revert a benchmark

Author: shuoweil

tests/benchmark/read_gbq_colab/aggregate_output.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 import pathlib
 
+import benchmark.utils as utils
+
 import bigframes.pandas as bpd
-import bigframes.session.execution_spec
-import tests.benchmark.utils as utils
 
 PAGE_SIZE = utils.READ_GBQ_COLAB_PAGE_SIZE
 
@@ -25,17 +25,9 @@ def aggregate_output(*, project_id, dataset_id, table_id):
     # e.g. "{local_inline}" or "{local_large}"
     df = bpd._read_gbq_colab(f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}")
 
-    # Call the executor directly to isolate the query execution time
-    # from other DataFrame overhead for this benchmark.
-    execute_result = df._block.session._executor.execute(
-        df._block.expr,
-        execution_spec=bigframes.session.execution_spec.ExecutionSpec(
-            ordered=True, promise_under_10gb=False
-        ),
-    )
-    assert execute_result.total_rows is not None and execute_result.total_rows >= 0
-    batches = execute_result.to_pandas_batches(page_size=PAGE_SIZE)
-    next(iter(batches))
+    # Simulate getting the first page, since we'll always do that first in the UI.
+    df.shape
+    next(iter(df.to_pandas_batches(page_size=PAGE_SIZE)))
 
     # To simulate very small rows that can only fit a boolean,
     # some tables don't have an integer column. If an integer column is available,
@@ -50,20 +42,9 @@ def aggregate_output(*, project_id, dataset_id, table_id):
         .groupby("rounded")
         .sum(numeric_only=True)
     )
-    execute_result_aggregated = df_aggregated._block.session._executor.execute(
-        df_aggregated._block.expr,
-        execution_spec=bigframes.session.execution_spec.ExecutionSpec(
-            ordered=True, promise_under_10gb=False
-        ),
-    )
-    assert (
-        execute_result_aggregated.total_rows is not None
-        and execute_result_aggregated.total_rows >= 0
-    )
-    batches_aggregated = execute_result_aggregated.to_pandas_batches(
-        page_size=PAGE_SIZE
-    )
-    next(iter(batches_aggregated))
+
+    df_aggregated.shape
+    next(iter(df_aggregated.to_pandas_batches(page_size=PAGE_SIZE)))
 
 
 if __name__ == "__main__":