datafusion_sql/unparser/
ast.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use core::fmt;
19use std::ops::ControlFlow;
20
21use sqlparser::ast::helpers::attached_token::AttachedToken;
22use sqlparser::ast::{
23    self, LimitClause, OrderByKind, SelectFlavor, visit_expressions_mut,
24};
25
26#[derive(Clone)]
27pub struct QueryBuilder {
28    with: Option<ast::With>,
29    body: Option<Box<ast::SetExpr>>,
30    order_by_kind: Option<OrderByKind>,
31    limit: Option<ast::Expr>,
32    limit_by: Vec<ast::Expr>,
33    offset: Option<ast::Offset>,
34    fetch: Option<ast::Fetch>,
35    locks: Vec<ast::LockClause>,
36    for_clause: Option<ast::ForClause>,
37    // If true, we need to unparse LogicalPlan::Union as a SQL `UNION` rather than a `UNION ALL`.
38    distinct_union: bool,
39}
40
41impl QueryBuilder {
42    pub fn with(&mut self, value: Option<ast::With>) -> &mut Self {
43        self.with = value;
44        self
45    }
46    pub fn body(&mut self, value: Box<ast::SetExpr>) -> &mut Self {
47        self.body = Some(value);
48        self
49    }
50    pub fn take_body(&mut self) -> Option<Box<ast::SetExpr>> {
51        self.body.take()
52    }
53    pub fn order_by(&mut self, value: OrderByKind) -> &mut Self {
54        self.order_by_kind = Some(value);
55        self
56    }
57    pub fn limit(&mut self, value: Option<ast::Expr>) -> &mut Self {
58        self.limit = value;
59        self
60    }
61    pub fn limit_by(&mut self, value: Vec<ast::Expr>) -> &mut Self {
62        self.limit_by = value;
63        self
64    }
65    pub fn offset(&mut self, value: Option<ast::Offset>) -> &mut Self {
66        self.offset = value;
67        self
68    }
69    pub fn fetch(&mut self, value: Option<ast::Fetch>) -> &mut Self {
70        self.fetch = value;
71        self
72    }
73    pub fn locks(&mut self, value: Vec<ast::LockClause>) -> &mut Self {
74        self.locks = value;
75        self
76    }
77    pub fn for_clause(&mut self, value: Option<ast::ForClause>) -> &mut Self {
78        self.for_clause = value;
79        self
80    }
81    pub fn distinct_union(&mut self) -> &mut Self {
82        self.distinct_union = true;
83        self
84    }
85    pub fn is_distinct_union(&self) -> bool {
86        self.distinct_union
87    }
88    pub fn build(&self) -> Result<ast::Query, BuilderError> {
89        let order_by = self
90            .order_by_kind
91            .as_ref()
92            .map(|order_by_kind| ast::OrderBy {
93                kind: order_by_kind.clone(),
94                interpolate: None,
95            });
96
97        Ok(ast::Query {
98            with: self.with.clone(),
99            body: match self.body {
100                Some(ref value) => value.clone(),
101                None => return Err(Into::into(UninitializedFieldError::from("body"))),
102            },
103            order_by,
104            limit_clause: Some(LimitClause::LimitOffset {
105                limit: self.limit.clone(),
106                offset: self.offset.clone(),
107                limit_by: self.limit_by.clone(),
108            }),
109            fetch: self.fetch.clone(),
110            locks: self.locks.clone(),
111            for_clause: self.for_clause.clone(),
112            settings: None,
113            format_clause: None,
114            pipe_operators: vec![],
115        })
116    }
117    fn create_empty() -> Self {
118        Self {
119            with: Default::default(),
120            body: Default::default(),
121            order_by_kind: Default::default(),
122            limit: Default::default(),
123            limit_by: Default::default(),
124            offset: Default::default(),
125            fetch: Default::default(),
126            locks: Default::default(),
127            for_clause: Default::default(),
128            distinct_union: false,
129        }
130    }
131}
132impl Default for QueryBuilder {
133    fn default() -> Self {
134        Self::create_empty()
135    }
136}
137
138#[derive(Clone)]
139pub struct SelectBuilder {
140    distinct: Option<ast::Distinct>,
141    top: Option<ast::Top>,
142    /// Projection items for the SELECT clause.
143    ///
144    /// This field uses `Option` to distinguish between three distinct states:
145    /// - `None`: No projection has been set (not yet initialized)
146    /// - `Some(vec![])`: Empty projection explicitly set (generates `SELECT FROM ...` or `SELECT 1 FROM ...`)
147    /// - `Some(vec![SelectItem::Wildcard(...)])`: Wildcard projection (generates `SELECT * FROM ...`)
148    /// - `Some(vec![...])`: Non-empty projection with specific columns/expressions
149    ///
150    /// Use `projection()` to set this field and `already_projected()` to check if it has been set.
151    projection: Option<Vec<ast::SelectItem>>,
152    into: Option<ast::SelectInto>,
153    from: Vec<TableWithJoinsBuilder>,
154    lateral_views: Vec<ast::LateralView>,
155    selection: Option<ast::Expr>,
156    group_by: Option<ast::GroupByExpr>,
157    cluster_by: Vec<ast::Expr>,
158    distribute_by: Vec<ast::Expr>,
159    sort_by: Vec<ast::OrderByExpr>,
160    having: Option<ast::Expr>,
161    named_window: Vec<ast::NamedWindowDefinition>,
162    qualify: Option<ast::Expr>,
163    value_table_mode: Option<ast::ValueTableMode>,
164    flavor: Option<SelectFlavor>,
165    /// Counter for generating unique LATERAL FLATTEN aliases within this SELECT.
166    flatten_alias_counter: usize,
167    /// Table aliases that correspond to LATERAL FLATTEN relations.
168    /// Column references into these aliases must use `VALUE` as the column name.
169    flatten_table_aliases: Vec<String>,
170}
171
172/// Prefix used for auto-generated LATERAL FLATTEN table aliases.
173const FLATTEN_ALIAS_PREFIX: &str = "_unnest";
174
175impl SelectBuilder {
176    /// Generate a unique alias for a LATERAL FLATTEN relation
177    /// (`_unnest_1`, `_unnest_2`, …). Each call returns a fresh name.
178    pub fn next_flatten_alias(&mut self) -> String {
179        self.flatten_alias_counter += 1;
180        format!("{FLATTEN_ALIAS_PREFIX}_{}", self.flatten_alias_counter)
181    }
182
183    /// Register a table alias as pointing to a LATERAL FLATTEN relation.
184    pub fn add_flatten_table_alias(&mut self, alias: String) {
185        self.flatten_table_aliases.push(alias);
186    }
187
188    /// Returns true if no FLATTEN table aliases have been registered.
189    pub fn flatten_table_aliases_empty(&self) -> bool {
190        self.flatten_table_aliases.is_empty()
191    }
192
193    /// Returns true if the given table alias refers to a FLATTEN relation.
194    pub fn is_flatten_table_alias(&self, alias: &str) -> bool {
195        self.flatten_table_aliases.iter().any(|a| a == alias)
196    }
197
198    /// Returns the most recently generated flatten alias, or `None` if
199    /// `next_flatten_alias` has not been called yet.
200    pub fn current_flatten_alias(&self) -> Option<String> {
201        if self.flatten_alias_counter > 0 {
202            Some(format!(
203                "{FLATTEN_ALIAS_PREFIX}_{}",
204                self.flatten_alias_counter
205            ))
206        } else {
207            None
208        }
209    }
210
211    pub fn distinct(&mut self, value: Option<ast::Distinct>) -> &mut Self {
212        self.distinct = value;
213        self
214    }
215    pub fn top(&mut self, value: Option<ast::Top>) -> &mut Self {
216        self.top = value;
217        self
218    }
219    pub fn projection(&mut self, value: Vec<ast::SelectItem>) -> &mut Self {
220        self.projection = Some(value);
221        self
222    }
223    pub fn pop_projections(&mut self) -> Vec<ast::SelectItem> {
224        self.projection.take().unwrap_or_default()
225    }
226    /// Returns true if a projection has been explicitly set via `projection()`.
227    ///
228    /// This method is used to determine whether the SELECT clause has already been
229    /// defined, which helps avoid creating duplicate projection nodes during query
230    /// unparsing. It returns `true` for both empty and non-empty projections.
231    ///
232    /// # Returns
233    ///
234    /// - `true` if `projection()` has been called (regardless of whether it was empty or not)
235    /// - `false` if no projection has been set yet
236    ///
237    /// # Example
238    ///
239    /// ```ignore
240    /// let mut builder = SelectBuilder::default();
241    /// assert!(!builder.already_projected());
242    ///
243    /// builder.projection(vec![]);
244    /// assert!(builder.already_projected()); // true even for empty projection
245    ///
246    /// builder.projection(vec![SelectItem::Wildcard(...)]);
247    /// assert!(builder.already_projected()); // true for non-empty projection
248    /// ```
249    pub fn already_projected(&self) -> bool {
250        self.projection.is_some()
251    }
252    pub fn into(&mut self, value: Option<ast::SelectInto>) -> &mut Self {
253        self.into = value;
254        self
255    }
256    pub fn from(&mut self, value: Vec<TableWithJoinsBuilder>) -> &mut Self {
257        self.from = value;
258        self
259    }
260    pub fn push_from(&mut self, value: TableWithJoinsBuilder) -> &mut Self {
261        self.from.push(value);
262        self
263    }
264    pub fn pop_from(&mut self) -> Option<TableWithJoinsBuilder> {
265        self.from.pop()
266    }
267    pub fn lateral_views(&mut self, value: Vec<ast::LateralView>) -> &mut Self {
268        self.lateral_views = value;
269        self
270    }
271
272    /// Replaces the selection with a new value.
273    ///
274    /// This function is used to replace a specific expression within the selection.
275    /// Unlike the `selection` method which combines existing and new selections with AND,
276    /// this method searches for and replaces occurrences of a specific expression.
277    ///
278    /// This method is primarily used to modify LEFT MARK JOIN expressions.
279    /// When processing a LEFT MARK JOIN, we need to replace the placeholder expression
280    /// with the actual join condition in the selection clause.
281    ///
282    /// # Arguments
283    ///
284    /// * `existing_expr` - The expression to replace
285    /// * `value` - The new expression to set as the selection
286    pub fn replace_mark(
287        &mut self,
288        existing_expr: &ast::Expr,
289        value: &ast::Expr,
290    ) -> &mut Self {
291        if let Some(selection) = &mut self.selection {
292            let _ = visit_expressions_mut(selection, |expr| {
293                if expr == existing_expr {
294                    *expr = value.clone();
295                }
296                ControlFlow::<()>::Continue(())
297            });
298        }
299        self
300    }
301
302    pub fn selection(&mut self, value: Option<ast::Expr>) -> &mut Self {
303        // With filter pushdown optimization, the LogicalPlan can have filters defined as part of `TableScan` and `Filter` nodes.
304        // To avoid overwriting one of the filters, we combine the existing filter with the additional filter.
305        // Example:                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
306        // |  Projection: customer.c_phone AS cntrycode, customer.c_acctbal                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      |
307        // |   Filter: CAST(customer.c_acctbal AS Decimal128(38, 6)) > (<subquery>)                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              |
308        // |     Subquery:
309        // |     ..                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 |
310        // |     TableScan: customer, full_filters=[customer.c_mktsegment = Utf8("BUILDING")]
311        match (&self.selection, value) {
312            (Some(existing_selection), Some(new_selection)) => {
313                self.selection = Some(ast::Expr::BinaryOp {
314                    left: Box::new(existing_selection.clone()),
315                    op: ast::BinaryOperator::And,
316                    right: Box::new(new_selection),
317                });
318            }
319            (None, Some(new_selection)) => {
320                self.selection = Some(new_selection);
321            }
322            (_, None) => (),
323        }
324
325        self
326    }
327    pub fn group_by(&mut self, value: ast::GroupByExpr) -> &mut Self {
328        self.group_by = Some(value);
329        self
330    }
331    pub fn cluster_by(&mut self, value: Vec<ast::Expr>) -> &mut Self {
332        self.cluster_by = value;
333        self
334    }
335    pub fn distribute_by(&mut self, value: Vec<ast::Expr>) -> &mut Self {
336        self.distribute_by = value;
337        self
338    }
339    pub fn sort_by(&mut self, value: Vec<ast::OrderByExpr>) -> &mut Self {
340        self.sort_by = value;
341        self
342    }
343    pub fn having(&mut self, value: Option<ast::Expr>) -> &mut Self {
344        self.having = value;
345        self
346    }
347    pub fn named_window(&mut self, value: Vec<ast::NamedWindowDefinition>) -> &mut Self {
348        self.named_window = value;
349        self
350    }
351    pub fn qualify(&mut self, value: Option<ast::Expr>) -> &mut Self {
352        self.qualify = value;
353        self
354    }
355    pub fn value_table_mode(&mut self, value: Option<ast::ValueTableMode>) -> &mut Self {
356        self.value_table_mode = value;
357        self
358    }
359    pub fn build(&self) -> Result<ast::Select, BuilderError> {
360        Ok(ast::Select {
361            optimizer_hints: vec![],
362            distinct: self.distinct.clone(),
363            select_modifiers: None,
364            top_before_distinct: false,
365            top: self.top.clone(),
366            projection: self.projection.clone().unwrap_or_default(),
367            into: self.into.clone(),
368            from: self
369                .from
370                .iter()
371                .filter_map(|b| b.build().transpose())
372                .collect::<Result<Vec<_>, BuilderError>>()?,
373            lateral_views: self.lateral_views.clone(),
374            selection: self.selection.clone(),
375            group_by: match self.group_by {
376                Some(ref value) => value.clone(),
377                None => {
378                    return Err(Into::into(UninitializedFieldError::from("group_by")));
379                }
380            },
381            cluster_by: self.cluster_by.clone(),
382            distribute_by: self.distribute_by.clone(),
383            sort_by: self.sort_by.clone(),
384            having: self.having.clone(),
385            named_window: self.named_window.clone(),
386            qualify: self.qualify.clone(),
387            value_table_mode: self.value_table_mode,
388            connect_by: Vec::new(),
389            window_before_qualify: false,
390            prewhere: None,
391            select_token: AttachedToken::empty(),
392            flavor: match self.flavor {
393                Some(ref value) => *value,
394                None => return Err(Into::into(UninitializedFieldError::from("flavor"))),
395            },
396            exclude: None,
397        })
398    }
399    fn create_empty() -> Self {
400        Self {
401            distinct: Default::default(),
402            top: Default::default(),
403            projection: None,
404            into: Default::default(),
405            from: Default::default(),
406            lateral_views: Default::default(),
407            selection: Default::default(),
408            group_by: Some(ast::GroupByExpr::Expressions(Vec::new(), Vec::new())),
409            cluster_by: Default::default(),
410            distribute_by: Default::default(),
411            sort_by: Default::default(),
412            having: Default::default(),
413            named_window: Default::default(),
414            qualify: Default::default(),
415            value_table_mode: Default::default(),
416            flavor: Some(SelectFlavor::Standard),
417            flatten_alias_counter: 0,
418            flatten_table_aliases: Vec::new(),
419        }
420    }
421}
422impl Default for SelectBuilder {
423    fn default() -> Self {
424        Self::create_empty()
425    }
426}
427
428#[derive(Clone)]
429pub struct TableWithJoinsBuilder {
430    relation: Option<RelationBuilder>,
431    joins: Vec<ast::Join>,
432}
433
434impl TableWithJoinsBuilder {
435    pub fn relation(&mut self, value: RelationBuilder) -> &mut Self {
436        self.relation = Some(value);
437        self
438    }
439
440    pub fn joins(&mut self, value: Vec<ast::Join>) -> &mut Self {
441        self.joins = value;
442        self
443    }
444    pub fn push_join(&mut self, value: ast::Join) -> &mut Self {
445        self.joins.push(value);
446        self
447    }
448
449    pub fn build(&self) -> Result<Option<ast::TableWithJoins>, BuilderError> {
450        match self.relation {
451            Some(ref value) => match value.build()? {
452                Some(relation) => Ok(Some(ast::TableWithJoins {
453                    relation,
454                    joins: self.joins.clone(),
455                })),
456                None => Ok(None),
457            },
458            None => Err(Into::into(UninitializedFieldError::from("relation"))),
459        }
460    }
461    fn create_empty() -> Self {
462        Self {
463            relation: Default::default(),
464            joins: Default::default(),
465        }
466    }
467}
468impl Default for TableWithJoinsBuilder {
469    fn default() -> Self {
470        Self::create_empty()
471    }
472}
473
474#[derive(Clone)]
475pub struct RelationBuilder {
476    relation: Option<TableFactorBuilder>,
477}
478
479#[derive(Clone)]
480// Boxing variants would penalize the common builder path; this enum is
481// constructed-then-consumed locally rather than stored at scale.
482#[expect(clippy::large_enum_variant)]
483enum TableFactorBuilder {
484    Table(TableRelationBuilder),
485    Derived(DerivedRelationBuilder),
486    Unnest(UnnestRelationBuilder),
487    Flatten(FlattenRelationBuilder),
488    Empty,
489}
490
491impl RelationBuilder {
492    pub fn has_relation(&self) -> bool {
493        self.relation.is_some()
494    }
495    pub fn table(&mut self, value: TableRelationBuilder) -> &mut Self {
496        self.relation = Some(TableFactorBuilder::Table(value));
497        self
498    }
499    pub fn derived(&mut self, value: DerivedRelationBuilder) -> &mut Self {
500        self.relation = Some(TableFactorBuilder::Derived(value));
501        self
502    }
503
504    pub fn unnest(&mut self, value: UnnestRelationBuilder) -> &mut Self {
505        self.relation = Some(TableFactorBuilder::Unnest(value));
506        self
507    }
508
509    pub fn flatten(&mut self, value: FlattenRelationBuilder) -> &mut Self {
510        self.relation = Some(TableFactorBuilder::Flatten(value));
511        self
512    }
513
514    pub fn empty(&mut self) -> &mut Self {
515        self.relation = Some(TableFactorBuilder::Empty);
516        self
517    }
518    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
519        let new = self;
520        match new.relation {
521            Some(TableFactorBuilder::Table(ref mut rel_builder)) => {
522                rel_builder.alias = value;
523            }
524            Some(TableFactorBuilder::Derived(ref mut rel_builder)) => {
525                rel_builder.alias = value;
526            }
527            Some(TableFactorBuilder::Unnest(ref mut rel_builder)) => {
528                rel_builder.alias = value;
529            }
530            Some(TableFactorBuilder::Flatten(ref mut rel_builder)) => {
531                rel_builder.alias = value;
532            }
533            Some(TableFactorBuilder::Empty) => (),
534            None => (),
535        }
536        new
537    }
538    pub fn build(&self) -> Result<Option<ast::TableFactor>, BuilderError> {
539        Ok(match self.relation {
540            Some(TableFactorBuilder::Table(ref value)) => Some(value.build()?),
541            Some(TableFactorBuilder::Derived(ref value)) => Some(value.build()?),
542            Some(TableFactorBuilder::Unnest(ref value)) => Some(value.build()?),
543            Some(TableFactorBuilder::Flatten(ref value)) => Some(value.build()?),
544            Some(TableFactorBuilder::Empty) => None,
545            None => return Err(Into::into(UninitializedFieldError::from("relation"))),
546        })
547    }
548    fn create_empty() -> Self {
549        Self {
550            relation: Default::default(),
551        }
552    }
553}
554impl Default for RelationBuilder {
555    fn default() -> Self {
556        Self::create_empty()
557    }
558}
559
560#[derive(Clone)]
561pub struct TableRelationBuilder {
562    name: Option<ast::ObjectName>,
563    alias: Option<ast::TableAlias>,
564    args: Option<Vec<ast::FunctionArg>>,
565    with_hints: Vec<ast::Expr>,
566    version: Option<ast::TableVersion>,
567    partitions: Vec<ast::Ident>,
568    index_hints: Vec<ast::TableIndexHints>,
569}
570
571impl TableRelationBuilder {
572    pub fn name(&mut self, value: ast::ObjectName) -> &mut Self {
573        self.name = Some(value);
574        self
575    }
576    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
577        self.alias = value;
578        self
579    }
580    pub fn args(&mut self, value: Option<Vec<ast::FunctionArg>>) -> &mut Self {
581        self.args = value;
582        self
583    }
584    pub fn with_hints(&mut self, value: Vec<ast::Expr>) -> &mut Self {
585        self.with_hints = value;
586        self
587    }
588    pub fn version(&mut self, value: Option<ast::TableVersion>) -> &mut Self {
589        self.version = value;
590        self
591    }
592    pub fn partitions(&mut self, value: Vec<ast::Ident>) -> &mut Self {
593        self.partitions = value;
594        self
595    }
596    pub fn index_hints(&mut self, value: Vec<ast::TableIndexHints>) -> &mut Self {
597        self.index_hints = value;
598        self
599    }
600    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
601        Ok(ast::TableFactor::Table {
602            name: match self.name {
603                Some(ref value) => value.clone(),
604                None => return Err(Into::into(UninitializedFieldError::from("name"))),
605            },
606            alias: self.alias.clone(),
607            args: self.args.clone().map(|args| ast::TableFunctionArgs {
608                args,
609                settings: None,
610            }),
611            with_hints: self.with_hints.clone(),
612            version: self.version.clone(),
613            partitions: self.partitions.clone(),
614            with_ordinality: false,
615            json_path: None,
616            sample: None,
617            index_hints: self.index_hints.clone(),
618        })
619    }
620    fn create_empty() -> Self {
621        Self {
622            name: Default::default(),
623            alias: Default::default(),
624            args: Default::default(),
625            with_hints: Default::default(),
626            version: Default::default(),
627            partitions: Default::default(),
628            index_hints: Default::default(),
629        }
630    }
631}
632impl Default for TableRelationBuilder {
633    fn default() -> Self {
634        Self::create_empty()
635    }
636}
637#[derive(Clone)]
638pub struct DerivedRelationBuilder {
639    lateral: Option<bool>,
640    subquery: Option<Box<ast::Query>>,
641    alias: Option<ast::TableAlias>,
642}
643
644impl DerivedRelationBuilder {
645    pub fn lateral(&mut self, value: bool) -> &mut Self {
646        self.lateral = Some(value);
647        self
648    }
649    pub fn subquery(&mut self, value: Box<ast::Query>) -> &mut Self {
650        self.subquery = Some(value);
651        self
652    }
653    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
654        self.alias = value;
655        self
656    }
657    fn build(&self) -> Result<ast::TableFactor, BuilderError> {
658        Ok(ast::TableFactor::Derived {
659            lateral: match self.lateral {
660                Some(ref value) => *value,
661                None => return Err(Into::into(UninitializedFieldError::from("lateral"))),
662            },
663            subquery: match self.subquery {
664                Some(ref value) => value.clone(),
665                None => {
666                    return Err(Into::into(UninitializedFieldError::from("subquery")));
667                }
668            },
669            alias: self.alias.clone(),
670            sample: None,
671        })
672    }
673    fn create_empty() -> Self {
674        Self {
675            lateral: Default::default(),
676            subquery: Default::default(),
677            alias: Default::default(),
678        }
679    }
680}
681impl Default for DerivedRelationBuilder {
682    fn default() -> Self {
683        Self::create_empty()
684    }
685}
686
687#[derive(Clone)]
688pub struct UnnestRelationBuilder {
689    pub alias: Option<ast::TableAlias>,
690    pub array_exprs: Vec<ast::Expr>,
691    with_offset: bool,
692    with_offset_alias: Option<ast::Ident>,
693    with_ordinality: bool,
694}
695
696impl UnnestRelationBuilder {
697    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
698        self.alias = value;
699        self
700    }
701    pub fn array_exprs(&mut self, value: Vec<ast::Expr>) -> &mut Self {
702        self.array_exprs = value;
703        self
704    }
705
706    pub fn with_offset(&mut self, value: bool) -> &mut Self {
707        self.with_offset = value;
708        self
709    }
710
711    pub fn with_offset_alias(&mut self, value: Option<ast::Ident>) -> &mut Self {
712        self.with_offset_alias = value;
713        self
714    }
715
716    pub fn with_ordinality(&mut self, value: bool) -> &mut Self {
717        self.with_ordinality = value;
718        self
719    }
720
721    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
722        Ok(ast::TableFactor::UNNEST {
723            alias: self.alias.clone(),
724            array_exprs: self.array_exprs.clone(),
725            with_offset: self.with_offset,
726            with_offset_alias: self.with_offset_alias.clone(),
727            with_ordinality: self.with_ordinality,
728        })
729    }
730
731    fn create_empty() -> Self {
732        Self {
733            alias: Default::default(),
734            array_exprs: Default::default(),
735            with_offset: Default::default(),
736            with_offset_alias: Default::default(),
737            with_ordinality: Default::default(),
738        }
739    }
740}
741
742impl Default for UnnestRelationBuilder {
743    fn default() -> Self {
744        Self::create_empty()
745    }
746}
747
748/// Builds a `LATERAL FLATTEN(INPUT => expr, OUTER => bool)` table factor
749/// for Snowflake-style unnesting.
750#[derive(Clone)]
751pub struct FlattenRelationBuilder {
752    pub alias: Option<ast::TableAlias>,
753    /// The input expression to flatten (e.g. a column reference).
754    pub input_expr: Option<ast::Expr>,
755    /// Whether to preserve rows for NULL/empty inputs (Snowflake `OUTER` param).
756    pub outer: bool,
757}
758
759impl FlattenRelationBuilder {
760    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
761        self.alias = value;
762        self
763    }
764
765    pub fn input_expr(&mut self, value: ast::Expr) -> &mut Self {
766        self.input_expr = Some(value);
767        self
768    }
769
770    pub fn outer(&mut self, value: bool) -> &mut Self {
771        self.outer = value;
772        self
773    }
774
775    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
776        let input = self.input_expr.clone().ok_or_else(|| {
777            BuilderError::from(UninitializedFieldError::from("input_expr"))
778        })?;
779
780        let mut args = vec![ast::FunctionArg::Named {
781            name: ast::Ident::new("INPUT"),
782            arg: ast::FunctionArgExpr::Expr(input),
783            operator: ast::FunctionArgOperator::RightArrow,
784        }];
785
786        if self.outer {
787            args.push(ast::FunctionArg::Named {
788                name: ast::Ident::new("OUTER"),
789                arg: ast::FunctionArgExpr::Expr(ast::Expr::Value(
790                    ast::Value::Boolean(true).into(),
791                )),
792                operator: ast::FunctionArgOperator::RightArrow,
793            });
794        }
795
796        Ok(ast::TableFactor::Function {
797            lateral: true,
798            name: ast::ObjectName::from(vec![ast::Ident::new("FLATTEN")]),
799            args,
800            with_ordinality: false,
801            alias: self.alias.clone(),
802        })
803    }
804
805    fn create_empty() -> Self {
806        Self {
807            alias: None,
808            input_expr: None,
809            outer: false,
810        }
811    }
812}
813
814impl Default for FlattenRelationBuilder {
815    fn default() -> Self {
816        Self::create_empty()
817    }
818}
819
820/// Runtime error when a `build()` method is called and one or more required fields
821/// do not have a value.
822#[derive(Debug, Clone)]
823pub struct UninitializedFieldError(&'static str);
824
825impl UninitializedFieldError {
826    /// Create a new `UninitializedFieldError` for the specified field name.
827    pub fn new(field_name: &'static str) -> Self {
828        UninitializedFieldError(field_name)
829    }
830
831    /// Get the name of the first-declared field that wasn't initialized
832    pub fn field_name(&self) -> &'static str {
833        self.0
834    }
835}
836
837impl fmt::Display for UninitializedFieldError {
838    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
839        write!(f, "Field not initialized: {}", self.0)
840    }
841}
842
843impl From<&'static str> for UninitializedFieldError {
844    fn from(field_name: &'static str) -> Self {
845        Self::new(field_name)
846    }
847}
848impl std::error::Error for UninitializedFieldError {}
849
850#[derive(Debug)]
851pub enum BuilderError {
852    UninitializedField(&'static str),
853    ValidationError(String),
854}
855impl From<UninitializedFieldError> for BuilderError {
856    fn from(s: UninitializedFieldError) -> Self {
857        Self::UninitializedField(s.field_name())
858    }
859}
860impl From<String> for BuilderError {
861    fn from(s: String) -> Self {
862        Self::ValidationError(s)
863    }
864}
865impl fmt::Display for BuilderError {
866    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
867        match self {
868            Self::UninitializedField(field) => {
869                write!(f, "`{field}` must be initialized")
870            }
871            Self::ValidationError(error) => write!(f, "{error}"),
872        }
873    }
874}
875impl std::error::Error for BuilderError {}
datafusion_sql/unparser/ast.rs

datafusion_sql/unparser/
ast.rs
