Skip to content

Commit c3113be

Browse files
committed
refactor: add agg_ops.ProductOp to the sqlglot compiler
1 parent b487cf1 commit c3113be

File tree

3 files changed

+32
-0
lines changed

3 files changed

+32
-0
lines changed

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -386,6 +386,17 @@ def _(
386386
return apply_window_if_present(sge.func("MIN", column.expr), window)
387387

388388

389+
@UNARY_OP_REGISTRATION.register(agg_ops.NuniqueOp)
390+
def _(
391+
op: agg_ops.NuniqueOp,
392+
column: typed_expr.TypedExpr,
393+
window: typing.Optional[window_spec.WindowSpec] = None,
394+
) -> sge.Expression:
395+
return apply_window_if_present(
396+
sge.func("COUNT", sge.Distinct(expressions=[column.expr])), window
397+
)
398+
399+
389400
@UNARY_OP_REGISTRATION.register(agg_ops.PopVarOp)
390401
def _(
391402
op: agg_ops.PopVarOp,
Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`int64_col`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
COUNT(DISTINCT `int64_col`) AS `bfcol_1`
8+
FROM `bfcte_0`
9+
)
10+
SELECT
11+
`bfcol_1` AS `int64_col`
12+
FROM `bfcte_1`

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -412,6 +412,15 @@ def test_min(scalar_types_df: bpd.DataFrame, snapshot):
412412
snapshot.assert_match(sql_window_partition, "window_partition_out.sql")
413413

414414

415+
def test_nunique(scalar_types_df: bpd.DataFrame, snapshot):
416+
col_name = "int64_col"
417+
bf_df = scalar_types_df[[col_name]]
418+
agg_expr = agg_ops.NuniqueOp().as_expr(col_name)
419+
sql = _apply_unary_agg_ops(bf_df, [agg_expr], [col_name])
420+
421+
snapshot.assert_match(sql, "out.sql")
422+
423+
415424
def test_pop_var(scalar_types_df: bpd.DataFrame, snapshot):
416425
col_names = ["int64_col", "bool_col"]
417426
bf_df = scalar_types_df[col_names]

0 commit comments

Comments
 (0)