Add integration test for scan rows with selection

parquet/src/arrow/array_reader/byte_array.rs

@@ -120,7 +120,7 @@ impl<I: OffsetSizeTrait + ScalarValue> ArrayReader for ByteArrayReader<I> {
     }
 
     fn skip_records(&mut self, num_records: usize) -> Result<usize> {
-        self.record_reader.skip_records(num_records)
+        self.record_reader.skip_records(num_records, self.pages.as_mut())
     }
 
     fn get_def_levels(&self) -> Option<&[i16]> {

parquet/src/arrow/array_reader/byte_array_dictionary.rs

@@ -181,7 +181,7 @@ where
     }
 
     fn skip_records(&mut self, num_records: usize) -> Result<usize> {
-        self.record_reader.skip_records(num_records)
+        self.record_reader.skip_records(num_records, self.pages.as_mut())
     }
 
     fn get_def_levels(&self) -> Option<&[i16]> {

parquet/src/arrow/array_reader/complex_object_array.rs

@@ -166,7 +166,13 @@ where
     fn skip_records(&mut self, num_records: usize) -> Result<usize> {
         match self.column_reader.as_mut() {
             Some(reader) => reader.skip_records(num_records),
-            None => Ok(0),
+            None => {
+                if self.next_column_reader()? {
+                    self.column_reader.as_mut().unwrap().skip_records(num_records)
+                }else {
+                    Ok(0)
+                }
+            }
         }
     }
 

parquet/src/arrow/array_reader/null_array.rs

@@ -97,7 +97,7 @@ where
     }
 
     fn skip_records(&mut self, num_records: usize) -> Result<usize> {
-        self.record_reader.skip_records(num_records)
+        self.record_reader.skip_records(num_records, self.pages.as_mut())
     }
 
     fn get_def_levels(&self) -> Option<&[i16]> {

parquet/src/arrow/array_reader/primitive_array.rs

@@ -222,7 +222,7 @@ where
     }
 
     fn skip_records(&mut self, num_records: usize) -> Result<usize> {
-        self.record_reader.skip_records(num_records)
+        self.record_reader.skip_records(num_records, self.pages.as_mut())
     }
 
     fn get_def_levels(&self) -> Option<&[i16]> {

parquet/src/arrow/arrow_reader.rs

@@ -33,6 +33,7 @@ use crate::arrow::ProjectionMask;
 use crate::errors::{ParquetError, Result};
 use crate::file::metadata::{KeyValue, ParquetMetaData};
 use crate::file::reader::{ChunkReader, FileReader, SerializedFileReader};
+use crate::file::serialized_reader::ReadOptionsBuilder;
 use crate::schema::types::SchemaDescriptor;
 
 /// Arrow reader api.
@@ -217,7 +218,15 @@ impl ParquetFileArrowReader {
         chunk_reader: R,
         options: ArrowReaderOptions,
     ) -> Result<Self> {
-        let file_reader = Arc::new(SerializedFileReader::new(chunk_reader)?);
+        let file_reader = if options.selection.is_some() {
+            let options = ReadOptionsBuilder::new().with_page_index().build();
+            Arc::new(SerializedFileReader::new_with_options(
+                chunk_reader,
+                options,
+            )?)
+        } else {
+            Arc::new(SerializedFileReader::new(chunk_reader)?)
+        };
         Ok(Self::new_with_options(file_reader, options))
     }
 
@@ -298,9 +307,14 @@ impl Iterator for ParquetRecordBatchReader {
                     continue;
                 }
 
+                // try to read record
                 let to_read = match front.row_count.checked_sub(self.batch_size) {
                     Some(remaining) => {
-                        selection.push_front(RowSelection::skip(remaining));
+                        // if page row count less than batch_size we must set batch size to page row count.
+                        // add check avoid dead loop
+                        if remaining != 0 {
+                            selection.push_front(RowSelection::select(remaining));
+                        }
                         self.batch_size
                     }
                     None => front.row_count,
@@ -390,6 +404,7 @@ mod tests {
 
     use crate::arrow::arrow_reader::{
         ArrowReader, ArrowReaderOptions, ParquetFileArrowReader,
+        ParquetRecordBatchReader, RowSelection,
     };
     use crate::arrow::buffer::converter::{
         BinaryArrayConverter, Converter, FixedSizeArrayConverter, FromConverter,
@@ -1586,4 +1601,206 @@ mod tests {
         test_row_group_batch(MIN_BATCH_SIZE, MIN_BATCH_SIZE - 1);
         test_row_group_batch(MIN_BATCH_SIZE - 1, MIN_BATCH_SIZE);
     }
+
+    #[test]
+    fn test_scan_row_with_selection() {
+        let testdata = arrow::util::test_util::parquet_test_data();
+        let path = format!("{}/alltypes_tiny_pages_plain.parquet", testdata);
+        let test_file = File::open(&path).unwrap();
+
+        // total row count 7300
+        // 1. test selection len more than one page row count
+        let batch_size = 1000;
+        let expected_data = create_expect_batch(&test_file, batch_size);
+
+        let selections = create_test_selection(batch_size, 7300, false);
+        let skip_reader = create_skip_reader(&test_file, batch_size, selections);
+        let mut total_row_count = 0;
+        let mut index = 0;
+        for batch in skip_reader {
+            let batch = batch.unwrap();
+            assert_eq!(batch, expected_data.get(index).unwrap().clone());
+            index += 2;
+            let num = batch.num_rows();
+            assert!(num == batch_size || num == 300);
+            total_row_count += num;
+        }
+        assert_eq!(total_row_count, 4000);
+
+        let selections = create_test_selection(batch_size, 7300, true);
+        let skip_reader = create_skip_reader(&test_file, batch_size, selections);
+        let mut total_row_count = 0;
+        let mut index = 1;
+        for batch in skip_reader {
+            let batch = batch.unwrap();
+            assert_eq!(batch, expected_data.get(index).unwrap().clone());
+            index += 2;
+            let num = batch.num_rows();
+            //the lase batch will be 300
+            assert!(num == batch_size || num == 300);
+            total_row_count += num;
+        }
+        assert_eq!(total_row_count, 3300);
+
+        // 2. test selection len less than one page row count
+        let batch_size = 20;
+        let expected_data = create_expect_batch(&test_file, batch_size);
+        let selections = create_test_selection(batch_size, 7300, false);
+
+        let skip_reader = create_skip_reader(&test_file, batch_size, selections);
+        let mut total_row_count = 0;
+        let mut index = 0;
+        for batch in skip_reader {
+            let batch = batch.unwrap();
+            assert_eq!(batch, expected_data.get(index).unwrap().clone());
+            index += 2;
+            let num = batch.num_rows();
+            assert_eq!(num, batch_size);
+            total_row_count += num;
+        }
+        assert_eq!(total_row_count, 3660);
+
+        let selections = create_test_selection(batch_size, 7300, true);
+        let skip_reader = create_skip_reader(&test_file, batch_size, selections);
+        let mut total_row_count = 0;
+        let mut index = 1;
+        for batch in skip_reader {
+            let batch = batch.unwrap();
+            assert_eq!(batch, expected_data.get(index).unwrap().clone());
+            index += 2;
+            let num = batch.num_rows();
+            assert_eq!(num, batch_size);
+            total_row_count += num;
+        }
+        assert_eq!(total_row_count, 3640);
+
+        // 3. test selection_len less than batch_size
+        let batch_size = 20;
+        let selection_len = 5;
+        let expected_data_batch = create_expect_batch(&test_file, batch_size);
+        let expected_data_selection = create_expect_batch(&test_file, selection_len);
+        let selections = create_test_selection(selection_len, 7300, false);
+        let skip_reader = create_skip_reader(&test_file, batch_size, selections);
+
+        let mut total_row_count = 0;
+
+        for batch in skip_reader {
+            let batch = batch.unwrap();
+            let num = batch.num_rows();
+            assert!(num == batch_size || num == selection_len);
+            if num == batch_size {
+                assert_eq!(
+                    batch,
+                    expected_data_batch
+                        .get(total_row_count / batch_size)
+                        .unwrap()
+                        .clone()
+                );
+                total_row_count += batch_size;
+            } else if num == selection_len {
+                assert_eq!(
+                    batch,
+                    expected_data_selection
+                        .get(total_row_count / selection_len)
+                        .unwrap()
+                        .clone()
+                );
+                total_row_count += selection_len;
+            }
+            // add skip offset
+            total_row_count += selection_len;
+        }
+
+        // 4. test selection_len more than batch_size
+        // If batch_size < selection_len will divide selection(50, read) ->
+        // selection(20, read), selection(20, read), selection(10, read)
+        let batch_size = 20;
+        let selection_len = 50;
+        let another_batch_size = 10;
+        let expected_data_batch = create_expect_batch(&test_file, batch_size);
+        let expected_data_batch2 = create_expect_batch(&test_file, another_batch_size);
+        let selections = create_test_selection(selection_len, 7300, false);
+        let skip_reader = create_skip_reader(&test_file, batch_size, selections);
+
+        let mut total_row_count = 0;
+
+        for batch in skip_reader {
+            let batch = batch.unwrap();
+            let num = batch.num_rows();
+            assert!(num == batch_size || num == another_batch_size);
+            if num == batch_size {
+                assert_eq!(
+                    batch,
+                    expected_data_batch
+                        .get(total_row_count / batch_size)
+                        .unwrap()
+                        .clone()
+                );
+                total_row_count += batch_size;
+            } else if num == another_batch_size {
+                assert_eq!(
+                    batch,
+                    expected_data_batch2
+                        .get(total_row_count / another_batch_size)
+                        .unwrap()
+                        .clone()
+                );
+                total_row_count += 10;
+                // add skip offset
+                total_row_count += selection_len;
+            }
+        }
+
+        fn create_skip_reader(
+            test_file: &File,
+            batch_size: usize,
+            selections: Vec<RowSelection>,
+        ) -> ParquetRecordBatchReader {
+            let arrow_reader_options =
+                ArrowReaderOptions::new().with_row_selection(selections);
+
+            let mut skip_arrow_reader = ParquetFileArrowReader::try_new_with_options(
+                test_file.try_clone().unwrap(),
+                arrow_reader_options,
+            )
+            .unwrap();
+            skip_arrow_reader.get_record_reader(batch_size).unwrap()
+        }
+
+        fn create_test_selection(
+            step_len: usize,
+            total_len: usize,
+            skip_first: bool,
+        ) -> Vec<RowSelection> {
+            let mut remaining = total_len;
+            let mut skip = skip_first;
+            let mut vec = vec![];
+            while remaining != 0 {
+                let step = if remaining > step_len {
+                    step_len
+                } else {
+                    remaining
+                };
+                vec.push(RowSelection {
+                    row_count: step,
+                    skip,
+                });
+                remaining -= step;
+                skip = !skip;
+            }
+            vec
+        }
+
+        fn create_expect_batch(test_file: &File, batch_size: usize) -> Vec<RecordBatch> {
+            let mut serial_arrow_reader =
+                ParquetFileArrowReader::try_new(test_file.try_clone().unwrap()).unwrap();
+            let serial_reader =
+                serial_arrow_reader.get_record_reader(batch_size).unwrap();
+            let mut expected_data = vec![];
+            for batch in serial_reader {
+                expected_data.push(batch.unwrap());
+            }
+            expected_data
+        }
+    }
 }

parquet/src/arrow/record_reader/mod.rs

@@ -24,6 +24,7 @@ use crate::arrow::record_reader::{
     buffer::{BufferQueue, ScalarBuffer, ValuesBuffer},
     definition_levels::{DefinitionLevelBuffer, DefinitionLevelBufferDecoder},
 };
+use crate::column::page::PageIterator;
 use crate::column::{
     page::PageReader,
     reader::{
@@ -184,11 +185,24 @@ where
     /// # Returns
     ///
     /// Number of records skipped
-    pub fn skip_records(&mut self, num_records: usize) -> Result<usize> {
+    pub fn skip_records(
+        &mut self,
+        num_records: usize,
+        pages: &mut dyn PageIterator,
+    ) -> Result<usize> {
         // First need to clear the buffer
         let end_of_column = match self.column_reader.as_mut() {
             Some(reader) => !reader.has_next()?,
-            None => return Ok(0),
+            None => {
+                // If we skip records before all read operation
+                // we need set `column_reader` by `set_page_reader`
+                if let Some(page_reader) = pages.next() {
+                    self.set_page_reader(page_reader?)?;
+                    false
+                } else {
+                    return Ok(0);
+                }
+            }
         };
 
         let (buffered_records, buffered_values) =

parquet/src/column/reader.rs

@@ -299,7 +299,7 @@ where
         let mut remaining = num_records;
         while remaining != 0 {
             if self.num_buffered_values == self.num_decoded_values {
-                let metadata = match self.page_reader.peek_next_page()? {
+                let mut metadata = match self.page_reader.peek_next_page()? {
                     None => return Ok(num_records - remaining),
                     Some(metadata) => metadata,
                 };
@@ -312,13 +312,20 @@ where
 
                 // If page has less rows than the remaining records to
                 // be skipped, skip entire page
-                if metadata.num_rows < remaining {
+                while metadata.num_rows < remaining {
                     self.page_reader.skip_next_page()?;
                     remaining -= metadata.num_rows;
-                    continue;
+                    metadata = match self.page_reader.peek_next_page()? {
+                        None => return Ok(num_records - remaining),
+                        Some(metadata) => metadata,
+                    };
                 }
+                // because self.num_buffered_values == self.num_decoded_values means
+                // we need reads a new page and set up the decoders for levels
+                self.read_new_page()?;
             }
 
+            // start skip values in page level
             let to_read = remaining
                 .min((self.num_buffered_values - self.num_decoded_values) as usize);