combine null ops into generic_ops files

Author: tswast

bigframes/core/compile/compiled.py

@@ -30,8 +30,8 @@
 import pyarrow as pa
 
 from bigframes.core import utils
-import bigframes.core.compile.aggregate_compiler as agg_compiler
 import bigframes.core.compile.googlesql
+import bigframes.core.compile.ibis_compiler.aggregate_compiler as agg_compiler
 import bigframes.core.compile.ibis_compiler.scalar_op_compiler as op_compilers
 import bigframes.core.compile.ibis_types
 import bigframes.core.expression as ex

bigframes/core/compile/ibis_compiler/__init__.py

@@ -20,6 +20,5 @@
 
 from __future__ import annotations
 
-import bigframes.core.compile.ibis_compiler.operations.generic_ops.isnull_op  # noqa: F401
-import bigframes.core.compile.ibis_compiler.operations.generic_ops.notnull_op  # noqa: F401
+import bigframes.core.compile.ibis_compiler.operations.generic_ops  # noqa: F401
 import bigframes.core.compile.ibis_compiler.scalar_op_registry  # noqa: F401

…rames/core/compile/aggregate_compiler.py …pile/ibis_compiler/aggregate_compiler.pybigframes/core/compile/aggregate_compiler.py renamed to bigframes/core/compile/ibis_compiler/aggregate_compiler.py bigframes/core/compile/aggregate_compiler.py renamed to bigframes/core/compile/ibis_compiler/aggregate_compiler.py

@@ -12,14 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""
+BigFrames -> Ibis compilation for the operations in bigframes.operations.generic_ops.
+
+Please keep implementations in sequential order by op name.
+"""
+
 from __future__ import annotations
 
 from bigframes_vendored.ibis.expr import types as ibis_types
 
 from bigframes.core.compile.ibis_compiler import scalar_op_compiler
-from bigframes.operations.generic_ops import isnull_op
+from bigframes.operations import generic_ops
+
+register_unary_op = scalar_op_compiler.scalar_op_compiler.register_unary_op
+
+
+@register_unary_op(generic_ops.notnull_op)
+def notnull_op_impl(x: ibis_types.Value):
+    return x.notnull()
 
 
-@scalar_op_compiler.scalar_op_compiler.register_unary_op(isnull_op.isnull_op)
-def _ibis_isnull_op_impl(x: ibis_types.Value):
+@register_unary_op(generic_ops.isnull_op)
+def isnull_op_impl(x: ibis_types.Value):
     return x.isnull()

…iler/operations/generic_ops/isnull_op.py …/ibis_compiler/operations/generic_ops.pybigframes/core/compile/ibis_compiler/operations/generic_ops/isnull_op.py renamed to bigframes/core/compile/ibis_compiler/operations/generic_ops.py bigframes/core/compile/ibis_compiler/operations/generic_ops/isnull_op.py renamed to bigframes/core/compile/ibis_compiler/operations/generic_ops.py

@@ -20,12 +20,15 @@
 
 import warnings
 
+# The ops imports appear first so that the implementations can be registered.
+# polars shouldn't be needed at import time, as register is a no-op if polars
+# isn't installed.
+import bigframes.core.compile.polars.operations.generic_ops  # noqa: F401
+
 try:
     import polars  # noqa
 
     from bigframes.core.compile.polars.compiler import PolarsCompiler
-    import bigframes.core.compile.polars.operations.generic_ops.isnull_op  # noqa: F401
-    import bigframes.core.compile.polars.operations.generic_ops.notnull_op  # noqa: F401
 
     __all__ = ["PolarsCompiler"]
 except Exception as exc:

bigframes/core/compile/ibis_compiler/operations/generic_ops/__init__.py

@@ -17,7 +17,7 @@
 import functools
 import itertools
 import operator
-from typing import cast, Literal, Optional, Sequence, Tuple, TYPE_CHECKING
+from typing import cast, Literal, Optional, Sequence, Tuple, Type, TYPE_CHECKING
 
 import pandas as pd
 
@@ -46,6 +46,27 @@
     except Exception:
         polars_installed = False
 
+
+def register_op(op: Type):
+    """Register a compilation from BigFrames to Ibis.
+
+    This decorator can be used, even if Polars is not installed.
+
+    Args:
+        op: The type of the operator the wrapped function compiles.
+    """
+
+    def decorator(func):
+        if polars_installed:
+            # Ignore the type because compile_op is a generic Callable, so
+            # register isn't available according to mypy.
+            return PolarsExpressionCompiler.compile_op.register(op)(func)  # type: ignore
+        else:
+            return func
+
+    return decorator
+
+
 if polars_installed:
     _DTYPE_MAPPING = {
         # Direct mappings

bigframes/core/compile/ibis_compiler/operations/generic_ops/notnull_op.py

@@ -12,28 +12,36 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""
+BigFrames -> Polars compilation for the operations in bigframes.operations.generic_ops.
+
+Please keep implementations in sequential order by op name.
+"""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
 import bigframes.core.compile.polars.compiler as polars_compiler
-from bigframes.operations.generic_ops import notnull_op
+from bigframes.operations import generic_ops
 
 if TYPE_CHECKING:
     import polars as pl
 
 
-def _polars_notnull_op_impl(
+@polars_compiler.register_op(generic_ops.NotNullOp)
+def notnull_op_impl(
     compiler: polars_compiler.PolarsExpressionCompiler,
-    op: notnull_op.NotNullOp,
+    op: generic_ops.NotNullOp,  # type: ignore
     input: pl.Expr,
 ) -> pl.Expr:
     return input.is_not_null()
 
 
-if hasattr(polars_compiler, "PolarsExpressionCompiler"):
-    # TODO(https://github.com/python/mypy/issues/13040): remove `type: ignore`
-    # when mypy can better handle singledispatch.
-    polars_compiler.PolarsExpressionCompiler.compile_op.register(  # type: ignore
-        notnull_op.NotNullOp, _polars_notnull_op_impl
-    )
+@polars_compiler.register_op(generic_ops.IsNullOp)
+def isnull_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: generic_ops.IsNullOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.is_null()