refactor(rust): implement ChunkArray::(try_)from_chunk_iter (pola-rs#10395)

Author: orlp

contribution/polars_ops_multiple_arguments/src/lib.rs

@@ -47,26 +47,21 @@ fn compute_chunked_array_2_args<T: PolarsNumericType>(
     ca_1: &ChunkedArray<T>,
     ca_2: &ChunkedArray<T>,
 ) -> ChunkedArray<T> {
-    // this ensures both ChunkedArrays have the same number of chunks with the same offset
-    // and the same length.
+    // This ensures both ChunkedArrays have the same number of chunks with the
+    // same offset and the same length.
     let (ca_1, ca_2) = align_chunks_binary(ca_1, ca_2);
-
     let chunks = ca_1
         .downcast_iter()
         .zip(ca_2.downcast_iter())
-        .map(|(arr_1, arr_2)| compute_kernel(arr_1, arr_2).boxed())
-        .collect::<Vec<_>>();
-
-    // Safety: we are sure the `ArrayRef` holds type `T`
-    unsafe { ChunkedArray::from_chunks(ca_1.name(), chunks) }
+        .map(|(arr_1, arr_2)| compute_kernel(arr_1, arr_2));
+    ChunkedArray::from_chunk_iter(ca_1.name(), chunks)
 }
 
 pub fn compute_expr_2_args(arg_1: &Series, arg_2: &Series) -> Series {
-    // dispatch the numerical series to `compute_chunked_array_2_args`
+    // Dispatch the numerical series to `compute_chunked_array_2_args`.
     with_match_physical_numeric_polars_type!(arg_1.dtype(), |$T| {
-            let ca_1: &ChunkedArray<$T> = arg_1.as_ref().as_ref().as_ref();
-            let ca_2: &ChunkedArray<$T> = arg_2.as_ref().as_ref().as_ref();
-
-    compute_chunked_array_2_args(ca_1, ca_2).into_series()
-        })
+        let ca_1: &ChunkedArray<$T> = arg_1.as_ref().as_ref().as_ref();
+        let ca_2: &ChunkedArray<$T> = arg_2.as_ref().as_ref().as_ref();
+        compute_chunked_array_2_args(ca_1, ca_2).into_series()
+    })
 }

crates/polars-arrow/src/array/utf8.rs

@@ -55,13 +55,13 @@ impl<T: AsRef<str>> AsRef<[u8]> for StrAsBytes<T> {
 
 pub trait Utf8FromIter {
     #[inline]
-    fn from_values_iter<I, S>(iter: I, len: usize, value_cap: usize) -> Utf8Array<i64>
+    fn from_values_iter<I, S>(iter: I, len: usize, size_hint: usize) -> Utf8Array<i64>
     where
         S: AsRef<str>,
         I: Iterator<Item = S>,
     {
         let iter = iter.map(StrAsBytes);
-        let (offsets, values) = unsafe { fill_offsets_and_values(iter, value_cap, len) };
+        let (offsets, values) = unsafe { fill_offsets_and_values(iter, size_hint, len) };
         unsafe {
             Utf8Array::new_unchecked(DataType::LargeUtf8, offsets.into(), values.into(), None)
         }

crates/polars-arrow/src/kernels/time.rs

@@ -5,16 +5,12 @@ use arrow::error::{Error as ArrowError, Result};
 use arrow::temporal_conversions::{
     timestamp_ms_to_datetime, timestamp_ns_to_datetime, timestamp_us_to_datetime,
 };
-#[cfg(feature = "timezones")]
 use chrono::{LocalResult, NaiveDateTime, TimeZone};
-#[cfg(feature = "timezones")]
 use chrono_tz::Tz;
 use polars_error::polars_bail;
 
 use crate::error::PolarsResult;
-use crate::prelude::ArrayRef;
 
-#[cfg(feature = "timezones")]
 fn convert_to_naive_local(
     from_tz: &Tz,
     to_tz: &Tz,
@@ -41,68 +37,55 @@ fn convert_to_naive_local(
     }
 }
 
-#[cfg(feature = "timezones")]
 fn convert_to_timestamp(
     from_tz: Tz,
     to_tz: Tz,
     arr: &PrimitiveArray<i64>,
     tu: TimeUnit,
     use_earliest: Option<bool>,
-) -> PolarsResult<ArrayRef> {
-    match tu {
-        TimeUnit::Millisecond => {
-            let data = try_unary(
-                arr,
-                |value| {
-                    let ndt = timestamp_ms_to_datetime(value);
-                    Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?
-                        .timestamp_millis())
-                },
-                ArrowDataType::Int64,
-            )?;
-            Ok(Box::new(data))
-        }
-        TimeUnit::Microsecond => {
-            let data = try_unary(
-                arr,
-                |value| {
-                    let ndt = timestamp_us_to_datetime(value);
-                    Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?
-                        .timestamp_micros())
-                },
-                ArrowDataType::Int64,
-            )?;
-            Ok(Box::new(data))
-        }
-        TimeUnit::Nanosecond => {
-            let data = try_unary(
-                arr,
-                |value| {
-                    let ndt = timestamp_ns_to_datetime(value);
-                    Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?
-                        .timestamp_nanos())
-                },
-                ArrowDataType::Int64,
-            )?;
-            Ok(Box::new(data))
-        }
+) -> PolarsResult<PrimitiveArray<i64>> {
+    let res = match tu {
+        TimeUnit::Millisecond => try_unary(
+            arr,
+            |value| {
+                let ndt = timestamp_ms_to_datetime(value);
+                Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?.timestamp_millis())
+            },
+            ArrowDataType::Int64,
+        ),
+        TimeUnit::Microsecond => try_unary(
+            arr,
+            |value| {
+                let ndt = timestamp_us_to_datetime(value);
+                Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?.timestamp_micros())
+            },
+            ArrowDataType::Int64,
+        ),
+        TimeUnit::Nanosecond => try_unary(
+            arr,
+            |value| {
+                let ndt = timestamp_ns_to_datetime(value);
+                Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?.timestamp_nanos())
+            },
+            ArrowDataType::Int64,
+        ),
         _ => unreachable!(),
-    }
+    };
+    Ok(res?)
 }
 
-#[cfg(feature = "timezones")]
 pub fn replace_time_zone(
     arr: &PrimitiveArray<i64>,
     tu: TimeUnit,
     from: &str,
     to: &str,
     use_earliest: Option<bool>,
-) -> PolarsResult<ArrayRef> {
-    Ok(match from.parse::<chrono_tz::Tz>() {
+) -> PolarsResult<PrimitiveArray<i64>> {
+    match from.parse::<chrono_tz::Tz>() {
         Ok(from_tz) => match to.parse::<chrono_tz::Tz>() {
-            Ok(to_tz) => convert_to_timestamp(from_tz, to_tz, arr, tu, use_earliest)?,
+            Ok(to_tz) => convert_to_timestamp(from_tz, to_tz, arr, tu, use_earliest),
             Err(_) => polars_bail!(ComputeError: "unable to parse time zone: '{}'", to),
         },
         Err(_) => polars_bail!(ComputeError: "unable to parse time zone: '{}'", from),
-    })
+    }
 }

crates/polars-arrow/src/lib.rs

@@ -2,7 +2,7 @@
 #![cfg_attr(
     feature = "nightly",
     allow(clippy::incorrect_partial_ord_impl_on_ord_type)
-)] // remove once stable
+)] // Remove once stable.
 pub mod array;
 pub mod bit_util;
 pub mod bitmap;

crates/polars-core/src/chunked_array/bitwise.rs

@@ -77,10 +77,8 @@ impl BitOr for &BooleanChunked {
         let chunks = lhs
             .downcast_iter()
             .zip(rhs.downcast_iter())
-            .map(|(lhs, rhs)| Box::new(compute::boolean_kleene::or(lhs, rhs)) as ArrayRef)
-            .collect();
-        // safety: same type
-        unsafe { BooleanChunked::from_chunks(self.name(), chunks) }
+            .map(|(lhs, rhs)| compute::boolean_kleene::or(lhs, rhs));
+        BooleanChunked::from_chunk_iter(self.name(), chunks)
     }
 }
 
@@ -132,14 +130,9 @@ impl BitXor for &BooleanChunked {
             .map(|(l_arr, r_arr)| {
                 let validity = combine_validities_and(l_arr.validity(), r_arr.validity());
                 let values = l_arr.values() ^ r_arr.values();
-
-                let arr = BooleanArray::from_data_default(values, validity);
-                Box::new(arr) as ArrayRef
-            })
-            .collect::<Vec<_>>();
-
-        // safety: same type
-        unsafe { ChunkedArray::from_chunks(self.name(), chunks) }
+                BooleanArray::from_data_default(values, validity)
+            });
+        ChunkedArray::from_chunk_iter(self.name(), chunks)
     }
 }
 
@@ -180,10 +173,8 @@ impl BitAnd for &BooleanChunked {
         let chunks = lhs
             .downcast_iter()
             .zip(rhs.downcast_iter())
-            .map(|(lhs, rhs)| Box::new(compute::boolean_kleene::and(lhs, rhs)) as ArrayRef)
-            .collect();
-        // safety: same type
-        unsafe { BooleanChunked::from_chunks(self.name(), chunks) }
+            .map(|(lhs, rhs)| compute::boolean_kleene::and(lhs, rhs));
+        BooleanChunked::from_chunk_iter(self.name(), chunks)
     }
 }
 

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

@@ -22,6 +22,7 @@ pub use list::*;
 pub use primitive::*;
 pub use utf8::*;
 
+use crate::chunked_array::to_primitive;
 use crate::prelude::*;
 use crate::utils::{get_iter_capacity, NoNull};
 
@@ -46,13 +47,10 @@ where
     T: PolarsNumericType,
 {
     fn from_iter<I: IntoIterator<Item = (Vec<T::Native>, Option<Bitmap>)>>(iter: I) -> Self {
-        let mut chunks = vec![];
-
-        for (values, opt_buffer) in iter {
-            chunks.push(to_array::<T>(values, opt_buffer))
-        }
-        // safety: same type
-        unsafe { ChunkedArray::from_chunks("from_iter", chunks) }
+        let chunks = iter
+            .into_iter()
+            .map(|(values, opt_buffer)| to_primitive::<T>(values, opt_buffer));
+        ChunkedArray::from_chunk_iter("from_iter", chunks)
     }
 }
 
@@ -72,9 +70,8 @@ where
     T: PolarsNumericType,
 {
     fn from_slice(name: &str, v: &[T::Native]) -> Self {
-        let arr = PrimitiveArray::<T::Native>::from_slice(v).to(T::get_dtype().to_arrow());
-        // safety: same type
-        unsafe { ChunkedArray::from_chunks(name, vec![Box::new(arr)]) }
+        let arr = PrimitiveArray::from_slice(v).to(T::get_dtype().to_arrow());
+        ChunkedArray::from_chunk_iter(name, [arr])
     }
 
     fn from_slice_options(name: &str, opt_v: &[Option<T::Native>]) -> Self {
@@ -131,13 +128,10 @@ where
 {
     fn from_slice(name: &str, v: &[S]) -> Self {
         let values_size = v.iter().fold(0, |acc, s| acc + s.as_ref().len());
-
         let mut builder = MutableUtf8Array::<i64>::with_capacities(v.len(), values_size);
         builder.extend_trusted_len_values(v.iter().map(|s| s.as_ref()));
-
-        let chunks = vec![builder.as_box()];
-        // safety: same type
-        unsafe { ChunkedArray::from_chunks(name, chunks) }
+        let imm: Utf8Array<i64> = builder.into();
+        ChunkedArray::from_chunk_iter(name, [imm])
     }
 
     fn from_slice_options(name: &str, opt_v: &[Option<S>]) -> Self {
@@ -147,10 +141,8 @@ where
         });
         let mut builder = MutableUtf8Array::<i64>::with_capacities(opt_v.len(), values_size);
         builder.extend_trusted_len(opt_v.iter().map(|s| s.as_ref()));
-
-        let chunks = vec![builder.as_box()];
-        // safety: same type
-        unsafe { ChunkedArray::from_chunks(name, chunks) }
+        let imm: Utf8Array<i64> = builder.into();
+        ChunkedArray::from_chunk_iter(name, [imm])
     }
 
     fn from_iter_options(name: &str, it: impl Iterator<Item = Option<S>>) -> Self {
@@ -175,13 +167,10 @@ where
 {
     fn from_slice(name: &str, v: &[B]) -> Self {
         let values_size = v.iter().fold(0, |acc, s| acc + s.as_ref().len());
-
         let mut builder = MutableBinaryArray::<i64>::with_capacities(v.len(), values_size);
         builder.extend_trusted_len_values(v.iter().map(|s| s.as_ref()));
-
-        let chunks = vec![builder.as_box()];
-        // safety: same type
-        unsafe { ChunkedArray::from_chunks(name, chunks) }
+        let imm: BinaryArray<i64> = builder.into();
+        ChunkedArray::from_chunk_iter(name, [imm])
     }
 
     fn from_slice_options(name: &str, opt_v: &[Option<B>]) -> Self {
@@ -191,10 +180,8 @@ where
         });
         let mut builder = MutableBinaryArray::<i64>::with_capacities(opt_v.len(), values_size);
         builder.extend_trusted_len(opt_v.iter().map(|s| s.as_ref()));
-
-        let chunks = vec![builder.as_box()];
-        // safety: same type
-        unsafe { ChunkedArray::from_chunks(name, chunks) }
+        let imm: BinaryArray<i64> = builder.into();
+        ChunkedArray::from_chunk_iter(name, [imm])
     }
 
     fn from_iter_options(name: &str, it: impl Iterator<Item = Option<B>>) -> Self {
@@ -233,7 +220,7 @@ mod test {
         let mut builder =
             ListPrimitiveChunkedBuilder::<Int32Type>::new("a", 10, 5, DataType::Int32);
 
-        // create a series containing two chunks
+        // Create a series containing two chunks.
         let mut s1 = Int32Chunked::from_slice("a", &[1, 2, 3]).into_series();
         let s2 = Int32Chunked::from_slice("b", &[4, 5, 6]).into_series();
         s1.append(&s2).unwrap();
@@ -252,7 +239,8 @@ mod test {
         } else {
             panic!()
         }
-        // test list collect
+
+        // Test list collect.
         let out = [&s1, &s2].iter().copied().collect::<ListChunked>();
         assert_eq!(out.get(0).unwrap().len(), 6);
         assert_eq!(out.get(1).unwrap().len(), 3);

crates/polars-core/src/chunked_array/from.rs

@@ -74,6 +74,30 @@ impl<T> ChunkedArray<T>
 where
     T: PolarsDataType,
 {
+    pub fn from_chunk_iter<I>(name: &str, iter: I) -> Self
+    where
+        I: IntoIterator,
+        <I as IntoIterator>::Item: StaticallyMatchesPolarsType<T> + Array,
+    {
+        let chunks = iter
+            .into_iter()
+            .map(|x| Box::new(x) as Box<dyn Array>)
+            .collect();
+        unsafe { Self::from_chunks(name, chunks) }
+    }
+
+    pub fn try_from_chunk_iter<I, A, E>(name: &str, iter: I) -> Result<Self, E>
+    where
+        I: IntoIterator<Item = Result<A, E>>,
+        A: StaticallyMatchesPolarsType<T> + Array,
+    {
+        let chunks: Result<_, _> = iter
+            .into_iter()
+            .map(|x| Ok(Box::new(x?) as Box<dyn Array>))
+            .collect();
+        unsafe { Ok(Self::from_chunks(name, chunks?)) }
+    }
+
     /// Create a new ChunkedArray from existing chunks.
     ///
     /// # Safety

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

@@ -68,7 +68,14 @@ impl ListChunked {
         let inner_dtype = self.inner_dtype().to_arrow();
 
         let chunks = ca.downcast_iter().map(|arr| {
-            let elements = unsafe { Series::try_from_arrow_unchecked(self.name(), vec![(*arr.values()).clone()], &inner_dtype).unwrap() } ;
+            let elements = unsafe {
+                Series::try_from_arrow_unchecked(
+                    self.name(),
+                    vec![(*arr.values()).clone()],
+                    &inner_dtype,
+                )
+                .unwrap()
+            };
 
             let expected_len = elements.len();
             let out: Series = func(elements)?;
@@ -86,9 +93,9 @@ impl ListChunked {
                 values,
                 arr.validity().cloned(),
             );
-            Ok(Box::new(arr) as ArrayRef)
-        }).collect::<PolarsResult<Vec<_>>>()?;
+            Ok(arr)
+        });
 
-        unsafe { Ok(ListChunked::from_chunks(self.name(), chunks)) }
+        ListChunked::try_from_chunk_iter(self.name(), chunks)
     }
 }