fix(rust, python): make weekday tz-aware (#5989)

pola-rs · Jan 2, 2023 · 935e04b · 935e04b
1 parent db0c741
commit 935e04b
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 41 deletions.
diff --git a/polars/polars-core/src/chunked_array/temporal/datetime.rs b/polars/polars-core/src/chunked_array/temporal/datetime.rs
@@ -42,14 +42,6 @@ impl DatetimeChunked {
         }
     }
 
-    #[cfg(feature = "timezones")]
-    pub fn apply_tz_offset(&self, tz: &str) -> PolarsResult<DatetimeChunked> {
-        let keep_tz = self.time_zone().as_deref().unwrap_or("UTC").to_string();
-        self.clone()
-            .with_time_zone(Some(tz.into()))
-            .cast_time_zone(&keep_tz)
-    }
-
     pub fn apply_on_tz_corrected<F>(&self, mut func: F) -> PolarsResult<DatetimeChunked>
     where
         F: FnMut(DatetimeChunked) -> PolarsResult<DatetimeChunked>,

diff --git a/polars/polars-time/src/chunkedarray/datetime.rs b/polars/polars-time/src/chunkedarray/datetime.rs
@@ -69,13 +69,7 @@ pub trait DatetimeMethods: AsDatetime {
     /// Extract ISO weekday from underlying NaiveDateTime representation.
     /// Returns the weekday number where monday = 1 and sunday = 7
     fn weekday(&self) -> UInt32Chunked {
-        let ca = self.as_datetime();
-        let f = match ca.time_unit() {
-            TimeUnit::Nanoseconds => datetime_to_weekday_ns,
-            TimeUnit::Microseconds => datetime_to_weekday_us,
-            TimeUnit::Milliseconds => datetime_to_weekday_ms,
-        };
-        ca.apply_kernel_cast::<UInt32Type>(&f)
+        cast_and_apply(self.as_datetime(), temporal::weekday)
     }
 
     /// Returns the ISO week number starting from 1.

diff --git a/polars/polars-time/src/chunkedarray/kernels.rs b/polars/polars-time/src/chunkedarray/kernels.rs
@@ -174,32 +174,6 @@ to_temporal_unit!(
     ArrowDataType::UInt32
 );
 
-#[cfg(feature = "dtype-datetime")]
-to_temporal_unit!(
-    datetime_to_weekday_ns,
-    p_weekday,
-    timestamp_ns_to_datetime,
-    i64,
-    ArrowDataType::UInt32
-);
-
-#[cfg(feature = "dtype-datetime")]
-to_temporal_unit!(
-    datetime_to_weekday_ms,
-    p_weekday,
-    timestamp_ms_to_datetime,
-    i64,
-    ArrowDataType::UInt32
-);
-#[cfg(feature = "dtype-datetime")]
-to_temporal_unit!(
-    datetime_to_weekday_us,
-    p_weekday,
-    timestamp_us_to_datetime,
-    i64,
-    ArrowDataType::UInt32
-);
-
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
     datetime_to_iso_year_ns,

diff --git a/py-polars/tests/unit/test_datelike.py b/py-polars/tests/unit/test_datelike.py
@@ -2306,3 +2306,34 @@ def test_cast_time_to_duration() -> None:
     assert pl.Series([time(hour=0, minute=0, second=2)]).cast(
         pl.Duration
     ).item() == timedelta(seconds=2)
+
+
+def test_tz_aware_day_weekday() -> None:
+    start = datetime(2001, 1, 1)
+    stop = datetime(2001, 1, 9)
+    df = pl.DataFrame({"date": pl.date_range(start, stop, timedelta(days=3))})
+
+    df = df.with_columns(
+        [
+            pl.col("date").dt.with_time_zone("Asia/Tokyo").alias("tyo_date"),
+            pl.col("date").dt.with_time_zone("America/New_York").alias("ny_date"),
+        ]
+    )
+
+    assert df.select(
+        [
+            pl.col("date").dt.day().alias("day"),
+            pl.col("tyo_date").dt.day().alias("tyo_day"),
+            pl.col("ny_date").dt.day().alias("ny_day"),
+            pl.col("date").dt.weekday().alias("weekday"),
+            pl.col("tyo_date").dt.weekday().alias("tyo_weekday"),
+            pl.col("ny_date").dt.weekday().alias("ny_weekday"),
+        ]
+    ).to_dict(False) == {
+        "day": [1, 4, 7],
+        "tyo_day": [1, 4, 7],
+        "ny_day": [31, 3, 6],
+        "weekday": [1, 4, 7],
+        "tyo_weekday": [1, 4, 7],
+        "ny_weekday": [7, 3, 6],
+    }