fix nested struct append (#4217)

pola-rs · Aug 2, 2022 · b2ff9fa · b2ff9fa
1 parent 53edb19
commit b2ff9fa
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 25 deletions.
diff --git a/polars/polars-core/src/chunked_array/logical/struct_/mod.rs b/polars/polars-core/src/chunked_array/logical/struct_/mod.rs
@@ -46,10 +46,6 @@ impl StructChunked {
         }
     }
 
-    pub(crate) fn arrow_array(&self) -> &ArrayRef {
-        &self.chunks[0]
-    }
-
     pub(crate) fn chunks(&self) -> &Vec<ArrayRef> {
         &self.chunks
     }

diff --git a/polars/polars-core/src/series/into.rs b/polars/polars-core/src/series/into.rs
@@ -10,14 +10,7 @@ use polars_arrow::compute::cast::cast;
 impl Series {
     /// Returns a reference to the Arrow ArrayRef
     pub fn array_ref(&self, chunk_idx: usize) -> &ArrayRef {
-        match self.dtype() {
-            #[cfg(feature = "dtype-struct")]
-            DataType::Struct(_) => {
-                let ca = self.struct_().unwrap();
-                ca.arrow_array()
-            }
-            _ => &self.chunks()[chunk_idx] as &ArrayRef,
-        }
+        &self.chunks()[chunk_idx] as &ArrayRef
     }
 
     /// Convert a chunk in the Series to the correct Arrow type.
@@ -78,8 +71,6 @@ impl Series {
             }
             #[cfg(feature = "dtype-time")]
             DataType::Time => cast(&*self.chunks()[chunk_idx], &DataType::Time.to_arrow()).unwrap(),
-            #[cfg(feature = "dtype-struct")]
-            DataType::Struct(_) => self.array_ref(chunk_idx).clone(),
             _ => self.array_ref(chunk_idx).clone(),
         }
     }

diff --git a/polars/polars-core/src/series/iterator.rs b/polars/polars-core/src/series/iterator.rs
@@ -59,19 +59,14 @@ impl FromIterator<String> for Series {
 
 #[cfg(any(feature = "rows", feature = "dtype-struct"))]
 impl Series {
+    /// iterate over [`Series`] as [`AnyValue`].
+    ///
+    /// # Panics
+    /// This will panic if the array is not rechunked first.
     pub fn iter(&self) -> SeriesIter<'_> {
         let dtype = self.dtype();
-        let arr = match dtype {
-            #[cfg(feature = "dtype-struct")]
-            DataType::Struct(_) => {
-                let ca = self.struct_().unwrap();
-                &**ca.arrow_array()
-            }
-            _ => {
-                assert_eq!(self.chunks().len(), 1, "impl error");
-                &*self.chunks()[0]
-            }
-        };
+        assert_eq!(self.chunks().len(), 1, "impl error");
+        let arr = &*self.chunks()[0];
         let len = arr.len();
         SeriesIter {
             arr,

diff --git a/py-polars/src/conversion.rs b/py-polars/src/conversion.rs
@@ -443,6 +443,9 @@ impl ToPyObject for Wrap<&Utf8Chunked> {
 impl ToPyObject for Wrap<&StructChunked> {
     fn to_object(&self, py: Python) -> PyObject {
         let s = self.0.clone().into_series();
+        // todo! iterate its chunks and flatten.
+        // make series::iter() accept a chunk index.
+        let s = s.rechunk();
         let iter = s.iter().map(|av| {
             if let AnyValue::Struct(vals, flds) = av {
                 struct_dict(py, vals, flds)

diff --git a/py-polars/tests/test_struct.py b/py-polars/tests/test_struct.py
@@ -566,3 +566,44 @@ def test_nested_explode_4026() -> None:
         ],
         "day": ["monday", "monday"],
     }
+
+
+def test_nested_struct_sliced_append() -> None:
+    s = pl.Series(
+        [
+            {
+                "_experience": {
+                    "aaid": {
+                        "id": "A",
+                        "namespace": {"code": "alpha"},
+                    }
+                }
+            },
+            {
+                "_experience": {
+                    "aaid": {
+                        "id": "B",
+                        "namespace": {"code": "bravo"},
+                    },
+                }
+            },
+            {
+                "_experience": {
+                    "aaid": {
+                        "id": "D",
+                        "namespace": {"code": "delta"},
+                    }
+                }
+            },
+        ]
+    )
+    s2 = s[1:]
+    s.append(s2)
+
+    assert s.to_list() == [
+        {"_experience": {"aaid": {"id": "A", "namespace": {"code": "alpha"}}}},
+        {"_experience": {"aaid": {"id": "B", "namespace": {"code": "bravo"}}}},
+        {"_experience": {"aaid": {"id": "D", "namespace": {"code": "delta"}}}},
+        {"_experience": {"aaid": {"id": "B", "namespace": {"code": "bravo"}}}},
+        {"_experience": {"aaid": {"id": "D", "namespace": {"code": "delta"}}}},
+    ]