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 3bc53b4d feat(python): support streaming SQL results (#1103)
3bc53b4d is described below
commit 3bc53b4d5f2fcd54d4f4bcbc2dec2aea1c4357a6
Author: Jiajia Li <[email protected]>
AuthorDate: Mon Oct 12 13:49:38 2026 +0800
feat(python): support streaming SQL results (#1103)
---
bindings/python/README.md | 26 ++++
.../python/python/pypaimon_rust/datafusion.pyi | 12 ++
bindings/python/src/context.rs | 106 ++++++++++++---
bindings/python/src/read.rs | 81 +++++++----
bindings/python/tests/test_datafusion.py | 15 ++-
bindings/python/tests/test_sql_stream.py | 149 +++++++++++++++++++++
6 files changed, 341 insertions(+), 48 deletions(-)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index 76ea23d8..9a55cf8b 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -56,6 +56,32 @@ ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp")
For the full SQL reference, see the [SQL Integration
docs](https://paimon.apache.org/docs/master/sql/).
+### Streaming SQL results
+
+`sql_stream()` executes one SQL query and returns a `RecordBatchReader`,
yielding
+PyArrow batches without collecting the entire result in memory. Close the
reader
+when stopping early to release the execution stream:
+
+```python
+from contextlib import closing
+
+with closing(ctx.sql_stream("SELECT id, name FROM paimon.my_db.users")) as
reader:
+ for batch in reader:
+ print(batch.to_pydict())
+```
+
+`read_next_batch()` returns `None` when the reader is exhausted or closed.
+`reader.schema` is known before the first batch, so
+`pa.Table.from_batches(reader, schema=reader.schema)` also handles an empty
result.
+Statements that write, such as `INSERT`, finish before `sql_stream()` returns.
+Errors during execution can surface while reading a batch and end the stream.
+Operators such as sorting and aggregation may still buffer data during
execution.
+
+`sql_stream_with_options(sql, options)` applies `SQLOptions` as
`sql_with_options()`
+does: they restrict statements delegated to DataFusion, but not statements
Paimon
+handles itself, such as `DELETE`, `UPDATE`, `MERGE INTO` or `DROP TABLE` on a
+Paimon table.
+
### Native Vector Search
A table with an existing vector column can search through the Rust core without
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index b14bbe5b..abc262e2 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -94,6 +94,10 @@ class StreamTableScan:
class RecordBatchReader:
+ @property
+ def schema(self) -> Optional[pyarrow.Schema]:
+ """Result schema of a SQL stream; None for a table read."""
+ ...
def __iter__(self) -> "RecordBatchReader": ...
def __next__(self) -> pyarrow.RecordBatch: ...
def read_next_batch(self) -> Optional[pyarrow.RecordBatch]: ...
@@ -589,6 +593,14 @@ class SQLContext:
def sql_with_options(
self, sql: str, options: SQLOptions
) -> List[pyarrow.RecordBatch]: ...
+ def sql_stream(self, sql: str) -> RecordBatchReader:
+ """Stream PyArrow batches; close the reader when stopping early."""
+ ...
+ def sql_stream_with_options(
+ self, sql: str, options: SQLOptions
+ ) -> RecordBatchReader:
+ """Stream results with SQLOptions restrictions."""
+ ...
class VectorSearchBuilder:
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 117dc05e..1b9c5397 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -20,16 +20,19 @@ use std::path::PathBuf;
use std::sync::Arc;
use arrow::compute::cast;
-use arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField};
+use arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField,
SchemaRef};
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};
+use datafusion::logical_expr::{LogicalPlan, Signature, TypeSignature,
Volatility};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::SendableRecordBatchStream;
use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use futures::StreamExt;
use paimon::catalog::Identifier;
use paimon::{Catalog, CatalogFactory, Options};
use paimon_datafusion::{PaimonCatalogProvider, SQLContext};
@@ -39,14 +42,15 @@ use pyo3::types::PyCapsule;
use crate::blob::PyBlobReaderRegistry;
use crate::error::{df_to_py_err, to_py_err};
+use crate::read::PyRecordBatchReader;
use crate::table::PyTable;
use crate::udf::{build_python_scalar_udf, udf, PyPythonScalarUDFObject};
use paimon_datafusion::runtime::runtime;
-fn pyarrow_compatible_batch(batch: &RecordBatch) ->
arrow::error::Result<RecordBatch> {
+/// Returns `None` when the schema is already PyArrow compatible.
+fn pyarrow_compatible_schema(schema: &arrow::datatypes::Schema) ->
Option<SchemaRef> {
let mut changed = false;
- let fields = batch
- .schema()
+ let fields = schema
.fields()
.iter()
.map(|field| {
@@ -60,13 +64,18 @@ fn pyarrow_compatible_batch(batch: &RecordBatch) ->
arrow::error::Result<RecordB
.collect::<Vec<_>>();
if !changed {
- return Ok(batch.clone());
+ return None;
}
-
- let schema = Arc::new(arrow::datatypes::Schema::new_with_metadata(
+ Some(Arc::new(arrow::datatypes::Schema::new_with_metadata(
fields,
- batch.schema().metadata().clone(),
- ));
+ schema.metadata().clone(),
+ )))
+}
+
+fn pyarrow_compatible_batch(batch: &RecordBatch) ->
arrow::error::Result<RecordBatch> {
+ let Some(schema) = pyarrow_compatible_schema(batch.schema_ref()) else {
+ return Ok(batch.clone());
+ };
let columns = batch
.columns()
.iter()
@@ -84,17 +93,28 @@ fn pyarrow_compatible_batch(batch: &RecordBatch) ->
arrow::error::Result<RecordB
RecordBatch::try_new_with_options(schema, columns, &options)
}
+/// DataFusion runs these plans only while their result stream is polled.
+/// `EXPLAIN ANALYZE` executes its input, so a wrapped write counts too.
+fn is_write_plan(plan: &LogicalPlan) -> bool {
+ match plan {
+ LogicalPlan::Dml(_) | LogicalPlan::Copy(_) => true,
+ LogicalPlan::Analyze(analyze) => is_write_plan(&analyze.input),
+ _ => false,
+ }
+}
+
+fn query_result_batch(batch: &RecordBatch) -> PyResult<RecordBatch> {
+ pyarrow_compatible_batch(batch).map_err(|err| {
+ PyValueError::new_err(format!(
+ "Failed to convert query result for PyArrow compatibility: {err}"
+ ))
+ })
+}
+
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())
- })
+ .map(|batch| Ok(query_result_batch(batch)?.to_pyarrow(py)?.unbind()))
.collect()
}
@@ -394,6 +414,41 @@ pub struct PySQLContext {
}
impl PySQLContext {
+ fn execute_sql_stream(
+ &self,
+ py: Python<'_>,
+ sql: &str,
+ options: SQLOptions,
+ ) -> PyResult<PyRecordBatchReader> {
+ let stream = py.detach(|| {
+ runtime().block_on(async {
+ let df = self
+ .inner
+ .sql_with_options(sql, options)
+ .await
+ .map_err(df_to_py_err)?;
+ if !is_write_plan(df.logical_plan()) {
+ return df.execute_stream().await.map_err(df_to_py_err);
+ }
+ // Finish the write before returning, so it does not depend on
+ // the caller reading or closing the result.
+ let schema = Arc::clone(df.schema().inner());
+ let batches = df.collect().await.map_err(df_to_py_err)?;
+ let stream =
futures::stream::iter(batches.into_iter().map(Ok));
+ let stream: SendableRecordBatchStream =
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream));
+ Ok(stream)
+ })
+ })?;
+ let schema = stream.schema();
+ let schema = pyarrow_compatible_schema(&schema).unwrap_or(schema);
+ let stream = stream.map(|batch|
query_result_batch(&batch.map_err(df_to_py_err)?));
+ Ok(PyRecordBatchReader::from_py_stream(
+ Some(schema),
+ Box::pin(stream),
+ ))
+ }
+
fn build_sql_context(
memory_pool_type: Option<String>,
memory_pool_bytes: Option<usize>,
@@ -720,6 +775,21 @@ impl PySQLContext {
})?;
batches_to_pyarrow(py, &batches)
}
+
+ /// Stream SQL results. Close the reader when stopping early.
+ fn sql_stream(&self, py: Python<'_>, sql: String) ->
PyResult<PyRecordBatchReader> {
+ self.execute_sql_stream(py, &sql, SQLOptions::new())
+ }
+
+ /// Stream SQL results using the supplied options.
+ fn sql_stream_with_options(
+ &self,
+ py: Python<'_>,
+ sql: String,
+ options: PyRef<'_, PySQLOptions>,
+ ) -> PyResult<PyRecordBatchReader> {
+ self.execute_sql_stream(py, &sql, options.inner)
+ }
}
pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) ->
PyResult<()> {
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
index e101e93b..0cec6518 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -19,8 +19,10 @@ use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
+use arrow::datatypes::SchemaRef;
use arrow::pyarrow::ToPyArrow;
use arrow::record_batch::RecordBatch;
+use futures::stream::BoxStream;
use futures::TryStreamExt;
use paimon::spec::{is_variant_extraction_row, DataField, DataType, Predicate,
RowType};
use paimon::table::{
@@ -746,11 +748,9 @@ impl PyTableRead {
py: Python<'_>,
splits: &Bound<'_, PyAny>,
) -> PyResult<PyRecordBatchReader> {
- Ok(PyRecordBatchReader {
- stream: Mutex::new(Some(self.read_stream(py, splits)?)),
- closed: AtomicBool::new(false),
- close_notify: Notify::new(),
- })
+ Ok(PyRecordBatchReader::from_stream(
+ self.read_stream(py, splits)?,
+ ))
}
/// Read the given splits into a list of PyArrow RecordBatches.
@@ -770,14 +770,23 @@ impl PyTableRead {
#[pyclass(name = "RecordBatchReader", module = "pypaimon_rust.datafusion")]
pub struct PyRecordBatchReader {
- stream: Mutex<Option<ArrowRecordBatchStream>>,
+ schema: Option<SchemaRef>,
+ stream: Mutex<Option<BoxStream<'static, PyResult<RecordBatch>>>>,
closed: AtomicBool,
close_notify: Notify,
}
impl PyRecordBatchReader {
pub(crate) fn from_stream(stream: ArrowRecordBatchStream) -> Self {
+ Self::from_py_stream(None, Box::pin(stream.map_err(to_py_err)))
+ }
+
+ pub(crate) fn from_py_stream(
+ schema: Option<SchemaRef>,
+ stream: BoxStream<'static, PyResult<RecordBatch>>,
+ ) -> Self {
Self {
+ schema,
stream: Mutex::new(Some(stream)),
closed: AtomicBool::new(false),
close_notify: Notify::new(),
@@ -794,18 +803,16 @@ impl PyRecordBatchReader {
return Ok(None);
}
let result = match stream.as_mut() {
- Some(stream) => runtime()
- .block_on(async {
- tokio::select! {
- biased;
- _ = self.close_notify.notified() => Ok(None),
- batch = stream.try_next() => batch,
- }
- })
- .map_err(to_py_err),
+ Some(stream) => runtime().block_on(async {
+ tokio::select! {
+ biased;
+ _ = self.close_notify.notified() => Ok(None),
+ batch = stream.try_next() => batch,
+ }
+ }),
None => Ok(None),
};
- if self.closed.load(Ordering::Acquire) {
+ if self.closed.load(Ordering::Acquire) || !matches!(result,
Ok(Some(_))) {
stream.take();
}
result
@@ -820,9 +827,7 @@ impl PyRecordBatchReader {
fn close_reader(&self) {
self.closed.store(true, Ordering::Release);
- // notify_one stores a permit when next_record_batch has not started
- // polling yet, avoiding a lost wake-up between its closed check and
- // the select.
+ // notify_one retains a permit if the read has not started polling.
self.close_notify.notify_one();
if let Ok(mut stream) = self.stream.try_lock() {
stream.take();
@@ -832,6 +837,15 @@ impl PyRecordBatchReader {
#[pymethods]
impl PyRecordBatchReader {
+ /// Result schema, known before the first batch for SQL streams.
+ #[getter]
+ fn schema(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
+ self.schema
+ .as_ref()
+ .map(|schema| Ok(schema.to_pyarrow(py)?.unbind()))
+ .transpose()
+ }
+
fn read_next_batch(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
self.next_batch(py)
}
@@ -967,6 +981,29 @@ mod tests {
use super::*;
+ #[test]
+ fn record_batch_reader_releases_stream_on_completion_or_error() {
+ for fail in [false, true] {
+ let resource = Arc::new(());
+ let weak = Arc::downgrade(&resource);
+ let stream: ArrowRecordBatchStream =
Box::pin(futures::stream::poll_fn(move |_cx| {
+ let _resource = &resource;
+ Poll::Ready(fail.then(|| {
+ Err(paimon::Error::Unsupported {
+ message: "read failed".to_string(),
+ })
+ }))
+ }));
+ let reader = PyRecordBatchReader::from_stream(stream);
+ assert!(weak.upgrade().is_some());
+
+ let result = reader.next_record_batch();
+ assert_eq!(result.is_err(), fail);
+ assert!(weak.upgrade().is_none());
+ assert!(matches!(reader.next_record_batch(), Ok(None)));
+ }
+ }
+
#[test]
fn record_batch_reader_close_interrupts_pending_next() {
let (polled_tx, polled_rx) = mpsc::channel();
@@ -978,11 +1015,7 @@ mod tests {
}
Poll::<Option<paimon::Result<RecordBatch>>>::Pending
}));
- let reader = Arc::new(PyRecordBatchReader {
- stream: Mutex::new(Some(stream)),
- closed: AtomicBool::new(false),
- close_notify: Notify::new(),
- });
+ let reader = Arc::new(PyRecordBatchReader::from_stream(stream));
let worker_reader = Arc::clone(&reader);
let (result_tx, result_rx) = mpsc::channel();
let worker = std::thread::spawn(move || {
diff --git a/bindings/python/tests/test_datafusion.py
b/bindings/python/tests/test_datafusion.py
index f9d09372..4fc8cfc8 100644
--- a/bindings/python/tests/test_datafusion.py
+++ b/bindings/python/tests/test_datafusion.py
@@ -132,30 +132,33 @@ def
test_sql_context_rejects_invalid_runtime_resource_configuration(kwargs, mess
SQLContext(**kwargs)
-def test_sql_with_options_disables_statements():
[email protected]("method", ["sql_with_options",
"sql_stream_with_options"])
+def test_sql_with_options_disables_statements(method):
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)
+ getattr(ctx, method)("SET datafusion.execution.batch_size = 1024",
options)
-def test_sql_with_options_disables_ddl():
[email protected]("method", ["sql_with_options",
"sql_stream_with_options"])
+def test_sql_with_options_disables_ddl(method):
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)
+ getattr(ctx, method)("CREATE VIEW blocked_view AS SELECT 1", options)
-def test_sql_with_options_disables_dml(tmp_path):
[email protected]("method", ["sql_with_options",
"sql_stream_with_options"])
+def test_sql_with_options_disables_dml(tmp_path, method):
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(
+ getattr(ctx, method)(
"INSERT INTO paimon.default.test_table VALUES (1)", options
)
diff --git a/bindings/python/tests/test_sql_stream.py
b/bindings/python/tests/test_sql_stream.py
new file mode 100644
index 00000000..cfa92d95
--- /dev/null
+++ b/bindings/python/tests/test_sql_stream.py
@@ -0,0 +1,149 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+import pytest
+
+from pypaimon_rust.datafusion import RecordBatchReader, SQLContext, SQLOptions
+
+
[email protected]
+def ctx(tmp_path):
+ ctx = SQLContext()
+ ctx.register_catalog("paimon", {"warehouse": str(tmp_path / "warehouse")})
+ ctx.sql("SET datafusion.execution.batch_size = 2")
+ ctx.sql("SET datafusion.execution.target_partitions = 1")
+ return ctx
+
+
[email protected](params=["sql_stream", "sql_stream_with_options"])
+def stream_sql(ctx, request):
+ if request.param == "sql_stream_with_options":
+ return lambda sql: ctx.sql_stream_with_options(sql, SQLOptions())
+ return ctx.sql_stream
+
+
+def test_sql_stream_returns_multiple_batches(ctx, stream_sql):
+ sql = "SELECT value FROM generate_series(1, 9) AS t(value)"
+ reader = stream_sql(sql)
+
+ assert isinstance(reader, RecordBatchReader)
+ assert iter(reader) is reader
+ first = reader.read_next_batch()
+ assert isinstance(first, pa.RecordBatch)
+ assert first.num_rows == 2
+ batches = [first, *reader]
+ assert len(batches) > 1
+ assert
pa.Table.from_batches(batches).equals(pa.Table.from_batches(ctx.sql(sql)))
+ assert reader.read_next_batch() is None
+ assert list(reader) == []
+ reader.close()
+
+
+def test_sql_stream_reads_paimon_table(ctx, stream_sql):
+ ctx.sql("CREATE TABLE paimon.default.users (id INT, name STRING)")
+ ctx.sql(
+ "INSERT INTO paimon.default.users VALUES "
+ "(1, 'alice'), (2, NULL), (3, 'charlie')"
+ )
+ sql = "SELECT id, name FROM paimon.default.users WHERE id >= 2 ORDER BY id"
+
+ result = pa.Table.from_batches(stream_sql(sql))
+
+ assert result.equals(pa.Table.from_batches(ctx.sql(sql)))
+ assert result.to_pydict() == {"id": [2, 3], "name": [None, "charlie"]}
+
+
+def test_sql_stream_preserves_pyarrow_string_compatibility(ctx, stream_sql):
+ ctx.register_batch(
+ "labels",
+ pa.record_batch(
+ [pa.array(["alice", None, "charlie"], type=pa.string_view())],
+ names=["name"],
+ ),
+ )
+ sql = "SELECT name FROM labels"
+
+ result = pa.Table.from_batches(stream_sql(sql))
+
+ assert result.schema.field("name").type == pa.string()
+ assert result.equals(pa.Table.from_batches(ctx.sql(sql)))
+
+
+def test_sql_stream_empty_result(stream_sql):
+ reader = stream_sql("SELECT 1 AS value WHERE FALSE")
+
+ assert sum(batch.num_rows for batch in reader) == 0
+ assert reader.read_next_batch() is None
+ assert list(reader) == []
+ reader.close()
+
+
+def test_sql_stream_schema_describes_an_empty_result(ctx, stream_sql):
+ ctx.sql("CREATE TABLE paimon.default.users (id INT, name STRING)")
+ reader = stream_sql("SELECT id, name FROM paimon.default.users WHERE id <
0")
+
+ result = pa.Table.from_batches(reader, schema=reader.schema)
+
+ assert result.num_rows == 0
+ assert result.schema.names == ["id", "name"]
+ assert result.schema.types == [pa.int32(), pa.string()]
+
+
[email protected]("prefix", ["", "EXPLAIN ANALYZE "])
+def test_sql_stream_finishes_insert_before_returning(ctx, stream_sql, prefix):
+ ctx.sql("CREATE TABLE paimon.default.events (id INT)")
+
+ sql = f"{prefix}INSERT INTO paimon.default.events VALUES (1), (2)"
+ stream_sql(sql).close()
+
+ result = pa.Table.from_batches(
+ ctx.sql("SELECT id FROM paimon.default.events ORDER BY id")
+ )
+ assert result.to_pydict() == {"id": [1, 2]}
+
+
+def test_sql_stream_reports_planning_errors(stream_sql):
+ with pytest.raises(ValueError, match="missing_table"):
+ stream_sql("SELECT * FROM missing_table")
+
+
+def test_sql_stream_yields_batches_before_execution_error(stream_sql):
+ reader = stream_sql(
+ "SELECT 10 / (3 - value) AS value FROM generate_series(1, 4) AS
t(value)"
+ )
+
+ assert next(reader).to_pydict() == {"value": [5, 10]}
+ with pytest.raises(ValueError, match="(?i)(divide|division) by zero"):
+ reader.read_next_batch()
+ assert reader.read_next_batch() is None
+ reader.close()
+
+
[email protected]("read_first", [False, True])
+def test_sql_stream_close_stops_iteration(stream_sql, read_first):
+ reader = stream_sql(
+ "SELECT value FROM generate_series(1, 1000000) AS t(value)"
+ )
+ if read_first:
+ assert next(reader).num_rows == 2
+
+ reader.close()
+
+ assert reader.read_next_batch() is None
+ assert list(reader) == []
+ reader.close()
