Separate ArrayReader::next_batch with read_records and consume_batch

parquet/src/arrow/array_reader/complex_object_array.rs

@@ -43,6 +43,9 @@ where
     column_desc: ColumnDescPtr,
     column_reader: Option<ColumnReaderImpl<T>>,
     converter: C,
+    need_consume_def_levels_buffer: Option<Vec<i16>>,
+    need_consume_rep_levels_buffer: Option<Vec<i16>>,
+    before_consume: bool,
     _parquet_type_marker: PhantomData<T>,
     _converter_marker: PhantomData<C>,
 }
@@ -61,6 +64,9 @@ where
     }
 
     fn read_records(&mut self, batch_size: usize) -> Result<usize> {
+        if !self.before_consume {
+            self.before_consume = true;
+        }
         // Try to initialize column reader
         if self.column_reader.is_none() {
             self.next_column_reader()?;
@@ -136,18 +142,18 @@ where
             .for_each(|buf| buf.truncate(num_read));
 
         if let Some(mut def_levels_buffer) = def_levels_buffer {
-            match &mut self.def_levels_buffer {
+            match &mut self.need_consume_def_levels_buffer {
                 None => {
-                    self.def_levels_buffer = Some(def_levels_buffer);
+                    self.need_consume_def_levels_buffer = Some(def_levels_buffer);
                 }
                 Some(buf) => buf.append(&mut def_levels_buffer),
             }
         }
 
         if let Some(mut rep_levels_buffer) = rep_levels_buffer {
-            match &mut self.rep_levels_buffer {
+            match &mut self.need_consume_rep_levels_buffer {
                 None => {
-                    self.rep_levels_buffer = Some(rep_levels_buffer);
+                    self.need_consume_rep_levels_buffer = Some(rep_levels_buffer);
                 }
                 Some(buf) => buf.append(&mut rep_levels_buffer),
             }
@@ -159,11 +165,11 @@ where
     }
 
     fn consume_batch(&mut self) -> Result<ArrayRef> {
-        let data: Vec<Option<T::T>> = if self.def_levels_buffer.is_some() {
+        let data: Vec<Option<T::T>> = if self.need_consume_def_levels_buffer.is_some() {
             let data_buffer = std::mem::take(&mut self.data_buffer);
             data_buffer
                 .into_iter()
-                .zip(self.def_levels_buffer.as_ref().unwrap().iter())
+                .zip(self.need_consume_def_levels_buffer.as_ref().unwrap().iter())
                 .map(|(t, def_level)| {
                     if *def_level == self.column_desc.max_def_level() {
                         Some(t)
@@ -183,8 +189,9 @@ where
         }
 
         self.data_buffer = vec![];
-        self.def_levels_buffer = None;
-        self.rep_levels_buffer = None;
+        self.def_levels_buffer = std::mem::take(&mut self.need_consume_def_levels_buffer);
+        self.rep_levels_buffer = std::mem::take(&mut self.need_consume_rep_levels_buffer);
+        self.before_consume = false;
 
         Ok(array)
     }
@@ -206,11 +213,19 @@ where
     }
 
     fn get_def_levels(&self) -> Option<&[i16]> {
-        self.def_levels_buffer.as_deref()
+        if self.before_consume {
+            self.need_consume_def_levels_buffer.as_deref()
+        } else {
+            self.def_levels_buffer.as_deref()
+        }
     }
 
     fn get_rep_levels(&self) -> Option<&[i16]> {
-        self.rep_levels_buffer.as_deref()
+        if self.before_consume {
+            self.need_consume_rep_levels_buffer.as_deref()
+        } else {
+            self.rep_levels_buffer.as_deref()
+        }
     }
 }
 
@@ -241,6 +256,9 @@ where
             column_desc,
             column_reader: None,
             converter,
+            need_consume_def_levels_buffer: None,
+            need_consume_rep_levels_buffer: None,
+            before_consume: true,
             _parquet_type_marker: PhantomData,
             _converter_marker: PhantomData,
         })

parquet/src/arrow/arrow_reader.rs

@@ -769,6 +769,44 @@ mod tests {
         assert_eq!(&written.slice(6, 2), &read[2]);
     }
 
+    #[test]
+    fn test_int32_nullable_struct() {
+        let int32 = Int32Array::from_iter_values([1, 2, 3, 4, 5, 6, 7, 8]);
+        let data = ArrayDataBuilder::new(ArrowDataType::Struct(vec![Field::new(
+            "int32",
+            int32.data_type().clone(),
+            false,
+        )]))
+        .len(8)
+        .null_bit_buffer(Some(Buffer::from(&[0b11101111])))
+        .child_data(vec![int32.into_data()])
+        .build()
+        .unwrap();
+
+        let written = RecordBatch::try_from_iter([(
+            "struct",
+            Arc::new(StructArray::from(data)) as ArrayRef,
+        )])
+        .unwrap();
+
+        let mut buffer = Vec::with_capacity(1024);
+        let mut writer =
+            ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap();
+        writer.write(&written).unwrap();
+        writer.close().unwrap();
+
+        let read = ParquetFileArrowReader::try_new(Bytes::from(buffer))
+            .unwrap()
+            .get_record_reader(3)
+            .unwrap()
+            .collect::<ArrowResult<Vec<_>>>()
+            .unwrap();
+
+        assert_eq!(&written.slice(0, 3), &read[0]);
+        assert_eq!(&written.slice(3, 3), &read[1]);
+        assert_eq!(&written.slice(6, 2), &read[2]);
+    }
+
     #[test]
     #[ignore] // https://github.com/apache/arrow-rs/issues/2253
     fn test_decimal_list() {

parquet/src/arrow/arrow_writer/mod.rs

@@ -1161,7 +1161,7 @@ mod tests {
             Some(props),
         )
         .expect("Unable to write file");
-        writer.write(&expected_batch).unwrap();
+        writer.write(expected_batch).unwrap();
         writer.close().unwrap();
 
         let mut arrow_reader =