Skip to content

Commit 30e1eb6

Browse files
Refactor IsNullOp and NotNullOp logic
This change consolidates the definition and compilation logic for IsNullOp, isnull_op, NotNullOp, and notnull_op into a new, dedicated file: `bigframes/operations/isnull_op.py`. Key changes include: - Moved operator definitions from `generic_ops.py` to `isnull_op.py`. - Moved Ibis scalar compilation logic from `scalar_op_compiler.py` to `isnull_op.py`. - Moved Polars expression compilation logic from `polars/compiler.py` to `isnull_op.py`. - Updated main compilers (`ScalarOpCompiler` and `PolarsExpressionCompiler`) to directly import and register the compilation functions from `isnull_op.py`. - Ensured all internal references and naming conventions (`IsNullOp`, `isnull_op`, `NotNullOp`, `notnull_op`) are consistent with the refactored structure. NOTE: I was unable to perform test validation (unit and system) due to missing project-specific dependencies, primarily `bigframes_vendored` and `test_utils.prefixer`. The changes are provided based on the completion of the refactoring steps as you requested.
1 parent f984381 commit 30e1eb6

File tree

6 files changed

+101
-34
lines changed

6 files changed

+101
-34
lines changed

bigframes/core/compile/polars/compiler.py

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -231,14 +231,6 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
231231
else:
232232
return input.is_in(op.values) or input.is_null()
233233

234-
@compile_op.register(gen_ops.IsNullOp)
235-
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
236-
return input.is_null()
237-
238-
@compile_op.register(gen_ops.NotNullOp)
239-
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
240-
return input.is_not_null()
241-
242234
@compile_op.register(gen_ops.FillNaOp)
243235
@compile_op.register(gen_ops.CoalesceOp)
244236
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:
268260
# eg. We want "True" instead of "true" for bool to strin
269261
return input.cast(_DTYPE_MAPPING[op.to_type], strict=not op.safe)
270262

263+
# Register ops from other modules
264+
from bigframes.operations import IsNullOp, NotNullOp
265+
from bigframes.operations.isnull_op import (
266+
_polars_isnull_op_impl,
267+
_polars_notnull_op_impl,
268+
)
269+
270+
PolarsExpressionCompiler.compile_op.register(IsNullOp, _polars_isnull_op_impl)
271+
PolarsExpressionCompiler.compile_op.register(NotNullOp, _polars_notnull_op_impl)
272+
271273
@dataclasses.dataclass(frozen=True)
272274
class PolarsAggregateCompiler:
273275
scalar_compiler = PolarsExpressionCompiler()

bigframes/core/compile/scalar_op_compiler.py

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -228,18 +228,23 @@ def _register(
228228
# Singleton compiler
229229
scalar_op_compiler = ScalarOpCompiler()
230230

231+
# Registrations for operations defined in isnull_op.py
232+
from bigframes.operations import isnull_op, notnull_op
233+
from bigframes.operations.isnull_op import (
234+
_ibis_isnull_op_impl,
235+
_ibis_notnull_op_impl,
236+
)
231237

232-
### Unary Ops
233-
@scalar_op_compiler.register_unary_op(ops.isnull_op)
234-
def isnull_op_impl(x: ibis_types.Value):
235-
return x.isnull()
236-
238+
@scalar_op_compiler.register_unary_op(isnull_op)
239+
def _scalar_isnull_op_impl_wrapper(x: ibis_types.Value):
240+
return _ibis_isnull_op_impl(x)
237241

238-
@scalar_op_compiler.register_unary_op(ops.notnull_op)
239-
def notnull_op_impl(x: ibis_types.Value):
240-
return x.notnull()
242+
@scalar_op_compiler.register_unary_op(notnull_op)
243+
def _scalar_notnull_op_impl_wrapper(x: ibis_types.Value):
244+
return _ibis_notnull_op_impl(x)
241245

242246

247+
### Unary Ops
243248
@scalar_op_compiler.register_unary_op(ops.hash_op)
244249
def hash_op_impl(x: ibis_types.Value):
245250
return typing.cast(ibis_types.IntegerValue, x).hash()

bigframes/dataframe.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2521,6 +2521,7 @@ def _filter_rows(
25212521
elif items is not None:
25222522
# Behavior matches pandas 2.1+, older pandas versions would reindex
25232523
block = self._block
2524+
block = self._block
25242525
block, mask_id = block.apply_unary_op(
25252526
self._block.index_columns[0], ops.IsInOp(values=tuple(items))
25262527
)

bigframes/operations/__init__.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,11 +81,9 @@
8181
hash_op,
8282
invert_op,
8383
IsInOp,
84-
isnull_op,
8584
MapOp,
8685
maximum_op,
8786
minimum_op,
88-
notnull_op,
8987
RowKey,
9088
SqlScalarOp,
9189
where_op,
@@ -104,6 +102,7 @@
104102
GeoStDistanceOp,
105103
GeoStLengthOp,
106104
)
105+
from bigframes.operations.isnull_op import isnull_op, notnull_op
107106
from bigframes.operations.json_ops import (
108107
JSONExtract,
109108
JSONExtractArray,

bigframes/operations/generic_ops.py

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -29,22 +29,6 @@
2929
)
3030
invert_op = InvertOp()
3131

32-
IsNullOp = base_ops.create_unary_op(
33-
name="isnull",
34-
type_signature=op_typing.FixedOutputType(
35-
lambda x: True, dtypes.BOOL_DTYPE, description="nullable"
36-
),
37-
)
38-
isnull_op = IsNullOp()
39-
40-
NotNullOp = base_ops.create_unary_op(
41-
name="notnull",
42-
type_signature=op_typing.FixedOutputType(
43-
lambda x: True, dtypes.BOOL_DTYPE, description="nullable"
44-
),
45-
)
46-
notnull_op = NotNullOp()
47-
4832
HashOp = base_ops.create_unary_op(
4933
name="hash",
5034
type_signature=op_typing.FixedOutputType(

bigframes/operations/isnull_op.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
# Copyright 2023 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import typing
18+
19+
# Direct imports from bigframes
20+
from bigframes import dtypes
21+
from bigframes.operations import base_ops
22+
import bigframes.operations.type as op_typing
23+
24+
# Imports for Ibis compilation
25+
from bigframes_vendored.ibis.expr import types as ibis_types
26+
27+
# Imports for Polars compilation
28+
try:
29+
import polars as pl
30+
except ImportError:
31+
# Polars is optional, error will be raised elsewhere if user tries to use it.
32+
pass
33+
34+
35+
# Definitions of IsNullOp and NotNullOp operations
36+
IsNullOp = base_ops.create_unary_op(
37+
name="isnull",
38+
type_signature=op_typing.FixedOutputType(
39+
lambda x: True, dtypes.BOOL_DTYPE, description="nullable"
40+
),
41+
)
42+
isnull_op = IsNullOp()
43+
44+
NotNullOp = base_ops.create_unary_op(
45+
name="notnull",
46+
type_signature=op_typing.FixedOutputType(
47+
lambda x: True, dtypes.BOOL_DTYPE, description="nullable"
48+
),
49+
)
50+
notnull_op = NotNullOp()
51+
52+
# Ibis Scalar Op Implementations
53+
def _ibis_isnull_op_impl(x: ibis_types.Value):
54+
return x.isnull()
55+
56+
def _ibis_notnull_op_impl(x: ibis_types.Value):
57+
return x.notnull()
58+
59+
60+
# Polars Expression Implementations
61+
def _polars_isnull_op_impl(op: IsNullOp, input: pl.Expr) -> pl.Expr:
62+
return input.is_null()
63+
64+
def _polars_notnull_op_impl(op: NotNullOp, input: pl.Expr) -> pl.Expr:
65+
return input.is_not_null()
66+
67+
__all__ = [
68+
"IsNullOp",
69+
"isnull_op",
70+
"NotNullOp",
71+
"notnull_op",
72+
"_ibis_isnull_op_impl",
73+
"_ibis_notnull_op_impl",
74+
"_polars_isnull_op_impl",
75+
"_polars_notnull_op_impl",
76+
]

0 commit comments

Comments
 (0)