@@ -21,14 +21,14 @@ use datafusion_expr::planner::{
2121 PlannerResult , RawBinaryExpr , RawDictionaryExpr , RawFieldAccessExpr ,
2222} ;
2323use sqlparser:: ast:: {
24- BinaryOperator , CastFormat , CastKind , DataType as SQLDataType , DictionaryField ,
25- Expr as SQLExpr , ExprWithAlias as SQLExprWithAlias , MapEntry , StructField , Subscript ,
26- TrimWhereField , Value ,
24+ AccessExpr , BinaryOperator , CastFormat , CastKind , DataType as SQLDataType ,
25+ DictionaryField , Expr as SQLExpr , ExprWithAlias as SQLExprWithAlias , MapEntry ,
26+ StructField , Subscript , TrimWhereField , Value ,
2727} ;
2828
2929use datafusion_common:: {
30- internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema , Result ,
31- ScalarValue ,
30+ internal_datafusion_err, internal_err, not_impl_err, plan_err, Column , DFSchema ,
31+ Result , ScalarValue ,
3232} ;
3333
3434use datafusion_expr:: expr:: ScalarFunction ;
@@ -238,14 +238,14 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
238238 self . sql_identifier_to_expr ( id, schema, planner_context)
239239 }
240240
241- SQLExpr :: MapAccess { .. } => {
242- not_impl_err ! ( "Map Access" )
243- }
244-
245241 // <expr>["foo"], <expr>[4] or <expr>[4:5]
246- SQLExpr :: Subscript { expr, subscript } => {
247- self . sql_subscript_to_expr ( * expr, subscript, schema, planner_context)
248- }
242+ SQLExpr :: CompoundFieldAccess { root, access_chain } => self
243+ . sql_compound_field_access_to_expr (
244+ * root,
245+ access_chain,
246+ schema,
247+ planner_context,
248+ ) ,
249249
250250 SQLExpr :: CompoundIdentifier ( ids) => {
251251 self . sql_compound_identifier_to_expr ( ids, schema, planner_context)
@@ -982,84 +982,146 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
982982 Ok ( Expr :: Cast ( Cast :: new ( Box :: new ( expr) , dt) ) )
983983 }
984984
985- fn sql_subscript_to_expr (
985+ fn sql_compound_field_access_to_expr (
986986 & self ,
987- expr : SQLExpr ,
988- subscript : Box < Subscript > ,
987+ root : SQLExpr ,
988+ access_chain : Vec < AccessExpr > ,
989989 schema : & DFSchema ,
990990 planner_context : & mut PlannerContext ,
991991 ) -> Result < Expr > {
992- let expr = self . sql_expr_to_logical_expr ( expr, schema, planner_context) ?;
993-
994- let field_access = match * subscript {
995- Subscript :: Index { index } => {
996- // index can be a name, in which case it is a named field access
997- match index {
998- SQLExpr :: Value (
999- Value :: SingleQuotedString ( s) | Value :: DoubleQuotedString ( s) ,
1000- ) => GetFieldAccess :: NamedStructField {
1001- name : ScalarValue :: from ( s) ,
1002- } ,
1003- SQLExpr :: JsonAccess { .. } => {
1004- return not_impl_err ! ( "JsonAccess" ) ;
992+ let mut root = self . sql_expr_to_logical_expr ( root, schema, planner_context) ?;
993+ let fields = access_chain
994+ . into_iter ( )
995+ . map ( |field| match field {
996+ AccessExpr :: Subscript ( subscript) => {
997+ match subscript {
998+ Subscript :: Index { index } => {
999+ // index can be a name, in which case it is a named field access
1000+ match index {
1001+ SQLExpr :: Value (
1002+ Value :: SingleQuotedString ( s)
1003+ | Value :: DoubleQuotedString ( s) ,
1004+ ) => Ok ( Some ( GetFieldAccess :: NamedStructField {
1005+ name : ScalarValue :: from ( s) ,
1006+ } ) ) ,
1007+ SQLExpr :: JsonAccess { .. } => {
1008+ not_impl_err ! ( "JsonAccess" )
1009+ }
1010+ // otherwise treat like a list index
1011+ _ => Ok ( Some ( GetFieldAccess :: ListIndex {
1012+ key : Box :: new ( self . sql_expr_to_logical_expr (
1013+ index,
1014+ schema,
1015+ planner_context,
1016+ ) ?) ,
1017+ } ) ) ,
1018+ }
1019+ }
1020+ Subscript :: Slice {
1021+ lower_bound,
1022+ upper_bound,
1023+ stride,
1024+ } => {
1025+ // Means access like [:2]
1026+ let lower_bound = if let Some ( lower_bound) = lower_bound {
1027+ self . sql_expr_to_logical_expr (
1028+ lower_bound,
1029+ schema,
1030+ planner_context,
1031+ )
1032+ } else {
1033+ not_impl_err ! ( "Slice subscript requires a lower bound" )
1034+ } ?;
1035+
1036+ // means access like [2:]
1037+ let upper_bound = if let Some ( upper_bound) = upper_bound {
1038+ self . sql_expr_to_logical_expr (
1039+ upper_bound,
1040+ schema,
1041+ planner_context,
1042+ )
1043+ } else {
1044+ not_impl_err ! ( "Slice subscript requires an upper bound" )
1045+ } ?;
1046+
1047+ // stride, default to 1
1048+ let stride = if let Some ( stride) = stride {
1049+ self . sql_expr_to_logical_expr (
1050+ stride,
1051+ schema,
1052+ planner_context,
1053+ ) ?
1054+ } else {
1055+ lit ( 1i64 )
1056+ } ;
1057+
1058+ Ok ( Some ( GetFieldAccess :: ListRange {
1059+ start : Box :: new ( lower_bound) ,
1060+ stop : Box :: new ( upper_bound) ,
1061+ stride : Box :: new ( stride) ,
1062+ } ) )
1063+ }
10051064 }
1006- // otherwise treat like a list index
1007- _ => GetFieldAccess :: ListIndex {
1008- key : Box :: new ( self . sql_expr_to_logical_expr (
1009- index,
1010- schema,
1011- planner_context,
1012- ) ?) ,
1013- } ,
10141065 }
1015- }
1016- Subscript :: Slice {
1017- lower_bound,
1018- upper_bound,
1019- stride,
1020- } => {
1021- // Means access like [:2]
1022- let lower_bound = if let Some ( lower_bound) = lower_bound {
1023- self . sql_expr_to_logical_expr ( lower_bound, schema, planner_context)
1024- } else {
1025- not_impl_err ! ( "Slice subscript requires a lower bound" )
1026- } ?;
1027-
1028- // means access like [2:]
1029- let upper_bound = if let Some ( upper_bound) = upper_bound {
1030- self . sql_expr_to_logical_expr ( upper_bound, schema, planner_context)
1031- } else {
1032- not_impl_err ! ( "Slice subscript requires an upper bound" )
1033- } ?;
1034-
1035- // stride, default to 1
1036- let stride = if let Some ( stride) = stride {
1037- self . sql_expr_to_logical_expr ( stride, schema, planner_context) ?
1038- } else {
1039- lit ( 1i64 )
1040- } ;
1041-
1042- GetFieldAccess :: ListRange {
1043- start : Box :: new ( lower_bound) ,
1044- stop : Box :: new ( upper_bound) ,
1045- stride : Box :: new ( stride) ,
1066+ AccessExpr :: Dot ( expr) => {
1067+ let expr =
1068+ self . sql_expr_to_logical_expr ( expr, schema, planner_context) ?;
1069+ match expr {
1070+ Expr :: Column ( Column {
1071+ name,
1072+ relation,
1073+ spans,
1074+ } ) => {
1075+ if let Some ( relation) = & relation {
1076+ // If the first part of the dot access is a column reference, we should
1077+ // check if the column is from the same table as the root expression.
1078+ // If it is, we should replace the root expression with the column reference.
1079+ // Otherwise, we should treat the dot access as a named field access.
1080+ if relation. table ( ) == root. schema_name ( ) . to_string ( ) {
1081+ root = Expr :: Column ( Column {
1082+ name,
1083+ relation : Some ( relation. clone ( ) ) ,
1084+ spans,
1085+ } ) ;
1086+ Ok ( None )
1087+ } else {
1088+ plan_err ! (
1089+ "table name mismatch: {} != {}" ,
1090+ relation. table( ) ,
1091+ root. schema_name( )
1092+ )
1093+ }
1094+ } else {
1095+ Ok ( Some ( GetFieldAccess :: NamedStructField {
1096+ name : ScalarValue :: from ( name) ,
1097+ } ) )
1098+ }
1099+ }
1100+ _ => not_impl_err ! (
1101+ "Dot access not supported for non-column expr: {expr:?}"
1102+ ) ,
1103+ }
10461104 }
1047- }
1048- } ;
1105+ } )
1106+ . collect :: < Result < Vec < _ > > > ( ) ? ;
10491107
1050- let mut field_access_expr = RawFieldAccessExpr { expr, field_access } ;
1051- for planner in self . context_provider . get_expr_planners ( ) {
1052- match planner. plan_field_access ( field_access_expr, schema) ? {
1053- PlannerResult :: Planned ( expr) => return Ok ( expr) ,
1054- PlannerResult :: Original ( expr) => {
1055- field_access_expr = expr;
1108+ fields
1109+ . into_iter ( )
1110+ . flatten ( )
1111+ . try_fold ( root, |expr, field_access| {
1112+ let mut field_access_expr = RawFieldAccessExpr { expr, field_access } ;
1113+ for planner in self . context_provider . get_expr_planners ( ) {
1114+ match planner. plan_field_access ( field_access_expr, schema) ? {
1115+ PlannerResult :: Planned ( expr) => return Ok ( expr) ,
1116+ PlannerResult :: Original ( expr) => {
1117+ field_access_expr = expr;
1118+ }
1119+ }
10561120 }
1057- }
1058- }
1059-
1060- not_impl_err ! (
1061- "GetFieldAccess not supported by ExprPlanner: {field_access_expr:?}"
1062- )
1121+ not_impl_err ! (
1122+ "GetFieldAccess not supported by ExprPlanner: {field_access_expr:?}"
1123+ )
1124+ } )
10631125 }
10641126}
10651127
0 commit comments