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

[adriangb]

@alamb sorry for the ping!

[alamb]

Sorry -- I am underwater with so many reviews.

I don't understand the nesting / reptition levels enough to review this code realistically as written.

Is there any chance you can write some "end to end" tests that show what this is doing from an end API perspective? Specifically, in https://github.com/apache/arrow-rs/blob/main/parquet/tests/arrow_reader/predicate_cache.rs

1. Create a file with nested structures / nulls
2. Create a read test with RowFilters on it
3. Use the cache statistics to show what is cached / not cached

Maybe you can get claude to help translate the datafusion repro into a arrow-rs only repo?

[adriangb]

Sorry for putting more on your plate…

I think the tests I added to parquet/tests/arrow_reader/predicate_cache.rs are exactly what you’re after?

[adriangb]

We've been able to work around this. @XiangpengHao @alamb do you think this is still valuable in its own right or should I close this?

[alamb]

I am sorry -- this pr is so large I am having trouble finding time to review it reasonably. Is it possibly to slim it down somehow? Perhaps breaking out tests (and/or reducing boiler plate, etc)?

[github-actions]

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.