feat: (Preview) Support arithmetics between dates and timedeltas (#1413)

* [WIP] Implement date_diff, date_add, and date_sub. Some tests are broken. * fix format * change implementation to emulate PyArrow * fix mypy * fix format * fix tests * fix more tests * bump pyarrow dependency to 15.0.2 * raise pandas version in system-3.12 tests * fix format * restore constraints-3.12 * raise pandas version to 2.2.0 in 3.12 testing env * remove pandas constraints for 3.12 * fix merge error * cast factor to int instead of calling floor()

Author: sycai

bigframes/core/compile/scalar_op_compiler.py

@@ -740,6 +740,21 @@ def timestamp_sub_op_impl(x: ibis_types.TimestampValue, y: ibis_types.IntegerVal
  return x - y.to_interval("us")
 
 
+@scalar_op_compiler.register_binary_op(ops.date_diff_op)
+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)
+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)
+def date_sub_op_impl(x: ibis_types.DateValue, y: ibis_types.IntegerValue):
+ return x.cast("timestamp") - y.to_interval("us") # type: ignore
+
+
 @scalar_op_compiler.register_unary_op(ops.FloorDtOp, pass_op=True)
 def floor_dt_op_impl(x: ibis_types.Value, op: ops.FloorDtOp):
  supported_freqs = ["Y", "Q", "M", "W", "D", "h", "min", "s", "ms", "us", "ns"]

bigframes/core/rewrite/timedeltas.py

@@ -151,6 +151,12 @@ def _rewrite_sub_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
  if dtypes.is_datetime_like(left.dtype) and right.dtype is dtypes.TIMEDELTA_DTYPE:
  return _TypedExpr.create_op_expr(ops.timestamp_sub_op, 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)
+
+ if left.dtype == dtypes.DATE_DTYPE and right.dtype is dtypes.TIMEDELTA_DTYPE:
+ return _TypedExpr.create_op_expr(ops.date_sub_op, left, right)
+
  return _TypedExpr.create_op_expr(ops.sub_op, left, right)
 
 
@@ -163,6 +169,14 @@ def _rewrite_add_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
  # always on the right.
  return _TypedExpr.create_op_expr(ops.timestamp_add_op, right, left)
 
+ if left.dtype == dtypes.DATE_DTYPE and right.dtype is dtypes.TIMEDELTA_DTYPE:
+ return _TypedExpr.create_op_expr(ops.date_add_op, left, right)
+
+ if left.dtype is 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.add_op, left, right)
 
 

bigframes/operations/__init__.py

@@ -39,6 +39,7 @@
  ne_op,
 )
 from bigframes.operations.date_ops import (
+ date_diff_op,
  day_op,
  dayofweek_op,
  month_op,
@@ -184,6 +185,8 @@
 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,
@@ -249,6 +252,7 @@
  "upper_op",
  "ZfillOp",
  # Date ops
+ "date_diff_op",
  "day_op",
  "month_op",
  "year_op",
@@ -260,6 +264,8 @@
  "second_op",
  "normalize_op",
  # Timedelta ops
+ "date_add_op",
+ "date_sub_op",
  "timedelta_floor_op",
  "timestamp_add_op",
  "timestamp_sub_op",

bigframes/operations/date_ops.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import dataclasses
+import typing
+
+from bigframes import dtypes
 from bigframes.operations import base_ops
 import bigframes.operations.type as op_typing
 
@@ -39,3 +43,22 @@
  name="quarter",
  type_signature=op_typing.DATELIKE_ACCESSOR,
 )
+
+
+@dataclasses.dataclass(frozen=True)
+class DateDiffOp(base_ops.BinaryOp):
+ name: typing.ClassVar[str] = "date_diff"
+
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ if input_types[0] is not input_types[1]:
+ raise TypeError(
+ f"two inputs have different types. left: {input_types[0]}, right: {input_types[1]}"
+ )
+
+ if input_types[0] != dtypes.DATE_DTYPE:
+ raise TypeError("expected date input")
+
+ return dtypes.TIMEDELTA_DTYPE
+
+
+date_diff_op = DateDiffOp()

bigframes/operations/numeric_ops.py

@@ -123,12 +123,18 @@ def output_type(self, *input_types):
  # String addition
  return input_types[0]
 
- # Timestamp addition.
+ # Temporal addition.
  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 dtypes.is_datetime_like(right_type):
  return right_type
 
+ if left_type == dtypes.DATE_DTYPE and right_type == dtypes.TIMEDELTA_DTYPE:
+ return dtypes.DATETIME_DTYPE
+
+ if left_type == dtypes.TIMEDELTA_DTYPE and right_type == dtypes.DATE_DTYPE:
+ return dtypes.DATETIME_DTYPE
+
  if left_type is dtypes.TIMEDELTA_DTYPE and right_type is dtypes.TIMEDELTA_DTYPE:
  return dtypes.TIMEDELTA_DTYPE
 
@@ -155,9 +161,15 @@ def output_type(self, *input_types):
  if dtypes.is_datetime_like(left_type) and dtypes.is_datetime_like(right_type):
  return dtypes.TIMEDELTA_DTYPE
 
+ if left_type == dtypes.DATE_DTYPE and right_type == dtypes.DATE_DTYPE:
+ return dtypes.TIMEDELTA_DTYPE
+
  if dtypes.is_datetime_like(left_type) and right_type is dtypes.TIMEDELTA_DTYPE:
  return left_type
 
+ if left_type == dtypes.DATE_DTYPE and right_type == dtypes.TIMEDELTA_DTYPE:
+ return dtypes.DATETIME_DTYPE
+
  if left_type is dtypes.TIMEDELTA_DTYPE and right_type is dtypes.TIMEDELTA_DTYPE:
  return dtypes.TIMEDELTA_DTYPE
 

bigframes/operations/timedelta_ops.py

@@ -79,6 +79,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 timestamp_add_op = TimestampAddOp()
 
 
+@dataclasses.dataclass(frozen=True)
 class TimestampSubOp(base_ops.BinaryOp):
  name: typing.ClassVar[str] = "timestamp_sub"
 
@@ -96,3 +97,49 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 
 
 timestamp_sub_op = TimestampSubOp()
+
+
+@dataclasses.dataclass(frozen=True)
+class DateAddOp(base_ops.BinaryOp):
+ name: typing.ClassVar[str] = "date_add"
+
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ # date + timedelta => timestamp without timezone
+ if (
+ input_types[0] == dtypes.DATE_DTYPE
+ and input_types[1] == dtypes.TIMEDELTA_DTYPE
+ ):
+ return dtypes.DATETIME_DTYPE
+ # timedelta + date => timestamp without timezone
+ if (
+ input_types[0] == dtypes.TIMEDELTA_DTYPE
+ and input_types[1] == dtypes.DATE_DTYPE
+ ):
+ return dtypes.DATETIME_DTYPE
+
+ raise TypeError(
+ f"unsupported types for date_add. left: {input_types[0]} right: {input_types[1]}"
+ )
+
+
+date_add_op = DateAddOp()
+
+
+@dataclasses.dataclass(frozen=True)
+class DateSubOp(base_ops.BinaryOp):
+ name: typing.ClassVar[str] = "date_sub"
+
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ # date - timedelta => timestamp without timezone
+ if (
+ input_types[0] == dtypes.DATE_DTYPE
+ and input_types[1] == dtypes.TIMEDELTA_DTYPE
+ ):
+ return dtypes.DATETIME_DTYPE
+
+ raise TypeError(
+ f"unsupported types for date_sub. left: {input_types[0]} right: {input_types[1]}"
+ )
+
+
+date_sub_op = DateSubOp()

setup.py

@@ -52,7 +52,7 @@
  "numpy >=1.24.0",
  "pandas >=1.5.3",
  "pandas-gbq >=0.26.0",
- "pyarrow >=10.0.1",
+ "pyarrow >=15.0.2",
  "pydata-google-auth >=1.8.2",
  "requests >=2.27.1",
  "sqlglot >=23.6.3",

testing/constraints-3.9.txt

@@ -16,7 +16,7 @@ jellyfish==0.8.9
 numpy==1.24.0
 pandas==1.5.3
 pandas-gbq==0.26.0
-pyarrow==10.0.1
+pyarrow==15.0.2
 pydata-google-auth==1.8.2
 requests==2.27.1
 scikit-learn==1.2.2

tests/system/small/operations/test_dates.py

@@ -12,11 +12,56 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+import datetime
+
+import pandas as pd
 import pandas.testing
 
 from bigframes import dtypes
 
 
+def test_date_diff_between_series(session):
+ pd_df = pd.DataFrame(
+ {
+ "col_1": [datetime.date(2025, 1, 2), datetime.date(2025, 2, 1)],
+ "col_2": [datetime.date(2024, 1, 2), datetime.date(2026, 1, 30)],
+ }
+ ).astype(dtypes.DATE_DTYPE)
+ bf_df = session.read_pandas(pd_df)
+
+ actual_result = (bf_df["col_1"] - bf_df["col_2"]).to_pandas()
+
+ expected_result = (pd_df["col_1"] - pd_df["col_2"]).astype(dtypes.TIMEDELTA_DTYPE)
+ pandas.testing.assert_series_equal(
+ actual_result, expected_result, check_index_type=False
+ )
+
+
+def test_date_diff_literal_sub_series(scalars_dfs):
+ bf_df, pd_df = scalars_dfs
+ literal = datetime.date(2030, 5, 20)
+
+ actual_result = (literal - bf_df["date_col"]).to_pandas()
+
+ expected_result = (literal - pd_df["date_col"]).astype(dtypes.TIMEDELTA_DTYPE)
+ pandas.testing.assert_series_equal(
+ actual_result, expected_result, check_index_type=False
+ )
+
+
+def test_date_diff_series_sub_literal(scalars_dfs):
+ bf_df, pd_df = scalars_dfs
+ literal = datetime.date(1980, 5, 20)
+
+ actual_result = (bf_df["date_col"] - literal).to_pandas()
+
+ expected_result = (pd_df["date_col"] - literal).astype(dtypes.TIMEDELTA_DTYPE)
+ pandas.testing.assert_series_equal(
+ actual_result, expected_result, check_index_type=False
+ )
+
+
 def test_date_series_diff_agg(scalars_dfs):
  bf_df, pd_df = scalars_dfs