revamp array building

pola-rs · Aug 22, 2023 · 109efbd · 109efbd
1 parent 69bd551
commit 109efbd
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 18 deletions.
diff --git a/crates/polars-core/src/chunked_array/builder/fixed_size_list.rs b/crates/polars-core/src/chunked_array/builder/fixed_size_list.rs
@@ -27,12 +27,24 @@ impl<T: NativeType> FixedSizeListNumericBuilder<T> {
 }
 
 pub(crate) trait FixedSizeListBuilder {
+    fn push(&mut self, arr: ArrayRef) -> PolarsResult<()>;
     unsafe fn push_unchecked(&mut self, arr: &dyn Array, offset: usize);
     unsafe fn push_null(&mut self);
     fn finish(&mut self) -> ArrayChunked;
 }
 
 impl<T: NativeType> FixedSizeListBuilder for FixedSizeListNumericBuilder<T> {
+    #[inline]
+    fn push(&mut self, arr: ArrayRef) -> PolarsResult<()> {
+        let inner = self.inner.as_mut().unwrap();
+        let arr = arr
+            .as_any()
+            .downcast_ref::<PrimitiveArray<T>>()
+            .ok_or_else(|| polars_err!(ComputeError : "failed to downcast array"))?;
+
+        Ok(inner.try_push(Some(arr.iter().map(|x| x.copied())))?)
+    }
+
     #[inline]
     unsafe fn push_unchecked(&mut self, arr: &dyn Array, offset: usize) {
         let start = offset * self.width;
@@ -95,6 +107,11 @@ impl AnonymousOwnedFixedSizeListBuilder {
 }
 
 impl FixedSizeListBuilder for AnonymousOwnedFixedSizeListBuilder {
+    #[inline]
+    fn push(&mut self, arr: ArrayRef) -> PolarsResult<()> {
+        self.inner.push(arr);
+        Ok(())
+    }
     #[inline]
     unsafe fn push_unchecked(&mut self, arr: &dyn Array, offset: usize) {
         let arr = arr.sliced_unchecked(offset * self.inner.width, self.inner.width);

diff --git a/crates/polars-core/src/serde/mod.rs b/crates/polars-core/src/serde/mod.rs
@@ -4,8 +4,6 @@ pub mod series;
 
 #[cfg(test)]
 mod test {
-    #[cfg(feature = "dtype-array")]
-    use crate::chunked_array::builder::get_fixed_size_list_builder;
     use crate::chunked_array::Settings;
     use crate::prelude::*;
     use crate::series::implementations::null::NullChunked;
@@ -117,17 +115,13 @@ mod test {
     #[test]
     #[cfg(feature = "dtype-array")]
     fn test_serde_array_owned_json() {
-        let arrays = vec![
-            UInt8Chunked::new("test", vec![1u8, 2, 3, 4, 5]),
-            UInt8Chunked::new("test2", vec![2u8, 3, 4, 5, 6]),
-        ];
-        let mut builder = get_fixed_size_list_builder(&DataType::UInt8, 2, 5, "array").unwrap();
-        for (i, ca) in arrays.iter().enumerate() {
-            unsafe {
-                builder.push_unchecked(ca.chunks().get_unchecked(0).as_ref(), i);
-            }
-        }
-        let s = builder.finish().into_series();
+        let input_s = Series::new("test", vec![1u8, 2, 3, 4, 5, 6, 7, 8, 9]).rechunk();
+        let dtype = ArrowDataType::FixedSizeList(
+            Box::new(ArrowField::new("test", ArrowDataType::UInt8, false)),
+            3,
+        );
+        let arr = FixedSizeListArray::new(dtype, input_s.to_arrow(0), None);
+        let s = Series::try_from(("test", arr.to_boxed())).unwrap();
         let json = serde_json::to_string(&s).unwrap();
         let out = serde_json::from_reader::<_, Series>(json.as_bytes()).unwrap();
         assert_eq!(out, s);

diff --git a/crates/polars-core/src/serde/series.rs b/crates/polars-core/src/serde/series.rs
@@ -256,13 +256,18 @@ impl<'de> Deserialize<'de> for Series {
                             .unwrap())
                     },
                     #[cfg(feature = "dtype-array")]
-                    DataType::Array(dt, size) => unsafe {
-                        let mut values: Vec<Series> = map.next_value()?;
+                    DataType::Array(dt, size) => {
+                        let mut values: Vec<Option<Series>> = map.next_value()?;
                         let mut builder =
                             get_fixed_size_list_builder(&dt, values.len(), size, &name).unwrap();
-                        for (i, s) in values.iter_mut().enumerate() {
-                            let s = s.rechunk();
-                            builder.push_unchecked(s.to_arrow(0).as_ref(), i)
+
+                        for s in values.iter_mut() {
+                            match s {
+                                None => unsafe { builder.push_null() },
+                                Some(s) => {
+                                    builder.push(s.rechunk().to_arrow(0)).unwrap();
+                                },
+                            }
                         }
                         Ok(builder.finish().into_series())
                     },