You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by pa...@apache.org on 2026/06/29 14:07:37 UTC
(sedona-db) branch main updated: chore(python/sedonadb): Upgrade PyO3 and enable free-threaded Python support (#1012)
This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new dfad028bf chore(python/sedonadb): Upgrade PyO3 and enable free-threaded Python support (#1012)
dfad028bf is described below
commit dfad028bf0cb18679c841393ea034290db0bf435
Author: ajaypadwal73 <51...@users.noreply.github.com>
AuthorDate: Mon Jun 29 19:37:31 2026 +0530
chore(python/sedonadb): Upgrade PyO3 and enable free-threaded Python support (#1012)
---
Cargo.lock | 21 +++++----
python/sedonadb-zarr/Cargo.toml | 2 +-
python/sedonadb-zarr/src/lib.rs | 7 ++-
python/sedonadb/Cargo.toml | 2 +-
python/sedonadb/src/context.rs | 78 +++++++++++++-------------------
python/sedonadb/src/dataframe.rs | 83 +++++++++++++++++++++--------------
python/sedonadb/src/datasource.rs | 22 +++++-----
python/sedonadb/src/lib.rs | 2 +-
python/sedonadb/src/raster_loader.rs | 22 +++++-----
python/sedonadb/src/runtime.rs | 6 +--
python/sedonadb/src/schema.rs | 6 +--
python/sedonadb/src/udf.rs | 46 +++++++++----------
python/sedonadb/tests/test_context.py | 18 ++++++++
rust/sedona/src/context.rs | 77 +++++++++++++++++++++++++++-----
rust/sedona/src/show.rs | 6 +--
15 files changed, 236 insertions(+), 162 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 95b434345..54fdd1e10 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4854,9 +4854,9 @@ dependencies = [
[[package]]
name = "pyo3"
-version = "0.25.1"
+version = "0.26.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8970a78afe0628a3e3430376fc5fd76b6b45c4d43360ffd6cdd40bdde72b682a"
+checksum = "7ba0117f4212101ee6544044dae45abe1083d30ce7b29c4b5cbdfa2354e07383"
dependencies = [
"indoc",
"libc",
@@ -4871,19 +4871,18 @@ dependencies = [
[[package]]
name = "pyo3-build-config"
-version = "0.25.1"
+version = "0.26.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "458eb0c55e7ece017adeba38f2248ff3ac615e53660d7c71a238d7d2a01c7598"
+checksum = "4fc6ddaf24947d12a9aa31ac65431fb1b851b8f4365426e182901eabfb87df5f"
dependencies = [
- "once_cell",
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
-version = "0.25.1"
+version = "0.26.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7114fe5457c61b276ab77c5055f206295b812608083644a5c5b2640c3102565c"
+checksum = "025474d3928738efb38ac36d4744a74a400c901c7596199e20e45d98eb194105"
dependencies = [
"libc",
"pyo3-build-config",
@@ -4891,9 +4890,9 @@ dependencies = [
[[package]]
name = "pyo3-macros"
-version = "0.25.1"
+version = "0.26.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a8725c0a622b374d6cb051d11a0983786448f7785336139c3c94f5aa6bef7e50"
+checksum = "2e64eb489f22fe1c95911b77c44cc41e7c19f3082fc81cce90f657cdc42ffded"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
@@ -4903,9 +4902,9 @@ dependencies = [
[[package]]
name = "pyo3-macros-backend"
-version = "0.25.1"
+version = "0.26.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4109984c22491085343c05b0dbc54ddc405c3cf7b4374fc533f5c3313a572ccc"
+checksum = "100246c0ecf400b475341b8455a9213344569af29a3c841d29270e53102e0fcf"
dependencies = [
"heck",
"proc-macro2",
diff --git a/python/sedonadb-zarr/Cargo.toml b/python/sedonadb-zarr/Cargo.toml
index aa35b68cc..f12f715f4 100644
--- a/python/sedonadb-zarr/Cargo.toml
+++ b/python/sedonadb-zarr/Cargo.toml
@@ -31,7 +31,7 @@ doc = false
arrow-array = { workspace = true, features = ["ffi"] }
arrow-buffer = { workspace = true }
object_store = { workspace = true, features = ["aws", "http"] }
-pyo3 = { version = "0.25.1" }
+pyo3 = { version = "0.26.0" }
sedona-raster = { workspace = true }
sedona-raster-zarr = { workspace = true }
sedona-schema = { workspace = true }
diff --git a/python/sedonadb-zarr/src/lib.rs b/python/sedonadb-zarr/src/lib.rs
index ca633bcda..f20ab4da0 100644
--- a/python/sedonadb-zarr/src/lib.rs
+++ b/python/sedonadb-zarr/src/lib.rs
@@ -372,9 +372,8 @@ impl PyZarrRasterLoader {
rust_requests.iter().collect();
// Run the async load on the shared runtime
- let results = py.allow_threads(|| {
- shared_runtime().block_on(async { self.loader.load(&request_refs).await })
- });
+ let results = py
+ .detach(|| shared_runtime().block_on(async { self.loader.load(&request_refs).await }));
let results = results.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
@@ -418,7 +417,7 @@ struct OwnedLoadRequest {
data_type: sedona_schema::raster::BandDataType,
}
-#[pymodule]
+#[pymodule(gil_used = false)]
fn _lib(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyZarrChunkReader>()?;
m.add_class::<PyZarrRasterLoader>()?;
diff --git a/python/sedonadb/Cargo.toml b/python/sedonadb/Cargo.toml
index b8df136b9..63f6259e9 100644
--- a/python/sedonadb/Cargo.toml
+++ b/python/sedonadb/Cargo.toml
@@ -46,7 +46,7 @@ datafusion-execution = { workspace = true }
datafusion-expr = { workspace = true }
datafusion-ffi = { workspace = true }
futures = { workspace = true }
-pyo3 = { version = "0.25.1" }
+pyo3 = { version = "0.26.0" }
sedona = { workspace = true }
sedona-adbc = { workspace = true }
sedona-datasource = { workspace = true }
diff --git a/python/sedonadb/src/context.rs b/python/sedonadb/src/context.rs
index e4e75b998..d935cd372 100644
--- a/python/sedonadb/src/context.rs
+++ b/python/sedonadb/src/context.rs
@@ -17,7 +17,6 @@
use std::{collections::HashMap, sync::Arc};
use arrow_schema::DataType;
-use datafusion_expr::{AggregateUDFImpl, ScalarUDFImpl};
use pyo3::prelude::*;
use sedona::context::SedonaContext;
use sedona::context_builder::SedonaContextBuilder;
@@ -87,7 +86,7 @@ impl InternalContext {
&self,
py: Python<'py>,
table_paths: Vec<String>,
- options: HashMap<String, PyObject>,
+ options: HashMap<String, Py<PyAny>>,
geometry_columns: Option<String>,
validate: bool,
partitioning: Option<Vec<String>>,
@@ -186,23 +185,28 @@ impl InternalContext {
name: &str,
) -> Result<Bound<'py, PyAny>, PySedonaError> {
let name_lower = name.to_lowercase();
- if let Some(sedona_scalar_udf) = self.inner.functions.scalar_udf(&name_lower) {
+ let inner = if let Some(sedona_scalar_udf) = self.inner.scalar_udf(&name_lower)? {
+ Some(ScalarUdfLookup::Sedona(sedona_scalar_udf))
+ } else {
+ self.inner
+ .ctx
+ .state()
+ .scalar_functions()
+ .get(&name_lower)
+ .cloned()
+ .map(ScalarUdfLookup::DataFusion)
+ };
+
+ if let Some(ScalarUdfLookup::Sedona(sedona_scalar_udf)) = inner {
Ok(Bound::new(
py,
PySedonaScalarUdf {
- inner: sedona_scalar_udf.clone(),
- },
- )?
- .into_any())
- } else if let Some(scalar_udf) = self.inner.ctx.state().scalar_functions().get(&name_lower)
- {
- Ok(Bound::new(
- py,
- PyScalarUdf {
- inner: scalar_udf.clone(),
+ inner: sedona_scalar_udf,
},
)?
.into_any())
+ } else if let Some(ScalarUdfLookup::DataFusion(scalar_udf)) = inner {
+ Ok(Bound::new(py, PyScalarUdf { inner: scalar_udf })?.into_any())
} else {
Err(PySedonaError::SedonaPython(format!(
"Scalar UDF with name {name} was not found"
@@ -216,17 +220,19 @@ impl InternalContext {
name: &str,
) -> Result<Bound<'py, PyAny>, PySedonaError> {
let name_lower = name.to_lowercase();
- if let Some(aggregate_udf) = self
+ let aggregate_udf = self
.inner
.ctx
.state()
.aggregate_functions()
.get(&name_lower)
- {
+ .cloned();
+
+ if let Some(aggregate_udf) = aggregate_udf {
Ok(Bound::new(
py,
PyAggregateUdf {
- inner: aggregate_udf.clone(),
+ inner: aggregate_udf,
},
)?
.into_any())
@@ -238,14 +244,7 @@ impl InternalContext {
}
pub fn list_scalar_udfs(&self) -> Result<Vec<String>, PySedonaError> {
- Ok(self
- .inner
- .ctx
- .state()
- .scalar_functions()
- .keys()
- .cloned()
- .collect())
+ Ok(self.inner.scalar_udf_names()?)
}
pub fn list_aggregate_udfs(&self) -> Result<Vec<String>, PySedonaError> {
@@ -259,42 +258,22 @@ impl InternalContext {
.collect())
}
- pub fn register_component(&mut self, component: Bound<PyAny>) -> Result<(), PySedonaError> {
+ pub fn register_component(&self, component: Bound<PyAny>) -> Result<(), PySedonaError> {
if component.hasattr("__sedonadb_internal_udf__")? {
let py_scalar_udf = component
.getattr("__sedonadb_internal_udf__")?
.call0()?
.extract::<PySedonaScalarUdf>()?;
- let name = py_scalar_udf.inner.name();
self.inner
- .functions
- .insert_scalar_udf(py_scalar_udf.inner.clone());
- self.inner.ctx.register_udf(
- self.inner
- .functions
- .scalar_udf(name)
- .unwrap()
- .clone()
- .into(),
- );
+ .register_sedona_scalar_udf(py_scalar_udf.inner.clone())?;
return Ok(());
} else if component.hasattr("__sedonadb_internal_aggregate_udf__")? {
let py_agg_udf = component
.getattr("__sedonadb_internal_aggregate_udf__")?
.call0()?
.extract::<PySedonaAggregateUdf>()?;
- let name = py_agg_udf.inner.name();
self.inner
- .functions
- .insert_aggregate_udf(py_agg_udf.inner.clone());
- self.inner.ctx.register_udaf(
- self.inner
- .functions
- .aggregate_udf(name)
- .unwrap()
- .clone()
- .into(),
- );
+ .register_sedona_aggregate_udf(py_agg_udf.inner.clone())?;
return Ok(());
} else if component.hasattr("__sedonadb_external_format__")? {
let spec = component
@@ -322,3 +301,8 @@ impl InternalContext {
))
}
}
+
+enum ScalarUdfLookup {
+ Sedona(sedona_expr::scalar_udf::SedonaScalarUDF),
+ DataFusion(Arc<datafusion_expr::ScalarUDF>),
+}
diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs
index f1bb0f6d1..3e5b4fed9 100644
--- a/python/sedonadb/src/dataframe.rs
+++ b/python/sedonadb/src/dataframe.rs
@@ -98,7 +98,7 @@ impl InternalDataFrame {
Ok(names)
}
- fn qualified_column_expr(&self, py: Python<'_>, key: PyObject) -> Result<PyExpr, PyErr> {
+ fn qualified_column_expr(&self, py: Python<'_>, key: Py<PyAny>) -> Result<PyExpr, PyErr> {
let num_fields = self.inner.schema().fields().len();
let all_names = || self.inner.schema().field_names();
@@ -414,12 +414,13 @@ impl InternalDataFrame {
let schema = self.inner.schema();
let partitions =
wait_for_future(py, &self.runtime, self.inner.clone().collect_partitioned())??;
- let provider = MemTable::try_new(schema.as_arrow().clone().into(), partitions)?;
+ let provider = Arc::new(MemTable::try_new(
+ schema.as_arrow().clone().into(),
+ partitions,
+ )?);
+ let df = ctx.inner.ctx.read_table(provider)?;
- Ok(Self::new(
- ctx.inner.ctx.read_table(Arc::new(provider))?,
- self.runtime.clone(),
- ))
+ Ok(Self::new(df, self.runtime.clone()))
}
fn to_batches<'py>(
@@ -489,36 +490,54 @@ impl InternalDataFrame {
#[cfg(not(feature = "s2geography"))]
let has_geography = false;
- let sort_by_expr = sort_by
+ let mut sort_by_expr = Vec::with_capacity(sort_by.len());
+ let mut needs_sd_order = false;
+ for name in sort_by {
+ let column = Expr::Column(Column::new_unqualified(name.clone()));
+ if geometry_column_names.contains(name.as_str()) {
+ if has_geography {
+ needs_sd_order = true;
+ sort_by_expr.push((Some(name), column));
+ } else {
+ return Err(PySedonaError::SedonaPython(
+ "Use maturin develop --features 's2geography,pyo3/extension-module' for dev geography support"
+ .to_string(),
+ ));
+ }
+ } else {
+ sort_by_expr.push((None, column));
+ }
+ }
+
+ let order_udf = if needs_sd_order {
+ Some(
+ ctx.inner
+ .ctx
+ .state()
+ .scalar_functions()
+ .get("sd_order")
+ .cloned()
+ .ok_or_else(|| {
+ PySedonaError::SedonaPython(
+ "Can't order by geometry field when sd_order() is not available"
+ .to_string(),
+ )
+ })?,
+ )
+ } else {
+ None
+ };
+
+ let sort_by_expr = sort_by_expr
.into_iter()
- .map(|name| {
- let column = Expr::Column(Column::new_unqualified(name.clone()));
- if geometry_column_names.contains(name.as_str()) {
- // Create the call sd_order(column). If we're ordering by geometry but don't have
- // the required feature for high quality sort output, give an error. This is mostly
- // an issue when using maturin develop because geography is not a default feature.
- if has_geography {
- let state = ctx.inner.ctx.state();
- let order_udf_opt = state.scalar_functions().get("sd_order");
- if let Some(order_udf) = order_udf_opt {
- Ok(SortExpr::new(order_udf.call(vec![column]), true, false))
- } else {
- Err(PySedonaError::SedonaPython(
- "Can't order by geometry field when sd_order() is not available"
- .to_string(),
- ))
- }
- } else {
- Err(PySedonaError::SedonaPython(
- "Use maturin develop --features 's2geography,pyo3/extension-module' for dev geography support"
- .to_string(),
- ))
- }
+ .map(|(geometry_name, column)| {
+ if geometry_name.is_some() {
+ SortExpr::new(order_udf.as_ref().unwrap().call(vec![column]), true, false)
} else {
- Ok(SortExpr::new(column, true, false))
+ SortExpr::new(column, true, false)
}
})
- .collect::<Result<Vec<_>, _>>()?;
+ .collect::<Vec<_>>();
let write_options = SedonaWriteOptions::new()
.with_partition_by(partition_by)
diff --git a/python/sedonadb/src/datasource.rs b/python/sedonadb/src/datasource.rs
index f3c254f27..12137729d 100644
--- a/python/sedonadb/src/datasource.rs
+++ b/python/sedonadb/src/datasource.rs
@@ -23,7 +23,7 @@ use async_trait::async_trait;
use datafusion::{physical_expr::conjunction, physical_plan::PhysicalExpr};
use datafusion_common::{DataFusionError, Result};
use pyo3::{
- exceptions::PyNotImplementedError, pyclass, pymethods, types::PyCapsule, Bound, PyObject,
+ exceptions::PyNotImplementedError, pyclass, pymethods, types::PyCapsule, Bound, Py, PyAny,
Python,
};
use sedona_datasource::{
@@ -52,12 +52,12 @@ pub struct PyExternalFormat {
/// `False`); we snapshot it once to avoid GIL traffic in
/// `list_single_object()`, which is called on hot paths.
list_single_object: bool,
- py_spec: PyObject,
+ py_spec: Py<PyAny>,
}
impl Clone for PyExternalFormat {
fn clone(&self) -> Self {
- Python::with_gil(|py| Self {
+ Python::attach(|py| Self {
extension: self.extension.clone(),
list_single_object: self.list_single_object,
py_spec: self.py_spec.clone_ref(py),
@@ -149,7 +149,7 @@ impl PyExternalFormat {
#[pymethods]
impl PyExternalFormat {
#[new]
- fn new<'py>(py: Python<'py>, py_spec: PyObject) -> Result<Self, PySedonaError> {
+ fn new<'py>(py: Python<'py>, py_spec: Py<PyAny>) -> Result<Self, PySedonaError> {
let extension = py_spec.getattr(py, "extension")?.extract::<String>(py)?;
let list_single_object = read_list_single_object(py, &py_spec)?;
Ok(Self {
@@ -168,7 +168,7 @@ impl PyExternalFormat {
/// — that's the intended failure mode.
fn read_list_single_object<'py>(
py: Python<'py>,
- py_spec: &PyObject,
+ py_spec: &Py<PyAny>,
) -> Result<bool, PySedonaError> {
Ok(py_spec
.getattr(py, "list_single_object")?
@@ -189,13 +189,13 @@ impl ExternalFormatSpec for PyExternalFormat {
&self,
options: &HashMap<String, String>,
) -> Result<Arc<dyn ExternalFormatSpec>> {
- let new_external_format = Python::with_gil(|py| self.with_options_impl(py, options))
+ let new_external_format = Python::attach(|py| self.with_options_impl(py, options))
.map_err(|e| DataFusionError::External(Box::new(e)))?;
Ok(Arc::new(new_external_format))
}
async fn infer_schema(&self, location: &Object) -> Result<Schema> {
- let schema = Python::with_gil(|py| self.infer_schema_impl(py, location))
+ let schema = Python::attach(|py| self.infer_schema_impl(py, location))
.map_err(|e| DataFusionError::External(Box::new(e)))?;
Ok(schema)
@@ -205,7 +205,7 @@ impl ExternalFormatSpec for PyExternalFormat {
&self,
args: &OpenReaderArgs,
) -> Result<Box<dyn RecordBatchReader + Send>> {
- let reader = Python::with_gil(|py| self.open_reader_impl(py, args))
+ let reader = Python::attach(|py| self.open_reader_impl(py, args))
.map_err(|e| DataFusionError::External(Box::new(e)))?;
Ok(reader)
@@ -344,7 +344,7 @@ impl PyFilter {
/// it to be constructed from Python using either a set of indices or a set of names.
#[pyclass]
pub struct PyProjectedRecordBatchReader {
- inner_object: PyObject,
+ inner_object: Py<PyAny>,
projection_indices: Option<Vec<usize>>,
projection_names: Option<Vec<String>>,
}
@@ -353,7 +353,7 @@ pub struct PyProjectedRecordBatchReader {
impl PyProjectedRecordBatchReader {
#[new]
fn new(
- inner_object: PyObject,
+ inner_object: Py<PyAny>,
projection_indices: Option<Vec<usize>>,
projection_names: Option<Vec<String>>,
) -> Self {
@@ -398,7 +398,7 @@ impl PyProjectedRecordBatchReader {
/// an ADBC statement/cursor).
struct WrappedRecordBatchReader {
pub inner: Box<dyn RecordBatchReader + Send>,
- pub shelter: Option<PyObject>,
+ pub shelter: Option<Py<PyAny>>,
}
impl RecordBatchReader for WrappedRecordBatchReader {
diff --git a/python/sedonadb/src/lib.rs b/python/sedonadb/src/lib.rs
index 1b6022108..6c4abb2e0 100644
--- a/python/sedonadb/src/lib.rs
+++ b/python/sedonadb/src/lib.rs
@@ -118,7 +118,7 @@ fn gdal_version() -> Result<Option<String>, PySedonaError> {
}
}
-#[pymodule]
+#[pymodule(gil_used = false)]
fn _lib(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
#[cfg(feature = "mimalloc")]
configure_tg_allocator();
diff --git a/python/sedonadb/src/raster_loader.rs b/python/sedonadb/src/raster_loader.rs
index 6ac04de86..42e798dc0 100644
--- a/python/sedonadb/src/raster_loader.rs
+++ b/python/sedonadb/src/raster_loader.rs
@@ -25,7 +25,7 @@ use pyo3::{
buffer::PyBuffer,
pyclass, pyfunction, pymethods,
types::{PyAnyMethods, PyList, PyListMethods},
- PyObject, Python,
+ Py, PyAny, Python,
};
use sedona_raster::{
raster_loader::{AsyncRasterLoader, RasterLoadRequest, RasterLoadResult},
@@ -45,15 +45,15 @@ pub struct PyRasterLoader {
/// Cached loader name (called once at construction)
name: String,
/// Python callable that checks if a format is supported
- py_supports_format: PyObject,
+ py_supports_format: Py<PyAny>,
/// Python callable that loads raster data
- py_load: PyObject,
+ py_load: Py<PyAny>,
}
impl PyRasterLoader {
- pub fn new(py_name: PyObject, py_supports_format: PyObject, py_load: PyObject) -> Self {
+ pub fn new(py_name: Py<PyAny>, py_supports_format: Py<PyAny>, py_load: Py<PyAny>) -> Self {
// Call py_name() once and cache the result
- let name = Python::with_gil(|py| {
+ let name = Python::attach(|py| {
py_name
.call(py, (), None)
.and_then(|obj| obj.extract::<String>(py))
@@ -75,7 +75,7 @@ impl AsyncRasterLoader for PyRasterLoader {
}
fn supports_format(&self, format: Option<&str>) -> bool {
- Python::with_gil(|py| {
+ Python::attach(|py| {
let result = self
.py_supports_format
.call(py, (format,), None)
@@ -97,11 +97,11 @@ impl AsyncRasterLoader for PyRasterLoader {
})
.collect();
- let py_load = Python::with_gil(|py| self.py_load.clone_ref(py));
+ let py_load = Python::attach(|py| self.py_load.clone_ref(py));
// Spawn a blocking task to hold the GIL and call Python
let results = tokio::task::spawn_blocking(move || {
- Python::with_gil(|py| -> Result<Vec<PyRasterLoadResultData>, ArrowError> {
+ Python::attach(|py| -> Result<Vec<PyRasterLoadResultData>, ArrowError> {
// Helper to convert pyo3 errors to ArrowError
fn py_err(e: impl std::fmt::Display) -> ArrowError {
ArrowError::InvalidArgumentError(format!("Python raster loader error: {e}"))
@@ -449,9 +449,9 @@ impl PyRasterLoadResult {
/// Create a Python-backed raster loader
#[pyfunction]
pub fn py_raster_loader(
- py_name: PyObject,
- py_supports_format: PyObject,
- py_load: PyObject,
+ py_name: Py<PyAny>,
+ py_supports_format: Py<PyAny>,
+ py_load: Py<PyAny>,
) -> PyRasterLoaderWrapper {
PyRasterLoaderWrapper {
inner: Arc::new(PyRasterLoader::new(py_name, py_supports_format, py_load)),
diff --git a/python/sedonadb/src/runtime.rs b/python/sedonadb/src/runtime.rs
index 772f90b1e..77af22760 100644
--- a/python/sedonadb/src/runtime.rs
+++ b/python/sedonadb/src/runtime.rs
@@ -33,14 +33,14 @@ where
py.run(cr"pass", None, None)?;
py.check_signals()?;
- py.allow_threads(|| {
+ py.detach(|| {
runtime.block_on(async {
tokio::pin!(fut);
loop {
tokio::select! {
res = &mut fut => break Ok(res),
_ = sleep(INTERVAL_CHECK_SIGNALS) => {
- Python::with_gil(|py| {
+ Python::attach(|py| {
py.run(cr"pass", None, None)?;
py.check_signals()
})?;
@@ -65,7 +65,7 @@ where
tokio::select! {
res = &mut fut => break Ok(res),
_ = sleep(INTERVAL_CHECK_SIGNALS) => {
- Python::with_gil(|py| {
+ Python::attach(|py| {
py.run(cr"pass", None, None)?;
py.check_signals()
})?;
diff --git a/python/sedonadb/src/schema.rs b/python/sedonadb/src/schema.rs
index 32466d175..73250498d 100644
--- a/python/sedonadb/src/schema.rs
+++ b/python/sedonadb/src/schema.rs
@@ -84,7 +84,7 @@ impl PySedonaSchema {
fn field<'py>(
&self,
py: Python<'py>,
- index_or_name: PyObject,
+ index_or_name: Py<PyAny>,
) -> Result<PySedonaField, PySedonaError> {
if let Ok(index) = index_or_name.extract::<usize>(py) {
if index < self.inner.fields().len() {
@@ -211,7 +211,7 @@ impl PySedonaType {
#[pymethods]
impl PySedonaType {
#[getter]
- fn crs<'py>(&self, py: Python<'py>) -> Result<Option<PyObject>, PySedonaError> {
+ fn crs<'py>(&self, py: Python<'py>) -> Result<Option<Py<PyAny>>, PySedonaError> {
if let Some(crs) = self.inner.crs() {
let json = py.import("json")?;
let geoarrow_types_crs = py.import("geoarrow.types.crs")?;
@@ -229,7 +229,7 @@ impl PySedonaType {
}
#[getter]
- fn edge_type<'py>(&self, py: Python<'py>) -> Result<Option<PyObject>, PySedonaError> {
+ fn edge_type<'py>(&self, py: Python<'py>) -> Result<Option<Py<PyAny>>, PySedonaError> {
match &self.inner {
SedonaType::Wkb(edges, _) | SedonaType::WkbView(edges, _) => {
let geoarrow_types = py.import("geoarrow.types")?;
diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs
index d0267a1a0..a679b03ef 100644
--- a/python/sedonadb/src/udf.rs
+++ b/python/sedonadb/src/udf.rs
@@ -32,7 +32,7 @@ use datafusion_ffi::udf::FFI_ScalarUDF;
use pyo3::{
pyclass, pyfunction, pymethods,
types::{PyAnyMethods, PyCapsule, PyTuple, PyTupleMethods},
- Bound, PyObject, Python,
+ Bound, Py, PyAny, Python,
};
use sedona_expr::aggregate_udf::{SedonaAccumulator, SedonaAccumulatorRef, SedonaAggregateUDF};
use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
@@ -185,9 +185,9 @@ fn validate_imported_type(
#[pyfunction]
pub fn sedona_scalar_udf<'py>(
py: Python<'py>,
- py_invoke_batch: PyObject,
- py_return_type: PyObject,
- py_input_types: Option<Vec<PyObject>>,
+ py_invoke_batch: Py<PyAny>,
+ py_return_type: Py<PyAny>,
+ py_input_types: Option<Vec<Py<PyAny>>>,
volatility: &str,
name: &str,
) -> Result<PySedonaScalarUdf, PySedonaError> {
@@ -203,9 +203,9 @@ pub fn sedona_scalar_udf<'py>(
fn sedona_scalar_kernel<'py>(
py: Python<'py>,
- input_types: Option<Vec<PyObject>>,
- py_return_field: PyObject,
- py_invoke_batch: PyObject,
+ input_types: Option<Vec<Py<PyAny>>>,
+ py_return_field: Py<PyAny>,
+ py_invoke_batch: Py<PyAny>,
) -> Result<PySedonaScalarKernel, PySedonaError> {
let matcher = if let Some(input_types) = input_types {
let arg_matchers = input_types
@@ -230,8 +230,8 @@ fn sedona_scalar_kernel<'py>(
#[derive(Debug)]
struct PySedonaScalarKernel {
matcher: Option<ArgMatcher>,
- py_return_field: PyObject,
- py_invoke_batch: PyObject,
+ py_return_field: Py<PyAny>,
+ py_invoke_batch: Py<PyAny>,
}
impl SedonaScalarKernel for PySedonaScalarKernel {
@@ -257,7 +257,7 @@ impl SedonaScalarKernel for PySedonaScalarKernel {
return Ok(return_type);
}
- let return_type = Python::with_gil(|py| -> Result<Option<SedonaType>, PySedonaError> {
+ let return_type = Python::attach(|py| -> Result<Option<SedonaType>, PySedonaError> {
let py_sedona_types = args
.iter()
.map(|arg| -> Result<_, PySedonaError> { Ok(PySedonaType::new(arg.clone())) })
@@ -294,7 +294,7 @@ impl SedonaScalarKernel for PySedonaScalarKernel {
num_rows: usize,
_config_options: Option<&ConfigOptions>,
) -> Result<ColumnarValue> {
- let result = Python::with_gil(|py| -> Result<ArrayRef, PySedonaError> {
+ let result = Python::attach(|py| -> Result<ArrayRef, PySedonaError> {
let py_values = zip(arg_types, args)
.map(|(sedona_type, arg)| PySedonaValue {
sedona_type: PySedonaType::new(sedona_type.clone()),
@@ -485,10 +485,10 @@ impl PySedonaAggregateUdf {
#[pyfunction]
pub fn sedona_aggregate_udf<'py>(
py: Python<'py>,
- py_factory: PyObject,
- py_return_type: PyObject,
- py_input_types: Vec<PyObject>,
- py_state_types: Vec<PyObject>,
+ py_factory: Py<PyAny>,
+ py_return_type: Py<PyAny>,
+ py_input_types: Vec<Py<PyAny>>,
+ py_state_types: Vec<Py<PyAny>>,
volatility: &str,
name: &str,
) -> Result<PySedonaAggregateUdf, PySedonaError> {
@@ -523,7 +523,7 @@ pub fn sedona_aggregate_udf<'py>(
struct PySedonaAggregateKernel {
matcher: ArgMatcher,
state_types: Vec<SedonaType>,
- py_factory: PyObject,
+ py_factory: Py<PyAny>,
}
impl SedonaAccumulator for PySedonaAggregateKernel {
@@ -536,7 +536,7 @@ impl SedonaAccumulator for PySedonaAggregateKernel {
args: &[SedonaType],
output_type: &SedonaType,
) -> Result<Box<dyn Accumulator>> {
- let instance = Python::with_gil(|py| -> Result<PyObject, PySedonaError> {
+ let instance = Python::attach(|py| -> Result<Py<PyAny>, PySedonaError> {
Ok(self.py_factory.call0(py)?)
})?;
Ok(Box::new(PySedonaAccumulator {
@@ -558,7 +558,7 @@ impl SedonaAccumulator for PySedonaAggregateKernel {
#[derive(Debug)]
struct PySedonaAccumulator {
- instance: PyObject,
+ instance: Py<PyAny>,
input_types: Vec<SedonaType>,
state_types: Vec<SedonaType>,
output_type: SedonaType,
@@ -594,7 +594,7 @@ impl PySedonaAccumulator {
/// SedonaType is used to validate the array's logical type.
fn import_scalar(
py: Python<'_>,
- result: PyObject,
+ result: Py<PyAny>,
expected: &SedonaType,
context: &str,
) -> Result<ScalarValue, PySedonaError> {
@@ -619,7 +619,7 @@ impl PySedonaAccumulator {
impl Accumulator for PySedonaAccumulator {
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- Python::with_gil(|py| -> Result<(), PySedonaError> {
+ Python::attach(|py| -> Result<(), PySedonaError> {
let args = Self::arrays_to_py_values(py, &self.input_types, values)?;
self.instance.call_method1(py, "update", (args,))?;
Ok(())
@@ -628,7 +628,7 @@ impl Accumulator for PySedonaAccumulator {
}
fn evaluate(&mut self) -> Result<ScalarValue> {
- let scalar = Python::with_gil(|py| -> Result<ScalarValue, PySedonaError> {
+ let scalar = Python::attach(|py| -> Result<ScalarValue, PySedonaError> {
let result = self.instance.call_method0(py, "evaluate")?;
Self::import_scalar(py, result, &self.output_type, "aggregate UDF evaluate()")
})?;
@@ -636,7 +636,7 @@ impl Accumulator for PySedonaAccumulator {
}
fn state(&mut self) -> Result<Vec<ScalarValue>> {
- let scalars = Python::with_gil(|py| -> Result<Vec<ScalarValue>, PySedonaError> {
+ let scalars = Python::attach(|py| -> Result<Vec<ScalarValue>, PySedonaError> {
let result = self.instance.call_method0(py, "state")?;
let result_bound = result.bind(py);
let tuple = result_bound.downcast::<PyTuple>().map_err(|_| {
@@ -667,7 +667,7 @@ impl Accumulator for PySedonaAccumulator {
}
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- Python::with_gil(|py| -> Result<(), PySedonaError> {
+ Python::attach(|py| -> Result<(), PySedonaError> {
let args = Self::arrays_to_py_values(py, &self.state_types, states)?;
self.instance.call_method1(py, "merge", (args,))?;
Ok(())
diff --git a/python/sedonadb/tests/test_context.py b/python/sedonadb/tests/test_context.py
index 1fe7811f6..e743fc644 100644
--- a/python/sedonadb/tests/test_context.py
+++ b/python/sedonadb/tests/test_context.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
import json
+from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any, Mapping
@@ -55,6 +56,23 @@ def test_options():
assert "DataFrame object at" not in repr(sd.sql("SELECT 1 as one"))
+def test_shared_context_and_dataframe_from_threads():
+ sd = sedonadb.connect()
+ df = sd.sql("SELECT * FROM (VALUES (1), (2), (3)) AS t(x)")
+
+ def count_shared_dataframe(_):
+ return df.count()
+
+ def create_and_read_view(i):
+ name = f"threaded_view_{i}"
+ sd.sql(f"SELECT {i} AS x").to_view(name, overwrite=True)
+ return sd.view(name).count()
+
+ with ThreadPoolExecutor(max_workers=8) as executor:
+ assert list(executor.map(count_shared_dataframe, range(16))) == [3] * 16
+ assert list(executor.map(create_and_read_view, range(16))) == [1] * 16
+
+
def test_read_parquet(con, geoarrow_data):
# Check one file
tab = con.read_parquet(
diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs
index dfc26fe09..7b9a2b349 100644
--- a/rust/sedona/src/context.rs
+++ b/rust/sedona/src/context.rs
@@ -16,7 +16,7 @@
// under the License.
use std::{
collections::{HashMap, VecDeque},
- sync::{Arc, RwLock},
+ sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
};
use crate::exec::create_plan_from_sql;
@@ -45,7 +45,7 @@ use datafusion::{dataframe::DataFrameWriteOptions, execution::memory_pool::Memor
use datafusion_common::not_impl_err;
use datafusion_expr::dml::InsertOp;
use datafusion_expr::sqlparser::dialect::{dialect_from_str, Dialect};
-use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, SortExpr};
+use datafusion_expr::{AggregateUDFImpl, LogicalPlan, LogicalPlanBuilder, ScalarUDFImpl, SortExpr};
use parking_lot::Mutex;
use sedona_common::{
option::add_sedona_option_extension, sedona_internal_datafusion_err, CrsProviderOption,
@@ -53,8 +53,11 @@ use sedona_common::{
};
use sedona_datasource::provider::external_table;
use sedona_datasource::spec::ExternalFormatSpec;
-use sedona_expr::scalar_udf::IntoScalarKernelRefs;
-use sedona_expr::{aggregate_udf::IntoSedonaAccumulatorRefs, function_set::FunctionSet};
+use sedona_expr::scalar_udf::{IntoScalarKernelRefs, SedonaScalarUDF};
+use sedona_expr::{
+ aggregate_udf::{IntoSedonaAccumulatorRefs, SedonaAggregateUDF},
+ function_set::FunctionSet,
+};
use sedona_geoparquet::options::TableGeoParquetOptions;
use sedona_geoparquet::{
format::GeoParquetFormatFactory,
@@ -83,7 +86,7 @@ use sedona_raster::raster_loader::{AsyncRasterLoader, RasterLoaderConfig, Raster
/// interface for configuring the behaviour of
pub struct SedonaContext {
pub ctx: SessionContext,
- pub functions: FunctionSet,
+ functions: RwLock<FunctionSet>,
/// Per-session registry of async raster byte loaders, keyed by
/// `outdb_format`. Held behind an `Arc<RwLock<…>>` so the registered
/// `RS_EnsureLoaded` UDF instance and any extension crates' `register(&ctx)`
@@ -235,7 +238,7 @@ impl SedonaContext {
pub fn new_from_context(ctx: SessionContext) -> Result<Self> {
let mut out = Self {
ctx,
- functions: FunctionSet::new(),
+ functions: RwLock::new(FunctionSet::new()),
raster_loader_registry: Arc::new(RwLock::new(RasterLoaderRegistry::new())),
};
@@ -367,15 +370,67 @@ impl SedonaContext {
}
}
+ fn functions(&self) -> Result<RwLockReadGuard<'_, FunctionSet>> {
+ self.functions
+ .read()
+ .map_err(|_| sedona_internal_datafusion_err!("Function registry lock poisoned"))
+ }
+
+ fn functions_mut(&self) -> Result<RwLockWriteGuard<'_, FunctionSet>> {
+ self.functions
+ .write()
+ .map_err(|_| sedona_internal_datafusion_err!("Function registry lock poisoned"))
+ }
+
+ pub fn scalar_udf(&self, name: &str) -> Result<Option<SedonaScalarUDF>> {
+ Ok(self.functions()?.scalar_udf(name).cloned())
+ }
+
+ pub fn scalar_udf_names(&self) -> Result<Vec<String>> {
+ Ok(self
+ .ctx
+ .state()
+ .scalar_functions()
+ .keys()
+ .cloned()
+ .collect())
+ }
+
+ pub fn register_sedona_scalar_udf(&self, udf: SedonaScalarUDF) -> Result<()> {
+ let name = udf.name().to_string();
+ let mut functions = self.functions_mut()?;
+ functions.insert_scalar_udf(udf);
+ let udf = functions.scalar_udf(&name).unwrap().clone();
+ drop(functions);
+
+ self.ctx.register_udf(udf.into());
+ Ok(())
+ }
+
+ pub fn register_sedona_aggregate_udf(&self, udf: SedonaAggregateUDF) -> Result<()> {
+ let name = udf.name().to_string();
+ let mut functions = self.functions_mut()?;
+ functions.insert_aggregate_udf(udf);
+ let udf = functions.aggregate_udf(&name).unwrap().clone();
+ drop(functions);
+
+ self.ctx.register_udaf(udf.into());
+ Ok(())
+ }
+
/// Register all functions in a [FunctionSet] with this context
pub fn register_function_set(&mut self, function_set: FunctionSet) {
+ let mut functions = self
+ .functions
+ .write()
+ .expect("fresh SedonaContext function registry lock should not be poisoned");
for udf in function_set.scalar_udfs() {
- self.functions.insert_scalar_udf(udf.clone());
+ functions.insert_scalar_udf(udf.clone());
self.ctx.register_udf(udf.clone().into());
}
for udf in function_set.aggregate_udfs() {
- self.functions.insert_aggregate_udf(udf.clone());
+ functions.insert_aggregate_udf(udf.clone());
self.ctx.register_udaf(udf.clone().into());
}
}
@@ -385,8 +440,9 @@ impl SedonaContext {
&mut self,
kernels: impl Iterator<Item = (&'a str, impl IntoScalarKernelRefs)>,
) -> Result<()> {
+ let mut functions = self.functions_mut()?;
for (name, kernel) in kernels {
- let udf = self.functions.add_scalar_udf_impl(name, kernel)?;
+ let udf = functions.add_scalar_udf_impl(name, kernel)?;
self.ctx.register_udf(udf.clone().into());
}
@@ -397,8 +453,9 @@ impl SedonaContext {
&mut self,
kernels: impl Iterator<Item = (&'a str, impl IntoSedonaAccumulatorRefs)>,
) -> Result<()> {
+ let mut functions = self.functions_mut()?;
for (name, kernel) in kernels {
- let udf = self.functions.add_aggregate_udf_kernel(name, kernel)?;
+ let udf = functions.add_aggregate_udf_kernel(name, kernel)?;
self.ctx.register_udaf(udf.clone().into());
}
diff --git a/rust/sedona/src/show.rs b/rust/sedona/src/show.rs
index 01376281f..0787e3820 100644
--- a/rust/sedona/src/show.rs
+++ b/rust/sedona/src/show.rs
@@ -44,12 +44,10 @@ pub fn show_batches<'a, W: std::io::Write>(
options: DisplayTableOptions<'a>,
) -> Result<()> {
let format_fn = ctx
- .functions
- .scalar_udf("sd_format")
+ .scalar_udf("sd_format")?
.ok_or(sedona_internal_datafusion_err!(
"sd_format UDF does not exist"
- ))?
- .clone();
+ ))?;
let session_config = ctx.ctx.copied_config();
let session_config_options = session_config.options();