feat: add dt.tz_localize()

Author: sycai

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -16,7 +16,7 @@
 
 import functools
 import typing
-from typing import cast
+from typing import cast, Union
 
 from bigframes_vendored import ibis
 import bigframes_vendored.ibis.expr.api as ibis_api
@@ -978,7 +978,7 @@ def isin_op_impl(x: ibis_types.Value, op: ops.IsInOp):
 
 @scalar_op_compiler.register_unary_op(ops.ToDatetimeOp, pass_op=True)
 def to_datetime_op_impl(x: ibis_types.Value, op: ops.ToDatetimeOp):
-    if x.type() == ibis_dtypes.str:
+    if x.type() in (ibis_dtypes.str, ibis_dtypes.Timestamp("UTC")):
         return x.try_cast(ibis_dtypes.Timestamp(None))  # type: ignore
     else:
         # Numerical inputs.
@@ -1001,6 +1001,8 @@ def to_timestamp_op_impl(x: ibis_types.Value, op: ops.ToTimestampOp):
             if op.format
             else timestamp(x)
         )
+    elif x.type() == ibis_dtypes.Timestamp(None):  # Datetime type
+        return timestamp(x)
     else:
         # Numerical inputs.
         if op.format:
@@ -2016,8 +2018,8 @@ def _ibis_num(number: float):
 
 
 @ibis_udf.scalar.builtin
-def timestamp(a: str) -> ibis_dtypes.timestamp:  # type: ignore
-    """Convert string to timestamp."""
+def timestamp(a) -> ibis_dtypes.timestamp:  # type: ignore
+    """Convert string or a datetime to timestamp."""
 
 
 @ibis_udf.scalar.builtin

bigframes/operations/datetime_ops.py

@@ -73,6 +73,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
             dtypes.INT_DTYPE,
             dtypes.STRING_DTYPE,
             dtypes.DATE_DTYPE,
+            dtypes.TIMESTAMP_DTYPE,
         ):
             raise TypeError("expected string or numeric input")
         return pd.ArrowDtype(pa.timestamp("us", tz=None))
@@ -91,6 +92,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
             dtypes.INT_DTYPE,
             dtypes.STRING_DTYPE,
             dtypes.DATE_DTYPE,
+            dtypes.DATETIME_DTYPE,
         ):
             raise TypeError("expected string or numeric input")
         return pd.ArrowDtype(pa.timestamp("us", tz="UTC"))

bigframes/operations/datetimes.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import datetime as dt
-from typing import Optional
+from typing import Literal, Optional
 
 import bigframes_vendored.pandas.core.arrays.datetimelike as vendored_pandas_datetimelike
 import bigframes_vendored.pandas.core.indexes.accessor as vendordt
@@ -147,6 +147,21 @@ def tz(self) -> Optional[dt.timezone]:
         else:
             raise ValueError(f"Unexpected timezone {tz_string}")
 
+    def tz_localize(self, tz: Literal["UTC"] | None) -> series.Series:
+        if tz == "UTC":
+            if self._data.dtype == dtypes.TIMESTAMP_DTYPE:
+                raise ValueError("Already tz-aware.")
+
+            return self._data._apply_unary_op(ops.ToTimestampOp())
+
+        if tz is None:
+            if self._data.dtype == dtypes.DATETIME_DTYPE:
+                return self._data  # no-op
+
+            return self._data._apply_unary_op(ops.ToDatetimeOp())
+
+        raise ValueError(f"Unsupported timezone {tz}")
+
     @property
     def unit(self) -> str:
         # Assumption: pyarrow dtype

tests/system/small/operations/test_datetimes.py

@@ -324,6 +324,42 @@ def test_dt_tz(scalars_dfs, col_name):
     assert bf_result == pd_result
 
 
+@pytest.mark.parametrize(
+    ("col_name", "tz"),
+    [
+        ("datetime_col", None),
+        ("timestamp_col", None),
+        ("datetime_col", "UTC"),
+    ],
+)
+def test_dt_tz_localize(scalars_dfs, col_name, tz):
+    pytest.importorskip("pandas", minversion="2.0.0")
+    scalars_df, scalars_pandas_df = scalars_dfs
+    bf_series = scalars_df[col_name]
+
+    bf_result = bf_series.dt.tz_localize(tz)
+    pd_result = scalars_pandas_df[col_name].dt.tz_localize(tz)
+
+    testing.assert_series_equal(
+        bf_result.to_pandas(), pd_result, check_index_type=False
+    )
+
+
+@pytest.mark.parametrize(
+    ("col_name", "tz"),
+    [
+        ("timestamp_col", "UTC"),
+        ("datetime_col", "US/Eastern"),
+    ],
+)
+def test_dt_tz_localize_invalid_inputs(scalars_dfs, col_name, tz):
+    pytest.importorskip("pandas", minversion="2.0.0")
+    scalars_df, _ = scalars_dfs
+
+    with pytest.raises(ValueError):
+        scalars_df[col_name].dt.tz_localize(tz)
+
+
 @pytest.mark.parametrize(
     ("col_name",),
     DATETIME_COL_NAMES,

third_party/bigframes_vendored/pandas/core/indexes/accessor.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from bigframes import constants
 
 
@@ -499,6 +501,34 @@ def tz(self):
 
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+    @property
+    def tz_localize(self, tz: Literal["UTC"] | None):
+        """Localize tz-naive Datetime Array/Index to tz-aware Datetime Array/Index.
+
+        This method takes a time zone (tz) naive Datetime Array/Index object and makes
+        this time zone aware. It does not move the time to another time zone. Only "UTC"
+        timezone is supported.
+
+        This method can also be used to do the inverse - to create a time zone unaware
+        object from an aware object. To that end, pass tz=None.
+
+        **Examples:**
+
+            >>> import bigframes.pandas as bpd
+            >>> s = bpd.Series([pd.Timestamp(year = 2026, month=1, day=1)])
+            >>> s
+            0    2026-01-01 00:00:00
+            dtype: timestamp[us][pyarrow]
+            >>> s.dt.tz_localize('UTC')
+            0    2026-01-01 00:00:00+00:00
+            dtype: timestamp[us, tz=UTC][pyarrow]
+
+        Returns:
+            A BigFrames series with the updated timezone.
+        """
+
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
     @property
     def unit(self) -> str:
         """Returns the unit of time precision.