fix: round-trip persistence for OBJ_REF_DTYPE

Author: shuoweil

bigframes/dataframe.py

@@ -4301,32 +4301,6 @@ def to_gbq(
         result_table = result.query_job.destination
         assert result_table is not None
 
-        obj_ref_dest_cols = []
-        for col_id in id_overrides.keys():
-            try:
-                if (
-                    export_array.get_column_type(col_id)
-                    == bigframes.dtypes.OBJ_REF_DTYPE
-                ):
-                    obj_ref_dest_cols.append(id_overrides[col_id])
-            except Exception:
-                pass
-
-        if obj_ref_dest_cols:
-            table = self._session.bqclient.get_table(result_table)
-            new_schema = []
-            for field in table.schema:
-                if field.name in obj_ref_dest_cols:
-                    field_dict = field.to_api_repr()
-                    field_dict["description"] = "bigframes_dtype: OBJ_REF_DTYPE"
-                    new_schema.append(
-                        google.cloud.bigquery.SchemaField.from_api_repr(field_dict)
-                    )
-                else:
-                    new_schema.append(field)
-            table.schema = new_schema
-            self._session.bqclient.update_table(table, ["schema"])
-
         if temp_table_ref:
             bigframes.session._io.bigquery.set_table_expiration(
                 self._session.bqclient,

bigframes/dtypes.py

@@ -772,7 +772,7 @@ def convert_schema_field(
 ) -> typing.Tuple[str, Dtype]:
     is_repeated = field.mode == "REPEATED"
     if field.field_type == "RECORD":
-        if field.description == "bigframes_dtype: OBJ_REF_DTYPE":
+        if field.description == OBJ_REF_DESCRIPTION_TAG:
             bf_dtype = OBJ_REF_DTYPE  # type: ignore
             if is_repeated:
                 pa_type = pa.list_(bigframes_dtype_to_arrow_dtype(bf_dtype))
@@ -834,8 +834,11 @@ def convert_to_schema_field(
                     convert_to_schema_field(field.name, inner_bf_type, overrides)
                 )
 
+            description = (
+                OBJ_REF_DESCRIPTION_TAG if bigframes_dtype == OBJ_REF_DTYPE else None
+            )
             return google.cloud.bigquery.SchemaField(
-                name, "RECORD", fields=inner_fields
+                name, "RECORD", fields=inner_fields, description=description
             )
         if bigframes_dtype.pyarrow_dtype == pa.duration("us"):
             # Timedeltas are represented as integers in microseconds.
@@ -978,6 +981,7 @@ def lcd_type_or_throw(dtype1: Dtype, dtype2: Dtype) -> Dtype:
 
 
 TIMEDELTA_DESCRIPTION_TAG = "#microseconds"
+OBJ_REF_DESCRIPTION_TAG = "bigframes_dtype: OBJ_REF_DTYPE"
 
 
 def contains_db_dtypes_json_arrow_type(type_):

bigframes/session/bq_caching_executor.py

@@ -334,13 +334,14 @@ def _export_gbq(
             session=array_value.session,
         )
 
-        has_timedelta_col = any(
-            t == bigframes.dtypes.TIMEDELTA_DTYPE for t in array_value.schema.dtypes
+        has_special_dtype_col = any(
+            t in (bigframes.dtypes.TIMEDELTA_DTYPE, bigframes.dtypes.OBJ_REF_DTYPE)
+            for t in array_value.schema.dtypes
         )
 
-        if spec.if_exists != "append" and has_timedelta_col:
+        if spec.if_exists != "append" and has_special_dtype_col:
             # Only update schema if this is not modifying an existing table, and the
-            # new table contains timedelta columns.
+            # new table contains special columns (like timedelta or obj_ref).
             table = self.bqclient.get_table(spec.table)
             table.schema = array_value.schema.to_bigquery()
             self.bqclient.update_table(table, ["schema"])

tests/system/small/test_dataframe_io.py

@@ -1012,16 +1012,13 @@ def test_to_gbq_obj_ref(session, dataset_id: str, bigquery_client):
     df["obj_ref_col"] = df["uri_col"].str.to_blob()
     df = df.drop(columns=["uri_col"])
 
-    # Save the dataframe to bigquery
     df.to_gbq(destination_table)
 
-    # Verify the table schema description is added
     table = bigquery_client.get_table(destination_table)
     obj_ref_field = next(f for f in table.schema if f.name == "obj_ref_col")
     assert obj_ref_field.field_type == "RECORD"
     assert obj_ref_field.description == "bigframes_dtype: OBJ_REF_DTYPE"
 
-    # Verify reloading it correctly restores the dtype
     reloaded_df = session.read_gbq(destination_table)
     assert reloaded_df["obj_ref_col"].dtype == dtypes.OBJ_REF_DTYPE
     assert len(reloaded_df) == 1