select some columns

Author: chelsea-lin

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

@@ -53,7 +53,10 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    return apply_window_if_present(sge.func("DENSE_RANK"), window)
+    # Ranking functions do not support window framing clauses.
+    return apply_window_if_present(
+        sge.func("DENSE_RANK"), window, include_framing_clauses=False
+    )
 
 
 @UNARY_OP_REGISTRATION.register(agg_ops.MaxOp)
@@ -121,7 +124,10 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    return apply_window_if_present(sge.func("RANK"), window)
+    # Ranking functions do not support window framing clauses.
+    return apply_window_if_present(
+        sge.func("RANK"), window, include_framing_clauses=False
+    )
 
 
 @UNARY_OP_REGISTRATION.register(agg_ops.SumOp)

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

@@ -25,6 +25,7 @@
 def apply_window_if_present(
     value: sge.Expression,
     window: typing.Optional[window_spec.WindowSpec] = None,
+    include_framing_clauses: bool = True,
 ) -> sge.Expression:
     if window is None:
         return value
@@ -64,23 +65,7 @@ def apply_window_if_present(
     if not window.bounds and not order:
         return sge.Window(this=value, partition_by=group_by)
 
-    # Ranking and navigation functions do not support window framing clauses.
-    # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts#window-frame-clause
-    no_frame_functions = {
-        "RANK",
-        "DENSE_RANK",
-        "NTILE",
-        "PERCENT_RANK",
-        "CUME_DIST",
-        "ROW_NUMBER",
-        "LAG",
-        "LEAD",
-        "FIRST_VALUE",
-        "LAST_VALUE",
-        "NTH_VALUE",
-    }
-    func_name = value.sql_name().upper() if isinstance(value, sge.Func) else ""
-    if not window.bounds and func_name in no_frame_functions:
+    if not window.bounds and not include_framing_clauses:
         return sge.Window(this=value, partition_by=group_by, order=order)
 
     kind = (

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

@@ -1,16 +1,13 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_0`,
-    `rowindex` AS `bfcol_1`
+    `int64_col` AS `bfcol_0`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    DENSE_RANK() OVER (ORDER BY `bfcol_0` IS NULL ASC NULLS LAST, `bfcol_0` ASC NULLS LAST) AS `bfcol_4`
+    DENSE_RANK() OVER (ORDER BY `bfcol_0` IS NULL ASC NULLS LAST, `bfcol_0` ASC NULLS LAST) AS `bfcol_1`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `bfuid_col_1`,
-  `bfcol_0` AS `int64_col`,
-  `bfcol_4` AS `agg_int64`
+  `bfcol_1` AS `agg_int64`
 FROM `bfcte_1`

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

@@ -1,16 +1,13 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_0`,
-    `rowindex` AS `bfcol_1`
+    `int64_col` AS `bfcol_0`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    RANK() OVER (ORDER BY `bfcol_0` IS NULL ASC NULLS LAST, `bfcol_0` ASC NULLS LAST) AS `bfcol_4`
+    RANK() OVER (ORDER BY `bfcol_0` IS NULL ASC NULLS LAST, `bfcol_0` ASC NULLS LAST) AS `bfcol_1`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `bfuid_col_1`,
-  `bfcol_0` AS `int64_col`,
-  `bfcol_4` AS `agg_int64`
+  `bfcol_1` AS `agg_int64`
 FROM `bfcte_1`

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

@@ -57,7 +57,7 @@ def _apply_unary_window_op(
         window_spec=window_spec,
         output_name=identifiers.ColumnId(new_name),
     )
-    result = array_value.ArrayValue(win_node)
+    result = array_value.ArrayValue(win_node).select_columns([new_name])
 
     sql = result.session._executor.to_sql(result, enable_cache=False)
     return sql