Skip to content

Commit

Permalink
generate new projection plan in inline_table_scan instead of discardi…
Browse files Browse the repository at this point in the history
…ng (apache#5371)

* fix inline_table_scan

* add test
  • Loading branch information
jackwener authored and jiangzhx committed Feb 24, 2023
1 parent 4b6e882 commit abfdb5d
Showing 1 changed file with 42 additions and 3 deletions.
45 changes: 42 additions & 3 deletions datafusion/optimizer/src/inline_table_scan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,11 +49,14 @@ impl OptimizerRule for InlineTableScan {
source,
table_name,
filters,
projection,
..
}) if filters.is_empty() => {
if let Some(sub_plan) = source.get_logical_plan() {
let projection_exprs =
generate_projection_expr(projection, sub_plan)?;
let plan = LogicalPlanBuilder::from(sub_plan.clone())
.project(vec![Expr::Wildcard])?
.project(projection_exprs)?
.alias(table_name)?;
Ok(Some(plan.build()?))
} else {
Expand All @@ -73,6 +76,23 @@ impl OptimizerRule for InlineTableScan {
}
}

fn generate_projection_expr(
projection: &Option<Vec<usize>>,
sub_plan: &LogicalPlan,
) -> Result<Vec<Expr>> {
let mut exprs = vec![];
if let Some(projection) = projection {
for i in projection {
exprs.push(Expr::Column(
sub_plan.schema().fields()[*i].qualified_column(),
));
}
} else {
exprs.push(Expr::Wildcard);
}
Ok(exprs)
}

#[cfg(test)]
mod tests {
use std::{sync::Arc, vec};
Expand All @@ -91,7 +111,10 @@ mod tests {
}

fn schema(&self) -> arrow::datatypes::SchemaRef {
Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]))
Arc::new(Schema::new(vec![
Field::new("a", DataType::Int64, false),
Field::new("b", DataType::Int64, false),
]))
}

fn supports_filter_pushdown(
Expand Down Expand Up @@ -150,9 +173,25 @@ mod tests {
let plan = scan.filter(col("x.a").eq(lit(1)))?.build()?;
let expected = "Filter: x.a = Int32(1)\
\n SubqueryAlias: x\
\n Projection: y.a\
\n Projection: y.a, y.b\
\n TableScan: y";

assert_optimized_plan_eq(Arc::new(InlineTableScan::new()), &plan, expected)
}

#[test]
fn inline_table_scan_with_projection() -> datafusion_common::Result<()> {
let scan = LogicalPlanBuilder::scan(
"x".to_string(),
Arc::new(CustomSource::new()),
Some(vec![0]),
)?;

let plan = scan.build()?;
let expected = "SubqueryAlias: x\
\n Projection: y.a\
\n TableScan: y";

assert_optimized_plan_eq(Arc::new(InlineTableScan::new()), &plan, expected)
}
}

0 comments on commit abfdb5d

Please sign in to comment.