datafusion_sql/
planner.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
18//! [`SqlToRel`]: SQL Query Planner (produces [`LogicalPlan`] from SQL AST)
19use std::collections::HashMap;
20use std::str::FromStr;
21use std::sync::{Arc, Mutex};
22use std::vec;
23
24use crate::utils::make_decimal_type;
25use arrow::datatypes::*;
26use datafusion_common::TableReference;
27use datafusion_common::config::SqlParserOptions;
28use datafusion_common::datatype::{DataTypeExt, FieldExt};
29use datafusion_common::error::add_possible_columns_to_diag;
30use datafusion_common::{DFSchema, DataFusionError, Result, not_impl_err, plan_err};
31use datafusion_common::{
32    DFSchemaRef, Diagnostic, SchemaError, field_not_found, internal_err,
33    plan_datafusion_err,
34};
35use datafusion_expr::Expr;
36use datafusion_expr::logical_plan::{LogicalPlan, LogicalPlanBuilder};
37pub use datafusion_expr::planner::ContextProvider;
38use datafusion_expr::utils::find_column_exprs;
39use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo, TimezoneInfo};
40use sqlparser::ast::{ColumnDef as SQLColumnDef, ColumnOption};
41use sqlparser::ast::{DataType as SQLDataType, Ident, ObjectName, TableAlias};
42
43/// SQL parser options
44#[derive(Debug, Clone, Copy)]
45pub struct ParserOptions {
46    /// Whether to parse float as decimal.
47    pub parse_float_as_decimal: bool,
48    /// Whether to normalize identifiers.
49    pub enable_ident_normalization: bool,
50    /// Whether to support varchar with length.
51    pub support_varchar_with_length: bool,
52    /// Whether to normalize options value.
53    pub enable_options_value_normalization: bool,
54    /// Whether to collect spans
55    pub collect_spans: bool,
56    /// Whether string types (VARCHAR, CHAR, Text, and String) are mapped to `Utf8View` during SQL planning.
57    pub map_string_types_to_utf8view: bool,
58    /// Default null ordering for sorting expressions.
59    pub default_null_ordering: NullOrdering,
60}
61
62impl ParserOptions {
63    /// Creates a new `ParserOptions` instance with default values.
64    ///
65    /// # Examples
66    ///
67    /// ```
68    /// use datafusion_sql::planner::ParserOptions;
69    /// let opts = ParserOptions::new();
70    /// assert_eq!(opts.parse_float_as_decimal, false);
71    /// assert_eq!(opts.enable_ident_normalization, true);
72    /// ```
73    pub fn new() -> Self {
74        Self {
75            parse_float_as_decimal: false,
76            enable_ident_normalization: true,
77            support_varchar_with_length: true,
78            map_string_types_to_utf8view: true,
79            enable_options_value_normalization: false,
80            collect_spans: false,
81            // By default, `nulls_max` is used to follow Postgres's behavior.
82            // postgres rule: https://www.postgresql.org/docs/current/queries-order.html
83            default_null_ordering: NullOrdering::NullsMax,
84        }
85    }
86
87    /// Sets the `parse_float_as_decimal` option.
88    ///
89    /// # Examples
90    ///
91    /// ```
92    /// use datafusion_sql::planner::ParserOptions;
93    /// let opts = ParserOptions::new().with_parse_float_as_decimal(true);
94    /// assert_eq!(opts.parse_float_as_decimal, true);
95    /// ```
96    pub fn with_parse_float_as_decimal(mut self, value: bool) -> Self {
97        self.parse_float_as_decimal = value;
98        self
99    }
100
101    /// Sets the `enable_ident_normalization` option.
102    ///
103    /// # Examples
104    ///
105    /// ```
106    /// use datafusion_sql::planner::ParserOptions;
107    /// let opts = ParserOptions::new().with_enable_ident_normalization(false);
108    /// assert_eq!(opts.enable_ident_normalization, false);
109    /// ```
110    pub fn with_enable_ident_normalization(mut self, value: bool) -> Self {
111        self.enable_ident_normalization = value;
112        self
113    }
114
115    /// Sets the `support_varchar_with_length` option.
116    pub fn with_support_varchar_with_length(mut self, value: bool) -> Self {
117        self.support_varchar_with_length = value;
118        self
119    }
120
121    /// Sets the `map_string_types_to_utf8view` option.
122    pub fn with_map_string_types_to_utf8view(mut self, value: bool) -> Self {
123        self.map_string_types_to_utf8view = value;
124        self
125    }
126
127    /// Sets the `enable_options_value_normalization` option.
128    pub fn with_enable_options_value_normalization(mut self, value: bool) -> Self {
129        self.enable_options_value_normalization = value;
130        self
131    }
132
133    /// Sets the `collect_spans` option.
134    pub fn with_collect_spans(mut self, value: bool) -> Self {
135        self.collect_spans = value;
136        self
137    }
138
139    /// Sets the `default_null_ordering` option.
140    pub fn with_default_null_ordering(mut self, value: NullOrdering) -> Self {
141        self.default_null_ordering = value;
142        self
143    }
144}
145
146impl Default for ParserOptions {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152impl From<&SqlParserOptions> for ParserOptions {
153    fn from(options: &SqlParserOptions) -> Self {
154        Self {
155            parse_float_as_decimal: options.parse_float_as_decimal,
156            enable_ident_normalization: options.enable_ident_normalization,
157            support_varchar_with_length: options.support_varchar_with_length,
158            map_string_types_to_utf8view: options.map_string_types_to_utf8view,
159            enable_options_value_normalization: options
160                .enable_options_value_normalization,
161            collect_spans: options.collect_spans,
162            default_null_ordering: options.default_null_ordering.as_str().into(),
163        }
164    }
165}
166
167/// Represents the null ordering for sorting expressions.
168#[derive(Debug, Clone, Copy)]
169pub enum NullOrdering {
170    /// Nulls appear last in ascending order.
171    NullsMax,
172    /// Nulls appear first in descending order.
173    NullsMin,
174    /// Nulls appear first.
175    NullsFirst,
176    /// Nulls appear last.
177    NullsLast,
178}
179
180impl NullOrdering {
181    /// Evaluates the null ordering based on the given ascending flag.
182    ///
183    /// # Returns
184    /// * `true` if nulls should appear first.
185    /// * `false` if nulls should appear last.
186    pub fn nulls_first(&self, asc: bool) -> bool {
187        match self {
188            Self::NullsMax => !asc,
189            Self::NullsMin => asc,
190            Self::NullsFirst => true,
191            Self::NullsLast => false,
192        }
193    }
194}
195
196impl FromStr for NullOrdering {
197    type Err = DataFusionError;
198
199    fn from_str(s: &str) -> Result<Self> {
200        match s {
201            "nulls_max" => Ok(Self::NullsMax),
202            "nulls_min" => Ok(Self::NullsMin),
203            "nulls_first" => Ok(Self::NullsFirst),
204            "nulls_last" => Ok(Self::NullsLast),
205            _ => plan_err!(
206                "Unknown null ordering: Expected one of 'nulls_first', 'nulls_last', 'nulls_min' or 'nulls_max'. Got {s}"
207            ),
208        }
209    }
210}
211
212impl From<&str> for NullOrdering {
213    fn from(s: &str) -> Self {
214        Self::from_str(s).unwrap_or(Self::NullsMax)
215    }
216}
217
218/// Ident Normalizer
219#[derive(Debug)]
220pub struct IdentNormalizer {
221    normalize: bool,
222}
223
224impl Default for IdentNormalizer {
225    fn default() -> Self {
226        Self { normalize: true }
227    }
228}
229
230impl IdentNormalizer {
231    pub fn new(normalize: bool) -> Self {
232        Self { normalize }
233    }
234
235    pub fn normalize(&self, ident: Ident) -> String {
236        if self.normalize {
237            crate::utils::normalize_ident(ident)
238        } else {
239            ident.value
240        }
241    }
242}
243
244/// Struct to store the states used by the Planner. The Planner will leverage the states
245/// to resolve CTEs, Views, subqueries and PREPARE statements. The states include
246/// Common Table Expression (CTE) provided with WITH clause and
247/// Parameter Data Types provided with PREPARE statement and the query schema of the
248/// outer query plan.
249///
250/// # Cloning
251///
252/// Only the `ctes` are truly cloned when the `PlannerContext` is cloned.
253/// This helps resolve scoping issues of CTEs.
254/// By using cloning, a subquery can inherit CTEs from the outer query
255/// and can also define its own private CTEs without affecting the outer query.
256#[derive(Debug, Clone)]
257pub struct PlannerContext {
258    /// Data types for numbered parameters ($1, $2, etc), if supplied
259    /// in `PREPARE` statement
260    prepare_param_data_types: Arc<Vec<Option<FieldRef>>>,
261    /// Map of CTE name to logical plan of the WITH clause.
262    /// Use `Arc<LogicalPlan>` to allow cheap cloning
263    ctes: HashMap<String, Arc<LogicalPlan>>,
264
265    /// The queries schemas of outer query relations, used to resolve the outer referenced
266    /// columns in subquery (recursive aware)
267    outer_queries_schemas_stack: Vec<DFSchemaRef>,
268    /// The joined schemas of all FROM clauses planned so far. When planning LATERAL
269    /// FROM clauses, this should become a suffix of the `outer_query_schema`.
270    outer_from_schema: Option<DFSchemaRef>,
271    /// The query schema defined by the table
272    create_table_schema: Option<DFSchemaRef>,
273    /// When planning non-first queries in a set expression
274    /// (UNION/INTERSECT/EXCEPT), holds the schema of the left-most query.
275    /// Used to alias duplicate expressions to match the left side's field names.
276    set_expr_left_schema: Option<DFSchemaRef>,
277    /// The parameters of all lambdas seen so far
278    lambda_parameters: HashMap<String, FieldRef>,
279}
280
281impl Default for PlannerContext {
282    fn default() -> Self {
283        Self::new()
284    }
285}
286
287impl PlannerContext {
288    /// Create an empty PlannerContext
289    pub fn new() -> Self {
290        Self {
291            prepare_param_data_types: Arc::new(vec![]),
292            ctes: HashMap::new(),
293            outer_queries_schemas_stack: vec![],
294            outer_from_schema: None,
295            create_table_schema: None,
296            set_expr_left_schema: None,
297            lambda_parameters: HashMap::new(),
298        }
299    }
300
301    /// Update the PlannerContext with provided prepare_param_data_types
302    pub fn with_prepare_param_data_types(
303        mut self,
304        prepare_param_data_types: Vec<Option<FieldRef>>,
305    ) -> Self {
306        self.prepare_param_data_types = prepare_param_data_types.into();
307        self
308    }
309
310    /// Return the stack of outer relations' schemas, the outer most
311    /// relation are at the first entry
312    pub fn outer_queries_schemas(&self) -> &[DFSchemaRef] {
313        &self.outer_queries_schemas_stack
314    }
315
316    /// Return an iterator of the subquery relations' schemas, innermost
317    /// relation is returned first.
318    ///
319    /// This order corresponds to the order of resolution when looking up column
320    /// references in subqueries, which start from the innermost relation and
321    /// then look up the outer relations one by one until a match is found or no
322    /// more outer relation exist.
323    ///
324    /// NOTE this is *REVERSED* order of [`Self::outer_queries_schemas`]
325    ///
326    /// This is useful to resolve the column reference in the subquery by
327    /// looking up the outer query schemas one by one.
328    pub fn outer_schemas_iter(&self) -> impl Iterator<Item = &DFSchemaRef> {
329        self.outer_queries_schemas_stack.iter().rev()
330    }
331
332    /// Sets the outer query schema, returning the existing one, if
333    /// any
334    pub fn append_outer_query_schema(&mut self, schema: DFSchemaRef) {
335        self.outer_queries_schemas_stack.push(schema);
336    }
337
338    /// The schema of the adjacent outer relation
339    pub fn latest_outer_query_schema(&self) -> Option<&DFSchemaRef> {
340        self.outer_queries_schemas_stack.last()
341    }
342
343    /// Remove the schema of the adjacent outer relation
344    pub fn pop_outer_query_schema(&mut self) -> Option<DFSchemaRef> {
345        self.outer_queries_schemas_stack.pop()
346    }
347
348    pub fn set_table_schema(
349        &mut self,
350        mut schema: Option<DFSchemaRef>,
351    ) -> Option<DFSchemaRef> {
352        std::mem::swap(&mut self.create_table_schema, &mut schema);
353        schema
354    }
355
356    pub fn table_schema(&self) -> Option<DFSchemaRef> {
357        self.create_table_schema.clone()
358    }
359
360    // Return a clone of the outer FROM schema
361    pub fn outer_from_schema(&self) -> Option<Arc<DFSchema>> {
362        self.outer_from_schema.clone()
363    }
364
365    /// Sets the outer FROM schema, returning the existing one, if any
366    pub fn set_outer_from_schema(
367        &mut self,
368        mut schema: Option<DFSchemaRef>,
369    ) -> Option<DFSchemaRef> {
370        std::mem::swap(&mut self.outer_from_schema, &mut schema);
371        schema
372    }
373
374    /// Extends the FROM schema, returning the existing one, if any
375    pub fn extend_outer_from_schema(&mut self, schema: &DFSchemaRef) -> Result<()> {
376        match self.outer_from_schema.as_mut() {
377            Some(from_schema) => Arc::make_mut(from_schema).merge(schema),
378            None => self.outer_from_schema = Some(Arc::clone(schema)),
379        };
380        Ok(())
381    }
382
383    /// Return the types of parameters (`$1`, `$2`, etc) if known
384    pub fn prepare_param_data_types(&self) -> &[Option<FieldRef>] {
385        &self.prepare_param_data_types
386    }
387
388    /// Returns true if there is a Common Table Expression (CTE) /
389    /// Subquery for the specified name
390    pub fn contains_cte(&self, cte_name: &str) -> bool {
391        self.ctes.contains_key(cte_name)
392    }
393
394    /// Inserts a LogicalPlan for the Common Table Expression (CTE) /
395    /// Subquery for the specified name
396    pub fn insert_cte(&mut self, cte_name: impl Into<String>, plan: LogicalPlan) {
397        let cte_name = cte_name.into();
398        self.ctes.insert(cte_name, Arc::new(plan));
399    }
400
401    /// Return a plan for the Common Table Expression (CTE) / Subquery for the
402    /// specified name
403    pub fn get_cte(&self, cte_name: &str) -> Option<&LogicalPlan> {
404        self.ctes.get(cte_name).map(|cte| cte.as_ref())
405    }
406
407    pub fn lambda_parameters(&self) -> &HashMap<String, FieldRef> {
408        &self.lambda_parameters
409    }
410
411    pub fn with_lambda_parameters(
412        mut self,
413        parameters: impl IntoIterator<Item = FieldRef>,
414    ) -> Self {
415        self.lambda_parameters
416            .extend(parameters.into_iter().map(|f| (f.name().clone(), f)));
417
418        self
419    }
420
421    /// Remove the plan of CTE / Subquery for the specified name
422    pub(super) fn remove_cte(&mut self, cte_name: &str) {
423        self.ctes.remove(cte_name);
424    }
425
426    /// Sets the left-most set expression schema, returning the previous value
427    pub(super) fn set_set_expr_left_schema(
428        &mut self,
429        schema: Option<DFSchemaRef>,
430    ) -> Option<DFSchemaRef> {
431        std::mem::replace(&mut self.set_expr_left_schema, schema)
432    }
433}
434
435/// SQL query planner and binder
436///
437/// This struct is used to convert a SQL AST into a [`LogicalPlan`].
438///
439/// You can control the behavior of the planner by providing [`ParserOptions`].
440///
441/// It performs the following tasks:
442///
443/// 1. Name and type resolution (called "binding" in other systems). This
444///    phase looks up table and column names using the [`ContextProvider`].
445/// 2. Mechanical translation of the AST into a [`LogicalPlan`].
446///
447/// It does not perform type coercion, or perform optimization, which are done
448/// by subsequent passes.
449///
450/// Key interfaces are:
451/// * [`Self::sql_statement_to_plan`]: Convert a statement
452///   (e.g. `SELECT ...`) into a [`LogicalPlan`]
453/// * [`Self::sql_to_expr`]: Convert an expression (e.g. `1 + 2`) into an [`Expr`]
454pub struct SqlToRel<'a, S: ContextProvider> {
455    pub(crate) context_provider: &'a S,
456    pub(crate) options: ParserOptions,
457    pub(crate) ident_normalizer: IdentNormalizer,
458    warnings: Mutex<Vec<Diagnostic>>,
459}
460
461impl<'a, S: ContextProvider> SqlToRel<'a, S> {
462    /// Create a new query planner.
463    ///
464    /// The query planner derives the parser options from the context provider.
465    pub fn new(context_provider: &'a S) -> Self {
466        let parser_options = ParserOptions::from(&context_provider.options().sql_parser);
467        Self::new_with_options(context_provider, parser_options)
468    }
469
470    /// Create a new query planner with the given parser options.
471    ///
472    /// The query planner ignores the parser options from the context provider
473    /// and uses the given parser options instead.
474    pub fn new_with_options(context_provider: &'a S, options: ParserOptions) -> Self {
475        let ident_normalize = options.enable_ident_normalization;
476
477        SqlToRel {
478            context_provider,
479            options,
480            ident_normalizer: IdentNormalizer::new(ident_normalize),
481            warnings: Mutex::new(vec![]),
482        }
483    }
484
485    pub(crate) fn add_warning(&self, warning: Diagnostic) {
486        self.warnings
487            .lock()
488            .expect("warning diagnostic lock poisoned")
489            .push(warning);
490    }
491
492    /// Drain and return non-fatal warnings collected during SQL planning.
493    pub fn take_warnings(&self) -> Vec<Diagnostic> {
494        std::mem::take(
495            &mut self
496                .warnings
497                .lock()
498                .expect("warning diagnostic lock poisoned"),
499        )
500    }
501
502    pub fn build_schema(&self, columns: Vec<SQLColumnDef>) -> Result<Schema> {
503        let mut fields = Vec::with_capacity(columns.len());
504
505        for column in columns {
506            let data_type = self.convert_data_type_to_field(&column.data_type)?;
507            let not_nullable = column
508                .options
509                .iter()
510                .any(|x| x.option == ColumnOption::NotNull);
511            fields.push(
512                data_type
513                    .as_ref()
514                    .clone()
515                    .with_name(self.ident_normalizer.normalize(column.name))
516                    .with_nullable(!not_nullable),
517            );
518        }
519
520        Ok(Schema::new(fields))
521    }
522
523    /// Returns a vector of (column_name, default_expr) pairs
524    pub(super) fn build_column_defaults(
525        &self,
526        columns: &Vec<SQLColumnDef>,
527        planner_context: &mut PlannerContext,
528    ) -> Result<Vec<(String, Expr)>> {
529        let mut column_defaults = vec![];
530        // Default expressions are restricted, column references are not allowed
531        let empty_schema = DFSchema::empty();
532        let error_desc = |e: DataFusionError| match e {
533            DataFusionError::SchemaError(ref err, _)
534                if matches!(**err, SchemaError::FieldNotFound { .. }) =>
535            {
536                plan_datafusion_err!(
537                    "Column reference is not allowed in the DEFAULT expression : {}",
538                    e
539                )
540            }
541            _ => e,
542        };
543
544        for column in columns {
545            if let Some(default_sql_expr) =
546                column.options.iter().find_map(|o| match &o.option {
547                    ColumnOption::Default(expr) => Some(expr),
548                    _ => None,
549                })
550            {
551                let default_expr = self
552                    .sql_to_expr(default_sql_expr.clone(), &empty_schema, planner_context)
553                    .map_err(error_desc)?;
554                column_defaults.push((
555                    self.ident_normalizer.normalize(column.name.clone()),
556                    default_expr,
557                ));
558            }
559        }
560        Ok(column_defaults)
561    }
562
563    /// Apply the given TableAlias to the input plan
564    pub(crate) fn apply_table_alias(
565        &self,
566        plan: LogicalPlan,
567        alias: TableAlias,
568    ) -> Result<LogicalPlan> {
569        let idents = alias.columns.into_iter().map(|c| c.name).collect();
570        let plan = self.apply_expr_alias(plan, idents)?;
571
572        LogicalPlanBuilder::from(plan)
573            .alias(TableReference::bare(
574                self.ident_normalizer.normalize(alias.name),
575            ))?
576            .build()
577    }
578
579    pub(crate) fn apply_expr_alias(
580        &self,
581        plan: LogicalPlan,
582        idents: Vec<Ident>,
583    ) -> Result<LogicalPlan> {
584        if idents.is_empty() {
585            Ok(plan)
586        } else if idents.len() != plan.schema().fields().len() {
587            plan_err!(
588                "Source table contains {} columns but only {} \
589                names given as column alias",
590                plan.schema().fields().len(),
591                idents.len()
592            )
593        } else {
594            let columns = plan.schema().columns();
595            LogicalPlanBuilder::from(plan)
596                .project(columns.into_iter().zip(idents).map(|(col, ident)| {
597                    Expr::Column(col).alias(self.ident_normalizer.normalize(ident))
598                }))?
599                .build()
600        }
601    }
602
603    /// Validate the schema provides all of the columns referenced in the expressions.
604    pub(crate) fn validate_schema_satisfies_exprs(
605        &self,
606        schema: &DFSchema,
607        exprs: &[Expr],
608    ) -> Result<()> {
609        find_column_exprs(exprs)
610            .iter()
611            .try_for_each(|col| match col {
612                Expr::Column(col) => match &col.relation {
613                    Some(r) => schema.field_with_qualified_name(r, &col.name).map(|_| ()),
614                    None => {
615                        if !schema.fields_with_unqualified_name(&col.name).is_empty() {
616                            Ok(())
617                        } else {
618                            Err(field_not_found(
619                                col.relation.clone(),
620                                col.name.as_str(),
621                                schema,
622                            ))
623                        }
624                    }
625                }
626                .map_err(|err: DataFusionError| match &err {
627                    DataFusionError::SchemaError(inner, _)
628                        if matches!(
629                            inner.as_ref(),
630                            SchemaError::FieldNotFound { .. }
631                        ) =>
632                    {
633                        let SchemaError::FieldNotFound {
634                            field,
635                            valid_fields,
636                        } = inner.as_ref()
637                        else {
638                            unreachable!()
639                        };
640                        let mut diagnostic = if let Some(relation) = &col.relation {
641                            Diagnostic::new_error(
642                                format!(
643                                    "column '{}' not found in '{}'",
644                                    col.name, relation
645                                ),
646                                col.spans().first(),
647                            )
648                        } else {
649                            Diagnostic::new_error(
650                                format!("column '{}' not found", col.name),
651                                col.spans().first(),
652                            )
653                        };
654                        add_possible_columns_to_diag(
655                            &mut diagnostic,
656                            field,
657                            valid_fields,
658                        );
659                        err.with_diagnostic(diagnostic)
660                    }
661                    _ => err,
662                }),
663                _ => internal_err!("Not a column"),
664            })
665    }
666
667    pub(crate) fn convert_data_type_to_field(
668        &self,
669        sql_type: &SQLDataType,
670    ) -> Result<FieldRef> {
671        // First check if any of the registered type_planner can handle this type
672        if let Some(type_planner) = self.context_provider.get_type_planner()
673            && let Some(data_type) = type_planner.plan_type_field(sql_type)?
674        {
675            return Ok(data_type);
676        }
677
678        // If no type_planner can handle this type, use the default conversion
679        match sql_type {
680            SQLDataType::Array(ArrayElemTypeDef::AngleBracket(inner_sql_type)) => {
681                // Arrays may be multi-dimensional.
682                Ok(self.convert_data_type_to_field(inner_sql_type)?.into_list())
683            }
684            SQLDataType::Array(ArrayElemTypeDef::SquareBracket(
685                inner_sql_type,
686                maybe_array_size,
687            )) => {
688                let inner_field = self.convert_data_type_to_field(inner_sql_type)?;
689                if let Some(array_size) = maybe_array_size {
690                    let array_size: i32 = (*array_size).try_into().map_err(|_| {
691                        plan_datafusion_err!(
692                            "Array size must be a positive 32 bit integer, got {array_size}"
693                        )
694                    })?;
695                    Ok(inner_field.into_fixed_size_list(array_size))
696                } else {
697                    Ok(inner_field.into_list())
698                }
699            }
700            SQLDataType::Array(ArrayElemTypeDef::None) => {
701                not_impl_err!("Arrays with unspecified type is not supported")
702            }
703            other => Ok(self
704                .convert_simple_data_type(other)?
705                .into_nullable_field_ref()),
706        }
707    }
708
709    fn convert_simple_data_type(&self, sql_type: &SQLDataType) -> Result<DataType> {
710        match sql_type {
711            SQLDataType::Boolean | SQLDataType::Bool => Ok(DataType::Boolean),
712            SQLDataType::TinyInt(_) => Ok(DataType::Int8),
713            SQLDataType::SmallInt(_) | SQLDataType::Int2(_) => Ok(DataType::Int16),
714            SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_) => {
715                Ok(DataType::Int32)
716            }
717            SQLDataType::BigInt(_) | SQLDataType::Int8(_) => Ok(DataType::Int64),
718            SQLDataType::TinyIntUnsigned(_) => Ok(DataType::UInt8),
719            SQLDataType::SmallIntUnsigned(_) | SQLDataType::Int2Unsigned(_) => {
720                Ok(DataType::UInt16)
721            }
722            SQLDataType::IntUnsigned(_)
723            | SQLDataType::IntegerUnsigned(_)
724            | SQLDataType::Int4Unsigned(_) => Ok(DataType::UInt32),
725            SQLDataType::Varchar(length) => {
726                match (length, self.options.support_varchar_with_length) {
727                    (Some(_), false) => plan_err!(
728                        "does not support Varchar with length, \
729                    please set `support_varchar_with_length` to be true"
730                    ),
731                    _ => {
732                        if self.options.map_string_types_to_utf8view {
733                            Ok(DataType::Utf8View)
734                        } else {
735                            Ok(DataType::Utf8)
736                        }
737                    }
738                }
739            }
740            SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => {
741                Ok(DataType::UInt64)
742            }
743            SQLDataType::Float(_) => Ok(DataType::Float32),
744            SQLDataType::Real | SQLDataType::Float4 => Ok(DataType::Float32),
745            SQLDataType::Double(ExactNumberInfo::None)
746            | SQLDataType::DoublePrecision
747            | SQLDataType::Float8 => Ok(DataType::Float64),
748            SQLDataType::Double(
749                ExactNumberInfo::Precision(_) | ExactNumberInfo::PrecisionAndScale(_, _),
750            ) => {
751                not_impl_err!(
752                    "Unsupported SQL type (precision/scale not supported) {sql_type}"
753                )
754            }
755            SQLDataType::Char(_) | SQLDataType::Text | SQLDataType::String(_) => {
756                if self.options.map_string_types_to_utf8view {
757                    Ok(DataType::Utf8View)
758                } else {
759                    Ok(DataType::Utf8)
760                }
761            }
762            SQLDataType::Timestamp(precision, tz_info)
763                if precision.is_none() || [0, 3, 6, 9].contains(&precision.unwrap()) =>
764            {
765                let tz = if *tz_info == TimezoneInfo::Tz
766                    || *tz_info == TimezoneInfo::WithTimeZone
767                {
768                    // Timestamp With Time Zone
769                    // INPUT : [SQLDataType]   TimestampTz + [Config] Time Zone
770                    // OUTPUT: [ArrowDataType] Timestamp<TimeUnit, Some(Time Zone)>
771                    self.context_provider.options().execution.time_zone.clone()
772                } else {
773                    // Timestamp Without Time zone
774                    None
775                };
776                let precision = match precision {
777                    Some(0) => TimeUnit::Second,
778                    Some(3) => TimeUnit::Millisecond,
779                    Some(6) => TimeUnit::Microsecond,
780                    None | Some(9) => TimeUnit::Nanosecond,
781                    _ => unreachable!(),
782                };
783                Ok(DataType::Timestamp(precision, tz.map(Into::into)))
784            }
785            SQLDataType::Date => Ok(DataType::Date32),
786            SQLDataType::Time(None, tz_info) => {
787                if *tz_info == TimezoneInfo::None
788                    || *tz_info == TimezoneInfo::WithoutTimeZone
789                {
790                    Ok(DataType::Time64(TimeUnit::Nanosecond))
791                } else {
792                    // We don't support TIMETZ and TIME WITH TIME ZONE for now
793                    not_impl_err!("Unsupported SQL type {sql_type}")
794                }
795            }
796            SQLDataType::Numeric(exact_number_info)
797            | SQLDataType::Decimal(exact_number_info) => {
798                let (precision, scale) = match *exact_number_info {
799                    ExactNumberInfo::None => (None, None),
800                    ExactNumberInfo::Precision(precision) => (Some(precision), None),
801                    ExactNumberInfo::PrecisionAndScale(precision, scale) => {
802                        (Some(precision), Some(scale))
803                    }
804                };
805                make_decimal_type(precision, scale.map(|s| s as u64))
806            }
807            SQLDataType::Bytea => Ok(DataType::Binary),
808            SQLDataType::Interval { fields, precision } => {
809                if fields.is_some() || precision.is_some() {
810                    return not_impl_err!("Unsupported SQL type {sql_type}");
811                }
812                Ok(DataType::Interval(IntervalUnit::MonthDayNano))
813            }
814            SQLDataType::Struct(fields, _) => {
815                let fields = fields
816                    .iter()
817                    .enumerate()
818                    .map(|(idx, sql_struct_field)| {
819                        let field = self.convert_data_type_to_field(&sql_struct_field.field_type)?;
820                        let field_name = match &sql_struct_field.field_name {
821                            Some(ident) => ident.clone(),
822                            None => Ident::new(format!("c{idx}")),
823                        };
824                        Ok(field.as_ref().clone().with_name(self.ident_normalizer.normalize(field_name)))
825                    })
826                    .collect::<Result<Vec<_>>>()?;
827                Ok(DataType::Struct(Fields::from(fields)))
828            }
829            SQLDataType::Nvarchar(_)
830            | SQLDataType::JSON
831            | SQLDataType::Uuid
832            | SQLDataType::Binary(_)
833            | SQLDataType::Varbinary(_)
834            | SQLDataType::Blob(_)
835            | SQLDataType::Datetime(_)
836            | SQLDataType::Regclass
837            | SQLDataType::Custom(_, _)
838            | SQLDataType::Array(_)
839            | SQLDataType::Enum(_, _)
840            | SQLDataType::Set(_)
841            | SQLDataType::MediumInt(_)
842            | SQLDataType::MediumIntUnsigned(_)
843            | SQLDataType::Character(_)
844            | SQLDataType::CharacterVarying(_)
845            | SQLDataType::CharVarying(_)
846            | SQLDataType::CharacterLargeObject(_)
847            | SQLDataType::CharLargeObject(_)
848            | SQLDataType::Timestamp(_, _)
849            | SQLDataType::Time(Some(_), _)
850            | SQLDataType::Dec(_)
851            | SQLDataType::BigNumeric(_)
852            | SQLDataType::BigDecimal(_)
853            | SQLDataType::Clob(_)
854            | SQLDataType::Bytes(_)
855            | SQLDataType::Int64
856            | SQLDataType::Float64
857            | SQLDataType::JSONB
858            | SQLDataType::Unspecified
859            | SQLDataType::Int16
860            | SQLDataType::Int32
861            | SQLDataType::Int128
862            | SQLDataType::Int256
863            | SQLDataType::UInt8
864            | SQLDataType::UInt16
865            | SQLDataType::UInt32
866            | SQLDataType::UInt64
867            | SQLDataType::UInt128
868            | SQLDataType::UInt256
869            | SQLDataType::Float32
870            | SQLDataType::Date32
871            | SQLDataType::Datetime64(_, _)
872            | SQLDataType::FixedString(_)
873            | SQLDataType::Map(_, _)
874            | SQLDataType::Tuple(_)
875            | SQLDataType::Nested(_)
876            | SQLDataType::Union(_)
877            | SQLDataType::Nullable(_)
878            | SQLDataType::LowCardinality(_)
879            | SQLDataType::Trigger
880            | SQLDataType::TinyBlob
881            | SQLDataType::MediumBlob
882            | SQLDataType::LongBlob
883            | SQLDataType::TinyText
884            | SQLDataType::MediumText
885            | SQLDataType::LongText
886            | SQLDataType::Bit(_)
887            | SQLDataType::BitVarying(_)
888            | SQLDataType::Signed
889            | SQLDataType::SignedInteger
890            | SQLDataType::Unsigned
891            | SQLDataType::UnsignedInteger
892            | SQLDataType::AnyType
893            | SQLDataType::Table(_)
894            | SQLDataType::VarBit(_)
895            | SQLDataType::UTinyInt
896            | SQLDataType::USmallInt
897            | SQLDataType::HugeInt
898            | SQLDataType::UHugeInt
899            | SQLDataType::UBigInt
900            | SQLDataType::TimestampNtz{..}
901            | SQLDataType::NamedTable { .. }
902            | SQLDataType::TsVector
903            | SQLDataType::TsQuery
904            | SQLDataType::GeometricType(_)
905            | SQLDataType::DecimalUnsigned(_) // deprecated mysql type
906            | SQLDataType::FloatUnsigned(_) // deprecated mysql type
907            | SQLDataType::RealUnsigned // deprecated mysql type
908            | SQLDataType::DecUnsigned(_) // deprecated mysql type
909            | SQLDataType::DoubleUnsigned(_) // deprecated mysql type
910            | SQLDataType::DoublePrecisionUnsigned // deprecated mysql type
911            => {
912                not_impl_err!("Unsupported SQL type {sql_type}")
913            }
914        }
915    }
916
917    pub(crate) fn object_name_to_table_reference(
918        &self,
919        object_name: ObjectName,
920    ) -> Result<TableReference> {
921        object_name_to_table_reference(
922            object_name,
923            self.options.enable_ident_normalization,
924        )
925    }
926}
927
928/// Create a [`TableReference`] after normalizing the specified ObjectName
929///
930/// Examples
931/// ```text
932/// ['foo']          -> Bare { table: "foo" }
933/// ['"foo.bar"]]    -> Bare { table: "foo.bar" }
934/// ['foo', 'Bar']   -> Partial { schema: "foo", table: "bar" } <-- note lower case "bar"
935/// ['foo', 'bar']   -> Partial { schema: "foo", table: "bar" }
936/// ['foo', '"Bar"'] -> Partial { schema: "foo", table: "Bar" }
937/// ```
938pub fn object_name_to_table_reference(
939    object_name: ObjectName,
940    enable_normalization: bool,
941) -> Result<TableReference> {
942    // Use destructure to make it clear no fields on ObjectName are ignored
943    let ObjectName(object_name_parts) = object_name;
944    let idents = object_name_parts
945        .into_iter()
946        .map(|object_name_part| {
947            object_name_part.as_ident().cloned().ok_or_else(|| {
948                plan_datafusion_err!(
949                    "Expected identifier, but found: {:?}",
950                    object_name_part
951                )
952            })
953        })
954        .collect::<Result<Vec<_>>>()?;
955    idents_to_table_reference(idents, enable_normalization)
956}
957
958struct IdentTaker {
959    normalizer: IdentNormalizer,
960    idents: Vec<Ident>,
961}
962
963/// Take the next identifier from the back of idents, panic'ing if
964/// there are none left
965impl IdentTaker {
966    fn new(idents: Vec<Ident>, enable_normalization: bool) -> Self {
967        Self {
968            normalizer: IdentNormalizer::new(enable_normalization),
969            idents,
970        }
971    }
972
973    fn take(&mut self) -> String {
974        let ident = self.idents.pop().expect("no more identifiers");
975        self.normalizer.normalize(ident)
976    }
977
978    /// Returns the number of remaining identifiers
979    fn len(&self) -> usize {
980        self.idents.len()
981    }
982}
983
984// impl Display for a nicer error message
985impl std::fmt::Display for IdentTaker {
986    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
987        let mut first = true;
988        for ident in self.idents.iter() {
989            if !first {
990                write!(f, ".")?;
991            }
992            write!(f, "{ident}")?;
993            first = false;
994        }
995
996        Ok(())
997    }
998}
999
1000/// Create a [`TableReference`] after normalizing the specified identifier
1001pub(crate) fn idents_to_table_reference(
1002    idents: Vec<Ident>,
1003    enable_normalization: bool,
1004) -> Result<TableReference> {
1005    let mut taker = IdentTaker::new(idents, enable_normalization);
1006
1007    match taker.len() {
1008        1 => {
1009            let table = taker.take();
1010            Ok(TableReference::bare(table))
1011        }
1012        2 => {
1013            let table = taker.take();
1014            let schema = taker.take();
1015            Ok(TableReference::partial(schema, table))
1016        }
1017        3 => {
1018            let table = taker.take();
1019            let schema = taker.take();
1020            let catalog = taker.take();
1021            Ok(TableReference::full(catalog, schema, table))
1022        }
1023        _ => plan_err!(
1024            "Unsupported compound identifier '{}'. Expected 1, 2 or 3 parts, got {}",
1025            taker,
1026            taker.len()
1027        ),
1028    }
1029}
1030
1031/// Construct a WHERE qualifier suitable for e.g. information_schema filtering
1032/// from the provided object identifiers (catalog, schema and table names).
1033pub fn object_name_to_qualifier(
1034    sql_table_name: &ObjectName,
1035    enable_normalization: bool,
1036) -> Result<String> {
1037    let columns = vec!["table_name", "table_schema", "table_catalog"].into_iter();
1038    let normalizer = IdentNormalizer::new(enable_normalization);
1039    sql_table_name
1040        .0
1041        .iter()
1042        .rev()
1043        .zip(columns)
1044        .map(|(object_name_part, column_name)| {
1045            object_name_part
1046                .as_ident()
1047                .map(|ident| {
1048                    format!(
1049                        r#"{} = '{}'"#,
1050                        column_name,
1051                        normalizer.normalize(ident.clone())
1052                    )
1053                })
1054                .ok_or_else(|| {
1055                    plan_datafusion_err!(
1056                        "Expected identifier, but found: {:?}",
1057                        object_name_part
1058                    )
1059                })
1060        })
1061        .collect::<Result<Vec<_>>>()
1062        .map(|parts| parts.join(" AND "))
1063}
datafusion_sql/planner.rs

datafusion_sql/
planner.rs
