Skip to content

Commit b9ce75a

Browse files
jmhainjmhain
committed
subscript
1 parent 9d15f7e commit b9ce75a

File tree

4 files changed

+216
-60
lines changed

4 files changed

+216
-60
lines changed

‎src/ast/mod.rs‎

Lines changed: 56 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -745,10 +745,10 @@ pub enum Expr {
745745
/// ```
746746
/// [1]: https://duckdb.org/docs/sql/data_types/struct#creating-structs
747747
Dictionary(Vec<DictionaryField>),
748-
/// An array index expression e.g. `(ARRAY[1, 2])[1]` or `(current_schemas(FALSE))[1]`
749-
ArrayIndex {
750-
obj: Box<Expr>,
751-
indexes: Vec<Expr>,
748+
/// An access of nested data using subscript syntax, for example `array[2]`.
749+
Subscript {
750+
expr: Box<Expr>,
751+
subscript: Box<Subscript>,
752752
},
753753
/// An array expression e.g. `ARRAY[1, 2]`
754754
Array(Array),
@@ -804,6 +804,53 @@ pub enum Expr {
804804
Lambda(LambdaFunction),
805805
}
806806

807+
/// The contents inside the `[` and `]` in a subscript expression.
808+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
809+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
810+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
811+
pub enum Subscript {
812+
/// Accesses the element of the array at the given index.
813+
Index { index: Expr },
814+
815+
/// Accesses a slice of an array on PostgreSQL, e.g.
816+
///
817+
/// ```plaintext
818+
/// => select (array[1,2,3,4,5,6])[2:5];
819+
/// -----------
820+
/// {2,3,4,5}
821+
/// ```
822+
///
823+
/// The lower and/or upper bound can be omitted to slice from the start or
824+
/// end of the array respectively.
825+
///
826+
/// See <https://www.postgresql.org/docs/current/arrays.html#ARRAYS-ACCESSING>.
827+
Slice {
828+
lower_bound: Option<Expr>,
829+
upper_bound: Option<Expr>,
830+
},
831+
}
832+
833+
impl fmt::Display for Subscript {
834+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
835+
match self {
836+
Subscript::Index { index } => write!(f, "{index}"),
837+
Subscript::Slice {
838+
lower_bound,
839+
upper_bound,
840+
} => {
841+
if let Some(lower) = lower_bound {
842+
write!(f, "{lower}")?;
843+
}
844+
write!(f, ":")?;
845+
if let Some(upper) = upper_bound {
846+
write!(f, "{upper}")?;
847+
}
848+
Ok(())
849+
}
850+
}
851+
}
852+
}
853+
807854
/// A lambda function.
808855
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
809856
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -1250,12 +1297,11 @@ impl fmt::Display for Expr {
12501297
Expr::Dictionary(fields) => {
12511298
write!(f, "{{{}}}", display_comma_separated(fields))
12521299
}
1253-
Expr::ArrayIndex { obj, indexes } => {
1254-
write!(f, "{obj}")?;
1255-
for i in indexes {
1256-
write!(f, "[{i}]")?;
1257-
}
1258-
Ok(())
1300+
Expr::Subscript {
1301+
expr,
1302+
subscript: key,
1303+
} => {
1304+
write!(f, "{expr}[{key}]")
12591305
}
12601306
Expr::Array(set) => {
12611307
write!(f, "{set}")

‎src/parser/mod.rs‎

Lines changed: 29 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2582,8 +2582,7 @@ impl<'a> Parser<'a> {
25822582
})
25832583
} else if Token::LBracket == tok {
25842584
if dialect_of!(self is PostgreSqlDialect | DuckDbDialect | GenericDialect) {
2585-
// parse index
2586-
self.parse_array_index(expr)
2585+
self.parse_subscript(expr)
25872586
} else if dialect_of!(self is SnowflakeDialect) {
25882587
self.prev_token();
25892588
self.parse_json_access(expr)
@@ -2611,18 +2610,34 @@ impl<'a> Parser<'a> {
26112610
}
26122611
}
26132612

2614-
pub fn parse_array_index(&mut self, expr: Expr) -> Result<Expr, ParserError> {
2615-
let index = self.parse_expr()?;
2613+
pub fn parse_subscript(&mut self, expr: Expr) -> Result<Expr, ParserError> {
2614+
let parse_upper_bound = |p: &mut Parser<'a>| {
2615+
if let Token::RBracket = p.peek_token().token {
2616+
Ok(None)
2617+
} else {
2618+
p.parse_expr().map(Some)
2619+
}
2620+
};
2621+
let subscript = if self.consume_token(&Token::Colon) {
2622+
Subscript::Slice {
2623+
lower_bound: None,
2624+
upper_bound: parse_upper_bound(self)?,
2625+
}
2626+
} else {
2627+
let expr = self.parse_expr()?;
2628+
if self.consume_token(&Token::Colon) {
2629+
Subscript::Slice {
2630+
lower_bound: Some(expr),
2631+
upper_bound: parse_upper_bound(self)?,
2632+
}
2633+
} else {
2634+
Subscript::Index { index: expr }
2635+
}
2636+
};
26162637
self.expect_token(&Token::RBracket)?;
2617-
let mut indexes: Vec<Expr> = vec![index];
2618-
while self.consume_token(&Token::LBracket) {
2619-
let index = self.parse_expr()?;
2620-
self.expect_token(&Token::RBracket)?;
2621-
indexes.push(index);
2622-
}
2623-
Ok(Expr::ArrayIndex {
2624-
obj: Box::new(expr),
2625-
indexes,
2638+
Ok(Expr::Subscript {
2639+
expr: Box::new(expr),
2640+
subscript: Box::new(subscript),
26262641
})
26272642
}
26282643

@@ -2872,7 +2887,7 @@ impl<'a> Parser<'a> {
28722887
Ok(Self::MUL_DIV_MOD_OP_PREC)
28732888
}
28742889
Token::DoubleColon => Ok(50),
2875-
Token::Colon => Ok(50),
2890+
Token::Colon if dialect_of!(self is SnowflakeDialect) => Ok(50),
28762891
Token::ExclamationMark => Ok(50),
28772892
Token::LBracket | Token::Overlap | Token::CaretAt => Ok(50),
28782893
Token::Arrow

‎tests/sqlparser_duckdb.rs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -528,16 +528,18 @@ fn test_array_index() {
528528
_ => panic!("Expected an expression with alias"),
529529
};
530530
assert_eq!(
531-
&Expr::ArrayIndex {
532-
obj: Box::new(Expr::Array(Array {
531+
&Expr::Subscript {
532+
expr: Box::new(Expr::Array(Array {
533533
elem: vec![
534534
Expr::Value(Value::SingleQuotedString("a".to_owned())),
535535
Expr::Value(Value::SingleQuotedString("b".to_owned())),
536536
Expr::Value(Value::SingleQuotedString("c".to_owned()))
537537
],
538538
named: false
539539
})),
540-
indexes: vec![Expr::Value(number("3"))]
540+
subscript: Box::new(Subscript::Index {
541+
index: Expr::Value(number("3"))
542+
})
541543
},
542544
expr
543545
);

‎tests/sqlparser_postgres.rs‎

Lines changed: 126 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1873,66 +1873,90 @@ fn parse_array_index_expr() {
18731873
let sql = "SELECT foo[0] FROM foos";
18741874
let select = pg_and_generic().verified_only_select(sql);
18751875
assert_eq!(
1876-
&Expr::ArrayIndex {
1877-
obj: Box::new(Expr::Identifier(Ident::new("foo"))),
1878-
indexes: vec![num[0].clone()],
1876+
&Expr::Subscript {
1877+
expr: Box::new(Expr::Identifier(Ident::new("foo"))),
1878+
subscript: Box::new(Subscript::Index {
1879+
index: num[0].clone()
1880+
}),
18791881
},
18801882
expr_from_projection(only(&select.projection)),
18811883
);
18821884

18831885
let sql = "SELECT foo[0][0] FROM foos";
18841886
let select = pg_and_generic().verified_only_select(sql);
18851887
assert_eq!(
1886-
&Expr::ArrayIndex {
1887-
obj: Box::new(Expr::Identifier(Ident::new("foo"))),
1888-
indexes: vec![num[0].clone(), num[0].clone()],
1888+
&Expr::Subscript {
1889+
expr: Box::new(Expr::Subscript {
1890+
expr: Box::new(Expr::Identifier(Ident::new("foo"))),
1891+
subscript: Box::new(Subscript::Index {
1892+
index: num[0].clone()
1893+
}),
1894+
}),
1895+
subscript: Box::new(Subscript::Index {
1896+
index: num[0].clone()
1897+
}),
18891898
},
18901899
expr_from_projection(only(&select.projection)),
18911900
);
18921901

18931902
let sql = r#"SELECT bar[0]["baz"]["fooz"] FROM foos"#;
18941903
let select = pg_and_generic().verified_only_select(sql);
18951904
assert_eq!(
1896-
&Expr::ArrayIndex {
1897-
obj: Box::new(Expr::Identifier(Ident::new("bar"))),
1898-
indexes: vec![
1899-
num[0].clone(),
1900-
Expr::Identifier(Ident {
1901-
value: "baz".to_string(),
1902-
quote_style: Some('"')
1905+
&Expr::Subscript {
1906+
expr: Box::new(Expr::Subscript {
1907+
expr: Box::new(Expr::Subscript {
1908+
expr: Box::new(Expr::Identifier(Ident::new("bar"))),
1909+
subscript: Box::new(Subscript::Index {
1910+
index: num[0].clone()
1911+
})
19031912
}),
1904-
Expr::Identifier(Ident {
1913+
subscript: Box::new(Subscript::Index {
1914+
index: Expr::Identifier(Ident {
1915+
value: "baz".to_string(),
1916+
quote_style: Some('"')
1917+
})
1918+
})
1919+
}),
1920+
subscript: Box::new(Subscript::Index {
1921+
index: Expr::Identifier(Ident {
19051922
value: "fooz".to_string(),
19061923
quote_style: Some('"')
19071924
})
1908-
],
1925+
})
19091926
},
19101927
expr_from_projection(only(&select.projection)),
19111928
);
19121929

19131930
let sql = "SELECT (CAST(ARRAY[ARRAY[2, 3]] AS INT[][]))[1][2]";
19141931
let select = pg_and_generic().verified_only_select(sql);
19151932
assert_eq!(
1916-
&Expr::ArrayIndex {
1917-
obj: Box::new(Expr::Nested(Box::new(Expr::Cast {
1918-
kind: CastKind::Cast,
1919-
expr: Box::new(Expr::Array(Array {
1920-
elem: vec![Expr::Array(Array {
1921-
elem: vec![num[2].clone(), num[3].clone(),],
1933+
&Expr::Subscript {
1934+
expr: Box::new(Expr::Subscript {
1935+
expr: Box::new(Expr::Nested(Box::new(Expr::Cast {
1936+
kind: CastKind::Cast,
1937+
expr: Box::new(Expr::Array(Array {
1938+
elem: vec![Expr::Array(Array {
1939+
elem: vec![num[2].clone(), num[3].clone(),],
1940+
named: true,
1941+
})],
19221942
named: true,
1923-
})],
1924-
named: true,
1925-
})),
1926-
data_type: DataType::Array(ArrayElemTypeDef::SquareBracket(
1927-
Box::new(DataType::Array(ArrayElemTypeDef::SquareBracket(
1928-
Box::new(DataType::Int(None)),
1943+
})),
1944+
data_type: DataType::Array(ArrayElemTypeDef::SquareBracket(
1945+
Box::new(DataType::Array(ArrayElemTypeDef::SquareBracket(
1946+
Box::new(DataType::Int(None)),
1947+
None
1948+
))),
19291949
None
1930-
))),
1931-
None
1932-
)),
1933-
format: None,
1934-
}))),
1935-
indexes: vec![num[1].clone(), num[2].clone()],
1950+
)),
1951+
format: None,
1952+
}))),
1953+
subscript: Box::new(Subscript::Index {
1954+
index: num[1].clone()
1955+
}),
1956+
}),
1957+
subscript: Box::new(Subscript::Index {
1958+
index: num[2].clone()
1959+
}),
19361960
},
19371961
expr_from_projection(only(&select.projection)),
19381962
);
@@ -1948,6 +1972,75 @@ fn parse_array_index_expr() {
19481972
);
19491973
}
19501974

1975+
#[test]
1976+
fn parse_array_subscript() {
1977+
let tests = [
1978+
(
1979+
"(ARRAY[1, 2, 3, 4, 5, 6])[2]",
1980+
Subscript::Index {
1981+
index: Expr::Value(number("2")),
1982+
},
1983+
),
1984+
(
1985+
"(ARRAY[1, 2, 3, 4, 5, 6])[foo]",
1986+
Subscript::Index {
1987+
index: Expr::Identifier(Ident::new("foo")),
1988+
},
1989+
),
1990+
(
1991+
"(ARRAY[1, 2, 3, 4, 5, 6])[2:5]",
1992+
Subscript::Slice {
1993+
lower_bound: Some(Expr::Value(number("2"))),
1994+
upper_bound: Some(Expr::Value(number("5"))),
1995+
},
1996+
),
1997+
(
1998+
"arr[array_length(arr) - 3:array_length(arr) - 1]",
1999+
Subscript::Slice {
2000+
lower_bound: Some(Expr::BinaryOp {
2001+
left: Box::new(call("array_length", [Expr::Identifier(Ident::new("arr"))])),
2002+
op: BinaryOperator::Minus,
2003+
right: Box::new(Expr::Value(number("3"))),
2004+
}),
2005+
upper_bound: Some(Expr::BinaryOp {
2006+
left: Box::new(call("array_length", [Expr::Identifier(Ident::new("arr"))])),
2007+
op: BinaryOperator::Minus,
2008+
right: Box::new(Expr::Value(number("1"))),
2009+
}),
2010+
},
2011+
),
2012+
(
2013+
"(ARRAY[1, 2, 3, 4, 5, 6])[:5]",
2014+
Subscript::Slice {
2015+
lower_bound: None,
2016+
upper_bound: Some(Expr::Value(number("5"))),
2017+
},
2018+
),
2019+
(
2020+
"(ARRAY[1, 2, 3, 4, 5, 6])[2:]",
2021+
Subscript::Slice {
2022+
lower_bound: Some(Expr::Value(number("2"))),
2023+
upper_bound: None,
2024+
},
2025+
),
2026+
(
2027+
"(ARRAY[1, 2, 3, 4, 5, 6])[:]",
2028+
Subscript::Slice {
2029+
lower_bound: None,
2030+
upper_bound: None,
2031+
},
2032+
),
2033+
];
2034+
for (sql, expect) in tests {
2035+
let Expr::Subscript { subscript, .. } = pg_and_generic().verified_expr(sql) else {
2036+
panic!("expected subscript expr");
2037+
};
2038+
assert_eq!(expect, *subscript);
2039+
}
2040+
2041+
// pg_and_generic().verified_expr("schedule[:2][2:]");
2042+
}
2043+
19512044
#[test]
19522045
fn parse_create_index() {
19532046
let sql = "CREATE INDEX IF NOT EXISTS my_index ON my_table(col1,col2)";

0 commit comments

Comments
 (0)