Skip to content

[Substrait] Change API to return LogicalPlan instead of DataFrame #4896

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Dandandan merged 4 commits into from
Jan 15, 2023
Merged

[Substrait] Change API to return LogicalPlan instead of DataFrame #4896

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
01fb4fd
Change API to return LogicalPlan instead of DataFrame
andygrove Jan 13, 2023
68d2f7c
optimize plans
andygrove Jan 13, 2023
158c37a
fix
andygrove Jan 13, 2023
f47a34d
ignore test
andygrove Jan 13, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 36 additions & 24 deletions datafusion/substrait/src/consumer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,17 @@

use async_recursion::async_recursion;
use datafusion::common::{DFField, DFSchema, DFSchemaRef};
use datafusion::logical_expr::build_join_schema;
use datafusion::logical_expr::expr;
use datafusion::logical_expr::{
aggregate_function, BinaryExpr, Case, Expr, LogicalPlan, Operator,
};
use datafusion::logical_expr::{build_join_schema, LogicalPlanBuilder};
use datafusion::prelude::JoinType;
use datafusion::sql::TableReference;
use datafusion::{
error::{DataFusionError, Result},
optimizer::utils::split_conjunction,
prelude::{Column, DataFrame, SessionContext},
prelude::{Column, SessionContext},
scalar::ScalarValue,
};
use substrait::protobuf::{
Expand Down Expand Up @@ -87,7 +87,7 @@ pub fn name_to_op(name: &str) -> Result<Operator> {
pub async fn from_substrait_plan(
ctx: &mut SessionContext,
plan: &Plan,
) -> Result<DataFrame> {
) -> Result<LogicalPlan> {
// Register function extension
let function_extension = plan
.extensions
Expand Down Expand Up @@ -135,17 +135,19 @@ pub async fn from_substrait_rel(
ctx: &mut SessionContext,
rel: &Rel,
extensions: &HashMap<u32, &String>,
) -> Result<DataFrame> {
) -> Result<LogicalPlan> {
match &rel.rel_type {
Some(RelType::Project(p)) => {
if let Some(input) = p.input.as_ref() {
let input = from_substrait_rel(ctx, input, extensions).await?;
let input = LogicalPlanBuilder::from(
from_substrait_rel(ctx, input, extensions).await?,
);
let mut exprs: Vec<Expr> = vec![];
for e in &p.expressions {
let x = from_substrait_rex(e, input.schema(), extensions).await?;
exprs.push(x.as_ref().clone());
}
input.select(exprs)
input.project(exprs)?.build()
} else {
Err(DataFusionError::NotImplemented(
"Projection without an input is not supported".to_string(),
Expand All @@ -154,11 +156,13 @@ pub async fn from_substrait_rel(
}
Some(RelType::Filter(filter)) => {
if let Some(input) = filter.input.as_ref() {
let input = from_substrait_rel(ctx, input, extensions).await?;
let input = LogicalPlanBuilder::from(
from_substrait_rel(ctx, input, extensions).await?,
);
if let Some(condition) = filter.condition.as_ref() {
let expr =
from_substrait_rex(condition, input.schema(), extensions).await?;
input.filter(expr.as_ref().clone())
input.filter(expr.as_ref().clone())?.build()
} else {
Err(DataFusionError::NotImplemented(
"Filter without an condition is not valid".to_string(),
Expand All @@ -172,10 +176,12 @@ pub async fn from_substrait_rel(
}
Some(RelType::Fetch(fetch)) => {
if let Some(input) = fetch.input.as_ref() {
let input = from_substrait_rel(ctx, input, extensions).await?;
let input = LogicalPlanBuilder::from(
from_substrait_rel(ctx, input, extensions).await?,
);
let offset = fetch.offset as usize;
let count = fetch.count as usize;
input.limit(offset, Some(count))
input.limit(offset, Some(count))?.build()
} else {
Err(DataFusionError::NotImplemented(
"Fetch without an input is not valid".to_string(),
Expand All @@ -184,7 +190,9 @@ pub async fn from_substrait_rel(
}
Some(RelType::Sort(sort)) => {
if let Some(input) = sort.input.as_ref() {
let input = from_substrait_rel(ctx, input, extensions).await?;
let input = LogicalPlanBuilder::from(
from_substrait_rel(ctx, input, extensions).await?,
);
let mut sorts: Vec<Expr> = vec![];
for s in &sort.sorts {
let expr = from_substrait_rex(
Expand Down Expand Up @@ -224,7 +232,7 @@ pub async fn from_substrait_rel(
nulls_first,
}));
}
input.sort(sorts)
input.sort(sorts)?.build()
} else {
Err(DataFusionError::NotImplemented(
"Sort without an input is not valid".to_string(),
Expand All @@ -233,7 +241,9 @@ pub async fn from_substrait_rel(
}
Some(RelType::Aggregate(agg)) => {
if let Some(input) = agg.input.as_ref() {
let input = from_substrait_rel(ctx, input, extensions).await?;
let input = LogicalPlanBuilder::from(
from_substrait_rel(ctx, input, extensions).await?,
);
let mut group_expr = vec![];
let mut aggr_expr = vec![];

Expand Down Expand Up @@ -292,18 +302,20 @@ pub async fn from_substrait_rel(
aggr_expr.push(agg_func?.as_ref().clone());
}

input.aggregate(group_expr, aggr_expr)
input.aggregate(group_expr, aggr_expr)?.build()
} else {
Err(DataFusionError::NotImplemented(
"Aggregate without an input is not valid".to_string(),
))
}
}
Some(RelType::Join(join)) => {
let left =
from_substrait_rel(ctx, join.left.as_ref().unwrap(), extensions).await?;
let right =
from_substrait_rel(ctx, join.right.as_ref().unwrap(), extensions).await?;
let left = LogicalPlanBuilder::from(
from_substrait_rel(ctx, join.left.as_ref().unwrap(), extensions).await?,
);
let right = LogicalPlanBuilder::from(
from_substrait_rel(ctx, join.right.as_ref().unwrap(), extensions).await?,
);
let join_type = match join.r#type {
1 => JoinType::Inner,
2 => JoinType::Left,
Expand Down Expand Up @@ -346,9 +358,9 @@ pub async fn from_substrait_rel(
)),
})
.collect::<Result<Vec<_>>>()?;
let left_cols: Vec<&str> = pairs.iter().map(|(l, _)| l.as_str()).collect();
let right_cols: Vec<&str> = pairs.iter().map(|(_, r)| r.as_str()).collect();
left.join(right, join_type, &left_cols, &right_cols, None)
let (left_cols, right_cols): (Vec<_>, Vec<_>) = pairs.iter().cloned().unzip();
left.join(right.build()?, join_type, (left_cols, right_cols), None)?
.build()
}
Some(RelType::Read(read)) => match &read.as_ref().read_type {
Some(ReadType::NamedTable(nt)) => {
Expand All @@ -372,6 +384,7 @@ pub async fn from_substrait_rel(
},
};
let t = ctx.table(table_reference).await?;
let t = t.into_optimized_plan()?;
match &read.projection {
Some(MaskExpression { select, .. }) => match &select.as_ref() {
Some(projection) => {
Expand All @@ -380,7 +393,7 @@ pub async fn from_substrait_rel(
.iter()
.map(|item| item.field as usize)
.collect();
match t.into_optimized_plan()? {
match &t {
LogicalPlan::TableScan(scan) => {
let fields: Vec<DFField> = column_indices
.iter()
Expand All @@ -395,8 +408,7 @@ pub async fn from_substrait_rel(
fields,
HashMap::new(),
)?);
let plan = LogicalPlan::TableScan(scan);
Ok(DataFrame::new(ctx.state(), plan))
Ok(LogicalPlan::TableScan(scan))
}
_ => Err(DataFusionError::Internal(
"unexpected plan for table".to_string(),
Expand Down
22 changes: 8 additions & 14 deletions datafusion/substrait/tests/roundtrip.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ mod tests {
roundtrip("SELECT * FROM data WHERE d AND a > 1").await
}

#[ignore] // tracked in https://github.com/apache/arrow-datafusion/issues/4897
#[tokio::test]
async fn select_with_limit() -> Result<()> {
roundtrip_fill_na("SELECT * FROM data LIMIT 100").await
Expand Down Expand Up @@ -186,8 +187,7 @@ mod tests {
let df = ctx.sql(sql).await?;
let plan = df.into_optimized_plan()?;
let proto = to_substrait_plan(&plan)?;
let df = from_substrait_plan(&mut ctx, &proto).await?;
let plan2 = df.into_optimized_plan()?;
let plan2 = from_substrait_plan(&mut ctx, &proto).await?;
let plan2str = format!("{:?}", plan2);
assert_eq!(expected_plan_str, &plan2str);
Ok(())
Expand All @@ -198,9 +198,7 @@ mod tests {
let df = ctx.sql(sql).await?;
let plan1 = df.into_optimized_plan()?;
let proto = to_substrait_plan(&plan1)?;

let df = from_substrait_plan(&mut ctx, &proto).await?;
let plan2 = df.into_optimized_plan()?;
let plan2 = from_substrait_plan(&mut ctx, &proto).await?;

// Format plan string and replace all None's with 0
let plan1str = format!("{:?}", plan1).replace("None", "0");
Expand All @@ -218,15 +216,11 @@ mod tests {

let df_a = ctx.sql(sql_with_alias).await?;
let proto_a = to_substrait_plan(&df_a.into_optimized_plan()?)?;
let plan_with_alias = from_substrait_plan(&mut ctx, &proto_a)
.await?
.into_optimized_plan()?;
let plan_with_alias = from_substrait_plan(&mut ctx, &proto_a).await?;

let df = ctx.sql(sql_no_alias).await?;
let proto = to_substrait_plan(&df.into_optimized_plan()?)?;
let plan = from_substrait_plan(&mut ctx, &proto)
.await?
.into_optimized_plan()?;
let plan = from_substrait_plan(&mut ctx, &proto).await?;

println!("{:#?}", plan_with_alias);
println!("{:#?}", plan);
Expand All @@ -237,14 +231,14 @@ mod tests {
Ok(())
}

#[allow(deprecated)]
async fn roundtrip(sql: &str) -> Result<()> {
let mut ctx = create_context().await?;
let df = ctx.sql(sql).await?;
let plan = df.into_optimized_plan()?;
let proto = to_substrait_plan(&plan)?;

let df = from_substrait_plan(&mut ctx, &proto).await?;
let plan2 = df.into_optimized_plan()?;
let plan2 = from_substrait_plan(&mut ctx, &proto).await?;
let plan2 = ctx.optimize(&plan2)?;

println!("{:#?}", plan);
println!("{:#?}", plan2);
Expand Down
5 changes: 3 additions & 2 deletions datafusion/substrait/tests/serialize.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,9 @@ mod tests {
// Read substrait plan from file
let proto = serializer::deserialize(path).await?;
// Check plan equality
let df = from_substrait_plan(&mut ctx, &proto).await?;
let plan = df.into_optimized_plan()?;
let plan = from_substrait_plan(&mut ctx, &proto).await?;
// #[allow(deprecated)]
// let plan = ctx.optimize(&plan)?;
Comment on lines +44 to +45

@waynexia waynexia Jan 14, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

some deadcode

let plan_str_ref = format!("{:?}", plan_ref);
let plan_str = format!("{:?}", plan);
assert_eq!(plan_str_ref, plan_str);
Expand Down