Cache the result of `Expr::data_type` / `Expr::nullable` / `Expr::to_field` to speed up planning

[alamb]

Is your feature request related to a problem or challenge?

Many parts of the planning process require knowing the type, nullability and recently the field information for each Expr

• Expr::to_field
• Expr::get_type
• Expr::nullable

The implementation of these methods recomputes the type information on demand, for example

datafusion/datafusion/expr/src/expr_schema.rs

Line 109 in 4f95e6d

fn get_type(&self, schema: &dyn ExprSchema) -> Result<DataType> {

This is slow for several reasons:

1. The computation is re-computed many times [during](repo:apache/datafusion nullable) planning
2. The computations themselves are sometimes non trivial (e.g. what @pepijnve added in #17801 Improve nullability reporting of case expressions #17813 for case expressions)
3. The calculations have to walk the entire expression tree, each time even when they don't change

I have some small evidence that this type computation takes a non trivial amount of planning time, as for example, when I removed some of the cloning in to_field in this PR, I saw a significant improvement in planning time (2-3%)

• Optimize planning / stop cloning Strings / Fields so much (2-3% faster planning time) #18415

Describe the solution you'd like

I would like to find some way to avoid recomputing the results of these calculations for Expressions

Describe alternatives you've considered

DataFusion already has a similar pattern for ExecutionPlan where the properties are computed once and then each ExecutionPlan must report the properties.

• https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.ExecutionPlan.html#tymethod.properties
• https://docs.rs/datafusion/latest/datafusion/physical_plan/struct.PlanProperties.html

We also need to be careful about rolling this change incrementally rather than as one giant PR out as it will be a potentially disruptive API change (as Exprs are used all over the place)

So one way to cache this would be to add a method like

/// Cached result of computing type for 
struct ExprProperties {
  field: FieldRef,
  table_reference: Option<TableReference>
}

Then add

trait ExprSchemable {
  /// return cached properties for the Expr, if available
  fn try_properties(&self) -> Option<&ExprProperties>;
}

impl Expr {
  fn get_type(&self) -> DataType {
    // first check if we have pre-computed results
    if let Some(properties) = self.try_properties() {
      return properties.field.data_type()
    }
    // otherwise compute as normal
  }

  // add similar checks for nullable and to_field
}

Then we could add an ExprProperties field to various Exprs over multiple PRs and roll out the cached values

Additional context

• See conversation with @pepijnve on #17801 Improve nullability reporting of case expressions #17813 (comment)

[pepijnve]

The main challenge I see here is that all ExprSchemable methods have signatures like

fn get_type(&self, schema: &dyn ExprSchema) -> Result<DataType>;

where the result of the method is dependent on the schema argument. As a consequence you need to use that argument to validate the usability/correctness of the memoised result. In the proposed ExprProperties solution, I'm not sure how you would be able to do that.

For ExecutionPlan this approach does works because plans know about their schema at construction time.

[alamb]

Note that @findepi suggested something similar on

• Proposal: introduced typed expressions, separate AST and IR #12604

@karlovnv proposed something as well on

• Bad performance on wide tables (1000+ columns) #7698

[alamb]

The main challenge I see here is that all ExprSchemable methods have signatures like

fn get_type(&self, schema: &dyn ExprSchema) -> Result<DataType>;

where the result of the method is dependent on the schema argument. As a consequence you need to use that argument to validate the usability/correctness of the memoised result. In the proposed ExprProperties solution, I'm not sure how you would be able to do that.

That is an excellent point. I think most times the schema should be the same, but I am not sure how to ensure that 🤔

Or maybe we need some way to "resolve" an Expr with a Schema and make some new type 🤔

Or I suppose we could put the memo table on the DFSchema itself maybe 🤔

[findepi]

There is no reason to "transplant" a typed (bound) expression from one schema to some other schema with different input types.
Doing so would be treating the expressions as syntactical rather than semantical operations.
Outside of initial SQL plan building (where expressions are constructed syntactical and then applied to a plan), the expression type does not change throughout the planning and optimization process.

That's why I think about this in the terms of typed expressions rather than "caching data type and nullability".
However, I can also see how caching may look more approachable to instill over existing code base.

[alamb]

Introducing the idea of typed / bound expressions is interesting 🤔 As you have said, we would then have to be much more explicit about bound vs unbound expressions in various logical and physical structures

[alamb]

In many ways a BoundExpr would basically be a Expr with the DataType and Nullability (and metadata)

DataFusion already does have an untyped IR (the sqlparser-rs AST) but it doesn't tightly bind the Exprs that come out of this

Perhaps one reason Expr isn't bound is that it is also part of DataFusion's public API -- users (e.g. DataFrame users) can create Exprs directly, they don't (just) come from SQL