enable engine tests

Author: chelsea-lin

tests/system/small/engines/test_aggregation.py

@@ -85,17 +85,23 @@ def test_engines_unary_aggregates(
     assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize(
-    "op",
-    [agg_ops.MedianOp],
-)
-def test_sql_engines_unary_aggregates(
+def test_sql_engines_median_op_aggregates(
     scalars_array_value: array_value.ArrayValue,
     bigquery_client: bigquery.Client,
-    op,
 ):
-    # TODO: this is not working??
-    node = apply_agg_to_all_valid(scalars_array_value, op).node
+    node = apply_agg_to_all_valid(
+        scalars_array_value,
+        agg_ops.MedianOp(),
+        # Exclude columns are not supported by Ibis.
+        excluded_cols=[
+            "bytes_col",
+            "date_col",
+            "datetime_col",
+            "time_col",
+            "timestamp_col",
+            "string_col",
+        ],
+    ).node
     left_engine = direct_gbq_execution.DirectGbqExecutor(bigquery_client)
     right_engine = direct_gbq_execution.DirectGbqExecutor(
         bigquery_client, compiler="sqlglot"

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

@@ -1,12 +1,30 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_0`
+    `bytes_col` AS `bfcol_0`,
+    `date_col` AS `bfcol_1`,
+    `datetime_col` AS `bfcol_2`,
+    `int64_col` AS `bfcol_3`,
+    `string_col` AS `bfcol_4`,
+    `time_col` AS `bfcol_5`,
+    `timestamp_col` AS `bfcol_6`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    APPROX_QUANTILES(`bfcol_0`, 2)[OFFSET(1)] AS `bfcol_1`
+    APPROX_QUANTILES(`bfcol_3`, 2)[OFFSET(1)] AS `bfcol_7`,
+    APPROX_QUANTILES(`bfcol_0`, 2)[OFFSET(1)] AS `bfcol_8`,
+    APPROX_QUANTILES(`bfcol_1`, 2)[OFFSET(1)] AS `bfcol_9`,
+    APPROX_QUANTILES(`bfcol_2`, 2)[OFFSET(1)] AS `bfcol_10`,
+    APPROX_QUANTILES(`bfcol_5`, 2)[OFFSET(1)] AS `bfcol_11`,
+    APPROX_QUANTILES(`bfcol_6`, 2)[OFFSET(1)] AS `bfcol_12`,
+    APPROX_QUANTILES(`bfcol_4`, 2)[OFFSET(1)] AS `bfcol_13`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `int64_col`
+  `bfcol_7` AS `int64_col`,
+  `bfcol_8` AS `bytes_col`,
+  `bfcol_9` AS `date_col`,
+  `bfcol_10` AS `datetime_col`,
+  `bfcol_11` AS `time_col`,
+  `bfcol_12` AS `timestamp_col`,
+  `bfcol_13` AS `string_col`
 FROM `bfcte_1`

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

@@ -57,10 +57,18 @@ def test_max(scalar_types_df: bpd.DataFrame, snapshot):
 
 
 def test_median(scalar_types_df: bpd.DataFrame, snapshot):
-    col_name = "int64_col"
-    bf_df = scalar_types_df[[col_name]]
-    agg_expr = agg_ops.MedianOp().as_expr(col_name)
-    sql = _apply_unary_agg_ops(bf_df, [agg_expr], [col_name])
+    bf_df = scalar_types_df
+    ops_map = {
+        "int64_col": agg_ops.MedianOp().as_expr("int64_col"),
+        # Includes columns are not supported by Ibis but supported by BigQuery.
+        "bytes_col": agg_ops.MedianOp().as_expr("bytes_col"),
+        "date_col": agg_ops.MedianOp().as_expr("date_col"),
+        "datetime_col": agg_ops.MedianOp().as_expr("datetime_col"),
+        "time_col": agg_ops.MedianOp().as_expr("time_col"),
+        "timestamp_col": agg_ops.MedianOp().as_expr("timestamp_col"),
+        "string_col": agg_ops.MedianOp().as_expr("string_col"),
+    }
+    sql = _apply_unary_agg_ops(bf_df, list(ops_map.values()), list(ops_map.keys()))
 
     snapshot.assert_match(sql, "out.sql")