Support planning Map literal

datafusion/sql/src/expr/mod.rs

@@ -21,13 +21,13 @@ use datafusion_expr::planner::PlannerResult;
 use datafusion_expr::planner::RawDictionaryExpr;
 use datafusion_expr::planner::RawFieldAccessExpr;
 use sqlparser::ast::{
-    CastKind, DictionaryField, Expr as SQLExpr, StructField, Subscript, TrimWhereField,
-    Value,
+    CastKind, DictionaryField, Expr as SQLExpr, MapEntry, StructField, Subscript,
+    TrimWhereField, Value,
 };
 
 use datafusion_common::{
-    internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema, Result,
-    ScalarValue,
+    internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema,
+    DataFusionError, Result, ScalarValue,
 };
 use datafusion_expr::expr::InList;
 use datafusion_expr::expr::ScalarFunction;
@@ -628,6 +628,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             SQLExpr::Dictionary(fields) => {
                 self.try_plan_dictionary_literal(fields, schema, planner_context)
             }
+            SQLExpr::Map(map) => {
+                self.try_plan_map_literal(map.entries, schema, planner_context)
+            }
             _ => not_impl_err!("Unsupported ast node in sqltorel: {sql:?}"),
         }
     }
@@ -714,6 +717,37 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         not_impl_err!("Unsupported dictionary literal: {raw_expr:?}")
     }
 
+    fn try_plan_map_literal(
+        &self,
+        entries: Vec<MapEntry>,
+        schema: &DFSchema,
+        planner_context: &mut PlannerContext,
+    ) -> Result<Expr> {
+        let mut exprs = vec![];
+        entries.into_iter().try_for_each(|entry| {
+            exprs.push(self.sql_expr_to_logical_expr(
+                *entry.key,
+                schema,
+                planner_context,
+            )?);
+            exprs.push(self.sql_expr_to_logical_expr(
+                *entry.value,
+                schema,
+                planner_context,
+            )?);
+            Ok::<_, DataFusionError>(())
+        })?;
+        for planner in self.context_provider.get_expr_planners() {
+            match planner.plan_make_map(exprs)? {
+                PlannerResult::Planned(expr) => {
+                    return Ok(expr);
+                }
+                PlannerResult::Original(expr) => exprs = expr,
+            }
+        }
+        not_impl_err!("Unsupported MAP literal: {exprs:?}")
+    }
+
     // Handles a call to struct(...) where the arguments are named. For example
     // `struct (v as foo, v2 as bar)` by creating a call to the `named_struct` function
     fn create_named_struct_expr(

datafusion/sqllogictest/test_files/map.slt

@@ -310,3 +310,152 @@ VALUES (MAP(['a'], [1])), (MAP(['b'], [2])), (MAP(['c', 'a'], [3, 1]))
 {a: 1}
 {b: 2}
 {c: 3, a: 1}
+
+query ?
+SELECT MAP {'a':1, 'b':2, 'c':3};
+----
+{a: 1, b: 2, c: 3}
+
+query ?
+SELECT MAP {'a':1, 'b':2, 'c':3 } FROM t;
+----
+{a: 1, b: 2, c: 3}
+{a: 1, b: 2, c: 3}
+{a: 1, b: 2, c: 3}
+
+query I
+SELECT MAP {'a':1, 'b':2, 'c':3}['a'];
+----
+1
+
+query I
+SELECT MAP {'a':1, 'b':2, 'c':3 }['a'] FROM t;
+----
+1
+1
+1
+
+# TODO: support parsing an empty map
+# query ?
+# SELECT MAP {};
+# ----
+# {}
+
+# values contain null
+query ?
+SELECT MAP {'a': 1, 'b': null};
+----
+{a: 1, b: }
+
+# keys contain null
+query error DataFusion error: Execution error: map key cannot be null
+SELECT MAP {'a': 1, null: 2}
+
+# array as key
+query ?
+SELECT MAP {[1,2,3]:1, [2,4]:2};
+----
+ {[1, 2, 3]: 1, [2, 4]: 2}
+
+# array with different type as key
+# expect to fail due to type coercion error
+query error
+SELECT MAP {[1,2,3]:1, ['a', 'b']:2};
+
+# array as value
+query ?
+SELECT MAP {'a':[1,2,3], 'b':[2,4]};
+----
+{a: [1, 2, 3], b: [2, 4]}
+
+# array with different type as value
+# expect to fail due to type coercion error
+query error
+SELECT MAP {'a':[1,2,3], 'b':['a', 'b']};
+
+# struct as key
+query ?
+SELECT MAP {{'a':1, 'b':2}:1, {'a':3, 'b':4}:2};
+----
+{{a: 1, b: 2}: 1, {a: 3, b: 4}: 2}
+
+# struct with different fields as key
+# expect to fail due to type coercion error
+query error
+SELECT MAP {{'a':1, 'b':2}:1, {'c':3, 'd':4}:2};
+
+# struct as value
+query ?
+SELECT MAP {'a':{'b':1, 'c':2}, 'b':{'b':3, 'c':4}};
+----
+{a: {b: 1, c: 2}, b: {b: 3, c: 4}}
+
+# struct with different fields as value
+# expect to fail due to type coercion error
+query error
+SELECT MAP {'a':{'b':1, 'c':2}, 'b':{'c':3, 'd':4}};
+
+# map as key
+query ?
+SELECT MAP { MAP {1:'a', 2:'b'}:1, MAP {1:'c', 2:'d'}:2 };
+----
+{{1: a, 2: b}: 1, {1: c, 2: d}: 2}
+
+# map with different keys as key
+query ?
+SELECT MAP { MAP {1:'a', 2:'b', 3:'c'}:1, MAP {2:'c', 4:'d'}:2 };
+----
+ {{1: a, 2: b, 3: c}: 1, {2: c, 4: d}: 2}
+
+# map as value
+query ?
+SELECT MAP {1: MAP {1:'a', 2:'b'}, 2: MAP {1:'c', 2:'d'} };
+----
+{1: {1: a, 2: b}, 2: {1: c, 2: d}}
+
+# map with different keys as value
+query ?
+SELECT MAP {'a': MAP {1:'a', 2:'b', 3:'c'}, 'b': MAP {2:'c', 4:'d'} };
+----
+{a: {1: a, 2: b, 3: c}, b: {2: c, 4: d}}
+
+# complex map for each row
+query ?
+SELECT MAP {'a': MAP {1:'a', 2:'b', 3:'c'}, 'b': MAP {2:'c', 4:'d'} } from t;
+----
+{a: {1: a, 2: b, 3: c}, b: {2: c, 4: d}}
+{a: {1: a, 2: b, 3: c}, b: {2: c, 4: d}}
+{a: {1: a, 2: b, 3: c}, b: {2: c, 4: d}}
+
+# access map with non-existent key
+query ?
+SELECT MAP {'a': MAP {1:'a', 2:'b', 3:'c'}, 'b': MAP {2:'c', 4:'d'} }['c'];
+----
+NULL
+
+# access map with null key
+query error
+SELECT MAP {'a': MAP {1:'a', 2:'b', 3:'c'}, 'b': MAP {2:'c', 4:'d'} }[NULL];
+
+query ?
+SELECT MAP { 'a': 1, 2: 3 };
+----
+{a: 1, 2: 3}
+
+# TODO: fix accessing map with non-string key
+# query ?
+# SELECT MAP { 1: 'a', 2: 'b', 3: 'c' }[1];
+# ----
+# a
+
+# TODO: fix accessing map with non-string key
+# query ?
+# SELECT MAP { MAP {1:'a', 2:'b'}:1, MAP {1:'c', 2:'d'}:2 }[MAP {1:'a', 2:'b'}];
+# ----
+# 1
+
+# TODO: fix accessing map with non-string key
+# query ?
+# SELECT MAKE_MAP(1, null, 2, 33, 3, null)[2];
+# ----
+# 33