Add lambda substrait support

[gstvg]

Which issue does this PR close?

Part of #21172

Rationale for this change

Substrait support wasn't implemented in the core lambda support to reduce PR size

What changes are included in this PR?

Substrait consuming and producing of higher-order functions, lambdas and lambda variables

Are these changes tested?

Unit tests added to datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

Are there any user-facing changes?

None

[benbellick]

👋 Hello from @substrait-io. Great to see the core lambda PR has gotten through! Once this PR is in a ready to review state and is rebased off of main, I will be more than happy to help review it 🙂

[gstvg]

Thanks @benbellick, I will open this tonight. Besides rebasing, I believe it misses some tests (I tested with sqllogictests only)

[rluvaton]

I updated your branch with main since the fix for detect breaking changes was resolved now, sorry for the trouble

[benbellick]

It overall looks good to me! I left a few comments on some stylistic things but the only thing that I would particularly like to see is just the tests for consumer / producer independently. Thanks!

[benbellick]

Is there a reason this is public? This feels like an implementation detail of the default lambda-handling logic. What about:

Suggested change

	pub struct DefaultSubstraitLambdaConsumer {
	struct LambdaConsumerState {

[gstvg]

The existing required methods for trait SubstraitConsumer are trivial to implement, but that's not the case for the newly added lambda methods. This is a just a convenience to custom implementations which don't want to customize the default lambda handling, should I remove it?

struct CustomSubstraitConsumer {
     extensions: Arc<Extensions>,
     state: Arc<SessionState>,
     // You can reuse existing consumer code related to lambdas
     lambda_consumer: DefaultSubstraitLambdaConsumer,
 }

 #[async_trait]
 impl SubstraitConsumer for CustomSubstraitConsumer {
     async fn resolve_table_ref(
         &self,
         table_ref: &TableReference,
     ) -> Result<Option<Arc<dyn TableProvider>>> {
         let table = table_ref.table().to_string();
         let schema = self.state.schema_for_ref(table_ref.clone())?;
         let table_provider = schema.table(&table).await?;
         Ok(table_provider)
     }

     fn get_extensions(&self) -> &Extensions {
         self.extensions.as_ref()
     }

     fn get_function_registry(&self) -> &impl FunctionRegistry {
         self.state.as_ref()
     }

     fn with_lambda_parameters(
        &self,
        lambda_parameters: &[Type],
        input_schema: &DFSchema,
    ) -> datafusion::common::Result<(Vec<String>, Self)> {
        let (names, lambda_consumer) = self.lambda_consumer.with_lambda_parameters(
            self,
            lambda_parameters,
            input_schema,
        )?;

        Ok((
            names,
            Self {
                extensions: self.extensions.clone(),
                state: self.state.clone(),
                lambda_consumer,
            },
        ))
    }

    fn lambda_variable(
        &self,
        steps_out: usize,
        field_idx: usize,
    ) -> datafusion::common::Result<Expr> {
        self.lambda_consumer.lambda_variable(steps_out, field_idx)
    }
 }

[benbellick]

Ah I see now, I appreciate the explanation. That makes a lot of sense then!

AFAICT the strategy for outer schemas is to provide default impls so that consumer implementors who don't care about the behavior can ignore their existence, but then will encounter a runtime error if they are used:

datafusion/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs

Lines 437 to 449 in 6cdfc05

	/// Push an outer schema onto the stack when entering a subquery.
	fn push_outer_schema(&self, _schema: Arc<DFSchema>) {}
	/// Pop an outer schema from the stack when leaving a subquery.
	fn pop_outer_schema(&self) {}
	/// Get the outer schema at the given nesting depth.
	/// `steps_out = 1` is the immediately enclosing query, `steps_out = 2`
	/// is two levels out, etc. Returns `None` if `steps_out` is 0 or
	/// exceeds the current nesting depth (the caller should treat this as
	/// an error in the Substrait plan).
	fn get_outer_schema(&self, _steps_out: usize) -> Option<Arc<DFSchema>> {
	None

I'm wondering if we should do the same thing for lambdas? Rather than enforce that implementors must implement the lambda-handling fns, they could instead optionally ignore them, resulting in an error if lambda expressions are encountered.

What do you think? I am not so particular here TBH. Ultimately my goal is to validate that the translation itself is correct, and I am happy to leave API concerns to the project maintainers :)

[gstvg]

Ah, I agree, they definetively should have a default impl returning an error to avoid a breaking change, thank you 2c23439
But even then, I still believe a public DefaultSubstraitLambdaConsumer is convenient, but yes, let's leave that for that maintainers, thanks

[benbellick]

Same comment as on the consumer side. I wonder if we can just keep this private, since its usage in implementing this producer is an implementation detail.

[gstvg]

Same reason for the consumer. What get's decided there (#21193 (comment)) I'll also apply here

struct CustomSubstraitProducer {
    extensions: Extensions,
    state: Arc<SessionState>,
    // You can reuse existing producer code related to lambdas
    lambda_producer: DefaultSubstraitLambdaProducer,
}

impl SubstraitProducer for CustomSubstraitProducer {
    fn register_function(&mut self, signature: String) -> u32 {
        self.extensions.register_function(&signature)
    }

    fn register_type(&mut self, type_name: String) -> u32 {
        self.extensions.register_type(&type_name)
    }

    fn get_extensions(self) -> Extensions {
        self.extensions
    }

    fn push_lambda_parameters(
        &mut self,
        lambda_parameters: Vec<FieldRef>,
    ) -> datafusion::common::Result<()> {
        let lambda_parameters_map = lambda_parameters_map(self, lambda_parameters)?;

        self.lambda_producer
            .push_lambda_parameters(lambda_parameters_map);

        Ok(())
    }

    fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> {
        self.lambda_producer.pop_lambda_parameters()
    }

    fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> {
        self.lambda_producer.lambda_variable(name)
    }

    fn lambda_parameter_type(
        &self,
        name: &str,
    ) -> datafusion::common::Result<substrait::proto::Type> {
        self.lambda_producer.lambda_parameter_type(name)
    }
}

[gstvg]

@benbellick I usually follows the pattern of existing methods, like push/pop_outer_schema from #20439. This is my first time dealing with substrait, so I may be wrong, but I didn't followed this pattern (using push/pop_lambda_parameter) because it modifies &self via a RwLock and I'm note sure this couldn't lead to conflicts if the same consumer is used to consume two different plans at the same time in different threads. If that's not the case I can change this to use push/pop_lambda_parameter as well.

[benbellick]

Ah, that is an interesting point... Is it expected/supported for the same SubstraitConsumer instance to be used concurrently?

If not, then I think it would be simpler and consistent to model lambda scope the same way as outer schemas.

If yes, then the scoped-consumer approach here makes sense, but it seems like the existing push_outer_schema / pop_outer_schema stack may have the same interleaving issue and should probably be addressed separately.

[gstvg]

I'll confirm with the maintainer who ends up reviewing this, but SubstraitConsumer is both Send + Sync and all it's methods take &self. Since the default consumer is cheap to create, and I expect most/all custom ones to be cheap as well, I guess it's mostly due to async and to easily embed it into other structures which should also implement Send + Sync than to allow efficient concurrent usage.
I won't expect any consumer to be used concurrently for performance, but, since it can be used, I think it's possible that it's/will be used concurrently incidentally as the easier/natural way within a given codebase

[benbellick]

Few more comments but on the whole this looks good to me! Thanks

[benbellick]

LGTM, great work! Excited to get this in 🚀

Feel free to summon the maintainers now.

[gstvg]

LGTM, great work! Excited to get this in 🚀

Thanks for the review @benbellick
cc @rluvaton @comphead this is ready for the final review

[rluvaton]

I'm unfamiliar with substrait except for the idea, so I don't know if I can have meaningful review

[rluvaton]

@gstvg according to the breaking change detector there are no breaking change api wise in this pr, is your commemt still correct?

Are there any user-facing changes?

Yes, there are breaking changes, new methods without default implementation have been added to SubstraitConsumer and SubstraitProducer

[gstvg]

@rluvaton, indeed, after applying @benbellick suggestions there's no breaking changes anymore, description updated, thanks