feat: Add having to group_by context (pola-rs#23550)

Co-authored-by: coastalwhite <me@gburghoorn.com>

Author: borchero

crates/polars-lazy/src/frame/mod.rs

@@ -1100,6 +1100,7 @@ impl LazyFrame {
                 logical_plan: self.logical_plan,
                 opt_state,
                 keys,
+                predicates: vec![],
                 maintain_order: false,
                 dynamic_options: None,
                 rolling_options: None,
@@ -1112,6 +1113,7 @@ impl LazyFrame {
                 logical_plan: self.logical_plan,
                 opt_state,
                 keys,
+                predicates: vec![],
                 maintain_order: false,
             }
         }
@@ -1147,6 +1149,7 @@ impl LazyFrame {
         LazyGroupBy {
             logical_plan: self.logical_plan,
             opt_state,
+            predicates: vec![],
             keys: group_by.as_ref().to_vec(),
             maintain_order: true,
             dynamic_options: None,
@@ -1192,6 +1195,7 @@ impl LazyFrame {
         LazyGroupBy {
             logical_plan: self.logical_plan,
             opt_state,
+            predicates: vec![],
             keys: group_by.as_ref().to_vec(),
             maintain_order: true,
             dynamic_options: Some(options),
@@ -1214,6 +1218,7 @@ impl LazyFrame {
                 logical_plan: self.logical_plan,
                 opt_state,
                 keys,
+                predicates: vec![],
                 maintain_order: true,
                 dynamic_options: None,
                 rolling_options: None,
@@ -1226,6 +1231,7 @@ impl LazyFrame {
                 logical_plan: self.logical_plan,
                 opt_state,
                 keys,
+                predicates: vec![],
                 maintain_order: true,
             }
         }
@@ -1975,6 +1981,7 @@ pub struct LazyGroupBy {
     pub logical_plan: DslPlan,
     opt_state: OptFlags,
     keys: Vec<Expr>,
+    predicates: Vec<Expr>,
     maintain_order: bool,
     #[cfg(feature = "dynamic_group_by")]
     dynamic_options: Option<DynamicGroupOptions>,
@@ -1993,6 +2000,31 @@ impl From<LazyGroupBy> for LazyFrame {
 }
 
 impl LazyGroupBy {
+    /// Filter groups with a predicate after aggregation.
+    ///
+    /// Similarly to the [LazyGroupBy::agg] method, the predicate must run an aggregation as it
+    /// is evaluated on the groups.
+    /// This method can be chained in which case all predicates must evaluate to `true` for a
+    /// group to be kept.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use polars_core::prelude::*;
+    /// use polars_lazy::prelude::*;
+    ///
+    /// fn example(df: DataFrame) -> LazyFrame {
+    ///       df.lazy()
+    ///        .group_by_stable([col("date")])
+    ///        .having(col("rain").sum().gt(lit(10)))
+    ///        .agg([col("rain").min().alias("min_rain")])
+    /// }
+    /// ```
+    pub fn having(mut self, predicate: Expr) -> Self {
+        self.predicates.push(predicate);
+        self
+    }
+
     /// Group by and aggregate.
     ///
     /// Select a column with [col] and choose an aggregation.
@@ -2019,6 +2051,7 @@ impl LazyGroupBy {
         let lp = DslBuilder::from(self.logical_plan)
             .group_by(
                 self.keys,
+                self.predicates,
                 aggs,
                 None,
                 self.maintain_order,
@@ -2029,7 +2062,7 @@ impl LazyGroupBy {
 
         #[cfg(not(feature = "dynamic_group_by"))]
         let lp = DslBuilder::from(self.logical_plan)
-            .group_by(self.keys, aggs, None, self.maintain_order)
+            .group_by(self.keys, self.predicates, aggs, None, self.maintain_order)
             .build();
         LazyFrame::from_logical_plan(lp, self.opt_state)
     }
@@ -2075,6 +2108,10 @@ impl LazyGroupBy {
     /// **It is not recommended that you use this as materializing the DataFrame is very
     /// expensive.**
     pub fn apply(self, f: PlanCallback<DataFrame, DataFrame>, schema: SchemaRef) -> LazyFrame {
+        if !self.predicates.is_empty() {
+            panic!("not yet implemented: `apply` cannot be used with `having` predicates");
+        }
+
         #[cfg(feature = "dynamic_group_by")]
         let options = GroupbyOptions {
             dynamic: self.dynamic_options,
@@ -2088,6 +2125,7 @@ impl LazyGroupBy {
         let lp = DslPlan::GroupBy {
             input: Arc::new(self.logical_plan),
             keys: self.keys,
+            predicates: vec![],
             aggs: vec![],
             apply: Some((f, schema)),
             maintain_order: self.maintain_order,

crates/polars-plan/dsl-schema-hashes.json

@@ -43,7 +43,7 @@
   "Dimension": "68880cdb10230df6c8c1632b073c80bd8ceb5c56a368c0cb438431ca9f3d3b31",
   "DistinctOptionsDSL": "41be5ec69ef9a614f2b36ac5deadfecdea5cca847ae1ada9d4bc626ff52a5b38",
   "DslFunction": "221f1a46a043c8ed54f57be981bf24509f04f5f91f0f08e0acc180d96f842ebf",
-  "DslPlan": "f9d1c8f632172a24dbaf76e56b2efab237937d2a2c615858d6f196932ada98f6",
+  "DslPlan": "2c776aa38f65fd85707d969967aa733e3f91a1002a200ca3a3ef7c1328c313eb",
   "Duration": "44999d59023085cbb592ce94b30d34f9b983081fc72bd6435a49bdf0869c0074",
   "DynListLiteralValue": "2266a553cb4a943f7097f24539eaa802453cf8742675996215235bd682dec0e8",
   "DynLiteralValue": "47dc404f42bef5ab71659b9e10a97413202a61bfa3ac9fc66fff4a176653f7fe",

crates/polars-plan/src/dsl/builder_dsl.rs

@@ -241,9 +241,11 @@ impl DslBuilder {
         .into()
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn group_by<E: AsRef<[Expr]>>(
         self,
         keys: Vec<Expr>,
+        predicates: Vec<Expr>,
         aggs: E,
         apply: Option<(PlanCallback<DataFrame, DataFrame>, SchemaRef)>,
         maintain_order: bool,
@@ -262,6 +264,7 @@ impl DslBuilder {
         DslPlan::GroupBy {
             input: Arc::new(self.0),
             keys,
+            predicates,
             aggs,
             apply,
             maintain_order,

crates/polars-plan/src/dsl/plan.rs

@@ -66,6 +66,7 @@ pub enum DslPlan {
     GroupBy {
         input: Arc<DslPlan>,
         keys: Vec<Expr>,
+        predicates: Vec<Expr>,
         aggs: Vec<Expr>,
         maintain_order: bool,
         options: Arc<GroupbyOptions>,
@@ -190,7 +191,7 @@ impl Clone for DslPlan {
             Self::Scan { sources,  unified_scan_args, scan_type, cached_ir } => Self::Scan { sources: sources.clone(), unified_scan_args: unified_scan_args.clone(), scan_type: scan_type.clone(), cached_ir: cached_ir.clone() },
             Self::DataFrameScan { df, schema, } => Self::DataFrameScan { df: df.clone(), schema: schema.clone(),  },
             Self::Select { expr, input, options } => Self::Select { expr: expr.clone(), input: input.clone(), options: options.clone() },
-            Self::GroupBy { input, keys, aggs,  apply, maintain_order, options } => Self::GroupBy { input: input.clone(), keys: keys.clone(), aggs: aggs.clone(), apply: apply.clone(), maintain_order: maintain_order.clone(), options: options.clone() },
+            Self::GroupBy { input, keys, predicates, aggs, apply, maintain_order, options } => Self::GroupBy { input: input.clone(), keys: keys.clone(), predicates: predicates.clone(), aggs: aggs.clone(), apply: apply.clone(), maintain_order: maintain_order.clone(), options: options.clone() },
             Self::Join { input_left, input_right, left_on, right_on, predicates, options } => Self::Join { input_left: input_left.clone(), input_right: input_right.clone(), left_on: left_on.clone(), right_on: right_on.clone(), options: options.clone(), predicates: predicates.clone() },
             Self::HStack { input, exprs, options } => Self::HStack { input: input.clone(), exprs: exprs.clone(),  options: options.clone() },
             Self::MatchToSchema { input, match_schema, per_column, extra_columns } => Self::MatchToSchema { input: input.clone(), match_schema: match_schema.clone(), per_column: per_column.clone(), extra_columns: *extra_columns },

crates/polars-plan/src/dsl/serializable_plan.rs

@@ -61,6 +61,7 @@ pub(crate) enum SerializableDslPlanNode {
         input: DslPlanKey,
         keys: Vec<Expr>,
         aggs: Vec<Expr>,
+        predicates: Vec<Expr>,
         maintain_order: bool,
         options: Arc<GroupbyOptions>,
         apply: Option<(PlanCallback<DataFrame, DataFrame>, SchemaRef)>,
@@ -218,13 +219,15 @@ fn convert_dsl_plan_to_serializable_plan(
             input,
             keys,
             aggs,
+            predicates,
             maintain_order,
             options,
             apply,
         } => SP::GroupBy {
             input: dsl_plan_key(input, arenas),
             keys: keys.clone(),
             aggs: aggs.clone(),
+            predicates: predicates.clone(),
             maintain_order: *maintain_order,
             options: options.clone(),
             apply: apply.clone(),
@@ -454,13 +457,15 @@ fn try_convert_serializable_plan_to_dsl_plan(
             input,
             keys,
             aggs,
+            predicates,
             maintain_order,
             options,
             apply,
         } => Ok(DP::GroupBy {
             input: get_dsl_plan(*input, ser_dsl_plan, arenas)?,
             keys: keys.clone(),
             aggs: aggs.clone(),
+            predicates: predicates.clone(),
             maintain_order: *maintain_order,
             options: options.clone(),
             apply: apply.clone(),

crates/polars-plan/src/plans/conversion/dsl_to_ir/mod.rs

@@ -5,6 +5,7 @@ use hive::hive_partitions_from_paths;
 use polars_core::chunked_array::cast::CastOptions;
 use polars_core::config::verbose;
 use polars_utils::format_pl_smallstr;
+use polars_utils::itertools::Itertools;
 use polars_utils::plpath::PlPath;
 use polars_utils::unique_id::UniqueId;
 
@@ -460,11 +461,51 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
         DslPlan::GroupBy {
             input,
             keys,
-            aggs,
+            predicates,
+            mut aggs,
             apply,
             maintain_order,
             options,
         } => {
+            // If the group by contains any predicates, we update the plan by turning the
+            // predicates into aggregations and filtering on them. Then, we recursively call
+            // this function.
+            if !predicates.is_empty() {
+                let predicate_names = (0..predicates.len())
+                    .map(|i| PlSmallStr::from_string(format!("__POLARS_HAVING_{i}")))
+                    .collect::<Arc<[_]>>();
+                let predicates = predicates
+                    .into_iter()
+                    .zip(predicate_names.iter())
+                    .map(|(p, name)| p.alias(name.clone()))
+                    .collect_vec();
+                aggs.extend(predicates);
+
+                let lp = DslPlan::GroupBy {
+                    input,
+                    keys,
+                    predicates: vec![],
+                    aggs,
+                    apply,
+                    maintain_order,
+                    options,
+                };
+                let lp = DslBuilder::from(lp)
+                    .filter(
+                        all_horizontal(
+                            predicate_names.iter().map(|n| col(n.clone())).collect_vec(),
+                        )
+                        .unwrap(),
+                    )
+                    .drop(Selector::ByName {
+                        names: predicate_names,
+                        strict: true,
+                    })
+                    .build();
+                return to_alp_impl(lp, ctxt);
+            }
+
+            // NOTE: As we went into this branch, we know that no predicates are provided.
             let input =
                 to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(group_by)))?;
 
@@ -504,7 +545,6 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
                 maintain_order,
                 options,
             };
-
             return run_conversion(lp, ctxt, "group_by")
                 .map_err(|e| e.context(failed_here!(group_by)));
         },
@@ -1476,6 +1516,7 @@ fn resolve_group_by(
 
     Ok((keys, aggs, Arc::new(output_schema)))
 }
+
 fn stats_helper<F, E>(condition: F, expr: E, schema: &Schema) -> Vec<Expr>
 where
     F: Fn(&DataType) -> bool,

crates/polars-python/src/lazygroupby.rs

@@ -19,6 +19,15 @@ pub struct PyLazyGroupBy {
 
 #[pymethods]
 impl PyLazyGroupBy {
+    fn having(&self, predicates: Vec<PyExpr>) -> PyLazyGroupBy {
+        let mut lgb = self.lgb.clone().unwrap();
+        let predicates = predicates.to_exprs();
+        for predicate in predicates.into_iter() {
+            lgb = lgb.having(predicate);
+        }
+        PyLazyGroupBy { lgb: Some(lgb) }
+    }
+
     fn agg(&self, aggs: Vec<PyExpr>) -> PyLazyFrame {
         let lgb = self.lgb.clone().unwrap();
         let aggs = aggs.to_exprs();

py-polars/src/polars/_plr.pyi

@@ -2282,6 +2282,7 @@ class PyLazyGroupBy:
     def agg(self, aggs: list[PyExpr]) -> PyLazyFrame: ...
     def head(self, n: int) -> PyLazyFrame: ...
     def tail(self, n: int) -> PyLazyFrame: ...
+    def having(self, predicates: list[PyExpr]) -> PyLazyGroupBy: ...
     def map_groups(
         self, lambda_function: Any, schema: Schema | None
     ) -> PyLazyFrame: ...

py-polars/src/polars/dataframe/frame.py

@@ -7063,7 +7063,9 @@ def group_by(
                     f"    group_by({value!r})"
                 )
                 raise TypeError(msg)
-        return GroupBy(self, *by, **named_by, maintain_order=maintain_order)
+        return GroupBy(
+            self, *by, **named_by, maintain_order=maintain_order, predicates=None
+        )
 
     @deprecate_renamed_parameter("by", "group_by", version="0.20.14")
     def rolling(
@@ -7220,6 +7222,7 @@ def rolling(
             offset=offset,
             closed=closed,
             group_by=group_by,
+            predicates=None,
         )
 
     @deprecate_renamed_parameter("by", "group_by", version="0.20.14")
@@ -7540,6 +7543,7 @@ def group_by_dynamic(
             closed=closed,
             group_by=group_by,
             start_by=start_by,
+            predicates=None,
         )
 
     @deprecate_renamed_parameter("by", "group_by", version="0.20.14")