fix[rust]: don't run parititioned groupby when computing mean of logi…

…cal column (#4412)
pola-rs · Aug 14, 2022 · cada9de · cada9de
1 parent beb802a
commit cada9de
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 10 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/aexpr.rs b/polars/polars-lazy/src/logical_plan/aexpr.rs
@@ -272,7 +272,7 @@ impl AExpr {
                     }
                     Mean(expr) => {
                         let mut field = arena.get(*expr).to_field(schema, ctxt, arena)?;
-                        field.coerce(DataType::Float64);
+                        coerce_numeric_aggregation(&mut field);
                         Ok(field)
                     }
                     List(expr) => {
@@ -284,12 +284,12 @@ impl AExpr {
                     }
                     Std(expr, _) => {
                         let mut field = arena.get(*expr).to_field(schema, ctxt, arena)?;
-                        field.coerce(DataType::Float64);
+                        coerce_numeric_aggregation(&mut field);
                         Ok(field)
                     }
                     Var(expr, _) => {
                         let mut field = arena.get(*expr).to_field(schema, ctxt, arena)?;
-                        field.coerce(DataType::Float64);
+                        coerce_numeric_aggregation(&mut field);
                         Ok(field)
                     }
                     NUnique(expr) => {
@@ -309,7 +309,7 @@ impl AExpr {
                     }
                     Quantile { expr, .. } => {
                         let mut field = arena.get(*expr).to_field(schema, ctxt, arena)?;
-                        field.coerce(DataType::Float64);
+                        coerce_numeric_aggregation(&mut field);
                         Ok(field)
                     }
                 }
@@ -359,3 +359,17 @@ impl AExpr {
         }
     }
 }
+
+fn coerce_numeric_aggregation(field: &mut Field) {
+    match field.dtype {
+        DataType::Duration(_) => {
+            // pass
+        }
+        DataType::Float32 => {
+            // pass
+        }
+        _ => {
+            field.coerce(DataType::Float64);
+        }
+    }
+}
diff --git a/polars/polars-lazy/src/physical_plan/planner/lp.rs b/polars/polars-lazy/src/physical_plan/planner/lp.rs
@@ -380,7 +380,12 @@ impl DefaultPlanner {
                                 match ae {
                                     // struct is needed to keep both states
                                     #[cfg(feature = "dtype-struct")]
-                                    Agg(AAggExpr::Mean(_)) => true,
+                                    Agg(AAggExpr::Mean(_)) => {
+                                        // only numeric means for now.
+                                        // logical types seem to break because of casts to float.
+                                        matches!(expr_arena.get(*agg).get_type(&input_schema, Context::Default, expr_arena).map(|dt| {
+                                        dt.is_numeric()}), Ok(true))
+                                    },
                                     // only allowed expressions
                                     Agg(agg_e) => {
                                         matches!(

diff --git a/polars/tests/it/lazy/groupby.rs b/polars/tests/it/lazy/groupby.rs
@@ -1,4 +1,5 @@
 use polars_core::series::ops::NullBehavior;
+use polars_core::SINGLE_LOCK;
 
 use super::*;
 
@@ -153,3 +154,30 @@ fn test_groupby_agg_list_with_not_aggregated() -> Result<()> {
     );
     Ok(())
 }
+
+#[test]
+#[cfg(all(feature = "dtype-duration", feature = "dtype-struct"))]
+fn test_logical_mean_partitioned_groupby_block() -> Result<()> {
+    let guard = SINGLE_LOCK.lock();
+    let df = df![
+        "a" => [1, 1, 2],
+        "duration" => [1000, 2000, 3000]
+    ]?;
+
+    let out = df
+        .lazy()
+        .with_column(col("duration").cast(DataType::Duration(TimeUnit::Microseconds)))
+        .groupby([col("a")])
+        .agg([col("duration").mean()])
+        .sort("duration", Default::default())
+        .collect()?;
+
+    let duration = out.column("duration")?;
+
+    assert_eq!(
+        duration.get(0),
+        AnyValue::Duration(1500, TimeUnit::Microseconds)
+    );
+
+    Ok(())
+}
diff --git a/py-polars/tests/db-benchmark/main.py b/py-polars/tests/db-benchmark/main.py
@@ -288,11 +288,6 @@
 print("total took:", total_time, "s")
 assert out.shape == (9999995, 8)
 
-if not ON_STRINGS:
-    if total_time > 11:
-        print("query took longer than 11s, may be noise")
-        exit(1)
-
 # Additional tests
 # the code below, does not belong to the db-benchmark
 # but it triggers other code paths so the checksums assertion
@@ -315,3 +310,8 @@
 
 assert out["id6"] == 2137755425
 assert np.isclose(out["v3"], 4.7040828499563754e8)
+
+if not ON_STRINGS:
+    if total_time > 12:
+        print("query took longer than 12s, may be noise")
+        exit(1)