feat: convert and replace timezone

daft/expressions/expressions.py

@@ -1808,6 +1808,26 @@ def to_datetime(self, format: builtins.str, timezone: builtins.str | None = None
 
         return to_datetime(self, format, timezone)
 
+    def convert_time_zone(self, to_timezone: builtins.str, from_timezone: builtins.str | None = None) -> Expression:
+        """Converts a timestamp to another timezone while preserving the instant in time.
+
+        Tip: See Also
+            [`daft.functions.convert_time_zone`](https://docs.daft.ai/en/stable/api/functions/convert_time_zone/)
+        """
+        from daft.functions import convert_time_zone
+
+        return convert_time_zone(self, to_timezone, from_timezone)
+
+    def replace_time_zone(self, timezone: builtins.str | None = None) -> Expression:
+        """Replaces the timezone of a timestamp while preserving the local time.
+
+        Tip: See Also
+            [`daft.functions.replace_time_zone`](https://docs.daft.ai/en/stable/api/functions/replace_time_zone/)
+        """
+        from daft.functions import replace_time_zone
+
+        return replace_time_zone(self, timezone)
+
     def contains(self, substr: builtins.str | Expression) -> Expression:
         """Checks whether each string contains the given pattern in a string column.
 

daft/functions/__init__.py

@@ -60,6 +60,8 @@
     total_days,
     to_date,
     to_datetime,
+    convert_time_zone,
+    replace_time_zone,
     date_trunc,
     to_unix_epoch,
 )
@@ -272,6 +274,7 @@
     "concat",
     "contains",
     "convert_image",
+    "convert_time_zone",
     "cos",
     "cosh",
     "cosine_distance",
@@ -399,6 +402,7 @@
     "regexp_split",
     "repeat",
     "replace",
+    "replace_time_zone",
     "resample",
     "resize",
     "reverse",

daft/functions/datetime.py

@@ -1121,6 +1121,37 @@ def to_datetime(expr: Expression, format: str, timezone: str | None = None) -> E
     return Expression._call_builtin_scalar_fn("to_datetime", expr, format=format, timezone=timezone)
 
 
+def convert_time_zone(expr: Expression, to_timezone: str, from_timezone: str | None = None) -> Expression:
+    """Converts a timestamp to another timezone while preserving the instant in time.
+
+    If the timestamp has no timezone, `from_timezone` must be provided to interpret the local time before converting to `to_timezone`.
+
+    Args:
+        expr: Timestamp expression to convert.
+        to_timezone: Target timezone (e.g. "UTC", "+02:00", "America/New_York").
+        from_timezone: Source timezone for timestamps without a timezone.
+
+    Returns:
+        Expression: Timestamp expression with the target timezone.
+    """
+    return Expression._call_builtin_scalar_fn("convert_time_zone", expr, to_timezone, from_timezone)
+
+
+def replace_time_zone(expr: Expression, timezone: str | None = None) -> Expression:
+    """Replaces the timezone of a timestamp while preserving the local time.
+
+    If `timezone` is not provided, the timezone is removed.
+
+    Args:
+        expr: Timestamp expression to update.
+        timezone: New timezone (e.g. "UTC", "+02:00", "America/New_York").
+
+    Returns:
+        Expression: Timestamp expression with the updated timezone.
+    """
+    return Expression._call_builtin_scalar_fn("replace_time_zone", expr, timezone)
+
+
 def date_trunc(interval: str, expr: Expression, relative_to: Expression | None = None) -> Expression:
     """Truncates the datetime column to the specified interval.
 

daft/series.py

@@ -1082,6 +1082,12 @@ def to_unix_epoch(self, time_unit: str | TimeUnit | None = None) -> Series:
     def strftime(self, fmt: str | None = None) -> Series:
         return self._eval_expressions("strftime", format=fmt)
 
+    def convert_time_zone(self, to_timezone: str, from_timezone: str | None = None) -> Series:
+        return self._eval_expressions("convert_time_zone", to_timezone, from_timezone)
+
+    def replace_time_zone(self, timezone: str | None = None) -> Series:
+        return self._eval_expressions("replace_time_zone", timezone)
+
     def total_seconds(self) -> Series:
         return self._eval_expressions("total_seconds")
 

src/daft-core/src/array/ops/cast.rs

@@ -291,39 +291,26 @@ impl TimestampArray {
                 let DataType::Timestamp(unit, timezone) = self.data_type() else {
                     panic!("Wrong dtype for TimestampArray: {}", self.data_type())
                 };
-
-                let str_array: daft_arrow::array::Utf8Array<i64> = timezone.as_ref().map_or_else(
-                    || {
-                        self.as_arrow2()
-                            .iter()
-                            .map(|val| val.map(|val| timestamp_to_str_naive(*val, unit)))
-                            .collect()
-                    },
-                    |timezone| {
-                        if let Ok(offset) = daft_schema::time_unit::parse_offset(timezone) {
-                            self.as_arrow2()
-                                .iter()
-                                .map(|val| {
-                                    val.map(|val| timestamp_to_str_offset(*val, unit, &offset))
-                                })
-                                .collect()
-                        } else if let Ok(tz) = daft_schema::time_unit::parse_offset_tz(timezone) {
-                            self.as_arrow2()
-                                .iter()
-                                .map(|val| val.map(|val| timestamp_to_str_tz(*val, unit, &tz)))
-                                .collect()
-                        } else {
-                            panic!("Unable to parse timezone string {}", timezone)
-                        }
-                    },
+                let tz_parsed = timezone.as_ref().map(|tz| {
+                    daft_schema::time_unit::parse_timezone(tz)
+                        .unwrap_or_else(|_| panic!("Unable to parse timezone string {tz}"))
+                });
+                let str_array = Utf8Array::from_iter(
+                    self.name(),
+                    self.physical.as_arrow()?.iter().map(|val| {
+                        val.map(|val| match &tz_parsed {
+                            Some(daft_schema::time_unit::ParsedTimezone::Fixed(offset)) => {
+                                timestamp_to_str_offset(val, unit, offset)
+                            }
+                            Some(daft_schema::time_unit::ParsedTimezone::Tz(tz)) => {
+                                timestamp_to_str_tz(val, unit, tz)
+                            }
+                            None => timestamp_to_str_naive(val, unit),
+                        })
+                    }),
                 );
 
-                Ok(Utf8Array::new(
-                    Field::new(self.name(), DataType::Utf8).into(),
-                    Box::new(str_array),
-                )
-                .unwrap()
-                .into_series())
+                Ok(str_array.into_series())
             }
             dtype if dtype.is_numeric() => self.physical.cast(dtype),
             #[cfg(feature = "python")]