Custom Predicates for ParquetExec and Parquet row indexes

[lightjacket]

Hello,

I've been looking at adding some indexing on parquet files by creating some auxiliary files and using those to create predicates for the ParquetExec struct to use. The short of it is I have an optimizer in place to replace existing ParquetExec instances with new instances using a custom PhysicalExpr for the predicate:

let new_exec = ParquetExec::new(
    exec.base_config().clone(),
    Some(Arc::new(MyCustomPhysicalExpr::new(/* args */))),
    None
).with_pushdown_filters(true);

My custom PhysicalExpr then has an evaluate function like the following:

fn evaluate(&self, batch: &RecordBatch) -> datafusion::common::Result<ColumnarValue> {
    let my_index: HashSet<usize> = self.run_my_index(/* args */);
    let indexes = (0..batch.num_rows()).map(|i| my_index.contains(i as usize)).collect::<Vec<_>>();
    Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(indexes))))
}

my_index currently has a set pf row numbers from the parquet file based on the args passed in. If there is only one RecordBatch, this happens to work (though I am not sure if I can assume the batch has not already been filtered somehow).

So my question is whether there is a way from a RecordBatch to load what the original record indexes in the parquet file were. I haven't seen anything yet, but I'd imagine I am looking in the wrong places. Alternatively, if the approach is just wrong and there's a better way to use indexes like this in DataFusion, I'd be very interested.

Thank you!

[lightjacket]

Just as some additional context on places I'd looked but without luck so far:

1. A PhysicalExpr that has the original row numbers (e.g. a Column expr on a hidden column like the rowid in sqlite)
2. Any info I could use passed to evaluate_bounds (I didn't actually see it get called in my testing, and hadn't dug into why yet)

Following up on (1), my current fallback plan is to use a column in the parquet as the "lookup" and use that in the index instead of a row number, but it would add additional constraints on the index usage (e.g. you need a unique column to index on).

[SteveLauC]

By row numbers, do you mean physical row numbers?

A ParquetExec, can have multiple partitions, assuming the underlying async runtime has sufficient threads (and cores), then those partitions will produce RecordBatches in parallel.

A pushed-down predicate, will be evaluated at partition levels, which means RecordBatches produced by different partitions will be filtered individually:

partition0 -> RecordBatch -> your custom predicate -> upper operators
partition1 -> RecordBatch -> your custom predicate -> upper operators
partition2 -> RecordBatch -> your custom predicate -> upper operators
...

So I think you need to ensure that the generated ParquetExec will only have 1 physical partition in order to make your custom predicate work.

By physical partition, I mean the parquet reader reads files in 1 physical thread. Changing it to 1 partition by operators like RepartitionExec won't work.

[lightjacket]

By row numbers, do you mean physical row numbers?

Yes.

Understood that the partitions will be run individually, but I don't think that conceptually prevents the original physical row numbers from being available, it would just mean you'd get a "random" set of row numbers in each partition the predicate is filtering. So instead of something like (0..batch.num_rows()) it'd be something like batch.physical_row_numbers().iter()....

But I would absolutely believe I am misunderstanding conceptually and the above is not correct.

[Ted-Jiang]

https://github.com/apache/arrow-datafusion/blob/4d389c2590370d85bfe3af77f5243d5b40f5a222/datafusion/core/src/datasource/physical_plan/parquet/mod.rs#L560

I think maybe you can use your indexes to build https://arrow.apache.org/rust/parquet/arrow/arrow_reader/struct.RowSelection.html
the selection is based on original record indexes

[lightjacket]

I'd missed this, and it looks promising! I think the only way to get the row selection set is via a page pruning predicate, which I think is created here from the ParquetExec:

https://github.com/apache/arrow-datafusion/blob/ec86acbc1fbc0da1e0bec9ad066a5177ec586c96/datafusion/core/src/physical_optimizer/pruning.rs#L1134

Looking at that code it doesn't look like there is a way to create the page pruning predicate based on physical row numbers (it looks like we can use a Column, InList, IsNull, etc.). I think the ultimately looks like the same issue where there's no access to the physical row numbers, but would love to be wrong about that!

But I'd imagine getting the execution to flow through this row selection code would result in better performance, so it's good to know about!

[alamb]

So my question is whether there is a way from a RecordBatch to load what the original record indexes in the parquet file were.

I don't think this is possible today in DataFusion (or in the parquet rust reader)

As @Ted-Jiang RowSelection can describe this concept, but I believe the row selections are per record group (not for the file as a whole) -- also I don't think this is exposed in some way you can provide a row selection to pass into the underlying reader

[lightjacket]

That looks to be the case as best I can tell, so I'm trying to move forward with a user-specified column for the index. Appreciate everyone taking a look and responding!