feat(rust, python): prefer streaming groupby if partitionable (#5580)

ritchie46 · web-flow · commit 7929246ffff6 · 2022-11-22T10:24:43.000+01:00
diff --git a/polars/polars-core/src/frame/from.rs b/polars/polars-core/src/frame/from.rs
@@ -24,3 +24,13 @@ impl TryFrom<StructArray> for DataFrame {
         DataFrame::new(columns)
     }
 }
+
+impl From<&Schema> for DataFrame {
+    fn from(schema: &Schema) -> Self {
+        let cols = schema
+            .iter()
+            .map(|(name, dtype)| Series::new_empty(name, dtype))
+            .collect();
+        DataFrame::new_no_checks(cols)
+    }
+}
diff --git a/polars/polars-core/src/utils/mod.rs b/polars/polars-core/src/utils/mod.rs
@@ -199,7 +199,7 @@ pub fn split_df_as_ref(df: &DataFrame, n: usize) -> PolarsResult<Vec<DataFrame>>
 #[doc(hidden)]
 /// Split a [`DataFrame`] into `n` parts. We take a `&mut` to be able to repartition/align chunks.
 pub fn split_df(df: &mut DataFrame, n: usize) -> PolarsResult<Vec<DataFrame>> {
-    if n == 0 {
+    if n == 0 || df.height() == 0 {
         return Ok(vec![df.clone()]);
     }
     // make sure that chunks are aligned.
diff --git a/polars/polars-lazy/Cargo.toml b/polars/polars-lazy/Cargo.toml
@@ -38,15 +38,15 @@ temporal = ["dtype-datetime", "dtype-date", "dtype-time", "dtype-duration", "pol
 fmt = ["polars-core/fmt", "polars-plan/fmt"]
 strings = ["polars-plan/strings"]
 future = []
-dtype-u8 = ["polars-plan/dtype-u8"]
-dtype-u16 = ["polars-plan/dtype-u16"]
-dtype-i8 = ["polars-plan/dtype-i8"]
-dtype-i16 = ["polars-plan/dtype-i16"]
+dtype-u8 = ["polars-plan/dtype-u8", "polars-pipe/dtype-u8"]
+dtype-u16 = ["polars-plan/dtype-u16", "polars-pipe/dtype-u16"]
+dtype-i8 = ["polars-plan/dtype-i8", "polars-pipe/dtype-i8"]
+dtype-i16 = ["polars-plan/dtype-i16", "polars-pipe/dtype-i16"]
 dtype-date = ["polars-plan/dtype-date", "polars-time/dtype-date", "temporal"]
 dtype-datetime = ["polars-plan/dtype-datetime", "polars-time/dtype-datetime", "temporal"]
 dtype-duration = ["polars-plan/dtype-duration", "polars-time/dtype-duration", "temporal"]
 dtype-time = ["polars-core/dtype-time", "temporal"]
-dtype-categorical = ["polars-plan/dtype-categorical"]
+dtype-categorical = ["polars-plan/dtype-categorical", "polars-pipe/dtype-categorical"]
 dtype-struct = ["polars-plan/dtype-struct"]
 dtype-binary = ["polars-plan/dtype-binary"]
 object = ["polars-plan/object"]
diff --git a/polars/polars-lazy/polars-pipe/Cargo.toml b/polars/polars-lazy/polars-pipe/Cargo.toml
@@ -26,3 +26,8 @@ csv-file = ["polars-plan/csv-file", "polars-io/csv-file"]
 parquet = ["polars-plan/parquet", "polars-io/parquet"]
 nightly = ["polars-core/nightly", "polars-utils/nightly", "hashbrown/nightly"]
 cross_join = ["polars-core/cross_join"]
+dtype-u8 = ["polars-core/dtype-u8"]
+dtype-u16 = ["polars-core/dtype-u16"]
+dtype-i8 = ["polars-core/dtype-i8"]
+dtype-i16 = ["polars-core/dtype-i16"]
+dtype-categorical = ["polars-core/dtype-categorical"]
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/convert.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/convert.rs
@@ -14,6 +14,7 @@ use crate::executors::sinks::groupby::aggregates::count::CountAgg;
 use crate::executors::sinks::groupby::aggregates::first::FirstAgg;
 use crate::executors::sinks::groupby::aggregates::last::LastAgg;
 use crate::executors::sinks::groupby::aggregates::mean::MeanAgg;
+use crate::executors::sinks::groupby::aggregates::null::NullAgg;
 use crate::executors::sinks::groupby::aggregates::{AggregateFunction, SumAgg};
 use crate::expressions::PhysicalPipedExpr;
 use crate::operators::DataChunk;
@@ -89,14 +90,24 @@ where
 {
     match expr_arena.get(node) {
         AExpr::Alias(input, _) => convert_to_hash_agg(*input, expr_arena, schema, to_physical),
-        AExpr::Count | AExpr::Agg(AAggExpr::Count(_)) => (
+        AExpr::Count => (
             Arc::new(Count {}),
             AggregateFunction::Count(CountAgg::new()),
         ),
         AExpr::Agg(agg) => match agg {
             AAggExpr::Sum(input) => {
                 let phys_expr = to_physical(*input, expr_arena).unwrap();
-                let agg_fn = match phys_expr.field(schema).unwrap().dtype.to_physical() {
+                let logical_dtype = phys_expr.field(schema).unwrap().dtype;
+
+                #[cfg(feature = "dtype-categorical")]
+                if matches!(logical_dtype, DataType::Categorical(_)) {
+                    return (
+                        phys_expr,
+                        AggregateFunction::Null(NullAgg::new(logical_dtype)),
+                    );
+                }
+
+                let agg_fn = match logical_dtype.to_physical() {
                     // Boolean is aggregated as the IDX type.
                     DataType::Boolean => {
                         if std::mem::size_of::<IdxSize>() == 4 {
@@ -117,19 +128,28 @@ where
                     DataType::Int64 => AggregateFunction::SumI64(SumAgg::<i64>::new()),
                     DataType::Float32 => AggregateFunction::SumF32(SumAgg::<f32>::new()),
                     DataType::Float64 => AggregateFunction::SumF64(SumAgg::<f64>::new()),
-                    _ => unreachable!(),
+                    dt => AggregateFunction::Null(NullAgg::new(dt)),
                 };
                 (phys_expr, agg_fn)
             }
             AAggExpr::Mean(input) => {
                 let phys_expr = to_physical(*input, expr_arena).unwrap();
-                let agg_fn = match phys_expr.field(schema).unwrap().dtype.to_physical() {
+
+                let logical_dtype = phys_expr.field(schema).unwrap().dtype;
+                #[cfg(feature = "dtype-categorical")]
+                if matches!(logical_dtype, DataType::Categorical(_)) {
+                    return (
+                        phys_expr,
+                        AggregateFunction::Null(NullAgg::new(logical_dtype)),
+                    );
+                }
+                let agg_fn = match logical_dtype.to_physical() {
                     dt if dt.is_integer() => AggregateFunction::MeanF64(MeanAgg::<f64>::new()),
                     // Boolean is aggregated as the IDX type.
                     DataType::Boolean => AggregateFunction::MeanF64(MeanAgg::<f64>::new()),
                     DataType::Float32 => AggregateFunction::MeanF32(MeanAgg::<f32>::new()),
                     DataType::Float64 => AggregateFunction::MeanF64(MeanAgg::<f64>::new()),
-                    _ => unreachable!(),
+                    dt => AggregateFunction::Null(NullAgg::new(dt)),
                 };
                 (phys_expr, agg_fn)
             }
@@ -143,6 +163,10 @@ where
                 let dtype = phys_expr.field(schema).unwrap().dtype;
                 (phys_expr, AggregateFunction::Last(LastAgg::new(dtype)))
             }
+            AAggExpr::Count(input) => {
+                let phys_expr = to_physical(*input, expr_arena).unwrap();
+                (phys_expr, AggregateFunction::Count(CountAgg::new()))
+            }
             agg => panic!("{:?} not yet implemented.", agg),
         },
         _ => todo!(),
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/count.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/count.rs
@@ -38,10 +38,6 @@ impl AggregateFn for CountAgg {
         self.count += other.count;
     }
 
-    fn split(&self) -> Box<dyn AggregateFn> {
-        Box::new(Self::new())
-    }
-
     fn finalize(&mut self) -> AnyValue<'static> {
         AnyValue::from(self.count)
     }
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/first.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/first.rs
@@ -44,10 +44,6 @@ impl AggregateFn for FirstAgg {
         };
     }
 
-    fn split(&self) -> Box<dyn AggregateFn> {
-        Box::new(Self::new(self.dtype.clone()))
-    }
-
     fn finalize(&mut self) -> AnyValue<'static> {
         std::mem::take(&mut self.first).unwrap_or(AnyValue::Null)
     }
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/interface.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/interface.rs
@@ -8,6 +8,7 @@ use crate::executors::sinks::groupby::aggregates::count::CountAgg;
 use crate::executors::sinks::groupby::aggregates::first::FirstAgg;
 use crate::executors::sinks::groupby::aggregates::last::LastAgg;
 use crate::executors::sinks::groupby::aggregates::mean::MeanAgg;
+use crate::executors::sinks::groupby::aggregates::null::NullAgg;
 use crate::executors::sinks::groupby::aggregates::SumAgg;
 use crate::operators::IdxSize;
 
@@ -54,8 +55,6 @@ pub trait AggregateFn: Send + Sync {
 
     fn combine(&mut self, other: &dyn Any);
 
-    fn split(&self) -> Box<dyn AggregateFn>;
-
     fn finalize(&mut self) -> AnyValue<'static>;
 
     fn as_any(&self) -> &dyn Any;
@@ -76,10 +75,10 @@ pub enum AggregateFunction {
     SumI64(SumAgg<i64>),
     MeanF32(MeanAgg<f32>),
     MeanF64(MeanAgg<f64>),
-    // place holder for any aggregate function
-    // this is not preferred because of the extra
-    // indirection
-    // Other(Box<dyn AggregateFn>)
+    Null(NullAgg), // place holder for any aggregate function
+                   // this is not preferred because of the extra
+                   // indirection
+                   // Other(Box<dyn AggregateFn>)
 }
 
 impl AggregateFunction {
@@ -97,6 +96,7 @@ impl AggregateFunction {
             MeanF32(_) => MeanF32(MeanAgg::new()),
             MeanF64(_) => MeanF64(MeanAgg::new()),
             Count(_) => Count(CountAgg::new()),
+            Null(a) => Null(a.clone()),
         }
     }
 }
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/last.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/last.rs
@@ -42,10 +42,6 @@ impl AggregateFn for LastAgg {
         };
     }
 
-    fn split(&self) -> Box<dyn AggregateFn> {
-        Box::new(Self::new(self.dtype.clone()))
-    }
-
     fn finalize(&mut self) -> AnyValue<'static> {
         std::mem::take(&mut self.last)
     }
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/mean.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/mean.rs
@@ -100,10 +100,6 @@ impl<K: NumericNative + Add<Output = K> + NumCast> AggregateFn for MeanAgg<K> {
         };
     }
 
-    fn split(&self) -> Box<dyn AggregateFn> {
-        Box::new(Self::new())
-    }
-
     fn finalize(&mut self) -> AnyValue<'static> {
         if let Some(val) = self.sum {
             unsafe {
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/mod.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/mod.rs
@@ -4,6 +4,7 @@ mod first;
 mod interface;
 mod last;
 mod mean;
+mod null;
 mod sum;
 
 pub use convert::*;
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/null.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/null.rs
@@ -0,0 +1,36 @@
+use std::any::Any;
+
+use polars_core::prelude::*;
+
+use crate::executors::sinks::groupby::aggregates::AggregateFn;
+
+#[derive(Clone)]
+pub struct NullAgg(DataType);
+
+impl NullAgg {
+    pub(crate) fn new(dt: DataType) -> Self {
+        Self(dt)
+    }
+}
+
+impl AggregateFn for NullAgg {
+    fn pre_agg(&mut self, _chunk_idx: IdxSize, _item: &mut dyn ExactSizeIterator<Item = AnyValue>) {
+        // no-op
+    }
+
+    fn dtype(&self) -> DataType {
+        self.0.clone()
+    }
+
+    fn combine(&mut self, _other: &dyn Any) {
+        // no-op
+    }
+
+    fn finalize(&mut self) -> AnyValue<'static> {
+        AnyValue::Null
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/sum.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/aggregates/sum.rs
@@ -79,10 +79,6 @@ impl<K: NumericNative + Add<Output = K> + NumCast> AggregateFn for SumAgg<K> {
         self.sum = sum;
     }
 
-    fn split(&self) -> Box<dyn AggregateFn> {
-        Box::new(Self::new())
-    }
-
     fn finalize(&mut self) -> AnyValue<'static> {
         if let Some(val) = self.sum {
             unsafe {
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/generic.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/generic.rs
@@ -15,6 +15,7 @@ use rayon::prelude::*;
 
 use super::aggregates::AggregateFn;
 use crate::executors::sinks::groupby::aggregates::AggregateFunction;
+use crate::executors::sinks::groupby::physical_agg_to_logical;
 use crate::executors::sinks::groupby::utils::compute_slices;
 use crate::executors::sinks::utils::{hash_series, load_vec};
 use crate::executors::sinks::HASHMAP_INIT_SIZE;
@@ -186,14 +187,7 @@ impl GenericGroupbySink {
                             cols.push(key_builder.into_series());
                         }
                         cols.extend(buffers.into_iter().map(|buf| buf.into_series()));
-                        for (s, (name, dtype)) in cols.iter_mut().zip(self.output_schema.iter()) {
-                            if s.name() != name {
-                                s.rename(name);
-                            }
-                            if s.dtype() != dtype {
-                                *s = s.cast(dtype).unwrap()
-                            }
-                        }
+                        physical_agg_to_logical(&mut cols, &self.output_schema);
                         Some(DataFrame::new_no_checks(cols))
                     })
                     .collect::<Vec<_>>();
@@ -416,6 +410,11 @@ impl Sink for GenericGroupbySink {
 
     fn finalize(&mut self, _context: &PExecutionContext) -> PolarsResult<FinalizedSink> {
         let dfs = self.pre_finalize()?;
+        if dfs.is_empty() {
+            return Ok(FinalizedSink::Finished(DataFrame::from(
+                self.output_schema.as_ref(),
+            )));
+        }
         let mut df = accumulate_dataframes_vertical_unchecked(dfs);
         DataFrame::new(std::mem::take(df.get_columns_mut())).map(FinalizedSink::Finished)
     }
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/mod.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/mod.rs
@@ -5,5 +5,21 @@ mod string;
 mod utils;
 
 pub(crate) use generic::*;
+use polars_core::prelude::*;
 pub(crate) use primitive::*;
 pub(crate) use string::*;
+
+pub(super) fn physical_agg_to_logical(cols: &mut [Series], output_schema: &Schema) {
+    for (s, (name, dtype)) in cols.iter_mut().zip(output_schema.iter()) {
+        if s.name() != name {
+            s.rename(name);
+        }
+        let dtype_left = s.dtype();
+        if dtype_left != dtype
+            && !matches!(dtype, DataType::Boolean)
+            && !(dtype.is_float() && dtype_left.is_float())
+        {
+            *s = s.cast(dtype).unwrap()
+        }
+    }
+}
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/primitive.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/primitive.rs
@@ -16,6 +16,7 @@ use rayon::prelude::*;
 
 use super::aggregates::AggregateFn;
 use crate::executors::sinks::groupby::aggregates::AggregateFunction;
+use crate::executors::sinks::groupby::physical_agg_to_logical;
 use crate::executors::sinks::groupby::string::{apply_aggregate, write_agg_idx};
 use crate::executors::sinks::groupby::utils::compute_slices;
 use crate::executors::sinks::utils::load_vec;
@@ -163,14 +164,7 @@ where
                         let mut cols = Vec::with_capacity(1 + self.number_of_aggs());
                         cols.push(key_builder.finish().into_series());
                         cols.extend(buffers.into_iter().map(|buf| buf.into_series()));
-                        for (s, (name, dtype)) in cols.iter_mut().zip(self.output_schema.iter()) {
-                            if s.name() != name {
-                                s.rename(name);
-                            }
-                            if s.dtype() != dtype {
-                                *s = s.cast(dtype).unwrap()
-                            }
-                        }
+                        physical_agg_to_logical(&mut cols, &self.output_schema);
                         Some(DataFrame::new_no_checks(cols))
                     })
                     .collect::<Vec<_>>();
@@ -308,6 +302,11 @@ where
 
     fn finalize(&mut self, _context: &PExecutionContext) -> PolarsResult<FinalizedSink> {
         let dfs = self.pre_finalize()?;
+        if dfs.is_empty() {
+            return Ok(FinalizedSink::Finished(DataFrame::from(
+                self.output_schema.as_ref(),
+            )));
+        }
         let mut df = accumulate_dataframes_vertical_unchecked(dfs);
         DataFrame::new(std::mem::take(df.get_columns_mut())).map(FinalizedSink::Finished)
     }
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/string.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/string.rs
diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/joins/inner_left.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/joins/inner_left.rs
diff --git a/polars/polars-lazy/polars-pipe/src/pipeline/convert.rs b/polars/polars-lazy/polars-pipe/src/pipeline/convert.rs
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/aexpr.rs b/polars/polars-lazy/polars-plan/src/logical_plan/aexpr.rs
diff --git a/polars/polars-lazy/src/physical_plan/planner/lp.rs b/polars/polars-lazy/src/physical_plan/planner/lp.rs
diff --git a/polars/polars-lazy/src/tests/queries.rs b/polars/polars-lazy/src/tests/queries.rs
diff --git a/py-polars/polars/internals/lazyframe/frame.py b/py-polars/polars/internals/lazyframe/frame.py

Original file line number	Diff line number	Diff line change
`@@ -24,3 +24,13 @@ impl TryFrom<StructArray> for DataFrame {`
`24`	`24`	`DataFrame::new(columns)`
`25`	`25`	`}`
`26`	`26`	`}`
	`27`	`+`
	`28`	`+impl From<&Schema> for DataFrame {`
	`29`	`+ fn from(schema: &Schema) -> Self {`
	`30`	`+ let cols = schema`
	`31`	`+ .iter()`
	`32`	`+ .map(\|(name, dtype)\| Series::new_empty(name, dtype))`
	`33`	`+ .collect();`
	`34`	`+ DataFrame::new_no_checks(cols)`
	`35`	`+ }`
	`36`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -199,7 +199,7 @@ pub fn split_df_as_ref(df: &DataFrame, n: usize) -> PolarsResult<Vec<DataFrame>>`
`199`	`199`	`#[doc(hidden)]`
`200`	`200`	/// Split a [`DataFrame`] into `n` parts. We take a `&mut` to be able to repartition/align chunks.
`201`	`201`	`pub fn split_df(df: &mut DataFrame, n: usize) -> PolarsResult<Vec<DataFrame>> {`
`202`		`- if n == 0 {`
	`202`	`+ if n == 0 \|\| df.height() == 0 {`
`203`	`203`	`return Ok(vec![df.clone()]);`
`204`	`204`	`}`
`205`	`205`	`// make sure that chunks are aligned.`
Original file line number	Diff line number	Diff line change
`@@ -38,10 +38,6 @@ impl AggregateFn for CountAgg {`
`38`	`38`	`self.count += other.count;`
`39`	`39`	`}`
`40`	`40`
`41`		`- fn split(&self) -> Box<dyn AggregateFn> {`
`42`		`- Box::new(Self::new())`
`43`		`- }`
`44`		`-`
`45`	`41`	`fn finalize(&mut self) -> AnyValue<'static> {`
`46`	`42`	`AnyValue::from(self.count)`
`47`	`43`	`}`
Original file line number	Diff line number	Diff line change
`@@ -44,10 +44,6 @@ impl AggregateFn for FirstAgg {`
`44`	`44`	`};`
`45`	`45`	`}`
`46`	`46`
`47`		`- fn split(&self) -> Box<dyn AggregateFn> {`
`48`		`- Box::new(Self::new(self.dtype.clone()))`
`49`		`- }`
`50`		`-`
`51`	`47`	`fn finalize(&mut self) -> AnyValue<'static> {`
`52`	`48`	`std::mem::take(&mut self.first).unwrap_or(AnyValue::Null)`
`53`	`49`	`}`
Original file line number	Diff line number	Diff line change
`@@ -42,10 +42,6 @@ impl AggregateFn for LastAgg {`
`42`	`42`	`};`
`43`	`43`	`}`
`44`	`44`
`45`		`- fn split(&self) -> Box<dyn AggregateFn> {`
`46`		`- Box::new(Self::new(self.dtype.clone()))`
`47`		`- }`
`48`		`-`
`49`	`45`	`fn finalize(&mut self) -> AnyValue<'static> {`
`50`	`46`	`std::mem::take(&mut self.last)`
`51`	`47`	`}`