Poc for adaptive parquet predicate pushdown(bitmap/range) with page cache(3 data pages)

parquet/src/arrow/array_reader/byte_view_array.rs

@@ -432,7 +432,8 @@ impl ByteViewArrayDecoderDictionary {
         }
     }
 
-    /// Reads the next indexes from self.decoder
+    /// Reads the next `len` indexes from self.decoder
+    ///
     /// the indexes are assumed to be indexes into `dict`
     /// the output values are written to output
     ///
@@ -464,6 +465,8 @@ impl ByteViewArrayDecoderDictionary {
             }
         }
 
+        output.views.reserve(len);
+
         // Calculate the offset of the dictionary buffers in the output buffers
         // For example if the 2nd buffer in the dictionary is the 5th buffer in the output buffers,
         // then the base_buffer_idx is 5 - 2 = 3

parquet/src/arrow/arrow_reader/mod.rs

@@ -17,16 +17,6 @@
 
 //! Contains reader which reads parquet data into arrow [`RecordBatch`]
 
-use arrow_array::cast::AsArray;
-use arrow_array::Array;
-use arrow_array::{RecordBatch, RecordBatchReader};
-use arrow_schema::{ArrowError, DataType as ArrowType, Schema, SchemaRef};
-use arrow_select::filter::prep_null_mask_filter;
-pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter};
-pub use selection::{RowSelection, RowSelector};
-use std::collections::VecDeque;
-use std::sync::Arc;
-
 pub use crate::arrow::array_reader::RowGroups;
 use crate::arrow::array_reader::{build_array_reader, ArrayReader};
 use crate::arrow::schema::{parquet_to_arrow_schema_and_fields, ParquetField};
@@ -38,6 +28,16 @@ use crate::errors::{ParquetError, Result};
 use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
 use crate::file::reader::{ChunkReader, SerializedPageReader};
 use crate::schema::types::SchemaDescriptor;
+use arrow_array::cast::AsArray;
+use arrow_array::{Array, BooleanArray};
+use arrow_array::{RecordBatch, RecordBatchReader};
+use arrow_schema::{ArrowError, DataType as ArrowType, Schema, SchemaRef};
+use arrow_select::filter::prep_null_mask_filter;
+pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter};
+pub use selection::{RowSelection, RowSelector};
+use std::collections::VecDeque;
+use std::mem::take;
+use std::sync::Arc;
 
 mod filter;
 mod selection;
@@ -761,6 +761,7 @@ impl<T: ChunkReader + 'static> ReaderPageIterator<T> {
             .map(|i| i[rg_idx][self.column_idx].page_locations.clone());
         let total_rows = rg.num_rows() as usize;
         let reader = self.reader.clone();
+        // todo: add cache???
 
         SerializedPageReader::new(reader, column_chunk_metadata, total_rows, page_locations)?
             .add_crypto_context(
@@ -792,18 +793,20 @@ pub struct ParquetRecordBatchReader {
     batch_size: usize,
     array_reader: Box<dyn ArrayReader>,
     schema: SchemaRef,
-    selection: Option<VecDeque<RowSelector>>,
+    selection: Option<RowSelection>,
 }
 
 impl Iterator for ParquetRecordBatchReader {
     type Item = Result<RecordBatch, ArrowError>;
 
     fn next(&mut self) -> Option<Self::Item> {
         let mut read_records = 0;
+        let mut filter: Option<BooleanArray> = None;
         match self.selection.as_mut() {
-            Some(selection) => {
-                while read_records < self.batch_size && !selection.is_empty() {
-                    let front = selection.pop_front().unwrap();
+            Some(RowSelection::Ranges(selectors)) => {
+                let mut selectors: VecDeque<RowSelector> = VecDeque::from_iter(take(selectors));
+                while read_records < self.batch_size && !selectors.is_empty() {
+                    let front = selectors.pop_front().unwrap();
                     if front.skip {
                         let skipped = match self.array_reader.skip_records(front.row_count) {
                             Ok(skipped) => skipped,
@@ -833,7 +836,7 @@ impl Iterator for ParquetRecordBatchReader {
                         Some(remaining) if remaining != 0 => {
                             // if page row count less than batch_size we must set batch size to page row count.
                             // add check avoid dead loop
-                            selection.push_front(RowSelector::select(remaining));
+                            selectors.push_front(RowSelector::select(remaining));
                             need_read
                         }
                         _ => front.row_count,
@@ -844,7 +847,14 @@ impl Iterator for ParquetRecordBatchReader {
                         Err(error) => return Some(Err(error.into())),
                     }
                 }
+                self.selection = Some(RowSelection::Ranges(selectors.into()));
             }
+            Some(RowSelection::BitMap(bitmap)) => {
+                self.array_reader.read_records(bitmap.len()).unwrap();
+                filter = Some(bitmap.clone());
+                self.selection = None;
+            }
+
             None => {
                 if let Err(error) = self.array_reader.read_records(self.batch_size) {
                     return Some(Err(error.into()));
@@ -861,9 +871,31 @@ impl Iterator for ParquetRecordBatchReader {
                     )
                 });
 
-                match struct_array {
-                    Err(err) => Some(Err(err)),
-                    Ok(e) => (e.len() > 0).then(|| Ok(RecordBatch::from(e))),
+                let batch: RecordBatch = match struct_array {
+                    Err(err) => return Some(Err(err)),
+                    Ok(e) => e.into(),
+                };
+
+                if let Some(filter) = filter.as_mut() {
+                    if batch.num_rows() == 0 {
+                        return None;
+                    }
+                    if filter.len() != batch.num_rows() {
+                        return Some(Err(ArrowError::ComputeError(format!(
+                            "Filter length ({}) does not match batch rows ({})",
+                            filter.len(),
+                            batch.num_rows()
+                        ))));
+                    }
+
+                    match arrow_select::filter::filter_record_batch(&batch, filter) {
+                        Ok(filtered_batch) => Some(Ok(filtered_batch)),
+                        Err(e) => Some(Err(e)),
+                    }
+                } else if batch.num_rows() > 0 {
+                    Some(Ok(batch))
+                } else {
+                    None
                 }
             }
         }
@@ -907,7 +939,7 @@ impl ParquetRecordBatchReader {
             batch_size,
             array_reader,
             schema: Arc::new(Schema::new(levels.fields.clone())),
-            selection: selection.map(|s| s.trim().into()),
+            selection: selection.map(|s| s.trim()),
         })
     }
 
@@ -928,7 +960,7 @@ impl ParquetRecordBatchReader {
             batch_size,
             array_reader,
             schema: Arc::new(schema),
-            selection: selection.map(|s| s.trim().into()),
+            selection: selection.map(|s| s.trim()),
         }
     }
 }
@@ -1002,6 +1034,7 @@ pub(crate) fn evaluate_predicate(
                 filter.len()
             ));
         }
+
         match filter.null_count() {
             0 => filters.push(filter),
             _ => filters.push(prep_null_mask_filter(&filter)),