lance_datafusion/
sql.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! SQL Parser utility
5
6use std::any::TypeId;
7
8use datafusion::sql::sqlparser::{
9    ast::{Expr, SelectItem, SetExpr, Statement},
10    dialect::{Dialect, GenericDialect},
11    parser::Parser,
12    tokenizer::{Token, Tokenizer},
13};
14
15use lance_core::{Error, Result};
16use snafu::location;
17#[derive(Debug, Default)]
18struct LanceDialect(GenericDialect);
19
20impl LanceDialect {
21    fn new() -> Self {
22        Self(GenericDialect {})
23    }
24}
25
26impl Dialect for LanceDialect {
27    fn dialect(&self) -> TypeId {
28        self.0.dialect()
29    }
30
31    fn is_identifier_start(&self, ch: char) -> bool {
32        self.0.is_identifier_start(ch)
33    }
34
35    fn is_identifier_part(&self, ch: char) -> bool {
36        self.0.is_identifier_part(ch)
37    }
38
39    fn is_delimited_identifier_start(&self, ch: char) -> bool {
40        ch == '`'
41    }
42}
43
44/// Parse sql filter to Expression.
45pub(crate) fn parse_sql_filter(filter: &str) -> Result<Expr> {
46    let sql = format!("SELECT 1 FROM t WHERE {filter}");
47    let statement = parse_statement(&sql)?;
48
49    let selection = if let Statement::Query(query) = &statement {
50        if let SetExpr::Select(s) = query.body.as_ref() {
51            s.selection.as_ref()
52        } else {
53            None
54        }
55    } else {
56        None
57    };
58    let expr = selection
59        .ok_or_else(|| Error::io(format!("Filter is not valid: {filter}"), location!()))?;
60    Ok(expr.clone())
61}
62
63/// Parse a SQL expression to Expression. This is more lenient than parse_sql_filter
64/// as it can be used for projection expressions as well.
65pub(crate) fn parse_sql_expr(expr: &str) -> Result<Expr> {
66    let sql = format!("SELECT {expr} FROM t");
67    let statement = parse_statement(&sql)?;
68
69    let selection = if let Statement::Query(query) = &statement {
70        if let SetExpr::Select(s) = query.body.as_ref() {
71            if let SelectItem::UnnamedExpr(expr) = &s.projection[0] {
72                Some(expr)
73            } else {
74                None
75            }
76        } else {
77            None
78        }
79    } else {
80        None
81    };
82    let expr = selection
83        .ok_or_else(|| Error::io(format!("Expression is not valid: {expr}"), location!()))?;
84    Ok(expr.clone())
85}
86
87fn parse_statement(statement: &str) -> Result<Statement> {
88    let dialect = LanceDialect::new();
89
90    // Hack to allow == as equals
91    // This is used to parse PyArrow expressions from strings.
92    // See: https://github.com/sqlparser-rs/sqlparser-rs/pull/815#issuecomment-1450714278
93    let mut tokenizer = Tokenizer::new(&dialect, statement);
94    let mut tokens = Vec::new();
95    let mut token_iter = tokenizer.tokenize()?.into_iter();
96    let mut prev_token = token_iter.next().unwrap();
97    for next_token in token_iter {
98        if let (Token::Eq, Token::Eq) = (&prev_token, &next_token) {
99            continue; // skip second equals
100        }
101        let token = std::mem::replace(&mut prev_token, next_token);
102        tokens.push(token);
103    }
104    tokens.push(prev_token);
105
106    Ok(Parser::new(&dialect)
107        .with_tokens(tokens)
108        .parse_statement()?)
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    use datafusion::sql::sqlparser::{
116        ast::{BinaryOperator, Ident, Value, ValueWithSpan},
117        tokenizer::Span,
118    };
119
120    #[test]
121    fn test_double_equal() {
122        let expr = parse_sql_filter("a == b").unwrap();
123        assert_eq!(
124            Expr::BinaryOp {
125                left: Box::new(Expr::Identifier(Ident::new("a"))),
126                op: BinaryOperator::Eq,
127                right: Box::new(Expr::Identifier(Ident::new("b")))
128            },
129            expr
130        );
131    }
132
133    #[test]
134    fn test_like() {
135        let expr = parse_sql_filter("a LIKE 'abc%'").unwrap();
136        assert_eq!(
137            Expr::Like {
138                negated: false,
139                expr: Box::new(Expr::Identifier(Ident::new("a"))),
140                pattern: Box::new(Expr::Value(ValueWithSpan {
141                    value: Value::SingleQuotedString("abc%".to_string()),
142                    span: Span::empty(),
143                })),
144                escape_char: None,
145                any: false,
146            },
147            expr
148        );
149    }
150
151    #[test]
152    fn test_quoted_ident() {
153        // CUBE is a SQL keyword, so it must be quoted.
154        let expr = parse_sql_filter("`a:Test_Something` == `CUBE`").unwrap();
155        assert_eq!(
156            Expr::BinaryOp {
157                left: Box::new(Expr::Identifier(Ident::with_quote('`', "a:Test_Something"))),
158                op: BinaryOperator::Eq,
159                right: Box::new(Expr::Identifier(Ident::with_quote('`', "CUBE")))
160            },
161            expr
162        );
163
164        let expr = parse_sql_filter("`outer field`.`inner field` == 1").unwrap();
165        assert_eq!(
166            Expr::BinaryOp {
167                left: Box::new(Expr::CompoundIdentifier(vec![
168                    Ident::with_quote('`', "outer field"),
169                    Ident::with_quote('`', "inner field")
170                ])),
171                op: BinaryOperator::Eq,
172                right: Box::new(Expr::Value(ValueWithSpan {
173                    value: Value::Number("1".to_string(), false),
174                    span: Span::empty(),
175                })),
176            },
177            expr
178        );
179    }
180}
lance_datafusion/sql.rs

lance_datafusion/
sql.rs
