Plan to replace `SchemaAdapter` with `PhysicalExprAdapter`

[adriangb]

As discussed in #16791 the long term plan in my mind (and that I would like to discuss with the community) is to replace SchemaAdapter with PhysicalExprAdapter.

There are multiple reasons for this:

• We can better optimize scenarios like missing columns or casts. For example, it's cheaper to cast a literal and evaluate it against the data as read from the file than it is to read the data from the file and cast that to the type of the literal. It is also cheaper to evaluate the expression 1 > col1 as 1 > null when col1 is missing than it is to create an array of nulls. Since we can also simplify PhysicalExpr we can even simplify 1 > null into just null.
• It's easier to manipulate PhysicalExprs than it is to manipulate arrays. We already have machinery (TreeNode APIs, etc.) to do so.
• This is necessary to be able to push down projections into file scans which we need for upcoming Variant work and will also allow us to read single fields in a struct without reading the entire struct into memory. Basically if we want to be able to customize how expressions are evaluated for a specific format, in particular how variant_get(column, 'field') or get_field(column, 'field') are executed in the context of a specific format (e.g. in parquet we can read single struct columns or use shredded variant) we need to have access to the expression in ParquetOpener in order to check if the file schema has the shredded variant field and generate the right ProjectionMask.
• Paves the path for any other advanced optimizations, e.g. we could do crazy stuff like only read the dictionary page from a parquet column for a filter col = 'a' and if 'a' is not in the dictionary don't even bother reading the keys.

We've already implemented a replacement system for predicate pushdown via PhysicalExprAdapter and have examples showing how to do some of the things a custom SchemaAdapter can do.
Once we implement #14993 we'll be able to deprecate SchemaAdapter for the most part.

[mbutrovich]

Comet makes increasing use of SchemaAdapter, but nothing you describe here sounds like a dealbreaker for Comet at first glance. I think we'd be able to make the necessary changes thin the Comet repo, and likely benefit from all of the stuff you're proposing from this change.

We'll be able to start making the migration with the DF 49.0 release?

[adriangb]

Thank you for chiming in!

We'll be able to start making the migration with the DF 49.0 release?

Yes that's the plan. I'm trying to figure out the way to make the migration as incremental and painless as possible. We'll also have to be doing this migration. The current plan as outlined in #16791 is that by default if you don't provide anything you use the new system. If you have a custom SchemaAdapter it will be used for everything. You can start migrating by implementing a PhysicalExprRewriter that has 1:1 pairity with your SchemaAdapter implementation and passing that in as well as your SchemaAdapter, then the SchemaAdapter will be used for projections and the PhysicalExprRewriter will be used for predicate pushdown.

Once we implement projection pushdown without any changes to your code you'd be fully migrated to the new system and can just remove the SchemaAdapter.

[viirya]

I came across this thread via @alamb’s post on dev and was curious to learn more. While I understand the general idea of what’s being proposed, it’s still unclear to me how it will work in practice or what steps are involved.

Could you provide a brief overview of the typical usage of SchemaAdapter, and explain how you plan to replace it with PhysicalExprAdapter? A few concrete examples would also be really helpful.

[adriangb]

Here's an example from Comet: https://github.com/vaibhawvipul/datafusion-comet/blob/main/native/core/src/parquet/schema_adapter.rs. As you can see it's a lot of code with a lot of duplication. The change will essentially be to re-implement that in terms of PhysicalExprAdapter which I think should be a lot less code and duplication. Maybe something like this (very much pseudocode):

/// A cast kernel that respect's Spark's casting rules.
pub struct SparkCast {
     /// The expression to cast
    pub expr: Arc<dyn PhysicalExpr>,
    /// The data type to cast to
    cast_type: DataType,
    /// Cast options
    cast_options: CastOptions<'static>,
}

impl PhysicalExpr for CastExpr {
    ...
    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
        let value = self.expr.evaluate(batch)?;
        spark_parquet_convert(value, &self.cast_type, Some(&self.cast_options))
    }
    ...
}

struct SparkPhysicalExprAdapter {
   inner: DefaultPhysicalExprAdapter,
}

impl PhysicalExprAdapter for SparkPhysicalExprAdapter {
  fn rewrite(&self, mut expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
    expr = self.inner.rewrite(expr);
    expr.transform(|expr| {
        if let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() {
            return Ok(Transformed::yes(SparkCastExpr::from(cast))
        }
        Ok(Transformed::no(expr))
   }).data()
 }
}

[adriangb]

Here's some more examples: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/json_shredding.rs, https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/default_column_values.rs

I think an example with custom casting rules like #16800 (comment) would be helpful.

[parthchandra]

I feel it may be a fair amount of work in Comet to move from SchemaAdapter to PhysicalExprAdapter but from the pseudocode example it appears tractable. I think we'll be willing to take the hit if the promise is improved predicate pushdown performance. :)

[adriangb]

Having looked at your implementation I think it may not be that bad! It seems like most of what your SchemaAdapter is doing is customizing casting rules, right?

[parthchandra]

Having looked at your implementation I think it may not be that bad! It seems like most of what your SchemaAdapter is doing is customizing casting rules, right?

All of it is customized casting rules :)

[adriangb]

@parthchandra @mbutrovich please take a look at #16803.
As per the comments in the example it looks like Comet already has a custom PhysicalExpr for casting so implementing a SchemaAdapter replacement might be <100 LOC because you can re-use your existing kernel / logic.

[viirya]

So the basic idea, as I understand it, is that instead of adapting the data batch using SchemaAdapter against the schema, the new approach involves rewriting or transforming the physical expressions using PhysicalExprAdapter.

There are multiple reasons for this:

Based on this reasoning, I think the first two points make sense.

Could you clarify the latter two? From your description, they sound like areas where PhysicalExprAdapter could bring benefits — but I'm not quite sure how SchemaAdapter fits into those cases. Are these scenarios where the current use of SchemaAdapter is actually blocking progress toward those goals, and replacing it with PhysicalExprAdapter is part of the solution?

[adriangb]

Could you clarify the latter two? From your description, they sound like areas where PhysicalExprAdapter could bring benefits — but I'm not quite sure how SchemaAdapter fits into those cases.

There's some discussion in #14993. Basically if we want to be able to customize how expressions are evaluated for a specific format, in particular how variant_get(column, 'field') or get_field(column, 'field') are executed in the context of a specific format (e.g. in parquet we can read single struct columns or use shredded variant) we need to have access to the expression in ParquetOpener in order to check if the file schema has the shredded variant field and generate the right ProjectionMask.

[parthchandra]

please take a look at #16803.

Thank you for this pointer to the example.