Skip to content

Support SubqueryAlias in optimizer, physcial planner. #4293

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
alamb merged 1 commit into from
Nov 23, 2022
Merged

Support SubqueryAlias in optimizer, physcial planner. #4293

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 1 addition & 6 deletions datafusion/core/src/physical_plan/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -969,12 +969,7 @@ impl DefaultPhysicalPlanner {
SchemaRef::new(schema.as_ref().to_owned().into()),
))),
LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => {
match input.as_ref() {
LogicalPlan::TableScan(..) => {
self.create_initial_plan(input, session_state).await
}
_ => Err(DataFusionError::Plan("SubqueryAlias should only wrap TableScan".to_string()))
}
self.create_initial_plan(input, session_state).await
}
LogicalPlan::Limit(Limit { input, skip, fetch, .. }) => {
let input = self.create_initial_plan(input, session_state).await?;
Expand Down
201 changes: 97 additions & 104 deletions datafusion/optimizer/src/projection_push_down.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
//! loaded into memory

use crate::{OptimizerConfig, OptimizerRule};
use arrow::datatypes::{Field, Schema};
use arrow::datatypes::Field;
use arrow::error::Result as ArrowResult;
use datafusion_common::{
Column, DFField, DFSchema, DFSchemaRef, DataFusionError, Result, ToDFSchema,
Expand All @@ -34,6 +34,7 @@ use datafusion_expr::{
utils::{expr_to_columns, exprlist_to_columns, find_sort_exprs, from_plan},
Expr,
};
use std::collections::HashMap;
use std::{
collections::{BTreeSet, HashSet},
sync::Arc,
Expand Down Expand Up @@ -72,60 +73,6 @@ impl ProjectionPushDown {
}
}

fn get_projected_schema(
table_name: Option<&String>,
schema: &Schema,
required_columns: &HashSet<Column>,
has_projection: bool,
) -> Result<(Vec<usize>, DFSchemaRef)> {
// once we reach the table scan, we can use the accumulated set of column
// names to construct the set of column indexes in the scan
//
// we discard non-existing columns because some column names are not part of the schema,
// e.g. when the column derives from an aggregation
//
// Use BTreeSet to remove potential duplicates (e.g. union) as
// well as to sort the projection to ensure deterministic behavior
let mut projection: BTreeSet<usize> = required_columns
.iter()
.filter(|c| c.relation.is_none() || c.relation.as_ref() == table_name)
.map(|c| schema.index_of(&c.name))
.filter_map(ArrowResult::ok)
.collect();

if projection.is_empty() {
if has_projection && !schema.fields().is_empty() {
// Ensure that we are reading at least one column from the table in case the query
// does not reference any columns directly such as "SELECT COUNT(1) FROM table",
// except when the table is empty (no column)
projection.insert(0);
} else {
// for table scan without projection, we default to return all columns
projection = schema
.fields()
.iter()
.enumerate()
.map(|(i, _)| i)
.collect::<BTreeSet<usize>>();
}
}

// create the projected schema
let projected_fields: Vec<DFField> = match table_name {
Some(qualifier) => projection
.iter()
.map(|i| DFField::from_qualified(qualifier, schema.fields()[*i].clone()))
.collect(),
None => projection
.iter()
.map(|i| DFField::from(schema.fields()[*i].clone()))
.collect(),
};

let projection = projection.into_iter().collect::<Vec<_>>();
Ok((projection, projected_fields.to_dfschema_ref()?))
}

/// Recursively transverses the logical plan removing expressions and that are not needed.
fn optimize_plan(
_optimizer: &ProjectionPushDown,
Expand Down Expand Up @@ -348,28 +295,8 @@ fn optimize_plan(
}
// scans:
// * remove un-used columns from the scan projection
LogicalPlan::TableScan(TableScan {
table_name,
source,
filters,
fetch: limit,
..
}) => {
let (projection, projected_schema) = get_projected_schema(
Some(table_name),
&source.schema(),
required_columns,
has_projection,
)?;
// return the table scan with projection
Ok(LogicalPlan::TableScan(TableScan {
table_name: table_name.clone(),
source: source.clone(),
projection: Some(projection),
projected_schema,
filters: filters.clone(),
fetch: *limit,
}))
LogicalPlan::TableScan(scan) => {
push_down_scan(scan, &new_required_columns, has_projection)
}
LogicalPlan::Explain { .. } => Err(DataFusionError::Internal(
"Unsupported logical plan: Explain must be root of the plan".to_string(),
Expand Down Expand Up @@ -441,32 +368,16 @@ fn optimize_plan(
}))
}
LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) => {
match input.as_ref() {
LogicalPlan::TableScan(TableScan { table_name, .. }) => {
let new_required_columns = new_required_columns
.iter()
.map(|c| match &c.relation {
Some(q) if q == alias => Column {
relation: Some(table_name.clone()),
name: c.name.clone(),
},
_ => c.clone(),
})
.collect();
let new_inputs = vec![optimize_plan(
_optimizer,
input,
&new_required_columns,
has_projection,
_optimizer_config,
)?];
let expr = vec![];
from_plan(plan, &expr, &new_inputs)
}
_ => Err(DataFusionError::Plan(
"SubqueryAlias should only wrap TableScan".to_string(),
)),
}
let new_required_columns =
Copy link
Contributor

@alamb alamb Nov 23, 2022

Choose a reason for hiding this comment

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

this is a nice cleanup 👍

replace_alias(required_columns, alias, input.schema());
let child = optimize_plan(
_optimizer,
input,
&new_required_columns,
has_projection,
_optimizer_config,
)?;
from_plan(plan, &plan.expressions(), &[child])
}
// all other nodes: Add any additional columns used by
// expressions in this node to the list of required columns
Expand Down Expand Up @@ -532,11 +443,93 @@ fn projection_equal(p: &Projection, p2: &Projection) -> bool {
&& p.expr.iter().zip(&p2.expr).all(|(l, r)| l == r)
}

fn replace_alias(
required_columns: &HashSet<Column>,
alias: &str,
input_schema: &DFSchemaRef,
) -> HashSet<Column> {
let mut map = HashMap::new();
for field in input_schema.fields() {
let col = field.qualified_column();
let alias_col = Column {
relation: Some(alias.to_owned()),
name: col.name.clone(),
};
map.insert(alias_col, col);
}
required_columns
.iter()
.map(|col| map.get(col).unwrap_or(col).clone())
.collect::<HashSet<_>>()
}

fn push_down_scan(
scan: &TableScan,
required_columns: &HashSet<Column>,
has_projection: bool,
) -> Result<LogicalPlan> {
// once we reach the table scan, we can use the accumulated set of column
// names to construct the set of column indexes in the scan
//
// we discard non-existing columns because some column names are not part of the schema,
// e.g. when the column derives from an aggregation
//
// Use BTreeSet to remove potential duplicates (e.g. union) as
// well as to sort the projection to ensure deterministic behavior
let schema = scan.source.schema();
let mut projection: BTreeSet<usize> = required_columns
.iter()
.filter(|c| {
c.relation.is_none() || c.relation.as_ref().unwrap() == &scan.table_name
})
.map(|c| schema.index_of(&c.name))
.filter_map(ArrowResult::ok)
.collect();

if projection.is_empty() {
if has_projection && !schema.fields().is_empty() {
// Ensure that we are reading at least one column from the table in case the query
// does not reference any columns directly such as "SELECT COUNT(1) FROM table",
// except when the table is empty (no column)
projection.insert(0);
} else {
// for table scan without projection, we default to return all columns
projection = scan
.source
.schema()
.fields()
.iter()
.enumerate()
.map(|(i, _)| i)
.collect::<BTreeSet<usize>>();
}
}

// create the projected schema
let projected_fields: Vec<DFField> = projection
.iter()
.map(|i| DFField::from_qualified(&scan.table_name, schema.fields()[*i].clone()))
.collect();

let projection = projection.into_iter().collect::<Vec<_>>();
let projected_schema = projected_fields.to_dfschema_ref()?;

// return the table scan with projection
Ok(LogicalPlan::TableScan(TableScan {
table_name: scan.table_name.clone(),
source: scan.source.clone(),
projection: Some(projection),
projected_schema,
filters: scan.filters.clone(),
fetch: scan.fetch,
}))
}

#[cfg(test)]
mod tests {
use super::*;
use crate::test::*;
use arrow::datatypes::DataType;
use arrow::datatypes::{DataType, Schema};
use datafusion_expr::expr::Cast;
use datafusion_expr::{
col, count, lit,
Expand Down