feat: support () operator between timedeltas (#1702)

Author: sycai

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

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):