From 2e126d4ea0107d9cdd94a8c2c77b493e27b88aa3 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Sun, 25 Dec 2022 13:27:48 +0100
Subject: [PATCH] fix(parquet): nested struct /list writing

---
 src/io/parquet/write/nested/rep.rs | 32 ++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/io/parquet/write/nested/rep.rs b/src/io/parquet/write/nested/rep.rs
index 68c7609acae..c7c5d15eb8e 100644
--- a/src/io/parquet/write/nested/rep.rs
+++ b/src/io/parquet/write/nested/rep.rs
@@ -8,14 +8,21 @@ impl<A: Iterator<Item = usize> + std::fmt::Debug> DebugIter for A {}
 fn iter<'a>(nested: &'a [Nested]) -> Vec<Box<dyn DebugIter + 'a>> {
     nested
         .iter()
-        .filter_map(|nested| match nested {
+        .enumerate()
+        .filter_map(|(i, nested)| match nested {
             Nested::Primitive(_, _, _) => None,
             Nested::List(nested) => Some(Box::new(to_length(nested.offsets)) as Box<dyn DebugIter>),
             Nested::LargeList(nested) => {
                 Some(Box::new(to_length(nested.offsets)) as Box<dyn DebugIter>)
             }
             Nested::Struct(_, _, length) => {
-                Some(Box::new(std::iter::repeat(1usize).take(*length)) as Box<dyn DebugIter>)
+                // only return 1, 1, 1, (x len) if struct is outer structure.
+                // otherwise treat as leaf
+                if i == 0 {
+                    Some(Box::new(std::iter::repeat(1usize).take(*length)) as Box<dyn DebugIter>)
+                } else {
+                    None
+                }
             }
         })
         .collect()
@@ -223,4 +230,25 @@ mod tests {
 
         test(nested, expected)
     }
+
+    #[test]
+    fn list_struct_list() {
+        let nested = vec![
+            Nested::List(ListNested {
+                is_optional: true,
+                offsets: &[0i32, 2, 3],
+                validity: None,
+            }),
+            Nested::Struct(None, true, 3),
+            Nested::List(ListNested {
+                is_optional: true,
+                offsets: &[0i32, 3, 6, 7],
+                validity: None,
+            }),
+            Nested::Primitive(None, true, 7),
+        ];
+        let expected = vec![0, 2, 2, 1, 2, 2, 0];
+
+        test(nested, expected)
+    }
 }