feat(rust, python): allow expression as quantile input (#5751)

pola-rs · Dec 8, 2022 · 3b7c0e5 · 3b7c0e5
1 parent 5c5cc58
commit 3b7c0e5
Show file tree

Hide file tree

Showing 20 changed files with 248 additions and 143 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/expr.rs b/polars/polars-lazy/polars-plan/src/dsl/expr.rs
@@ -240,7 +240,7 @@ pub enum AggExpr {
     Count(Box<Expr>),
     Quantile {
         expr: Box<Expr>,
-        quantile: f64,
+        quantile: Box<Expr>,
         interpol: QuantileInterpolOptions,
     },
     Sum(Box<Expr>),

diff --git a/polars/polars-lazy/polars-plan/src/dsl/functions.rs b/polars/polars-lazy/polars-plan/src/dsl/functions.rs
@@ -711,7 +711,7 @@ pub fn median(name: &str) -> Expr {
 }
 
 /// Find a specific quantile of all the values in this Expression.
-pub fn quantile(name: &str, quantile: f64, interpol: QuantileInterpolOptions) -> Expr {
+pub fn quantile(name: &str, quantile: Expr, interpol: QuantileInterpolOptions) -> Expr {
     col(name).quantile(quantile, interpol)
 }
 

diff --git a/polars/polars-lazy/polars-plan/src/dsl/mod.rs b/polars/polars-lazy/polars-plan/src/dsl/mod.rs
@@ -368,10 +368,10 @@ impl Expr {
     }
 
     /// Compute the quantile per group.
-    pub fn quantile(self, quantile: f64, interpol: QuantileInterpolOptions) -> Self {
+    pub fn quantile(self, quantile: Expr, interpol: QuantileInterpolOptions) -> Self {
         AggExpr::Quantile {
             expr: Box::new(self),
-            quantile,
+            quantile: Box::new(quantile),
             interpol,
         }
         .into()

diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/aexpr.rs b/polars/polars-lazy/polars-plan/src/logical_plan/aexpr.rs
@@ -28,7 +28,7 @@ pub enum AAggExpr {
     List(Node),
     Quantile {
         expr: Node,
-        quantile: f64,
+        quantile: Node,
         interpol: QuantileInterpolOptions,
     },
     Sum(Node),

diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/conversion.rs b/polars/polars-lazy/polars-plan/src/logical_plan/conversion.rs
@@ -77,7 +77,7 @@ pub fn to_aexpr(expr: Expr, arena: &mut Arena<AExpr>) -> Node {
                     interpol,
                 } => AAggExpr::Quantile {
                     expr: to_aexpr(*expr, arena),
-                    quantile,
+                    quantile: to_aexpr(*quantile, arena),
                     interpol,
                 },
                 AggExpr::Sum(expr) => AAggExpr::Sum(to_aexpr(*expr, arena)),
@@ -540,10 +540,11 @@ pub fn node_to_expr(node: Node, expr_arena: &Arena<AExpr>) -> Expr {
                 quantile,
                 interpol,
             } => {
-                let exp = node_to_expr(expr, expr_arena);
+                let expr = node_to_expr(expr, expr_arena);
+                let quantile = node_to_expr(quantile, expr_arena);
                 AggExpr::Quantile {
-                    expr: Box::new(exp),
-                    quantile,
+                    expr: Box::new(expr),
+                    quantile: Box::new(quantile),
                     interpol,
                 }
                 .into()

diff --git a/polars/polars-lazy/src/frame/mod.rs b/polars/polars-lazy/src/frame/mod.rs
@@ -684,7 +684,7 @@ impl LazyFrame {
     ///        .agg([
     ///            col("rain").min(),
     ///            col("rain").sum(),
-    ///            col("rain").quantile(0.5, QuantileInterpolOptions::Nearest).alias("median_rain"),
+    ///            col("rain").quantile(lit(0.5), QuantileInterpolOptions::Nearest).alias("median_rain"),
     ///        ])
     /// }
     /// ```
@@ -982,7 +982,7 @@ impl LazyFrame {
     }
 
     /// Aggregate all the columns as their quantile values.
-    pub fn quantile(self, quantile: f64, interpol: QuantileInterpolOptions) -> LazyFrame {
+    pub fn quantile(self, quantile: Expr, interpol: QuantileInterpolOptions) -> LazyFrame {
         self.select_local(vec![col("*").quantile(quantile, interpol)])
     }
 
@@ -1237,7 +1237,7 @@ impl LazyGroupBy {
     ///        .agg([
     ///            col("rain").min(),
     ///            col("rain").sum(),
-    ///            col("rain").quantile(0.5, QuantileInterpolOptions::Nearest).alias("median_rain"),
+    ///            col("rain").quantile(lit(0.5), QuantileInterpolOptions::Nearest).alias("median_rain"),
     ///        ])
     /// }
     /// ```

diff --git a/polars/polars-lazy/src/lib.rs b/polars/polars-lazy/src/lib.rs
@@ -107,7 +107,7 @@
 //!     .agg([
 //!         col("rain").min(),
 //!         col("rain").sum(),
-//!         col("rain").quantile(0.5, QuantileInterpolOptions::Nearest).alias("median_rain"),
+//!         col("rain").quantile(lit(0.5), QuantileInterpolOptions::Nearest).alias("median_rain"),
 //!     ])
 //!     .sort("date", Default::default())
 //!     .collect()

diff --git a/polars/polars-lazy/src/physical_plan/expressions/aggregation.rs b/polars/polars-lazy/src/physical_plan/expressions/aggregation.rs
@@ -457,32 +457,46 @@ impl PartitionedAggregation for AggregationExpr {
 }
 
 pub struct AggQuantileExpr {
-    pub(crate) expr: Arc<dyn PhysicalExpr>,
-    pub(crate) quantile: f64,
+    pub(crate) input: Arc<dyn PhysicalExpr>,
+    pub(crate) quantile: Arc<dyn PhysicalExpr>,
     pub(crate) interpol: QuantileInterpolOptions,
 }
 
 impl AggQuantileExpr {
     pub fn new(
-        expr: Arc<dyn PhysicalExpr>,
-        quantile: f64,
+        input: Arc<dyn PhysicalExpr>,
+        quantile: Arc<dyn PhysicalExpr>,
         interpol: QuantileInterpolOptions,
     ) -> Self {
         Self {
-            expr,
+            input,
             quantile,
             interpol,
         }
     }
+
+    fn get_quantile(&self, df: &DataFrame, state: &ExecutionState) -> PolarsResult<f64> {
+        let quantile = self.quantile.evaluate(df, state)?;
+        if quantile.len() > 1 {
+            return Err(PolarsError::ComputeError(
+                "Polars only supports computing a single quantile. \
+            Make sure the 'quantile' expression input produces a single quantile."
+                    .into(),
+            ));
+        }
+        quantile.get(0).try_extract::<f64>()
+    }
 }
 
 impl PhysicalExpr for AggQuantileExpr {
     fn as_expression(&self) -> Option<&Expr> {
         None
     }
 
-    fn evaluate(&self, _df: &DataFrame, _state: &ExecutionState) -> PolarsResult<Series> {
-        unimplemented!()
+    fn evaluate(&self, df: &DataFrame, state: &ExecutionState) -> PolarsResult<Series> {
+        let input = self.input.evaluate(df, state)?;
+        let quantile = self.get_quantile(df, state)?;
+        input.quantile_as_series(quantile, self.interpol)
     }
     #[allow(clippy::ptr_arg)]
     fn evaluate_on_groups<'a>(
@@ -491,23 +505,25 @@ impl PhysicalExpr for AggQuantileExpr {
         groups: &'a GroupsProxy,
         state: &ExecutionState,
     ) -> PolarsResult<AggregationContext<'a>> {
-        let mut ac = self.expr.evaluate_on_groups(df, groups, state)?;
+        let mut ac = self.input.evaluate_on_groups(df, groups, state)?;
         // don't change names by aggregations as is done in polars-core
         let keep_name = ac.series().name().to_string();
 
+        let quantile = self.get_quantile(df, state)?;
+
         // safety:
         // groups are in bounds
         let mut agg = unsafe {
             ac.flat_naive()
                 .into_owned()
-                .agg_quantile(ac.groups(), self.quantile, self.interpol)
+                .agg_quantile(ac.groups(), quantile, self.interpol)
         };
         agg.rename(&keep_name);
         Ok(AggregationContext::new(agg, Cow::Borrowed(groups), true))
     }
 
     fn to_field(&self, input_schema: &Schema) -> PolarsResult<Field> {
-        self.expr.to_field(input_schema)
+        self.input.to_field(input_schema)
     }
 
     fn is_valid_aggregation(&self) -> bool {

diff --git a/polars/polars-lazy/src/physical_plan/planner/expr.rs b/polars/polars-lazy/src/physical_plan/planner/expr.rs
@@ -452,24 +452,28 @@ pub(crate) fn create_physical_expr(
                 } => {
                     // todo! add schema to get correct output type
                     let input = create_physical_expr(expr, ctxt, expr_arena, schema)?;
-                    match ctxt {
-                        Context::Aggregation => {
-                            Ok(Arc::new(AggQuantileExpr::new(input, quantile, interpol)))
-                        }
-                        Context::Default => {
-                            let function = SpecialEq::new(Arc::new(move |s: &mut [Series]| {
-                                let s = std::mem::take(&mut s[0]);
-                                s.quantile_as_series(quantile, interpol)
-                            })
-                                as Arc<dyn SeriesUdf>);
-                            Ok(Arc::new(ApplyExpr::new_minimal(
-                                vec![input],
-                                function,
-                                node_to_expr(expression, expr_arena),
-                                ApplyOptions::ApplyFlat,
-                            )))
-                        }
-                    }
+                    let quantile = create_physical_expr(quantile, ctxt, expr_arena, schema)?;
+                    Ok(Arc::new(AggQuantileExpr::new(input, quantile, interpol)))
+                    //
+                    // match ctxt {
+                    //     Context::Aggregation => {
+                    //
+                    //         Ok(Arc::new(AggQuantileExpr::new(input, quantile, interpol)))
+                    //     }
+                    //     Context::Default => {
+                    //         let function = SpecialEq::new(Arc::new(move |s: &mut [Series]| {
+                    //             let s = std::mem::take(&mut s[0]);
+                    //             s.quantile_as_series(quantile, interpol)
+                    //         })
+                    //             as Arc<dyn SeriesUdf>);
+                    //         Ok(Arc::new(ApplyExpr::new_minimal(
+                    //             vec![input],
+                    //             function,
+                    //             node_to_expr(expression, expr_arena),
+                    //             ApplyOptions::ApplyFlat,
+                    //         )))
+                    //     }
+                    // }
                 }
                 AAggExpr::AggGroups(expr) => {
                     if let Context::Default = ctxt {

diff --git a/polars/polars-lazy/src/tests/aggregations.rs b/polars/polars-lazy/src/tests/aggregations.rs
@@ -64,36 +64,6 @@ fn test_lazy_agg_scan() {
     assert!(df.frame_equal_missing(&lf().collect().unwrap().mean()));
 }
 
-#[test]
-fn test_lazy_df_aggregations() {
-    let df = load_df();
-
-    assert!(df
-        .clone()
-        .lazy()
-        .min()
-        .collect()
-        .unwrap()
-        .frame_equal_missing(&df.min()));
-    assert!(df
-        .clone()
-        .lazy()
-        .median()
-        .collect()
-        .unwrap()
-        .frame_equal_missing(&df.median()));
-    assert!(df
-        .clone()
-        .lazy()
-        .quantile(0.5, QuantileInterpolOptions::default())
-        .collect()
-        .unwrap()
-        .frame_equal_missing(
-            &df.quantile(0.5, QuantileInterpolOptions::default())
-                .unwrap()
-        ));
-}
-
 #[test]
 fn test_cumsum_agg_as_key() -> PolarsResult<()> {
     let df = df![

diff --git a/polars/polars-lazy/src/tests/queries.rs b/polars/polars-lazy/src/tests/queries.rs
@@ -151,42 +151,6 @@ fn test_lazy_pushdown_through_agg() {
     assert_eq!(bar.get(0), AnyValue::Float64(1.3));
 }
 
-#[test]
-#[cfg(feature = "temporal")]
-fn test_lazy_agg() {
-    let s0 = DateChunked::parse_from_str_slice(
-        "date",
-        &[
-            "2020-08-21",
-            "2020-08-21",
-            "2020-08-22",
-            "2020-08-23",
-            "2020-08-22",
-        ],
-        "%Y-%m-%d",
-    )
-    .into_series();
-    let s1 = Series::new("temp", [20, 10, 7, 9, 1].as_ref());
-    let s2 = Series::new("rain", [0.2, 0.1, 0.3, 0.1, 0.01].as_ref());
-    let df = DataFrame::new(vec![s0, s1, s2]).unwrap();
-
-    let lf = df
-        .lazy()
-        .groupby([col("date")])
-        .agg([
-            col("rain").min().alias("min"),
-            col("rain").sum().alias("sum"),
-            col("rain")
-                .quantile(0.5, QuantileInterpolOptions::default())
-                .alias("median_rain"),
-        ])
-        .sort("date", Default::default());
-
-    let new = lf.collect().unwrap();
-    let min = new.column("min").unwrap();
-    assert_eq!(min, &Series::new("min", [0.1f64, 0.01, 0.1]));
-}
-
 #[test]
 fn test_lazy_shift() {
     let df = get_df();