Repository: modin-project/modin

Issue:
HDK: Add support for the quantile group by aggregation.



Discussion:

Resolving patch:
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/calcite_builder.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/calcite_builder.py
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/calcite_builder.py
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/calcite_builder.py
@@ -46,6 +46,7 @@
     InputRefExpr,
     LiteralExpr,
     OpExpr,
+    _quantile_agg_dtype,
     build_if_then_else,
     build_row_idx_filter_expr,
 )
@@ -377,11 +378,42 @@ def gen_reduce_expr(self):
                 self._dtype,
             )
 
+    class QuantileAggregate(CompoundAggregateWithColArg):
+        """
+        A QUANTILE aggregate generator.
+
+        Parameters
+        ----------
+        builder : CalciteBuilder
+            A builder to use for translation.
+        arg : List of BaseExpr
+            A list of 3 values:
+                0. InputRefExpr - the column to compute the quantiles for.
+                1. LiteralExpr - the quantile value.
+                2. str - the interpolation method to use.
+        """
+
+        def __init__(self, builder, arg):
+            super().__init__(
+                "QUANTILE",
+                builder,
+                arg,
+                _quantile_agg_dtype(arg[0]._dtype),
+            )
+            self._interpolation = arg[2].val.upper()
+
+        def gen_agg_exprs(self):
+            exprs = super().gen_agg_exprs()
+            for expr in exprs.values():
+                expr.interpolation = self._interpolation
+            return exprs
+
     _compound_aggregates = {
         "std": StdAggregate,
         "skew": SkewAggregate,
         "nlargest": TopkAggregate,
         "nsmallest": TopkAggregate,
+        "quantile": QuantileAggregate,
     }
 
     class InputContext:
@@ -419,7 +451,6 @@ class InputContext:
             "min": "MIN",
             "size": "COUNT",
             "count": "COUNT",
-            "median": "APPROX_QUANTILE",
         }
         _no_arg_aggregates = {"size"}
 
@@ -662,7 +693,7 @@ def __exit__(self, type, value, traceback):
     _bool_cast_aggregates = {
         "sum": _get_dtype(int),
         "mean": _get_dtype(float),
-        "median": _get_dtype(float),
+        "quantile": _get_dtype(float),
     }
 
     def __init__(self):
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
@@ -646,26 +646,19 @@ def generate_by_name(by):
 
         agg_exprs = OrderedDict()
         if isinstance(agg, str):
-            if agg == "nlargest" or agg == "nsmallest":
-                n = kwargs["agg_kwargs"]["n"]
-                if agg == "nsmallest":
-                    n = -n
-                n = LiteralExpr(n)
-                for col in agg_cols:
-                    agg_exprs[col] = AggregateExpr(agg, [base.ref(col), n])
-            else:
-                for col in agg_cols:
-                    agg_exprs[col] = AggregateExpr(agg, base.ref(col))
+            col_to_ref = {col: base.ref(col) for col in agg_cols}
+            self._add_agg_exprs(agg, col_to_ref, kwargs, agg_exprs)
         else:
             assert isinstance(agg, dict), "unsupported aggregate type"
             multiindex = any(isinstance(v, list) for v in agg.values())
-            for k, v in agg.items():
-                if isinstance(v, list):
-                    for item in v:
-                        agg_exprs[(k, item)] = AggregateExpr(item, base.ref(k))
+            for col, aggs in agg.items():
+                if isinstance(aggs, list):
+                    for a in aggs:
+                        col_to_ref = {(col, a): base.ref(col)}
+                        self._add_agg_exprs(a, col_to_ref, kwargs, agg_exprs)
                 else:
-                    col_name = (k, v) if multiindex else k
-                    agg_exprs[col_name] = AggregateExpr(v, base.ref(k))
+                    col_to_ref = {((col, aggs) if multiindex else col): base.ref(col)}
+                    self._add_agg_exprs(aggs, col_to_ref, kwargs, agg_exprs)
         new_columns.extend(agg_exprs.keys())
         new_dtypes.extend((x._dtype for x in agg_exprs.values()))
         new_columns = Index.__new__(Index, data=new_columns, dtype=self.columns.dtype)
@@ -692,6 +685,37 @@ def generate_by_name(by):
                 )
         return new_frame
 
+    def _add_agg_exprs(self, agg, col_to_ref, kwargs, agg_exprs):
+        """
+        Add `AggregateExpr`s for each column to `agg_exprs`.
+
+        Parameters
+        ----------
+        agg : str
+        col_to_ref : dict
+        kwargs : dict
+        agg_exprs : dict
+        """
+        if agg == "nlargest" or agg == "nsmallest":
+            n = kwargs["agg_kwargs"]["n"]
+            if agg == "nsmallest":
+                n = -n
+            n = LiteralExpr(n)
+            for col, ref in col_to_ref.items():
+                agg_exprs[col] = AggregateExpr(agg, [ref, n])
+        elif agg == "median" or agg == "quantile":
+            agg_kwargs = kwargs["agg_kwargs"]
+            q = agg_kwargs.get("q", 0.5)
+            if not isinstance(q, float):
+                raise NotImplementedError("Non-float quantile")
+            q = LiteralExpr(q)
+            interpolation = LiteralExpr(agg_kwargs.get("interpolation", "linear"))
+            for col, ref in col_to_ref.items():
+                agg_exprs[col] = AggregateExpr("quantile", [ref, q, interpolation])
+        else:
+            for col, ref in col_to_ref.items():
+                agg_exprs[col] = AggregateExpr(agg, ref)
+
     def _groupby_head_tail(
         self, agg: str, n: int, cols: Iterable[str]
     ) -> "HdkOnNativeDataframe":
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/expr.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/expr.py
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/expr.py
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/expr.py
@@ -82,6 +82,21 @@ def _get_common_dtype(lhs_dtype, rhs_dtype):
 _aggs_with_float_result = {"mean", "median", "std", "skew"}
 
 
+def _quantile_agg_dtype(dtype):
+    """
+    Compute the quantile aggregate data type.
+
+    Parameters
+    ----------
+    dtype : dtype
+
+    Returns
+    -------
+    dtype
+    """
+    return dtype if is_datetime64_any_dtype(dtype) else _get_dtype(float)
+
+
 def _agg_dtype(agg, dtype):
     """
     Compute aggregate data type.
@@ -104,6 +119,8 @@ def _agg_dtype(agg, dtype):
         return _get_dtype(int)
     elif agg in _aggs_with_float_result:
         return _get_dtype(float)
+    elif agg == "quantile":
+        return _quantile_agg_dtype(dtype)
     else:
         raise NotImplementedError(f"unsupported aggregate {agg}")
 


Test patch:
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
@@ -1287,6 +1287,19 @@ def topk(df, **kwargs):
 
         run_and_compare(topk, data=self.h2o_data)
 
+    @pytest.mark.parametrize("time", [False, True])
+    @pytest.mark.parametrize("q", [0.1, 0.5, 1.0])
+    @pytest.mark.parametrize(
+        "interpolation", ["linear", "lower", "higher", "midpoint", "nearest"]
+    )
+    def test_quantile(self, time, q, interpolation):
+        def quantile(df, **kwargs):
+            if time:
+                df["v1"] = df["v1"].astype("datetime64[ns]")
+            return df.groupby("id4")[["v1", "v2", "v3"]].quantile(q, interpolation)
+
+        run_and_compare(quantile, data=self.h2o_data)
+
 
 class TestAgg:
     data = {