From 1388ae718b902353bf3d93dd5c3102b4489925fc Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 22 Oct 2025 18:17:44 +0000 Subject: [PATCH 01/10] feat: Implement isnan and isfinite in bigframes This commit implements the `isnan` and `isfinite` numpy ufuncs in bigframes. The following changes were made: - Added `IsNanOp` and `IsFiniteOp` to `bigframes/operations/numeric_ops.py` - Mapped `np.isnan` and `np.isfinite` to the new ops in `bigframes/operations/numpy_op_maps.py` - Added compilation logic for the new ops to the ibis, polars, and sqlglot compilers - Added tests for the new ops in `tests/system/small/test_numpy.py` --- .../compile/ibis_compiler/scalar_op_compiler.py | 11 +++++++++++ bigframes/core/compile/polars/compiler.py | 8 ++++++++ .../core/compile/sqlglot/scalar_compiler.py | 11 +++++++++++ bigframes/operations/numeric_ops.py | 16 ++++++++++++++++ bigframes/operations/numpy_op_maps.py | 2 ++ tests/system/small/test_numpy.py | 2 ++ 6 files changed, 50 insertions(+) diff --git a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py index 1197f6b9da..7f74a1b7d9 100644 --- a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py +++ b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py @@ -26,6 +26,7 @@ from bigframes.core import agg_expressions, ordering import bigframes.core.compile.ibis_types import bigframes.core.expression as ex +from bigframes.operations import numeric_ops if TYPE_CHECKING: import bigframes.operations as ops @@ -267,3 +268,13 @@ def _convert_range_ordering_to_table_value( # Singleton compiler scalar_op_compiler = ExpressionCompiler() + + +@scalar_op_compiler.register_unary_op(numeric_ops.isnan_op) +def isnan(arg): + return arg.isnan() + + +@scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) +def isfinite(arg): + return arg.isinf().not_() & arg.isnan().not_() diff --git a/bigframes/core/compile/polars/compiler.py b/bigframes/core/compile/polars/compiler.py index acaf1b8f22..7f1f42631c 100644 --- a/bigframes/core/compile/polars/compiler.py +++ b/bigframes/core/compile/polars/compiler.py @@ -199,6 +199,14 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: return input.ceil() + @compile_op.register(num_ops.IsNanOp) + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: + return input.is_nan() + + @compile_op.register(num_ops.IsFiniteOp) + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: + return input.is_finite() + @compile_op.register(num_ops.PosOp) def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: return input.__pos__() diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index 8167f40fc3..c6e052e3cd 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -22,6 +22,7 @@ import bigframes.core.compile.sqlglot.sqlglot_ir as ir import bigframes.core.expression as ex import bigframes.operations as ops +from bigframes.operations import numeric_ops class ScalarOpCompiler: @@ -180,3 +181,13 @@ def _register( # Singleton compiler scalar_op_compiler = ScalarOpCompiler() + + +@scalar_op_compiler.register_unary_op(numeric_ops.isnan_op) +def isnan(arg: TypedExpr) -> sge.Expression: + return sge.IsNan(this=arg.expr) + + +@scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) +def isfinite(arg: TypedExpr) -> sge.Expression: + return sge.Not(this=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr))) diff --git a/bigframes/operations/numeric_ops.py b/bigframes/operations/numeric_ops.py index afdc924c0b..83e2078c88 100644 --- a/bigframes/operations/numeric_ops.py +++ b/bigframes/operations/numeric_ops.py @@ -348,3 +348,19 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT name="unsafe_pow_op", type_signature=op_typing.BINARY_REAL_NUMERIC ) unsafe_pow_op = UnsafePowOp() + +IsNanOp = base_ops.create_unary_op( + name="isnan", + type_signature=op_typing.FixedOutputType( + dtypes.is_numeric, dtypes.BOOL_DTYPE, "numeric" + ), +) +isnan_op = IsNanOp() + +IsFiniteOp = base_ops.create_unary_op( + name="isfinite", + type_signature=op_typing.FixedOutputType( + dtypes.is_numeric, dtypes.BOOL_DTYPE, "numeric" + ), +) +isfinite_op = IsFiniteOp() diff --git a/bigframes/operations/numpy_op_maps.py b/bigframes/operations/numpy_op_maps.py index 7f3decdfa0..791e2eb890 100644 --- a/bigframes/operations/numpy_op_maps.py +++ b/bigframes/operations/numpy_op_maps.py @@ -40,6 +40,8 @@ np.ceil: numeric_ops.ceil_op, np.log1p: numeric_ops.log1p_op, np.expm1: numeric_ops.expm1_op, + np.isnan: numeric_ops.isnan_op, + np.isfinite: numeric_ops.isfinite_op, } diff --git a/tests/system/small/test_numpy.py b/tests/system/small/test_numpy.py index 37a707b9d0..490f927114 100644 --- a/tests/system/small/test_numpy.py +++ b/tests/system/small/test_numpy.py @@ -37,6 +37,8 @@ ("log10",), ("sqrt",), ("abs",), + ("isnan",), + ("isfinite",), ], ) def test_series_ufuncs(floats_pd, floats_bf, opname): From 2f1331141675a830e4de85164cb5107c043dc29a Mon Sep 17 00:00:00 2001 From: Owl Bot Date: Wed, 22 Oct 2025 18:21:15 +0000 Subject: [PATCH 02/10] =?UTF-8?q?=F0=9F=A6=89=20Updates=20from=20OwlBot=20?= =?UTF-8?q?post-processor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --- bigframes/core/compile/sqlglot/scalar_compiler.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index c6e052e3cd..496ea3ea26 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -21,8 +21,8 @@ from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr import bigframes.core.compile.sqlglot.sqlglot_ir as ir import bigframes.core.expression as ex -import bigframes.operations as ops from bigframes.operations import numeric_ops +import bigframes.operations as ops class ScalarOpCompiler: @@ -190,4 +190,6 @@ def isnan(arg: TypedExpr) -> sge.Expression: @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) def isfinite(arg: TypedExpr) -> sge.Expression: - return sge.Not(this=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr))) + return sge.Not( + this=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr)) + ) From 07e677bf693bd1f2e525197798fcc174e37dd160 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 22 Oct 2025 18:31:51 +0000 Subject: [PATCH 03/10] feat: Implement isnan and isfinite in bigframes This commit implements the `isnan` and `isfinite` numpy ufuncs in bigframes. The following changes were made: - Added `IsNanOrNullOp` and `IsFiniteOp` to `bigframes/operations/numeric_ops.py` - Mapped `np.isnan` and `np.isfinite` to the new ops in `bigframes/operations/numpy_op_maps.py` - Added compilation logic for the new ops to the ibis, polars, and sqlglot compilers - Added tests for the new ops in `tests/system/small/test_numpy.py` - Renamed `IsNanOp` to `IsNanOrNullOp` to match numpy semantics and updated compilers accordingly. --- .../core/compile/ibis_compiler/scalar_op_compiler.py | 6 +++--- bigframes/core/compile/polars/compiler.py | 4 ++-- bigframes/core/compile/sqlglot/scalar_compiler.py | 12 +++++------- bigframes/operations/numeric_ops.py | 6 +++--- bigframes/operations/numpy_op_maps.py | 2 +- 5 files changed, 14 insertions(+), 16 deletions(-) diff --git a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py index 7f74a1b7d9..35cdbc0cc7 100644 --- a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py +++ b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py @@ -270,9 +270,9 @@ def _convert_range_ordering_to_table_value( scalar_op_compiler = ExpressionCompiler() -@scalar_op_compiler.register_unary_op(numeric_ops.isnan_op) -def isnan(arg): - return arg.isnan() +@scalar_op_compiler.register_unary_op(numeric_ops.isnanornull_op) +def isnanornull(arg): + return arg.isnan() | arg.isnull() @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) diff --git a/bigframes/core/compile/polars/compiler.py b/bigframes/core/compile/polars/compiler.py index 7f1f42631c..10299ecf6d 100644 --- a/bigframes/core/compile/polars/compiler.py +++ b/bigframes/core/compile/polars/compiler.py @@ -199,9 +199,9 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: return input.ceil() - @compile_op.register(num_ops.IsNanOp) + @compile_op.register(num_ops.IsNanOrNullOp) def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: - return input.is_nan() + return input.is_nan() | input.is_null() @compile_op.register(num_ops.IsFiniteOp) def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index 496ea3ea26..7e11d55616 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -21,8 +21,8 @@ from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr import bigframes.core.compile.sqlglot.sqlglot_ir as ir import bigframes.core.expression as ex -from bigframes.operations import numeric_ops import bigframes.operations as ops +from bigframes.operations import numeric_ops class ScalarOpCompiler: @@ -183,13 +183,11 @@ def _register( scalar_op_compiler = ScalarOpCompiler() -@scalar_op_compiler.register_unary_op(numeric_ops.isnan_op) -def isnan(arg: TypedExpr) -> sge.Expression: - return sge.IsNan(this=arg.expr) +@scalar_op_compiler.register_unary_op(numeric_ops.isnanornull_op) +def isnanornull(arg: TypedExpr) -> sge.Expression: + return sge.Or(this=sge.IsNan(this=arg.expr), right=sge.IsNull(this=arg.expr)) @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) def isfinite(arg: TypedExpr) -> sge.Expression: - return sge.Not( - this=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr)) - ) + return sge.Not(this=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr))) diff --git a/bigframes/operations/numeric_ops.py b/bigframes/operations/numeric_ops.py index 83e2078c88..649289ce74 100644 --- a/bigframes/operations/numeric_ops.py +++ b/bigframes/operations/numeric_ops.py @@ -349,13 +349,13 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT ) unsafe_pow_op = UnsafePowOp() -IsNanOp = base_ops.create_unary_op( - name="isnan", +IsNanOrNullOp = base_ops.create_unary_op( + name="isnanornull", type_signature=op_typing.FixedOutputType( dtypes.is_numeric, dtypes.BOOL_DTYPE, "numeric" ), ) -isnan_op = IsNanOp() +isnanornull_op = IsNanOrNullOp() IsFiniteOp = base_ops.create_unary_op( name="isfinite", diff --git a/bigframes/operations/numpy_op_maps.py b/bigframes/operations/numpy_op_maps.py index 791e2eb890..57f7c422b8 100644 --- a/bigframes/operations/numpy_op_maps.py +++ b/bigframes/operations/numpy_op_maps.py @@ -40,7 +40,7 @@ np.ceil: numeric_ops.ceil_op, np.log1p: numeric_ops.log1p_op, np.expm1: numeric_ops.expm1_op, - np.isnan: numeric_ops.isnan_op, + np.isnan: numeric_ops.isnanornull_op, np.isfinite: numeric_ops.isfinite_op, } From 99afbdab7ccfeb76d671bf7a3dcf09ca0702ad02 Mon Sep 17 00:00:00 2001 From: Owl Bot Date: Wed, 22 Oct 2025 18:34:13 +0000 Subject: [PATCH 04/10] =?UTF-8?q?=F0=9F=A6=89=20Updates=20from=20OwlBot=20?= =?UTF-8?q?post-processor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --- bigframes/core/compile/sqlglot/scalar_compiler.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index 7e11d55616..9a2b34e79c 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -21,8 +21,8 @@ from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr import bigframes.core.compile.sqlglot.sqlglot_ir as ir import bigframes.core.expression as ex -import bigframes.operations as ops from bigframes.operations import numeric_ops +import bigframes.operations as ops class ScalarOpCompiler: @@ -190,4 +190,6 @@ def isnanornull(arg: TypedExpr) -> sge.Expression: @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) def isfinite(arg: TypedExpr) -> sge.Expression: - return sge.Not(this=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr))) + return sge.Not( + this=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr)) + ) From 3b0af20ca82a1f971440b518f0967e1df142b32c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a=20=28Swast=29?= Date: Wed, 22 Oct 2025 13:37:24 -0500 Subject: [PATCH 05/10] Apply suggestions from code review --- bigframes/core/compile/ibis_compiler/scalar_op_compiler.py | 2 +- bigframes/core/compile/sqlglot/scalar_compiler.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py index 35cdbc0cc7..213ff10c29 100644 --- a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py +++ b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py @@ -277,4 +277,4 @@ def isnanornull(arg): @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) def isfinite(arg): - return arg.isinf().not_() & arg.isnan().not_() + return arg.isinf().not_() & arg.isnan().not_() & arg.isnull().not_() diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index 9a2b34e79c..6234143037 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -191,5 +191,5 @@ def isnanornull(arg: TypedExpr) -> sge.Expression: @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) def isfinite(arg: TypedExpr) -> sge.Expression: return sge.Not( - this=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr)) + this=sge.Or(this=sge.Is(this=expr.expr, expression=sge.Null()), right=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr)) ) From 8f67cb28349aebafe2e36e2fd780999af4a389cd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a=20=28Swast=29?= Date: Wed, 22 Oct 2025 13:41:58 -0500 Subject: [PATCH 06/10] Update bigframes/core/compile/sqlglot/scalar_compiler.py --- bigframes/core/compile/sqlglot/scalar_compiler.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index 6234143037..63cc93b4ca 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -191,5 +191,11 @@ def isnanornull(arg: TypedExpr) -> sge.Expression: @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) def isfinite(arg: TypedExpr) -> sge.Expression: return sge.Not( - this=sge.Or(this=sge.Is(this=expr.expr, expression=sge.Null()), right=sge.Or(this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr)) + this=sge.Or( + this=sge.Is(this=expr.expr, expression=sge.Null()), + right=sge.Or( + this=sge.IsInf(this=arg.expr), + right=sge.IsNan(this=arg.expr), + ), + ), ) From aae29026e9ba66d44f42e39ba5ec09fb85da4bed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a=20=28Swast=29?= Date: Wed, 22 Oct 2025 13:44:16 -0500 Subject: [PATCH 07/10] Update bigframes/core/compile/sqlglot/scalar_compiler.py --- bigframes/core/compile/sqlglot/scalar_compiler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index 63cc93b4ca..96a4a875ff 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -192,7 +192,7 @@ def isnanornull(arg: TypedExpr) -> sge.Expression: def isfinite(arg: TypedExpr) -> sge.Expression: return sge.Not( this=sge.Or( - this=sge.Is(this=expr.expr, expression=sge.Null()), + this=sge.Is(this=arg.expr, expression=sge.Null()), right=sge.Or( this=sge.IsInf(this=arg.expr), right=sge.IsNan(this=arg.expr), From 41ce94b5e61bb34b9a7fdb6e8019c9532432602a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a=20=28Swast=29?= Date: Wed, 22 Oct 2025 13:48:46 -0500 Subject: [PATCH 08/10] Update bigframes/core/compile/sqlglot/scalar_compiler.py --- bigframes/core/compile/sqlglot/scalar_compiler.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index 96a4a875ff..1345593402 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -185,7 +185,10 @@ def _register( @scalar_op_compiler.register_unary_op(numeric_ops.isnanornull_op) def isnanornull(arg: TypedExpr) -> sge.Expression: - return sge.Or(this=sge.IsNan(this=arg.expr), right=sge.IsNull(this=arg.expr)) + return sge.Or( + this=sge.IsNan(this=arg.expr), + right=sge.Is(this=arg.expr, expression=sge.Null()), + ) @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) From c0b5e685a53c242e2071a63d6d52845526d54956 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 22 Oct 2025 20:28:40 +0000 Subject: [PATCH 09/10] preserve NA --- .../ibis_compiler/scalar_op_compiler.py | 6 ++--- bigframes/core/compile/polars/compiler.py | 8 ------- .../compile/polars/operations/numeric_ops.py | 18 +++++++++++++++ .../sqlglot/expressions/numeric_ops.py | 16 ++++++++++++++ .../core/compile/sqlglot/scalar_compiler.py | 22 ------------------- bigframes/operations/numeric_ops.py | 4 ++-- bigframes/operations/numpy_op_maps.py | 2 +- 7 files changed, 40 insertions(+), 36 deletions(-) diff --git a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py index 213ff10c29..8a027ca296 100644 --- a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py +++ b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py @@ -270,11 +270,11 @@ def _convert_range_ordering_to_table_value( scalar_op_compiler = ExpressionCompiler() -@scalar_op_compiler.register_unary_op(numeric_ops.isnanornull_op) +@scalar_op_compiler.register_unary_op(numeric_ops.isnan_op) def isnanornull(arg): - return arg.isnan() | arg.isnull() + return arg.isnan() @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) def isfinite(arg): - return arg.isinf().not_() & arg.isnan().not_() & arg.isnull().not_() + return arg.isinf().negate() & arg.isnan().negate() diff --git a/bigframes/core/compile/polars/compiler.py b/bigframes/core/compile/polars/compiler.py index 10299ecf6d..acaf1b8f22 100644 --- a/bigframes/core/compile/polars/compiler.py +++ b/bigframes/core/compile/polars/compiler.py @@ -199,14 +199,6 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: return input.ceil() - @compile_op.register(num_ops.IsNanOrNullOp) - def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: - return input.is_nan() | input.is_null() - - @compile_op.register(num_ops.IsFiniteOp) - def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: - return input.is_finite() - @compile_op.register(num_ops.PosOp) def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: return input.__pos__() diff --git a/bigframes/core/compile/polars/operations/numeric_ops.py b/bigframes/core/compile/polars/operations/numeric_ops.py index 2572d862e3..89bcb75bf5 100644 --- a/bigframes/core/compile/polars/operations/numeric_ops.py +++ b/bigframes/core/compile/polars/operations/numeric_ops.py @@ -89,3 +89,21 @@ def sqrt_op_impl( import polars as pl return pl.when(input < 0).then(float("nan")).otherwise(input.sqrt()) + + +@polars_compiler.register_op(numeric_ops.IsNanOp) +def is_nan_op_impl( + compiler: polars_compiler.PolarsExpressionCompiler, + op: numeric_ops.SqrtOp, # type: ignore + input: pl.Expr, +) -> pl.Expr: + return input.is_nan() + + +@polars_compiler.register_op(numeric_ops.IsFiniteOp) +def is_finite_op_impl( + compiler: polars_compiler.PolarsExpressionCompiler, + op: numeric_ops.SqrtOp, # type: ignore + input: pl.Expr, +) -> pl.Expr: + return input.is_finite() diff --git a/bigframes/core/compile/sqlglot/expressions/numeric_ops.py b/bigframes/core/compile/sqlglot/expressions/numeric_ops.py index 3bbe2623ea..64e27ebe79 100644 --- a/bigframes/core/compile/sqlglot/expressions/numeric_ops.py +++ b/bigframes/core/compile/sqlglot/expressions/numeric_ops.py @@ -22,6 +22,7 @@ import bigframes.core.compile.sqlglot.expressions.constants as constants from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler +from bigframes.operations import numeric_ops register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op @@ -408,6 +409,21 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: ) +@register_unary_op(numeric_ops.isnan_op) +def isnan(arg: TypedExpr) -> sge.Expression: + return sge.IsNan(this=arg.expr) + + +@register_unary_op(numeric_ops.isfinite_op) +def isfinite(arg: TypedExpr) -> sge.Expression: + return sge.Not( + this=sge.Or( + this=sge.IsInf(this=arg.expr), + right=sge.IsNan(this=arg.expr), + ), + ) + + def _coerce_bool_to_int(typed_expr: TypedExpr) -> sge.Expression: """Coerce boolean expression to integer.""" if typed_expr.dtype == dtypes.BOOL_DTYPE: diff --git a/bigframes/core/compile/sqlglot/scalar_compiler.py b/bigframes/core/compile/sqlglot/scalar_compiler.py index 1345593402..8167f40fc3 100644 --- a/bigframes/core/compile/sqlglot/scalar_compiler.py +++ b/bigframes/core/compile/sqlglot/scalar_compiler.py @@ -21,7 +21,6 @@ from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr import bigframes.core.compile.sqlglot.sqlglot_ir as ir import bigframes.core.expression as ex -from bigframes.operations import numeric_ops import bigframes.operations as ops @@ -181,24 +180,3 @@ def _register( # Singleton compiler scalar_op_compiler = ScalarOpCompiler() - - -@scalar_op_compiler.register_unary_op(numeric_ops.isnanornull_op) -def isnanornull(arg: TypedExpr) -> sge.Expression: - return sge.Or( - this=sge.IsNan(this=arg.expr), - right=sge.Is(this=arg.expr, expression=sge.Null()), - ) - - -@scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) -def isfinite(arg: TypedExpr) -> sge.Expression: - return sge.Not( - this=sge.Or( - this=sge.Is(this=arg.expr, expression=sge.Null()), - right=sge.Or( - this=sge.IsInf(this=arg.expr), - right=sge.IsNan(this=arg.expr), - ), - ), - ) diff --git a/bigframes/operations/numeric_ops.py b/bigframes/operations/numeric_ops.py index 649289ce74..aee71e7bdb 100644 --- a/bigframes/operations/numeric_ops.py +++ b/bigframes/operations/numeric_ops.py @@ -349,13 +349,13 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT ) unsafe_pow_op = UnsafePowOp() -IsNanOrNullOp = base_ops.create_unary_op( +IsNanOp = base_ops.create_unary_op( name="isnanornull", type_signature=op_typing.FixedOutputType( dtypes.is_numeric, dtypes.BOOL_DTYPE, "numeric" ), ) -isnanornull_op = IsNanOrNullOp() +isnan_op = IsNanOp() IsFiniteOp = base_ops.create_unary_op( name="isfinite", diff --git a/bigframes/operations/numpy_op_maps.py b/bigframes/operations/numpy_op_maps.py index 57f7c422b8..791e2eb890 100644 --- a/bigframes/operations/numpy_op_maps.py +++ b/bigframes/operations/numpy_op_maps.py @@ -40,7 +40,7 @@ np.ceil: numeric_ops.ceil_op, np.log1p: numeric_ops.log1p_op, np.expm1: numeric_ops.expm1_op, - np.isnan: numeric_ops.isnanornull_op, + np.isnan: numeric_ops.isnan_op, np.isfinite: numeric_ops.isfinite_op, } From 7f1bca4424ca2ff978aea03ddb9576c6d7e17221 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 22 Oct 2025 21:41:54 +0000 Subject: [PATCH 10/10] fix annotations --- bigframes/core/compile/polars/operations/numeric_ops.py | 4 ++-- bigframes/operations/numeric_ops.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/bigframes/core/compile/polars/operations/numeric_ops.py b/bigframes/core/compile/polars/operations/numeric_ops.py index 89bcb75bf5..8e44f15955 100644 --- a/bigframes/core/compile/polars/operations/numeric_ops.py +++ b/bigframes/core/compile/polars/operations/numeric_ops.py @@ -94,7 +94,7 @@ def sqrt_op_impl( @polars_compiler.register_op(numeric_ops.IsNanOp) def is_nan_op_impl( compiler: polars_compiler.PolarsExpressionCompiler, - op: numeric_ops.SqrtOp, # type: ignore + op: numeric_ops.IsNanOp, # type: ignore input: pl.Expr, ) -> pl.Expr: return input.is_nan() @@ -103,7 +103,7 @@ def is_nan_op_impl( @polars_compiler.register_op(numeric_ops.IsFiniteOp) def is_finite_op_impl( compiler: polars_compiler.PolarsExpressionCompiler, - op: numeric_ops.SqrtOp, # type: ignore + op: numeric_ops.IsFiniteOp, # type: ignore input: pl.Expr, ) -> pl.Expr: return input.is_finite() diff --git a/bigframes/operations/numeric_ops.py b/bigframes/operations/numeric_ops.py index aee71e7bdb..83e2078c88 100644 --- a/bigframes/operations/numeric_ops.py +++ b/bigframes/operations/numeric_ops.py @@ -350,7 +350,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT unsafe_pow_op = UnsafePowOp() IsNanOp = base_ops.create_unary_op( - name="isnanornull", + name="isnan", type_signature=op_typing.FixedOutputType( dtypes.is_numeric, dtypes.BOOL_DTYPE, "numeric" ),