fix: Fix split_chunks for nested dtypes (#16493)

crates/polars-arrow/src/record_batch.rs

@@ -34,7 +34,7 @@ impl<A: AsRef<dyn Array>> RecordBatchT<A> {
                 .any(|array| array.len() != len)
             {
                 polars_bail!(ComputeError:
-                    "Chunk require all its arrays to have an equal number of rows".to_string(),
+                    "RecordBatch requires all its arrays to have an equal number of rows".to_string(),
                 );
             }
         }

crates/polars-core/src/frame/chunks.rs

@@ -30,11 +30,11 @@ impl DataFrame {
                 .get_columns()
                 .iter()
                 .map(|s| {
-                    Series::from_chunks_and_dtype_unchecked(
-                        s.name(),
-                        vec![s.chunks()[i].clone()],
-                        s.dtype(),
-                    )
+                    if s.n_chunks() == 1 {
+                        s.clone()
+                    } else {
+                        s.replace_with_chunk(s.chunks()[i].clone())
+                    }
                 })
                 .collect::<Vec<_>>();
 

crates/polars-core/src/series/mod.rs

@@ -189,6 +189,32 @@ impl Series {
         ca.chunks_mut()
     }
 
+    /// Create a `Series` of the same data type with all chunks replaced.
+    /// # Safety
+    /// These chunks should align with the data-type
+    pub unsafe fn replace_chunks(&self, chunks: Vec<ArrayRef>) -> Self {
+        let mut new = self.clear();
+        // Assign mut so we go through arc only once.
+        let mut_new = new._get_inner_mut();
+        *mut_new.chunks_mut() = chunks;
+        mut_new.compute_len();
+        new
+    }
+
+    /// Create a `Series` of the same data type with all chunks replaced.
+    /// # Safety
+    /// This chunk should align with the data-type
+    pub unsafe fn replace_with_chunk(&self, chunk: ArrayRef) -> Self {
+        let mut new = self.clear();
+        // Assign mut so we go through arc only once.
+        let mut_new = new._get_inner_mut();
+        let chunks = mut_new.chunks_mut();
+        chunks.clear();
+        chunks.push(chunk);
+        mut_new.compute_len();
+        new
+    }
+
     pub fn is_sorted_flag(&self) -> IsSorted {
         if self.len() <= 1 {
             return IsSorted::Ascending;