Retrieve array from RecordBatch for a leaf column

[viirya]

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

While working filter pushdown for iceberg-rs: apache/iceberg-rust#295, I am going to use the APIs like ArrowPredicateFn and RowFilter.

When constructing ArrowPredicateFn for iceberg predicate, we provide a filtering function that takes RecordBatch based on the given projection.

The RecordBatch contains the columns specified in the projection. And we need to access correct column in the batch to evaluate the predicate.

For top-level column, it should be straightforward. But for nested column, seems no way to access the particular array from the RecordBatch.

We only have the projection (i.e., ProjectionMask) which contains indices of leaf columns in the batch.

For example, if the schema has [a, b, c] top columns. b is a struct column with [aa, bb, cc] columns. Give a predicate like cc > 1, and we know the leaf indices of the nested column cc is 3.

Is there API we can use to access the array of cc in the RecordBatch?

Describe the solution you'd like

Describe alternatives you've considered

Additional context

[tustvold]

I don't believe we currently have a mechanism for nested projection of RecordBatch but this is something that I think would be generally useful and a worthwhile addition

[adriangb]

I'll note that I think this will end up being a blocker for #6736

[alamb]

The ProjectionMask includes all leaf columns (so the [aa, bb, cc] in the example above can be selected in the predicate)

However, I think they would come back as their own columns (and I am not sure how nullability would work, etc)

The first step for this ticket is probably to work up an example:

1. Write a struct array with multiple fields into a parquet file
2. Write a test that tries to apply an ArrowPredicate to only one field and see what happens

[superserious-dev]

take

[superserious-dev]

Hello, I did some investigation and came up with the following example. It looks like it's possible to apply a ProjectionMask targeting the nested column cc and to access the column in the batch for filtering.

use std::{error::Error, fs::File, sync::Arc};

use arrow::array::{Array, AsArray, BooleanArray, Int32Array, RecordBatch, Scalar, StructArray};
use arrow::{
    compute::kernels::cmp::gt,
    datatypes::{DataType, Field, Schema},
    error::ArrowError,
};
use parquet::{
    arrow::{
        ArrowWriter, ProjectionMask,
        arrow_reader::{ArrowPredicateFn, ParquetRecordBatchReaderBuilder, RowFilter},
    },
    file::metadata::{ParquetMetaData, ParquetMetaDataReader},
};

fn main() -> Result<(), Box<dyn Error>> {
    let path = "sample.parquet";

    // Generate sample data w/ a StructArray
    generate_parquet(path)?;

    // Fetch metadata
    let metadata = read_metadata(path)?;

    // Construct filter
    let mask = ProjectionMask::leaves(metadata.file_metadata().schema_descr(), vec![3]); // Only look at cc (leaf = 3) in the filter
    let threshold = Int32Array::from(vec![1]);
    let cc_filter = ArrowPredicateFn::new(mask.clone(), move |batch| {
        // Assert only 1 top-level column is present: "b"
        assert_eq!(batch.num_columns(), 1);

        // Downcast the column to a struct
        let cc_struct = batch.column(0).as_struct();

        // Assert only 1 struct column is present: "cc"
        assert_eq!(batch.num_columns(), 1);

        // Filter column
        let cc = cc_struct.column(0);
        gt(cc, &Scalar::new(&threshold))
    });

    // Execute the filter w/ reader mask (same as filter mask)
    let filtered_row_count = filter_parquet(path, cc_filter, mask)?;

    // Assert that filtering was successful
    let original_row_count = metadata
        .row_groups()
        .iter()
        .fold(0, |total, rg| total + rg.num_rows());

    assert_eq!(original_row_count, 10);
    assert_eq!(filtered_row_count, 3);

    Ok(())
}

fn generate_parquet(path: &str) -> Result<(), Box<dyn Error>> {
    let a_col = Int32Array::from((0..10).collect::<Vec<i32>>());
    let a_field = Field::new("a", DataType::Int32, false);

    let b_col = StructArray::from(vec![
        (
            Arc::new(Field::new("aa", DataType::Boolean, false)),
            Arc::new(BooleanArray::from(vec![
                false, false, true, true, false, false, true, true, false, true,
            ])) as Arc<dyn Array>,
        ),
        (
            Arc::new(Field::new("bb", DataType::Boolean, false)),
            Arc::new(BooleanArray::from(vec![
                true, true, false, false, true, true, false, false, true, false,
            ])) as Arc<dyn Array>,
        ),
        (
            Arc::new(Field::new("cc", DataType::Int32, false)),
            Arc::new(Int32Array::from((-5..5).collect::<Vec<i32>>())),
        ),
    ]);
    let b_field = Field::new("b", b_col.data_type().clone(), false);

    let c_col = Int32Array::from((10..20).collect::<Vec<i32>>());
    let c_field = Field::new("c", DataType::Int32, true);

    let schema = Schema::new(vec![a_field, b_field, c_field]);
    let batch = RecordBatch::try_new(
        Arc::new(schema),
        vec![Arc::new(a_col), Arc::new(b_col), Arc::new(c_col)],
    )?;

    let file = File::create(path)?;
    let mut writer = ArrowWriter::try_new(file, batch.schema(), None)?;
    writer.write(&batch)?;
    writer.close()?;

    Ok(())
}

fn read_metadata(path: &str) -> Result<ParquetMetaData, Box<dyn Error>> {
    let parquet_file = File::open(path)?;
    ParquetMetaDataReader::new()
        .with_page_indexes(true)
        .parse_and_finish(&parquet_file)
        .map_err(|e| e.into())
}

fn filter_parquet<F>(
    path: &str,
    filter: ArrowPredicateFn<F>,
    mask: ProjectionMask,
) -> Result<usize, Box<dyn Error>>
where
    F: FnMut(RecordBatch) -> Result<BooleanArray, ArrowError> + Send + 'static,
{
    let parquet_file = File::open(path)?;

    let filter = RowFilter::new(vec![Box::new(filter)]);
    let reader = ParquetRecordBatchReaderBuilder::try_new(parquet_file)?
        .with_batch_size(10)
        .with_projection(mask)
        .with_row_filter(filter)
        .build()?;

    let mut row_count = 0;
    for b in reader {
        let num_rows = b?.num_rows();
        row_count += num_rows;
    }

    Ok(row_count)
}

[adriangb]

That's great news! I think then we just need to get apache/datafusion#14993 in DataFusion across the line so that ParquetOpener can inspect get_field(column, 'field') and generate the right ProjectionMask from that!

[superserious-dev]

Awesome. I'd be happy to help out on the DataFusion side if there are any tickets that need attention.

[adriangb]

I'll let you know! Right now I think we're going to need to do some design work and refactoring my to move from pushing down projections as column indices to expressions

[viirya]

The ProjectionMask includes all leaf columns (so the [aa, bb, cc] in the example above can be selected in the predicate)
However, I think they would come back as their own columns (and I am not sure how nullability would work, etc)
The first step for this ticket is probably to work up an example:
1. Write a struct array with multiple fields into a parquet file
2. Write a test that tries to apply an ArrowPredicate to only one field and see what happens

Hmm, I thought that I provided enough context in the description, but seems there is still confusion on this issue.

This issue originated from some works on iceberg-rs, specially this is the discussion around nested type apache/iceberg-rust#295 (comment).

To summarize the discussion, in short, although we can use ProjectionMask to access leaf columns, you cannot get "flatten" values in RowFilter API. In other words, you still get nested values where the nested struct is pruned to include the leaf columns. But it is not the "projected" column needed for filter pushdown.

To verify how the API works and if it fits for the filter pushdown work, I wrote a test #5600 at that time. I think the test does what @alamb asked.