feat: Allow (non-)coalescing in join_asof (pola-rs#17066)

Wouittone · Jun 22, 2024 · 1e86dc9 · 1e86dc9
1 parent e0965b9
commit 1e86dc9
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py
@@ -6216,6 +6216,7 @@ def join_asof(
         tolerance: str | int | float | timedelta | None = None,
         allow_parallel: bool = True,
         force_parallel: bool = False,
+        coalesce: bool | None = None,
     ) -> DataFrame:
         """
         Perform an asof join.
@@ -6292,6 +6293,15 @@ def join_asof(
         force_parallel
             Force the physical plan to evaluate the computation of both DataFrames up to
             the join in parallel.
+        coalesce
+            Coalescing behavior (merging of join columns).
+
+            - None: -> join specific.
+            - True: -> Always coalesce join columns.
+            - False: -> Never coalesce join columns.
+
+            Note that joining on any other expressions than `col`
+            will turn off coalescing.
 
         Examples
         --------
@@ -6509,6 +6519,7 @@ def join_asof(
                 tolerance=tolerance,
                 allow_parallel=allow_parallel,
                 force_parallel=force_parallel,
+                coalesce=coalesce,
             )
             .collect(_eager=True)
         )

diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
@@ -3733,6 +3733,7 @@ def join_asof(
         tolerance: str | int | float | timedelta | None = None,
         allow_parallel: bool = True,
         force_parallel: bool = False,
+        coalesce: bool | None = None,
     ) -> Self:
         """
         Perform an asof join.
@@ -3809,6 +3810,15 @@ def join_asof(
         force_parallel
             Force the physical plan to evaluate the computation of both DataFrames up to
             the join in parallel.
+        coalesce
+            Coalescing behavior (merging of join columns).
+
+            - None: -> join specific.
+            - True: -> Always coalesce join columns.
+            - False: -> Never coalesce join columns.
+
+            Note that joining on any other expressions than `col`
+            will turn off coalescing.
 
 
         Examples
@@ -3899,6 +3909,7 @@ def join_asof(
                 strategy,
                 tolerance_num,
                 tolerance_str,
+                coalesce=coalesce,
             )
         )
 

diff --git a/py-polars/src/lazyframe/mod.rs b/py-polars/src/lazyframe/mod.rs
@@ -892,7 +892,7 @@ impl PyLazyFrame {
     }
 
     #[cfg(feature = "asof_join")]
-    #[pyo3(signature = (other, left_on, right_on, left_by, right_by, allow_parallel, force_parallel, suffix, strategy, tolerance, tolerance_str))]
+    #[pyo3(signature = (other, left_on, right_on, left_by, right_by, allow_parallel, force_parallel, suffix, strategy, tolerance, tolerance_str, coalesce))]
     fn join_asof(
         &self,
         other: Self,
@@ -906,7 +906,13 @@ impl PyLazyFrame {
         strategy: Wrap<AsofStrategy>,
         tolerance: Option<Wrap<AnyValue<'_>>>,
         tolerance_str: Option<String>,
+        coalesce: Option<bool>,
     ) -> PyResult<Self> {
+        let coalesce = match coalesce {
+            None => JoinCoalesce::JoinSpecific,
+            Some(true) => JoinCoalesce::CoalesceColumns,
+            Some(false) => JoinCoalesce::KeepColumns,
+        };
         let ldf = self.ldf.clone();
         let other = other.ldf;
         let left_on = left_on.inner;
@@ -918,6 +924,7 @@ impl PyLazyFrame {
             .right_on([right_on])
             .allow_parallel(allow_parallel)
             .force_parallel(force_parallel)
+            .coalesce(coalesce)
             .how(JoinType::AsOf(AsOfOptions {
                 strategy: strategy.0,
                 left_by: left_by.map(strings_to_smartstrings),