chore: import upstream snapshot with attribution

python/ray/data/namespace_expressions/dt_namespace.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal
+
+import pyarrow.compute as pc
+
+from ray.data.datatype import DataType
+from ray.data.expressions import _create_pyarrow_compute_udf
+
+if TYPE_CHECKING:
+    from ray.data.expressions import Expr, PyArrowComputeUDFExpr
+
+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"
+
+    # extractors
+
+    def year(self) -> "PyArrowComputeUDFExpr":
+        """Extract year component."""
+        return _create_pyarrow_compute_udf(pc.year, DataType.int32())(self._expr)
+
+    def month(self) -> "PyArrowComputeUDFExpr":
+        """Extract month component."""
+        return _create_pyarrow_compute_udf(pc.month, DataType.int32())(self._expr)
+
+    def day(self) -> "PyArrowComputeUDFExpr":
+        """Extract day component."""
+        return _create_pyarrow_compute_udf(pc.day, DataType.int32())(self._expr)
+
+    def hour(self) -> "PyArrowComputeUDFExpr":
+        """Extract hour component."""
+        return _create_pyarrow_compute_udf(pc.hour, DataType.int32())(self._expr)
+
+    def minute(self) -> "PyArrowComputeUDFExpr":
+        """Extract minute component."""
+        return _create_pyarrow_compute_udf(pc.minute, DataType.int32())(self._expr)
+
+    def second(self) -> "PyArrowComputeUDFExpr":
+        """Extract second component."""
+        return _create_pyarrow_compute_udf(pc.second, DataType.int32())(self._expr)
+
+    # formatting
+
+    def strftime(self, fmt: str) -> "PyArrowComputeUDFExpr":
+        """Format timestamps with a strftime pattern."""
+        return _create_pyarrow_compute_udf(pc.strftime, DataType.string())(
+            self._expr, format=fmt
+        )
+
+    # rounding
+
+    def ceil(self, unit: TemporalUnit) -> "PyArrowComputeUDFExpr":
+        """Ceil timestamps to the next multiple of the given unit."""
+        return _create_pyarrow_compute_udf(pc.ceil_temporal, self._expr.data_type)(
+            self._expr, multiple=1, unit=unit
+        )
+
+    def floor(self, unit: TemporalUnit) -> "PyArrowComputeUDFExpr":
+        """Floor timestamps to the previous multiple of the given unit."""
+        return _create_pyarrow_compute_udf(pc.floor_temporal, self._expr.data_type)(
+            self._expr, multiple=1, unit=unit
+        )
+
+    def round(self, unit: TemporalUnit) -> "PyArrowComputeUDFExpr":
+        """Round timestamps to the nearest multiple of the given unit."""
+        return _create_pyarrow_compute_udf(pc.round_temporal, self._expr.data_type)(
+            self._expr, multiple=1, unit=unit
+        )