From 3eb28d2dcc16ebe999811e4ee42d53d42881a352 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Wed, 7 May 2025 20:18:27 +0000 Subject: [PATCH] feat: support () operator between timedeltas --- bigframes/operations/numeric_ops.py | 24 +++++++++++++++++-- .../small/operations/test_timedeltas.py | 3 +++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/bigframes/operations/numeric_ops.py b/bigframes/operations/numeric_ops.py index d06d6eb336..9d6749a169 100644 --- a/bigframes/operations/numeric_ops.py +++ b/bigframes/operations/numeric_ops.py @@ -260,9 +260,29 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT floordiv_op = FloorDivOp() -pow_op = base_ops.create_binary_op(name="pow", type_signature=op_typing.BINARY_NUMERIC) -mod_op = base_ops.create_binary_op(name="mod", type_signature=op_typing.BINARY_NUMERIC) +@dataclasses.dataclass(frozen=True) +class ModOp(base_ops.BinaryOp): + name: typing.ClassVar[str] = "mod" + + def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType: + left_type = input_types[0] + right_type = input_types[1] + + if left_type == dtypes.TIMEDELTA_DTYPE and right_type == dtypes.TIMEDELTA_DTYPE: + return dtypes.TIMEDELTA_DTYPE + + if (left_type is None or dtypes.is_numeric(left_type)) and ( + right_type is None or dtypes.is_numeric(right_type) + ): + return dtypes.coerce_to_common(left_type, right_type) + + raise TypeError(f"Cannot mod dtypes {left_type} and {right_type}") + + +mod_op = ModOp() + +pow_op = base_ops.create_binary_op(name="pow", type_signature=op_typing.BINARY_NUMERIC) arctan2_op = base_ops.create_binary_op( name="arctan2", type_signature=op_typing.BINARY_REAL_NUMERIC diff --git a/tests/system/small/operations/test_timedeltas.py b/tests/system/small/operations/test_timedeltas.py index d6b32a3508..18c88db8eb 100644 --- a/tests/system/small/operations/test_timedeltas.py +++ b/tests/system/small/operations/test_timedeltas.py @@ -98,6 +98,7 @@ def _assert_series_equal(actual: pd.Series, expected: pd.Series): (operator.floordiv, "timedelta_col_1", "float_col"), (operator.mul, "timedelta_col_1", "float_col"), (operator.mul, "float_col", "timedelta_col_1"), + (operator.mod, "timedelta_col_1", "timedelta_col_2"), ], ) def test_timedelta_binary_ops_between_series(temporal_dfs, op, col_1, col_2): @@ -120,6 +121,7 @@ def test_timedelta_binary_ops_between_series(temporal_dfs, op, col_1, col_2): (operator.floordiv, "timedelta_col_1", 3), (operator.mul, "timedelta_col_1", 3), (operator.mul, "float_col", pd.Timedelta(1, "s")), + (operator.mod, "timedelta_col_1", pd.Timedelta(7, "s")), ], ) def test_timedelta_binary_ops_series_and_literal(temporal_dfs, op, col, literal): @@ -142,6 +144,7 @@ def test_timedelta_binary_ops_series_and_literal(temporal_dfs, op, col, literal) (operator.floordiv, "float_col", pd.Timedelta(2, "s")), (operator.mul, "timedelta_col_1", 3), (operator.mul, "float_col", pd.Timedelta(1, "s")), + (operator.mod, "timedelta_col_1", pd.Timedelta(7, "s")), ], ) def test_timedelta_binary_ops_literal_and_series(temporal_dfs, op, col, literal):