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`

Author: google-labs-jules[bot]

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

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

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

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