Switching to explicit param_values or named parameters that will perform string replacement via parsed tokens

Author: timsaucer

python/datafusion/context.py

@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+import uuid
 import warnings
 from typing import TYPE_CHECKING, Any, Protocol
 
@@ -594,10 +595,18 @@ def register_listing_table(
         )
 
     def sql(
-        self, query: str, options: SQLOptions | None = None, **named_params: Any
+        self,
+        query: str,
+        options: SQLOptions | None = None,
+        param_values: dict[str, Any] | None = None,
+        **named_params: Any,
     ) -> DataFrame:
         """Create a :py:class:`~datafusion.DataFrame` from SQL query text.
 
+        See the online documentation for a description of how to perform
+        parameterized substitution via either the param_values option
+        or passing in named parameters.
+
         Note: This API implements DDL statements such as ``CREATE TABLE`` and
         ``CREATE VIEW`` and DML statements such as ``INSERT INTO`` with in-memory
         default implementation.See
@@ -606,45 +615,56 @@ def sql(
         Args:
             query: SQL query text.
             options: If provided, the query will be validated against these options.
-            named_params: Provides substitution in the query string.
+            param_values: Provides substitution of scalar values in the query
+                after parsing.
+            named_params: Provides string or DataFrame substitution in the query string.
 
         Returns:
             DataFrame representation of the SQL query.
         """
 
-        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)
-            ]
-
-        def dataframe_params(**p: Any) -> list[tuple[str, DataFrame]]:
-            if p is None:
-                return []
-
-            return [
-                (name, value.df)
-                for (name, value) in p.items()
-                if isinstance(value, DataFrame)
-            ]
+        def value_to_scalar(value) -> pa.Scalar:
+            if isinstance(value, pa.Scalar):
+                return value
+            return pa.scalar(value)
+
+        def value_to_string(value) -> str:
+            if isinstance(value, DataFrame):
+                view_name = str(uuid.uuid4()).replace("-", "_")
+                view_name = f"view_{view_name}"
+                view = value.df.into_view(temporary=True)
+                self.ctx.register_table(view_name, view)
+                return view_name
+            return str(value)
+
+        param_values = (
+            {name: value_to_scalar(value) for (name, value) in param_values}
+            if param_values is not None
+            else {}
+        )
+        param_strings = (
+            {name: value_to_string(value) for (name, value) in named_params.items()}
+            if named_params is not None
+            else {}
+        )
 
         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),
+                param_values=param_values,
+                param_strings=param_strings,
             )
         )
 
     def sql_with_options(
-        self, query: str, options: SQLOptions, **named_params: Any
+        self,
+        query: str,
+        options: SQLOptions,
+        param_values: dict[str, Any] | None = None,
+        **named_params: Any,
     ) -> DataFrame:
         """Create a :py:class:`~datafusion.dataframe.DataFrame` from SQL query text.
 
@@ -654,12 +674,16 @@ def sql_with_options(
         Args:
             query: SQL query text.
             options: SQL options.
-            named_params: Provides substitution in the query string.
+            param_values: Provides substitution of scalar values in the query
+                after parsing.
+            named_params: Provides string or DataFrame substitution in the query string.
 
         Returns:
             DataFrame representation of the SQL query.
         """
-        return self.sql(query, options, **named_params)
+        return self.sql(
+            query, options=options, param_values=param_values, **named_params
+        )
 
     def create_dataframe(
         self,

src/context.rs

@@ -64,9 +64,9 @@ 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::sql::util::replace_placeholders_with_strings;
 use crate::store::StorageContexts;
-use crate::table::{PyTable, TempViewTable};
+use crate::table::PyTable;
 use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
@@ -424,54 +424,41 @@ impl PySessionContext {
         self.ctx.register_udtf(&name, func);
     }
 
-    #[pyo3(signature = (query, options=None, scalar_params=vec![], dataframe_params=vec![]))]
+    #[pyo3(signature = (query, options=None, param_values=HashMap::default(), param_strings=HashMap::default()))]
     pub fn sql_with_options(
         &self,
         py: Python,
-        query: &str,
+        mut query: String,
         options: Option<PySQLOptions>,
-        scalar_params: Vec<(String, PyScalarValue)>,
-        dataframe_params: Vec<(String, PyDataFrame)>,
+        param_values: HashMap<String, PyScalarValue>,
+        param_strings: HashMap<String, String>,
     ) -> PyDataFusionResult<PyDataFrame> {
         let options = if let Some(options) = options {
             options.options
         } else {
             SQLOptions::new()
         };
 
-        let scalar_params = scalar_params
+        let param_values = param_values
             .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>>()?;
+            .collect::<HashMap<_, _>>();
 
         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)
+        if !param_strings.is_empty() {
+            query = replace_placeholders_with_strings(&query, dialect, param_strings)?;
+        }
+
+        let mut df = wait_for_future(py, async {
+            self.ctx.sql_with_options(&query, options).await
         })??;
 
+        if !param_values.is_empty() {
+            df = df.with_param_values(param_values)?;
+        }
+
         Ok(PyDataFrame::new(df))
     }
 
@@ -516,29 +503,6 @@ impl PySessionContext {
         PyDataFrame::new(DataFrame::new(self.ctx.state(), plan.plan.as_ref().clone()))
     }
 
-    pub fn create_temporary_view(
-        &self,
-        py: Python,
-        name: &str,
-        df: PyDataFrame,
-        replace_if_exists: bool,
-    ) -> PyDataFusionResult<()> {
-        if self.table(name, py).is_ok() {
-            if replace_if_exists {
-                let _ = self.deregister_table(name);
-            } else {
-                exec_err!(
-                    "Unable to create temporary view. Table with name {name} already exists."
-                )?;
-            }
-        }
-
-        let table = Arc::new(TempViewTable::new(df.inner_df()));
-        self.ctx.register_table(name, table)?;
-
-        Ok(())
-    }
-
     /// Construct datafusion dataframe from Python list
     #[pyo3(signature = (data, name=None))]
     pub fn from_pylist(

src/sql/util.rs

@@ -1,93 +1,69 @@
-use datafusion::common::{internal_err, plan_datafusion_err, DataFusionError};
+use datafusion::common::{exec_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 datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer};
 use std::collections::HashMap;
 
-fn value_from_replacements(
+fn tokens_from_replacements(
     placeholder: &str,
-    replacements: &HashMap<String, String>,
-) -> Option<Token> {
+    replacements: &HashMap<String, Vec<Token>>,
+) -> Option<Vec<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,
-            })
-        })
+        replacements.get(pattern).cloned()
     } 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
+fn get_tokens_for_string_replacement(
+    dialect: &dyn Dialect,
+    replacements: HashMap<String, String>,
+) -> Result<HashMap<String, Vec<Token>>, DataFusionError> {
+    replacements
+        .into_iter()
+        .map(|(name, value)| {
+            let tokens = Tokenizer::new(dialect, &value)
+                .tokenize()
+                .map_err(|err| DataFusionError::External(err.into()))?;
+            Ok((name, tokens))
+        })
+        .collect()
 }
 
-pub(crate) fn replace_placeholders_with_table_names(
+pub(crate) fn replace_placeholders_with_strings(
     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."
-        )
-    })?;
+    let dialect = dialect_from_str(dialect)
+        .ok_or_else(|| plan_datafusion_err!("Unsupported SQL dialect: {dialect}."))?;
 
-    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 replacements = get_tokens_for_string_replacement(dialect.as_ref(), replacements)?;
+
+    let tokens = Tokenizer::new(dialect.as_ref(), query)
+        .tokenize()
+        .map_err(|err| DataFusionError::External(err.into()))?;
 
     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)
+        .flat_map(|token| {
+            if let Token::Placeholder(placeholder) = &token {
+                tokens_from_replacements(placeholder, &replacements).unwrap_or(vec![token])
             } else {
-                token
+                vec![token]
             }
         })
         .collect::<Vec<Token>>();
 
-    Ok(Parser::new(dialect.as_ref())
+    let statement = 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(" "))
+        .map_err(|err| DataFusionError::External(Box::new(err)))?;
+
+    if statement.len() != 1 {
+        return exec_err!("placeholder replacement should return exactly one statement");
+    }
+
+    Ok(statement[0].to_string())
 }