refactor: Remove col_id_overrides from executor interfaces (#1438)

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -18,7 +18,7 @@
 import functools
 import io
 import typing
-from typing import Iterable, List, Optional, Sequence, Tuple
+from typing import Iterable, List, Mapping, Optional, Sequence, Tuple
 import warnings
 
 import google.cloud.bigquery
@@ -349,6 +349,20 @@ def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue:
             )
         )
 
+    def rename_columns(self, col_id_overrides: Mapping[str, str]) -> ArrayValue:
+        if not col_id_overrides:
+            return self
+        output_ids = [col_id_overrides.get(id, id) for id in self.node.schema.names]
+        return ArrayValue(
+            nodes.SelectionNode(
+                self.node,
+                tuple(
+                    nodes.AliasedRef(ex.DerefOp(old_id), ids.ColumnId(out_id))
+                    for old_id, out_id in zip(self.node.ids, output_ids)
+                ),
+            )
+        )
+
     def drop_columns(self, columns: Iterable[str]) -> ArrayValue:
         return self.select_columns(
             [col_id for col_id in self.column_ids if col_id not in columns]

bigframes/core/blocks.py

@@ -2431,7 +2431,7 @@ def to_sql_query(
         # implementaton. It will reference cached tables instead of original data sources.
         # Maybe should just compile raw BFET? Depends on user intent.
         sql = self.session._executor.to_sql(
-            array_value, col_id_overrides=substitutions, enable_cache=enable_cache
+            array_value.rename_columns(substitutions), enable_cache=enable_cache
         )
         return (
             sql,

bigframes/dataframe.py

@@ -3606,8 +3606,7 @@ def to_csv(
             "header": header,
         }
         query_job = self._session._executor.export_gcs(
-            export_array,
-            id_overrides,
+            export_array.rename_columns(id_overrides),
             path_or_buf,
             format="csv",
             export_options=options,
@@ -3656,8 +3655,7 @@ def to_json(
             ordering_id=bigframes.session._io.bigquery.IO_ORDERING_ID,
         )
         query_job = self._session._executor.export_gcs(
-            export_array,
-            id_overrides,
+            export_array.rename_columns(id_overrides),
             path_or_buf,
             format="json",
             export_options={},
@@ -3736,9 +3734,8 @@ def to_gbq(
             )
         )
         query_job = self._session._executor.export_gbq(
-            export_array,
+            export_array.rename_columns(id_overrides),
             destination=destination,
-            col_id_overrides=id_overrides,
             cluster_cols=clustering_fields,
             if_exists=if_exists,
         )
@@ -3814,8 +3811,7 @@ def to_parquet(
             ordering_id=bigframes.session._io.bigquery.IO_ORDERING_ID,
         )
         query_job = self._session._executor.export_gcs(
-            export_array,
-            id_overrides,
+            export_array.rename_columns(id_overrides),
             path,
             format="parquet",
             export_options=export_options,
@@ -4070,7 +4066,9 @@ def _prepare_export(
         # the arbitrary unicode column labels feature in BigQuery, which is
         # currently (June 2023) in preview.
         id_overrides = {
-            col_id: col_label for col_id, col_label in zip(columns, column_labels)
+            col_id: col_label
+            for col_id, col_label in zip(columns, column_labels)
+            if (col_id != col_label)
         }
 
         if ordering_id is not None:

bigframes/session/executor.py

@@ -40,7 +40,6 @@
 import pyarrow
 
 import bigframes.core
-from bigframes.core import expression
 import bigframes.core.compile
 import bigframes.core.guid
 import bigframes.core.identifiers
@@ -91,7 +90,6 @@ def to_sql(
         self,
         array_value: bigframes.core.ArrayValue,
         offset_column: Optional[str] = None,
-        col_id_overrides: Mapping[str, str] = {},
         ordered: bool = False,
         enable_cache: bool = True,
     ) -> str:
@@ -105,7 +103,6 @@ def execute(
         array_value: bigframes.core.ArrayValue,
         *,
         ordered: bool = True,
-        col_id_overrides: Mapping[str, str] = {},
         use_explicit_destination: Optional[bool] = False,
         get_size_bytes: bool = False,
         page_size: Optional[int] = None,
@@ -119,7 +116,6 @@ def execute(
     def export_gbq(
         self,
         array_value: bigframes.core.ArrayValue,
-        col_id_overrides: Mapping[str, str],
         destination: bigquery.TableReference,
         if_exists: Literal["fail", "replace", "append"] = "fail",
         cluster_cols: Sequence[str] = [],
@@ -132,7 +128,6 @@ def export_gbq(
     def export_gcs(
         self,
         array_value: bigframes.core.ArrayValue,
-        col_id_overrides: Mapping[str, str],
         uri: str,
         format: Literal["json", "csv", "parquet"],
         export_options: Mapping[str, Union[bool, str]],
@@ -220,29 +215,23 @@ def to_sql(
         self,
         array_value: bigframes.core.ArrayValue,
         offset_column: Optional[str] = None,
-        col_id_overrides: Mapping[str, str] = {},
         ordered: bool = False,
         enable_cache: bool = True,
     ) -> str:
         if offset_column:
             array_value, internal_offset_col = array_value.promote_offsets()
-            col_id_overrides = dict(col_id_overrides)
-            col_id_overrides[internal_offset_col] = offset_column
         node = (
             self.replace_cached_subtrees(array_value.node)
             if enable_cache
             else array_value.node
         )
-        if col_id_overrides:
-            node = override_ids(node, col_id_overrides)
         return self.compiler.compile(node, ordered=ordered)
 
     def execute(
         self,
         array_value: bigframes.core.ArrayValue,
         *,
         ordered: bool = True,
-        col_id_overrides: Mapping[str, str] = {},
         use_explicit_destination: Optional[bool] = False,
         get_size_bytes: bool = False,
         page_size: Optional[int] = None,
@@ -254,15 +243,12 @@ def execute(
         if bigframes.options.compute.enable_multi_query_execution:
             self._simplify_with_caching(array_value)
 
-        sql = self.to_sql(
-            array_value, ordered=ordered, col_id_overrides=col_id_overrides
-        )
-        adjusted_schema = array_value.schema.rename(col_id_overrides)
+        sql = self.to_sql(array_value, ordered=ordered)
         job_config = bigquery.QueryJobConfig()
         # Use explicit destination to avoid 10GB limit of temporary table
         if use_explicit_destination:
             destination_table = self.storage_manager.create_temp_table(
-                adjusted_schema.to_bigquery(), cluster_cols=[]
+                array_value.schema.to_bigquery(), cluster_cols=[]
             )
             job_config.destination = destination_table
         # TODO(swast): plumb through the api_name of the user-facing api that
@@ -293,12 +279,12 @@ def iterator_supplier():
             )
         # Runs strict validations to ensure internal type predictions and ibis are completely in sync
         # Do not execute these validations outside of testing suite.
-        if "PYTEST_CURRENT_TEST" in os.environ and len(col_id_overrides) == 0:
+        if "PYTEST_CURRENT_TEST" in os.environ:
             self._validate_result_schema(array_value, iterator.schema)
 
         return ExecuteResult(
             arrow_batches=iterator_supplier,
-            schema=adjusted_schema,
+            schema=array_value.schema,
             query_job=query_job,
             total_bytes=size_bytes,
             total_rows=iterator.total_rows,
@@ -307,7 +293,6 @@ def iterator_supplier():
     def export_gbq(
         self,
         array_value: bigframes.core.ArrayValue,
-        col_id_overrides: Mapping[str, str],
         destination: bigquery.TableReference,
         if_exists: Literal["fail", "replace", "append"] = "fail",
         cluster_cols: Sequence[str] = [],
@@ -323,7 +308,7 @@ def export_gbq(
             "replace": bigquery.WriteDisposition.WRITE_TRUNCATE,
             "append": bigquery.WriteDisposition.WRITE_APPEND,
         }
-        sql = self.to_sql(array_value, ordered=False, col_id_overrides=col_id_overrides)
+        sql = self.to_sql(array_value, ordered=False)
         job_config = bigquery.QueryJobConfig(
             write_disposition=dispositions[if_exists],
             destination=destination,
@@ -340,15 +325,13 @@ def export_gbq(
     def export_gcs(
         self,
         array_value: bigframes.core.ArrayValue,
-        col_id_overrides: Mapping[str, str],
         uri: str,
         format: Literal["json", "csv", "parquet"],
         export_options: Mapping[str, Union[bool, str]],
     ):
         query_job = self.execute(
             array_value,
             ordered=False,
-            col_id_overrides=col_id_overrides,
             use_explicit_destination=True,
         ).query_job
         result_table = query_job.destination
@@ -678,18 +661,3 @@ def generate_head_plan(node: nodes.BigFrameNode, n: int):
 
 def generate_row_count_plan(node: nodes.BigFrameNode):
     return nodes.RowCountNode(node)
-
-
-def override_ids(
-    node: nodes.BigFrameNode, col_id_overrides: Mapping[str, str]
-) -> nodes.SelectionNode:
-    output_ids = [col_id_overrides.get(id, id) for id in node.schema.names]
-    return nodes.SelectionNode(
-        node,
-        tuple(
-            nodes.AliasedRef(
-                expression.DerefOp(old_id), bigframes.core.identifiers.ColumnId(out_id)
-            )
-            for old_id, out_id in zip(node.ids, output_ids)
-        ),
-    )

tests/unit/polars_session.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import dataclasses
-from typing import Mapping, Optional, Union
+from typing import Optional, Union
 import weakref
 
 import polars
@@ -39,7 +39,6 @@ def execute(
         array_value: bigframes.core.ArrayValue,
         *,
         ordered: bool = True,
-        col_id_overrides: Mapping[str, str] = {},
         use_explicit_destination: Optional[bool] = False,
         get_size_bytes: bool = False,
         page_size: Optional[int] = None,