diff --git a/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py b/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py index 1197f6b9da..8a027ca296 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 isnanornull(arg): + return arg.isnan() + + +@scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op) +def isfinite(arg): + return arg.isinf().negate() & arg.isnan().negate() diff --git a/bigframes/core/compile/polars/operations/numeric_ops.py b/bigframes/core/compile/polars/operations/numeric_ops.py index 2572d862e3..8e44f15955 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.IsNanOp, # 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.IsFiniteOp, # 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/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):