diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index 71ff7c03bea2..b80ffb6aed3f 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -21,8 +21,8 @@ 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::{ @@ -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:?}"), } } @@ -711,7 +714,29 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { PlannerResult::Original(expr) => raw_expr = expr, } } - not_impl_err!("Unsupported dictionary literal: {raw_expr:?}") + not_impl_err!("Dictionary not supported by ExprPlanner: {raw_expr:?}") + } + + fn try_plan_map_literal( + &self, + entries: Vec, + schema: &DFSchema, + planner_context: &mut PlannerContext, + ) -> Result { + let mut exprs: Vec<_> = entries + .into_iter() + .flat_map(|entry| vec![entry.key, entry.value].into_iter()) + .map(|expr| self.sql_expr_to_logical_expr(*expr, schema, planner_context)) + .collect::>>()?; + 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!("MAP not supported by ExprPlanner: {exprs:?}") } // Handles a call to struct(...) where the arguments are named. For example diff --git a/datafusion/sqllogictest/test_files/map.slt b/datafusion/sqllogictest/test_files/map.slt index e530e14df66e..11998eea9044 100644 --- a/datafusion/sqllogictest/test_files/map.slt +++ b/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(https://github.com/sqlparser-rs/sqlparser-rs/pull/1361): support parsing an empty map. Enable this after upgrading sqlparser-rs. +# 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(https://github.com/apache/datafusion/issues/11785): fix accessing map with non-string key +# query ? +# SELECT MAP { 1: 'a', 2: 'b', 3: 'c' }[1]; +# ---- +# a + +# TODO(https://github.com/apache/datafusion/issues/11785): 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(https://github.com/apache/datafusion/issues/11785): fix accessing map with non-string key +# query ? +# SELECT MAKE_MAP(1, null, 2, 33, 3, null)[2]; +# ---- +# 33