update

parquet/src/arrow/array_reader/builder.rs

@@ -41,6 +41,8 @@ pub struct ArrayReaderBuilder<'a> {
     row_groups: &'a dyn RowGroups,
 }
 
+type CacheContext<'a> = (&'a ProjectionMask, Arc<Mutex<RowGroupCache>>);
+
 impl<'a> ArrayReaderBuilder<'a> {
     pub fn new(row_groups: &'a dyn RowGroups) -> Self {
         Self { row_groups }
@@ -51,12 +53,25 @@ impl<'a> ArrayReaderBuilder<'a> {
         &self,
         field: Option<&ParquetField>,
         mask: &ProjectionMask,
-        cache_mask: &ProjectionMask,
-        cache: Arc<Mutex<RowGroupCache>>,
+    ) -> Result<Box<dyn ArrayReader>> {
+        let reader = field
+            .and_then(|field| self.build_reader(field, mask, None).transpose())
+            .transpose()?
+            .unwrap_or_else(|| make_empty_array_reader(self.num_rows()));
+
+        Ok(reader)
+    }
+
+    /// Create [`ArrayReader`] from parquet schema, projection mask, and parquet file reader.
+    pub fn build_array_reader_with_cache(
+        &self,
+        field: Option<&ParquetField>,
+        mask: &ProjectionMask,
+        cache_context: CacheContext,
     ) -> Result<Box<dyn ArrayReader>> {
         let reader = field
             .and_then(|field| {
-                self.build_reader(field, mask, cache_mask, cache)
+                self.build_reader(field, mask, Some(cache_context))
                     .transpose()
             })
             .transpose()?
@@ -74,32 +89,32 @@ impl<'a> ArrayReaderBuilder<'a> {
         &self,
         field: &ParquetField,
         mask: &ProjectionMask,
-        cache_mask: &ProjectionMask,
-        cache: Arc<Mutex<RowGroupCache>>,
+        cache_context: Option<CacheContext>,
     ) -> Result<Option<Box<dyn ArrayReader>>> {
         match field.field_type {
             ParquetFieldType::Primitive { col_idx, .. } => {
-                if let Some(reader) = self.build_primitive_reader(field, mask)? {
-                    if cache_mask.leaf_included(col_idx) {
-                        Ok(Some(Box::new(CachedArrayReader::new(
-                            reader, cache, col_idx,
-                        ))))
-                    } else {
-                        Ok(Some(reader))
-                    }
+                let Some(reader) = self.build_primitive_reader(field, mask)? else {
+                    return Ok(None);
+                };
+                let Some((cache_mask, cache)) = cache_context else {
+                    return Ok(Some(reader));
+                };
+
+                if cache_mask.leaf_included(col_idx) {
+                    Ok(Some(Box::new(CachedArrayReader::new(
+                        reader, cache, col_idx,
+                    ))))
                 } else {
-                    Ok(None)
+                    Ok(Some(reader))
                 }
             }
             ParquetFieldType::Group { .. } => match &field.arrow_type {
-                DataType::Map(_, _) => self.build_map_reader(field, mask, cache_mask, cache),
-                DataType::Struct(_) => self.build_struct_reader(field, mask, cache_mask, cache),
-                DataType::List(_) => self.build_list_reader(field, mask, cache_mask, cache, false),
-                DataType::LargeList(_) => {
-                    self.build_list_reader(field, mask, cache_mask, cache, true)
-                }
+                DataType::Map(_, _) => self.build_map_reader(field, mask, cache_context),
+                DataType::Struct(_) => self.build_struct_reader(field, mask, cache_context),
+                DataType::List(_) => self.build_list_reader(field, mask, cache_context, false),
+                DataType::LargeList(_) => self.build_list_reader(field, mask, cache_context, true),
                 DataType::FixedSizeList(_, _) => {
-                    self.build_fixed_size_list_reader(field, mask, cache_mask, cache)
+                    self.build_fixed_size_list_reader(field, mask, cache_context)
                 }
                 d => unimplemented!("reading group type {} not implemented", d),
             },
@@ -111,14 +126,13 @@ impl<'a> ArrayReaderBuilder<'a> {
         &self,
         field: &ParquetField,
         mask: &ProjectionMask,
-        cache_mask: &ProjectionMask,
-        cache: Arc<Mutex<RowGroupCache>>,
+        cache_context: Option<CacheContext>,
     ) -> Result<Option<Box<dyn ArrayReader>>> {
         let children = field.children().unwrap();
         assert_eq!(children.len(), 2);
 
-        let key_reader = self.build_reader(&children[0], mask, cache_mask, cache.clone())?;
-        let value_reader = self.build_reader(&children[1], mask, cache_mask, cache)?;
+        let key_reader = self.build_reader(&children[0], mask, cache_context.clone())?;
+        let value_reader = self.build_reader(&children[1], mask, cache_context)?;
 
         match (key_reader, value_reader) {
             (Some(key_reader), Some(value_reader)) => {
@@ -164,14 +178,13 @@ impl<'a> ArrayReaderBuilder<'a> {
         &self,
         field: &ParquetField,
         mask: &ProjectionMask,
-        cache_mask: &ProjectionMask,
-        cache: Arc<Mutex<RowGroupCache>>,
+        cache_context: Option<CacheContext>,
         is_large: bool,
     ) -> Result<Option<Box<dyn ArrayReader>>> {
         let children = field.children().unwrap();
         assert_eq!(children.len(), 1);
 
-        let reader = match self.build_reader(&children[0], mask, cache_mask, cache)? {
+        let reader = match self.build_reader(&children[0], mask, cache_context)? {
             Some(item_reader) => {
                 // Need to retrieve underlying data type to handle projection
                 let item_type = item_reader.get_data_type().clone();
@@ -213,13 +226,12 @@ impl<'a> ArrayReaderBuilder<'a> {
         &self,
         field: &ParquetField,
         mask: &ProjectionMask,
-        cache_mask: &ProjectionMask,
-        cache: Arc<Mutex<RowGroupCache>>,
+        cache_context: Option<CacheContext>,
     ) -> Result<Option<Box<dyn ArrayReader>>> {
         let children = field.children().unwrap();
         assert_eq!(children.len(), 1);
 
-        let reader = match self.build_reader(&children[0], mask, cache_mask, cache)? {
+        let reader = match self.build_reader(&children[0], mask, cache_context)? {
             Some(item_reader) => {
                 let item_type = item_reader.get_data_type().clone();
                 let reader = match &field.arrow_type {
@@ -349,8 +361,7 @@ impl<'a> ArrayReaderBuilder<'a> {
         &self,
         field: &ParquetField,
         mask: &ProjectionMask,
-        cache_mask: &ProjectionMask,
-        cache: Arc<Mutex<RowGroupCache>>,
+        cache_context: Option<CacheContext>,
     ) -> Result<Option<Box<dyn ArrayReader>>> {
         let arrow_fields = match &field.arrow_type {
             DataType::Struct(children) => children,
@@ -363,7 +374,7 @@ impl<'a> ArrayReaderBuilder<'a> {
         let mut builder = SchemaBuilder::with_capacity(children.len());
 
         for (arrow, parquet) in arrow_fields.iter().zip(children) {
-            if let Some(reader) = self.build_reader(parquet, mask, cache_mask, cache.clone())? {
+            if let Some(reader) = self.build_reader(parquet, mask, cache_context.clone())? {
                 // Need to retrieve underlying data type to handle projection
                 let child_type = reader.get_data_type().clone();
                 builder.push(arrow.as_ref().clone().with_data_type(child_type));
@@ -407,10 +418,9 @@ mod tests {
             file_metadata.key_value_metadata(),
         )
         .unwrap();
-        let cache = Arc::new(Mutex::new(RowGroupCache::new(1000)));
 
         let array_reader = ArrayReaderBuilder::new(&file_reader)
-            .build_array_reader(fields.as_ref(), &mask, &ProjectionMask::all(), cache)
+            .build_array_reader(fields.as_ref(), &mask)
             .unwrap();
 
         // Create arrow types

parquet/src/arrow/array_reader/cached_array_reader.rs

@@ -204,11 +204,11 @@ impl ArrayReader for CachedArrayReader {
     }
 
     fn get_def_levels(&self) -> Option<&[i16]> {
-        self.inner.get_def_levels()
+        None // we don't allow nullable parent for now.
     }
 
     fn get_rep_levels(&self) -> Option<&[i16]> {
-        self.inner.get_rep_levels()
+        None
     }
 }
 

parquet/src/arrow/array_reader/list_array.rs

@@ -247,7 +247,6 @@ impl<OffsetSize: OffsetSizeTrait> ArrayReader for ListArrayReader<OffsetSize> {
 mod tests {
     use super::*;
     use crate::arrow::array_reader::list_array::ListArrayReader;
-    use crate::arrow::array_reader::row_group_cache::RowGroupCache;
     use crate::arrow::array_reader::test_util::InMemoryArrayReader;
     use crate::arrow::array_reader::ArrayReaderBuilder;
     use crate::arrow::schema::parquet_to_arrow_schema_and_fields;
@@ -260,7 +259,7 @@ mod tests {
     use arrow_array::{Array, PrimitiveArray};
     use arrow_data::ArrayDataBuilder;
     use arrow_schema::Fields;
-    use std::sync::{Arc, Mutex};
+    use std::sync::Arc;
 
     fn list_type<OffsetSize: OffsetSizeTrait>(
         data_type: ArrowType,
@@ -563,10 +562,9 @@ mod tests {
             file_metadata.key_value_metadata(),
         )
         .unwrap();
-        let cache = Arc::new(Mutex::new(RowGroupCache::new(1000)));
 
         let mut array_reader = ArrayReaderBuilder::new(&file_reader)
-            .build_array_reader(fields.as_ref(), &mask, &ProjectionMask::all(), cache)
+            .build_array_reader(fields.as_ref(), &mask)
             .unwrap();
 
         let batch = array_reader.next_batch(100).unwrap();

parquet/src/arrow/arrow_reader/mod.rs

@@ -24,10 +24,10 @@ use arrow_schema::{ArrowError, DataType as ArrowType, Schema, SchemaRef};
 pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter};
 pub use selection::{RowSelection, RowSelector};
 use std::fmt::{Debug, Formatter};
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 
 pub use crate::arrow::array_reader::RowGroups;
-use crate::arrow::array_reader::{ArrayReader, ArrayReaderBuilder, RowGroupCache};
+use crate::arrow::array_reader::{ArrayReader, ArrayReaderBuilder};
 use crate::arrow::schema::{parquet_to_arrow_schema_and_fields, ParquetField};
 use crate::arrow::{parquet_to_arrow_field_levels, FieldLevels, ProjectionMask};
 use crate::column::page::{PageIterator, PageReader};
@@ -711,10 +711,6 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
         let batch_size = self
             .batch_size
             .min(self.metadata.file_metadata().num_rows() as usize);
-        let cache_projection = match self.compute_cache_projection(&self.projection) {
-            Some(projection) => projection,
-            None => ProjectionMask::all(),
-        };
 
         let row_groups = self
             .row_groups
@@ -725,7 +721,6 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
             metadata: self.metadata,
             row_groups,
         };
-        let row_group_cache = Arc::new(Mutex::new(RowGroupCache::new(batch_size)));
 
         let mut filter = self.filter;
         let mut plan_builder = ReadPlanBuilder::new(batch_size).with_selection(self.selection);
@@ -741,23 +736,15 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
                 let mut cache_projection = predicate.projection().clone();
                 cache_projection.intersect(&self.projection);
 
-                let array_reader = ArrayReaderBuilder::new(&reader).build_array_reader(
-                    self.fields.as_deref(),
-                    predicate.projection(),
-                    &cache_projection,
-                    row_group_cache.clone(),
-                )?;
+                let array_reader = ArrayReaderBuilder::new(&reader)
+                    .build_array_reader(self.fields.as_deref(), predicate.projection())?;
 
                 plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?;
             }
         }
 
-        let array_reader = ArrayReaderBuilder::new(&reader).build_array_reader(
-            self.fields.as_deref(),
-            &self.projection,
-            &cache_projection,
-            row_group_cache.clone(),
-        )?;
+        let array_reader = ArrayReaderBuilder::new(&reader)
+            .build_array_reader(self.fields.as_deref(), &self.projection)?;
 
         let read_plan = plan_builder
             .limited(reader.num_rows())
@@ -968,12 +955,7 @@ impl ParquetRecordBatchReader {
         selection: Option<RowSelection>,
     ) -> Result<Self> {
         let array_reader = ArrayReaderBuilder::new(row_groups)
-            .build_array_reader(
-                levels.levels.as_ref(),
-                &ProjectionMask::all(),
-                &ProjectionMask::all(),
-                Arc::new(Mutex::new(RowGroupCache::new(batch_size))),
-            )?;
+            .build_array_reader(levels.levels.as_ref(), &ProjectionMask::all())?;
 
         let read_plan = ReadPlanBuilder::new(batch_size)
             .with_selection(selection)

parquet/src/arrow/async_reader/mod.rs

@@ -621,12 +621,12 @@ where
 
                 let mut cache_projection = predicate.projection().clone();
                 cache_projection.intersect(&projection);
-                let array_reader = ArrayReaderBuilder::new(&row_group).build_array_reader(
-                    self.fields.as_deref(),
-                    predicate.projection(),
-                    &cache_projection,
-                    row_group_cache.clone(),
-                )?;
+                let array_reader = ArrayReaderBuilder::new(&row_group)
+                    .build_array_reader_with_cache(
+                        self.fields.as_deref(),
+                        predicate.projection(),
+                        (&cache_projection, row_group_cache.clone()),
+                    )?;
 
                 plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?;
             }
@@ -673,13 +673,11 @@ where
 
         let plan = plan_builder.build();
 
-        let array_reader = ArrayReaderBuilder::new(&row_group)
-            .build_array_reader(
-                self.fields.as_deref(),
-                &projection,
-                &cache_projection,
-                row_group_cache.clone(),
-            )?;
+        let array_reader = ArrayReaderBuilder::new(&row_group).build_array_reader_with_cache(
+            self.fields.as_deref(),
+            &projection,
+            (&cache_projection, row_group_cache.clone()),
+        )?;
 
         let reader = ParquetRecordBatchReader::new(array_reader, plan);