Skip to content

Commit 08119e6

Browse files
goldmedalgoldmedal
authored
Support unparsing UNNEST plan to UNNEST table factor SQL (#13660)
* add `unnest_as_table_factor` and `UnnestRelationBuilder` * unparse unnest as table factor * fix typo * add tests for the default configs * add a static const for unnest_placeholder * fix tests * fix tests
1 parent 28e4c64 commit 08119e6

12 files changed

Lines changed: 313 additions & 74 deletions

File tree

‎datafusion/sql/src/unparser/ast.rs‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -353,6 +353,7 @@ pub(super) struct RelationBuilder {
353353
enum TableFactorBuilder {
354354
Table(TableRelationBuilder),
355355
Derived(DerivedRelationBuilder),
356+
Unnest(UnnestRelationBuilder),
356357
Empty,
357358
}
358359

@@ -369,6 +370,12 @@ impl RelationBuilder {
369370
self.relation = Some(TableFactorBuilder::Derived(value));
370371
self
371372
}
373+
374+
pub fn unnest(&mut self, value: UnnestRelationBuilder) -> &mut Self {
375+
self.relation = Some(TableFactorBuilder::Unnest(value));
376+
self
377+
}
378+
372379
pub fn empty(&mut self) -> &mut Self {
373380
self.relation = Some(TableFactorBuilder::Empty);
374381
self
@@ -382,6 +389,9 @@ impl RelationBuilder {
382389
Some(TableFactorBuilder::Derived(ref mut rel_builder)) => {
383390
rel_builder.alias = value;
384391
}
392+
Some(TableFactorBuilder::Unnest(ref mut rel_builder)) => {
393+
rel_builder.alias = value;
394+
}
385395
Some(TableFactorBuilder::Empty) => (),
386396
None => (),
387397
}
@@ -391,6 +401,7 @@ impl RelationBuilder {
391401
Ok(match self.relation {
392402
Some(TableFactorBuilder::Table(ref value)) => Some(value.build()?),
393403
Some(TableFactorBuilder::Derived(ref value)) => Some(value.build()?),
404+
Some(TableFactorBuilder::Unnest(ref value)) => Some(value.build()?),
394405
Some(TableFactorBuilder::Empty) => None,
395406
None => return Err(Into::into(UninitializedFieldError::from("relation"))),
396407
})
@@ -526,6 +537,68 @@ impl Default for DerivedRelationBuilder {
526537
}
527538
}
528539

540+
#[derive(Clone)]
541+
pub(super) struct UnnestRelationBuilder {
542+
pub alias: Option<ast::TableAlias>,
543+
pub array_exprs: Vec<ast::Expr>,
544+
with_offset: bool,
545+
with_offset_alias: Option<ast::Ident>,
546+
with_ordinality: bool,
547+
}
548+
549+
#[allow(dead_code)]
550+
impl UnnestRelationBuilder {
551+
pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
552+
self.alias = value;
553+
self
554+
}
555+
pub fn array_exprs(&mut self, value: Vec<ast::Expr>) -> &mut Self {
556+
self.array_exprs = value;
557+
self
558+
}
559+
560+
pub fn with_offset(&mut self, value: bool) -> &mut Self {
561+
self.with_offset = value;
562+
self
563+
}
564+
565+
pub fn with_offset_alias(&mut self, value: Option<ast::Ident>) -> &mut Self {
566+
self.with_offset_alias = value;
567+
self
568+
}
569+
570+
pub fn with_ordinality(&mut self, value: bool) -> &mut Self {
571+
self.with_ordinality = value;
572+
self
573+
}
574+
575+
pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
576+
Ok(ast::TableFactor::UNNEST {
577+
alias: self.alias.clone(),
578+
array_exprs: self.array_exprs.clone(),
579+
with_offset: self.with_offset,
580+
with_offset_alias: self.with_offset_alias.clone(),
581+
with_ordinality: self.with_ordinality,
582+
})
583+
}
584+
585+
fn create_empty() -> Self {
586+
Self {
587+
alias: Default::default(),
588+
array_exprs: Default::default(),
589+
with_offset: Default::default(),
590+
with_offset_alias: Default::default(),
591+
with_ordinality: Default::default(),
592+
}
593+
}
594+
}
595+
596+
impl Default for UnnestRelationBuilder {
597+
fn default() -> Self {
598+
Self::create_empty()
599+
}
600+
}
601+
529602
/// Runtime error when a `build()` method is called and one or more required fields
530603
/// do not have a value.
531604
#[derive(Debug, Clone)]

‎datafusion/sql/src/unparser/dialect.rs‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,15 @@ pub trait Dialect: Send + Sync {
157157
fn full_qualified_col(&self) -> bool {
158158
false
159159
}
160+
161+
/// Allow to unparse the unnest plan as [ast::TableFactor::UNNEST].
162+
///
163+
/// Some dialects like BigQuery require UNNEST to be used in the FROM clause but
164+
/// the LogicalPlan planner always puts UNNEST in the SELECT clause. This flag allows
165+
/// to unparse the UNNEST plan as [ast::TableFactor::UNNEST] instead of a subquery.
166+
fn unnest_as_table_factor(&self) -> bool {
167+
false
168+
}
160169
}
161170

162171
/// `IntervalStyle` to use for unparsing
@@ -448,6 +457,7 @@ pub struct CustomDialect {
448457
requires_derived_table_alias: bool,
449458
division_operator: BinaryOperator,
450459
full_qualified_col: bool,
460+
unnest_as_table_factor: bool,
451461
}
452462

453463
impl Default for CustomDialect {
@@ -474,6 +484,7 @@ impl Default for CustomDialect {
474484
requires_derived_table_alias: false,
475485
division_operator: BinaryOperator::Divide,
476486
full_qualified_col: false,
487+
unnest_as_table_factor: false,
477488
}
478489
}
479490
}
@@ -582,6 +593,10 @@ impl Dialect for CustomDialect {
582593
fn full_qualified_col(&self) -> bool {
583594
self.full_qualified_col
584595
}
596+
597+
fn unnest_as_table_factor(&self) -> bool {
598+
self.unnest_as_table_factor
599+
}
585600
}
586601

587602
/// `CustomDialectBuilder` to build `CustomDialect` using builder pattern
@@ -617,6 +632,7 @@ pub struct CustomDialectBuilder {
617632
requires_derived_table_alias: bool,
618633
division_operator: BinaryOperator,
619634
full_qualified_col: bool,
635+
unnest_as_table_factor: bool,
620636
}
621637

622638
impl Default for CustomDialectBuilder {
@@ -649,6 +665,7 @@ impl CustomDialectBuilder {
649665
requires_derived_table_alias: false,
650666
division_operator: BinaryOperator::Divide,
651667
full_qualified_col: false,
668+
unnest_as_table_factor: false,
652669
}
653670
}
654671

@@ -673,6 +690,7 @@ impl CustomDialectBuilder {
673690
requires_derived_table_alias: self.requires_derived_table_alias,
674691
division_operator: self.division_operator,
675692
full_qualified_col: self.full_qualified_col,
693+
unnest_as_table_factor: self.unnest_as_table_factor,
676694
}
677695
}
678696

@@ -800,4 +818,9 @@ impl CustomDialectBuilder {
800818
self.full_qualified_col = full_qualified_col;
801819
self
802820
}
821+
822+
pub fn with_unnest_as_table_factor(mut self, _unnest_as_table_factor: bool) -> Self {
823+
self.unnest_as_table_factor = _unnest_as_table_factor;
824+
self
825+
}
803826
}

‎datafusion/sql/src/unparser/plan.rs‎

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,15 +32,17 @@ use super::{
3232
},
3333
Unparser,
3434
};
35+
use crate::unparser::ast::UnnestRelationBuilder;
3536
use crate::unparser::utils::unproject_agg_exprs;
37+
use crate::utils::UNNEST_PLACEHOLDER;
3638
use datafusion_common::{
3739
internal_err, not_impl_err,
3840
tree_node::{TransformedResult, TreeNode},
3941
Column, DataFusionError, Result, TableReference,
4042
};
4143
use datafusion_expr::{
4244
expr::Alias, BinaryExpr, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan,
43-
LogicalPlanBuilder, Operator, Projection, SortExpr, TableScan,
45+
LogicalPlanBuilder, Operator, Projection, SortExpr, TableScan, Unnest,
4446
};
4547
use sqlparser::ast::{self, Ident, SetExpr};
4648
use std::sync::Arc;
@@ -312,6 +314,19 @@ impl Unparser<'_> {
312314
.select_to_sql_recursively(&new_plan, query, select, relation);
313315
}
314316

317+
// Projection can be top-level plan for unnest relation
318+
// The projection generated by the `RecursiveUnnestRewriter` from a UNNEST relation will have
319+
// only one expression, which is the placeholder column generated by the rewriter.
320+
if self.dialect.unnest_as_table_factor()
321+
&& p.expr.len() == 1
322+
&& Self::is_unnest_placeholder(&p.expr[0])
323+
{
324+
if let LogicalPlan::Unnest(unnest) = &p.input.as_ref() {
325+
return self
326+
.unnest_to_table_factor_sql(unnest, query, select, relation);
327+
}
328+
}
329+
315330
// Projection can be top-level plan for derived table
316331
if select.already_projected() {
317332
return self.derive_with_dialect_alias(
@@ -678,7 +693,11 @@ impl Unparser<'_> {
678693
)
679694
}
680695
LogicalPlan::EmptyRelation(_) => {
681-
relation.empty();
696+
// An EmptyRelation could be behind an UNNEST node. If the dialect supports UNNEST as a table factor,
697+
// a TableRelationBuilder will be created for the UNNEST node first.
698+
if !relation.has_relation() {
699+
relation.empty();
700+
}
682701
Ok(())
683702
}
684703
LogicalPlan::Extension(_) => not_impl_err!("Unsupported operator: {plan:?}"),
@@ -708,6 +727,38 @@ impl Unparser<'_> {
708727
}
709728
}
710729

730+
/// Try to find the placeholder column name generated by `RecursiveUnnestRewriter`
731+
/// Only match the pattern `Expr::Alias(Expr::Column("__unnest_placeholder(...)"))`
732+
fn is_unnest_placeholder(expr: &Expr) -> bool {
733+
if let Expr::Alias(Alias { expr, .. }) = expr {
734+
if let Expr::Column(Column { name, .. }) = expr.as_ref() {
735+
return name.starts_with(UNNEST_PLACEHOLDER);
736+
}
737+
}
738+
false
739+
}
740+
741+
fn unnest_to_table_factor_sql(
742+
&self,
743+
unnest: &Unnest,
744+
query: &mut Option<QueryBuilder>,
745+
select: &mut SelectBuilder,
746+
relation: &mut RelationBuilder,
747+
) -> Result<()> {
748+
let mut unnest_relation = UnnestRelationBuilder::default();
749+
let LogicalPlan::Projection(p) = unnest.input.as_ref() else {
750+
return internal_err!("Unnest input is not a Projection: {unnest:?}");
751+
};
752+
let exprs = p
753+
.expr
754+
.iter()
755+
.map(|e| self.expr_to_sql(e))
756+
.collect::<Result<Vec<_>>>()?;
757+
unnest_relation.array_exprs(exprs);
758+
relation.unnest(unnest_relation);
759+
self.select_to_sql_recursively(p.input.as_ref(), query, select, relation)
760+
}
761+
711762
fn is_scan_with_pushdown(scan: &TableScan) -> bool {
712763
scan.projection.is_some() || !scan.filters.is_empty() || scan.fetch.is_some()
713764
}

‎datafusion/sql/src/unparser/utils.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ pub(crate) fn find_window_nodes_within_select<'a>(
133133

134134
/// Recursively identify Column expressions and transform them into the appropriate unnest expression
135135
///
136-
/// For example, if expr contains the column expr "unnest_placeholder(make_array(Int64(1),Int64(2),Int64(2),Int64(5),NULL),depth=1)"
136+
/// For example, if expr contains the column expr "__unnest_placeholder(make_array(Int64(1),Int64(2),Int64(2),Int64(5),NULL),depth=1)"
137137
/// it will be transformed into an actual unnest expression UNNEST([1, 2, 2, 5, NULL])
138138
pub(crate) fn unproject_unnest_expr(expr: Expr, unnest: &Unnest) -> Result<Expr> {
139139
expr.transform(|sub_expr| {

0 commit comments

Comments
 (0)