feat: add support for np.isnan and np.isfinite ufuncs (#2188)

* 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`

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* 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.

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* Apply suggestions from code review

* Update bigframes/core/compile/sqlglot/scalar_compiler.py

* Update bigframes/core/compile/sqlglot/scalar_compiler.py

* Update bigframes/core/compile/sqlglot/scalar_compiler.py

* preserve NA

* fix annotations

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: 3 people

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()

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()

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:

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()

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,
 }
 
 

tests/system/small/test_numpy.py

@@ -37,6 +37,8 @@
         ("log10",),
         ("sqrt",),
         ("abs",),
+        ("isnan",),
+        ("isfinite",),
     ],
 )
 def test_series_ufuncs(floats_pd, floats_bf, opname):