perf: Fold row count ops when known

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -204,7 +204,17 @@ def get_column_type(self, key: str) -> bigframes.dtypes.Dtype:
 
     def row_count(self) -> ArrayValue:
         """Get number of rows in ArrayValue as a single-entry ArrayValue."""
-        return ArrayValue(nodes.RowCountNode(child=self.node))
+        return ArrayValue(
+            nodes.AggregateNode(
+                child=self.node,
+                aggregations=(
+                    (
+                        ex.NullaryAggregation(agg_ops.size_op),
+                        ids.ColumnId(bigframes.core.guid.generate_guid()),
+                    ),
+                ),
+            )
+        )
 
     # Operations
     def filter_by_id(self, predicate_id: str, keep_null: bool = False) -> ArrayValue:

bigframes/core/blocks.py

@@ -221,7 +221,14 @@ def shape(self) -> typing.Tuple[int, int]:
             except Exception:
                 pass
 
-        row_count = self.session._executor.get_row_count(self.expr)
+        row_count = next(
+            iter(
+                self.session._executor.execute(self.expr.row_count())
+                .to_arrow_table()
+                .to_pydict()
+                .values()
+            )
+        )[0]
         return (row_count, len(self.value_columns))
 
     @property
@@ -485,7 +492,7 @@ def to_arrow(
         *,
         ordered: bool = True,
         allow_large_results: Optional[bool] = None,
-    ) -> Tuple[pa.Table, bigquery.QueryJob]:
+    ) -> Tuple[pa.Table, Optional[bigquery.QueryJob]]:
         """Run query and download results as a pyarrow Table."""
         execute_result = self.session._executor.execute(
             self.expr, ordered=ordered, use_explicit_destination=allow_large_results
@@ -659,7 +666,7 @@ def _materialize_local(
 
         # TODO: Maybe materialize before downsampling
         # Some downsampling methods
-        if fraction < 1:
+        if fraction < 1 and (execute_result.total_rows is not None):
             if not sample_config.enable_downsampling:
                 raise RuntimeError(
                     f"The data size ({table_mb:.2f} MB) exceeds the maximum download limit of "
@@ -690,7 +697,6 @@ def _materialize_local(
                 MaterializationOptions(ordered=materialize_options.ordered)
             )
         else:
-            total_rows = execute_result.total_rows
             arrow = execute_result.to_arrow_table()
             df = io_pandas.arrow_to_pandas(arrow, schema=self.expr.schema)
             self._copy_index_to_pandas(df)
@@ -1570,12 +1576,19 @@ def retrieve_repr_request_results(
 
         # head caches full underlying expression, so row_count will be free after
         head_result = self.session._executor.head(self.expr, max_results)
-        count = self.session._executor.get_row_count(self.expr)
+        row_count = next(
+            iter(
+                self.session._executor.execute(self.expr.row_count())
+                .to_arrow_table()
+                .to_pydict()
+                .values()
+            )
+        )[0]
 
         arrow = head_result.to_arrow_table()
         df = io_pandas.arrow_to_pandas(arrow, schema=self.expr.schema)
         self._copy_index_to_pandas(df)
-        return df, count, head_result.query_job
+        return df, row_count, head_result.query_job
 
     def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]:
         expr, result_id = self._expr.promote_offsets()

bigframes/core/compile/compiler.py

@@ -270,11 +270,6 @@ def compile_concat(node: nodes.ConcatNode, *children: compiled.UnorderedIR):
     return concat_impl.concat_unordered(children, output_ids)
 
 
-@_compile_node.register
-def compile_rowcount(node: nodes.RowCountNode, child: compiled.UnorderedIR):
-    return child.row_count(name=node.col_id.sql)
-
-
 @_compile_node.register
 def compile_aggregate(node: nodes.AggregateNode, child: compiled.UnorderedIR):
     aggs = tuple((agg, id.sql) for agg, id in node.aggregations)

bigframes/core/compile/polars/compiler.py

@@ -252,11 +252,6 @@ def compile_projection(self, node: nodes.ProjectionNode):
         ]
         return self.compile_node(node.child).with_columns(new_cols)
 
-    @compile_node.register
-    def compile_rowcount(self, node: nodes.RowCountNode):
-        df = cast(pl.LazyFrame, self.compile_node(node.child))
-        return df.select(pl.len().alias(node.col_id.sql))
-
     @compile_node.register
     def compile_offsets(self, node: nodes.PromoteOffsetsNode):
         return self.compile_node(node.child).with_columns(

bigframes/core/nodes.py

@@ -1256,55 +1256,6 @@ def remap_refs(
         return dataclasses.replace(self, assignments=new_fields)
 
 
-# TODO: Merge RowCount into Aggregate Node?
-# Row count can be compute from table metadata sometimes, so it is a bit special.
-@dataclasses.dataclass(frozen=True, eq=False)
-class RowCountNode(UnaryNode):
-    col_id: identifiers.ColumnId = identifiers.ColumnId("count")
-
-    @property
-    def row_preserving(self) -> bool:
-        return False
-
-    @property
-    def non_local(self) -> bool:
-        return True
-
-    @property
-    def fields(self) -> Sequence[Field]:
-        return (Field(self.col_id, bigframes.dtypes.INT_DTYPE, nullable=False),)
-
-    @property
-    def variables_introduced(self) -> int:
-        return 1
-
-    @property
-    def defines_namespace(self) -> bool:
-        return True
-
-    @property
-    def row_count(self) -> Optional[int]:
-        return 1
-
-    @property
-    def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
-        return (self.col_id,)
-
-    @property
-    def consumed_ids(self) -> COLUMN_SET:
-        return frozenset()
-
-    def remap_vars(
-        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
-    ) -> RowCountNode:
-        return dataclasses.replace(self, col_id=mappings.get(self.col_id, self.col_id))
-
-    def remap_refs(
-        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
-    ) -> RowCountNode:
-        return self
-
-
 @dataclasses.dataclass(frozen=True, eq=False)
 class AggregateNode(UnaryNode):
     aggregations: typing.Tuple[typing.Tuple[ex.Aggregation, identifiers.ColumnId], ...]

bigframes/core/rewrite/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from bigframes.core.rewrite.fold_row_count import fold_row_counts
 from bigframes.core.rewrite.identifiers import remap_variables
 from bigframes.core.rewrite.implicit_align import try_row_join
 from bigframes.core.rewrite.legacy_align import legacy_join_as_projection
@@ -38,4 +39,5 @@
     "try_reduce_to_table_scan",
     "bake_order",
     "try_reduce_to_local_scan",
+    "fold_row_counts",
 ]

bigframes/core/rewrite/fold_row_count.py

@@ -0,0 +1,44 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import pyarrow as pa
+
+from bigframes import dtypes
+from bigframes.core import local_data, nodes
+from bigframes.operations import aggregations
+
+
+def fold_row_counts(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+    if not isinstance(node, nodes.AggregateNode):
+        return node
+    if len(node.by_column_ids) > 0:
+        return node
+    if node.child.row_count is None:
+        return node
+    for agg, _ in node.aggregations:
+        if agg.op != aggregations.size_op:
+            return node
+    local_data_source = local_data.ManagedArrowTable.from_pyarrow(
+        pa.table({"count": pa.array([node.child.row_count], type=pa.int64())})
+    )
+    scan_list = nodes.ScanList(
+        tuple(
+            nodes.ScanItem(out_id, dtypes.INT_DTYPE, "count")
+            for _, out_id in node.aggregations
+        )
+    )
+    return nodes.ReadLocalNode(
+        local_data_source=local_data_source, scan_list=scan_list, session=node.session
+    )

bigframes/core/rewrite/order.py

@@ -211,11 +211,6 @@ def pull_up_order_inner(
             )
             new_order = child_order.remap_column_refs(new_select_node.get_id_mapping())
             return new_select_node, new_order
-        elif isinstance(node, bigframes.core.nodes.RowCountNode):
-            child_result = remove_order(node.child)
-            return node.replace_child(
-                child_result
-            ), bigframes.core.ordering.TotalOrdering.from_primary_key([node.col_id])
         elif isinstance(node, bigframes.core.nodes.AggregateNode):
             if node.has_ordered_ops:
                 child_result, child_order = pull_up_order_inner(node.child)