feat(rust): Implement strict AnyValue construction for temporal types (…

…#15146)
pola-rs · Mar 19, 2024 · 83c3fb0 · 83c3fb0
1 parent 223f281
commit 83c3fb0
Show file tree

Hide file tree

Showing 2 changed files with 243 additions and 63 deletions.
diff --git a/crates/polars-core/src/series/any_value.rs b/crates/polars-core/src/series/any_value.rs
@@ -103,21 +103,15 @@ impl Series {
             DataType::String => any_values_to_string(values, strict)?.into_series(),
             DataType::Binary => any_values_to_binary(values, strict)?.into_series(),
             #[cfg(feature = "dtype-date")]
-            DataType::Date => any_values_to_primitive_nonstrict::<Int32Type>(values)
-                .into_date()
-                .into_series(),
-            #[cfg(feature = "dtype-datetime")]
-            DataType::Datetime(tu, tz) => any_values_to_primitive_nonstrict::<Int64Type>(values)
-                .into_datetime(*tu, (*tz).clone())
-                .into_series(),
+            DataType::Date => any_values_to_date(values, strict)?.into_series(),
             #[cfg(feature = "dtype-time")]
-            DataType::Time => any_values_to_primitive_nonstrict::<Int64Type>(values)
-                .into_time()
-                .into_series(),
+            DataType::Time => any_values_to_time(values, strict)?.into_series(),
+            #[cfg(feature = "dtype-datetime")]
+            DataType::Datetime(tu, tz) => {
+                any_values_to_datetime(values, *tu, (*tz).clone(), strict)?.into_series()
+            },
             #[cfg(feature = "dtype-duration")]
-            DataType::Duration(tu) => any_values_to_primitive_nonstrict::<Int64Type>(values)
-                .into_duration(*tu)
-                .into_series(),
+            DataType::Duration(tu) => any_values_to_duration(values, *tu, strict)?.into_series(),
             #[cfg(feature = "dtype-categorical")]
             dt @ (DataType::Categorical(_, _) | DataType::Enum(_, _)) => {
                 any_values_to_categorical(values, dt, strict)?
@@ -225,33 +219,23 @@ fn any_values_to_f64(values: &[AnyValue], strict: bool) -> PolarsResult<Float64C
 }
 
 fn any_values_to_bool(values: &[AnyValue], strict: bool) -> PolarsResult<BooleanChunked> {
-    fn any_values_to_bool_strict(values: &[AnyValue]) -> PolarsResult<BooleanChunked> {
-        let mut builder = BooleanChunkedBuilder::new("", values.len());
-        for av in values {
-            match av {
-                AnyValue::Boolean(b) => builder.append_value(*b),
-                AnyValue::Null => builder.append_null(),
-                av => return Err(invalid_value_error(&DataType::Boolean, av)),
-            }
-        }
-        Ok(builder.finish())
-    }
-    fn any_values_to_bool_nonstrict(values: &[AnyValue]) -> BooleanChunked {
-        let mapper = |av: &AnyValue| match av {
-            AnyValue::Boolean(b) => Some(*b),
-            AnyValue::Null => None,
-            av => match av.cast(&DataType::Boolean) {
-                AnyValue::Boolean(b) => Some(b),
-                _ => None,
+    let mut builder = BooleanChunkedBuilder::new("", values.len());
+    for av in values {
+        match av {
+            AnyValue::Boolean(b) => builder.append_value(*b),
+            AnyValue::Null => builder.append_null(),
+            av => {
+                if strict {
+                    return Err(invalid_value_error(&DataType::Boolean, av));
+                }
+                match av.cast(&DataType::Boolean) {
+                    AnyValue::Boolean(b) => builder.append_value(b),
+                    _ => builder.append_null(),
+                }
             },
-        };
-        values.iter().map(mapper).collect_trusted()
-    }
-    if strict {
-        any_values_to_bool_strict(values)
-    } else {
-        Ok(any_values_to_bool_nonstrict(values))
+        }
     }
+    Ok(builder.finish())
 }
 
 fn any_values_to_string(values: &[AnyValue], strict: bool) -> PolarsResult<StringChunked> {
@@ -324,6 +308,101 @@ fn any_values_to_binary(values: &[AnyValue], strict: bool) -> PolarsResult<Binar
     }
 }
 
+#[cfg(feature = "dtype-date")]
+fn any_values_to_date(values: &[AnyValue], strict: bool) -> PolarsResult<DateChunked> {
+    let mut builder = PrimitiveChunkedBuilder::<Int32Type>::new("", values.len());
+    for av in values {
+        match av {
+            AnyValue::Date(i) => builder.append_value(*i),
+            AnyValue::Null => builder.append_null(),
+            av => {
+                if strict {
+                    return Err(invalid_value_error(&DataType::Date, av));
+                }
+                match av.cast(&DataType::Date) {
+                    AnyValue::Date(i) => builder.append_value(i),
+                    _ => builder.append_null(),
+                }
+            },
+        }
+    }
+    Ok(builder.finish().into())
+}
+
+#[cfg(feature = "dtype-time")]
+fn any_values_to_time(values: &[AnyValue], strict: bool) -> PolarsResult<TimeChunked> {
+    let mut builder = PrimitiveChunkedBuilder::<Int64Type>::new("", values.len());
+    for av in values {
+        match av {
+            AnyValue::Time(i) => builder.append_value(*i),
+            AnyValue::Null => builder.append_null(),
+            av => {
+                if strict {
+                    return Err(invalid_value_error(&DataType::Time, av));
+                }
+                match av.cast(&DataType::Time) {
+                    AnyValue::Time(i) => builder.append_value(i),
+                    _ => builder.append_null(),
+                }
+            },
+        }
+    }
+    Ok(builder.finish().into())
+}
+
+#[cfg(feature = "dtype-datetime")]
+fn any_values_to_datetime(
+    values: &[AnyValue],
+    time_unit: TimeUnit,
+    time_zone: Option<TimeZone>,
+    strict: bool,
+) -> PolarsResult<DatetimeChunked> {
+    let mut builder = PrimitiveChunkedBuilder::<Int64Type>::new("", values.len());
+    let target_dtype = DataType::Datetime(time_unit, time_zone.clone());
+    for av in values {
+        match av {
+            AnyValue::Datetime(i, tu, _) if *tu == time_unit => builder.append_value(*i),
+            AnyValue::Null => builder.append_null(),
+            av => {
+                if strict {
+                    return Err(invalid_value_error(&target_dtype, av));
+                }
+                match av.cast(&target_dtype) {
+                    AnyValue::Datetime(i, _, _) => builder.append_value(i),
+                    _ => builder.append_null(),
+                }
+            },
+        }
+    }
+    Ok(builder.finish().into_datetime(time_unit, time_zone))
+}
+
+#[cfg(feature = "dtype-duration")]
+fn any_values_to_duration(
+    values: &[AnyValue],
+    time_unit: TimeUnit,
+    strict: bool,
+) -> PolarsResult<DurationChunked> {
+    let mut builder = PrimitiveChunkedBuilder::<Int64Type>::new("", values.len());
+    let target_dtype = DataType::Duration(time_unit);
+    for av in values {
+        match av {
+            AnyValue::Duration(i, tu) if *tu == time_unit => builder.append_value(*i),
+            AnyValue::Null => builder.append_null(),
+            av => {
+                if strict {
+                    return Err(invalid_value_error(&target_dtype, av));
+                }
+                match av.cast(&target_dtype) {
+                    AnyValue::Duration(i, _) => builder.append_value(i),
+                    _ => builder.append_null(),
+                }
+            },
+        }
+    }
+    Ok(builder.finish().into_duration(time_unit))
+}
+
 #[cfg(feature = "dtype-categorical")]
 fn any_values_to_categorical(
     values: &[AnyValue],