Skip to content

Commit 4ae25fd

Browse files
Samyak2Samyak2
committed
GenericDialect: support colon operator for JsonAccess
- Port JsonAccess colon operator from Snowflake to Generic dialect - This will be used in variant data type support in Datafusion - see discussion in datafusion-contrib/datafusion-variant#2
1 parent ce74e7f commit 4ae25fd

6 files changed

Lines changed: 152 additions & 116 deletions

File tree

‎src/dialect/mod.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -759,6 +759,13 @@ pub trait Dialect: Debug + Any {
759759
Token::DoubleColon | Token::ExclamationMark | Token::LBracket | Token::CaretAt => {
760760
Ok(p!(DoubleColon))
761761
}
762+
Token::Colon => match parser.peek_nth_token(1).token {
763+
// When colon is followed by a string or a number, it's usually in MAP syntax.
764+
Token::SingleQuotedString(_) | Token::Number(_, _) => Ok(self.prec_unknown()),
765+
// In other cases, it's used in semi-structured data traversal like in variant or JSON
766+
// string columns. See `JsonAccess`.
767+
_ => Ok(p!(Colon)),
768+
},
762769
Token::Arrow
763770
| Token::LongArrow
764771
| Token::HashArrow
@@ -812,6 +819,7 @@ pub trait Dialect: Debug + Any {
812819
Precedence::Ampersand => 23,
813820
Precedence::Caret => 22,
814821
Precedence::Pipe => 21,
822+
Precedence::Colon => 21,
815823
Precedence::Between => 20,
816824
Precedence::Eq => 20,
817825
Precedence::Like => 19,
@@ -1232,6 +1240,7 @@ pub enum Precedence {
12321240
Ampersand,
12331241
Caret,
12341242
Pipe,
1243+
Colon,
12351244
Between,
12361245
Eq,
12371246
Like,

‎src/dialect/mssql.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,15 @@ impl Dialect for MsSqlDialect {
160160
None
161161
}
162162
}
163+
164+
fn get_next_precedence(&self, parser: &Parser) -> Option<Result<u8, ParserError>> {
165+
let token = parser.peek_token();
166+
match token.token {
167+
// lowest prec to prevent it from turning into a binary op
168+
Token::Colon => Some(Ok(self.prec_unknown())),
169+
_ => None,
170+
}
171+
}
163172
}
164173

165174
impl MsSqlDialect {

‎src/dialect/postgresql.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,8 @@ impl Dialect for PostgreSqlDialect {
136136
| Token::ShiftRight
137137
| Token::ShiftLeft
138138
| Token::CustomBinaryOperator(_) => Some(Ok(PG_OTHER_PREC)),
139+
// lowest prec to prevent it from turning into a binary op
140+
Token::Colon => Some(Ok(self.prec_unknown())),
139141
_ => None,
140142
}
141143
}
@@ -159,6 +161,7 @@ impl Dialect for PostgreSqlDialect {
159161
Precedence::Ampersand => PG_OTHER_PREC,
160162
Precedence::Caret => CARET_PREC,
161163
Precedence::Pipe => PG_OTHER_PREC,
164+
Precedence::Colon => PG_OTHER_PREC,
162165
Precedence::Between => BETWEEN_LIKE_PREC,
163166
Precedence::Eq => EQ_PREC,
164167
Precedence::Like => BETWEEN_LIKE_PREC,

‎src/parser/mod.rs‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3847,7 +3847,8 @@ impl<'a> Parser<'a> {
38473847
let lower_bound = if self.consume_token(&Token::Colon) {
38483848
None
38493849
} else {
3850-
Some(self.parse_expr()?)
3850+
// parse expr until we hit a colon (or any token with lower precedence)
3851+
Some(self.parse_subexpr(self.dialect.prec_value(Precedence::Colon))?)
38513852
};
38523853

38533854
// check for end
@@ -3875,7 +3876,8 @@ impl<'a> Parser<'a> {
38753876
stride: None,
38763877
});
38773878
} else {
3878-
Some(self.parse_expr()?)
3879+
// parse expr until we hit a colon (or any token with lower precedence)
3880+
Some(self.parse_subexpr(self.dialect.prec_value(Precedence::Colon))?)
38793881
};
38803882

38813883
// check for end
@@ -3892,7 +3894,8 @@ impl<'a> Parser<'a> {
38923894
let stride = if self.consume_token(&Token::RBracket) {
38933895
None
38943896
} else {
3895-
Some(self.parse_expr()?)
3897+
// parse expr until we hit a colon (or any token with lower precedence)
3898+
Some(self.parse_subexpr(self.dialect.prec_value(Precedence::Colon))?)
38963899
};
38973900

38983901
if stride.is_some() {

‎tests/sqlparser_common.rs‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17972,3 +17972,126 @@ fn parse_select_parenthesized_wildcard() {
1797217972
assert_eq!(select2.projection.len(), 1);
1797317973
assert!(matches!(select2.projection[0], SelectItem::Wildcard(_)));
1797417974
}
17975+
17976+
// https://docs.snowflake.com/en/user-guide/querying-semistructured
17977+
#[test]
17978+
fn parse_semi_structured_data_traversal() {
17979+
let dialects = TestedDialects::new(vec![
17980+
Box::new(GenericDialect {}),
17981+
Box::new(SnowflakeDialect {}),
17982+
]);
17983+
17984+
// most basic case
17985+
let sql = "SELECT a:b FROM t";
17986+
let select = dialects.verified_only_select(sql);
17987+
assert_eq!(
17988+
SelectItem::UnnamedExpr(Expr::JsonAccess {
17989+
value: Box::new(Expr::Identifier(Ident::new("a"))),
17990+
path: JsonPath {
17991+
path: vec![JsonPathElem::Dot {
17992+
key: "b".to_owned(),
17993+
quoted: false
17994+
}]
17995+
},
17996+
}),
17997+
select.projection[0]
17998+
);
17999+
18000+
// identifier can be quoted
18001+
let sql = r#"SELECT a:"my long object key name" FROM t"#;
18002+
let select = dialects.verified_only_select(sql);
18003+
assert_eq!(
18004+
SelectItem::UnnamedExpr(Expr::JsonAccess {
18005+
value: Box::new(Expr::Identifier(Ident::new("a"))),
18006+
path: JsonPath {
18007+
path: vec![JsonPathElem::Dot {
18008+
key: "my long object key name".to_owned(),
18009+
quoted: true
18010+
}]
18011+
},
18012+
}),
18013+
select.projection[0]
18014+
);
18015+
18016+
dialects.verified_stmt("SELECT a:b::INT FROM t");
18017+
18018+
// unquoted keywords are permitted in the object key
18019+
let sql = "SELECT a:select, a:from FROM t";
18020+
let select = dialects.verified_only_select(sql);
18021+
assert_eq!(
18022+
vec![
18023+
SelectItem::UnnamedExpr(Expr::JsonAccess {
18024+
value: Box::new(Expr::Identifier(Ident::new("a"))),
18025+
path: JsonPath {
18026+
path: vec![JsonPathElem::Dot {
18027+
key: "select".to_owned(),
18028+
quoted: false
18029+
}]
18030+
},
18031+
}),
18032+
SelectItem::UnnamedExpr(Expr::JsonAccess {
18033+
value: Box::new(Expr::Identifier(Ident::new("a"))),
18034+
path: JsonPath {
18035+
path: vec![JsonPathElem::Dot {
18036+
key: "from".to_owned(),
18037+
quoted: false
18038+
}]
18039+
},
18040+
})
18041+
],
18042+
select.projection
18043+
);
18044+
18045+
// multiple levels can be traversed
18046+
// https://docs.snowflake.com/en/user-guide/querying-semistructured#dot-notation
18047+
let sql = r#"SELECT a:foo."bar".baz"#;
18048+
let select = dialects.verified_only_select(sql);
18049+
assert_eq!(
18050+
vec![SelectItem::UnnamedExpr(Expr::JsonAccess {
18051+
value: Box::new(Expr::Identifier(Ident::new("a"))),
18052+
path: JsonPath {
18053+
path: vec![
18054+
JsonPathElem::Dot {
18055+
key: "foo".to_owned(),
18056+
quoted: false,
18057+
},
18058+
JsonPathElem::Dot {
18059+
key: "bar".to_owned(),
18060+
quoted: true,
18061+
},
18062+
JsonPathElem::Dot {
18063+
key: "baz".to_owned(),
18064+
quoted: false,
18065+
}
18066+
]
18067+
},
18068+
})],
18069+
select.projection
18070+
);
18071+
18072+
// dot and bracket notation can be mixed (starting with : case)
18073+
// https://docs.snowflake.com/en/user-guide/querying-semistructured#dot-notation
18074+
let sql = r#"SELECT a:foo[0].bar"#;
18075+
let select = dialects.verified_only_select(sql);
18076+
assert_eq!(
18077+
vec![SelectItem::UnnamedExpr(Expr::JsonAccess {
18078+
value: Box::new(Expr::Identifier(Ident::new("a"))),
18079+
path: JsonPath {
18080+
path: vec![
18081+
JsonPathElem::Dot {
18082+
key: "foo".to_owned(),
18083+
quoted: false,
18084+
},
18085+
JsonPathElem::Bracket {
18086+
key: Expr::value(number("0")),
18087+
},
18088+
JsonPathElem::Dot {
18089+
key: "bar".to_owned(),
18090+
quoted: false,
18091+
}
18092+
]
18093+
},
18094+
})],
18095+
select.projection
18096+
);
18097+
}

‎tests/sqlparser_snowflake.rs‎

Lines changed: 2 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -1265,37 +1265,8 @@ fn parse_lateral_flatten() {
12651265
// https://docs.snowflake.com/en/user-guide/querying-semistructured
12661266
#[test]
12671267
fn parse_semi_structured_data_traversal() {
1268-
// most basic case
1269-
let sql = "SELECT a:b FROM t";
1270-
let select = snowflake().verified_only_select(sql);
1271-
assert_eq!(
1272-
SelectItem::UnnamedExpr(Expr::JsonAccess {
1273-
value: Box::new(Expr::Identifier(Ident::new("a"))),
1274-
path: JsonPath {
1275-
path: vec![JsonPathElem::Dot {
1276-
key: "b".to_owned(),
1277-
quoted: false
1278-
}]
1279-
},
1280-
}),
1281-
select.projection[0]
1282-
);
1283-
1284-
// identifier can be quoted
1285-
let sql = r#"SELECT a:"my long object key name" FROM t"#;
1286-
let select = snowflake().verified_only_select(sql);
1287-
assert_eq!(
1288-
SelectItem::UnnamedExpr(Expr::JsonAccess {
1289-
value: Box::new(Expr::Identifier(Ident::new("a"))),
1290-
path: JsonPath {
1291-
path: vec![JsonPathElem::Dot {
1292-
key: "my long object key name".to_owned(),
1293-
quoted: true
1294-
}]
1295-
},
1296-
}),
1297-
select.projection[0]
1298-
);
1268+
// see `tests/sqlparser_common.rs` -> `parse_semi_structured_data_traversal` for more test
1269+
// cases. This test only has Snowflake-specific syntax like array access.
12991270

13001271
// expressions are allowed in bracket notation
13011272
let sql = r#"SELECT a[2 + 2] FROM t"#;
@@ -1316,88 +1287,6 @@ fn parse_semi_structured_data_traversal() {
13161287
select.projection[0]
13171288
);
13181289

1319-
snowflake().verified_stmt("SELECT a:b::INT FROM t");
1320-
1321-
// unquoted keywords are permitted in the object key
1322-
let sql = "SELECT a:select, a:from FROM t";
1323-
let select = snowflake().verified_only_select(sql);
1324-
assert_eq!(
1325-
vec![
1326-
SelectItem::UnnamedExpr(Expr::JsonAccess {
1327-
value: Box::new(Expr::Identifier(Ident::new("a"))),
1328-
path: JsonPath {
1329-
path: vec![JsonPathElem::Dot {
1330-
key: "select".to_owned(),
1331-
quoted: false
1332-
}]
1333-
},
1334-
}),
1335-
SelectItem::UnnamedExpr(Expr::JsonAccess {
1336-
value: Box::new(Expr::Identifier(Ident::new("a"))),
1337-
path: JsonPath {
1338-
path: vec![JsonPathElem::Dot {
1339-
key: "from".to_owned(),
1340-
quoted: false
1341-
}]
1342-
},
1343-
})
1344-
],
1345-
select.projection
1346-
);
1347-
1348-
// multiple levels can be traversed
1349-
// https://docs.snowflake.com/en/user-guide/querying-semistructured#dot-notation
1350-
let sql = r#"SELECT a:foo."bar".baz"#;
1351-
let select = snowflake().verified_only_select(sql);
1352-
assert_eq!(
1353-
vec![SelectItem::UnnamedExpr(Expr::JsonAccess {
1354-
value: Box::new(Expr::Identifier(Ident::new("a"))),
1355-
path: JsonPath {
1356-
path: vec![
1357-
JsonPathElem::Dot {
1358-
key: "foo".to_owned(),
1359-
quoted: false,
1360-
},
1361-
JsonPathElem::Dot {
1362-
key: "bar".to_owned(),
1363-
quoted: true,
1364-
},
1365-
JsonPathElem::Dot {
1366-
key: "baz".to_owned(),
1367-
quoted: false,
1368-
}
1369-
]
1370-
},
1371-
})],
1372-
select.projection
1373-
);
1374-
1375-
// dot and bracket notation can be mixed (starting with : case)
1376-
// https://docs.snowflake.com/en/user-guide/querying-semistructured#dot-notation
1377-
let sql = r#"SELECT a:foo[0].bar"#;
1378-
let select = snowflake().verified_only_select(sql);
1379-
assert_eq!(
1380-
vec![SelectItem::UnnamedExpr(Expr::JsonAccess {
1381-
value: Box::new(Expr::Identifier(Ident::new("a"))),
1382-
path: JsonPath {
1383-
path: vec![
1384-
JsonPathElem::Dot {
1385-
key: "foo".to_owned(),
1386-
quoted: false,
1387-
},
1388-
JsonPathElem::Bracket {
1389-
key: Expr::value(number("0")),
1390-
},
1391-
JsonPathElem::Dot {
1392-
key: "bar".to_owned(),
1393-
quoted: false,
1394-
}
1395-
]
1396-
},
1397-
})],
1398-
select.projection
1399-
);
1400-
14011290
// dot and bracket notation can be mixed (starting with bracket case)
14021291
// https://docs.snowflake.com/en/user-guide/querying-semistructured#dot-notation
14031292
let sql = r#"SELECT a[0].foo.bar"#;

0 commit comments

Comments
 (0)