feat: allow setting table labels in to_gbq (#941)

* chore: allow setting table labels in `to_gbq`

---------

Co-authored-by: Matthew Laurence Chen <[email protected]>
Co-authored-by: Chelsea Lin <[email protected]>

Author: 3 people

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(

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"),
     [

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