Visitar URL original
Feat/parameterized sql queries (#964) · voidstackloop/datafusion-python@c141dd3 · GitHub
Skip to content

Commit c141dd3

Browse files
timsaucertimsaucer
authored
Feat/parameterized sql queries (apache#964)
* Intermediate work on parameterizing queries * Reworking to do token parsing of sql query instead of string manipulation * Switching to explicit param_values or named parameters that will perform string replacement via parsed tokens * Add additional unit tests for parameterized queries * merge conflict * license text * Add documentation * cargo clippy and fmt * Need at least pyarrow 16 now * add type hints * Minor docstring update
1 parent 276dc6a commit c141dd3

8 files changed

Lines changed: 319 additions & 28 deletions

File tree

‎docs/source/user-guide/configuration.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
.. specific language governing permissions and limitations
1616
.. under the License.
1717
18+
.. _configuration:
19+
1820
Configuration
1921
=============
2022

‎docs/source/user-guide/sql.rst‎

Lines changed: 87 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,17 +23,100 @@ DataFusion also offers a SQL API, read the full reference `here <https://arrow.a
2323
.. ipython:: python
2424
2525
import datafusion
26-
from datafusion import col
27-
import pyarrow
26+
from datafusion import DataFrame, SessionContext
2827
2928
# create a context
3029
ctx = datafusion.SessionContext()
3130
3231
# register a CSV
33-
ctx.register_csv('pokemon', 'pokemon.csv')
32+
ctx.register_csv("pokemon", "pokemon.csv")
3433
3534
# create a new statement via SQL
3635
df = ctx.sql('SELECT "Attack"+"Defense", "Attack"-"Defense" FROM pokemon')
3736
3837
# collect and convert to pandas DataFrame
39-
df.to_pandas()
38+
df.to_pandas()
39+
40+
Parameterized queries
41+
---------------------
42+
43+
In DataFusion-Python 51.0.0 we introduced the ability to pass parameters
44+
in a SQL query. These are similar in concept to
45+
`prepared statements <https://datafusion.apache.org/user-guide/sql/prepared_statements.html>`_,
46+
but allow passing named parameters into a SQL query. Consider this simple
47+
example.
48+
49+
.. ipython:: python
50+
51+
def show_attacks(ctx: SessionContext, threshold: int) -> None:
52+
ctx.sql(
53+
'SELECT "Name", "Attack" FROM pokemon WHERE "Attack" > $val', val=threshold
54+
).show(num=5)
55+
show_attacks(ctx, 75)
56+
57+
When passing parameters like the example above we convert the Python objects
58+
into their string representation. We also have special case handling
59+
for :py:class:`~datafusion.dataframe.DataFrame` objects, since they cannot simply
60+
be turned into string representations for an SQL query. In these cases we
61+
will register a temporary view in the :py:class:`~datafusion.context.SessionContext`
62+
using a generated table name.
63+
64+
The formatting for passing string replacement objects is to precede the
65+
variable name with a single ``$``. This works for all dialects in
66+
the SQL parser except ``hive`` and ``mysql``. Since these dialects do not
67+
support named placeholders, we are unable to do this type of replacement.
68+
We recommend either switching to another dialect or using Python
69+
f-string style replacement.
70+
71+
.. warning::
72+
73+
To support DataFrame parameterized queries, your session must support
74+
registration of temporary views. The default
75+
:py:class:`~datafusion.catalog.CatalogProvider` and
76+
:py:class:`~datafusion.catalog.SchemaProvider` do have this capability.
77+
If you have implemented custom providers, it is important that temporary
78+
views do not persist across :py:class:`~datafusion.context.SessionContext`
79+
or you may get unintended consequences.
80+
81+
The following example shows passing in both a :py:class:`~datafusion.dataframe.DataFrame`
82+
object as well as a Python object to be used in parameterized replacement.
83+
84+
.. ipython:: python
85+
86+
def show_column(
87+
ctx: SessionContext, column: str, df: DataFrame, threshold: int
88+
) -> None:
89+
ctx.sql(
90+
'SELECT "Name", $col FROM $df WHERE $col > $val',
91+
col=column,
92+
df=df,
93+
val=threshold,
94+
).show(num=5)
95+
df = ctx.table("pokemon")
96+
show_column(ctx, '"Defense"', df, 75)
97+
98+
The approach implemented for conversion of variables into a SQL query
99+
relies on string conversion. This has the potential for data loss,
100+
specifically for cases like floating point numbers. If you need to pass
101+
variables into a parameterized query and it is important to maintain the
102+
original value without conversion to a string, then you can use the
103+
optional parameter ``param_values`` to specify these. This parameter
104+
expects a dictionary mapping from the parameter name to a Python
105+
object. Those objects will be cast into a
106+
`PyArrow Scalar Value <https://arrow.apache.org/docs/python/generated/pyarrow.Scalar.html>`_.
107+
108+
Using ``param_values`` will rely on the SQL dialect you have configured
109+
for your session. This can be set using the :ref:`configuration options <configuration>`
110+
of your :py:class:`~datafusion.context.SessionContext`. Similar to how
111+
`prepared statements <https://datafusion.apache.org/user-guide/sql/prepared_statements.html>`_
112+
work, these parameters are limited to places where you would pass in a
113+
scalar value, such as a comparison.
114+
115+
.. ipython:: python
116+
117+
def param_attacks(ctx: SessionContext, threshold: int) -> None:
118+
ctx.sql(
119+
'SELECT "Name", "Attack" FROM pokemon WHERE "Attack" > $val',
120+
param_values={"val": threshold},
121+
).show(num=5)
122+
param_attacks(ctx, 75)

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ classifiers = [
4444
"Programming Language :: Rust",
4545
]
4646
dependencies = [
47-
"pyarrow>=11.0.0;python_version<'3.14'",
47+
"pyarrow>=16.0.0;python_version<'3.14'",
4848
"pyarrow>=22.0.0;python_version>='3.14'",
4949
"typing-extensions;python_version<'3.13'"
5050
]

‎python/datafusion/context.py‎

Lines changed: 65 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
from __future__ import annotations
2121

22+
import uuid
2223
import warnings
2324
from typing import TYPE_CHECKING, Any, Protocol
2425

@@ -27,6 +28,7 @@
2728
except ImportError:
2829
from typing_extensions import deprecated # Python 3.12
2930

31+
3032
import pyarrow as pa
3133

3234
from datafusion.catalog import Catalog
@@ -592,9 +594,19 @@ def register_listing_table(
592594
self._convert_file_sort_order(file_sort_order),
593595
)
594596

595-
def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame:
597+
def sql(
598+
self,
599+
query: str,
600+
options: SQLOptions | None = None,
601+
param_values: dict[str, Any] | None = None,
602+
**named_params: Any,
603+
) -> DataFrame:
596604
"""Create a :py:class:`~datafusion.DataFrame` from SQL query text.
597605
606+
See the online documentation for a description of how to perform
607+
parameterized substitution via either the ``param_values`` option
608+
or passing in ``named_params``.
609+
598610
Note: This API implements DDL statements such as ``CREATE TABLE`` and
599611
``CREATE VIEW`` and DML statements such as ``INSERT INTO`` with in-memory
600612
default implementation.See
@@ -603,15 +615,57 @@ def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame:
603615
Args:
604616
query: SQL query text.
605617
options: If provided, the query will be validated against these options.
618+
param_values: Provides substitution of scalar values in the query
619+
after parsing.
620+
named_params: Provides string or DataFrame substitution in the query string.
606621
607622
Returns:
608623
DataFrame representation of the SQL query.
609624
"""
610-
if options is None:
611-
return DataFrame(self.ctx.sql(query))
612-
return DataFrame(self.ctx.sql_with_options(query, options.options_internal))
613625

614-
def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame:
626+
def value_to_scalar(value: Any) -> pa.Scalar:
627+
if isinstance(value, pa.Scalar):
628+
return value
629+
return pa.scalar(value)
630+
631+
def value_to_string(value: Any) -> str:
632+
if isinstance(value, DataFrame):
633+
view_name = str(uuid.uuid4()).replace("-", "_")
634+
view_name = f"view_{view_name}"
635+
view = value.df.into_view(temporary=True)
636+
self.ctx.register_table(view_name, view)
637+
return view_name
638+
return str(value)
639+
640+
param_values = (
641+
{name: value_to_scalar(value) for (name, value) in param_values.items()}
642+
if param_values is not None
643+
else {}
644+
)
645+
param_strings = (
646+
{name: value_to_string(value) for (name, value) in named_params.items()}
647+
if named_params is not None
648+
else {}
649+
)
650+
651+
options_raw = options.options_internal if options is not None else None
652+
653+
return DataFrame(
654+
self.ctx.sql_with_options(
655+
query,
656+
options=options_raw,
657+
param_values=param_values,
658+
param_strings=param_strings,
659+
)
660+
)
661+
662+
def sql_with_options(
663+
self,
664+
query: str,
665+
options: SQLOptions,
666+
param_values: dict[str, Any] | None = None,
667+
**named_params: Any,
668+
) -> DataFrame:
615669
"""Create a :py:class:`~datafusion.dataframe.DataFrame` from SQL query text.
616670
617671
This function will first validate that the query is allowed by the
@@ -620,11 +674,16 @@ def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame:
620674
Args:
621675
query: SQL query text.
622676
options: SQL options.
677+
param_values: Provides substitution of scalar values in the query
678+
after parsing.
679+
named_params: Provides string or DataFrame substitution in the query string.
623680
624681
Returns:
625682
DataFrame representation of the SQL query.
626683
"""
627-
return self.sql(query, options)
684+
return self.sql(
685+
query, options=options, param_values=param_values, **named_params
686+
)
628687

629688
def create_dataframe(
630689
self,

‎python/tests/test_sql.py‎

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
import pyarrow as pa
2222
import pyarrow.dataset as ds
2323
import pytest
24-
from datafusion import col, udf
24+
from datafusion import SessionContext, col, udf
2525
from datafusion.object_store import Http
2626
from pyarrow.csv import write_csv
2727

@@ -550,3 +550,46 @@ def test_register_listing_table(
550550

551551
rd = result.to_pydict()
552552
assert dict(zip(rd["grp"], rd["count"], strict=False)) == {"a": 3, "b": 2}
553+
554+
555+
def test_parameterized_named_params(ctx, tmp_path) -> None:
556+
path = helpers.write_parquet(tmp_path / "a.parquet", helpers.data())
557+
558+
df = ctx.read_parquet(path)
559+
result = ctx.sql(
560+
"SELECT COUNT(a) AS cnt, $lit_val as lit_val FROM $replaced_df",
561+
lit_val=3,
562+
replaced_df=df,
563+
).collect()
564+
result = pa.Table.from_batches(result)
565+
assert result.to_pydict() == {"cnt": [100], "lit_val": [3]}
566+
567+
568+
def test_parameterized_param_values(ctx: SessionContext) -> None:
569+
# Test the parameters that should be handled by the parser rather
570+
# than our manipulation of the query string by searching for tokens
571+
batch = pa.RecordBatch.from_arrays(
572+
[pa.array([1, 2, 3, 4])],
573+
names=["a"],
574+
)
575+
576+
ctx.register_record_batches("t", [[batch]])
577+
result = ctx.sql("SELECT a FROM t WHERE a < $val", param_values={"val": 3})
578+
assert result.to_pydict() == {"a": [1, 2]}
579+
580+
581+
def test_parameterized_mixed_query(ctx: SessionContext) -> None:
582+
batch = pa.RecordBatch.from_arrays(
583+
[pa.array([1, 2, 3, 4])],
584+
names=["a"],
585+
)
586+
ctx.register_record_batches("t", [[batch]])
587+
registered_df = ctx.table("t")
588+
589+
result = ctx.sql(
590+
"SELECT $col_name FROM $df WHERE a < $val",
591+
param_values={"val": 3},
592+
df=registered_df,
593+
col_name="a",
594+
)
595+
assert result.to_pydict() == {"a": [1, 2]}

‎src/context.rs‎

Lines changed: 32 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -55,14 +55,16 @@ use url::Url;
5555
use uuid::Uuid;
5656

5757
use crate::catalog::{PyCatalog, RustWrappedPyCatalogProvider};
58+
use crate::common::data_type::PyScalarValue;
5859
use crate::dataframe::PyDataFrame;
5960
use crate::dataset::Dataset;
60-
use crate::errors::{py_datafusion_err, PyDataFusionResult};
61+
use crate::errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult};
6162
use crate::expr::sort_expr::PySortExpr;
6263
use crate::physical_plan::PyExecutionPlan;
6364
use crate::record_batch::PyRecordBatchStream;
6465
use crate::sql::exceptions::py_value_err;
6566
use crate::sql::logical::PyLogicalPlan;
67+
use crate::sql::util::replace_placeholders_with_strings;
6668
use crate::store::StorageContexts;
6769
use crate::table::PyTable;
6870
use crate::udaf::PyAggregateUDF;
@@ -422,27 +424,41 @@ impl PySessionContext {
422424
self.ctx.register_udtf(&name, func);
423425
}
424426

425-
/// Returns a PyDataFrame whose plan corresponds to the SQL statement.
426-
pub fn sql(&self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
427-
let result = self.ctx.sql(query);
428-
let df = wait_for_future(py, result)??;
429-
Ok(PyDataFrame::new(df))
430-
}
431-
432-
#[pyo3(signature = (query, options=None))]
427+
#[pyo3(signature = (query, options=None, param_values=HashMap::default(), param_strings=HashMap::default()))]
433428
pub fn sql_with_options(
434429
&self,
435-
query: &str,
436-
options: Option<PySQLOptions>,
437430
py: Python,
431+
mut query: String,
432+
options: Option<PySQLOptions>,
433+
param_values: HashMap<String, PyScalarValue>,
434+
param_strings: HashMap<String, String>,
438435
) -> PyDataFusionResult<PyDataFrame> {
439436
let options = if let Some(options) = options {
440437
options.options
441438
} else {
442439
SQLOptions::new()
443440
};
444-
let result = self.ctx.sql_with_options(query, options);
445-
let df = wait_for_future(py, result)??;
441+
442+
let param_values = param_values
443+
.into_iter()
444+
.map(|(name, value)| (name, ScalarValue::from(value)))
445+
.collect::<HashMap<_, _>>();
446+
447+
let state = self.ctx.state();
448+
let dialect = state.config().options().sql_parser.dialect.as_str();
449+
450+
if !param_strings.is_empty() {
451+
query = replace_placeholders_with_strings(&query, dialect, param_strings)?;
452+
}
453+
454+
let mut df = wait_for_future(py, async {
455+
self.ctx.sql_with_options(&query, options).await
456+
})??;
457+
458+
if !param_values.is_empty() {
459+
df = df.with_param_values(param_values)?;
460+
}
461+
446462
Ok(PyDataFrame::new(df))
447463
}
448464

@@ -550,7 +566,7 @@ impl PySessionContext {
550566

551567
(array.schema().as_ref().to_owned(), vec![array])
552568
} else {
553-
return Err(crate::errors::PyDataFusionError::Common(
569+
return Err(PyDataFusionError::Common(
554570
"Expected either a Arrow Array or Arrow Stream in from_arrow().".to_string(),
555571
));
556572
};
@@ -714,7 +730,7 @@ impl PySessionContext {
714730
) -> PyDataFusionResult<()> {
715731
let delimiter = delimiter.as_bytes();
716732
if delimiter.len() != 1 {
717-
return Err(crate::errors::PyDataFusionError::PythonError(py_value_err(
733+
return Err(PyDataFusionError::PythonError(py_value_err(
718734
"Delimiter must be a single character",
719735
)));
720736
}
@@ -968,7 +984,7 @@ impl PySessionContext {
968984
) -> PyDataFusionResult<PyDataFrame> {
969985
let delimiter = delimiter.as_bytes();
970986
if delimiter.len() != 1 {
971-
return Err(crate::errors::PyDataFusionError::PythonError(py_value_err(
987+
return Err(PyDataFusionError::PythonError(py_value_err(
972988
"Delimiter must be a single character",
973989
)));
974990
};

0 commit comments

Comments
 (0)