add nested nullable field support for CachedArrayReader

Author: adriangb

parquet/src/arrow/array_reader/cached_array_reader.rs

@@ -281,8 +281,15 @@ impl ArrayReader for CachedArrayReader {
     fn consume_batch(&mut self) -> Result<ArrayRef> {
         let row_count = self.selections.len();
         if row_count == 0 {
-            self.def_levels_buffer = None;
-            self.rep_levels_buffer = None;
+            // When there's no data to consume, set empty level buffers if we
+            // previously had levels. This ensures the levels match the empty array.
+            // We keep Some([]) rather than None to indicate this reader provides levels.
+            if self.def_levels_buffer.is_some() {
+                self.def_levels_buffer = Some(Vec::new());
+            }
+            if self.rep_levels_buffer.is_some() {
+                self.rep_levels_buffer = Some(Vec::new());
+            }
             return Ok(new_empty_array(self.inner.get_data_type()));
         }
 

parquet/src/arrow/mod.rs

@@ -419,51 +419,6 @@ impl ProjectionMask {
             }
         }
     }
-
-    /// Return a new [`ProjectionMask`] that excludes any leaf columns that are
-    /// part of a nested type, such as struct, list, or map
-    ///
-    /// If there are no non-nested columns in the mask, returns `None`
-    pub(crate) fn without_nested_types(&self, schema: &SchemaDescriptor) -> Option<Self> {
-        let num_leaves = schema.num_columns();
-
-        // Count how many leaves each root column has
-        let num_roots = schema.root_schema().get_fields().len();
-        let mut root_leaf_counts = vec![0usize; num_roots];
-        for leaf_idx in 0..num_leaves {
-            let root_idx = schema.get_column_root_idx(leaf_idx);
-            root_leaf_counts[root_idx] += 1;
-        }
-
-        // Keep only leaves whose root has exactly one leaf (non-nested) and is not a
-        // LIST. LIST is encoded as a wrapped logical type with a single leaf, e.g.
-        // https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#lists
-        //
-        // ```text
-        // // List<String> (list non-null, elements nullable)
-        // required group my_list (LIST) {
-        //   repeated group list {
-        //     optional binary element (STRING);
-        //   }
-        // }
-        // ```
-        let mut included_leaves = Vec::new();
-        for leaf_idx in 0..num_leaves {
-            if self.leaf_included(leaf_idx) {
-                let root = schema.get_column_root(leaf_idx);
-                let root_idx = schema.get_column_root_idx(leaf_idx);
-                if root_leaf_counts[root_idx] == 1 && !root.is_list() {
-                    included_leaves.push(leaf_idx);
-                }
-            }
-        }
-
-        if included_leaves.is_empty() {
-            None
-        } else {
-            Some(ProjectionMask::leaves(schema, included_leaves))
-        }
-    }
 }
 
 /// Lookups up the parquet column by name

parquet/src/arrow/push_decoder/reader_builder/mod.rs

@@ -636,12 +636,7 @@ impl RowGroupReaderBuilder {
             cache_projection.union(predicate.projection());
         }
         cache_projection.intersect(&self.projection);
-        self.exclude_nested_columns_from_cache(&cache_projection)
-    }
-
-    /// Exclude leaves belonging to roots that span multiple parquet leaves (i.e. nested columns)
-    fn exclude_nested_columns_from_cache(&self, mask: &ProjectionMask) -> Option<ProjectionMask> {
-        mask.without_nested_types(self.metadata.file_metadata().schema_descr())
+        Some(cache_projection)
     }
 
     /// Get the offset index for the specified row group, if any

parquet/tests/arrow_reader/predicate_cache.rs

@@ -85,13 +85,11 @@ async fn test_cache_disabled_with_filters() {
 }
 
 #[tokio::test]
-async fn test_cache_projection_excludes_nested_columns() {
-    let test = ParquetPredicateCacheTest::new_nested().with_expected_records_read_from_cache(0);
-
-    let sync_builder = test.sync_builder().add_nested_filter();
-    test.run_sync(sync_builder);
-
-    let async_builder = test.async_builder().await.add_nested_filter();
+async fn test_async_cache_with_nested_columns() {
+    // Nested columns now work with cache - expect records from cache
+    // 100 rows × 2 leaf columns (b.aa, b.bb) = 200 records
+    let test = ParquetPredicateCacheTest::new_nested().with_expected_records_read_from_cache(200);
+    let async_builder = test.async_builder().await.add_nested_root_filter();
     test.run_async(async_builder).await;
 }
 
@@ -282,9 +280,9 @@ trait ArrowReaderBuilderExt {
     /// 2. a row_filter applied to "b": 575 < "b" < 625 (select 1 data page from each row group)
     fn add_project_ab_and_filter_b(self) -> Self;
 
-    /// Adds a row filter that projects the nested leaf column "b.aa" and
+    /// Adds a row filter that projects the nested ROOT column "b" and
     /// returns true for all rows.
-    fn add_nested_filter(self) -> Self;
+    fn add_nested_root_filter(self) -> Self;
 }
 
 impl<T> ArrowReaderBuilderExt for ArrowReaderBuilder<T> {
@@ -306,23 +304,21 @@ impl<T> ArrowReaderBuilderExt for ArrowReaderBuilder<T> {
             .with_row_filter(RowFilter::new(vec![Box::new(row_filter)]))
     }
 
-    fn add_nested_filter(self) -> Self {
+    fn add_nested_root_filter(self) -> Self {
         let schema_descr = self.metadata().file_metadata().schema_descr_ptr();
 
-        // Build a RowFilter whose predicate projects a leaf under the nested root `b`
-        // Leaf indices are depth-first; with schema [a, b.aa, b.bb] we pick index 1 (b.aa)
-        let nested_leaf_mask = ProjectionMask::leaves(&schema_descr, vec![1]);
+        // Project the ROOT struct column "b", not just leaf "b.aa"
+        let root_mask = ProjectionMask::roots(&schema_descr, [1]); // column index 1 = "b"
 
-        let always_true = ArrowPredicateFn::new(nested_leaf_mask.clone(), |batch: RecordBatch| {
+        let always_true = ArrowPredicateFn::new(root_mask.clone(), |batch: RecordBatch| {
             Ok(arrow_array::BooleanArray::from(vec![
                 true;
                 batch.num_rows()
             ]))
         });
         let row_filter = RowFilter::new(vec![Box::new(always_true)]);
 
-        self.with_projection(nested_leaf_mask)
-            .with_row_filter(row_filter)
+        self.with_projection(root_mask).with_row_filter(row_filter)
     }
 }