Separate ArrayReader::next_batch with read_records and consume_batch

parquet/src/arrow/array_reader/byte_array.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::arrow::array_reader::{read_records, skip_records, ArrayReader};
+use crate::arrow::array_reader::{read_records_inner, skip_records, ArrayReader};
 use crate::arrow::buffer::offset_buffer::OffsetBuffer;
 use crate::arrow::record_reader::buffer::ScalarValue;
 use crate::arrow::record_reader::GenericRecordReader;
@@ -109,7 +109,23 @@ impl<I: OffsetSizeTrait + ScalarValue> ArrayReader for ByteArrayReader<I> {
     }
 
     fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
-        read_records(&mut self.record_reader, self.pages.as_mut(), batch_size)?;
+        let size = self.read_records(batch_size)?;
+        self.consume_batch(size)
+    }
+
+    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
+        read_records_inner(&mut self.record_reader, self.pages.as_mut(), batch_size)
+    }
+
+    fn consume_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
+        if self.record_reader.num_records() < batch_size {
+            return Err(general_err!(
+                "Invalid batch_size: {}, current consume: {} records in buffer.",
+                batch_size,
+                self.record_reader.num_records()
+            ));
+        }
+
         let buffer = self.record_reader.consume_record_data();
         let null_buffer = self.record_reader.consume_bitmap_buffer();
         self.def_levels_buffer = self.record_reader.consume_def_levels();

parquet/src/arrow/array_reader/byte_array_dictionary.rs

@@ -25,7 +25,7 @@ use arrow::buffer::Buffer;
 use arrow::datatypes::{ArrowNativeType, DataType as ArrowType};
 
 use crate::arrow::array_reader::byte_array::{ByteArrayDecoder, ByteArrayDecoderPlain};
-use crate::arrow::array_reader::{read_records, ArrayReader, skip_records};
+use crate::arrow::array_reader::{read_records_inner, skip_records, ArrayReader};
 use crate::arrow::buffer::{
     dictionary_buffer::DictionaryBuffer, offset_buffer::OffsetBuffer,
 };
@@ -168,7 +168,24 @@ where
     }
 
     fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
-        read_records(&mut self.record_reader, self.pages.as_mut(), batch_size)?;
+        let size = self.read_records(batch_size)?;
+        self.consume_batch(size)
+    }
+
+    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
+        read_records_inner(&mut self.record_reader, self.pages.as_mut(), batch_size)
+    }
+
+    fn consume_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
+        let num_records = self.record_reader.num_records();
+        if num_records < batch_size {
+            return Err(general_err!(
+                "Invalid batch_size: {}, current consume: {} records in buffer.",
+                batch_size,
+                num_records
+            ));
+        }
+
         let buffer = self.record_reader.consume_record_data();
         let null_buffer = self.record_reader.consume_bitmap_buffer();
         let array = buffer.into_array(null_buffer, &self.data_type)?;

parquet/src/arrow/array_reader/complex_object_array.rs

@@ -21,7 +21,7 @@ use crate::arrow::schema::parquet_to_arrow_field;
 use crate::column::page::PageIterator;
 use crate::column::reader::ColumnReaderImpl;
 use crate::data_type::DataType;
-use crate::errors::Result;
+use crate::errors::{ParquetError, Result};
 use crate::schema::types::ColumnDescPtr;
 use arrow::array::ArrayRef;
 use arrow::datatypes::DataType as ArrowType;
@@ -39,6 +39,7 @@ where
     pages: Box<dyn PageIterator>,
     def_levels_buffer: Option<Vec<i16>>,
     rep_levels_buffer: Option<Vec<i16>>,
+    data_buffer: Vec<T::T>,
     column_desc: ColumnDescPtr,
     column_reader: Option<ColumnReaderImpl<T>>,
     converter: C,
@@ -60,6 +61,11 @@ where
     }
 
     fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
+        let size = self.read_records(batch_size)?;
+        self.consume_batch(size)
+    }
+
+    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
         // Try to initialize column reader
         if self.column_reader.is_none() {
             self.next_column_reader()?;
@@ -126,7 +132,6 @@ where
                 break;
             }
         }
-
         data_buffer.truncate(num_read);
         def_levels_buffer
             .iter_mut()
@@ -135,23 +140,54 @@ where
             .iter_mut()
             .for_each(|buf| buf.truncate(num_read));
 
-        self.def_levels_buffer = def_levels_buffer;
-        self.rep_levels_buffer = rep_levels_buffer;
+        if let Some(mut def_levels_buffer) = def_levels_buffer {
+            match &mut self.def_levels_buffer {
+                None => {
+                    self.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 {
+                None => {
+                    self.rep_levels_buffer = Some(rep_levels_buffer);
+                }
+                Some(buf) => buf.append(&mut rep_levels_buffer),
+            }
+        }
+
+        self.data_buffer.append(&mut data_buffer);
+
+        Ok(num_read)
+    }
+
+    fn consume_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
+        // uncheck null count
+        let len = self.data_buffer.len();
+        if len < batch_size {
+            return Err(general_err!(
+                "Invalid batch_size: {}, current consume: {} records in buffer.",
+                batch_size,
+                len
+            ));
+        }
 
         let data: Vec<Option<T::T>> = if self.def_levels_buffer.is_some() {
-            data_buffer
-                .into_iter()
+            self.data_buffer
+                .iter()
                 .zip(self.def_levels_buffer.as_ref().unwrap().iter())
                 .map(|(t, def_level)| {
                     if *def_level == self.column_desc.max_def_level() {
-                        Some(t)
+                        Some(t.clone())
                     } else {
                         None
                     }
                 })
                 .collect()
         } else {
-            data_buffer.into_iter().map(Some).collect()
+            self.data_buffer.iter().map(|x| Some(x.clone())).collect()
         };
 
         let mut array = self.converter.convert(data)?;
@@ -160,6 +196,15 @@ where
             array = arrow::compute::cast(&array, &self.data_type)?;
         }
 
+        self.data_buffer = self.data_buffer.split_off(batch_size);
+        if let Some(buf) = &mut self.def_levels_buffer {
+            let rest = buf.split_off(batch_size);
+            self.def_levels_buffer = Some(rest);
+        }
+        if let Some(buf) = &mut self.rep_levels_buffer {
+            let rest = buf.split_off(batch_size);
+            self.rep_levels_buffer = Some(rest);
+        }
         Ok(array)
     }
 
@@ -168,8 +213,11 @@ where
             Some(reader) => reader.skip_records(num_records),
             None => {
                 if self.next_column_reader()? {
-                    self.column_reader.as_mut().unwrap().skip_records(num_records)
-                }else {
+                    self.column_reader
+                        .as_mut()
+                        .unwrap()
+                        .skip_records(num_records)
+                } else {
                     Ok(0)
                 }
             }
@@ -208,6 +256,7 @@ where
             pages,
             def_levels_buffer: None,
             rep_levels_buffer: None,
+            data_buffer: vec![],
             column_desc,
             column_reader: None,
             converter,
@@ -349,30 +398,32 @@ mod tests {
 
         let mut accu_len: usize = 0;
 
-        let array = array_reader.next_batch(values_per_page / 2).unwrap();
-        assert_eq!(array.len(), values_per_page / 2);
+        let len = array_reader.read_records(values_per_page / 2).unwrap();
+        assert_eq!(len, values_per_page / 2);
         assert_eq!(
-            Some(&def_levels[accu_len..(accu_len + array.len())]),
+            Some(&def_levels[accu_len..(accu_len + len)]),
             array_reader.get_def_levels()
         );
         assert_eq!(
-            Some(&rep_levels[accu_len..(accu_len + array.len())]),
+            Some(&rep_levels[accu_len..(accu_len + len)]),
             array_reader.get_rep_levels()
         );
-        accu_len += array.len();
+        accu_len += len;
+        array_reader.consume_batch(values_per_page / 2).unwrap();
 
         // Read next values_per_page values, the first values_per_page/2 ones are from the first column chunk,
         // and the last values_per_page/2 ones are from the second column chunk
-        let array = array_reader.next_batch(values_per_page).unwrap();
-        assert_eq!(array.len(), values_per_page);
+        let len = array_reader.read_records(values_per_page).unwrap();
+        assert_eq!(len, values_per_page);
         assert_eq!(
-            Some(&def_levels[accu_len..(accu_len + array.len())]),
+            Some(&def_levels[accu_len..(accu_len + len)]),
             array_reader.get_def_levels()
         );
         assert_eq!(
-            Some(&rep_levels[accu_len..(accu_len + array.len())]),
+            Some(&rep_levels[accu_len..(accu_len + len)]),
             array_reader.get_rep_levels()
         );
+        let array = array_reader.consume_batch(values_per_page).unwrap();
         let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
         for i in 0..array.len() {
             if array.is_valid(i) {
@@ -384,19 +435,20 @@ mod tests {
                 assert_eq!(all_values[i + accu_len], None)
             }
         }
-        accu_len += array.len();
+        accu_len += len;
 
         // Try to read values_per_page values, however there are only values_per_page/2 values
-        let array = array_reader.next_batch(values_per_page).unwrap();
-        assert_eq!(array.len(), values_per_page / 2);
+        let len = array_reader.read_records(values_per_page).unwrap();
+        assert_eq!(len, values_per_page / 2);
         assert_eq!(
-            Some(&def_levels[accu_len..(accu_len + array.len())]),
+            Some(&def_levels[accu_len..(accu_len + len)]),
             array_reader.get_def_levels()
         );
         assert_eq!(
-            Some(&rep_levels[accu_len..(accu_len + array.len())]),
+            Some(&rep_levels[accu_len..(accu_len + len)]),
             array_reader.get_rep_levels()
         );
+        array_reader.consume_batch(len).unwrap();
     }
 
     #[test]
@@ -491,31 +543,34 @@ mod tests {
         let mut accu_len: usize = 0;
 
         // println!("---------- reading a batch of {} values ----------", values_per_page / 2);
-        let array = array_reader.next_batch(values_per_page / 2).unwrap();
-        assert_eq!(array.len(), values_per_page / 2);
+        let len = array_reader.read_records(values_per_page / 2).unwrap();
+        assert_eq!(len, values_per_page / 2);
         assert_eq!(
-            Some(&def_levels[accu_len..(accu_len + array.len())]),
+            Some(&def_levels[accu_len..(accu_len + len)]),
             array_reader.get_def_levels()
         );
         assert_eq!(
-            Some(&rep_levels[accu_len..(accu_len + array.len())]),
+            Some(&rep_levels[accu_len..(accu_len + len)]),
             array_reader.get_rep_levels()
         );
-        accu_len += array.len();
+        accu_len += len;
+        array_reader.consume_batch(len).unwrap();
 
         // Read next values_per_page values, the first values_per_page/2 ones are from the first column chunk,
         // and the last values_per_page/2 ones are from the second column chunk
         // println!("---------- reading a batch of {} values ----------", values_per_page);
-        let array = array_reader.next_batch(values_per_page).unwrap();
-        assert_eq!(array.len(), values_per_page);
+        //let array = array_reader.next_batch(values_per_page).unwrap();
+        let len = array_reader.read_records(values_per_page).unwrap();
+        assert_eq!(len, values_per_page);
         assert_eq!(
-            Some(&def_levels[accu_len..(accu_len + array.len())]),
+            Some(&def_levels[accu_len..(accu_len + len)]),
             array_reader.get_def_levels()
         );
         assert_eq!(
-            Some(&rep_levels[accu_len..(accu_len + array.len())]),
+            Some(&rep_levels[accu_len..(accu_len + len)]),
             array_reader.get_rep_levels()
         );
+        let array = array_reader.consume_batch(len).unwrap();
         let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
         for i in 0..array.len() {
             if array.is_valid(i) {
@@ -527,19 +582,20 @@ mod tests {
                 assert_eq!(all_values[i + accu_len], None)
             }
         }
-        accu_len += array.len();
+        accu_len += len;
 
         // Try to read values_per_page values, however there are only values_per_page/2 values
         // println!("---------- reading a batch of {} values ----------", values_per_page);
-        let array = array_reader.next_batch(values_per_page).unwrap();
-        assert_eq!(array.len(), values_per_page / 2);
+        let len = array_reader.read_records(values_per_page).unwrap();
+        assert_eq!(len, values_per_page / 2);
         assert_eq!(
-            Some(&def_levels[accu_len..(accu_len + array.len())]),
+            Some(&def_levels[accu_len..(accu_len + len)]),
             array_reader.get_def_levels()
         );
         assert_eq!(
-            Some(&rep_levels[accu_len..(accu_len + array.len())]),
+            Some(&rep_levels[accu_len..(accu_len + len)]),
             array_reader.get_rep_levels()
         );
+        array_reader.consume_batch(len).unwrap();
     }
 }

parquet/src/arrow/array_reader/empty_array.rs

@@ -54,14 +54,21 @@ impl ArrayReader for EmptyArrayReader {
     }
 
     fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
+        let size = self.read_records(batch_size)?;
+        self.consume_batch(size)
+    }
+
+    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
         let len = self.remaining_rows.min(batch_size);
         self.remaining_rows -= len;
+        Ok(len)
+    }
 
+    fn consume_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
         let data = ArrayDataBuilder::new(self.data_type.clone())
-            .len(len)
+            .len(batch_size)
             .build()
             .unwrap();
-
         Ok(Arc::new(StructArray::from(data)))
     }
 

parquet/src/arrow/array_reader/list_array.rs

@@ -79,8 +79,17 @@ impl<OffsetSize: OffsetSizeTrait> ArrayReader for ListArrayReader<OffsetSize> {
     }
 
     fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
-        let next_batch_array = self.item_reader.next_batch(batch_size)?;
+        let size = self.read_records(batch_size)?;
+        self.consume_batch(size)
+    }
+
+    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
+        let size = self.item_reader.read_records(batch_size)?;
+        Ok(size)
+    }
 
+    fn consume_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
+        let next_batch_array = self.item_reader.consume_batch(batch_size)?;
         if next_batch_array.len() == 0 {
             return Ok(new_empty_array(&self.data_type));
         }