fix more tests

Author: TrevorBergeron

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -169,6 +169,8 @@ def arctanh_op_impl(x: ibis_types.Value):
 @scalar_op_compiler.register_unary_op(ops.floor_op)
 def floor_op_impl(x: ibis_types.Value):
     x_numeric = typing.cast(ibis_types.NumericValue, x)
+    if x_numeric.type().is_boolean():
+        return x_numeric.cast(ibis_dtypes.Int64()).cast(ibis_dtypes.Float64())
     if x_numeric.type().is_integer():
         return x_numeric.cast(ibis_dtypes.Float64())
     if x_numeric.type().is_floating():
@@ -181,6 +183,8 @@ def floor_op_impl(x: ibis_types.Value):
 @scalar_op_compiler.register_unary_op(ops.ceil_op)
 def ceil_op_impl(x: ibis_types.Value):
     x_numeric = typing.cast(ibis_types.NumericValue, x)
+    if x_numeric.type().is_boolean():
+        return x_numeric.cast(ibis_dtypes.Int64()).cast(ibis_dtypes.Float64())
     if x_numeric.type().is_integer():
         return x_numeric.cast(ibis_dtypes.Float64())
     if x_numeric.type().is_floating():

bigframes/core/compile/polars/lowering.py

@@ -174,12 +174,10 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
             divisor.output_type
         ):
             # exact same as floordiv impl for timedelta
-            numeric_result = ops.floordiv_op.as_expr(
+            numeric_result = ops.div_op.as_expr(
                 ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(dividend), divisor
             )
-            int_result = ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(numeric_result)
-            return ops.AsTypeOp(to_type=dtypes.TIMEDELTA_DTYPE).as_expr(int_result)
-
+            return _numeric_to_timedelta(numeric_result)
         if (
             dividend.output_type == dtypes.BOOL_DTYPE
             and divisor.output_type == dtypes.BOOL_DTYPE
@@ -226,11 +224,10 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
             divisor.output_type
         ):
             # this is pretty fragile as zero will break it, and must fit back into int
-            numeric_result = expr.op.as_expr(
+            numeric_result = ops.div_op.as_expr(
                 ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(dividend), divisor
             )
-            int_result = ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(numeric_result)
-            return ops.AsTypeOp(to_type=dtypes.TIMEDELTA_DTYPE).as_expr(int_result)
+            return _numeric_to_timedelta(numeric_result)
 
         if dividend.output_type == dtypes.BOOL_DTYPE:
             dividend = ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(dividend)
@@ -319,6 +316,32 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
         return expr
 
 
+class LowerCeilOp(op_lowering.OpLoweringRule):
+    @property
+    def op(self) -> type[ops.ScalarOp]:
+        return numeric_ops.CeilOp
+
+    def lower(self, expr: expression.OpExpression) -> expression.Expression:
+        assert isinstance(expr.op, numeric_ops.CeilOp)
+        arg = expr.children[0]
+        if arg.output_type in (dtypes.INT_DTYPE, dtypes.BOOL_DTYPE):
+            return expr.op.as_expr(ops.AsTypeOp(dtypes.FLOAT_DTYPE).as_expr(arg))
+        return expr
+
+
+class LowerFloorOp(op_lowering.OpLoweringRule):
+    @property
+    def op(self) -> type[ops.ScalarOp]:
+        return numeric_ops.FloorOp
+
+    def lower(self, expr: expression.OpExpression) -> expression.Expression:
+        assert isinstance(expr.op, numeric_ops.FloorOp)
+        arg = expr.children[0]
+        if arg.output_type in (dtypes.INT_DTYPE, dtypes.BOOL_DTYPE):
+            return expr.op.as_expr(ops.AsTypeOp(dtypes.FLOAT_DTYPE).as_expr(arg))
+        return expr
+
+
 class LowerIsinOp(op_lowering.OpLoweringRule):
     @property
     def op(self) -> type[ops.ScalarOp]:
@@ -465,8 +488,21 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
     LowerInvertOp(),
     LowerIsinOp(),
     LowerLenOp(),
+    LowerCeilOp(),
+    LowerFloorOp(),
 )
 
 
 def lower_ops_to_polars(root: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:
     return op_lowering.lower_ops(root, rules=POLARS_LOWERING_RULES)
+
+
+def _numeric_to_timedelta(expr: expression.Expression) -> expression.Expression:
+    """rounding logic used for emulating timedelta ops"""
+    rounded_value = ops.where_op.as_expr(
+        ops.floor_op.as_expr(expr),
+        ops.gt_op.as_expr(expr, expression.const(0)),
+        ops.ceil_op.as_expr(expr),
+    )
+    int_value = ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(rounded_value)
+    return ops.AsTypeOp(to_type=dtypes.TIMEDELTA_DTYPE).as_expr(int_value)

bigframes/core/compile/sqlglot/expressions/common.py

@@ -0,0 +1,33 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import bigframes_vendored.sqlglot.expressions as sge
+
+
+def round_towards_zero(expr: sge.Expression):
+    """
+    Round a float value to to an integer, always rounding towards zero.
+
+    This is used to handle duration/timedelta emulation mostly.
+    """
+    return sge.Cast(
+        this=sge.If(
+            this=sge.GT(this=expr, expression=sge.convert(0)),
+            true=sge.Floor(this=expr),
+            false=sge.Ceil(this=expr),
+        ),
+        to="INT64",
+    )

bigframes/core/compile/sqlglot/expressions/numeric_ops.py

@@ -20,6 +20,7 @@
 from bigframes import dtypes
 from bigframes import operations as ops
 import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
+from bigframes.core.compile.sqlglot.expressions.common import round_towards_zero
 import bigframes.core.compile.sqlglot.expressions.constants as constants
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 from bigframes.operations import numeric_ops
@@ -467,7 +468,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
     result = sge.func("IEEE_DIVIDE", left_expr, right_expr)
     if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
-        return sge.Cast(this=sge.Floor(this=result), to="INT64")
+        return round_towards_zero(result)
     else:
         return result
 
@@ -510,7 +511,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     )
 
     if dtypes.is_numeric(right.dtype) and left.dtype == dtypes.TIMEDELTA_DTYPE:
-        result = sge.Cast(this=sge.Floor(this=result), to="INT64")
+        result = round_towards_zero(sge.func("IEEE_DIVIDE", left_expr, right_expr))
 
     return result
 
@@ -578,7 +579,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     if (dtypes.is_numeric(left.dtype) and right.dtype == dtypes.TIMEDELTA_DTYPE) or (
         left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype)
     ):
-        return sge.Cast(this=sge.Floor(this=result), to="INT64")
+        return round_towards_zero(result)
     else:
         return result
 

bigframes/session/polars_executor.py

@@ -79,6 +79,8 @@
     numeric_ops.SubOp,
     numeric_ops.MulOp,
     numeric_ops.DivOp,
+    numeric_ops.CeilOp,
+    numeric_ops.FloorOp,
     numeric_ops.FloorDivOp,
     numeric_ops.ModOp,
     generic_ops.AsTypeOp,

tests/system/large/functions/test_remote_function.py

@@ -1354,7 +1354,7 @@ def double_num(x):
         bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).to_pandas()
 
         pd_int64_col = scalars_pandas_df["int64_col"]
-        pd_result_col = pd_int64_col.apply(lambda x: x if x is None else x * x)
+        pd_result_col = pd_int64_col.apply(lambda x: x if x is None else x + x)
         pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
 
         assert_frame_equal(bf_result, pd_result, check_dtype=False)

tests/system/small/engines/test_numeric_ops.py

@@ -53,6 +53,47 @@ def apply_op_pairwise(
     return new_arr
 
 
+def apply_op(
+    array: array_value.ArrayValue, op: ops.UnaryOp, excluded_cols=[]
+) -> array_value.ArrayValue:
+    exprs = []
+    labels = []
+    for arg in array.column_ids:
+        if arg in excluded_cols:
+            continue
+        try:
+            _ = op.output_type(array.get_column_type(arg))
+            expr = op.as_expr(arg)
+            exprs.append(expr)
+            labels.append(f"{arg}_{op.name}")
+        except TypeError:
+            continue
+    assert len(exprs) > 0
+    new_arr, ids = array.compute_values(exprs)
+    new_arr = new_arr.rename_columns(
+        {new_col: label for new_col, label in zip(ids, labels)}
+    )
+    return new_arr
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_project_ceil(
+    scalars_array_value: array_value.ArrayValue,
+    engine,
+):
+    arr = apply_op(scalars_array_value, ops.ceil_op)
+    assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_project_floor(
+    scalars_array_value: array_value.ArrayValue,
+    engine,
+):
+    arr = apply_op(scalars_array_value, ops.floor_op)
+    assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
 @pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 def test_engines_project_add(
     scalars_array_value: array_value.ArrayValue,