test: use jobless query api for faster tests (#1478)

* use fixture to reset mutated global sessions

* fix tests

* fix progress bar tests by allowing large results

* remove undesired location setting in tests

* use fixture for test_df_construct_inline_respects_location

* remove redundant fixture

Author: TrevorBergeron

tests/system/conftest.py

@@ -139,9 +139,7 @@ def resourcemanager_client(
 
 @pytest.fixture(scope="session")
 def session() -> Generator[bigframes.Session, None, None]:
-    context = bigframes.BigQueryOptions(
-        location="US",
-    )
+    context = bigframes.BigQueryOptions(location="US", allow_large_results=False)
     session = bigframes.Session(context=context)
     yield session
     session.close()  # close generated session at cleanup time
@@ -157,15 +155,19 @@ def session_load() -> Generator[bigframes.Session, None, None]:
 
 @pytest.fixture(scope="session", params=["strict", "partial"])
 def maybe_ordered_session(request) -> Generator[bigframes.Session, None, None]:
-    context = bigframes.BigQueryOptions(location="US", ordering_mode=request.param)
+    context = bigframes.BigQueryOptions(
+        location="US", ordering_mode=request.param, allow_large_results=False
+    )
     session = bigframes.Session(context=context)
     yield session
     session.close()  # close generated session at cleanup type
 
 
 @pytest.fixture(scope="session")
 def unordered_session() -> Generator[bigframes.Session, None, None]:
-    context = bigframes.BigQueryOptions(location="US", ordering_mode="partial")
+    context = bigframes.BigQueryOptions(
+        location="US", ordering_mode="partial", allow_large_results=False
+    )
     session = bigframes.Session(context=context)
     yield session
     session.close()  # close generated session at cleanup type
@@ -1378,6 +1380,12 @@ def floats_product_bf(session, floats_product_pd):
     return session.read_pandas(floats_product_pd)
 
 
+@pytest.fixture(scope="session", autouse=True)
+def use_fast_query_path():
+    with bpd.option_context("bigquery.allow_large_results", False):
+        yield
+
+
 @pytest.fixture(scope="session", autouse=True)
 def cleanup_cloud_functions(session, cloudfunctions_client, dataset_id_permanent):
     """Clean up stale cloud functions."""
@@ -1460,3 +1468,12 @@ def images_mm_df(
     return session.from_glob_path(
         images_gcs_path, name="blob_col", connection=bq_connection
     )
+
+
+@pytest.fixture()
+def reset_default_session_and_location():
+    bpd.close_session()
+    with bpd.option_context("bigquery.location", None):
+        yield
+    bpd.close_session()
+    bpd.options.bigquery.location = None

tests/system/small/bigquery/test_json.py

@@ -22,6 +22,13 @@
 import bigframes.pandas as bpd
 
 
+@pytest.fixture(scope="module", autouse=True)
+def use_large_query_path():
+    # b/401630655
+    with bpd.option_context("bigquery.allow_large_results", True):
+        yield
+
+
 @pytest.mark.parametrize(
     ("json_path", "expected_json"),
     [

tests/system/small/bigquery/test_struct.py

@@ -53,9 +53,10 @@ def test_struct_from_dataframe(columns_arg):
     srs = series.Series(
         columns_arg,
     )
+    # Use allow_large_results=True, due to b/403028465
     pd.testing.assert_series_equal(
-        srs.to_pandas(),
-        bbq.struct(srs.struct.explode()).to_pandas(),
+        srs.to_pandas(allow_large_results=True),
+        bbq.struct(srs.struct.explode()).to_pandas(allow_large_results=True),
         check_index_type=False,
         check_dtype=False,
     )

tests/system/small/blob/test_properties.py

@@ -55,31 +55,33 @@ def test_blob_version(images_mm_df: bpd.DataFrame):
 
 
 def test_blob_metadata(images_mm_df: bpd.DataFrame):
-    bigframes.options.experiments.blob = True
-
-    actual = images_mm_df["blob_col"].blob.metadata().to_pandas()
-    expected = pd.Series(
-        [
-            {
-                "content_type": "image/jpeg",
-                "md5_hash": "e130ad042261a1883cd2cc06831cf748",
-                "size": 338390,
-                "updated": 1739574332000000,
-            },
-            {
-                "content_type": "image/jpeg",
-                "md5_hash": "e2ae3191ff2b809fd0935f01a537c650",
-                "size": 43333,
-                "updated": 1739574332000000,
-            },
-        ],
-        name="metadata",
-        dtype=db_dtypes.JSONDtype(),
-    )
-
-    pd.testing.assert_series_equal(
-        actual, expected, check_dtype=False, check_index_type=False
-    )
+    # allow_large_result=False incompatible with json b/401630655
+    with bigframes.option_context(
+        "bigquery.allow_large_results", True, "experiments.blob", True
+    ):
+        actual = images_mm_df["blob_col"].blob.metadata().to_pandas()
+        expected = pd.Series(
+            [
+                {
+                    "content_type": "image/jpeg",
+                    "md5_hash": "e130ad042261a1883cd2cc06831cf748",
+                    "size": 338390,
+                    "updated": 1739574332000000,
+                },
+                {
+                    "content_type": "image/jpeg",
+                    "md5_hash": "e2ae3191ff2b809fd0935f01a537c650",
+                    "size": 43333,
+                    "updated": 1739574332000000,
+                },
+            ],
+            name="metadata",
+            dtype=db_dtypes.JSONDtype(),
+        )
+
+        pd.testing.assert_series_equal(
+            actual, expected, check_dtype=False, check_index_type=False
+        )
 
 
 def test_blob_content_type(images_mm_df: bpd.DataFrame):

tests/system/small/ml/test_core.py

@@ -401,9 +401,10 @@ def test_model_generate_text(
         "top_p": 0.5,
         "flatten_json_output": True,
     }
+    # Until b/401630655 is resolved, json not compatible with allow_large_results=False
     df = bqml_palm2_text_generator_model.generate_text(
         llm_text_df, options=options
-    ).to_pandas()
+    ).to_pandas(allow_large_results=True)
 
     utils.check_pandas_df_schema_and_index(
         df, columns=utils.ML_GENERATE_TEXT_OUTPUT, index=3, col_exact=False

tests/system/small/ml/test_llm.py

@@ -24,6 +24,13 @@
 from tests.system import utils
 
 
+# Until b/401630655 is resolved, ML apis return json, not compatible with allow_large_results=False
+@pytest.fixture(scope="module", autouse=True)
+def always_create_table():
+    with bigframes.option_context("bigquery.allow_large_results", True):
+        yield
+
+
 def test_create_load_text_generator_model(
     palm2_text_generator_model, dataset_id, bq_connection
 ):

tests/system/small/test_dataframe.py

@@ -163,11 +163,11 @@ def test_df_construct_from_dict():
     )
 
 
-def test_df_construct_inline_respects_location():
+def test_df_construct_inline_respects_location(reset_default_session_and_location):
     # Note: This starts a thread-local session.
     with bpd.option_context("bigquery.location", "europe-west1"):
         df = bpd.DataFrame([[1, 2, 3], [4, 5, 6]])
-        repr(df)
+        df.to_gbq()
         assert df.query_job is not None
         table = bpd.get_global_session().bqclient.get_table(df.query_job.destination)
 
@@ -666,7 +666,8 @@ def test_df_peek(scalars_dfs_maybe_ordered):
 
     session = scalars_df._block.session
     slot_millis_sum = session.slot_millis_sum
-    peek_result = scalars_df.peek(n=3, force=False)
+    # allow_large_results=False needed to get slot_millis_sum statistics only
+    peek_result = scalars_df.peek(n=3, force=False, allow_large_results=True)
 
     assert session.slot_millis_sum - slot_millis_sum > 1000
     pd.testing.assert_index_equal(scalars_pandas_df.columns, peek_result.columns)
@@ -4584,12 +4585,13 @@ def test_df_drop_duplicates(scalars_df_index, scalars_pandas_df_index, keep, sub
     ],
 )
 def test_df_drop_duplicates_w_json(json_df, keep):
-    bf_df = json_df.drop_duplicates(keep=keep).to_pandas()
+    bf_df = json_df.drop_duplicates(keep=keep).to_pandas(allow_large_results=True)
 
     # drop_duplicates relies on pa.compute.dictionary_encode, which is incompatible
     # with Arrow string extension types. Temporary conversion to standard Pandas
     # strings is required.
-    json_pandas_df = json_df.to_pandas()
+    # allow_large_results=True for b/401630655
+    json_pandas_df = json_df.to_pandas(allow_large_results=True)
     json_pandas_df["json_col"] = json_pandas_df["json_col"].astype(
         pd.StringDtype(storage="pyarrow")
     )
@@ -4951,14 +4953,16 @@ def test_df_bool_interpretation_error(scalars_df_index):
 
 
 def test_query_job_setters(scalars_df_default_index: dataframe.DataFrame):
-    job_ids = set()
-    repr(scalars_df_default_index)
-    assert scalars_df_default_index.query_job is not None
-    job_ids.add(scalars_df_default_index.query_job.job_id)
-    scalars_df_default_index.to_pandas()
-    job_ids.add(scalars_df_default_index.query_job.job_id)
+    # if allow_large_results=False, might not create query job
+    with bigframes.option_context("bigquery.allow_large_results", True):
+        job_ids = set()
+        repr(scalars_df_default_index)
+        assert scalars_df_default_index.query_job is not None
+        job_ids.add(scalars_df_default_index.query_job.job_id)
+        scalars_df_default_index.to_pandas(allow_large_results=True)
+        job_ids.add(scalars_df_default_index.query_job.job_id)
 
-    assert len(job_ids) == 2
+        assert len(job_ids) == 2
 
 
 def test_df_cached(scalars_df_index):
@@ -5196,7 +5200,12 @@ def test_to_pandas_downsampling_option_override(session):
     df = session.read_gbq("bigframes-dev.bigframes_tests_sys.batting")
     download_size = 1
 
-    df = df.to_pandas(max_download_size=download_size, sampling_method="head")
+    # limits only apply for allow_large_result=True
+    df = df.to_pandas(
+        max_download_size=download_size,
+        sampling_method="head",
+        allow_large_results=True,
+    )
 
     total_memory_bytes = df.memory_usage(deep=True).sum()
     total_memory_mb = total_memory_bytes / (1024 * 1024)

tests/system/small/test_dataframe_io.py

@@ -254,27 +254,31 @@ def test_to_pandas_array_struct_correct_result(session):
 def test_to_pandas_override_global_option(scalars_df_index):
     # Direct call to_pandas uses global default setting (allow_large_results=True),
     # table has 'bqdf' prefix.
-    scalars_df_index.to_pandas()
-    table_id = scalars_df_index._query_job.destination.table_id
-    assert table_id.startswith("bqdf")
+    with bigframes.option_context("bigquery.allow_large_results", True):
 
-    # When allow_large_results=False, a query_job object should not be created.
-    # Therefore, the table_id should remain unchanged.
-    scalars_df_index.to_pandas(allow_large_results=False)
-    assert scalars_df_index._query_job.destination.table_id == table_id
+        scalars_df_index.to_pandas()
+        table_id = scalars_df_index._query_job.destination.table_id
+        assert table_id.startswith("bqdf")
+
+        # When allow_large_results=False, a query_job object should not be created.
+        # Therefore, the table_id should remain unchanged.
+        scalars_df_index.to_pandas(allow_large_results=False)
+        assert scalars_df_index._query_job.destination.table_id == table_id
 
 
 def test_to_arrow_override_global_option(scalars_df_index):
     # Direct call to_arrow uses global default setting (allow_large_results=True),
     # table has 'bqdf' prefix.
-    scalars_df_index.to_arrow()
-    table_id = scalars_df_index._query_job.destination.table_id
-    assert table_id.startswith("bqdf")
-
-    # When allow_large_results=False, a query_job object should not be created.
-    # Therefore, the table_id should remain unchanged.
-    scalars_df_index.to_arrow(allow_large_results=False)
-    assert scalars_df_index._query_job.destination.table_id == table_id
+    with bigframes.option_context("bigquery.allow_large_results", True):
+
+        scalars_df_index.to_arrow()
+        table_id = scalars_df_index._query_job.destination.table_id
+        assert table_id.startswith("bqdf")
+
+        # When allow_large_results=False, a query_job object should not be created.
+        # Therefore, the table_id should remain unchanged.
+        scalars_df_index.to_arrow(allow_large_results=False)
+        assert scalars_df_index._query_job.destination.table_id == table_id
 
 
 def test_load_json_w_json_string_items(session):

tests/system/small/test_encryption.py

@@ -41,6 +41,7 @@ def bq_cmek() -> str:
 
 @pytest.fixture(scope="module")
 def session_with_bq_cmek(bq_cmek) -> bigframes.Session:
+    # allow_large_results = False might not create table, and therefore no encryption config
     session = bigframes.Session(bigframes.BigQueryOptions(kms_key_name=bq_cmek))
 
     return session
@@ -52,7 +53,7 @@ def _assert_bq_table_is_encrypted(
     session: bigframes.Session,
 ):
     # Materialize the data in BQ
-    repr(df)
+    df.to_gbq()
 
     # The df should be backed by a query job with intended encryption on the result table
     assert df.query_job is not None

tests/system/small/test_groupby.py

@@ -546,8 +546,7 @@ def test_dataframe_groupby_nonnumeric_with_mean():
     )
     pd_result = df.groupby(["key1", "key2"]).mean()
 
-    with bpd.option_context("bigquery.location", "US"):
-        bf_result = bpd.DataFrame(df).groupby(["key1", "key2"]).mean().to_pandas()
+    bf_result = bpd.DataFrame(df).groupby(["key1", "key2"]).mean().to_pandas()
 
     pd.testing.assert_frame_equal(
         pd_result, bf_result, check_index_type=False, check_dtype=False