WIP: TimeUnits and TimeZones init (#2209)

.gitignore

@@ -10,4 +10,5 @@ polars/vendor
 .ENV
 .env
 AUTO_CHANGELOG.md
-node_modules/
+node_modules/
+.coverage

polars/polars-core/src/chunked_array/cast.rs

@@ -19,7 +19,8 @@ fn cast_impl(name: &str, chunks: &[ArrayRef], dtype: &DataType) -> Result<Series
     let out = Series::try_from((name, chunks))?;
     use DataType::*;
     let out = match dtype {
-        Date | Datetime => out.into_date(),
+        Date => out.into_date(),
+        Datetime(tu, tz) => out.into_datetime(*tu, tz.clone()),
         #[cfg(feature = "dtype-time")]
         Time => out.into_time(),
         _ => out,

polars/polars-core/src/chunked_array/kernels/temporal.rs

@@ -5,7 +5,9 @@ use arrow::array::{ArrayRef, PrimitiveArray};
 use arrow::compute::arity::unary;
 #[cfg(feature = "dtype-time")]
 use arrow::temporal_conversions::time64ns_to_time;
-use arrow::temporal_conversions::{date32_to_datetime, timestamp_ns_to_datetime};
+use arrow::temporal_conversions::{
+    date32_to_datetime, timestamp_ms_to_datetime, timestamp_ns_to_datetime,
+};
 use polars_time::export::chrono::{Datelike, NaiveDate, NaiveDateTime, Timelike};
 use std::sync::Arc;
 
@@ -46,6 +48,7 @@ macro_rules! to_temporal_unit {
         }
     };
 }
+// Dates
 #[cfg(feature = "dtype-date")]
 to_temporal_unit!(
     date_to_week,
@@ -94,116 +97,202 @@ to_temporal_unit!(
     i32,
     ArrowDataType::UInt32
 );
+
+// Times
+#[cfg(feature = "dtype-time")]
+to_temporal_unit!(
+    time_to_hour,
+    hour,
+    time64ns_to_time,
+    i64,
+    ArrowDataType::UInt32
+);
+#[cfg(feature = "dtype-time")]
+to_temporal_unit!(
+    time_to_minute,
+    minute,
+    time64ns_to_time,
+    i64,
+    ArrowDataType::UInt32
+);
+#[cfg(feature = "dtype-time")]
+to_temporal_unit!(
+    time_to_second,
+    second,
+    time64ns_to_time,
+    i64,
+    ArrowDataType::UInt32
+);
+#[cfg(feature = "dtype-time")]
+to_temporal_unit!(
+    time_to_nanosecond,
+    nanosecond,
+    time64ns_to_time,
+    i64,
+    ArrowDataType::UInt32
+);
+
+// Datetimes nanoseconds
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_week,
+    datetime_to_week_ns,
     week,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_weekday,
+    datetime_to_weekday_ns,
     p_weekday,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_year,
+    datetime_to_year_ns,
     year,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::Int32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_month,
+    datetime_to_month_ns,
     month,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_day,
+    datetime_to_day_ns,
     day,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_hour,
+    datetime_to_hour_ns,
     hour,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_minute,
+    datetime_to_minute_ns,
     minute,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_second,
+    datetime_to_second_ns,
     second,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_nanosecond,
+    datetime_to_nanosecond_ns,
     nanosecond,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 #[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    datetime_to_ordinal,
+    datetime_to_ordinal_ns,
     ordinal,
     timestamp_ns_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
 
-#[cfg(feature = "dtype-time")]
+// Datetimes milliseconds
+
+#[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    time_to_hour,
+    datetime_to_week_ms,
+    week,
+    timestamp_ms_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_year_ms,
+    year,
+    timestamp_ms_to_datetime,
+    i64,
+    ArrowDataType::Int32
+);
+#[cfg(feature = "dtype-datetime")]
+to_temporal_unit!(
+    datetime_to_month_ms,
+    month,
+    timestamp_ms_to_datetime,
+    i64,
+    ArrowDataType::UInt32
+);
+#[cfg(feature = "dtype-datetime")]
+to_temporal_unit!(
+    datetime_to_day_ms,
+    day,
+    timestamp_ms_to_datetime,
+    i64,
+    ArrowDataType::UInt32
+);
+#[cfg(feature = "dtype-datetime")]
+to_temporal_unit!(
+    datetime_to_hour_ms,
     hour,
-    time64ns_to_time,
+    timestamp_ms_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
-#[cfg(feature = "dtype-time")]
+#[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    time_to_minute,
+    datetime_to_minute_ms,
     minute,
-    time64ns_to_time,
+    timestamp_ms_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
-#[cfg(feature = "dtype-time")]
+#[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    time_to_second,
+    datetime_to_second_ms,
     second,
-    time64ns_to_time,
+    timestamp_ms_to_datetime,
     i64,
     ArrowDataType::UInt32
 );
-#[cfg(feature = "dtype-time")]
+#[cfg(feature = "dtype-datetime")]
 to_temporal_unit!(
-    time_to_nanosecond,
+    datetime_to_nanosecond_ms,
     nanosecond,
-    time64ns_to_time,
+    timestamp_ms_to_datetime,
+    i64,
+    ArrowDataType::UInt32
+);
+#[cfg(feature = "dtype-datetime")]
+to_temporal_unit!(
+    datetime_to_ordinal_ms,
+    ordinal,
+    timestamp_ms_to_datetime,
     i64,
     ArrowDataType::UInt32
 );

polars/polars-core/src/chunked_array/logical/date.rs

@@ -16,7 +16,7 @@ impl Int32Chunked {
 }
 
 impl LogicalType for DateChunked {
-    fn dtype(&self) -> &'static DataType {
+    fn dtype(&self) -> &DataType {
         &DataType::Date
     }
 

polars/polars-core/src/chunked_array/logical/datetime.rs

@@ -3,25 +3,23 @@ use crate::prelude::*;
 
 pub type DatetimeChunked = Logical<DatetimeType, Int64Type>;
 
-impl From<Int64Chunked> for DatetimeChunked {
-    fn from(ca: Int64Chunked) -> Self {
-        DatetimeChunked::new(ca)
-    }
-}
-
 impl Int64Chunked {
-    pub fn into_date(self) -> DatetimeChunked {
-        DatetimeChunked::new(self)
+    pub fn into_datetime(self, timeunit: TimeUnit, tz: Option<TimeZone>) -> DatetimeChunked {
+        let mut dt = DatetimeChunked::new(self);
+        dt.2 = Some(DataType::Datetime(timeunit, tz));
+        dt
     }
 }
 
 impl LogicalType for DatetimeChunked {
-    fn dtype(&self) -> &'static DataType {
-        &DataType::Datetime
+    fn dtype(&self) -> &DataType {
+        self.2.as_ref().unwrap()
     }
 
-    #[cfg(feature = "dtype-date")]
+    #[cfg(feature = "dtype-datetime")]
     fn get_any_value(&self, i: usize) -> AnyValue<'_> {
-        self.0.get_any_value(i).into_date()
+        self.0
+            .get_any_value(i)
+            .into_datetime(self.time_unit(), self.time_zone())
     }
 }

polars/polars-core/src/chunked_array/logical/mod.rs

@@ -10,11 +10,17 @@ use std::ops::{Deref, DerefMut};
 
 /// Maps a logical type to a a chunked array implementation of the physical type.
 /// This saves a lot of compiler bloat and allows us to reuse functionality.
-pub struct Logical<K: PolarsDataType, T: PolarsDataType>(pub ChunkedArray<T>, PhantomData<K>);
+pub struct Logical<K: PolarsDataType, T: PolarsDataType>(
+    pub ChunkedArray<T>,
+    PhantomData<K>,
+    pub Option<DataType>,
+);
 
 impl<K: PolarsDataType, T: PolarsDataType> Clone for Logical<K, T> {
     fn clone(&self) -> Self {
-        Logical::<K, _>::new(self.0.clone())
+        let mut new = Logical::<K, _>::new(self.0.clone());
+        new.2 = self.2.clone();
+        new
     }
 }
 
@@ -34,13 +40,13 @@ impl<K: PolarsDataType, T: PolarsDataType> DerefMut for Logical<K, T> {
 
 impl<K: PolarsDataType, T: PolarsDataType> Logical<K, T> {
     pub fn new<J: PolarsDataType>(ca: ChunkedArray<T>) -> Logical<J, T> {
-        Logical(ca, PhantomData)
+        Logical(ca, PhantomData, None)
     }
 }
 
 pub trait LogicalType {
     /// Get data type of ChunkedArray.
-    fn dtype(&self) -> &'static DataType;
+    fn dtype(&self) -> &DataType;
 
     fn get_any_value(&self, _i: usize) -> AnyValue<'_> {
         unimplemented!()

polars/polars-core/src/chunked_array/mod.rs

@@ -237,7 +237,7 @@ impl<T> ChunkedArray<T> {
         } else {
             use DataType::*;
             match (self.dtype(), series.dtype()) {
-                (Int64, Datetime) | (Int32, Date) => {
+                (Int64, Datetime(_, _)) | (Int32, Date) => {
                     &*(series_trait as *const dyn SeriesTrait as *const ChunkedArray<T>)
                 }
                 _ => panic!(

polars/polars-core/src/chunked_array/ops/any_value.rs

@@ -11,7 +11,7 @@ unsafe fn arr_to_any_value<'a>(
     arr: &'a dyn Array,
     idx: usize,
     categorical_map: &'a Option<Arc<RevMapping>>,
-    dtype: &DataType,
+    dtype: &'a DataType,
 ) -> AnyValue<'a> {
     if arr.is_null(idx) {
         return AnyValue::Null;
@@ -47,7 +47,10 @@ unsafe fn arr_to_any_value<'a>(
         #[cfg(feature = "dtype-date")]
         DataType::Date => downcast_and_pack!(Int32Array, Date),
         #[cfg(feature = "dtype-datetime")]
-        DataType::Datetime => downcast_and_pack!(Int64Array, Datetime),
+        DataType::Datetime(tu, tz) => {
+            let ts: i64 = downcast!(Int64Array);
+            AnyValue::Datetime(ts, *tu, tz)
+        }
         DataType::List(dt) => {
             let v: ArrayRef = downcast!(LargeListArray).into();
             let mut s = Series::try_from(("", v)).unwrap();