[Data] Compute Expressions-datetime

python/ray/data/expressions.py

@@ -23,6 +23,7 @@
 from ray.util.annotations import DeveloperAPI, PublicAPI
 
 if TYPE_CHECKING:
+    from ray.data.namespace_expressions.dt_namespace import _DatetimeNamespace
     from ray.data.namespace_expressions.list_namespace import _ListNamespace
     from ray.data.namespace_expressions.string_namespace import _StringNamespace
     from ray.data.namespace_expressions.struct_namespace import _StructNamespace
@@ -486,6 +487,13 @@ def struct(self) -> "_StructNamespace":
 
         return _StructNamespace(self)
 
+    @property
+    def dt(self) -> "_DatetimeNamespace":
+        """Access datetime operations for this expression."""
+        from ray.data.namespace_expressions.dt_namespace import _DatetimeNamespace
+
+        return _DatetimeNamespace(self)
+
     def _unalias(self) -> "Expr":
         return self
 
@@ -1061,6 +1069,7 @@ def download(uri_column_name: str) -> DownloadExpr:
     "_ListNamespace",
     "_StringNamespace",
     "_StructNamespace",
+    "_DatetimeNamespace",
 ]
 
 
@@ -1078,4 +1087,8 @@ def __getattr__(name: str):
         from ray.data.namespace_expressions.struct_namespace import _StructNamespace
 
         return _StructNamespace
+    elif name == "_DatetimeNamespace":
+        from ray.data.namespace_expressions.dt_namespace import _DatetimeNamespace
+
+        return _DatetimeNamespace
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

python/ray/data/namespace_expressions/dt_namespace.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Literal
+
+import pyarrow
+import pyarrow.compute as pc
+
+from ray.data.datatype import DataType
+from ray.data.expressions import pyarrow_udf
+
+if TYPE_CHECKING:
+    from ray.data.expressions import Expr, UDFExpr
+
+TemporalUnit = Literal[
+    "year",
+    "quarter",
+    "month",
+    "week",
+    "day",
+    "hour",
+    "minute",
+    "second",
+    "millisecond",
+    "microsecond",
+    "nanosecond",
+]
+
+
+@dataclass
+class _DatetimeNamespace:
+    """Datetime namespace for operations on datetime-typed expression columns."""
+
+    _expr: "Expr"
+
+    def _unary_temporal_int(
+        self, func: Callable[[pyarrow.Array], pyarrow.Array]
+    ) -> "UDFExpr":
+        """Helper for year/month/… that return int32."""
+
+        @pyarrow_udf(return_dtype=DataType.int32())
+        def _udf(arr: pyarrow.Array) -> pyarrow.Array:
+            return func(arr)
+
+        return _udf(self._expr)
+
+    # extractors
+
+    def year(self) -> "UDFExpr":
+        """Extract year component."""
+        return self._unary_temporal_int(pc.year)
+
+    def month(self) -> "UDFExpr":
+        """Extract month component."""
+        return self._unary_temporal_int(pc.month)
+
+    def day(self) -> "UDFExpr":
+        """Extract day component."""
+        return self._unary_temporal_int(pc.day)
+
+    def hour(self) -> "UDFExpr":
+        """Extract hour component."""
+        return self._unary_temporal_int(pc.hour)
+
+    def minute(self) -> "UDFExpr":
+        """Extract minute component."""
+        return self._unary_temporal_int(pc.minute)
+
+    def second(self) -> "UDFExpr":
+        """Extract second component."""
+        return self._unary_temporal_int(pc.second)
+
+    # formatting
+
+    def strftime(self, fmt: str) -> "UDFExpr":
+        """Format timestamps with a strftime pattern."""
+
+        @pyarrow_udf(return_dtype=DataType.string())
+        def _format(arr: pyarrow.Array) -> pyarrow.Array:
+            return pc.strftime(arr, format=fmt)
+
+        return _format(self._expr)
+
+    # rounding
+
+    def ceil(self, unit: TemporalUnit) -> "UDFExpr":
+        """Ceil timestamps to the next multiple of the given unit."""
+        return_dtype = DataType.temporal()
+
+        @pyarrow_udf(return_dtype=return_dtype)
+        def _ceil(arr: pyarrow.Array) -> pyarrow.Array:
+            return pc.ceil_temporal(arr, multiple=1, unit=unit)
+
+        return _ceil(self._expr)
+
+    def floor(self, unit: TemporalUnit) -> "UDFExpr":
+        """Floor timestamps to the previous multiple of the given unit."""
+        return_dtype = DataType.temporal()
+
+        @pyarrow_udf(return_dtype=return_dtype)
+        def _floor(arr: pyarrow.Array) -> pyarrow.Array:
+            return pc.floor_temporal(arr, multiple=1, unit=unit)
+
+        return _floor(self._expr)
+
+    def round(self, unit: TemporalUnit) -> "UDFExpr":
+        """Round timestamps to the nearest multiple of the given unit."""
+        return_dtype = DataType.temporal()
+
+        @pyarrow_udf(return_dtype=return_dtype)
+        def _round(arr: pyarrow.Array) -> pyarrow.Array:
+
+            return pc.round_temporal(arr, multiple=1, unit=unit)
+
+        return _round(self._expr)

python/ray/data/tests/test_namespace_expressions.py

@@ -4,6 +4,7 @@
 convenient access to PyArrow compute functions through the expression API.
 """
 
+import datetime
 from typing import Any
 
 import pandas as pd
@@ -521,6 +522,75 @@ def test_struct_nested_bracket(self, dataset_format):
         assert_df_equal(result, expected)
 
 
+# ──────────────────────────────────────
+# Datatime Namespace Tests
+# ──────────────────────────────────────
+
+
+def test_dt_namespace_extractors(ray_start_regular):
+    ds = ray.data.from_items(
+        [
+            {
+                "ts": datetime.datetime(2024, 1, 2, 3, 4, 5),
+            }
+        ]
+    )
+
+    result_ds = ds.select(
+        [
+            col("ts").dt.year().alias("year"),
+            col("ts").dt.month().alias("month"),
+            col("ts").dt.day().alias("day"),
+            col("ts").dt.hour().alias("hour"),
+            col("ts").dt.minute().alias("minute"),
+            col("ts").dt.second().alias("second"),
+        ]
+    )
+
+    row = result_ds.take(1)[0]
+    assert row["year"] == 2024
+    assert row["month"] == 1
+    assert row["day"] == 2
+    assert row["hour"] == 3
+    assert row["minute"] == 4
+    assert row["second"] == 5
+
+
+def test_dt_namespace_strftime(ray_start_regular):
+    ds = ray.data.from_items(
+        [
+            {
+                "ts": datetime.datetime(2024, 1, 2, 3, 4, 5),
+            }
+        ]
+    )
+
+    result_ds = ds.select([col("ts").dt.strftime("%Y-%m-%d").alias("date_str")])
+
+    row = result_ds.take(1)[0]
+    assert row["date_str"] == "2024-01-02"
+
+
+def test_dt_namespace_rounding(ray_start_regular):
+    ts = datetime.datetime(2024, 1, 2, 10, 30, 0)
+
+    ds = ray.data.from_items([{"ts": ts}])
+
+    floored = ds.select([col("ts").dt.floor("day").alias("ts_floor")]).take(1)[0][
+        "ts_floor"
+    ]
+    ceiled = ds.select([col("ts").dt.ceil("day").alias("ts_ceil")]).take(1)[0][
+        "ts_ceil"
+    ]
+    rounded = ds.select([col("ts").dt.round("day").alias("ts_round")]).take(1)[0][
+        "ts_round"
+    ]
+
+    assert floored == datetime.datetime(2024, 1, 2, 0, 0, 0)
+    assert ceiled == datetime.datetime(2024, 1, 3, 0, 0, 0)
+    assert rounded == datetime.datetime(2024, 1, 3, 0, 0, 0)
+
+
 # ──────────────────────────────────────
 # Integration Tests
 # ──────────────────────────────────────