Intermediate work on parameterizing queries

Author: timsaucer

python/datafusion/context.py

@@ -27,6 +27,8 @@
 except ImportError:
     from typing_extensions import deprecated  # Python 3.12
 
+import uuid
+
 import pyarrow as pa
 
 from datafusion.catalog import Catalog
@@ -592,7 +594,9 @@ def register_listing_table(
             self._convert_file_sort_order(file_sort_order),
         )
 
-    def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame:
+    def sql(
+        self, query: str, options: SQLOptions | None = None, **named_params: Any
+    ) -> DataFrame:
         """Create a :py:class:`~datafusion.DataFrame` from SQL query text.
 
         Note: This API implements DDL statements such as ``CREATE TABLE`` and
@@ -603,10 +607,25 @@ def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame:
         Args:
             query: SQL query text.
             options: If provided, the query will be validated against these options.
+            named_params: Provides substitution in the query string.
 
         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)
+
         if options is None:
             return DataFrame(self.ctx.sql(query))
         return DataFrame(self.ctx.sql_with_options(query, options.options_internal))

python/tests/test_sql.py

@@ -550,3 +550,13 @@ def test_register_listing_table(
 
     rd = result.to_pydict()
     assert dict(zip(rd["grp"], rd["count"], strict=False)) == {"a": 3, "b": 2}
+
+
+def test_parameterized_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
+    ).collect()
+    result = pa.Table.from_batches(result)
+    assert result.to_pydict() == {"cnt": [100]}

src/context.rs

@@ -64,7 +64,7 @@ use crate::record_batch::PyRecordBatchStream;
 use crate::sql::exceptions::py_value_err;
 use crate::sql::logical::PyLogicalPlan;
 use crate::store::StorageContexts;
-use crate::table::PyTable;
+use crate::table::{PyTable, TempViewTable};
 use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
@@ -487,6 +487,29 @@ 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/table.rs

@@ -24,7 +24,11 @@ use async_trait::async_trait;
 use datafusion::catalog::Session;
 use datafusion::common::Column;
 use datafusion::datasource::{TableProvider, TableType};
+<<<<<<< HEAD
 use datafusion::logical_expr::{Expr, LogicalPlanBuilder, TableProviderFilterPushDown};
+=======
+use datafusion::logical_expr::{Expr, LogicalPlanBuilder};
+>>>>>>> 3c9a96c (Intermediate work on temp views and parameterizing queries)
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::DataFrame;
 use pyo3::prelude::*;