Consolidate MapAccess, and Subscript into CompoundExpr to handle the complex field access chain (#1551)

Author: goldmedal

src/ast/mod.rs

@@ -459,40 +459,6 @@ pub enum CastFormat {
     ValueAtTimeZone(Value, Value),
 }
 
-/// Represents the syntax/style used in a map access.
-#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-pub enum MapAccessSyntax {
-    /// Access using bracket notation. `mymap[mykey]`
-    Bracket,
-    /// Access using period notation. `mymap.mykey`
-    Period,
-}
-
-/// Expression used to access a value in a nested structure.
-///
-/// Example: `SAFE_OFFSET(0)` in
-/// ```sql
-/// SELECT mymap[SAFE_OFFSET(0)];
-/// ```
-#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-pub struct MapAccessKey {
-    pub key: Expr,
-    pub syntax: MapAccessSyntax,
-}
-
-impl fmt::Display for MapAccessKey {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self.syntax {
-            MapAccessSyntax::Bracket => write!(f, "[{}]", self.key),
-            MapAccessSyntax::Period => write!(f, ".{}", self.key),
-        }
-    }
-}
-
 /// An element of a JSON path.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -629,6 +595,28 @@ pub enum Expr {
     Identifier(Ident),
     /// Multi-part identifier, e.g. `table_alias.column` or `schema.table.col`
     CompoundIdentifier(Vec<Ident>),
+    /// Multi-part expression access.
+    ///
+    /// This structure represents an access chain in structured / nested types
+    /// such as maps, arrays, and lists:
+    /// - Array
+    ///     - A 1-dim array `a[1]` will be represented like:
+    ///         `CompoundFieldAccess(Ident('a'), vec![Subscript(1)]`
+    ///     - A 2-dim array `a[1][2]` will be represented like:
+    ///         `CompoundFieldAccess(Ident('a'), vec![Subscript(1), Subscript(2)]`
+    /// - Map or Struct (Bracket-style)
+    ///     - A map `a['field1']` will be represented like:
+    ///         `CompoundFieldAccess(Ident('a'), vec![Subscript('field')]`
+    ///     - A 2-dim map `a['field1']['field2']` will be represented like:
+    ///         `CompoundFieldAccess(Ident('a'), vec![Subscript('field1'), Subscript('field2')]`
+    /// - Struct (Dot-style) (only effect when the chain contains both subscript and expr)
+    ///     - A struct access `a[field1].field2` will be represented like:
+    ///         `CompoundFieldAccess(Ident('a'), vec![Subscript('field1'), Ident('field2')]`
+    /// - If a struct access likes `a.field1.field2`, it will be represented by CompoundIdentifier([a, field1, field2])
+    CompoundFieldAccess {
+        root: Box<Expr>,
+        access_chain: Vec<AccessExpr>,
+    },
     /// Access data nested in a value containing semi-structured data, such as
     /// the `VARIANT` type on Snowflake. for example `src:customer[0].name`.
     ///
@@ -882,14 +870,6 @@ pub enum Expr {
         data_type: DataType,
         value: String,
     },
-    /// Access a map-like object by field (e.g. `column['field']` or `column[4]`
-    /// Note that depending on the dialect, struct like accesses may be
-    /// parsed as [`Subscript`](Self::Subscript) or [`MapAccess`](Self::MapAccess)
-    /// <https://clickhouse.com/docs/en/sql-reference/data-types/map/>
-    MapAccess {
-        column: Box<Expr>,
-        keys: Vec<MapAccessKey>,
-    },
     /// Scalar function call e.g. `LEFT(foo, 5)`
     Function(Function),
     /// Arbitrary expr method call
@@ -978,11 +958,6 @@ pub enum Expr {
     /// ```
     /// [1]: https://duckdb.org/docs/sql/data_types/map#creating-maps
     Map(Map),
-    /// An access of nested data using subscript syntax, for example `array[2]`.
-    Subscript {
-        expr: Box<Expr>,
-        subscript: Box<Subscript>,
-    },
     /// An array expression e.g. `ARRAY[1, 2]`
     Array(Array),
     /// An interval expression e.g. `INTERVAL '1' YEAR`
@@ -1099,6 +1074,27 @@ impl fmt::Display for Subscript {
     }
 }
 
+/// An element of a [`Expr::CompoundFieldAccess`].
+/// It can be an expression or a subscript.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum AccessExpr {
+    /// Accesses a field using dot notation, e.g. `foo.bar.baz`.
+    Dot(Expr),
+    /// Accesses a field or array element using bracket notation, e.g. `foo['bar']`.
+    Subscript(Subscript),
+}
+
+impl fmt::Display for AccessExpr {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            AccessExpr::Dot(expr) => write!(f, ".{}", expr),
+            AccessExpr::Subscript(subscript) => write!(f, "[{}]", subscript),
+        }
+    }
+}
+
 /// A lambda function.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -1295,12 +1291,16 @@ impl fmt::Display for Expr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Expr::Identifier(s) => write!(f, "{s}"),
-            Expr::MapAccess { column, keys } => {
-                write!(f, "{column}{}", display_separated(keys, ""))
-            }
             Expr::Wildcard(_) => f.write_str("*"),
             Expr::QualifiedWildcard(prefix, _) => write!(f, "{}.*", prefix),
             Expr::CompoundIdentifier(s) => write!(f, "{}", display_separated(s, ".")),
+            Expr::CompoundFieldAccess { root, access_chain } => {
+                write!(f, "{}", root)?;
+                for field in access_chain {
+                    write!(f, "{}", field)?;
+                }
+                Ok(())
+            }
             Expr::IsTrue(ast) => write!(f, "{ast} IS TRUE"),
             Expr::IsNotTrue(ast) => write!(f, "{ast} IS NOT TRUE"),
             Expr::IsFalse(ast) => write!(f, "{ast} IS FALSE"),
@@ -1720,12 +1720,6 @@ impl fmt::Display for Expr {
             Expr::Map(map) => {
                 write!(f, "{map}")
             }
-            Expr::Subscript {
-                expr,
-                subscript: key,
-            } => {
-                write!(f, "{expr}[{key}]")
-            }
             Expr::Array(set) => {
                 write!(f, "{set}")
             }

src/ast/spans.rs

@@ -20,20 +20,20 @@ use core::iter;
 use crate::tokenizer::Span;
 
 use super::{
-    dcl::SecondaryRoles, AlterColumnOperation, AlterIndexOperation, AlterTableOperation, Array,
-    Assignment, AssignmentTarget, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption,
-    ColumnOptionDef, ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex,
-    CreateTable, CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem,
-    Expr, ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
-    FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
-    IlikeSelectItem, Insert, Interpolate, InterpolateExpr, Join, JoinConstraint, JoinOperator,
-    JsonPath, JsonPathElem, LateralView, MatchRecognizePattern, Measure, NamedWindowDefinition,
-    ObjectName, Offset, OnConflict, OnConflictAction, OnInsert, OrderBy, OrderByExpr, Partition,
-    PivotValueSource, ProjectionSelect, Query, ReferentialAction, RenameSelectItem,
-    ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption,
-    Statement, Subscript, SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint,
-    TableFactor, TableOptionsClustered, TableWithJoins, Use, Value, Values, ViewColumnDef,
-    WildcardAdditionalOptions, With, WithFill,
+    dcl::SecondaryRoles, AccessExpr, AlterColumnOperation, AlterIndexOperation,
+    AlterTableOperation, Array, Assignment, AssignmentTarget, CloseCursor, ClusteredIndex,
+    ColumnDef, ColumnOption, ColumnOptionDef, ConflictTarget, ConnectBy, ConstraintCharacteristics,
+    CopySource, CreateIndex, CreateTable, CreateTableOptions, Cte, Delete, DoUpdate,
+    ExceptSelectItem, ExcludeSelectItem, Expr, ExprWithAlias, Fetch, FromTable, Function,
+    FunctionArg, FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList, FunctionArguments,
+    GroupByExpr, HavingBound, IlikeSelectItem, Insert, Interpolate, InterpolateExpr, Join,
+    JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView, MatchRecognizePattern,
+    Measure, NamedWindowDefinition, ObjectName, Offset, OnConflict, OnConflictAction, OnInsert,
+    OrderBy, OrderByExpr, Partition, PivotValueSource, ProjectionSelect, Query, ReferentialAction,
+    RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem,
+    SetExpr, SqlOption, Statement, Subscript, SymbolDefinition, TableAlias, TableAliasColumnDef,
+    TableConstraint, TableFactor, TableOptionsClustered, TableWithJoins, Use, Value, Values,
+    ViewColumnDef, WildcardAdditionalOptions, With, WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -1262,6 +1262,9 @@ impl Spanned for Expr {
             Expr::Identifier(ident) => ident.span,
             Expr::CompoundIdentifier(vec) => union_spans(vec.iter().map(|i| i.span)),
             Expr::CompositeAccess { expr, key } => expr.span().union(&key.span),
+            Expr::CompoundFieldAccess { root, access_chain } => {
+                union_spans(iter::once(root.span()).chain(access_chain.iter().map(|i| i.span())))
+            }
             Expr::IsFalse(expr) => expr.span(),
             Expr::IsNotFalse(expr) => expr.span(),
             Expr::IsTrue(expr) => expr.span(),
@@ -1336,9 +1339,6 @@ impl Spanned for Expr {
             Expr::Nested(expr) => expr.span(),
             Expr::Value(value) => value.span(),
             Expr::TypedString { .. } => Span::empty(),
-            Expr::MapAccess { column, keys } => column
-                .span()
-                .union(&union_spans(keys.iter().map(|i| i.key.span()))),
             Expr::Function(function) => function.span(),
             Expr::GroupingSets(vec) => {
                 union_spans(vec.iter().flat_map(|i| i.iter().map(|k| k.span())))
@@ -1434,7 +1434,6 @@ impl Spanned for Expr {
             Expr::Named { .. } => Span::empty(),
             Expr::Dictionary(_) => Span::empty(),
             Expr::Map(_) => Span::empty(),
-            Expr::Subscript { expr, subscript } => expr.span().union(&subscript.span()),
             Expr::Interval(interval) => interval.value.span(),
             Expr::Wildcard(token) => token.0.span,
             Expr::QualifiedWildcard(object_name, token) => union_spans(
@@ -1473,6 +1472,15 @@ impl Spanned for Subscript {
     }
 }
 
+impl Spanned for AccessExpr {
+    fn span(&self) -> Span {
+        match self {
+            AccessExpr::Dot(ident) => ident.span(),
+            AccessExpr::Subscript(subscript) => subscript.span(),
+        }
+    }
+}
+
 impl Spanned for ObjectName {
     fn span(&self) -> Span {
         let ObjectName(segments) = self;

src/dialect/snowflake.rs

@@ -234,6 +234,10 @@ impl Dialect for SnowflakeDialect {
             RESERVED_FOR_IDENTIFIER.contains(&kw)
         }
     }
+
+    fn supports_partiql(&self) -> bool {
+        true
+    }
 }
 
 /// Parse snowflake create table statement.