preserve NA

Author: tswast

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

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

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

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

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

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