refactor: Simplify expression generation for some block ops (#1298)

Author: TrevorBergeron

bigframes/core/block_transforms.py

@@ -43,19 +43,16 @@ def equals(block1: blocks.Block, block2: blocks.Block) -> bool:
 
     joined_block, (lmap, rmap) = block1.join(block2, how="outer")
 
-    equality_ids = []
+    exprs = []
     for lcol, rcol in zip(block1.value_columns, block2.value_columns):
-        lcolmapped = lmap[lcol]
-        rcolmapped = rmap[rcol]
-        joined_block, result_id = joined_block.project_expr(
+        exprs.append(
             ops.fillna_op.as_expr(
-                ops.eq_null_match_op.as_expr(lcolmapped, rcolmapped), ex.const(False)
+                ops.eq_null_match_op.as_expr(lmap[lcol], rmap[rcol]), ex.const(False)
             )
         )
-        equality_ids.append(result_id)
 
-    joined_block = joined_block.select_columns(equality_ids).with_column_labels(
-        list(range(len(equality_ids)))
+    joined_block = joined_block.project_exprs(
+        exprs, labels=list(range(len(exprs))), drop=True
     )
     stacked_block = joined_block.stack()
     result = stacked_block.get_stat(stacked_block.value_columns[0], agg_ops.all_op)
@@ -395,12 +392,12 @@ def pct_change(block: blocks.Block, periods: int = 1) -> blocks.Block:
     block, shift_columns = block.multi_apply_window_op(
         original_columns, agg_ops.ShiftOp(periods), window_spec=window_spec
     )
-    result_ids = []
+    exprs = []
     for original_col, shifted_col in zip(original_columns, shift_columns):
-        block, change_id = block.apply_binary_op(original_col, shifted_col, ops.sub_op)
-        block, pct_change_id = block.apply_binary_op(change_id, shifted_col, ops.div_op)
-        result_ids.append(pct_change_id)
-    return block.select_columns(result_ids).with_column_labels(column_labels)
+        change_expr = ops.sub_op.as_expr(original_col, shifted_col)
+        pct_change_expr = ops.div_op.as_expr(change_expr, shifted_col)
+        exprs.append(pct_change_expr)
+    return block.project_exprs(exprs, labels=column_labels, drop=True)
 
 
 def rank(
@@ -470,16 +467,23 @@ def rank(
     # Step 3: post processing: mask null values and cast to float
     if method in ["min", "max", "first", "dense"]:
         # Pandas rank always produces Float64, so must cast for aggregation types that produce ints
-        block = block.multi_apply_unary_op(
-            rownum_col_ids, ops.AsTypeOp(pd.Float64Dtype())
+        return (
+            block.select_columns(rownum_col_ids)
+            .multi_apply_unary_op(ops.AsTypeOp(pd.Float64Dtype()))
+            .with_column_labels(labels)
         )
     if na_option == "keep":
         # For na_option "keep", null inputs must produce null outputs
+        exprs = []
         for i in range(len(columns)):
-            block, null_const = block.create_constant(pd.NA, dtype=pd.Float64Dtype())
-            block, rownum_col_ids[i] = block.apply_ternary_op(
-                null_const, nullity_col_ids[i], rownum_col_ids[i], ops.where_op
+            exprs.append(
+                ops.where_op.as_expr(
+                    ex.const(pd.NA, dtype=pd.Float64Dtype()),
+                    nullity_col_ids[i],
+                    rownum_col_ids[i],
+                )
             )
+        return block.project_exprs(exprs, labels=labels, drop=True)
 
     return block.select_columns(rownum_col_ids).with_column_labels(labels)
 

bigframes/core/blocks.py

@@ -897,7 +897,6 @@ def multi_apply_window_op(
 
     def multi_apply_unary_op(
         self,
-        columns: typing.Sequence[str],
         op: Union[ops.UnaryOp, ex.Expression],
     ) -> Block:
         if isinstance(op, ops.UnaryOp):
@@ -911,27 +910,37 @@ def multi_apply_unary_op(
 
         block = self
 
-        result_ids = []
-        for col_id in columns:
-            label = self.col_id_to_label[col_id]
-            block, result_id = block.project_expr(
-                expr.bind_variables({input_varname: ex.deref(col_id)}),
-                label=label,
-            )
-            block = block.copy_values(result_id, col_id)
-            result_ids.append(result_id)
-        block = block.drop_columns(result_ids)
+        exprs = [
+            expr.bind_variables({input_varname: ex.deref(col_id)})
+            for col_id in self.value_columns
+        ]
+        block = self.project_exprs(exprs, labels=self.column_labels, drop=True)
+
         # Special case, we can preserve transpose cache for full-frame unary ops
-        if (self._transpose_cache is not None) and set(self.value_columns) == set(
-            columns
-        ):
-            transpose_columns = self._transpose_cache.value_columns
-            new_transpose_cache = self._transpose_cache.multi_apply_unary_op(
-                transpose_columns, op
-            )
+        if self._transpose_cache is not None:
+            new_transpose_cache = self._transpose_cache.multi_apply_unary_op(op)
             block = block.with_transpose_cache(new_transpose_cache)
         return block
 
+    def project_exprs(
+        self,
+        exprs: Sequence[ex.Expression],
+        labels: Union[Sequence[Label], pd.Index],
+        drop=False,
+    ) -> Block:
+        new_array, _ = self.expr.compute_values(exprs)
+        if drop:
+            new_array = new_array.drop_columns(self.value_columns)
+
+        return Block(
+            new_array,
+            index_columns=self.index_columns,
+            column_labels=labels
+            if drop
+            else self.column_labels.append(pd.Index(labels)),
+            index_labels=self._index_labels,
+        )
+
     def apply_window_op(
         self,
         column: str,
@@ -2279,18 +2288,15 @@ def _apply_binop(
         labels: pd.Index,
         reverse: bool = False,
     ) -> Block:
-        block = self
-        binop_result_ids = []
+        exprs = []
         for left_input, right_input in inputs:
-            expr = (
+            exprs.append(
                 op.as_expr(right_input, left_input)
                 if reverse
                 else op.as_expr(left_input, right_input)
             )
-            block, result_col_id = block.project_expr(expr)
-            binop_result_ids.append(result_col_id)
 
-        return block.select_columns(binop_result_ids).with_column_labels(labels)
+        return self.project_exprs(exprs, labels=labels, drop=True)
 
     def join(
         self,

bigframes/dataframe.py

@@ -179,9 +179,7 @@ def __init__(
             if columns:
                 block = block.select_columns(list(columns))  # type:ignore
             if dtype:
-                block = block.multi_apply_unary_op(
-                    block.value_columns, ops.AsTypeOp(to_type=dtype)
-                )
+                block = block.multi_apply_unary_op(ops.AsTypeOp(to_type=dtype))
             self._block = block
 
         else:
@@ -845,9 +843,7 @@ def _apply_scalar_binop(
                 left_input=ex.free_var("var1"),
                 right_input=ex.const(other),
             )
-        return DataFrame(
-            self._block.multi_apply_unary_op(self._block.value_columns, expr)
-        )
+        return DataFrame(self._block.multi_apply_unary_op(expr))
 
     def _apply_series_binop_axis_0(
         self,
@@ -2400,9 +2396,7 @@ def dropna(
                 result = result.reset_index()
             return DataFrame(result)
         else:
-            isnull_block = self._block.multi_apply_unary_op(
-                self._block.value_columns, ops.isnull_op
-            )
+            isnull_block = self._block.multi_apply_unary_op(ops.isnull_op)
             if how == "any":
                 null_locations = DataFrame(isnull_block).any().to_pandas()
             else:  # 'all'
@@ -3828,7 +3822,7 @@ def to_orc(self, path=None, **kwargs) -> bytes | None:
         return as_pandas_default_index.to_orc(path, **kwargs)
 
     def _apply_unary_op(self, operation: ops.UnaryOp) -> DataFrame:
-        block = self._block.multi_apply_unary_op(self._block.value_columns, operation)
+        block = self._block.multi_apply_unary_op(operation)
         return DataFrame(block)
 
     def _map_clustering_columns(