feat: Add where, coalesce, fillna, casewhen, invert local impl (#1976)

Author: TrevorBergeron

bigframes/core/compile/polars/compiler.py

@@ -168,7 +168,7 @@ def compile_op(self, op: ops.ScalarOp, *args: pl.Expr) -> pl.Expr:
 
  @compile_op.register(gen_ops.InvertOp)
  def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
- return ~input
+ return input.not_()
 
  @compile_op.register(num_ops.AbsOp)
  def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:

bigframes/core/compile/polars/lowering.py

@@ -14,10 +14,18 @@
 
 import dataclasses
 
+import numpy as np
+
 from bigframes import dtypes
 from bigframes.core import bigframe_node, expression
 from bigframes.core.rewrite import op_lowering
-from bigframes.operations import comparison_ops, datetime_ops, json_ops, numeric_ops
+from bigframes.operations import (
+ comparison_ops,
+ datetime_ops,
+ generic_ops,
+ json_ops,
+ numeric_ops,
+)
 import bigframes.operations as ops
 
 # TODO: Would be more precise to actually have separate op set for polars ops (where they diverge from the original ops)
@@ -288,6 +296,26 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
  return _lower_cast(expr.op, expr.inputs[0])
 
 
+def invert_bytes(byte_string):
+ inverted_bytes = ~np.frombuffer(byte_string, dtype=np.uint8)
+ return inverted_bytes.tobytes()
+
+
+class LowerInvertOp(op_lowering.OpLoweringRule):
+ @property
+ def op(self) -> type[ops.ScalarOp]:
+ return generic_ops.InvertOp
+
+ def lower(self, expr: expression.OpExpression) -> expression.Expression:
+ assert isinstance(expr.op, generic_ops.InvertOp)
+ arg = expr.children[0]
+ if arg.output_type == dtypes.BYTES_DTYPE:
+ return generic_ops.PyUdfOp(invert_bytes, dtypes.BYTES_DTYPE).as_expr(
+ expr.inputs[0]
+ )
+ return expr
+
+
 def _coerce_comparables(
  expr1: expression.Expression,
  expr2: expression.Expression,
@@ -385,6 +413,7 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
  LowerFloorDivRule(),
  LowerModRule(),
  LowerAsTypeRule(),
+ LowerInvertOp(),
 )
 
 

bigframes/core/compile/polars/operations/generic_ops.py

@@ -45,3 +45,14 @@ def isnull_op_impl(
  input: pl.Expr,
 ) -> pl.Expr:
  return input.is_null()
+
+
+@polars_compiler.register_op(generic_ops.PyUdfOp)
+def py_udf_op_impl(
+ compiler: polars_compiler.PolarsExpressionCompiler,
+ op: generic_ops.PyUdfOp, # type: ignore
+ input: pl.Expr,
+) -> pl.Expr:
+ return input.map_elements(
+ op.fn, return_dtype=polars_compiler._DTYPE_MAPPING[op._output_type]
+ )

bigframes/operations/generic_ops.py

@@ -446,3 +446,15 @@ class SqlScalarOp(base_ops.NaryOp):
 
  def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
  return self._output_type
+
+
+@dataclasses.dataclass(frozen=True)
+class PyUdfOp(base_ops.NaryOp):
+ """Represents a local UDF."""
+
+ name: typing.ClassVar[str] = "py_udf"
+ fn: typing.Callable
+ _output_type: dtypes.ExpressionType
+
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ return self._output_type

bigframes/session/polars_executor.py

@@ -58,6 +58,11 @@
  numeric_ops.FloorDivOp,
  numeric_ops.ModOp,
  generic_ops.AsTypeOp,
+ generic_ops.WhereOp,
+ generic_ops.CoalesceOp,
+ generic_ops.FillNaOp,
+ generic_ops.CaseWhenOp,
+ generic_ops.InvertOp,
 )
 _COMPATIBLE_AGG_OPS = (
  agg_ops.SizeOp,

tests/system/small/engines/test_generic_ops.py

@@ -59,6 +59,7 @@ def test_engines_astype_int(scalars_array_value: array_value.ArrayValue, engine)
  ops.AsTypeOp(to_type=bigframes.dtypes.INT_DTYPE),
  excluded_cols=["string_col"],
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -73,6 +74,7 @@ def test_engines_astype_string_int(scalars_array_value: array_value.ArrayValue,
  for val in vals
  ]
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -83,6 +85,7 @@ def test_engines_astype_float(scalars_array_value: array_value.ArrayValue, engin
  ops.AsTypeOp(to_type=bigframes.dtypes.FLOAT_DTYPE),
  excluded_cols=["string_col"],
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -99,6 +102,7 @@ def test_engines_astype_string_float(
  for val in vals
  ]
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -107,6 +111,7 @@ def test_engines_astype_bool(scalars_array_value: array_value.ArrayValue, engine
  arr = apply_op(
  scalars_array_value, ops.AsTypeOp(to_type=bigframes.dtypes.BOOL_DTYPE)
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -118,6 +123,7 @@ def test_engines_astype_string(scalars_array_value: array_value.ArrayValue, engi
  ops.AsTypeOp(to_type=bigframes.dtypes.STRING_DTYPE),
  excluded_cols=["float64_col"],
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -128,6 +134,7 @@ def test_engines_astype_numeric(scalars_array_value: array_value.ArrayValue, eng
  ops.AsTypeOp(to_type=bigframes.dtypes.NUMERIC_DTYPE),
  excluded_cols=["string_col"],
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -144,6 +151,7 @@ def test_engines_astype_string_numeric(
  for val in vals
  ]
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -154,6 +162,7 @@ def test_engines_astype_date(scalars_array_value: array_value.ArrayValue, engine
  ops.AsTypeOp(to_type=bigframes.dtypes.DATE_DTYPE),
  excluded_cols=["string_col"],
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -170,6 +179,7 @@ def test_engines_astype_string_date(
  for val in vals
  ]
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -180,6 +190,7 @@ def test_engines_astype_datetime(scalars_array_value: array_value.ArrayValue, en
  ops.AsTypeOp(to_type=bigframes.dtypes.DATETIME_DTYPE),
  excluded_cols=["string_col"],
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -196,6 +207,7 @@ def test_engines_astype_string_datetime(
  for val in vals
  ]
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -206,6 +218,7 @@ def test_engines_astype_timestamp(scalars_array_value: array_value.ArrayValue, e
  ops.AsTypeOp(to_type=bigframes.dtypes.TIMESTAMP_DTYPE),
  excluded_cols=["string_col"],
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -226,6 +239,7 @@ def test_engines_astype_string_timestamp(
  for val in vals
  ]
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -236,6 +250,7 @@ def test_engines_astype_time(scalars_array_value: array_value.ArrayValue, engine
  ops.AsTypeOp(to_type=bigframes.dtypes.TIME_DTYPE),
  excluded_cols=["string_col", "int64_col", "int64_too"],
  )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -256,6 +271,7 @@ def test_engines_astype_from_json(scalars_array_value: array_value.ArrayValue, e
  ),
  ]
  arr, _ = scalars_array_value.compute_values(exprs)
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
@@ -265,4 +281,112 @@ def test_engines_astype_timedelta(scalars_array_value: array_value.ArrayValue, e
  scalars_array_value,
  ops.AsTypeOp(to_type=bigframes.dtypes.TIMEDELTA_DTYPE),
  )
+
+ assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_where_op(scalars_array_value: array_value.ArrayValue, engine):
+ arr, _ = scalars_array_value.compute_values(
+ [
+ ops.where_op.as_expr(
+ expression.deref("int64_col"),
+ expression.deref("bool_col"),
+ expression.deref("float64_col"),
+ )
+ ]
+ )
+
+ assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_coalesce_op(scalars_array_value: array_value.ArrayValue, engine):
+ arr, _ = scalars_array_value.compute_values(
+ [
+ ops.coalesce_op.as_expr(
+ expression.deref("int64_col"),
+ expression.deref("float64_col"),
+ )
+ ]
+ )
+
+ assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_fillna_op(scalars_array_value: array_value.ArrayValue, engine):
+ arr, _ = scalars_array_value.compute_values(
+ [
+ ops.fillna_op.as_expr(
+ expression.deref("int64_col"),
+ expression.deref("float64_col"),
+ )
+ ]
+ )
+
+ assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_casewhen_op_single_case(
+ scalars_array_value: array_value.ArrayValue, engine
+):
+ arr, _ = scalars_array_value.compute_values(
+ [
+ ops.case_when_op.as_expr(
+ expression.deref("bool_col"),
+ expression.deref("int64_col"),
+ )
+ ]
+ )
+
+ assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_casewhen_op_double_case(
+ scalars_array_value: array_value.ArrayValue, engine
+):
+ arr, _ = scalars_array_value.compute_values(
+ [
+ ops.case_when_op.as_expr(
+ ops.gt_op.as_expr(expression.deref("int64_col"), expression.const(3)),
+ expression.deref("int64_col"),
+ ops.lt_op.as_expr(expression.deref("int64_col"), expression.const(-3)),
+ expression.deref("int64_too"),
+ )
+ ]
+ )
+
+ assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_isnull_op(scalars_array_value: array_value.ArrayValue, engine):
+ arr, _ = scalars_array_value.compute_values(
+ [ops.isnull_op.as_expr(expression.deref("string_col"))]
+ )
+
+ assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_notnull_op(scalars_array_value: array_value.ArrayValue, engine):
+ arr, _ = scalars_array_value.compute_values(
+ [ops.notnull_op.as_expr(expression.deref("string_col"))]
+ )
+
+ assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_invert_op(scalars_array_value: array_value.ArrayValue, engine):
+ arr, _ = scalars_array_value.compute_values(
+ [
+ ops.invert_op.as_expr(expression.deref("bytes_col")),
+ ops.invert_op.as_expr(expression.deref("bool_col")),
+ ]
+ )
+
  assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)

third_party/bigframes_vendored/ibis/expr/operations/numeric.py

@@ -326,7 +326,7 @@ class Tan(TrigonometricUnary):
 class BitwiseNot(Unary):
  """Bitwise NOT operation."""
 
- arg: Integer
+ arg: Value[dt.Integer | dt.Binary]
 
  dtype = rlz.numeric_like("args", operator.invert)
 

third_party/bigframes_vendored/ibis/expr/types/binary.py

@@ -32,6 +32,9 @@ def hashbytes(
  """
  return ops.HashBytes(self, how).to_expr()
 
+ def __invert__(self) -> BinaryValue:
+ return ops.BitwiseNot(self).to_expr()
+
 
 @public
 class BinaryScalar(Scalar, BinaryValue):