perf: Defer some data uploads to execution time

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -133,8 +133,17 @@ def from_table(
             ordering=ordering,
             n_rows=n_rows,
         )
+        return cls.from_bq_data_source(source_def, scan_list, session)
+
+    @classmethod
+    def from_bq_data_source(
+        cls,
+        source: nodes.BigqueryDataSource,
+        scan_list: nodes.ScanList,
+        session: Session,
+    ):
         node = nodes.ReadTableNode(
-            source=source_def,
+            source=source,
             scan_list=scan_list,
             table_session=session,
         )

bigframes/session/loader.py

@@ -43,7 +43,7 @@
 import pandas
 import pyarrow as pa
 
-from bigframes.core import guid, local_data, utils
+from bigframes.core import guid, identifiers, local_data, nodes, ordering, utils
 import bigframes.core as core
 import bigframes.core.blocks as blocks
 import bigframes.core.schema as schemata
@@ -183,35 +183,59 @@ def read_pandas(
         )
         managed_data = local_data.ManagedArrowTable.from_pandas(prepared_df)
 
+        block = blocks.Block(
+            self.read_managed_data(managed_data, method=method, api_name=api_name),
+            index_columns=idx_cols,
+            column_labels=pandas_dataframe.columns,
+            index_labels=pandas_dataframe.index.names,
+        )
+        return dataframe.DataFrame(block)
+
+    def read_managed_data(
+        self,
+        data: local_data.ManagedArrowTable,
+        method: Literal["load", "stream", "write"],
+        api_name: str,
+    ) -> core.ArrayValue:
+        offsets_col = guid.generate_guid("upload_offsets_")
         if method == "load":
-            array_value = self.load_data(managed_data, api_name=api_name)
+            gbq_source = self.load_data(
+                data, offsets_col=offsets_col, api_name=api_name
+            )
         elif method == "stream":
-            array_value = self.stream_data(managed_data)
+            gbq_source = self.stream_data(data, offsets_col=offsets_col)
         elif method == "write":
-            array_value = self.write_data(managed_data)
+            gbq_source = self.write_data(data, offsets_col=offsets_col)
         else:
             raise ValueError(f"Unsupported read method {method}")
 
-        block = blocks.Block(
-            array_value,
-            index_columns=idx_cols,
-            column_labels=pandas_dataframe.columns,
-            index_labels=pandas_dataframe.index.names,
+        return core.ArrayValue.from_bq_data_source(
+            source=gbq_source,
+            scan_list=nodes.ScanList(
+                tuple(
+                    nodes.ScanItem(
+                        identifiers.ColumnId(item.column), item.dtype, item.column
+                    )
+                    for item in data.schema.items
+                )
+            ),
+            session=self._session,
         )
-        return dataframe.DataFrame(block)
 
     def load_data(
-        self, data: local_data.ManagedArrowTable, api_name: Optional[str] = None
-    ) -> core.ArrayValue:
+        self,
+        data: local_data.ManagedArrowTable,
+        offsets_col: str,
+        api_name: Optional[str] = None,
+    ) -> nodes.BigqueryDataSource:
         """Load managed data into bigquery"""
-        ordering_col = guid.generate_guid("load_offsets_")
 
         # JSON support incomplete
         for item in data.schema.items:
             _validate_dtype_can_load(item.column, item.dtype)
 
         schema_w_offsets = data.schema.append(
-            schemata.SchemaItem(ordering_col, bigframes.dtypes.INT_DTYPE)
+            schemata.SchemaItem(offsets_col, bigframes.dtypes.INT_DTYPE)
         )
         bq_schema = schema_w_offsets.to_bigquery(_LOAD_JOB_TYPE_OVERRIDES)
 
@@ -222,13 +246,13 @@ def load_data(
             job_config.labels = {"bigframes-api": api_name}
 
         load_table_destination = self._storage_manager.create_temp_table(
-            bq_schema, [ordering_col]
+            bq_schema, [offsets_col]
         )
 
         buffer = io.BytesIO()
         data.to_parquet(
             buffer,
-            offsets_col=ordering_col,
+            offsets_col=offsets_col,
             geo_format="wkt",
             duration_type="duration",
             json_type="string",
@@ -240,23 +264,24 @@ def load_data(
         self._start_generic_job(load_job)
         # must get table metadata after load job for accurate metadata
         destination_table = self._bqclient.get_table(load_table_destination)
-        return core.ArrayValue.from_table(
-            table=destination_table,
-            schema=schema_w_offsets,
-            session=self._session,
-            offsets_col=ordering_col,
-            n_rows=data.data.num_rows,
-        ).drop_columns([ordering_col])
+        return nodes.BigqueryDataSource(
+            destination_table,
+            ordering=ordering.TotalOrdering.from_offset_col(offsets_col),
+            n_rows=destination_table.num_rows,
+        )
 
-    def stream_data(self, data: local_data.ManagedArrowTable) -> core.ArrayValue:
+    def stream_data(
+        self,
+        data: local_data.ManagedArrowTable,
+        offsets_col: str,
+    ) -> nodes.BigqueryDataSource:
         """Load managed data into bigquery"""
-        ordering_col = guid.generate_guid("stream_offsets_")
         schema_w_offsets = data.schema.append(
-            schemata.SchemaItem(ordering_col, bigframes.dtypes.INT_DTYPE)
+            schemata.SchemaItem(offsets_col, bigframes.dtypes.INT_DTYPE)
         )
         bq_schema = schema_w_offsets.to_bigquery(_STREAM_JOB_TYPE_OVERRIDES)
         load_table_destination = self._storage_manager.create_temp_table(
-            bq_schema, [ordering_col]
+            bq_schema, [offsets_col]
         )
 
         rows = data.itertuples(
@@ -279,20 +304,21 @@ def stream_data(self, data: local_data.ManagedArrowTable) -> core.ArrayValue:
             table=destination_table,
             schema=schema_w_offsets,
             session=self._session,
-            offsets_col=ordering_col,
+            offsets_col=offsets_col,
             n_rows=data.data.num_rows,
-        ).drop_columns([ordering_col])
+        ).drop_columns([offsets_col])
 
-    def write_data(self, data: local_data.ManagedArrowTable) -> core.ArrayValue:
+    def write_data(
+        self,
+        data: local_data.ManagedArrowTable,
+        offsets_col: str,
+    ) -> nodes.BigqueryDataSource:
         """Load managed data into bigquery"""
-        ordering_col = guid.generate_guid("stream_offsets_")
         schema_w_offsets = data.schema.append(
-            schemata.SchemaItem(ordering_col, bigframes.dtypes.INT_DTYPE)
+            schemata.SchemaItem(offsets_col, bigframes.dtypes.INT_DTYPE)
         )
         bq_schema = schema_w_offsets.to_bigquery(_STREAM_JOB_TYPE_OVERRIDES)
-        bq_table_ref = self._storage_manager.create_temp_table(
-            bq_schema, [ordering_col]
-        )
+        bq_table_ref = self._storage_manager.create_temp_table(bq_schema, [offsets_col])
 
         requested_stream = bq_storage_types.stream.WriteStream()
         requested_stream.type_ = bq_storage_types.stream.WriteStream.Type.COMMITTED  # type: ignore
@@ -304,7 +330,7 @@ def write_data(self, data: local_data.ManagedArrowTable) -> core.ArrayValue:
 
         def request_gen() -> Generator[bq_storage_types.AppendRowsRequest, None, None]:
             schema, batches = data.to_arrow(
-                offsets_col=ordering_col, duration_type="int"
+                offsets_col=offsets_col, duration_type="int"
             )
             offset = 0
             for batch in batches:
@@ -334,9 +360,9 @@ def request_gen() -> Generator[bq_storage_types.AppendRowsRequest, None, None]:
             table=destination_table,
             schema=schema_w_offsets,
             session=self._session,
-            offsets_col=ordering_col,
+            offsets_col=offsets_col,
             n_rows=data.data.num_rows,
-        ).drop_columns([ordering_col])
+        ).drop_columns([offsets_col])
 
     def _start_generic_job(self, job: formatting_helpers.GenericJob):
         if bigframes.options.display.progress_bar is not None:
@@ -533,7 +559,7 @@ def read_gbq_table(
             if not primary_key:
                 array_value = array_value.order_by(
                     [
-                        bigframes.core.ordering.OrderingExpression(
+                        ordering.OrderingExpression(
                             bigframes.operations.RowKey().as_expr(
                                 *(id for id in array_value.column_ids)
                             ),