diff --git a/bigframes/core/compile/scalar_op_compiler.py b/bigframes/core/compile/scalar_op_compiler.py index 3e5f10eca4..d5ce6e9e09 100644 --- a/bigframes/core/compile/scalar_op_compiler.py +++ b/bigframes/core/compile/scalar_op_compiler.py @@ -747,6 +747,11 @@ def timestamp_add_op_impl(x: ibis_types.TimestampValue, y: ibis_types.IntegerVal return x + y.to_interval("us") +@scalar_op_compiler.register_binary_op(ops.timestamp_sub_op) +def timestamp_sub_op_impl(x: ibis_types.TimestampValue, y: ibis_types.IntegerValue): + return x - y.to_interval("us") + + @scalar_op_compiler.register_unary_op(ops.FloorDtOp, pass_op=True) def floor_dt_op_impl(x: ibis_types.Value, op: ops.FloorDtOp): supported_freqs = ["Y", "Q", "M", "W", "D", "h", "min", "s", "ms", "us", "ns"] diff --git a/bigframes/core/rewrite/timedeltas.py b/bigframes/core/rewrite/timedeltas.py index 990aca1f18..db3a426635 100644 --- a/bigframes/core/rewrite/timedeltas.py +++ b/bigframes/core/rewrite/timedeltas.py @@ -110,6 +110,9 @@ def _rewrite_sub_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr: if dtypes.is_datetime_like(left.dtype) and dtypes.is_datetime_like(right.dtype): return _TypedExpr.create_op_expr(ops.timestamp_diff_op, left, right) + if dtypes.is_datetime_like(left.dtype) and right.dtype is dtypes.TIMEDELTA_DTYPE: + return _TypedExpr.create_op_expr(ops.timestamp_sub_op, left, right) + return _TypedExpr.create_op_expr(ops.sub_op, left, right) diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index 88406317fe..21a1171ddc 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -178,7 +178,11 @@ ) from bigframes.operations.struct_ops import StructFieldOp, StructOp from bigframes.operations.time_ops import hour_op, minute_op, normalize_op, second_op -from bigframes.operations.timedelta_ops import timestamp_add_op, ToTimedeltaOp +from bigframes.operations.timedelta_ops import ( + timestamp_add_op, + timestamp_sub_op, + ToTimedeltaOp, +) __all__ = [ # Base ops @@ -251,6 +255,7 @@ "normalize_op", # Timedelta ops "timestamp_add_op", + "timestamp_sub_op", "ToTimedeltaOp", # Datetime ops "date_op", diff --git a/bigframes/operations/numeric_ops.py b/bigframes/operations/numeric_ops.py index 5183e5c4c5..61544984fb 100644 --- a/bigframes/operations/numeric_ops.py +++ b/bigframes/operations/numeric_ops.py @@ -151,6 +151,9 @@ def output_type(self, *input_types): if dtypes.is_datetime_like(left_type) and dtypes.is_datetime_like(right_type): return dtypes.TIMEDELTA_DTYPE + if dtypes.is_datetime_like(left_type) and right_type is dtypes.TIMEDELTA_DTYPE: + return left_type + raise TypeError(f"Cannot subtract dtypes {left_type} and {right_type}") diff --git a/bigframes/operations/timedelta_ops.py b/bigframes/operations/timedelta_ops.py index 69e054fa5c..3d3c3bfeeb 100644 --- a/bigframes/operations/timedelta_ops.py +++ b/bigframes/operations/timedelta_ops.py @@ -54,3 +54,23 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT timestamp_add_op = TimestampAdd() + + +@dataclasses.dataclass(frozen=True) +class TimestampSub(base_ops.BinaryOp): + name: typing.ClassVar[str] = "timestamp_sub" + + def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType: + # timestamp - timedelta => timestamp + if ( + dtypes.is_datetime_like(input_types[0]) + and input_types[1] is dtypes.TIMEDELTA_DTYPE + ): + return input_types[0] + + raise TypeError( + f"unsupported types for timestamp_sub. left: {input_types[0]} right: {input_types[1]}" + ) + + +timestamp_sub_op = TimestampSub() diff --git a/tests/system/small/operations/test_timedeltas.py b/tests/system/small/operations/test_timedeltas.py index fe779a8524..9dc889f619 100644 --- a/tests/system/small/operations/test_timedeltas.py +++ b/tests/system/small/operations/test_timedeltas.py @@ -178,6 +178,96 @@ def test_timestamp_add_dataframes(temporal_dfs): ) +@pytest.mark.parametrize( + ("column", "pd_dtype"), + [ + ("datetime_col", "