This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 6f10dbc124 Support parsing SQL strings to Exprs (#10995)
6f10dbc124 is described below
commit 6f10dbc124de6b999f47a7f45e5dc1c486f8a4dd
Author: Xin Li <[email protected]>
AuthorDate: Sun Jun 23 18:11:36 2024 +0800
Support parsing SQL strings to Exprs (#10995)
* Draft parse_sql
* Allow stirng pass
* Complete sql to expr support
* Add examples
* Add unit tests
* Fix format
* Remove async for trival operation and add parquet demo
* Fix comments
* fix comments
* fix comments
* Fix doc link
---
datafusion-examples/README.md | 1 +
datafusion-examples/examples/parse_sql_expr.rs | 157 +++++++++++++++++++++
datafusion/core/src/dataframe/mod.rs | 27 ++++
.../datasource/physical_plan/parquet/statistics.rs | 2 +-
datafusion/core/src/execution/context/mod.rs | 26 ++++
datafusion/core/src/execution/session_state.rs | 67 +++++++--
datafusion/core/tests/expr_api/mod.rs | 1 +
datafusion/core/tests/expr_api/parse_sql_expr.rs | 93 ++++++++++++
datafusion/sql/src/parser.rs | 23 ++-
9 files changed, 383 insertions(+), 14 deletions(-)
diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md
index 0eef9091b2..0446a00818 100644
--- a/datafusion-examples/README.md
+++ b/datafusion-examples/README.md
@@ -66,6 +66,7 @@ cargo run --example csv_sql
- [`parquet_sql.rs`](examples/parquet_sql.rs): Build and run a query plan from
a SQL statement against a local Parquet file
- [`parquet_sql_multiple_files.rs`](examples/parquet_sql_multiple_files.rs):
Build and run a query plan from a SQL statement against multiple local Parquet
files
- ['parquet_exec_visitor.rs'](examples/parquet_exec_visitor.rs): Extract
statistics by visiting an ExecutionPlan after execution
+- [`parse_sql_expr.rs`](examples/parse_sql_expr.rs): Parse SQL text into
Datafusion `Expr`.
- [`plan_to_sql.rs`](examples/plan_to_sql.rs): Generate SQL from Datafusion
`Expr` and `LogicalPlan`
- [`pruning.rs`](examples/parquet_sql.rs): Use pruning to rule out files based
on statistics
- [`query-aws-s3.rs`](examples/external_dependency/query-aws-s3.rs): Configure
`object_store` and run a query against files stored in AWS S3
diff --git a/datafusion-examples/examples/parse_sql_expr.rs
b/datafusion-examples/examples/parse_sql_expr.rs
new file mode 100644
index 0000000000..6444eb68b6
--- /dev/null
+++ b/datafusion-examples/examples/parse_sql_expr.rs
@@ -0,0 +1,157 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::{
+ assert_batches_eq,
+ error::Result,
+ prelude::{ParquetReadOptions, SessionContext},
+};
+use datafusion_common::DFSchema;
+use datafusion_expr::{col, lit};
+use datafusion_sql::unparser::Unparser;
+
+/// This example demonstrates the programmatic parsing of SQL expressions using
+/// the DataFusion [`SessionContext::parse_sql_expr`] API or the
[`DataFrame::parse_sql_expr`] API.
+///
+///
+/// The code in this example shows how to:
+///
+/// 1. [`simple_session_context_parse_sql_expr_demo`]: Parse a simple SQL text
into a logical
+/// expression using a schema at [`SessionContext`].
+///
+/// 2. [`simple_dataframe_parse_sql_expr_demo`]: Parse a simple SQL text into
a logical expression
+/// using a schema at [`DataFrame`].
+///
+/// 3. [`query_parquet_demo`]: Query a parquet file using the parsed_sql_expr
from a DataFrame.
+///
+/// 4. [`round_trip_parse_sql_expr_demo`]: Parse a SQL text and convert it
back to SQL using [`Unparser`].
+
+#[tokio::main]
+async fn main() -> Result<()> {
+ // See how to evaluate expressions
+ simple_session_context_parse_sql_expr_demo()?;
+ simple_dataframe_parse_sql_expr_demo().await?;
+ query_parquet_demo().await?;
+ round_trip_parse_sql_expr_demo().await?;
+ Ok(())
+}
+
+/// DataFusion can parse a SQL text to a logical expression against a schema
at [`SessionContext`].
+fn simple_session_context_parse_sql_expr_demo() -> Result<()> {
+ let sql = "a < 5 OR a = 8";
+ let expr = col("a").lt(lit(5_i64)).or(col("a").eq(lit(8_i64)));
+
+ // provide type information that `a` is an Int32
+ let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
+ let df_schema = DFSchema::try_from(schema).unwrap();
+ let ctx = SessionContext::new();
+
+ let parsed_expr = ctx.parse_sql_expr(sql, &df_schema)?;
+
+ assert_eq!(parsed_expr, expr);
+
+ Ok(())
+}
+
+/// DataFusion can parse a SQL text to an logical expression using schema at
[`DataFrame`].
+async fn simple_dataframe_parse_sql_expr_demo() -> Result<()> {
+ let sql = "int_col < 5 OR double_col = 8.0";
+ let expr = col("int_col")
+ .lt(lit(5_i64))
+ .or(col("double_col").eq(lit(8.0_f64)));
+
+ let ctx = SessionContext::new();
+ let testdata = datafusion::test_util::parquet_test_data();
+ let df = ctx
+ .read_parquet(
+ &format!("{testdata}/alltypes_plain.parquet"),
+ ParquetReadOptions::default(),
+ )
+ .await?;
+
+ let parsed_expr = df.parse_sql_expr(sql)?;
+
+ assert_eq!(parsed_expr, expr);
+
+ Ok(())
+}
+
+async fn query_parquet_demo() -> Result<()> {
+ let ctx = SessionContext::new();
+ let testdata = datafusion::test_util::parquet_test_data();
+ let df = ctx
+ .read_parquet(
+ &format!("{testdata}/alltypes_plain.parquet"),
+ ParquetReadOptions::default(),
+ )
+ .await?;
+
+ let df = df
+ .clone()
+ .select(vec![
+ df.parse_sql_expr("int_col")?,
+ df.parse_sql_expr("double_col")?,
+ ])?
+ .filter(df.parse_sql_expr("int_col < 5 OR double_col = 8.0")?)?
+ .aggregate(
+ vec![df.parse_sql_expr("double_col")?],
+ vec![df.parse_sql_expr("SUM(int_col) as sum_int_col")?],
+ )?
+ // Directly parsing the SQL text into a sort expression is not
supported yet, so
+ // construct it programatically
+ .sort(vec![col("double_col").sort(false, false)])?
+ .limit(0, Some(1))?;
+
+ let result = df.collect().await?;
+
+ assert_batches_eq!(
+ &[
+ "+------------+----------------------+",
+ "| double_col | sum(?table?.int_col) |",
+ "+------------+----------------------+",
+ "| 10.1 | 4 |",
+ "+------------+----------------------+",
+ ],
+ &result
+ );
+
+ Ok(())
+}
+
+/// DataFusion can parse a SQL text and convert it back to SQL using
[`Unparser`].
+async fn round_trip_parse_sql_expr_demo() -> Result<()> {
+ let sql = "((int_col < 5) OR (double_col = 8))";
+
+ let ctx = SessionContext::new();
+ let testdata = datafusion::test_util::parquet_test_data();
+ let df = ctx
+ .read_parquet(
+ &format!("{testdata}/alltypes_plain.parquet"),
+ ParquetReadOptions::default(),
+ )
+ .await?;
+
+ let parsed_expr = df.parse_sql_expr(sql)?;
+
+ let unparser = Unparser::default();
+ let round_trip_sql = unparser.expr_to_sql(&parsed_expr)?.to_string();
+
+ assert_eq!(sql, round_trip_sql);
+
+ Ok(())
+}
diff --git a/datafusion/core/src/dataframe/mod.rs
b/datafusion/core/src/dataframe/mod.rs
index d55caf1404..90cde08768 100644
--- a/datafusion/core/src/dataframe/mod.rs
+++ b/datafusion/core/src/dataframe/mod.rs
@@ -177,6 +177,33 @@ impl DataFrame {
}
}
+ /// Creates logical expression from a SQL query text.
+ /// The expression is created and processed againt the current schema.
+ ///
+ /// # Example: Parsing SQL queries
+ /// ```
+ /// # use arrow::datatypes::{DataType, Field, Schema};
+ /// # use datafusion::prelude::*;
+ /// # use datafusion_common::{DFSchema, Result};
+ /// # #[tokio::main]
+ /// # async fn main() -> Result<()> {
+ /// // datafusion will parse number as i64 first.
+ /// let sql = "a > 1 and b in (1, 10)";
+ /// let expected = col("a").gt(lit(1 as i64))
+ /// .and(col("b").in_list(vec![lit(1 as i64), lit(10 as i64)], false));
+ /// let ctx = SessionContext::new();
+ /// let df = ctx.read_csv("tests/data/example.csv",
CsvReadOptions::new()).await?;
+ /// let expr = df.parse_sql_expr(sql)?;
+ /// assert_eq!(expected, expr);
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub fn parse_sql_expr(&self, sql: &str) -> Result<Expr> {
+ let df_schema = self.schema();
+
+ self.session_state.create_logical_expr(sql, df_schema)
+ }
+
/// Consume the DataFrame and produce a physical plan
pub async fn create_physical_plan(self) -> Result<Arc<dyn ExecutionPlan>> {
self.session_state.create_physical_plan(&self.plan).await
diff --git a/datafusion/core/src/datasource/physical_plan/parquet/statistics.rs
b/datafusion/core/src/datasource/physical_plan/parquet/statistics.rs
index 43483a0b66..17bfe72dbd 100644
--- a/datafusion/core/src/datasource/physical_plan/parquet/statistics.rs
+++ b/datafusion/core/src/datasource/physical_plan/parquet/statistics.rs
@@ -67,7 +67,7 @@ pub(crate) fn from_bytes_to_f16(b: &[u8]) -> Option<f16> {
// Copy from arrow-rs
//
https://github.com/apache/arrow-rs/blob/198af7a3f4aa20f9bd003209d9f04b0f37bb120e/parquet/src/arrow/buffer/bit_util.rs#L54
// Convert the byte slice to fixed length byte array with the length of N.
-pub fn sign_extend_be<const N: usize>(b: &[u8]) -> [u8; N] {
+fn sign_extend_be<const N: usize>(b: &[u8]) -> [u8; N] {
assert!(b.len() <= N, "Array too large, expected less than {N}");
let is_negative = (b[0] & 128u8) == 128u8;
let mut result = if is_negative { [255u8; N] } else { [0u8; N] };
diff --git a/datafusion/core/src/execution/context/mod.rs
b/datafusion/core/src/execution/context/mod.rs
index c44e974260..57e94cad06 100644
--- a/datafusion/core/src/execution/context/mod.rs
+++ b/datafusion/core/src/execution/context/mod.rs
@@ -486,6 +486,32 @@ impl SessionContext {
self.execute_logical_plan(plan).await
}
+ /// Creates logical expresssions from SQL query text.
+ ///
+ /// # Example: Parsing SQL queries
+ ///
+ /// ```
+ /// # use arrow::datatypes::{DataType, Field, Schema};
+ /// # use datafusion::prelude::*;
+ /// # use datafusion_common::{DFSchema, Result};
+ /// # #[tokio::main]
+ /// # async fn main() -> Result<()> {
+ /// // datafusion will parse number as i64 first.
+ /// let sql = "a > 10";
+ /// let expected = col("a").gt(lit(10 as i64));
+ /// // provide type information that `a` is an Int32
+ /// let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
+ /// let df_schema = DFSchema::try_from(schema).unwrap();
+ /// let expr = SessionContext::new()
+ /// .parse_sql_expr(sql, &df_schema)?;
+ /// assert_eq!(expected, expr);
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub fn parse_sql_expr(&self, sql: &str, df_schema: &DFSchema) ->
Result<Expr> {
+ self.state.read().create_logical_expr(sql, df_schema)
+ }
+
/// Execute the [`LogicalPlan`], return a [`DataFrame`]. This API
/// is not featured limited (so all SQL such as `CREATE TABLE` and
/// `COPY` will be run).
diff --git a/datafusion/core/src/execution/session_state.rs
b/datafusion/core/src/execution/session_state.rs
index e9441a89cd..4173a8dcc4 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -67,7 +67,8 @@ use datafusion_physical_expr::create_physical_expr;
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_sql::parser::{DFParser, Statement};
-use datafusion_sql::planner::{ContextProvider, ParserOptions, SqlToRel};
+use datafusion_sql::planner::{ContextProvider, ParserOptions, PlannerContext,
SqlToRel};
+use sqlparser::ast::Expr as SQLExpr;
use sqlparser::dialect::dialect_from_str;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
@@ -490,6 +491,27 @@ impl SessionState {
Ok(statement)
}
+ /// parse a sql string into a sqlparser-rs AST [`SQLExpr`].
+ ///
+ /// See [`Self::create_logical_expr`] for parsing sql to [`Expr`].
+ pub fn sql_to_expr(
+ &self,
+ sql: &str,
+ dialect: &str,
+ ) -> datafusion_common::Result<SQLExpr> {
+ let dialect = dialect_from_str(dialect).ok_or_else(|| {
+ plan_datafusion_err!(
+ "Unsupported SQL dialect: {dialect}. Available dialects: \
+ Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake,
Redshift, \
+ MsSQL, ClickHouse, BigQuery, Ansi."
+ )
+ })?;
+
+ let expr = DFParser::parse_sql_into_expr_with_dialect(sql,
dialect.as_ref())?;
+
+ Ok(expr)
+ }
+
/// Resolve all table references in the SQL statement. Does not include
CTE references.
///
/// See [`catalog::resolve_table_references`] for more information.
@@ -520,10 +542,6 @@ impl SessionState {
tables: HashMap::with_capacity(references.len()),
};
- let enable_ident_normalization =
- self.config.options().sql_parser.enable_ident_normalization;
- let parse_float_as_decimal =
- self.config.options().sql_parser.parse_float_as_decimal;
for reference in references {
let resolved = &self.resolve_table_ref(reference);
if let Entry::Vacant(v) =
provider.tables.entry(resolved.to_string()) {
@@ -535,16 +553,19 @@ impl SessionState {
}
}
- let query = SqlToRel::new_with_options(
- &provider,
- ParserOptions {
- parse_float_as_decimal,
- enable_ident_normalization,
- },
- );
+ let query = SqlToRel::new_with_options(&provider,
self.get_parser_options());
query.statement_to_plan(statement)
}
+ fn get_parser_options(&self) -> ParserOptions {
+ let sql_parser_options = &self.config.options().sql_parser;
+
+ ParserOptions {
+ parse_float_as_decimal: sql_parser_options.parse_float_as_decimal,
+ enable_ident_normalization:
sql_parser_options.enable_ident_normalization,
+ }
+ }
+
/// Creates a [`LogicalPlan`] from the provided SQL string. This
/// interface will plan any SQL DataFusion supports, including DML
/// like `CREATE TABLE`, and `COPY` (which can write to local
@@ -567,6 +588,28 @@ impl SessionState {
Ok(plan)
}
+ /// Creates a datafusion style AST [`Expr`] from a SQL string.
+ ///
+ /// See example on
[SessionContext::parse_sql_expr](crate::execution::context::SessionContext::parse_sql_expr)
+ pub fn create_logical_expr(
+ &self,
+ sql: &str,
+ df_schema: &DFSchema,
+ ) -> datafusion_common::Result<Expr> {
+ let dialect = self.config.options().sql_parser.dialect.as_str();
+
+ let sql_expr = self.sql_to_expr(sql, dialect)?;
+
+ let provider = SessionContextProvider {
+ state: self,
+ tables: HashMap::new(),
+ };
+
+ let query = SqlToRel::new_with_options(&provider,
self.get_parser_options());
+
+ query.sql_to_expr(sql_expr, df_schema, &mut PlannerContext::new())
+ }
+
/// Optimizes the logical plan by applying optimizer rules.
pub fn optimize(&self, plan: &LogicalPlan) ->
datafusion_common::Result<LogicalPlan> {
if let LogicalPlan::Explain(e) = plan {
diff --git a/datafusion/core/tests/expr_api/mod.rs
b/datafusion/core/tests/expr_api/mod.rs
index 7085333bee..f36f2d5398 100644
--- a/datafusion/core/tests/expr_api/mod.rs
+++ b/datafusion/core/tests/expr_api/mod.rs
@@ -30,6 +30,7 @@ use sqlparser::ast::NullTreatment;
/// Tests of using and evaluating `Expr`s outside the context of a LogicalPlan
use std::sync::{Arc, OnceLock};
+mod parse_sql_expr;
mod simplification;
#[test]
diff --git a/datafusion/core/tests/expr_api/parse_sql_expr.rs
b/datafusion/core/tests/expr_api/parse_sql_expr.rs
new file mode 100644
index 0000000000..991579b5a3
--- /dev/null
+++ b/datafusion/core/tests/expr_api/parse_sql_expr.rs
@@ -0,0 +1,93 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow_schema::{DataType, Field, Schema};
+use datafusion::prelude::{CsvReadOptions, SessionContext};
+use datafusion_common::{DFSchemaRef, Result, ToDFSchema};
+use datafusion_expr::Expr;
+use datafusion_sql::unparser::Unparser;
+
+/// A schema like:
+///
+/// a: Int32 (possibly with nulls)
+/// b: Int32
+/// s: Float32
+fn schema() -> DFSchemaRef {
+ Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, false),
+ Field::new("c", DataType::Float32, false),
+ ])
+ .to_dfschema_ref()
+ .unwrap()
+}
+
+#[tokio::test]
+async fn round_trip_parse_sql_expr() -> Result<()> {
+ let tests = vec![
+ "(a = 10)",
+ "((a = 10) AND (b <> 20))",
+ "((a = 10) OR (b <> 20))",
+ "(((a = 10) AND (b <> 20)) OR (c = a))",
+ "((a = 10) AND b IN (20, 30))",
+ "((a = 10) AND b NOT IN (20, 30))",
+ "sum(a)",
+ "(sum(a) + 1)",
+ "(MIN(a) + MAX(b))",
+ "(MIN(a) + (MAX(b) * sum(c)))",
+ "(MIN(a) + ((MAX(b) * sum(c)) / 10))",
+ ];
+
+ for test in tests {
+ round_trip_session_context(test)?;
+ round_trip_dataframe(test).await?;
+ }
+
+ Ok(())
+}
+
+fn round_trip_session_context(sql: &str) -> Result<()> {
+ let ctx = SessionContext::new();
+ let df_schema = schema();
+ let expr = ctx.parse_sql_expr(sql, &df_schema)?;
+ let sql2 = unparse_sql_expr(&expr)?;
+ assert_eq!(sql, sql2);
+
+ Ok(())
+}
+
+async fn round_trip_dataframe(sql: &str) -> Result<()> {
+ let ctx = SessionContext::new();
+ let df = ctx
+ .read_csv(
+ &"tests/data/example.csv".to_string(),
+ CsvReadOptions::default(),
+ )
+ .await?;
+ let expr = df.parse_sql_expr(sql)?;
+ let sql2 = unparse_sql_expr(&expr)?;
+ assert_eq!(sql, sql2);
+
+ Ok(())
+}
+
+fn unparse_sql_expr(expr: &Expr) -> Result<String> {
+ let unparser = Unparser::default();
+
+ let round_trip_sql = unparser.expr_to_sql(expr)?.to_string();
+ Ok(round_trip_sql)
+}
diff --git a/datafusion/sql/src/parser.rs b/datafusion/sql/src/parser.rs
index bbc3a52f07..d2f3f508a3 100644
--- a/datafusion/sql/src/parser.rs
+++ b/datafusion/sql/src/parser.rs
@@ -22,7 +22,7 @@ use std::fmt;
use sqlparser::{
ast::{
- ColumnDef, ColumnOptionDef, ObjectName, OrderByExpr, Query,
+ ColumnDef, ColumnOptionDef, Expr, ObjectName, OrderByExpr, Query,
Statement as SQLStatement, TableConstraint, Value,
},
dialect::{keywords::Keyword, Dialect, GenericDialect},
@@ -323,6 +323,14 @@ impl<'a> DFParser<'a> {
Ok(stmts)
}
+ pub fn parse_sql_into_expr_with_dialect(
+ sql: &str,
+ dialect: &dyn Dialect,
+ ) -> Result<Expr, ParserError> {
+ let mut parser = DFParser::new_with_dialect(sql, dialect)?;
+ parser.parse_expr()
+ }
+
/// Report an unexpected token
fn expected<T>(
&self,
@@ -367,6 +375,19 @@ impl<'a> DFParser<'a> {
}
}
+ pub fn parse_expr(&mut self) -> Result<Expr, ParserError> {
+ if let Token::Word(w) = self.parser.peek_token().token {
+ match w.keyword {
+ Keyword::CREATE | Keyword::COPY | Keyword::EXPLAIN => {
+ return parser_err!("Unsupported command in expression");
+ }
+ _ => {}
+ }
+ }
+
+ self.parser.parse_expr()
+ }
+
/// Parse a SQL `COPY TO` statement
pub fn parse_copy(&mut self) -> Result<Statement, ParserError> {
// parse as a query
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
