chore: support comparison, ordering, and filtering for timedeltas (#1387)

* [WIP] support timedelta ordering and filtering

* chore: support comparison, ordering, and filtering for timedeltas

* fix format

* some cleanups

* use operator package for testing

* fix test error

Author: sycai

bigframes/core/rewrite/timedeltas.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import dataclasses
 import functools
 import typing
@@ -27,6 +29,14 @@ class _TypedExpr:
     expr: ex.Expression
     dtype: dtypes.Dtype
 
+    @classmethod
+    def create_op_expr(
+        cls, op: typing.Union[ops.ScalarOp, ops.RowOp], *inputs: _TypedExpr
+    ) -> _TypedExpr:
+        expr = op.as_expr(*tuple(x.expr for x in inputs))  # type: ignore
+        dtype = op.output_type(*tuple(x.dtype for x in inputs))
+        return cls(expr, dtype)
+
 
 def rewrite_timedelta_expressions(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
     """
@@ -38,12 +48,27 @@ def rewrite_timedelta_expressions(root: nodes.BigFrameNode) -> nodes.BigFrameNod
             (_rewrite_expressions(expr, root.schema).expr, column_id)
             for expr, column_id in root.assignments
         )
-        root = nodes.ProjectionNode(root.child, updated_assignments)
+        return nodes.ProjectionNode(root.child, updated_assignments)
+
+    if isinstance(root, nodes.FilterNode):
+        return nodes.FilterNode(
+            root.child, _rewrite_expressions(root.predicate, root.schema).expr
+        )
+
+    if isinstance(root, nodes.OrderByNode):
+        by = tuple(_rewrite_ordering_expr(x, root.schema) for x in root.by)
+        return nodes.OrderByNode(root.child, by)
 
-    # TODO(b/394354614): FilterByNode and OrderNode also contain expressions. Need to update them too.
     return root
 
 
+def _rewrite_ordering_expr(
+    expr: nodes.OrderingExpression, schema: schema.ArraySchema
+) -> nodes.OrderingExpression:
+    by = _rewrite_expressions(expr.scalar_expression, schema).expr
+    return nodes.OrderingExpression(by, expr.direction, expr.na_last)
+
+
 @functools.cache
 def _rewrite_expressions(expr: ex.Expression, schema: schema.ArraySchema) -> _TypedExpr:
     if isinstance(expr, ex.DerefOp):
@@ -78,37 +103,23 @@ def _rewrite_op_expr(
     if isinstance(expr.op, ops.AddOp):
         return _rewrite_add_op(inputs[0], inputs[1])
 
-    input_types = tuple(map(lambda x: x.dtype, inputs))
-    return _TypedExpr(expr, expr.op.output_type(*input_types))
+    return _TypedExpr.create_op_expr(expr.op, *inputs)
 
 
 def _rewrite_sub_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
-    result_op: ops.BinaryOp = ops.sub_op
     if dtypes.is_datetime_like(left.dtype) and dtypes.is_datetime_like(right.dtype):
-        result_op = ops.timestamp_diff_op
+        return _TypedExpr.create_op_expr(ops.timestamp_diff_op, left, right)
 
-    return _TypedExpr(
-        result_op.as_expr(left.expr, right.expr),
-        result_op.output_type(left.dtype, right.dtype),
-    )
+    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 is dtypes.TIMEDELTA_DTYPE:
-        return _TypedExpr(
-            ops.timestamp_add_op.as_expr(left.expr, right.expr),
-            ops.timestamp_add_op.output_type(left.dtype, right.dtype),
-        )
+        return _TypedExpr.create_op_expr(ops.timestamp_add_op, left, right)
 
     if left.dtype is 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(
-            ops.timestamp_add_op.as_expr(right.expr, left.expr),
-            ops.timestamp_add_op.output_type(right.dtype, left.dtype),
-        )
+        return _TypedExpr.create_op_expr(ops.timestamp_add_op, right, left)
 
-    return _TypedExpr(
-        ops.add_op.as_expr(left.expr, right.expr),
-        ops.add_op.output_type(left.dtype, right.dtype),
-    )
+    return _TypedExpr.create_op_expr(ops.add_op, left, right)

bigframes/dtypes.py

@@ -358,7 +358,7 @@ def is_comparable(type_: ExpressionType) -> bool:
 
 def is_orderable(type_: ExpressionType) -> bool:
     # On BQ side, ARRAY, STRUCT, GEOGRAPHY, JSON are not orderable
-    return type_ in _ORDERABLE_SIMPLE_TYPES
+    return type_ in _ORDERABLE_SIMPLE_TYPES or type_ is TIMEDELTA_DTYPE
 
 
 _CLUSTERABLE_SIMPLE_TYPES = set(

tests/system/small/operations/test_timedeltas.py

@@ -14,6 +14,7 @@
 
 
 import datetime
+import operator
 
 import numpy as np
 import pandas as pd
@@ -28,12 +29,23 @@ def temporal_dfs(session):
             "datetime_col": [
                 pd.Timestamp("2025-02-01 01:00:01"),
                 pd.Timestamp("2019-01-02 02:00:00"),
+                pd.Timestamp("1997-01-01 19:00:00"),
             ],
             "timestamp_col": [
                 pd.Timestamp("2023-01-01 01:00:01", tz="UTC"),
                 pd.Timestamp("2024-01-02 02:00:00", tz="UTC"),
+                pd.Timestamp("2005-03-05 02:00:00", tz="UTC"),
+            ],
+            "timedelta_col_1": [
+                pd.Timedelta(3, "s"),
+                pd.Timedelta(-4, "d"),
+                pd.Timedelta(5, "h"),
+            ],
+            "timedelta_col_2": [
+                pd.Timedelta(2, "s"),
+                pd.Timedelta(-4, "d"),
+                pd.Timedelta(6, "h"),
             ],
-            "timedelta_col": [pd.Timedelta(3, "s"), pd.Timedelta(-4, "d")],
         }
     )
 
@@ -53,10 +65,10 @@ def test_timestamp_add__ts_series_plus_td_series(temporal_dfs, column, pd_dtype)
     bf_df, pd_df = temporal_dfs
 
     actual_result = (
-        (bf_df[column] + bf_df["timedelta_col"]).to_pandas().astype(pd_dtype)
+        (bf_df[column] + bf_df["timedelta_col_1"]).to_pandas().astype(pd_dtype)
     )
 
-    expected_result = pd_df[column] + pd_df["timedelta_col"]
+    expected_result = pd_df[column] + pd_df["timedelta_col_1"]
     pandas.testing.assert_series_equal(
         actual_result, expected_result, check_index_type=False
     )
@@ -94,10 +106,10 @@ def test_timestamp_add__td_series_plus_ts_series(temporal_dfs, column, pd_dtype)
     bf_df, pd_df = temporal_dfs
 
     actual_result = (
-        (bf_df["timedelta_col"] + bf_df[column]).to_pandas().astype(pd_dtype)
+        (bf_df["timedelta_col_1"] + bf_df[column]).to_pandas().astype(pd_dtype)
     )
 
-    expected_result = pd_df["timedelta_col"] + pd_df[column]
+    expected_result = pd_df["timedelta_col_1"] + pd_df[column]
     pandas.testing.assert_series_equal(
         actual_result, expected_result, check_index_type=False
     )
@@ -120,10 +132,10 @@ def test_timestamp_add__ts_literal_plus_td_series(temporal_dfs):
     timestamp = pd.Timestamp("2025-01-01", tz="UTC")
 
     actual_result = (
-        (timestamp + bf_df["timedelta_col"]).to_pandas().astype("datetime64[ns, UTC]")
+        (timestamp + bf_df["timedelta_col_1"]).to_pandas().astype("datetime64[ns, UTC]")
     )
 
-    expected_result = timestamp + pd_df["timedelta_col"]
+    expected_result = timestamp + pd_df["timedelta_col_1"]
     pandas.testing.assert_series_equal(
         actual_result, expected_result, check_index_type=False
     )
@@ -140,10 +152,10 @@ def test_timestamp_add_with_numpy_op(temporal_dfs, column, pd_dtype):
     bf_df, pd_df = temporal_dfs
 
     actual_result = (
-        np.add(bf_df[column], bf_df["timedelta_col"]).to_pandas().astype(pd_dtype)
+        np.add(bf_df[column], bf_df["timedelta_col_1"]).to_pandas().astype(pd_dtype)
     )
 
-    expected_result = np.add(pd_df[column], pd_df["timedelta_col"])
+    expected_result = np.add(pd_df[column], pd_df["timedelta_col_1"])
     pandas.testing.assert_series_equal(
         actual_result, expected_result, check_index_type=False
     )
@@ -164,3 +176,105 @@ def test_timestamp_add_dataframes(temporal_dfs):
     pandas.testing.assert_frame_equal(
         actual_result, expected_result, check_index_type=False
     )
+
+
+@pytest.mark.parametrize(
+    "compare_func",
+    [
+        pytest.param(operator.gt, id="gt"),
+        pytest.param(operator.ge, id="ge"),
+        pytest.param(operator.eq, id="eq"),
+        pytest.param(operator.ne, id="ne"),
+        pytest.param(operator.lt, id="lt"),
+        pytest.param(operator.le, id="le"),
+    ],
+)
+def test_timedelta_series_comparison(temporal_dfs, compare_func):
+    bf_df, pd_df = temporal_dfs
+
+    actual_result = compare_func(
+        bf_df["timedelta_col_1"], bf_df["timedelta_col_2"]
+    ).to_pandas()
+
+    expected_result = compare_func(
+        pd_df["timedelta_col_1"], pd_df["timedelta_col_2"]
+    ).astype("boolean")
+    pandas.testing.assert_series_equal(
+        actual_result, expected_result, check_index_type=False
+    )
+
+
+@pytest.mark.parametrize(
+    "compare_func",
+    [
+        pytest.param(operator.gt, id="gt"),
+        pytest.param(operator.ge, id="ge"),
+        pytest.param(operator.eq, id="eq"),
+        pytest.param(operator.ne, id="ne"),
+        pytest.param(operator.lt, id="lt"),
+        pytest.param(operator.le, id="le"),
+    ],
+)
+def test_timedelta_series_and_literal_comparison(temporal_dfs, compare_func):
+    bf_df, pd_df = temporal_dfs
+    literal = pd.Timedelta(3, "s")
+
+    actual_result = compare_func(literal, bf_df["timedelta_col_2"]).to_pandas()
+
+    expected_result = compare_func(literal, pd_df["timedelta_col_2"]).astype("boolean")
+    pandas.testing.assert_series_equal(
+        actual_result, expected_result, check_index_type=False
+    )
+
+
+def test_timedelta_filtering(session):
+    pd_series = pd.Series(
+        [
+            pd.Timestamp("2025-01-01 01:00:00"),
+            pd.Timestamp("2025-01-01 02:00:00"),
+            pd.Timestamp("2025-01-01 03:00:00"),
+        ]
+    )
+    bf_series = session.read_pandas(pd_series)
+    timestamp = pd.Timestamp("2025-01-01, 00:00:01")
+
+    actual_result = (
+        bf_series[((bf_series - timestamp) > pd.Timedelta(1, "h"))]
+        .to_pandas()
+        .astype("<M8[ns]")
+    )
+
+    expected_result = pd_series[(pd_series - timestamp) > pd.Timedelta(1, "h")]
+    pandas.testing.assert_series_equal(
+        actual_result, expected_result, check_index_type=False
+    )
+
+
+def test_timedelta_ordering(session):
+    pd_df = pd.DataFrame(
+        {
+            "col_1": [
+                pd.Timestamp("2025-01-01 01:00:00"),
+                pd.Timestamp("2025-01-01 02:00:00"),
+                pd.Timestamp("2025-01-01 03:00:00"),
+            ],
+            "col_2": [
+                pd.Timestamp("2025-01-01 01:00:02"),
+                pd.Timestamp("2025-01-01 02:00:01"),
+                pd.Timestamp("2025-01-01 02:59:59"),
+            ],
+        }
+    )
+    bf_df = session.read_pandas(pd_df)
+
+    actual_result = (
+        (bf_df["col_2"] - bf_df["col_1"])
+        .sort_values()
+        .to_pandas()
+        .astype("timedelta64[ns]")
+    )
+
+    expected_result = (pd_df["col_2"] - pd_df["col_1"]).sort_values()
+    pandas.testing.assert_series_equal(
+        actual_result, expected_result, check_index_type=False
+    )