refactor: fix pull_up_select disorder the columns of readtable nodes

chelsea-lin · chelsea-lin · commit 56c28791072e · 2026-02-04T20:48:53.000Z
diff --git a/bigframes/core/rewrite/select_pullup.py b/bigframes/core/rewrite/select_pullup.py
@@ -54,13 +54,12 @@ def pull_up_source_ids(node: nodes.ReadTableNode) -> nodes.BigFrameNode:
     if all(id.sql == source_id for id, source_id in node.scan_list.items):
         return node
     else:
-        source_ids = sorted(
-            set(scan_item.source_id for scan_item in node.scan_list.items)
-        )
         new_scan_list = nodes.ScanList.from_items(
             [
-                nodes.ScanItem(identifiers.ColumnId(source_id), source_id)
-                for source_id in source_ids
+                nodes.ScanItem(
+                    identifiers.ColumnId(scan_item.source_id), scan_item.source_id
+                )
+                for scan_item in node.scan_list.items
             ]
         )
         new_source = dataclasses.replace(node, scan_list=new_scan_list)
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_corr/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_corr/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_cov/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_cov/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_row_number/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_row_number/out.sql
@@ -4,17 +4,17 @@ WITH `bfcte_0` AS (
     `bytes_col`,
     `date_col`,
     `datetime_col`,
-    `duration_col`,
-    `float64_col`,
     `geography_col`,
     `int64_col`,
     `int64_too`,
     `numeric_col`,
+    `float64_col`,
     `rowindex`,
     `rowindex_2`,
     `string_col`,
     `time_col`,
-    `timestamp_col`
+    `timestamp_col`,
+    `duration_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_size/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_size/out.sql
@@ -4,17 +4,17 @@ WITH `bfcte_0` AS (
     `bytes_col`,
     `date_col`,
     `datetime_col`,
-    `duration_col`,
-    `float64_col`,
     `geography_col`,
     `int64_col`,
     `int64_too`,
     `numeric_col`,
+    `float64_col`,
     `rowindex`,
     `rowindex_2`,
     `string_col`,
     `time_col`,
-    `timestamp_col`
+    `timestamp_col`,
+    `duration_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_mean/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_mean/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `duration_col`,
-    `int64_col`
+    `int64_col`,
+    `duration_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_std/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_std/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `duration_col`,
-    `int64_col`
+    `int64_col`,
+    `duration_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_array_ops/test_to_array_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_array_ops/test_to_array_op/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `string_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_is_in/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_is_in/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_maximum_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_maximum_op/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_minimum_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_minimum_op/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_sub_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_sub_timedelta/out.sql
@@ -1,9 +1,9 @@
 WITH `bfcte_0` AS (
   SELECT
     `date_col`,
-    `duration_col`,
     `rowindex`,
-    `timestamp_col`
+    `timestamp_col`,
+    `duration_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_to_datetime/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_to_datetime/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `string_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_to_timestamp/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_to_timestamp/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_astype_int/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_astype_int/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `datetime_col`,
-    `float64_col`,
     `numeric_col`,
+    `float64_col`,
     `time_col`,
     `timestamp_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_astype_json/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_astype_json/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `string_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_binary_remote_function_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_binary_remote_function_op/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_case_when_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_case_when_op/out.sql
@@ -1,9 +1,9 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
     `int64_col`,
-    `int64_too`
+    `int64_too`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_fillna/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_fillna/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_nary_remote_function_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_nary_remote_function_op/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `string_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_row_key/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_row_key/out.sql
@@ -4,17 +4,17 @@ WITH `bfcte_0` AS (
     `bytes_col`,
     `date_col`,
     `datetime_col`,
-    `duration_col`,
-    `float64_col`,
     `geography_col`,
     `int64_col`,
     `int64_too`,
     `numeric_col`,
+    `float64_col`,
     `rowindex`,
     `rowindex_2`,
     `string_col`,
     `time_col`,
-    `timestamp_col`
+    `timestamp_col`,
+    `duration_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_where/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_where/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_arctan2/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_arctan2/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_cosine_distance/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_cosine_distance/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float_list_col`,
-    `int_list_col`
+    `int_list_col`,
+    `float_list_col`
   FROM `bigframes-dev`.`sqlglot_test`.`repeated_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_div_numeric/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_div_numeric/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `rowindex`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mod_numeric/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mod_numeric/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `rowindex`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mul_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mul_timedelta/out.sql
@@ -1,9 +1,9 @@
 WITH `bfcte_0` AS (
   SELECT
-    `duration_col`,
     `int64_col`,
     `rowindex`,
-    `timestamp_col`
+    `timestamp_col`,
+    `duration_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_pow/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_pow/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `rowindex`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_round/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_round/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `rowindex`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_sub_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_sub_timedelta/out.sql
@@ -1,9 +1,9 @@
 WITH `bfcte_0` AS (
   SELECT
     `date_col`,
-    `duration_col`,
     `rowindex`,
-    `timestamp_col`
+    `timestamp_col`,
+    `duration_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_unsafe_pow_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_unsafe_pow_op/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
-    `int64_col`
+    `int64_col`,
+    `float64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_struct_ops/test_struct_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_struct_ops/test_struct_op/out.sql
@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `string_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_timedelta_ops/test_to_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_timedelta_ops/test_to_timedelta/out.sql
@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col`,
     `int64_col`,
+    `float64_col`,
     `rowindex`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat_filter_sorted/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat_filter_sorted/out.sql
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_explode/test_compile_explode_dataframe/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_explode/test_compile_explode_dataframe/out.sql
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_explode/test_compile_explode_series/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_explode/test_compile_explode_series/out.sql
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable/out.sql
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_columns_filters/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_columns_filters/out.sql
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_json_types/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_json_types/out.sql
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_repeated_types/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_repeated_types/out.sql
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_system_time/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_system_time/out.sql

Original file line number	Diff line number	Diff line change
`@@ -54,13 +54,12 @@ def pull_up_source_ids(node: nodes.ReadTableNode) -> nodes.BigFrameNode:`
`54`	`54`	`if all(id.sql == source_id for id, source_id in node.scan_list.items):`
`55`	`55`	`return node`
`56`	`56`	`else:`
`57`		`- source_ids = sorted(`
`58`		`- set(scan_item.source_id for scan_item in node.scan_list.items)`
`59`		`- )`
`60`	`57`	`new_scan_list = nodes.ScanList.from_items(`
`61`	`58`	`[`
`62`		`- nodes.ScanItem(identifiers.ColumnId(source_id), source_id)`
`63`		`- for source_id in source_ids`
	`59`	`+ nodes.ScanItem(`
	`60`	`+ identifiers.ColumnId(scan_item.source_id), scan_item.source_id`
	`61`	`+ )`
	`62`	`+ for scan_item in node.scan_list.items`
`64`	`63`	`]`
`65`	`64`	`)`
`66`	`65`	`new_source = dataclasses.replace(node, scan_list=new_scan_list)`