googleapis · sycai · Feb 19, 2025 · Feb 13, 2025 · Feb 13, 2025 · Feb 14, 2025
@@ -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
@@ -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",

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


@@ -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"


@@ -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()

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