You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "vbarua (via GitHub)" <gi...@apache.org> on 2024/12/17 02:32:03 UTC
[PR] feat(substrait): modular substrait consumer [datafusion]
vbarua opened a new pull request, #13803:
URL: https://github.com/apache/datafusion/pull/13803
Note: When reviewing this PR I would recommend hiding whitespace changes.
## Which issue does this PR close?
Initial work for https://github.com/apache/datafusion/issues/13318
## Rationale for this change
Improves the reusability of the Substrait consumer for users that utilize user-defined extensions and types.
## What changes are included in this PR?
* Relation and expression handling code has been extracted into a series of `from_substrait_*` functions (i.e. `from_filter_rel`, `from_scalar_function`) to aid re-use.
* A `SubstraitConsumer` trait has been introduced with default methods that handle relations and expression using the above functions.
* The `SubstraitConsumer` trait also includes methods for handling extension relations (i.e `consume_extension_single`) along with user-defined types (`consume_user_defined_type`, `consume_user_defined_literal`).
* All conversion methods take as their first argument an `&impl SubstraitConsumer`.
* The role of `SubstraitPlanningState` has been subsumed into the `SubstraitConsumer`.
## Are these changes tested?
These changes refactor existing code and leverage their tests.
A doc comment has been added to show how the `SubstraitConsumer` trait can be implemented by users.
## Are there any user-facing changes?
Yes.
The top-level `from_substrait_plan(state, plan)` function now takes a `&SessionState` as its first argument, instead of `&dyn SubstraitPlanningState`. The same is true for `from_substrait_extended_expr`. These two functions are the primary external facing APIs for intiation a conversion.
A number of functions like `from_substrait_rel`, `from_substrait_rex`, etc have had their API changed to make the first argument `&impl SubstraitConsumer`, and remove all `state: &dyn SubstraitPlanningState` and `extensions: &Extensions` arguments.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1888625076
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -1503,370 +1858,416 @@ pub async fn from_substrait_func_args(
/// Convert Substrait AggregateFunction to DataFusion Expr
pub async fn from_substrait_agg_func(
- state: &dyn SubstraitPlanningState,
+ consumer: &impl SubstraitConsumer,
f: &AggregateFunction,
input_schema: &DFSchema,
- extensions: &Extensions,
filter: Option<Box<Expr>>,
order_by: Option<Vec<SortExpr>>,
distinct: bool,
) -> Result<Arc<Expr>> {
- let args =
- from_substrait_func_args(state, &f.arguments, input_schema, extensions).await?;
-
- let Some(function_name) = extensions.functions.get(&f.function_reference) else {
+ let Some(fn_signature) = consumer
+ .get_extensions()
Review Comment:
You mentioned you'd prefer to not have the get_extensions(), I guess you could replace it with something like `get_function` which in DefaultSubstraitConsumer would just do this `extension.functions.get(..)`? Or did you have something else in mind?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894475748
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -262,7 +718,25 @@ async fn except_rels(
/// Convert Substrait Plan to DataFusion LogicalPlan
pub async fn from_substrait_plan(
- state: &dyn SubstraitPlanningState,
+ state: &SessionState,
+ plan: &Plan,
+) -> Result<LogicalPlan> {
+ // Register function extension
+ let extensions = Extensions::try_from(&plan.extensions)?;
+ if !extensions.type_variations.is_empty() {
+ return not_impl_err!("Type variation extensions are not supported");
+ }
+
+ let consumer = DefaultSubstraitConsumer {
+ extensions: &extensions,
+ state,
+ };
+ from_substrait_plan_with_consumer(&consumer, plan).await
+}
+
+/// Convert Substrait Plan to DataFusion LogicalPlan using the given consumer
+pub async fn from_substrait_plan_with_consumer(
+ consumer: &impl SubstraitConsumer,
Review Comment:
That is redundant now isn't it. Removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894475383
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -262,7 +718,25 @@ async fn except_rels(
/// Convert Substrait Plan to DataFusion LogicalPlan
pub async fn from_substrait_plan(
- state: &dyn SubstraitPlanningState,
+ state: &SessionState,
+ plan: &Plan,
+) -> Result<LogicalPlan> {
+ // Register function extension
+ let extensions = Extensions::try_from(&plan.extensions)?;
+ if !extensions.type_variations.is_empty() {
+ return not_impl_err!("Type variation extensions are not supported");
+ }
+
+ let consumer = DefaultSubstraitConsumer {
+ extensions: &extensions,
Review Comment:
I do want to streamline the API further, but in the future 😅
Initialising the consumer directly from the plan makes sense to me, because we always need the extensions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894445457
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -262,7 +718,25 @@ async fn except_rels(
/// Convert Substrait Plan to DataFusion LogicalPlan
pub async fn from_substrait_plan(
- state: &dyn SubstraitPlanningState,
+ state: &SessionState,
+ plan: &Plan,
+) -> Result<LogicalPlan> {
+ // Register function extension
+ let extensions = Extensions::try_from(&plan.extensions)?;
+ if !extensions.type_variations.is_empty() {
+ return not_impl_err!("Type variation extensions are not supported");
+ }
+
+ let consumer = DefaultSubstraitConsumer {
+ extensions: &extensions,
+ state,
+ };
+ from_substrait_plan_with_consumer(&consumer, plan).await
+}
+
+/// Convert Substrait Plan to DataFusion LogicalPlan using the given consumer
+pub async fn from_substrait_plan_with_consumer(
+ consumer: &impl SubstraitConsumer,
Review Comment:
nit: I think if you want to, you could remove the Extension lines from below - no need to duplicate them there given they're no longer used
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -262,7 +718,25 @@ async fn except_rels(
/// Convert Substrait Plan to DataFusion LogicalPlan
pub async fn from_substrait_plan(
- state: &dyn SubstraitPlanningState,
+ state: &SessionState,
+ plan: &Plan,
+) -> Result<LogicalPlan> {
+ // Register function extension
+ let extensions = Extensions::try_from(&plan.extensions)?;
+ if !extensions.type_variations.is_empty() {
+ return not_impl_err!("Type variation extensions are not supported");
+ }
+
+ let consumer = DefaultSubstraitConsumer {
+ extensions: &extensions,
Review Comment:
I started thinking whether it actually really make sense to pass the Extensions into the SubstraitConsumer, or if there should be some `initialize(plan: &Plan)` method or something in the SubstraitConsumer that loads them. But I didn't think yet much about it, maybe there is a reason to rather have it like this, and anyways no need to change anything for this PR anymore in any case 😄
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1893185369
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +101,471 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::{FunctionRegistry, SessionState};
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[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()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and FunctionRegistry is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_function_registry(&self) -> &impl FunctionRegistry;
+
+ // Relation Methods
Review Comment:
Since these `//` comments aren't displayed in the documentationI think we should also add some documentation to them eventually in `///` style
https://docs.rs/datafusion-substrait/41.0.0/datafusion_substrait/
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +101,471 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::{FunctionRegistry, SessionState};
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
Review Comment:
You can make examples like this potentially read more nicely by putting `#` in front of the import lines so they aren't rendered int he output.
I think it is also fine the way it is too
```suggestion
/// # use async_trait::async_trait;
/// # use datafusion::catalog::TableProvider;
/// # use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
/// # use datafusion::error::Result;
/// # use datafusion::execution::{FunctionRegistry, SessionState};
/// # use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
/// # use std::sync::Arc;
/// # use substrait::proto;
/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
/// # use datafusion::arrow::datatypes::DataType;
/// # use datafusion::logical_expr::expr::ScalarFunction;
/// # use datafusion_substrait::extensions::Extensions;
/// # use datafusion_substrait::logical_plan::consumer::{
/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
/// # };
```
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +101,471 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::{FunctionRegistry, SessionState};
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[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()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
Review Comment:
Do you intend to handle this TODO in this PR or is it for some future PR? If it is for a future PR perhaps you can file a ticket to track the extra work and link to it here
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -989,8 +992,13 @@ async fn extension_logical_plan() -> Result<()> {
}),
});
+ let consumer = DefaultSubstraitConsumer::new(
+ Arc::new(Extensions::default()),
+ Arc::new(ctx.state()),
+ );
+
let proto = to_substrait_plan(&ext_plan, &ctx.state())?;
- let plan2 = from_substrait_plan(&ctx.state(), &proto).await?;
+ let plan2 = from_substrait_plan_with_consumer(&consumer, &proto).await?;
Review Comment:
why doesn't this call `from_substrait_plan` ?
I ask because I was initially confused and thought that the main public APIs for substrait had changed.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +101,471 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
Review Comment:
😍
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -262,7 +724,25 @@ async fn except_rels(
/// Convert Substrait Plan to DataFusion LogicalPlan
pub async fn from_substrait_plan(
Review Comment:
💯 for not changing the API
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +101,471 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::{FunctionRegistry, SessionState};
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[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()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and FunctionRegistry is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_function_registry(&self) -> &impl FunctionRegistry;
+
+ // Relation Methods
+ // There is one method per Substrait relation to allow for easy overriding of consumer behaviour.
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_read(&self, rel: &ReadRel) -> Result<LogicalPlan> {
+ from_read_rel(self, rel).await
+ }
+
+ async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+ from_filter_rel(self, rel).await
+ }
+
+ async fn consume_fetch(&self, rel: &FetchRel) -> Result<LogicalPlan> {
+ from_fetch_rel(self, rel).await
+ }
+
+ async fn consume_aggregate(&self, rel: &AggregateRel) -> Result<LogicalPlan> {
+ from_aggregate_rel(self, rel).await
+ }
+
+ async fn consume_sort(&self, rel: &SortRel) -> Result<LogicalPlan> {
+ from_sort_rel(self, rel).await
+ }
+
+ async fn consume_join(&self, rel: &JoinRel) -> Result<LogicalPlan> {
+ from_join_rel(self, rel).await
+ }
+
+ async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+ from_project_rel(self, rel).await
+ }
+
+ async fn consume_set(&self, rel: &SetRel) -> Result<LogicalPlan> {
+ from_set_rel(self, rel).await
+ }
+
+ async fn consume_cross(&self, rel: &CrossRel) -> Result<LogicalPlan> {
+ from_cross_rel(self, rel).await
+ }
+
+ async fn consume_consistent_partition_window(
+ &self,
+ _rel: &ConsistentPartitionWindowRel,
+ ) -> Result<LogicalPlan> {
+ not_impl_err!("Consistent Partition Window Rel not supported")
+ }
+
+ async fn consume_exchange(&self, rel: &ExchangeRel) -> Result<LogicalPlan> {
+ from_exchange_rel(self, rel).await
+ }
+
+ // Expression Methods
+ // There is one method per Substrait expression to allow for easy overriding of consumer behaviour
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_literal(&self, expr: &Literal) -> Result<Expr> {
+ from_literal(self, expr).await
+ }
+
+ async fn consume_field_reference(
+ &self,
+ expr: &FieldReference,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_field_reference(self, expr, input_schema).await
+ }
+
+ async fn consume_scalar_function(
+ &self,
+ expr: &ScalarFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_scalar_function(self, expr, input_schema).await
+ }
+
+ async fn consume_window_function(
+ &self,
+ expr: &WindowFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_window_function(self, expr, input_schema).await
+ }
+
+ async fn consume_if_then(
+ &self,
+ expr: &IfThen,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_if_then(self, expr, input_schema).await
+ }
+
+ async fn consume_switch(
+ &self,
+ _expr: &SwitchExpression,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Switch expression not supported")
+ }
+
+ async fn consume_singular_or_list(
+ &self,
+ expr: &SingularOrList,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_singular_or_list(self, expr, input_schema).await
+ }
+
+ async fn consume_multi_or_list(
+ &self,
+ _expr: &MultiOrList,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Multi Or List expression not supported")
+ }
+
+ async fn consume_cast(
+ &self,
+ expr: &substrait_expression::Cast,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_cast(self, expr, input_schema).await
+ }
+
+ async fn consume_subquery(
+ &self,
+ expr: &substrait_expression::Subquery,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_subquery(self, expr, input_schema).await
+ }
+
+ async fn consume_nested(
+ &self,
+ _expr: &Nested,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Nested expression not supported")
+ }
+
+ async fn consume_enum(&self, _expr: &Enum, _input_schema: &DFSchema) -> Result<Expr> {
+ not_impl_err!("Enum expression not supported")
+ }
+
+ // User-Defined Functionality
+
+ // The details of extension relations, and how to handle them, are fully up to users to specify.
+ // The following methods allow users to customize the consumer behaviour
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionLeafRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionLeafRel")
+ }
+
+ async fn consume_extension_single(
+ &self,
+ rel: &ExtensionSingleRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionSingleRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionSingleRel")
+ }
+
+ async fn consume_extension_multi(
+ &self,
+ rel: &ExtensionMultiRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionMultiRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionMultiRel")
+ }
+
+ // Users can bring their own types to Substrait which require custom handling
+
+ fn consume_user_defined_type(
+ &self,
+ user_defined_type: &r#type::UserDefined,
+ ) -> Result<DataType> {
+ substrait_err!(
+ "Missing handler for user-defined type: {}",
+ user_defined_type.type_reference
+ )
+ }
+
+ fn consume_user_defined_literal(
+ &self,
+ user_defined_literal: &proto::expression::literal::UserDefined,
+ ) -> Result<ScalarValue> {
+ substrait_err!(
+ "Missing handler for user-defined literals {}",
+ user_defined_literal.type_reference
+ )
+ }
+}
+
+/// Convert Substrait Rel to DataFusion DataFrame
+#[async_recursion]
+pub async fn from_substrait_rel(
+ consumer: &impl SubstraitConsumer,
+ relation: &Rel,
+) -> Result<LogicalPlan> {
+ let plan: Result<LogicalPlan> = match &relation.rel_type {
+ Some(rel_type) => match rel_type {
+ RelType::Read(rel) => consumer.consume_read(rel).await,
+ RelType::Filter(rel) => consumer.consume_filter(rel).await,
+ RelType::Fetch(rel) => consumer.consume_fetch(rel).await,
+ RelType::Aggregate(rel) => consumer.consume_aggregate(rel).await,
+ RelType::Sort(rel) => consumer.consume_sort(rel).await,
+ RelType::Join(rel) => consumer.consume_join(rel).await,
+ RelType::Project(rel) => consumer.consume_project(rel).await,
+ RelType::Set(rel) => consumer.consume_set(rel).await,
+ RelType::ExtensionSingle(rel) => consumer.consume_extension_single(rel).await,
+ RelType::ExtensionMulti(rel) => consumer.consume_extension_multi(rel).await,
+ RelType::ExtensionLeaf(rel) => consumer.consume_extension_leaf(rel).await,
+ RelType::Cross(rel) => consumer.consume_cross(rel).await,
+ RelType::Window(rel) => {
+ consumer.consume_consistent_partition_window(rel).await
+ }
+ RelType::Exchange(rel) => consumer.consume_exchange(rel).await,
+ rt => not_impl_err!("{rt:?} rel not supported yet"),
+ },
+ None => return substrait_err!("rel must set rel_type"),
+ };
+ apply_emit_kind(retrieve_rel_common(relation), plan?)
+}
+
+/// Can be used to consume standard Substrait without user-defined extensions
Review Comment:
Perhaps you could create an example for this one or point at the existing documentation examples
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -1503,370 +1858,416 @@ pub async fn from_substrait_func_args(
/// Convert Substrait AggregateFunction to DataFusion Expr
pub async fn from_substrait_agg_func(
- state: &dyn SubstraitPlanningState,
+ consumer: &impl SubstraitConsumer,
f: &AggregateFunction,
input_schema: &DFSchema,
- extensions: &Extensions,
filter: Option<Box<Expr>>,
order_by: Option<Vec<SortExpr>>,
distinct: bool,
) -> Result<Arc<Expr>> {
- let args =
- from_substrait_func_args(state, &f.arguments, input_schema, extensions).await?;
-
- let Some(function_name) = extensions.functions.get(&f.function_reference) else {
+ let Some(fn_signature) = consumer
+ .get_extensions()
Review Comment:
It seems like perhaps we could make a specific API in the trait for resolving functions or something 🤔
I also think this is fine to leave in the initial PR
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1888608145
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +977,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
+ consumer: &impl SubstraitConsumer,
+ p: &ProjectRel,
) -> Result<LogicalPlan> {
- let plan: Result<LogicalPlan> = match &rel.rel_type {
- Some(RelType::Project(p)) => {
- if let Some(input) = p.input.as_ref() {
- let mut input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let original_schema = input.schema().clone();
-
- // Ensure that all expressions have a unique display name, so that
- // validate_unique_names does not fail when constructing the project.
- let mut name_tracker = NameTracker::new();
-
- // By default, a Substrait Project emits all inputs fields followed by all expressions.
- // We build the explicit expressions first, and then the input expressions to avoid
- // adding aliases to the explicit expressions (as part of ensuring unique names).
- //
- // This is helpful for plan visualization and tests, because when DataFusion produces
- // Substrait Projects it adds an output mapping that excludes all input columns
- // leaving only explicit expressions.
-
- let mut explicit_exprs: Vec<Expr> = vec![];
- for expr in &p.expressions {
- let e = from_substrait_rex(
- state,
- expr,
- input.clone().schema(),
- extensions,
- )
- .await?;
- // if the expression is WindowFunction, wrap in a Window relation
- if let Expr::WindowFunction(_) = &e {
- // Adding the same expression here and in the project below
- // works because the project's builder uses columnize_expr(..)
- // to transform it into a column reference
- input = input.window(vec![e.clone()])?
- }
- explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
-
- let mut final_exprs: Vec<Expr> = vec![];
- for index in 0..original_schema.fields().len() {
- let e = Expr::Column(Column::from(
- original_schema.qualified_field(index),
- ));
- final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
- final_exprs.append(&mut explicit_exprs);
-
- input.project(final_exprs)?.build()
- } else {
- not_impl_err!("Projection without an input is not supported")
+ if let Some(input) = p.input.as_ref() {
+ let mut input =
+ LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let original_schema = input.schema().clone();
+
+ // Ensure that all expressions have a unique display name, so that
+ // validate_unique_names does not fail when constructing the project.
+ let mut name_tracker = NameTracker::new();
+
+ // By default, a Substrait Project emits all inputs fields followed by all expressions.
+ // We build the explicit expressions first, and then the input expressions to avoid
+ // adding aliases to the explicit expressions (as part of ensuring unique names).
+ //
+ // This is helpful for plan visualization and tests, because when DataFusion produces
+ // Substrait Projects it adds an output mapping that excludes all input columns
+ // leaving only explicit expressions.
+
+ let mut explicit_exprs: Vec<Expr> = vec![];
+ for expr in &p.expressions {
+ let e = from_substrait_rex(consumer, expr, input.clone().schema()).await?;
+ // if the expression is WindowFunction, wrap in a Window relation
+ if let Expr::WindowFunction(_) = &e {
+ // Adding the same expression here and in the project below
+ // works because the project's builder uses columnize_expr(..)
+ // to transform it into a column reference
+ input = input.window(vec![e.clone()])?
}
+ explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Filter(filter)) => {
- if let Some(input) = filter.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- if let Some(condition) = filter.condition.as_ref() {
- let expr =
- from_substrait_rex(state, condition, input.schema(), extensions)
- .await?;
- input.filter(expr)?.build()
- } else {
- not_impl_err!("Filter without an condition is not valid")
- }
- } else {
- not_impl_err!("Filter without an input is not valid")
- }
+
+ let mut final_exprs: Vec<Expr> = vec![];
+ for index in 0..original_schema.fields().len() {
+ let e = Expr::Column(Column::from(original_schema.qualified_field(index)));
+ final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Fetch(fetch)) => {
- if let Some(input) = fetch.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let offset = fetch.offset as usize;
- // -1 means that ALL records should be returned
- let count = if fetch.count == -1 {
- None
- } else {
- Some(fetch.count as usize)
- };
- input.limit(offset, count)?.build()
- } else {
- not_impl_err!("Fetch without an input is not valid")
- }
+ final_exprs.append(&mut explicit_exprs);
+ input.project(final_exprs)?.build()
+ } else {
+ not_impl_err!("Projection without an input is not supported")
+ }
+}
+
+#[async_recursion]
+pub async fn from_filter_rel(
+ consumer: &impl SubstraitConsumer,
+ filter: &FilterRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = filter.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ if let Some(condition) = filter.condition.as_ref() {
+ let expr = from_substrait_rex(consumer, condition, input.schema()).await?;
+ input.filter(expr)?.build()
+ } else {
+ not_impl_err!("Filter without an condition is not valid")
}
- Some(RelType::Sort(sort)) => {
- if let Some(input) = sort.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let sorts =
- from_substrait_sorts(state, &sort.sorts, input.schema(), extensions)
- .await?;
- input.sort(sorts)?.build()
- } else {
- not_impl_err!("Sort without an input is not valid")
- }
+ } else {
+ not_impl_err!("Filter without an input is not valid")
+ }
+}
+
+#[async_recursion]
+pub async fn from_fetch_rel(
+ consumer: &impl SubstraitConsumer,
+ fetch: &FetchRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = fetch.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let offset = fetch.offset as usize;
+ // -1 means that ALL records should be returned
+ let count = if fetch.count == -1 {
+ None
+ } else {
+ Some(fetch.count as usize)
+ };
+ input.limit(offset, count)?.build()
+ } else {
+ not_impl_err!("Fetch without an input is not valid")
+ }
+}
+
+pub async fn from_sort_rel(
+ consumer: &impl SubstraitConsumer,
+ sort: &SortRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = sort.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?;
+ input.sort(sorts)?.build()
+ } else {
+ not_impl_err!("Sort without an input is not valid")
+ }
+}
+
+pub async fn from_aggregate_rel(
+ consumer: &impl SubstraitConsumer,
+ agg: &AggregateRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = agg.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let mut ref_group_exprs = vec![];
+
+ for e in &agg.grouping_expressions {
+ let x = from_substrait_rex(consumer, e, input.schema()).await?;
+ ref_group_exprs.push(x);
}
- Some(RelType::Aggregate(agg)) => {
- if let Some(input) = agg.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let mut ref_group_exprs = vec![];
- for e in &agg.grouping_expressions {
- let x =
- from_substrait_rex(state, e, input.schema(), extensions).await?;
- ref_group_exprs.push(x);
+ let mut group_exprs = vec![];
+ let mut aggr_exprs = vec![];
+
+ match agg.groupings.len() {
+ 1 => {
+ group_exprs.extend_from_slice(
+ &from_substrait_grouping(
+ consumer,
+ &agg.groupings[0],
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?,
+ );
+ }
+ _ => {
+ let mut grouping_sets = vec![];
+ for grouping in &agg.groupings {
+ let grouping_set = from_substrait_grouping(
+ consumer,
+ grouping,
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?;
+ grouping_sets.push(grouping_set);
}
+ // Single-element grouping expression of type Expr::GroupingSet.
+ // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+ // parsed by the producer and consumer, since Substrait does not have a type dedicated
+ // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+ group_exprs
+ .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+ }
+ };
- let mut group_exprs = vec![];
- let mut aggr_exprs = vec![];
-
- match agg.groupings.len() {
- 1 => {
- group_exprs.extend_from_slice(
- &from_substrait_grouping(
- state,
- &agg.groupings[0],
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?,
- );
- }
- _ => {
- let mut grouping_sets = vec![];
- for grouping in &agg.groupings {
- let grouping_set = from_substrait_grouping(
- state,
- grouping,
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?;
- grouping_sets.push(grouping_set);
+ for m in &agg.measures {
+ let filter = match &m.filter {
+ Some(fil) => Some(Box::new(
+ from_substrait_rex(consumer, fil, input.schema()).await?,
+ )),
+ None => None,
+ };
+ let agg_func = match &m.measure {
+ Some(f) => {
+ let distinct = match f.invocation {
+ _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+ true
}
- // Single-element grouping expression of type Expr::GroupingSet.
- // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
- // parsed by the producer and consumer, since Substrait does not have a type dedicated
- // to ROLLUP. Only vector of Groupings (grouping sets) is available.
- group_exprs.push(Expr::GroupingSet(GroupingSet::GroupingSets(
- grouping_sets,
- )));
- }
- };
-
- for m in &agg.measures {
- let filter = match &m.filter {
- Some(fil) => Some(Box::new(
- from_substrait_rex(state, fil, input.schema(), extensions)
- .await?,
- )),
- None => None,
+ _ if f.invocation == AggregationInvocation::All as i32 => false,
+ _ => false,
};
- let agg_func = match &m.measure {
- Some(f) => {
- let distinct = match f.invocation {
- _ if f.invocation
- == AggregationInvocation::Distinct as i32 =>
- {
- true
- }
- _ if f.invocation
- == AggregationInvocation::All as i32 =>
- {
- false
- }
- _ => false,
- };
- let order_by = if !f.sorts.is_empty() {
- Some(
- from_substrait_sorts(
- state,
- &f.sorts,
- input.schema(),
- extensions,
- )
- .await?,
- )
- } else {
- None
- };
-
- from_substrait_agg_func(
- state,
- f,
- input.schema(),
- extensions,
- filter,
- order_by,
- distinct,
- )
- .await
- }
- None => not_impl_err!(
- "Aggregate without aggregate function is not supported"
- ),
+ let order_by = if !f.sorts.is_empty() {
+ Some(
+ from_substrait_sorts(consumer, &f.sorts, input.schema())
+ .await?,
+ )
+ } else {
+ None
};
- aggr_exprs.push(agg_func?.as_ref().clone());
- }
- input.aggregate(group_exprs, aggr_exprs)?.build()
- } else {
- not_impl_err!("Aggregate without an input is not valid")
- }
- }
- Some(RelType::Join(join)) => {
- if join.post_join_filter.is_some() {
- return not_impl_err!(
- "JoinRel with post_join_filter is not yet supported"
- );
- }
- let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
-
- let join_type = from_substrait_jointype(join.r#type)?;
- // The join condition expression needs full input schema and not the output schema from join since we lose columns from
- // certain join types such as semi and anti joins
- let in_join_schema = left.schema().join(right.schema())?;
-
- // If join expression exists, parse the `on` condition expression, build join and return
- // Otherwise, build join with only the filter, without join keys
- match &join.expression.as_ref() {
- Some(expr) => {
- let on = from_substrait_rex(state, expr, &in_join_schema, extensions)
- .await?;
- // The join expression can contain both equal and non-equal ops.
- // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
- // So we extract each part as follows:
- // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
- // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
- let (join_ons, nulls_equal_nulls, join_filter) =
- split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
- let (left_cols, right_cols): (Vec<_>, Vec<_>) =
- itertools::multiunzip(join_ons);
- left.join_detailed(
- right.build()?,
- join_type,
- (left_cols, right_cols),
- join_filter,
- nulls_equal_nulls,
- )?
- .build()
+ from_substrait_agg_func(
+ consumer,
+ f,
+ input.schema(),
+ filter,
+ order_by,
+ distinct,
+ )
+ .await
}
None => {
- let on: Vec<String> = vec![];
- left.join_detailed(
- right.build()?,
- join_type,
- (on.clone(), on),
- None,
- false,
- )?
- .build()
+ not_impl_err!("Aggregate without aggregate function is not supported")
}
- }
+ };
+ aggr_exprs.push(agg_func?.as_ref().clone());
}
- Some(RelType::Cross(cross)) => {
- let left = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
- left.cross_join(right.build()?)?.build()
+ input.aggregate(group_exprs, aggr_exprs)?.build()
+ } else {
+ not_impl_err!("Aggregate without an input is not valid")
+ }
+}
+
+pub async fn from_join_rel(
+ consumer: &impl SubstraitConsumer,
+ join: &JoinRel,
+) -> Result<LogicalPlan> {
+ if join.post_join_filter.is_some() {
+ return not_impl_err!("JoinRel with post_join_filter is not yet supported");
+ }
+
+ let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+
+ let join_type = from_substrait_jointype(join.r#type)?;
+ // The join condition expression needs full input schema and not the output schema from join since we lose columns from
+ // certain join types such as semi and anti joins
+ let in_join_schema = left.schema().join(right.schema())?;
+
+ // If join expression exists, parse the `on` condition expression, build join and return
+ // Otherwise, build join with only the filter, without join keys
+ match &join.expression.as_ref() {
+ Some(expr) => {
+ let on = from_substrait_rex(consumer, expr, &in_join_schema).await?;
+ // The join expression can contain both equal and non-equal ops.
+ // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
+ // So we extract each part as follows:
+ // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
+ // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
+ let (join_ons, nulls_equal_nulls, join_filter) =
+ split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) =
+ itertools::multiunzip(join_ons);
+ left.join_detailed(
+ right.build()?,
+ join_type,
+ (left_cols, right_cols),
+ join_filter,
+ nulls_equal_nulls,
+ )?
+ .build()
}
- Some(RelType::Read(read)) => {
- async fn read_with_schema(
- state: &dyn SubstraitPlanningState,
- table_ref: TableReference,
- schema: DFSchema,
- projection: &Option<MaskExpression>,
- ) -> Result<LogicalPlan> {
- let schema = schema.replace_qualifier(table_ref.clone());
-
- let plan = {
- let provider = match state.table(&table_ref).await? {
- Some(ref provider) => Arc::clone(provider),
- _ => return plan_err!("No table named '{table_ref}'"),
- };
+ None => {
+ let on: Vec<String> = vec![];
+ left.join_detailed(right.build()?, join_type, (on.clone(), on), None, false)?
+ .build()
+ }
+ }
+}
- LogicalPlanBuilder::scan(
- table_ref,
- provider_as_source(Arc::clone(&provider)),
- None,
- )?
- .build()?
- };
+pub async fn from_cross_rel(
+ consumer: &impl SubstraitConsumer,
+ cross: &CrossRel,
+) -> Result<LogicalPlan> {
+ let left = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+ left.cross_join(right.build()?)?.build()
+}
- ensure_schema_compatability(plan.schema(), schema.clone())?;
+#[allow(deprecated)]
+pub async fn from_read_rel(
+ consumer: &impl SubstraitConsumer,
+ read: &ReadRel,
+) -> Result<LogicalPlan> {
+ async fn read_with_schema(
+ consumer: &impl SubstraitConsumer,
+ table_ref: TableReference,
+ schema: DFSchema,
+ projection: &Option<MaskExpression>,
+ ) -> Result<LogicalPlan> {
+ let schema = schema.replace_qualifier(table_ref.clone());
+
+ let plan = {
+ let provider = match consumer.resolve_table_ref(&table_ref).await? {
+ Some(ref provider) => Arc::clone(provider),
+ _ => return plan_err!("No table named '{table_ref}'"),
+ };
- let schema = apply_masking(schema, projection)?;
+ LogicalPlanBuilder::scan(
+ table_ref,
+ provider_as_source(Arc::clone(&provider)),
+ None,
+ )?
+ .build()?
+ };
- apply_projection(plan, schema)
- }
+ ensure_schema_compatability(plan.schema(), schema.clone())?;
- let named_struct = read.base_schema.as_ref().ok_or_else(|| {
- substrait_datafusion_err!("No base schema provided for Read Relation")
- })?;
+ let schema = apply_masking(schema, projection)?;
- let substrait_schema = from_substrait_named_struct(named_struct, extensions)?;
+ apply_projection(plan, schema)
+ }
- match &read.as_ref().read_type {
- Some(ReadType::NamedTable(nt)) => {
- let table_reference = match nt.names.len() {
- 0 => {
- return plan_err!("No table name found in NamedTable");
- }
- 1 => TableReference::Bare {
- table: nt.names[0].clone().into(),
- },
- 2 => TableReference::Partial {
- schema: nt.names[0].clone().into(),
- table: nt.names[1].clone().into(),
- },
- _ => TableReference::Full {
- catalog: nt.names[0].clone().into(),
- schema: nt.names[1].clone().into(),
- table: nt.names[2].clone().into(),
- },
- };
+ let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+ substrait_datafusion_err!("No base schema provided for Read Relation")
+ })?;
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
+ let substrait_schema = from_substrait_named_struct(consumer, named_struct)?;
+
+ match &read.read_type {
+ Some(ReadType::NamedTable(nt)) => {
+ let table_reference = match nt.names.len() {
+ 0 => {
+ return plan_err!("No table name found in NamedTable");
}
- Some(ReadType::VirtualTable(vt)) => {
- if vt.values.is_empty() {
- return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: DFSchemaRef::new(substrait_schema),
- }));
+ 1 => TableReference::Bare {
+ table: nt.names[0].clone().into(),
+ },
+ 2 => TableReference::Partial {
+ schema: nt.names[0].clone().into(),
+ table: nt.names[1].clone().into(),
+ },
+ _ => TableReference::Full {
+ catalog: nt.names[0].clone().into(),
+ schema: nt.names[1].clone().into(),
+ table: nt.names[2].clone().into(),
+ },
+ };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ Some(ReadType::VirtualTable(vt)) => {
+ if vt.values.is_empty() {
+ return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row: false,
+ schema: DFSchemaRef::new(substrait_schema),
+ }));
+ }
+
+ let values = vt
+ .values
+ .iter()
+ .map(|row| {
+ let mut name_idx = 0;
+ let lits = row
+ .fields
+ .iter()
+ .map(|lit| {
+ name_idx += 1; // top-level names are provided through schema
+ Ok(Expr::Literal(from_substrait_literal(
+ consumer,
+ lit,
+ &named_struct.names,
+ &mut name_idx,
+ )?))
+ })
+ .collect::<Result<_>>()?;
+ if name_idx != named_struct.names.len() {
+ return substrait_err!(
+ "Names list must match exactly to nested schema, but found {} uses for {} names",
+ name_idx,
+ named_struct.names.len()
+ );
}
+ Ok(lits)
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(LogicalPlan::Values(Values {
+ schema: DFSchemaRef::new(substrait_schema),
+ values,
+ }))
+ }
+ Some(ReadType::LocalFiles(lf)) => {
+ fn extract_filename(name: &str) -> Option<String> {
+ let corrected_url =
+ if name.starts_with("file://") && !name.starts_with("file:///") {
+ name.replacen("file://", "file:///", 1)
+ } else {
+ name.to_string()
+ };
+
+ Url::parse(&corrected_url).ok().and_then(|url| {
+ let path = url.path();
+ std::path::Path::new(path)
+ .file_name()
+ .map(|filename| filename.to_string_lossy().to_string())
+ })
+ }
+
+ // we could use the file name to check the original table provider
+ // TODO: currently does not support multiple local files
+ let filename: Option<String> =
+ lf.items.first().and_then(|x| match x.path_type.as_ref() {
+ Some(UriFile(name)) => extract_filename(name),
+ _ => None,
+ });
+
+ if lf.items.len() > 1 || filename.is_none() {
+ return not_impl_err!("Only single file reads are supported");
+ }
+ let name = filename.unwrap();
+ // directly use unwrap here since we could determine it is a valid one
+ let table_reference = TableReference::Bare { table: name.into() };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ _ => {
+ not_impl_err!("Unsupported ReadType: {:?}", read.read_type)
+ }
+ }
+}
+
+pub async fn from_set_rel(
+ consumer: &impl SubstraitConsumer,
+ set: &SetRel,
+) -> Result<LogicalPlan> {
+ if set.inputs.len() < 2 {
+ substrait_err!("Set operation requires at least two inputs")
+ } else {
+ match set.op() {
+ SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
+ SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
+ SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect(
+ from_substrait_rel(consumer, &set.inputs[0]).await?,
+ union_rels(consumer, &set.inputs[1..], true).await?,
+ false,
+ ),
+ SetOp::IntersectionMultiset => {
+ intersect_rels(consumer, &set.inputs, false).await
+ }
+ SetOp::IntersectionMultisetAll => {
+ intersect_rels(consumer, &set.inputs, true).await
+ }
+ SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
+ SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+ set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+ }
+ }
+}
- let values = vt
- .values
- .iter()
- .map(|row| {
- let mut name_idx = 0;
- let lits = row
- .fields
- .iter()
- .map(|lit| {
- name_idx += 1; // top-level names are provided through schema
- Ok(Expr::Literal(from_substrait_literal(
- lit,
- extensions,
- &named_struct.names,
- &mut name_idx,
- )?))
- })
- .collect::<Result<_>>()?;
- if name_idx != named_struct.names.len() {
- return substrait_err!(
- "Names list must match exactly to nested schema, but found {} uses for {} names",
- name_idx,
- named_struct.names.len()
- );
- }
- Ok(lits)
- })
- .collect::<Result<_>>()?;
+pub async fn from_extension_leaf_rel(
Review Comment:
Hm, these aren't used anywhere, and aren't they kinda weirdly opposite of what the other functions do (ie these call consumer.consume.., instead of consumer.consume.. calling these), maybe let's not add them?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1890180571
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +102,467 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::SessionState;
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// use datafusion_substrait::logical_plan::state::SubstraitPlanningState;
Review Comment:
do you need this? SubstraitPlanningState doesn't seem to be used anywhere in the code anymore afaict. Actually, if it's not used anymore after this, maybe remove the whole state.rs file as part of this PR?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#issuecomment-2551239991
> I think it makes sense to do that, and I'm happy to own it, but I would prefer to do that as a followup. I structured the `consumer.rs` changes in order to make it easy to see that the functions I added were extracted from the big `from_substrait_rel` and `from_substrait_rex` switch statements.
Yup, I totally agree!
I looked through the latest changes, looks even better now, just some more comments re SubstraitPlanningState but otherwise this looks very good by me!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1888615020
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +971,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
+ consumer: &impl SubstraitConsumer,
+ p: &ProjectRel,
) -> Result<LogicalPlan> {
- let plan: Result<LogicalPlan> = match &rel.rel_type {
- Some(RelType::Project(p)) => {
- if let Some(input) = p.input.as_ref() {
- let mut input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let original_schema = input.schema().clone();
-
- // Ensure that all expressions have a unique display name, so that
- // validate_unique_names does not fail when constructing the project.
- let mut name_tracker = NameTracker::new();
-
- // By default, a Substrait Project emits all inputs fields followed by all expressions.
- // We build the explicit expressions first, and then the input expressions to avoid
- // adding aliases to the explicit expressions (as part of ensuring unique names).
- //
- // This is helpful for plan visualization and tests, because when DataFusion produces
- // Substrait Projects it adds an output mapping that excludes all input columns
- // leaving only explicit expressions.
-
- let mut explicit_exprs: Vec<Expr> = vec![];
- for expr in &p.expressions {
- let e = from_substrait_rex(
- state,
- expr,
- input.clone().schema(),
- extensions,
- )
- .await?;
- // if the expression is WindowFunction, wrap in a Window relation
- if let Expr::WindowFunction(_) = &e {
- // Adding the same expression here and in the project below
- // works because the project's builder uses columnize_expr(..)
- // to transform it into a column reference
- input = input.window(vec![e.clone()])?
- }
- explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
-
- let mut final_exprs: Vec<Expr> = vec![];
- for index in 0..original_schema.fields().len() {
- let e = Expr::Column(Column::from(
- original_schema.qualified_field(index),
- ));
- final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
- final_exprs.append(&mut explicit_exprs);
-
- input.project(final_exprs)?.build()
- } else {
- not_impl_err!("Projection without an input is not supported")
+ if let Some(input) = p.input.as_ref() {
+ let mut input =
+ LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let original_schema = input.schema().clone();
+
+ // Ensure that all expressions have a unique display name, so that
+ // validate_unique_names does not fail when constructing the project.
+ let mut name_tracker = NameTracker::new();
+
+ // By default, a Substrait Project emits all inputs fields followed by all expressions.
+ // We build the explicit expressions first, and then the input expressions to avoid
+ // adding aliases to the explicit expressions (as part of ensuring unique names).
+ //
+ // This is helpful for plan visualization and tests, because when DataFusion produces
+ // Substrait Projects it adds an output mapping that excludes all input columns
+ // leaving only explicit expressions.
+
+ let mut explicit_exprs: Vec<Expr> = vec![];
+ for expr in &p.expressions {
+ let e = from_substrait_rex(consumer, expr, input.clone().schema()).await?;
+ // if the expression is WindowFunction, wrap in a Window relation
+ if let Expr::WindowFunction(_) = &e {
+ // Adding the same expression here and in the project below
+ // works because the project's builder uses columnize_expr(..)
+ // to transform it into a column reference
+ input = input.window(vec![e.clone()])?
}
+ explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Filter(filter)) => {
- if let Some(input) = filter.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- if let Some(condition) = filter.condition.as_ref() {
- let expr =
- from_substrait_rex(state, condition, input.schema(), extensions)
- .await?;
- input.filter(expr)?.build()
- } else {
- not_impl_err!("Filter without an condition is not valid")
- }
- } else {
- not_impl_err!("Filter without an input is not valid")
- }
+
+ let mut final_exprs: Vec<Expr> = vec![];
+ for index in 0..original_schema.fields().len() {
+ let e = Expr::Column(Column::from(original_schema.qualified_field(index)));
+ final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Fetch(fetch)) => {
- if let Some(input) = fetch.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let offset = fetch.offset as usize;
- // -1 means that ALL records should be returned
- let count = if fetch.count == -1 {
- None
- } else {
- Some(fetch.count as usize)
- };
- input.limit(offset, count)?.build()
- } else {
- not_impl_err!("Fetch without an input is not valid")
- }
+ final_exprs.append(&mut explicit_exprs);
+ input.project(final_exprs)?.build()
+ } else {
+ not_impl_err!("Projection without an input is not supported")
+ }
+}
+
+#[async_recursion]
+pub async fn from_filter_rel(
+ consumer: &impl SubstraitConsumer,
+ filter: &FilterRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = filter.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ if let Some(condition) = filter.condition.as_ref() {
+ let expr = from_substrait_rex(consumer, condition, input.schema()).await?;
+ input.filter(expr)?.build()
+ } else {
+ not_impl_err!("Filter without an condition is not valid")
}
- Some(RelType::Sort(sort)) => {
- if let Some(input) = sort.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let sorts =
- from_substrait_sorts(state, &sort.sorts, input.schema(), extensions)
- .await?;
- input.sort(sorts)?.build()
- } else {
- not_impl_err!("Sort without an input is not valid")
- }
+ } else {
+ not_impl_err!("Filter without an input is not valid")
+ }
+}
+
+#[async_recursion]
+pub async fn from_fetch_rel(
+ consumer: &impl SubstraitConsumer,
+ fetch: &FetchRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = fetch.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let offset = fetch.offset as usize;
+ // -1 means that ALL records should be returned
+ let count = if fetch.count == -1 {
+ None
+ } else {
+ Some(fetch.count as usize)
+ };
+ input.limit(offset, count)?.build()
+ } else {
+ not_impl_err!("Fetch without an input is not valid")
+ }
+}
+
+pub async fn from_sort_rel(
+ consumer: &impl SubstraitConsumer,
+ sort: &SortRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = sort.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?;
+ input.sort(sorts)?.build()
+ } else {
+ not_impl_err!("Sort without an input is not valid")
+ }
+}
+
+pub async fn from_aggregate_rel(
+ consumer: &impl SubstraitConsumer,
+ agg: &AggregateRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = agg.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let mut ref_group_exprs = vec![];
+
+ for e in &agg.grouping_expressions {
+ let x = from_substrait_rex(consumer, e, input.schema()).await?;
+ ref_group_exprs.push(x);
}
- Some(RelType::Aggregate(agg)) => {
- if let Some(input) = agg.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let mut ref_group_exprs = vec![];
- for e in &agg.grouping_expressions {
- let x =
- from_substrait_rex(state, e, input.schema(), extensions).await?;
- ref_group_exprs.push(x);
+ let mut group_exprs = vec![];
+ let mut aggr_exprs = vec![];
+
+ match agg.groupings.len() {
+ 1 => {
+ group_exprs.extend_from_slice(
+ &from_substrait_grouping(
+ consumer,
+ &agg.groupings[0],
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?,
+ );
+ }
+ _ => {
+ let mut grouping_sets = vec![];
+ for grouping in &agg.groupings {
+ let grouping_set = from_substrait_grouping(
+ consumer,
+ grouping,
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?;
+ grouping_sets.push(grouping_set);
}
+ // Single-element grouping expression of type Expr::GroupingSet.
+ // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+ // parsed by the producer and consumer, since Substrait does not have a type dedicated
+ // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+ group_exprs
+ .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+ }
+ };
- let mut group_exprs = vec![];
- let mut aggr_exprs = vec![];
-
- match agg.groupings.len() {
- 1 => {
- group_exprs.extend_from_slice(
- &from_substrait_grouping(
- state,
- &agg.groupings[0],
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?,
- );
- }
- _ => {
- let mut grouping_sets = vec![];
- for grouping in &agg.groupings {
- let grouping_set = from_substrait_grouping(
- state,
- grouping,
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?;
- grouping_sets.push(grouping_set);
+ for m in &agg.measures {
+ let filter = match &m.filter {
+ Some(fil) => Some(Box::new(
+ from_substrait_rex(consumer, fil, input.schema()).await?,
+ )),
+ None => None,
+ };
+ let agg_func = match &m.measure {
+ Some(f) => {
+ let distinct = match f.invocation {
+ _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+ true
}
- // Single-element grouping expression of type Expr::GroupingSet.
- // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
- // parsed by the producer and consumer, since Substrait does not have a type dedicated
- // to ROLLUP. Only vector of Groupings (grouping sets) is available.
- group_exprs.push(Expr::GroupingSet(GroupingSet::GroupingSets(
- grouping_sets,
- )));
- }
- };
-
- for m in &agg.measures {
- let filter = match &m.filter {
- Some(fil) => Some(Box::new(
- from_substrait_rex(state, fil, input.schema(), extensions)
- .await?,
- )),
- None => None,
+ _ if f.invocation == AggregationInvocation::All as i32 => false,
+ _ => false,
};
- let agg_func = match &m.measure {
- Some(f) => {
- let distinct = match f.invocation {
- _ if f.invocation
- == AggregationInvocation::Distinct as i32 =>
- {
- true
- }
- _ if f.invocation
- == AggregationInvocation::All as i32 =>
- {
- false
- }
- _ => false,
- };
- let order_by = if !f.sorts.is_empty() {
- Some(
- from_substrait_sorts(
- state,
- &f.sorts,
- input.schema(),
- extensions,
- )
- .await?,
- )
- } else {
- None
- };
-
- from_substrait_agg_func(
- state,
- f,
- input.schema(),
- extensions,
- filter,
- order_by,
- distinct,
- )
- .await
- }
- None => not_impl_err!(
- "Aggregate without aggregate function is not supported"
- ),
+ let order_by = if !f.sorts.is_empty() {
+ Some(
+ from_substrait_sorts(consumer, &f.sorts, input.schema())
+ .await?,
+ )
+ } else {
+ None
};
- aggr_exprs.push(agg_func?.as_ref().clone());
- }
- input.aggregate(group_exprs, aggr_exprs)?.build()
- } else {
- not_impl_err!("Aggregate without an input is not valid")
- }
- }
- Some(RelType::Join(join)) => {
- if join.post_join_filter.is_some() {
- return not_impl_err!(
- "JoinRel with post_join_filter is not yet supported"
- );
- }
- let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
-
- let join_type = from_substrait_jointype(join.r#type)?;
- // The join condition expression needs full input schema and not the output schema from join since we lose columns from
- // certain join types such as semi and anti joins
- let in_join_schema = left.schema().join(right.schema())?;
-
- // If join expression exists, parse the `on` condition expression, build join and return
- // Otherwise, build join with only the filter, without join keys
- match &join.expression.as_ref() {
- Some(expr) => {
- let on = from_substrait_rex(state, expr, &in_join_schema, extensions)
- .await?;
- // The join expression can contain both equal and non-equal ops.
- // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
- // So we extract each part as follows:
- // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
- // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
- let (join_ons, nulls_equal_nulls, join_filter) =
- split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
- let (left_cols, right_cols): (Vec<_>, Vec<_>) =
- itertools::multiunzip(join_ons);
- left.join_detailed(
- right.build()?,
- join_type,
- (left_cols, right_cols),
- join_filter,
- nulls_equal_nulls,
- )?
- .build()
+ from_substrait_agg_func(
+ consumer,
+ f,
+ input.schema(),
+ filter,
+ order_by,
+ distinct,
+ )
+ .await
}
None => {
- let on: Vec<String> = vec![];
- left.join_detailed(
- right.build()?,
- join_type,
- (on.clone(), on),
- None,
- false,
- )?
- .build()
+ not_impl_err!("Aggregate without aggregate function is not supported")
}
- }
+ };
+ aggr_exprs.push(agg_func?.as_ref().clone());
}
- Some(RelType::Cross(cross)) => {
- let left = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
- left.cross_join(right.build()?)?.build()
+ input.aggregate(group_exprs, aggr_exprs)?.build()
+ } else {
+ not_impl_err!("Aggregate without an input is not valid")
+ }
+}
+
+pub async fn from_join_rel(
+ consumer: &impl SubstraitConsumer,
+ join: &JoinRel,
+) -> Result<LogicalPlan> {
+ if join.post_join_filter.is_some() {
+ return not_impl_err!("JoinRel with post_join_filter is not yet supported");
+ }
+
+ let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+
+ let join_type = from_substrait_jointype(join.r#type)?;
+ // The join condition expression needs full input schema and not the output schema from join since we lose columns from
+ // certain join types such as semi and anti joins
+ let in_join_schema = left.schema().join(right.schema())?;
+
+ // If join expression exists, parse the `on` condition expression, build join and return
+ // Otherwise, build join with only the filter, without join keys
+ match &join.expression.as_ref() {
+ Some(expr) => {
+ let on = from_substrait_rex(consumer, expr, &in_join_schema).await?;
+ // The join expression can contain both equal and non-equal ops.
+ // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
+ // So we extract each part as follows:
+ // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
+ // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
+ let (join_ons, nulls_equal_nulls, join_filter) =
+ split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) =
+ itertools::multiunzip(join_ons);
+ left.join_detailed(
+ right.build()?,
+ join_type,
+ (left_cols, right_cols),
+ join_filter,
+ nulls_equal_nulls,
+ )?
+ .build()
}
- Some(RelType::Read(read)) => {
- async fn read_with_schema(
- state: &dyn SubstraitPlanningState,
- table_ref: TableReference,
- schema: DFSchema,
- projection: &Option<MaskExpression>,
- ) -> Result<LogicalPlan> {
- let schema = schema.replace_qualifier(table_ref.clone());
-
- let plan = {
- let provider = match state.table(&table_ref).await? {
- Some(ref provider) => Arc::clone(provider),
- _ => return plan_err!("No table named '{table_ref}'"),
- };
+ None => {
+ let on: Vec<String> = vec![];
+ left.join_detailed(right.build()?, join_type, (on.clone(), on), None, false)?
+ .build()
+ }
+ }
+}
- LogicalPlanBuilder::scan(
- table_ref,
- provider_as_source(Arc::clone(&provider)),
- None,
- )?
- .build()?
- };
+pub async fn from_cross_rel(
+ consumer: &impl SubstraitConsumer,
+ cross: &CrossRel,
+) -> Result<LogicalPlan> {
+ let left = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+ left.cross_join(right.build()?)?.build()
+}
- ensure_schema_compatability(plan.schema(), schema.clone())?;
+#[allow(deprecated)]
+pub async fn from_read_rel(
+ consumer: &impl SubstraitConsumer,
+ read: &ReadRel,
+) -> Result<LogicalPlan> {
+ async fn read_with_schema(
+ consumer: &impl SubstraitConsumer,
+ table_ref: TableReference,
+ schema: DFSchema,
+ projection: &Option<MaskExpression>,
+ ) -> Result<LogicalPlan> {
+ let schema = schema.replace_qualifier(table_ref.clone());
+
+ let plan = {
+ let provider = match consumer.resolve_table_ref(&table_ref).await? {
+ Some(ref provider) => Arc::clone(provider),
+ _ => return plan_err!("No table named '{table_ref}'"),
+ };
- let schema = apply_masking(schema, projection)?;
+ LogicalPlanBuilder::scan(
+ table_ref,
+ provider_as_source(Arc::clone(&provider)),
+ None,
+ )?
+ .build()?
+ };
- apply_projection(plan, schema)
- }
+ ensure_schema_compatability(plan.schema(), schema.clone())?;
- let named_struct = read.base_schema.as_ref().ok_or_else(|| {
- substrait_datafusion_err!("No base schema provided for Read Relation")
- })?;
+ let schema = apply_masking(schema, projection)?;
- let substrait_schema = from_substrait_named_struct(named_struct, extensions)?;
+ apply_projection(plan, schema)
+ }
- match &read.as_ref().read_type {
- Some(ReadType::NamedTable(nt)) => {
- let table_reference = match nt.names.len() {
- 0 => {
- return plan_err!("No table name found in NamedTable");
- }
- 1 => TableReference::Bare {
- table: nt.names[0].clone().into(),
- },
- 2 => TableReference::Partial {
- schema: nt.names[0].clone().into(),
- table: nt.names[1].clone().into(),
- },
- _ => TableReference::Full {
- catalog: nt.names[0].clone().into(),
- schema: nt.names[1].clone().into(),
- table: nt.names[2].clone().into(),
- },
- };
+ let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+ substrait_datafusion_err!("No base schema provided for Read Relation")
+ })?;
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
+ let substrait_schema = from_substrait_named_struct(consumer, named_struct)?;
+
+ match &read.read_type {
+ Some(ReadType::NamedTable(nt)) => {
+ let table_reference = match nt.names.len() {
+ 0 => {
+ return plan_err!("No table name found in NamedTable");
}
- Some(ReadType::VirtualTable(vt)) => {
- if vt.values.is_empty() {
- return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: DFSchemaRef::new(substrait_schema),
- }));
+ 1 => TableReference::Bare {
+ table: nt.names[0].clone().into(),
+ },
+ 2 => TableReference::Partial {
+ schema: nt.names[0].clone().into(),
+ table: nt.names[1].clone().into(),
+ },
+ _ => TableReference::Full {
+ catalog: nt.names[0].clone().into(),
+ schema: nt.names[1].clone().into(),
+ table: nt.names[2].clone().into(),
+ },
+ };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ Some(ReadType::VirtualTable(vt)) => {
+ if vt.values.is_empty() {
+ return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row: false,
+ schema: DFSchemaRef::new(substrait_schema),
+ }));
+ }
+
+ let values = vt
+ .values
+ .iter()
+ .map(|row| {
+ let mut name_idx = 0;
+ let lits = row
+ .fields
+ .iter()
+ .map(|lit| {
+ name_idx += 1; // top-level names are provided through schema
+ Ok(Expr::Literal(from_substrait_literal(
+ consumer,
+ lit,
+ &named_struct.names,
+ &mut name_idx,
+ )?))
+ })
+ .collect::<Result<_>>()?;
+ if name_idx != named_struct.names.len() {
+ return substrait_err!(
+ "Names list must match exactly to nested schema, but found {} uses for {} names",
+ name_idx,
+ named_struct.names.len()
+ );
}
+ Ok(lits)
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(LogicalPlan::Values(Values {
+ schema: DFSchemaRef::new(substrait_schema),
+ values,
+ }))
+ }
+ Some(ReadType::LocalFiles(lf)) => {
+ fn extract_filename(name: &str) -> Option<String> {
+ let corrected_url =
+ if name.starts_with("file://") && !name.starts_with("file:///") {
+ name.replacen("file://", "file:///", 1)
+ } else {
+ name.to_string()
+ };
+
+ Url::parse(&corrected_url).ok().and_then(|url| {
+ let path = url.path();
+ std::path::Path::new(path)
+ .file_name()
+ .map(|filename| filename.to_string_lossy().to_string())
+ })
+ }
+
+ // we could use the file name to check the original table provider
+ // TODO: currently does not support multiple local files
+ let filename: Option<String> =
+ lf.items.first().and_then(|x| match x.path_type.as_ref() {
+ Some(UriFile(name)) => extract_filename(name),
+ _ => None,
+ });
+
+ if lf.items.len() > 1 || filename.is_none() {
+ return not_impl_err!("Only single file reads are supported");
+ }
+ let name = filename.unwrap();
+ // directly use unwrap here since we could determine it is a valid one
+ let table_reference = TableReference::Bare { table: name.into() };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ _ => {
+ not_impl_err!("Unsupported ReadType: {:?}", read.read_type)
+ }
+ }
+}
+
+pub async fn from_set_rel(
+ consumer: &impl SubstraitConsumer,
+ set: &SetRel,
+) -> Result<LogicalPlan> {
+ if set.inputs.len() < 2 {
+ substrait_err!("Set operation requires at least two inputs")
+ } else {
+ match set.op() {
+ SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
+ SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
+ SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect(
+ from_substrait_rel(consumer, &set.inputs[0]).await?,
+ union_rels(consumer, &set.inputs[1..], true).await?,
+ false,
+ ),
+ SetOp::IntersectionMultiset => {
+ intersect_rels(consumer, &set.inputs, false).await
+ }
+ SetOp::IntersectionMultisetAll => {
+ intersect_rels(consumer, &set.inputs, true).await
+ }
+ SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
+ SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+ set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+ }
+ }
+}
- let values = vt
- .values
- .iter()
- .map(|row| {
- let mut name_idx = 0;
- let lits = row
- .fields
- .iter()
- .map(|lit| {
- name_idx += 1; // top-level names are provided through schema
- Ok(Expr::Literal(from_substrait_literal(
- lit,
- extensions,
- &named_struct.names,
- &mut name_idx,
- )?))
- })
- .collect::<Result<_>>()?;
- if name_idx != named_struct.names.len() {
- return substrait_err!(
- "Names list must match exactly to nested schema, but found {} uses for {} names",
- name_idx,
- named_struct.names.len()
- );
- }
- Ok(lits)
- })
- .collect::<Result<_>>()?;
+pub async fn from_extension_leaf_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_leaf_rel: &ExtensionLeafRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_leaf(extension_leaf_rel).await
+}
- Ok(LogicalPlan::Values(Values {
- schema: DFSchemaRef::new(substrait_schema),
- values,
- }))
- }
- Some(ReadType::LocalFiles(lf)) => {
- fn extract_filename(name: &str) -> Option<String> {
- let corrected_url = if name.starts_with("file://")
- && !name.starts_with("file:///")
- {
- name.replacen("file://", "file:///", 1)
- } else {
- name.to_string()
- };
+pub async fn from_extension_single_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_single_rel: &ExtensionSingleRel,
+) -> Result<LogicalPlan> {
+ consumer
+ .consume_extension_single(extension_single_rel)
+ .await
+}
- Url::parse(&corrected_url).ok().and_then(|url| {
- let path = url.path();
- std::path::Path::new(path)
- .file_name()
- .map(|filename| filename.to_string_lossy().to_string())
- })
- }
+pub async fn from_extension_multi_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_multi_rel: &ExtensionMultiRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_multi(extension_multi_rel).await
+}
- // we could use the file name to check the original table provider
- // TODO: currently does not support multiple local files
- let filename: Option<String> =
- lf.items.first().and_then(|x| match x.path_type.as_ref() {
- Some(UriFile(name)) => extract_filename(name),
- _ => None,
- });
+pub async fn from_exchange_rel(
+ consumer: &impl SubstraitConsumer,
+ exchange: &ExchangeRel,
+) -> Result<LogicalPlan> {
+ let Some(input) = exchange.input.as_ref() else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+ let input = Arc::new(from_substrait_rel(consumer, input).await?);
- if lf.items.len() > 1 || filename.is_none() {
- return not_impl_err!("Only single file reads are supported");
- }
- let name = filename.unwrap();
- // directly use unwrap here since we could determine it is a valid one
- let table_reference = TableReference::Bare { table: name.into() };
-
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
- }
- _ => {
- not_impl_err!("Unsupported ReadType: {:?}", &read.as_ref().read_type)
- }
- }
- }
- Some(RelType::Set(set)) => match set_rel::SetOp::try_from(set.op) {
- Ok(set_op) => {
- if set.inputs.len() < 2 {
- substrait_err!("Set operation requires at least two inputs")
- } else {
- match set_op {
- set_rel::SetOp::UnionAll => {
- union_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::UnionDistinct => {
- union_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionPrimary => {
- LogicalPlanBuilder::intersect(
- from_substrait_rel(state, &set.inputs[0], extensions)
- .await?,
- union_rels(&set.inputs[1..], state, extensions, true)
- .await?,
- false,
- )
- }
- set_rel::SetOp::IntersectionMultiset => {
- intersect_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionMultisetAll => {
- intersect_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::MinusPrimary => {
- except_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::MinusPrimaryAll => {
- except_rels(&set.inputs, state, extensions, true).await
- }
- _ => not_impl_err!("Unsupported set operator: {set_op:?}"),
- }
- }
+ let Some(exchange_kind) = &exchange.exchange_kind else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+
+ // ref: https://substrait.io/relations/physical_relations/#exchange-types
+ let partitioning_scheme = match exchange_kind {
+ ExchangeKind::ScatterByFields(scatter_fields) => {
+ let mut partition_columns = vec![];
+ let input_schema = input.schema();
+ for field_ref in &scatter_fields.fields {
+ let column = from_substrait_field_reference(field_ref, input_schema)?;
+ partition_columns.push(column);
}
- Err(e) => not_impl_err!("Invalid set operation type {}: {e}", set.op),
- },
- Some(RelType::ExtensionLeaf(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
- }
- Some(RelType::ExtensionSingle(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- let Some(input_rel) = &extension.input else {
- return substrait_err!(
- "ExtensionSingleRel doesn't contains input rel. Try use ExtensionLeafRel instead"
- );
- };
- let input_plan = from_substrait_rel(state, input_rel, extensions).await?;
- let plan =
- plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
+ Partitioning::Hash(partition_columns, exchange.partition_count as usize)
}
- Some(RelType::ExtensionMulti(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- let mut inputs = Vec::with_capacity(extension.inputs.len());
- for input in &extension.inputs {
- let input_plan = from_substrait_rel(state, input, extensions).await?;
- inputs.push(input_plan);
- }
- let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
+ ExchangeKind::RoundRobin(_) => {
+ Partitioning::RoundRobinBatch(exchange.partition_count as usize)
}
Review Comment:
Would it make sense to keep these in the default implementation though? Just as in the producer we do produce these types (https://github.com/apache/datafusion/blob/b8e435f3e38f1f8d2d8f24c3da3c2b21e76b700b/datafusion/substrait/src/logical_plan/producer.rs#L596), and I guess it'd be nice if the roundtrip tests that don't require special setup on producer wouldn't require special setup for consumer either.
That's not a strictly held opinion.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +971,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
+ consumer: &impl SubstraitConsumer,
+ p: &ProjectRel,
) -> Result<LogicalPlan> {
- let plan: Result<LogicalPlan> = match &rel.rel_type {
- Some(RelType::Project(p)) => {
- if let Some(input) = p.input.as_ref() {
- let mut input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let original_schema = input.schema().clone();
-
- // Ensure that all expressions have a unique display name, so that
- // validate_unique_names does not fail when constructing the project.
- let mut name_tracker = NameTracker::new();
-
- // By default, a Substrait Project emits all inputs fields followed by all expressions.
- // We build the explicit expressions first, and then the input expressions to avoid
- // adding aliases to the explicit expressions (as part of ensuring unique names).
- //
- // This is helpful for plan visualization and tests, because when DataFusion produces
- // Substrait Projects it adds an output mapping that excludes all input columns
- // leaving only explicit expressions.
-
- let mut explicit_exprs: Vec<Expr> = vec![];
- for expr in &p.expressions {
- let e = from_substrait_rex(
- state,
- expr,
- input.clone().schema(),
- extensions,
- )
- .await?;
- // if the expression is WindowFunction, wrap in a Window relation
- if let Expr::WindowFunction(_) = &e {
- // Adding the same expression here and in the project below
- // works because the project's builder uses columnize_expr(..)
- // to transform it into a column reference
- input = input.window(vec![e.clone()])?
- }
- explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
-
- let mut final_exprs: Vec<Expr> = vec![];
- for index in 0..original_schema.fields().len() {
- let e = Expr::Column(Column::from(
- original_schema.qualified_field(index),
- ));
- final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
- final_exprs.append(&mut explicit_exprs);
-
- input.project(final_exprs)?.build()
- } else {
- not_impl_err!("Projection without an input is not supported")
+ if let Some(input) = p.input.as_ref() {
+ let mut input =
+ LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let original_schema = input.schema().clone();
+
+ // Ensure that all expressions have a unique display name, so that
+ // validate_unique_names does not fail when constructing the project.
+ let mut name_tracker = NameTracker::new();
+
+ // By default, a Substrait Project emits all inputs fields followed by all expressions.
+ // We build the explicit expressions first, and then the input expressions to avoid
+ // adding aliases to the explicit expressions (as part of ensuring unique names).
+ //
+ // This is helpful for plan visualization and tests, because when DataFusion produces
+ // Substrait Projects it adds an output mapping that excludes all input columns
+ // leaving only explicit expressions.
+
+ let mut explicit_exprs: Vec<Expr> = vec![];
+ for expr in &p.expressions {
+ let e = from_substrait_rex(consumer, expr, input.clone().schema()).await?;
+ // if the expression is WindowFunction, wrap in a Window relation
+ if let Expr::WindowFunction(_) = &e {
+ // Adding the same expression here and in the project below
+ // works because the project's builder uses columnize_expr(..)
+ // to transform it into a column reference
+ input = input.window(vec![e.clone()])?
}
+ explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Filter(filter)) => {
- if let Some(input) = filter.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- if let Some(condition) = filter.condition.as_ref() {
- let expr =
- from_substrait_rex(state, condition, input.schema(), extensions)
- .await?;
- input.filter(expr)?.build()
- } else {
- not_impl_err!("Filter without an condition is not valid")
- }
- } else {
- not_impl_err!("Filter without an input is not valid")
- }
+
+ let mut final_exprs: Vec<Expr> = vec![];
+ for index in 0..original_schema.fields().len() {
+ let e = Expr::Column(Column::from(original_schema.qualified_field(index)));
+ final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Fetch(fetch)) => {
- if let Some(input) = fetch.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let offset = fetch.offset as usize;
- // -1 means that ALL records should be returned
- let count = if fetch.count == -1 {
- None
- } else {
- Some(fetch.count as usize)
- };
- input.limit(offset, count)?.build()
- } else {
- not_impl_err!("Fetch without an input is not valid")
- }
+ final_exprs.append(&mut explicit_exprs);
+ input.project(final_exprs)?.build()
+ } else {
+ not_impl_err!("Projection without an input is not supported")
+ }
+}
+
+#[async_recursion]
+pub async fn from_filter_rel(
+ consumer: &impl SubstraitConsumer,
+ filter: &FilterRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = filter.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ if let Some(condition) = filter.condition.as_ref() {
+ let expr = from_substrait_rex(consumer, condition, input.schema()).await?;
+ input.filter(expr)?.build()
+ } else {
+ not_impl_err!("Filter without an condition is not valid")
}
- Some(RelType::Sort(sort)) => {
- if let Some(input) = sort.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let sorts =
- from_substrait_sorts(state, &sort.sorts, input.schema(), extensions)
- .await?;
- input.sort(sorts)?.build()
- } else {
- not_impl_err!("Sort without an input is not valid")
- }
+ } else {
+ not_impl_err!("Filter without an input is not valid")
+ }
+}
+
+#[async_recursion]
+pub async fn from_fetch_rel(
+ consumer: &impl SubstraitConsumer,
+ fetch: &FetchRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = fetch.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let offset = fetch.offset as usize;
+ // -1 means that ALL records should be returned
+ let count = if fetch.count == -1 {
+ None
+ } else {
+ Some(fetch.count as usize)
+ };
+ input.limit(offset, count)?.build()
+ } else {
+ not_impl_err!("Fetch without an input is not valid")
+ }
+}
+
+pub async fn from_sort_rel(
+ consumer: &impl SubstraitConsumer,
+ sort: &SortRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = sort.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?;
+ input.sort(sorts)?.build()
+ } else {
+ not_impl_err!("Sort without an input is not valid")
+ }
+}
+
+pub async fn from_aggregate_rel(
+ consumer: &impl SubstraitConsumer,
+ agg: &AggregateRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = agg.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let mut ref_group_exprs = vec![];
+
+ for e in &agg.grouping_expressions {
+ let x = from_substrait_rex(consumer, e, input.schema()).await?;
+ ref_group_exprs.push(x);
}
- Some(RelType::Aggregate(agg)) => {
- if let Some(input) = agg.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let mut ref_group_exprs = vec![];
- for e in &agg.grouping_expressions {
- let x =
- from_substrait_rex(state, e, input.schema(), extensions).await?;
- ref_group_exprs.push(x);
+ let mut group_exprs = vec![];
+ let mut aggr_exprs = vec![];
+
+ match agg.groupings.len() {
+ 1 => {
+ group_exprs.extend_from_slice(
+ &from_substrait_grouping(
+ consumer,
+ &agg.groupings[0],
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?,
+ );
+ }
+ _ => {
+ let mut grouping_sets = vec![];
+ for grouping in &agg.groupings {
+ let grouping_set = from_substrait_grouping(
+ consumer,
+ grouping,
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?;
+ grouping_sets.push(grouping_set);
}
+ // Single-element grouping expression of type Expr::GroupingSet.
+ // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+ // parsed by the producer and consumer, since Substrait does not have a type dedicated
+ // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+ group_exprs
+ .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+ }
+ };
- let mut group_exprs = vec![];
- let mut aggr_exprs = vec![];
-
- match agg.groupings.len() {
- 1 => {
- group_exprs.extend_from_slice(
- &from_substrait_grouping(
- state,
- &agg.groupings[0],
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?,
- );
- }
- _ => {
- let mut grouping_sets = vec![];
- for grouping in &agg.groupings {
- let grouping_set = from_substrait_grouping(
- state,
- grouping,
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?;
- grouping_sets.push(grouping_set);
+ for m in &agg.measures {
+ let filter = match &m.filter {
+ Some(fil) => Some(Box::new(
+ from_substrait_rex(consumer, fil, input.schema()).await?,
+ )),
+ None => None,
+ };
+ let agg_func = match &m.measure {
+ Some(f) => {
+ let distinct = match f.invocation {
+ _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+ true
}
- // Single-element grouping expression of type Expr::GroupingSet.
- // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
- // parsed by the producer and consumer, since Substrait does not have a type dedicated
- // to ROLLUP. Only vector of Groupings (grouping sets) is available.
- group_exprs.push(Expr::GroupingSet(GroupingSet::GroupingSets(
- grouping_sets,
- )));
- }
- };
-
- for m in &agg.measures {
- let filter = match &m.filter {
- Some(fil) => Some(Box::new(
- from_substrait_rex(state, fil, input.schema(), extensions)
- .await?,
- )),
- None => None,
+ _ if f.invocation == AggregationInvocation::All as i32 => false,
+ _ => false,
};
- let agg_func = match &m.measure {
- Some(f) => {
- let distinct = match f.invocation {
- _ if f.invocation
- == AggregationInvocation::Distinct as i32 =>
- {
- true
- }
- _ if f.invocation
- == AggregationInvocation::All as i32 =>
- {
- false
- }
- _ => false,
- };
- let order_by = if !f.sorts.is_empty() {
- Some(
- from_substrait_sorts(
- state,
- &f.sorts,
- input.schema(),
- extensions,
- )
- .await?,
- )
- } else {
- None
- };
-
- from_substrait_agg_func(
- state,
- f,
- input.schema(),
- extensions,
- filter,
- order_by,
- distinct,
- )
- .await
- }
- None => not_impl_err!(
- "Aggregate without aggregate function is not supported"
- ),
+ let order_by = if !f.sorts.is_empty() {
+ Some(
+ from_substrait_sorts(consumer, &f.sorts, input.schema())
+ .await?,
+ )
+ } else {
+ None
};
- aggr_exprs.push(agg_func?.as_ref().clone());
- }
- input.aggregate(group_exprs, aggr_exprs)?.build()
- } else {
- not_impl_err!("Aggregate without an input is not valid")
- }
- }
- Some(RelType::Join(join)) => {
- if join.post_join_filter.is_some() {
- return not_impl_err!(
- "JoinRel with post_join_filter is not yet supported"
- );
- }
- let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
-
- let join_type = from_substrait_jointype(join.r#type)?;
- // The join condition expression needs full input schema and not the output schema from join since we lose columns from
- // certain join types such as semi and anti joins
- let in_join_schema = left.schema().join(right.schema())?;
-
- // If join expression exists, parse the `on` condition expression, build join and return
- // Otherwise, build join with only the filter, without join keys
- match &join.expression.as_ref() {
- Some(expr) => {
- let on = from_substrait_rex(state, expr, &in_join_schema, extensions)
- .await?;
- // The join expression can contain both equal and non-equal ops.
- // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
- // So we extract each part as follows:
- // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
- // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
- let (join_ons, nulls_equal_nulls, join_filter) =
- split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
- let (left_cols, right_cols): (Vec<_>, Vec<_>) =
- itertools::multiunzip(join_ons);
- left.join_detailed(
- right.build()?,
- join_type,
- (left_cols, right_cols),
- join_filter,
- nulls_equal_nulls,
- )?
- .build()
+ from_substrait_agg_func(
+ consumer,
+ f,
+ input.schema(),
+ filter,
+ order_by,
+ distinct,
+ )
+ .await
}
None => {
- let on: Vec<String> = vec![];
- left.join_detailed(
- right.build()?,
- join_type,
- (on.clone(), on),
- None,
- false,
- )?
- .build()
+ not_impl_err!("Aggregate without aggregate function is not supported")
}
- }
+ };
+ aggr_exprs.push(agg_func?.as_ref().clone());
}
- Some(RelType::Cross(cross)) => {
- let left = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
- left.cross_join(right.build()?)?.build()
+ input.aggregate(group_exprs, aggr_exprs)?.build()
+ } else {
+ not_impl_err!("Aggregate without an input is not valid")
+ }
+}
+
+pub async fn from_join_rel(
+ consumer: &impl SubstraitConsumer,
+ join: &JoinRel,
+) -> Result<LogicalPlan> {
+ if join.post_join_filter.is_some() {
+ return not_impl_err!("JoinRel with post_join_filter is not yet supported");
+ }
+
+ let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+
+ let join_type = from_substrait_jointype(join.r#type)?;
+ // The join condition expression needs full input schema and not the output schema from join since we lose columns from
+ // certain join types such as semi and anti joins
+ let in_join_schema = left.schema().join(right.schema())?;
+
+ // If join expression exists, parse the `on` condition expression, build join and return
+ // Otherwise, build join with only the filter, without join keys
+ match &join.expression.as_ref() {
+ Some(expr) => {
+ let on = from_substrait_rex(consumer, expr, &in_join_schema).await?;
+ // The join expression can contain both equal and non-equal ops.
+ // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
+ // So we extract each part as follows:
+ // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
+ // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
+ let (join_ons, nulls_equal_nulls, join_filter) =
+ split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) =
+ itertools::multiunzip(join_ons);
+ left.join_detailed(
+ right.build()?,
+ join_type,
+ (left_cols, right_cols),
+ join_filter,
+ nulls_equal_nulls,
+ )?
+ .build()
}
- Some(RelType::Read(read)) => {
- async fn read_with_schema(
- state: &dyn SubstraitPlanningState,
- table_ref: TableReference,
- schema: DFSchema,
- projection: &Option<MaskExpression>,
- ) -> Result<LogicalPlan> {
- let schema = schema.replace_qualifier(table_ref.clone());
-
- let plan = {
- let provider = match state.table(&table_ref).await? {
- Some(ref provider) => Arc::clone(provider),
- _ => return plan_err!("No table named '{table_ref}'"),
- };
+ None => {
+ let on: Vec<String> = vec![];
+ left.join_detailed(right.build()?, join_type, (on.clone(), on), None, false)?
+ .build()
+ }
+ }
+}
- LogicalPlanBuilder::scan(
- table_ref,
- provider_as_source(Arc::clone(&provider)),
- None,
- )?
- .build()?
- };
+pub async fn from_cross_rel(
+ consumer: &impl SubstraitConsumer,
+ cross: &CrossRel,
+) -> Result<LogicalPlan> {
+ let left = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+ left.cross_join(right.build()?)?.build()
+}
- ensure_schema_compatability(plan.schema(), schema.clone())?;
+#[allow(deprecated)]
+pub async fn from_read_rel(
+ consumer: &impl SubstraitConsumer,
+ read: &ReadRel,
+) -> Result<LogicalPlan> {
+ async fn read_with_schema(
+ consumer: &impl SubstraitConsumer,
+ table_ref: TableReference,
+ schema: DFSchema,
+ projection: &Option<MaskExpression>,
+ ) -> Result<LogicalPlan> {
+ let schema = schema.replace_qualifier(table_ref.clone());
+
+ let plan = {
+ let provider = match consumer.resolve_table_ref(&table_ref).await? {
+ Some(ref provider) => Arc::clone(provider),
+ _ => return plan_err!("No table named '{table_ref}'"),
+ };
- let schema = apply_masking(schema, projection)?;
+ LogicalPlanBuilder::scan(
+ table_ref,
+ provider_as_source(Arc::clone(&provider)),
+ None,
+ )?
+ .build()?
+ };
- apply_projection(plan, schema)
- }
+ ensure_schema_compatability(plan.schema(), schema.clone())?;
- let named_struct = read.base_schema.as_ref().ok_or_else(|| {
- substrait_datafusion_err!("No base schema provided for Read Relation")
- })?;
+ let schema = apply_masking(schema, projection)?;
- let substrait_schema = from_substrait_named_struct(named_struct, extensions)?;
+ apply_projection(plan, schema)
+ }
- match &read.as_ref().read_type {
- Some(ReadType::NamedTable(nt)) => {
- let table_reference = match nt.names.len() {
- 0 => {
- return plan_err!("No table name found in NamedTable");
- }
- 1 => TableReference::Bare {
- table: nt.names[0].clone().into(),
- },
- 2 => TableReference::Partial {
- schema: nt.names[0].clone().into(),
- table: nt.names[1].clone().into(),
- },
- _ => TableReference::Full {
- catalog: nt.names[0].clone().into(),
- schema: nt.names[1].clone().into(),
- table: nt.names[2].clone().into(),
- },
- };
+ let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+ substrait_datafusion_err!("No base schema provided for Read Relation")
+ })?;
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
+ let substrait_schema = from_substrait_named_struct(consumer, named_struct)?;
+
+ match &read.read_type {
+ Some(ReadType::NamedTable(nt)) => {
+ let table_reference = match nt.names.len() {
+ 0 => {
+ return plan_err!("No table name found in NamedTable");
}
- Some(ReadType::VirtualTable(vt)) => {
- if vt.values.is_empty() {
- return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: DFSchemaRef::new(substrait_schema),
- }));
+ 1 => TableReference::Bare {
+ table: nt.names[0].clone().into(),
+ },
+ 2 => TableReference::Partial {
+ schema: nt.names[0].clone().into(),
+ table: nt.names[1].clone().into(),
+ },
+ _ => TableReference::Full {
+ catalog: nt.names[0].clone().into(),
+ schema: nt.names[1].clone().into(),
+ table: nt.names[2].clone().into(),
+ },
+ };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ Some(ReadType::VirtualTable(vt)) => {
+ if vt.values.is_empty() {
+ return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row: false,
+ schema: DFSchemaRef::new(substrait_schema),
+ }));
+ }
+
+ let values = vt
+ .values
+ .iter()
+ .map(|row| {
+ let mut name_idx = 0;
+ let lits = row
+ .fields
+ .iter()
+ .map(|lit| {
+ name_idx += 1; // top-level names are provided through schema
+ Ok(Expr::Literal(from_substrait_literal(
+ consumer,
+ lit,
+ &named_struct.names,
+ &mut name_idx,
+ )?))
+ })
+ .collect::<Result<_>>()?;
+ if name_idx != named_struct.names.len() {
+ return substrait_err!(
+ "Names list must match exactly to nested schema, but found {} uses for {} names",
+ name_idx,
+ named_struct.names.len()
+ );
}
+ Ok(lits)
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(LogicalPlan::Values(Values {
+ schema: DFSchemaRef::new(substrait_schema),
+ values,
+ }))
+ }
+ Some(ReadType::LocalFiles(lf)) => {
+ fn extract_filename(name: &str) -> Option<String> {
+ let corrected_url =
+ if name.starts_with("file://") && !name.starts_with("file:///") {
+ name.replacen("file://", "file:///", 1)
+ } else {
+ name.to_string()
+ };
+
+ Url::parse(&corrected_url).ok().and_then(|url| {
+ let path = url.path();
+ std::path::Path::new(path)
+ .file_name()
+ .map(|filename| filename.to_string_lossy().to_string())
+ })
+ }
+
+ // we could use the file name to check the original table provider
+ // TODO: currently does not support multiple local files
+ let filename: Option<String> =
+ lf.items.first().and_then(|x| match x.path_type.as_ref() {
+ Some(UriFile(name)) => extract_filename(name),
+ _ => None,
+ });
+
+ if lf.items.len() > 1 || filename.is_none() {
+ return not_impl_err!("Only single file reads are supported");
+ }
+ let name = filename.unwrap();
+ // directly use unwrap here since we could determine it is a valid one
+ let table_reference = TableReference::Bare { table: name.into() };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ _ => {
+ not_impl_err!("Unsupported ReadType: {:?}", read.read_type)
+ }
+ }
+}
+
+pub async fn from_set_rel(
+ consumer: &impl SubstraitConsumer,
+ set: &SetRel,
+) -> Result<LogicalPlan> {
+ if set.inputs.len() < 2 {
+ substrait_err!("Set operation requires at least two inputs")
+ } else {
+ match set.op() {
+ SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
+ SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
+ SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect(
+ from_substrait_rel(consumer, &set.inputs[0]).await?,
+ union_rels(consumer, &set.inputs[1..], true).await?,
+ false,
+ ),
+ SetOp::IntersectionMultiset => {
+ intersect_rels(consumer, &set.inputs, false).await
+ }
+ SetOp::IntersectionMultisetAll => {
+ intersect_rels(consumer, &set.inputs, true).await
+ }
+ SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
+ SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+ set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+ }
+ }
+}
- let values = vt
- .values
- .iter()
- .map(|row| {
- let mut name_idx = 0;
- let lits = row
- .fields
- .iter()
- .map(|lit| {
- name_idx += 1; // top-level names are provided through schema
- Ok(Expr::Literal(from_substrait_literal(
- lit,
- extensions,
- &named_struct.names,
- &mut name_idx,
- )?))
- })
- .collect::<Result<_>>()?;
- if name_idx != named_struct.names.len() {
- return substrait_err!(
- "Names list must match exactly to nested schema, but found {} uses for {} names",
- name_idx,
- named_struct.names.len()
- );
- }
- Ok(lits)
- })
- .collect::<Result<_>>()?;
+pub async fn from_extension_leaf_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_leaf_rel: &ExtensionLeafRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_leaf(extension_leaf_rel).await
+}
- Ok(LogicalPlan::Values(Values {
- schema: DFSchemaRef::new(substrait_schema),
- values,
- }))
- }
- Some(ReadType::LocalFiles(lf)) => {
- fn extract_filename(name: &str) -> Option<String> {
- let corrected_url = if name.starts_with("file://")
- && !name.starts_with("file:///")
- {
- name.replacen("file://", "file:///", 1)
- } else {
- name.to_string()
- };
+pub async fn from_extension_single_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_single_rel: &ExtensionSingleRel,
+) -> Result<LogicalPlan> {
+ consumer
+ .consume_extension_single(extension_single_rel)
+ .await
+}
- Url::parse(&corrected_url).ok().and_then(|url| {
- let path = url.path();
- std::path::Path::new(path)
- .file_name()
- .map(|filename| filename.to_string_lossy().to_string())
- })
- }
+pub async fn from_extension_multi_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_multi_rel: &ExtensionMultiRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_multi(extension_multi_rel).await
+}
- // we could use the file name to check the original table provider
- // TODO: currently does not support multiple local files
- let filename: Option<String> =
- lf.items.first().and_then(|x| match x.path_type.as_ref() {
- Some(UriFile(name)) => extract_filename(name),
- _ => None,
- });
+pub async fn from_exchange_rel(
+ consumer: &impl SubstraitConsumer,
+ exchange: &ExchangeRel,
+) -> Result<LogicalPlan> {
+ let Some(input) = exchange.input.as_ref() else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+ let input = Arc::new(from_substrait_rel(consumer, input).await?);
- if lf.items.len() > 1 || filename.is_none() {
- return not_impl_err!("Only single file reads are supported");
- }
- let name = filename.unwrap();
- // directly use unwrap here since we could determine it is a valid one
- let table_reference = TableReference::Bare { table: name.into() };
-
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
- }
- _ => {
- not_impl_err!("Unsupported ReadType: {:?}", &read.as_ref().read_type)
- }
- }
- }
- Some(RelType::Set(set)) => match set_rel::SetOp::try_from(set.op) {
- Ok(set_op) => {
- if set.inputs.len() < 2 {
- substrait_err!("Set operation requires at least two inputs")
- } else {
- match set_op {
- set_rel::SetOp::UnionAll => {
- union_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::UnionDistinct => {
- union_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionPrimary => {
- LogicalPlanBuilder::intersect(
- from_substrait_rel(state, &set.inputs[0], extensions)
- .await?,
- union_rels(&set.inputs[1..], state, extensions, true)
- .await?,
- false,
- )
- }
- set_rel::SetOp::IntersectionMultiset => {
- intersect_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionMultisetAll => {
- intersect_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::MinusPrimary => {
- except_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::MinusPrimaryAll => {
- except_rels(&set.inputs, state, extensions, true).await
- }
- _ => not_impl_err!("Unsupported set operator: {set_op:?}"),
- }
- }
+ let Some(exchange_kind) = &exchange.exchange_kind else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+
+ // ref: https://substrait.io/relations/physical_relations/#exchange-types
+ let partitioning_scheme = match exchange_kind {
+ ExchangeKind::ScatterByFields(scatter_fields) => {
+ let mut partition_columns = vec![];
+ let input_schema = input.schema();
+ for field_ref in &scatter_fields.fields {
+ let column = from_substrait_field_reference(field_ref, input_schema)?;
+ partition_columns.push(column);
}
- Err(e) => not_impl_err!("Invalid set operation type {}: {e}", set.op),
- },
- Some(RelType::ExtensionLeaf(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
- }
- Some(RelType::ExtensionSingle(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- let Some(input_rel) = &extension.input else {
- return substrait_err!(
- "ExtensionSingleRel doesn't contains input rel. Try use ExtensionLeafRel instead"
- );
- };
- let input_plan = from_substrait_rel(state, input_rel, extensions).await?;
- let plan =
- plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
+ Partitioning::Hash(partition_columns, exchange.partition_count as usize)
}
- Some(RelType::ExtensionMulti(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- let mut inputs = Vec::with_capacity(extension.inputs.len());
- for input in &extension.inputs {
- let input_plan = from_substrait_rel(state, input, extensions).await?;
- inputs.push(input_plan);
- }
- let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
+ ExchangeKind::RoundRobin(_) => {
+ Partitioning::RoundRobinBatch(exchange.partition_count as usize)
}
Review Comment:
Would it make sense to keep these in the default implementation though? Just as in the producer we do produce these types (https://github.com/apache/datafusion/blob/b8e435f3e38f1f8d2d8f24c3da3c2b21e76b700b/datafusion/substrait/src/logical_plan/producer.rs#L596), and I guess it'd be nice if the roundtrip tests that don't require special setup on producer wouldn't require special setup for consumer either.
That's not a strictly held opinion though.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1888729805
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +971,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
+ consumer: &impl SubstraitConsumer,
+ p: &ProjectRel,
) -> Result<LogicalPlan> {
- let plan: Result<LogicalPlan> = match &rel.rel_type {
- Some(RelType::Project(p)) => {
- if let Some(input) = p.input.as_ref() {
- let mut input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let original_schema = input.schema().clone();
-
- // Ensure that all expressions have a unique display name, so that
- // validate_unique_names does not fail when constructing the project.
- let mut name_tracker = NameTracker::new();
-
- // By default, a Substrait Project emits all inputs fields followed by all expressions.
- // We build the explicit expressions first, and then the input expressions to avoid
- // adding aliases to the explicit expressions (as part of ensuring unique names).
- //
- // This is helpful for plan visualization and tests, because when DataFusion produces
- // Substrait Projects it adds an output mapping that excludes all input columns
- // leaving only explicit expressions.
-
- let mut explicit_exprs: Vec<Expr> = vec![];
- for expr in &p.expressions {
- let e = from_substrait_rex(
- state,
- expr,
- input.clone().schema(),
- extensions,
- )
- .await?;
- // if the expression is WindowFunction, wrap in a Window relation
- if let Expr::WindowFunction(_) = &e {
- // Adding the same expression here and in the project below
- // works because the project's builder uses columnize_expr(..)
- // to transform it into a column reference
- input = input.window(vec![e.clone()])?
- }
- explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
-
- let mut final_exprs: Vec<Expr> = vec![];
- for index in 0..original_schema.fields().len() {
- let e = Expr::Column(Column::from(
- original_schema.qualified_field(index),
- ));
- final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
- final_exprs.append(&mut explicit_exprs);
-
- input.project(final_exprs)?.build()
- } else {
- not_impl_err!("Projection without an input is not supported")
+ if let Some(input) = p.input.as_ref() {
+ let mut input =
+ LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let original_schema = input.schema().clone();
+
+ // Ensure that all expressions have a unique display name, so that
+ // validate_unique_names does not fail when constructing the project.
+ let mut name_tracker = NameTracker::new();
+
+ // By default, a Substrait Project emits all inputs fields followed by all expressions.
+ // We build the explicit expressions first, and then the input expressions to avoid
+ // adding aliases to the explicit expressions (as part of ensuring unique names).
+ //
+ // This is helpful for plan visualization and tests, because when DataFusion produces
+ // Substrait Projects it adds an output mapping that excludes all input columns
+ // leaving only explicit expressions.
+
+ let mut explicit_exprs: Vec<Expr> = vec![];
+ for expr in &p.expressions {
+ let e = from_substrait_rex(consumer, expr, input.clone().schema()).await?;
+ // if the expression is WindowFunction, wrap in a Window relation
+ if let Expr::WindowFunction(_) = &e {
+ // Adding the same expression here and in the project below
+ // works because the project's builder uses columnize_expr(..)
+ // to transform it into a column reference
+ input = input.window(vec![e.clone()])?
}
+ explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Filter(filter)) => {
- if let Some(input) = filter.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- if let Some(condition) = filter.condition.as_ref() {
- let expr =
- from_substrait_rex(state, condition, input.schema(), extensions)
- .await?;
- input.filter(expr)?.build()
- } else {
- not_impl_err!("Filter without an condition is not valid")
- }
- } else {
- not_impl_err!("Filter without an input is not valid")
- }
+
+ let mut final_exprs: Vec<Expr> = vec![];
+ for index in 0..original_schema.fields().len() {
+ let e = Expr::Column(Column::from(original_schema.qualified_field(index)));
+ final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Fetch(fetch)) => {
- if let Some(input) = fetch.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let offset = fetch.offset as usize;
- // -1 means that ALL records should be returned
- let count = if fetch.count == -1 {
- None
- } else {
- Some(fetch.count as usize)
- };
- input.limit(offset, count)?.build()
- } else {
- not_impl_err!("Fetch without an input is not valid")
- }
+ final_exprs.append(&mut explicit_exprs);
+ input.project(final_exprs)?.build()
+ } else {
+ not_impl_err!("Projection without an input is not supported")
+ }
+}
+
+#[async_recursion]
+pub async fn from_filter_rel(
+ consumer: &impl SubstraitConsumer,
+ filter: &FilterRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = filter.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ if let Some(condition) = filter.condition.as_ref() {
+ let expr = from_substrait_rex(consumer, condition, input.schema()).await?;
+ input.filter(expr)?.build()
+ } else {
+ not_impl_err!("Filter without an condition is not valid")
}
- Some(RelType::Sort(sort)) => {
- if let Some(input) = sort.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let sorts =
- from_substrait_sorts(state, &sort.sorts, input.schema(), extensions)
- .await?;
- input.sort(sorts)?.build()
- } else {
- not_impl_err!("Sort without an input is not valid")
- }
+ } else {
+ not_impl_err!("Filter without an input is not valid")
+ }
+}
+
+#[async_recursion]
+pub async fn from_fetch_rel(
+ consumer: &impl SubstraitConsumer,
+ fetch: &FetchRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = fetch.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let offset = fetch.offset as usize;
+ // -1 means that ALL records should be returned
+ let count = if fetch.count == -1 {
+ None
+ } else {
+ Some(fetch.count as usize)
+ };
+ input.limit(offset, count)?.build()
+ } else {
+ not_impl_err!("Fetch without an input is not valid")
+ }
+}
+
+pub async fn from_sort_rel(
+ consumer: &impl SubstraitConsumer,
+ sort: &SortRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = sort.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?;
+ input.sort(sorts)?.build()
+ } else {
+ not_impl_err!("Sort without an input is not valid")
+ }
+}
+
+pub async fn from_aggregate_rel(
+ consumer: &impl SubstraitConsumer,
+ agg: &AggregateRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = agg.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let mut ref_group_exprs = vec![];
+
+ for e in &agg.grouping_expressions {
+ let x = from_substrait_rex(consumer, e, input.schema()).await?;
+ ref_group_exprs.push(x);
}
- Some(RelType::Aggregate(agg)) => {
- if let Some(input) = agg.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let mut ref_group_exprs = vec![];
- for e in &agg.grouping_expressions {
- let x =
- from_substrait_rex(state, e, input.schema(), extensions).await?;
- ref_group_exprs.push(x);
+ let mut group_exprs = vec![];
+ let mut aggr_exprs = vec![];
+
+ match agg.groupings.len() {
+ 1 => {
+ group_exprs.extend_from_slice(
+ &from_substrait_grouping(
+ consumer,
+ &agg.groupings[0],
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?,
+ );
+ }
+ _ => {
+ let mut grouping_sets = vec![];
+ for grouping in &agg.groupings {
+ let grouping_set = from_substrait_grouping(
+ consumer,
+ grouping,
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?;
+ grouping_sets.push(grouping_set);
}
+ // Single-element grouping expression of type Expr::GroupingSet.
+ // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+ // parsed by the producer and consumer, since Substrait does not have a type dedicated
+ // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+ group_exprs
+ .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+ }
+ };
- let mut group_exprs = vec![];
- let mut aggr_exprs = vec![];
-
- match agg.groupings.len() {
- 1 => {
- group_exprs.extend_from_slice(
- &from_substrait_grouping(
- state,
- &agg.groupings[0],
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?,
- );
- }
- _ => {
- let mut grouping_sets = vec![];
- for grouping in &agg.groupings {
- let grouping_set = from_substrait_grouping(
- state,
- grouping,
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?;
- grouping_sets.push(grouping_set);
+ for m in &agg.measures {
+ let filter = match &m.filter {
+ Some(fil) => Some(Box::new(
+ from_substrait_rex(consumer, fil, input.schema()).await?,
+ )),
+ None => None,
+ };
+ let agg_func = match &m.measure {
+ Some(f) => {
+ let distinct = match f.invocation {
+ _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+ true
}
- // Single-element grouping expression of type Expr::GroupingSet.
- // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
- // parsed by the producer and consumer, since Substrait does not have a type dedicated
- // to ROLLUP. Only vector of Groupings (grouping sets) is available.
- group_exprs.push(Expr::GroupingSet(GroupingSet::GroupingSets(
- grouping_sets,
- )));
- }
- };
-
- for m in &agg.measures {
- let filter = match &m.filter {
- Some(fil) => Some(Box::new(
- from_substrait_rex(state, fil, input.schema(), extensions)
- .await?,
- )),
- None => None,
+ _ if f.invocation == AggregationInvocation::All as i32 => false,
+ _ => false,
};
- let agg_func = match &m.measure {
- Some(f) => {
- let distinct = match f.invocation {
- _ if f.invocation
- == AggregationInvocation::Distinct as i32 =>
- {
- true
- }
- _ if f.invocation
- == AggregationInvocation::All as i32 =>
- {
- false
- }
- _ => false,
- };
- let order_by = if !f.sorts.is_empty() {
- Some(
- from_substrait_sorts(
- state,
- &f.sorts,
- input.schema(),
- extensions,
- )
- .await?,
- )
- } else {
- None
- };
-
- from_substrait_agg_func(
- state,
- f,
- input.schema(),
- extensions,
- filter,
- order_by,
- distinct,
- )
- .await
- }
- None => not_impl_err!(
- "Aggregate without aggregate function is not supported"
- ),
+ let order_by = if !f.sorts.is_empty() {
+ Some(
+ from_substrait_sorts(consumer, &f.sorts, input.schema())
+ .await?,
+ )
+ } else {
+ None
};
- aggr_exprs.push(agg_func?.as_ref().clone());
- }
- input.aggregate(group_exprs, aggr_exprs)?.build()
- } else {
- not_impl_err!("Aggregate without an input is not valid")
- }
- }
- Some(RelType::Join(join)) => {
- if join.post_join_filter.is_some() {
- return not_impl_err!(
- "JoinRel with post_join_filter is not yet supported"
- );
- }
- let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
-
- let join_type = from_substrait_jointype(join.r#type)?;
- // The join condition expression needs full input schema and not the output schema from join since we lose columns from
- // certain join types such as semi and anti joins
- let in_join_schema = left.schema().join(right.schema())?;
-
- // If join expression exists, parse the `on` condition expression, build join and return
- // Otherwise, build join with only the filter, without join keys
- match &join.expression.as_ref() {
- Some(expr) => {
- let on = from_substrait_rex(state, expr, &in_join_schema, extensions)
- .await?;
- // The join expression can contain both equal and non-equal ops.
- // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
- // So we extract each part as follows:
- // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
- // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
- let (join_ons, nulls_equal_nulls, join_filter) =
- split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
- let (left_cols, right_cols): (Vec<_>, Vec<_>) =
- itertools::multiunzip(join_ons);
- left.join_detailed(
- right.build()?,
- join_type,
- (left_cols, right_cols),
- join_filter,
- nulls_equal_nulls,
- )?
- .build()
+ from_substrait_agg_func(
+ consumer,
+ f,
+ input.schema(),
+ filter,
+ order_by,
+ distinct,
+ )
+ .await
}
None => {
- let on: Vec<String> = vec![];
- left.join_detailed(
- right.build()?,
- join_type,
- (on.clone(), on),
- None,
- false,
- )?
- .build()
+ not_impl_err!("Aggregate without aggregate function is not supported")
}
- }
+ };
+ aggr_exprs.push(agg_func?.as_ref().clone());
}
- Some(RelType::Cross(cross)) => {
- let left = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
- left.cross_join(right.build()?)?.build()
+ input.aggregate(group_exprs, aggr_exprs)?.build()
+ } else {
+ not_impl_err!("Aggregate without an input is not valid")
+ }
+}
+
+pub async fn from_join_rel(
+ consumer: &impl SubstraitConsumer,
+ join: &JoinRel,
+) -> Result<LogicalPlan> {
+ if join.post_join_filter.is_some() {
+ return not_impl_err!("JoinRel with post_join_filter is not yet supported");
+ }
+
+ let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+
+ let join_type = from_substrait_jointype(join.r#type)?;
+ // The join condition expression needs full input schema and not the output schema from join since we lose columns from
+ // certain join types such as semi and anti joins
+ let in_join_schema = left.schema().join(right.schema())?;
+
+ // If join expression exists, parse the `on` condition expression, build join and return
+ // Otherwise, build join with only the filter, without join keys
+ match &join.expression.as_ref() {
+ Some(expr) => {
+ let on = from_substrait_rex(consumer, expr, &in_join_schema).await?;
+ // The join expression can contain both equal and non-equal ops.
+ // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
+ // So we extract each part as follows:
+ // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
+ // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
+ let (join_ons, nulls_equal_nulls, join_filter) =
+ split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) =
+ itertools::multiunzip(join_ons);
+ left.join_detailed(
+ right.build()?,
+ join_type,
+ (left_cols, right_cols),
+ join_filter,
+ nulls_equal_nulls,
+ )?
+ .build()
}
- Some(RelType::Read(read)) => {
- async fn read_with_schema(
- state: &dyn SubstraitPlanningState,
- table_ref: TableReference,
- schema: DFSchema,
- projection: &Option<MaskExpression>,
- ) -> Result<LogicalPlan> {
- let schema = schema.replace_qualifier(table_ref.clone());
-
- let plan = {
- let provider = match state.table(&table_ref).await? {
- Some(ref provider) => Arc::clone(provider),
- _ => return plan_err!("No table named '{table_ref}'"),
- };
+ None => {
+ let on: Vec<String> = vec![];
+ left.join_detailed(right.build()?, join_type, (on.clone(), on), None, false)?
+ .build()
+ }
+ }
+}
- LogicalPlanBuilder::scan(
- table_ref,
- provider_as_source(Arc::clone(&provider)),
- None,
- )?
- .build()?
- };
+pub async fn from_cross_rel(
+ consumer: &impl SubstraitConsumer,
+ cross: &CrossRel,
+) -> Result<LogicalPlan> {
+ let left = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+ left.cross_join(right.build()?)?.build()
+}
- ensure_schema_compatability(plan.schema(), schema.clone())?;
+#[allow(deprecated)]
+pub async fn from_read_rel(
+ consumer: &impl SubstraitConsumer,
+ read: &ReadRel,
+) -> Result<LogicalPlan> {
+ async fn read_with_schema(
+ consumer: &impl SubstraitConsumer,
+ table_ref: TableReference,
+ schema: DFSchema,
+ projection: &Option<MaskExpression>,
+ ) -> Result<LogicalPlan> {
+ let schema = schema.replace_qualifier(table_ref.clone());
+
+ let plan = {
+ let provider = match consumer.resolve_table_ref(&table_ref).await? {
+ Some(ref provider) => Arc::clone(provider),
+ _ => return plan_err!("No table named '{table_ref}'"),
+ };
- let schema = apply_masking(schema, projection)?;
+ LogicalPlanBuilder::scan(
+ table_ref,
+ provider_as_source(Arc::clone(&provider)),
+ None,
+ )?
+ .build()?
+ };
- apply_projection(plan, schema)
- }
+ ensure_schema_compatability(plan.schema(), schema.clone())?;
- let named_struct = read.base_schema.as_ref().ok_or_else(|| {
- substrait_datafusion_err!("No base schema provided for Read Relation")
- })?;
+ let schema = apply_masking(schema, projection)?;
- let substrait_schema = from_substrait_named_struct(named_struct, extensions)?;
+ apply_projection(plan, schema)
+ }
- match &read.as_ref().read_type {
- Some(ReadType::NamedTable(nt)) => {
- let table_reference = match nt.names.len() {
- 0 => {
- return plan_err!("No table name found in NamedTable");
- }
- 1 => TableReference::Bare {
- table: nt.names[0].clone().into(),
- },
- 2 => TableReference::Partial {
- schema: nt.names[0].clone().into(),
- table: nt.names[1].clone().into(),
- },
- _ => TableReference::Full {
- catalog: nt.names[0].clone().into(),
- schema: nt.names[1].clone().into(),
- table: nt.names[2].clone().into(),
- },
- };
+ let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+ substrait_datafusion_err!("No base schema provided for Read Relation")
+ })?;
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
+ let substrait_schema = from_substrait_named_struct(consumer, named_struct)?;
+
+ match &read.read_type {
+ Some(ReadType::NamedTable(nt)) => {
+ let table_reference = match nt.names.len() {
+ 0 => {
+ return plan_err!("No table name found in NamedTable");
}
- Some(ReadType::VirtualTable(vt)) => {
- if vt.values.is_empty() {
- return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: DFSchemaRef::new(substrait_schema),
- }));
+ 1 => TableReference::Bare {
+ table: nt.names[0].clone().into(),
+ },
+ 2 => TableReference::Partial {
+ schema: nt.names[0].clone().into(),
+ table: nt.names[1].clone().into(),
+ },
+ _ => TableReference::Full {
+ catalog: nt.names[0].clone().into(),
+ schema: nt.names[1].clone().into(),
+ table: nt.names[2].clone().into(),
+ },
+ };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ Some(ReadType::VirtualTable(vt)) => {
+ if vt.values.is_empty() {
+ return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row: false,
+ schema: DFSchemaRef::new(substrait_schema),
+ }));
+ }
+
+ let values = vt
+ .values
+ .iter()
+ .map(|row| {
+ let mut name_idx = 0;
+ let lits = row
+ .fields
+ .iter()
+ .map(|lit| {
+ name_idx += 1; // top-level names are provided through schema
+ Ok(Expr::Literal(from_substrait_literal(
+ consumer,
+ lit,
+ &named_struct.names,
+ &mut name_idx,
+ )?))
+ })
+ .collect::<Result<_>>()?;
+ if name_idx != named_struct.names.len() {
+ return substrait_err!(
+ "Names list must match exactly to nested schema, but found {} uses for {} names",
+ name_idx,
+ named_struct.names.len()
+ );
}
+ Ok(lits)
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(LogicalPlan::Values(Values {
+ schema: DFSchemaRef::new(substrait_schema),
+ values,
+ }))
+ }
+ Some(ReadType::LocalFiles(lf)) => {
+ fn extract_filename(name: &str) -> Option<String> {
+ let corrected_url =
+ if name.starts_with("file://") && !name.starts_with("file:///") {
+ name.replacen("file://", "file:///", 1)
+ } else {
+ name.to_string()
+ };
+
+ Url::parse(&corrected_url).ok().and_then(|url| {
+ let path = url.path();
+ std::path::Path::new(path)
+ .file_name()
+ .map(|filename| filename.to_string_lossy().to_string())
+ })
+ }
+
+ // we could use the file name to check the original table provider
+ // TODO: currently does not support multiple local files
+ let filename: Option<String> =
+ lf.items.first().and_then(|x| match x.path_type.as_ref() {
+ Some(UriFile(name)) => extract_filename(name),
+ _ => None,
+ });
+
+ if lf.items.len() > 1 || filename.is_none() {
+ return not_impl_err!("Only single file reads are supported");
+ }
+ let name = filename.unwrap();
+ // directly use unwrap here since we could determine it is a valid one
+ let table_reference = TableReference::Bare { table: name.into() };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ _ => {
+ not_impl_err!("Unsupported ReadType: {:?}", read.read_type)
+ }
+ }
+}
+
+pub async fn from_set_rel(
+ consumer: &impl SubstraitConsumer,
+ set: &SetRel,
+) -> Result<LogicalPlan> {
+ if set.inputs.len() < 2 {
+ substrait_err!("Set operation requires at least two inputs")
+ } else {
+ match set.op() {
+ SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
+ SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
+ SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect(
+ from_substrait_rel(consumer, &set.inputs[0]).await?,
+ union_rels(consumer, &set.inputs[1..], true).await?,
+ false,
+ ),
+ SetOp::IntersectionMultiset => {
+ intersect_rels(consumer, &set.inputs, false).await
+ }
+ SetOp::IntersectionMultisetAll => {
+ intersect_rels(consumer, &set.inputs, true).await
+ }
+ SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
+ SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+ set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+ }
+ }
+}
- let values = vt
- .values
- .iter()
- .map(|row| {
- let mut name_idx = 0;
- let lits = row
- .fields
- .iter()
- .map(|lit| {
- name_idx += 1; // top-level names are provided through schema
- Ok(Expr::Literal(from_substrait_literal(
- lit,
- extensions,
- &named_struct.names,
- &mut name_idx,
- )?))
- })
- .collect::<Result<_>>()?;
- if name_idx != named_struct.names.len() {
- return substrait_err!(
- "Names list must match exactly to nested schema, but found {} uses for {} names",
- name_idx,
- named_struct.names.len()
- );
- }
- Ok(lits)
- })
- .collect::<Result<_>>()?;
+pub async fn from_extension_leaf_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_leaf_rel: &ExtensionLeafRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_leaf(extension_leaf_rel).await
+}
- Ok(LogicalPlan::Values(Values {
- schema: DFSchemaRef::new(substrait_schema),
- values,
- }))
- }
- Some(ReadType::LocalFiles(lf)) => {
- fn extract_filename(name: &str) -> Option<String> {
- let corrected_url = if name.starts_with("file://")
- && !name.starts_with("file:///")
- {
- name.replacen("file://", "file:///", 1)
- } else {
- name.to_string()
- };
+pub async fn from_extension_single_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_single_rel: &ExtensionSingleRel,
+) -> Result<LogicalPlan> {
+ consumer
+ .consume_extension_single(extension_single_rel)
+ .await
+}
- Url::parse(&corrected_url).ok().and_then(|url| {
- let path = url.path();
- std::path::Path::new(path)
- .file_name()
- .map(|filename| filename.to_string_lossy().to_string())
- })
- }
+pub async fn from_extension_multi_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_multi_rel: &ExtensionMultiRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_multi(extension_multi_rel).await
+}
- // we could use the file name to check the original table provider
- // TODO: currently does not support multiple local files
- let filename: Option<String> =
- lf.items.first().and_then(|x| match x.path_type.as_ref() {
- Some(UriFile(name)) => extract_filename(name),
- _ => None,
- });
+pub async fn from_exchange_rel(
+ consumer: &impl SubstraitConsumer,
+ exchange: &ExchangeRel,
+) -> Result<LogicalPlan> {
+ let Some(input) = exchange.input.as_ref() else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+ let input = Arc::new(from_substrait_rel(consumer, input).await?);
- if lf.items.len() > 1 || filename.is_none() {
- return not_impl_err!("Only single file reads are supported");
- }
- let name = filename.unwrap();
- // directly use unwrap here since we could determine it is a valid one
- let table_reference = TableReference::Bare { table: name.into() };
-
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
- }
- _ => {
- not_impl_err!("Unsupported ReadType: {:?}", &read.as_ref().read_type)
- }
- }
- }
- Some(RelType::Set(set)) => match set_rel::SetOp::try_from(set.op) {
- Ok(set_op) => {
- if set.inputs.len() < 2 {
- substrait_err!("Set operation requires at least two inputs")
- } else {
- match set_op {
- set_rel::SetOp::UnionAll => {
- union_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::UnionDistinct => {
- union_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionPrimary => {
- LogicalPlanBuilder::intersect(
- from_substrait_rel(state, &set.inputs[0], extensions)
- .await?,
- union_rels(&set.inputs[1..], state, extensions, true)
- .await?,
- false,
- )
- }
- set_rel::SetOp::IntersectionMultiset => {
- intersect_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionMultisetAll => {
- intersect_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::MinusPrimary => {
- except_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::MinusPrimaryAll => {
- except_rels(&set.inputs, state, extensions, true).await
- }
- _ => not_impl_err!("Unsupported set operator: {set_op:?}"),
- }
- }
+ let Some(exchange_kind) = &exchange.exchange_kind else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+
+ // ref: https://substrait.io/relations/physical_relations/#exchange-types
+ let partitioning_scheme = match exchange_kind {
+ ExchangeKind::ScatterByFields(scatter_fields) => {
+ let mut partition_columns = vec![];
+ let input_schema = input.schema();
+ for field_ref in &scatter_fields.fields {
+ let column = from_substrait_field_reference(field_ref, input_schema)?;
+ partition_columns.push(column);
}
- Err(e) => not_impl_err!("Invalid set operation type {}: {e}", set.op),
- },
- Some(RelType::ExtensionLeaf(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
- }
- Some(RelType::ExtensionSingle(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- let Some(input_rel) = &extension.input else {
- return substrait_err!(
- "ExtensionSingleRel doesn't contains input rel. Try use ExtensionLeafRel instead"
- );
- };
- let input_plan = from_substrait_rel(state, input_rel, extensions).await?;
- let plan =
- plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
+ Partitioning::Hash(partition_columns, exchange.partition_count as usize)
}
- Some(RelType::ExtensionMulti(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- let mut inputs = Vec::with_capacity(extension.inputs.len());
- for input in &extension.inputs {
- let input_plan = from_substrait_rel(state, input, extensions).await?;
- inputs.push(input_plan);
- }
- let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
+ ExchangeKind::RoundRobin(_) => {
+ Partitioning::RoundRobinBatch(exchange.partition_count as usize)
}
Review Comment:
> Would it make sense to keep these in the default implementation though?
I think it would be reasonable to include these in the `DefaultSubstraitConsumer`. That would play nice with roundtrip tests, and it might be a reasonable default to assume that extension relations are `UserDefinedLogicalNode`s.
I'll add these back in.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894407409
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +101,471 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::{FunctionRegistry, SessionState};
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[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()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and FunctionRegistry is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_function_registry(&self) -> &impl FunctionRegistry;
+
+ // Relation Methods
Review Comment:
@/Blizzara also suggested splitting the code into modules. I've combined both of those ideas into a future issue https://github.com/apache/datafusion/issues/13864
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894401312
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +101,471 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::{FunctionRegistry, SessionState};
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[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()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and FunctionRegistry is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_function_registry(&self) -> &impl FunctionRegistry;
+
+ // Relation Methods
+ // There is one method per Substrait relation to allow for easy overriding of consumer behaviour.
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_read(&self, rel: &ReadRel) -> Result<LogicalPlan> {
+ from_read_rel(self, rel).await
+ }
+
+ async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+ from_filter_rel(self, rel).await
+ }
+
+ async fn consume_fetch(&self, rel: &FetchRel) -> Result<LogicalPlan> {
+ from_fetch_rel(self, rel).await
+ }
+
+ async fn consume_aggregate(&self, rel: &AggregateRel) -> Result<LogicalPlan> {
+ from_aggregate_rel(self, rel).await
+ }
+
+ async fn consume_sort(&self, rel: &SortRel) -> Result<LogicalPlan> {
+ from_sort_rel(self, rel).await
+ }
+
+ async fn consume_join(&self, rel: &JoinRel) -> Result<LogicalPlan> {
+ from_join_rel(self, rel).await
+ }
+
+ async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+ from_project_rel(self, rel).await
+ }
+
+ async fn consume_set(&self, rel: &SetRel) -> Result<LogicalPlan> {
+ from_set_rel(self, rel).await
+ }
+
+ async fn consume_cross(&self, rel: &CrossRel) -> Result<LogicalPlan> {
+ from_cross_rel(self, rel).await
+ }
+
+ async fn consume_consistent_partition_window(
+ &self,
+ _rel: &ConsistentPartitionWindowRel,
+ ) -> Result<LogicalPlan> {
+ not_impl_err!("Consistent Partition Window Rel not supported")
+ }
+
+ async fn consume_exchange(&self, rel: &ExchangeRel) -> Result<LogicalPlan> {
+ from_exchange_rel(self, rel).await
+ }
+
+ // Expression Methods
+ // There is one method per Substrait expression to allow for easy overriding of consumer behaviour
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_literal(&self, expr: &Literal) -> Result<Expr> {
+ from_literal(self, expr).await
+ }
+
+ async fn consume_field_reference(
+ &self,
+ expr: &FieldReference,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_field_reference(self, expr, input_schema).await
+ }
+
+ async fn consume_scalar_function(
+ &self,
+ expr: &ScalarFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_scalar_function(self, expr, input_schema).await
+ }
+
+ async fn consume_window_function(
+ &self,
+ expr: &WindowFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_window_function(self, expr, input_schema).await
+ }
+
+ async fn consume_if_then(
+ &self,
+ expr: &IfThen,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_if_then(self, expr, input_schema).await
+ }
+
+ async fn consume_switch(
+ &self,
+ _expr: &SwitchExpression,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Switch expression not supported")
+ }
+
+ async fn consume_singular_or_list(
+ &self,
+ expr: &SingularOrList,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_singular_or_list(self, expr, input_schema).await
+ }
+
+ async fn consume_multi_or_list(
+ &self,
+ _expr: &MultiOrList,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Multi Or List expression not supported")
+ }
+
+ async fn consume_cast(
+ &self,
+ expr: &substrait_expression::Cast,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_cast(self, expr, input_schema).await
+ }
+
+ async fn consume_subquery(
+ &self,
+ expr: &substrait_expression::Subquery,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_subquery(self, expr, input_schema).await
+ }
+
+ async fn consume_nested(
+ &self,
+ _expr: &Nested,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Nested expression not supported")
+ }
+
+ async fn consume_enum(&self, _expr: &Enum, _input_schema: &DFSchema) -> Result<Expr> {
+ not_impl_err!("Enum expression not supported")
+ }
+
+ // User-Defined Functionality
+
+ // The details of extension relations, and how to handle them, are fully up to users to specify.
+ // The following methods allow users to customize the consumer behaviour
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionLeafRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionLeafRel")
+ }
+
+ async fn consume_extension_single(
+ &self,
+ rel: &ExtensionSingleRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionSingleRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionSingleRel")
+ }
+
+ async fn consume_extension_multi(
+ &self,
+ rel: &ExtensionMultiRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionMultiRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionMultiRel")
+ }
+
+ // Users can bring their own types to Substrait which require custom handling
+
+ fn consume_user_defined_type(
+ &self,
+ user_defined_type: &r#type::UserDefined,
+ ) -> Result<DataType> {
+ substrait_err!(
+ "Missing handler for user-defined type: {}",
+ user_defined_type.type_reference
+ )
+ }
+
+ fn consume_user_defined_literal(
+ &self,
+ user_defined_literal: &proto::expression::literal::UserDefined,
+ ) -> Result<ScalarValue> {
+ substrait_err!(
+ "Missing handler for user-defined literals {}",
+ user_defined_literal.type_reference
+ )
+ }
+}
+
+/// Convert Substrait Rel to DataFusion DataFrame
+#[async_recursion]
+pub async fn from_substrait_rel(
+ consumer: &impl SubstraitConsumer,
+ relation: &Rel,
+) -> Result<LogicalPlan> {
+ let plan: Result<LogicalPlan> = match &relation.rel_type {
+ Some(rel_type) => match rel_type {
+ RelType::Read(rel) => consumer.consume_read(rel).await,
+ RelType::Filter(rel) => consumer.consume_filter(rel).await,
+ RelType::Fetch(rel) => consumer.consume_fetch(rel).await,
+ RelType::Aggregate(rel) => consumer.consume_aggregate(rel).await,
+ RelType::Sort(rel) => consumer.consume_sort(rel).await,
+ RelType::Join(rel) => consumer.consume_join(rel).await,
+ RelType::Project(rel) => consumer.consume_project(rel).await,
+ RelType::Set(rel) => consumer.consume_set(rel).await,
+ RelType::ExtensionSingle(rel) => consumer.consume_extension_single(rel).await,
+ RelType::ExtensionMulti(rel) => consumer.consume_extension_multi(rel).await,
+ RelType::ExtensionLeaf(rel) => consumer.consume_extension_leaf(rel).await,
+ RelType::Cross(rel) => consumer.consume_cross(rel).await,
+ RelType::Window(rel) => {
+ consumer.consume_consistent_partition_window(rel).await
+ }
+ RelType::Exchange(rel) => consumer.consume_exchange(rel).await,
+ rt => not_impl_err!("{rt:?} rel not supported yet"),
+ },
+ None => return substrait_err!("rel must set rel_type"),
+ };
+ apply_emit_kind(retrieve_rel_common(relation), plan?)
+}
+
+/// Can be used to consume standard Substrait without user-defined extensions
Review Comment:
In the docs, I've pointed users towards `from_substrait_plan` which uses this.
This is really only `pub` so it can be used in integration tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894400506
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +101,471 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::{FunctionRegistry, SessionState};
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[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()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
Review Comment:
Created https://github.com/apache/datafusion/issues/13863 for this linked it in the code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "ccciudatu (via GitHub)" <gi...@apache.org>.
ccciudatu commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1888961320
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -977,6 +984,43 @@ async fn new_test_grammar() -> Result<()> {
.await
}
+struct MockSubstraitConsumer {
+ extensions: Extensions,
+ state: SessionState,
+}
+
+#[async_trait]
+impl SubstraitConsumer for MockSubstraitConsumer {
+ async fn resolve_table_ref(
+ &self,
+ _: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>> {
+ not_impl_err!("MockSubstraitConsumer does not implement resolve_table_ref")
+ }
+
+ fn get_extensions(&self) -> &Extensions {
+ &self.extensions
+ }
+
+ fn get_state(&self) -> &SessionState {
+ &self.state
+ }
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ let Some(ext_detail) = &rel.detail else {
+ return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
+ };
+ let plan = self
+ .state
+ .serializer_registry()
+ .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
Review Comment:
One drawback of the current implementation is that it's abusing the `type_url` field in `proto::Any` (which is quite [clearly specified](https://protobuf.dev/reference/java/api-docs/com/google/protobuf/AnyOrBuilder.html#getTypeUrl--)). Working with Substrait objects directly would eliminate the need for hacks like encoding node names as Any::type_url.
As a side note, this may sound hypocritical, given that I made it even worse [with my PR](https://github.com/apache/datafusion/pull/13772) by setting it to table names as well. But I only did it for consistency. :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1889320195
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -1503,370 +1858,416 @@ pub async fn from_substrait_func_args(
/// Convert Substrait AggregateFunction to DataFusion Expr
pub async fn from_substrait_agg_func(
- state: &dyn SubstraitPlanningState,
+ consumer: &impl SubstraitConsumer,
f: &AggregateFunction,
input_schema: &DFSchema,
- extensions: &Extensions,
filter: Option<Box<Expr>>,
order_by: Option<Vec<SortExpr>>,
distinct: bool,
) -> Result<Arc<Expr>> {
- let args =
- from_substrait_func_args(state, &f.arguments, input_schema, extensions).await?;
-
- let Some(function_name) = extensions.functions.get(&f.function_reference) else {
+ let Some(fn_signature) = consumer
+ .get_extensions()
+ .functions
+ .get(&f.function_reference)
+ else {
return plan_err!(
"Aggregate function not registered: function anchor = {:?}",
f.function_reference
);
};
- let function_name = substrait_fun_name(function_name);
- // try udaf first, then built-in aggr fn.
- if let Ok(fun) = state.udaf(function_name) {
- // deal with situation that count(*) got no arguments
- let args = if fun.name() == "count" && args.is_empty() {
- vec![Expr::Literal(ScalarValue::Int64(Some(1)))]
- } else {
- args
- };
-
- Ok(Arc::new(Expr::AggregateFunction(
- expr::AggregateFunction::new_udf(fun, args, distinct, filter, order_by, None),
- )))
- } else {
- not_impl_err!(
+ let fn_name = substrait_fun_name(fn_signature);
+ let udaf = consumer.get_state().udaf(fn_name);
+ let udaf = udaf.map_err(|_| {
+ not_impl_datafusion_err!(
"Aggregate function {} is not supported: function anchor = {:?}",
- function_name,
+ fn_signature,
f.function_reference
)
- }
+ })?;
+
+ let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?;
+
+ // deal with situation that count(*) got no arguments
+ let args = if udaf.name() == "count" && args.is_empty() {
+ vec![Expr::Literal(ScalarValue::Int64(Some(1)))]
+ } else {
+ args
+ };
+
+ Ok(Arc::new(Expr::AggregateFunction(
+ expr::AggregateFunction::new_udf(udaf, args, distinct, filter, order_by, None),
+ )))
}
/// Convert Substrait Rex to DataFusion Expr
-#[async_recursion]
pub async fn from_substrait_rex(
- state: &dyn SubstraitPlanningState,
- e: &Expression,
+ consumer: &impl SubstraitConsumer,
+ expression: &Expression,
input_schema: &DFSchema,
- extensions: &Extensions,
) -> Result<Expr> {
- match &e.rex_type {
- Some(RexType::SingularOrList(s)) => {
- let substrait_expr = s.value.as_ref().unwrap();
- let substrait_list = s.options.as_ref();
- Ok(Expr::InList(InList {
- expr: Box::new(
- from_substrait_rex(state, substrait_expr, input_schema, extensions)
- .await?,
- ),
- list: from_substrait_rex_vec(
- state,
- substrait_list,
- input_schema,
- extensions,
- )
- .await?,
- negated: false,
- }))
- }
- Some(RexType::Selection(field_ref)) => {
- Ok(from_substrait_field_reference(field_ref, input_schema)?)
- }
- Some(RexType::IfThen(if_then)) => {
- // Parse `ifs`
- // If the first element does not have a `then` part, then we can assume it's a base expression
- let mut when_then_expr: Vec<(Box<Expr>, Box<Expr>)> = vec![];
- let mut expr = None;
- for (i, if_expr) in if_then.ifs.iter().enumerate() {
- if i == 0 {
- // Check if the first element is type base expression
- if if_expr.then.is_none() {
- expr = Some(Box::new(
- from_substrait_rex(
- state,
- if_expr.r#if.as_ref().unwrap(),
- input_schema,
- extensions,
- )
- .await?,
- ));
- continue;
- }
- }
- when_then_expr.push((
- Box::new(
- from_substrait_rex(
- state,
- if_expr.r#if.as_ref().unwrap(),
- input_schema,
- extensions,
- )
- .await?,
- ),
- Box::new(
- from_substrait_rex(
- state,
- if_expr.then.as_ref().unwrap(),
- input_schema,
- extensions,
- )
- .await?,
- ),
- ));
+ match &expression.rex_type {
+ Some(t) => match t {
+ RexType::Literal(expr) => consumer.consume_literal(expr).await,
+ RexType::Selection(expr) => {
+ consumer.consume_selection(expr, input_schema).await
+ }
+ RexType::ScalarFunction(expr) => {
+ consumer.consume_scalar_function(expr, input_schema).await
+ }
+ RexType::WindowFunction(expr) => {
+ consumer.consume_window_function(expr, input_schema).await
+ }
+ RexType::IfThen(expr) => consumer.consume_if_then(expr, input_schema).await,
+ RexType::SwitchExpression(expr) => {
+ consumer.consume_switch(expr, input_schema).await
+ }
+ RexType::SingularOrList(expr) => {
+ consumer.consume_singular_or_list(expr, input_schema).await
}
- // Parse `else`
- let else_expr = match &if_then.r#else {
- Some(e) => Some(Box::new(
- from_substrait_rex(state, e, input_schema, extensions).await?,
- )),
- None => None,
- };
- Ok(Expr::Case(Case {
- expr,
- when_then_expr,
- else_expr,
- }))
- }
- Some(RexType::ScalarFunction(f)) => {
- let Some(fn_name) = extensions.functions.get(&f.function_reference) else {
- return plan_err!(
- "Scalar function not found: function reference = {:?}",
- f.function_reference
- );
- };
- let fn_name = substrait_fun_name(fn_name);
- let args =
- from_substrait_func_args(state, &f.arguments, input_schema, extensions)
- .await?;
+ RexType::MultiOrList(expr) => {
+ consumer.consume_multi_or_list(expr, input_schema).await
+ }
- // try to first match the requested function into registered udfs, then built-in ops
- // and finally built-in expressions
- if let Ok(func) = state.udf(fn_name) {
- Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf(
- func.to_owned(),
- args,
- )))
- } else if let Some(op) = name_to_op(fn_name) {
- if f.arguments.len() < 2 {
- return not_impl_err!(
- "Expect at least two arguments for binary operator {op:?}, the provided number of operators is {:?}",
- f.arguments.len()
- );
- }
- // Some expressions are binary in DataFusion but take in a variadic number of args in Substrait.
- // In those cases we iterate through all the arguments, applying the binary expression against them all
- let combined_expr = args
- .into_iter()
- .fold(None, |combined_expr: Option<Expr>, arg: Expr| {
- Some(match combined_expr {
- Some(expr) => Expr::BinaryExpr(BinaryExpr {
- left: Box::new(expr),
- op,
- right: Box::new(arg),
- }),
- None => arg,
- })
- })
- .unwrap();
+ RexType::Cast(expr) => {
+ consumer.consume_cast(expr.as_ref(), input_schema).await
+ }
- Ok(combined_expr)
- } else if let Some(builder) = BuiltinExprBuilder::try_from_name(fn_name) {
- builder.build(state, f, input_schema, extensions).await
- } else {
- not_impl_err!("Unsupported function name: {fn_name:?}")
+ RexType::Subquery(expr) => {
+ consumer.consume_subquery(expr.as_ref(), input_schema).await
}
- }
- Some(RexType::Literal(lit)) => {
- let scalar_value = from_substrait_literal_without_names(lit, extensions)?;
- Ok(Expr::Literal(scalar_value))
- }
- Some(RexType::Cast(cast)) => match cast.as_ref().r#type.as_ref() {
- Some(output_type) => {
- let input_expr = Box::new(
+ RexType::Nested(expr) => consumer.consume_nested(expr, input_schema).await,
+ RexType::Enum(expr) => consumer.consume_enum(expr, input_schema).await,
+ },
+ None => substrait_err!("Expression must set rex_type: {:?}", expression),
+ }
+}
+
+pub async fn from_singular_or_list(
+ consumer: &impl SubstraitConsumer,
+ expr: &SingularOrList,
+ input_schema: &DFSchema,
+) -> Result<Expr> {
+ let substrait_expr = expr.value.as_ref().unwrap();
+ let substrait_list = expr.options.as_ref();
+ Ok(Expr::InList(InList {
+ expr: Box::new(from_substrait_rex(consumer, substrait_expr, input_schema).await?),
+ list: from_substrait_rex_vec(consumer, substrait_list, input_schema).await?,
+ negated: false,
+ }))
+}
+
+pub async fn from_selection(
Review Comment:
I think that's a better name. I updated both `consume_selection` and `from_selection`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1889342502
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -1503,370 +1858,416 @@ pub async fn from_substrait_func_args(
/// Convert Substrait AggregateFunction to DataFusion Expr
pub async fn from_substrait_agg_func(
- state: &dyn SubstraitPlanningState,
+ consumer: &impl SubstraitConsumer,
f: &AggregateFunction,
input_schema: &DFSchema,
- extensions: &Extensions,
filter: Option<Box<Expr>>,
order_by: Option<Vec<SortExpr>>,
distinct: bool,
) -> Result<Arc<Expr>> {
- let args =
- from_substrait_func_args(state, &f.arguments, input_schema, extensions).await?;
-
- let Some(function_name) = extensions.functions.get(&f.function_reference) else {
+ let Some(fn_signature) = consumer
+ .get_extensions()
Review Comment:
I was trying to figure out the right abstraction not just for resolving the function via the extension, but also mapping it into DataFusion. However I found the existing logic a little unwieldy to reason about. Specifically scalar functions first try and lookup the function by name, then look it up as an operator and then try and treat is a BuiltinExpr.
I want to think about function mapping in isolation, because I suspect it will be easier to figure out the right abstraction if we simplify the lookup logic first.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#issuecomment-2549828856
@Blizzara
> One thing is we should consider splitting the consumer.rs into smaller files. It could be a separate PR to keep the diff smaller, but maybe moving all the static "consumer_X" functions into like consumer/relations.rs, consumer/expressions.rs, consumer/literals.rs and consumer/types.rs (and one more for utils etc) or similar? Feels like now could be a good moment to do that.
I think it makes sense to do that, and I'm happy to own it, but I would prefer to do that as a followup. I structured the `consumer.rs` changes in order to make it easy to see that the functions I added were extracted from the big `from_substrait_rel` and `from_substrait_rex` switch statements.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1889330504
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -2037,42 +2438,48 @@ fn from_substrait_type(
Ok(DataType::Interval(IntervalUnit::MonthDayNano))
}
r#type::Kind::UserDefined(u) => {
- if let Some(name) = extensions.types.get(&u.type_reference) {
+ if let Ok(data_type) = consumer.consume_user_defined_type(u) {
+ return Ok(data_type);
+ }
+
+ // TODO: remove the code below once the producer has been updated
Review Comment:
I thought about removing it as part of this PR, but figured that would be better handled as a separate follow-up change. I'm already breaking enough API with this as is 😅
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1890181294
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -21,23 +21,24 @@ use datafusion::arrow::array::{GenericListArray, MapArray};
use datafusion::arrow::datatypes::{
DataType, Field, FieldRef, Fields, IntervalUnit, Schema, TimeUnit,
};
-use datafusion::common::plan_err;
use datafusion::common::{
- not_impl_err, plan_datafusion_err, substrait_datafusion_err, substrait_err, DFSchema,
- DFSchemaRef,
+ not_impl_datafusion_err, not_impl_err, plan_datafusion_err, plan_err,
+ substrait_datafusion_err, substrait_err, DFSchema, DFSchemaRef,
};
use datafusion::datasource::provider_as_source;
use datafusion::logical_expr::expr::{Exists, InSubquery, Sort};
use datafusion::logical_expr::{
- Aggregate, BinaryExpr, Case, EmptyRelation, Expr, ExprSchemable, LogicalPlan,
- Operator, Projection, SortExpr, TryCast, Values,
+ Aggregate, BinaryExpr, Case, Cast, EmptyRelation, Expr, ExprSchemable, Extension,
+ LogicalPlan, Operator, Projection, SortExpr, Subquery, TryCast, Values,
};
use substrait::proto::aggregate_rel::Grouping;
+use substrait::proto::expression as substrait_expression;
use substrait::proto::expression::subquery::set_predicate::PredicateOp;
use substrait::proto::expression_reference::ExprType;
use url::Url;
+use super::state::SubstraitPlanningState;
Review Comment:
here as well, is this actually used? I don't see any references to it anymore
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #13803:
URL: https://github.com/apache/datafusion/pull/13803
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1888633370
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -1503,370 +1858,416 @@ pub async fn from_substrait_func_args(
/// Convert Substrait AggregateFunction to DataFusion Expr
pub async fn from_substrait_agg_func(
- state: &dyn SubstraitPlanningState,
+ consumer: &impl SubstraitConsumer,
f: &AggregateFunction,
input_schema: &DFSchema,
- extensions: &Extensions,
filter: Option<Box<Expr>>,
order_by: Option<Vec<SortExpr>>,
distinct: bool,
) -> Result<Arc<Expr>> {
- let args =
- from_substrait_func_args(state, &f.arguments, input_schema, extensions).await?;
-
- let Some(function_name) = extensions.functions.get(&f.function_reference) else {
+ let Some(fn_signature) = consumer
+ .get_extensions()
+ .functions
+ .get(&f.function_reference)
+ else {
return plan_err!(
"Aggregate function not registered: function anchor = {:?}",
f.function_reference
);
};
- let function_name = substrait_fun_name(function_name);
- // try udaf first, then built-in aggr fn.
- if let Ok(fun) = state.udaf(function_name) {
- // deal with situation that count(*) got no arguments
- let args = if fun.name() == "count" && args.is_empty() {
- vec![Expr::Literal(ScalarValue::Int64(Some(1)))]
- } else {
- args
- };
-
- Ok(Arc::new(Expr::AggregateFunction(
- expr::AggregateFunction::new_udf(fun, args, distinct, filter, order_by, None),
- )))
- } else {
- not_impl_err!(
+ let fn_name = substrait_fun_name(fn_signature);
+ let udaf = consumer.get_state().udaf(fn_name);
+ let udaf = udaf.map_err(|_| {
+ not_impl_datafusion_err!(
"Aggregate function {} is not supported: function anchor = {:?}",
- function_name,
+ fn_signature,
f.function_reference
)
- }
+ })?;
+
+ let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?;
+
+ // deal with situation that count(*) got no arguments
+ let args = if udaf.name() == "count" && args.is_empty() {
+ vec![Expr::Literal(ScalarValue::Int64(Some(1)))]
+ } else {
+ args
+ };
+
+ Ok(Arc::new(Expr::AggregateFunction(
+ expr::AggregateFunction::new_udf(udaf, args, distinct, filter, order_by, None),
+ )))
}
/// Convert Substrait Rex to DataFusion Expr
-#[async_recursion]
pub async fn from_substrait_rex(
- state: &dyn SubstraitPlanningState,
- e: &Expression,
+ consumer: &impl SubstraitConsumer,
+ expression: &Expression,
input_schema: &DFSchema,
- extensions: &Extensions,
) -> Result<Expr> {
- match &e.rex_type {
- Some(RexType::SingularOrList(s)) => {
- let substrait_expr = s.value.as_ref().unwrap();
- let substrait_list = s.options.as_ref();
- Ok(Expr::InList(InList {
- expr: Box::new(
- from_substrait_rex(state, substrait_expr, input_schema, extensions)
- .await?,
- ),
- list: from_substrait_rex_vec(
- state,
- substrait_list,
- input_schema,
- extensions,
- )
- .await?,
- negated: false,
- }))
- }
- Some(RexType::Selection(field_ref)) => {
- Ok(from_substrait_field_reference(field_ref, input_schema)?)
- }
- Some(RexType::IfThen(if_then)) => {
- // Parse `ifs`
- // If the first element does not have a `then` part, then we can assume it's a base expression
- let mut when_then_expr: Vec<(Box<Expr>, Box<Expr>)> = vec![];
- let mut expr = None;
- for (i, if_expr) in if_then.ifs.iter().enumerate() {
- if i == 0 {
- // Check if the first element is type base expression
- if if_expr.then.is_none() {
- expr = Some(Box::new(
- from_substrait_rex(
- state,
- if_expr.r#if.as_ref().unwrap(),
- input_schema,
- extensions,
- )
- .await?,
- ));
- continue;
- }
- }
- when_then_expr.push((
- Box::new(
- from_substrait_rex(
- state,
- if_expr.r#if.as_ref().unwrap(),
- input_schema,
- extensions,
- )
- .await?,
- ),
- Box::new(
- from_substrait_rex(
- state,
- if_expr.then.as_ref().unwrap(),
- input_schema,
- extensions,
- )
- .await?,
- ),
- ));
+ match &expression.rex_type {
+ Some(t) => match t {
+ RexType::Literal(expr) => consumer.consume_literal(expr).await,
+ RexType::Selection(expr) => {
+ consumer.consume_selection(expr, input_schema).await
+ }
+ RexType::ScalarFunction(expr) => {
+ consumer.consume_scalar_function(expr, input_schema).await
+ }
+ RexType::WindowFunction(expr) => {
+ consumer.consume_window_function(expr, input_schema).await
+ }
+ RexType::IfThen(expr) => consumer.consume_if_then(expr, input_schema).await,
+ RexType::SwitchExpression(expr) => {
+ consumer.consume_switch(expr, input_schema).await
+ }
+ RexType::SingularOrList(expr) => {
+ consumer.consume_singular_or_list(expr, input_schema).await
}
- // Parse `else`
- let else_expr = match &if_then.r#else {
- Some(e) => Some(Box::new(
- from_substrait_rex(state, e, input_schema, extensions).await?,
- )),
- None => None,
- };
- Ok(Expr::Case(Case {
- expr,
- when_then_expr,
- else_expr,
- }))
- }
- Some(RexType::ScalarFunction(f)) => {
- let Some(fn_name) = extensions.functions.get(&f.function_reference) else {
- return plan_err!(
- "Scalar function not found: function reference = {:?}",
- f.function_reference
- );
- };
- let fn_name = substrait_fun_name(fn_name);
- let args =
- from_substrait_func_args(state, &f.arguments, input_schema, extensions)
- .await?;
+ RexType::MultiOrList(expr) => {
+ consumer.consume_multi_or_list(expr, input_schema).await
+ }
- // try to first match the requested function into registered udfs, then built-in ops
- // and finally built-in expressions
- if let Ok(func) = state.udf(fn_name) {
- Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf(
- func.to_owned(),
- args,
- )))
- } else if let Some(op) = name_to_op(fn_name) {
- if f.arguments.len() < 2 {
- return not_impl_err!(
- "Expect at least two arguments for binary operator {op:?}, the provided number of operators is {:?}",
- f.arguments.len()
- );
- }
- // Some expressions are binary in DataFusion but take in a variadic number of args in Substrait.
- // In those cases we iterate through all the arguments, applying the binary expression against them all
- let combined_expr = args
- .into_iter()
- .fold(None, |combined_expr: Option<Expr>, arg: Expr| {
- Some(match combined_expr {
- Some(expr) => Expr::BinaryExpr(BinaryExpr {
- left: Box::new(expr),
- op,
- right: Box::new(arg),
- }),
- None => arg,
- })
- })
- .unwrap();
+ RexType::Cast(expr) => {
+ consumer.consume_cast(expr.as_ref(), input_schema).await
+ }
- Ok(combined_expr)
- } else if let Some(builder) = BuiltinExprBuilder::try_from_name(fn_name) {
- builder.build(state, f, input_schema, extensions).await
- } else {
- not_impl_err!("Unsupported function name: {fn_name:?}")
+ RexType::Subquery(expr) => {
+ consumer.consume_subquery(expr.as_ref(), input_schema).await
}
- }
- Some(RexType::Literal(lit)) => {
- let scalar_value = from_substrait_literal_without_names(lit, extensions)?;
- Ok(Expr::Literal(scalar_value))
- }
- Some(RexType::Cast(cast)) => match cast.as_ref().r#type.as_ref() {
- Some(output_type) => {
- let input_expr = Box::new(
+ RexType::Nested(expr) => consumer.consume_nested(expr, input_schema).await,
+ RexType::Enum(expr) => consumer.consume_enum(expr, input_schema).await,
+ },
+ None => substrait_err!("Expression must set rex_type: {:?}", expression),
+ }
+}
+
+pub async fn from_singular_or_list(
+ consumer: &impl SubstraitConsumer,
+ expr: &SingularOrList,
+ input_schema: &DFSchema,
+) -> Result<Expr> {
+ let substrait_expr = expr.value.as_ref().unwrap();
+ let substrait_list = expr.options.as_ref();
+ Ok(Expr::InList(InList {
+ expr: Box::new(from_substrait_rex(consumer, substrait_expr, input_schema).await?),
+ list: from_substrait_rex_vec(consumer, substrait_list, input_schema).await?,
+ negated: false,
+ }))
+}
+
+pub async fn from_selection(
Review Comment:
should this be renamed to `from_field_reference`? I saw this is used for RexType::Selection, but maybe it'd still make more sense to name it after the parameter type..?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1887815430
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -977,6 +984,43 @@ async fn new_test_grammar() -> Result<()> {
.await
}
+struct MockSubstraitConsumer {
+ extensions: Extensions,
+ state: SessionState,
+}
+
+#[async_trait]
+impl SubstraitConsumer for MockSubstraitConsumer {
+ async fn resolve_table_ref(
+ &self,
+ _: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>> {
+ not_impl_err!("MockSubstraitConsumer does not implement resolve_table_ref")
+ }
+
+ fn get_extensions(&self) -> &Extensions {
+ &self.extensions
+ }
+
+ fn get_state(&self) -> &SessionState {
+ &self.state
+ }
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ let Some(ext_detail) = &rel.detail else {
+ return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
+ };
+ let plan = self
+ .state
+ .serializer_registry()
+ .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
+ Ok(LogicalPlan::Extension(Extension { node: plan }))
+ }
Review Comment:
This is required because the DefaultSubstraitConsumer does not handle extension relations by default.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -202,16 +597,15 @@ fn split_eq_and_noneq_join_predicate_with_nulls_equality(
}
async fn union_rels(
+ consumer: &impl SubstraitConsumer,
rels: &[Rel],
- state: &dyn SubstraitPlanningState,
- extensions: &Extensions,
is_all: bool,
Review Comment:
Starting from here, a lot of the changes look like this. Replacing `state` and `extensions` with a single `consumer` argument that subsumes them both, and then threading the `consumer` into calls.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +971,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
Review Comment:
All of the `from_*` are extractions of existing functionality. It is easier to see this if you hide whitespaces changes in the diff.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +102,400 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::SessionState;
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// use datafusion_substrait::logical_plan::state::SubstraitPlanningState;
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[async_trait]
+/// impl SubstraitConsumer for CustomSubstraitConsumer {
+/// async fn resolve_table_ref(
+/// &self,
+/// table_ref: &TableReference,
+/// ) -> Result<Option<Arc<dyn TableProvider>>> {
+/// self.state.table(table_ref).await
+/// }
+///
+/// fn get_extensions(&self) -> &Extensions {
+/// self.extensions.as_ref()
+/// }
+///
+/// fn get_state(&self) -> &SessionState {
+/// self.state.as_ref()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// async fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// async fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and SessionState is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_state(&self) -> &SessionState;
Review Comment:
I couldn't figure out a clean way to remove these methods as part of this PR that wouldn't make the review very challenging. I've chosen to include these two methods as part of the trait for now.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +971,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
+ consumer: &impl SubstraitConsumer,
+ p: &ProjectRel,
) -> Result<LogicalPlan> {
- let plan: Result<LogicalPlan> = match &rel.rel_type {
- Some(RelType::Project(p)) => {
- if let Some(input) = p.input.as_ref() {
- let mut input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let original_schema = input.schema().clone();
-
- // Ensure that all expressions have a unique display name, so that
- // validate_unique_names does not fail when constructing the project.
- let mut name_tracker = NameTracker::new();
-
- // By default, a Substrait Project emits all inputs fields followed by all expressions.
- // We build the explicit expressions first, and then the input expressions to avoid
- // adding aliases to the explicit expressions (as part of ensuring unique names).
- //
- // This is helpful for plan visualization and tests, because when DataFusion produces
- // Substrait Projects it adds an output mapping that excludes all input columns
- // leaving only explicit expressions.
-
- let mut explicit_exprs: Vec<Expr> = vec![];
- for expr in &p.expressions {
- let e = from_substrait_rex(
- state,
- expr,
- input.clone().schema(),
- extensions,
- )
- .await?;
- // if the expression is WindowFunction, wrap in a Window relation
- if let Expr::WindowFunction(_) = &e {
- // Adding the same expression here and in the project below
- // works because the project's builder uses columnize_expr(..)
- // to transform it into a column reference
- input = input.window(vec![e.clone()])?
- }
- explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
-
- let mut final_exprs: Vec<Expr> = vec![];
- for index in 0..original_schema.fields().len() {
- let e = Expr::Column(Column::from(
- original_schema.qualified_field(index),
- ));
- final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
- final_exprs.append(&mut explicit_exprs);
-
- input.project(final_exprs)?.build()
- } else {
- not_impl_err!("Projection without an input is not supported")
+ if let Some(input) = p.input.as_ref() {
+ let mut input =
+ LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let original_schema = input.schema().clone();
+
+ // Ensure that all expressions have a unique display name, so that
+ // validate_unique_names does not fail when constructing the project.
+ let mut name_tracker = NameTracker::new();
+
+ // By default, a Substrait Project emits all inputs fields followed by all expressions.
+ // We build the explicit expressions first, and then the input expressions to avoid
+ // adding aliases to the explicit expressions (as part of ensuring unique names).
+ //
+ // This is helpful for plan visualization and tests, because when DataFusion produces
+ // Substrait Projects it adds an output mapping that excludes all input columns
+ // leaving only explicit expressions.
+
+ let mut explicit_exprs: Vec<Expr> = vec![];
+ for expr in &p.expressions {
+ let e = from_substrait_rex(consumer, expr, input.clone().schema()).await?;
+ // if the expression is WindowFunction, wrap in a Window relation
+ if let Expr::WindowFunction(_) = &e {
+ // Adding the same expression here and in the project below
+ // works because the project's builder uses columnize_expr(..)
+ // to transform it into a column reference
+ input = input.window(vec![e.clone()])?
}
+ explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Filter(filter)) => {
- if let Some(input) = filter.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- if let Some(condition) = filter.condition.as_ref() {
- let expr =
- from_substrait_rex(state, condition, input.schema(), extensions)
- .await?;
- input.filter(expr)?.build()
- } else {
- not_impl_err!("Filter without an condition is not valid")
- }
- } else {
- not_impl_err!("Filter without an input is not valid")
- }
+
+ let mut final_exprs: Vec<Expr> = vec![];
+ for index in 0..original_schema.fields().len() {
+ let e = Expr::Column(Column::from(original_schema.qualified_field(index)));
+ final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Fetch(fetch)) => {
- if let Some(input) = fetch.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let offset = fetch.offset as usize;
- // -1 means that ALL records should be returned
- let count = if fetch.count == -1 {
- None
- } else {
- Some(fetch.count as usize)
- };
- input.limit(offset, count)?.build()
- } else {
- not_impl_err!("Fetch without an input is not valid")
- }
+ final_exprs.append(&mut explicit_exprs);
+ input.project(final_exprs)?.build()
+ } else {
+ not_impl_err!("Projection without an input is not supported")
+ }
+}
+
+#[async_recursion]
+pub async fn from_filter_rel(
+ consumer: &impl SubstraitConsumer,
+ filter: &FilterRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = filter.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ if let Some(condition) = filter.condition.as_ref() {
+ let expr = from_substrait_rex(consumer, condition, input.schema()).await?;
+ input.filter(expr)?.build()
+ } else {
+ not_impl_err!("Filter without an condition is not valid")
}
- Some(RelType::Sort(sort)) => {
- if let Some(input) = sort.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let sorts =
- from_substrait_sorts(state, &sort.sorts, input.schema(), extensions)
- .await?;
- input.sort(sorts)?.build()
- } else {
- not_impl_err!("Sort without an input is not valid")
- }
+ } else {
+ not_impl_err!("Filter without an input is not valid")
+ }
+}
+
+#[async_recursion]
+pub async fn from_fetch_rel(
+ consumer: &impl SubstraitConsumer,
+ fetch: &FetchRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = fetch.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let offset = fetch.offset as usize;
+ // -1 means that ALL records should be returned
+ let count = if fetch.count == -1 {
+ None
+ } else {
+ Some(fetch.count as usize)
+ };
+ input.limit(offset, count)?.build()
+ } else {
+ not_impl_err!("Fetch without an input is not valid")
+ }
+}
+
+pub async fn from_sort_rel(
+ consumer: &impl SubstraitConsumer,
+ sort: &SortRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = sort.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?;
+ input.sort(sorts)?.build()
+ } else {
+ not_impl_err!("Sort without an input is not valid")
+ }
+}
+
+pub async fn from_aggregate_rel(
+ consumer: &impl SubstraitConsumer,
+ agg: &AggregateRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = agg.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let mut ref_group_exprs = vec![];
+
+ for e in &agg.grouping_expressions {
+ let x = from_substrait_rex(consumer, e, input.schema()).await?;
+ ref_group_exprs.push(x);
}
- Some(RelType::Aggregate(agg)) => {
- if let Some(input) = agg.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let mut ref_group_exprs = vec![];
- for e in &agg.grouping_expressions {
- let x =
- from_substrait_rex(state, e, input.schema(), extensions).await?;
- ref_group_exprs.push(x);
+ let mut group_exprs = vec![];
+ let mut aggr_exprs = vec![];
+
+ match agg.groupings.len() {
+ 1 => {
+ group_exprs.extend_from_slice(
+ &from_substrait_grouping(
+ consumer,
+ &agg.groupings[0],
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?,
+ );
+ }
+ _ => {
+ let mut grouping_sets = vec![];
+ for grouping in &agg.groupings {
+ let grouping_set = from_substrait_grouping(
+ consumer,
+ grouping,
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?;
+ grouping_sets.push(grouping_set);
}
+ // Single-element grouping expression of type Expr::GroupingSet.
+ // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+ // parsed by the producer and consumer, since Substrait does not have a type dedicated
+ // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+ group_exprs
+ .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+ }
+ };
- let mut group_exprs = vec![];
- let mut aggr_exprs = vec![];
-
- match agg.groupings.len() {
- 1 => {
- group_exprs.extend_from_slice(
- &from_substrait_grouping(
- state,
- &agg.groupings[0],
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?,
- );
- }
- _ => {
- let mut grouping_sets = vec![];
- for grouping in &agg.groupings {
- let grouping_set = from_substrait_grouping(
- state,
- grouping,
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?;
- grouping_sets.push(grouping_set);
+ for m in &agg.measures {
+ let filter = match &m.filter {
+ Some(fil) => Some(Box::new(
+ from_substrait_rex(consumer, fil, input.schema()).await?,
+ )),
+ None => None,
+ };
+ let agg_func = match &m.measure {
+ Some(f) => {
+ let distinct = match f.invocation {
+ _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+ true
}
- // Single-element grouping expression of type Expr::GroupingSet.
- // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
- // parsed by the producer and consumer, since Substrait does not have a type dedicated
- // to ROLLUP. Only vector of Groupings (grouping sets) is available.
- group_exprs.push(Expr::GroupingSet(GroupingSet::GroupingSets(
- grouping_sets,
- )));
- }
- };
-
- for m in &agg.measures {
- let filter = match &m.filter {
- Some(fil) => Some(Box::new(
- from_substrait_rex(state, fil, input.schema(), extensions)
- .await?,
- )),
- None => None,
+ _ if f.invocation == AggregationInvocation::All as i32 => false,
+ _ => false,
};
- let agg_func = match &m.measure {
- Some(f) => {
- let distinct = match f.invocation {
- _ if f.invocation
- == AggregationInvocation::Distinct as i32 =>
- {
- true
- }
- _ if f.invocation
- == AggregationInvocation::All as i32 =>
- {
- false
- }
- _ => false,
- };
- let order_by = if !f.sorts.is_empty() {
- Some(
- from_substrait_sorts(
- state,
- &f.sorts,
- input.schema(),
- extensions,
- )
- .await?,
- )
- } else {
- None
- };
-
- from_substrait_agg_func(
- state,
- f,
- input.schema(),
- extensions,
- filter,
- order_by,
- distinct,
- )
- .await
- }
- None => not_impl_err!(
- "Aggregate without aggregate function is not supported"
- ),
+ let order_by = if !f.sorts.is_empty() {
+ Some(
+ from_substrait_sorts(consumer, &f.sorts, input.schema())
+ .await?,
+ )
+ } else {
+ None
};
- aggr_exprs.push(agg_func?.as_ref().clone());
- }
- input.aggregate(group_exprs, aggr_exprs)?.build()
- } else {
- not_impl_err!("Aggregate without an input is not valid")
- }
- }
- Some(RelType::Join(join)) => {
- if join.post_join_filter.is_some() {
- return not_impl_err!(
- "JoinRel with post_join_filter is not yet supported"
- );
- }
- let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
-
- let join_type = from_substrait_jointype(join.r#type)?;
- // The join condition expression needs full input schema and not the output schema from join since we lose columns from
- // certain join types such as semi and anti joins
- let in_join_schema = left.schema().join(right.schema())?;
-
- // If join expression exists, parse the `on` condition expression, build join and return
- // Otherwise, build join with only the filter, without join keys
- match &join.expression.as_ref() {
- Some(expr) => {
- let on = from_substrait_rex(state, expr, &in_join_schema, extensions)
- .await?;
- // The join expression can contain both equal and non-equal ops.
- // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
- // So we extract each part as follows:
- // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
- // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
- let (join_ons, nulls_equal_nulls, join_filter) =
- split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
- let (left_cols, right_cols): (Vec<_>, Vec<_>) =
- itertools::multiunzip(join_ons);
- left.join_detailed(
- right.build()?,
- join_type,
- (left_cols, right_cols),
- join_filter,
- nulls_equal_nulls,
- )?
- .build()
+ from_substrait_agg_func(
+ consumer,
+ f,
+ input.schema(),
+ filter,
+ order_by,
+ distinct,
+ )
+ .await
}
None => {
- let on: Vec<String> = vec![];
- left.join_detailed(
- right.build()?,
- join_type,
- (on.clone(), on),
- None,
- false,
- )?
- .build()
+ not_impl_err!("Aggregate without aggregate function is not supported")
}
- }
+ };
+ aggr_exprs.push(agg_func?.as_ref().clone());
}
- Some(RelType::Cross(cross)) => {
- let left = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
- left.cross_join(right.build()?)?.build()
+ input.aggregate(group_exprs, aggr_exprs)?.build()
+ } else {
+ not_impl_err!("Aggregate without an input is not valid")
+ }
+}
+
+pub async fn from_join_rel(
+ consumer: &impl SubstraitConsumer,
+ join: &JoinRel,
+) -> Result<LogicalPlan> {
+ if join.post_join_filter.is_some() {
+ return not_impl_err!("JoinRel with post_join_filter is not yet supported");
+ }
+
+ let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+
+ let join_type = from_substrait_jointype(join.r#type)?;
+ // The join condition expression needs full input schema and not the output schema from join since we lose columns from
+ // certain join types such as semi and anti joins
+ let in_join_schema = left.schema().join(right.schema())?;
+
+ // If join expression exists, parse the `on` condition expression, build join and return
+ // Otherwise, build join with only the filter, without join keys
+ match &join.expression.as_ref() {
+ Some(expr) => {
+ let on = from_substrait_rex(consumer, expr, &in_join_schema).await?;
+ // The join expression can contain both equal and non-equal ops.
+ // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
+ // So we extract each part as follows:
+ // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
+ // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
+ let (join_ons, nulls_equal_nulls, join_filter) =
+ split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) =
+ itertools::multiunzip(join_ons);
+ left.join_detailed(
+ right.build()?,
+ join_type,
+ (left_cols, right_cols),
+ join_filter,
+ nulls_equal_nulls,
+ )?
+ .build()
}
- Some(RelType::Read(read)) => {
- async fn read_with_schema(
- state: &dyn SubstraitPlanningState,
- table_ref: TableReference,
- schema: DFSchema,
- projection: &Option<MaskExpression>,
- ) -> Result<LogicalPlan> {
- let schema = schema.replace_qualifier(table_ref.clone());
-
- let plan = {
- let provider = match state.table(&table_ref).await? {
- Some(ref provider) => Arc::clone(provider),
- _ => return plan_err!("No table named '{table_ref}'"),
- };
+ None => {
+ let on: Vec<String> = vec![];
+ left.join_detailed(right.build()?, join_type, (on.clone(), on), None, false)?
+ .build()
+ }
+ }
+}
- LogicalPlanBuilder::scan(
- table_ref,
- provider_as_source(Arc::clone(&provider)),
- None,
- )?
- .build()?
- };
+pub async fn from_cross_rel(
+ consumer: &impl SubstraitConsumer,
+ cross: &CrossRel,
+) -> Result<LogicalPlan> {
+ let left = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+ left.cross_join(right.build()?)?.build()
+}
- ensure_schema_compatability(plan.schema(), schema.clone())?;
+#[allow(deprecated)]
+pub async fn from_read_rel(
+ consumer: &impl SubstraitConsumer,
+ read: &ReadRel,
+) -> Result<LogicalPlan> {
+ async fn read_with_schema(
+ consumer: &impl SubstraitConsumer,
+ table_ref: TableReference,
+ schema: DFSchema,
+ projection: &Option<MaskExpression>,
+ ) -> Result<LogicalPlan> {
+ let schema = schema.replace_qualifier(table_ref.clone());
+
+ let plan = {
+ let provider = match consumer.resolve_table_ref(&table_ref).await? {
+ Some(ref provider) => Arc::clone(provider),
+ _ => return plan_err!("No table named '{table_ref}'"),
+ };
- let schema = apply_masking(schema, projection)?;
+ LogicalPlanBuilder::scan(
+ table_ref,
+ provider_as_source(Arc::clone(&provider)),
+ None,
+ )?
+ .build()?
+ };
- apply_projection(plan, schema)
- }
+ ensure_schema_compatability(plan.schema(), schema.clone())?;
- let named_struct = read.base_schema.as_ref().ok_or_else(|| {
- substrait_datafusion_err!("No base schema provided for Read Relation")
- })?;
+ let schema = apply_masking(schema, projection)?;
- let substrait_schema = from_substrait_named_struct(named_struct, extensions)?;
+ apply_projection(plan, schema)
+ }
- match &read.as_ref().read_type {
- Some(ReadType::NamedTable(nt)) => {
- let table_reference = match nt.names.len() {
- 0 => {
- return plan_err!("No table name found in NamedTable");
- }
- 1 => TableReference::Bare {
- table: nt.names[0].clone().into(),
- },
- 2 => TableReference::Partial {
- schema: nt.names[0].clone().into(),
- table: nt.names[1].clone().into(),
- },
- _ => TableReference::Full {
- catalog: nt.names[0].clone().into(),
- schema: nt.names[1].clone().into(),
- table: nt.names[2].clone().into(),
- },
- };
+ let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+ substrait_datafusion_err!("No base schema provided for Read Relation")
+ })?;
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
+ let substrait_schema = from_substrait_named_struct(consumer, named_struct)?;
+
+ match &read.read_type {
+ Some(ReadType::NamedTable(nt)) => {
+ let table_reference = match nt.names.len() {
+ 0 => {
+ return plan_err!("No table name found in NamedTable");
}
- Some(ReadType::VirtualTable(vt)) => {
- if vt.values.is_empty() {
- return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: DFSchemaRef::new(substrait_schema),
- }));
+ 1 => TableReference::Bare {
+ table: nt.names[0].clone().into(),
+ },
+ 2 => TableReference::Partial {
+ schema: nt.names[0].clone().into(),
+ table: nt.names[1].clone().into(),
+ },
+ _ => TableReference::Full {
+ catalog: nt.names[0].clone().into(),
+ schema: nt.names[1].clone().into(),
+ table: nt.names[2].clone().into(),
+ },
+ };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ Some(ReadType::VirtualTable(vt)) => {
+ if vt.values.is_empty() {
+ return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row: false,
+ schema: DFSchemaRef::new(substrait_schema),
+ }));
+ }
+
+ let values = vt
+ .values
+ .iter()
+ .map(|row| {
+ let mut name_idx = 0;
+ let lits = row
+ .fields
+ .iter()
+ .map(|lit| {
+ name_idx += 1; // top-level names are provided through schema
+ Ok(Expr::Literal(from_substrait_literal(
+ consumer,
+ lit,
+ &named_struct.names,
+ &mut name_idx,
+ )?))
+ })
+ .collect::<Result<_>>()?;
+ if name_idx != named_struct.names.len() {
+ return substrait_err!(
+ "Names list must match exactly to nested schema, but found {} uses for {} names",
+ name_idx,
+ named_struct.names.len()
+ );
}
+ Ok(lits)
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(LogicalPlan::Values(Values {
+ schema: DFSchemaRef::new(substrait_schema),
+ values,
+ }))
+ }
+ Some(ReadType::LocalFiles(lf)) => {
+ fn extract_filename(name: &str) -> Option<String> {
+ let corrected_url =
+ if name.starts_with("file://") && !name.starts_with("file:///") {
+ name.replacen("file://", "file:///", 1)
+ } else {
+ name.to_string()
+ };
+
+ Url::parse(&corrected_url).ok().and_then(|url| {
+ let path = url.path();
+ std::path::Path::new(path)
+ .file_name()
+ .map(|filename| filename.to_string_lossy().to_string())
+ })
+ }
+
+ // we could use the file name to check the original table provider
+ // TODO: currently does not support multiple local files
+ let filename: Option<String> =
+ lf.items.first().and_then(|x| match x.path_type.as_ref() {
+ Some(UriFile(name)) => extract_filename(name),
+ _ => None,
+ });
+
+ if lf.items.len() > 1 || filename.is_none() {
+ return not_impl_err!("Only single file reads are supported");
+ }
+ let name = filename.unwrap();
+ // directly use unwrap here since we could determine it is a valid one
+ let table_reference = TableReference::Bare { table: name.into() };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ _ => {
+ not_impl_err!("Unsupported ReadType: {:?}", read.read_type)
+ }
+ }
+}
+
+pub async fn from_set_rel(
+ consumer: &impl SubstraitConsumer,
+ set: &SetRel,
+) -> Result<LogicalPlan> {
+ if set.inputs.len() < 2 {
+ substrait_err!("Set operation requires at least two inputs")
+ } else {
+ match set.op() {
+ SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
+ SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
+ SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect(
+ from_substrait_rel(consumer, &set.inputs[0]).await?,
+ union_rels(consumer, &set.inputs[1..], true).await?,
+ false,
+ ),
+ SetOp::IntersectionMultiset => {
+ intersect_rels(consumer, &set.inputs, false).await
+ }
+ SetOp::IntersectionMultisetAll => {
+ intersect_rels(consumer, &set.inputs, true).await
+ }
+ SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
+ SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+ set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+ }
+ }
+}
- let values = vt
- .values
- .iter()
- .map(|row| {
- let mut name_idx = 0;
- let lits = row
- .fields
- .iter()
- .map(|lit| {
- name_idx += 1; // top-level names are provided through schema
- Ok(Expr::Literal(from_substrait_literal(
- lit,
- extensions,
- &named_struct.names,
- &mut name_idx,
- )?))
- })
- .collect::<Result<_>>()?;
- if name_idx != named_struct.names.len() {
- return substrait_err!(
- "Names list must match exactly to nested schema, but found {} uses for {} names",
- name_idx,
- named_struct.names.len()
- );
- }
- Ok(lits)
- })
- .collect::<Result<_>>()?;
+pub async fn from_extension_leaf_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_leaf_rel: &ExtensionLeafRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_leaf(extension_leaf_rel).await
+}
- Ok(LogicalPlan::Values(Values {
- schema: DFSchemaRef::new(substrait_schema),
- values,
- }))
- }
- Some(ReadType::LocalFiles(lf)) => {
- fn extract_filename(name: &str) -> Option<String> {
- let corrected_url = if name.starts_with("file://")
- && !name.starts_with("file:///")
- {
- name.replacen("file://", "file:///", 1)
- } else {
- name.to_string()
- };
+pub async fn from_extension_single_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_single_rel: &ExtensionSingleRel,
+) -> Result<LogicalPlan> {
+ consumer
+ .consume_extension_single(extension_single_rel)
+ .await
+}
- Url::parse(&corrected_url).ok().and_then(|url| {
- let path = url.path();
- std::path::Path::new(path)
- .file_name()
- .map(|filename| filename.to_string_lossy().to_string())
- })
- }
+pub async fn from_extension_multi_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_multi_rel: &ExtensionMultiRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_multi(extension_multi_rel).await
+}
- // we could use the file name to check the original table provider
- // TODO: currently does not support multiple local files
- let filename: Option<String> =
- lf.items.first().and_then(|x| match x.path_type.as_ref() {
- Some(UriFile(name)) => extract_filename(name),
- _ => None,
- });
+pub async fn from_exchange_rel(
+ consumer: &impl SubstraitConsumer,
+ exchange: &ExchangeRel,
+) -> Result<LogicalPlan> {
+ let Some(input) = exchange.input.as_ref() else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+ let input = Arc::new(from_substrait_rel(consumer, input).await?);
- if lf.items.len() > 1 || filename.is_none() {
- return not_impl_err!("Only single file reads are supported");
- }
- let name = filename.unwrap();
- // directly use unwrap here since we could determine it is a valid one
- let table_reference = TableReference::Bare { table: name.into() };
-
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
- }
- _ => {
- not_impl_err!("Unsupported ReadType: {:?}", &read.as_ref().read_type)
- }
- }
- }
- Some(RelType::Set(set)) => match set_rel::SetOp::try_from(set.op) {
- Ok(set_op) => {
- if set.inputs.len() < 2 {
- substrait_err!("Set operation requires at least two inputs")
- } else {
- match set_op {
- set_rel::SetOp::UnionAll => {
- union_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::UnionDistinct => {
- union_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionPrimary => {
- LogicalPlanBuilder::intersect(
- from_substrait_rel(state, &set.inputs[0], extensions)
- .await?,
- union_rels(&set.inputs[1..], state, extensions, true)
- .await?,
- false,
- )
- }
- set_rel::SetOp::IntersectionMultiset => {
- intersect_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionMultisetAll => {
- intersect_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::MinusPrimary => {
- except_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::MinusPrimaryAll => {
- except_rels(&set.inputs, state, extensions, true).await
- }
- _ => not_impl_err!("Unsupported set operator: {set_op:?}"),
- }
- }
+ let Some(exchange_kind) = &exchange.exchange_kind else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+
+ // ref: https://substrait.io/relations/physical_relations/#exchange-types
+ let partitioning_scheme = match exchange_kind {
+ ExchangeKind::ScatterByFields(scatter_fields) => {
+ let mut partition_columns = vec![];
+ let input_schema = input.schema();
+ for field_ref in &scatter_fields.fields {
+ let column = from_substrait_field_reference(field_ref, input_schema)?;
+ partition_columns.push(column);
}
- Err(e) => not_impl_err!("Invalid set operation type {}: {e}", set.op),
- },
- Some(RelType::ExtensionLeaf(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
- }
- Some(RelType::ExtensionSingle(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- let Some(input_rel) = &extension.input else {
- return substrait_err!(
- "ExtensionSingleRel doesn't contains input rel. Try use ExtensionLeafRel instead"
- );
- };
- let input_plan = from_substrait_rel(state, input_rel, extensions).await?;
- let plan =
- plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
+ Partitioning::Hash(partition_columns, exchange.partition_count as usize)
}
- Some(RelType::ExtensionMulti(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- let mut inputs = Vec::with_capacity(extension.inputs.len());
- for input in &extension.inputs {
- let input_plan = from_substrait_rel(state, input, extensions).await?;
- inputs.push(input_plan);
- }
- let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
+ ExchangeKind::RoundRobin(_) => {
+ Partitioning::RoundRobinBatch(exchange.partition_count as usize)
}
Review Comment:
These Extension relation handling code has been fully removed as that must be provided by users.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +102,400 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::SessionState;
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// use datafusion_substrait::logical_plan::state::SubstraitPlanningState;
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[async_trait]
+/// impl SubstraitConsumer for CustomSubstraitConsumer {
+/// async fn resolve_table_ref(
+/// &self,
+/// table_ref: &TableReference,
+/// ) -> Result<Option<Arc<dyn TableProvider>>> {
+/// self.state.table(table_ref).await
+/// }
+///
+/// fn get_extensions(&self) -> &Extensions {
+/// self.extensions.as_ref()
+/// }
+///
+/// fn get_state(&self) -> &SessionState {
+/// self.state.as_ref()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// async fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// async fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and SessionState is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_state(&self) -> &SessionState;
+
+ // Relation Methods
+ // There is one method per Substrait relation to allow for easy overriding of consumer behaviour.
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_read(&self, rel: &ReadRel) -> Result<LogicalPlan> {
+ from_read_rel(self, rel).await
+ }
+
+ async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+ from_filter_rel(self, rel).await
+ }
+
+ async fn consume_fetch(&self, rel: &FetchRel) -> Result<LogicalPlan> {
+ from_fetch_rel(self, rel).await
+ }
+
+ async fn consume_aggregate(&self, rel: &AggregateRel) -> Result<LogicalPlan> {
+ from_aggregate_rel(self, rel).await
+ }
+
+ async fn consume_sort(&self, rel: &SortRel) -> Result<LogicalPlan> {
+ from_sort_rel(self, rel).await
+ }
+
+ async fn consume_join(&self, rel: &JoinRel) -> Result<LogicalPlan> {
+ from_join_rel(self, rel).await
+ }
+
+ async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+ from_project_rel(self, rel).await
+ }
+
+ async fn consume_set(&self, rel: &SetRel) -> Result<LogicalPlan> {
+ from_set_rel(self, rel).await
+ }
+
+ async fn consume_cross(&self, rel: &CrossRel) -> Result<LogicalPlan> {
+ from_cross_rel(self, rel).await
+ }
+
+ async fn consume_consistent_partition_window(
+ &self,
+ _rel: &ConsistentPartitionWindowRel,
+ ) -> Result<LogicalPlan> {
+ not_impl_err!("Consistent Partition Window Rel not supported")
+ }
+
+ async fn consume_exchange(&self, rel: &ExchangeRel) -> Result<LogicalPlan> {
+ from_exchange_rel(self, rel).await
+ }
+
+ // Expression Methods
+ // There is one method per Substrait expression to allow for easy overriding of consumer behaviour
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_literal(&self, expr: &Literal) -> Result<Expr> {
+ from_literal(self, expr).await
+ }
+
+ async fn consume_selection(
+ &self,
+ expr: &FieldReference,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_selection(self, expr, input_schema).await
+ }
+
+ async fn consume_scalar_function(
+ &self,
+ expr: &ScalarFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_scalar_function(self, expr, input_schema).await
+ }
+
+ async fn consume_window_function(
+ &self,
+ expr: &WindowFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_window_function(self, expr, input_schema).await
+ }
+
+ async fn consume_if_then(
+ &self,
+ expr: &IfThen,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_if_then(self, expr, input_schema).await
+ }
+
+ async fn consume_switch(
+ &self,
+ _expr: &SwitchExpression,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Switch expression not supported")
+ }
+
+ async fn consume_singular_or_list(
+ &self,
+ expr: &SingularOrList,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_singular_or_list(self, expr, input_schema).await
+ }
+
+ async fn consume_multi_or_list(
+ &self,
+ _expr: &MultiOrList,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Multi Or List expression not supported")
+ }
+
+ async fn consume_cast(
+ &self,
+ expr: &substrait_expression::Cast,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_cast(self, expr, input_schema).await
+ }
+
+ async fn consume_subquery(
+ &self,
+ expr: &substrait_expression::Subquery,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_subquery(self, expr, input_schema).await
+ }
+
+ async fn consume_nested(
+ &self,
+ _expr: &Nested,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Nested expression not supported")
+ }
+
+ async fn consume_enum(&self, _expr: &Enum, _input_schema: &DFSchema) -> Result<Expr> {
+ not_impl_err!("Enum expression not supported")
+ }
+
+ // User-Defined Functionality
+
+ // The details of extension relations, and how to handle them, are fully up to users to specify.
+ // The following methods allow users to customize the consumer behaviour
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionLeafRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionLeafRel")
+ }
+
+ async fn consume_extension_single(
+ &self,
+ rel: &ExtensionSingleRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionSingleRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionSingleRel")
+ }
+
+ async fn consume_extension_multi(
+ &self,
+ rel: &ExtensionMultiRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionMultiRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionMultiRel")
+ }
+
+ // Users can bring their own types to Substrait which require custom handling
+
+ fn consume_user_defined_type(
+ &self,
+ user_defined_type: &r#type::UserDefined
+ ) -> Result<DataType> {
+ substrait_err!("Missing handler for user-defined type: {}", user_defined_type.type_reference)
+ }
+
+ fn consume_user_defined_literal(
+ &self,
+ user_defined_literal: &proto::expression::literal::UserDefined
+ ) -> Result<ScalarValue> {
+ substrait_err!("Missing handler for user-defined literals {}", user_defined_literal.type_reference)
+ }
+}
+
+/// Convert Substrait Rel to DataFusion DataFrame
+#[async_recursion]
+pub async fn from_substrait_rel(
+ consumer: &impl SubstraitConsumer,
+ relation: &Rel,
+) -> Result<LogicalPlan> {
+ let plan: Result<LogicalPlan> = match &relation.rel_type {
+ Some(rel_type) => match rel_type {
+ RelType::Read(rel) => consumer.consume_read(rel).await,
+ RelType::Filter(rel) => consumer.consume_filter(rel).await,
+ RelType::Fetch(rel) => consumer.consume_fetch(rel).await,
+ RelType::Aggregate(rel) => consumer.consume_aggregate(rel).await,
+ RelType::Sort(rel) => consumer.consume_sort(rel).await,
+ RelType::Join(rel) => consumer.consume_join(rel).await,
+ RelType::Project(rel) => consumer.consume_project(rel).await,
+ RelType::Set(rel) => consumer.consume_set(rel).await,
+ RelType::ExtensionSingle(rel) => consumer.consume_extension_single(rel).await,
+ RelType::ExtensionMulti(rel) => consumer.consume_extension_multi(rel).await,
+ RelType::ExtensionLeaf(rel) => consumer.consume_extension_leaf(rel).await,
+ RelType::Cross(rel) => consumer.consume_cross(rel).await,
+ RelType::Window(rel) => {
+ consumer.consume_consistent_partition_window(rel).await
+ }
+ RelType::Exchange(rel) => consumer.consume_exchange(rel).await,
+ rt => not_impl_err!("{rt:?} rel not supported yet"),
+ },
+ None => return substrait_err!("rel must set rel_type"),
+ };
+ apply_emit_kind(retrieve_rel_common(relation), plan?)
+}
+
+/// Can be used to consume standard Substrait without user-defined extensions
+pub struct DefaultSubstraitConsumer {
+ extensions: Arc<Extensions>,
+ state: Arc<SessionState>,
Review Comment:
I'm not sure if Arc is actually needed here, I just got tired of fighting the borrow-checker. I'm open to suggestions, as all we really need is to be able to get references to Extensions and SessionState.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1889343110
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -977,6 +984,43 @@ async fn new_test_grammar() -> Result<()> {
.await
}
+struct MockSubstraitConsumer {
+ extensions: Extensions,
+ state: SessionState,
+}
+
+#[async_trait]
+impl SubstraitConsumer for MockSubstraitConsumer {
+ async fn resolve_table_ref(
+ &self,
+ _: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>> {
+ not_impl_err!("MockSubstraitConsumer does not implement resolve_table_ref")
+ }
+
+ fn get_extensions(&self) -> &Extensions {
+ &self.extensions
+ }
+
+ fn get_state(&self) -> &SessionState {
+ &self.state
+ }
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ let Some(ext_detail) = &rel.detail else {
+ return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
+ };
+ let plan = self
+ .state
+ .serializer_registry()
+ .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
Review Comment:
> As a side note, this may sound hypocritical, given that I made it even worse https://github.com/apache/datafusion/pull/13772
I don't think it's hypocritical, you were aiming for consistency and I'm just being much more aggressive with my changes and breaking the API.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "ccciudatu (via GitHub)" <gi...@apache.org>.
ccciudatu commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1888961320
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -977,6 +984,43 @@ async fn new_test_grammar() -> Result<()> {
.await
}
+struct MockSubstraitConsumer {
+ extensions: Extensions,
+ state: SessionState,
+}
+
+#[async_trait]
+impl SubstraitConsumer for MockSubstraitConsumer {
+ async fn resolve_table_ref(
+ &self,
+ _: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>> {
+ not_impl_err!("MockSubstraitConsumer does not implement resolve_table_ref")
+ }
+
+ fn get_extensions(&self) -> &Extensions {
+ &self.extensions
+ }
+
+ fn get_state(&self) -> &SessionState {
+ &self.state
+ }
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ let Some(ext_detail) = &rel.detail else {
+ return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
+ };
+ let plan = self
+ .state
+ .serializer_registry()
+ .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
Review Comment:
One drawback of the current implementation is that it's abusing the `type_url` field in `proto::Any` (which is quite [clearly specified](https://protobuf.dev/reference/java/api-docs/com/google/protobuf/AnyOrBuilder.html#getTypeUrl--). Working with Substrait objects directly would eliminate the need for hacks like encoding node names as Any::type_url.
As a side note, this may sound hypocritical, given that I made it even worse [with my PR](https://github.com/apache/datafusion/pull/13772) by setting it to table names as well. But I only did it for consistency. :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1890511596
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -21,23 +21,24 @@ use datafusion::arrow::array::{GenericListArray, MapArray};
use datafusion::arrow::datatypes::{
DataType, Field, FieldRef, Fields, IntervalUnit, Schema, TimeUnit,
};
-use datafusion::common::plan_err;
use datafusion::common::{
- not_impl_err, plan_datafusion_err, substrait_datafusion_err, substrait_err, DFSchema,
- DFSchemaRef,
+ not_impl_datafusion_err, not_impl_err, plan_datafusion_err, plan_err,
+ substrait_datafusion_err, substrait_err, DFSchema, DFSchemaRef,
};
use datafusion::datasource::provider_as_source;
use datafusion::logical_expr::expr::{Exists, InSubquery, Sort};
use datafusion::logical_expr::{
- Aggregate, BinaryExpr, Case, EmptyRelation, Expr, ExprSchemable, LogicalPlan,
- Operator, Projection, SortExpr, TryCast, Values,
+ Aggregate, BinaryExpr, Case, Cast, EmptyRelation, Expr, ExprSchemable, Extension,
+ LogicalPlan, Operator, Projection, SortExpr, Subquery, TryCast, Values,
};
use substrait::proto::aggregate_rel::Grouping;
+use substrait::proto::expression as substrait_expression;
use substrait::proto::expression::subquery::set_predicate::PredicateOp;
use substrait::proto::expression_reference::ExprType;
use url::Url;
+use super::state::SubstraitPlanningState;
Review Comment:
Ah I see, yep I'd be for inlining it here, but not a critical thing then!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894403299
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +102,400 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::SessionState;
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// use datafusion_substrait::logical_plan::state::SubstraitPlanningState;
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[async_trait]
+/// impl SubstraitConsumer for CustomSubstraitConsumer {
+/// async fn resolve_table_ref(
+/// &self,
+/// table_ref: &TableReference,
+/// ) -> Result<Option<Arc<dyn TableProvider>>> {
+/// self.state.table(table_ref).await
+/// }
+///
+/// fn get_extensions(&self) -> &Extensions {
+/// self.extensions.as_ref()
+/// }
+///
+/// fn get_state(&self) -> &SessionState {
+/// self.state.as_ref()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// async fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// async fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and SessionState is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_state(&self) -> &SessionState;
+
+ // Relation Methods
+ // There is one method per Substrait relation to allow for easy overriding of consumer behaviour.
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_read(&self, rel: &ReadRel) -> Result<LogicalPlan> {
+ from_read_rel(self, rel).await
+ }
+
+ async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+ from_filter_rel(self, rel).await
+ }
+
+ async fn consume_fetch(&self, rel: &FetchRel) -> Result<LogicalPlan> {
+ from_fetch_rel(self, rel).await
+ }
+
+ async fn consume_aggregate(&self, rel: &AggregateRel) -> Result<LogicalPlan> {
+ from_aggregate_rel(self, rel).await
+ }
+
+ async fn consume_sort(&self, rel: &SortRel) -> Result<LogicalPlan> {
+ from_sort_rel(self, rel).await
+ }
+
+ async fn consume_join(&self, rel: &JoinRel) -> Result<LogicalPlan> {
+ from_join_rel(self, rel).await
+ }
+
+ async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+ from_project_rel(self, rel).await
+ }
+
+ async fn consume_set(&self, rel: &SetRel) -> Result<LogicalPlan> {
+ from_set_rel(self, rel).await
+ }
+
+ async fn consume_cross(&self, rel: &CrossRel) -> Result<LogicalPlan> {
+ from_cross_rel(self, rel).await
+ }
+
+ async fn consume_consistent_partition_window(
+ &self,
+ _rel: &ConsistentPartitionWindowRel,
+ ) -> Result<LogicalPlan> {
+ not_impl_err!("Consistent Partition Window Rel not supported")
+ }
+
+ async fn consume_exchange(&self, rel: &ExchangeRel) -> Result<LogicalPlan> {
+ from_exchange_rel(self, rel).await
+ }
+
+ // Expression Methods
+ // There is one method per Substrait expression to allow for easy overriding of consumer behaviour
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_literal(&self, expr: &Literal) -> Result<Expr> {
+ from_literal(self, expr).await
+ }
+
+ async fn consume_selection(
+ &self,
+ expr: &FieldReference,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_selection(self, expr, input_schema).await
+ }
+
+ async fn consume_scalar_function(
+ &self,
+ expr: &ScalarFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_scalar_function(self, expr, input_schema).await
+ }
+
+ async fn consume_window_function(
+ &self,
+ expr: &WindowFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_window_function(self, expr, input_schema).await
+ }
+
+ async fn consume_if_then(
+ &self,
+ expr: &IfThen,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_if_then(self, expr, input_schema).await
+ }
+
+ async fn consume_switch(
+ &self,
+ _expr: &SwitchExpression,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Switch expression not supported")
+ }
+
+ async fn consume_singular_or_list(
+ &self,
+ expr: &SingularOrList,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_singular_or_list(self, expr, input_schema).await
+ }
+
+ async fn consume_multi_or_list(
+ &self,
+ _expr: &MultiOrList,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Multi Or List expression not supported")
+ }
+
+ async fn consume_cast(
+ &self,
+ expr: &substrait_expression::Cast,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_cast(self, expr, input_schema).await
+ }
+
+ async fn consume_subquery(
+ &self,
+ expr: &substrait_expression::Subquery,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_subquery(self, expr, input_schema).await
+ }
+
+ async fn consume_nested(
+ &self,
+ _expr: &Nested,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Nested expression not supported")
+ }
+
+ async fn consume_enum(&self, _expr: &Enum, _input_schema: &DFSchema) -> Result<Expr> {
+ not_impl_err!("Enum expression not supported")
+ }
+
+ // User-Defined Functionality
+
+ // The details of extension relations, and how to handle them, are fully up to users to specify.
+ // The following methods allow users to customize the consumer behaviour
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionLeafRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionLeafRel")
+ }
+
+ async fn consume_extension_single(
+ &self,
+ rel: &ExtensionSingleRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionSingleRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionSingleRel")
+ }
+
+ async fn consume_extension_multi(
+ &self,
+ rel: &ExtensionMultiRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionMultiRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionMultiRel")
+ }
+
+ // Users can bring their own types to Substrait which require custom handling
+
+ fn consume_user_defined_type(
+ &self,
+ user_defined_type: &r#type::UserDefined
+ ) -> Result<DataType> {
+ substrait_err!("Missing handler for user-defined type: {}", user_defined_type.type_reference)
+ }
+
+ fn consume_user_defined_literal(
+ &self,
+ user_defined_literal: &proto::expression::literal::UserDefined
+ ) -> Result<ScalarValue> {
+ substrait_err!("Missing handler for user-defined literals {}", user_defined_literal.type_reference)
+ }
+}
+
+/// Convert Substrait Rel to DataFusion DataFrame
+#[async_recursion]
+pub async fn from_substrait_rel(
+ consumer: &impl SubstraitConsumer,
+ relation: &Rel,
+) -> Result<LogicalPlan> {
+ let plan: Result<LogicalPlan> = match &relation.rel_type {
+ Some(rel_type) => match rel_type {
+ RelType::Read(rel) => consumer.consume_read(rel).await,
+ RelType::Filter(rel) => consumer.consume_filter(rel).await,
+ RelType::Fetch(rel) => consumer.consume_fetch(rel).await,
+ RelType::Aggregate(rel) => consumer.consume_aggregate(rel).await,
+ RelType::Sort(rel) => consumer.consume_sort(rel).await,
+ RelType::Join(rel) => consumer.consume_join(rel).await,
+ RelType::Project(rel) => consumer.consume_project(rel).await,
+ RelType::Set(rel) => consumer.consume_set(rel).await,
+ RelType::ExtensionSingle(rel) => consumer.consume_extension_single(rel).await,
+ RelType::ExtensionMulti(rel) => consumer.consume_extension_multi(rel).await,
+ RelType::ExtensionLeaf(rel) => consumer.consume_extension_leaf(rel).await,
+ RelType::Cross(rel) => consumer.consume_cross(rel).await,
+ RelType::Window(rel) => {
+ consumer.consume_consistent_partition_window(rel).await
+ }
+ RelType::Exchange(rel) => consumer.consume_exchange(rel).await,
+ rt => not_impl_err!("{rt:?} rel not supported yet"),
+ },
+ None => return substrait_err!("rel must set rel_type"),
+ };
+ apply_emit_kind(retrieve_rel_common(relation), plan?)
+}
+
+/// Can be used to consume standard Substrait without user-defined extensions
+pub struct DefaultSubstraitConsumer {
+ extensions: Arc<Extensions>,
+ state: Arc<SessionState>,
Review Comment:
Replaced with lifetime annotations based on @alamb proposals in https://github.com/vbarua/datafusion/pull/1
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894401830
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -989,8 +992,13 @@ async fn extension_logical_plan() -> Result<()> {
}),
});
+ let consumer = DefaultSubstraitConsumer::new(
+ Arc::new(Extensions::default()),
+ Arc::new(ctx.state()),
+ );
+
let proto = to_substrait_plan(&ext_plan, &ctx.state())?;
- let plan2 = from_substrait_plan(&ctx.state(), &proto).await?;
+ let plan2 = from_substrait_plan_with_consumer(&consumer, &proto).await?;
Review Comment:
I've updated this to call `from_substrait_plan`. This is a holdover from an earlier design in which the DefaultSubstraitConsumer didn't handle extension relations
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894424867
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -90,17 +97,469 @@ use substrait::proto::{
window_function::bound::Kind as BoundKind, window_function::Bound,
window_function::BoundsType, MaskExpression, RexType,
},
+ fetch_rel,
function_argument::ArgType,
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtendedExpression, ExtensionLeafRel, ExtensionMultiRel,
+ ExtensionSingleRel, FetchRel, FilterRel, FunctionArgument, JoinRel, NamedStruct,
+ Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel, SortField, SortRel, Type,
};
-use substrait::proto::{fetch_rel, ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// # use async_trait::async_trait;
+/// # use datafusion::catalog::TableProvider;
+/// # use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// # use datafusion::error::Result;
+/// # use datafusion::execution::{FunctionRegistry, SessionState};
+/// # use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// # use std::sync::Arc;
+/// # use substrait::proto;
+/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// # use datafusion::arrow::datatypes::DataType;
+/// # use datafusion::logical_expr::expr::ScalarFunction;
+/// # use datafusion_substrait::extensions::Extensions;
+/// # use datafusion_substrait::logical_plan::consumer::{
+/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// # };
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[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()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and FunctionRegistry is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ // See: https://github.com/apache/datafusion/issues/13863
+ fn get_extensions(&self) -> &Extensions;
+ fn get_function_registry(&self) -> &impl FunctionRegistry;
+
+ // Relation Methods
+ // There is one method per Substrait relation to allow for easy overriding of consumer behaviour.
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_read(&self, rel: &ReadRel) -> Result<LogicalPlan> {
+ from_read_rel(self, rel).await
+ }
+
+ async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+ from_filter_rel(self, rel).await
+ }
+
+ async fn consume_fetch(&self, rel: &FetchRel) -> Result<LogicalPlan> {
+ from_fetch_rel(self, rel).await
+ }
+
+ async fn consume_aggregate(&self, rel: &AggregateRel) -> Result<LogicalPlan> {
+ from_aggregate_rel(self, rel).await
+ }
+
+ async fn consume_sort(&self, rel: &SortRel) -> Result<LogicalPlan> {
+ from_sort_rel(self, rel).await
+ }
+
+ async fn consume_join(&self, rel: &JoinRel) -> Result<LogicalPlan> {
+ from_join_rel(self, rel).await
+ }
+
+ async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+ from_project_rel(self, rel).await
+ }
+
+ async fn consume_set(&self, rel: &SetRel) -> Result<LogicalPlan> {
+ from_set_rel(self, rel).await
+ }
+
+ async fn consume_cross(&self, rel: &CrossRel) -> Result<LogicalPlan> {
+ from_cross_rel(self, rel).await
+ }
+
+ async fn consume_consistent_partition_window(
+ &self,
+ _rel: &ConsistentPartitionWindowRel,
+ ) -> Result<LogicalPlan> {
+ not_impl_err!("Consistent Partition Window Rel not supported")
+ }
+
+ async fn consume_exchange(&self, rel: &ExchangeRel) -> Result<LogicalPlan> {
+ from_exchange_rel(self, rel).await
+ }
+
+ // Expression Methods
+ // There is one method per Substrait expression to allow for easy overriding of consumer behaviour
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_literal(&self, expr: &Literal) -> Result<Expr> {
+ from_literal(self, expr).await
+ }
+
+ async fn consume_field_reference(
+ &self,
+ expr: &FieldReference,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_field_reference(self, expr, input_schema).await
+ }
+
+ async fn consume_scalar_function(
+ &self,
+ expr: &ScalarFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_scalar_function(self, expr, input_schema).await
+ }
+
+ async fn consume_window_function(
+ &self,
+ expr: &WindowFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_window_function(self, expr, input_schema).await
+ }
+
+ async fn consume_if_then(
+ &self,
+ expr: &IfThen,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_if_then(self, expr, input_schema).await
+ }
+
+ async fn consume_switch(
+ &self,
+ _expr: &SwitchExpression,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Switch expression not supported")
+ }
+
+ async fn consume_singular_or_list(
+ &self,
+ expr: &SingularOrList,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_singular_or_list(self, expr, input_schema).await
+ }
+
+ async fn consume_multi_or_list(
+ &self,
+ _expr: &MultiOrList,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Multi Or List expression not supported")
+ }
+
+ async fn consume_cast(
+ &self,
+ expr: &substrait_expression::Cast,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_cast(self, expr, input_schema).await
+ }
+
+ async fn consume_subquery(
+ &self,
+ expr: &substrait_expression::Subquery,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_subquery(self, expr, input_schema).await
+ }
+
+ async fn consume_nested(
+ &self,
+ _expr: &Nested,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Nested expression not supported")
+ }
+
+ async fn consume_enum(&self, _expr: &Enum, _input_schema: &DFSchema) -> Result<Expr> {
+ not_impl_err!("Enum expression not supported")
+ }
+
+ // User-Defined Functionality
+
+ // The details of extension relations, and how to handle them, are fully up to users to specify.
+ // The following methods allow users to customize the consumer behaviour
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionLeafRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionLeafRel")
+ }
+
+ async fn consume_extension_single(
+ &self,
+ rel: &ExtensionSingleRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionSingleRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionSingleRel")
+ }
+
+ async fn consume_extension_multi(
+ &self,
+ rel: &ExtensionMultiRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionMultiRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionMultiRel")
+ }
+
+ // Users can bring their own types to Substrait which require custom handling
+
+ fn consume_user_defined_type(
+ &self,
+ user_defined_type: &r#type::UserDefined,
+ ) -> Result<DataType> {
+ substrait_err!(
+ "Missing handler for user-defined type: {}",
+ user_defined_type.type_reference
+ )
+ }
+
+ fn consume_user_defined_literal(
+ &self,
+ user_defined_literal: &proto::expression::literal::UserDefined,
+ ) -> Result<ScalarValue> {
+ substrait_err!(
+ "Missing handler for user-defined literals {}",
+ user_defined_literal.type_reference
+ )
+ }
+}
+
+/// Convert Substrait Rel to DataFusion DataFrame
+#[async_recursion]
+pub async fn from_substrait_rel(
+ consumer: &impl SubstraitConsumer,
+ relation: &Rel,
+) -> Result<LogicalPlan> {
+ let plan: Result<LogicalPlan> = match &relation.rel_type {
+ Some(rel_type) => match rel_type {
+ RelType::Read(rel) => consumer.consume_read(rel).await,
+ RelType::Filter(rel) => consumer.consume_filter(rel).await,
+ RelType::Fetch(rel) => consumer.consume_fetch(rel).await,
+ RelType::Aggregate(rel) => consumer.consume_aggregate(rel).await,
+ RelType::Sort(rel) => consumer.consume_sort(rel).await,
+ RelType::Join(rel) => consumer.consume_join(rel).await,
+ RelType::Project(rel) => consumer.consume_project(rel).await,
+ RelType::Set(rel) => consumer.consume_set(rel).await,
+ RelType::ExtensionSingle(rel) => consumer.consume_extension_single(rel).await,
+ RelType::ExtensionMulti(rel) => consumer.consume_extension_multi(rel).await,
+ RelType::ExtensionLeaf(rel) => consumer.consume_extension_leaf(rel).await,
+ RelType::Cross(rel) => consumer.consume_cross(rel).await,
+ RelType::Window(rel) => {
+ consumer.consume_consistent_partition_window(rel).await
+ }
+ RelType::Exchange(rel) => consumer.consume_exchange(rel).await,
+ rt => not_impl_err!("{rt:?} rel not supported yet"),
+ },
+ None => return substrait_err!("rel must set rel_type"),
+ };
+ apply_emit_kind(retrieve_rel_common(relation), plan?)
+}
+
+/// Default SubstraitConsumer for converting standard Substrait without user-defined extensions.
+///
+/// Used as the consumer in [from_substrait_plan]
+pub struct DefaultSubstraitConsumer<'a> {
+ extensions: &'a Extensions,
+ state: &'a SessionState,
+}
+
+impl<'a> DefaultSubstraitConsumer<'a> {
+ pub fn new(extensions: &'a Extensions, state: &'a SessionState) -> Self {
+ DefaultSubstraitConsumer { extensions, state }
+ }
+}
Review Comment:
In a prior versions this used to implement `Default`. I've removed it as it was intended mostly for testing. External users should just call `from_substrait_plan` as before.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -90,17 +97,469 @@ use substrait::proto::{
window_function::bound::Kind as BoundKind, window_function::Bound,
window_function::BoundsType, MaskExpression, RexType,
},
+ fetch_rel,
function_argument::ArgType,
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtendedExpression, ExtensionLeafRel, ExtensionMultiRel,
+ ExtensionSingleRel, FetchRel, FilterRel, FunctionArgument, JoinRel, NamedStruct,
+ Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel, SortField, SortRel, Type,
};
-use substrait::proto::{fetch_rel, ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// # use async_trait::async_trait;
+/// # use datafusion::catalog::TableProvider;
+/// # use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// # use datafusion::error::Result;
+/// # use datafusion::execution::{FunctionRegistry, SessionState};
+/// # use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// # use std::sync::Arc;
+/// # use substrait::proto;
+/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// # use datafusion::arrow::datatypes::DataType;
+/// # use datafusion::logical_expr::expr::ScalarFunction;
+/// # use datafusion_substrait::extensions::Extensions;
+/// # use datafusion_substrait::logical_plan::consumer::{
+/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// # };
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[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()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and FunctionRegistry is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ // See: https://github.com/apache/datafusion/issues/13863
+ fn get_extensions(&self) -> &Extensions;
+ fn get_function_registry(&self) -> &impl FunctionRegistry;
+
+ // Relation Methods
+ // There is one method per Substrait relation to allow for easy overriding of consumer behaviour.
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_read(&self, rel: &ReadRel) -> Result<LogicalPlan> {
+ from_read_rel(self, rel).await
+ }
+
+ async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+ from_filter_rel(self, rel).await
+ }
+
+ async fn consume_fetch(&self, rel: &FetchRel) -> Result<LogicalPlan> {
+ from_fetch_rel(self, rel).await
+ }
+
+ async fn consume_aggregate(&self, rel: &AggregateRel) -> Result<LogicalPlan> {
+ from_aggregate_rel(self, rel).await
+ }
+
+ async fn consume_sort(&self, rel: &SortRel) -> Result<LogicalPlan> {
+ from_sort_rel(self, rel).await
+ }
+
+ async fn consume_join(&self, rel: &JoinRel) -> Result<LogicalPlan> {
+ from_join_rel(self, rel).await
+ }
+
+ async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+ from_project_rel(self, rel).await
+ }
+
+ async fn consume_set(&self, rel: &SetRel) -> Result<LogicalPlan> {
+ from_set_rel(self, rel).await
+ }
+
+ async fn consume_cross(&self, rel: &CrossRel) -> Result<LogicalPlan> {
+ from_cross_rel(self, rel).await
+ }
+
+ async fn consume_consistent_partition_window(
+ &self,
+ _rel: &ConsistentPartitionWindowRel,
+ ) -> Result<LogicalPlan> {
+ not_impl_err!("Consistent Partition Window Rel not supported")
+ }
+
+ async fn consume_exchange(&self, rel: &ExchangeRel) -> Result<LogicalPlan> {
+ from_exchange_rel(self, rel).await
+ }
+
+ // Expression Methods
+ // There is one method per Substrait expression to allow for easy overriding of consumer behaviour
+ // These methods have default implementations calling the common handler code, to allow for users
+ // to re-use common handling logic.
+
+ async fn consume_literal(&self, expr: &Literal) -> Result<Expr> {
+ from_literal(self, expr).await
+ }
+
+ async fn consume_field_reference(
+ &self,
+ expr: &FieldReference,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_field_reference(self, expr, input_schema).await
+ }
+
+ async fn consume_scalar_function(
+ &self,
+ expr: &ScalarFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_scalar_function(self, expr, input_schema).await
+ }
+
+ async fn consume_window_function(
+ &self,
+ expr: &WindowFunction,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_window_function(self, expr, input_schema).await
+ }
+
+ async fn consume_if_then(
+ &self,
+ expr: &IfThen,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_if_then(self, expr, input_schema).await
+ }
+
+ async fn consume_switch(
+ &self,
+ _expr: &SwitchExpression,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Switch expression not supported")
+ }
+
+ async fn consume_singular_or_list(
+ &self,
+ expr: &SingularOrList,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_singular_or_list(self, expr, input_schema).await
+ }
+
+ async fn consume_multi_or_list(
+ &self,
+ _expr: &MultiOrList,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Multi Or List expression not supported")
+ }
+
+ async fn consume_cast(
+ &self,
+ expr: &substrait_expression::Cast,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_cast(self, expr, input_schema).await
+ }
+
+ async fn consume_subquery(
+ &self,
+ expr: &substrait_expression::Subquery,
+ input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ from_subquery(self, expr, input_schema).await
+ }
+
+ async fn consume_nested(
+ &self,
+ _expr: &Nested,
+ _input_schema: &DFSchema,
+ ) -> Result<Expr> {
+ not_impl_err!("Nested expression not supported")
+ }
+
+ async fn consume_enum(&self, _expr: &Enum, _input_schema: &DFSchema) -> Result<Expr> {
+ not_impl_err!("Enum expression not supported")
+ }
+
+ // User-Defined Functionality
+
+ // The details of extension relations, and how to handle them, are fully up to users to specify.
+ // The following methods allow users to customize the consumer behaviour
+
+ async fn consume_extension_leaf(
+ &self,
+ rel: &ExtensionLeafRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionLeafRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionLeafRel")
+ }
+
+ async fn consume_extension_single(
+ &self,
+ rel: &ExtensionSingleRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionSingleRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionSingleRel")
+ }
+
+ async fn consume_extension_multi(
+ &self,
+ rel: &ExtensionMultiRel,
+ ) -> Result<LogicalPlan> {
+ if let Some(detail) = rel.detail.as_ref() {
+ return substrait_err!(
+ "Missing handler for ExtensionMultiRel: {}",
+ detail.type_url
+ );
+ }
+ substrait_err!("Missing handler for ExtensionMultiRel")
+ }
+
+ // Users can bring their own types to Substrait which require custom handling
+
+ fn consume_user_defined_type(
+ &self,
+ user_defined_type: &r#type::UserDefined,
+ ) -> Result<DataType> {
+ substrait_err!(
+ "Missing handler for user-defined type: {}",
+ user_defined_type.type_reference
+ )
+ }
+
+ fn consume_user_defined_literal(
+ &self,
+ user_defined_literal: &proto::expression::literal::UserDefined,
+ ) -> Result<ScalarValue> {
+ substrait_err!(
+ "Missing handler for user-defined literals {}",
+ user_defined_literal.type_reference
+ )
+ }
+}
+
+/// Convert Substrait Rel to DataFusion DataFrame
+#[async_recursion]
+pub async fn from_substrait_rel(
+ consumer: &impl SubstraitConsumer,
+ relation: &Rel,
+) -> Result<LogicalPlan> {
+ let plan: Result<LogicalPlan> = match &relation.rel_type {
+ Some(rel_type) => match rel_type {
+ RelType::Read(rel) => consumer.consume_read(rel).await,
+ RelType::Filter(rel) => consumer.consume_filter(rel).await,
+ RelType::Fetch(rel) => consumer.consume_fetch(rel).await,
+ RelType::Aggregate(rel) => consumer.consume_aggregate(rel).await,
+ RelType::Sort(rel) => consumer.consume_sort(rel).await,
+ RelType::Join(rel) => consumer.consume_join(rel).await,
+ RelType::Project(rel) => consumer.consume_project(rel).await,
+ RelType::Set(rel) => consumer.consume_set(rel).await,
+ RelType::ExtensionSingle(rel) => consumer.consume_extension_single(rel).await,
+ RelType::ExtensionMulti(rel) => consumer.consume_extension_multi(rel).await,
+ RelType::ExtensionLeaf(rel) => consumer.consume_extension_leaf(rel).await,
+ RelType::Cross(rel) => consumer.consume_cross(rel).await,
+ RelType::Window(rel) => {
+ consumer.consume_consistent_partition_window(rel).await
+ }
+ RelType::Exchange(rel) => consumer.consume_exchange(rel).await,
+ rt => not_impl_err!("{rt:?} rel not supported yet"),
+ },
+ None => return substrait_err!("rel must set rel_type"),
+ };
+ apply_emit_kind(retrieve_rel_common(relation), plan?)
+}
+
+/// Default SubstraitConsumer for converting standard Substrait without user-defined extensions.
+///
+/// Used as the consumer in [from_substrait_plan]
+pub struct DefaultSubstraitConsumer<'a> {
+ extensions: &'a Extensions,
+ state: &'a SessionState,
+}
+
+impl<'a> DefaultSubstraitConsumer<'a> {
+ pub fn new(extensions: &'a Extensions, state: &'a SessionState) -> Self {
+ DefaultSubstraitConsumer { extensions, state }
+ }
+}
Review Comment:
In a prior versions this used to implement `Default`. I've removed it as it was intended mostly for testing. External users should just call `from_substrait_plan` as before.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#issuecomment-2557775943
I'll plan to merge this tomorrow unless anyone would like some more time to review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1890594630
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -21,23 +21,24 @@ use datafusion::arrow::array::{GenericListArray, MapArray};
use datafusion::arrow::datatypes::{
DataType, Field, FieldRef, Fields, IntervalUnit, Schema, TimeUnit,
};
-use datafusion::common::plan_err;
use datafusion::common::{
- not_impl_err, plan_datafusion_err, substrait_datafusion_err, substrait_err, DFSchema,
- DFSchemaRef,
+ not_impl_datafusion_err, not_impl_err, plan_datafusion_err, plan_err,
+ substrait_datafusion_err, substrait_err, DFSchema, DFSchemaRef,
};
use datafusion::datasource::provider_as_source;
use datafusion::logical_expr::expr::{Exists, InSubquery, Sort};
use datafusion::logical_expr::{
- Aggregate, BinaryExpr, Case, EmptyRelation, Expr, ExprSchemable, LogicalPlan,
- Operator, Projection, SortExpr, TryCast, Values,
+ Aggregate, BinaryExpr, Case, Cast, EmptyRelation, Expr, ExprSchemable, Extension,
+ LogicalPlan, Operator, Projection, SortExpr, Subquery, TryCast, Values,
};
use substrait::proto::aggregate_rel::Grouping;
+use substrait::proto::expression as substrait_expression;
use substrait::proto::expression::subquery::set_predicate::PredicateOp;
use substrait::proto::expression_reference::ExprType;
use url::Url;
+use super::state::SubstraitPlanningState;
Review Comment:
Went ahead and inlined it to be able to remove all usages from the consumer. Will make it easier to remove fully once we updated the producer as well.
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +102,467 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::SessionState;
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// use datafusion_substrait::logical_plan::state::SubstraitPlanningState;
Review Comment:
Removed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#issuecomment-2548987977
> @Blizzara I would appreciate if you could take a look and let me know if this approach seems reasonable, or if I've [architecture astronauted](https://en.wikipedia.org/wiki/Architecture_astronaut) myself.
Thanks for the ping - I took a look, and overall looks really good to me! I left couple thoughts (I see you already saw some), but nothing major. I already have usecase in mind for this internally, so looking forward to getting it merged 😄
One thing is we should consider splitting the consumer.rs into smaller files. It could be a separate PR to keep the diff smaller, but maybe moving all the static "consumer_X" functions into like `consumer/relations.rs`, `consumer/expressions.rs`, `consumer/literals.rs` and `consumer/types.rs` (and one more for utils etc) or similar? Feels like now could be a good moment to do that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1889326909
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +977,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
+ consumer: &impl SubstraitConsumer,
+ p: &ProjectRel,
) -> Result<LogicalPlan> {
- let plan: Result<LogicalPlan> = match &rel.rel_type {
- Some(RelType::Project(p)) => {
- if let Some(input) = p.input.as_ref() {
- let mut input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let original_schema = input.schema().clone();
-
- // Ensure that all expressions have a unique display name, so that
- // validate_unique_names does not fail when constructing the project.
- let mut name_tracker = NameTracker::new();
-
- // By default, a Substrait Project emits all inputs fields followed by all expressions.
- // We build the explicit expressions first, and then the input expressions to avoid
- // adding aliases to the explicit expressions (as part of ensuring unique names).
- //
- // This is helpful for plan visualization and tests, because when DataFusion produces
- // Substrait Projects it adds an output mapping that excludes all input columns
- // leaving only explicit expressions.
-
- let mut explicit_exprs: Vec<Expr> = vec![];
- for expr in &p.expressions {
- let e = from_substrait_rex(
- state,
- expr,
- input.clone().schema(),
- extensions,
- )
- .await?;
- // if the expression is WindowFunction, wrap in a Window relation
- if let Expr::WindowFunction(_) = &e {
- // Adding the same expression here and in the project below
- // works because the project's builder uses columnize_expr(..)
- // to transform it into a column reference
- input = input.window(vec![e.clone()])?
- }
- explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
-
- let mut final_exprs: Vec<Expr> = vec![];
- for index in 0..original_schema.fields().len() {
- let e = Expr::Column(Column::from(
- original_schema.qualified_field(index),
- ));
- final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
- final_exprs.append(&mut explicit_exprs);
-
- input.project(final_exprs)?.build()
- } else {
- not_impl_err!("Projection without an input is not supported")
+ if let Some(input) = p.input.as_ref() {
+ let mut input =
+ LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let original_schema = input.schema().clone();
+
+ // Ensure that all expressions have a unique display name, so that
+ // validate_unique_names does not fail when constructing the project.
+ let mut name_tracker = NameTracker::new();
+
+ // By default, a Substrait Project emits all inputs fields followed by all expressions.
+ // We build the explicit expressions first, and then the input expressions to avoid
+ // adding aliases to the explicit expressions (as part of ensuring unique names).
+ //
+ // This is helpful for plan visualization and tests, because when DataFusion produces
+ // Substrait Projects it adds an output mapping that excludes all input columns
+ // leaving only explicit expressions.
+
+ let mut explicit_exprs: Vec<Expr> = vec![];
+ for expr in &p.expressions {
+ let e = from_substrait_rex(consumer, expr, input.clone().schema()).await?;
+ // if the expression is WindowFunction, wrap in a Window relation
+ if let Expr::WindowFunction(_) = &e {
+ // Adding the same expression here and in the project below
+ // works because the project's builder uses columnize_expr(..)
+ // to transform it into a column reference
+ input = input.window(vec![e.clone()])?
}
+ explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Filter(filter)) => {
- if let Some(input) = filter.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- if let Some(condition) = filter.condition.as_ref() {
- let expr =
- from_substrait_rex(state, condition, input.schema(), extensions)
- .await?;
- input.filter(expr)?.build()
- } else {
- not_impl_err!("Filter without an condition is not valid")
- }
- } else {
- not_impl_err!("Filter without an input is not valid")
- }
+
+ let mut final_exprs: Vec<Expr> = vec![];
+ for index in 0..original_schema.fields().len() {
+ let e = Expr::Column(Column::from(original_schema.qualified_field(index)));
+ final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Fetch(fetch)) => {
- if let Some(input) = fetch.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let offset = fetch.offset as usize;
- // -1 means that ALL records should be returned
- let count = if fetch.count == -1 {
- None
- } else {
- Some(fetch.count as usize)
- };
- input.limit(offset, count)?.build()
- } else {
- not_impl_err!("Fetch without an input is not valid")
- }
+ final_exprs.append(&mut explicit_exprs);
+ input.project(final_exprs)?.build()
+ } else {
+ not_impl_err!("Projection without an input is not supported")
+ }
+}
+
+#[async_recursion]
+pub async fn from_filter_rel(
+ consumer: &impl SubstraitConsumer,
+ filter: &FilterRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = filter.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ if let Some(condition) = filter.condition.as_ref() {
+ let expr = from_substrait_rex(consumer, condition, input.schema()).await?;
+ input.filter(expr)?.build()
+ } else {
+ not_impl_err!("Filter without an condition is not valid")
}
- Some(RelType::Sort(sort)) => {
- if let Some(input) = sort.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let sorts =
- from_substrait_sorts(state, &sort.sorts, input.schema(), extensions)
- .await?;
- input.sort(sorts)?.build()
- } else {
- not_impl_err!("Sort without an input is not valid")
- }
+ } else {
+ not_impl_err!("Filter without an input is not valid")
+ }
+}
+
+#[async_recursion]
+pub async fn from_fetch_rel(
+ consumer: &impl SubstraitConsumer,
+ fetch: &FetchRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = fetch.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let offset = fetch.offset as usize;
+ // -1 means that ALL records should be returned
+ let count = if fetch.count == -1 {
+ None
+ } else {
+ Some(fetch.count as usize)
+ };
+ input.limit(offset, count)?.build()
+ } else {
+ not_impl_err!("Fetch without an input is not valid")
+ }
+}
+
+pub async fn from_sort_rel(
+ consumer: &impl SubstraitConsumer,
+ sort: &SortRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = sort.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?;
+ input.sort(sorts)?.build()
+ } else {
+ not_impl_err!("Sort without an input is not valid")
+ }
+}
+
+pub async fn from_aggregate_rel(
+ consumer: &impl SubstraitConsumer,
+ agg: &AggregateRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = agg.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let mut ref_group_exprs = vec![];
+
+ for e in &agg.grouping_expressions {
+ let x = from_substrait_rex(consumer, e, input.schema()).await?;
+ ref_group_exprs.push(x);
}
- Some(RelType::Aggregate(agg)) => {
- if let Some(input) = agg.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let mut ref_group_exprs = vec![];
- for e in &agg.grouping_expressions {
- let x =
- from_substrait_rex(state, e, input.schema(), extensions).await?;
- ref_group_exprs.push(x);
+ let mut group_exprs = vec![];
+ let mut aggr_exprs = vec![];
+
+ match agg.groupings.len() {
+ 1 => {
+ group_exprs.extend_from_slice(
+ &from_substrait_grouping(
+ consumer,
+ &agg.groupings[0],
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?,
+ );
+ }
+ _ => {
+ let mut grouping_sets = vec![];
+ for grouping in &agg.groupings {
+ let grouping_set = from_substrait_grouping(
+ consumer,
+ grouping,
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?;
+ grouping_sets.push(grouping_set);
}
+ // Single-element grouping expression of type Expr::GroupingSet.
+ // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+ // parsed by the producer and consumer, since Substrait does not have a type dedicated
+ // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+ group_exprs
+ .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+ }
+ };
- let mut group_exprs = vec![];
- let mut aggr_exprs = vec![];
-
- match agg.groupings.len() {
- 1 => {
- group_exprs.extend_from_slice(
- &from_substrait_grouping(
- state,
- &agg.groupings[0],
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?,
- );
- }
- _ => {
- let mut grouping_sets = vec![];
- for grouping in &agg.groupings {
- let grouping_set = from_substrait_grouping(
- state,
- grouping,
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?;
- grouping_sets.push(grouping_set);
+ for m in &agg.measures {
+ let filter = match &m.filter {
+ Some(fil) => Some(Box::new(
+ from_substrait_rex(consumer, fil, input.schema()).await?,
+ )),
+ None => None,
+ };
+ let agg_func = match &m.measure {
+ Some(f) => {
+ let distinct = match f.invocation {
+ _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+ true
}
- // Single-element grouping expression of type Expr::GroupingSet.
- // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
- // parsed by the producer and consumer, since Substrait does not have a type dedicated
- // to ROLLUP. Only vector of Groupings (grouping sets) is available.
- group_exprs.push(Expr::GroupingSet(GroupingSet::GroupingSets(
- grouping_sets,
- )));
- }
- };
-
- for m in &agg.measures {
- let filter = match &m.filter {
- Some(fil) => Some(Box::new(
- from_substrait_rex(state, fil, input.schema(), extensions)
- .await?,
- )),
- None => None,
+ _ if f.invocation == AggregationInvocation::All as i32 => false,
+ _ => false,
};
- let agg_func = match &m.measure {
- Some(f) => {
- let distinct = match f.invocation {
- _ if f.invocation
- == AggregationInvocation::Distinct as i32 =>
- {
- true
- }
- _ if f.invocation
- == AggregationInvocation::All as i32 =>
- {
- false
- }
- _ => false,
- };
- let order_by = if !f.sorts.is_empty() {
- Some(
- from_substrait_sorts(
- state,
- &f.sorts,
- input.schema(),
- extensions,
- )
- .await?,
- )
- } else {
- None
- };
-
- from_substrait_agg_func(
- state,
- f,
- input.schema(),
- extensions,
- filter,
- order_by,
- distinct,
- )
- .await
- }
- None => not_impl_err!(
- "Aggregate without aggregate function is not supported"
- ),
+ let order_by = if !f.sorts.is_empty() {
+ Some(
+ from_substrait_sorts(consumer, &f.sorts, input.schema())
+ .await?,
+ )
+ } else {
+ None
};
- aggr_exprs.push(agg_func?.as_ref().clone());
- }
- input.aggregate(group_exprs, aggr_exprs)?.build()
- } else {
- not_impl_err!("Aggregate without an input is not valid")
- }
- }
- Some(RelType::Join(join)) => {
- if join.post_join_filter.is_some() {
- return not_impl_err!(
- "JoinRel with post_join_filter is not yet supported"
- );
- }
- let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
-
- let join_type = from_substrait_jointype(join.r#type)?;
- // The join condition expression needs full input schema and not the output schema from join since we lose columns from
- // certain join types such as semi and anti joins
- let in_join_schema = left.schema().join(right.schema())?;
-
- // If join expression exists, parse the `on` condition expression, build join and return
- // Otherwise, build join with only the filter, without join keys
- match &join.expression.as_ref() {
- Some(expr) => {
- let on = from_substrait_rex(state, expr, &in_join_schema, extensions)
- .await?;
- // The join expression can contain both equal and non-equal ops.
- // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
- // So we extract each part as follows:
- // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
- // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
- let (join_ons, nulls_equal_nulls, join_filter) =
- split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
- let (left_cols, right_cols): (Vec<_>, Vec<_>) =
- itertools::multiunzip(join_ons);
- left.join_detailed(
- right.build()?,
- join_type,
- (left_cols, right_cols),
- join_filter,
- nulls_equal_nulls,
- )?
- .build()
+ from_substrait_agg_func(
+ consumer,
+ f,
+ input.schema(),
+ filter,
+ order_by,
+ distinct,
+ )
+ .await
}
None => {
- let on: Vec<String> = vec![];
- left.join_detailed(
- right.build()?,
- join_type,
- (on.clone(), on),
- None,
- false,
- )?
- .build()
+ not_impl_err!("Aggregate without aggregate function is not supported")
}
- }
+ };
+ aggr_exprs.push(agg_func?.as_ref().clone());
}
- Some(RelType::Cross(cross)) => {
- let left = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
- left.cross_join(right.build()?)?.build()
+ input.aggregate(group_exprs, aggr_exprs)?.build()
+ } else {
+ not_impl_err!("Aggregate without an input is not valid")
+ }
+}
+
+pub async fn from_join_rel(
+ consumer: &impl SubstraitConsumer,
+ join: &JoinRel,
+) -> Result<LogicalPlan> {
+ if join.post_join_filter.is_some() {
+ return not_impl_err!("JoinRel with post_join_filter is not yet supported");
+ }
+
+ let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+
+ let join_type = from_substrait_jointype(join.r#type)?;
+ // The join condition expression needs full input schema and not the output schema from join since we lose columns from
+ // certain join types such as semi and anti joins
+ let in_join_schema = left.schema().join(right.schema())?;
+
+ // If join expression exists, parse the `on` condition expression, build join and return
+ // Otherwise, build join with only the filter, without join keys
+ match &join.expression.as_ref() {
+ Some(expr) => {
+ let on = from_substrait_rex(consumer, expr, &in_join_schema).await?;
+ // The join expression can contain both equal and non-equal ops.
+ // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
+ // So we extract each part as follows:
+ // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
+ // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
+ let (join_ons, nulls_equal_nulls, join_filter) =
+ split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) =
+ itertools::multiunzip(join_ons);
+ left.join_detailed(
+ right.build()?,
+ join_type,
+ (left_cols, right_cols),
+ join_filter,
+ nulls_equal_nulls,
+ )?
+ .build()
}
- Some(RelType::Read(read)) => {
- async fn read_with_schema(
- state: &dyn SubstraitPlanningState,
- table_ref: TableReference,
- schema: DFSchema,
- projection: &Option<MaskExpression>,
- ) -> Result<LogicalPlan> {
- let schema = schema.replace_qualifier(table_ref.clone());
-
- let plan = {
- let provider = match state.table(&table_ref).await? {
- Some(ref provider) => Arc::clone(provider),
- _ => return plan_err!("No table named '{table_ref}'"),
- };
+ None => {
+ let on: Vec<String> = vec![];
+ left.join_detailed(right.build()?, join_type, (on.clone(), on), None, false)?
+ .build()
+ }
+ }
+}
- LogicalPlanBuilder::scan(
- table_ref,
- provider_as_source(Arc::clone(&provider)),
- None,
- )?
- .build()?
- };
+pub async fn from_cross_rel(
+ consumer: &impl SubstraitConsumer,
+ cross: &CrossRel,
+) -> Result<LogicalPlan> {
+ let left = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+ left.cross_join(right.build()?)?.build()
+}
- ensure_schema_compatability(plan.schema(), schema.clone())?;
+#[allow(deprecated)]
+pub async fn from_read_rel(
+ consumer: &impl SubstraitConsumer,
+ read: &ReadRel,
+) -> Result<LogicalPlan> {
+ async fn read_with_schema(
+ consumer: &impl SubstraitConsumer,
+ table_ref: TableReference,
+ schema: DFSchema,
+ projection: &Option<MaskExpression>,
+ ) -> Result<LogicalPlan> {
+ let schema = schema.replace_qualifier(table_ref.clone());
+
+ let plan = {
+ let provider = match consumer.resolve_table_ref(&table_ref).await? {
+ Some(ref provider) => Arc::clone(provider),
+ _ => return plan_err!("No table named '{table_ref}'"),
+ };
- let schema = apply_masking(schema, projection)?;
+ LogicalPlanBuilder::scan(
+ table_ref,
+ provider_as_source(Arc::clone(&provider)),
+ None,
+ )?
+ .build()?
+ };
- apply_projection(plan, schema)
- }
+ ensure_schema_compatability(plan.schema(), schema.clone())?;
- let named_struct = read.base_schema.as_ref().ok_or_else(|| {
- substrait_datafusion_err!("No base schema provided for Read Relation")
- })?;
+ let schema = apply_masking(schema, projection)?;
- let substrait_schema = from_substrait_named_struct(named_struct, extensions)?;
+ apply_projection(plan, schema)
+ }
- match &read.as_ref().read_type {
- Some(ReadType::NamedTable(nt)) => {
- let table_reference = match nt.names.len() {
- 0 => {
- return plan_err!("No table name found in NamedTable");
- }
- 1 => TableReference::Bare {
- table: nt.names[0].clone().into(),
- },
- 2 => TableReference::Partial {
- schema: nt.names[0].clone().into(),
- table: nt.names[1].clone().into(),
- },
- _ => TableReference::Full {
- catalog: nt.names[0].clone().into(),
- schema: nt.names[1].clone().into(),
- table: nt.names[2].clone().into(),
- },
- };
+ let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+ substrait_datafusion_err!("No base schema provided for Read Relation")
+ })?;
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
+ let substrait_schema = from_substrait_named_struct(consumer, named_struct)?;
+
+ match &read.read_type {
+ Some(ReadType::NamedTable(nt)) => {
+ let table_reference = match nt.names.len() {
+ 0 => {
+ return plan_err!("No table name found in NamedTable");
}
- Some(ReadType::VirtualTable(vt)) => {
- if vt.values.is_empty() {
- return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: DFSchemaRef::new(substrait_schema),
- }));
+ 1 => TableReference::Bare {
+ table: nt.names[0].clone().into(),
+ },
+ 2 => TableReference::Partial {
+ schema: nt.names[0].clone().into(),
+ table: nt.names[1].clone().into(),
+ },
+ _ => TableReference::Full {
+ catalog: nt.names[0].clone().into(),
+ schema: nt.names[1].clone().into(),
+ table: nt.names[2].clone().into(),
+ },
+ };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ Some(ReadType::VirtualTable(vt)) => {
+ if vt.values.is_empty() {
+ return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row: false,
+ schema: DFSchemaRef::new(substrait_schema),
+ }));
+ }
+
+ let values = vt
+ .values
+ .iter()
+ .map(|row| {
+ let mut name_idx = 0;
+ let lits = row
+ .fields
+ .iter()
+ .map(|lit| {
+ name_idx += 1; // top-level names are provided through schema
+ Ok(Expr::Literal(from_substrait_literal(
+ consumer,
+ lit,
+ &named_struct.names,
+ &mut name_idx,
+ )?))
+ })
+ .collect::<Result<_>>()?;
+ if name_idx != named_struct.names.len() {
+ return substrait_err!(
+ "Names list must match exactly to nested schema, but found {} uses for {} names",
+ name_idx,
+ named_struct.names.len()
+ );
}
+ Ok(lits)
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(LogicalPlan::Values(Values {
+ schema: DFSchemaRef::new(substrait_schema),
+ values,
+ }))
+ }
+ Some(ReadType::LocalFiles(lf)) => {
+ fn extract_filename(name: &str) -> Option<String> {
+ let corrected_url =
+ if name.starts_with("file://") && !name.starts_with("file:///") {
+ name.replacen("file://", "file:///", 1)
+ } else {
+ name.to_string()
+ };
+
+ Url::parse(&corrected_url).ok().and_then(|url| {
+ let path = url.path();
+ std::path::Path::new(path)
+ .file_name()
+ .map(|filename| filename.to_string_lossy().to_string())
+ })
+ }
+
+ // we could use the file name to check the original table provider
+ // TODO: currently does not support multiple local files
+ let filename: Option<String> =
+ lf.items.first().and_then(|x| match x.path_type.as_ref() {
+ Some(UriFile(name)) => extract_filename(name),
+ _ => None,
+ });
+
+ if lf.items.len() > 1 || filename.is_none() {
+ return not_impl_err!("Only single file reads are supported");
+ }
+ let name = filename.unwrap();
+ // directly use unwrap here since we could determine it is a valid one
+ let table_reference = TableReference::Bare { table: name.into() };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ _ => {
+ not_impl_err!("Unsupported ReadType: {:?}", read.read_type)
+ }
+ }
+}
+
+pub async fn from_set_rel(
+ consumer: &impl SubstraitConsumer,
+ set: &SetRel,
+) -> Result<LogicalPlan> {
+ if set.inputs.len() < 2 {
+ substrait_err!("Set operation requires at least two inputs")
+ } else {
+ match set.op() {
+ SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
+ SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
+ SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect(
+ from_substrait_rel(consumer, &set.inputs[0]).await?,
+ union_rels(consumer, &set.inputs[1..], true).await?,
+ false,
+ ),
+ SetOp::IntersectionMultiset => {
+ intersect_rels(consumer, &set.inputs, false).await
+ }
+ SetOp::IntersectionMultisetAll => {
+ intersect_rels(consumer, &set.inputs, true).await
+ }
+ SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
+ SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+ set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+ }
+ }
+}
- let values = vt
- .values
- .iter()
- .map(|row| {
- let mut name_idx = 0;
- let lits = row
- .fields
- .iter()
- .map(|lit| {
- name_idx += 1; // top-level names are provided through schema
- Ok(Expr::Literal(from_substrait_literal(
- lit,
- extensions,
- &named_struct.names,
- &mut name_idx,
- )?))
- })
- .collect::<Result<_>>()?;
- if name_idx != named_struct.names.len() {
- return substrait_err!(
- "Names list must match exactly to nested schema, but found {} uses for {} names",
- name_idx,
- named_struct.names.len()
- );
- }
- Ok(lits)
- })
- .collect::<Result<_>>()?;
+pub async fn from_extension_leaf_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_leaf_rel: &ExtensionLeafRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_leaf(extension_leaf_rel).await
+}
- Ok(LogicalPlan::Values(Values {
- schema: DFSchemaRef::new(substrait_schema),
- values,
- }))
- }
- Some(ReadType::LocalFiles(lf)) => {
- fn extract_filename(name: &str) -> Option<String> {
- let corrected_url = if name.starts_with("file://")
- && !name.starts_with("file:///")
- {
- name.replacen("file://", "file:///", 1)
- } else {
- name.to_string()
- };
+pub async fn from_extension_single_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_single_rel: &ExtensionSingleRel,
+) -> Result<LogicalPlan> {
+ consumer
+ .consume_extension_single(extension_single_rel)
+ .await
+}
- Url::parse(&corrected_url).ok().and_then(|url| {
- let path = url.path();
- std::path::Path::new(path)
- .file_name()
- .map(|filename| filename.to_string_lossy().to_string())
- })
- }
+pub async fn from_extension_multi_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_multi_rel: &ExtensionMultiRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_multi(extension_multi_rel).await
+}
- // we could use the file name to check the original table provider
- // TODO: currently does not support multiple local files
- let filename: Option<String> =
- lf.items.first().and_then(|x| match x.path_type.as_ref() {
- Some(UriFile(name)) => extract_filename(name),
- _ => None,
- });
+pub async fn from_exchange_rel(
+ consumer: &impl SubstraitConsumer,
+ exchange: &ExchangeRel,
+) -> Result<LogicalPlan> {
+ let Some(input) = exchange.input.as_ref() else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+ let input = Arc::new(from_substrait_rel(consumer, input).await?);
- if lf.items.len() > 1 || filename.is_none() {
- return not_impl_err!("Only single file reads are supported");
- }
- let name = filename.unwrap();
- // directly use unwrap here since we could determine it is a valid one
- let table_reference = TableReference::Bare { table: name.into() };
-
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
- }
- _ => {
- not_impl_err!("Unsupported ReadType: {:?}", &read.as_ref().read_type)
- }
- }
- }
- Some(RelType::Set(set)) => match set_rel::SetOp::try_from(set.op) {
- Ok(set_op) => {
- if set.inputs.len() < 2 {
- substrait_err!("Set operation requires at least two inputs")
- } else {
- match set_op {
- set_rel::SetOp::UnionAll => {
- union_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::UnionDistinct => {
- union_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionPrimary => {
- LogicalPlanBuilder::intersect(
- from_substrait_rel(state, &set.inputs[0], extensions)
- .await?,
- union_rels(&set.inputs[1..], state, extensions, true)
- .await?,
- false,
- )
- }
- set_rel::SetOp::IntersectionMultiset => {
- intersect_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionMultisetAll => {
- intersect_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::MinusPrimary => {
- except_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::MinusPrimaryAll => {
- except_rels(&set.inputs, state, extensions, true).await
- }
- _ => not_impl_err!("Unsupported set operator: {set_op:?}"),
- }
- }
+ let Some(exchange_kind) = &exchange.exchange_kind else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+
+ // ref: https://substrait.io/relations/physical_relations/#exchange-types
+ let partitioning_scheme = match exchange_kind {
+ ExchangeKind::ScatterByFields(scatter_fields) => {
+ let mut partition_columns = vec![];
+ let input_schema = input.schema();
+ for field_ref in &scatter_fields.fields {
+ let column = from_substrait_field_reference(field_ref, input_schema)?;
+ partition_columns.push(column);
}
- Err(e) => not_impl_err!("Invalid set operation type {}: {e}", set.op),
- },
- Some(RelType::ExtensionLeaf(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
- }
- Some(RelType::ExtensionSingle(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
Review Comment:
I added the SerializerRegistry based code back into the DefaultSubstraitConsumer for now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#issuecomment-2547414143
@Blizzara I would appreciate if you could take a look and let me know if this approach seems reasonable, or if I've [architecture astronauted](https://en.wikipedia.org/wiki/Architecture_astronaut) myself.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1890403551
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -21,23 +21,24 @@ use datafusion::arrow::array::{GenericListArray, MapArray};
use datafusion::arrow::datatypes::{
DataType, Field, FieldRef, Fields, IntervalUnit, Schema, TimeUnit,
};
-use datafusion::common::plan_err;
use datafusion::common::{
- not_impl_err, plan_datafusion_err, substrait_datafusion_err, substrait_err, DFSchema,
- DFSchemaRef,
+ not_impl_datafusion_err, not_impl_err, plan_datafusion_err, plan_err,
+ substrait_datafusion_err, substrait_err, DFSchema, DFSchemaRef,
};
use datafusion::datasource::provider_as_source;
use datafusion::logical_expr::expr::{Exists, InSubquery, Sort};
use datafusion::logical_expr::{
- Aggregate, BinaryExpr, Case, EmptyRelation, Expr, ExprSchemable, LogicalPlan,
- Operator, Projection, SortExpr, TryCast, Values,
+ Aggregate, BinaryExpr, Case, Cast, EmptyRelation, Expr, ExprSchemable, Extension,
+ LogicalPlan, Operator, Projection, SortExpr, Subquery, TryCast, Values,
};
use substrait::proto::aggregate_rel::Grouping;
+use substrait::proto::expression as substrait_expression;
use substrait::proto::expression::subquery::set_predicate::PredicateOp;
use substrait::proto::expression_reference::ExprType;
use url::Url;
+use super::state::SubstraitPlanningState;
Review Comment:
`SubstraitPlanningState` is used in `producer.rs` still, so we can't remove it yet.
The import here brings the `table()` fn here https://github.com/apache/datafusion/blob/82a40f3edf264834c28f3a1bf2a3143951802634/datafusion/substrait/src/logical_plan/state.rs#L54-L62
into scope. It could make sense to inline it to remove it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1894445457
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -262,7 +718,25 @@ async fn except_rels(
/// Convert Substrait Plan to DataFusion LogicalPlan
pub async fn from_substrait_plan(
- state: &dyn SubstraitPlanningState,
+ state: &SessionState,
+ plan: &Plan,
+) -> Result<LogicalPlan> {
+ // Register function extension
+ let extensions = Extensions::try_from(&plan.extensions)?;
+ if !extensions.type_variations.is_empty() {
+ return not_impl_err!("Type variation extensions are not supported");
+ }
+
+ let consumer = DefaultSubstraitConsumer {
+ extensions: &extensions,
+ state,
+ };
+ from_substrait_plan_with_consumer(&consumer, plan).await
+}
+
+/// Convert Substrait Plan to DataFusion LogicalPlan using the given consumer
+pub async fn from_substrait_plan_with_consumer(
+ consumer: &impl SubstraitConsumer,
Review Comment:
nit: I think if you want to, you could remove the Extension lines from below - no need to duplicate them there
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1888753594
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -2037,42 +2438,48 @@ fn from_substrait_type(
Ok(DataType::Interval(IntervalUnit::MonthDayNano))
}
r#type::Kind::UserDefined(u) => {
- if let Some(name) = extensions.types.get(&u.type_reference) {
+ if let Ok(data_type) = consumer.consume_user_defined_type(u) {
+ return Ok(data_type);
+ }
+
+ // TODO: remove the code below once the producer has been updated
Review Comment:
I think the producer doesn't produce any of these anymore, we just didn't remove them yet in case someone is producing plans with an older version and trying to run them with a newer one. Or like, run a plan that has been stored somewhere. But maybe it's time to remove them from the next version for example.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1887852717
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +977,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
+ consumer: &impl SubstraitConsumer,
+ p: &ProjectRel,
) -> Result<LogicalPlan> {
- let plan: Result<LogicalPlan> = match &rel.rel_type {
- Some(RelType::Project(p)) => {
- if let Some(input) = p.input.as_ref() {
- let mut input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let original_schema = input.schema().clone();
-
- // Ensure that all expressions have a unique display name, so that
- // validate_unique_names does not fail when constructing the project.
- let mut name_tracker = NameTracker::new();
-
- // By default, a Substrait Project emits all inputs fields followed by all expressions.
- // We build the explicit expressions first, and then the input expressions to avoid
- // adding aliases to the explicit expressions (as part of ensuring unique names).
- //
- // This is helpful for plan visualization and tests, because when DataFusion produces
- // Substrait Projects it adds an output mapping that excludes all input columns
- // leaving only explicit expressions.
-
- let mut explicit_exprs: Vec<Expr> = vec![];
- for expr in &p.expressions {
- let e = from_substrait_rex(
- state,
- expr,
- input.clone().schema(),
- extensions,
- )
- .await?;
- // if the expression is WindowFunction, wrap in a Window relation
- if let Expr::WindowFunction(_) = &e {
- // Adding the same expression here and in the project below
- // works because the project's builder uses columnize_expr(..)
- // to transform it into a column reference
- input = input.window(vec![e.clone()])?
- }
- explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
-
- let mut final_exprs: Vec<Expr> = vec![];
- for index in 0..original_schema.fields().len() {
- let e = Expr::Column(Column::from(
- original_schema.qualified_field(index),
- ));
- final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
- final_exprs.append(&mut explicit_exprs);
-
- input.project(final_exprs)?.build()
- } else {
- not_impl_err!("Projection without an input is not supported")
+ if let Some(input) = p.input.as_ref() {
+ let mut input =
+ LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let original_schema = input.schema().clone();
+
+ // Ensure that all expressions have a unique display name, so that
+ // validate_unique_names does not fail when constructing the project.
+ let mut name_tracker = NameTracker::new();
+
+ // By default, a Substrait Project emits all inputs fields followed by all expressions.
+ // We build the explicit expressions first, and then the input expressions to avoid
+ // adding aliases to the explicit expressions (as part of ensuring unique names).
+ //
+ // This is helpful for plan visualization and tests, because when DataFusion produces
+ // Substrait Projects it adds an output mapping that excludes all input columns
+ // leaving only explicit expressions.
+
+ let mut explicit_exprs: Vec<Expr> = vec![];
+ for expr in &p.expressions {
+ let e = from_substrait_rex(consumer, expr, input.clone().schema()).await?;
+ // if the expression is WindowFunction, wrap in a Window relation
+ if let Expr::WindowFunction(_) = &e {
+ // Adding the same expression here and in the project below
+ // works because the project's builder uses columnize_expr(..)
+ // to transform it into a column reference
+ input = input.window(vec![e.clone()])?
}
+ explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Filter(filter)) => {
- if let Some(input) = filter.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- if let Some(condition) = filter.condition.as_ref() {
- let expr =
- from_substrait_rex(state, condition, input.schema(), extensions)
- .await?;
- input.filter(expr)?.build()
- } else {
- not_impl_err!("Filter without an condition is not valid")
- }
- } else {
- not_impl_err!("Filter without an input is not valid")
- }
+
+ let mut final_exprs: Vec<Expr> = vec![];
+ for index in 0..original_schema.fields().len() {
+ let e = Expr::Column(Column::from(original_schema.qualified_field(index)));
+ final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Fetch(fetch)) => {
- if let Some(input) = fetch.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let offset = fetch.offset as usize;
- // -1 means that ALL records should be returned
- let count = if fetch.count == -1 {
- None
- } else {
- Some(fetch.count as usize)
- };
- input.limit(offset, count)?.build()
- } else {
- not_impl_err!("Fetch without an input is not valid")
- }
+ final_exprs.append(&mut explicit_exprs);
+ input.project(final_exprs)?.build()
+ } else {
+ not_impl_err!("Projection without an input is not supported")
+ }
+}
+
+#[async_recursion]
+pub async fn from_filter_rel(
+ consumer: &impl SubstraitConsumer,
+ filter: &FilterRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = filter.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ if let Some(condition) = filter.condition.as_ref() {
+ let expr = from_substrait_rex(consumer, condition, input.schema()).await?;
+ input.filter(expr)?.build()
+ } else {
+ not_impl_err!("Filter without an condition is not valid")
}
- Some(RelType::Sort(sort)) => {
- if let Some(input) = sort.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let sorts =
- from_substrait_sorts(state, &sort.sorts, input.schema(), extensions)
- .await?;
- input.sort(sorts)?.build()
- } else {
- not_impl_err!("Sort without an input is not valid")
- }
+ } else {
+ not_impl_err!("Filter without an input is not valid")
+ }
+}
+
+#[async_recursion]
+pub async fn from_fetch_rel(
+ consumer: &impl SubstraitConsumer,
+ fetch: &FetchRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = fetch.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let offset = fetch.offset as usize;
+ // -1 means that ALL records should be returned
+ let count = if fetch.count == -1 {
+ None
+ } else {
+ Some(fetch.count as usize)
+ };
+ input.limit(offset, count)?.build()
+ } else {
+ not_impl_err!("Fetch without an input is not valid")
+ }
+}
+
+pub async fn from_sort_rel(
+ consumer: &impl SubstraitConsumer,
+ sort: &SortRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = sort.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?;
+ input.sort(sorts)?.build()
+ } else {
+ not_impl_err!("Sort without an input is not valid")
+ }
+}
+
+pub async fn from_aggregate_rel(
+ consumer: &impl SubstraitConsumer,
+ agg: &AggregateRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = agg.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let mut ref_group_exprs = vec![];
+
+ for e in &agg.grouping_expressions {
+ let x = from_substrait_rex(consumer, e, input.schema()).await?;
+ ref_group_exprs.push(x);
}
- Some(RelType::Aggregate(agg)) => {
- if let Some(input) = agg.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let mut ref_group_exprs = vec![];
- for e in &agg.grouping_expressions {
- let x =
- from_substrait_rex(state, e, input.schema(), extensions).await?;
- ref_group_exprs.push(x);
+ let mut group_exprs = vec![];
+ let mut aggr_exprs = vec![];
+
+ match agg.groupings.len() {
+ 1 => {
+ group_exprs.extend_from_slice(
+ &from_substrait_grouping(
+ consumer,
+ &agg.groupings[0],
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?,
+ );
+ }
+ _ => {
+ let mut grouping_sets = vec![];
+ for grouping in &agg.groupings {
+ let grouping_set = from_substrait_grouping(
+ consumer,
+ grouping,
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?;
+ grouping_sets.push(grouping_set);
}
+ // Single-element grouping expression of type Expr::GroupingSet.
+ // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+ // parsed by the producer and consumer, since Substrait does not have a type dedicated
+ // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+ group_exprs
+ .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+ }
+ };
- let mut group_exprs = vec![];
- let mut aggr_exprs = vec![];
-
- match agg.groupings.len() {
- 1 => {
- group_exprs.extend_from_slice(
- &from_substrait_grouping(
- state,
- &agg.groupings[0],
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?,
- );
- }
- _ => {
- let mut grouping_sets = vec![];
- for grouping in &agg.groupings {
- let grouping_set = from_substrait_grouping(
- state,
- grouping,
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?;
- grouping_sets.push(grouping_set);
+ for m in &agg.measures {
+ let filter = match &m.filter {
+ Some(fil) => Some(Box::new(
+ from_substrait_rex(consumer, fil, input.schema()).await?,
+ )),
+ None => None,
+ };
+ let agg_func = match &m.measure {
+ Some(f) => {
+ let distinct = match f.invocation {
+ _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+ true
}
- // Single-element grouping expression of type Expr::GroupingSet.
- // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
- // parsed by the producer and consumer, since Substrait does not have a type dedicated
- // to ROLLUP. Only vector of Groupings (grouping sets) is available.
- group_exprs.push(Expr::GroupingSet(GroupingSet::GroupingSets(
- grouping_sets,
- )));
- }
- };
-
- for m in &agg.measures {
- let filter = match &m.filter {
- Some(fil) => Some(Box::new(
- from_substrait_rex(state, fil, input.schema(), extensions)
- .await?,
- )),
- None => None,
+ _ if f.invocation == AggregationInvocation::All as i32 => false,
+ _ => false,
};
- let agg_func = match &m.measure {
- Some(f) => {
- let distinct = match f.invocation {
- _ if f.invocation
- == AggregationInvocation::Distinct as i32 =>
- {
- true
- }
- _ if f.invocation
- == AggregationInvocation::All as i32 =>
- {
- false
- }
- _ => false,
- };
- let order_by = if !f.sorts.is_empty() {
- Some(
- from_substrait_sorts(
- state,
- &f.sorts,
- input.schema(),
- extensions,
- )
- .await?,
- )
- } else {
- None
- };
-
- from_substrait_agg_func(
- state,
- f,
- input.schema(),
- extensions,
- filter,
- order_by,
- distinct,
- )
- .await
- }
- None => not_impl_err!(
- "Aggregate without aggregate function is not supported"
- ),
+ let order_by = if !f.sorts.is_empty() {
+ Some(
+ from_substrait_sorts(consumer, &f.sorts, input.schema())
+ .await?,
+ )
+ } else {
+ None
};
- aggr_exprs.push(agg_func?.as_ref().clone());
- }
- input.aggregate(group_exprs, aggr_exprs)?.build()
- } else {
- not_impl_err!("Aggregate without an input is not valid")
- }
- }
- Some(RelType::Join(join)) => {
- if join.post_join_filter.is_some() {
- return not_impl_err!(
- "JoinRel with post_join_filter is not yet supported"
- );
- }
- let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
-
- let join_type = from_substrait_jointype(join.r#type)?;
- // The join condition expression needs full input schema and not the output schema from join since we lose columns from
- // certain join types such as semi and anti joins
- let in_join_schema = left.schema().join(right.schema())?;
-
- // If join expression exists, parse the `on` condition expression, build join and return
- // Otherwise, build join with only the filter, without join keys
- match &join.expression.as_ref() {
- Some(expr) => {
- let on = from_substrait_rex(state, expr, &in_join_schema, extensions)
- .await?;
- // The join expression can contain both equal and non-equal ops.
- // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
- // So we extract each part as follows:
- // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
- // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
- let (join_ons, nulls_equal_nulls, join_filter) =
- split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
- let (left_cols, right_cols): (Vec<_>, Vec<_>) =
- itertools::multiunzip(join_ons);
- left.join_detailed(
- right.build()?,
- join_type,
- (left_cols, right_cols),
- join_filter,
- nulls_equal_nulls,
- )?
- .build()
+ from_substrait_agg_func(
+ consumer,
+ f,
+ input.schema(),
+ filter,
+ order_by,
+ distinct,
+ )
+ .await
}
None => {
- let on: Vec<String> = vec![];
- left.join_detailed(
- right.build()?,
- join_type,
- (on.clone(), on),
- None,
- false,
- )?
- .build()
+ not_impl_err!("Aggregate without aggregate function is not supported")
}
- }
+ };
+ aggr_exprs.push(agg_func?.as_ref().clone());
}
- Some(RelType::Cross(cross)) => {
- let left = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
- left.cross_join(right.build()?)?.build()
+ input.aggregate(group_exprs, aggr_exprs)?.build()
+ } else {
+ not_impl_err!("Aggregate without an input is not valid")
+ }
+}
+
+pub async fn from_join_rel(
+ consumer: &impl SubstraitConsumer,
+ join: &JoinRel,
+) -> Result<LogicalPlan> {
+ if join.post_join_filter.is_some() {
+ return not_impl_err!("JoinRel with post_join_filter is not yet supported");
+ }
+
+ let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+
+ let join_type = from_substrait_jointype(join.r#type)?;
+ // The join condition expression needs full input schema and not the output schema from join since we lose columns from
+ // certain join types such as semi and anti joins
+ let in_join_schema = left.schema().join(right.schema())?;
+
+ // If join expression exists, parse the `on` condition expression, build join and return
+ // Otherwise, build join with only the filter, without join keys
+ match &join.expression.as_ref() {
+ Some(expr) => {
+ let on = from_substrait_rex(consumer, expr, &in_join_schema).await?;
+ // The join expression can contain both equal and non-equal ops.
+ // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
+ // So we extract each part as follows:
+ // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
+ // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
+ let (join_ons, nulls_equal_nulls, join_filter) =
+ split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) =
+ itertools::multiunzip(join_ons);
+ left.join_detailed(
+ right.build()?,
+ join_type,
+ (left_cols, right_cols),
+ join_filter,
+ nulls_equal_nulls,
+ )?
+ .build()
}
- Some(RelType::Read(read)) => {
- async fn read_with_schema(
- state: &dyn SubstraitPlanningState,
- table_ref: TableReference,
- schema: DFSchema,
- projection: &Option<MaskExpression>,
- ) -> Result<LogicalPlan> {
- let schema = schema.replace_qualifier(table_ref.clone());
-
- let plan = {
- let provider = match state.table(&table_ref).await? {
- Some(ref provider) => Arc::clone(provider),
- _ => return plan_err!("No table named '{table_ref}'"),
- };
+ None => {
+ let on: Vec<String> = vec![];
+ left.join_detailed(right.build()?, join_type, (on.clone(), on), None, false)?
+ .build()
+ }
+ }
+}
- LogicalPlanBuilder::scan(
- table_ref,
- provider_as_source(Arc::clone(&provider)),
- None,
- )?
- .build()?
- };
+pub async fn from_cross_rel(
+ consumer: &impl SubstraitConsumer,
+ cross: &CrossRel,
+) -> Result<LogicalPlan> {
+ let left = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+ left.cross_join(right.build()?)?.build()
+}
- ensure_schema_compatability(plan.schema(), schema.clone())?;
+#[allow(deprecated)]
+pub async fn from_read_rel(
+ consumer: &impl SubstraitConsumer,
+ read: &ReadRel,
+) -> Result<LogicalPlan> {
+ async fn read_with_schema(
+ consumer: &impl SubstraitConsumer,
+ table_ref: TableReference,
+ schema: DFSchema,
+ projection: &Option<MaskExpression>,
+ ) -> Result<LogicalPlan> {
+ let schema = schema.replace_qualifier(table_ref.clone());
+
+ let plan = {
+ let provider = match consumer.resolve_table_ref(&table_ref).await? {
+ Some(ref provider) => Arc::clone(provider),
+ _ => return plan_err!("No table named '{table_ref}'"),
+ };
- let schema = apply_masking(schema, projection)?;
+ LogicalPlanBuilder::scan(
+ table_ref,
+ provider_as_source(Arc::clone(&provider)),
+ None,
+ )?
+ .build()?
+ };
- apply_projection(plan, schema)
- }
+ ensure_schema_compatability(plan.schema(), schema.clone())?;
- let named_struct = read.base_schema.as_ref().ok_or_else(|| {
- substrait_datafusion_err!("No base schema provided for Read Relation")
- })?;
+ let schema = apply_masking(schema, projection)?;
- let substrait_schema = from_substrait_named_struct(named_struct, extensions)?;
+ apply_projection(plan, schema)
+ }
- match &read.as_ref().read_type {
- Some(ReadType::NamedTable(nt)) => {
- let table_reference = match nt.names.len() {
- 0 => {
- return plan_err!("No table name found in NamedTable");
- }
- 1 => TableReference::Bare {
- table: nt.names[0].clone().into(),
- },
- 2 => TableReference::Partial {
- schema: nt.names[0].clone().into(),
- table: nt.names[1].clone().into(),
- },
- _ => TableReference::Full {
- catalog: nt.names[0].clone().into(),
- schema: nt.names[1].clone().into(),
- table: nt.names[2].clone().into(),
- },
- };
+ let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+ substrait_datafusion_err!("No base schema provided for Read Relation")
+ })?;
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
+ let substrait_schema = from_substrait_named_struct(consumer, named_struct)?;
+
+ match &read.read_type {
+ Some(ReadType::NamedTable(nt)) => {
+ let table_reference = match nt.names.len() {
+ 0 => {
+ return plan_err!("No table name found in NamedTable");
}
- Some(ReadType::VirtualTable(vt)) => {
- if vt.values.is_empty() {
- return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: DFSchemaRef::new(substrait_schema),
- }));
+ 1 => TableReference::Bare {
+ table: nt.names[0].clone().into(),
+ },
+ 2 => TableReference::Partial {
+ schema: nt.names[0].clone().into(),
+ table: nt.names[1].clone().into(),
+ },
+ _ => TableReference::Full {
+ catalog: nt.names[0].clone().into(),
+ schema: nt.names[1].clone().into(),
+ table: nt.names[2].clone().into(),
+ },
+ };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ Some(ReadType::VirtualTable(vt)) => {
+ if vt.values.is_empty() {
+ return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row: false,
+ schema: DFSchemaRef::new(substrait_schema),
+ }));
+ }
+
+ let values = vt
+ .values
+ .iter()
+ .map(|row| {
+ let mut name_idx = 0;
+ let lits = row
+ .fields
+ .iter()
+ .map(|lit| {
+ name_idx += 1; // top-level names are provided through schema
+ Ok(Expr::Literal(from_substrait_literal(
+ consumer,
+ lit,
+ &named_struct.names,
+ &mut name_idx,
+ )?))
+ })
+ .collect::<Result<_>>()?;
+ if name_idx != named_struct.names.len() {
+ return substrait_err!(
+ "Names list must match exactly to nested schema, but found {} uses for {} names",
+ name_idx,
+ named_struct.names.len()
+ );
}
+ Ok(lits)
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(LogicalPlan::Values(Values {
+ schema: DFSchemaRef::new(substrait_schema),
+ values,
+ }))
+ }
+ Some(ReadType::LocalFiles(lf)) => {
+ fn extract_filename(name: &str) -> Option<String> {
+ let corrected_url =
+ if name.starts_with("file://") && !name.starts_with("file:///") {
+ name.replacen("file://", "file:///", 1)
+ } else {
+ name.to_string()
+ };
+
+ Url::parse(&corrected_url).ok().and_then(|url| {
+ let path = url.path();
+ std::path::Path::new(path)
+ .file_name()
+ .map(|filename| filename.to_string_lossy().to_string())
+ })
+ }
+
+ // we could use the file name to check the original table provider
+ // TODO: currently does not support multiple local files
+ let filename: Option<String> =
+ lf.items.first().and_then(|x| match x.path_type.as_ref() {
+ Some(UriFile(name)) => extract_filename(name),
+ _ => None,
+ });
+
+ if lf.items.len() > 1 || filename.is_none() {
+ return not_impl_err!("Only single file reads are supported");
+ }
+ let name = filename.unwrap();
+ // directly use unwrap here since we could determine it is a valid one
+ let table_reference = TableReference::Bare { table: name.into() };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ _ => {
+ not_impl_err!("Unsupported ReadType: {:?}", read.read_type)
+ }
+ }
+}
+
+pub async fn from_set_rel(
+ consumer: &impl SubstraitConsumer,
+ set: &SetRel,
+) -> Result<LogicalPlan> {
+ if set.inputs.len() < 2 {
+ substrait_err!("Set operation requires at least two inputs")
+ } else {
+ match set.op() {
+ SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
+ SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
+ SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect(
+ from_substrait_rel(consumer, &set.inputs[0]).await?,
+ union_rels(consumer, &set.inputs[1..], true).await?,
+ false,
+ ),
+ SetOp::IntersectionMultiset => {
+ intersect_rels(consumer, &set.inputs, false).await
+ }
+ SetOp::IntersectionMultisetAll => {
+ intersect_rels(consumer, &set.inputs, true).await
+ }
+ SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
+ SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+ set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+ }
+ }
+}
- let values = vt
- .values
- .iter()
- .map(|row| {
- let mut name_idx = 0;
- let lits = row
- .fields
- .iter()
- .map(|lit| {
- name_idx += 1; // top-level names are provided through schema
- Ok(Expr::Literal(from_substrait_literal(
- lit,
- extensions,
- &named_struct.names,
- &mut name_idx,
- )?))
- })
- .collect::<Result<_>>()?;
- if name_idx != named_struct.names.len() {
- return substrait_err!(
- "Names list must match exactly to nested schema, but found {} uses for {} names",
- name_idx,
- named_struct.names.len()
- );
- }
- Ok(lits)
- })
- .collect::<Result<_>>()?;
+pub async fn from_extension_leaf_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_leaf_rel: &ExtensionLeafRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_leaf(extension_leaf_rel).await
+}
- Ok(LogicalPlan::Values(Values {
- schema: DFSchemaRef::new(substrait_schema),
- values,
- }))
- }
- Some(ReadType::LocalFiles(lf)) => {
- fn extract_filename(name: &str) -> Option<String> {
- let corrected_url = if name.starts_with("file://")
- && !name.starts_with("file:///")
- {
- name.replacen("file://", "file:///", 1)
- } else {
- name.to_string()
- };
+pub async fn from_extension_single_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_single_rel: &ExtensionSingleRel,
+) -> Result<LogicalPlan> {
+ consumer
+ .consume_extension_single(extension_single_rel)
+ .await
+}
- Url::parse(&corrected_url).ok().and_then(|url| {
- let path = url.path();
- std::path::Path::new(path)
- .file_name()
- .map(|filename| filename.to_string_lossy().to_string())
- })
- }
+pub async fn from_extension_multi_rel(
+ consumer: &impl SubstraitConsumer,
+ extension_multi_rel: &ExtensionMultiRel,
+) -> Result<LogicalPlan> {
+ consumer.consume_extension_multi(extension_multi_rel).await
+}
- // we could use the file name to check the original table provider
- // TODO: currently does not support multiple local files
- let filename: Option<String> =
- lf.items.first().and_then(|x| match x.path_type.as_ref() {
- Some(UriFile(name)) => extract_filename(name),
- _ => None,
- });
+pub async fn from_exchange_rel(
+ consumer: &impl SubstraitConsumer,
+ exchange: &ExchangeRel,
+) -> Result<LogicalPlan> {
+ let Some(input) = exchange.input.as_ref() else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+ let input = Arc::new(from_substrait_rel(consumer, input).await?);
- if lf.items.len() > 1 || filename.is_none() {
- return not_impl_err!("Only single file reads are supported");
- }
- let name = filename.unwrap();
- // directly use unwrap here since we could determine it is a valid one
- let table_reference = TableReference::Bare { table: name.into() };
-
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
- }
- _ => {
- not_impl_err!("Unsupported ReadType: {:?}", &read.as_ref().read_type)
- }
- }
- }
- Some(RelType::Set(set)) => match set_rel::SetOp::try_from(set.op) {
- Ok(set_op) => {
- if set.inputs.len() < 2 {
- substrait_err!("Set operation requires at least two inputs")
- } else {
- match set_op {
- set_rel::SetOp::UnionAll => {
- union_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::UnionDistinct => {
- union_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionPrimary => {
- LogicalPlanBuilder::intersect(
- from_substrait_rel(state, &set.inputs[0], extensions)
- .await?,
- union_rels(&set.inputs[1..], state, extensions, true)
- .await?,
- false,
- )
- }
- set_rel::SetOp::IntersectionMultiset => {
- intersect_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::IntersectionMultisetAll => {
- intersect_rels(&set.inputs, state, extensions, true).await
- }
- set_rel::SetOp::MinusPrimary => {
- except_rels(&set.inputs, state, extensions, false).await
- }
- set_rel::SetOp::MinusPrimaryAll => {
- except_rels(&set.inputs, state, extensions, true).await
- }
- _ => not_impl_err!("Unsupported set operator: {set_op:?}"),
- }
- }
+ let Some(exchange_kind) = &exchange.exchange_kind else {
+ return substrait_err!("Unexpected empty input in ExchangeRel");
+ };
+
+ // ref: https://substrait.io/relations/physical_relations/#exchange-types
+ let partitioning_scheme = match exchange_kind {
+ ExchangeKind::ScatterByFields(scatter_fields) => {
+ let mut partition_columns = vec![];
+ let input_schema = input.schema();
+ for field_ref in &scatter_fields.fields {
+ let column = from_substrait_field_reference(field_ref, input_schema)?;
+ partition_columns.push(column);
}
- Err(e) => not_impl_err!("Invalid set operation type {}: {e}", set.op),
- },
- Some(RelType::ExtensionLeaf(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionLeafRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
- Ok(LogicalPlan::Extension(Extension { node: plan }))
- }
- Some(RelType::ExtensionSingle(extension)) => {
- let Some(ext_detail) = &extension.detail else {
- return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
- };
- let plan = state
- .serializer_registry()
- .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
Review Comment:
The usage of `SerializerRegistry` is problematic here because `deserialize_logical_plan` assumes that the extension details are a fully formed `UserDefinedLogicalNode`. This _may_ sometimes be the case, but the purpose of the Substrait Extension relations is to allow users to encoded arbitrary relations into Substrait in a manner they see fit.
The only really foolproof way of handling this is to provide a hook like `consume_extension_single` which users can use to provide their deserialisation logic.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1889313376
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -557,570 +977,510 @@ fn make_renamed_schema(
)
}
-/// Convert Substrait Rel to DataFusion DataFrame
-#[allow(deprecated)]
#[async_recursion]
-pub async fn from_substrait_rel(
- state: &dyn SubstraitPlanningState,
- rel: &Rel,
- extensions: &Extensions,
+pub async fn from_project_rel(
+ consumer: &impl SubstraitConsumer,
+ p: &ProjectRel,
) -> Result<LogicalPlan> {
- let plan: Result<LogicalPlan> = match &rel.rel_type {
- Some(RelType::Project(p)) => {
- if let Some(input) = p.input.as_ref() {
- let mut input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let original_schema = input.schema().clone();
-
- // Ensure that all expressions have a unique display name, so that
- // validate_unique_names does not fail when constructing the project.
- let mut name_tracker = NameTracker::new();
-
- // By default, a Substrait Project emits all inputs fields followed by all expressions.
- // We build the explicit expressions first, and then the input expressions to avoid
- // adding aliases to the explicit expressions (as part of ensuring unique names).
- //
- // This is helpful for plan visualization and tests, because when DataFusion produces
- // Substrait Projects it adds an output mapping that excludes all input columns
- // leaving only explicit expressions.
-
- let mut explicit_exprs: Vec<Expr> = vec![];
- for expr in &p.expressions {
- let e = from_substrait_rex(
- state,
- expr,
- input.clone().schema(),
- extensions,
- )
- .await?;
- // if the expression is WindowFunction, wrap in a Window relation
- if let Expr::WindowFunction(_) = &e {
- // Adding the same expression here and in the project below
- // works because the project's builder uses columnize_expr(..)
- // to transform it into a column reference
- input = input.window(vec![e.clone()])?
- }
- explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
-
- let mut final_exprs: Vec<Expr> = vec![];
- for index in 0..original_schema.fields().len() {
- let e = Expr::Column(Column::from(
- original_schema.qualified_field(index),
- ));
- final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
- }
- final_exprs.append(&mut explicit_exprs);
-
- input.project(final_exprs)?.build()
- } else {
- not_impl_err!("Projection without an input is not supported")
+ if let Some(input) = p.input.as_ref() {
+ let mut input =
+ LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let original_schema = input.schema().clone();
+
+ // Ensure that all expressions have a unique display name, so that
+ // validate_unique_names does not fail when constructing the project.
+ let mut name_tracker = NameTracker::new();
+
+ // By default, a Substrait Project emits all inputs fields followed by all expressions.
+ // We build the explicit expressions first, and then the input expressions to avoid
+ // adding aliases to the explicit expressions (as part of ensuring unique names).
+ //
+ // This is helpful for plan visualization and tests, because when DataFusion produces
+ // Substrait Projects it adds an output mapping that excludes all input columns
+ // leaving only explicit expressions.
+
+ let mut explicit_exprs: Vec<Expr> = vec![];
+ for expr in &p.expressions {
+ let e = from_substrait_rex(consumer, expr, input.clone().schema()).await?;
+ // if the expression is WindowFunction, wrap in a Window relation
+ if let Expr::WindowFunction(_) = &e {
+ // Adding the same expression here and in the project below
+ // works because the project's builder uses columnize_expr(..)
+ // to transform it into a column reference
+ input = input.window(vec![e.clone()])?
}
+ explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Filter(filter)) => {
- if let Some(input) = filter.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- if let Some(condition) = filter.condition.as_ref() {
- let expr =
- from_substrait_rex(state, condition, input.schema(), extensions)
- .await?;
- input.filter(expr)?.build()
- } else {
- not_impl_err!("Filter without an condition is not valid")
- }
- } else {
- not_impl_err!("Filter without an input is not valid")
- }
+
+ let mut final_exprs: Vec<Expr> = vec![];
+ for index in 0..original_schema.fields().len() {
+ let e = Expr::Column(Column::from(original_schema.qualified_field(index)));
+ final_exprs.push(name_tracker.get_uniquely_named_expr(e)?);
}
- Some(RelType::Fetch(fetch)) => {
- if let Some(input) = fetch.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let offset = fetch.offset as usize;
- // -1 means that ALL records should be returned
- let count = if fetch.count == -1 {
- None
- } else {
- Some(fetch.count as usize)
- };
- input.limit(offset, count)?.build()
- } else {
- not_impl_err!("Fetch without an input is not valid")
- }
+ final_exprs.append(&mut explicit_exprs);
+ input.project(final_exprs)?.build()
+ } else {
+ not_impl_err!("Projection without an input is not supported")
+ }
+}
+
+#[async_recursion]
+pub async fn from_filter_rel(
+ consumer: &impl SubstraitConsumer,
+ filter: &FilterRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = filter.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ if let Some(condition) = filter.condition.as_ref() {
+ let expr = from_substrait_rex(consumer, condition, input.schema()).await?;
+ input.filter(expr)?.build()
+ } else {
+ not_impl_err!("Filter without an condition is not valid")
}
- Some(RelType::Sort(sort)) => {
- if let Some(input) = sort.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let sorts =
- from_substrait_sorts(state, &sort.sorts, input.schema(), extensions)
- .await?;
- input.sort(sorts)?.build()
- } else {
- not_impl_err!("Sort without an input is not valid")
- }
+ } else {
+ not_impl_err!("Filter without an input is not valid")
+ }
+}
+
+#[async_recursion]
+pub async fn from_fetch_rel(
+ consumer: &impl SubstraitConsumer,
+ fetch: &FetchRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = fetch.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let offset = fetch.offset as usize;
+ // -1 means that ALL records should be returned
+ let count = if fetch.count == -1 {
+ None
+ } else {
+ Some(fetch.count as usize)
+ };
+ input.limit(offset, count)?.build()
+ } else {
+ not_impl_err!("Fetch without an input is not valid")
+ }
+}
+
+pub async fn from_sort_rel(
+ consumer: &impl SubstraitConsumer,
+ sort: &SortRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = sort.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?;
+ input.sort(sorts)?.build()
+ } else {
+ not_impl_err!("Sort without an input is not valid")
+ }
+}
+
+pub async fn from_aggregate_rel(
+ consumer: &impl SubstraitConsumer,
+ agg: &AggregateRel,
+) -> Result<LogicalPlan> {
+ if let Some(input) = agg.input.as_ref() {
+ let input = LogicalPlanBuilder::from(from_substrait_rel(consumer, input).await?);
+ let mut ref_group_exprs = vec![];
+
+ for e in &agg.grouping_expressions {
+ let x = from_substrait_rex(consumer, e, input.schema()).await?;
+ ref_group_exprs.push(x);
}
- Some(RelType::Aggregate(agg)) => {
- if let Some(input) = agg.input.as_ref() {
- let input = LogicalPlanBuilder::from(
- from_substrait_rel(state, input, extensions).await?,
- );
- let mut ref_group_exprs = vec![];
- for e in &agg.grouping_expressions {
- let x =
- from_substrait_rex(state, e, input.schema(), extensions).await?;
- ref_group_exprs.push(x);
+ let mut group_exprs = vec![];
+ let mut aggr_exprs = vec![];
+
+ match agg.groupings.len() {
+ 1 => {
+ group_exprs.extend_from_slice(
+ &from_substrait_grouping(
+ consumer,
+ &agg.groupings[0],
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?,
+ );
+ }
+ _ => {
+ let mut grouping_sets = vec![];
+ for grouping in &agg.groupings {
+ let grouping_set = from_substrait_grouping(
+ consumer,
+ grouping,
+ &ref_group_exprs,
+ input.schema(),
+ )
+ .await?;
+ grouping_sets.push(grouping_set);
}
+ // Single-element grouping expression of type Expr::GroupingSet.
+ // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+ // parsed by the producer and consumer, since Substrait does not have a type dedicated
+ // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+ group_exprs
+ .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+ }
+ };
- let mut group_exprs = vec![];
- let mut aggr_exprs = vec![];
-
- match agg.groupings.len() {
- 1 => {
- group_exprs.extend_from_slice(
- &from_substrait_grouping(
- state,
- &agg.groupings[0],
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?,
- );
- }
- _ => {
- let mut grouping_sets = vec![];
- for grouping in &agg.groupings {
- let grouping_set = from_substrait_grouping(
- state,
- grouping,
- &ref_group_exprs,
- input.schema(),
- extensions,
- )
- .await?;
- grouping_sets.push(grouping_set);
+ for m in &agg.measures {
+ let filter = match &m.filter {
+ Some(fil) => Some(Box::new(
+ from_substrait_rex(consumer, fil, input.schema()).await?,
+ )),
+ None => None,
+ };
+ let agg_func = match &m.measure {
+ Some(f) => {
+ let distinct = match f.invocation {
+ _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+ true
}
- // Single-element grouping expression of type Expr::GroupingSet.
- // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
- // parsed by the producer and consumer, since Substrait does not have a type dedicated
- // to ROLLUP. Only vector of Groupings (grouping sets) is available.
- group_exprs.push(Expr::GroupingSet(GroupingSet::GroupingSets(
- grouping_sets,
- )));
- }
- };
-
- for m in &agg.measures {
- let filter = match &m.filter {
- Some(fil) => Some(Box::new(
- from_substrait_rex(state, fil, input.schema(), extensions)
- .await?,
- )),
- None => None,
+ _ if f.invocation == AggregationInvocation::All as i32 => false,
+ _ => false,
};
- let agg_func = match &m.measure {
- Some(f) => {
- let distinct = match f.invocation {
- _ if f.invocation
- == AggregationInvocation::Distinct as i32 =>
- {
- true
- }
- _ if f.invocation
- == AggregationInvocation::All as i32 =>
- {
- false
- }
- _ => false,
- };
- let order_by = if !f.sorts.is_empty() {
- Some(
- from_substrait_sorts(
- state,
- &f.sorts,
- input.schema(),
- extensions,
- )
- .await?,
- )
- } else {
- None
- };
-
- from_substrait_agg_func(
- state,
- f,
- input.schema(),
- extensions,
- filter,
- order_by,
- distinct,
- )
- .await
- }
- None => not_impl_err!(
- "Aggregate without aggregate function is not supported"
- ),
+ let order_by = if !f.sorts.is_empty() {
+ Some(
+ from_substrait_sorts(consumer, &f.sorts, input.schema())
+ .await?,
+ )
+ } else {
+ None
};
- aggr_exprs.push(agg_func?.as_ref().clone());
- }
- input.aggregate(group_exprs, aggr_exprs)?.build()
- } else {
- not_impl_err!("Aggregate without an input is not valid")
- }
- }
- Some(RelType::Join(join)) => {
- if join.post_join_filter.is_some() {
- return not_impl_err!(
- "JoinRel with post_join_filter is not yet supported"
- );
- }
- let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, join.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
-
- let join_type = from_substrait_jointype(join.r#type)?;
- // The join condition expression needs full input schema and not the output schema from join since we lose columns from
- // certain join types such as semi and anti joins
- let in_join_schema = left.schema().join(right.schema())?;
-
- // If join expression exists, parse the `on` condition expression, build join and return
- // Otherwise, build join with only the filter, without join keys
- match &join.expression.as_ref() {
- Some(expr) => {
- let on = from_substrait_rex(state, expr, &in_join_schema, extensions)
- .await?;
- // The join expression can contain both equal and non-equal ops.
- // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
- // So we extract each part as follows:
- // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
- // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
- let (join_ons, nulls_equal_nulls, join_filter) =
- split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
- let (left_cols, right_cols): (Vec<_>, Vec<_>) =
- itertools::multiunzip(join_ons);
- left.join_detailed(
- right.build()?,
- join_type,
- (left_cols, right_cols),
- join_filter,
- nulls_equal_nulls,
- )?
- .build()
+ from_substrait_agg_func(
+ consumer,
+ f,
+ input.schema(),
+ filter,
+ order_by,
+ distinct,
+ )
+ .await
}
None => {
- let on: Vec<String> = vec![];
- left.join_detailed(
- right.build()?,
- join_type,
- (on.clone(), on),
- None,
- false,
- )?
- .build()
+ not_impl_err!("Aggregate without aggregate function is not supported")
}
- }
+ };
+ aggr_exprs.push(agg_func?.as_ref().clone());
}
- Some(RelType::Cross(cross)) => {
- let left = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.left.as_ref().unwrap(), extensions)
- .await?,
- );
- let right = LogicalPlanBuilder::from(
- from_substrait_rel(state, cross.right.as_ref().unwrap(), extensions)
- .await?,
- );
- let (left, right) = requalify_sides_if_needed(left, right)?;
- left.cross_join(right.build()?)?.build()
+ input.aggregate(group_exprs, aggr_exprs)?.build()
+ } else {
+ not_impl_err!("Aggregate without an input is not valid")
+ }
+}
+
+pub async fn from_join_rel(
+ consumer: &impl SubstraitConsumer,
+ join: &JoinRel,
+) -> Result<LogicalPlan> {
+ if join.post_join_filter.is_some() {
+ return not_impl_err!("JoinRel with post_join_filter is not yet supported");
+ }
+
+ let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, join.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+
+ let join_type = from_substrait_jointype(join.r#type)?;
+ // The join condition expression needs full input schema and not the output schema from join since we lose columns from
+ // certain join types such as semi and anti joins
+ let in_join_schema = left.schema().join(right.schema())?;
+
+ // If join expression exists, parse the `on` condition expression, build join and return
+ // Otherwise, build join with only the filter, without join keys
+ match &join.expression.as_ref() {
+ Some(expr) => {
+ let on = from_substrait_rex(consumer, expr, &in_join_schema).await?;
+ // The join expression can contain both equal and non-equal ops.
+ // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields.
+ // So we extract each part as follows:
+ // - If an Eq or IsNotDistinctFrom op is encountered, add the left column, right column and is_null_equal_nulls to `join_ons` vector
+ // - Otherwise we add the expression to join_filter (use conjunction if filter already exists)
+ let (join_ons, nulls_equal_nulls, join_filter) =
+ split_eq_and_noneq_join_predicate_with_nulls_equality(&on);
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) =
+ itertools::multiunzip(join_ons);
+ left.join_detailed(
+ right.build()?,
+ join_type,
+ (left_cols, right_cols),
+ join_filter,
+ nulls_equal_nulls,
+ )?
+ .build()
}
- Some(RelType::Read(read)) => {
- async fn read_with_schema(
- state: &dyn SubstraitPlanningState,
- table_ref: TableReference,
- schema: DFSchema,
- projection: &Option<MaskExpression>,
- ) -> Result<LogicalPlan> {
- let schema = schema.replace_qualifier(table_ref.clone());
-
- let plan = {
- let provider = match state.table(&table_ref).await? {
- Some(ref provider) => Arc::clone(provider),
- _ => return plan_err!("No table named '{table_ref}'"),
- };
+ None => {
+ let on: Vec<String> = vec![];
+ left.join_detailed(right.build()?, join_type, (on.clone(), on), None, false)?
+ .build()
+ }
+ }
+}
- LogicalPlanBuilder::scan(
- table_ref,
- provider_as_source(Arc::clone(&provider)),
- None,
- )?
- .build()?
- };
+pub async fn from_cross_rel(
+ consumer: &impl SubstraitConsumer,
+ cross: &CrossRel,
+) -> Result<LogicalPlan> {
+ let left = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.left.as_ref().unwrap()).await?,
+ );
+ let right = LogicalPlanBuilder::from(
+ from_substrait_rel(consumer, cross.right.as_ref().unwrap()).await?,
+ );
+ let (left, right) = requalify_sides_if_needed(left, right)?;
+ left.cross_join(right.build()?)?.build()
+}
- ensure_schema_compatability(plan.schema(), schema.clone())?;
+#[allow(deprecated)]
+pub async fn from_read_rel(
+ consumer: &impl SubstraitConsumer,
+ read: &ReadRel,
+) -> Result<LogicalPlan> {
+ async fn read_with_schema(
+ consumer: &impl SubstraitConsumer,
+ table_ref: TableReference,
+ schema: DFSchema,
+ projection: &Option<MaskExpression>,
+ ) -> Result<LogicalPlan> {
+ let schema = schema.replace_qualifier(table_ref.clone());
+
+ let plan = {
+ let provider = match consumer.resolve_table_ref(&table_ref).await? {
+ Some(ref provider) => Arc::clone(provider),
+ _ => return plan_err!("No table named '{table_ref}'"),
+ };
- let schema = apply_masking(schema, projection)?;
+ LogicalPlanBuilder::scan(
+ table_ref,
+ provider_as_source(Arc::clone(&provider)),
+ None,
+ )?
+ .build()?
+ };
- apply_projection(plan, schema)
- }
+ ensure_schema_compatability(plan.schema(), schema.clone())?;
- let named_struct = read.base_schema.as_ref().ok_or_else(|| {
- substrait_datafusion_err!("No base schema provided for Read Relation")
- })?;
+ let schema = apply_masking(schema, projection)?;
- let substrait_schema = from_substrait_named_struct(named_struct, extensions)?;
+ apply_projection(plan, schema)
+ }
- match &read.as_ref().read_type {
- Some(ReadType::NamedTable(nt)) => {
- let table_reference = match nt.names.len() {
- 0 => {
- return plan_err!("No table name found in NamedTable");
- }
- 1 => TableReference::Bare {
- table: nt.names[0].clone().into(),
- },
- 2 => TableReference::Partial {
- schema: nt.names[0].clone().into(),
- table: nt.names[1].clone().into(),
- },
- _ => TableReference::Full {
- catalog: nt.names[0].clone().into(),
- schema: nt.names[1].clone().into(),
- table: nt.names[2].clone().into(),
- },
- };
+ let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+ substrait_datafusion_err!("No base schema provided for Read Relation")
+ })?;
- read_with_schema(
- state,
- table_reference,
- substrait_schema,
- &read.projection,
- )
- .await
+ let substrait_schema = from_substrait_named_struct(consumer, named_struct)?;
+
+ match &read.read_type {
+ Some(ReadType::NamedTable(nt)) => {
+ let table_reference = match nt.names.len() {
+ 0 => {
+ return plan_err!("No table name found in NamedTable");
}
- Some(ReadType::VirtualTable(vt)) => {
- if vt.values.is_empty() {
- return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: DFSchemaRef::new(substrait_schema),
- }));
+ 1 => TableReference::Bare {
+ table: nt.names[0].clone().into(),
+ },
+ 2 => TableReference::Partial {
+ schema: nt.names[0].clone().into(),
+ table: nt.names[1].clone().into(),
+ },
+ _ => TableReference::Full {
+ catalog: nt.names[0].clone().into(),
+ schema: nt.names[1].clone().into(),
+ table: nt.names[2].clone().into(),
+ },
+ };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ Some(ReadType::VirtualTable(vt)) => {
+ if vt.values.is_empty() {
+ return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row: false,
+ schema: DFSchemaRef::new(substrait_schema),
+ }));
+ }
+
+ let values = vt
+ .values
+ .iter()
+ .map(|row| {
+ let mut name_idx = 0;
+ let lits = row
+ .fields
+ .iter()
+ .map(|lit| {
+ name_idx += 1; // top-level names are provided through schema
+ Ok(Expr::Literal(from_substrait_literal(
+ consumer,
+ lit,
+ &named_struct.names,
+ &mut name_idx,
+ )?))
+ })
+ .collect::<Result<_>>()?;
+ if name_idx != named_struct.names.len() {
+ return substrait_err!(
+ "Names list must match exactly to nested schema, but found {} uses for {} names",
+ name_idx,
+ named_struct.names.len()
+ );
}
+ Ok(lits)
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(LogicalPlan::Values(Values {
+ schema: DFSchemaRef::new(substrait_schema),
+ values,
+ }))
+ }
+ Some(ReadType::LocalFiles(lf)) => {
+ fn extract_filename(name: &str) -> Option<String> {
+ let corrected_url =
+ if name.starts_with("file://") && !name.starts_with("file:///") {
+ name.replacen("file://", "file:///", 1)
+ } else {
+ name.to_string()
+ };
+
+ Url::parse(&corrected_url).ok().and_then(|url| {
+ let path = url.path();
+ std::path::Path::new(path)
+ .file_name()
+ .map(|filename| filename.to_string_lossy().to_string())
+ })
+ }
+
+ // we could use the file name to check the original table provider
+ // TODO: currently does not support multiple local files
+ let filename: Option<String> =
+ lf.items.first().and_then(|x| match x.path_type.as_ref() {
+ Some(UriFile(name)) => extract_filename(name),
+ _ => None,
+ });
+
+ if lf.items.len() > 1 || filename.is_none() {
+ return not_impl_err!("Only single file reads are supported");
+ }
+ let name = filename.unwrap();
+ // directly use unwrap here since we could determine it is a valid one
+ let table_reference = TableReference::Bare { table: name.into() };
+
+ read_with_schema(
+ consumer,
+ table_reference,
+ substrait_schema,
+ &read.projection,
+ )
+ .await
+ }
+ _ => {
+ not_impl_err!("Unsupported ReadType: {:?}", read.read_type)
+ }
+ }
+}
+
+pub async fn from_set_rel(
+ consumer: &impl SubstraitConsumer,
+ set: &SetRel,
+) -> Result<LogicalPlan> {
+ if set.inputs.len() < 2 {
+ substrait_err!("Set operation requires at least two inputs")
+ } else {
+ match set.op() {
+ SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
+ SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
+ SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect(
+ from_substrait_rel(consumer, &set.inputs[0]).await?,
+ union_rels(consumer, &set.inputs[1..], true).await?,
+ false,
+ ),
+ SetOp::IntersectionMultiset => {
+ intersect_rels(consumer, &set.inputs, false).await
+ }
+ SetOp::IntersectionMultisetAll => {
+ intersect_rels(consumer, &set.inputs, true).await
+ }
+ SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
+ SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+ set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+ }
+ }
+}
- let values = vt
- .values
- .iter()
- .map(|row| {
- let mut name_idx = 0;
- let lits = row
- .fields
- .iter()
- .map(|lit| {
- name_idx += 1; // top-level names are provided through schema
- Ok(Expr::Literal(from_substrait_literal(
- lit,
- extensions,
- &named_struct.names,
- &mut name_idx,
- )?))
- })
- .collect::<Result<_>>()?;
- if name_idx != named_struct.names.len() {
- return substrait_err!(
- "Names list must match exactly to nested schema, but found {} uses for {} names",
- name_idx,
- named_struct.names.len()
- );
- }
- Ok(lits)
- })
- .collect::<Result<_>>()?;
+pub async fn from_extension_leaf_rel(
Review Comment:
These were leftover from an earlier approach I took to this and aren't needed anymore. I've removed them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1890179391
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +102,400 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::SessionState;
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// use datafusion_substrait::logical_plan::state::SubstraitPlanningState;
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[async_trait]
+/// impl SubstraitConsumer for CustomSubstraitConsumer {
+/// async fn resolve_table_ref(
+/// &self,
+/// table_ref: &TableReference,
+/// ) -> Result<Option<Arc<dyn TableProvider>>> {
+/// self.state.table(table_ref).await
+/// }
+///
+/// fn get_extensions(&self) -> &Extensions {
+/// self.extensions.as_ref()
+/// }
+///
+/// fn get_state(&self) -> &SessionState {
+/// self.state.as_ref()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// async fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// async fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and SessionState is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_state(&self) -> &SessionState;
Review Comment:
It looks like `get_state()` is only used for getting udf/udaf/udwfs? So if you'd wish, you could replace SessionState with FunctionRegistry. Then the change would be compatible with anyone using the current `SubstraitPlanningState` rather than a SessionState, since SubstraitPlanningState requires FunctionRegistry.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1890823837
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -94,13 +102,400 @@ use substrait::proto::{
join_rel, plan_rel, r#type,
read_rel::ReadType,
rel::RelType,
- rel_common, set_rel,
+ rel_common,
sort_field::{SortDirection, SortKind::*},
- AggregateFunction, Expression, NamedStruct, Plan, Rel, RelCommon, Type,
+ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel,
+ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
+ FilterRel, JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel,
+ SortRel, Type,
};
use substrait::proto::{ExtendedExpression, FunctionArgument, SortField};
-use super::state::SubstraitPlanningState;
+#[async_trait]
+/// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
+/// It can be implemented by users to allow for custom handling of relations, expressions, etc.
+///
+/// # Example Usage
+///
+/// ```
+/// use async_trait::async_trait;
+/// use datafusion::catalog::TableProvider;
+/// use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference};
+/// use datafusion::error::Result;
+/// use datafusion::execution::SessionState;
+/// use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder};
+/// use std::sync::Arc;
+/// use substrait::proto;
+/// use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel};
+/// use datafusion::arrow::datatypes::DataType;
+/// use datafusion::logical_expr::expr::ScalarFunction;
+/// use datafusion_substrait::extensions::Extensions;
+/// use datafusion_substrait::logical_plan::consumer::{
+/// from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// };
+///
+/// use datafusion_substrait::logical_plan::state::SubstraitPlanningState;
+///
+/// struct CustomSubstraitConsumer {
+/// extensions: Arc<Extensions>,
+/// state: Arc<SessionState>,
+/// }
+///
+/// #[async_trait]
+/// impl SubstraitConsumer for CustomSubstraitConsumer {
+/// async fn resolve_table_ref(
+/// &self,
+/// table_ref: &TableReference,
+/// ) -> Result<Option<Arc<dyn TableProvider>>> {
+/// self.state.table(table_ref).await
+/// }
+///
+/// fn get_extensions(&self) -> &Extensions {
+/// self.extensions.as_ref()
+/// }
+///
+/// fn get_state(&self) -> &SessionState {
+/// self.state.as_ref()
+/// }
+///
+/// // You can reuse existing consumer code to assist in handling advanced extensions
+/// async fn consume_project(&self, rel: &ProjectRel) -> Result<LogicalPlan> {
+/// let df_plan = from_project_rel(self, rel).await?;
+/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() {
+/// not_impl_err!(
+/// "decode and handle an advanced extension: {:?}",
+/// advanced_extension
+/// )
+/// } else {
+/// Ok(df_plan)
+/// }
+/// }
+///
+/// // You can implement a fully custom consumer method if you need special handling
+/// async fn consume_filter(&self, rel: &FilterRel) -> Result<LogicalPlan> {
+/// let input = from_substrait_rel(self, rel.input.as_ref().unwrap()).await?;
+/// let expression =
+/// from_substrait_rex(self, rel.condition.as_ref().unwrap(), input.schema())
+/// .await?;
+/// // though this one is quite boring
+/// LogicalPlanBuilder::from(input).filter(expression)?.build()
+/// }
+///
+/// // You can add handlers for extension relations
+/// async fn consume_extension_leaf(
+/// &self,
+/// rel: &ExtensionLeafRel,
+/// ) -> Result<LogicalPlan> {
+/// not_impl_err!(
+/// "handle protobuf Any {} as you need",
+/// rel.detail.as_ref().unwrap().type_url
+/// )
+/// }
+///
+/// // and handlers for user-define types
+/// async fn consume_user_defined_type(&self, typ: &proto::r#type::UserDefined) -> Result<DataType> {
+/// let type_string = self.extensions.types.get(&typ.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+///
+/// // and user-defined literals
+/// async fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result<ScalarValue> {
+/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap();
+/// match type_string.as_str() {
+/// "u!foo" => not_impl_err!("handle foo conversion"),
+/// "u!bar" => not_impl_err!("handle bar conversion"),
+/// _ => substrait_err!("unexpected type")
+/// }
+/// }
+/// }
+/// ```
+///
+pub trait SubstraitConsumer: Send + Sync + Sized {
+ async fn resolve_table_ref(
+ &self,
+ table_ref: &TableReference,
+ ) -> Result<Option<Arc<dyn TableProvider>>>;
+
+ // TODO: Remove these two methods
+ // Ideally, the abstract consumer should not place any constraints on implementations.
+ // The functionality for which the Extensions and SessionState is needed should be abstracted
+ // out into methods on the trait. As an example, resolve_table_reference is such a method.
+ fn get_extensions(&self) -> &Extensions;
+ fn get_state(&self) -> &SessionState;
Review Comment:
Good idea, it's a more precise constraint to have.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#issuecomment-2557967567
FYI @ccciudatu
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(substrait): modular substrait consumer [datafusion]
Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #13803:
URL: https://github.com/apache/datafusion/pull/13803#discussion_r1890163208
##########
datafusion/substrait/src/logical_plan/consumer.rs:
##########
@@ -1503,370 +1858,416 @@ pub async fn from_substrait_func_args(
/// Convert Substrait AggregateFunction to DataFusion Expr
pub async fn from_substrait_agg_func(
- state: &dyn SubstraitPlanningState,
+ consumer: &impl SubstraitConsumer,
f: &AggregateFunction,
input_schema: &DFSchema,
- extensions: &Extensions,
filter: Option<Box<Expr>>,
order_by: Option<Vec<SortExpr>>,
distinct: bool,
) -> Result<Arc<Expr>> {
- let args =
- from_substrait_func_args(state, &f.arguments, input_schema, extensions).await?;
-
- let Some(function_name) = extensions.functions.get(&f.function_reference) else {
+ let Some(fn_signature) = consumer
+ .get_extensions()
Review Comment:
> scalar functions first try and lookup the function by name, then look it up as an operator and then try and treat is a BuiltinExpr.
Yeah, I think this is a side-effect of DF not having (yet?) converged fully on UDFs. I think for aggregates and windows the built-in modes were already removed in favor of UDAF/UDWF but not for scalars.
But yep, totally good to do as a followup someday!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org