perf: rechunk before grouping on multiple keys (pola-rs#11711)

cmdlineluser · Oct 13, 2023 · abd7586 · abd7586
1 parent f18583c
commit abd7586
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 7 deletions.
diff --git a/crates/polars-plan/src/dsl/string.rs b/crates/polars-plan/src/dsl/string.rs
@@ -395,7 +395,7 @@ impl StringNameSpace {
     }
 
     /// Slice the string values.
-    pub fn str_slice(self, start: i64, length: Option<u64>) -> Expr {
+    pub fn slice(self, start: i64, length: Option<u64>) -> Expr {
         self.0
             .map_private(FunctionExpr::StringExpr(StringFunction::Slice(
                 start, length,

diff --git a/crates/polars-plan/src/logical_plan/optimizer/delay_rechunk.rs b/crates/polars-plan/src/logical_plan/optimizer/delay_rechunk.rs
@@ -25,8 +25,9 @@ impl OptimizationRule for DelayRechunk {
         match lp_arena.get(node) {
             // An aggregation can be partitioned, its wasteful to rechunk before that partition.
             #[allow(unused_mut)]
-            ALogicalPlan::Aggregate { input, .. } => {
-                if !self.processed.insert(node.0) {
+            ALogicalPlan::Aggregate { input, keys, .. } => {
+                // Multiple keys on multiple chunks is much slower, so rechunk.
+                if !self.processed.insert(node.0) || keys.len() > 1 {
                     return None;
                 };
 

diff --git a/crates/polars-sql/src/functions.rs b/crates/polars-sql/src/functions.rs
@@ -683,7 +683,7 @@ impl SqlFunctionVisitor<'_> {
             EndsWith => self.visit_binary(|e, s| e.str().ends_with(s)),
             InitCap => self.visit_unary(|e| e.str().to_titlecase()),
             Left => self.try_visit_binary(|e, length| {
-                Ok(e.str().str_slice(0, match length {
+                Ok(e.str().slice(0, match length {
                     Expr::Literal(LiteralValue::Int64(n)) => Some(n as u64),
                     _ => {
                         polars_bail!(InvalidOperation: "Invalid 'length' for Left: {}", function.args[1]);
@@ -729,15 +729,15 @@ impl SqlFunctionVisitor<'_> {
             StartsWith => self.visit_binary(|e, s| e.str().starts_with(s)),
             Substring => match function.args.len() {
                 2 => self.try_visit_binary(|e, start| {
-                    Ok(e.str().str_slice(match start {
+                    Ok(e.str().slice(match start {
                         Expr::Literal(LiteralValue::Int64(n)) => n,
                         _ => {
                             polars_bail!(InvalidOperation: "Invalid 'start' for Substring: {}", function.args[1]);
                         }
                     }, None))
                 }),
                 3 => self.try_visit_ternary(|e, start, length| {
-                    Ok(e.str().str_slice(
+                    Ok(e.str().slice(
                         match start {
                             Expr::Literal(LiteralValue::Int64(n)) => n,
                             _ => {

diff --git a/py-polars/src/expr/string.rs b/py-polars/src/expr/string.rs
@@ -91,7 +91,7 @@ impl PyExpr {
     }
 
     fn str_slice(&self, start: i64, length: Option<u64>) -> Self {
-        self.inner.clone().str().str_slice(start, length).into()
+        self.inner.clone().str().slice(start, length).into()
     }
 
     fn str_explode(&self) -> Self {