more fixes

Author: TrevorBergeron

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -108,7 +108,7 @@ def __init__(
         self._upload_lock = threading.Lock()
         self._gbq_executor = direct_gbq_execution.DirectGbqExecutor(
             bqclient,
-            compiler=compile.compiler,
+            compiler=compile.compiler().compile_sql,
             bqstoragereadclient=bqstoragereadclient,
             metrics=self.metrics,
             publisher=self._publisher,

packages/bigframes/bigframes/testing/engine_utils.py

@@ -15,16 +15,21 @@
 import pandas.testing
 
 from bigframes.core import nodes
-from bigframes.session import semi_executor
+from bigframes.session import semi_executor, execution_spec
+
+
+SPEC = execution_spec.ExecutionSpec(
+    ordered=True,
+)
 
 
 def assert_equivalence_execution(
     node: nodes.BigFrameNode,
     engine1: semi_executor.SemiExecutor,
     engine2: semi_executor.SemiExecutor,
 ):
-    e1_result = engine1.execute(node, ordered=True)
-    e2_result = engine2.execute(node, ordered=True)
+    e1_result = engine1.execute(node, SPEC)
+    e2_result = engine2.execute(node, SPEC)
     assert e1_result is not None
     assert e2_result is not None
     # Convert to pandas, as pandas has better comparison utils than arrow

packages/bigframes/tests/system/conftest.py

@@ -27,6 +27,7 @@
 import google.cloud.bigquery_connection_v1 as bigquery_connection_v1
 import google.cloud.exceptions
 import google.cloud.functions_v2 as functions_v2
+import google.cloud.bigquery_storage_v1
 import google.cloud.resourcemanager_v3 as resourcemanager_v3
 import google.cloud.storage as storage  # type: ignore
 import numpy as np
@@ -103,6 +104,13 @@ def bigquery_client(session: bigframes.Session) -> bigquery.Client:
     return session.bqclient
 
 
+@pytest.fixture(scope="session")
+def bigquery_storage_read_client(
+    session: bigframes.Session,
+) -> google.cloud.bigquery_storage_v1.BigQueryReadClient:
+    return session.bqstoragereadclient
+
+
 @pytest.fixture(scope="session")
 def bigquery_client_tokyo(session_tokyo: bigframes.Session) -> bigquery.Client:
     return session_tokyo.bqclient

packages/bigframes/tests/system/small/engines/conftest.py

@@ -19,6 +19,7 @@
 from google.cloud import bigquery
 
 import bigframes
+import google.cloud.bigquery_storage_v1
 from bigframes.core import ArrayValue, events, local_data
 from bigframes.session import (
     direct_gbq_execution,
@@ -45,19 +46,28 @@ def fake_session() -> Generator[bigframes.Session, None, None]:
 
 
 @pytest.fixture(scope="session", params=["pyarrow", "polars", "bq", "bq-sqlglot"])
-def engine(request, bigquery_client: bigquery.Client) -> semi_executor.SemiExecutor:
+def engine(
+    request,
+    bigquery_client: bigquery.Client,
+    bigquery_storage_read_client: google.cloud.bigquery_storage_v1.BigQueryReadClient,
+) -> semi_executor.SemiExecutor:
     if request.param == "pyarrow":
         return local_scan_executor.LocalScanExecutor()
     if request.param == "polars":
         return polars_executor.PolarsExecutor()
     publisher = events.Publisher()
     if request.param == "bq":
         return direct_gbq_execution.DirectGbqExecutor(
-            bigquery_client, publisher=publisher
+            bigquery_client,
+            bqstoragereadclient=bigquery_storage_read_client,
+            publisher=publisher,
         )
     if request.param == "bq-sqlglot":
         return direct_gbq_execution.DirectGbqExecutor(
-            bigquery_client, compiler="sqlglot", publisher=publisher
+            bigquery_client,
+            bqstoragereadclient=bigquery_storage_read_client,
+            compiler="sqlglot",
+            publisher=publisher,
         )
     raise ValueError(f"Unrecognized param: {request.param}")
 

packages/bigframes/tests/unit/session/test_local_scan_executor.py

@@ -17,9 +17,13 @@
 import pytest
 
 from bigframes.core import identifiers, local_data, nodes
-from bigframes.session import local_scan_executor
+from bigframes.session import local_scan_executor, execution_spec
 from bigframes.testing import mocks
 
+SPEC = execution_spec.ExecutionSpec(
+    ordered=True,
+)
+
 
 @pytest.fixture
 def object_under_test():
@@ -72,7 +76,7 @@ def test_local_scan_executor_with_slice(start, stop, expected_rows, object_under
         stop=stop,
     )
 
-    result = object_under_test.execute(plan, ordered=True)
+    result = object_under_test.execute(plan, SPEC)
     result_table = pyarrow.Table.from_batches(result.batches().arrow_batches)
     assert result_table.num_rows == expected_rows
 
@@ -98,4 +102,4 @@ def test_local_scan_executor_with_slice_unsupported_inputs(
         stop=stop,
         step=step,
     )
-    assert object_under_test.execute(plan, ordered=True) is None
+    assert object_under_test.execute(plan, SPEC) is None