refactor: move operator definitions from __init__.py to separate files (#1290)

* refactor: move operator definitions from __init__.py to separate files

* import geo operators

* use Union instead of |

* update more Union types

* expose CaseWhenOp

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* fix lint error

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: sycai

bigframes/operations/__init__.py

@@ -0,0 +1,65 @@
+# 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.
+
+import dataclasses
+import typing
+
+from bigframes import dtypes
+from bigframes.operations import base_ops
+
+
+@dataclasses.dataclass(frozen=True)
+class ArrayToStringOp(base_ops.UnaryOp):
+    name: typing.ClassVar[str] = "array_to_string"
+    delimiter: str
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if not dtypes.is_array_string_like(input_type):
+            raise TypeError("Input type must be an array of string type.")
+        return dtypes.STRING_DTYPE
+
+
+@dataclasses.dataclass(frozen=True)
+class ArrayIndexOp(base_ops.UnaryOp):
+    name: typing.ClassVar[str] = "array_index"
+    index: int
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if dtypes.is_string_like(input_type):
+            return dtypes.STRING_DTYPE
+        elif dtypes.is_array_like(input_type):
+            return dtypes.arrow_dtype_to_bigframes_dtype(
+                input_type.pyarrow_dtype.value_type
+            )
+        else:
+            raise TypeError("Input type must be an array or string-like type.")
+
+
+@dataclasses.dataclass(frozen=True)
+class ArraySliceOp(base_ops.UnaryOp):
+    name: typing.ClassVar[str] = "array_slice"
+    start: int
+    stop: typing.Optional[int] = None
+    step: typing.Optional[int] = None
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if dtypes.is_string_like(input_type):
+            return dtypes.STRING_DTYPE
+        elif dtypes.is_array_like(input_type):
+            return input_type
+        else:
+            raise TypeError("Input type must be an array or string-like type.")

bigframes/operations/array_ops.py

@@ -0,0 +1,197 @@
+# 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 dataclasses
+import typing
+
+from bigframes import dtypes
+import bigframes.operations.type as op_typing
+
+if typing.TYPE_CHECKING:
+    # Avoids circular dependency
+    import bigframes.core.expression
+
+
+class RowOp(typing.Protocol):
+    @property
+    def name(self) -> str:
+        ...
+
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        ...
+
+    @property
+    def is_monotonic(self) -> bool:
+        """Whether the row operation preserves total ordering. Can be pruned from ordering expressions."""
+        ...
+
+    @property
+    def is_bijective(self) -> bool:
+        """Whether the operation has a 1:1 mapping between inputs and outputs"""
+        ...
+
+    @property
+    def deterministic(self) -> bool:
+        """Whether the operation is deterministic" (given deterministic inputs)"""
+        ...
+
+
+@dataclasses.dataclass(frozen=True)
+class ScalarOp:
+    @property
+    def name(self) -> str:
+        raise NotImplementedError("RowOp abstract base class has no implementation")
+
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        raise NotImplementedError("Abstract operation has no output type")
+
+    @property
+    def is_monotonic(self) -> bool:
+        """Whether the row operation preserves total ordering. Can be pruned from ordering expressions."""
+        return False
+
+    @property
+    def is_bijective(self) -> bool:
+        """Whether the operation has a 1:1 mapping between inputs and outputs"""
+        return False
+
+    @property
+    def deterministic(self) -> bool:
+        """Whether the operation is deterministic" (given deterministic inputs)"""
+        return True
+
+
+@dataclasses.dataclass(frozen=True)
+class NaryOp(ScalarOp):
+    def as_expr(
+        self,
+        *exprs: typing.Union[str, bigframes.core.expression.Expression],
+    ) -> bigframes.core.expression.Expression:
+        import bigframes.core.expression
+
+        # Keep this in sync with output_type and compilers
+        inputs: list[bigframes.core.expression.Expression] = []
+
+        for expr in exprs:
+            inputs.append(_convert_expr_input(expr))
+
+        return bigframes.core.expression.OpExpression(
+            self,
+            tuple(inputs),
+        )
+
+
+# These classes can be used to create simple ops that don't take local parameters
+# All is needed is a unique name, and to register an implementation in ibis_mappings.py
+@dataclasses.dataclass(frozen=True)
+class UnaryOp(ScalarOp):
+    @property
+    def arguments(self) -> int:
+        return 1
+
+    def as_expr(
+        self, input_id: typing.Union[str, bigframes.core.expression.Expression] = "arg"
+    ) -> bigframes.core.expression.Expression:
+        import bigframes.core.expression
+
+        return bigframes.core.expression.OpExpression(
+            self, (_convert_expr_input(input_id),)
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BinaryOp(ScalarOp):
+    @property
+    def arguments(self) -> int:
+        return 2
+
+    def as_expr(
+        self,
+        left_input: typing.Union[str, bigframes.core.expression.Expression] = "arg1",
+        right_input: typing.Union[str, bigframes.core.expression.Expression] = "arg2",
+    ) -> bigframes.core.expression.Expression:
+        import bigframes.core.expression
+
+        return bigframes.core.expression.OpExpression(
+            self,
+            (
+                _convert_expr_input(left_input),
+                _convert_expr_input(right_input),
+            ),
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class TernaryOp(ScalarOp):
+    @property
+    def arguments(self) -> int:
+        return 3
+
+    def as_expr(
+        self,
+        input1: typing.Union[str, bigframes.core.expression.Expression] = "arg1",
+        input2: typing.Union[str, bigframes.core.expression.Expression] = "arg2",
+        input3: typing.Union[str, bigframes.core.expression.Expression] = "arg3",
+    ) -> bigframes.core.expression.Expression:
+        import bigframes.core.expression
+
+        return bigframes.core.expression.OpExpression(
+            self,
+            (
+                _convert_expr_input(input1),
+                _convert_expr_input(input2),
+                _convert_expr_input(input3),
+            ),
+        )
+
+
+def _convert_expr_input(
+    input: typing.Union[str, bigframes.core.expression.Expression]
+) -> bigframes.core.expression.Expression:
+    """Allows creating column references with just a string"""
+    import bigframes.core.expression
+
+    if isinstance(input, str):
+        return bigframes.core.expression.deref(input)
+    else:
+        return input
+
+
+# Operation Factories
+def create_unary_op(name: str, type_signature: op_typing.UnaryTypeSignature) -> UnaryOp:
+    return dataclasses.make_dataclass(
+        name,
+        [
+            ("name", typing.ClassVar[str], name),
+            ("output_type", typing.ClassVar[typing.Callable], type_signature.as_method),
+        ],
+        bases=(UnaryOp,),
+        frozen=True,
+    )()
+
+
+def create_binary_op(
+    name: str, type_signature: op_typing.BinaryTypeSignature
+) -> BinaryOp:
+    return dataclasses.make_dataclass(
+        name,
+        [
+            ("name", typing.ClassVar[str], name),
+            ("output_type", typing.ClassVar[typing.Callable], type_signature.as_method),
+        ],
+        bases=(BinaryOp,),
+        frozen=True,
+    )()

bigframes/operations/base_ops.py

@@ -0,0 +1,47 @@
+# 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.
+
+import dataclasses
+import typing
+
+from bigframes import dtypes
+from bigframes.operations import base_ops
+import bigframes.operations.type as op_typing
+
+obj_fetch_metadata_op = base_ops.create_unary_op(
+    name="obj_fetch_metadata", type_signature=op_typing.BLOB_TRANSFORM
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class ObjGetAccessUrl(base_ops.UnaryOp):
+    name: typing.ClassVar[str] = "obj_get_access_url"
+    mode: str  # access mode, e.g. R read, W write, RW read & write
+
+    def output_type(self, *input_types):
+        return dtypes.JSON_DTYPE
+
+
+@dataclasses.dataclass(frozen=True)
+class ObjMakeRef(base_ops.BinaryOp):
+    name: typing.ClassVar[str] = "obj.make_ref"
+
+    def output_type(self, *input_types):
+        if not all(map(dtypes.is_string_like, input_types)):
+            raise TypeError("obj.make_ref requires string-like arguments")
+
+        return dtypes.OBJ_REF_DTYPE
+
+
+obj_make_ref_op = ObjMakeRef()

bigframes/operations/blob_ops.py

@@ -0,0 +1,23 @@
+# 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 bigframes.operations import base_ops
+import bigframes.operations.type as op_typing
+
+and_op = base_ops.create_binary_op(name="and", type_signature=op_typing.LOGICAL)
+
+or_op = base_ops.create_binary_op(name="or", type_signature=op_typing.LOGICAL)
+
+xor_op = base_ops.create_binary_op(name="xor", type_signature=op_typing.LOGICAL)

bigframes/operations/bool_ops.py

@@ -0,0 +1,33 @@
+# 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 bigframes.operations import base_ops
+import bigframes.operations.type as op_typing
+
+eq_op = base_ops.create_binary_op(name="eq", type_signature=op_typing.COMPARISON)
+
+eq_null_match_op = base_ops.create_binary_op(
+    name="eq_nulls_match", type_signature=op_typing.COMPARISON
+)
+
+ne_op = base_ops.create_binary_op(name="ne", type_signature=op_typing.COMPARISON)
+
+lt_op = base_ops.create_binary_op(name="lt", type_signature=op_typing.COMPARISON)
+
+gt_op = base_ops.create_binary_op(name="gt", type_signature=op_typing.COMPARISON)
+
+le_op = base_ops.create_binary_op(name="le", type_signature=op_typing.COMPARISON)
+
+ge_op = base_ops.create_binary_op(name="ge", type_signature=op_typing.COMPARISON)