Fixed reading arrays from parquet with required children

parquet_integration/bench_compute.py

@@ -26,7 +26,7 @@ def get_f32(size, null_density):
 
 
 def bench_add_f32_pyarrow(log2_size, null_density):
-    size = 2 ** log2_size
+    size = 2**log2_size
 
     values, validity = get_f32(size, null_density)
     array1 = pa.array(values, pa.float32(), mask=validity)
@@ -50,7 +50,7 @@ def f():
 
 
 def bench_add_f32_numpy(log2_size):
-    size = 2 ** log2_size
+    size = 2**log2_size
 
     array1, _ = get_f32(size, 0)
     array2, _ = get_f32(size, 0)
@@ -71,7 +71,7 @@ def f():
 
 
 def _bench_unary_f32_pyarrow(log2_size, null_density, name, op):
-    size = 2 ** log2_size
+    size = 2**log2_size
 
     values, validity = get_f32(size, null_density)
     array = pa.array(values, pa.float32(), mask=validity)
@@ -100,7 +100,7 @@ def bench_min_f32_pyarrow(log2_size, null_density):
 
 
 def _bench_unary_f32_numpy(log2_size, name, op):
-    size = 2 ** log2_size
+    size = 2**log2_size
 
     values, _ = get_f32(size, 0)
 
@@ -128,7 +128,7 @@ def bench_min_f32_numpy(log2_size):
 
 
 def bench_sort_f32_pyarrow(log2_size, null_density):
-    size = 2 ** log2_size
+    size = 2**log2_size
 
     values, validity = get_f32(size, null_density)
     array = pa.array(values, pa.float32(), mask=validity)
@@ -151,7 +151,7 @@ def f():
 
 def bench_sort_f32_numpy(log2_size):
     null_density = 0
-    size = 2 ** log2_size
+    size = 2**log2_size
 
     array, _ = get_f32(size, null_density)
 
@@ -171,7 +171,7 @@ def f():
 
 
 def bench_filter_f32_pyarrow(log2_size, null_density):
-    size = 2 ** log2_size
+    size = 2**log2_size
 
     values, validity = get_f32(size, null_density)
     _, mask = get_f32(size, 0.9)
@@ -198,7 +198,7 @@ def f():
 
 def bench_filter_f32_numpy(log2_size):
     null_density = 0
-    size = 2 ** log2_size
+    size = 2**log2_size
 
     array, _ = get_f32(size, null_density)
     _, mask = get_f32(size, 0.1)

parquet_integration/write_parquet.py

@@ -223,20 +223,31 @@ def case_struct() -> Tuple[dict, pa.Schema, str]:
                     ]
                 ),
             ),
+            pa.field(
+                "struct_nullable",
+                pa.struct(struct_fields),
+            ),
         ]
     )
 
     struct = pa.StructArray.from_arrays(
         [pa.array(string), pa.array(boolean)],
         fields=struct_fields,
     )
+    struct_nullable = pa.StructArray.from_arrays(
+        [pa.array(string), pa.array(boolean)],
+        fields=struct_fields,
+        mask=pa.array([False, False, True, False, False, False, False, False, False, False]),
+    )
+
     return (
         {
             "struct": struct,
             "struct_struct": pa.StructArray.from_arrays(
                 [struct, pa.array(boolean)],
                 names=["f1", "f2"],
             ),
+            "struct_nullable": struct_nullable,
         },
         schema,
         f"struct_nullable_10.parquet",
@@ -307,16 +318,16 @@ def write_pyarrow(
         base_path = f"{base_path}/{compression}"
 
     if multiple_pages:
-        data_page_size = 2 ** 10  # i.e. a small number to ensure multiple pages
+        data_page_size = 2**10  # i.e. a small number to ensure multiple pages
     else:
-        data_page_size = 2 ** 40  # i.e. a large number to ensure a single page
+        data_page_size = 2**40  # i.e. a large number to ensure a single page
 
     t = pa.table(data, schema=schema)
     os.makedirs(base_path, exist_ok=True)
     pa.parquet.write_table(
         t,
         f"{base_path}/{path}",
-        row_group_size=2 ** 40,
+        row_group_size=2**40,
         use_dictionary=use_dictionary,
         compression=compression,
         write_statistics=True,
@@ -325,7 +336,14 @@ def write_pyarrow(
     )
 
 
-for case in [case_basic_nullable, case_basic_required, case_nested, case_struct, case_nested_edge, case_map]:
+for case in [
+    case_basic_nullable,
+    case_basic_required,
+    case_nested,
+    case_struct,
+    case_nested_edge,
+    case_map,
+]:
     for version in [1, 2]:
         for use_dict in [True, False]:
             for compression in ["lz4", None, "snappy"]:
@@ -351,14 +369,14 @@ def case_benches_required(size):
 # for read benchmarks
 for i in range(10, 22, 2):
     # two pages (dict)
-    write_pyarrow(case_benches(2 ** i), 1, True, False, None)
+    write_pyarrow(case_benches(2**i), 1, True, False, None)
     # single page
-    write_pyarrow(case_benches(2 ** i), 1, False, False, None)
+    write_pyarrow(case_benches(2**i), 1, False, False, None)
     # single page required
-    write_pyarrow(case_benches_required(2 ** i), 1, False, False, None)
+    write_pyarrow(case_benches_required(2**i), 1, False, False, None)
     # multiple pages
-    write_pyarrow(case_benches(2 ** i), 1, False, True, None)
+    write_pyarrow(case_benches(2**i), 1, False, True, None)
     # multiple compressed pages
-    write_pyarrow(case_benches(2 ** i), 1, False, True, "snappy")
+    write_pyarrow(case_benches(2**i), 1, False, True, "snappy")
     # single compressed page
-    write_pyarrow(case_benches(2 ** i), 1, False, False, "snappy")
+    write_pyarrow(case_benches(2**i), 1, False, False, "snappy")

src/io/parquet/read/deserialize/binary/nested.rs

@@ -17,25 +17,24 @@ use super::basic::ValuesDictionary;
 use super::utils::*;
 use super::{
     super::utils,
-    basic::{finish, Required, TraitBinaryArray},
+    basic::{finish, TraitBinaryArray},
 };
 
-#[allow(clippy::large_enum_variant)]
 #[derive(Debug)]
 enum State<'a> {
-    Optional(Optional<'a>, BinaryIter<'a>),
-    Required(Required<'a>),
+    Optional(BinaryIter<'a>),
+    Required(BinaryIter<'a>),
     RequiredDictionary(ValuesDictionary<'a>),
-    OptionalDictionary(Optional<'a>, ValuesDictionary<'a>),
+    OptionalDictionary(ValuesDictionary<'a>),
 }
 
 impl<'a> utils::PageState<'a> for State<'a> {
     fn len(&self) -> usize {
         match self {
-            State::Optional(validity, _) => validity.len(),
-            State::Required(state) => state.len(),
+            State::Optional(validity) => validity.size_hint().0,
+            State::Required(state) => state.size_hint().0,
             State::RequiredDictionary(required) => required.len(),
-            State::OptionalDictionary(optional, _) => optional.len(),
+            State::OptionalDictionary(optional) => optional.len(),
         }
     }
 }
@@ -45,7 +44,7 @@ struct BinaryDecoder<O: Offset> {
     phantom_o: std::marker::PhantomData<O>,
 }
 
-impl<'a, O: Offset> utils::Decoder<'a> for BinaryDecoder<O> {
+impl<'a, O: Offset> NestedDecoder<'a> for BinaryDecoder<O> {
     type State = State<'a>;
     type DecodedState = (Binary<O>, MutableBitmap);
 
@@ -62,25 +61,26 @@ impl<'a, O: Offset> utils::Decoder<'a> for BinaryDecoder<O> {
         ) {
             (Encoding::PlainDictionary | Encoding::RleDictionary, Some(dict), false, false) => {
                 let dict = dict.as_any().downcast_ref().unwrap();
-                Ok(State::RequiredDictionary(ValuesDictionary::try_new(
-                    page, dict,
-                )?))
+                ValuesDictionary::try_new(page, dict).map(State::RequiredDictionary)
             }
             (Encoding::PlainDictionary | Encoding::RleDictionary, Some(dict), true, false) => {
                 let dict = dict.as_any().downcast_ref().unwrap();
-                Ok(State::OptionalDictionary(
-                    Optional::try_new(page)?,
-                    ValuesDictionary::try_new(page, dict)?,
-                ))
+                ValuesDictionary::try_new(page, dict).map(State::OptionalDictionary)
             }
             (Encoding::Plain, _, true, false) => {
                 let (_, _, values) = split_buffer(page)?;
 
                 let values = BinaryIter::new(values);
 
-                Ok(State::Optional(Optional::try_new(page)?, values))
+                Ok(State::Optional(values))
+            }
+            (Encoding::Plain, _, false, false) => {
+                let (_, _, values) = split_buffer(page)?;
+
+                let values = BinaryIter::new(values);
+
+                Ok(State::Required(values))
             }
-            (Encoding::Plain, _, false, false) => Ok(State::Required(Required::try_new(page)?)),
             _ => Err(utils::not_implemented(page)),
         }
     }
@@ -92,24 +92,17 @@ impl<'a, O: Offset> utils::Decoder<'a> for BinaryDecoder<O> {
         )
     }
 
-    fn extend_from_state(
-        &self,
-        state: &mut Self::State,
-        decoded: &mut Self::DecodedState,
-        additional: usize,
-    ) {
+    fn push_valid(&self, state: &mut Self::State, decoded: &mut Self::DecodedState) {
         let (values, validity) = decoded;
         match state {
-            State::Optional(page_validity, page_values) => {
-                let items = page_validity.by_ref().take(additional);
-                let items = Zip::new(items, page_values.by_ref());
-
-                read_optional_values(items, values, validity)
+            State::Optional(page) => {
+                let value = page.next().unwrap_or_default();
+                values.push(value);
+                validity.push(true);
             }
             State::Required(page) => {
-                for x in page.values.by_ref().take(additional) {
-                    values.push(x)
-                }
+                let value = page.next().unwrap_or_default();
+                values.push(value);
             }
             State::RequiredDictionary(page) => {
                 let dict_values = page.dict.values();
@@ -121,36 +114,38 @@ impl<'a, O: Offset> utils::Decoder<'a> for BinaryDecoder<O> {
                     let dict_offset_ip1 = dict_offsets[index + 1] as usize;
                     &dict_values[dict_offset_i..dict_offset_ip1]
                 };
-                for x in page.values.by_ref().map(op).take(additional) {
-                    values.push(x)
-                }
+                let item = page.values.next().map(op).unwrap_or_default();
+                values.push(item);
             }
-            State::OptionalDictionary(page_validity, page_values) => {
-                let dict_values = page_values.dict.values();
-                let dict_offsets = page_values.dict.offsets();
+            State::OptionalDictionary(page) => {
+                let dict_values = page.dict.values();
+                let dict_offsets = page.dict.offsets();
 
                 let op = move |index: u32| {
                     let index = index as usize;
                     let dict_offset_i = dict_offsets[index] as usize;
                     let dict_offset_ip1 = dict_offsets[index + 1] as usize;
                     &dict_values[dict_offset_i..dict_offset_ip1]
                 };
-
-                let items = page_validity.by_ref().take(additional);
-                let items = Zip::new(items, page_values.values.by_ref().map(op));
-
-                read_optional_values(items, values, validity)
+                let item = page.values.next().map(op).unwrap_or_default();
+                values.push(item);
+                validity.push(true);
             }
         }
     }
+
+    fn push_null(&self, decoded: &mut Self::DecodedState) {
+        let (values, validity) = decoded;
+        values.push(&[]);
+        validity.push(false);
+    }
 }
 
 pub struct ArrayIterator<O: Offset, A: TraitBinaryArray<O>, I: DataPages> {
     iter: I,
     data_type: DataType,
     init: Vec<InitNested>,
-    items: VecDeque<(Binary<O>, MutableBitmap)>,
-    nested: VecDeque<NestedState>,
+    items: VecDeque<(NestedState, (Binary<O>, MutableBitmap))>,
     chunk_size: Option<usize>,
     phantom_a: std::marker::PhantomData<A>,
 }
@@ -167,7 +162,6 @@ impl<O: Offset, A: TraitBinaryArray<O>, I: DataPages> ArrayIterator<O, A, I> {
             data_type,
             init,
             items: VecDeque::new(),
-            nested: VecDeque::new(),
             chunk_size,
             phantom_a: Default::default(),
         }
@@ -181,7 +175,6 @@ impl<O: Offset, A: TraitBinaryArray<O>, I: DataPages> Iterator for ArrayIterator
         let maybe_state = next(
             &mut self.iter,
             &mut self.items,
-            &mut self.nested,
             &self.init,
             self.chunk_size,
             &BinaryDecoder::<O>::default(),