refactor: fix test_dataframe_aggregates_median on agg_ops.QuantileOp

Author: chelsea-lin

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

@@ -515,14 +515,18 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    # TODO: Support interpolation argument
-    # TODO: Support percentile_disc
-    result: sge.Expression = sge.func("PERCENTILE_CONT", column.expr, sge.convert(op.q))
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=expr, to="INT64")
+
+    result: sge.Expression = sge.func("PERCENTILE_CONT", 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.
+        # 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

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

@@ -1,14 +1,17 @@
 WITH `bfcte_0` AS (
   SELECT
+    `bool_col`,
     `int64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    PERCENTILE_CONT(`int64_col`, 0.5) OVER () AS `bfcol_1`,
-    CAST(FLOOR(PERCENTILE_CONT(`int64_col`, 0.5) OVER ()) AS INT64) AS `bfcol_2`
+    PERCENTILE_CONT(`int64_col`, 0.5) OVER () AS `bfcol_4`,
+    PERCENTILE_CONT(CAST(`bool_col` AS INT64), 0.5) OVER () AS `bfcol_5`,
+    CAST(FLOOR(PERCENTILE_CONT(`int64_col`, 0.5) OVER ()) AS INT64) AS `bfcol_6`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `quantile`,
-  `bfcol_2` AS `quantile_floor`
+  `bfcol_4` AS `int64`,
+  `bfcol_5` AS `bool`,
+  `bfcol_6` AS `int64_w_floor`
 FROM `bfcte_1`

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

@@ -491,12 +491,12 @@ def test_qcut(scalar_types_df: bpd.DataFrame, snapshot):
 
 
 def test_quantile(scalar_types_df: bpd.DataFrame, snapshot):
-    col_name = "int64_col"
-    bf_df = scalar_types_df[[col_name]]
+    bf_df = scalar_types_df[["int64_col", "bool_col"]]
     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
+        "int64": agg_ops.QuantileOp(q=0.5).as_expr("int64_col"),
+        "bool": agg_ops.QuantileOp(q=0.5).as_expr("bool_col"),
+        "int64_w_floor": agg_ops.QuantileOp(q=0.5, should_floor_result=True).as_expr(
+            "int64_col"
         ),
     }
     sql = _apply_unary_agg_ops(