Add field to dataframe join to indicate if we should keep duplicate keys

Author: timsaucer

python/datafusion/dataframe.py

@@ -643,6 +643,7 @@ def join(
         left_on: None = None,
         right_on: None = None,
         join_keys: None = None,
+        keep_duplicate_keys: bool = False,
     ) -> DataFrame: ...
 
     @overload
@@ -655,6 +656,7 @@ def join(
         left_on: str | Sequence[str],
         right_on: str | Sequence[str],
         join_keys: tuple[list[str], list[str]] | None = None,
+        keep_duplicate_keys: bool = False,
     ) -> DataFrame: ...
 
     @overload
@@ -667,6 +669,7 @@ def join(
         join_keys: tuple[list[str], list[str]],
         left_on: None = None,
         right_on: None = None,
+        keep_duplicate_keys: bool = False,
     ) -> DataFrame: ...
 
     def join(
@@ -678,6 +681,7 @@ def join(
         left_on: str | Sequence[str] | None = None,
         right_on: str | Sequence[str] | None = None,
         join_keys: tuple[list[str], list[str]] | None = None,
+        keep_duplicate_keys: bool = False,
     ) -> DataFrame:
         """Join this :py:class:`DataFrame` with another :py:class:`DataFrame`.
 
@@ -690,11 +694,23 @@ def join(
                 "right", "full", "semi", "anti".
             left_on: Join column of the left dataframe.
             right_on: Join column of the right dataframe.
+            keep_duplicate_keys: When False, the columns from the right DataFrame
+                that have identical names in the ``on`` fields to the left DataFrame
+                will be dropped.
             join_keys: Tuple of two lists of column names to join on. [Deprecated]
 
         Returns:
             DataFrame after join.
         """
+        if join_keys is not None:
+            warnings.warn(
+                "`join_keys` is deprecated, use `on` or `left_on` with `right_on`",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
+            left_on = join_keys[0]
+            right_on = join_keys[1]
+
         # This check is to prevent breaking API changes where users prior to
         # DF 43.0.0 would  pass the join_keys as a positional argument instead
         # of a keyword argument.
@@ -705,18 +721,10 @@ def join(
             and isinstance(on[1], list)
         ):
             # We know this is safe because we've checked the types
-            join_keys = on  # type: ignore[assignment]
+            left_on = on[0]
+            right_on = on[1]
             on = None
 
-        if join_keys is not None:
-            warnings.warn(
-                "`join_keys` is deprecated, use `on` or `left_on` with `right_on`",
-                category=DeprecationWarning,
-                stacklevel=2,
-            )
-            left_on = join_keys[0]
-            right_on = join_keys[1]
-
         if on is not None:
             if left_on is not None or right_on is not None:
                 error_msg = "`left_on` or `right_on` should not provided with `on`"
@@ -735,7 +743,9 @@ def join(
         if isinstance(right_on, str):
             right_on = [right_on]
 
-        return DataFrame(self.df.join(right.df, how, left_on, right_on))
+        return DataFrame(
+            self.df.join(right.df, how, left_on, right_on, keep_duplicate_keys)
+        )
 
     def join_on(
         self,

python/tests/test_dataframe.py

@@ -400,7 +400,6 @@ def test_unnest_without_nulls(nested_df):
     assert result.column(1) == pa.array([7, 8, 8, 9, 9, 9])
 
 
-@pytest.mark.filterwarnings("ignore:`join_keys`:DeprecationWarning")
 def test_join():
     ctx = SessionContext()
 

src/dataframe.rs

@@ -566,6 +566,7 @@ impl PyDataFrame {
         how: &str,
         left_on: Vec<PyBackedStr>,
         right_on: Vec<PyBackedStr>,
+        keep_duplicate_keys: bool,
     ) -> PyDataFusionResult<Self> {
         let join_type = match how {
             "inner" => JoinType::Inner,
@@ -584,13 +585,46 @@ impl PyDataFrame {
         let left_keys = left_on.iter().map(|s| s.as_ref()).collect::<Vec<&str>>();
         let right_keys = right_on.iter().map(|s| s.as_ref()).collect::<Vec<&str>>();
 
-        let df = self.df.as_ref().clone().join(
+        let mut df = self.df.as_ref().clone().join(
             right.df.as_ref().clone(),
             join_type,
             &left_keys,
             &right_keys,
             None,
         )?;
+
+        if !keep_duplicate_keys {
+            let mutual_keys = left_keys
+                .iter()
+                .zip(right_keys.iter())
+                .filter(|(l, r)| l == r)
+                .map(|(key, _)| *key)
+                .collect::<Vec<_>>();
+
+            let fields_to_drop = mutual_keys
+                .iter()
+                .map(|name| {
+                    df.logical_plan()
+                        .schema()
+                        .qualified_fields_with_unqualified_name(name)
+                })
+                .filter(|r| r.len() == 2)
+                .map(|r| r[1])
+                .collect::<Vec<_>>();
+
+            let expr: Vec<Expr> = df
+                .logical_plan()
+                .schema()
+                .fields()
+                .into_iter()
+                .enumerate()
+                .map(|(idx, _)| df.logical_plan().schema().qualified_field(idx))
+                .filter(|(qualifier, f)| !fields_to_drop.contains(&(*qualifier, f)))
+                .map(|(qualifier, field)| Expr::Column(Column::from((qualifier, field))))
+                .collect();
+            df = df.select(expr)?;
+        }
+
         Ok(Self::new(df))
     }