perf: Optimize repr for unordered gbq table (#1778)

Author: TrevorBergeron

bigframes/core/rewrite/__init__.py

@@ -22,7 +22,7 @@
     try_reduce_to_local_scan,
     try_reduce_to_table_scan,
 )
-from bigframes.core.rewrite.slices import pull_up_limits, rewrite_slice
+from bigframes.core.rewrite.slices import pull_out_limit, pull_up_limits, rewrite_slice
 from bigframes.core.rewrite.timedeltas import rewrite_timedelta_expressions
 from bigframes.core.rewrite.windows import rewrite_range_rolling
 
@@ -32,6 +32,7 @@
     "rewrite_slice",
     "rewrite_timedelta_expressions",
     "pull_up_limits",
+    "pull_out_limit",
     "remap_variables",
     "defer_order",
     "column_pruning",

bigframes/core/rewrite/slices.py

@@ -26,7 +26,7 @@
 
 
 def pull_up_limits(root: nodes.ResultNode) -> nodes.ResultNode:
-    new_child, pulled_limit = _pullup_slice_inner(root.child)
+    new_child, pulled_limit = pull_out_limit(root.child)
     if new_child == root.child:
         return root
     elif pulled_limit is None:
@@ -37,7 +37,7 @@ def pull_up_limits(root: nodes.ResultNode) -> nodes.ResultNode:
         return dataclasses.replace(root, child=new_child, limit=new_limit)
 
 
-def _pullup_slice_inner(
+def pull_out_limit(
     root: nodes.BigFrameNode,
 ) -> Tuple[nodes.BigFrameNode, Optional[int]]:
     """
@@ -53,15 +53,15 @@ def _pullup_slice_inner(
             assert root.step == 1
             assert root.stop is not None
             limit = root.stop
-            new_root, prior_limit = _pullup_slice_inner(root.child)
+            new_root, prior_limit = pull_out_limit(root.child)
             if (prior_limit is not None) and (prior_limit < limit):
                 limit = prior_limit
             return new_root, limit
     elif (
         isinstance(root, (nodes.SelectionNode, nodes.ProjectionNode))
         and root.row_preserving
     ):
-        new_child, prior_limit = _pullup_slice_inner(root.child)
+        new_child, prior_limit = pull_out_limit(root.child)
         if prior_limit is not None:
             return root.transform_children(lambda _: new_child), prior_limit
     # Most ops don't support pulling up slice, like filter, agg, join, etc.

bigframes/core/tree_properties.py

@@ -45,26 +45,13 @@ def can_fast_head(node: nodes.BigFrameNode) -> bool:
     # To do fast head operation:
     # (1) the underlying data must be arranged/indexed according to the logical ordering
     # (2) transformations must support pushing down LIMIT or a filter on row numbers
-    return has_fast_offset_address(node) or has_fast_offset_address(node)
-
-
-def has_fast_orderby_limit(node: nodes.BigFrameNode) -> bool:
-    """True iff ORDER BY LIMIT can be performed without a large full table scan."""
-    # TODO: In theory compatible with some Slice nodes, potentially by adding OFFSET
-    if isinstance(node, nodes.LeafNode):
-        return node.fast_ordered_limit
-    if isinstance(node, (nodes.ProjectionNode, nodes.SelectionNode)):
-        return has_fast_orderby_limit(node.child)
-    return False
-
-
-def has_fast_offset_address(node: nodes.BigFrameNode) -> bool:
-    """True iff specific offsets can be scanned without a large full table scan."""
-    # TODO: In theory can push offset lookups through slice operators by translating indices
-    if isinstance(node, nodes.LeafNode):
-        return node.fast_offsets
+    if isinstance(node, nodes.ReadLocalNode):
+        # always cheap to push slice into local data
+        return True
+    if isinstance(node, nodes.ReadTableNode):
+        return (node.source.ordering is None) or (node.fast_ordered_limit)
     if isinstance(node, (nodes.ProjectionNode, nodes.SelectionNode)):
-        return has_fast_offset_address(node.child)
+        return can_fast_head(node.child)
     return False
 
 

bigframes/session/loader.py

@@ -387,6 +387,7 @@ def read_gbq_table(  # type: ignore[overload-overlap]
         enable_snapshot: bool = ...,
         dry_run: Literal[False] = ...,
         force_total_order: Optional[bool] = ...,
+        n_rows: Optional[int] = None,
     ) -> dataframe.DataFrame:
         ...
 
@@ -408,6 +409,7 @@ def read_gbq_table(
         enable_snapshot: bool = ...,
         dry_run: Literal[True] = ...,
         force_total_order: Optional[bool] = ...,
+        n_rows: Optional[int] = None,
     ) -> pandas.Series:
         ...
 
@@ -428,6 +430,7 @@ def read_gbq_table(
         enable_snapshot: bool = True,
         dry_run: bool = False,
         force_total_order: Optional[bool] = None,
+        n_rows: Optional[int] = None,
     ) -> dataframe.DataFrame | pandas.Series:
         import bigframes._tools.strings
         import bigframes.dataframe as dataframe
@@ -618,6 +621,7 @@ def read_gbq_table(
             at_time=time_travel_timestamp if enable_snapshot else None,
             primary_key=primary_key,
             session=self._session,
+            n_rows=n_rows,
         )
         # if we don't have a unique index, we order by row hash if we are in strict mode
         if (
@@ -852,6 +856,7 @@ def read_gbq_query(
             columns=columns,
             use_cache=configuration["query"]["useQueryCache"],
             force_total_order=force_total_order,
+            n_rows=query_job.result().total_rows,
             # max_results and filters are omitted because they are already
             # handled by to_query(), above.
         )

bigframes/session/read_api_execution.py

@@ -39,12 +39,17 @@ def execute(
         ordered: bool,
         peek: Optional[int] = None,
     ) -> Optional[executor.ExecuteResult]:
-        node = self._try_adapt_plan(plan, ordered)
-        if not node:
+        adapt_result = self._try_adapt_plan(plan, ordered)
+        if not adapt_result:
             return None
+        node, limit = adapt_result
         if node.explicitly_ordered and ordered:
             return None
 
+        if limit is not None:
+            if peek is None or limit < peek:
+                peek = limit
+
         import google.cloud.bigquery_storage_v1.types as bq_storage_types
         from google.protobuf import timestamp_pb2
 
@@ -117,11 +122,20 @@ def _try_adapt_plan(
         self,
         plan: bigframe_node.BigFrameNode,
         ordered: bool,
-    ) -> Optional[nodes.ReadTableNode]:
+    ) -> Optional[tuple[nodes.ReadTableNode, Optional[int]]]:
         """
-        Tries to simplify the plan to an equivalent single ReadTableNode. Otherwise, returns None.
+        Tries to simplify the plan to an equivalent single ReadTableNode and a limit. Otherwise, returns None.
         """
+        plan, limit = rewrite.pull_out_limit(plan)
+        # bake_order does not allow slice ops
+        plan = plan.bottom_up(rewrite.rewrite_slice)
         if not ordered:
             # gets rid of order_by ops
             plan = rewrite.bake_order(plan)
-        return rewrite.try_reduce_to_table_scan(plan)
+        read_table_node = rewrite.try_reduce_to_table_scan(plan)
+        if read_table_node is None:
+            return None
+        if (limit is not None) and (read_table_node.source.ordering is not None):
+            # read api can only use physical ordering to limit, not a logical ordering
+            return None
+        return (read_table_node, limit)

tests/system/small/session/test_read_gbq_colab.py

@@ -19,6 +19,7 @@
 
 
 def test_read_gbq_colab_to_pandas_batches_preserves_order_by(maybe_ordered_session):
+    executions_before_sql = maybe_ordered_session._metrics.execution_count
     df = maybe_ordered_session._read_gbq_colab(
         """
         SELECT
@@ -32,16 +33,67 @@ def test_read_gbq_colab_to_pandas_batches_preserves_order_by(maybe_ordered_sessi
         LIMIT 300
         """
     )
+    executions_before_python = maybe_ordered_session._metrics.execution_count
     batches = df.to_pandas_batches(
         page_size=100,
     )
+    executions_after = maybe_ordered_session._metrics.execution_count
 
     total_rows = 0
     for batch in batches:
         assert batch["total"].is_monotonic_decreasing
         total_rows += len(batch.index)
 
     assert total_rows > 0
+    assert executions_after == executions_before_python == executions_before_sql + 1
+
+
+def test_read_gbq_colab_peek_avoids_requery(maybe_ordered_session):
+    executions_before_sql = maybe_ordered_session._metrics.execution_count
+    df = maybe_ordered_session._read_gbq_colab(
+        """
+        SELECT
+            name,
+            SUM(number) AS total
+        FROM
+            `bigquery-public-data.usa_names.usa_1910_2013`
+        WHERE state LIKE 'W%'
+        GROUP BY name
+        ORDER BY total DESC
+        LIMIT 300
+        """
+    )
+    executions_before_python = maybe_ordered_session._metrics.execution_count
+    result = df.peek(100)
+    executions_after = maybe_ordered_session._metrics.execution_count
+
+    # Ok, this isn't guaranteed by peek, but should happen with read api based impl
+    # if starts failing, maybe stopped using read api?
+    assert result["total"].is_monotonic_decreasing
+
+    assert len(result) == 100
+    assert executions_after == executions_before_python == executions_before_sql + 1
+
+
+def test_read_gbq_colab_repr_avoids_requery(maybe_ordered_session):
+    executions_before_sql = maybe_ordered_session._metrics.execution_count
+    df = maybe_ordered_session._read_gbq_colab(
+        """
+        SELECT
+            name,
+            SUM(number) AS total
+        FROM
+            `bigquery-public-data.usa_names.usa_1910_2013`
+        WHERE state LIKE 'W%'
+        GROUP BY name
+        ORDER BY total DESC
+        LIMIT 300
+        """
+    )
+    executions_before_python = maybe_ordered_session._metrics.execution_count
+    _ = repr(df)
+    executions_after = maybe_ordered_session._metrics.execution_count
+    assert executions_after == executions_before_python == executions_before_sql + 1
 
 
 def test_read_gbq_colab_includes_formatted_scalars(session):