fix: Support varying quantile values per group in group_by aggregation

Co-authored-by: Claude <noreply@anthropic.com>

Author: wtn

crates/polars-expr/src/expressions/aggregation.rs

@@ -483,6 +483,57 @@ impl AggQuantileExpr {
         );
         quantile.get(0).unwrap().try_extract()
     }
+
+    fn get_quantile_from_scalar(&self, quantile_ac: &AggregationContext) -> PolarsResult<f64> {
+        let quantile_col = quantile_ac.get_values();
+        polars_ensure!(quantile_col.len() <= 1, ComputeError:
+            "polars only supports computing a single quantile; \
+            make sure the 'quantile' expression input produces a single quantile"
+        );
+        quantile_col.get(0).unwrap().try_extract()
+    }
+
+    /// Compute quantile per group when quantile values vary by group.
+    fn agg_quantile_per_group(
+        &self,
+        ac: &mut AggregationContext,
+        quantile_ac: &mut AggregationContext,
+        keep_name: PlSmallStr,
+    ) -> PolarsResult<Column> {
+        let ac_list = ac.aggregated_as_list();
+        let quantile_values = quantile_ac.flat_naive();
+
+        // Get quantile values as f64
+        let quantile_values = quantile_values.cast(&DataType::Float64)?;
+        let quantile_arr = quantile_values.f64()?;
+
+        let method = self.method;
+        let result = ac_list
+            .amortized_iter()
+            .zip(quantile_arr.iter())
+            .map(|(opt_s, opt_q)| match (opt_s, opt_q) {
+                (Some(s), Some(q)) => {
+                    let s = s.as_ref();
+                    if !(0.0..=1.0).contains(&q) {
+                        Ok(None)
+                    } else {
+                        let scalar = s.quantile_reduce(q, method)?;
+                        // Extract f64 from the scalar, returning None if null
+                        if scalar.is_null() {
+                            Ok(None)
+                        } else {
+                            scalar.value().try_extract::<f64>().map(Some)
+                        }
+                    }
+                },
+                _ => Ok(None),
+            })
+            .collect::<PolarsResult<Float64Chunked>>()?
+            .with_name(keep_name)
+            .into_column();
+
+        Ok(result)
+    }
 }
 
 impl PhysicalExpr for AggQuantileExpr {
@@ -505,6 +556,7 @@ impl PhysicalExpr for AggQuantileExpr {
         state: &ExecutionState,
     ) -> PolarsResult<AggregationContext<'a>> {
         let mut ac = self.input.evaluate_on_groups(df, groups, state)?;
+        let mut quantile_ac = self.quantile.evaluate_on_groups(df, groups, state)?;
 
         // AggregatedScalar has no defined group structure. We fix it up here, so that we can
         // reliably call `agg_quantile` functions with the groups.
@@ -513,27 +565,41 @@ impl PhysicalExpr for AggQuantileExpr {
         // don't change names by aggregations as is done in polars-core
         let keep_name = ac.get_values().name().clone();
 
-        let quantile = self.get_quantile(df, state)?;
+        // Check if quantile is a literal/scalar (same for all groups) or varies per group
+        let is_uniform_quantile = quantile_ac.is_literal()
+            || matches!(quantile_ac.agg_state(), AggState::LiteralScalar(_));
 
-        if let AggState::LiteralScalar(c) = &mut ac.state {
-            *c = c
-                .quantile_reduce(quantile, self.method)?
-                .into_column(keep_name);
-            return Ok(ac);
-        }
+        if is_uniform_quantile {
+            // Fast path: single quantile value for all groups
+            let quantile = self.get_quantile_from_scalar(&quantile_ac)?;
 
-        // SAFETY:
-        // groups are in bounds
-        let mut agg = unsafe {
-            ac.flat_naive()
-                .into_owned()
-                .agg_quantile(ac.groups(), quantile, self.method)
-        };
-        agg.rename(keep_name);
-        Ok(AggregationContext::from_agg_state(
-            AggregatedScalar(agg),
-            Cow::Borrowed(groups),
-        ))
+            if let AggState::LiteralScalar(c) = &mut ac.state {
+                *c = c
+                    .quantile_reduce(quantile, self.method)?
+                    .into_column(keep_name);
+                return Ok(ac);
+            }
+
+            // SAFETY:
+            // groups are in bounds
+            let mut agg = unsafe {
+                ac.flat_naive()
+                    .into_owned()
+                    .agg_quantile(ac.groups(), quantile, self.method)
+            };
+            agg.rename(keep_name);
+            Ok(AggregationContext::from_agg_state(
+                AggregatedScalar(agg),
+                Cow::Borrowed(groups),
+            ))
+        } else {
+            // Slow path: different quantile value per group
+            let agg = self.agg_quantile_per_group(&mut ac, &mut quantile_ac, keep_name)?;
+            Ok(AggregationContext::from_agg_state(
+                AggregatedScalar(agg),
+                Cow::Borrowed(groups),
+            ))
+        }
     }
 
     fn to_field(&self, input_schema: &Schema) -> PolarsResult<Field> {

py-polars/tests/unit/operations/aggregation/test_aggregations.py

@@ -27,6 +27,45 @@ def test_quantile_expr_input() -> None:
     )
 
 
+def test_quantile_varying_by_group_20951() -> None:
+    df = pl.DataFrame(
+        {
+            "value": [1, 2, 1, 2],
+            "quantile": [0, 0, 1, 1],
+        }
+    )
+    result = df.group_by(pl.col.quantile).agg(
+        pl.col.value.quantile(pl.col.quantile.first())
+    ).sort("quantile")
+
+    expected = pl.DataFrame(
+        {
+            "quantile": [0, 1],
+            "value": [1.0, 2.0],
+        }
+    )
+    assert_frame_equal(result, expected)
+
+    df = pl.DataFrame(
+        {
+            "group": ["a", "a", "a", "b", "b", "b"],
+            "value": [1, 2, 3, 10, 20, 30],
+            "q": [0.0, 0.0, 0.0, 0.5, 0.5, 0.5],
+        }
+    )
+    result = df.group_by("group").agg(
+        pl.col.value.quantile(pl.col.q.first())
+    ).sort("group")
+
+    expected = pl.DataFrame(
+        {
+            "group": ["a", "b"],
+            "value": [1.0, 20.0],
+        }
+    )
+    assert_frame_equal(result, expected)
+
+
 def test_boolean_aggs() -> None:
     df = pl.DataFrame({"bool": [True, False, None, True]})