googleapis · mattyopl · Sep 5, 2024 · Aug 5, 2024 · Aug 5, 2024 · Aug 6, 2024
@@ -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(

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

@@ -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
-Original file line number
+Diff line change
@@ -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"
             Copy link

  
      
    
  

  
      

  
    Contributor


      

  

  
    
      

      
            shobsi
  

      

      

      


        Sep 6, 2024


      
    

  


        
      
  
  
  
    

    There was a problem hiding this comment.


  

 
  
    

    Choose a reason for hiding this comment

    
      The reason will be displayed to describe this comment to others. Learn more.
    

    
      
      


  


  
    
      Ideally the test should not be hard coupled with a project/dataset. Do we have to use a fixed table name? I think the test would still serve its purpose if we did to_gbq  without an explicit table arg and asserted labels on the returned anonymous table.
    
  
  


    

        
      
  
  
    
  
  
  
    
    Sorry, something went wrong.
  

  
    
  
    
      

              Uh oh!

              
There was an error while loading. Please reload this page.


  
  


          
      
  
    
    
      
        
            
    All reactions
  


          
          
        
      
    

    



    
        
  
    
        
    
  


      
          
  
      
            Copy link

  
      
    
  

  
      

  
    Contributor


      

  Author


  

  
    
      

      
            mattyopl
  

      

      

      


        Sep 9, 2024


      
    

  


        
      
  
  
  
    

    There was a problem hiding this comment.


  

 
  
    

    Choose a reason for hiding this comment

    
      The reason will be displayed to describe this comment to others. Learn more.
    

    
      
      


  


  
    
      done in #968!
    
  
  


    

        
      
  
  
    
  
  
  
    
    Sorry, something went wrong.
  

  
    
  
    
      

              Uh oh!

              
There was an error while loading. Please reload this page.


  
  


          
      
  
    
    
      
        
            
    All reactions
+        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"),
         [
-          Expand Down