use limit rewrite for slice support

Author: tswast

bigframes/core/nodes.py

@@ -154,6 +154,16 @@ def is_limit(self) -> bool:
             and (self.stop > 0)
         )
 
+    @property
+    def is_noop(self) -> bool:
+        """Returns whether this node doesn't actually change the results."""
+        # TODO: Handle tail case.
+        return (
+            ((not self.start) or (self.start == 0))
+            and (self.step == 1)
+            and ((self.stop is None) or (self.stop == self.row_count))
+        )
+
     @property
     def row_count(self) -> typing.Optional[int]:
         child_length = self.child.row_count

bigframes/core/rewrite/scan_reduction.py

@@ -16,6 +16,7 @@
 from typing import Optional
 
 from bigframes.core import nodes
+import bigframes.core.rewrite.slices
 
 
 def try_reduce_to_table_scan(root: nodes.BigFrameNode) -> Optional[nodes.ReadTableNode]:
@@ -28,7 +29,15 @@ def try_reduce_to_table_scan(root: nodes.BigFrameNode) -> Optional[nodes.ReadTab
     return None
 
 
-def try_reduce_to_local_scan(node: nodes.BigFrameNode) -> Optional[nodes.ReadLocalNode]:
+def try_reduce_to_local_scan(
+    node: nodes.BigFrameNode,
+) -> Optional[tuple[nodes.ReadLocalNode, Optional[int]]]:
+    """Create a ReadLocalNode with optional limit, if possible.
+
+    Similar to ReadApiSemiExecutor._try_adapt_plan.
+    """
+    node, limit = bigframes.core.rewrite.slices.pull_out_limit(node)
+
     if not all(
         map(
             lambda x: isinstance(x, (nodes.ReadLocalNode, nodes.SelectionNode)),
@@ -38,7 +47,7 @@ def try_reduce_to_local_scan(node: nodes.BigFrameNode) -> Optional[nodes.ReadLoc
         return None
     result = node.bottom_up(merge_scan)
     if isinstance(result, nodes.ReadLocalNode):
-        return result
+        return result, limit
     return None
 
 

bigframes/core/rewrite/slices.py

@@ -57,6 +57,9 @@ def pull_out_limit(
             if (prior_limit is not None) and (prior_limit < limit):
                 limit = prior_limit
             return new_root, limit
+        if root.is_noop:
+            new_root, prior_limit = pull_out_limit(root.child)
+            return new_root, prior_limit
     elif (
         isinstance(root, (nodes.SelectionNode, nodes.ProjectionNode))
         and root.row_preserving

bigframes/session/local_scan_executor.py

@@ -15,7 +15,7 @@
 
 from typing import Optional
 
-from bigframes.core import bigframe_node, nodes, rewrite
+from bigframes.core import bigframe_node, rewrite
 from bigframes.session import executor, semi_executor
 
 
@@ -30,30 +30,15 @@ def execute(
         ordered: bool,
         peek: Optional[int] = None,
     ) -> Optional[executor.ExecuteResult]:
-        rewrite_node = plan
-
-        # Implement top-level slice here so that we don't have to implement
-        # slice in all ReadLocalNode compilers.
-        slice_start: Optional[int] = None
-        slice_stop: Optional[int] = None
-
-        if isinstance(plan, nodes.SliceNode):
-            # These slice features are not supported by pyarrow.Table.slice. Must
-            # have a non-negative start and stop and no custom step size.
-            if (
-                (plan.step is not None and plan.step != 1)
-                or (plan.start is not None and plan.start < 0)
-                or (plan.stop is not None and plan.stop < 0)
-            ):
-                return None
+        reduced_result = rewrite.try_reduce_to_local_scan(plan)
+        if not reduced_result:
+            return None
 
-            slice_start = plan.start
-            slice_stop = plan.stop
-            rewrite_node = plan.child
+        node, limit = reduced_result
 
-        node = rewrite.try_reduce_to_local_scan(rewrite_node)
-        if not node:
-            return None
+        if limit is not None:
+            if peek is None or limit < peek:
+                peek = limit
 
         # TODO: Can support some sorting
         offsets_col = node.offsets_col.sql if (node.offsets_col is not None) else None
@@ -67,23 +52,7 @@ def execute(
 
         arrow_table = arrow_table.select(needed_cols)
         arrow_table = arrow_table.rename_columns([id.sql for id in node.ids])
-
-        if slice_start is not None or slice_stop is not None:
-            slice_length: Optional[int] = None
-
-            if slice_stop is not None:
-                if slice_start is None:
-                    slice_length = slice_stop
-                else:
-                    slice_length = slice_stop - slice_start
-
-            arrow_table = arrow_table.slice(
-                offset=slice_start if slice_start is not None else 0,
-                length=slice_length,
-            )
-            total_rows = arrow_table.num_rows
-        else:
-            total_rows = node.row_count
+        total_rows = node.row_count
 
         if (peek is not None) and (total_rows is not None):
             total_rows = min(peek, total_rows)

tests/unit/session/test_local_scan_executor.py

@@ -51,12 +51,13 @@ def create_read_local_node(arrow_table: pyarrow.Table):
 @pytest.mark.parametrize(
     ("start", "stop", "expected_rows"),
     (
+        # No-op slices.
         (None, None, 10),
         (0, None, 10),
-        (4, None, 6),
         (None, 10, 10),
+        # Slices equivalent to limits.
         (None, 7, 7),
-        (1, 9, 8),
+        (0, 3, 3),
     ),
 )
 def test_local_scan_executor_with_slice(start, stop, expected_rows, object_under_test):
@@ -87,6 +88,8 @@ def test_local_scan_executor_with_slice(start, stop, expected_rows, object_under
         (None, -1, 1),
         (None, None, 2),
         (None, None, -1),
+        (4, None, 6),
+        (1, 9, 8),
     ),
 )
 def test_local_scan_executor_with_slice_unsupported_inputs(