chore: refactor IsNullOp and NotNullOp logic to make scalar ops generation easier

bigframes/core/compile/polars/compiler.py

@@ -231,14 +231,6 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             else:
                 return input.is_in(op.values) or input.is_null()
 
-        @compile_op.register(gen_ops.IsNullOp)
-        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
-            return input.is_null()
-
-        @compile_op.register(gen_ops.NotNullOp)
-        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
-            return input.is_not_null()
-
         @compile_op.register(gen_ops.FillNaOp)
         @compile_op.register(gen_ops.CoalesceOp)
         def _(self, op: ops.ScalarOp, l_input: pl.Expr, r_input: pl.Expr) -> pl.Expr:
@@ -268,6 +260,16 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             # eg. We want "True" instead of "true" for bool to strin
             return input.cast(_DTYPE_MAPPING[op.to_type], strict=not op.safe)
 
+    # Register ops from other modules
+    from bigframes.operations import IsNullOp, NotNullOp
+    from bigframes.operations.isnull_op import (
+        _polars_isnull_op_impl,
+        _polars_notnull_op_impl,
+    )
+
+    PolarsExpressionCompiler.compile_op.register(IsNullOp, _polars_isnull_op_impl)
+    PolarsExpressionCompiler.compile_op.register(NotNullOp, _polars_notnull_op_impl)
+
     @dataclasses.dataclass(frozen=True)
     class PolarsAggregateCompiler:
         scalar_compiler = PolarsExpressionCompiler()

bigframes/core/compile/scalar_op_compiler.py

@@ -228,18 +228,23 @@ def _register(
 # Singleton compiler
 scalar_op_compiler = ScalarOpCompiler()
 
+# Registrations for operations defined in isnull_op.py
+from bigframes.operations import isnull_op, notnull_op
+from bigframes.operations.isnull_op import (
+    _ibis_isnull_op_impl,
+    _ibis_notnull_op_impl,
+)
 
-### Unary Ops
-@scalar_op_compiler.register_unary_op(ops.isnull_op)
-def isnull_op_impl(x: ibis_types.Value):
-    return x.isnull()
-
+@scalar_op_compiler.register_unary_op(isnull_op)
+def _scalar_isnull_op_impl_wrapper(x: ibis_types.Value):
+    return _ibis_isnull_op_impl(x)
 
-@scalar_op_compiler.register_unary_op(ops.notnull_op)
-def notnull_op_impl(x: ibis_types.Value):
-    return x.notnull()
+@scalar_op_compiler.register_unary_op(notnull_op)
+def _scalar_notnull_op_impl_wrapper(x: ibis_types.Value):
+    return _ibis_notnull_op_impl(x)
 
 
+### Unary Ops
 @scalar_op_compiler.register_unary_op(ops.hash_op)
 def hash_op_impl(x: ibis_types.Value):
     return typing.cast(ibis_types.IntegerValue, x).hash()

bigframes/dataframe.py

@@ -2521,6 +2521,7 @@ def _filter_rows(
         elif items is not None:
             # Behavior matches pandas 2.1+, older pandas versions would reindex
             block = self._block
+            block = self._block
             block, mask_id = block.apply_unary_op(
                 self._block.index_columns[0], ops.IsInOp(values=tuple(items))
             )

bigframes/operations/__init__.py

@@ -81,11 +81,9 @@
     hash_op,
     invert_op,
     IsInOp,
-    isnull_op,
     MapOp,
     maximum_op,
     minimum_op,
-    notnull_op,
     RowKey,
     SqlScalarOp,
     where_op,
@@ -104,6 +102,7 @@
     GeoStDistanceOp,
     GeoStLengthOp,
 )
+from bigframes.operations.isnull_op import isnull_op, notnull_op
 from bigframes.operations.json_ops import (
     JSONExtract,
     JSONExtractArray,

bigframes/operations/generic_ops.py

@@ -29,22 +29,6 @@
 )
 invert_op = InvertOp()
 
-IsNullOp = base_ops.create_unary_op(
-    name="isnull",
-    type_signature=op_typing.FixedOutputType(
-        lambda x: True, dtypes.BOOL_DTYPE, description="nullable"
-    ),
-)
-isnull_op = IsNullOp()
-
-NotNullOp = base_ops.create_unary_op(
-    name="notnull",
-    type_signature=op_typing.FixedOutputType(
-        lambda x: True, dtypes.BOOL_DTYPE, description="nullable"
-    ),
-)
-notnull_op = NotNullOp()
-
 HashOp = base_ops.create_unary_op(
     name="hash",
     type_signature=op_typing.FixedOutputType(

bigframes/operations/isnull_op.py

@@ -0,0 +1,76 @@
+# Copyright 2023 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 typing
+
+# Direct imports from bigframes
+from bigframes import dtypes
+from bigframes.operations import base_ops
+import bigframes.operations.type as op_typing
+
+# Imports for Ibis compilation
+from bigframes_vendored.ibis.expr import types as ibis_types
+
+# Imports for Polars compilation
+try:
+    import polars as pl
+except ImportError:
+    # Polars is optional, error will be raised elsewhere if user tries to use it.
+    pass
+
+
+# Definitions of IsNullOp and NotNullOp operations
+IsNullOp = base_ops.create_unary_op(
+    name="isnull",
+    type_signature=op_typing.FixedOutputType(
+        lambda x: True, dtypes.BOOL_DTYPE, description="nullable"
+    ),
+)
+isnull_op = IsNullOp()
+
+NotNullOp = base_ops.create_unary_op(
+    name="notnull",
+    type_signature=op_typing.FixedOutputType(
+        lambda x: True, dtypes.BOOL_DTYPE, description="nullable"
+    ),
+)
+notnull_op = NotNullOp()
+
+# Ibis Scalar Op Implementations
+def _ibis_isnull_op_impl(x: ibis_types.Value):
+    return x.isnull()
+
+def _ibis_notnull_op_impl(x: ibis_types.Value):
+    return x.notnull()
+
+
+# Polars Expression Implementations
+def _polars_isnull_op_impl(op: IsNullOp, input: pl.Expr) -> pl.Expr:
+    return input.is_null()
+
+def _polars_notnull_op_impl(op: NotNullOp, input: pl.Expr) -> pl.Expr:
+    return input.is_not_null()
+
+__all__ = [
+    "IsNullOp",
+    "isnull_op",
+    "NotNullOp",
+    "notnull_op",
+    "_ibis_isnull_op_impl",
+    "_ibis_notnull_op_impl",
+    "_polars_isnull_op_impl",
+    "_polars_notnull_op_impl",
+]