[Data] Fix datetime namespace review comments

Author: 400Ping

python/ray/data/namespace_expressions/dt_namespace.py

@@ -1,27 +1,7 @@
-"""
-Datetime namespace for expression operations on datetime-typed columns.
-
-This module defines the ``_DatetimeNamespace`` class which exposes a set of
-convenience methods for working with timestamp and date columns in Ray Data
-expressions.  The API mirrors pandas' ``Series.dt`` accessor and is backed by
-PyArrow compute functions for efficient execution.
-
-Example
--------
-
->>> from ray.data.expressions import col
->>> # Extract year, month and day from a timestamp column
->>> expr_year = col("timestamp").dt.year()
->>> expr_month = col("timestamp").dt.month()
->>> expr_day = col("timestamp").dt.day()
->>> # Format the timestamp as a string
->>> expr_fmt = col("timestamp").dt.strftime("%Y-%m-%d")
-"""
-
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Callable, Literal
 
 import pyarrow
 import pyarrow.compute as pc
@@ -32,75 +12,104 @@
 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:
-    """Namespace for datetime operations on expression columns."""
+    """Datetime namespace for operations on datetime-typed expression columns."""
+
     _expr: "Expr"
 
-    def year(self) -> "UDFExpr":
+    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 _year(arr: pyarrow.Array) -> pyarrow.Array:
-            return pc.year(arr)
-        return _year(self._expr)
+        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":
-        @pyarrow_udf(return_dtype=DataType.int32())
-        def _month(arr: pyarrow.Array) -> pyarrow.Array:
-            return pc.month(arr)
-        return _month(self._expr)
+        """Extract month component."""
+        return self._unary_temporal_int(pc.month)
 
     def day(self) -> "UDFExpr":
-        @pyarrow_udf(return_dtype=DataType.int32())
-        def _day(arr: pyarrow.Array) -> pyarrow.Array:
-            return pc.day(arr)
-        return _day(self._expr)
+        """Extract day component."""
+        return self._unary_temporal_int(pc.day)
 
     def hour(self) -> "UDFExpr":
-        @pyarrow_udf(return_dtype=DataType.int32())
-        def _hour(arr: pyarrow.Array) -> pyarrow.Array:
-            return pc.hour(arr)
-        return _hour(self._expr)
+        """Extract hour component."""
+        return self._unary_temporal_int(pc.hour)
 
     def minute(self) -> "UDFExpr":
-        @pyarrow_udf(return_dtype=DataType.int32())
-        def _minute(arr: pyarrow.Array) -> pyarrow.Array:
-            return pc.minute(arr)
-        return _minute(self._expr)
+        """Extract minute component."""
+        return self._unary_temporal_int(pc.minute)
 
     def second(self) -> "UDFExpr":
-        @pyarrow_udf(return_dtype=DataType.int32())
-        def _second(arr: pyarrow.Array) -> pyarrow.Array:
-            return pc.second(arr)
-        return _second(self._expr)
+        """Extract second component."""
+        return self._unary_temporal_int(pc.second)
+
+    # formatting
 
     def strftime(self, fmt: str) -> "UDFExpr":
-        """Format each timestamp using a strftime format string."""
+        """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)
 
-    def ceil(self, unit: str) -> "UDFExpr":
-        """Ceil timestamps up to the nearest boundary of the given unit."""
+    # rounding
+
+    def ceil(self, unit: TemporalUnit) -> "UDFExpr":
+        """Ceil timestamps to the next multiple of the given unit."""
         return_dtype = self._expr.data_type
+
         @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: str) -> "UDFExpr":
-        """Floor timestamps down to the previous boundary of the given unit."""
+    def floor(self, unit: TemporalUnit) -> "UDFExpr":
+        """Floor timestamps to the previous multiple of the given unit."""
         return_dtype = self._expr.data_type
+
         @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: str, tie_breaker: str = "half_to_even") -> "UDFExpr":
-        """Round timestamps to the nearest boundary of the given unit."""
+    def round(self, unit: TemporalUnit) -> "UDFExpr":
+        """Round timestamps to the nearest multiple of the given unit."""
         return_dtype = self._expr.data_type
+
         @pyarrow_udf(return_dtype=return_dtype)
         def _round(arr: pyarrow.Array) -> pyarrow.Array:
-            return pc.round_temporal(arr, multiple=1, unit=unit,
-                                      tie_breaker=tie_breaker)
+
+            return pc.round_temporal(arr, multiple=1, unit=unit)
+
         return _round(self._expr)

python/ray/data/tests/test_namespace_expressions.py

@@ -5,6 +5,7 @@
 """
 
 from typing import Any
+import datetime
 
 import pandas as pd
 import pyarrow as pa
@@ -521,6 +522,77 @@ def test_struct_nested_bracket(self, dataset_format):
         assert_df_equal(result, expected)
 
 
+# ──────────────────────────────────────
+# Struct 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
 # ──────────────────────────────────────