[WIP] Log __init__ calls of timedelta-related ops

bigframes/core/compile/scalar_op_compiler.py

@@ -725,32 +725,32 @@ def unix_millis_op_impl(x: ibis_types.TimestampValue):
     return unix_millis(x)
 
 
-@scalar_op_compiler.register_binary_op(ops.timestamp_diff_op)
+@scalar_op_compiler.register_binary_op(ops.TimestampDiffOp)
 def timestamp_diff_op_impl(x: ibis_types.TimestampValue, y: ibis_types.TimestampValue):
     return x.delta(y, "microsecond")
 
 
-@scalar_op_compiler.register_binary_op(ops.timestamp_add_op)
+@scalar_op_compiler.register_binary_op(ops.TimestampAddOp)
 def timestamp_add_op_impl(x: ibis_types.TimestampValue, y: ibis_types.IntegerValue):
     return x + y.to_interval("us")
 
 
-@scalar_op_compiler.register_binary_op(ops.timestamp_sub_op)
+@scalar_op_compiler.register_binary_op(ops.TimestampSubOp)
 def timestamp_sub_op_impl(x: ibis_types.TimestampValue, y: ibis_types.IntegerValue):
     return x - y.to_interval("us")
 
 
-@scalar_op_compiler.register_binary_op(ops.date_diff_op)
+@scalar_op_compiler.register_binary_op(ops.DateDiffOp)
 def date_diff_op_impl(x: ibis_types.DateValue, y: ibis_types.DateValue):
     return x.delta(y, "day") * int(UNIT_TO_US_CONVERSION_FACTORS["d"])  # type: ignore
 
 
-@scalar_op_compiler.register_binary_op(ops.date_add_op)
+@scalar_op_compiler.register_binary_op(ops.DateAddOp)
 def date_add_op_impl(x: ibis_types.DateValue, y: ibis_types.IntegerValue):
     return x.cast("timestamp") + y.to_interval("us")  # type: ignore
 
 
-@scalar_op_compiler.register_binary_op(ops.date_sub_op)
+@scalar_op_compiler.register_binary_op(ops.DateSubOp)
 def date_sub_op_impl(x: ibis_types.DateValue, y: ibis_types.IntegerValue):
     return x.cast("timestamp") - y.to_interval("us")  # type: ignore
 
@@ -1208,7 +1208,7 @@ def to_timedelta_op_impl(x: ibis_types.Value, op: ops.ToTimedeltaOp):
     ).floor()
 
 
-@scalar_op_compiler.register_unary_op(ops.timedelta_floor_op)
+@scalar_op_compiler.register_unary_op(ops.TimedeltaFloorOp)
 def timedelta_floor_op_impl(x: ibis_types.NumericValue):
     return x.floor()
 

bigframes/core/rewrite/timedeltas.py

@@ -28,7 +28,7 @@
 @dataclasses.dataclass
 class _TypedExpr:
     expr: ex.Expression
-    dtype: dtypes.Dtype
+    dtype: dtypes.ExpressionType
 
     @classmethod
     def create_op_expr(
@@ -146,36 +146,36 @@ def _rewrite_op_expr(
 
 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)
+        return _TypedExpr.create_op_expr(ops.TimestampDiffOp(), left, right)
 
     if dtypes.is_datetime_like(left.dtype) and right.dtype == dtypes.TIMEDELTA_DTYPE:
-        return _TypedExpr.create_op_expr(ops.timestamp_sub_op, left, right)
+        return _TypedExpr.create_op_expr(ops.TimestampSubOp(), left, right)
 
     if left.dtype == dtypes.DATE_DTYPE and right.dtype == dtypes.DATE_DTYPE:
-        return _TypedExpr.create_op_expr(ops.date_diff_op, left, right)
+        return _TypedExpr.create_op_expr(ops.DateDiffOp(), left, right)
 
     if left.dtype == dtypes.DATE_DTYPE and right.dtype == dtypes.TIMEDELTA_DTYPE:
-        return _TypedExpr.create_op_expr(ops.date_sub_op, left, right)
+        return _TypedExpr.create_op_expr(ops.DateSubOp(), left, right)
 
     return _TypedExpr.create_op_expr(ops.sub_op, left, right)
 
 
 def _rewrite_add_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
     if dtypes.is_datetime_like(left.dtype) and right.dtype == dtypes.TIMEDELTA_DTYPE:
-        return _TypedExpr.create_op_expr(ops.timestamp_add_op, left, right)
+        return _TypedExpr.create_op_expr(ops.TimestampAddOp(), left, right)
 
     if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_datetime_like(right.dtype):
         # Re-arrange operands such that timestamp is always on the left and timedelta is
         # always on the right.
-        return _TypedExpr.create_op_expr(ops.timestamp_add_op, right, left)
+        return _TypedExpr.create_op_expr(ops.TimestampAddOp(), right, left)
 
     if left.dtype == dtypes.DATE_DTYPE and right.dtype == dtypes.TIMEDELTA_DTYPE:
-        return _TypedExpr.create_op_expr(ops.date_add_op, left, right)
+        return _TypedExpr.create_op_expr(ops.DateAddOp(), left, right)
 
     if left.dtype == dtypes.TIMEDELTA_DTYPE and right.dtype == dtypes.DATE_DTYPE:
         # Re-arrange operands such that date is always on the left and timedelta is
         # always on the right.
-        return _TypedExpr.create_op_expr(ops.date_add_op, right, left)
+        return _TypedExpr.create_op_expr(ops.DateAddOp(), right, left)
 
     return _TypedExpr.create_op_expr(ops.add_op, left, right)
 
@@ -184,9 +184,9 @@ def _rewrite_mul_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
     result = _TypedExpr.create_op_expr(ops.mul_op, left, right)
 
     if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
-        return _TypedExpr.create_op_expr(ops.timedelta_floor_op, result)
+        return _TypedExpr.create_op_expr(ops.TimedeltaFloorOp(), result)
     if dtypes.is_numeric(left.dtype) and right.dtype == dtypes.TIMEDELTA_DTYPE:
-        return _TypedExpr.create_op_expr(ops.timedelta_floor_op, result)
+        return _TypedExpr.create_op_expr(ops.TimedeltaFloorOp(), result)
 
     return result
 
@@ -195,7 +195,7 @@ def _rewrite_div_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
     result = _TypedExpr.create_op_expr(ops.div_op, left, right)
 
     if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
-        return _TypedExpr.create_op_expr(ops.timedelta_floor_op, result)
+        return _TypedExpr.create_op_expr(ops.TimedeltaFloorOp(), result)
 
     return result
 
@@ -204,7 +204,7 @@ def _rewrite_floordiv_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
     result = _TypedExpr.create_op_expr(ops.floordiv_op, left, right)
 
     if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
-        return _TypedExpr.create_op_expr(ops.timedelta_floor_op, result)
+        return _TypedExpr.create_op_expr(ops.TimedeltaFloorOp(), result)
 
     return result
 

bigframes/operations/__init__.py

@@ -39,7 +39,7 @@
     ne_op,
 )
 from bigframes.operations.date_ops import (
-    date_diff_op,
+    DateDiffOp,
     day_op,
     dayofweek_op,
     month_op,
@@ -50,7 +50,7 @@
     date_op,
     StrftimeOp,
     time_op,
-    timestamp_diff_op,
+    TimestampDiffOp,
     ToDatetimeOp,
     ToTimestampOp,
     UnixMicros,
@@ -188,11 +188,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 (
-    date_add_op,
-    date_sub_op,
-    timedelta_floor_op,
-    timestamp_add_op,
-    timestamp_sub_op,
+    DateAddOp,
+    DateSubOp,
+    TimedeltaFloorOp,
+    TimestampAddOp,
+    TimestampSubOp,
     ToTimedeltaOp,
 )
 
@@ -255,7 +255,7 @@
     "upper_op",
     "ZfillOp",
     # Date ops
-    "date_diff_op",
+    "DateDiffOp",
     "day_op",
     "month_op",
     "year_op",
@@ -267,16 +267,16 @@
     "second_op",
     "normalize_op",
     # Timedelta ops
-    "date_add_op",
-    "date_sub_op",
-    "timedelta_floor_op",
-    "timestamp_add_op",
-    "timestamp_sub_op",
+    "DateAddOp",
+    "DateSubOp",
+    "TimedeltaFloorOp",
+    "TimestampAddOp",
+    "TimestampSubOp",
     "ToTimedeltaOp",
     # Datetime ops
     "date_op",
     "time_op",
-    "timestamp_diff_op",
+    "TimestampDiffOp",
     "ToDatetimeOp",
     "ToTimestampOp",
     "StrftimeOp",

bigframes/operations/date_ops.py

@@ -16,6 +16,7 @@
 import typing
 
 from bigframes import dtypes
+from bigframes.core import log_adapter
 from bigframes.operations import base_ops
 import bigframes.operations.type as op_typing
 
@@ -45,6 +46,7 @@
 )
 
 
+@log_adapter.class_logger(include_internal_calls=True)
 @dataclasses.dataclass(frozen=True)
 class DateDiffOp(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "date_diff"
@@ -59,6 +61,3 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
             raise TypeError("expected date input")
 
         return dtypes.TIMEDELTA_DTYPE
-
-
-date_diff_op = DateDiffOp()

bigframes/operations/datetime_ops.py

@@ -19,6 +19,7 @@
 import pyarrow as pa
 
 from bigframes import dtypes
+from bigframes.core import log_adapter
 from bigframes.operations import base_ops
 import bigframes.operations.type as op_typing
 
@@ -109,8 +110,9 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         return dtypes.INT_DTYPE
 
 
+@log_adapter.class_logger(include_internal_calls=True)
 @dataclasses.dataclass(frozen=True)
-class TimestampDiff(base_ops.BinaryOp):
+class TimestampDiffOp(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "timestamp_diff"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
@@ -123,6 +125,3 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
             raise TypeError("expected timestamp input")
 
         return dtypes.TIMEDELTA_DTYPE
-
-
-timestamp_diff_op = TimestampDiff()

bigframes/operations/timedelta_ops.py

@@ -17,9 +17,11 @@
 import typing
 
 from bigframes import dtypes
+from bigframes.core import log_adapter
 from bigframes.operations import base_ops
 
 
+@log_adapter.class_logger(include_internal_calls=True)
 @dataclasses.dataclass(frozen=True)
 class ToTimedeltaOp(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "to_timedelta"
@@ -35,6 +37,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         raise TypeError("expected integer or float input")
 
 
+@log_adapter.class_logger(include_internal_calls=True)
 @dataclasses.dataclass(frozen=True)
 class TimedeltaFloorOp(base_ops.UnaryOp):
     """Floors the numeric value to the nearest integer and use it to represent a timedelta.
@@ -51,9 +54,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         raise TypeError(f"unsupported type: {input_type}")
 
 
-timedelta_floor_op = TimedeltaFloorOp()
-
-
+@log_adapter.class_logger(include_internal_calls=True)
 @dataclasses.dataclass(frozen=True)
 class TimestampAddOp(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "timestamp_add"
@@ -76,9 +77,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         )
 
 
-timestamp_add_op = TimestampAddOp()
-
-
+@log_adapter.class_logger(include_internal_calls=True)
 @dataclasses.dataclass(frozen=True)
 class TimestampSubOp(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "timestamp_sub"
@@ -96,9 +95,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         )
 
 
-timestamp_sub_op = TimestampSubOp()
-
-
+@log_adapter.class_logger(include_internal_calls=True)
 @dataclasses.dataclass(frozen=True)
 class DateAddOp(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "date_add"
@@ -122,9 +119,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         )
 
 
-date_add_op = DateAddOp()
-
-
+@log_adapter.class_logger(include_internal_calls=True)
 @dataclasses.dataclass(frozen=True)
 class DateSubOp(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "date_sub"
@@ -140,6 +135,3 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         raise TypeError(
             f"unsupported types for date_sub. left: {input_types[0]} right: {input_types[1]}"
         )
-
-
-date_sub_op = DateSubOp()