Support quantile in cudf-polars grouped aggregations (#18634)

- Closes #17123

Authors:
  - Lawrence Mitchell (https://github.com/wence-)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: #18634

Author: wence-

python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py

@@ -66,9 +66,13 @@ def __init__(
                 else plc.types.NullPolicy.INCLUDE
             )
         elif name == "quantile":
-            _, quantile = self.children
+            child, quantile = self.children
             if not isinstance(quantile, Literal):
                 raise NotImplementedError("Only support literal quantile values")
+            if options == "equiprobable":
+                raise NotImplementedError("Quantile with equiprobable interpolation")
+            if plc.traits.is_duration(child.dtype):
+                raise NotImplementedError("Quantile with duration data type")
             req = plc.aggregation.quantile(
                 quantiles=[quantile.value.as_py()], interp=Agg.interp_mapping[options]
             )

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -1012,7 +1012,10 @@ def do_evaluate(
                 # A count aggregation, we need a column so use a key column
                 col = keys[0].obj
             elif isinstance(value, expr.Agg):
-                (child,) = value.children
+                if value.name == "quantile":
+                    child = value.children[0]
+                else:
+                    (child,) = value.children
                 col = child.evaluate(df).obj
             else:
                 # Anything else, we pre-evaluate

python/cudf_polars/cudf_polars/dsl/utils/aggregations.py

@@ -72,7 +72,12 @@ def decompose_single_agg(
     if isinstance(agg, (expr.Literal, expr.LiteralColumn)):
         return [], named_expr, False
     if isinstance(agg, expr.Agg):
-        (child,) = agg.children
+        if agg.name == "quantile":
+            # Second child the requested quantile (which is asserted
+            # to be a literal on construction)
+            child = agg.children[0]
+        else:
+            (child,) = agg.children
         needs_masking = agg.name in {"min", "max"} and plc.traits.is_floating_point(
             child.dtype
         )

python/cudf_polars/cudf_polars/testing/plugin.py

@@ -160,6 +160,7 @@ def pytest_configure(config: pytest.Config) -> None:
     "tests/unit/operations/arithmetic/test_list.py::test_list_arithmetic_values[exec_op_with_expr_no_type_coercion-broadcast_both-none]": "cudf-polars doesn't nullify division by zero",
     "tests/unit/operations/arithmetic/test_list.py::test_list_arithmetic_values[exec_op_with_expr_no_type_coercion-broadcast_none-none]": "cudf-polars doesn't nullify division by zero",
     "tests/unit/operations/test_abs.py::test_abs_duration": "Need to raise for unsupported uops on timelike values",
+    "tests/unit/operations/test_group_by.py::test_group_by_shorthand_quantile": "libcudf quantiles are round to nearest ties to even, polars quantiles are round to nearest ties away from zero",
     "tests/unit/operations/test_group_by.py::test_group_by_mean_by_dtype[input7-expected7-Float32-Float32]": "Mismatching dtypes, needs cudf#15852",
     "tests/unit/operations/test_group_by.py::test_group_by_mean_by_dtype[input10-expected10-Date-output_dtype10]": "Unsupported groupby-agg for a particular dtype",
     "tests/unit/operations/test_group_by.py::test_group_by_mean_by_dtype[input11-expected11-input_dtype11-output_dtype11]": "Unsupported groupby-agg for a particular dtype",

python/cudf_polars/tests/expressions/test_agg.py

@@ -125,6 +125,18 @@ def test_quantile_invalid_q(df):
     assert_ir_translation_raises(q, NotImplementedError)
 
 
+def test_quantile_equiprobable_unsupported(df):
+    expr = pl.col("a").quantile(0.5, interpolation="equiprobable")
+    q = df.select(expr)
+    assert_ir_translation_raises(q, NotImplementedError)
+
+
+def test_quantile_duration_unsupported():
+    df = pl.LazyFrame({"a": pl.Series([1, 2, 3, 4], dtype=pl.Duration("ns"))})
+    q = df.select(pl.col("a").quantile(0.5))
+    assert_ir_translation_raises(q, NotImplementedError)
+
+
 @pytest.mark.parametrize(
     "op", [pl.Expr.min, pl.Expr.nan_min, pl.Expr.max, pl.Expr.nan_max]
 )

python/cudf_polars/tests/test_groupby.py

@@ -81,6 +81,11 @@ def keys(request):
         [pl.col("float").round(decimals=1).sum()],
         [pl.col("int").first(), pl.col("float").last()],
         [pl.col("int").sum(), pl.col("string").str.replace("h", "foo", literal=True)],
+        [pl.col("float").quantile(0.3, interpolation="nearest")],
+        [pl.col("float").quantile(0.3, interpolation="higher")],
+        [pl.col("float").quantile(0.3, interpolation="lower")],
+        [pl.col("float").quantile(0.3, interpolation="midpoint")],
+        [pl.col("float").quantile(0.3, interpolation="linear")],
         [
             pl.col("datetime").max(),
             pl.col("datetime").max().dt.is_leap_year().alias("leapyear"),