diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 6b782b4692..2ae6aefe1b 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -3027,6 +3027,7 @@ def to_gbq( index: bool = True, ordering_id: Optional[str] = None, clustering_columns: Union[pandas.Index, Iterable[typing.Hashable]] = (), + labels: dict[str, str] = {}, ) -> str: temp_table_ref = None @@ -3081,9 +3082,11 @@ def to_gbq( export_array, id_overrides = self._prepare_export( index=index and self._has_index, ordering_id=ordering_id ) - destination = bigquery.table.TableReference.from_string( - destination_table, - default_project=default_project, + destination: bigquery.table.TableReference = ( + bigquery.table.TableReference.from_string( + destination_table, + default_project=default_project, + ) ) _, query_job = self._session._export( export_array, @@ -3106,6 +3109,11 @@ def to_gbq( + constants.DEFAULT_EXPIRATION, ) + if len(labels) != 0: + table = bigquery.Table(result_table) + table.labels = labels + self._session.bqclient.update_table(table, ["labels"]) + return destination_table def to_numpy( diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index ddcf044911..f51b597650 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -4657,6 +4657,17 @@ def test_to_gbq_and_create_dataset(session, scalars_df_index, dataset_id_not_cre assert not loaded_scalars_df_index.empty +def test_to_gbq_table_labels(scalars_df_index): + destination_table = "bigframes-dev.bigframes_tests_sys.table_labels" + result_table = scalars_df_index.to_gbq( + destination_table, labels={"test": "labels"}, if_exists="replace" + ) + client = scalars_df_index._session.bqclient + table = client.get_table(result_table) + assert table.labels + assert table.labels["test"] == "labels" + + @pytest.mark.parametrize( ("col_names", "ignore_index"), [ diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 10565a2552..fe1c8a12ff 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -390,6 +390,7 @@ def to_gbq( index: bool = True, ordering_id: Optional[str] = None, clustering_columns: Union[pd.Index, Iterable[Hashable]] = (), + labels: dict[str, str] = {}, ) -> str: """Write a DataFrame to a BigQuery table. @@ -467,6 +468,9 @@ def to_gbq( clustering order within the Index/DataFrame columns follows the order specified in `clustering_columns`. + labels (dict[str, str], default None): + Specifies table labels within BigQuery + Returns: str: The fully-qualified ID for the written table, in the form