fix order by, window key issues

Author: TrevorBergeron

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

@@ -193,13 +193,13 @@ def _get_window_bounds(
 def _compile_group_by_key(key: ex.Expression) -> sge.Expression:
     expr = expression_compiler.expression_compiler.compile_expression(key)
     # The group_by keys has been rewritten by bind_schema_to_node
-    assert isinstance(key, ex.ResolvedDerefOp)
+    assert key.is_scalar_expr and key.is_resolved
 
     # Some types need to be converted to another type to enable groupby
-    if key.dtype == dtypes.FLOAT_DTYPE:
+    if key.output_type == dtypes.FLOAT_DTYPE:
         expr = sge.Cast(this=expr, to="STRING")
-    elif key.dtype == dtypes.GEO_DTYPE:
+    elif key.output_type == dtypes.GEO_DTYPE:
         expr = sge.func("ST_ASBINARY", expr)
-    elif key.dtype == dtypes.JSON_DTYPE:
+    elif key.output_type == dtypes.JSON_DTYPE:
         expr = sge.func("TO_JSON_STRING", expr)
     return expr

bigframes/core/compile/sqlglot/compiler.py

@@ -108,6 +108,7 @@ def _compile_result_node(root: nodes.ResultNode) -> str:
     root = typing.cast(nodes.ResultNode, schema_binding.bind_schema_to_tree(root))
 
     sqlglot_ir = compile_node(rewrite.as_sql_nodes(root), uid_gen)
+    print(sqlglot_ir.sql)
     return sqlglot_ir.sql
 
 

bigframes/core/rewrite/as_sql.py

@@ -202,7 +202,7 @@ def _as_sql_node(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         if node.order_by is not None:
             result = _sort(result, node.order_by.all_ordering_columns)
         result = _remap_select_cols(
-            node.child,
+            result,
             [
                 nodes.AliasedRef(ref, identifiers.ColumnId(name))
                 for ref, name in node.output_cols

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

@@ -1,61 +1,35 @@
-WITH `bfcte_0` AS (
-  SELECT
-    `int64_col`,
-    `rowindex`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
-), `bfcte_1` AS (
-  SELECT
-    *,
-    NOT `int64_col` IS NULL AS `bfcol_4`
-  FROM `bfcte_0`
-), `bfcte_2` AS (
-  SELECT
-    *,
+SELECT
+  `rowindex`,
+  `int64_col`,
+  IF(
+    NOT `int64_col` IS NULL,
     IF(
       `int64_col` IS NULL,
       NULL,
       CAST(GREATEST(
-        CEIL(PERCENT_RANK() OVER (PARTITION BY `bfcol_4` ORDER BY `int64_col` ASC) * 4) - 1,
+        CEIL(
+          PERCENT_RANK() OVER (PARTITION BY NOT `int64_col` IS NULL ORDER BY `int64_col` ASC) * 4
+        ) - 1,
         0
       ) AS INT64)
-    ) AS `bfcol_5`
-  FROM `bfcte_1`
-), `bfcte_3` AS (
-  SELECT
-    *,
-    IF(`bfcol_4`, `bfcol_5`, NULL) AS `bfcol_6`
-  FROM `bfcte_2`
-), `bfcte_4` AS (
-  SELECT
-    *,
-    NOT `int64_col` IS NULL AS `bfcol_10`
-  FROM `bfcte_3`
-), `bfcte_5` AS (
-  SELECT
-    *,
+    ),
+    NULL
+  ) AS `qcut_w_int`,
+  IF(
+    NOT `int64_col` IS NULL,
     CASE
-      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) < 0
+      WHEN PERCENT_RANK() OVER (PARTITION BY NOT `int64_col` IS NULL ORDER BY `int64_col` ASC) < 0
       THEN NULL
-      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) <= 0.25
+      WHEN PERCENT_RANK() OVER (PARTITION BY NOT `int64_col` IS NULL ORDER BY `int64_col` ASC) <= 0.25
       THEN 0
-      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) <= 0.5
+      WHEN PERCENT_RANK() OVER (PARTITION BY NOT `int64_col` IS NULL ORDER BY `int64_col` ASC) <= 0.5
       THEN 1
-      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) <= 0.75
+      WHEN PERCENT_RANK() OVER (PARTITION BY NOT `int64_col` IS NULL ORDER BY `int64_col` ASC) <= 0.75
       THEN 2
-      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) <= 1
+      WHEN PERCENT_RANK() OVER (PARTITION BY NOT `int64_col` IS NULL ORDER BY `int64_col` ASC) <= 1
       THEN 3
       ELSE NULL
-    END AS `bfcol_11`
-  FROM `bfcte_4`
-), `bfcte_6` AS (
-  SELECT
-    *,
-    IF(`bfcol_10`, `bfcol_11`, NULL) AS `bfcol_12`
-  FROM `bfcte_5`
-)
-SELECT
-  `rowindex`,
-  `int64_col`,
-  `bfcol_6` AS `qcut_w_int`,
-  `bfcol_12` AS `qcut_w_list`
-FROM `bfcte_6`
+    END,
+    NULL
+  ) AS `qcut_w_list`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`

tests/unit/core/compile/sqlglot/snapshots/test_compile_aggregate/test_compile_aggregate/out.sql

@@ -18,4 +18,6 @@ WITH `bfcte_0` AS (
 SELECT
   `bfcol_3` AS `bool_col`,
   `bfcol_6` AS `int64_too`
-FROM `bfcte_1`
+FROM `bfcte_1`
+ORDER BY
+  `bfcol_3` ASC NULLS LAST

tests/unit/core/compile/sqlglot/snapshots/test_compile_aggregate/test_compile_aggregate_wo_dropna/out.sql

@@ -16,4 +16,6 @@ WITH `bfcte_0` AS (
 SELECT
   `bfcol_3` AS `bool_col`,
   `bfcol_6` AS `int64_too`
-FROM `bfcte_1`
+FROM `bfcte_1`
+ORDER BY
+  `bfcol_3` ASC NULLS LAST

tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat/out.sql

@@ -35,4 +35,7 @@ SELECT
   `bfcol_31` AS `rowindex_1`,
   `bfcol_32` AS `int64_col`,
   `bfcol_33` AS `string_col`
-FROM `bfcte_0`
+FROM `bfcte_0`
+ORDER BY
+  `bfcol_34` ASC NULLS LAST,
+  `bfcol_35` ASC NULLS LAST

tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat_filter_sorted/out.sql

@@ -49,4 +49,7 @@ WITH `bfcte_0` AS (
 SELECT
   `bfcol_42` AS `float64_col`,
   `bfcol_43` AS `int64_col`
-FROM `bfcte_0`
+FROM `bfcte_0`
+ORDER BY
+  `bfcol_44` ASC NULLS LAST,
+  `bfcol_45` ASC NULLS LAST

tests/unit/core/compile/sqlglot/snapshots/test_compile_explode/test_compile_explode_dataframe/out.sql

@@ -16,4 +16,6 @@ SELECT
   `rowindex` AS `rowindex_1`,
   `int_list_col`,
   `string_list_col`
-FROM `bfcte_1`
+FROM `bfcte_1`
+ORDER BY
+  `bfcol_7` ASC NULLS LAST

tests/unit/core/compile/sqlglot/snapshots/test_compile_explode/test_compile_explode_series/out.sql

@@ -13,4 +13,6 @@ WITH `bfcte_0` AS (
 SELECT
   `rowindex`,
   `int_list_col`
-FROM `bfcte_1`
+FROM `bfcte_1`
+ORDER BY
+  `bfcol_4` ASC NULLS LAST