Reworking to do token parsing of sql query instead of string manipulation

Author: timsaucer

python/datafusion/context.py

@@ -27,7 +27,6 @@
 except ImportError:
     from typing_extensions import deprecated  # Python 3.12
 
-import uuid
 
 import pyarrow as pa
 
@@ -612,25 +611,41 @@ def sql(
         Returns:
             DataFrame representation of the SQL query.
         """
-        if named_params:
-            for alias, param in named_params.items():
-                if isinstance(param, DataFrame):
-                    view_name = str(uuid.uuid4()).replace("-", "_")
-                    view_name = f"view_{view_name}"
-                    self.ctx.create_temporary_view(
-                        view_name, param.df, replace_if_exists=True
-                    )
-                    replace_str = view_name
-                else:
-                    replace_str = str(param)
 
-                query = query.replace(f"{{{alias}}}", replace_str)
+        def scalar_params(**p: Any) -> list[tuple[str, pa.Scalar]]:
+            if p is None:
+                return []
+
+            return [
+                (name, pa.scalar(value))
+                for (name, value) in p.items()
+                if not isinstance(value, DataFrame)
+            ]
 
-        if options is None:
-            return DataFrame(self.ctx.sql(query))
-        return DataFrame(self.ctx.sql_with_options(query, options.options_internal))
+        def dataframe_params(**p: Any) -> list[tuple[str, DataFrame]]:
+            if p is None:
+                return []
 
-    def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame:
+            return [
+                (name, value.df)
+                for (name, value) in p.items()
+                if isinstance(value, DataFrame)
+            ]
+
+        options_raw = options.options_internal if options is not None else None
+
+        return DataFrame(
+            self.ctx.sql_with_options(
+                query,
+                options=options_raw,
+                scalar_params=scalar_params(**named_params),
+                dataframe_params=dataframe_params(**named_params),
+            )
+        )
+
+    def sql_with_options(
+        self, query: str, options: SQLOptions, **named_params: Any
+    ) -> DataFrame:
         """Create a :py:class:`~datafusion.dataframe.DataFrame` from SQL query text.
 
         This function will first validate that the query is allowed by the
@@ -639,11 +654,12 @@ def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame:
         Args:
             query: SQL query text.
             options: SQL options.
+            named_params: Provides substitution in the query string.
 
         Returns:
             DataFrame representation of the SQL query.
         """
-        return self.sql(query, options)
+        return self.sql(query, options, **named_params)
 
     def create_dataframe(
         self,

python/tests/test_sql.py

@@ -21,7 +21,7 @@
 import pyarrow as pa
 import pyarrow.dataset as ds
 import pytest
-from datafusion import col, udf
+from datafusion import SessionContext, col, udf
 from datafusion.object_store import Http
 from pyarrow.csv import write_csv
 
@@ -535,11 +535,25 @@ def test_register_listing_table(
     assert dict(zip(rd["grp"], rd["count"])) == {"a": 3, "b": 2}
 
 
-def test_parameterized_sql(ctx, tmp_path) -> None:
+def test_parameterized_df_in_sql(ctx, tmp_path) -> None:
     path = helpers.write_parquet(tmp_path / "a.parquet", helpers.data())
+
     df = ctx.read_parquet(path)
     result = ctx.sql(
-        "SELECT COUNT(a) AS cnt FROM {replaced_df}", replaced_df=df
+        "SELECT COUNT(a) AS cnt FROM $replaced_df", replaced_df=df
     ).collect()
     result = pa.Table.from_batches(result)
     assert result.to_pydict() == {"cnt": [100]}
+
+
+def test_parameterized_pass_through_in_sql(ctx: SessionContext) -> None:
+    # Test the parameters that should be handled by the parser rather
+    # than our manipulation of the query string by searching for tokens
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3, 4])],
+        names=["a"],
+    )
+
+    ctx.register_record_batches("t", [[batch]])
+    result = ctx.sql("SELECT a FROM t WHERE a < $val", val=3)
+    assert result.to_pydict() == {"a": [1, 2]}

src/context.rs

@@ -32,14 +32,16 @@ use pyo3::exceptions::{PyKeyError, PyValueError};
 use pyo3::prelude::*;
 
 use crate::catalog::{PyCatalog, RustWrappedPyCatalogProvider};
+use crate::common::data_type::PyScalarValue;
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
-use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionResult};
+use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::expr::sort_expr::PySortExpr;
 use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::PyRecordBatchStream;
 use crate::sql::exceptions::py_value_err;
 use crate::sql::logical::PyLogicalPlan;
+use crate::sql::util::replace_placeholders_with_table_names;
 use crate::store::StorageContexts;
 use crate::table::{PyTable, TempViewTable};
 use crate::udaf::PyAggregateUDF;
@@ -427,27 +429,54 @@ impl PySessionContext {
         self.ctx.register_udtf(&name, func);
     }
 
-    /// Returns a PyDataFrame whose plan corresponds to the SQL statement.
-    pub fn sql(&self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
-        let result = self.ctx.sql(query);
-        let df = wait_for_future(py, result)??;
-        Ok(PyDataFrame::new(df))
-    }
-
-    #[pyo3(signature = (query, options=None))]
+    #[pyo3(signature = (query, options=None, scalar_params=vec![], dataframe_params=vec![]))]
     pub fn sql_with_options(
         &self,
+        py: Python,
         query: &str,
         options: Option<PySQLOptions>,
-        py: Python,
+        scalar_params: Vec<(String, PyScalarValue)>,
+        dataframe_params: Vec<(String, PyDataFrame)>,
     ) -> PyDataFusionResult<PyDataFrame> {
         let options = if let Some(options) = options {
             options.options
         } else {
             SQLOptions::new()
         };
-        let result = self.ctx.sql_with_options(query, options);
-        let df = wait_for_future(py, result)??;
+
+        let scalar_params = scalar_params
+            .into_iter()
+            .map(|(name, value)| (name, ScalarValue::from(value)))
+            .collect::<Vec<_>>();
+
+        let dataframe_params = dataframe_params
+            .into_iter()
+            .map(|(name, df)| {
+                let uuid = Uuid::new_v4().to_string().replace("-", "");
+                let view_name = format!("view_{uuid}");
+
+                self.create_temporary_view(py, view_name.as_str(), df, true)?;
+                Ok((name, view_name))
+            })
+            .collect::<Result<HashMap<_, _>, PyDataFusionError>>()?;
+
+        let state = self.ctx.state();
+        let dialect = state.config().options().sql_parser.dialect.as_str();
+
+        let query = replace_placeholders_with_table_names(query, dialect, dataframe_params)?;
+
+        println!("using scalar params: {scalar_params:?}");
+        let df = wait_for_future(py, async {
+            self.ctx
+                .sql_with_options(&query, options)
+                .await
+                .map_err(|err| {
+                    println!("error before param replacement: {}", err);
+                    err
+                })?
+                .with_param_values(scalar_params)
+        })??;
+
         Ok(PyDataFrame::new(df))
     }
 

src/sql.rs

@@ -17,3 +17,4 @@
 
 pub mod exceptions;
 pub mod logical;
+pub(crate) mod util;

src/sql/util.rs

@@ -0,0 +1,93 @@
+use datafusion::common::{internal_err, plan_datafusion_err, DataFusionError};
+use datafusion::logical_expr::sqlparser::dialect::dialect_from_str;
+use datafusion::sql::sqlparser::dialect::Dialect;
+use datafusion::sql::sqlparser::keywords::Keyword;
+use datafusion::sql::sqlparser::parser::Parser;
+use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer, Word};
+use std::collections::HashMap;
+
+fn value_from_replacements(
+    placeholder: &str,
+    replacements: &HashMap<String, String>,
+) -> Option<Token> {
+    if let Some(pattern) = placeholder.strip_prefix("$") {
+        replacements.get(pattern).map(|replacement| {
+            Token::Word(Word {
+                value: replacement.to_owned(),
+                quote_style: None,
+                keyword: Keyword::NoKeyword,
+            })
+        })
+    } else {
+        None
+    }
+}
+
+fn table_names_are_valid(dialect: &dyn Dialect, replacements: &HashMap<String, String>) -> bool {
+    for name in replacements.values() {
+        let tokens = Tokenizer::new(dialect, name).tokenize().unwrap();
+        if tokens.len() != 1 {
+            // We should get exactly one token for our temporary table name
+            return false;
+        }
+
+        if let Token::Word(word) = &tokens[0] {
+            // Generated table names should be not quoted or have keywords
+            if word.quote_style.is_some() || word.keyword != Keyword::NoKeyword {
+                return false;
+            }
+        } else {
+            // We should always parse table names to a Word
+            return false;
+        }
+    }
+
+    true
+}
+
+pub(crate) fn replace_placeholders_with_table_names(
+    query: &str,
+    dialect: &str,
+    replacements: HashMap<String, String>,
+) -> Result<String, DataFusionError> {
+    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, DuckDB, Databricks."
+        )
+    })?;
+
+    if !table_names_are_valid(dialect.as_ref(), &replacements) {
+        return internal_err!("Invalid generated table name when replacing placeholders");
+    }
+    let tokens = Tokenizer::new(dialect.as_ref(), query).tokenize().unwrap();
+
+    let replaced_tokens = tokens
+        .into_iter()
+        .map(|token| {
+            if let Token::Word(word) = &token {
+                let Word {
+                    value,
+                    quote_style: _,
+                    keyword: _,
+                } = word;
+
+                value_from_replacements(value, &replacements).unwrap_or(token)
+            } else if let Token::Placeholder(placeholder) = &token {
+                value_from_replacements(placeholder, &replacements).unwrap_or(token)
+            } else {
+                token
+            }
+        })
+        .collect::<Vec<Token>>();
+
+    Ok(Parser::new(dialect.as_ref())
+        .with_tokens(replaced_tokens)
+        .parse_statements()
+        .map_err(|err| DataFusionError::External(Box::new(err)))?
+        .into_iter()
+        .map(|s| s.to_string())
+        .collect::<Vec<_>>()
+        .join(" "))
+}