Skip to content

Commit 11a3743

Browse files
refactor: export operations use sesssion.execute path (#418)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent 36173b0 commit 11a3743

File tree

2 files changed

+22
-13
lines changed

2 files changed

+22
-13
lines changed

bigframes/dataframe.py

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2932,8 +2932,9 @@ def map_columns_on_occurrence(columns):
29322932

29332933
return clustering_columns_for_index + clustering_columns_for_df
29342934

2935-
def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str:
2936-
"""Create query text representing this dataframe for I/O."""
2935+
def _prepare_export(
2936+
self, index: bool, ordering_id: Optional[str]
2937+
) -> Tuple[bigframes.core.ArrayValue, Dict[str, str]]:
29372938
array_value = self._block.expr
29382939

29392940
new_col_labels, new_idx_labels = utils.get_standardized_ids(
@@ -2961,10 +2962,7 @@ def _create_io_query(self, index: bool, ordering_id: Optional[str]) -> str:
29612962

29622963
if ordering_id is not None:
29632964
array_value = array_value.promote_offsets(ordering_id)
2964-
return self._block.session._to_sql(
2965-
array_value=array_value,
2966-
col_id_overrides=id_overrides,
2967-
)
2965+
return array_value, id_overrides
29682966

29692967
def _run_io_query(
29702968
self,
@@ -2974,11 +2972,16 @@ def _run_io_query(
29742972
) -> bigquery.TableReference:
29752973
"""Executes a query job presenting this dataframe and returns the destination
29762974
table."""
2977-
expr = self._block.expr
2978-
session = expr.session
2979-
sql = self._create_io_query(index=index, ordering_id=ordering_id)
2980-
_, query_job = session._start_query(
2981-
sql=sql, job_config=job_config # type: ignore
2975+
session = self._block.expr.session
2976+
export_array, id_overrides = self._prepare_export(
2977+
index=index, ordering_id=ordering_id
2978+
)
2979+
2980+
_, query_job = session._execute(
2981+
export_array,
2982+
job_config=job_config,
2983+
sorted=False,
2984+
col_id_overrides=id_overrides,
29822985
)
29832986
self._set_internal_query_job(query_job)
29842987

bigframes/session/__init__.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1626,9 +1626,15 @@ def _execute(
16261626
*,
16271627
sorted: bool = True,
16281628
dry_run=False,
1629+
col_id_overrides: Mapping[str, str] = {},
16291630
) -> tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
1630-
sql = self._to_sql(array_value, sorted=sorted) # type:ignore
1631-
job_config = bigquery.QueryJobConfig(dry_run=dry_run)
1631+
sql = self._to_sql(
1632+
array_value, sorted=sorted, col_id_overrides=col_id_overrides
1633+
) # type:ignore
1634+
if job_config is None:
1635+
job_config = bigquery.QueryJobConfig(dry_run=dry_run)
1636+
else:
1637+
job_config.dry_run = dry_run
16321638
return self._start_query(
16331639
sql=sql,
16341640
job_config=job_config,

0 commit comments

Comments
 (0)