add nested nullable field support for CachedArrayReader

[adriangb]

Needed to unblock apache/datafusion#19556

This was mostly written by Claude with guidance on to fix the following MRE (reproducible on the branch for apache/datafusion#19556):

COPY (
  SELECT *
  FROM VALUES ({field: [{nested: 1}]}, 1), ({field: [{nested: 2}]}, 2) AS t(struct_col, id)
)
TO 'test.parquet';

set datafusion.execution.parquet.pushdown_filters = true;

CREATE EXTERNAL TABLE t1 STORED AS PARQUET LOCATION 'test.parquet';

-- Works, no issues
select struct_col['field'][1]['nested']
from t1
where id = 1;

-- Works, no issues
select id, struct_col
from t1
where struct_col['field'][1]['nested']  = 1;

-- Error
select id, struct_col['field'][1]['nested']
from t1
where struct_col['field'][1]['nested'] = 1;

[Weijun-H]

LGTM, Thanks @adriangb

better add this test

    #[test]
    fn test_level_propagation_empty_after_skip() {
        let metrics = ArrowReaderMetrics::disabled();
        let cache = Arc::new(Mutex::new(RowGroupCache::new(4, usize::MAX)));

        // Producer populates cache with levels
        let data = vec![1, 2, 3, 4];
        let def_levels = vec![1, 0, 1, 1];
        let rep_levels = vec![0, 1, 1, 0];
        let mock_reader =
            MockArrayReaderWithLevels::new(data, def_levels.clone(), rep_levels.clone());
        let mut producer = CachedArrayReader::new(
            Box::new(mock_reader),
            cache.clone(),
            0,
            CacheRole::Producer,
            metrics.clone(),
        );

        producer.read_records(4).unwrap();
        producer.consume_batch().unwrap();

        // Consumer skips all rows, resulting in an empty output batch
        let mock_reader2 = MockArrayReaderWithLevels::new(
            vec![10, 20, 30, 40],
            vec![0, 0, 0, 0],
            vec![0, 0, 0, 0],
        );
        let mut consumer = CachedArrayReader::new(
            Box::new(mock_reader2),
            cache,
            0,
            CacheRole::Consumer,
            metrics,
        );

        let skipped = consumer.skip_records(4).unwrap();
        assert_eq!(skipped, 4);

        let array = consumer.consume_batch().unwrap();
        assert_eq!(array.len(), 0);

        assert_eq!(consumer.get_def_levels().unwrap(), &[]);
        assert_eq!(consumer.get_rep_levels().unwrap(), &[]);
    }

[adriangb]

Added in a495d4d

[adriangb]

@XiangpengHao any chance you could take a look at this?

[XiangpengHao]

Looks good to me, thank you @adriangb !

[adriangb]

Looks good to me, thank you @adriangb !

Thank you for reviewing! I think this needs approval from someone with write permissions and for them to hit merge, I don't have these permissions in this repo.

[XiangpengHao]

me neither, asking for @alamb's help!

[alamb]

Thanks -- the high level idea of this looks good. However, I think we need some end to end tests prior to merging this

Specifically an end to end test that actually reads a StructAray that uses repetition and definition levels, and then a row filter that filters out some of the rows

I looked around in

arrow-rs/parquet/tests/arrow_reader/mod.rs

Line 1136 in 02fa779

builder = builder.set_statistics_truncate_length(DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH);

and didn't find any clear example

I am imagining something like

arrow-rs/parquet/tests/arrow_reader/io/async_reader.rs

Lines 295 to 297 in 02fa779

	let builder = builder
	.with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"]))
	.with_row_filter(filter_b_false(&schema_descr));

where we setup a reader and read the results back

[alamb]

this copies the levels, right? It seems like we should:

1. Only copy the levels when needed (caching nested types)
2. Look into potentially being able to take the levels rather than having to copy them

[adriangb]

@alamb I've reduced some allocations and added pretty extensive integration tests

[adriangb]

@alamb this was the best I could easily do: get_def_levels() returns an Option<&[u16]> and we're storing this on structures that can't easily have a lifetime (without a major refactor at least) so we have to allocate an owned Vec

[adriangb]

Alternatively could pass field: &ParquetField or def_level: i16 and have CachedArrayReader::new figure out how to use that