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 f07703a6 feat(datafusion): configure SQLContext runtime resources
(#675)
f07703a6 is described below
commit f07703a6f3d7edbe3410f4efd3038e23f775b42f
Author: shyjsarah <[email protected]>
AuthorDate: Wed Aug 5 11:44:49 2026 +0800
feat(datafusion): configure SQLContext runtime resources (#675)
---
.../python/python/pypaimon_rust/datafusion.pyi | 12 ++-
bindings/python/src/context.rs | 81 ++++++++++++++++++-
bindings/python/tests/test_datafusion.py | 42 ++++++++++
crates/integrations/datafusion/src/lib.rs | 2 +-
crates/integrations/datafusion/src/sql_context.rs | 93 +++++++++++++++++++---
docs/src/python-binding.md | 24 ++++++
6 files changed, 239 insertions(+), 15 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 07586e63..a43e57f2 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -15,7 +15,8 @@
# specific language governing permissions and limitations
# under the License.
-from typing import Any, Callable, Dict, List, Optional, Sequence, TypeAlias,
Union
+from os import PathLike
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence,
TypeAlias, Union
import pyarrow
@@ -185,7 +186,14 @@ def udf(
...
class SQLContext:
- def __init__(self) -> None: ...
+ def __init__(
+ self,
+ *,
+ memory_pool_type: Optional[Literal["fair", "greedy"]] = None,
+ memory_pool_bytes: Optional[int] = None,
+ temp_directory: Optional[Union[str, PathLike[str]]] = None,
+ max_temp_directory_size_bytes: Optional[int] = None,
+ ) -> None: ...
def register_catalog(
self, catalog_name: str, catalog_options: Dict[str, str]
) -> None: ...
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 2536c5d5..6ab04ad7 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -16,6 +16,7 @@
// under the License.
use std::collections::HashMap;
+use std::path::PathBuf;
use std::sync::Arc;
use arrow::compute::cast;
@@ -23,6 +24,8 @@ 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::memory_pool::{FairSpillPool, GreedyMemoryPool,
MemoryPool};
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::logical_expr::{Signature, TypeSignature, Volatility};
use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
@@ -181,6 +184,62 @@ pub struct PySQLContext {
}
impl PySQLContext {
+ fn build_sql_context(
+ memory_pool_type: Option<String>,
+ memory_pool_bytes: Option<usize>,
+ temp_directory: Option<PathBuf>,
+ max_temp_directory_size_bytes: Option<u64>,
+ ) -> PyResult<SQLContext> {
+ if memory_pool_type.is_some() && memory_pool_bytes.is_none() {
+ return Err(PyValueError::new_err(
+ "memory_pool_type requires memory_pool_bytes",
+ ));
+ }
+ if memory_pool_bytes == Some(0) {
+ return Err(PyValueError::new_err(
+ "memory_pool_bytes must be greater than zero",
+ ));
+ }
+ if temp_directory
+ .as_ref()
+ .is_some_and(|path| path.as_os_str().is_empty())
+ {
+ return Err(PyValueError::new_err("temp_directory must not be
empty"));
+ }
+
+ let has_runtime_config = memory_pool_bytes.is_some()
+ || temp_directory.is_some()
+ || max_temp_directory_size_bytes.is_some();
+ if !has_runtime_config {
+ return Ok(SQLContext::new());
+ }
+
+ let mut runtime_builder = RuntimeEnvBuilder::new();
+ if let Some(memory_pool_bytes) = memory_pool_bytes {
+ let memory_pool: Arc<dyn MemoryPool> =
+ match memory_pool_type.as_deref().unwrap_or("greedy") {
+ "fair" => Arc::new(FairSpillPool::new(memory_pool_bytes)),
+ "greedy" =>
Arc::new(GreedyMemoryPool::new(memory_pool_bytes)),
+ other => {
+ return Err(PyValueError::new_err(format!(
+ "unsupported memory_pool_type '{other}'; expected
'fair' or 'greedy'"
+ )));
+ }
+ };
+ runtime_builder = runtime_builder.with_memory_pool(memory_pool);
+ }
+ if let Some(temp_directory) = temp_directory {
+ runtime_builder =
runtime_builder.with_temp_file_path(temp_directory);
+ }
+ if let Some(max_temp_directory_size_bytes) =
max_temp_directory_size_bytes {
+ runtime_builder =
+
runtime_builder.with_max_temp_directory_size(max_temp_directory_size_bytes);
+ }
+
+ let runtime_env = runtime_builder.build_arc().map_err(df_to_py_err)?;
+ Ok(SQLContext::builder().with_runtime_env(runtime_env).build())
+ }
+
fn vector_float32_type() -> ArrowDataType {
ArrowDataType::List(Arc::new(ArrowField::new(
"item",
@@ -325,9 +384,27 @@ impl PySQLContext {
#[pymethods]
impl PySQLContext {
#[new]
- fn new(py: Python<'_>) -> PyResult<Self> {
+ /// Creates a Paimon SQL context.
+ ///
+ /// `memory_pool_type` accepts `"greedy"` or `"fair"` and requires
+ /// `memory_pool_bytes`. When only `memory_pool_bytes` is provided, the
+ /// greedy pool is used. Temporary files can be directed to
+ /// `temp_directory` and bounded with `max_temp_directory_size_bytes`.
+ #[pyo3(signature = (*, memory_pool_type=None, memory_pool_bytes=None,
temp_directory=None, max_temp_directory_size_bytes=None))]
+ fn new(
+ py: Python<'_>,
+ memory_pool_type: Option<String>,
+ memory_pool_bytes: Option<usize>,
+ temp_directory: Option<PathBuf>,
+ max_temp_directory_size_bytes: Option<u64>,
+ ) -> PyResult<Self> {
let ctx = Self {
- inner: SQLContext::new(),
+ inner: Self::build_sql_context(
+ memory_pool_type,
+ memory_pool_bytes,
+ temp_directory,
+ max_temp_directory_size_bytes,
+ )?,
};
if let Err(err) = ctx.register_multimodal_builtins(py) {
Self::warn_multimodal_builtin_registration_failure(py, err);
diff --git a/bindings/python/tests/test_datafusion.py
b/bindings/python/tests/test_datafusion.py
index 916e1cc3..cfaa1a0e 100644
--- a/bindings/python/tests/test_datafusion.py
+++ b/bindings/python/tests/test_datafusion.py
@@ -18,6 +18,7 @@
import io
import json
import os
+import re
import struct
import sys
import tempfile
@@ -83,6 +84,47 @@ def extract_rows(batches):
return sorted(zip(table["id"].to_pylist(), table["name"].to_pylist()))
+def test_sql_context_accepts_runtime_resource_configuration():
+ with tempfile.TemporaryDirectory() as temp_directory:
+ ctx = SQLContext(
+ memory_pool_type="fair",
+ memory_pool_bytes=16 * 1024 * 1024,
+ temp_directory=temp_directory,
+ max_temp_directory_size_bytes=256 * 1024 * 1024,
+ )
+
+ batches = ctx.sql(
+ """
+ EXPLAIN ANALYZE
+ SELECT value
+ FROM generate_series(1, 1000000) AS t(value)
+ ORDER BY value DESC
+ """
+ )
+ plan = pa.Table.from_batches(batches)["plan"].to_pylist()[0]
+ spill_count = re.search(r"spill_count=(\d+)", plan)
+
+ assert spill_count is not None
+ assert int(spill_count.group(1)) > 0
+ assert any(Path(temp_directory).iterdir())
+
+
[email protected](
+ ("kwargs", "message"),
+ [
+ ({"memory_pool_type": "fair"}, "requires memory_pool_bytes"),
+ ({"memory_pool_bytes": 0}, "must be greater than zero"),
+ (
+ {"memory_pool_type": "unknown", "memory_pool_bytes": 1024},
+ "expected 'fair' or 'greedy'",
+ ),
+ ],
+)
+def test_sql_context_rejects_invalid_runtime_resource_configuration(kwargs,
message):
+ with pytest.raises(ValueError, match=message):
+ SQLContext(**kwargs)
+
+
def test_video_snapshot_builtin_registered_on_context_init():
ctx = SQLContext()
diff --git a/crates/integrations/datafusion/src/lib.rs
b/crates/integrations/datafusion/src/lib.rs
index 05e4037a..d9044fbc 100644
--- a/crates/integrations/datafusion/src/lib.rs
+++ b/crates/integrations/datafusion/src/lib.rs
@@ -81,7 +81,7 @@ pub use full_text_search::{register_full_text_search,
FullTextSearchFunction};
pub use hybrid_search::{register_hybrid_search, HybridSearchFunction};
pub use physical_plan::PaimonTableScan;
pub use relation_planner::PaimonRelationPlanner;
-pub use sql_context::SQLContext;
+pub use sql_context::{SQLContext, SQLContextBuilder};
pub use table::PaimonTableProvider;
pub use variant_functions::register_variant_functions;
pub use vector_search::{register_vector_search, VectorSearchFunction};
diff --git a/crates/integrations/datafusion/src/sql_context.rs
b/crates/integrations/datafusion/src/sql_context.rs
index cd54a08e..4c39c49a 100644
--- a/crates/integrations/datafusion/src/sql_context.rs
+++ b/crates/integrations/datafusion/src/sql_context.rs
@@ -55,6 +55,7 @@ use datafusion::common::tree_node::{TreeNode,
TreeNodeRecursion};
use datafusion::common::TableReference;
use datafusion::datasource::{MemTable, TableProvider};
use datafusion::error::{DataFusionError, Result as DFResult};
+use datafusion::execution::runtime_env::RuntimeEnv;
use datafusion::execution::SessionStateBuilder;
use datafusion::logical_expr::{Expr as LogicalExpr, LogicalPlan, Volatility};
use datafusion::prelude::{DataFrame, SessionContext};
@@ -101,16 +102,50 @@ pub struct SQLContext {
blob_reader_registry: BlobReaderRegistry,
}
-impl Default for SQLContext {
- fn default() -> Self {
- Self::new()
- }
+/// Builder for [`SQLContext`].
+///
+/// The builder preserves Paimon's session configuration while allowing callers
+/// to customize DataFusion runtime resources such as memory pools, temporary
+/// directories, and object store registries.
+///
+/// # Example
+///
+/// ```ignore
+/// use std::sync::Arc;
+///
+/// use datafusion::execution::memory_pool::FairSpillPool;
+/// use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+/// use paimon_datafusion::SQLContext;
+///
+/// let runtime_env = RuntimeEnvBuilder::new()
+/// .with_memory_pool(Arc::new(FairSpillPool::new(512 * 1024 * 1024)))
+/// .with_temp_file_path("/tmp/paimon-spill")
+/// .with_max_temp_directory_size(4 * 1024 * 1024 * 1024)
+/// .build_arc()?;
+/// let ctx = SQLContext::builder()
+/// .with_runtime_env(runtime_env)
+/// .build();
+/// ```
+#[derive(Default)]
+pub struct SQLContextBuilder {
+ runtime_env: Option<Arc<RuntimeEnv>>,
}
-impl SQLContext {
- /// Creates a new empty SQL context.
+impl SQLContextBuilder {
+ /// Creates a builder with DataFusion's default runtime environment.
pub fn new() -> Self {
- let state = SessionStateBuilder::new()
+ Self::default()
+ }
+
+ /// Uses the provided DataFusion runtime environment.
+ pub fn with_runtime_env(mut self, runtime_env: Arc<RuntimeEnv>) -> Self {
+ self.runtime_env = Some(runtime_env);
+ self
+ }
+
+ /// Builds a [`SQLContext`].
+ pub fn build(self) -> SQLContext {
+ let mut state_builder = SessionStateBuilder::new()
.with_config(crate::lateral_vector_search::session_config())
.with_default_features()
.with_relation_planners(vec![Arc::new(
@@ -119,18 +154,39 @@ impl SQLContext {
.with_optimizer_rules(crate::lateral_vector_search::optimizer_rules())
.with_query_planner(Arc::new(
crate::lateral_vector_search::PaimonQueryPlanner::new(),
- ))
- .build();
+ ));
+ if let Some(runtime_env) = self.runtime_env {
+ state_builder = state_builder.with_runtime_env(runtime_env);
+ }
+ let state = state_builder.build();
let ctx = SessionContext::new_with_state(state);
crate::blob_descriptor_functions::register_blob_descriptor_functions(&ctx);
crate::variant_functions::register_variant_functions(&ctx);
- Self {
+ SQLContext {
ctx,
catalogs: HashMap::new(),
dynamic_options: Default::default(),
blob_reader_registry: BlobReaderRegistry::default(),
}
}
+}
+
+impl Default for SQLContext {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SQLContext {
+ /// Creates a new empty SQL context.
+ pub fn new() -> Self {
+ Self::builder().build()
+ }
+
+ /// Creates a builder for customizing the underlying DataFusion session.
+ pub fn builder() -> SQLContextBuilder {
+ SQLContextBuilder::new()
+ }
pub fn blob_reader_registry(&self) -> BlobReaderRegistry {
self.blob_reader_registry.clone()
@@ -3601,6 +3657,8 @@ mod tests {
use async_trait::async_trait;
use datafusion::arrow::array::StringViewArray;
+ use datafusion::execution::memory_pool::FairSpillPool;
+ use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::sql::sqlparser::dialect::PostgreSqlDialect;
use paimon::catalog::Database;
use paimon::spec::{
@@ -3608,6 +3666,21 @@ mod tests {
};
use paimon::table::Table;
+ #[test]
+ fn test_sql_context_builder_uses_custom_runtime_env() {
+ let runtime_env = RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::new(FairSpillPool::new(1024)))
+ .build_arc()
+ .unwrap();
+
+ let ctx = SQLContext::builder()
+ .with_runtime_env(Arc::clone(&runtime_env))
+ .build();
+
+ assert!(Arc::ptr_eq(&runtime_env, &ctx.ctx().runtime_env()));
+ assert_eq!("fair", ctx.ctx().runtime_env().memory_pool.name());
+ }
+
// ==================== Mock Catalog ====================
#[allow(clippy::enum_variant_names)]
diff --git a/docs/src/python-binding.md b/docs/src/python-binding.md
index 57dc657f..bc4aa3ee 100644
--- a/docs/src/python-binding.md
+++ b/docs/src/python-binding.md
@@ -94,6 +94,30 @@ for batch in batches:
print(batch)
```
+### Runtime Resource Configuration
+
+`SQLContext` can use a bounded DataFusion memory pool and a dedicated temporary
+directory. This allows spill-capable operators, such as sorts and aggregations,
+to move intermediate data to disk when execution memory is constrained.
+
+```python
+ctx = SQLContext(
+ memory_pool_type="fair",
+ memory_pool_bytes=512 * 1024 * 1024,
+ temp_directory="/tmp/paimon-spill",
+ max_temp_directory_size_bytes=4 * 1024 * 1024 * 1024,
+)
+```
+
+`memory_pool_type` accepts `"fair"` or `"greedy"` and requires
+`memory_pool_bytes`. If only `memory_pool_bytes` is provided, the greedy pool
is
+used. All arguments are optional, and `SQLContext()` continues to use
+DataFusion's default runtime environment.
+
+Memory limits apply to allocations tracked by DataFusion's memory pool. They do
+not account for every allocation made by the host application or by external
+libraries.
+
## Reading a Table
Paimon Python uses a **scan-then-read** pattern: first scan the table to
produce splits, then read data from those splits as PyArrow RecordBatches.
