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::ast::{BinaryOperator, Ident, Value};
116
117    #[test]
118    fn test_double_equal() {
119        let expr = parse_sql_filter("a == b").unwrap();
120        assert_eq!(
121            Expr::BinaryOp {
122                left: Box::new(Expr::Identifier(Ident::new("a"))),
123                op: BinaryOperator::Eq,
124                right: Box::new(Expr::Identifier(Ident::new("b")))
125            },
126            expr
127        );
128    }
129
130    #[test]
131    fn test_like() {
132        let expr = parse_sql_filter("a LIKE 'abc%'").unwrap();
133        assert_eq!(
134            Expr::Like {
135                negated: false,
136                expr: Box::new(Expr::Identifier(Ident::new("a"))),
137                pattern: Box::new(Expr::Value(Value::SingleQuotedString("abc%".to_string()))),
138                escape_char: None,
139                any: false,
140            },
141            expr
142        );
143    }
144
145    #[test]
146    fn test_quoted_ident() {
147        // CUBE is a SQL keyword, so it must be quoted.
148        let expr = parse_sql_filter("`a:Test_Something` == `CUBE`").unwrap();
149        assert_eq!(
150            Expr::BinaryOp {
151                left: Box::new(Expr::Identifier(Ident::with_quote('`', "a:Test_Something"))),
152                op: BinaryOperator::Eq,
153                right: Box::new(Expr::Identifier(Ident::with_quote('`', "CUBE")))
154            },
155            expr
156        );
157
158        let expr = parse_sql_filter("`outer field`.`inner field` == 1").unwrap();
159        assert_eq!(
160            Expr::BinaryOp {
161                left: Box::new(Expr::CompoundIdentifier(vec![
162                    Ident::with_quote('`', "outer field"),
163                    Ident::with_quote('`', "inner field")
164                ])),
165                op: BinaryOperator::Eq,
166                right: Box::new(Expr::Value(Value::Number("1".to_string(), false))),
167            },
168            expr
169        );
170    }
171}
lance_datafusion/sql.rs

lance_datafusion/
sql.rs
