feat: support add, sub, mult, div, and more between timedeltas (#1396)

bigframes/core/rewrite/timedeltas.py

@@ -103,6 +103,17 @@ def _rewrite_op_expr(
  if isinstance(expr.op, ops.AddOp):
  return _rewrite_add_op(inputs[0], inputs[1])
 
+ if isinstance(expr.op, ops.MulOp):
+ return _rewrite_mul_op(inputs[0], inputs[1])
+
+ if isinstance(expr.op, ops.DivOp):
+ return _rewrite_div_op(inputs[0], inputs[1])
+
+ if isinstance(expr.op, ops.FloorDivOp):
+ # We need to re-write floor div because for numerics: int // float => float
+ # but for timedeltas: int(timedelta) // float => int(timedelta)
+ return _rewrite_floordiv_op(inputs[0], inputs[1])
+
  return _TypedExpr.create_op_expr(expr.op, *inputs)
 
 
@@ -126,3 +137,32 @@ def _rewrite_add_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
  return _TypedExpr.create_op_expr(ops.timestamp_add_op, right, left)
 
  return _TypedExpr.create_op_expr(ops.add_op, left, right)
+
+
+def _rewrite_mul_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
+ result = _TypedExpr.create_op_expr(ops.mul_op, left, right)
+
+ if left.dtype is dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
+ return _TypedExpr.create_op_expr(ops.ToTimedeltaOp("us"), result)
+ if dtypes.is_numeric(left.dtype) and right.dtype is dtypes.TIMEDELTA_DTYPE:
+ return _TypedExpr.create_op_expr(ops.ToTimedeltaOp("us"), result)
+
+ return result
+
+
+def _rewrite_div_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
+ result = _TypedExpr.create_op_expr(ops.div_op, left, right)
+
+ if left.dtype is dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
+ return _TypedExpr.create_op_expr(ops.ToTimedeltaOp("us"), 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 is dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
+ return _TypedExpr.create_op_expr(ops.ToTimedeltaOp("us"), result)
+
+ return result

bigframes/operations/__init__.py

@@ -115,15 +115,18 @@
  cos_op,
  cosh_op,
  div_op,
+ DivOp,
  exp_op,
  expm1_op,
  floor_op,
  floordiv_op,
+ FloorDivOp,
  ln_op,
  log1p_op,
  log10_op,
  mod_op,
  mul_op,
+ MulOp,
  neg_op,
  pos_op,
  pow_op,
@@ -282,15 +285,18 @@
  "cos_op",
  "cosh_op",
  "div_op",
+ "DivOp",
  "exp_op",
  "expm1_op",
  "floor_op",
  "floordiv_op",
+ "FloorDivOp",
  "ln_op",
  "log1p_op",
  "log10_op",
  "mod_op",
  "mul_op",
+ "MulOp",
  "neg_op",
  "pos_op",
  "pow_op",

bigframes/operations/numeric_ops.py

@@ -75,11 +75,17 @@
  name="ceil", type_signature=op_typing.UNARY_REAL_NUMERIC
 )
 
-abs_op = base_ops.create_unary_op(name="abs", type_signature=op_typing.UNARY_NUMERIC)
+abs_op = base_ops.create_unary_op(
+ name="abs", type_signature=op_typing.UNARY_NUMERIC_AND_TIMEDELTA
+)
 
-pos_op = base_ops.create_unary_op(name="pos", type_signature=op_typing.UNARY_NUMERIC)
+pos_op = base_ops.create_unary_op(
+ name="pos", type_signature=op_typing.UNARY_NUMERIC_AND_TIMEDELTA
+)
 
-neg_op = base_ops.create_unary_op(name="neg", type_signature=op_typing.UNARY_NUMERIC)
+neg_op = base_ops.create_unary_op(
+ name="neg", type_signature=op_typing.UNARY_NUMERIC_AND_TIMEDELTA
+)
 
 exp_op = base_ops.create_unary_op(
  name="exp", type_signature=op_typing.UNARY_REAL_NUMERIC
@@ -123,6 +129,9 @@ def output_type(self, *input_types):
  if left_type is dtypes.TIMEDELTA_DTYPE and dtypes.is_datetime_like(right_type):
  return right_type
 
+ if left_type is dtypes.TIMEDELTA_DTYPE and right_type is 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)
  ):
@@ -142,32 +151,102 @@ class SubOp(base_ops.BinaryOp):
  def output_type(self, *input_types):
  left_type = input_types[0]
  right_type = input_types[1]
- if (left_type is None or dtypes.is_numeric(left_type)) and (
- right_type is None or dtypes.is_numeric(right_type)
- ):
- # Numeric subtraction
- return dtypes.coerce_to_common(left_type, right_type)
 
  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
 
+ if left_type is dtypes.TIMEDELTA_DTYPE and right_type is 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)
+ ):
+ # Numeric subtraction
+ return dtypes.coerce_to_common(left_type, right_type)
+
  raise TypeError(f"Cannot subtract dtypes {left_type} and {right_type}")
 
 
 sub_op = SubOp()
 
-mul_op = base_ops.create_binary_op(name="mul", type_signature=op_typing.BINARY_NUMERIC)
 
-div_op = base_ops.create_binary_op(
- name="div", type_signature=op_typing.BINARY_REAL_NUMERIC
-)
+@dataclasses.dataclass(frozen=True)
+class MulOp(base_ops.BinaryOp):
+ name: typing.ClassVar[str] = "mul"
 
-floordiv_op = base_ops.create_binary_op(
- name="floordiv", type_signature=op_typing.BINARY_NUMERIC
-)
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ left_type = input_types[0]
+ right_type = input_types[1]
+
+ if left_type is dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right_type):
+ return dtypes.TIMEDELTA_DTYPE
+ if dtypes.is_numeric(left_type) and right_type is 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 multiply dtypes {left_type} and {right_type}")
+
+
+mul_op = MulOp()
+
+
+@dataclasses.dataclass(frozen=True)
+class DivOp(base_ops.BinaryOp):
+ name: typing.ClassVar[str] = "div"
+
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ left_type = input_types[0]
+ right_type = input_types[1]
+
+ if left_type is dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right_type):
+ return dtypes.TIMEDELTA_DTYPE
+
+ if left_type is dtypes.TIMEDELTA_DTYPE and right_type is dtypes.TIMEDELTA_DTYPE:
+ return dtypes.FLOAT_DTYPE
+
+ if (left_type is None or dtypes.is_numeric(left_type)) and (
+ right_type is None or dtypes.is_numeric(right_type)
+ ):
+ lcd_type = dtypes.coerce_to_common(left_type, right_type)
+ # Real numeric ops produce floats on int input
+ return dtypes.FLOAT_DTYPE if lcd_type == dtypes.INT_DTYPE else lcd_type
+
+ raise TypeError(f"Cannot divide dtypes {left_type} and {right_type}")
+
+
+div_op = DivOp()
+
+
+@dataclasses.dataclass(frozen=True)
+class FloorDivOp(base_ops.BinaryOp):
+ name: typing.ClassVar[str] = "floordiv"
+
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ left_type = input_types[0]
+ right_type = input_types[1]
+
+ if left_type is dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right_type):
+ return dtypes.TIMEDELTA_DTYPE
+
+ if left_type is dtypes.TIMEDELTA_DTYPE and right_type is dtypes.TIMEDELTA_DTYPE:
+ return dtypes.INT_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 floor divide dtypes {left_type} and {right_type}")
+
+
+floordiv_op = FloorDivOp()
 
 pow_op = base_ops.create_binary_op(name="pow", type_signature=op_typing.BINARY_NUMERIC)
 

bigframes/operations/timedelta_ops.py

@@ -26,7 +26,11 @@ class ToTimedeltaOp(base_ops.UnaryOp):
  unit: typing.Literal["us", "ms", "s", "m", "h", "d", "W"]
 
  def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
- if input_types[0] in (dtypes.INT_DTYPE, dtypes.FLOAT_DTYPE):
+ if input_types[0] in (
+ dtypes.INT_DTYPE,
+ dtypes.FLOAT_DTYPE,
+ dtypes.TIMEDELTA_DTYPE,
+ ):
  return dtypes.TIMEDELTA_DTYPE
  raise TypeError("expected integer or float input")
 
@@ -56,7 +60,6 @@ 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"
 

bigframes/operations/type.py

@@ -224,6 +224,10 @@ def output_type(
 
 # Common type signatures
 UNARY_NUMERIC = TypePreserving(bigframes.dtypes.is_numeric, description="numeric")
+UNARY_NUMERIC_AND_TIMEDELTA = TypePreserving(
+ lambda x: bigframes.dtypes.is_numeric(x) or x is bigframes.dtypes.TIMEDELTA_DTYPE,
+ description="numeric_and_timedelta",
+)
 UNARY_REAL_NUMERIC = UnaryRealNumeric()
 BINARY_NUMERIC = BinaryNumeric()
 BINARY_REAL_NUMERIC = BinaryRealNumeric()

bigframes/series.py

@@ -964,6 +964,9 @@ def update(self, other: Union[Series, Sequence, Mapping]) -> None:
  )
  self._set_block(result._get_block())
 
+ def __abs__(self) -> Series:
+ return self.abs()
+
  def abs(self) -> Series:
  return self._apply_unary_op(ops.abs_op)