add fast path for binary/ternary exprs in case when all groups are un…

…ique
pola-rs · Nov 26, 2021 · 5eec42c · 5eec42c
1 parent fb57767
commit 5eec42c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/binary.rs b/polars/polars-lazy/src/physical_plan/expressions/binary.rs
@@ -98,7 +98,9 @@ impl PhysicalExpr for BinaryExpr {
 
         match (ac_l.agg_state(), ac_r.agg_state()) {
             // One of the two exprs is aggregated with flat aggregation, e.g. `e.min(), e.max(), e.first()`
-            (AggState::AggregatedFlat(_), AggState::NotAggregated(_)) => {
+
+            // if the groups_len == df.len we can just apply all flat.
+            (AggState::AggregatedFlat(s), AggState::NotAggregated(_)) if s.len() != df.height() => {
                 // this is a flat series of len eq to group tuples
                 let l = ac_l.aggregated();
                 let l = l.as_ref();
@@ -141,7 +143,8 @@ impl PhysicalExpr for BinaryExpr {
                 ac_l.with_series(ca.into_series(), true);
                 Ok(ac_l)
             }
-            (AggState::NotAggregated(_), AggState::AggregatedFlat(_)) => {
+            // if the groups_len == df.len we can just apply all flat.
+            (AggState::NotAggregated(_), AggState::AggregatedFlat(s)) if s.len() != df.height() => {
                 // this is now a list
                 let l = ac_l.aggregated();
                 let l = l.list().unwrap();

diff --git a/polars/polars-lazy/src/physical_plan/expressions/ternary.rs b/polars/polars-lazy/src/physical_plan/expressions/ternary.rs
@@ -64,7 +64,8 @@ The predicate produced {} values. Where the original DataFrame has {} values",
         );
 
         match (ac_truthy.agg_state(), ac_falsy.agg_state()) {
-            (AggState::AggregatedFlat(_), AggState::NotAggregated(_)) => {
+            // if the groups_len == df.len we can just apply all flat.
+            (AggState::AggregatedFlat(s), AggState::NotAggregated(_)) if s.len() != df.height() => {
                 // this is a flat series of len eq to group tuples
                 let truthy = ac_truthy.aggregated();
                 let truthy = truthy.as_ref();
@@ -120,7 +121,8 @@ The predicate produced {} values. Where the original DataFrame has {} values",
                 ac_truthy.with_series(ca.into_series(), true);
                 Ok(ac_truthy)
             }
-            (AggState::NotAggregated(_), AggState::AggregatedFlat(_)) => {
+            // if the groups_len == df.len we can just apply all flat.
+            (AggState::NotAggregated(_), AggState::AggregatedFlat(s)) if s.len() != df.height() => {
                 // this is now a list
                 let truthy = ac_truthy.aggregated();
                 let truthy = truthy.as_ref();