This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new ddbdf4bf7f Unparse TableScan with projections, filters or fetch to SQL
string (#12158)
ddbdf4bf7f is described below
commit ddbdf4bf7f215c01aca435f8edd00412e84aab60
Author: Jax Liu <[email protected]>
AuthorDate: Sat Sep 7 18:18:22 2024 +0800
Unparse TableScan with projections, filters or fetch to SQL string (#12158)
* support to unparse table scan with projection
* support to unparse table_scan with filter and fetch
* add the doc for public function
* clippy
* remove pub get method for table scan
* fix merge conflict
* add issue reference
* fix typo
---
datafusion/expr/src/logical_plan/builder.rs | 36 +++++++
datafusion/sql/src/unparser/plan.rs | 99 ++++++++++++++++--
datafusion/sql/tests/cases/plan_to_sql.rs | 154 +++++++++++++++++++++++++++-
3 files changed, 278 insertions(+), 11 deletions(-)
diff --git a/datafusion/expr/src/logical_plan/builder.rs
b/datafusion/expr/src/logical_plan/builder.rs
index fc961b83f7..f4dad08542 100644
--- a/datafusion/expr/src/logical_plan/builder.rs
+++ b/datafusion/expr/src/logical_plan/builder.rs
@@ -333,6 +333,19 @@ impl LogicalPlanBuilder {
.map(Self::new)
}
+ /// Convert a table provider into a builder with a TableScan with filter
and fetch
+ pub fn scan_with_filters_fetch(
+ table_name: impl Into<TableReference>,
+ table_source: Arc<dyn TableSource>,
+ projection: Option<Vec<usize>>,
+ filters: Vec<Expr>,
+ fetch: Option<usize>,
+ ) -> Result<Self> {
+ TableScan::try_new(table_name, table_source, projection, filters,
fetch)
+ .map(LogicalPlan::TableScan)
+ .map(Self::new)
+ }
+
/// Wrap a plan in a window
pub fn window_plan(
input: LogicalPlan,
@@ -1424,6 +1437,29 @@ pub fn table_scan_with_filters(
LogicalPlanBuilder::scan_with_filters(name, table_source, projection,
filters)
}
+/// Create a LogicalPlanBuilder representing a scan of a table with the
provided name and schema,
+/// filters, and inlined fetch.
+/// This is mostly used for testing and documentation.
+pub fn table_scan_with_filter_and_fetch(
+ name: Option<impl Into<TableReference>>,
+ table_schema: &Schema,
+ projection: Option<Vec<usize>>,
+ filters: Vec<Expr>,
+ fetch: Option<usize>,
+) -> Result<LogicalPlanBuilder> {
+ let table_source = table_source(table_schema);
+ let name = name
+ .map(|n| n.into())
+ .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
+ LogicalPlanBuilder::scan_with_filters_fetch(
+ name,
+ table_source,
+ projection,
+ filters,
+ fetch,
+ )
+}
+
fn table_source(table_schema: &Schema) -> Arc<dyn TableSource> {
let table_schema = Arc::new(table_schema.clone());
Arc::new(LogicalTableSource { table_schema })
diff --git a/datafusion/sql/src/unparser/plan.rs
b/datafusion/sql/src/unparser/plan.rs
index 478cfd0ecd..802d476257 100644
--- a/datafusion/sql/src/unparser/plan.rs
+++ b/datafusion/sql/src/unparser/plan.rs
@@ -15,16 +15,16 @@
// specific language governing permissions and limitations
// under the License.
+use crate::unparser::utils::unproject_agg_exprs;
use datafusion_common::{
- internal_err, not_impl_err, plan_err, Column, DataFusionError, Result,
+ internal_err, not_impl_err, plan_err, Column, DataFusionError, Result,
TableReference,
};
use datafusion_expr::{
- expr::Alias, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan,
Projection,
- SortExpr,
+ expr::Alias, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan,
+ LogicalPlanBuilder, Projection, SortExpr,
};
use sqlparser::ast::{self, Ident, SetExpr};
-
-use crate::unparser::utils::unproject_agg_exprs;
+use std::sync::Arc;
use super::{
ast::{
@@ -240,6 +240,19 @@ impl Unparser<'_> {
) -> Result<()> {
match plan {
LogicalPlan::TableScan(scan) => {
+ if scan.projection.is_some()
+ || !scan.filters.is_empty()
+ || scan.fetch.is_some()
+ {
+ let unparsed_table_scan =
+ Self::unparse_table_scan_pushdown(plan, None)?;
+ return self.select_to_sql_recursively(
+ &unparsed_table_scan,
+ query,
+ select,
+ relation,
+ );
+ }
let mut builder = TableRelationBuilder::default();
let mut table_parts = vec![];
if let Some(catalog_name) = scan.table_name.catalog() {
@@ -455,7 +468,10 @@ impl Unparser<'_> {
LogicalPlan::SubqueryAlias(plan_alias) => {
let (plan, mut columns) =
subquery_alias_inner_query_and_columns(plan_alias);
-
+ let plan = Self::unparse_table_scan_pushdown(
+ plan,
+ Some(plan_alias.alias.clone()),
+ )?;
if !columns.is_empty()
&& !self.dialect.supports_column_alias_in_table_alias()
{
@@ -467,7 +483,7 @@ impl Unparser<'_> {
};
// Instead of specifying column aliases as part of the
outer table, inject them directly into the inner projection
- let rewritten_plan = inject_column_aliases(inner_p,
columns);
+ let rewritten_plan = inject_column_aliases(&inner_p,
columns);
columns = vec![];
self.select_to_sql_recursively(
@@ -477,7 +493,7 @@ impl Unparser<'_> {
relation,
)?;
} else {
- self.select_to_sql_recursively(plan, query, select,
relation)?;
+ self.select_to_sql_recursively(&plan, query, select,
relation)?;
}
relation.alias(Some(
@@ -532,6 +548,73 @@ impl Unparser<'_> {
}
}
+ fn unparse_table_scan_pushdown(
+ plan: &LogicalPlan,
+ alias: Option<TableReference>,
+ ) -> Result<LogicalPlan> {
+ match plan {
+ LogicalPlan::TableScan(table_scan) => {
+ // TODO: support filters for table scan with alias. Remove
this check after #12368 issue.
+ // see the issue:
https://github.com/apache/datafusion/issues/12368
+ if alias.is_some() && !table_scan.filters.is_empty() {
+ return not_impl_err!(
+ "Subquery alias is not supported for table scan with
pushdown filters"
+ );
+ }
+
+ let mut builder = LogicalPlanBuilder::scan(
+ table_scan.table_name.clone(),
+ Arc::clone(&table_scan.source),
+ None,
+ )?;
+ if let Some(project_vec) = &table_scan.projection {
+ let project_columns = project_vec
+ .iter()
+ .cloned()
+ .map(|i| {
+ let (qualifier, field) =
+ table_scan.projected_schema.qualified_field(i);
+ if alias.is_some() {
+ Column::new(alias.clone(),
field.name().clone())
+ } else {
+ Column::new(qualifier.cloned(),
field.name().clone())
+ }
+ })
+ .collect::<Vec<_>>();
+ if let Some(alias) = alias {
+ builder = builder.alias(alias)?;
+ }
+ builder = builder.project(project_columns)?;
+ }
+
+ let filter_expr = table_scan
+ .filters
+ .iter()
+ .cloned()
+ .reduce(|acc, expr| acc.and(expr));
+ if let Some(filter) = filter_expr {
+ builder = builder.filter(filter)?;
+ }
+
+ if let Some(fetch) = table_scan.fetch {
+ builder = builder.limit(0, Some(fetch))?;
+ }
+
+ builder.build()
+ }
+ LogicalPlan::SubqueryAlias(subquery_alias) => {
+ let new_plan = Self::unparse_table_scan_pushdown(
+ &subquery_alias.input,
+ Some(subquery_alias.alias.clone()),
+ )?;
+ LogicalPlanBuilder::from(new_plan)
+ .alias(subquery_alias.alias.clone())?
+ .build()
+ }
+ _ => Ok(plan.clone()),
+ }
+ }
+
fn select_item_to_sql(&self, expr: &Expr) -> Result<ast::SelectItem> {
match expr {
Expr::Alias(Alias { expr, name, .. }) => {
diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs
b/datafusion/sql/tests/cases/plan_to_sql.rs
index 4a6923593f..fa95d05c32 100644
--- a/datafusion/sql/tests/cases/plan_to_sql.rs
+++ b/datafusion/sql/tests/cases/plan_to_sql.rs
@@ -21,7 +21,7 @@ use std::vec;
use arrow_schema::*;
use datafusion_common::{DFSchema, Result, TableReference};
use datafusion_expr::test::function_stub::{count_udaf, max_udaf, min_udaf,
sum_udaf};
-use datafusion_expr::{col, table_scan};
+use datafusion_expr::{col, lit, table_scan, wildcard, LogicalPlanBuilder};
use datafusion_sql::planner::{ContextProvider, PlannerContext, SqlToRel};
use datafusion_sql::unparser::dialect::{
DefaultDialect as UnparserDefaultDialect, Dialect as UnparserDialect,
@@ -29,12 +29,14 @@ use datafusion_sql::unparser::dialect::{
};
use datafusion_sql::unparser::{expr_to_sql, plan_to_sql, Unparser};
+use crate::common::{MockContextProvider, MockSessionState};
+use datafusion_expr::builder::{
+ table_scan_with_filter_and_fetch, table_scan_with_filters,
+};
use datafusion_functions::core::planner::CoreFunctionPlanner;
use sqlparser::dialect::{Dialect, GenericDialect, MySqlDialect};
use sqlparser::parser::Parser;
-use crate::common::{MockContextProvider, MockSessionState};
-
#[test]
fn roundtrip_expr() {
let tests: Vec<(TableReference, &str, &str)> = vec![
@@ -619,6 +621,152 @@ fn sql_round_trip(query: &str, expect: &str) {
assert_eq!(roundtrip_statement.to_string(), expect);
}
+#[test]
+fn test_table_scan_pushdown() -> Result<()> {
+ let schema = Schema::new(vec![
+ Field::new("id", DataType::Utf8, false),
+ Field::new("age", DataType::Utf8, false),
+ ]);
+
+ let scan_with_projection =
+ table_scan(Some("t1"), &schema, Some(vec![0, 1]))?.build()?;
+ let scan_with_projection = plan_to_sql(&scan_with_projection)?;
+ assert_eq!(
+ format!("{}", scan_with_projection),
+ "SELECT t1.id, t1.age FROM t1"
+ );
+
+ let scan_with_no_projection = table_scan(Some("t1"), &schema,
None)?.build()?;
+ let scan_with_no_projection = plan_to_sql(&scan_with_no_projection)?;
+ assert_eq!(format!("{}", scan_with_no_projection), "SELECT * FROM t1");
+
+ let table_scan_with_projection_alias =
+ table_scan(Some("t1"), &schema, Some(vec![0, 1]))?
+ .alias("ta")?
+ .build()?;
+ let table_scan_with_projection_alias =
+ plan_to_sql(&table_scan_with_projection_alias)?;
+ assert_eq!(
+ format!("{}", table_scan_with_projection_alias),
+ "SELECT ta.id, ta.age FROM t1 AS ta"
+ );
+
+ let table_scan_with_no_projection_alias = table_scan(Some("t1"), &schema,
None)?
+ .alias("ta")?
+ .build()?;
+ let table_scan_with_no_projection_alias =
+ plan_to_sql(&table_scan_with_no_projection_alias)?;
+ assert_eq!(
+ format!("{}", table_scan_with_no_projection_alias),
+ "SELECT * FROM t1 AS ta"
+ );
+
+ let query_from_table_scan_with_projection = LogicalPlanBuilder::from(
+ table_scan(Some("t1"), &schema, Some(vec![0, 1]))?.build()?,
+ )
+ .project(vec![wildcard()])?
+ .build()?;
+ let query_from_table_scan_with_projection =
+ plan_to_sql(&query_from_table_scan_with_projection)?;
+ assert_eq!(
+ format!("{}", query_from_table_scan_with_projection),
+ "SELECT * FROM (SELECT t1.id, t1.age FROM t1)"
+ );
+
+ let table_scan_with_filter = table_scan_with_filters(
+ Some("t1"),
+ &schema,
+ None,
+ vec![col("id").gt(col("age"))],
+ )?
+ .build()?;
+ let table_scan_with_filter = plan_to_sql(&table_scan_with_filter)?;
+ assert_eq!(
+ format!("{}", table_scan_with_filter),
+ "SELECT * FROM t1 WHERE (t1.id > t1.age)"
+ );
+
+ let table_scan_with_two_filter = table_scan_with_filters(
+ Some("t1"),
+ &schema,
+ None,
+ vec![col("id").gt(lit(1)), col("age").lt(lit(2))],
+ )?
+ .build()?;
+ let table_scan_with_two_filter = plan_to_sql(&table_scan_with_two_filter)?;
+ assert_eq!(
+ format!("{}", table_scan_with_two_filter),
+ "SELECT * FROM t1 WHERE ((t1.id > 1) AND (t1.age < 2))"
+ );
+
+ // TODO: support filters for table scan with alias. Enable this test after
#12368 issue is fixed
+ // see the issue: https://github.com/apache/datafusion/issues/12368
+ // let table_scan_with_filter_alias = table_scan_with_filters(
+ // Some("t1"),
+ // &schema,
+ // None,
+ // vec![col("id").gt(col("age"))],
+ // )?.alias("ta")?.build()?;
+ // let table_scan_with_filter_alias =
plan_to_sql(&table_scan_with_filter_alias)?;
+ // assert_eq!(
+ // format!("{}", table_scan_with_filter_alias),
+ // "SELECT * FROM t1 AS ta WHERE (ta.id > ta.age)"
+ // );
+
+ let table_scan_with_projection_and_filter = table_scan_with_filters(
+ Some("t1"),
+ &schema,
+ Some(vec![0, 1]),
+ vec![col("id").gt(col("age"))],
+ )?
+ .build()?;
+ let table_scan_with_projection_and_filter =
+ plan_to_sql(&table_scan_with_projection_and_filter)?;
+ assert_eq!(
+ format!("{}", table_scan_with_projection_and_filter),
+ "SELECT t1.id, t1.age FROM t1 WHERE (t1.id > t1.age)"
+ );
+
+ let table_scan_with_inline_fetch =
+ table_scan_with_filter_and_fetch(Some("t1"), &schema, None, vec![],
Some(10))?
+ .build()?;
+ let table_scan_with_inline_fetch =
plan_to_sql(&table_scan_with_inline_fetch)?;
+ assert_eq!(
+ format!("{}", table_scan_with_inline_fetch),
+ "SELECT * FROM t1 LIMIT 10"
+ );
+
+ let table_scan_with_projection_and_inline_fetch =
table_scan_with_filter_and_fetch(
+ Some("t1"),
+ &schema,
+ Some(vec![0, 1]),
+ vec![],
+ Some(10),
+ )?
+ .build()?;
+ let table_scan_with_projection_and_inline_fetch =
+ plan_to_sql(&table_scan_with_projection_and_inline_fetch)?;
+ assert_eq!(
+ format!("{}", table_scan_with_projection_and_inline_fetch),
+ "SELECT t1.id, t1.age FROM t1 LIMIT 10"
+ );
+
+ let table_scan_with_all = table_scan_with_filter_and_fetch(
+ Some("t1"),
+ &schema,
+ Some(vec![0, 1]),
+ vec![col("id").gt(col("age"))],
+ Some(10),
+ )?
+ .build()?;
+ let table_scan_with_all = plan_to_sql(&table_scan_with_all)?;
+ assert_eq!(
+ format!("{}", table_scan_with_all),
+ "SELECT t1.id, t1.age FROM t1 WHERE (t1.id > t1.age) LIMIT 10"
+ );
+ Ok(())
+}
+
#[test]
fn test_interval_lhs_eq() {
sql_round_trip(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
