refactor: add agg_ops.QuantileOp to sqlglot compiler

Author: chelsea-lin

bigframes/core/compile/sqlglot/aggregations/op_registration.py

@@ -58,5 +58,7 @@ def __getitem__(self, op: str | agg_ops.WindowOp) -> CompilationFunc:
                 raise ValueError(f"The operator must have a 'name' attribute. Got {op}")
             else:
                 key = typing.cast(str, op.name)
-                return self._registered_ops[key]
+                if key in self._registered_ops:
+                    return self._registered_ops[key]
+                return self._registered_ops[type(op)]
         return self._registered_ops[op]

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -109,13 +109,23 @@ def _(
     return apply_window_if_present(sge.func("MIN", column.expr), window)
 
 
-@UNARY_OP_REGISTRATION.register(agg_ops.SizeUnaryOp)
+@UNARY_OP_REGISTRATION.register(agg_ops.QuantileOp)
 def _(
-    op: agg_ops.SizeUnaryOp,
-    _,
+    op: agg_ops.QuantileOp,
+    column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    return apply_window_if_present(sge.func("COUNT", sge.convert(1)), window)
+    # TODO: Support interpolation argument
+    # TODO: Support percentile_disc
+    result = sge.func("PERCENTILE_CONT", column.expr, sge.convert(op.q))
+    if window is None:
+        # PERCENTILE_CONT is a navigation function, not an aggregate function, so it always needs an OVER clause.
+        result = sge.Window(this=result)
+    else:
+        result = apply_window_if_present(result, window)
+    if op.should_floor_result:
+        result = sge.Cast(this=sge.func("FLOOR", result), to="INT64")
+    return result
 
 
 @UNARY_OP_REGISTRATION.register(agg_ops.RankOp)
@@ -130,6 +140,15 @@ def _(
     )
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.SizeUnaryOp)
+def _(
+    op: agg_ops.SizeUnaryOp,
+    _,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    return apply_window_if_present(sge.func("COUNT", sge.convert(1)), window)
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.SumOp)
 def _(
     op: agg_ops.SumOp,

bigframes/operations/aggregations.py

@@ -223,13 +223,10 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 
 @dataclasses.dataclass(frozen=True)
 class QuantileOp(UnaryAggregateOp):
+    name: typing.ClassVar[str] = "quantile"
     q: float
     should_floor_result: bool = False
 
-    @property
-    def name(self):
-        return f"{int(self.q * 100)}%"
-
     @property
     def order_independent(self) -> bool:
         return True

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_quantile/out.sql

@@ -0,0 +1,14 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    PERCENTILE_CONT(`bfcol_0`, 0.5) OVER () AS `bfcol_1`,
+    CAST(FLOOR(PERCENTILE_CONT(`bfcol_0`, 0.5) OVER ()) AS INT64) AS `bfcol_2`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `quantile`,
+  `bfcol_2` AS `quantile_floor`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/test_unary_compiler.py

@@ -141,6 +141,22 @@ def test_min(scalar_types_df: bpd.DataFrame, snapshot):
     snapshot.assert_match(sql, "out.sql")
 
 
+def test_quantile(scalar_types_df: bpd.DataFrame, snapshot):
+    col_name = "int64_col"
+    bf_df = scalar_types_df[[col_name]]
+    agg_ops_map = {
+        "quantile": agg_ops.QuantileOp(q=0.5).as_expr(col_name),
+        "quantile_floor": agg_ops.QuantileOp(q=0.5, should_floor_result=True).as_expr(
+            col_name
+        ),
+    }
+    sql = _apply_unary_agg_ops(
+        bf_df, list(agg_ops_map.values()), list(agg_ops_map.keys())
+    )
+
+    snapshot.assert_match(sql, "out.sql")
+
+
 def test_rank(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "int64_col"
     bf_df = scalar_types_df[[col_name]]