This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 13cc3199 python: expose SQL execution options (#1084)
13cc3199 is described below
commit 13cc3199ad420972b6a5090c9110f8c0b55eae92
Author: shyjsarah <[email protected]>
AuthorDate: Fri Oct 9 15:59:59 2026 +0800
python: expose SQL execution options (#1084)
---
.../python/python/pypaimon_rust/datafusion.pyi | 9 +++
bindings/python/src/context.rs | 80 +++++++++++++++++++---
bindings/python/tests/test_datafusion.py | 36 +++++++++-
3 files changed, 114 insertions(+), 11 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 0b701ded..5cce8692 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -502,6 +502,12 @@ def udf(
"""
...
+class SQLOptions:
+ def __init__(self) -> None: ...
+ def with_allow_ddl(self, allow: bool = True) -> "SQLOptions": ...
+ def with_allow_dml(self, allow: bool = True) -> "SQLOptions": ...
+ def with_allow_statements(self, allow: bool = True) -> "SQLOptions": ...
+
class SQLContext:
def __init__(
self,
@@ -525,3 +531,6 @@ class SQLContext:
def register_batch(self, name: str, batch: pyarrow.RecordBatch) -> None:
...
def register_udf(self, udf: PythonScalarUDF) -> None: ...
def sql(self, sql: str) -> List[pyarrow.RecordBatch]: ...
+ def sql_with_options(
+ self, sql: str, options: SQLOptions
+ ) -> List[pyarrow.RecordBatch]: ...
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 2bc67988..b6a10ff8 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -24,6 +24,7 @@ use arrow::datatypes::{DataType as ArrowDataType, Field as
ArrowField};
use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use arrow::record_batch::{RecordBatch, RecordBatchOptions};
use datafusion::catalog::CatalogProvider;
+use datafusion::execution::context::SQLOptions;
use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool,
MemoryPool};
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::logical_expr::{Signature, TypeSignature, Volatility};
@@ -83,6 +84,20 @@ fn pyarrow_compatible_batch(batch: &RecordBatch) ->
arrow::error::Result<RecordB
RecordBatch::try_new_with_options(schema, columns, &options)
}
+fn batches_to_pyarrow(py: Python<'_>, batches: &[RecordBatch]) ->
PyResult<Vec<Py<PyAny>>> {
+ batches
+ .iter()
+ .map(|batch| {
+ let batch = pyarrow_compatible_batch(batch).map_err(|err| {
+ PyValueError::new_err(format!(
+ "Failed to convert query result for PyArrow compatibility:
{err}"
+ ))
+ })?;
+ Ok(batch.to_pyarrow(py)?.unbind())
+ })
+ .collect()
+}
+
const OSS_IMPL: &str = "fs.oss.impl";
const JINDO_LIBRARY_PATH: &str = "fs.jindo.library.path";
const JINDOSDK_LIBRARY_PATH: &str = "JINDOSDK_LIBRARY_PATH";
@@ -339,6 +354,39 @@ impl PaimonCatalog {
}
}
+#[pyclass(name = "SQLOptions", module = "pypaimon_rust.datafusion")]
+pub struct PySQLOptions {
+ inner: SQLOptions,
+}
+
+#[pymethods]
+impl PySQLOptions {
+ #[new]
+ fn new() -> Self {
+ Self {
+ inner: SQLOptions::new(),
+ }
+ }
+
+ #[pyo3(signature = (allow = true))]
+ fn with_allow_ddl(mut slf: PyRefMut<'_, Self>, allow: bool) ->
PyRefMut<'_, Self> {
+ slf.inner = slf.inner.with_allow_ddl(allow);
+ slf
+ }
+
+ #[pyo3(signature = (allow = true))]
+ fn with_allow_dml(mut slf: PyRefMut<'_, Self>, allow: bool) ->
PyRefMut<'_, Self> {
+ slf.inner = slf.inner.with_allow_dml(allow);
+ slf
+ }
+
+ #[pyo3(signature = (allow = true))]
+ fn with_allow_statements(mut slf: PyRefMut<'_, Self>, allow: bool) ->
PyRefMut<'_, Self> {
+ slf.inner = slf.inner.with_allow_statements(allow);
+ slf
+ }
+}
+
/// A SQL context that supports registering multiple Paimon catalogs and
executing SQL.
#[pyclass(name = "SQLContext")]
pub struct PySQLContext {
@@ -649,17 +697,28 @@ impl PySQLContext {
df.collect().await.map_err(df_to_py_err)
})
})?;
- batches
- .iter()
- .map(|batch| {
- let batch = pyarrow_compatible_batch(batch).map_err(|err| {
- PyValueError::new_err(format!(
- "Failed to convert query result for PyArrow
compatibility: {err}"
- ))
- })?;
- Ok(batch.to_pyarrow(py)?.unbind())
+ batches_to_pyarrow(py, &batches)
+ }
+
+ fn sql_with_options(
+ &self,
+ py: Python<'_>,
+ sql: String,
+ options: PyRef<'_, PySQLOptions>,
+ ) -> PyResult<Vec<Py<PyAny>>> {
+ let rt = runtime();
+ let options = options.inner;
+ let batches = py.detach(|| {
+ rt.block_on(async {
+ let df = self
+ .inner
+ .sql_with_options(&sql, options)
+ .await
+ .map_err(df_to_py_err)?;
+ df.collect().await.map_err(df_to_py_err)
})
- .collect()
+ })?;
+ batches_to_pyarrow(py, &batches)
}
}
@@ -677,6 +736,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
this.add_class::<crate::schema::PyTableSchema>()?;
this.add_class::<crate::schema::PyDataField>()?;
this.add_class::<PyPythonScalarUDFObject>()?;
+ this.add_class::<PySQLOptions>()?;
this.add_class::<PySQLContext>()?;
this.add_class::<crate::write::PyBatchWriteBuilder>()?;
this.add_class::<crate::write::PyPostponeFixedBucketWriteBuilder>()?;
diff --git a/bindings/python/tests/test_datafusion.py
b/bindings/python/tests/test_datafusion.py
index 94e63e0b..f9d09372 100644
--- a/bindings/python/tests/test_datafusion.py
+++ b/bindings/python/tests/test_datafusion.py
@@ -30,7 +30,13 @@ import pyarrow as pa
import pytest
from datafusion import SessionContext
-from pypaimon_rust.datafusion import PaimonCatalog, PythonScalarUDF,
SQLContext, udf
+from pypaimon_rust.datafusion import (
+ PaimonCatalog,
+ PythonScalarUDF,
+ SQLContext,
+ SQLOptions,
+ udf,
+)
WAREHOUSE = os.environ.get("PAIMON_TEST_WAREHOUSE", "/tmp/paimon-warehouse")
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
@@ -126,6 +132,34 @@ def
test_sql_context_rejects_invalid_runtime_resource_configuration(kwargs, mess
SQLContext(**kwargs)
+def test_sql_with_options_disables_statements():
+ ctx = SQLContext()
+ options = SQLOptions().with_allow_statements(False)
+
+ with pytest.raises(ValueError, match="Statement not supported"):
+ ctx.sql_with_options("SET datafusion.execution.batch_size = 1024",
options)
+
+
+def test_sql_with_options_disables_ddl():
+ ctx = SQLContext()
+ options = SQLOptions().with_allow_ddl(False)
+
+ with pytest.raises(ValueError, match="DDL not supported"):
+ ctx.sql_with_options("CREATE VIEW blocked_view AS SELECT 1", options)
+
+
+def test_sql_with_options_disables_dml(tmp_path):
+ ctx = SQLContext()
+ ctx.register_catalog("paimon", {"warehouse": str(tmp_path / "warehouse")})
+ ctx.sql("CREATE TABLE paimon.default.test_table (id INT)")
+ options = SQLOptions().with_allow_dml(False)
+
+ with pytest.raises(ValueError, match="DML not supported"):
+ ctx.sql_with_options(
+ "INSERT INTO paimon.default.test_table VALUES (1)", options
+ )
+
+
def test_video_snapshot_builtin_registered_on_context_init():
ctx = SQLContext()
