perf: Prevent inlining of remote ops (#1347)

Author: TrevorBergeron

bigframes/core/compile/compiled.py

@@ -110,12 +110,19 @@ def projection(
         expression_id_pairs: typing.Tuple[typing.Tuple[ex.Expression, str], ...],
     ) -> UnorderedIR:
         """Apply an expression to the ArrayValue and assign the output to a column."""
+        cannot_inline = any(expr.expensive for expr, _ in expression_id_pairs)
+
         bindings = {col: self._get_ibis_column(col) for col in self.column_ids}
         new_values = [
             op_compiler.compile_expression(expression, bindings).name(id)
             for expression, id in expression_id_pairs
         ]
-        return UnorderedIR(self._table, (*self._columns, *new_values))
+        result = UnorderedIR(self._table, (*self._columns, *new_values))
+        if cannot_inline:
+            return result._reproject_to_table()
+        else:
+            # Cheap ops can defer "SELECT" and inline into later ops
+            return result
 
     def selection(
         self,
@@ -174,13 +181,12 @@ def _to_ibis_expr(
         Returns:
             An ibis expression representing the data help by the ArrayValue object.
         """
-        columns = list(self._columns)
         # Special case for empty tables, since we can't create an empty
         # projection.
-        if not columns:
+        if not self._columns:
             return bigframes_vendored.ibis.memtable([])
 
-        table = self._table.select(columns)
+        table = self._table.select(self._columns)
         if fraction is not None:
             table = table.filter(
                 bigframes_vendored.ibis.random() < ibis_types.literal(fraction)

bigframes/core/expression.py

@@ -18,7 +18,7 @@
 import dataclasses
 import itertools
 import typing
-from typing import Mapping, TypeVar, Union
+from typing import Generator, Mapping, TypeVar, Union
 
 import pandas as pd
 
@@ -155,6 +155,16 @@ class Expression(abc.ABC):
     def free_variables(self) -> typing.Tuple[str, ...]:
         return ()
 
+    @property
+    def children(self) -> typing.Tuple[Expression, ...]:
+        return ()
+
+    @property
+    def expensive(self) -> bool:
+        return any(
+            isinstance(ex, OpExpression) and ex.op.expensive for ex in self.walk()
+        )
+
     @property
     @abc.abstractmethod
     def column_references(self) -> typing.Tuple[ids.ColumnId, ...]:
@@ -216,6 +226,11 @@ def is_identity(self) -> bool:
         """True for identity operation that does not transform input."""
         return False
 
+    def walk(self) -> Generator[Expression, None, None]:
+        yield self
+        for child in self.children:
+            yield from child.children
+
 
 @dataclasses.dataclass(frozen=True)
 class ScalarConstantExpression(Expression):
@@ -389,6 +404,10 @@ def free_variables(self) -> typing.Tuple[str, ...]:
     def is_const(self) -> bool:
         return all(child.is_const for child in self.inputs)
 
+    @property
+    def children(self):
+        return self.inputs
+
     def output_type(
         self, input_types: dict[ids.ColumnId, dtypes.ExpressionType]
     ) -> dtypes.ExpressionType:

bigframes/operations/base_ops.py

@@ -48,6 +48,11 @@ def deterministic(self) -> bool:
         """Whether the operation is deterministic" (given deterministic inputs)"""
         ...
 
+    @property
+    def expensive(self) -> bool:
+        """Whether the operation is expensive to calculate. Such ops shouldn't be inlined if referenced multiple places."""
+        ...
+
 
 @dataclasses.dataclass(frozen=True)
 class ScalarOp:
@@ -73,6 +78,10 @@ def deterministic(self) -> bool:
         """Whether the operation is deterministic" (given deterministic inputs)"""
         return True
 
+    @property
+    def expensive(self) -> bool:
+        return False
+
 
 @dataclasses.dataclass(frozen=True)
 class NaryOp(ScalarOp):

bigframes/operations/remote_function_ops.py

@@ -25,6 +25,10 @@ class RemoteFunctionOp(base_ops.UnaryOp):
     func: typing.Callable
     apply_on_null: bool
 
+    @property
+    def expensive(self) -> bool:
+        return True
+
     def output_type(self, *input_types):
         # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
         if hasattr(self.func, "output_dtype"):
@@ -45,6 +49,10 @@ class BinaryRemoteFunctionOp(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "binary_remote_function"
     func: typing.Callable
 
+    @property
+    def expensive(self) -> bool:
+        return True
+
     def output_type(self, *input_types):
         # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
         if hasattr(self.func, "output_dtype"):
@@ -65,6 +73,10 @@ class NaryRemoteFunctionOp(base_ops.NaryOp):
     name: typing.ClassVar[str] = "nary_remote_function"
     func: typing.Callable
 
+    @property
+    def expensive(self) -> bool:
+        return True
+
     def output_type(self, *input_types):
         # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
         if hasattr(self.func, "output_dtype"):

third_party/bigframes_vendored/ibis/backends/sql/rewrites.py

@@ -245,6 +245,8 @@ def merge_select_select(_, **kwargs):
         ops.InSubquery,
         ops.Unnest,
         ops.Impure,
+        # This is used for remote functions, which we don't want to copy
+        ops.ScalarUDF,
     )
     if _.find_below(blocking, filter=ops.Value):
         return _