fix tests related to blob api deprecation

Author: shuoweil

packages/bigframes/tests/system/small/bigquery/test_ai.py

@@ -22,6 +22,31 @@
 import bigframes.pandas as bpd
 from bigframes import dataframe, dtypes, series
 from bigframes.testing import utils as test_utils
+import uuid
+import google.cloud.bigquery
+
+
+def _create_mock_obj_ref_df(session, uris, name="image"):
+    df = bpd.DataFrame({name: uris}, session=session)
+    table_id = f"bigframes-dev.bigframes_tests_sys.tmp_obj_ref_{uuid.uuid4().hex}"
+    df.to_gbq(table_id, if_exists="replace")
+
+    client = session.bqclient
+    table = client.get_table(table_id)
+    schema = list(table.schema)
+    for i, field in enumerate(schema):
+        if field.name == name:
+            schema[i] = google.cloud.bigquery.SchemaField(
+                name=field.name,
+                field_type=field.field_type,
+                mode=field.mode,
+                description="bigframes_dtype: OBJ_REF_DTYPE",
+            )
+            break
+    table.schema = schema
+    client.update_table(table, ["schema"])
+
+    return session.read_gbq(table_id)
 
 
 def test_ai_function_pandas_input(session):
@@ -159,8 +184,8 @@ def test_ai_generate_bool(session):
 
 
 def test_ai_generate_bool_multi_model(session):
-    df = session.from_glob_path(
-        "gs://bigframes-dev-testing/a_multimodel/images/*", name="image"
+    df = _create_mock_obj_ref_df(
+        session, ["gs://cloud-samples-data/vision/ocr/sign.jpg"], name="image"
     )
 
     result = bbq.ai.generate_bool((df["image"], " contains an animal"))
@@ -196,8 +221,8 @@ def test_ai_generate_int(session):
 
 
 def test_ai_generate_int_multi_model(session):
-    df = session.from_glob_path(
-        "gs://bigframes-dev-testing/a_multimodel/images/*", name="image"
+    df = _create_mock_obj_ref_df(
+        session, ["gs://cloud-samples-data/vision/ocr/sign.jpg"], name="image"
     )
 
     result = bbq.ai.generate_int(
@@ -235,8 +260,8 @@ def test_ai_generate_double(session):
 
 
 def test_ai_generate_double_multi_model(session):
-    df = session.from_glob_path(
-        "gs://bigframes-dev-testing/a_multimodel/images/*", name="image"
+    df = _create_mock_obj_ref_df(
+        session, ["gs://cloud-samples-data/vision/ocr/sign.jpg"], name="image"
     )
 
     result = bbq.ai.generate_double(
@@ -267,10 +292,8 @@ def test_ai_if(session):
 
 
 def test_ai_if_multi_model(session, bq_connection):
-    df = session.from_glob_path(
-        "gs://bigframes-dev-testing/a_multimodel/images/*",
-        name="image",
-        connection=bq_connection,
+    df = _create_mock_obj_ref_df(
+        session, ["gs://cloud-samples-data/vision/ocr/sign.jpg"], name="image"
     )
 
     result = bbq.ai.if_((df["image"], " contains an animal"))
@@ -289,10 +312,8 @@ def test_ai_classify(session):
 
 
 def test_ai_classify_multi_model(session, bq_connection):
-    df = session.from_glob_path(
-        "gs://bigframes-dev-testing/a_multimodel/images/*",
-        name="image",
-        connection=bq_connection,
+    df = _create_mock_obj_ref_df(
+        session, ["gs://cloud-samples-data/vision/ocr/sign.jpg"], name="image"
     )
 
     result = bbq.ai.classify(df["image"], ["photo", "cartoon"])
@@ -312,8 +333,8 @@ def test_ai_score(session):
 
 
 def test_ai_score_multi_model(session):
-    df = session.from_glob_path(
-        "gs://bigframes-dev-testing/a_multimodel/images/*", name="image"
+    df = _create_mock_obj_ref_df(
+        session, ["gs://cloud-samples-data/vision/ocr/sign.jpg"], name="image"
     )
     prompt = ("Rank the liveliness of ", df["image"], "on the scale from 1 to 3")
 

packages/bigframes/tests/system/small/pandas/test_describe.py

@@ -358,15 +358,36 @@ def test_series_groupby_describe(scalars_dfs):
 
 def test_describe_json_and_obj_ref_returns_count(session):
     # Test describe() works on JSON and OBJ_REF types (without nunique, which fails)
+    import uuid
+    import google.cloud.bigquery
+
     sql = """
         SELECT
             PARSE_JSON('{"a": 1}') AS json_col,
             'gs://cloud-samples-data/vision/ocr/sign.jpg' AS uri_col
     """
-    df = session.read_gbq(sql)
+    df_init = session.read_gbq(sql)
+
+    table_id = f"bigframes-dev.bigframes_tests_sys.tmp_obj_ref_{uuid.uuid4().hex}"
+    df_init.to_gbq(table_id, if_exists="replace")
+
+    client = session.bqclient
+    table = client.get_table(table_id)
+    schema = list(table.schema)
+    for i, field in enumerate(schema):
+        if field.name == "uri_col":
+            schema[i] = google.cloud.bigquery.SchemaField(
+                name=field.name,
+                field_type=field.field_type,
+                mode=field.mode,
+                description="bigframes_dtype: OBJ_REF_DTYPE",
+            )
+            break
+    table.schema = schema
+    client.update_table(table, ["schema"])
 
-    df["obj_ref_col"] = df["uri_col"].str.to_blob()
-    df = df.drop(columns=["uri_col"])
+    df = session.read_gbq(table_id)
+    df = df.rename(columns={"uri_col": "obj_ref_col"})
 
     res = df.describe(include="all").to_pandas()
 

packages/bigframes/tests/system/small/test_dataframe.py

@@ -5919,9 +5919,34 @@ def test_to_gbq_table_labels(scalars_df_index):
 
 def test_to_gbq_obj_ref_persists(session):
     # Test that saving and loading an Object Reference retains its dtype
-    bdf = session.from_glob_path(
-        "gs://cloud-samples-data/vision/ocr/*.jpg", name="uris"
-    ).head(1)
+    import uuid
+    import google.cloud.bigquery
+
+    sql = """
+        SELECT STRUCT('gs://cloud-samples-data/vision/ocr/sign.jpg' AS uri, CAST(NULL AS STRING) AS version, CAST(NULL AS STRING) AS authorizer, PARSE_JSON('{}') AS details) AS uris
+    """
+    df_init = session.read_gbq(sql)
+
+    tmp_table_id = f"bigframes-dev.bigframes_tests_sys.tmp_obj_ref_{uuid.uuid4().hex}"
+    df_init.to_gbq(tmp_table_id, if_exists="replace")
+
+    client = session.bqclient
+    table = client.get_table(tmp_table_id)
+    schema = list(table.schema)
+    for i, field in enumerate(schema):
+        if field.name == "uris":
+            schema[i] = google.cloud.bigquery.SchemaField(
+                name=field.name,
+                field_type=field.field_type,
+                mode=field.mode,
+                description="bigframes_dtype: OBJ_REF_DTYPE",
+                fields=field.fields,
+            )
+            break
+    table.schema = schema
+    client.update_table(table, ["schema"])
+
+    bdf = session.read_gbq(tmp_table_id)
 
     destination_table = "bigframes-dev.bigframes_tests_sys.test_obj_ref_persistence"
     bdf.to_gbq(destination_table, if_exists="replace")

packages/bigframes/tests/system/small/test_dataframe_io.py

@@ -1010,16 +1010,38 @@ def test_to_gbq_timedelta_tag_ignored_when_appending(bigquery_client, dataset_id
 
 
 def test_to_gbq_obj_ref(session, dataset_id: str, bigquery_client):
+    import uuid
+    import google.cloud.bigquery
+
     destination_table = f"{dataset_id}.test_to_gbq_obj_ref"
     sql = """
-        SELECT
-            'gs://cloud-samples-data/vision/ocr/sign.jpg' AS uri_col
+        SELECT STRUCT('gs://cloud-samples-data/vision/ocr/sign.jpg' AS uri, CAST(NULL AS STRING) AS version, CAST(NULL AS STRING) AS authorizer, PARSE_JSON('{}') AS details) AS uri_col
     """
-    df = session.read_gbq(sql)
-    df["obj_ref_col"] = df["uri_col"].str.to_blob()
-    df = df.drop(columns=["uri_col"])
-
-    df.to_gbq(destination_table)
+    df_init = session.read_gbq(sql)
+
+    tmp_table_id = f"{dataset_id}.tmp_obj_ref_{uuid.uuid4().hex}"
+    df_init.to_gbq(tmp_table_id, if_exists="replace")
+
+    client = session.bqclient
+    table = client.get_table(tmp_table_id)
+    schema = list(table.schema)
+    for i, field in enumerate(schema):
+        if field.name == "uri_col":
+            schema[i] = google.cloud.bigquery.SchemaField(
+                name=field.name,
+                field_type=field.field_type,
+                mode=field.mode,
+                description="bigframes_dtype: OBJ_REF_DTYPE",
+                fields=field.fields,
+            )
+            break
+    table.schema = schema
+    client.update_table(table, ["schema"])
+
+    df = session.read_gbq(tmp_table_id)
+    df = df.rename(columns={"uri_col": "obj_ref_col"})
+
+    df.to_gbq(destination_table, if_exists="replace")
 
     table = bigquery_client.get_table(destination_table)
     obj_ref_field = next(f for f in table.schema if f.name == "obj_ref_col")

packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_blob_ops.py

@@ -38,19 +38,3 @@ def test_blob_get_access_url_with_duration(scalar_types_df: bpd.DataFrame, snaps
         [col_name],
     )
     snapshot.assert_match(sql, "out.sql")
-
-
-def test_blob_make_ref(scalar_types_df: bpd.DataFrame, snapshot):
-    ref_s = obj.make_ref(scalar_types_df["string_col"], authorizer="my-connection")
-    snapshot.assert_match(ref_s.to_frame().sql, "out.sql")
-
-
-def test_blob_make_ref_json(scalar_types_df: bpd.DataFrame, snapshot):
-    col_name = "string_col"
-    bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_ops_to_sql(
-        bf_df,
-        [ops.obj_make_ref_json_op.as_expr(col_name)],
-        [col_name],
-    )
-    snapshot.assert_match(sql, "out.sql")