refactor: unify row operators to same interface

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -40,6 +40,7 @@
 import bigframes.dtypes
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
+import bigframes.operations.ibis_mappings as ops_ibis
 
 if typing.TYPE_CHECKING:
     from bigframes.session import Session
@@ -509,9 +510,10 @@ def project_unary_op(
         self, column_name: str, op: ops.UnaryOp, output_name=None
     ) -> ArrayValue:
         """Creates a new expression based on this expression with unary operation applied to one column."""
-        value = op._as_ibis(self._get_ibis_column(column_name)).name(
-            output_name or column_name
+        ibis_compiled = ops_ibis.ibis_compiler.compile_row_op(
+            op, (self._get_ibis_column(column_name),)
         )
+        value = ibis_compiled.name(output_name or column_name)
         return self._set_or_replace_by_id(output_name or column_name, value)
 
     def project_binary_op(
@@ -522,10 +524,12 @@ def project_binary_op(
         output_column_id: str,
     ) -> ArrayValue:
         """Creates a new expression based on this expression with binary operation applied to two columns."""
-        value = op(
-            self._get_ibis_column(left_column_id),
-            self._get_ibis_column(right_column_id),
-        ).name(output_column_id)
+        left_input = self._get_ibis_column(left_column_id)
+        right_input = self._get_ibis_column(right_column_id)
+        ibis_compiled = ops_ibis.ibis_compiler.compile_row_op(
+            op, (left_input, right_input)
+        )
+        value = ibis_compiled.name(output_column_id)
         return self._set_or_replace_by_id(output_column_id, value)
 
     def project_ternary_op(
@@ -537,11 +541,14 @@ def project_ternary_op(
         output_column_id: str,
     ) -> ArrayValue:
         """Creates a new expression based on this expression with ternary operation applied to three columns."""
-        value = op(
+        ibis_inputs = (
             self._get_ibis_column(col_id_1),
             self._get_ibis_column(col_id_2),
             self._get_ibis_column(col_id_3),
-        ).name(output_column_id)
+        )
+        value = ops_ibis.ibis_compiler.compile_row_op(op, ibis_inputs).name(
+            output_column_id
+        )
         return self._set_or_replace_by_id(output_column_id, value)
 
     def aggregate(
@@ -574,10 +581,11 @@ def aggregate(
             )
             columns = tuple(result[key] for key in result.columns)
             expr = ArrayValue(self._session, result, columns=columns, ordering=ordering)
+
             if dropna:
                 for column_id in by_column_ids:
                     expr = expr._filter(
-                        ops.notnull_op._as_ibis(expr._get_ibis_column(column_id))
+                        ops_ibis.notnull_op_impl(expr._get_ibis_column(column_id))
                     )
             # Can maybe remove this as Ordering id is redundant as by_column is unique after aggregation
             return expr._project_offsets()
@@ -1023,8 +1031,9 @@ def unpivot(
             null_value = bigframes.dtypes.literal_to_ibis_scalar(
                 None, force_dtype=col_dtype
             )
+            cast_op = ops.AsTypeOp(to_type=col_dtype)
             ibis_values = [
-                ops.AsTypeOp(col_dtype)._as_ibis(unpivot_table[col])
+                ops_ibis.astype_op_impl(unpivot_table[col], cast_op)
                 if col is not None
                 else null_value
                 for col in source_cols

bigframes/core/block_transforms.py

@@ -43,7 +43,7 @@ def equals(block1: blocks.Block, block2: blocks.Block) -> bool:
         lcolmapped = lmap[lcol]
         rcolmapped = rmap[rcol]
         joined_block, result_id = joined_block.apply_binary_op(
-            lcolmapped, rcolmapped, ops.eq_nulls_match_op
+            lcolmapped, rcolmapped, ops.eq_null_match_op
         )
         joined_block, result_id = joined_block.apply_unary_op(
             result_id, ops.partial_right(ops.fillna_op, False)
@@ -239,7 +239,7 @@ def rank(
     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())
+            rownum_col_ids, ops.AsTypeOp(to_type=pd.Float64Dtype())
         )
     if na_option == "keep":
         # For na_option "keep", null inputs must produce null outputs

bigframes/core/blocks.py

@@ -554,12 +554,12 @@ def _split(
         # Create an ordering col and convert to string
         block, ordering_col = block.promote_offsets()
         block, string_ordering_col = block.apply_unary_op(
-            ordering_col, ops.AsTypeOp("string[pyarrow]")
+            ordering_col, ops.AsTypeOp(to_type="string[pyarrow]")
         )
 
         # Apply hash method to sum col and order by it.
         block, string_sum_col = block.apply_binary_op(
-            string_ordering_col, random_state_col, ops.concat_op
+            string_ordering_col, random_state_col, ops.strconcat_op
         )
         block, hash_string_sum_col = block.apply_unary_op(string_sum_col, ops.hash_op)
         block = block.order_by([ordering.OrderingColumnReference(hash_string_sum_col)])
@@ -1223,8 +1223,8 @@ def add_prefix(self, prefix: str, axis: str | int | None = None) -> Block:
         if axis_number == 0:
             expr = self._expr
             for index_col in self._index_columns:
-                expr = expr.project_unary_op(index_col, ops.AsTypeOp("string"))
-                prefix_op = ops.BinopPartialLeft(ops.add_op, prefix)
+                expr = expr.project_unary_op(index_col, ops.AsTypeOp(to_type="string"))
+                prefix_op = ops.ApplyLeft(base_op=ops.add_op, left_scalar=prefix)
                 expr = expr.project_unary_op(index_col, prefix_op)
             return Block(
                 expr,
@@ -1242,8 +1242,8 @@ def add_suffix(self, suffix: str, axis: str | int | None = None) -> Block:
         if axis_number == 0:
             expr = self._expr
             for index_col in self._index_columns:
-                expr = expr.project_unary_op(index_col, ops.AsTypeOp("string"))
-                prefix_op = ops.BinopPartialRight(ops.add_op, suffix)
+                expr = expr.project_unary_op(index_col, ops.AsTypeOp(to_type="string"))
+                prefix_op = ops.ApplyRight(base_op=ops.add_op, right_scalar=suffix)
                 expr = expr.project_unary_op(index_col, prefix_op)
             return Block(
                 expr,

bigframes/core/indexes/index.py

@@ -179,7 +179,7 @@ def astype(
     ) -> Index:
         if self.nlevels > 1:
             raise TypeError("Multiindex does not support 'astype'")
-        return self._apply_unary_op(ops.AsTypeOp(dtype))
+        return self._apply_unary_op(ops.AsTypeOp(to_type=dtype))
 
     def all(self) -> bool:
         if self.nlevels > 1:
@@ -271,7 +271,7 @@ def drop(
         level_id = self._block.index_columns[0]
         if utils.is_list_like(labels):
             block, inverse_condition_id = block.apply_unary_op(
-                level_id, ops.IsInOp(labels, match_nulls=True)
+                level_id, ops.IsInOp(values=tuple(labels), match_nulls=True)
             )
             block, condition_id = block.apply_unary_op(
                 inverse_condition_id, ops.invert_op
@@ -301,9 +301,9 @@ def isin(self, values) -> Index:
                 f"isin(), you passed a [{type(values).__name__}]"
             )
 
-        return self._apply_unary_op(ops.IsInOp(values, match_nulls=True)).fillna(
-            value=False
-        )
+        return self._apply_unary_op(
+            ops.IsInOp(values=tuple(values), match_nulls=True)
+        ).fillna(value=False)
 
     def _apply_unary_op(
         self,

bigframes/dataframe.py

@@ -148,7 +148,7 @@ def __init__(
                 block = block.select_columns(list(columns))  # type:ignore
             if dtype:
                 block = block.multi_apply_unary_op(
-                    block.value_columns, ops.AsTypeOp(dtype)
+                    block.value_columns, ops.AsTypeOp(to_type=dtype)
                 )
             self._block = block
 
@@ -309,7 +309,7 @@ def astype(
         self,
         dtype: Union[bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype],
     ) -> DataFrame:
-        return self._apply_unary_op(ops.AsTypeOp(dtype))
+        return self._apply_unary_op(ops.AsTypeOp(to_type=dtype))
 
     def _to_sql_query(
         self, include_index: bool
@@ -536,7 +536,7 @@ def _apply_binop(
 
     def _apply_scalar_binop(self, other: float | int, op: ops.BinaryOp) -> DataFrame:
         block = self._block
-        partial_op = ops.BinopPartialRight(op, other)
+        partial_op = ops.ApplyRight(base_op=op, right_scalar=other)
         for column_id, label in zip(
             self._block.value_columns, self._block.column_labels
         ):
@@ -951,7 +951,7 @@ def drop(
 
             if utils.is_list_like(index):
                 block, inverse_condition_id = block.apply_unary_op(
-                    level_id, ops.IsInOp(index, match_nulls=True)
+                    level_id, ops.IsInOp(values=tuple(index), match_nulls=True)
                 )
                 block, condition_id = block.apply_unary_op(
                     inverse_condition_id, ops.invert_op
@@ -1318,16 +1318,16 @@ def _filter_rows(
             block = self._block
             block, label_string_id = block.apply_unary_op(
                 self._block.index_columns[0],
-                ops.AsTypeOp(pandas.StringDtype(storage="pyarrow")),
+                ops.AsTypeOp(to_type=pandas.StringDtype(storage="pyarrow")),
             )
             if like is not None:
                 block, mask_id = block.apply_unary_op(
-                    label_string_id, ops.ContainsStringOp(pat=like)
+                    label_string_id, ops.StrContainsOp(pat=like)
                 )
             else:  # regex
                 assert regex is not None
                 block, mask_id = block.apply_unary_op(
-                    label_string_id, ops.ContainsRegexOp(pat=regex)
+                    label_string_id, ops.StrContainsRegexOp(pat=regex)
                 )
 
             block = block.filter(mask_id)
@@ -1337,7 +1337,7 @@ def _filter_rows(
             # Behavior matches pandas 2.1+, older pandas versions would reindex
             block = self._block
             block, mask_id = block.apply_unary_op(
-                self._block.index_columns[0], ops.IsInOp(values=list(items))
+                self._block.index_columns[0], ops.IsInOp(values=tuple(items))
             )
             block = block.filter(mask_id)
             block = block.select_columns(self._block.value_columns)
@@ -1467,7 +1467,9 @@ def isin(self, values) -> DataFrame:
                 if label in values.keys():
                     value_for_key = values[label]
                     block, result_id = block.apply_unary_op(
-                        col, ops.IsInOp(value_for_key, match_nulls=True), label
+                        col,
+                        ops.IsInOp(values=tuple(value_for_key), match_nulls=True),
+                        label,
                     )
                     result_ids.append(result_id)
                 else:
@@ -1477,9 +1479,9 @@ def isin(self, values) -> DataFrame:
                     result_ids.append(result_id)
             return DataFrame(block.select_columns(result_ids)).fillna(value=False)
         elif utils.is_list_like(values):
-            return self._apply_unary_op(ops.IsInOp(values, match_nulls=True)).fillna(
-                value=False
-            )
+            return self._apply_unary_op(
+                ops.IsInOp(values=tuple(values), match_nulls=True)
+            ).fillna(value=False)
         else:
             raise TypeError(
                 "only list-like objects are allowed to be passed to "
@@ -2478,7 +2480,7 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame:
         # inputs causing errors.
         reprojected_df = DataFrame(self._block._force_reproject())
         return reprojected_df._apply_unary_op(
-            ops.RemoteFunctionOp(func, apply_on_null=(na_action is None))
+            ops.RemoteFunctionOp(func=func, apply_on_null=(na_action is None))
         )
 
     def apply(self, func, *, args: typing.Tuple = (), **kwargs):