handle strictness lower

pola-rs · Jun 7, 2024 · 1f4845c · 1f4845c
1 parent 8c1f923
commit 1f4845c
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/crates/polars-core/src/chunked_array/cast.rs b/crates/polars-core/src/chunked_array/cast.rs
@@ -48,12 +48,24 @@ pub(crate) fn cast_chunks(
     dtype: &DataType,
     options: CastOptions,
 ) -> PolarsResult<Vec<ArrayRef>> {
+    let check_nulls = matches!(options, CastOptions::Strict);
     let options = options.into();
 
     let arrow_dtype = dtype.to_arrow(true);
     chunks
         .iter()
-        .map(|arr| arrow::compute::cast::cast(arr.as_ref(), &arrow_dtype, options))
+        .map(|arr| {
+            let out = arrow::compute::cast::cast(arr.as_ref(), &arrow_dtype, options);
+            if check_nulls {
+                out.and_then(|new| {
+                    polars_ensure!(arr.null_count() == new.null_count(), ComputeError: "strict cast failed");
+                    Ok(new)
+                })
+
+            } else {
+                out
+            }
+        })
         .collect::<PolarsResult<Vec<_>>>()
 }
 

diff --git a/crates/polars-core/src/series/mod.rs b/crates/polars-core/src/series/mod.rs
@@ -439,7 +439,13 @@ impl Series {
             return Ok(Series::full_null(self.name(), len, dtype));
         }
 
-        let ret = self.0.cast(dtype, options);
+        let new_options = match options {
+            // Strictness is handled on this level to improve error messages.
+            CastOptions::Strict => CastOptions::NonStrict,
+            opt => opt,
+        };
+
+        let ret = self.0.cast(dtype, new_options);
 
         match options {
             CastOptions::NonStrict | CastOptions::Overflowing => ret,