fix time roundings issues

Author: TrevorBergeron

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1026,7 +1026,12 @@ def to_timedelta_op_impl(x: ibis_types.Value, op: ops.ToTimedeltaOp):
 
 @scalar_op_compiler.register_unary_op(ops.timedelta_floor_op)
 def timedelta_floor_op_impl(x: ibis_types.NumericValue):
-    return x.floor()
+    return ibis_api.case().when(x > 0, x.floor()).else_(x.ceil()).end()
+
+
+@scalar_op_compiler.register_unary_op(ops.timedelta_round_op)
+def timedelta_round_op_impl(x: ibis_types.NumericValue):
+    return ibis_api.case().when(x > 0, x.floor()).else_(x.ceil()).end()
 
 
 @scalar_op_compiler.register_unary_op(ops.RemoteFunctionOp, pass_op=True)

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

@@ -30,6 +30,11 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.Floor(this=expr.expr)
 
 
+@register_unary_op(ops.timedelta_round_op)
+def _(expr: TypedExpr) -> sge.Expression:
+    return sge.Round(this=expr.expr).cast()
+
+
 @register_unary_op(ops.ToTimedeltaOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.ToTimedeltaOp) -> sge.Expression:
     value = expr.expr

bigframes/core/rewrite/timedeltas.py

@@ -189,9 +189,9 @@ def _rewrite_mul_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
     result = _TypedExpr.create_op_expr(ops.mul_op, left, right)
 
     if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
-        return _TypedExpr.create_op_expr(ops.timedelta_floor_op, result)
+        return _TypedExpr.create_op_expr(ops.timedelta_round_op, result)
     if dtypes.is_numeric(left.dtype) and right.dtype == dtypes.TIMEDELTA_DTYPE:
-        return _TypedExpr.create_op_expr(ops.timedelta_floor_op, result)
+        return _TypedExpr.create_op_expr(ops.timedelta_round_op, result)
 
     return result
 
@@ -200,18 +200,18 @@ def _rewrite_div_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
     result = _TypedExpr.create_op_expr(ops.div_op, left, right)
 
     if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
-        return _TypedExpr.create_op_expr(ops.timedelta_floor_op, result)
+        return _TypedExpr.create_op_expr(ops.timedelta_round_op, result)
 
     return result
 
 
 def _rewrite_floordiv_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
-    result = _TypedExpr.create_op_expr(ops.floordiv_op, left, right)
-
     if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
-        return _TypedExpr.create_op_expr(ops.timedelta_floor_op, result)
+        return _TypedExpr.create_op_expr(
+            ops.timedelta_round_op, _TypedExpr.create_op_expr(ops.div_op, left, right)
+        )
 
-    return result
+    return _TypedExpr.create_op_expr(ops.floordiv_op, left, right)
 
 
 def _rewrite_to_timedelta_op(op: ops.ToTimedeltaOp, arg: _TypedExpr):

bigframes/operations/__init__.py

@@ -224,6 +224,7 @@
     date_add_op,
     date_sub_op,
     timedelta_floor_op,
+    timedelta_round_op,
     timestamp_add_op,
     timestamp_sub_op,
     ToTimedeltaOp,
@@ -308,6 +309,7 @@
     "date_add_op",
     "date_sub_op",
     "timedelta_floor_op",
+    "timedelta_round_op",
     "timestamp_add_op",
     "timestamp_sub_op",
     "ToTimedeltaOp",

bigframes/operations/timedelta_ops.py

@@ -54,6 +54,25 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 timedelta_floor_op = TimedeltaFloorOp()
 
 
+@dataclasses.dataclass(frozen=True)
+class TimedeltaRoundOp(base_ops.UnaryOp):
+    """Rounds the numeric value to the nearest integer and use it to represent a timedelta.
+
+    This operator is only meant to be used during expression tree rewrites. Do not use it anywhere else!
+    """
+
+    name: typing.ClassVar[str] = "timedelta_round"
+
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        input_type = input_types[0]
+        if dtypes.is_numeric(input_type) or input_type == dtypes.TIMEDELTA_DTYPE:
+            return dtypes.TIMEDELTA_DTYPE
+        raise TypeError(f"unsupported type: {input_type}")
+
+
+timedelta_round_op = TimedeltaRoundOp()
+
+
 @dataclasses.dataclass(frozen=True)
 class TimestampAddOp(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "timestamp_add"

tests/system/small/operations/test_timedeltas.py

@@ -87,10 +87,9 @@ def temporal_dfs(session):
 
 def _assert_series_equal(actual: pd.Series, expected: pd.Series):
     """Helper function specifically for timedelta testing. Don't use it outside of this module."""
-    # expected[expected.select_dtypes('timedelta64').columns] = expected.select_dtypes('timedelta64').astype(dtypes.TIMEDELTA_DTYPE)
     bigframes.testing.assert_series_equal(
         actual,
-        expected,  # .convert_dtypes(dtype_backend="pyarrow"),
+        expected,
         check_index_type=False,
         check_dtype=False,
     )
@@ -158,25 +157,35 @@ def test_timedelta_binary_ops_series_and_literal(
 
 
 @pytest.mark.parametrize(
-    ("op", "col", "literal"),
+    ("op", "col", "literal", "arrow_supported"),
     [
-        (operator.add, "timedelta_col_1", pd.Timedelta(2, "s").as_unit("us")),
-        (operator.sub, "timedelta_col_1", pd.Timedelta(2, "s").as_unit("us")),
-        (operator.truediv, "timedelta_col_1", pd.Timedelta(2, "s").as_unit("us")),
-        (operator.floordiv, "timedelta_col_1", pd.Timedelta(2, "s").as_unit("us")),
-        (operator.truediv, "float_col", pd.Timedelta(2, "s").as_unit("us")),
-        (operator.floordiv, "float_col", pd.Timedelta(2, "s").as_unit("us")),
-        (operator.mul, "timedelta_col_1", 3),
-        (operator.mul, "float_col", pd.Timedelta(1, "s").as_unit("us")),
-        (operator.mod, "timedelta_col_1", pd.Timedelta(7, "s").as_unit("us")),
+        (operator.add, "timedelta_col_1", pd.Timedelta(2, "s").as_unit("us"), True),
+        (operator.sub, "timedelta_col_1", pd.Timedelta(2, "s").as_unit("us"), True),
+        (operator.truediv, "timedelta_col_1", pd.Timedelta(2, "s").as_unit("us"), True),
+        (
+            operator.floordiv,
+            "timedelta_col_1",
+            pd.Timedelta(2, "s").as_unit("us"),
+            True,
+        ),
+        (operator.truediv, "float_col", pd.Timedelta(2, "s").as_unit("us"), True),
+        (operator.floordiv, "float_col", pd.Timedelta(2, "s").as_unit("us"), True),
+        (operator.mul, "timedelta_col_1", 3, True),
+        (operator.mul, "float_col", pd.Timedelta(1, "s").as_unit("us"), False),
+        (operator.mod, "timedelta_col_1", pd.Timedelta(7, "s").as_unit("us"), False),
     ],
 )
-def test_timedelta_binary_ops_literal_and_series(temporal_dfs, op, col, literal):
+def test_timedelta_binary_ops_literal_and_series(
+    temporal_dfs, op, col, literal, arrow_supported
+):
     bf_df, pd_df = temporal_dfs
 
     actual_result = op(literal, bf_df[col]).to_pandas()
 
-    expected_result = op(literal, pd_df[col])
+    if not arrow_supported:
+        expected_result = pd_df[col].map(lambda x: op(literal, x))
+    else:
+        expected_result = op(literal, pd_df[col])
     _assert_series_equal(actual_result, expected_result)
 
 

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

@@ -158,8 +158,6 @@ class Round(Value):
     def dtype(self):
         if self.arg.dtype.is_decimal():
             return self.arg.dtype
-        elif self.digits is None:
-            return dt.int64
         else:
             return dt.double