From c76091d354c18b8f7ab5f5f19946f988bbd752aa Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 22 Jun 2026 22:13:10 -0500 Subject: [PATCH 01/55] sketch structure --- Cargo.lock | 4 + c/sedona-extension/Cargo.toml | 4 + c/sedona-extension/src/execution_plan.rs | 95 ++++++++++++++++++++++++ c/sedona-extension/src/lib.rs | 2 + c/sedona-extension/src/table_provider.rs | 93 +++++++++++++++++++++++ 5 files changed, 198 insertions(+) create mode 100644 c/sedona-extension/src/execution_plan.rs create mode 100644 c/sedona-extension/src/table_provider.rs diff --git a/Cargo.lock b/Cargo.lock index 95b4343453..70590fcbff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5781,8 +5781,12 @@ version = "0.4.0" dependencies = [ "arrow-array", "arrow-schema", + "async-trait", + "datafusion-catalog", "datafusion-common", + "datafusion-execution", "datafusion-expr", + "datafusion-physical-plan", "libc", "sedona-common", "sedona-expr", diff --git a/c/sedona-extension/Cargo.toml b/c/sedona-extension/Cargo.toml index e8f8c17c56..b0ea8225b2 100644 --- a/c/sedona-extension/Cargo.toml +++ b/c/sedona-extension/Cargo.toml @@ -30,8 +30,12 @@ rust-version.workspace = true [dependencies] arrow-array = { workspace = true, features = ["ffi"]} arrow-schema = { workspace = true, features = ["ffi"]} +async-trait = { workspace = true } datafusion-common = { workspace = true } +datafusion-catalog = { workspace = true } +datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } +datafusion-physical-plan = { workspace = true } libc = "0.2.178" sedona-common = { workspace = true } sedona-expr = { workspace = true } diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs new file mode 100644 index 0000000000..6486462900 --- /dev/null +++ b/c/sedona-extension/src/execution_plan.rs @@ -0,0 +1,95 @@ +// 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. + +use std::{any::Any, sync::Arc}; + +use datafusion_common::{Result, Statistics}; +use datafusion_execution::TaskContext; +use datafusion_physical_plan::{ + metrics::MetricsSet, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, +}; + +#[derive(Debug)] +struct ImportedTableProviderExec { + name: String, + properties: Arc, +} + +impl DisplayAs for ImportedTableProviderExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + todo!() + } +} + +impl ExecutionPlan for ImportedTableProviderExec { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "MyExecPlan" + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn partition_statistics(&self, partition: Option) -> Result { + todo!() + } + + fn cardinality_effect(&self) -> datafusion_physical_plan::execution_plan::CardinalityEffect { + todo!() + } + + fn maintains_input_order(&self) -> Vec { + todo!() + } + + fn metrics(&self) -> Option { + todo!() + } + + fn supports_limit_pushdown(&self) -> bool { + todo!() + } + + fn statistics(&self) -> Result { + todo!() + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + todo!() + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert!(children.is_empty()); + Ok(self) + } +} diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs index 073ed939bc..0358a28707 100644 --- a/c/sedona-extension/src/lib.rs +++ b/c/sedona-extension/src/lib.rs @@ -17,3 +17,5 @@ pub mod extension; pub mod scalar_kernel; +pub mod execution_plan; +pub mod table_provider; diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs new file mode 100644 index 0000000000..a064c327b0 --- /dev/null +++ b/c/sedona-extension/src/table_provider.rs @@ -0,0 +1,93 @@ +// 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. + +use std::{any::Any, sync::Arc}; + +use arrow_schema::SchemaRef; +use async_trait::async_trait; +use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider}; +use datafusion_common::{Result, Statistics}; +use datafusion_expr::{dml::InsertOp, Expr, TableType}; +use datafusion_physical_plan::ExecutionPlan; + +#[derive(Debug)] +struct ImportedTableProvider { + schema: SchemaRef, +} + +#[async_trait] +impl TableProvider for ImportedTableProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + todo!() + } + + fn statistics(&self) -> Option { + todo!() + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> Result> { + todo!() + } + + async fn scan_with_args<'a>( + &self, + state: &dyn Session, + args: ScanArgs<'a>, + ) -> Result { + todo!() + } + + async fn insert_into( + &self, + _state: &dyn Session, + _input: Arc, + _insert_op: InsertOp, + ) -> Result> { + todo!() + } + + async fn delete_from( + &self, + _state: &dyn Session, + _filters: Vec, + ) -> Result> { + todo!() + } + + async fn update( + &self, + _state: &dyn Session, + _assignments: Vec<(String, Expr)>, + _filters: Vec, + ) -> Result> { + todo!() + } +} From aacba44ea9c8afea9dcfa4d09c917c098065e4b5 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 22 Jun 2026 23:16:41 -0500 Subject: [PATCH 02/55] sketch --- c/sedona-extension/src/execution_plan.rs | 13 ++-- c/sedona-extension/src/sedona_extension.h | 80 +++++++++++++++++++++++ 2 files changed, 86 insertions(+), 7 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 6486462900..0fb4f16766 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -20,14 +20,14 @@ use std::{any::Any, sync::Arc}; use datafusion_common::{Result, Statistics}; use datafusion_execution::TaskContext; use datafusion_physical_plan::{ - metrics::MetricsSet, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, - SendableRecordBatchStream, + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, execution_plan::CardinalityEffect, metrics::MetricsSet, }; #[derive(Debug)] struct ImportedTableProviderExec { name: String, properties: Arc, + children: Vec> } impl DisplayAs for ImportedTableProviderExec { @@ -42,7 +42,7 @@ impl ExecutionPlan for ImportedTableProviderExec { } fn name(&self) -> &str { - "MyExecPlan" + &self.name } fn properties(&self) -> &PlanProperties { @@ -53,7 +53,7 @@ impl ExecutionPlan for ImportedTableProviderExec { todo!() } - fn cardinality_effect(&self) -> datafusion_physical_plan::execution_plan::CardinalityEffect { + fn cardinality_effect(&self) -> CardinalityEffect { todo!() } @@ -82,14 +82,13 @@ impl ExecutionPlan for ImportedTableProviderExec { } fn children(&self) -> Vec<&Arc> { - vec![] + self.children.iter().collect() } fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - assert!(children.is_empty()); - Ok(self) + todo!() } } diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 7191af21f6..4485fe2603 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -18,6 +18,7 @@ #ifndef SEDONA_EXTENSION_H #define SEDONA_EXTENSION_H +#include #include #ifdef __cplusplus @@ -203,6 +204,85 @@ struct SedonaCScalarKernel { void* private_data; }; +struct SedonaCError { + const char* err; + + uint32_t err_len; + + uint32_t reserved; + + /// \brief Release this instance + /// + /// Implementations of this callback must set self->release to NULL. + void (*release)(struct SedonaCError* self); +}; + +struct SedonaCExpr { + // Get a property of this expression (e.g., serialize, extract bbox) + int (*get_property_schema)(const struct SedonaCExpr* self, const char* property, + struct SedonaCError* err); + int (*get_property)(const struct SedonaCExpr* self, const char* property, + const char* args, struct ArrowArray* out, struct SedonaCError* err); + + void* reserved; + + /// \brief Release this instance + /// + /// Implementations of this callback must set self->release to NULL. + void (*release)(struct SedonaCExecutionPlan* self); + + /// \brief Opaque implementation-specific data + void* private_data; +}; + +struct SedonaCExecutionPlan; + +struct SedonaCExecutionPlanArgs { + const uint8_t* args; + size_t args_len; + struct SedonaCExecutionPlan** children; + size_t num_children; + void* reserved; +}; + +struct SedonaCExecutionPlan { + void (*get_schema)(const struct SedonaCExecutionPlan* self, struct ArrowSchema* out); + + // Extract some serializable property from this plan (e.g., plan properties) + int (*get_property_schema)(const struct SedonaCExecutionPlan* self, + const char* property, struct ArrowSchema* out, + struct SedonaCError* err); + int (*get_property)(const struct SedonaCExecutionPlan* self, const char* property, + struct SedonaCExecutionPlanArgs* args, struct ArrowArray* out, + struct SedonaCError* err); + + // Clone this plan based on some new information (e.g., try pushdown filters) + int (*with_property)(const struct SedonaCExecutionPlan* self, const char* property, + struct SedonaCExecutionPlanArgs* args, + struct SedonaCExecutionPlan* out, struct SedonaCError* err); + + // Resolve a synchronous stream for one partition from this plan + int (*execute)(const struct SedonaCExecutionPlan* self, + struct SedonaCExecutionPlanArgs* args, struct ArrowArrayStream* out, + struct SedonaCError* err); + + // Future implemenatation with async streams + int (*execute_async)(const struct SedonaCExecutionPlan* self, + struct SedonaCExecutionPlanArgs* args, void* out, + struct SedonaCError* err); + + // Reserved for future use + void* reserved; + + /// \brief Release this instance + /// + /// Implementations of this callback must set self->release to NULL. + void (*release)(struct SedonaCExecutionPlan* self); + + /// \brief Opaque implementation-specific data + void* private_data; +}; + #ifdef __cplusplus } #endif From 6c0037b2232d2cff32bb599bbfee3e9415743bb9 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 22 Jun 2026 23:36:02 -0500 Subject: [PATCH 03/55] rust translation --- c/sedona-extension/src/extension.rs | 192 +++++++++++++++++++++++++++- 1 file changed, 190 insertions(+), 2 deletions(-) diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs index 7d957a57c2..3dd7cbc9c1 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -18,10 +18,13 @@ use std::{ ffi::c_int, os::raw::{c_char, c_void}, - ptr::null_mut, + ptr::{null_mut}, }; -use arrow_array::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; +use arrow_array::{ + ffi::{FFI_ArrowArray, FFI_ArrowSchema}, + ffi_stream::FFI_ArrowArrayStream, +}; /// Raw FFI representation of the SedonaCScalarKernel /// @@ -120,3 +123,188 @@ struct ArrowSchemaInternal { release: Option, private_data: *mut c_void, } + +/// Raw FFI representation of the SedonaCError +#[derive(Default)] +#[repr(C)] +pub struct SedonaCError { + pub err: *const c_char, + pub err_len: u32, + pub reserved: u32, + pub release: Option, +} + +impl SedonaCError { + pub fn new(message: &str) -> Self { + use std::ffi::CString; + + match CString::new(message) { + Ok(c_string) => { + let len = message.len() as u32; + let ptr = c_string.into_raw(); + + extern "C" fn release_error(self_: *mut SedonaCError) { + unsafe { + if !(*self_).err.is_null() { + let _ = std::ffi::CString::from_raw((*self_).err as *mut c_char); + } + (*self_).release = None; + } + } + + SedonaCError { + err: ptr, + err_len: len, + reserved: 0, + release: Some(release_error), + } + } + Err(_) => UNKNOWN_SEDONA_C_ERROR + } + } +} + +impl Drop for SedonaCError { + fn drop(&mut self) { + if let Some(releaser) = self.release { + unsafe { releaser(self) } + self.release = None; + } + } +} + +extern "C" fn sedona_c_noop_release(_self: *mut SedonaCError) {} + +pub const UNKNOWN_SEDONA_C_ERROR: SedonaCError = SedonaCError { + err: c"Unknown error".as_ptr(), + err_len: "Unknown error".len() as u32, + reserved: 0, + release: Some(sedona_c_noop_release), +}; + +/// Raw FFI representation of the SedonaCExpr +#[derive(Default)] +#[repr(C)] +pub struct SedonaCExpr { + pub get_property_schema: Option< + unsafe extern "C" fn( + self_: *const SedonaCExpr, + property: *const c_char, + err: *mut SedonaCError, + ) -> c_int, + >, + + pub get_property: Option< + unsafe extern "C" fn( + self_: *const SedonaCExpr, + property: *const c_char, + args: *const c_char, + out: *mut FFI_ArrowArray, + err: *mut SedonaCError, + ) -> c_int, + >, + + pub reserved: *mut c_void, + + pub release: Option, + + pub private_data: *mut c_void, +} + +unsafe impl Send for SedonaCExpr {} +unsafe impl Sync for SedonaCExpr {} + +impl Drop for SedonaCExpr { + fn drop(&mut self) { + if let Some(releaser) = self.release { + unsafe { releaser(self) } + self.release = None; + self.private_data = null_mut(); + } + } +} + +/// Raw FFI representation of the SedonaCExecutionPlanArgs +#[derive(Default)] +#[repr(C)] +pub struct SedonaCExecutionPlanArgs { + pub args: *const u8, + pub args_len: usize, + pub children: *mut *mut SedonaCExecutionPlan, + pub num_children: usize, + pub reserved: *mut c_void, +} + +/// Raw FFI representation of the SedonaCExecutionPlan +#[derive(Default)] +#[repr(C)] +pub struct SedonaCExecutionPlan { + pub get_schema: + Option, + + pub get_property_schema: Option< + unsafe extern "C" fn( + self_: *const SedonaCExecutionPlan, + property: *const c_char, + out: *mut FFI_ArrowSchema, + err: *mut SedonaCError, + ) -> c_int, + >, + + pub get_property: Option< + unsafe extern "C" fn( + self_: *const SedonaCExecutionPlan, + property: *const c_char, + args: *mut SedonaCExecutionPlanArgs, + out: *mut FFI_ArrowArray, + err: *mut SedonaCError, + ) -> c_int, + >, + + pub with_property: Option< + unsafe extern "C" fn( + self_: *const SedonaCExecutionPlan, + property: *const c_char, + args: *mut SedonaCExecutionPlanArgs, + out: *mut SedonaCExecutionPlan, + err: *mut SedonaCError, + ) -> c_int, + >, + + pub execute: Option< + unsafe extern "C" fn( + self_: *const SedonaCExecutionPlan, + args: *mut SedonaCExecutionPlanArgs, + out: *mut FFI_ArrowArrayStream, + err: *mut SedonaCError, + ) -> c_int, + >, + + pub execute_async: Option< + unsafe extern "C" fn( + self_: *const SedonaCExecutionPlan, + args: *mut SedonaCExecutionPlanArgs, + out: *mut c_void, + err: *mut SedonaCError, + ) -> c_int, + >, + + pub reserved: *mut c_void, + + pub release: Option, + + pub private_data: *mut c_void, +} + +unsafe impl Send for SedonaCExecutionPlan {} +unsafe impl Sync for SedonaCExecutionPlan {} + +impl Drop for SedonaCExecutionPlan { + fn drop(&mut self) { + if let Some(releaser) = self.release { + unsafe { releaser(self) } + self.release = None; + self.private_data = null_mut(); + } + } +} From 57906b0478d95e6bf319559051466c28ef1db6e0 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 22 Jun 2026 23:44:23 -0500 Subject: [PATCH 04/55] error stuff --- c/sedona-extension/src/extension.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs index 3dd7cbc9c1..fea947a801 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -18,7 +18,7 @@ use std::{ ffi::c_int, os::raw::{c_char, c_void}, - ptr::{null_mut}, + ptr::null_mut, }; use arrow_array::{ @@ -159,7 +159,19 @@ impl SedonaCError { release: Some(release_error), } } - Err(_) => UNKNOWN_SEDONA_C_ERROR + Err(_) => UNKNOWN_SEDONA_C_ERROR, + } + } +} + +impl std::fmt::Display for SedonaCError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.err.is_null() || self.err_len == 0 { + return write!(f, ""); + } + unsafe { + let bytes = std::slice::from_raw_parts(self.err as *const u8, self.err_len as usize); + write!(f, "{}", String::from_utf8_lossy(bytes)) } } } From 90c2fd7eafcfba970216d2b435486a0c37d4fe5e Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 23 Jun 2026 00:07:22 -0500 Subject: [PATCH 05/55] try the expr piece --- Cargo.lock | 4 + c/sedona-extension/Cargo.toml | 4 + c/sedona-extension/src/expr.rs | 353 ++++++++++++++++++++++ c/sedona-extension/src/extension.rs | 3 +- c/sedona-extension/src/lib.rs | 1 + c/sedona-extension/src/sedona_extension.h | 3 +- 6 files changed, 366 insertions(+), 2 deletions(-) create mode 100644 c/sedona-extension/src/expr.rs diff --git a/Cargo.lock b/Cargo.lock index 70590fcbff..deabb83025 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5786,12 +5786,16 @@ dependencies = [ "datafusion-common", "datafusion-execution", "datafusion-expr", + "datafusion-physical-expr", "datafusion-physical-plan", "libc", "sedona-common", "sedona-expr", + "sedona-geometry", "sedona-schema", "sedona-testing", + "serde", + "serde_json", ] [[package]] diff --git a/c/sedona-extension/Cargo.toml b/c/sedona-extension/Cargo.toml index b0ea8225b2..15c7109326 100644 --- a/c/sedona-extension/Cargo.toml +++ b/c/sedona-extension/Cargo.toml @@ -35,9 +35,13 @@ datafusion-common = { workspace = true } datafusion-catalog = { workspace = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } +datafusion-physical-expr = { workspace = true } datafusion-physical-plan = { workspace = true } libc = "0.2.178" sedona-common = { workspace = true } sedona-expr = { workspace = true } +sedona-geometry = { workspace = true } sedona-schema = { workspace = true } sedona-testing = { path = "../../rust/sedona-testing" } +serde = { workspace = true } +serde_json = { workspace = true } diff --git a/c/sedona-extension/src/expr.rs b/c/sedona-extension/src/expr.rs new file mode 100644 index 0000000000..6694a392ef --- /dev/null +++ b/c/sedona-extension/src/expr.rs @@ -0,0 +1,353 @@ +// 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. + +use std::{ + ffi::{c_char, c_int, c_void, CStr}, + fmt::Debug, + ptr::null_mut, + sync::Arc, +}; + +use arrow_array::{builder::StringBuilder, ffi::FFI_ArrowArray, Array}; +use arrow_schema::SchemaRef; +use datafusion_common::{DataFusionError, Result}; +use datafusion_physical_expr::PhysicalExpr; +use sedona_expr::spatial_filter::SpatialFilterFactory; +use serde::Deserialize; + +use crate::extension::{SedonaCError, SedonaCExpr}; + +/// A wrapper around a DataFusion [PhysicalExpr] with its associated schema. +/// +/// This struct bundles a physical expression with the schema context needed to +/// interpret it, which is required for operations like type resolution +/// and extracting spatial filter information. +pub struct PhysicalExprWithSchema { + expr: Arc, + schema: SchemaRef, +} + +impl Debug for PhysicalExprWithSchema { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PhysicalExprWithSchema") + .field("expr", &self.expr.to_string()) + .field("schema", &self.schema) + .finish() + } +} + +impl PhysicalExprWithSchema { + /// Create a new PhysicalExprWithSchema + pub fn new(expr: Arc, schema: SchemaRef) -> Self { + Self { expr, schema } + } + + /// Get a reference to the expression + pub fn expr(&self) -> &Arc { + &self.expr + } + + /// Get a reference to the schema + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// Consume self and return the inner expression and schema + pub fn into_parts(self) -> (Arc, SchemaRef) { + (self.expr, self.schema) + } + + /// Extract the bounding box for a spatial filter on the given column + /// + /// Returns the bounding box as a JSON string, or an error if extraction fails. + pub fn filter_bbox(&self, column_name: &str) -> Result { + let factory = SpatialFilterFactory::default(); + let spatial_filter = factory.try_from_expr(&self.expr)?; + let bbox = spatial_filter.filter_bbox(column_name); + serde_json::to_string(&bbox).map_err(|e| { + DataFusionError::Internal(format!("Failed to serialize bounding box: {}", e)) + }) + } +} + +/// Arguments for the bbox property +#[derive(Debug, Deserialize)] +struct BboxArgs { + column: String, +} + +/// Wrapper around a [SedonaCExpr] that can be used to import an expression +/// from a C implementation. +pub struct ImportedExpr { + inner: SedonaCExpr, +} + +impl Debug for ImportedExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ImportedExpr") + .field("inner", &"") + .finish() + } +} + +impl TryFrom for ImportedExpr { + type Error = DataFusionError; + + fn try_from(value: SedonaCExpr) -> Result { + match (&value.get_property, &value.release) { + (Some(_), Some(_)) => Ok(Self { inner: value }), + _ => Err(DataFusionError::Internal( + "Can't import released or uninitialized SedonaCExpr".to_string(), + )), + } + } +} + +impl ImportedExpr { + /// Get a property from this expression + /// + /// # Safety + /// + /// The caller must ensure that the property name and args are valid + /// null-terminated C strings if provided. + pub unsafe fn get_property( + &self, + property: &CStr, + args: Option<&CStr>, + ) -> Result { + let get_property = self.inner.get_property.ok_or_else(|| { + DataFusionError::Internal("get_property callback is null".to_string()) + })?; + + let args_ptr = args.map(|a| a.as_ptr()).unwrap_or(std::ptr::null()); + let mut out = FFI_ArrowArray::empty(); + let mut err = SedonaCError::default(); + + let result = get_property( + &self.inner, + property.as_ptr(), + args_ptr, + &mut out, + &mut err, + ); + + if result != 0 { + return Err(DataFusionError::Internal(format!( + "get_property failed: {}", + err + ))); + } + + Ok(out) + } +} + +/// Export a [PhysicalExprWithSchema] as a [SedonaCExpr] for use across FFI boundaries. +pub struct ExportedExpr { + inner: Arc, +} + +impl ExportedExpr { + /// Create a new ExportedExpr from a PhysicalExprWithSchema + pub fn new(expr_with_schema: PhysicalExprWithSchema) -> Self { + Self { + inner: Arc::new(expr_with_schema), + } + } + + /// Export this expression as a SedonaCExpr + /// + /// The returned SedonaCExpr takes ownership of the Arc and will release + /// it when the release callback is called. + pub fn export(self) -> SedonaCExpr { + let boxed = Box::new(self.inner); + + SedonaCExpr { + get_property_schema: Some(exported_expr_get_property_schema), + get_property: Some(exported_expr_get_property), + reserved: null_mut(), + release: Some(exported_expr_release), + private_data: Box::into_raw(boxed) as *mut c_void, + } + } +} + +unsafe extern "C" fn exported_expr_get_property_schema( + _self_: *const SedonaCExpr, + _property: *const c_char, + err: *mut SedonaCError, +) -> c_int { + // TODO: Implement property schema retrieval + if !err.is_null() { + *err = SedonaCError::new("get_property_schema not implemented"); + } + -1 +} + +unsafe extern "C" fn exported_expr_get_property( + self_: *const SedonaCExpr, + property: *const c_char, + args: *const c_char, + out: *mut FFI_ArrowArray, + err: *mut SedonaCError, +) -> c_int { + if self_.is_null() || (*self_).private_data.is_null() { + if !err.is_null() { + *err = SedonaCError::new("null self pointer"); + } + return -1; + } + + let expr_arc = &*((*self_).private_data as *const Arc); + + if property.is_null() { + if !err.is_null() { + *err = SedonaCError::new("null property pointer"); + } + return -1; + } + + let property_str = match CStr::from_ptr(property).to_str() { + Ok(s) => s, + Err(_) => { + if !err.is_null() { + *err = SedonaCError::new("invalid UTF-8 in property name"); + } + return -1; + } + }; + + match property_str { + "bbox" => { + // Parse args JSON to get column name + if args.is_null() { + if !err.is_null() { + *err = SedonaCError::new("bbox property requires args with column name"); + } + return -1; + } + + let args_str = match CStr::from_ptr(args).to_str() { + Ok(s) => s, + Err(_) => { + if !err.is_null() { + *err = SedonaCError::new("invalid UTF-8 in args"); + } + return -1; + } + }; + + let bbox_args: BboxArgs = match serde_json::from_str(args_str) { + Ok(a) => a, + Err(e) => { + if !err.is_null() { + *err = SedonaCError::new(&format!("failed to parse args: {}", e)); + } + return -1; + } + }; + + // Extract bbox using SpatialFilterFactory + let bbox_json = match expr_arc.filter_bbox(&bbox_args.column) { + Ok(json) => json, + Err(e) => { + if !err.is_null() { + *err = SedonaCError::new(&format!("failed to extract bbox: {}", e)); + } + return -1; + } + }; + + // Create a length-1 string array with the JSON result + let mut builder = StringBuilder::new(); + builder.append_value(&bbox_json); + let array = builder.finish(); + + // Export the array via FFI + std::ptr::write(out, FFI_ArrowArray::new(&array.to_data())); + 0 + } + _ => { + if !err.is_null() { + *err = SedonaCError::new(&format!("unknown property: {}", property_str)); + } + -1 + } + } +} + +unsafe extern "C" fn exported_expr_release(self_: *mut SedonaCExpr) { + if self_.is_null() { + return; + } + + let expr = &mut *self_; + if !expr.private_data.is_null() { + let _ = Box::from_raw(expr.private_data as *mut Arc); + expr.private_data = null_mut(); + } + + expr.get_property_schema = None; + expr.get_property = None; + expr.release = None; +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::Literal; + use datafusion_common::ScalarValue; + + #[test] + fn test_physical_expr_with_schema_new() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let expr: Arc = Arc::new(Literal::new(ScalarValue::Int32(Some(42)))); + let expr_with_schema = PhysicalExprWithSchema::new(expr.clone(), schema.clone()); + + assert!(Arc::ptr_eq(&expr_with_schema.schema, &schema)); + } + + #[test] + fn test_physical_expr_with_schema_into_parts() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let expr: Arc = Arc::new(Literal::new(ScalarValue::Int32(Some(42)))); + let expr_with_schema = PhysicalExprWithSchema::new(expr.clone(), schema.clone()); + + let (returned_expr, returned_schema) = expr_with_schema.into_parts(); + assert_eq!(returned_expr.to_string(), expr.to_string()); + assert!(Arc::ptr_eq(&returned_schema, &schema)); + } + + #[test] + fn test_exported_expr_roundtrip() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let expr: Arc = Arc::new(Literal::new(ScalarValue::Int32(Some(42)))); + let expr_with_schema = PhysicalExprWithSchema::new(expr, schema); + + let exported = ExportedExpr::new(expr_with_schema); + let mut c_expr = exported.export(); + + // Verify release callback works + assert!(c_expr.release.is_some()); + unsafe { + (c_expr.release.unwrap())(&mut c_expr); + } + assert!(c_expr.release.is_none()); + assert!(c_expr.private_data.is_null()); + } +} diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs index fea947a801..61c1cb54b3 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -242,8 +242,9 @@ impl Drop for SedonaCExpr { pub struct SedonaCExecutionPlanArgs { pub args: *const u8, pub args_len: usize, - pub children: *mut *mut SedonaCExecutionPlan, + pub children: *const *const SedonaCExecutionPlan, pub num_children: usize, + pub expr: *const SedonaCExpr, pub reserved: *mut c_void, } diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs index 0358a28707..03757c5f46 100644 --- a/c/sedona-extension/src/lib.rs +++ b/c/sedona-extension/src/lib.rs @@ -16,6 +16,7 @@ // under the License. pub mod extension; +pub mod expr; pub mod scalar_kernel; pub mod execution_plan; pub mod table_provider; diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 4485fe2603..939208b9b4 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -240,8 +240,9 @@ struct SedonaCExecutionPlan; struct SedonaCExecutionPlanArgs { const uint8_t* args; size_t args_len; - struct SedonaCExecutionPlan** children; + const struct SedonaCExecutionPlan** children; size_t num_children; + const struct SedonaCExpr* expr; void* reserved; }; From 56f2ed6f6aa1604007aaf33808b55ae52d114b45 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 24 Jun 2026 11:59:55 -0500 Subject: [PATCH 06/55] move dataframe stuff to dataframe --- python/sedonadb/src/dataframe.rs | 2 +- r/sedonadb/src/rust/src/dataframe.rs | 2 +- rust/sedona/src/context.rs | 264 +----------------------- rust/sedona/src/dataframe.rs | 289 +++++++++++++++++++++++++++ rust/sedona/src/lib.rs | 1 + 5 files changed, 298 insertions(+), 260 deletions(-) create mode 100644 rust/sedona/src/dataframe.rs diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs index 761b4f171f..cb81311bc6 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -34,7 +34,7 @@ use futures::lock::Mutex; use futures::TryStreamExt; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict, PyList}; -use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; +use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; use sedona::projected_reader::simplify_record_batch_reader; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_geoparquet::options::TableGeoParquetOptions; diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index 20af9bf5bc..7d22ed4596 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -28,7 +28,7 @@ use datafusion_expr::JoinType; use datafusion_expr::{select_expr::SelectExpr, Expr, SortExpr}; use datafusion_ffi::table_provider::FFI_TableProvider; use savvy::{savvy, savvy_err, sexp, IntoExtPtrSexp, Result}; -use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; +use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; use sedona::reader::SedonaStreamReader; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_geoparquet::options::TableGeoParquetOptions; diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index dfc26fe09b..d2b2d55e28 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -24,12 +24,8 @@ use crate::object_storage::ensure_object_store_registered_with_options; use crate::{ catalog::DynamicObjectStoreCatalog, random_geometry_provider::RandomGeometryFunction, - show::{show_batches, DisplayTableOptions}, }; -use arrow_array::RecordBatch; use arrow_schema::DataType; -use async_trait::async_trait; -use datafusion::datasource::file_format::format_as_file_type; use datafusion::{ common::plan_err, error::{DataFusionError, Result}, @@ -41,11 +37,8 @@ use datafusion::{ prelude::{DataFrame, SessionConfig, SessionContext}, sql::parser::{DFParser, Statement}, }; -use datafusion::{dataframe::DataFrameWriteOptions, execution::memory_pool::MemoryLimit}; -use datafusion_common::not_impl_err; -use datafusion_expr::dml::InsertOp; +use datafusion::{execution::memory_pool::MemoryLimit}; use datafusion_expr::sqlparser::dialect::{dialect_from_str, Dialect}; -use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, SortExpr}; use parking_lot::Mutex; use sedona_common::{ option::add_sedona_option_extension, sedona_internal_datafusion_err, CrsProviderOption, @@ -55,7 +48,6 @@ 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_geoparquet::options::TableGeoParquetOptions; use sedona_geoparquet::{ format::GeoParquetFormatFactory, provider::{geoparquet_listing_table, GeoParquetReadOptions}, @@ -517,179 +509,6 @@ impl Default for SedonaContext { } } -/// Sedona-specific [`DataFrame`] actions -/// -/// This trait, implemented for [`DataFrame`], extends the DataFrame API to make it -/// ergonomic to work with dataframes that contain geometry columns. Currently these -/// are limited to output functions, as geometry columns currently require special -/// handling when written or exported to an external system. -#[async_trait] -pub trait SedonaDataFrame { - /// Build a table of the first `limit` results in this DataFrame - /// - /// This will limit and execute the query and build a table using [show_batches]. - async fn show_sedona<'a>( - self, - ctx: &SedonaContext, - limit: Option, - options: DisplayTableOptions<'a>, - ) -> Result; - - async fn write_geoparquet( - self, - ctx: &SedonaContext, - path: &str, - options: SedonaWriteOptions, - writer_options: Option, - ) -> Result>; -} - -#[async_trait] -impl SedonaDataFrame for DataFrame { - async fn show_sedona<'a>( - self, - ctx: &SedonaContext, - limit: Option, - mut options: DisplayTableOptions<'a>, - ) -> Result { - let df = if matches!( - self.logical_plan(), - LogicalPlan::Explain(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::Analyze(_) - ) { - // Show multi-line output without truncation for plans like `EXPLAIN` - options.max_row_height = usize::MAX; - - // We don't want to apply an additional .limit() to plans like `Explain` - // as that will trigger an internal error: Unsupported logical plan: Explain must be root of the plan - self - } else { - // Apply limit if specified - self.limit(0, limit)? - }; - - let schema_without_qualifiers = df.schema().clone().strip_qualifiers(); - let schema = schema_without_qualifiers.as_arrow(); - let batches = df.collect().await?; - let mut out = Vec::new(); - show_batches(ctx, &mut out, schema, batches, options)?; - String::from_utf8(out).map_err(|e| DataFusionError::External(Box::new(e))) - } - - async fn write_geoparquet( - self, - ctx: &SedonaContext, - path: &str, - options: SedonaWriteOptions, - writer_options: Option, - ) -> Result, DataFusionError> { - if options.insert_op != InsertOp::Append { - return not_impl_err!( - "{} is not implemented for DataFrame::write_geoparquet.", - options.insert_op - ); - } - - let format = if let Some(parquet_opts) = writer_options { - Arc::new(GeoParquetFormatFactory::new_with_options(parquet_opts)) - } else { - Arc::new(GeoParquetFormatFactory::new()) - }; - - let file_type = format_as_file_type(format); - - let plan = if options.sort_by.is_empty() { - self.into_unoptimized_plan() - } else { - LogicalPlanBuilder::from(self.into_unoptimized_plan()) - .sort(options.sort_by)? - .build()? - }; - - let plan = LogicalPlanBuilder::copy_to( - plan, - path.into(), - file_type, - Default::default(), - options.partition_by, - )? - .build()?; - - DataFrame::new(ctx.ctx.state(), plan).collect().await - } -} - -/// A Sedona-specific copy of [DataFrameWriteOptions] -/// -/// This is needed because [DataFrameWriteOptions] has private fields, so we -/// can't use it in our interfaces. This object can be converted to a -/// [DataFrameWriteOptions] using `.into()`. -pub struct SedonaWriteOptions { - /// Controls how new data should be written to the table, determining whether - /// to append, overwrite, or replace existing data. - pub insert_op: InsertOp, - /// Controls if all partitions should be coalesced into a single output file - /// Generally will have slower performance when set to true. - pub single_file_output: bool, - /// Sets which columns should be used for hive-style partitioned writes by name. - /// Can be set to empty vec![] for non-partitioned writes. - pub partition_by: Vec, - /// Sets which columns should be used for sorting the output by name. - /// Can be set to empty vec![] for non-sorted writes. - pub sort_by: Vec, -} - -impl From for DataFrameWriteOptions { - fn from(value: SedonaWriteOptions) -> Self { - DataFrameWriteOptions::new() - .with_insert_operation(value.insert_op) - .with_single_file_output(value.single_file_output) - .with_partition_by(value.partition_by) - .with_sort_by(value.sort_by) - } -} - -impl SedonaWriteOptions { - /// Create a new SedonaWriteOptions with default values - pub fn new() -> Self { - SedonaWriteOptions { - insert_op: InsertOp::Append, - single_file_output: false, - partition_by: vec![], - sort_by: vec![], - } - } - - /// Set the insert operation - pub fn with_insert_operation(mut self, insert_op: InsertOp) -> Self { - self.insert_op = insert_op; - self - } - - /// Set the single_file_output value to true or false - pub fn with_single_file_output(mut self, single_file_output: bool) -> Self { - self.single_file_output = single_file_output; - self - } - - /// Sets the partition_by columns for output partitioning - pub fn with_partition_by(mut self, partition_by: Vec) -> Self { - self.partition_by = partition_by; - self - } - - /// Sets the sort_by columns for output sorting - pub fn with_sort_by(mut self, sort_by: Vec) -> Self { - self.sort_by = sort_by; - self - } -} - -impl Default for SedonaWriteOptions { - fn default() -> Self { - Self::new() - } -} - // Because Dialect/dialect_from_str is not marked as Send, using the async // function in certain contexts will fail to compile. Here we use a wrapper // to ensure that that the Dialect can be specified and parsed in any async @@ -720,18 +539,18 @@ impl ThreadSafeDialect { #[cfg(test)] mod tests { - - use arrow_array::{create_array, ArrayRef, RecordBatchIterator, RecordBatchReader}; + use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, create_array}; use arrow_schema::{DataType, Field, Schema}; - use datafusion::assert_batches_eq; - use sedona_datasource::spec::{Object, OpenReaderArgs}; + use async_trait::async_trait; +use datafusion::assert_batches_eq; + use datafusion_common::not_impl_err; +use sedona_datasource::spec::{Object, OpenReaderArgs}; use sedona_schema::{ crs::{deserialize_crs, lnglat}, datatypes::{Edges, SedonaType}, schema::SedonaSchema, }; use sedona_testing::data::test_geoparquet; - use tempfile::tempdir; use super::*; @@ -931,77 +750,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn show() { - let ctx = SedonaContext::new(); - let tbl = ctx - .sql("SELECT 1 as one") - .await - .unwrap() - .show_sedona(&ctx, None, DisplayTableOptions::default()) - .await - .unwrap(); - - #[rustfmt::skip] - assert_eq!( - tbl.lines().collect::>(), - vec![ - "+-----+", - "| one |", - "+-----+", - "| 1 |", - "+-----+" - ] - ); - } - - #[tokio::test] - async fn show_explain() { - let ctx = SedonaContext::new(); - for limit in [None, Some(10)] { - let tbl = ctx - .sql("EXPLAIN SELECT 1 as one") - .await - .unwrap() - .show_sedona(&ctx, limit, DisplayTableOptions::default()) - .await - .unwrap(); - - #[rustfmt::skip] - assert_eq!( - tbl.lines().collect::>(), - vec![ - "+---------------+---------------------------------+", - "| plan_type | plan |", - "+---------------+---------------------------------+", - "| logical_plan | Projection: Int64(1) AS one |", - "| | EmptyRelation: rows=1 |", - "| physical_plan | ProjectionExec: expr=[1 as one] |", - "| | PlaceholderRowExec |", - "| | |", - "+---------------+---------------------------------+", - ] - ); - } - } - - #[tokio::test] - async fn write_geoparquet() { - let tmpdir = tempdir().unwrap(); - let tmp_parquet = tmpdir.path().join("tmp.parquet"); - let ctx = SedonaContext::new(); - ctx.sql("SELECT 1 as one") - .await - .unwrap() - .write_parquet( - &tmp_parquet.to_string_lossy(), - DataFrameWriteOptions::default(), - None, - ) - .await - .unwrap(); - } - #[tokio::test] async fn geoparquet_format() { // Make sure that our context can be set up to identify and read diff --git a/rust/sedona/src/dataframe.rs b/rust/sedona/src/dataframe.rs new file mode 100644 index 0000000000..0bb9ea1ae1 --- /dev/null +++ b/rust/sedona/src/dataframe.rs @@ -0,0 +1,289 @@ +// 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. + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use async_trait::async_trait; +use datafusion::{ + dataframe::DataFrameWriteOptions, + datasource::file_format::format_as_file_type, + error::{DataFusionError, Result}, + prelude::DataFrame, +}; +use datafusion_common::not_impl_err; +use datafusion_expr::{dml::InsertOp, LogicalPlan, LogicalPlanBuilder, SortExpr}; +use sedona_geoparquet::{format::GeoParquetFormatFactory, options::TableGeoParquetOptions}; + +use crate::{ + context::SedonaContext, + show::{show_batches, DisplayTableOptions}, +}; + +/// Sedona-specific [`DataFrame`] actions +/// +/// This trait, implemented for [`DataFrame`], extends the DataFrame API to make it +/// ergonomic to work with dataframes that contain geometry columns. Currently these +/// are limited to output functions, as geometry columns currently require special +/// handling when written or exported to an external system. +#[async_trait] +pub trait SedonaDataFrame { + /// Build a table of the first `limit` results in this DataFrame + /// + /// This will limit and execute the query and build a table using [show_batches]. + async fn show_sedona<'a>( + self, + ctx: &SedonaContext, + limit: Option, + options: DisplayTableOptions<'a>, + ) -> Result; + + async fn write_geoparquet( + self, + ctx: &SedonaContext, + path: &str, + options: SedonaWriteOptions, + writer_options: Option, + ) -> Result>; +} + +#[async_trait] +impl SedonaDataFrame for DataFrame { + async fn show_sedona<'a>( + self, + ctx: &SedonaContext, + limit: Option, + mut options: DisplayTableOptions<'a>, + ) -> Result { + let df = if matches!( + self.logical_plan(), + LogicalPlan::Explain(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::Analyze(_) + ) { + // Show multi-line output without truncation for plans like `EXPLAIN` + options.max_row_height = usize::MAX; + + // We don't want to apply an additional .limit() to plans like `Explain` + // as that will trigger an internal error: Unsupported logical plan: Explain must be root of the plan + self + } else { + // Apply limit if specified + self.limit(0, limit)? + }; + + let schema_without_qualifiers = df.schema().clone().strip_qualifiers(); + let schema = schema_without_qualifiers.as_arrow(); + let batches = df.collect().await?; + let mut out = Vec::new(); + show_batches(ctx, &mut out, schema, batches, options)?; + String::from_utf8(out).map_err(|e| DataFusionError::External(Box::new(e))) + } + + async fn write_geoparquet( + self, + ctx: &SedonaContext, + path: &str, + options: SedonaWriteOptions, + writer_options: Option, + ) -> Result, DataFusionError> { + if options.insert_op != InsertOp::Append { + return not_impl_err!( + "{} is not implemented for DataFrame::write_geoparquet.", + options.insert_op + ); + } + + let format = if let Some(parquet_opts) = writer_options { + Arc::new(GeoParquetFormatFactory::new_with_options(parquet_opts)) + } else { + Arc::new(GeoParquetFormatFactory::new()) + }; + + let file_type = format_as_file_type(format); + + let plan = if options.sort_by.is_empty() { + self.into_unoptimized_plan() + } else { + LogicalPlanBuilder::from(self.into_unoptimized_plan()) + .sort(options.sort_by)? + .build()? + }; + + let plan = LogicalPlanBuilder::copy_to( + plan, + path.into(), + file_type, + Default::default(), + options.partition_by, + )? + .build()?; + + DataFrame::new(ctx.ctx.state(), plan).collect().await + } +} + +/// A Sedona-specific copy of [DataFrameWriteOptions] +/// +/// This is needed because [DataFrameWriteOptions] has private fields, so we +/// can't use it in our interfaces. This object can be converted to a +/// [DataFrameWriteOptions] using `.into()`. +pub struct SedonaWriteOptions { + /// Controls how new data should be written to the table, determining whether + /// to append, overwrite, or replace existing data. + pub insert_op: InsertOp, + /// Controls if all partitions should be coalesced into a single output file + /// Generally will have slower performance when set to true. + pub single_file_output: bool, + /// Sets which columns should be used for hive-style partitioned writes by name. + /// Can be set to empty vec![] for non-partitioned writes. + pub partition_by: Vec, + /// Sets which columns should be used for sorting the output by name. + /// Can be set to empty vec![] for non-sorted writes. + pub sort_by: Vec, +} + +impl From for DataFrameWriteOptions { + fn from(value: SedonaWriteOptions) -> Self { + DataFrameWriteOptions::new() + .with_insert_operation(value.insert_op) + .with_single_file_output(value.single_file_output) + .with_partition_by(value.partition_by) + .with_sort_by(value.sort_by) + } +} + +impl SedonaWriteOptions { + /// Create a new SedonaWriteOptions with default values + pub fn new() -> Self { + SedonaWriteOptions { + insert_op: InsertOp::Append, + single_file_output: false, + partition_by: vec![], + sort_by: vec![], + } + } + + /// Set the insert operation + pub fn with_insert_operation(mut self, insert_op: InsertOp) -> Self { + self.insert_op = insert_op; + self + } + + /// Set the single_file_output value to true or false + pub fn with_single_file_output(mut self, single_file_output: bool) -> Self { + self.single_file_output = single_file_output; + self + } + + /// Sets the partition_by columns for output partitioning + pub fn with_partition_by(mut self, partition_by: Vec) -> Self { + self.partition_by = partition_by; + self + } + + /// Sets the sort_by columns for output sorting + pub fn with_sort_by(mut self, sort_by: Vec) -> Self { + self.sort_by = sort_by; + self + } +} + +impl Default for SedonaWriteOptions { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod test { + use tempfile::tempdir; + +use crate::context::SedonaContext; + + use super::*; + + + #[tokio::test] + async fn show() { + let ctx = SedonaContext::new(); + let tbl = ctx + .sql("SELECT 1 as one") + .await + .unwrap() + .show_sedona(&ctx, None, DisplayTableOptions::default()) + .await + .unwrap(); + + #[rustfmt::skip] + assert_eq!( + tbl.lines().collect::>(), + vec![ + "+-----+", + "| one |", + "+-----+", + "| 1 |", + "+-----+" + ] + ); + } + + #[tokio::test] + async fn show_explain() { + let ctx = SedonaContext::new(); + for limit in [None, Some(10)] { + let tbl = ctx + .sql("EXPLAIN SELECT 1 as one") + .await + .unwrap() + .show_sedona(&ctx, limit, DisplayTableOptions::default()) + .await + .unwrap(); + + #[rustfmt::skip] + assert_eq!( + tbl.lines().collect::>(), + vec![ + "+---------------+---------------------------------+", + "| plan_type | plan |", + "+---------------+---------------------------------+", + "| logical_plan | Projection: Int64(1) AS one |", + "| | EmptyRelation: rows=1 |", + "| physical_plan | ProjectionExec: expr=[1 as one] |", + "| | PlaceholderRowExec |", + "| | |", + "+---------------+---------------------------------+", + ] + ); + } + } + + #[tokio::test] + async fn write_geoparquet() { + let tmpdir = tempdir().unwrap(); + let tmp_parquet = tmpdir.path().join("tmp.parquet"); + let ctx = SedonaContext::new(); + ctx.sql("SELECT 1 as one") + .await + .unwrap() + .write_parquet( + &tmp_parquet.to_string_lossy(), + DataFrameWriteOptions::default(), + None, + ) + .await + .unwrap(); + } +} diff --git a/rust/sedona/src/lib.rs b/rust/sedona/src/lib.rs index 81c103f9f4..e29bbe2ecb 100644 --- a/rust/sedona/src/lib.rs +++ b/rust/sedona/src/lib.rs @@ -17,6 +17,7 @@ mod catalog; pub mod context; pub mod context_builder; +pub mod dataframe; mod exec; pub mod memory_pool; mod object_storage; From bc7d9c5a7271bcb707e385fab49e5f0b9b949d69 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 24 Jun 2026 12:01:07 -0500 Subject: [PATCH 07/55] formatting --- c/sedona-extension/src/execution_plan.rs | 5 +++-- c/sedona-extension/src/expr.rs | 10 ++-------- c/sedona-extension/src/lib.rs | 4 ++-- c/sedona-extension/src/sedona_extension.h | 2 +- rust/sedona/src/context.rs | 15 +++++++-------- rust/sedona/src/dataframe.rs | 3 +-- 6 files changed, 16 insertions(+), 23 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 0fb4f16766..3e01ca53d1 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -20,14 +20,15 @@ use std::{any::Any, sync::Arc}; use datafusion_common::{Result, Statistics}; use datafusion_execution::TaskContext; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, execution_plan::CardinalityEffect, metrics::MetricsSet, + execution_plan::CardinalityEffect, metrics::MetricsSet, DisplayAs, DisplayFormatType, + ExecutionPlan, PlanProperties, SendableRecordBatchStream, }; #[derive(Debug)] struct ImportedTableProviderExec { name: String, properties: Arc, - children: Vec> + children: Vec>, } impl DisplayAs for ImportedTableProviderExec { diff --git a/c/sedona-extension/src/expr.rs b/c/sedona-extension/src/expr.rs index 6694a392ef..7f1474f439 100644 --- a/c/sedona-extension/src/expr.rs +++ b/c/sedona-extension/src/expr.rs @@ -137,13 +137,7 @@ impl ImportedExpr { let mut out = FFI_ArrowArray::empty(); let mut err = SedonaCError::default(); - let result = get_property( - &self.inner, - property.as_ptr(), - args_ptr, - &mut out, - &mut err, - ); + let result = get_property(&self.inner, property.as_ptr(), args_ptr, &mut out, &mut err); if result != 0 { return Err(DataFusionError::Internal(format!( @@ -310,8 +304,8 @@ unsafe extern "C" fn exported_expr_release(self_: *mut SedonaCExpr) { mod tests { use super::*; use arrow_schema::{DataType, Field, Schema}; - use datafusion_physical_expr::expressions::Literal; use datafusion_common::ScalarValue; + use datafusion_physical_expr::expressions::Literal; #[test] fn test_physical_expr_with_schema_new() { diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs index 03757c5f46..6d019079c3 100644 --- a/c/sedona-extension/src/lib.rs +++ b/c/sedona-extension/src/lib.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -pub mod extension; +pub mod execution_plan; pub mod expr; +pub mod extension; pub mod scalar_kernel; -pub mod execution_plan; pub mod table_provider; diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 939208b9b4..a78d2f2452 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -267,7 +267,7 @@ struct SedonaCExecutionPlan { struct SedonaCExecutionPlanArgs* args, struct ArrowArrayStream* out, struct SedonaCError* err); - // Future implemenatation with async streams + // Future implementation with async streams int (*execute_async)(const struct SedonaCExecutionPlan* self, struct SedonaCExecutionPlanArgs* args, void* out, struct SedonaCError* err); diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index d2b2d55e28..526d727aa2 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -21,11 +21,9 @@ use std::{ use crate::exec::create_plan_from_sql; use crate::object_storage::ensure_object_store_registered_with_options; -use crate::{ - catalog::DynamicObjectStoreCatalog, - random_geometry_provider::RandomGeometryFunction, -}; +use crate::{catalog::DynamicObjectStoreCatalog, random_geometry_provider::RandomGeometryFunction}; use arrow_schema::DataType; +use datafusion::execution::memory_pool::MemoryLimit; use datafusion::{ common::plan_err, error::{DataFusionError, Result}, @@ -37,7 +35,6 @@ use datafusion::{ prelude::{DataFrame, SessionConfig, SessionContext}, sql::parser::{DFParser, Statement}, }; -use datafusion::{execution::memory_pool::MemoryLimit}; use datafusion_expr::sqlparser::dialect::{dialect_from_str, Dialect}; use parking_lot::Mutex; use sedona_common::{ @@ -539,12 +536,14 @@ impl ThreadSafeDialect { #[cfg(test)] mod tests { - use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, create_array}; + use arrow_array::{ + create_array, ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, + }; use arrow_schema::{DataType, Field, Schema}; use async_trait::async_trait; -use datafusion::assert_batches_eq; + use datafusion::assert_batches_eq; use datafusion_common::not_impl_err; -use sedona_datasource::spec::{Object, OpenReaderArgs}; + use sedona_datasource::spec::{Object, OpenReaderArgs}; use sedona_schema::{ crs::{deserialize_crs, lnglat}, datatypes::{Edges, SedonaType}, diff --git a/rust/sedona/src/dataframe.rs b/rust/sedona/src/dataframe.rs index 0bb9ea1ae1..5d21ed811a 100644 --- a/rust/sedona/src/dataframe.rs +++ b/rust/sedona/src/dataframe.rs @@ -211,11 +211,10 @@ impl Default for SedonaWriteOptions { mod test { use tempfile::tempdir; -use crate::context::SedonaContext; + use crate::context::SedonaContext; use super::*; - #[tokio::test] async fn show() { let ctx = SedonaContext::new(); From b5eb184eb01b5c2cb9e93dd7682bec6f018a46ca Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 24 Jun 2026 15:34:25 -0500 Subject: [PATCH 08/55] in decent shape --- Cargo.lock | 3 + c/sedona-extension/Cargo.toml | 5 + c/sedona-extension/src/execution_plan.rs | 376 ++++++++++++- c/sedona-extension/src/expr.rs | 45 +- c/sedona-extension/src/extension.rs | 120 +++- c/sedona-extension/src/lib.rs | 1 + c/sedona-extension/src/sedona_extension.h | 78 ++- c/sedona-extension/src/table_provider.rs | 642 ++++++++++++++++++++-- c/sedona-extension/src/utils.rs | 301 ++++++++++ 9 files changed, 1482 insertions(+), 89 deletions(-) create mode 100644 c/sedona-extension/src/utils.rs diff --git a/Cargo.lock b/Cargo.lock index deabb83025..dce2da1700 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5782,12 +5782,14 @@ dependencies = [ "arrow-array", "arrow-schema", "async-trait", + "datafusion", "datafusion-catalog", "datafusion-common", "datafusion-execution", "datafusion-expr", "datafusion-physical-expr", "datafusion-physical-plan", + "futures", "libc", "sedona-common", "sedona-expr", @@ -5796,6 +5798,7 @@ dependencies = [ "sedona-testing", "serde", "serde_json", + "tokio", ] [[package]] diff --git a/c/sedona-extension/Cargo.toml b/c/sedona-extension/Cargo.toml index 15c7109326..a416d4fb45 100644 --- a/c/sedona-extension/Cargo.toml +++ b/c/sedona-extension/Cargo.toml @@ -37,6 +37,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr = { workspace = true } datafusion-physical-plan = { workspace = true } +futures = { workspace = true } libc = "0.2.178" sedona-common = { workspace = true } sedona-expr = { workspace = true } @@ -45,3 +46,7 @@ sedona-schema = { workspace = true } sedona-testing = { path = "../../rust/sedona-testing" } serde = { workspace = true } serde_json = { workspace = true } +tokio = { workspace = true, features = ["rt-multi-thread"] } + +[dev-dependencies] +datafusion = { workspace = true, features = ["sql"] } diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 3e01ca53d1..b45f66efc4 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -15,29 +15,329 @@ // specific language governing permissions and limitations // under the License. -use std::{any::Any, sync::Arc}; +use std::{ + any::Any, + ffi::{c_int, c_void, CStr}, + fmt::Debug, + ptr::null_mut, + sync::Arc, +}; -use datafusion_common::{Result, Statistics}; +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_schema::{ffi::FFI_ArrowSchema, Schema, SchemaRef}; +use datafusion_common::{exec_err, Result, Statistics}; use datafusion_execution::TaskContext; use datafusion_physical_plan::{ execution_plan::CardinalityEffect, metrics::MetricsSet, DisplayAs, DisplayFormatType, - ExecutionPlan, PlanProperties, SendableRecordBatchStream, + ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream, +}; +use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; +use serde::{Deserialize, Serialize}; + +use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs}; +use crate::utils::{ + ffi_stream_to_sendable, get_plan_property, get_plan_string_property, StreamingRecordBatchReader, + ERRNO_OK, }; -#[derive(Debug)] -struct ImportedTableProviderExec { +/// Arguments for executing a partition of an execution plan. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecuteArgs { + pub partition: usize, +} + +/// Properties of an execution plan serialized across FFI. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PlanPropertiesArgs { + pub num_partitions: usize, + pub supports_limit_pushdown: bool, +} + +impl PlanPropertiesArgs { + /// Extract plan properties from a SedonaCExecutionPlan via FFI. + pub fn from_ffi_plan(plan: &SedonaCExecutionPlan) -> Result { + get_plan_property::(plan, "plan_properties", None) + } + + /// Convert to DataFusion PlanProperties. + pub fn into_plan_properties(self, schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + datafusion_physical_expr::EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(self.num_partitions), + datafusion_physical_plan::execution_plan::EmissionType::Incremental, + datafusion_physical_plan::execution_plan::Boundedness::Bounded, + ) + } +} + +/// Helper wrapper to format an ExecutionPlan with a specific DisplayFormatType. +struct DisplayAsWrapper<'a>(&'a Arc, DisplayFormatType); + +impl std::fmt::Display for DisplayAsWrapper<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt_as(self.1, f) + } +} + +/// Wrapper around an [ExecutionPlan] that can be exported across FFI. +pub struct ExportedExecutionPlan { + plan: Arc, + task_context: Arc, + runtime: tokio::runtime::Handle, +} + +impl Debug for ExportedExecutionPlan { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExportedExecutionPlan") + .field("plan", &self.plan) + .finish() + } +} + +impl ExportedExecutionPlan { + /// Create a new ExportedExecutionPlan from an ExecutionPlan. + pub fn new( + plan: Arc, + task_context: Arc, + runtime: tokio::runtime::Handle, + ) -> Self { + Self { + plan, + task_context, + runtime, + } + } + + fn schema(&self) -> SchemaRef { + self.plan.schema() + } + + fn get_property(&self, property: &str) -> Result> { + match property { + "plan_properties" => { + let props = PlanPropertiesArgs { + num_partitions: self + .plan + .properties() + .output_partitioning() + .partition_count(), + supports_limit_pushdown: self.plan.supports_limit_pushdown(), + }; + serde_json::to_vec(&props).map_err(|e| { + sedona_internal_datafusion_err!("Failed to serialize plan properties: {}", e) + }) + } + "debug_string" => Ok(format!("{:?}", self.plan).into_bytes()), + "display_default" => { + use std::fmt::Write; + let mut s = String::new(); + let _ = write!(s, "{}", DisplayAsWrapper(&self.plan, DisplayFormatType::Default)); + Ok(s.into_bytes()) + } + "display_verbose" => { + use std::fmt::Write; + let mut s = String::new(); + let _ = write!(s, "{}", DisplayAsWrapper(&self.plan, DisplayFormatType::Verbose)); + Ok(s.into_bytes()) + } + "display_tree_render" => { + use std::fmt::Write; + let mut s = String::new(); + let _ = write!(s, "{}", DisplayAsWrapper(&self.plan, DisplayFormatType::TreeRender)); + Ok(s.into_bytes()) + } + "name" => Ok(self.plan.name().as_bytes().to_vec()), + _ => exec_err!("Unknown property: {}", property), + } + } + + fn execute(&self, partition: usize) -> Result { + self.plan.execute(partition, self.task_context.clone()) + } +} + +impl From for SedonaCExecutionPlan { + fn from(value: ExportedExecutionPlan) -> Self { + let boxed = Box::new(value); + Self { + get_schema: Some(c_exec_plan_get_schema), + get_property_schema: None, + get_property: Some(c_exec_plan_get_property), + with_property: None, + execute: Some(c_exec_plan_execute), + execute_async: None, + reserved: null_mut(), + release: Some(c_exec_plan_release), + private_data: Box::into_raw(boxed) as *mut c_void, + } + } +} + +unsafe extern "C" fn c_exec_plan_get_schema( + self_: *const SedonaCExecutionPlan, + out: *mut FFI_ArrowSchema, +) { + let plan = &*((*self_).private_data as *const ExportedExecutionPlan); + let schema = plan.schema(); + if let Ok(ffi_schema) = FFI_ArrowSchema::try_from(schema.as_ref()) { + std::ptr::write(out, ffi_schema); + } +} + +unsafe extern "C" fn c_exec_plan_get_property( + self_: *const SedonaCExecutionPlan, + property: *const std::ffi::c_char, + _args: *mut SedonaCExecutionPlanArgs, + out: *mut arrow_array::ffi::FFI_ArrowArray, + err: *mut SedonaCError, +) -> c_int { + let plan = &*((*self_).private_data as *const ExportedExecutionPlan); + let property_str = CStr::from_ptr(property).to_string_lossy(); + + match plan.get_property(&property_str) { + Ok(bytes) => { + // Return the bytes as a single-element binary array + use arrow_array::{builder::BinaryBuilder, Array}; + let mut builder = BinaryBuilder::new(); + builder.append_value(&bytes); + let array = builder.finish(); + let ffi_array = arrow_array::ffi::FFI_ArrowArray::new(&array.to_data()); + std::ptr::write(out, ffi_array); + ERRNO_OK + } + Err(e) => { + *err = SedonaCError::new(&e.to_string()); + libc::EINVAL + } + } +} + +unsafe extern "C" fn c_exec_plan_execute( + self_: *const SedonaCExecutionPlan, + args: *mut SedonaCExecutionPlanArgs, + out: *mut FFI_ArrowArrayStream, + err: *mut SedonaCError, +) -> c_int { + let plan = &*((*self_).private_data as *const ExportedExecutionPlan); + + // Parse the execute args + let args_ref = &*args; + let args_slice = if args_ref.args.is_null() || args_ref.args_len == 0 { + &[] + } else { + std::slice::from_raw_parts(args_ref.args, args_ref.args_len) + }; + + let execute_args: ExecuteArgs = match serde_json::from_slice(args_slice) { + Ok(a) => a, + Err(e) => { + *err = SedonaCError::new(&format!("Failed to parse execute args: {}", e)); + return libc::EINVAL; + } + }; + + match plan.execute(execute_args.partition) { + Ok(stream) => { + // Create a streaming reader that polls the stream lazily + let reader = StreamingRecordBatchReader::new(stream, plan.runtime.clone()); + let ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); + std::ptr::write(out, ffi_stream); + ERRNO_OK + } + Err(e) => { + *err = SedonaCError::new(&e.to_string()); + libc::EINVAL + } + } +} + +unsafe extern "C" fn c_exec_plan_release(self_: *mut SedonaCExecutionPlan) { + if !(*self_).private_data.is_null() { + let _ = Box::from_raw((*self_).private_data as *mut ExportedExecutionPlan); + (*self_).private_data = null_mut(); + } + (*self_).release = None; +} + +/// An [ExecutionPlan] that wraps an imported [SedonaCExecutionPlan]. +/// +/// This allows plans to be imported from across an FFI boundary and executed +/// within DataFusion. +pub struct ImportedSedonaCExec { + inner: SedonaCExecutionPlan, + schema: SchemaRef, + properties: PlanProperties, + supports_limit_pushdown: bool, name: String, - properties: Arc, - children: Vec>, } -impl DisplayAs for ImportedTableProviderExec { +impl Debug for ImportedSedonaCExec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Try to get debug string from the FFI plan + if let Ok(debug_str) = self.get_debug_string() { + f.debug_struct("ImportedTableProviderExec") + .field("inner", &debug_str) + .finish() + } else { + f.debug_struct("ImportedTableProviderExec").finish() + } + } +} + +impl ImportedSedonaCExec { + /// Create a new ImportedTableProviderExec from a SedonaCExecutionPlan. + /// + /// This will query the plan for its schema and properties. + pub fn try_new(inner: SedonaCExecutionPlan) -> Result { + // Get schema + let Some(get_schema) = inner.get_schema else { + return sedona_internal_err!("SedonaCExecutionPlan does not have get_schema"); + }; + + let mut ffi_schema = FFI_ArrowSchema::empty(); + unsafe { get_schema(&inner, &mut ffi_schema) }; + let schema = Arc::new(Schema::try_from(&ffi_schema)?); + + // Get plan properties + let props_args = PlanPropertiesArgs::from_ffi_plan(&inner)?; + let supports_limit_pushdown = props_args.supports_limit_pushdown; + let properties = props_args.into_plan_properties(schema.clone()); + + // Get the inner plan's name + let inner_name = get_plan_string_property(&inner, "name").unwrap_or_default(); + let name = format!("ImportedSedonaCExec<{}>", inner_name); + + Ok(Self { + inner, + schema, + properties, + supports_limit_pushdown, + name, + }) + } + + fn get_debug_string(&self) -> Result { + get_plan_string_property(&self.inner, "debug_string") + } +} + +impl DisplayAs for ImportedSedonaCExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { - todo!() + let property = match t { + DisplayFormatType::Default => "display_default", + DisplayFormatType::Verbose => "display_verbose", + DisplayFormatType::TreeRender => "display_tree_render", + }; + + if let Ok(display_str) = get_plan_string_property(&self.inner, property) { + write!(f, "{}", display_str) + } else { + write!(f, "ImportedSedonaCExec") + } } } -impl ExecutionPlan for ImportedTableProviderExec { +impl ExecutionPlan for ImportedSedonaCExec { fn as_any(&self) -> &dyn Any { self } @@ -50,46 +350,80 @@ impl ExecutionPlan for ImportedTableProviderExec { &self.properties } - fn partition_statistics(&self, partition: Option) -> Result { - todo!() + fn partition_statistics(&self, _partition: Option) -> Result { + Ok(Statistics::new_unknown(&self.schema)) } fn cardinality_effect(&self) -> CardinalityEffect { - todo!() + CardinalityEffect::Unknown } fn maintains_input_order(&self) -> Vec { - todo!() + vec![] } fn metrics(&self) -> Option { - todo!() + None } fn supports_limit_pushdown(&self) -> bool { - todo!() + self.supports_limit_pushdown } fn statistics(&self) -> Result { - todo!() + Ok(Statistics::new_unknown(&self.schema)) } fn execute( &self, partition: usize, - context: Arc, + _context: Arc, ) -> Result { - todo!() + let Some(execute) = self.inner.execute else { + return sedona_internal_err!("SedonaCExecutionPlan does not have execute"); + }; + + let args = ExecuteArgs { partition }; + let args_bytes = serde_json::to_vec(&args).map_err(|e| { + sedona_internal_datafusion_err!("Failed to serialize execute args: {}", e) + })?; + + let mut ffi_args = SedonaCExecutionPlanArgs { + args: args_bytes.as_ptr(), + args_len: args_bytes.len(), + exec_plans: std::ptr::null(), + num_exec_plans: 0, + exprs: std::ptr::null(), + num_exprs: 0, + reserved: null_mut(), + }; + + let mut ffi_stream = FFI_ArrowArrayStream::empty(); + let mut err = SedonaCError::default(); + + let code = unsafe { execute(&self.inner, &mut ffi_args, &mut ffi_stream, &mut err) }; + + if code != ERRNO_OK { + return exec_err!("Failed to execute plan: {}", err); + } + + // Convert FFI stream to SendableRecordBatchStream + unsafe { ffi_stream_to_sendable(&mut ffi_stream) } } fn children(&self) -> Vec<&Arc> { - self.children.iter().collect() + vec![] } fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - todo!() + if !children.is_empty() { + return exec_err!("ImportedTableProviderExec does not support children"); + } + Ok(self) } } + + diff --git a/c/sedona-extension/src/expr.rs b/c/sedona-extension/src/expr.rs index 7f1474f439..cc25b283ed 100644 --- a/c/sedona-extension/src/expr.rs +++ b/c/sedona-extension/src/expr.rs @@ -26,10 +26,12 @@ use arrow_array::{builder::StringBuilder, ffi::FFI_ArrowArray, Array}; use arrow_schema::SchemaRef; use datafusion_common::{DataFusionError, Result}; use datafusion_physical_expr::PhysicalExpr; +use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; use sedona_expr::spatial_filter::SpatialFilterFactory; use serde::Deserialize; use crate::extension::{SedonaCError, SedonaCExpr}; +use crate::utils::ERRNO_OK; /// A wrapper around a DataFusion [PhysicalExpr] with its associated schema. /// @@ -78,9 +80,8 @@ impl PhysicalExprWithSchema { let factory = SpatialFilterFactory::default(); let spatial_filter = factory.try_from_expr(&self.expr)?; let bbox = spatial_filter.filter_bbox(column_name); - serde_json::to_string(&bbox).map_err(|e| { - DataFusionError::Internal(format!("Failed to serialize bounding box: {}", e)) - }) + serde_json::to_string(&bbox) + .map_err(|e| sedona_internal_datafusion_err!("Failed to serialize bounding box: {}", e)) } } @@ -110,8 +111,8 @@ impl TryFrom for ImportedExpr { fn try_from(value: SedonaCExpr) -> Result { match (&value.get_property, &value.release) { (Some(_), Some(_)) => Ok(Self { inner: value }), - _ => Err(DataFusionError::Internal( - "Can't import released or uninitialized SedonaCExpr".to_string(), + _ => Err(sedona_internal_datafusion_err!( + "Can't import released or uninitialized SedonaCExpr" )), } } @@ -129,9 +130,10 @@ impl ImportedExpr { property: &CStr, args: Option<&CStr>, ) -> Result { - let get_property = self.inner.get_property.ok_or_else(|| { - DataFusionError::Internal("get_property callback is null".to_string()) - })?; + let get_property = self + .inner + .get_property + .ok_or_else(|| sedona_internal_datafusion_err!("get_property callback is null"))?; let args_ptr = args.map(|a| a.as_ptr()).unwrap_or(std::ptr::null()); let mut out = FFI_ArrowArray::empty(); @@ -139,11 +141,8 @@ impl ImportedExpr { let result = get_property(&self.inner, property.as_ptr(), args_ptr, &mut out, &mut err); - if result != 0 { - return Err(DataFusionError::Internal(format!( - "get_property failed: {}", - err - ))); + if result != ERRNO_OK { + return sedona_internal_err!("get_property failed: {}", err); } Ok(out) @@ -189,7 +188,7 @@ unsafe extern "C" fn exported_expr_get_property_schema( if !err.is_null() { *err = SedonaCError::new("get_property_schema not implemented"); } - -1 + libc::EINVAL } unsafe extern "C" fn exported_expr_get_property( @@ -203,7 +202,7 @@ unsafe extern "C" fn exported_expr_get_property( if !err.is_null() { *err = SedonaCError::new("null self pointer"); } - return -1; + return libc::EINVAL; } let expr_arc = &*((*self_).private_data as *const Arc); @@ -212,7 +211,7 @@ unsafe extern "C" fn exported_expr_get_property( if !err.is_null() { *err = SedonaCError::new("null property pointer"); } - return -1; + return libc::EINVAL; } let property_str = match CStr::from_ptr(property).to_str() { @@ -221,7 +220,7 @@ unsafe extern "C" fn exported_expr_get_property( if !err.is_null() { *err = SedonaCError::new("invalid UTF-8 in property name"); } - return -1; + return libc::EINVAL; } }; @@ -232,7 +231,7 @@ unsafe extern "C" fn exported_expr_get_property( if !err.is_null() { *err = SedonaCError::new("bbox property requires args with column name"); } - return -1; + return libc::EINVAL; } let args_str = match CStr::from_ptr(args).to_str() { @@ -241,7 +240,7 @@ unsafe extern "C" fn exported_expr_get_property( if !err.is_null() { *err = SedonaCError::new("invalid UTF-8 in args"); } - return -1; + return libc::EINVAL; } }; @@ -251,7 +250,7 @@ unsafe extern "C" fn exported_expr_get_property( if !err.is_null() { *err = SedonaCError::new(&format!("failed to parse args: {}", e)); } - return -1; + return libc::EINVAL; } }; @@ -262,7 +261,7 @@ unsafe extern "C" fn exported_expr_get_property( if !err.is_null() { *err = SedonaCError::new(&format!("failed to extract bbox: {}", e)); } - return -1; + return libc::EINVAL; } }; @@ -273,13 +272,13 @@ unsafe extern "C" fn exported_expr_get_property( // Export the array via FFI std::ptr::write(out, FFI_ArrowArray::new(&array.to_data())); - 0 + ERRNO_OK } _ => { if !err.is_null() { *err = SedonaCError::new(&format!("unknown property: {}", property_str)); } - -1 + libc::EINVAL } } } diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs index 61c1cb54b3..309c2c2ab1 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -237,14 +237,21 @@ impl Drop for SedonaCExpr { } /// Raw FFI representation of the SedonaCExecutionPlanArgs +/// +/// This structure is passed to methods that need JSON-serialized arguments, +/// optional execution plans, and/or expressions. #[derive(Default)] #[repr(C)] pub struct SedonaCExecutionPlanArgs { + /// JSON-serialized arguments pub args: *const u8, pub args_len: usize, - pub children: *const *const SedonaCExecutionPlan, - pub num_children: usize, - pub expr: *const SedonaCExpr, + /// Optional array of execution plans + pub exec_plans: *const *const SedonaCExecutionPlan, + pub num_exec_plans: usize, + /// Optional array of expressions + pub exprs: *const *const SedonaCExpr, + pub num_exprs: usize, pub reserved: *mut c_void, } @@ -321,3 +328,110 @@ impl Drop for SedonaCExecutionPlan { } } } + +/// Raw FFI representation of a TableProvider. +/// +/// This provides a minimal interface for importing a TableProvider +/// across an FFI boundary in a version-agnostic manner. +#[derive(Default)] +#[repr(C)] +pub struct SedonaCTableProvider { + /// Get the schema of this table provider + pub get_schema: + Option, + + /// Get the schema of a property + pub get_property_schema: Option< + unsafe extern "C" fn( + self_: *const SedonaCTableProvider, + property: *const c_char, + out: *mut FFI_ArrowSchema, + err: *mut SedonaCError, + ) -> c_int, + >, + + /// Get a property value as an Arrow array + pub get_property: Option< + unsafe extern "C" fn( + self_: *const SedonaCTableProvider, + property: *const c_char, + args: *mut SedonaCExecutionPlanArgs, + out: *mut FFI_ArrowArray, + err: *mut SedonaCError, + ) -> c_int, + >, + + /// Perform a scan operation and return an execution plan + /// + /// The args parameter contains JSON-serialized scan arguments. + /// Returns an execution plan that can be used to read the data. + pub scan: Option< + unsafe extern "C" fn( + self_: *const SedonaCTableProvider, + args: *mut SedonaCExecutionPlanArgs, + out: *mut SedonaCExecutionPlan, + err: *mut SedonaCError, + ) -> c_int, + >, + + /// Perform an insert operation + /// + /// The args parameter contains JSON-serialized insert arguments. + /// The exec_plans field should contain the plan providing rows to insert. + /// Returns an execution plan that performs the insert. + pub insert: Option< + unsafe extern "C" fn( + self_: *const SedonaCTableProvider, + args: *mut SedonaCExecutionPlanArgs, + out: *mut SedonaCExecutionPlan, + err: *mut SedonaCError, + ) -> c_int, + >, + + /// Perform an update operation + /// + /// The args parameter contains JSON-serialized update arguments + /// (filters, column assignments, etc.). + /// Returns an execution plan that performs the update. + pub update: Option< + unsafe extern "C" fn( + self_: *const SedonaCTableProvider, + args: *mut SedonaCExecutionPlanArgs, + out: *mut SedonaCExecutionPlan, + err: *mut SedonaCError, + ) -> c_int, + >, + + /// Perform a delete operation + /// + /// The args parameter contains JSON-serialized delete arguments + /// (filters, etc.). + /// Returns an execution plan that performs the delete. + pub delete_rows: Option< + unsafe extern "C" fn( + self_: *const SedonaCTableProvider, + args: *mut SedonaCExecutionPlanArgs, + out: *mut SedonaCExecutionPlan, + err: *mut SedonaCError, + ) -> c_int, + >, + + pub reserved: *mut c_void, + + pub release: Option, + + pub private_data: *mut c_void, +} + +unsafe impl Send for SedonaCTableProvider {} +unsafe impl Sync for SedonaCTableProvider {} + +impl Drop for SedonaCTableProvider { + fn drop(&mut self) { + if let Some(releaser) = self.release { + unsafe { releaser(self) } + self.release = None; + self.private_data = null_mut(); + } + } +} diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs index 6d019079c3..2d101a1d1a 100644 --- a/c/sedona-extension/src/lib.rs +++ b/c/sedona-extension/src/lib.rs @@ -20,3 +20,4 @@ pub mod expr; pub mod extension; pub mod scalar_kernel; pub mod table_provider; +mod utils; diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index a78d2f2452..b87cf7e787 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -237,12 +237,20 @@ struct SedonaCExpr { struct SedonaCExecutionPlan; +/// \brief Arguments for execution plan and table provider operations +/// +/// This structure is passed to methods that need JSON-serialized arguments, +/// optional execution plans, and/or expressions. struct SedonaCExecutionPlanArgs { + /// JSON-serialized arguments const uint8_t* args; size_t args_len; - const struct SedonaCExecutionPlan** children; - size_t num_children; - const struct SedonaCExpr* expr; + /// Optional array of execution plans + const struct SedonaCExecutionPlan** exec_plans; + size_t num_exec_plans; + /// Optional array of expressions + const struct SedonaCExpr** exprs; + size_t num_exprs; void* reserved; }; @@ -267,7 +275,7 @@ struct SedonaCExecutionPlan { struct SedonaCExecutionPlanArgs* args, struct ArrowArrayStream* out, struct SedonaCError* err); - // Future implementation with async streams + // Future implementation with async streams (don't implement now) int (*execute_async)(const struct SedonaCExecutionPlan* self, struct SedonaCExecutionPlanArgs* args, void* out, struct SedonaCError* err); @@ -284,6 +292,68 @@ struct SedonaCExecutionPlan { void* private_data; }; +/// \brief ABI-stable table provider interface +/// +/// This provides a minimal interface for importing a table provider +/// across an FFI boundary in a version-agnostic manner. +struct SedonaCTableProvider { + /// Get the schema of this table provider + void (*get_schema)(const struct SedonaCTableProvider* self, struct ArrowSchema* out); + + // Extract some serializable property from this table provider + int (*get_property_schema)(const struct SedonaCTableProvider* self, + const char* property, struct ArrowSchema* out, + struct SedonaCError* err); + int (*get_property)(const struct SedonaCTableProvider* self, const char* property, + struct SedonaCExecutionPlanArgs* args, struct ArrowArray* out, + struct SedonaCError* err); + + /// Perform a scan operation and return an execution plan + /// + /// The args parameter contains JSON-serialized scan arguments. + /// Returns an execution plan that can be used to read the data. + int (*scan)(const struct SedonaCTableProvider* self, + struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, + struct SedonaCError* err); + + /// Perform an insert operation + /// + /// The args parameter contains JSON-serialized insert arguments. + /// The exec_plans field should contain the plan providing rows to insert. + /// Returns an execution plan that performs the insert. + int (*insert)(const struct SedonaCTableProvider* self, + struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, + struct SedonaCError* err); + + /// Perform an update operation + /// + /// The args parameter contains JSON-serialized update arguments + /// (filters, column assignments, etc.). + /// Returns an execution plan that performs the update. + int (*update)(const struct SedonaCTableProvider* self, + struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, + struct SedonaCError* err); + + /// Perform a delete operation + /// + /// The args parameter contains JSON-serialized delete arguments + /// (filters, etc.). + /// Returns an execution plan that performs the delete. + int (*delete_rows)(const struct SedonaCTableProvider* self, + struct SedonaCExecutionPlanArgs* args, + struct SedonaCExecutionPlan* out, struct SedonaCError* err); + + void* reserved; + + /// \brief Release this instance + /// + /// Implementations of this callback must set self->release to NULL. + void (*release)(struct SedonaCTableProvider* self); + + /// \brief Opaque implementation-specific data + void* private_data; +}; + #ifdef __cplusplus } #endif diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index a064c327b0..c208f6cbff 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -15,18 +15,266 @@ // specific language governing permissions and limitations // under the License. -use std::{any::Any, sync::Arc}; +use std::{ + any::Any, + ffi::{c_int, c_void, CStr}, + fmt::Debug, + ptr::null_mut, + sync::Arc, +}; -use arrow_schema::SchemaRef; +use arrow_array::ffi::FFI_ArrowArray; +use arrow_schema::{ffi::FFI_ArrowSchema, Schema, SchemaRef}; use async_trait::async_trait; -use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider}; -use datafusion_common::{Result, Statistics}; -use datafusion_expr::{dml::InsertOp, Expr, TableType}; +use datafusion_catalog::{Session, TableProvider}; +use datafusion_common::{exec_err, Result, Statistics}; +use datafusion_expr::{Expr, TableType}; use datafusion_physical_plan::ExecutionPlan; +use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; +use serde::{Deserialize, Serialize}; -#[derive(Debug)] -struct ImportedTableProvider { +use crate::execution_plan::{ExportedExecutionPlan, ImportedSedonaCExec}; +use crate::extension::{ + SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs, SedonaCTableProvider, +}; +use crate::utils::{get_table_provider_string_property, ERRNO_OK}; + +/// Arguments for a scan operation, serialized as JSON across FFI. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ScanArgs { + /// Column indices to project, or None for all columns. + pub projection: Option>, + /// Maximum number of rows to return, or None for unlimited. + pub limit: Option, +} + +/// A TableProvider wrapper that can be exported across FFI. +/// +/// This wraps an inner TableProvider and exposes it via the SedonaCTableProvider +/// FFI interface. +pub struct ExportedTableProvider { + inner: Arc, + session: Arc, + runtime: tokio::runtime::Handle, +} + +impl Debug for ExportedTableProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExportedTableProvider") + .field("inner", &self.inner) + .finish() + } +} + +impl ExportedTableProvider { + /// Create a new ExportedTableProvider from a TableProvider. + pub fn new( + inner: Arc, + session: Arc, + runtime: tokio::runtime::Handle, + ) -> Self { + Self { + inner, + session, + runtime, + } + } + + fn scan( + &self, + projection: Option>, + limit: Option, + ) -> Result> { + // Convert projection to the expected format + let projection_ref = projection.as_ref(); + + // Execute the scan - we use block_in_place to allow blocking within an async context + let inner = self.inner.clone(); + let session = self.session.as_ref(); + tokio::task::block_in_place(|| { + self.runtime + .block_on(async { inner.scan(session, projection_ref, &[], limit).await }) + }) + } + + fn get_property(&self, property: &str) -> Result> { + match property { + "table_type" => { + let table_type = self.inner.table_type(); + let type_str = match table_type { + TableType::Base => "Base", + TableType::View => "View", + TableType::Temporary => "Temporary", + }; + Ok(type_str.as_bytes().to_vec()) + } + _ => exec_err!("Unknown property: {}", property), + } + } +} + +impl From for SedonaCTableProvider { + fn from(value: ExportedTableProvider) -> Self { + let boxed = Box::new(value); + Self { + get_schema: Some(c_table_provider_get_schema), + get_property_schema: None, + get_property: Some(c_table_provider_get_property), + scan: Some(c_table_provider_scan), + insert: None, + update: None, + delete_rows: None, + reserved: null_mut(), + release: Some(c_table_provider_release), + private_data: Box::into_raw(boxed) as *mut c_void, + } + } +} + +unsafe extern "C" fn c_table_provider_get_schema( + self_: *const SedonaCTableProvider, + out: *mut FFI_ArrowSchema, +) { + let provider = &*((*self_).private_data as *const ExportedTableProvider); + let schema = provider.inner.schema(); + if let Ok(ffi_schema) = FFI_ArrowSchema::try_from(schema.as_ref()) { + std::ptr::write(out, ffi_schema); + } +} + +unsafe extern "C" fn c_table_provider_get_property( + self_: *const SedonaCTableProvider, + property: *const std::ffi::c_char, + _args: *mut SedonaCExecutionPlanArgs, + out: *mut FFI_ArrowArray, + err: *mut SedonaCError, +) -> c_int { + let provider = &*((*self_).private_data as *const ExportedTableProvider); + let property_str = CStr::from_ptr(property).to_string_lossy(); + + match provider.get_property(&property_str) { + Ok(bytes) => { + // Return the bytes as a single-element binary array + use arrow_array::{builder::BinaryBuilder, Array}; + let mut builder = BinaryBuilder::new(); + builder.append_value(&bytes); + let array = builder.finish(); + let ffi_array = FFI_ArrowArray::new(&array.to_data()); + std::ptr::write(out, ffi_array); + ERRNO_OK + } + Err(e) => { + *err = SedonaCError::new(&e.to_string()); + libc::EINVAL + } + } +} + +unsafe extern "C" fn c_table_provider_scan( + self_: *const SedonaCTableProvider, + args: *mut SedonaCExecutionPlanArgs, + out: *mut SedonaCExecutionPlan, + err: *mut SedonaCError, +) -> c_int { + let provider = &*((*self_).private_data as *const ExportedTableProvider); + + // Parse scan args + let args_ref = &*args; + let args_slice = if args_ref.args.is_null() || args_ref.args_len == 0 { + &[] + } else { + std::slice::from_raw_parts(args_ref.args, args_ref.args_len) + }; + + let scan_args: ScanArgs = if args_slice.is_empty() { + ScanArgs { + projection: None, + limit: None, + } + } else { + match serde_json::from_slice(args_slice) { + Ok(a) => a, + Err(e) => { + *err = SedonaCError::new(&format!("Failed to parse scan args: {}", e)); + return libc::EINVAL; + } + } + }; + + match provider.scan(scan_args.projection, scan_args.limit) { + Ok(plan) => { + let task_ctx = provider.session.task_ctx(); + let exported = ExportedExecutionPlan::new(plan, task_ctx, provider.runtime.clone()); + let ffi_plan: SedonaCExecutionPlan = exported.into(); + std::ptr::write(out, ffi_plan); + ERRNO_OK + } + Err(e) => { + *err = SedonaCError::new(&e.to_string()); + libc::EINVAL + } + } +} + +unsafe extern "C" fn c_table_provider_release(self_: *mut SedonaCTableProvider) { + if !(*self_).private_data.is_null() { + let _ = Box::from_raw((*self_).private_data as *mut ExportedTableProvider); + (*self_).private_data = null_mut(); + } + (*self_).release = None; +} + +/// A TableProvider that wraps an imported SedonaCTableProvider. +/// +/// This allows table providers from across an FFI boundary to be used +/// within DataFusion. +pub struct ImportedTableProvider { + inner: SedonaCTableProvider, schema: SchemaRef, + table_type: TableType, +} + +impl Debug for ImportedTableProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ImportedTableProvider").finish() + } +} + +impl ImportedTableProvider { + /// Create a new ImportedTableProvider from a SedonaCTableProvider. + pub fn try_new(inner: SedonaCTableProvider) -> Result { + // Get schema via Arrow C Data Interface + let Some(get_schema) = inner.get_schema else { + return sedona_internal_err!("SedonaCTableProvider does not have get_schema"); + }; + + let mut ffi_schema = FFI_ArrowSchema::empty(); + unsafe { get_schema(&inner, &mut ffi_schema) }; + let schema = Arc::new(Schema::try_from(&ffi_schema)?); + + // Get table type via get_property + let table_type = Self::get_table_type(&inner)?; + + Ok(Self { + inner, + schema, + table_type, + }) + } + + fn get_table_type(provider: &SedonaCTableProvider) -> Result { + let type_str = match get_table_provider_string_property(provider, "table_type") { + Ok(s) => s, + Err(_) => return Ok(TableType::Base), // Default to Base if property fetch fails + }; + + match type_str.as_str() { + "Base" => Ok(TableType::Base), + "View" => Ok(TableType::View), + "Temporary" => Ok(TableType::Temporary), + _ => Ok(TableType::Base), // Default to Base for unknown types + } + } } #[async_trait] @@ -40,54 +288,372 @@ impl TableProvider for ImportedTableProvider { } fn table_type(&self) -> TableType { - todo!() + self.table_type } fn statistics(&self) -> Option { - todo!() + None } async fn scan( &self, - state: &dyn Session, + _state: &dyn Session, projection: Option<&Vec>, - filters: &[Expr], + _filters: &[Expr], limit: Option, ) -> Result> { - todo!() + let Some(scan) = self.inner.scan else { + return sedona_internal_err!("SedonaCTableProvider does not have scan"); + }; + + let args = ScanArgs { + projection: projection.cloned(), + limit, + }; + let args_bytes = serde_json::to_vec(&args) + .map_err(|e| sedona_internal_datafusion_err!("Failed to serialize scan args: {}", e))?; + + let mut ffi_args = SedonaCExecutionPlanArgs { + args: args_bytes.as_ptr(), + args_len: args_bytes.len(), + exec_plans: std::ptr::null(), + num_exec_plans: 0, + exprs: std::ptr::null(), + num_exprs: 0, + reserved: null_mut(), + }; + + let mut ffi_plan = SedonaCExecutionPlan::default(); + let mut err = SedonaCError::default(); + + let code = unsafe { scan(&self.inner, &mut ffi_args, &mut ffi_plan, &mut err) }; + + if code != ERRNO_OK { + return exec_err!("Failed to scan table: {}", err); + } + + let exec = ImportedSedonaCExec::try_new(ffi_plan)?; + Ok(Arc::new(exec)) } +} - async fn scan_with_args<'a>( - &self, - state: &dyn Session, - args: ScanArgs<'a>, - ) -> Result { - todo!() +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{Float64Array, Int32Array, Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion::prelude::SessionContext; + use datafusion_common::assert_batches_eq; + + /// Create a SessionContext with a test table containing 5 numeric columns and multiple batches. + async fn create_test_context() -> Result { + let ctx = SessionContext::new(); + + // Create schema with 5 numeric columns + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value_a", DataType::Int64, false), + Field::new("value_b", DataType::Float64, false), + Field::new("value_c", DataType::Int32, false), + Field::new("value_d", DataType::Int64, false), + ])); + + // Create 5 batches with different data + let mut batches = Vec::new(); + for batch_num in 0..5 { + let offset = batch_num * 10; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![ + offset + 1, + offset + 2, + offset + 3, + offset + 4, + offset + 5, + ])), + Arc::new(Int64Array::from(vec![ + (offset + 1) as i64 * 100, + (offset + 2) as i64 * 100, + (offset + 3) as i64 * 100, + (offset + 4) as i64 * 100, + (offset + 5) as i64 * 100, + ])), + Arc::new(Float64Array::from(vec![ + (offset + 1) as f64 * 1.5, + (offset + 2) as f64 * 1.5, + (offset + 3) as f64 * 1.5, + (offset + 4) as f64 * 1.5, + (offset + 5) as f64 * 1.5, + ])), + Arc::new(Int32Array::from(vec![ + (offset + 1) * 2, + (offset + 2) * 2, + (offset + 3) * 2, + (offset + 4) * 2, + (offset + 5) * 2, + ])), + Arc::new(Int64Array::from(vec![ + (offset + 1) as i64 * 1000, + (offset + 2) as i64 * 1000, + (offset + 3) as i64 * 1000, + (offset + 4) as i64 * 1000, + (offset + 5) as i64 * 1000, + ])), + ], + )?; + batches.push(batch); + } + + // Use a MemTable with all batches + let provider = datafusion::datasource::MemTable::try_new(schema, vec![batches])?; + ctx.register_table("test_data", Arc::new(provider))?; + + Ok(ctx) } - async fn insert_into( - &self, - _state: &dyn Session, - _input: Arc, - _insert_op: InsertOp, - ) -> Result> { - todo!() + #[tokio::test(flavor = "multi_thread")] + async fn test_roundtrip_simple_select() -> Result<()> { + let ctx = create_test_context().await?; + + // Get the table provider from the context + let table = ctx.table_provider("test_data").await?; + + // Export the table provider + let runtime = tokio::runtime::Handle::current(); + let session = Arc::new(ctx.state()); + let exported = ExportedTableProvider::new(table, session, runtime); + let ffi_provider: SedonaCTableProvider = exported.into(); + + // Import the table provider + let imported = ImportedTableProvider::try_new(ffi_provider)?; + + // Create a new context and register the imported table + let ctx2 = SessionContext::new(); + ctx2.register_table("imported_data", Arc::new(imported))?; + + // Query and verify + let result = ctx2 + .sql("SELECT id, value_a FROM imported_data ORDER BY id LIMIT 5") + .await? + .collect() + .await?; + + let expected = [ + "+----+---------+", + "| id | value_a |", + "+----+---------+", + "| 1 | 100 |", + "| 2 | 200 |", + "| 3 | 300 |", + "| 4 | 400 |", + "| 5 | 500 |", + "+----+---------+", + ]; + + assert_batches_eq!(expected, &result); + Ok(()) } - async fn delete_from( - &self, - _state: &dyn Session, - _filters: Vec, - ) -> Result> { - todo!() + #[tokio::test(flavor = "multi_thread")] + async fn test_roundtrip_projection() -> Result<()> { + let ctx = create_test_context().await?; + + // Get the table provider from the context + let table = ctx.table_provider("test_data").await?; + + // Export the table provider + let runtime = tokio::runtime::Handle::current(); + let session = Arc::new(ctx.state()); + let exported = ExportedTableProvider::new(table, session, runtime); + let ffi_provider: SedonaCTableProvider = exported.into(); + + // Import the table provider + let imported = ImportedTableProvider::try_new(ffi_provider)?; + + // Create a new context and register the imported table + let ctx2 = SessionContext::new(); + ctx2.register_table("imported_data", Arc::new(imported))?; + + // Test projection with only specific columns + let result = ctx2 + .sql("SELECT value_b, value_d FROM imported_data ORDER BY value_b LIMIT 3") + .await? + .collect() + .await?; + + let expected = [ + "+---------+---------+", + "| value_b | value_d |", + "+---------+---------+", + "| 1.5 | 1000 |", + "| 3.0 | 2000 |", + "| 4.5 | 3000 |", + "+---------+---------+", + ]; + + assert_batches_eq!(expected, &result); + Ok(()) } - async fn update( - &self, - _state: &dyn Session, - _assignments: Vec<(String, Expr)>, - _filters: Vec, - ) -> Result> { - todo!() + #[tokio::test(flavor = "multi_thread")] + async fn test_roundtrip_filter() -> Result<()> { + let ctx = create_test_context().await?; + + // Get the table provider from the context + let table = ctx.table_provider("test_data").await?; + + // Export the table provider + let runtime = tokio::runtime::Handle::current(); + let session = Arc::new(ctx.state()); + let exported = ExportedTableProvider::new(table, session, runtime); + let ffi_provider: SedonaCTableProvider = exported.into(); + + // Import the table provider + let imported = ImportedTableProvider::try_new(ffi_provider)?; + + // Create a new context and register the imported table + let ctx2 = SessionContext::new(); + ctx2.register_table("imported_data", Arc::new(imported))?; + + // Test filter (filter is applied on the DataFusion side, not pushed to FFI yet) + let result = ctx2 + .sql("SELECT id, value_a FROM imported_data WHERE id > 20 ORDER BY id LIMIT 5") + .await? + .collect() + .await?; + + let expected = [ + "+----+---------+", + "| id | value_a |", + "+----+---------+", + "| 21 | 2100 |", + "| 22 | 2200 |", + "| 23 | 2300 |", + "| 24 | 2400 |", + "| 25 | 2500 |", + "+----+---------+", + ]; + + assert_batches_eq!(expected, &result); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_roundtrip_sort() -> Result<()> { + let ctx = create_test_context().await?; + + // Get the table provider from the context + let table = ctx.table_provider("test_data").await?; + + // Export the table provider + let runtime = tokio::runtime::Handle::current(); + let session = Arc::new(ctx.state()); + let exported = ExportedTableProvider::new(table, session, runtime); + let ffi_provider: SedonaCTableProvider = exported.into(); + + // Import the table provider + let imported = ImportedTableProvider::try_new(ffi_provider)?; + + // Create a new context and register the imported table + let ctx2 = SessionContext::new(); + ctx2.register_table("imported_data", Arc::new(imported))?; + + // Test sorting in descending order + let result = ctx2 + .sql("SELECT id, value_c FROM imported_data ORDER BY id DESC LIMIT 5") + .await? + .collect() + .await?; + + let expected = [ + "+----+---------+", + "| id | value_c |", + "+----+---------+", + "| 45 | 90 |", + "| 44 | 88 |", + "| 43 | 86 |", + "| 42 | 84 |", + "| 41 | 82 |", + "+----+---------+", + ]; + + assert_batches_eq!(expected, &result); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_roundtrip_limit() -> Result<()> { + let ctx = create_test_context().await?; + + // Get the table provider from the context + let table = ctx.table_provider("test_data").await?; + + // Export the table provider + let runtime = tokio::runtime::Handle::current(); + let session = Arc::new(ctx.state()); + let exported = ExportedTableProvider::new(table, session, runtime); + let ffi_provider: SedonaCTableProvider = exported.into(); + + // Import the table provider + let imported = ImportedTableProvider::try_new(ffi_provider)?; + + // Create a new context and register the imported table + let ctx2 = SessionContext::new(); + ctx2.register_table("imported_data", Arc::new(imported))?; + + // Test limit + let result = ctx2 + .sql("SELECT id FROM imported_data ORDER BY id LIMIT 3") + .await? + .collect() + .await?; + + let expected = [ + "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 3 |", "+----+", + ]; + + assert_batches_eq!(expected, &result); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_roundtrip_all_columns() -> Result<()> { + let ctx = create_test_context().await?; + + // Get the table provider from the context + let table = ctx.table_provider("test_data").await?; + + // Export the table provider + let runtime = tokio::runtime::Handle::current(); + let session = Arc::new(ctx.state()); + let exported = ExportedTableProvider::new(table, session, runtime); + let ffi_provider: SedonaCTableProvider = exported.into(); + + // Import the table provider + let imported = ImportedTableProvider::try_new(ffi_provider)?; + + // Create a new context and register the imported table + let ctx2 = SessionContext::new(); + ctx2.register_table("imported_data", Arc::new(imported))?; + + // Test selecting all columns + let result = ctx2 + .sql("SELECT * FROM imported_data ORDER BY id LIMIT 2") + .await? + .collect() + .await?; + + let expected = [ + "+----+---------+---------+---------+---------+", + "| id | value_a | value_b | value_c | value_d |", + "+----+---------+---------+---------+---------+", + "| 1 | 100 | 1.5 | 2 | 1000 |", + "| 2 | 200 | 3.0 | 4 | 2000 |", + "+----+---------+---------+---------+---------+", + ]; + + assert_batches_eq!(expected, &result); + Ok(()) } } diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs new file mode 100644 index 0000000000..bb27eb2013 --- /dev/null +++ b/c/sedona-extension/src/utils.rs @@ -0,0 +1,301 @@ +// 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. + +//! Utilities for converting between Arrow FFI streams and DataFusion streams. + +use std::ffi::{c_int, CString}; +use std::ptr::null_mut; +use std::sync::Mutex; + +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::{RecordBatch, RecordBatchReader}; +use arrow_schema::{ArrowError, SchemaRef}; +use datafusion_common::Result; +use datafusion_physical_plan::SendableRecordBatchStream; +use futures::StreamExt; +use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; +use serde::de::DeserializeOwned; +use serde::Serialize; + +use crate::extension::{ + SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs, SedonaCTableProvider, +}; + +/// Success return code for FFI functions. +pub const ERRNO_OK: c_int = 0; + +/// Call `get_property` on a [SedonaCExecutionPlan] and deserialize the result. +/// +/// This handles the common pattern of: +/// 1. Extracting the get_property function pointer +/// 2. Calling it with a property name and optional serializable args +/// 3. Parsing the binary array result +/// 4. Deserializing from JSON to the target type +/// +/// # Arguments +/// +/// * `plan` - The execution plan to query +/// * `property` - The property name to retrieve +/// * `args` - Optional arguments to pass (will be serialized to JSON bytes) +/// +/// # Errors +/// +/// Returns an error if: +/// - The plan does not have a `get_property` callback +/// - The FFI call fails +/// - The result cannot be parsed as a binary array +/// - The JSON deserialization fails +pub fn get_plan_property( + plan: &SedonaCExecutionPlan, + property: &str, + args: Option<&A>, +) -> Result +where + T: DeserializeOwned, + A: Serialize, +{ + let Some(get_property) = plan.get_property else { + return sedona_internal_err!("SedonaCExecutionPlan does not have get_property"); + }; + + let property_cstr = CString::new(property) + .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; + + // Serialize args if provided + let args_bytes = match args { + Some(a) => serde_json::to_vec(a) + .map_err(|e| sedona_internal_datafusion_err!("Failed to serialize args: {}", e))?, + None => Vec::new(), + }; + + let mut ffi_args = SedonaCExecutionPlanArgs { + args: if args_bytes.is_empty() { + std::ptr::null() + } else { + args_bytes.as_ptr() + }, + args_len: args_bytes.len(), + exec_plans: std::ptr::null(), + num_exec_plans: 0, + exprs: std::ptr::null(), + num_exprs: 0, + reserved: null_mut(), + }; + + let mut ffi_array = arrow_array::ffi::FFI_ArrowArray::empty(); + let mut err = SedonaCError::default(); + + let code = unsafe { + get_property( + plan, + property_cstr.as_ptr(), + &mut ffi_args, + &mut ffi_array, + &mut err, + ) + }; + + if code != ERRNO_OK { + return sedona_internal_err!("Failed to get property '{}': {}", property, err); + } + + // Parse the binary array to get the JSON bytes + parse_binary_ffi_array(ffi_array) +} + +/// Parse a binary FFI array containing JSON and deserialize to the target type. +fn parse_binary_ffi_array( + ffi_array: arrow_array::ffi::FFI_ArrowArray, +) -> Result { + let bytes = parse_binary_ffi_array_to_bytes(ffi_array)?; + serde_json::from_slice::(&bytes) + .map_err(|e| sedona_internal_datafusion_err!("Failed to deserialize property: {}", e)) +} + +/// Parse a binary FFI array and return the raw bytes. +fn parse_binary_ffi_array_to_bytes( + ffi_array: arrow_array::ffi::FFI_ArrowArray, +) -> Result> { + let data = unsafe { + arrow_array::ffi::from_ffi_and_data_type(ffi_array, arrow_schema::DataType::Binary)? + }; + let array = arrow_array::make_array(data); + let binary_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + sedona_internal_datafusion_err!("Expected binary array from get_property") + })?; + + Ok(binary_array.value(0).to_vec()) +} + +/// Get a string property from a [SedonaCExecutionPlan]. +/// +/// This is a convenience wrapper around [get_plan_property] for string values. +pub fn get_plan_string_property(plan: &SedonaCExecutionPlan, property: &str) -> Result { + let Some(get_property) = plan.get_property else { + return sedona_internal_err!("SedonaCExecutionPlan does not have get_property"); + }; + + let property_cstr = CString::new(property) + .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; + + let mut ffi_args = SedonaCExecutionPlanArgs { + args: std::ptr::null(), + args_len: 0, + exec_plans: std::ptr::null(), + num_exec_plans: 0, + exprs: std::ptr::null(), + num_exprs: 0, + reserved: null_mut(), + }; + + let mut ffi_array = arrow_array::ffi::FFI_ArrowArray::empty(); + let mut err = SedonaCError::default(); + + let code = unsafe { + get_property( + plan, + property_cstr.as_ptr(), + &mut ffi_args, + &mut ffi_array, + &mut err, + ) + }; + + if code != ERRNO_OK { + return sedona_internal_err!("Failed to get '{}': {}", property, err); + } + + let bytes = parse_binary_ffi_array_to_bytes(ffi_array)?; + String::from_utf8(bytes) + .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) +} + +/// Get a string property from a [SedonaCTableProvider]. +pub fn get_table_provider_string_property( + provider: &SedonaCTableProvider, + property: &str, +) -> Result { + let Some(get_property) = provider.get_property else { + return sedona_internal_err!("SedonaCTableProvider does not have get_property"); + }; + + let property_cstr = CString::new(property) + .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; + + let mut ffi_args = SedonaCExecutionPlanArgs { + args: std::ptr::null(), + args_len: 0, + exec_plans: std::ptr::null(), + num_exec_plans: 0, + exprs: std::ptr::null(), + num_exprs: 0, + reserved: null_mut(), + }; + + let mut ffi_array = arrow_array::ffi::FFI_ArrowArray::empty(); + let mut err = SedonaCError::default(); + + let code = unsafe { + get_property( + provider, + property_cstr.as_ptr(), + &mut ffi_args, + &mut ffi_array, + &mut err, + ) + }; + + if code != ERRNO_OK { + return sedona_internal_err!("Failed to get '{}': {}", property, err); + } + + let bytes = parse_binary_ffi_array_to_bytes(ffi_array)?; + String::from_utf8(bytes) + .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) +} + +/// A RecordBatchReader that lazily polls a SendableRecordBatchStream. +/// +/// This allows exporting a DataFusion stream over FFI without collecting all +/// batches upfront. Used when exporting execution plans across FFI boundaries. +pub struct StreamingRecordBatchReader { + schema: SchemaRef, + stream: Mutex, + runtime: tokio::runtime::Handle, +} + +impl StreamingRecordBatchReader { + /// Create a new StreamingRecordBatchReader from a SendableRecordBatchStream. + pub fn new(stream: SendableRecordBatchStream, runtime: tokio::runtime::Handle) -> Self { + Self { + schema: stream.schema(), + stream: Mutex::new(stream), + runtime, + } + } +} + +impl Iterator for StreamingRecordBatchReader { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + let mut stream = self.stream.lock().unwrap(); + tokio::task::block_in_place(|| { + self.runtime.block_on(async { + match stream.next().await { + Some(Ok(batch)) => Some(Ok(batch)), + Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), + None => None, + } + }) + }) + } +} + +impl RecordBatchReader for StreamingRecordBatchReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +/// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream. +/// +/// This is the inverse of StreamingRecordBatchReader - it takes an FFI stream +/// and converts it back to a DataFusion stream. Used when importing execution +/// plans from across FFI boundaries. +/// +/// # Safety +/// +/// The caller must ensure that the FFI stream pointer is valid and properly +/// initialized. +pub unsafe fn ffi_stream_to_sendable( + ffi_stream: &mut FFI_ArrowArrayStream, +) -> Result { + let reader = arrow_array::ffi_stream::ArrowArrayStreamReader::from_raw(ffi_stream)?; + + let schema = reader.schema(); + let stream = futures::stream::iter(reader).map(|result| { + result.map_err(|e| datafusion_common::DataFusionError::ArrowError(Box::new(e), None)) + }); + + Ok(Box::pin( + datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), + )) +} From 205a804390dd1e80e45b151834f1c048c2325d8d Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 24 Jun 2026 16:26:44 -0500 Subject: [PATCH 09/55] more cases --- c/sedona-extension/src/execution_plan.rs | 630 ++++++++++++++++++++--- c/sedona-extension/src/table_provider.rs | 37 +- c/sedona-extension/src/utils.rs | 148 +++++- 3 files changed, 715 insertions(+), 100 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index b45f66efc4..9eef7683ae 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -18,7 +18,7 @@ use std::{ any::Any, ffi::{c_int, c_void, CStr}, - fmt::Debug, + fmt::{Debug, Display, Formatter}, ptr::null_mut, sync::Arc, }; @@ -28,57 +28,20 @@ use arrow_schema::{ffi::FFI_ArrowSchema, Schema, SchemaRef}; use datafusion_common::{exec_err, Result, Statistics}; use datafusion_execution::TaskContext; use datafusion_physical_plan::{ - execution_plan::CardinalityEffect, metrics::MetricsSet, DisplayAs, DisplayFormatType, - ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream, + execution_plan::{Boundedness, CardinalityEffect, EmissionType}, + metrics::{CustomMetricValue, Metric, MetricValue, MetricsSet}, + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + SendableRecordBatchStream, }; use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; use serde::{Deserialize, Serialize}; use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs}; use crate::utils::{ - ffi_stream_to_sendable, get_plan_property, get_plan_string_property, StreamingRecordBatchReader, - ERRNO_OK, + ffi_stream_to_sendable, get_plan_property, get_plan_string_property, + StreamingRecordBatchReader, ERRNO_OK, }; -/// Arguments for executing a partition of an execution plan. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ExecuteArgs { - pub partition: usize, -} - -/// Properties of an execution plan serialized across FFI. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct PlanPropertiesArgs { - pub num_partitions: usize, - pub supports_limit_pushdown: bool, -} - -impl PlanPropertiesArgs { - /// Extract plan properties from a SedonaCExecutionPlan via FFI. - pub fn from_ffi_plan(plan: &SedonaCExecutionPlan) -> Result { - get_plan_property::(plan, "plan_properties", None) - } - - /// Convert to DataFusion PlanProperties. - pub fn into_plan_properties(self, schema: SchemaRef) -> PlanProperties { - PlanProperties::new( - datafusion_physical_expr::EquivalenceProperties::new(schema), - Partitioning::UnknownPartitioning(self.num_partitions), - datafusion_physical_plan::execution_plan::EmissionType::Incremental, - datafusion_physical_plan::execution_plan::Boundedness::Bounded, - ) - } -} - -/// Helper wrapper to format an ExecutionPlan with a specific DisplayFormatType. -struct DisplayAsWrapper<'a>(&'a Arc, DisplayFormatType); - -impl std::fmt::Display for DisplayAsWrapper<'_> { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - self.0.fmt_as(self.1, f) - } -} - /// Wrapper around an [ExecutionPlan] that can be exported across FFI. pub struct ExportedExecutionPlan { plan: Arc, @@ -112,41 +75,75 @@ impl ExportedExecutionPlan { self.plan.schema() } - fn get_property(&self, property: &str) -> Result> { + fn get_property(&self, property: &str) -> Result { match property { "plan_properties" => { - let props = PlanPropertiesArgs { - num_partitions: self - .plan - .properties() - .output_partitioning() - .partition_count(), - supports_limit_pushdown: self.plan.supports_limit_pushdown(), - }; - serde_json::to_vec(&props).map_err(|e| { + let props = PlanPropertiesArgs::from_plan(self.plan.as_ref()); + serde_json::to_string(&props).map_err(|e| { sedona_internal_datafusion_err!("Failed to serialize plan properties: {}", e) }) } - "debug_string" => Ok(format!("{:?}", self.plan).into_bytes()), + "debug_string" => Ok(format!("{:?}", self.plan)), "display_default" => { use std::fmt::Write; let mut s = String::new(); - let _ = write!(s, "{}", DisplayAsWrapper(&self.plan, DisplayFormatType::Default)); - Ok(s.into_bytes()) + let _ = write!( + s, + "{}", + DisplayAsWrapper(&self.plan, DisplayFormatType::Default) + ); + Ok(s) } "display_verbose" => { use std::fmt::Write; let mut s = String::new(); - let _ = write!(s, "{}", DisplayAsWrapper(&self.plan, DisplayFormatType::Verbose)); - Ok(s.into_bytes()) + let _ = write!( + s, + "{}", + DisplayAsWrapper(&self.plan, DisplayFormatType::Verbose) + ); + Ok(s) } "display_tree_render" => { use std::fmt::Write; let mut s = String::new(); - let _ = write!(s, "{}", DisplayAsWrapper(&self.plan, DisplayFormatType::TreeRender)); - Ok(s.into_bytes()) + let _ = write!( + s, + "{}", + DisplayAsWrapper(&self.plan, DisplayFormatType::TreeRender) + ); + Ok(s) + } + "name" => Ok(self.plan.name().to_string()), + "cardinality_effect" => { + let effect = match self.plan.cardinality_effect() { + CardinalityEffect::Unknown => "Unknown", + CardinalityEffect::Equal => "Equal", + CardinalityEffect::LowerEqual => "LowerEqual", + CardinalityEffect::GreaterEqual => "GreaterEqual", + }; + Ok(effect.to_string()) + } + "maintains_input_order" => { + let order = self.plan.maintains_input_order(); + serde_json::to_string(&order).map_err(|e| { + sedona_internal_datafusion_err!( + "Failed to serialize maintains_input_order: {}", + e + ) + }) + } + "metrics" => { + // Serialize metrics as JSON with aggregated values + if let Some(metrics) = self.plan.metrics() { + let serialized = SerializedMetrics::from_metrics_set(&metrics); + serde_json::to_string(&serialized).map_err(|e| { + sedona_internal_datafusion_err!("Failed to serialize metrics: {}", e) + }) + } else { + Ok("null".to_string()) + } } - "name" => Ok(self.plan.name().as_bytes().to_vec()), _ => exec_err!("Unknown property: {}", property), } } @@ -161,7 +158,7 @@ impl From for SedonaCExecutionPlan { let boxed = Box::new(value); Self { get_schema: Some(c_exec_plan_get_schema), - get_property_schema: None, + get_property_schema: Some(c_exec_plan_get_property_schema), get_property: Some(c_exec_plan_get_property), with_property: None, execute: Some(c_exec_plan_execute), @@ -184,6 +181,27 @@ unsafe extern "C" fn c_exec_plan_get_schema( } } +unsafe extern "C" fn c_exec_plan_get_property_schema( + _self_: *const SedonaCExecutionPlan, + _property: *const std::ffi::c_char, + out: *mut FFI_ArrowSchema, + err: *mut SedonaCError, +) -> c_int { + // All properties are returned as Utf8 strings (including JSON) + use arrow_schema::{DataType, Field}; + let field = Field::new("value", DataType::Utf8, false); + match FFI_ArrowSchema::try_from(&field) { + Ok(ffi_schema) => { + std::ptr::write(out, ffi_schema); + ERRNO_OK + } + Err(e) => { + *err = SedonaCError::new(&format!("Failed to convert field to FFI schema: {}", e)); + libc::EINVAL + } + } +} + unsafe extern "C" fn c_exec_plan_get_property( self_: *const SedonaCExecutionPlan, property: *const std::ffi::c_char, @@ -195,11 +213,11 @@ unsafe extern "C" fn c_exec_plan_get_property( let property_str = CStr::from_ptr(property).to_string_lossy(); match plan.get_property(&property_str) { - Ok(bytes) => { - // Return the bytes as a single-element binary array - use arrow_array::{builder::BinaryBuilder, Array}; - let mut builder = BinaryBuilder::new(); - builder.append_value(&bytes); + Ok(value) => { + // Return the string as a single-element string array + use arrow_array::{builder::StringBuilder, Array}; + let mut builder = StringBuilder::new(); + builder.append_value(&value); let array = builder.finish(); let ffi_array = arrow_array::ffi::FFI_ArrowArray::new(&array.to_data()); std::ptr::write(out, ffi_array); @@ -355,15 +373,27 @@ impl ExecutionPlan for ImportedSedonaCExec { } fn cardinality_effect(&self) -> CardinalityEffect { - CardinalityEffect::Unknown + get_plan_string_property(&self.inner, "cardinality_effect") + .ok() + .and_then(|s| match s.as_str() { + "Equal" => Some(CardinalityEffect::Equal), + "LowerEqual" => Some(CardinalityEffect::LowerEqual), + "GreaterEqual" => Some(CardinalityEffect::GreaterEqual), + _ => None, + }) + .unwrap_or(CardinalityEffect::Unknown) } fn maintains_input_order(&self) -> Vec { - vec![] + get_plan_property::, ()>(&self.inner, "maintains_input_order", None) + .unwrap_or_default() } fn metrics(&self) -> Option { - None + get_plan_property::, ()>(&self.inner, "metrics", None) + .ok() + .flatten() + .map(|sm| sm.into_metrics_set()) } fn supports_limit_pushdown(&self) -> bool { @@ -426,4 +456,462 @@ impl ExecutionPlan for ImportedSedonaCExec { } } +/// Arguments for executing a partition of an execution plan. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecuteArgs { + pub partition: usize, +} + +/// Metrics serialized for FFI transfer. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SerializedMetrics { + pub display: String, +} + +impl SerializedMetrics { + /// Create from a MetricsSet by capturing its display string. + pub fn from_metrics_set(set: &MetricsSet) -> Self { + Self { + display: set.to_string(), + } + } + + /// Convert back to a MetricsSet with a custom metric containing the display string. + pub fn into_metrics_set(self) -> MetricsSet { + let mut set = MetricsSet::new(); + let custom = ImportedMetrics::new(self.display); + let metric = Metric::new( + MetricValue::Custom { + name: "imported_metrics".into(), + value: Arc::new(custom), + }, + None, + ); + set.push(Arc::new(metric)); + set + } +} + +/// Custom metric value that holds imported metrics display string. +#[derive(Debug, Clone)] +pub struct ImportedMetrics { + display: String, +} + +impl ImportedMetrics { + pub fn new(display: String) -> Self { + Self { display } + } +} + +impl Display for ImportedMetrics { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.display) + } +} + +impl CustomMetricValue for ImportedMetrics { + fn new_empty(&self) -> Arc { + Arc::new(Self { + display: String::new(), + }) + } + + fn aggregate(&self, other: Arc) { + // No aggregation for imported metrics - they're read-only snapshots + let _ = other; + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn is_eq(&self, other: &Arc) -> bool { + other + .as_any() + .downcast_ref::() + .is_some_and(|o| o.display == self.display) + } +} + +/// Properties of an execution plan serialized across FFI. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PlanPropertiesArgs { + pub num_partitions: usize, + pub supports_limit_pushdown: bool, + pub emission_type: String, + pub boundedness: String, +} + +impl PlanPropertiesArgs { + /// Create from an ExecutionPlan. + pub fn from_plan(plan: &dyn ExecutionPlan) -> Self { + let plan_props = plan.properties(); + let emission_type = match plan_props.emission_type { + EmissionType::Incremental => "Incremental", + EmissionType::Final => "Final", + EmissionType::Both => "Both", + }; + let boundedness = match plan_props.boundedness { + Boundedness::Bounded => "Bounded", + Boundedness::Unbounded { + requires_infinite_memory: false, + } => "Unbounded", + Boundedness::Unbounded { + requires_infinite_memory: true, + } => "UnboundedInfiniteMemory", + }; + Self { + num_partitions: plan_props.output_partitioning().partition_count(), + supports_limit_pushdown: plan.supports_limit_pushdown(), + emission_type: emission_type.to_string(), + boundedness: boundedness.to_string(), + } + } + + /// Extract plan properties from a SedonaCExecutionPlan via FFI. + pub fn from_ffi_plan(plan: &SedonaCExecutionPlan) -> Result { + get_plan_property::(plan, "plan_properties", None) + } + + /// Convert to DataFusion PlanProperties. + pub fn into_plan_properties(self, schema: SchemaRef) -> PlanProperties { + let emission_type = match self.emission_type.as_str() { + "Final" => EmissionType::Final, + "Both" => EmissionType::Both, + _ => EmissionType::Incremental, + }; + let boundedness = match self.boundedness.as_str() { + "Unbounded" => Boundedness::Unbounded { + requires_infinite_memory: false, + }, + "UnboundedInfiniteMemory" => Boundedness::Unbounded { + requires_infinite_memory: true, + }, + _ => Boundedness::Bounded, + }; + PlanProperties::new( + datafusion_physical_expr::EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(self.num_partitions), + emission_type, + boundedness, + ) + } +} + +/// Helper wrapper to format an ExecutionPlan with a specific DisplayFormatType. +struct DisplayAsWrapper<'a>(&'a Arc, DisplayFormatType); + +impl std::fmt::Display for DisplayAsWrapper<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt_as(self.1, f) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field}; + use datafusion_common::assert_batches_eq; + use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + use futures::{stream, StreamExt}; + use std::fmt::Formatter; + + /// A dummy ExecutionPlan with fixed, predictable values for testing FFI roundtrip. + #[derive(Debug)] + struct DummyExec { + schema: SchemaRef, + properties: PlanProperties, + limit_pushdown: bool, + } + + impl DummyExec { + fn new() -> Self { + Self::with_properties(EmissionType::Incremental, Boundedness::Bounded, true) + } + + fn with_properties( + emission_type: EmissionType, + boundedness: Boundedness, + supports_limit_pushdown: bool, + ) -> Self { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + let properties = PlanProperties::new( + datafusion_physical_expr::EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(3), + emission_type, + boundedness, + ); + Self { + schema, + properties, + limit_pushdown: supports_limit_pushdown, + } + } + } + + impl DisplayAs for DummyExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default => write!(f, "DummyExec: default format"), + DisplayFormatType::Verbose => write!(f, "DummyExec: verbose format with schema"), + DisplayFormatType::TreeRender => write!(f, "DummyExec: tree render format"), + } + } + } + + impl ExecutionPlan for DummyExec { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "DummyExec" + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if !children.is_empty() { + return exec_err!("DummyExec does not support children"); + } + Ok(self) + } + + fn supports_limit_pushdown(&self) -> bool { + self.limit_pushdown + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + // Return a batch with partition-specific data + let ids = Int32Array::from(vec![ + partition as i32 * 10 + 1, + partition as i32 * 10 + 2, + partition as i32 * 10 + 3, + ]); + let values = Int32Array::from(vec![100, 200, 300]); + let batch = + RecordBatch::try_new(self.schema.clone(), vec![Arc::new(ids), Arc::new(values)])?; + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream::iter(vec![Ok(batch)]), + ))) + } + } + + /// Helper to set up an imported plan from a DummyExec through FFI roundtrip. + fn setup_imported_plan() -> (ImportedSedonaCExec, Arc) { + let dummy = Arc::new(DummyExec::new()); + let runtime = tokio::runtime::Handle::current(); + let task_ctx = Arc::new(TaskContext::default()); + + let exported = ExportedExecutionPlan::new(dummy, task_ctx.clone(), runtime); + let ffi_plan: SedonaCExecutionPlan = exported.into(); + let imported = ImportedSedonaCExec::try_new(ffi_plan).unwrap(); + (imported, task_ctx) + } + + fn setup_imported_plan_with( + emission_type: EmissionType, + boundedness: Boundedness, + supports_limit_pushdown: bool, + ) -> (ImportedSedonaCExec, Arc) { + let dummy = Arc::new(DummyExec::with_properties( + emission_type, + boundedness, + supports_limit_pushdown, + )); + let runtime = tokio::runtime::Handle::current(); + let task_ctx = Arc::new(TaskContext::default()); + + let exported = ExportedExecutionPlan::new(dummy, task_ctx.clone(), runtime); + let ffi_plan: SedonaCExecutionPlan = exported.into(); + let imported = ImportedSedonaCExec::try_new(ffi_plan).unwrap(); + + (imported, task_ctx) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_execution_plan_roundtrip_schema() { + let (imported, _) = setup_imported_plan(); + + // Verify schema matches + assert_eq!(imported.schema().fields().len(), 2); + assert_eq!(imported.schema().field(0).name(), "id"); + assert_eq!(imported.schema().field(1).name(), "value"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_execution_plan_roundtrip_name() { + let (imported, _) = setup_imported_plan(); + assert_eq!(imported.name(), "ImportedSedonaCExec"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_execution_plan_roundtrip_properties() { + // Test all EmissionType variants + for (emission_type, expected_str) in [ + (EmissionType::Incremental, "Incremental"), + (EmissionType::Final, "Final"), + (EmissionType::Both, "Both"), + ] { + let (imported, _) = setup_imported_plan_with(emission_type, Boundedness::Bounded, true); + let props = PlanPropertiesArgs::from_plan(&imported); + assert_eq!( + props.emission_type, expected_str, + "emission_type mismatch for {:?}", + emission_type + ); + } + + // Test all Boundedness variants + for (boundedness, expected_str) in [ + (Boundedness::Bounded, "Bounded"), + ( + Boundedness::Unbounded { + requires_infinite_memory: false, + }, + "Unbounded", + ), + ( + Boundedness::Unbounded { + requires_infinite_memory: true, + }, + "UnboundedInfiniteMemory", + ), + ] { + let (imported, _) = + setup_imported_plan_with(EmissionType::Incremental, boundedness, true); + let props = PlanPropertiesArgs::from_plan(&imported); + assert_eq!( + props.boundedness, expected_str, + "boundedness mismatch for {:?}", + boundedness + ); + } + + // Test supports_limit_pushdown + let (imported_with, _) = + setup_imported_plan_with(EmissionType::Incremental, Boundedness::Bounded, true); + assert!(imported_with.supports_limit_pushdown()); + + let (imported_without, _) = + setup_imported_plan_with(EmissionType::Incremental, Boundedness::Bounded, false); + assert!(!imported_without.supports_limit_pushdown()); + + // Test partition count + let (imported, _) = setup_imported_plan(); + assert_eq!( + imported + .properties() + .output_partitioning() + .partition_count(), + 3 + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_execution_plan_roundtrip_display_as() { + let (imported, _) = setup_imported_plan(); + + // Helper struct to format DisplayAs with a specific format type + struct DisplayAsFormat<'a, T: DisplayAs>(&'a T, DisplayFormatType); + impl std::fmt::Display for DisplayAsFormat<'_, T> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.0.fmt_as(self.1, f) + } + } + + assert_eq!( + format!("{}", DisplayAsFormat(&imported, DisplayFormatType::Default)), + "DummyExec: default format" + ); + assert_eq!( + format!("{}", DisplayAsFormat(&imported, DisplayFormatType::Verbose)), + "DummyExec: verbose format with schema" + ); + assert_eq!( + format!( + "{}", + DisplayAsFormat(&imported, DisplayFormatType::TreeRender) + ), + "DummyExec: tree render format" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_execution_plan_roundtrip_debug_string() { + let (imported, _) = setup_imported_plan(); + + let debug_str = imported.get_debug_string().unwrap(); + assert!( + debug_str.contains("DummyExec"), + "debug_string should contain 'DummyExec', got: {}", + debug_str + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_execution_plan_roundtrip_execute() { + let (imported, task_ctx) = setup_imported_plan(); + + // Execute partition 0 + let stream = imported.execute(0, task_ctx.clone()).unwrap(); + let batches: Vec = stream + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap(); + + let expected = [ + "+----+-------+", + "| id | value |", + "+----+-------+", + "| 1 | 100 |", + "| 2 | 200 |", + "| 3 | 300 |", + "+----+-------+", + ]; + assert_batches_eq!(expected, &batches); + + // Execute partition 1 - needs a fresh import since we consumed the first + let stream2 = imported.execute(1, task_ctx).unwrap(); + let batches2: Vec = stream2 + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap(); + + let expected2 = [ + "+----+-------+", + "| id | value |", + "+----+-------+", + "| 11 | 100 |", + "| 12 | 200 |", + "| 13 | 300 |", + "+----+-------+", + ]; + assert_batches_eq!(expected2, &batches2); + } +} diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index c208f6cbff..79e70c9d00 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -97,7 +97,7 @@ impl ExportedTableProvider { }) } - fn get_property(&self, property: &str) -> Result> { + fn get_property(&self, property: &str) -> Result { match property { "table_type" => { let table_type = self.inner.table_type(); @@ -106,7 +106,7 @@ impl ExportedTableProvider { TableType::View => "View", TableType::Temporary => "Temporary", }; - Ok(type_str.as_bytes().to_vec()) + Ok(type_str.to_string()) } _ => exec_err!("Unknown property: {}", property), } @@ -118,7 +118,7 @@ impl From for SedonaCTableProvider { let boxed = Box::new(value); Self { get_schema: Some(c_table_provider_get_schema), - get_property_schema: None, + get_property_schema: Some(c_table_provider_get_property_schema), get_property: Some(c_table_provider_get_property), scan: Some(c_table_provider_scan), insert: None, @@ -142,6 +142,27 @@ unsafe extern "C" fn c_table_provider_get_schema( } } +unsafe extern "C" fn c_table_provider_get_property_schema( + _self_: *const SedonaCTableProvider, + _property: *const std::ffi::c_char, + out: *mut FFI_ArrowSchema, + err: *mut SedonaCError, +) -> c_int { + // All properties are returned as Utf8 strings + use arrow_schema::{DataType, Field}; + let field = Field::new("value", DataType::Utf8, false); + match FFI_ArrowSchema::try_from(&field) { + Ok(ffi_schema) => { + std::ptr::write(out, ffi_schema); + ERRNO_OK + } + Err(e) => { + *err = SedonaCError::new(&format!("Failed to convert field to FFI schema: {}", e)); + libc::EINVAL + } + } +} + unsafe extern "C" fn c_table_provider_get_property( self_: *const SedonaCTableProvider, property: *const std::ffi::c_char, @@ -153,11 +174,11 @@ unsafe extern "C" fn c_table_provider_get_property( let property_str = CStr::from_ptr(property).to_string_lossy(); match provider.get_property(&property_str) { - Ok(bytes) => { - // Return the bytes as a single-element binary array - use arrow_array::{builder::BinaryBuilder, Array}; - let mut builder = BinaryBuilder::new(); - builder.append_value(&bytes); + Ok(value) => { + // Return the string as a single-element string array + use arrow_array::{builder::StringBuilder, Array}; + let mut builder = StringBuilder::new(); + builder.append_value(&value); let array = builder.finish(); let ffi_array = FFI_ArrowArray::new(&array.to_data()); std::ptr::write(out, ffi_array); diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index bb27eb2013..093b86e7c9 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -23,7 +23,8 @@ use std::sync::Mutex; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{RecordBatch, RecordBatchReader}; -use arrow_schema::{ArrowError, SchemaRef}; +use arrow_schema::ffi::FFI_ArrowSchema; +use arrow_schema::{ArrowError, DataType, Field, SchemaRef}; use datafusion_common::Result; use datafusion_physical_plan::SendableRecordBatchStream; use futures::StreamExt; @@ -38,6 +39,36 @@ use crate::extension::{ /// Success return code for FFI functions. pub const ERRNO_OK: c_int = 0; +/// Get the schema for a property from a [SedonaCExecutionPlan]. +/// +/// Returns the DataType describing the property's data type. +/// If `get_property_schema` is not implemented, defaults to Binary. +fn get_plan_property_data_type(plan: &SedonaCExecutionPlan, property: &str) -> Result { + let Some(get_property_schema) = plan.get_property_schema else { + // Default to Binary if get_property_schema is not implemented + return Ok(DataType::Binary); + }; + + let property_cstr = CString::new(property) + .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; + + let mut ffi_schema = FFI_ArrowSchema::empty(); + let mut err = SedonaCError::default(); + + let code = + unsafe { get_property_schema(plan, property_cstr.as_ptr(), &mut ffi_schema, &mut err) }; + + if code != ERRNO_OK { + return sedona_internal_err!("Failed to get property schema for '{}': {}", property, err); + } + + // Try to convert the FFI schema to a Field + let field = Field::try_from(&ffi_schema).map_err(|e| { + sedona_internal_datafusion_err!("Failed to parse property schema for '{}': {}", property, e) + })?; + Ok(field.data_type().clone()) +} + /// Call `get_property` on a [SedonaCExecutionPlan] and deserialize the result. /// /// This handles the common pattern of: @@ -113,35 +144,71 @@ where return sedona_internal_err!("Failed to get property '{}': {}", property, err); } - // Parse the binary array to get the JSON bytes - parse_binary_ffi_array(ffi_array) + // Get the property schema to know how to interpret the array + let data_type = get_plan_property_data_type(plan, property)?; + + // Parse the array to get the JSON bytes + parse_ffi_array(ffi_array, &data_type) } -/// Parse a binary FFI array containing JSON and deserialize to the target type. -fn parse_binary_ffi_array( +/// Parse an FFI array containing JSON and deserialize to the target type. +fn parse_ffi_array( ffi_array: arrow_array::ffi::FFI_ArrowArray, + data_type: &DataType, ) -> Result { - let bytes = parse_binary_ffi_array_to_bytes(ffi_array)?; + let bytes = parse_ffi_array_to_bytes(ffi_array, data_type)?; serde_json::from_slice::(&bytes) .map_err(|e| sedona_internal_datafusion_err!("Failed to deserialize property: {}", e)) } -/// Parse a binary FFI array and return the raw bytes. -fn parse_binary_ffi_array_to_bytes( +/// Parse an FFI array and return the raw bytes. +fn parse_ffi_array_to_bytes( ffi_array: arrow_array::ffi::FFI_ArrowArray, + data_type: &DataType, ) -> Result> { - let data = unsafe { - arrow_array::ffi::from_ffi_and_data_type(ffi_array, arrow_schema::DataType::Binary)? - }; + let data = unsafe { arrow_array::ffi::from_ffi_and_data_type(ffi_array, data_type.clone())? }; let array = arrow_array::make_array(data); - let binary_array = array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - sedona_internal_datafusion_err!("Expected binary array from get_property") - })?; - - Ok(binary_array.value(0).to_vec()) + + // Handle different array types + match data_type { + DataType::Binary => { + let binary_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + sedona_internal_datafusion_err!("Expected binary array from get_property") + })?; + Ok(binary_array.value(0).to_vec()) + } + DataType::LargeBinary => { + let binary_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + sedona_internal_datafusion_err!("Expected large binary array from get_property") + })?; + Ok(binary_array.value(0).to_vec()) + } + DataType::Utf8 => { + let string_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + sedona_internal_datafusion_err!("Expected string array from get_property") + })?; + Ok(string_array.value(0).as_bytes().to_vec()) + } + DataType::LargeUtf8 => { + let string_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + sedona_internal_datafusion_err!("Expected large string array from get_property") + })?; + Ok(string_array.value(0).as_bytes().to_vec()) + } + _ => sedona_internal_err!("Unsupported data type for property: {:?}", data_type), + } } /// Get a string property from a [SedonaCExecutionPlan]. @@ -182,11 +249,47 @@ pub fn get_plan_string_property(plan: &SedonaCExecutionPlan, property: &str) -> return sedona_internal_err!("Failed to get '{}': {}", property, err); } - let bytes = parse_binary_ffi_array_to_bytes(ffi_array)?; + // Get the property schema to know how to interpret the array + let data_type = get_plan_property_data_type(plan, property)?; + + let bytes = parse_ffi_array_to_bytes(ffi_array, &data_type)?; String::from_utf8(bytes) .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) } +/// Get the schema for a property from a [SedonaCTableProvider]. +/// +/// Returns the DataType describing the property's data type. +/// If `get_property_schema` is not implemented, defaults to Binary. +fn get_table_provider_property_data_type( + provider: &SedonaCTableProvider, + property: &str, +) -> Result { + let Some(get_property_schema) = provider.get_property_schema else { + // Default to Binary if get_property_schema is not implemented + return Ok(DataType::Binary); + }; + + let property_cstr = CString::new(property) + .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; + + let mut ffi_schema = FFI_ArrowSchema::empty(); + let mut err = SedonaCError::default(); + + let code = + unsafe { get_property_schema(provider, property_cstr.as_ptr(), &mut ffi_schema, &mut err) }; + + if code != ERRNO_OK { + return sedona_internal_err!("Failed to get property schema for '{}': {}", property, err); + } + + // Try to convert the FFI schema to a Field + let field = Field::try_from(&ffi_schema).map_err(|e| { + sedona_internal_datafusion_err!("Failed to parse property schema for '{}': {}", property, e) + })?; + Ok(field.data_type().clone()) +} + /// Get a string property from a [SedonaCTableProvider]. pub fn get_table_provider_string_property( provider: &SedonaCTableProvider, @@ -226,7 +329,10 @@ pub fn get_table_provider_string_property( return sedona_internal_err!("Failed to get '{}': {}", property, err); } - let bytes = parse_binary_ffi_array_to_bytes(ffi_array)?; + // Get the property schema to know how to interpret the array + let data_type = get_table_provider_property_data_type(provider, property)?; + + let bytes = parse_ffi_array_to_bytes(ffi_array, &data_type)?; String::from_utf8(bytes) .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) } From aa6293552557cc6da29506ce7876e785191b8fce Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 24 Jun 2026 16:30:43 -0500 Subject: [PATCH 10/55] a few more spots --- c/sedona-extension/src/execution_plan.rs | 16 ++++++++--- c/sedona-extension/src/table_provider.rs | 34 +++++++++++++++--------- 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 9eef7683ae..cfb5363469 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -196,7 +196,9 @@ unsafe extern "C" fn c_exec_plan_get_property_schema( ERRNO_OK } Err(e) => { - *err = SedonaCError::new(&format!("Failed to convert field to FFI schema: {}", e)); + if !err.is_null() { + *err = SedonaCError::new(&format!("Failed to convert field to FFI schema: {}", e)); + } libc::EINVAL } } @@ -224,7 +226,9 @@ unsafe extern "C" fn c_exec_plan_get_property( ERRNO_OK } Err(e) => { - *err = SedonaCError::new(&e.to_string()); + if !err.is_null() { + *err = SedonaCError::new(&e.to_string()); + } libc::EINVAL } } @@ -249,7 +253,9 @@ unsafe extern "C" fn c_exec_plan_execute( let execute_args: ExecuteArgs = match serde_json::from_slice(args_slice) { Ok(a) => a, Err(e) => { - *err = SedonaCError::new(&format!("Failed to parse execute args: {}", e)); + if !err.is_null() { + *err = SedonaCError::new(&format!("Failed to parse execute args: {}", e)); + } return libc::EINVAL; } }; @@ -263,7 +269,9 @@ unsafe extern "C" fn c_exec_plan_execute( ERRNO_OK } Err(e) => { - *err = SedonaCError::new(&e.to_string()); + if !err.is_null() { + *err = SedonaCError::new(&e.to_string()); + } libc::EINVAL } } diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 79e70c9d00..0ce3e22bbe 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -39,15 +39,6 @@ use crate::extension::{ }; use crate::utils::{get_table_provider_string_property, ERRNO_OK}; -/// Arguments for a scan operation, serialized as JSON across FFI. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ScanArgs { - /// Column indices to project, or None for all columns. - pub projection: Option>, - /// Maximum number of rows to return, or None for unlimited. - pub limit: Option, -} - /// A TableProvider wrapper that can be exported across FFI. /// /// This wraps an inner TableProvider and exposes it via the SedonaCTableProvider @@ -157,7 +148,9 @@ unsafe extern "C" fn c_table_provider_get_property_schema( ERRNO_OK } Err(e) => { - *err = SedonaCError::new(&format!("Failed to convert field to FFI schema: {}", e)); + if !err.is_null() { + *err = SedonaCError::new(&format!("Failed to convert field to FFI schema: {}", e)); + } libc::EINVAL } } @@ -185,7 +178,9 @@ unsafe extern "C" fn c_table_provider_get_property( ERRNO_OK } Err(e) => { - *err = SedonaCError::new(&e.to_string()); + if !err.is_null() { + *err = SedonaCError::new(&e.to_string()); + } libc::EINVAL } } @@ -216,7 +211,9 @@ unsafe extern "C" fn c_table_provider_scan( match serde_json::from_slice(args_slice) { Ok(a) => a, Err(e) => { - *err = SedonaCError::new(&format!("Failed to parse scan args: {}", e)); + if !err.is_null() { + *err = SedonaCError::new(&format!("Failed to parse scan args: {}", e)); + } return libc::EINVAL; } } @@ -231,7 +228,9 @@ unsafe extern "C" fn c_table_provider_scan( ERRNO_OK } Err(e) => { - *err = SedonaCError::new(&e.to_string()); + if !err.is_null() { + *err = SedonaCError::new(&e.to_string()); + } libc::EINVAL } } @@ -358,6 +357,15 @@ impl TableProvider for ImportedTableProvider { } } +/// Arguments for a scan operation, serialized as JSON across FFI. +#[derive(Debug, Clone, Serialize, Deserialize)] +struct ScanArgs { + /// Column indices to project, or None for all columns. + pub projection: Option>, + /// Maximum number of rows to return, or None for unlimited. + pub limit: Option, +} + #[cfg(test)] mod tests { use super::*; From 878c799ab540e091810a9e277b5f92158507f5d0 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 24 Jun 2026 16:45:00 -0500 Subject: [PATCH 11/55] add a lower-level test for the tabler provider --- c/sedona-extension/src/table_provider.rs | 80 ++++++++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 0ce3e22bbe..2d8de2c854 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -372,7 +372,10 @@ mod tests { use arrow_array::{Float64Array, Int32Array, Int64Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; use datafusion::prelude::SessionContext; + use datafusion_catalog::Session; use datafusion_common::assert_batches_eq; + use datafusion_expr::Expr; + use datafusion_physical_plan::ExecutionPlan; /// Create a SessionContext with a test table containing 5 numeric columns and multiple batches. async fn create_test_context() -> Result { @@ -685,4 +688,81 @@ mod tests { assert_batches_eq!(expected, &result); Ok(()) } + + /// A dummy TableProvider with configurable table_type for testing FFI roundtrip. + #[derive(Debug)] + struct DummyTableProvider { + schema: SchemaRef, + table_type: TableType, + } + + impl DummyTableProvider { + fn with_table_type(table_type: TableType) -> Self { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + Self { schema, table_type } + } + } + + #[async_trait] + impl TableProvider for DummyTableProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + self.table_type + } + + async fn scan( + &self, + _state: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + exec_err!("DummyTableProvider does not support scan") + } + } + + /// Helper to set up an imported table provider from a DummyTableProvider through FFI roundtrip. + fn setup_imported_provider_with(table_type: TableType) -> ImportedTableProvider { + let dummy = Arc::new(DummyTableProvider::with_table_type(table_type)); + let ctx = SessionContext::new(); + let runtime = tokio::runtime::Handle::current(); + let session = Arc::new(ctx.state()); + let exported = ExportedTableProvider::new(dummy, session, runtime); + let ffi_provider: SedonaCTableProvider = exported.into(); + ImportedTableProvider::try_new(ffi_provider).expect("Failed to import table provider") + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_table_provider_roundtrip_schema() -> Result<()> { + let imported = setup_imported_provider_with(TableType::Base); + + // Check schema roundtrip + let schema = imported.schema(); + assert_eq!(schema.fields().len(), 2); + assert_eq!(schema.field(0).name(), "id"); + assert_eq!(schema.field(0).data_type(), &DataType::Int32); + assert_eq!(schema.field(1).name(), "value"); + assert_eq!(schema.field(1).data_type(), &DataType::Int32); + + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_table_provider_roundtrip_table_type() -> Result<()> { + for table_type in [TableType::Base, TableType::View, TableType::Temporary] { + let imported = setup_imported_provider_with(table_type); + assert_eq!(imported.table_type(), table_type); + } + Ok(()) + } } From ff433c7f0237f49f0b161ab211508c8e1b6a5794 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 24 Jun 2026 16:55:42 -0500 Subject: [PATCH 12/55] don't require a specific runtime --- c/sedona-extension/src/execution_plan.rs | 12 ++++----- c/sedona-extension/src/table_provider.rs | 34 +++++++++++++----------- c/sedona-extension/src/utils.rs | 27 +++++++++++++------ 3 files changed, 43 insertions(+), 30 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index cfb5363469..25c94318db 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -758,7 +758,7 @@ mod tests { (imported, task_ctx) } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_execution_plan_roundtrip_schema() { let (imported, _) = setup_imported_plan(); @@ -768,13 +768,13 @@ mod tests { assert_eq!(imported.schema().field(1).name(), "value"); } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_execution_plan_roundtrip_name() { let (imported, _) = setup_imported_plan(); assert_eq!(imported.name(), "ImportedSedonaCExec"); } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_execution_plan_roundtrip_properties() { // Test all EmissionType variants for (emission_type, expected_str) in [ @@ -837,7 +837,7 @@ mod tests { ); } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_execution_plan_roundtrip_display_as() { let (imported, _) = setup_imported_plan(); @@ -866,7 +866,7 @@ mod tests { ); } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_execution_plan_roundtrip_debug_string() { let (imported, _) = setup_imported_plan(); @@ -878,7 +878,7 @@ mod tests { ); } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_execution_plan_roundtrip_execute() { let (imported, task_ctx) = setup_imported_plan(); diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 2d8de2c854..7201626ee3 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -76,16 +76,18 @@ impl ExportedTableProvider { projection: Option>, limit: Option, ) -> Result> { - // Convert projection to the expected format - let projection_ref = projection.as_ref(); - - // Execute the scan - we use block_in_place to allow blocking within an async context let inner = self.inner.clone(); - let session = self.session.as_ref(); - tokio::task::block_in_place(|| { - self.runtime - .block_on(async { inner.scan(session, projection_ref, &[], limit).await }) + let session = self.session.clone(); + let runtime = self.runtime.clone(); + + std::thread::spawn(move || { + let projection_ref = projection.as_ref(); + runtime.block_on(inner.scan(session.as_ref(), projection_ref, &[], limit)) }) + .join() + .map_err(|_| { + datafusion_common::DataFusionError::Internal("Scan thread panicked".to_string()) + })? } fn get_property(&self, property: &str) -> Result { @@ -444,7 +446,7 @@ mod tests { Ok(ctx) } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_roundtrip_simple_select() -> Result<()> { let ctx = create_test_context().await?; @@ -487,7 +489,7 @@ mod tests { Ok(()) } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_roundtrip_projection() -> Result<()> { let ctx = create_test_context().await?; @@ -528,7 +530,7 @@ mod tests { Ok(()) } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_roundtrip_filter() -> Result<()> { let ctx = create_test_context().await?; @@ -571,7 +573,7 @@ mod tests { Ok(()) } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_roundtrip_sort() -> Result<()> { let ctx = create_test_context().await?; @@ -614,7 +616,7 @@ mod tests { Ok(()) } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_roundtrip_limit() -> Result<()> { let ctx = create_test_context().await?; @@ -649,7 +651,7 @@ mod tests { Ok(()) } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_roundtrip_all_columns() -> Result<()> { let ctx = create_test_context().await?; @@ -742,7 +744,7 @@ mod tests { ImportedTableProvider::try_new(ffi_provider).expect("Failed to import table provider") } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_table_provider_roundtrip_schema() -> Result<()> { let imported = setup_imported_provider_with(TableType::Base); @@ -757,7 +759,7 @@ mod tests { Ok(()) } - #[tokio::test(flavor = "multi_thread")] + #[tokio::test] async fn test_table_provider_roundtrip_table_type() -> Result<()> { for table_type in [TableType::Base, TableType::View, TableType::Temporary] { let imported = setup_imported_provider_with(table_type); diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index 093b86e7c9..723552bf67 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -362,14 +362,25 @@ impl Iterator for StreamingRecordBatchReader { type Item = std::result::Result; fn next(&mut self) -> Option { - let mut stream = self.stream.lock().unwrap(); - tokio::task::block_in_place(|| { - self.runtime.block_on(async { - match stream.next().await { - Some(Ok(batch)) => Some(Ok(batch)), - Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), - None => None, - } + let stream = &self.stream; + let runtime = &self.runtime; + + std::thread::scope(|s| { + s.spawn(|| { + let mut guard = stream.lock().unwrap(); + runtime.block_on(async { + match guard.next().await { + Some(Ok(batch)) => Some(Ok(batch)), + Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), + None => None, + } + }) + }) + .join() + .unwrap_or_else(|_| { + Some(Err(ArrowError::InvalidArgumentError( + "Iterator thread panicked".to_string(), + ))) }) }) } From 2a8db881daac4b739d140072d1d337f3f1d5cb8f Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 24 Jun 2026 17:21:48 -0500 Subject: [PATCH 13/55] add a test with cancels --- c/sedona-extension/src/lib.rs | 2 +- c/sedona-extension/src/utils.rs | 258 +++++++++++++++++++++++++++++++- 2 files changed, 251 insertions(+), 9 deletions(-) diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs index 2d101a1d1a..d4f94155b4 100644 --- a/c/sedona-extension/src/lib.rs +++ b/c/sedona-extension/src/lib.rs @@ -20,4 +20,4 @@ pub mod expr; pub mod extension; pub mod scalar_kernel; pub mod table_provider; -mod utils; +pub mod utils; diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index 723552bf67..0f022a152d 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -19,13 +19,12 @@ use std::ffi::{c_int, CString}; use std::ptr::null_mut; -use std::sync::Mutex; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::ffi::FFI_ArrowSchema; use arrow_schema::{ArrowError, DataType, Field, SchemaRef}; -use datafusion_common::Result; +use datafusion_common::{Result, exec_err}; use datafusion_physical_plan::SendableRecordBatchStream; use futures::StreamExt; use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; @@ -337,14 +336,23 @@ pub fn get_table_provider_string_property( .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) } +/// A cancellation check callback for FFI operations. +/// +/// Returns `true` if the operation should be cancelled, `false` to continue. +/// This allows Python, R, ADBC, and other runtimes to integrate their own +/// cancellation mechanisms. +pub type CancelChecker = Box bool + Send + Sync>; + /// A RecordBatchReader that lazily polls a SendableRecordBatchStream. /// /// This allows exporting a DataFusion stream over FFI without collecting all /// batches upfront. Used when exporting execution plans across FFI boundaries. pub struct StreamingRecordBatchReader { schema: SchemaRef, - stream: Mutex, + stream: SendableRecordBatchStream, runtime: tokio::runtime::Handle, + cancel_checker: Option, + cancelled: bool, } impl StreamingRecordBatchReader { @@ -352,8 +360,29 @@ impl StreamingRecordBatchReader { pub fn new(stream: SendableRecordBatchStream, runtime: tokio::runtime::Handle) -> Self { Self { schema: stream.schema(), - stream: Mutex::new(stream), + stream, + runtime, + cancel_checker: None, + cancelled: false, + } + } + + /// Create a new StreamingRecordBatchReader with a cancellation checker. + /// + /// The cancellation checker is called before each batch is fetched. If it + /// returns `true`, iteration stops with a cancellation error on the next + /// call and `None` on subsequent calls. + pub fn with_cancel_checker( + stream: SendableRecordBatchStream, + runtime: tokio::runtime::Handle, + cancel_checker: CancelChecker, + ) -> Self { + Self { + schema: stream.schema(), + stream, runtime, + cancel_checker: Some(cancel_checker), + cancelled: false, } } } @@ -362,14 +391,28 @@ impl Iterator for StreamingRecordBatchReader { type Item = std::result::Result; fn next(&mut self) -> Option { - let stream = &self.stream; + // If already cancelled, return None to stop iteration + if self.cancelled { + return None; + } + + // Check for cancellation before fetching the next batch + if let Some(ref checker) = self.cancel_checker { + if checker() { + self.cancelled = true; + return Some(Err(ArrowError::ExternalError(Box::new( + std::io::Error::new(std::io::ErrorKind::Interrupted, "Operation cancelled"), + )))); + } + } + + let stream = &mut self.stream; let runtime = &self.runtime; std::thread::scope(|s| { s.spawn(|| { - let mut guard = stream.lock().unwrap(); runtime.block_on(async { - match guard.next().await { + match stream.next().await { Some(Ok(batch)) => Some(Ok(batch)), Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), None => None, @@ -404,11 +447,33 @@ impl RecordBatchReader for StreamingRecordBatchReader { /// initialized. pub unsafe fn ffi_stream_to_sendable( ffi_stream: &mut FFI_ArrowArrayStream, +) -> Result { + ffi_stream_to_sendable_with_cancel(ffi_stream, None) +} + +/// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream with cancellation support. +/// +/// The cancellation checker is called before each batch is read. If it returns +/// `true`, the stream yields a cancellation error. +/// +/// # Safety +/// +/// The caller must ensure that the FFI stream pointer is valid and properly +/// initialized. +pub unsafe fn ffi_stream_to_sendable_with_cancel( + ffi_stream: &mut FFI_ArrowArrayStream, + cancel_checker: Option, ) -> Result { let reader = arrow_array::ffi_stream::ArrowArrayStreamReader::from_raw(ffi_stream)?; let schema = reader.schema(); - let stream = futures::stream::iter(reader).map(|result| { + let stream = futures::stream::iter(reader).map(move |result| { + // Check for cancellation before yielding each batch + if let Some(ref checker) = cancel_checker { + if checker() { + return exec_err!("Operation cancelled"); + } + } result.map_err(|e| datafusion_common::DataFusionError::ArrowError(Box::new(e), None)) }); @@ -416,3 +481,180 @@ pub unsafe fn ffi_stream_to_sendable( datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), )) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::Int32Array; + use arrow_schema::{Field, Schema}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::sync::Arc; + use std::time::Duration; + + /// Create a slow stream that yields batches with a configurable delay. + fn create_slow_stream( + num_batches: usize, + delay_ms: u64, + ) -> (SchemaRef, SendableRecordBatchStream) { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Int32, false)])); + let schema_clone = schema.clone(); + + let stream = futures::stream::iter(0..num_batches).then(move |i| { + let schema = schema_clone.clone(); + async move { + if delay_ms > 0 { + // Use std::thread::sleep because tokio::time::sleep requires + // the runtime to poll, which can deadlock with block_on + std::thread::sleep(Duration::from_millis(delay_ms)); + } + let array = Int32Array::from(vec![i as i32]); + Ok(RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()) + } + }); + + ( + schema.clone(), + Box::pin(datafusion_physical_plan::stream::RecordBatchStreamAdapter::new( + schema, stream, + )), + ) + } + + #[tokio::test] + async fn test_streaming_reader_basic() { + let runtime = tokio::runtime::Handle::current(); + let (_schema, stream) = create_slow_stream(5, 10); + + let reader = StreamingRecordBatchReader::new(stream, runtime); + let batches: Vec<_> = reader.collect(); + + assert_eq!(batches.len(), 5); + for (i, batch) in batches.iter().enumerate() { + let batch = batch.as_ref().unwrap(); + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(array.value(0), i as i32); + } + } + + #[tokio::test] + async fn test_streaming_reader_cancel() { + let runtime = tokio::runtime::Handle::current(); + let (_schema, stream) = create_slow_stream(10, 50); + + // Cancel after reading 3 batches + let counter = Arc::new(AtomicUsize::new(0)); + let counter_clone = counter.clone(); + let cancel_checker: CancelChecker = Box::new(move || { + let count = counter_clone.fetch_add(1, Ordering::SeqCst); + count >= 3 + }); + + let reader = StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker); + let batches: Vec<_> = reader.collect(); + + // Should have 3 successful batches + 1 cancellation error + assert_eq!(batches.len(), 4); + + // First 3 should be Ok + for i in 0..3 { + assert!(batches[i].is_ok(), "batch {} should be Ok", i); + } + + // Last one should be a cancellation error + let last = batches.last().unwrap(); + assert!(last.is_err()); + let err = last.as_ref().unwrap_err(); + assert!( + err.to_string().contains("cancelled"), + "error should mention cancellation: {}", + err + ); + } + + #[tokio::test] + async fn test_ffi_stream_to_sendable_basic() { + use arrow_array::ffi_stream::FFI_ArrowArrayStream; + + let runtime = tokio::runtime::Handle::current(); + let (_schema, stream) = create_slow_stream(5, 10); + + // Export to FFI stream + let reader = StreamingRecordBatchReader::new(stream, runtime.clone()); + let mut ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); + + // Import back + let imported = unsafe { ffi_stream_to_sendable(&mut ffi_stream).unwrap() }; + + // Collect results + let batches: Vec<_> = imported.collect::>().await; + + assert_eq!(batches.len(), 5); + for (i, batch) in batches.iter().enumerate() { + let batch = batch.as_ref().unwrap(); + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(array.value(0), i as i32); + } + } + + #[tokio::test] + async fn test_ffi_stream_to_sendable_cancel() { + use arrow_array::ffi_stream::FFI_ArrowArrayStream; + + let runtime = tokio::runtime::Handle::current(); + let (_schema, stream) = create_slow_stream(10, 50); + + // Export to FFI stream (no cancellation on export side) + let reader = StreamingRecordBatchReader::new(stream, runtime.clone()); + let mut ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); + + // Cancel after reading 3 batches on import side + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_clone = cancelled.clone(); + let cancel_checker: CancelChecker = + Box::new(move || cancelled_clone.load(Ordering::SeqCst)); + + let imported = unsafe { + ffi_stream_to_sendable_with_cancel(&mut ffi_stream, Some(cancel_checker)).unwrap() + }; + + // Collect with cancellation after 3 batches, stop on first error + let mut batches = Vec::new(); + let mut stream = imported; + while let Some(result) = stream.next().await { + let is_err = result.is_err(); + batches.push(result); + if batches.len() == 3 { + cancelled.store(true, Ordering::SeqCst); + } + if is_err { + break; // Stop on first error + } + } + + // Should have 3 successful batches + 1 cancellation error + assert_eq!(batches.len(), 4); + + // First 3 should be Ok + for i in 0..3 { + assert!(batches[i].is_ok(), "batch {} should be Ok", i); + } + + // Last one should be a cancellation error + let last = batches.last().unwrap(); + assert!(last.is_err()); + let err = last.as_ref().unwrap_err(); + assert!( + err.to_string().contains("cancelled"), + "error should mention cancellation: {}", + err + ); + } +} From 1fdef0c8050cbdd82952f4e61d98c5741f0d4ef5 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 26 Jun 2026 16:56:08 -0500 Subject: [PATCH 14/55] fmt --- c/sedona-extension/src/utils.rs | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index 0f022a152d..d50991c932 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -24,7 +24,7 @@ use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::ffi::FFI_ArrowSchema; use arrow_schema::{ArrowError, DataType, Field, SchemaRef}; -use datafusion_common::{Result, exec_err}; +use datafusion_common::{exec_err, Result}; use datafusion_physical_plan::SendableRecordBatchStream; use futures::StreamExt; use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; @@ -496,7 +496,11 @@ mod tests { num_batches: usize, delay_ms: u64, ) -> (SchemaRef, SendableRecordBatchStream) { - let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Int32, false)])); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); let schema_clone = schema.clone(); let stream = futures::stream::iter(0..num_batches).then(move |i| { @@ -514,9 +518,9 @@ mod tests { ( schema.clone(), - Box::pin(datafusion_physical_plan::stream::RecordBatchStreamAdapter::new( - schema, stream, - )), + Box::pin( + datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), + ), ) } @@ -553,7 +557,8 @@ mod tests { count >= 3 }); - let reader = StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker); + let reader = + StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker); let batches: Vec<_> = reader.collect(); // Should have 3 successful batches + 1 cancellation error From 9b567a31e6b849e6c21dd8e0dd33a566b7f566f7 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 26 Jun 2026 17:11:43 -0500 Subject: [PATCH 15/55] use enumerate in tests --- c/sedona-extension/src/utils.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index d50991c932..2cf0cea724 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -565,8 +565,8 @@ mod tests { assert_eq!(batches.len(), 4); // First 3 should be Ok - for i in 0..3 { - assert!(batches[i].is_ok(), "batch {} should be Ok", i); + for (i, batch) in batches.iter().enumerate().take(3) { + assert!(batch.is_ok(), "batch {} should be Ok", i); } // Last one should be a cancellation error @@ -648,8 +648,8 @@ mod tests { assert_eq!(batches.len(), 4); // First 3 should be Ok - for i in 0..3 { - assert!(batches[i].is_ok(), "batch {} should be Ok", i); + for (i, batch) in batches.iter().enumerate().take(3) { + assert!(batch.is_ok(), "batch {} should be Ok", i); } // Last one should be a cancellation error From 83b0025a58c646c44b7d752343432e2c4b14fa05 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 26 Jun 2026 20:58:49 -0500 Subject: [PATCH 16/55] Revert "move dataframe stuff to dataframe" This reverts commit 56f2ed6f6aa1604007aaf33808b55ae52d114b45. --- python/sedonadb/src/dataframe.rs | 2 +- r/sedonadb/src/rust/src/dataframe.rs | 2 +- rust/sedona/src/context.rs | 267 ++++++++++++++++++++++++++- rust/sedona/src/lib.rs | 1 - 4 files changed, 262 insertions(+), 10 deletions(-) diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs index cb81311bc6..761b4f171f 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -34,7 +34,7 @@ use futures::lock::Mutex; use futures::TryStreamExt; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict, PyList}; -use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; +use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; use sedona::projected_reader::simplify_record_batch_reader; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_geoparquet::options::TableGeoParquetOptions; diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index 7d22ed4596..20af9bf5bc 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -28,7 +28,7 @@ use datafusion_expr::JoinType; use datafusion_expr::{select_expr::SelectExpr, Expr, SortExpr}; use datafusion_ffi::table_provider::FFI_TableProvider; use savvy::{savvy, savvy_err, sexp, IntoExtPtrSexp, Result}; -use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; +use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; use sedona::reader::SedonaStreamReader; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_geoparquet::options::TableGeoParquetOptions; diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index 526d727aa2..dfc26fe09b 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -21,9 +21,15 @@ use std::{ use crate::exec::create_plan_from_sql; use crate::object_storage::ensure_object_store_registered_with_options; -use crate::{catalog::DynamicObjectStoreCatalog, random_geometry_provider::RandomGeometryFunction}; +use crate::{ + catalog::DynamicObjectStoreCatalog, + random_geometry_provider::RandomGeometryFunction, + show::{show_batches, DisplayTableOptions}, +}; +use arrow_array::RecordBatch; use arrow_schema::DataType; -use datafusion::execution::memory_pool::MemoryLimit; +use async_trait::async_trait; +use datafusion::datasource::file_format::format_as_file_type; use datafusion::{ common::plan_err, error::{DataFusionError, Result}, @@ -35,7 +41,11 @@ use datafusion::{ prelude::{DataFrame, SessionConfig, SessionContext}, sql::parser::{DFParser, Statement}, }; +use datafusion::{dataframe::DataFrameWriteOptions, execution::memory_pool::MemoryLimit}; +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 parking_lot::Mutex; use sedona_common::{ option::add_sedona_option_extension, sedona_internal_datafusion_err, CrsProviderOption, @@ -45,6 +55,7 @@ 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_geoparquet::options::TableGeoParquetOptions; use sedona_geoparquet::{ format::GeoParquetFormatFactory, provider::{geoparquet_listing_table, GeoParquetReadOptions}, @@ -506,6 +517,179 @@ impl Default for SedonaContext { } } +/// Sedona-specific [`DataFrame`] actions +/// +/// This trait, implemented for [`DataFrame`], extends the DataFrame API to make it +/// ergonomic to work with dataframes that contain geometry columns. Currently these +/// are limited to output functions, as geometry columns currently require special +/// handling when written or exported to an external system. +#[async_trait] +pub trait SedonaDataFrame { + /// Build a table of the first `limit` results in this DataFrame + /// + /// This will limit and execute the query and build a table using [show_batches]. + async fn show_sedona<'a>( + self, + ctx: &SedonaContext, + limit: Option, + options: DisplayTableOptions<'a>, + ) -> Result; + + async fn write_geoparquet( + self, + ctx: &SedonaContext, + path: &str, + options: SedonaWriteOptions, + writer_options: Option, + ) -> Result>; +} + +#[async_trait] +impl SedonaDataFrame for DataFrame { + async fn show_sedona<'a>( + self, + ctx: &SedonaContext, + limit: Option, + mut options: DisplayTableOptions<'a>, + ) -> Result { + let df = if matches!( + self.logical_plan(), + LogicalPlan::Explain(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::Analyze(_) + ) { + // Show multi-line output without truncation for plans like `EXPLAIN` + options.max_row_height = usize::MAX; + + // We don't want to apply an additional .limit() to plans like `Explain` + // as that will trigger an internal error: Unsupported logical plan: Explain must be root of the plan + self + } else { + // Apply limit if specified + self.limit(0, limit)? + }; + + let schema_without_qualifiers = df.schema().clone().strip_qualifiers(); + let schema = schema_without_qualifiers.as_arrow(); + let batches = df.collect().await?; + let mut out = Vec::new(); + show_batches(ctx, &mut out, schema, batches, options)?; + String::from_utf8(out).map_err(|e| DataFusionError::External(Box::new(e))) + } + + async fn write_geoparquet( + self, + ctx: &SedonaContext, + path: &str, + options: SedonaWriteOptions, + writer_options: Option, + ) -> Result, DataFusionError> { + if options.insert_op != InsertOp::Append { + return not_impl_err!( + "{} is not implemented for DataFrame::write_geoparquet.", + options.insert_op + ); + } + + let format = if let Some(parquet_opts) = writer_options { + Arc::new(GeoParquetFormatFactory::new_with_options(parquet_opts)) + } else { + Arc::new(GeoParquetFormatFactory::new()) + }; + + let file_type = format_as_file_type(format); + + let plan = if options.sort_by.is_empty() { + self.into_unoptimized_plan() + } else { + LogicalPlanBuilder::from(self.into_unoptimized_plan()) + .sort(options.sort_by)? + .build()? + }; + + let plan = LogicalPlanBuilder::copy_to( + plan, + path.into(), + file_type, + Default::default(), + options.partition_by, + )? + .build()?; + + DataFrame::new(ctx.ctx.state(), plan).collect().await + } +} + +/// A Sedona-specific copy of [DataFrameWriteOptions] +/// +/// This is needed because [DataFrameWriteOptions] has private fields, so we +/// can't use it in our interfaces. This object can be converted to a +/// [DataFrameWriteOptions] using `.into()`. +pub struct SedonaWriteOptions { + /// Controls how new data should be written to the table, determining whether + /// to append, overwrite, or replace existing data. + pub insert_op: InsertOp, + /// Controls if all partitions should be coalesced into a single output file + /// Generally will have slower performance when set to true. + pub single_file_output: bool, + /// Sets which columns should be used for hive-style partitioned writes by name. + /// Can be set to empty vec![] for non-partitioned writes. + pub partition_by: Vec, + /// Sets which columns should be used for sorting the output by name. + /// Can be set to empty vec![] for non-sorted writes. + pub sort_by: Vec, +} + +impl From for DataFrameWriteOptions { + fn from(value: SedonaWriteOptions) -> Self { + DataFrameWriteOptions::new() + .with_insert_operation(value.insert_op) + .with_single_file_output(value.single_file_output) + .with_partition_by(value.partition_by) + .with_sort_by(value.sort_by) + } +} + +impl SedonaWriteOptions { + /// Create a new SedonaWriteOptions with default values + pub fn new() -> Self { + SedonaWriteOptions { + insert_op: InsertOp::Append, + single_file_output: false, + partition_by: vec![], + sort_by: vec![], + } + } + + /// Set the insert operation + pub fn with_insert_operation(mut self, insert_op: InsertOp) -> Self { + self.insert_op = insert_op; + self + } + + /// Set the single_file_output value to true or false + pub fn with_single_file_output(mut self, single_file_output: bool) -> Self { + self.single_file_output = single_file_output; + self + } + + /// Sets the partition_by columns for output partitioning + pub fn with_partition_by(mut self, partition_by: Vec) -> Self { + self.partition_by = partition_by; + self + } + + /// Sets the sort_by columns for output sorting + pub fn with_sort_by(mut self, sort_by: Vec) -> Self { + self.sort_by = sort_by; + self + } +} + +impl Default for SedonaWriteOptions { + fn default() -> Self { + Self::new() + } +} + // Because Dialect/dialect_from_str is not marked as Send, using the async // function in certain contexts will fail to compile. Here we use a wrapper // to ensure that that the Dialect can be specified and parsed in any async @@ -536,13 +720,10 @@ impl ThreadSafeDialect { #[cfg(test)] mod tests { - use arrow_array::{ - create_array, ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, - }; + + use arrow_array::{create_array, ArrayRef, RecordBatchIterator, RecordBatchReader}; use arrow_schema::{DataType, Field, Schema}; - use async_trait::async_trait; use datafusion::assert_batches_eq; - use datafusion_common::not_impl_err; use sedona_datasource::spec::{Object, OpenReaderArgs}; use sedona_schema::{ crs::{deserialize_crs, lnglat}, @@ -550,6 +731,7 @@ mod tests { schema::SedonaSchema, }; use sedona_testing::data::test_geoparquet; + use tempfile::tempdir; use super::*; @@ -749,6 +931,77 @@ mod tests { Ok(()) } + #[tokio::test] + async fn show() { + let ctx = SedonaContext::new(); + let tbl = ctx + .sql("SELECT 1 as one") + .await + .unwrap() + .show_sedona(&ctx, None, DisplayTableOptions::default()) + .await + .unwrap(); + + #[rustfmt::skip] + assert_eq!( + tbl.lines().collect::>(), + vec![ + "+-----+", + "| one |", + "+-----+", + "| 1 |", + "+-----+" + ] + ); + } + + #[tokio::test] + async fn show_explain() { + let ctx = SedonaContext::new(); + for limit in [None, Some(10)] { + let tbl = ctx + .sql("EXPLAIN SELECT 1 as one") + .await + .unwrap() + .show_sedona(&ctx, limit, DisplayTableOptions::default()) + .await + .unwrap(); + + #[rustfmt::skip] + assert_eq!( + tbl.lines().collect::>(), + vec![ + "+---------------+---------------------------------+", + "| plan_type | plan |", + "+---------------+---------------------------------+", + "| logical_plan | Projection: Int64(1) AS one |", + "| | EmptyRelation: rows=1 |", + "| physical_plan | ProjectionExec: expr=[1 as one] |", + "| | PlaceholderRowExec |", + "| | |", + "+---------------+---------------------------------+", + ] + ); + } + } + + #[tokio::test] + async fn write_geoparquet() { + let tmpdir = tempdir().unwrap(); + let tmp_parquet = tmpdir.path().join("tmp.parquet"); + let ctx = SedonaContext::new(); + ctx.sql("SELECT 1 as one") + .await + .unwrap() + .write_parquet( + &tmp_parquet.to_string_lossy(), + DataFrameWriteOptions::default(), + None, + ) + .await + .unwrap(); + } + #[tokio::test] async fn geoparquet_format() { // Make sure that our context can be set up to identify and read diff --git a/rust/sedona/src/lib.rs b/rust/sedona/src/lib.rs index e29bbe2ecb..81c103f9f4 100644 --- a/rust/sedona/src/lib.rs +++ b/rust/sedona/src/lib.rs @@ -17,7 +17,6 @@ mod catalog; pub mod context; pub mod context_builder; -pub mod dataframe; mod exec; pub mod memory_pool; mod object_storage; From bb6b1f9219aa2ff1011907e4a77a1ccbb8cfe9d3 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 26 Jun 2026 21:48:28 -0500 Subject: [PATCH 17/55] flip the switch on Python --- Cargo.lock | 1 + c/sedona-extension/src/execution_plan.rs | 6 +++-- c/sedona-extension/src/table_provider.rs | 3 +++ python/sedonadb/Cargo.toml | 1 + python/sedonadb/python/sedonadb/dataframe.py | 8 +++---- python/sedonadb/src/dataframe.rs | 20 ++++++++--------- python/sedonadb/src/import_from.rs | 23 +++++++++++++------- python/sedonadb/tests/test_dataframe.py | 2 +- 8 files changed, 39 insertions(+), 25 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dce2da1700..aef794658d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6415,6 +6415,7 @@ dependencies = [ "sedona-adbc", "sedona-datasource", "sedona-expr", + "sedona-extension", "sedona-gdal", "sedona-geometry", "sedona-geoparquet", diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 25c94318db..f0e2c53257 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -393,8 +393,10 @@ impl ExecutionPlan for ImportedSedonaCExec { } fn maintains_input_order(&self) -> Vec { - get_plan_property::, ()>(&self.inner, "maintains_input_order", None) - .unwrap_or_default() + // ImportedSedonaCExec has no children (it's a leaf node), so we return + // an empty Vec. The inner plan's maintains_input_order doesn't apply + // since we don't expose the inner plan's children. + vec![] } fn metrics(&self) -> Option { diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 7201626ee3..ad1dac73a6 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -59,6 +59,9 @@ impl Debug for ExportedTableProvider { impl ExportedTableProvider { /// Create a new ExportedTableProvider from a TableProvider. + /// + /// The session is used during scan operations and must support physical planning + /// if the inner TableProvider requires it (e.g., for Views). pub fn new( inner: Arc, session: Arc, diff --git a/python/sedonadb/Cargo.toml b/python/sedonadb/Cargo.toml index b8df136b93..cd64fbd2ce 100644 --- a/python/sedonadb/Cargo.toml +++ b/python/sedonadb/Cargo.toml @@ -50,6 +50,7 @@ pyo3 = { version = "0.25.1" } sedona = { workspace = true } sedona-adbc = { workspace = true } sedona-datasource = { workspace = true } +sedona-extension = { workspace = true } sedona-geometry = { workspace = true } sedona-expr = { workspace = true } sedona-geoparquet = { workspace = true } diff --git a/python/sedonadb/python/sedonadb/dataframe.py b/python/sedonadb/python/sedonadb/dataframe.py index a0e6675028..c2ea812dad 100644 --- a/python/sedonadb/python/sedonadb/dataframe.py +++ b/python/sedonadb/python/sedonadb/dataframe.py @@ -1215,8 +1215,8 @@ def to_memtable(self) -> "DataFrame": """ return DataFrame(self._ctx, self._impl.to_memtable(self._ctx._impl)) - def __datafusion_table_provider__(self): - return self._impl.__datafusion_table_provider__() + def __sedonadb_table_provider__(self): + return self._impl.__sedonadb_table_provider__() def to_arrow_table(self, schema: Any = None) -> "pa.Table": """Execute and collect results as a PyArrow Table @@ -1601,7 +1601,7 @@ def _create_data_frame(ctx, obj, schema) -> DataFrame: """ # If we're dealing with an anonymous data frame on the same context, # just return it. Otherwise, fall back to the default interpretation - # (which uses __datafusion_table_provider__). + # (which uses __sedonadb_table_provider__). if isinstance(obj, DataFrame) and obj._ctx is ctx and schema is None: return obj @@ -1614,7 +1614,7 @@ def _create_data_frame(ctx, obj, schema) -> DataFrame: return SPECIAL_CASED_SCANS[type_name](ctx, obj, schema) # The default implementation handles objects that implement - # __datafusion_table_provider__ or __arrow_c_stream__. For objects implementing + # __sedonadb_table_provider__ or __arrow_c_stream__. For objects implementing # __arrow_c_stream__, this currently will only work for a single scan (i.e., # the returned data frame can't be previewed before the query is computed). return _scan_default(ctx, obj, schema) diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs index 761b4f171f..db4be30df5 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -27,9 +27,7 @@ use datafusion::config::ConfigField; use datafusion::logical_expr::SortExpr; use datafusion::prelude::{DataFrame, SessionContext}; use datafusion_common::{Column, DataFusionError, ParamValues}; -use datafusion_execution::TaskContextProvider; use datafusion_expr::{ExplainFormat, ExplainOption, Expr, JoinType, LogicalPlanBuilder}; -use datafusion_ffi::table_provider::FFI_TableProvider; use futures::lock::Mutex; use futures::TryStreamExt; use pyo3::prelude::*; @@ -621,20 +619,22 @@ impl InternalDataFrame { Ok(InternalDataFrame::new(df, self.runtime.clone())) } - fn __datafusion_table_provider__<'py>( + fn __sedonadb_table_provider__<'py>( &self, py: Python<'py>, ) -> Result, PySedonaError> { - let name = cr"datafusion_table_provider".into(); + let name = cr"sedonadb_table_provider".into(); let provider = self.inner.clone().into_view(); - let ctx = Arc::new(SessionContext::new()) as Arc; - let ffi_provider = FFI_TableProvider::new( + // Create a session context for FFI - the consuming side will use its own + // session for actual execution, this is just needed for the FFI interface. + let ctx = SessionContext::new(); + let session = Arc::new(ctx.state()); + let exported = sedona_extension::table_provider::ExportedTableProvider::new( provider, - true, - Some(self.runtime.handle().clone()), - &ctx, - None, + session, + self.runtime.handle().clone(), ); + let ffi_provider: sedona_extension::extension::SedonaCTableProvider = exported.into(); Ok(PyCapsule::new(py, ffi_provider, Some(name))?) } } diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs index 530efa3a8e..ff837353e1 100644 --- a/python/sedonadb/src/import_from.rs +++ b/python/sedonadb/src/import_from.rs @@ -28,12 +28,12 @@ use arrow_schema::{Field, Schema}; use datafusion::catalog::TableProvider; use datafusion_common::{metadata::ScalarAndMetadata, ScalarValue}; use datafusion_expr::expr::FieldMetadata; -use datafusion_ffi::table_provider::FFI_TableProvider; use pyo3::{ types::{PyAnyMethods, PyCapsule, PyCapsuleMethods}, Bound, PyAny, Python, }; use sedona::record_batch_reader_provider::RecordBatchReaderProvider; +use sedona_extension::{extension::SedonaCTableProvider, table_provider::ImportedTableProvider}; use sedona_schema::{ datatypes::SedonaType, matchers::{ArgMatcher, TypeMatcher}, @@ -46,8 +46,8 @@ pub fn import_table_provider_from_any<'py>( obj: &Bound, requested_schema: Option<&Bound>, ) -> Result, PySedonaError> { - if obj.hasattr("__datafusion_table_provider__")? { - let provider = import_ffi_table_provider(obj)?; + if obj.hasattr("__sedonadb_table_provider__")? { + let provider = import_sedona_ffi_table_provider(obj)?; Ok(provider) } else if obj.hasattr("__arrow_c_stream__")? { let reader = import_arrow_array_stream(py, obj, requested_schema)?; @@ -59,14 +59,21 @@ pub fn import_table_provider_from_any<'py>( } } -pub fn import_ffi_table_provider( +pub fn import_sedona_ffi_table_provider( obj: &Bound, ) -> Result, PySedonaError> { - let capsule = obj.getattr("__datafusion_table_provider__")?.call0()?; + let capsule = obj.getattr("__sedonadb_table_provider__")?.call0()?; let contents = - check_pycapsule(&capsule, "datafusion_table_provider")? as *mut FFI_TableProvider; - let provider = Arc::::from(unsafe { contents.as_ref().unwrap() }); - Ok(provider) + check_pycapsule(&capsule, "sedonadb_table_provider")? as *mut SedonaCTableProvider; + // Move the SedonaCTableProvider out of the capsule into our ImportedTableProvider. + // We null out the release function to prevent double-free when the capsule is dropped. + let ffi_provider = unsafe { + let provider = std::ptr::read(contents); + (*contents).release = None; + provider + }; + let provider = ImportedTableProvider::try_new(ffi_provider)?; + Ok(Arc::new(provider)) } pub fn import_arrow_array_stream<'py>( diff --git a/python/sedonadb/tests/test_dataframe.py b/python/sedonadb/tests/test_dataframe.py index 768ca9e5d2..6f183ca999 100644 --- a/python/sedonadb/tests/test_dataframe.py +++ b/python/sedonadb/tests/test_dataframe.py @@ -36,7 +36,7 @@ def test_dataframe_from_dataframe(con): # On a separate context the table should still be collected the same # but should be a separate Python reference. This also has the effect - # of testing the __datafusion_table_provider__ interface. + # of testing the __sedonadb_table_provider__ interface. new_con = sedonadb.connect() new_df = new_con.create_data_frame(df) assert new_df is not df From 593af3584a50c21513348fd568ddce4cab92eeec Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 26 Jun 2026 22:43:36 -0500 Subject: [PATCH 18/55] add non-trivial test --- python/sedonadb/tests/test_dataframe_ffi.py | 123 ++++++++++++++++++++ 1 file changed, 123 insertions(+) create mode 100644 python/sedonadb/tests/test_dataframe_ffi.py diff --git a/python/sedonadb/tests/test_dataframe_ffi.py b/python/sedonadb/tests/test_dataframe_ffi.py new file mode 100644 index 0000000000..9d0c54eee3 --- /dev/null +++ b/python/sedonadb/tests/test_dataframe_ffi.py @@ -0,0 +1,123 @@ +# 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 pandas as pd +import pytest + +import sedonadb +from sedonadb.testing import skip_if_not_exists + + +# Test cases: (producer_sql, consumer_sql, description) +FFI_TEST_CASES = [ + # Simple select all + pytest.param( + "SELECT * FROM water_point", + "SELECT * FROM df_producer", + id="select_all", + ), + # Projection - select specific columns + pytest.param( + 'SELECT "OBJECTID", "FEAT_CODE" FROM water_point', + 'SELECT * FROM df_producer ORDER BY "OBJECTID" LIMIT 10', + id="projection", + ), + # Filter on producer side + pytest.param( + 'SELECT "OBJECTID", "FEAT_CODE" FROM water_point WHERE "OBJECTID" > 100', + 'SELECT * FROM df_producer ORDER BY "OBJECTID" LIMIT 5', + id="filter_producer", + ), + # Filter on consumer side + pytest.param( + 'SELECT "OBJECTID", "FEAT_CODE" FROM water_point', + 'SELECT * FROM df_producer WHERE "OBJECTID" < 50 ORDER BY "OBJECTID"', + id="filter_consumer", + ), + # Sort on producer side (descending) + pytest.param( + 'SELECT "OBJECTID", "FEAT_CODE" FROM water_point ORDER BY "OBJECTID" DESC LIMIT 20', + "SELECT * FROM df_producer", + id="sort_producer_desc", + ), + # Sort on consumer side + pytest.param( + 'SELECT "OBJECTID", "FEAT_CODE" FROM water_point', + 'SELECT * FROM df_producer ORDER BY "FEAT_CODE", "OBJECTID" DESC LIMIT 10', + id="sort_consumer", + ), + # Limit on producer + pytest.param( + "SELECT * FROM water_point LIMIT 5", + "SELECT * FROM df_producer", + id="limit_producer", + ), + # Aggregate on consumer + pytest.param( + 'SELECT "OBJECTID", "FEAT_CODE" FROM water_point', + 'SELECT "FEAT_CODE", COUNT(*) as cnt FROM df_producer GROUP BY "FEAT_CODE" ORDER BY cnt DESC LIMIT 5', + id="aggregate_consumer", + ), + # Aggregate on producer side + pytest.param( + 'SELECT "FEAT_CODE", COUNT(*) as cnt, MIN("OBJECTID") as min_id, MAX("OBJECTID") as max_id FROM water_point GROUP BY "FEAT_CODE"', + 'SELECT * FROM df_producer ORDER BY cnt DESC, "FEAT_CODE" LIMIT 10', + id="aggregate_producer", + ), + # Spatial distance join on producer side (with pre-filtered data to keep it fast) + pytest.param( + """ + WITH subset AS (SELECT * FROM water_point WHERE "OBJECTID" < 1000) + SELECT a."OBJECTID" as id_a, b."OBJECTID" as id_b, ST_Distance(a.geometry, b.geometry) as dist + FROM subset a + JOIN subset b ON ST_DWithin(a.geometry, b.geometry, 100) + ORDER BY a."OBJECTID", b."OBJECTID" + """, + "SELECT * FROM df_producer ORDER BY id_a, id_b LIMIT 10", + id="spatial_join_producer", + ), + # Multiple operations chained + pytest.param( + 'SELECT "OBJECTID", "FEAT_CODE" FROM water_point WHERE "OBJECTID" BETWEEN 10 AND 100 ORDER BY "OBJECTID"', + 'SELECT "FEAT_CODE", COUNT(*) as cnt FROM df_producer GROUP BY "FEAT_CODE" HAVING COUNT(*) > 1 ORDER BY cnt DESC', + id="filter_then_aggregate", + ), +] + + +@pytest.mark.parametrize("producer_sql,consumer_sql", FFI_TEST_CASES) +def test_ffi_roundtrip(geoarrow_data, producer_sql, consumer_sql): + # Use a real file with a reasonable number of rows so that parallelism + # and multiple batches are invoked + path = geoarrow_data / "ns-water" / "files" / "ns-water_water-point_geo.parquet" + skip_if_not_exists(path) + + sd_producer = sedonadb.connect() + sd_consumer = sedonadb.connect() + + sd_producer.read_parquet(path).to_view("water_point") + df_producer = sd_producer.sql(producer_sql) + + # Run the query on the producer (no FFI) + df_producer.to_view("df_producer") + result_no_ffi = sd_producer.sql(consumer_sql).to_pandas() + + # Run the query on the consumer (separated via FFI) + sd_consumer.create_data_frame(df_producer).to_view("df_producer") + result_over_ffi = sd_consumer.sql(consumer_sql).to_pandas() + + pd.testing.assert_frame_equal(result_no_ffi, result_over_ffi) From 9aa06bd50cdd4d0b36f9c9092e82a1cd620a475f Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 26 Jun 2026 23:01:21 -0500 Subject: [PATCH 19/55] add the importedcexec to the explain --- c/sedona-extension/src/execution_plan.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index f0e2c53257..52ba24a948 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -355,8 +355,9 @@ impl DisplayAs for ImportedSedonaCExec { DisplayFormatType::TreeRender => "display_tree_render", }; + // Always show the wrapper name, with the inner plan's display info if let Ok(display_str) = get_plan_string_property(&self.inner, property) { - write!(f, "{}", display_str) + write!(f, "ImportedSedonaCExec: {}", display_str) } else { write!(f, "ImportedSedonaCExec") } @@ -851,20 +852,21 @@ mod tests { } } + // ImportedSedonaCExec shows itself with the inner plan's display assert_eq!( format!("{}", DisplayAsFormat(&imported, DisplayFormatType::Default)), - "DummyExec: default format" + "ImportedSedonaCExec: DummyExec: default format" ); assert_eq!( format!("{}", DisplayAsFormat(&imported, DisplayFormatType::Verbose)), - "DummyExec: verbose format with schema" + "ImportedSedonaCExec: DummyExec: verbose format with schema" ); assert_eq!( format!( "{}", DisplayAsFormat(&imported, DisplayFormatType::TreeRender) ), - "DummyExec: tree render format" + "ImportedSedonaCExec: DummyExec: tree render format" ); } From 47166d63511a01c2bce0bd86f894271c8a395e10 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Sun, 28 Jun 2026 21:40:51 -0500 Subject: [PATCH 20/55] migrate the other reader --- Cargo.lock | 1 + r/sedonadb/src/rust/Cargo.toml | 1 + r/sedonadb/src/rust/src/dataframe.rs | 4 ++-- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index aef794658d..8786b4edf2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6459,6 +6459,7 @@ dependencies = [ "sedona", "sedona-adbc", "sedona-expr", + "sedona-extension", "sedona-geoparquet", "sedona-proj", "sedona-schema", diff --git a/r/sedonadb/src/rust/Cargo.toml b/r/sedonadb/src/rust/Cargo.toml index 3c1cadbdcd..5dc64c9a81 100644 --- a/r/sedonadb/src/rust/Cargo.toml +++ b/r/sedonadb/src/rust/Cargo.toml @@ -37,6 +37,7 @@ savvy-ffi = "*" sedona = { workspace = true } sedona-adbc = { workspace = true } sedona-expr = { workspace = true } +sedona-extension = { workspace = true } sedona-geoparquet = { workspace = true } sedona-proj = { workspace = true } sedona-schema = { workspace = true } diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index 20af9bf5bc..913c13e39c 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -29,7 +29,7 @@ use datafusion_expr::{select_expr::SelectExpr, Expr, SortExpr}; use datafusion_ffi::table_provider::FFI_TableProvider; use savvy::{savvy, savvy_err, sexp, IntoExtPtrSexp, Result}; use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; -use sedona::reader::SedonaStreamReader; +use sedona_extension::utils::StreamingRecordBatchReader; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_geoparquet::options::TableGeoParquetOptions; use sedona_schema::schema::SedonaSchema; @@ -115,7 +115,7 @@ impl InternalDataFrame { async move { inner.execute_stream().await }, )??; - let reader = SedonaStreamReader::new(self.runtime.clone(), stream); + let reader = StreamingRecordBatchReader::new(stream, self.runtime.handle().clone()); let reader: Box = Box::new(reader); let mut ffi_stream = FFI_ArrowArrayStream::new(reader); From 7a43fbadb3c325466cd9dbd28c61d1397d7aac78 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Sun, 28 Jun 2026 21:52:50 -0500 Subject: [PATCH 21/55] replace the adbc reader --- Cargo.lock | 1 + rust/sedona-adbc/Cargo.toml | 1 + rust/sedona-adbc/src/statement.rs | 9 +-- rust/sedona/src/lib.rs | 1 - rust/sedona/src/reader.rs | 126 ------------------------------ 5 files changed, 6 insertions(+), 132 deletions(-) delete mode 100644 rust/sedona/src/reader.rs diff --git a/Cargo.lock b/Cargo.lock index 8786b4edf2..4a345edce1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5693,6 +5693,7 @@ dependencies = [ "datafusion", "futures", "sedona", + "sedona-extension", "tokio", ] diff --git a/rust/sedona-adbc/Cargo.toml b/rust/sedona-adbc/Cargo.toml index 2ebdac2f9d..e25460d823 100644 --- a/rust/sedona-adbc/Cargo.toml +++ b/rust/sedona-adbc/Cargo.toml @@ -38,4 +38,5 @@ arrow-schema = { workspace = true } datafusion = { workspace = true } futures = { workspace = true } sedona = { workspace = true } +sedona-extension = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread"] } diff --git a/rust/sedona-adbc/src/statement.rs b/rust/sedona-adbc/src/statement.rs index 159cb661a7..59927a24d0 100644 --- a/rust/sedona-adbc/src/statement.rs +++ b/rust/sedona-adbc/src/statement.rs @@ -17,7 +17,8 @@ use adbc_core::PartitionedResult; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::Schema; -use sedona::{context::SedonaContext, reader::SedonaStreamReader}; +use sedona::context::SedonaContext; +use sedona_extension::utils::StreamingRecordBatchReader; use std::sync::Arc; use tokio::runtime::Runtime; @@ -98,10 +99,8 @@ impl Statement for SedonaStatement { self.runtime.block_on(async { let df = self.ctx.sql(&query).await.map_err(from_datafusion_error)?; let stream = df.execute_stream().await.map_err(from_datafusion_error)?; - Ok( - Box::new(SedonaStreamReader::new(self.runtime.clone(), stream)) - as Box, - ) + let reader = StreamingRecordBatchReader::new(stream, self.runtime.handle().clone()); + Ok(Box::new(reader) as Box) }) } else { Err(Error::with_message_and_status( diff --git a/rust/sedona/src/lib.rs b/rust/sedona/src/lib.rs index 81c103f9f4..dab1169a04 100644 --- a/rust/sedona/src/lib.rs +++ b/rust/sedona/src/lib.rs @@ -23,7 +23,6 @@ mod object_storage; pub mod pool_type; pub mod projected_reader; pub mod random_geometry_provider; -pub mod reader; pub mod record_batch_reader_provider; pub mod show; pub mod size_parser; diff --git a/rust/sedona/src/reader.rs b/rust/sedona/src/reader.rs deleted file mode 100644 index 2b4673fe15..0000000000 --- a/rust/sedona/src/reader.rs +++ /dev/null @@ -1,126 +0,0 @@ -// 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. -use futures::TryStreamExt; -use std::sync::Arc; - -use arrow_array::{RecordBatch, RecordBatchReader}; -use arrow_schema::{ArrowError, SchemaRef}; -use datafusion::execution::SendableRecordBatchStream; -use tokio::runtime::Runtime; - -/// Utility to convert a [SendableRecordBatchStream] into a [RecordBatchReader] -/// -/// This is needed for clients like ADBC and Python where this is the format -/// that is required for export. -pub struct SedonaStreamReader { - runtime: Arc, - stream: SendableRecordBatchStream, -} - -impl SedonaStreamReader { - pub fn new(runtime: Arc, stream: SendableRecordBatchStream) -> Self { - Self { runtime, stream } - } -} - -impl Iterator for SedonaStreamReader { - type Item = std::result::Result; - - fn next(&mut self) -> Option { - loop { - match self.runtime.block_on(self.stream.try_next()) { - Ok(maybe_batch) => match maybe_batch { - Some(batch) => { - if batch.num_rows() == 0 { - continue; - } - - return Some(Ok(batch)); - } - None => return None, - }, - Err(err) => return Some(Err(ArrowError::ExternalError(Box::new(err)))), - } - } - } -} - -impl RecordBatchReader for SedonaStreamReader { - fn schema(&self) -> SchemaRef { - self.stream.schema() - } -} - -#[cfg(test)] -mod test { - - use arrow_array::record_batch; - use arrow_schema::{DataType, Field, Schema}; - use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - - use crate::context::SedonaContext; - - use super::*; - - #[test] - fn reader() { - let runtime = Arc::new(tokio::runtime::Runtime::new().unwrap()); - let ctx = SedonaContext::new(); - let df = runtime.block_on(ctx.sql("SELECT 1 as one")).unwrap(); - let expected_batches = runtime.block_on(df.clone().collect()).unwrap(); - assert_eq!(expected_batches.len(), 1); - - let stream = runtime.block_on(df.execute_stream()).unwrap(); - let mut reader = SedonaStreamReader::new(runtime, stream); - - let expected_schema = Arc::new(Schema::new([ - Field::new("one", DataType::Int64, false).into() - ])); - assert_eq!(reader.schema(), expected_schema); - - assert_eq!(reader.next().unwrap().unwrap(), expected_batches[0]); - assert!(reader.next().is_none()); - } - - #[test] - fn reader_empty_chunks() { - let runtime = Arc::new(tokio::runtime::Runtime::new().unwrap()); - - let batch0 = record_batch!( - ("a", Int32, [1, 2, 3]), - ("b", Float64, [Some(4.0), None, Some(5.0)]) - ) - .expect("created batch"); - let schema = batch0.schema(); - - let batch1 = RecordBatch::new_empty(schema.clone()); - let batch2 = batch0.clone(); - - let stream = futures::stream::iter(vec![ - Ok(batch0.clone()), - Ok(batch1.clone()), - Ok(batch2.clone()), - ]); - let adapter = RecordBatchStreamAdapter::new(schema, stream); - let batch_stream: SendableRecordBatchStream = Box::pin(adapter); - - let mut reader = SedonaStreamReader::new(runtime, batch_stream); - assert_eq!(reader.next().unwrap().unwrap(), batch0); - assert_eq!(reader.next().unwrap().unwrap(), batch2); - assert!(reader.next().is_none()); - } -} From d02afad99e96a46be454390280ea01807a7f817e Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Sun, 28 Jun 2026 22:24:47 -0500 Subject: [PATCH 22/55] replace Python reader --- c/sedona-extension/src/utils.rs | 245 ++++++++++++++++++++++++++++--- python/sedonadb/src/dataframe.rs | 4 +- python/sedonadb/src/reader.rs | 69 ++++----- python/sedonadb/src/runtime.rs | 24 --- 4 files changed, 253 insertions(+), 89 deletions(-) diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index 2cf0cea724..aae79d6de2 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -353,6 +353,8 @@ pub struct StreamingRecordBatchReader { runtime: tokio::runtime::Handle, cancel_checker: Option, cancelled: bool, + skip_empty_batches: bool, + periodic_check_interval: Option, } impl StreamingRecordBatchReader { @@ -364,6 +366,8 @@ impl StreamingRecordBatchReader { runtime, cancel_checker: None, cancelled: false, + skip_empty_batches: false, + periodic_check_interval: None, } } @@ -383,6 +387,102 @@ impl StreamingRecordBatchReader { runtime, cancel_checker: Some(cancel_checker), cancelled: false, + skip_empty_batches: false, + periodic_check_interval: None, + } + } + + /// Set whether to skip empty batches (batches with 0 rows). + /// + /// When enabled, the iterator will automatically skip over any batches + /// that have no rows and continue to the next batch. + pub fn with_skip_empty_batches(mut self, skip: bool) -> Self { + self.skip_empty_batches = skip; + self + } + + /// Set a periodic interval for checking cancellation during batch fetches. + /// + /// When set, the cancel checker will be called periodically at this interval + /// even while waiting for a single batch to be fetched. This is useful for + /// Python where we need to periodically check for signals (Ctrl+C) during + /// long-running operations. + /// + /// Without this, the cancel checker is only called between batch fetches. + pub fn with_periodic_check_interval(mut self, interval: std::time::Duration) -> Self { + self.periodic_check_interval = Some(interval); + self + } + + fn fetch_next_batch(&mut self) -> Option> { + let stream = &mut self.stream; + let runtime = &self.runtime; + + match &self.periodic_check_interval { + Some(interval) => { + // Use tokio::select! to periodically check for cancellation + let interval = *interval; + let cancel_checker = &self.cancel_checker; + + std::thread::scope(|s| { + s.spawn(|| { + runtime.block_on(async { + use futures::StreamExt; + tokio::pin!(stream); + loop { + tokio::select! { + res = stream.next() => { + return match res { + Some(Ok(batch)) => Some(Ok(batch)), + Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), + None => None, + }; + } + _ = tokio::time::sleep(interval) => { + if let Some(ref checker) = cancel_checker { + if checker() { + return Some(Err(ArrowError::ExternalError(Box::new( + std::io::Error::new( + std::io::ErrorKind::Interrupted, + "Operation cancelled", + ), + )))); + } + } + // Continue waiting for the batch + } + } + } + }) + }) + .join() + .unwrap_or_else(|_| { + Some(Err(ArrowError::InvalidArgumentError( + "Iterator thread panicked".to_string(), + ))) + }) + }) + } + None => { + // Simple blocking fetch without periodic checking + std::thread::scope(|s| { + s.spawn(|| { + runtime.block_on(async { + match stream.next().await { + Some(Ok(batch)) => Some(Ok(batch)), + Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), + None => None, + } + }) + }) + .join() + .unwrap_or_else(|_| { + Some(Err(ArrowError::InvalidArgumentError( + "Iterator thread panicked".to_string(), + ))) + }) + }) + } } } } @@ -397,6 +497,7 @@ impl Iterator for StreamingRecordBatchReader { } // Check for cancellation before fetching the next batch + // (periodic checking during fetch is handled by fetch_next_batch if configured) if let Some(ref checker) = self.cancel_checker { if checker() { self.cancelled = true; @@ -406,26 +507,24 @@ impl Iterator for StreamingRecordBatchReader { } } - let stream = &mut self.stream; - let runtime = &self.runtime; - - std::thread::scope(|s| { - s.spawn(|| { - runtime.block_on(async { - match stream.next().await { - Some(Ok(batch)) => Some(Ok(batch)), - Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), - None => None, + loop { + match self.fetch_next_batch() { + Some(Ok(batch)) => { + if self.skip_empty_batches && batch.num_rows() == 0 { + continue; } - }) - }) - .join() - .unwrap_or_else(|_| { - Some(Err(ArrowError::InvalidArgumentError( - "Iterator thread panicked".to_string(), - ))) - }) - }) + return Some(Ok(batch)); + } + Some(Err(e)) => { + // Check if this was a cancellation error from periodic checking + if e.to_string().contains("Operation cancelled") { + self.cancelled = true; + } + return Some(Err(e)); + } + None => return None, + } + } } } @@ -662,4 +761,112 @@ mod tests { err ); } + + /// Create a stream that yields some empty batches. + fn create_stream_with_empty_batches( + batch_row_counts: Vec, + ) -> (SchemaRef, SendableRecordBatchStream) { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let schema_clone = schema.clone(); + + let stream = futures::stream::iter(batch_row_counts.into_iter().enumerate()).map( + move |(i, row_count)| { + let schema = schema_clone.clone(); + let array: Int32Array = (0..row_count).map(|_| i as i32).collect(); + Ok(RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()) + }, + ); + + ( + schema.clone(), + Box::pin( + datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), + ), + ) + } + + #[tokio::test] + async fn test_streaming_reader_skip_empty_batches() { + let runtime = tokio::runtime::Handle::current(); + // Create stream with: 2 rows, 0 rows, 3 rows, 0 rows, 0 rows, 1 row + let (_schema, stream) = create_stream_with_empty_batches(vec![2, 0, 3, 0, 0, 1]); + + let reader = StreamingRecordBatchReader::new(stream, runtime).with_skip_empty_batches(true); + let batches: Vec<_> = reader.collect(); + + // Should only get 3 batches (the non-empty ones) + assert_eq!(batches.len(), 3); + + // Check row counts: 2, 3, 1 + assert_eq!(batches[0].as_ref().unwrap().num_rows(), 2); + assert_eq!(batches[1].as_ref().unwrap().num_rows(), 3); + assert_eq!(batches[2].as_ref().unwrap().num_rows(), 1); + } + + #[tokio::test] + async fn test_streaming_reader_no_skip_empty_batches() { + let runtime = tokio::runtime::Handle::current(); + // Create stream with: 2 rows, 0 rows, 3 rows + let (_schema, stream) = create_stream_with_empty_batches(vec![2, 0, 3]); + + // Default: skip_empty_batches is false + let reader = StreamingRecordBatchReader::new(stream, runtime); + let batches: Vec<_> = reader.collect(); + + // Should get all 3 batches including the empty one + assert_eq!(batches.len(), 3); + assert_eq!(batches[0].as_ref().unwrap().num_rows(), 2); + assert_eq!(batches[1].as_ref().unwrap().num_rows(), 0); + assert_eq!(batches[2].as_ref().unwrap().num_rows(), 3); + } + + #[tokio::test] + async fn test_streaming_reader_periodic_check_interval() { + let runtime = tokio::runtime::Handle::current(); + // Create a slow stream where each batch takes 200ms + let (_schema, stream) = create_slow_stream(5, 200); + + // Cancel flag - will be set after 150ms + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_clone = cancelled.clone(); + let cancel_checker: CancelChecker = + Box::new(move || cancelled_clone.load(Ordering::SeqCst)); + + let reader = + StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker) + .with_periodic_check_interval(Duration::from_millis(50)); + + // Spawn a task to set cancelled after 150ms + let cancelled_setter = cancelled.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(150)); + cancelled_setter.store(true, Ordering::SeqCst); + }); + + let batches: Vec<_> = reader.collect(); + + // The first batch takes 200ms but we cancel at 150ms during the fetch + // With periodic_check_interval of 50ms, the check happens at 50ms, 100ms, 150ms + // At 150ms the cancel check should trigger + // So we should get 0 successful batches + 1 cancellation error + assert!( + batches.len() <= 2, + "expected at most 1 batch + 1 error, got {}", + batches.len() + ); + + // At least one should be an error + let has_cancel_error = batches + .iter() + .any(|b| b.is_err() && b.as_ref().unwrap_err().to_string().contains("cancelled")); + assert!( + has_cancel_error, + "should have a cancellation error in batches: {:?}", + batches + ); + } } diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs index db4be30df5..17030f9848 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -43,7 +43,7 @@ use crate::context::InternalContext; use crate::error::PySedonaError; use crate::expr::{PyExpr, PySortExpr}; use crate::import_from::{import_arrow_scalar, import_arrow_schema}; -use crate::reader::PySedonaStreamReader; +use crate::reader::new_py_streaming_reader; use crate::runtime::wait_for_future; use crate::schema::PySedonaSchema; @@ -422,7 +422,7 @@ impl InternalDataFrame { simplify: Option, ) -> Result { let stream = wait_for_future(py, &self.runtime, self.inner.clone().execute_stream())??; - let reader = PySedonaStreamReader::new(self.runtime.clone(), stream); + let reader = new_py_streaming_reader(stream, self.runtime.handle().clone()); let mut reader: Box = Box::new(reader); if simplify.unwrap_or(false) { diff --git a/python/sedonadb/src/reader.rs b/python/sedonadb/src/reader.rs index 4c13e7d042..804c4a4b01 100644 --- a/python/sedonadb/src/reader.rs +++ b/python/sedonadb/src/reader.rs @@ -14,56 +14,37 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. -use futures::TryStreamExt; -use std::sync::Arc; +use std::time::Duration; -use arrow_array::{RecordBatch, RecordBatchReader}; -use arrow_schema::{ArrowError, SchemaRef}; use datafusion::execution::SendableRecordBatchStream; -use tokio::runtime::Runtime; +use pyo3::Python; +use sedona_extension::utils::StreamingRecordBatchReader; -use crate::runtime::wait_for_future_from_rust; +/// Interval for checking Python signals during batch fetches. +const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(2_000); -/// Utility to convert a [SendableRecordBatchStream] into a [RecordBatchReader] +/// Create a Python-aware [StreamingRecordBatchReader] that: +/// - Skips empty batches +/// - Periodically checks for Python cancellation signals (Ctrl+C) /// -/// This is like the SedonaStreamReader except it checks for Python signals such -/// as cancellation. -pub struct PySedonaStreamReader { - runtime: Arc, +/// The reader will check Python signals every 2 seconds during batch fetches, +/// allowing users to interrupt long-running queries. +pub fn new_py_streaming_reader( stream: SendableRecordBatchStream, -} - -impl PySedonaStreamReader { - pub fn new(runtime: Arc, stream: SendableRecordBatchStream) -> Self { - Self { runtime, stream } - } -} - -impl Iterator for PySedonaStreamReader { - type Item = std::result::Result; - - fn next(&mut self) -> Option { - loop { - match wait_for_future_from_rust(&self.runtime, self.stream.try_next()) { - Ok(Ok(maybe_batch)) => match maybe_batch { - Some(batch) => { - if batch.num_rows() == 0 { - continue; - } - - return Some(Ok(batch)); - } - None => return None, - }, - Ok(Err(df_err)) => return Some(Err(ArrowError::ExternalError(Box::new(df_err)))), - Err(py_err) => return Some(Err(ArrowError::ExternalError(Box::new(py_err)))), + runtime: tokio::runtime::Handle, +) -> StreamingRecordBatchReader { + // Create a cancel checker that checks Python signals + let cancel_checker: Box bool + Send + Sync> = Box::new(|| { + Python::with_gil(|py| { + // Run `pass` to process any pending signals, then check for errors + if py.run(cr"pass", None, None).is_err() { + return true; } - } - } -} + py.check_signals().is_err() + }) + }); -impl RecordBatchReader for PySedonaStreamReader { - fn schema(&self) -> SchemaRef { - self.stream.schema() - } + StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker) + .with_skip_empty_batches(true) + .with_periodic_check_interval(INTERVAL_CHECK_SIGNALS) } diff --git a/python/sedonadb/src/runtime.rs b/python/sedonadb/src/runtime.rs index 772f90b1e4..d2ca9a1223 100644 --- a/python/sedonadb/src/runtime.rs +++ b/python/sedonadb/src/runtime.rs @@ -50,27 +50,3 @@ where }) }) } - -// A version of the above except for use from an arbitrary Rust function instead -// of from somewhere that had already acquired the GIL. -pub fn wait_for_future_from_rust(runtime: &Runtime, fut: F) -> Result -where - F: Future + Send, - F::Output: Send, -{ - const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(2_000); - runtime.block_on(async { - tokio::pin!(fut); - loop { - tokio::select! { - res = &mut fut => break Ok(res), - _ = sleep(INTERVAL_CHECK_SIGNALS) => { - Python::with_gil(|py| { - py.run(cr"pass", None, None)?; - py.check_signals() - })?; - } - } - } - }) -} From 47ad3a726cf00b781cdb1cf4914b8596603d93db Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Sun, 28 Jun 2026 22:29:56 -0500 Subject: [PATCH 23/55] fmt --- r/sedonadb/src/rust/src/dataframe.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index 913c13e39c..e1390507dd 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -29,8 +29,8 @@ use datafusion_expr::{select_expr::SelectExpr, Expr, SortExpr}; use datafusion_ffi::table_provider::FFI_TableProvider; use savvy::{savvy, savvy_err, sexp, IntoExtPtrSexp, Result}; use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; -use sedona_extension::utils::StreamingRecordBatchReader; use sedona::show::{DisplayMode, DisplayTableOptions}; +use sedona_extension::utils::StreamingRecordBatchReader; use sedona_geoparquet::options::TableGeoParquetOptions; use sedona_schema::schema::SedonaSchema; use std::{iter::zip, ptr::swap_nonoverlapping, sync::Arc}; @@ -115,7 +115,8 @@ impl InternalDataFrame { async move { inner.execute_stream().await }, )??; - let reader = StreamingRecordBatchReader::new(stream, self.runtime.handle().clone()); + let reader = StreamingRecordBatchReader::new(stream, self.runtime.handle().clone()) + .with_skip_empty_batches(true); let reader: Box = Box::new(reader); let mut ffi_stream = FFI_ArrowArrayStream::new(reader); From b8d962046d8a6ff6eed9a4a30c03e4c98781fcf2 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 14:02:47 -0500 Subject: [PATCH 24/55] address easier comments --- c/sedona-extension/src/execution_plan.rs | 5 + c/sedona-extension/src/sedona_extension.h | 2 +- c/sedona-extension/src/table_provider.rs | 5 + c/sedona-extension/src/utils.rs | 8 +- python/sedonadb/src/dataframe.rs | 2 +- python/sedonadb/src/import_from.rs | 7 +- r/sedonadb/src/rust/src/dataframe.rs | 2 +- rust/sedona/src/context.rs | 198 ++-------------------- rust/sedona/src/lib.rs | 1 + 9 files changed, 35 insertions(+), 195 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 52ba24a948..4bc53bf560 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -315,6 +315,11 @@ impl ImportedSedonaCExec { /// /// This will query the plan for its schema and properties. pub fn try_new(inner: SedonaCExecutionPlan) -> Result { + // Refuse to import a structure without a valid release callback + if inner.release.is_none() { + return sedona_internal_err!("SedonaCExecutionPlan does not have a release callback"); + } + // Get schema let Some(get_schema) = inner.get_schema else { return sedona_internal_err!("SedonaCExecutionPlan does not have get_schema"); diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index b87cf7e787..c55ccbd6a8 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -229,7 +229,7 @@ struct SedonaCExpr { /// \brief Release this instance /// /// Implementations of this callback must set self->release to NULL. - void (*release)(struct SedonaCExecutionPlan* self); + void (*release)(struct SedonaCExpr* self); /// \brief Opaque implementation-specific data void* private_data; diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index ad1dac73a6..f0f371eb22 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -268,6 +268,11 @@ impl Debug for ImportedTableProvider { impl ImportedTableProvider { /// Create a new ImportedTableProvider from a SedonaCTableProvider. pub fn try_new(inner: SedonaCTableProvider) -> Result { + // Refuse to import a structure without a valid release callback + if inner.release.is_none() { + return sedona_internal_err!("SedonaCTableProvider does not have a release callback"); + } + // Get schema via Arrow C Data Interface let Some(get_schema) = inner.get_schema else { return sedona_internal_err!("SedonaCTableProvider does not have get_schema"); diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index aae79d6de2..c388aa6d8d 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -517,8 +517,12 @@ impl Iterator for StreamingRecordBatchReader { } Some(Err(e)) => { // Check if this was a cancellation error from periodic checking - if e.to_string().contains("Operation cancelled") { - self.cancelled = true; + if let ArrowError::ExternalError(box_err) = &e { + if let Some(io_err) = box_err.downcast_ref::() { + if io_err.kind() == std::io::ErrorKind::Interrupted { + self.cancelled = true; + } + } } return Some(Err(e)); } diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs index 17030f9848..7c5272c1d9 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -32,7 +32,7 @@ use futures::lock::Mutex; use futures::TryStreamExt; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict, PyList}; -use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; +use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; use sedona::projected_reader::simplify_record_batch_reader; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_geoparquet::options::TableGeoParquetOptions; diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs index ff837353e1..8ff094d5ec 100644 --- a/python/sedonadb/src/import_from.rs +++ b/python/sedonadb/src/import_from.rs @@ -65,13 +65,16 @@ pub fn import_sedona_ffi_table_provider( let capsule = obj.getattr("__sedonadb_table_provider__")?.call0()?; let contents = check_pycapsule(&capsule, "sedonadb_table_provider")? as *mut SedonaCTableProvider; + // Move the SedonaCTableProvider out of the capsule into our ImportedTableProvider. - // We null out the release function to prevent double-free when the capsule is dropped. + // Clear the structure after reading to prevent double-free when the capsule is dropped. let ffi_provider = unsafe { let provider = std::ptr::read(contents); - (*contents).release = None; + // Clear the entire structure to prevent any accidental use + std::ptr::write_bytes(contents, 0, 1); provider }; + // try_new validates the release callback let provider = ImportedTableProvider::try_new(ffi_provider)?; Ok(Arc::new(provider)) } diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index e1390507dd..e914f4176e 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -28,7 +28,7 @@ use datafusion_expr::JoinType; use datafusion_expr::{select_expr::SelectExpr, Expr, SortExpr}; use datafusion_ffi::table_provider::FFI_TableProvider; use savvy::{savvy, savvy_err, sexp, IntoExtPtrSexp, Result}; -use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; +use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_extension::utils::StreamingRecordBatchReader; use sedona_geoparquet::options::TableGeoParquetOptions; diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index dfc26fe09b..63f06d0300 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -21,15 +21,9 @@ use std::{ use crate::exec::create_plan_from_sql; use crate::object_storage::ensure_object_store_registered_with_options; -use crate::{ - catalog::DynamicObjectStoreCatalog, - random_geometry_provider::RandomGeometryFunction, - show::{show_batches, DisplayTableOptions}, -}; -use arrow_array::RecordBatch; +use crate::{catalog::DynamicObjectStoreCatalog, random_geometry_provider::RandomGeometryFunction}; use arrow_schema::DataType; -use async_trait::async_trait; -use datafusion::datasource::file_format::format_as_file_type; +use datafusion::execution::memory_pool::MemoryLimit; use datafusion::{ common::plan_err, error::{DataFusionError, Result}, @@ -41,11 +35,7 @@ use datafusion::{ prelude::{DataFrame, SessionConfig, SessionContext}, sql::parser::{DFParser, Statement}, }; -use datafusion::{dataframe::DataFrameWriteOptions, execution::memory_pool::MemoryLimit}; -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 parking_lot::Mutex; use sedona_common::{ option::add_sedona_option_extension, sedona_internal_datafusion_err, CrsProviderOption, @@ -55,7 +45,6 @@ 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_geoparquet::options::TableGeoParquetOptions; use sedona_geoparquet::{ format::GeoParquetFormatFactory, provider::{geoparquet_listing_table, GeoParquetReadOptions}, @@ -517,179 +506,6 @@ impl Default for SedonaContext { } } -/// Sedona-specific [`DataFrame`] actions -/// -/// This trait, implemented for [`DataFrame`], extends the DataFrame API to make it -/// ergonomic to work with dataframes that contain geometry columns. Currently these -/// are limited to output functions, as geometry columns currently require special -/// handling when written or exported to an external system. -#[async_trait] -pub trait SedonaDataFrame { - /// Build a table of the first `limit` results in this DataFrame - /// - /// This will limit and execute the query and build a table using [show_batches]. - async fn show_sedona<'a>( - self, - ctx: &SedonaContext, - limit: Option, - options: DisplayTableOptions<'a>, - ) -> Result; - - async fn write_geoparquet( - self, - ctx: &SedonaContext, - path: &str, - options: SedonaWriteOptions, - writer_options: Option, - ) -> Result>; -} - -#[async_trait] -impl SedonaDataFrame for DataFrame { - async fn show_sedona<'a>( - self, - ctx: &SedonaContext, - limit: Option, - mut options: DisplayTableOptions<'a>, - ) -> Result { - let df = if matches!( - self.logical_plan(), - LogicalPlan::Explain(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::Analyze(_) - ) { - // Show multi-line output without truncation for plans like `EXPLAIN` - options.max_row_height = usize::MAX; - - // We don't want to apply an additional .limit() to plans like `Explain` - // as that will trigger an internal error: Unsupported logical plan: Explain must be root of the plan - self - } else { - // Apply limit if specified - self.limit(0, limit)? - }; - - let schema_without_qualifiers = df.schema().clone().strip_qualifiers(); - let schema = schema_without_qualifiers.as_arrow(); - let batches = df.collect().await?; - let mut out = Vec::new(); - show_batches(ctx, &mut out, schema, batches, options)?; - String::from_utf8(out).map_err(|e| DataFusionError::External(Box::new(e))) - } - - async fn write_geoparquet( - self, - ctx: &SedonaContext, - path: &str, - options: SedonaWriteOptions, - writer_options: Option, - ) -> Result, DataFusionError> { - if options.insert_op != InsertOp::Append { - return not_impl_err!( - "{} is not implemented for DataFrame::write_geoparquet.", - options.insert_op - ); - } - - let format = if let Some(parquet_opts) = writer_options { - Arc::new(GeoParquetFormatFactory::new_with_options(parquet_opts)) - } else { - Arc::new(GeoParquetFormatFactory::new()) - }; - - let file_type = format_as_file_type(format); - - let plan = if options.sort_by.is_empty() { - self.into_unoptimized_plan() - } else { - LogicalPlanBuilder::from(self.into_unoptimized_plan()) - .sort(options.sort_by)? - .build()? - }; - - let plan = LogicalPlanBuilder::copy_to( - plan, - path.into(), - file_type, - Default::default(), - options.partition_by, - )? - .build()?; - - DataFrame::new(ctx.ctx.state(), plan).collect().await - } -} - -/// A Sedona-specific copy of [DataFrameWriteOptions] -/// -/// This is needed because [DataFrameWriteOptions] has private fields, so we -/// can't use it in our interfaces. This object can be converted to a -/// [DataFrameWriteOptions] using `.into()`. -pub struct SedonaWriteOptions { - /// Controls how new data should be written to the table, determining whether - /// to append, overwrite, or replace existing data. - pub insert_op: InsertOp, - /// Controls if all partitions should be coalesced into a single output file - /// Generally will have slower performance when set to true. - pub single_file_output: bool, - /// Sets which columns should be used for hive-style partitioned writes by name. - /// Can be set to empty vec![] for non-partitioned writes. - pub partition_by: Vec, - /// Sets which columns should be used for sorting the output by name. - /// Can be set to empty vec![] for non-sorted writes. - pub sort_by: Vec, -} - -impl From for DataFrameWriteOptions { - fn from(value: SedonaWriteOptions) -> Self { - DataFrameWriteOptions::new() - .with_insert_operation(value.insert_op) - .with_single_file_output(value.single_file_output) - .with_partition_by(value.partition_by) - .with_sort_by(value.sort_by) - } -} - -impl SedonaWriteOptions { - /// Create a new SedonaWriteOptions with default values - pub fn new() -> Self { - SedonaWriteOptions { - insert_op: InsertOp::Append, - single_file_output: false, - partition_by: vec![], - sort_by: vec![], - } - } - - /// Set the insert operation - pub fn with_insert_operation(mut self, insert_op: InsertOp) -> Self { - self.insert_op = insert_op; - self - } - - /// Set the single_file_output value to true or false - pub fn with_single_file_output(mut self, single_file_output: bool) -> Self { - self.single_file_output = single_file_output; - self - } - - /// Sets the partition_by columns for output partitioning - pub fn with_partition_by(mut self, partition_by: Vec) -> Self { - self.partition_by = partition_by; - self - } - - /// Sets the sort_by columns for output sorting - pub fn with_sort_by(mut self, sort_by: Vec) -> Self { - self.sort_by = sort_by; - self - } -} - -impl Default for SedonaWriteOptions { - fn default() -> Self { - Self::new() - } -} - // Because Dialect/dialect_from_str is not marked as Send, using the async // function in certain contexts will fail to compile. Here we use a wrapper // to ensure that that the Dialect can be specified and parsed in any async @@ -721,9 +537,13 @@ impl ThreadSafeDialect { #[cfg(test)] mod tests { - use arrow_array::{create_array, ArrayRef, RecordBatchIterator, RecordBatchReader}; + use arrow_array::{ + create_array, ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, + }; use arrow_schema::{DataType, Field, Schema}; - use datafusion::assert_batches_eq; + use async_trait::async_trait; + use datafusion::{assert_batches_eq, dataframe::DataFrameWriteOptions}; + use datafusion_common::not_impl_err; use sedona_datasource::spec::{Object, OpenReaderArgs}; use sedona_schema::{ crs::{deserialize_crs, lnglat}, @@ -733,6 +553,8 @@ mod tests { use sedona_testing::data::test_geoparquet; use tempfile::tempdir; + use crate::{dataframe::SedonaDataFrame, show::DisplayTableOptions}; + use super::*; #[tokio::test] diff --git a/rust/sedona/src/lib.rs b/rust/sedona/src/lib.rs index dab1169a04..7db7d8d0ca 100644 --- a/rust/sedona/src/lib.rs +++ b/rust/sedona/src/lib.rs @@ -17,6 +17,7 @@ mod catalog; pub mod context; pub mod context_builder; +pub mod dataframe; mod exec; pub mod memory_pool; mod object_storage; From 7ac82e897e7f7f78ea0db6afbb540e78adb3f1b5 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 14:18:52 -0500 Subject: [PATCH 25/55] maybe fix streaming --- c/sedona-extension/src/lib.rs | 1 + c/sedona-extension/src/streaming.rs | 613 ++++++++++++++++++++++++++++ c/sedona-extension/src/utils.rs | 549 +------------------------ 3 files changed, 617 insertions(+), 546 deletions(-) create mode 100644 c/sedona-extension/src/streaming.rs diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs index d4f94155b4..68b82ca87c 100644 --- a/c/sedona-extension/src/lib.rs +++ b/c/sedona-extension/src/lib.rs @@ -19,5 +19,6 @@ pub mod execution_plan; pub mod expr; pub mod extension; pub mod scalar_kernel; +pub mod streaming; pub mod table_provider; pub mod utils; diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs new file mode 100644 index 0000000000..2ef03552c7 --- /dev/null +++ b/c/sedona-extension/src/streaming.rs @@ -0,0 +1,613 @@ +// 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. + +//! Utilities for streaming record batches across FFI boundaries. +//! +//! This module provides [`StreamingRecordBatchReader`] for exporting DataFusion streams +//! over FFI, and [`ffi_stream_to_sendable`] for importing FFI streams back into DataFusion. + +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::{RecordBatch, RecordBatchReader}; +use arrow_schema::{ArrowError, SchemaRef}; +use datafusion_common::{exec_err, Result}; +use datafusion_physical_plan::SendableRecordBatchStream; +use futures::StreamExt; + +/// A cancellation check callback for FFI operations. +/// +/// Returns `true` if the operation should be cancelled, `false` to continue. +/// This allows Python, R, ADBC, and other runtimes to integrate their own +/// cancellation mechanisms. +pub type CancelChecker = Box bool + Send + Sync>; + +/// Result type for batch fetching from the worker thread. +type BatchResult = Option>; + +/// Worker thread state for streaming record batch reads. +struct StreamWorker { + /// Channel to request the next batch (send () to request). + request_tx: std::sync::mpsc::SyncSender<()>, + /// Channel to receive batch results. + response_rx: std::sync::mpsc::Receiver, + /// Worker thread handle for cleanup. + _handle: std::thread::JoinHandle<()>, +} + +/// A RecordBatchReader that lazily polls a SendableRecordBatchStream. +/// +/// This allows exporting a DataFusion stream over FFI without collecting all +/// batches upfront. Used when exporting execution plans across FFI boundaries. +/// +/// Internally uses a dedicated worker thread to avoid blocking issues when +/// called from within a tokio runtime context. +pub struct StreamingRecordBatchReader { + schema: SchemaRef, + /// Stream and runtime, wrapped in Option so they can be moved to the worker. + stream_and_runtime: Option<(SendableRecordBatchStream, tokio::runtime::Handle)>, + /// Worker thread state, lazily initialized on first fetch. + worker: Option, + cancel_checker: Option, + cancelled: bool, + skip_empty_batches: bool, + periodic_check_interval: Option, +} + +impl StreamingRecordBatchReader { + /// Create a new StreamingRecordBatchReader from a SendableRecordBatchStream. + pub fn new(stream: SendableRecordBatchStream, runtime: tokio::runtime::Handle) -> Self { + Self { + schema: stream.schema(), + stream_and_runtime: Some((stream, runtime)), + worker: None, + cancel_checker: None, + cancelled: false, + skip_empty_batches: false, + periodic_check_interval: None, + } + } + + /// Create a new StreamingRecordBatchReader with a cancellation checker. + /// + /// The cancellation checker is called before each batch is fetched. If it + /// returns `true`, iteration stops with a cancellation error on the next + /// call and `None` on subsequent calls. + pub fn with_cancel_checker( + stream: SendableRecordBatchStream, + runtime: tokio::runtime::Handle, + cancel_checker: CancelChecker, + ) -> Self { + Self { + schema: stream.schema(), + stream_and_runtime: Some((stream, runtime)), + worker: None, + cancel_checker: Some(cancel_checker), + cancelled: false, + skip_empty_batches: false, + periodic_check_interval: None, + } + } + + /// Set whether to skip empty batches (batches with 0 rows). + /// + /// When enabled, the iterator will automatically skip over any batches + /// that have no rows and continue to the next batch. + pub fn with_skip_empty_batches(mut self, skip: bool) -> Self { + self.skip_empty_batches = skip; + self + } + + /// Set a periodic interval for checking cancellation during batch fetches. + /// + /// When set, the cancel checker will be called periodically at this interval + /// even while waiting for a single batch to be fetched. This is useful for + /// Python where we need to periodically check for signals (Ctrl+C) during + /// long-running operations. + /// + /// Without this, the cancel checker is only called between batch fetches. + pub fn with_periodic_check_interval(mut self, interval: std::time::Duration) -> Self { + self.periodic_check_interval = Some(interval); + self + } + + /// Ensure the worker thread is running, spawning it if necessary. + fn ensure_worker(&mut self) { + if self.worker.is_some() { + return; + } + + // Take ownership of stream and runtime + let Some((stream, runtime)) = self.stream_and_runtime.take() else { + return; + }; + + // Create channels for communication + // Use bounded channel with size 0 (rendezvous) for backpressure + let (request_tx, request_rx) = std::sync::mpsc::sync_channel::<()>(0); + let (response_tx, response_rx) = std::sync::mpsc::sync_channel::(0); + + // Spawn the worker thread + let handle = std::thread::spawn(move || { + let mut stream = stream; + // Process requests until the channel is closed + while request_rx.recv().is_ok() { + let result = runtime.block_on(async { + match stream.next().await { + Some(Ok(batch)) => Some(Ok(batch)), + Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), + None => None, + } + }); + + // Send the result back; if send fails, the reader was dropped + if response_tx.send(result).is_err() { + break; + } + } + }); + + self.worker = Some(StreamWorker { + request_tx, + response_rx, + _handle: handle, + }); + } + + fn fetch_next_batch(&mut self) -> Option> { + self.ensure_worker(); + + let Some(worker) = &self.worker else { + return Some(Err(ArrowError::InvalidArgumentError( + "Worker not initialized".to_string(), + ))); + }; + + // Request the next batch + if worker.request_tx.send(()).is_err() { + return Some(Err(ArrowError::InvalidArgumentError( + "Worker thread terminated".to_string(), + ))); + } + + // Wait for the response, with optional periodic cancellation checking + match &self.periodic_check_interval { + Some(interval) => { + let interval = *interval; + loop { + match worker.response_rx.recv_timeout(interval) { + Ok(result) => return result, + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { + // Check for cancellation + if let Some(ref checker) = self.cancel_checker { + if checker() { + return Some(Err(ArrowError::ExternalError(Box::new( + std::io::Error::new( + std::io::ErrorKind::Interrupted, + "Operation cancelled", + ), + )))); + } + } + // Continue waiting + } + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + return Some(Err(ArrowError::InvalidArgumentError( + "Worker thread terminated".to_string(), + ))); + } + } + } + } + None => { + // Simple blocking receive + match worker.response_rx.recv() { + Ok(result) => result, + Err(_) => Some(Err(ArrowError::InvalidArgumentError( + "Worker thread terminated".to_string(), + ))), + } + } + } + } +} + +impl Iterator for StreamingRecordBatchReader { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + // If already cancelled, return None to stop iteration + if self.cancelled { + return None; + } + + // Check for cancellation before fetching the next batch + // (periodic checking during fetch is handled by fetch_next_batch if configured) + if let Some(ref checker) = self.cancel_checker { + if checker() { + self.cancelled = true; + return Some(Err(ArrowError::ExternalError(Box::new( + std::io::Error::new(std::io::ErrorKind::Interrupted, "Operation cancelled"), + )))); + } + } + + loop { + match self.fetch_next_batch() { + Some(Ok(batch)) => { + if self.skip_empty_batches && batch.num_rows() == 0 { + continue; + } + return Some(Ok(batch)); + } + Some(Err(e)) => { + // Check if this was a cancellation error from periodic checking + if let ArrowError::ExternalError(box_err) = &e { + if let Some(io_err) = box_err.downcast_ref::() { + if io_err.kind() == std::io::ErrorKind::Interrupted { + self.cancelled = true; + } + } + } + return Some(Err(e)); + } + None => return None, + } + } + } +} + +impl RecordBatchReader for StreamingRecordBatchReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +/// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream. +/// +/// This is the inverse of StreamingRecordBatchReader - it takes an FFI stream +/// and converts it back to a DataFusion stream. Used when importing execution +/// plans from across FFI boundaries. +/// +/// # Safety +/// +/// The caller must ensure that the FFI stream pointer is valid and properly +/// initialized. +pub unsafe fn ffi_stream_to_sendable( + ffi_stream: &mut FFI_ArrowArrayStream, +) -> Result { + ffi_stream_to_sendable_with_cancel(ffi_stream, None) +} + +/// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream with cancellation support. +/// +/// The cancellation checker is called before each batch is read. If it returns +/// `true`, the stream yields a cancellation error. +/// +/// # Safety +/// +/// The caller must ensure that the FFI stream pointer is valid and properly +/// initialized. +pub unsafe fn ffi_stream_to_sendable_with_cancel( + ffi_stream: &mut FFI_ArrowArrayStream, + cancel_checker: Option, +) -> Result { + let reader = arrow_array::ffi_stream::ArrowArrayStreamReader::from_raw(ffi_stream)?; + + let schema = reader.schema(); + let stream = futures::stream::iter(reader).map(move |result| { + // Check for cancellation before yielding each batch + if let Some(ref checker) = cancel_checker { + if checker() { + return exec_err!("Operation cancelled"); + } + } + result.map_err(|e| datafusion_common::DataFusionError::ArrowError(Box::new(e), None)) + }); + + Ok(Box::pin( + datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::Int32Array; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::sync::Arc; + use std::time::Duration; + + /// Create a slow stream that yields batches with a configurable delay. + fn create_slow_stream( + num_batches: usize, + delay_ms: u64, + ) -> (SchemaRef, SendableRecordBatchStream) { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let schema_clone = schema.clone(); + + let stream = futures::stream::iter(0..num_batches).then(move |i| { + let schema = schema_clone.clone(); + async move { + if delay_ms > 0 { + // Use std::thread::sleep because tokio::time::sleep requires + // the runtime to poll, which can deadlock with block_on + std::thread::sleep(Duration::from_millis(delay_ms)); + } + let array = Int32Array::from(vec![i as i32]); + Ok(RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()) + } + }); + + ( + schema.clone(), + Box::pin( + datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), + ), + ) + } + + #[tokio::test] + async fn test_streaming_reader_basic() { + let runtime = tokio::runtime::Handle::current(); + let (_schema, stream) = create_slow_stream(5, 10); + + let reader = StreamingRecordBatchReader::new(stream, runtime); + let batches: Vec<_> = reader.collect(); + + assert_eq!(batches.len(), 5); + for (i, batch) in batches.iter().enumerate() { + let batch = batch.as_ref().unwrap(); + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(array.value(0), i as i32); + } + } + + #[tokio::test] + async fn test_streaming_reader_cancel() { + let runtime = tokio::runtime::Handle::current(); + let (_schema, stream) = create_slow_stream(10, 50); + + // Cancel after reading 3 batches + let counter = Arc::new(AtomicUsize::new(0)); + let counter_clone = counter.clone(); + let cancel_checker: CancelChecker = Box::new(move || { + let count = counter_clone.fetch_add(1, Ordering::SeqCst); + count >= 3 + }); + + let reader = + StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker); + let batches: Vec<_> = reader.collect(); + + // Should have 3 successful batches + 1 cancellation error + assert_eq!(batches.len(), 4); + + // First 3 should be Ok + for (i, batch) in batches.iter().enumerate().take(3) { + assert!(batch.is_ok(), "batch {} should be Ok", i); + } + + // Last one should be a cancellation error + let last = batches.last().unwrap(); + assert!(last.is_err()); + let err = last.as_ref().unwrap_err(); + assert!( + err.to_string().contains("cancelled"), + "error should mention cancellation: {}", + err + ); + } + + #[tokio::test] + async fn test_ffi_stream_to_sendable_basic() { + use arrow_array::ffi_stream::FFI_ArrowArrayStream; + + let runtime = tokio::runtime::Handle::current(); + let (_schema, stream) = create_slow_stream(5, 10); + + // Export to FFI stream + let reader = StreamingRecordBatchReader::new(stream, runtime.clone()); + let mut ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); + + // Import back + let imported = unsafe { ffi_stream_to_sendable(&mut ffi_stream).unwrap() }; + + // Collect results + let batches: Vec<_> = imported.collect::>().await; + + assert_eq!(batches.len(), 5); + for (i, batch) in batches.iter().enumerate() { + let batch = batch.as_ref().unwrap(); + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(array.value(0), i as i32); + } + } + + #[tokio::test] + async fn test_ffi_stream_to_sendable_cancel() { + use arrow_array::ffi_stream::FFI_ArrowArrayStream; + + let runtime = tokio::runtime::Handle::current(); + let (_schema, stream) = create_slow_stream(10, 50); + + // Export to FFI stream (no cancellation on export side) + let reader = StreamingRecordBatchReader::new(stream, runtime.clone()); + let mut ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); + + // Cancel after reading 3 batches on import side + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_clone = cancelled.clone(); + let cancel_checker: CancelChecker = + Box::new(move || cancelled_clone.load(Ordering::SeqCst)); + + let imported = unsafe { + ffi_stream_to_sendable_with_cancel(&mut ffi_stream, Some(cancel_checker)).unwrap() + }; + + // Collect with cancellation after 3 batches, stop on first error + let mut batches = Vec::new(); + let mut stream = imported; + while let Some(result) = stream.next().await { + let is_err = result.is_err(); + batches.push(result); + if batches.len() == 3 { + cancelled.store(true, Ordering::SeqCst); + } + if is_err { + break; // Stop on first error + } + } + + // Should have 3 successful batches + 1 cancellation error + assert_eq!(batches.len(), 4); + + // First 3 should be Ok + for (i, batch) in batches.iter().enumerate().take(3) { + assert!(batch.is_ok(), "batch {} should be Ok", i); + } + + // Last one should be a cancellation error + let last = batches.last().unwrap(); + assert!(last.is_err()); + let err = last.as_ref().unwrap_err(); + assert!( + err.to_string().contains("cancelled"), + "error should mention cancellation: {}", + err + ); + } + + /// Create a stream that yields some empty batches. + fn create_stream_with_empty_batches( + batch_row_counts: Vec, + ) -> (SchemaRef, SendableRecordBatchStream) { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let schema_clone = schema.clone(); + + let stream = futures::stream::iter(batch_row_counts.into_iter().enumerate()).map( + move |(i, row_count)| { + let schema = schema_clone.clone(); + let array: Int32Array = (0..row_count).map(|_| i as i32).collect(); + Ok(RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()) + }, + ); + + ( + schema.clone(), + Box::pin( + datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), + ), + ) + } + + #[tokio::test] + async fn test_streaming_reader_skip_empty_batches() { + let runtime = tokio::runtime::Handle::current(); + // Create stream with: 2 rows, 0 rows, 3 rows, 0 rows, 0 rows, 1 row + let (_schema, stream) = create_stream_with_empty_batches(vec![2, 0, 3, 0, 0, 1]); + + let reader = StreamingRecordBatchReader::new(stream, runtime).with_skip_empty_batches(true); + let batches: Vec<_> = reader.collect(); + + // Should only get 3 batches (the non-empty ones) + assert_eq!(batches.len(), 3); + + // Check row counts: 2, 3, 1 + assert_eq!(batches[0].as_ref().unwrap().num_rows(), 2); + assert_eq!(batches[1].as_ref().unwrap().num_rows(), 3); + assert_eq!(batches[2].as_ref().unwrap().num_rows(), 1); + } + + #[tokio::test] + async fn test_streaming_reader_no_skip_empty_batches() { + let runtime = tokio::runtime::Handle::current(); + // Create stream with: 2 rows, 0 rows, 3 rows + let (_schema, stream) = create_stream_with_empty_batches(vec![2, 0, 3]); + + // Default: skip_empty_batches is false + let reader = StreamingRecordBatchReader::new(stream, runtime); + let batches: Vec<_> = reader.collect(); + + // Should get all 3 batches including the empty one + assert_eq!(batches.len(), 3); + assert_eq!(batches[0].as_ref().unwrap().num_rows(), 2); + assert_eq!(batches[1].as_ref().unwrap().num_rows(), 0); + assert_eq!(batches[2].as_ref().unwrap().num_rows(), 3); + } + + #[tokio::test] + async fn test_streaming_reader_periodic_check_interval() { + let runtime = tokio::runtime::Handle::current(); + // Create a slow stream where each batch takes 200ms + let (_schema, stream) = create_slow_stream(5, 200); + + // Cancel flag - will be set after 150ms + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_clone = cancelled.clone(); + let cancel_checker: CancelChecker = + Box::new(move || cancelled_clone.load(Ordering::SeqCst)); + + let reader = + StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker) + .with_periodic_check_interval(Duration::from_millis(50)); + + // Spawn a task to set cancelled after 150ms + let cancelled_setter = cancelled.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(150)); + cancelled_setter.store(true, Ordering::SeqCst); + }); + + let batches: Vec<_> = reader.collect(); + + // The first batch takes 200ms but we cancel at 150ms during the fetch + // With periodic_check_interval of 50ms, the check happens at 50ms, 100ms, 150ms + // At 150ms the cancel check should trigger + // So we should get 0 successful batches + 1 cancellation error + assert!( + batches.len() <= 2, + "expected at most 1 batch + 1 error, got {}", + batches.len() + ); + + // At least one should be an error + let has_cancel_error = batches + .iter() + .any(|b| b.is_err() && b.as_ref().unwrap_err().to_string().contains("cancelled")); + assert!( + has_cancel_error, + "should have a cancellation error in batches: {:?}", + batches + ); + } +} diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index c388aa6d8d..d7a91a8965 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -15,18 +15,14 @@ // specific language governing permissions and limitations // under the License. -//! Utilities for converting between Arrow FFI streams and DataFusion streams. +//! Utilities for FFI property access and conversion. use std::ffi::{c_int, CString}; use std::ptr::null_mut; -use arrow_array::ffi_stream::FFI_ArrowArrayStream; -use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::ffi::FFI_ArrowSchema; -use arrow_schema::{ArrowError, DataType, Field, SchemaRef}; -use datafusion_common::{exec_err, Result}; -use datafusion_physical_plan::SendableRecordBatchStream; -use futures::StreamExt; +use arrow_schema::{DataType, Field}; +use datafusion_common::Result; use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; use serde::de::DeserializeOwned; use serde::Serialize; @@ -335,542 +331,3 @@ pub fn get_table_provider_string_property( String::from_utf8(bytes) .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) } - -/// A cancellation check callback for FFI operations. -/// -/// Returns `true` if the operation should be cancelled, `false` to continue. -/// This allows Python, R, ADBC, and other runtimes to integrate their own -/// cancellation mechanisms. -pub type CancelChecker = Box bool + Send + Sync>; - -/// A RecordBatchReader that lazily polls a SendableRecordBatchStream. -/// -/// This allows exporting a DataFusion stream over FFI without collecting all -/// batches upfront. Used when exporting execution plans across FFI boundaries. -pub struct StreamingRecordBatchReader { - schema: SchemaRef, - stream: SendableRecordBatchStream, - runtime: tokio::runtime::Handle, - cancel_checker: Option, - cancelled: bool, - skip_empty_batches: bool, - periodic_check_interval: Option, -} - -impl StreamingRecordBatchReader { - /// Create a new StreamingRecordBatchReader from a SendableRecordBatchStream. - pub fn new(stream: SendableRecordBatchStream, runtime: tokio::runtime::Handle) -> Self { - Self { - schema: stream.schema(), - stream, - runtime, - cancel_checker: None, - cancelled: false, - skip_empty_batches: false, - periodic_check_interval: None, - } - } - - /// Create a new StreamingRecordBatchReader with a cancellation checker. - /// - /// The cancellation checker is called before each batch is fetched. If it - /// returns `true`, iteration stops with a cancellation error on the next - /// call and `None` on subsequent calls. - pub fn with_cancel_checker( - stream: SendableRecordBatchStream, - runtime: tokio::runtime::Handle, - cancel_checker: CancelChecker, - ) -> Self { - Self { - schema: stream.schema(), - stream, - runtime, - cancel_checker: Some(cancel_checker), - cancelled: false, - skip_empty_batches: false, - periodic_check_interval: None, - } - } - - /// Set whether to skip empty batches (batches with 0 rows). - /// - /// When enabled, the iterator will automatically skip over any batches - /// that have no rows and continue to the next batch. - pub fn with_skip_empty_batches(mut self, skip: bool) -> Self { - self.skip_empty_batches = skip; - self - } - - /// Set a periodic interval for checking cancellation during batch fetches. - /// - /// When set, the cancel checker will be called periodically at this interval - /// even while waiting for a single batch to be fetched. This is useful for - /// Python where we need to periodically check for signals (Ctrl+C) during - /// long-running operations. - /// - /// Without this, the cancel checker is only called between batch fetches. - pub fn with_periodic_check_interval(mut self, interval: std::time::Duration) -> Self { - self.periodic_check_interval = Some(interval); - self - } - - fn fetch_next_batch(&mut self) -> Option> { - let stream = &mut self.stream; - let runtime = &self.runtime; - - match &self.periodic_check_interval { - Some(interval) => { - // Use tokio::select! to periodically check for cancellation - let interval = *interval; - let cancel_checker = &self.cancel_checker; - - std::thread::scope(|s| { - s.spawn(|| { - runtime.block_on(async { - use futures::StreamExt; - tokio::pin!(stream); - loop { - tokio::select! { - res = stream.next() => { - return match res { - Some(Ok(batch)) => Some(Ok(batch)), - Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), - None => None, - }; - } - _ = tokio::time::sleep(interval) => { - if let Some(ref checker) = cancel_checker { - if checker() { - return Some(Err(ArrowError::ExternalError(Box::new( - std::io::Error::new( - std::io::ErrorKind::Interrupted, - "Operation cancelled", - ), - )))); - } - } - // Continue waiting for the batch - } - } - } - }) - }) - .join() - .unwrap_or_else(|_| { - Some(Err(ArrowError::InvalidArgumentError( - "Iterator thread panicked".to_string(), - ))) - }) - }) - } - None => { - // Simple blocking fetch without periodic checking - std::thread::scope(|s| { - s.spawn(|| { - runtime.block_on(async { - match stream.next().await { - Some(Ok(batch)) => Some(Ok(batch)), - Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), - None => None, - } - }) - }) - .join() - .unwrap_or_else(|_| { - Some(Err(ArrowError::InvalidArgumentError( - "Iterator thread panicked".to_string(), - ))) - }) - }) - } - } - } -} - -impl Iterator for StreamingRecordBatchReader { - type Item = std::result::Result; - - fn next(&mut self) -> Option { - // If already cancelled, return None to stop iteration - if self.cancelled { - return None; - } - - // Check for cancellation before fetching the next batch - // (periodic checking during fetch is handled by fetch_next_batch if configured) - if let Some(ref checker) = self.cancel_checker { - if checker() { - self.cancelled = true; - return Some(Err(ArrowError::ExternalError(Box::new( - std::io::Error::new(std::io::ErrorKind::Interrupted, "Operation cancelled"), - )))); - } - } - - loop { - match self.fetch_next_batch() { - Some(Ok(batch)) => { - if self.skip_empty_batches && batch.num_rows() == 0 { - continue; - } - return Some(Ok(batch)); - } - Some(Err(e)) => { - // Check if this was a cancellation error from periodic checking - if let ArrowError::ExternalError(box_err) = &e { - if let Some(io_err) = box_err.downcast_ref::() { - if io_err.kind() == std::io::ErrorKind::Interrupted { - self.cancelled = true; - } - } - } - return Some(Err(e)); - } - None => return None, - } - } - } -} - -impl RecordBatchReader for StreamingRecordBatchReader { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} - -/// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream. -/// -/// This is the inverse of StreamingRecordBatchReader - it takes an FFI stream -/// and converts it back to a DataFusion stream. Used when importing execution -/// plans from across FFI boundaries. -/// -/// # Safety -/// -/// The caller must ensure that the FFI stream pointer is valid and properly -/// initialized. -pub unsafe fn ffi_stream_to_sendable( - ffi_stream: &mut FFI_ArrowArrayStream, -) -> Result { - ffi_stream_to_sendable_with_cancel(ffi_stream, None) -} - -/// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream with cancellation support. -/// -/// The cancellation checker is called before each batch is read. If it returns -/// `true`, the stream yields a cancellation error. -/// -/// # Safety -/// -/// The caller must ensure that the FFI stream pointer is valid and properly -/// initialized. -pub unsafe fn ffi_stream_to_sendable_with_cancel( - ffi_stream: &mut FFI_ArrowArrayStream, - cancel_checker: Option, -) -> Result { - let reader = arrow_array::ffi_stream::ArrowArrayStreamReader::from_raw(ffi_stream)?; - - let schema = reader.schema(); - let stream = futures::stream::iter(reader).map(move |result| { - // Check for cancellation before yielding each batch - if let Some(ref checker) = cancel_checker { - if checker() { - return exec_err!("Operation cancelled"); - } - } - result.map_err(|e| datafusion_common::DataFusionError::ArrowError(Box::new(e), None)) - }); - - Ok(Box::pin( - datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), - )) -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow_array::Int32Array; - use arrow_schema::{Field, Schema}; - use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; - use std::sync::Arc; - use std::time::Duration; - - /// Create a slow stream that yields batches with a configurable delay. - fn create_slow_stream( - num_batches: usize, - delay_ms: u64, - ) -> (SchemaRef, SendableRecordBatchStream) { - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Int32, - false, - )])); - let schema_clone = schema.clone(); - - let stream = futures::stream::iter(0..num_batches).then(move |i| { - let schema = schema_clone.clone(); - async move { - if delay_ms > 0 { - // Use std::thread::sleep because tokio::time::sleep requires - // the runtime to poll, which can deadlock with block_on - std::thread::sleep(Duration::from_millis(delay_ms)); - } - let array = Int32Array::from(vec![i as i32]); - Ok(RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()) - } - }); - - ( - schema.clone(), - Box::pin( - datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), - ), - ) - } - - #[tokio::test] - async fn test_streaming_reader_basic() { - let runtime = tokio::runtime::Handle::current(); - let (_schema, stream) = create_slow_stream(5, 10); - - let reader = StreamingRecordBatchReader::new(stream, runtime); - let batches: Vec<_> = reader.collect(); - - assert_eq!(batches.len(), 5); - for (i, batch) in batches.iter().enumerate() { - let batch = batch.as_ref().unwrap(); - let array = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(array.value(0), i as i32); - } - } - - #[tokio::test] - async fn test_streaming_reader_cancel() { - let runtime = tokio::runtime::Handle::current(); - let (_schema, stream) = create_slow_stream(10, 50); - - // Cancel after reading 3 batches - let counter = Arc::new(AtomicUsize::new(0)); - let counter_clone = counter.clone(); - let cancel_checker: CancelChecker = Box::new(move || { - let count = counter_clone.fetch_add(1, Ordering::SeqCst); - count >= 3 - }); - - let reader = - StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker); - let batches: Vec<_> = reader.collect(); - - // Should have 3 successful batches + 1 cancellation error - assert_eq!(batches.len(), 4); - - // First 3 should be Ok - for (i, batch) in batches.iter().enumerate().take(3) { - assert!(batch.is_ok(), "batch {} should be Ok", i); - } - - // Last one should be a cancellation error - let last = batches.last().unwrap(); - assert!(last.is_err()); - let err = last.as_ref().unwrap_err(); - assert!( - err.to_string().contains("cancelled"), - "error should mention cancellation: {}", - err - ); - } - - #[tokio::test] - async fn test_ffi_stream_to_sendable_basic() { - use arrow_array::ffi_stream::FFI_ArrowArrayStream; - - let runtime = tokio::runtime::Handle::current(); - let (_schema, stream) = create_slow_stream(5, 10); - - // Export to FFI stream - let reader = StreamingRecordBatchReader::new(stream, runtime.clone()); - let mut ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); - - // Import back - let imported = unsafe { ffi_stream_to_sendable(&mut ffi_stream).unwrap() }; - - // Collect results - let batches: Vec<_> = imported.collect::>().await; - - assert_eq!(batches.len(), 5); - for (i, batch) in batches.iter().enumerate() { - let batch = batch.as_ref().unwrap(); - let array = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(array.value(0), i as i32); - } - } - - #[tokio::test] - async fn test_ffi_stream_to_sendable_cancel() { - use arrow_array::ffi_stream::FFI_ArrowArrayStream; - - let runtime = tokio::runtime::Handle::current(); - let (_schema, stream) = create_slow_stream(10, 50); - - // Export to FFI stream (no cancellation on export side) - let reader = StreamingRecordBatchReader::new(stream, runtime.clone()); - let mut ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); - - // Cancel after reading 3 batches on import side - let cancelled = Arc::new(AtomicBool::new(false)); - let cancelled_clone = cancelled.clone(); - let cancel_checker: CancelChecker = - Box::new(move || cancelled_clone.load(Ordering::SeqCst)); - - let imported = unsafe { - ffi_stream_to_sendable_with_cancel(&mut ffi_stream, Some(cancel_checker)).unwrap() - }; - - // Collect with cancellation after 3 batches, stop on first error - let mut batches = Vec::new(); - let mut stream = imported; - while let Some(result) = stream.next().await { - let is_err = result.is_err(); - batches.push(result); - if batches.len() == 3 { - cancelled.store(true, Ordering::SeqCst); - } - if is_err { - break; // Stop on first error - } - } - - // Should have 3 successful batches + 1 cancellation error - assert_eq!(batches.len(), 4); - - // First 3 should be Ok - for (i, batch) in batches.iter().enumerate().take(3) { - assert!(batch.is_ok(), "batch {} should be Ok", i); - } - - // Last one should be a cancellation error - let last = batches.last().unwrap(); - assert!(last.is_err()); - let err = last.as_ref().unwrap_err(); - assert!( - err.to_string().contains("cancelled"), - "error should mention cancellation: {}", - err - ); - } - - /// Create a stream that yields some empty batches. - fn create_stream_with_empty_batches( - batch_row_counts: Vec, - ) -> (SchemaRef, SendableRecordBatchStream) { - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Int32, - false, - )])); - let schema_clone = schema.clone(); - - let stream = futures::stream::iter(batch_row_counts.into_iter().enumerate()).map( - move |(i, row_count)| { - let schema = schema_clone.clone(); - let array: Int32Array = (0..row_count).map(|_| i as i32).collect(); - Ok(RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()) - }, - ); - - ( - schema.clone(), - Box::pin( - datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(schema, stream), - ), - ) - } - - #[tokio::test] - async fn test_streaming_reader_skip_empty_batches() { - let runtime = tokio::runtime::Handle::current(); - // Create stream with: 2 rows, 0 rows, 3 rows, 0 rows, 0 rows, 1 row - let (_schema, stream) = create_stream_with_empty_batches(vec![2, 0, 3, 0, 0, 1]); - - let reader = StreamingRecordBatchReader::new(stream, runtime).with_skip_empty_batches(true); - let batches: Vec<_> = reader.collect(); - - // Should only get 3 batches (the non-empty ones) - assert_eq!(batches.len(), 3); - - // Check row counts: 2, 3, 1 - assert_eq!(batches[0].as_ref().unwrap().num_rows(), 2); - assert_eq!(batches[1].as_ref().unwrap().num_rows(), 3); - assert_eq!(batches[2].as_ref().unwrap().num_rows(), 1); - } - - #[tokio::test] - async fn test_streaming_reader_no_skip_empty_batches() { - let runtime = tokio::runtime::Handle::current(); - // Create stream with: 2 rows, 0 rows, 3 rows - let (_schema, stream) = create_stream_with_empty_batches(vec![2, 0, 3]); - - // Default: skip_empty_batches is false - let reader = StreamingRecordBatchReader::new(stream, runtime); - let batches: Vec<_> = reader.collect(); - - // Should get all 3 batches including the empty one - assert_eq!(batches.len(), 3); - assert_eq!(batches[0].as_ref().unwrap().num_rows(), 2); - assert_eq!(batches[1].as_ref().unwrap().num_rows(), 0); - assert_eq!(batches[2].as_ref().unwrap().num_rows(), 3); - } - - #[tokio::test] - async fn test_streaming_reader_periodic_check_interval() { - let runtime = tokio::runtime::Handle::current(); - // Create a slow stream where each batch takes 200ms - let (_schema, stream) = create_slow_stream(5, 200); - - // Cancel flag - will be set after 150ms - let cancelled = Arc::new(AtomicBool::new(false)); - let cancelled_clone = cancelled.clone(); - let cancel_checker: CancelChecker = - Box::new(move || cancelled_clone.load(Ordering::SeqCst)); - - let reader = - StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker) - .with_periodic_check_interval(Duration::from_millis(50)); - - // Spawn a task to set cancelled after 150ms - let cancelled_setter = cancelled.clone(); - std::thread::spawn(move || { - std::thread::sleep(Duration::from_millis(150)); - cancelled_setter.store(true, Ordering::SeqCst); - }); - - let batches: Vec<_> = reader.collect(); - - // The first batch takes 200ms but we cancel at 150ms during the fetch - // With periodic_check_interval of 50ms, the check happens at 50ms, 100ms, 150ms - // At 150ms the cancel check should trigger - // So we should get 0 successful batches + 1 cancellation error - assert!( - batches.len() <= 2, - "expected at most 1 batch + 1 error, got {}", - batches.len() - ); - - // At least one should be an error - let has_cancel_error = batches - .iter() - .any(|b| b.is_err() && b.as_ref().unwrap_err().to_string().contains("cancelled")); - assert!( - has_cancel_error, - "should have a cancellation error in batches: {:?}", - batches - ); - } -} From 86d2b3445d3d52831bf4bccbf5f760fbcff3fb28 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 14:58:19 -0500 Subject: [PATCH 26/55] merge formatting --- rust/sedona/src/context.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index 403ec80673..6831491881 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -45,9 +45,9 @@ use sedona_common::{ use sedona_datasource::provider::external_table; use sedona_datasource::spec::ExternalFormatSpec; use sedona_expr::{ - scalar_udf::{IntoScalarKernelRefs, SedonaScalarUDF}, aggregate_udf::{IntoSedonaAccumulatorRefs, SedonaAggregateUDF}, function_set::FunctionSet, + scalar_udf::{IntoScalarKernelRefs, SedonaScalarUDF}, }; use sedona_geoparquet::options::TableGeoParquetOptions; use sedona_geoparquet::{ From 271681c8431209e1ecd9517e9ad0c43b84addd87 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 15:51:31 -0500 Subject: [PATCH 27/55] fix imports --- c/sedona-extension/src/execution_plan.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 4bc53bf560..9bf09595d2 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -37,10 +37,8 @@ use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; use serde::{Deserialize, Serialize}; use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs}; -use crate::utils::{ - ffi_stream_to_sendable, get_plan_property, get_plan_string_property, - StreamingRecordBatchReader, ERRNO_OK, -}; +use crate::streaming::{ffi_stream_to_sendable, StreamingRecordBatchReader}; +use crate::utils::{get_plan_property, get_plan_string_property, ERRNO_OK}; /// Wrapper around an [ExecutionPlan] that can be exported across FFI. pub struct ExportedExecutionPlan { From 78042b7c1084492e8a03cb8990274e5cc9fee2af Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 16:18:42 -0500 Subject: [PATCH 28/55] maybe actually fix imports --- python/sedonadb/src/reader.rs | 4 ++-- r/sedonadb/src/rust/src/dataframe.rs | 2 +- rust/sedona-adbc/src/statement.rs | 2 +- rust/sedona/src/context.rs | 3 +-- 4 files changed, 5 insertions(+), 6 deletions(-) diff --git a/python/sedonadb/src/reader.rs b/python/sedonadb/src/reader.rs index 804c4a4b01..fc0b54815b 100644 --- a/python/sedonadb/src/reader.rs +++ b/python/sedonadb/src/reader.rs @@ -18,7 +18,7 @@ use std::time::Duration; use datafusion::execution::SendableRecordBatchStream; use pyo3::Python; -use sedona_extension::utils::StreamingRecordBatchReader; +use sedona_extension::streaming::StreamingRecordBatchReader; /// Interval for checking Python signals during batch fetches. const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(2_000); @@ -35,7 +35,7 @@ pub fn new_py_streaming_reader( ) -> StreamingRecordBatchReader { // Create a cancel checker that checks Python signals let cancel_checker: Box bool + Send + Sync> = Box::new(|| { - Python::with_gil(|py| { + Python::attach(|py| { // Run `pass` to process any pending signals, then check for errors if py.run(cr"pass", None, None).is_err() { return true; diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index e914f4176e..6aa1b0e23e 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -30,7 +30,7 @@ use datafusion_ffi::table_provider::FFI_TableProvider; use savvy::{savvy, savvy_err, sexp, IntoExtPtrSexp, Result}; use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; use sedona::show::{DisplayMode, DisplayTableOptions}; -use sedona_extension::utils::StreamingRecordBatchReader; +use sedona_extension::streaming::StreamingRecordBatchReader; use sedona_geoparquet::options::TableGeoParquetOptions; use sedona_schema::schema::SedonaSchema; use std::{iter::zip, ptr::swap_nonoverlapping, sync::Arc}; diff --git a/rust/sedona-adbc/src/statement.rs b/rust/sedona-adbc/src/statement.rs index 59927a24d0..f8bb9344ba 100644 --- a/rust/sedona-adbc/src/statement.rs +++ b/rust/sedona-adbc/src/statement.rs @@ -18,7 +18,7 @@ use adbc_core::PartitionedResult; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::Schema; use sedona::context::SedonaContext; -use sedona_extension::utils::StreamingRecordBatchReader; +use sedona_extension::streaming::StreamingRecordBatchReader; use std::sync::Arc; use tokio::runtime::Runtime; diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index 6831491881..2582f0ca66 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -36,7 +36,7 @@ use datafusion::{ sql::parser::{DFParser, Statement}, }; use datafusion_expr::sqlparser::dialect::{dialect_from_str, Dialect}; -use datafusion_expr::{AggregateUDFImpl, LogicalPlan, LogicalPlanBuilder, ScalarUDFImpl, SortExpr}; +use datafusion_expr::{AggregateUDFImpl, ScalarUDFImpl}; use parking_lot::Mutex; use sedona_common::{ option::add_sedona_option_extension, sedona_internal_datafusion_err, CrsProviderOption, @@ -49,7 +49,6 @@ use sedona_expr::{ function_set::FunctionSet, scalar_udf::{IntoScalarKernelRefs, SedonaScalarUDF}, }; -use sedona_geoparquet::options::TableGeoParquetOptions; use sedona_geoparquet::{ format::GeoParquetFormatFactory, provider::{geoparquet_listing_table, GeoParquetReadOptions}, From 49edb5b182f64cd172dfdba3a31e0ff2b8e97e18 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 16:39:40 -0500 Subject: [PATCH 29/55] refix dataframe export --- python/sedonadb/src/dataframe.rs | 2 +- r/sedonadb/src/rust/src/dataframe.rs | 2 +- rust/sedona/src/context.rs | 207 +++++++++++++++++-- rust/sedona/src/dataframe.rs | 288 --------------------------- rust/sedona/src/lib.rs | 1 - 5 files changed, 196 insertions(+), 304 deletions(-) delete mode 100644 rust/sedona/src/dataframe.rs diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs index 83cb7fb6ec..1f62763ff4 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -32,7 +32,7 @@ use futures::lock::Mutex; use futures::TryStreamExt; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict, PyList}; -use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; +use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; use sedona::projected_reader::simplify_record_batch_reader; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_geoparquet::options::TableGeoParquetOptions; diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index 6aa1b0e23e..cc150b3f9b 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -28,7 +28,7 @@ use datafusion_expr::JoinType; use datafusion_expr::{select_expr::SelectExpr, Expr, SortExpr}; use datafusion_ffi::table_provider::FFI_TableProvider; use savvy::{savvy, savvy_err, sexp, IntoExtPtrSexp, Result}; -use sedona::dataframe::{SedonaDataFrame, SedonaWriteOptions}; +use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_extension::streaming::StreamingRecordBatchReader; use sedona_geoparquet::options::TableGeoParquetOptions; diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index 2582f0ca66..d44c9beb18 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -21,22 +21,35 @@ use std::{ use crate::exec::create_plan_from_sql; use crate::object_storage::ensure_object_store_registered_with_options; -use crate::{catalog::DynamicObjectStoreCatalog, random_geometry_provider::RandomGeometryFunction}; +use crate::{ + catalog::DynamicObjectStoreCatalog, + random_geometry_provider::RandomGeometryFunction, + show::{show_batches, DisplayTableOptions}, +}; +use arrow_array::RecordBatch; use arrow_schema::DataType; -use datafusion::execution::memory_pool::MemoryLimit; +use async_trait::async_trait; use datafusion::{ common::plan_err, + dataframe::DataFrameWriteOptions, + datasource::file_format::format_as_file_type, error::{DataFusionError, Result}, execution::{ context::DataFilePaths, + memory_pool::MemoryLimit, runtime_env::{RuntimeEnv, RuntimeEnvBuilder}, SessionStateBuilder, }, prelude::{DataFrame, SessionConfig, SessionContext}, sql::parser::{DFParser, Statement}, }; -use datafusion_expr::sqlparser::dialect::{dialect_from_str, Dialect}; -use datafusion_expr::{AggregateUDFImpl, ScalarUDFImpl}; + +use datafusion_common::not_impl_err; +use datafusion_expr::{ + dml::InsertOp, + sqlparser::dialect::{dialect_from_str, Dialect}, + AggregateUDFImpl, LogicalPlan, LogicalPlanBuilder, ScalarUDFImpl, SortExpr, +}; use parking_lot::Mutex; use sedona_common::{ option::add_sedona_option_extension, sedona_internal_datafusion_err, CrsProviderOption, @@ -49,6 +62,7 @@ use sedona_expr::{ function_set::FunctionSet, scalar_udf::{IntoScalarKernelRefs, SedonaScalarUDF}, }; +use sedona_geoparquet::options::TableGeoParquetOptions; use sedona_geoparquet::{ format::GeoParquetFormatFactory, provider::{geoparquet_listing_table, GeoParquetReadOptions}, @@ -76,7 +90,7 @@ use sedona_raster::raster_loader::{AsyncRasterLoader, RasterLoaderConfig, Raster /// interface for configuring the behaviour of pub struct SedonaContext { pub ctx: SessionContext, - functions: RwLock, + pub functions: RwLock, /// Per-session registry of async raster byte loaders, keyed by /// `outdb_format`. Held behind an `Arc>` so the registered /// `RS_EnsureLoaded` UDF instance and any extension crates' `register(&ctx)` @@ -564,6 +578,179 @@ impl Default for SedonaContext { } } +/// Sedona-specific [`DataFrame`] actions +/// +/// This trait, implemented for [`DataFrame`], extends the DataFrame API to make it +/// ergonomic to work with dataframes that contain geometry columns. Currently these +/// are limited to output functions, as geometry columns currently require special +/// handling when written or exported to an external system. +#[async_trait] +pub trait SedonaDataFrame { + /// Build a table of the first `limit` results in this DataFrame + /// + /// This will limit and execute the query and build a table using [show_batches]. + async fn show_sedona<'a>( + self, + ctx: &SedonaContext, + limit: Option, + options: DisplayTableOptions<'a>, + ) -> Result; + + async fn write_geoparquet( + self, + ctx: &SedonaContext, + path: &str, + options: SedonaWriteOptions, + writer_options: Option, + ) -> Result>; +} + +#[async_trait] +impl SedonaDataFrame for DataFrame { + async fn show_sedona<'a>( + self, + ctx: &SedonaContext, + limit: Option, + mut options: DisplayTableOptions<'a>, + ) -> Result { + let df = if matches!( + self.logical_plan(), + LogicalPlan::Explain(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::Analyze(_) + ) { + // Show multi-line output without truncation for plans like `EXPLAIN` + options.max_row_height = usize::MAX; + + // We don't want to apply an additional .limit() to plans like `Explain` + // as that will trigger an internal error: Unsupported logical plan: Explain must be root of the plan + self + } else { + // Apply limit if specified + self.limit(0, limit)? + }; + + let schema_without_qualifiers = df.schema().clone().strip_qualifiers(); + let schema = schema_without_qualifiers.as_arrow(); + let batches = df.collect().await?; + let mut out = Vec::new(); + show_batches(ctx, &mut out, schema, batches, options)?; + String::from_utf8(out).map_err(|e| DataFusionError::External(Box::new(e))) + } + + async fn write_geoparquet( + self, + ctx: &SedonaContext, + path: &str, + options: SedonaWriteOptions, + writer_options: Option, + ) -> Result, DataFusionError> { + if options.insert_op != InsertOp::Append { + return not_impl_err!( + "{} is not implemented for DataFrame::write_geoparquet.", + options.insert_op + ); + } + + let format = if let Some(parquet_opts) = writer_options { + Arc::new(GeoParquetFormatFactory::new_with_options(parquet_opts)) + } else { + Arc::new(GeoParquetFormatFactory::new()) + }; + + let file_type = format_as_file_type(format); + + let plan = if options.sort_by.is_empty() { + self.into_unoptimized_plan() + } else { + LogicalPlanBuilder::from(self.into_unoptimized_plan()) + .sort(options.sort_by)? + .build()? + }; + + let plan = LogicalPlanBuilder::copy_to( + plan, + path.into(), + file_type, + Default::default(), + options.partition_by, + )? + .build()?; + + DataFrame::new(ctx.ctx.state(), plan).collect().await + } +} + +/// A Sedona-specific copy of [DataFrameWriteOptions] +/// +/// This is needed because [DataFrameWriteOptions] has private fields, so we +/// can't use it in our interfaces. This object can be converted to a +/// [DataFrameWriteOptions] using `.into()`. +pub struct SedonaWriteOptions { + /// Controls how new data should be written to the table, determining whether + /// to append, overwrite, or replace existing data. + pub insert_op: InsertOp, + /// Controls if all partitions should be coalesced into a single output file + /// Generally will have slower performance when set to true. + pub single_file_output: bool, + /// Sets which columns should be used for hive-style partitioned writes by name. + /// Can be set to empty vec![] for non-partitioned writes. + pub partition_by: Vec, + /// Sets which columns should be used for sorting the output by name. + /// Can be set to empty vec![] for non-sorted writes. + pub sort_by: Vec, +} + +impl From for DataFrameWriteOptions { + fn from(value: SedonaWriteOptions) -> Self { + DataFrameWriteOptions::new() + .with_insert_operation(value.insert_op) + .with_single_file_output(value.single_file_output) + .with_partition_by(value.partition_by) + .with_sort_by(value.sort_by) + } +} + +impl SedonaWriteOptions { + /// Create a new SedonaWriteOptions with default values + pub fn new() -> Self { + SedonaWriteOptions { + insert_op: InsertOp::Append, + single_file_output: false, + partition_by: vec![], + sort_by: vec![], + } + } + + /// Set the insert operation + pub fn with_insert_operation(mut self, insert_op: InsertOp) -> Self { + self.insert_op = insert_op; + self + } + + /// Set the single_file_output value to true or false + pub fn with_single_file_output(mut self, single_file_output: bool) -> Self { + self.single_file_output = single_file_output; + self + } + + /// Sets the partition_by columns for output partitioning + pub fn with_partition_by(mut self, partition_by: Vec) -> Self { + self.partition_by = partition_by; + self + } + + /// Sets the sort_by columns for output sorting + pub fn with_sort_by(mut self, sort_by: Vec) -> Self { + self.sort_by = sort_by; + self + } +} + +impl Default for SedonaWriteOptions { + fn default() -> Self { + Self::new() + } +} + // Because Dialect/dialect_from_str is not marked as Send, using the async // function in certain contexts will fail to compile. Here we use a wrapper // to ensure that that the Dialect can be specified and parsed in any async @@ -595,13 +782,9 @@ impl ThreadSafeDialect { #[cfg(test)] mod tests { - use arrow_array::{ - create_array, ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, - }; + use arrow_array::{create_array, ArrayRef, RecordBatchIterator, RecordBatchReader}; use arrow_schema::{DataType, Field, Schema}; - use async_trait::async_trait; - use datafusion::{assert_batches_eq, dataframe::DataFrameWriteOptions}; - use datafusion_common::not_impl_err; + use datafusion::assert_batches_eq; use sedona_datasource::spec::{Object, OpenReaderArgs}; use sedona_schema::{ crs::{deserialize_crs, lnglat}, @@ -611,8 +794,6 @@ mod tests { use sedona_testing::data::test_geoparquet; use tempfile::tempdir; - use crate::{dataframe::SedonaDataFrame, show::DisplayTableOptions}; - use super::*; #[tokio::test] diff --git a/rust/sedona/src/dataframe.rs b/rust/sedona/src/dataframe.rs deleted file mode 100644 index 5d21ed811a..0000000000 --- a/rust/sedona/src/dataframe.rs +++ /dev/null @@ -1,288 +0,0 @@ -// 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. - -use std::sync::Arc; - -use arrow_array::RecordBatch; -use async_trait::async_trait; -use datafusion::{ - dataframe::DataFrameWriteOptions, - datasource::file_format::format_as_file_type, - error::{DataFusionError, Result}, - prelude::DataFrame, -}; -use datafusion_common::not_impl_err; -use datafusion_expr::{dml::InsertOp, LogicalPlan, LogicalPlanBuilder, SortExpr}; -use sedona_geoparquet::{format::GeoParquetFormatFactory, options::TableGeoParquetOptions}; - -use crate::{ - context::SedonaContext, - show::{show_batches, DisplayTableOptions}, -}; - -/// Sedona-specific [`DataFrame`] actions -/// -/// This trait, implemented for [`DataFrame`], extends the DataFrame API to make it -/// ergonomic to work with dataframes that contain geometry columns. Currently these -/// are limited to output functions, as geometry columns currently require special -/// handling when written or exported to an external system. -#[async_trait] -pub trait SedonaDataFrame { - /// Build a table of the first `limit` results in this DataFrame - /// - /// This will limit and execute the query and build a table using [show_batches]. - async fn show_sedona<'a>( - self, - ctx: &SedonaContext, - limit: Option, - options: DisplayTableOptions<'a>, - ) -> Result; - - async fn write_geoparquet( - self, - ctx: &SedonaContext, - path: &str, - options: SedonaWriteOptions, - writer_options: Option, - ) -> Result>; -} - -#[async_trait] -impl SedonaDataFrame for DataFrame { - async fn show_sedona<'a>( - self, - ctx: &SedonaContext, - limit: Option, - mut options: DisplayTableOptions<'a>, - ) -> Result { - let df = if matches!( - self.logical_plan(), - LogicalPlan::Explain(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::Analyze(_) - ) { - // Show multi-line output without truncation for plans like `EXPLAIN` - options.max_row_height = usize::MAX; - - // We don't want to apply an additional .limit() to plans like `Explain` - // as that will trigger an internal error: Unsupported logical plan: Explain must be root of the plan - self - } else { - // Apply limit if specified - self.limit(0, limit)? - }; - - let schema_without_qualifiers = df.schema().clone().strip_qualifiers(); - let schema = schema_without_qualifiers.as_arrow(); - let batches = df.collect().await?; - let mut out = Vec::new(); - show_batches(ctx, &mut out, schema, batches, options)?; - String::from_utf8(out).map_err(|e| DataFusionError::External(Box::new(e))) - } - - async fn write_geoparquet( - self, - ctx: &SedonaContext, - path: &str, - options: SedonaWriteOptions, - writer_options: Option, - ) -> Result, DataFusionError> { - if options.insert_op != InsertOp::Append { - return not_impl_err!( - "{} is not implemented for DataFrame::write_geoparquet.", - options.insert_op - ); - } - - let format = if let Some(parquet_opts) = writer_options { - Arc::new(GeoParquetFormatFactory::new_with_options(parquet_opts)) - } else { - Arc::new(GeoParquetFormatFactory::new()) - }; - - let file_type = format_as_file_type(format); - - let plan = if options.sort_by.is_empty() { - self.into_unoptimized_plan() - } else { - LogicalPlanBuilder::from(self.into_unoptimized_plan()) - .sort(options.sort_by)? - .build()? - }; - - let plan = LogicalPlanBuilder::copy_to( - plan, - path.into(), - file_type, - Default::default(), - options.partition_by, - )? - .build()?; - - DataFrame::new(ctx.ctx.state(), plan).collect().await - } -} - -/// A Sedona-specific copy of [DataFrameWriteOptions] -/// -/// This is needed because [DataFrameWriteOptions] has private fields, so we -/// can't use it in our interfaces. This object can be converted to a -/// [DataFrameWriteOptions] using `.into()`. -pub struct SedonaWriteOptions { - /// Controls how new data should be written to the table, determining whether - /// to append, overwrite, or replace existing data. - pub insert_op: InsertOp, - /// Controls if all partitions should be coalesced into a single output file - /// Generally will have slower performance when set to true. - pub single_file_output: bool, - /// Sets which columns should be used for hive-style partitioned writes by name. - /// Can be set to empty vec![] for non-partitioned writes. - pub partition_by: Vec, - /// Sets which columns should be used for sorting the output by name. - /// Can be set to empty vec![] for non-sorted writes. - pub sort_by: Vec, -} - -impl From for DataFrameWriteOptions { - fn from(value: SedonaWriteOptions) -> Self { - DataFrameWriteOptions::new() - .with_insert_operation(value.insert_op) - .with_single_file_output(value.single_file_output) - .with_partition_by(value.partition_by) - .with_sort_by(value.sort_by) - } -} - -impl SedonaWriteOptions { - /// Create a new SedonaWriteOptions with default values - pub fn new() -> Self { - SedonaWriteOptions { - insert_op: InsertOp::Append, - single_file_output: false, - partition_by: vec![], - sort_by: vec![], - } - } - - /// Set the insert operation - pub fn with_insert_operation(mut self, insert_op: InsertOp) -> Self { - self.insert_op = insert_op; - self - } - - /// Set the single_file_output value to true or false - pub fn with_single_file_output(mut self, single_file_output: bool) -> Self { - self.single_file_output = single_file_output; - self - } - - /// Sets the partition_by columns for output partitioning - pub fn with_partition_by(mut self, partition_by: Vec) -> Self { - self.partition_by = partition_by; - self - } - - /// Sets the sort_by columns for output sorting - pub fn with_sort_by(mut self, sort_by: Vec) -> Self { - self.sort_by = sort_by; - self - } -} - -impl Default for SedonaWriteOptions { - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -mod test { - use tempfile::tempdir; - - use crate::context::SedonaContext; - - use super::*; - - #[tokio::test] - async fn show() { - let ctx = SedonaContext::new(); - let tbl = ctx - .sql("SELECT 1 as one") - .await - .unwrap() - .show_sedona(&ctx, None, DisplayTableOptions::default()) - .await - .unwrap(); - - #[rustfmt::skip] - assert_eq!( - tbl.lines().collect::>(), - vec![ - "+-----+", - "| one |", - "+-----+", - "| 1 |", - "+-----+" - ] - ); - } - - #[tokio::test] - async fn show_explain() { - let ctx = SedonaContext::new(); - for limit in [None, Some(10)] { - let tbl = ctx - .sql("EXPLAIN SELECT 1 as one") - .await - .unwrap() - .show_sedona(&ctx, limit, DisplayTableOptions::default()) - .await - .unwrap(); - - #[rustfmt::skip] - assert_eq!( - tbl.lines().collect::>(), - vec![ - "+---------------+---------------------------------+", - "| plan_type | plan |", - "+---------------+---------------------------------+", - "| logical_plan | Projection: Int64(1) AS one |", - "| | EmptyRelation: rows=1 |", - "| physical_plan | ProjectionExec: expr=[1 as one] |", - "| | PlaceholderRowExec |", - "| | |", - "+---------------+---------------------------------+", - ] - ); - } - } - - #[tokio::test] - async fn write_geoparquet() { - let tmpdir = tempdir().unwrap(); - let tmp_parquet = tmpdir.path().join("tmp.parquet"); - let ctx = SedonaContext::new(); - ctx.sql("SELECT 1 as one") - .await - .unwrap() - .write_parquet( - &tmp_parquet.to_string_lossy(), - DataFrameWriteOptions::default(), - None, - ) - .await - .unwrap(); - } -} diff --git a/rust/sedona/src/lib.rs b/rust/sedona/src/lib.rs index 7db7d8d0ca..dab1169a04 100644 --- a/rust/sedona/src/lib.rs +++ b/rust/sedona/src/lib.rs @@ -17,7 +17,6 @@ mod catalog; pub mod context; pub mod context_builder; -pub mod dataframe; mod exec; pub mod memory_pool; mod object_storage; From e24d3f9977fafe2524ed8ea4967af89659540fc0 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 16:46:55 -0500 Subject: [PATCH 30/55] smaller self review fixes --- c/sedona-extension/Cargo.toml | 2 +- rust/sedona-adbc/src/statement.rs | 3 ++- rust/sedona/src/context.rs | 18 +++++++----------- 3 files changed, 10 insertions(+), 13 deletions(-) diff --git a/c/sedona-extension/Cargo.toml b/c/sedona-extension/Cargo.toml index a416d4fb45..f0ca89a2f2 100644 --- a/c/sedona-extension/Cargo.toml +++ b/c/sedona-extension/Cargo.toml @@ -46,7 +46,7 @@ sedona-schema = { workspace = true } sedona-testing = { path = "../../rust/sedona-testing" } serde = { workspace = true } serde_json = { workspace = true } -tokio = { workspace = true, features = ["rt-multi-thread"] } +tokio = { workspace = true } [dev-dependencies] datafusion = { workspace = true, features = ["sql"] } diff --git a/rust/sedona-adbc/src/statement.rs b/rust/sedona-adbc/src/statement.rs index f8bb9344ba..74d4f83472 100644 --- a/rust/sedona-adbc/src/statement.rs +++ b/rust/sedona-adbc/src/statement.rs @@ -99,7 +99,8 @@ impl Statement for SedonaStatement { self.runtime.block_on(async { let df = self.ctx.sql(&query).await.map_err(from_datafusion_error)?; let stream = df.execute_stream().await.map_err(from_datafusion_error)?; - let reader = StreamingRecordBatchReader::new(stream, self.runtime.handle().clone()); + let reader = StreamingRecordBatchReader::new(stream, self.runtime.handle().clone()) + .with_skip_empty_batches(true); Ok(Box::new(reader) as Box) }) } else { diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index d44c9beb18..7b9a2b349b 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -29,27 +29,23 @@ use crate::{ use arrow_array::RecordBatch; use arrow_schema::DataType; use async_trait::async_trait; +use datafusion::datasource::file_format::format_as_file_type; use datafusion::{ common::plan_err, - dataframe::DataFrameWriteOptions, - datasource::file_format::format_as_file_type, error::{DataFusionError, Result}, execution::{ context::DataFilePaths, - memory_pool::MemoryLimit, runtime_env::{RuntimeEnv, RuntimeEnvBuilder}, SessionStateBuilder, }, prelude::{DataFrame, SessionConfig, SessionContext}, sql::parser::{DFParser, Statement}, }; - +use datafusion::{dataframe::DataFrameWriteOptions, execution::memory_pool::MemoryLimit}; use datafusion_common::not_impl_err; -use datafusion_expr::{ - dml::InsertOp, - sqlparser::dialect::{dialect_from_str, Dialect}, - AggregateUDFImpl, LogicalPlan, LogicalPlanBuilder, ScalarUDFImpl, SortExpr, -}; +use datafusion_expr::dml::InsertOp; +use datafusion_expr::sqlparser::dialect::{dialect_from_str, Dialect}; +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, @@ -57,10 +53,10 @@ use sedona_common::{ }; use sedona_datasource::provider::external_table; use sedona_datasource::spec::ExternalFormatSpec; +use sedona_expr::scalar_udf::{IntoScalarKernelRefs, SedonaScalarUDF}; use sedona_expr::{ aggregate_udf::{IntoSedonaAccumulatorRefs, SedonaAggregateUDF}, function_set::FunctionSet, - scalar_udf::{IntoScalarKernelRefs, SedonaScalarUDF}, }; use sedona_geoparquet::options::TableGeoParquetOptions; use sedona_geoparquet::{ @@ -90,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: RwLock, + functions: RwLock, /// Per-session registry of async raster byte loaders, keyed by /// `outdb_format`. Held behind an `Arc>` so the registered /// `RS_EnsureLoaded` UDF instance and any extension crates' `register(&ctx)` From 72bdedf6a4616ed8e7f841146ca0c7c69c76c175 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 17:18:01 -0500 Subject: [PATCH 31/55] document the ffi --- c/sedona-extension/src/sedona_extension.h | 74 ++++++++++++++++++----- 1 file changed, 59 insertions(+), 15 deletions(-) diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index c55ccbd6a8..94e6ad0686 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -205,10 +205,16 @@ struct SedonaCScalarKernel { }; struct SedonaCError { + /// \brief A UTF-8 encoded error message + /// + /// May be NULL if err_len is 0. This string is not necessarily null + /// terminated. const char* err; + /// \brief The number of bytes pointed to be err. uint32_t err_len; + /// \brief Reserved for future use. Must be 0. uint32_t reserved; /// \brief Release this instance @@ -218,9 +224,15 @@ struct SedonaCError { }; struct SedonaCExpr { - // Get a property of this expression (e.g., serialize, extract bbox) + /// \brief Get the data type of a property int (*get_property_schema)(const struct SedonaCExpr* self, const char* property, struct SedonaCError* err); + + /// \brief Extract a serializable property from this expression + /// + /// This is used to implement PlanProperties and other values that can be + /// easily retreived and serialized. The data type associated with the out + /// array may be retrieved with the get_property_schema callback. int (*get_property)(const struct SedonaCExpr* self, const char* property, const char* args, struct ArrowArray* out, struct SedonaCError* err); @@ -235,6 +247,7 @@ struct SedonaCExpr { void* private_data; }; +/// Forward declaration of the execution plan struct SedonaCExecutionPlan; /// \brief Arguments for execution plan and table provider operations @@ -242,45 +255,66 @@ struct SedonaCExecutionPlan; /// This structure is passed to methods that need JSON-serialized arguments, /// optional execution plans, and/or expressions. struct SedonaCExecutionPlanArgs { - /// JSON-serialized arguments + /// \brief JSON-serialized arguments const uint8_t* args; size_t args_len; - /// Optional array of execution plans + + /// \brief Optional array of execution plans const struct SedonaCExecutionPlan** exec_plans; size_t num_exec_plans; - /// Optional array of expressions + + /// \brief Optional array of expressions const struct SedonaCExpr** exprs; size_t num_exprs; + + /// \brief Reserved for future use. Must be NULL. void* reserved; }; +/// \brief FFI interface for a physical execution operator +/// +/// Before using this structure, the release callback MUST be checked. +/// Instances with a NULL release callback are not valid and must not be used. struct SedonaCExecutionPlan { + /// \brief Get the schema associated with the output of this plan void (*get_schema)(const struct SedonaCExecutionPlan* self, struct ArrowSchema* out); - // Extract some serializable property from this plan (e.g., plan properties) + /// \brief Get the data type of a property int (*get_property_schema)(const struct SedonaCExecutionPlan* self, const char* property, struct ArrowSchema* out, struct SedonaCError* err); + + /// \brief Extract a serializable property from this plan + /// + /// This is used to implement PlanProperties and other values that can be + /// easily retreived and serialized. The data type associated with the out + /// array may be retrieved with the get_property_schema callback. int (*get_property)(const struct SedonaCExecutionPlan* self, const char* property, struct SedonaCExecutionPlanArgs* args, struct ArrowArray* out, struct SedonaCError* err); - // Clone this plan based on some new information (e.g., try pushdown filters) + /// \brief Clone this plan based on information about a property + /// + /// This can used to implement operations that require modifying a plan. int (*with_property)(const struct SedonaCExecutionPlan* self, const char* property, struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, struct SedonaCError* err); - // Resolve a synchronous stream for one partition from this plan + /// \brief Resolve a synchronous stream for one partition from this plan int (*execute)(const struct SedonaCExecutionPlan* self, struct SedonaCExecutionPlanArgs* args, struct ArrowArrayStream* out, struct SedonaCError* err); - // Future implementation with async streams (don't implement now) + /// \brief Resolve an asynchronous stream for one partition from this plan + /// + /// This is not currently implemented and must be NULL. In the future, + /// out must point to a caller-supplied struct ArrowAsyncDeviceStreamHandler + /// as specified in the Arrow C Device Async Stream specification. int (*execute_async)(const struct SedonaCExecutionPlan* self, struct SedonaCExecutionPlanArgs* args, void* out, struct SedonaCError* err); - // Reserved for future use + /// \brief Reserved for future use (must be NULL). void* reserved; /// \brief Release this instance @@ -296,19 +330,28 @@ struct SedonaCExecutionPlan { /// /// This provides a minimal interface for importing a table provider /// across an FFI boundary in a version-agnostic manner. +/// +/// Before using this structure, the release callback MUST be checked. +/// Instances with a NULL release callback are not valid and must not be used. struct SedonaCTableProvider { - /// Get the schema of this table provider + /// \brief Get the schema of this table provider void (*get_schema)(const struct SedonaCTableProvider* self, struct ArrowSchema* out); - // Extract some serializable property from this table provider + /// \brief Get the data type of a property int (*get_property_schema)(const struct SedonaCTableProvider* self, const char* property, struct ArrowSchema* out, struct SedonaCError* err); + + /// \brief Extract a serializable property from this table provider + /// + /// This is used to implement PlanProperties and other values that can be + /// easily retreived and serialized. The data type associated with the out + /// array may be retrieved with the get_property_schema callback. int (*get_property)(const struct SedonaCTableProvider* self, const char* property, struct SedonaCExecutionPlanArgs* args, struct ArrowArray* out, struct SedonaCError* err); - /// Perform a scan operation and return an execution plan + /// \brief Perform a scan operation and return an execution plan /// /// The args parameter contains JSON-serialized scan arguments. /// Returns an execution plan that can be used to read the data. @@ -316,7 +359,7 @@ struct SedonaCTableProvider { struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, struct SedonaCError* err); - /// Perform an insert operation + /// \brief Perform an insert operation /// /// The args parameter contains JSON-serialized insert arguments. /// The exec_plans field should contain the plan providing rows to insert. @@ -325,7 +368,7 @@ struct SedonaCTableProvider { struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, struct SedonaCError* err); - /// Perform an update operation + /// \brief Perform an update operation /// /// The args parameter contains JSON-serialized update arguments /// (filters, column assignments, etc.). @@ -334,7 +377,7 @@ struct SedonaCTableProvider { struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, struct SedonaCError* err); - /// Perform a delete operation + /// \brief Perform a delete operation /// /// The args parameter contains JSON-serialized delete arguments /// (filters, etc.). @@ -343,6 +386,7 @@ struct SedonaCTableProvider { struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, struct SedonaCError* err); + /// \brief Reserved for future use. Must be NULL. void* reserved; /// \brief Release this instance From 5b0c58b41c6304cccdf00d1a53359cdb043691fb Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 29 Jun 2026 22:23:15 -0500 Subject: [PATCH 32/55] review comments --- c/sedona-extension/src/execution_plan.rs | 40 ++- c/sedona-extension/src/expr.rs | 346 -------------------- c/sedona-extension/src/lib.rs | 1 - c/sedona-extension/src/sedona_extension.h | 6 +- c/sedona-extension/src/table_provider.rs | 39 ++- c/sedona-extension/src/utils.rs | 24 +- python/sedonadb/tests/test_dataframe_ffi.py | 2 +- 7 files changed, 83 insertions(+), 375 deletions(-) delete mode 100644 c/sedona-extension/src/expr.rs diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 9bf09595d2..3259a22fa1 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -17,7 +17,7 @@ use std::{ any::Any, - ffi::{c_int, c_void, CStr}, + ffi::{c_int, c_void}, fmt::{Debug, Display, Formatter}, ptr::null_mut, sync::Arc, @@ -38,7 +38,7 @@ use serde::{Deserialize, Serialize}; use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs}; use crate::streaming::{ffi_stream_to_sendable, StreamingRecordBatchReader}; -use crate::utils::{get_plan_property, get_plan_string_property, ERRNO_OK}; +use crate::utils::{cstr_from_ptr_or_empty, get_plan_property, get_plan_string_property, ERRNO_OK}; /// Wrapper around an [ExecutionPlan] that can be exported across FFI. pub struct ExportedExecutionPlan { @@ -172,7 +172,12 @@ unsafe extern "C" fn c_exec_plan_get_schema( self_: *const SedonaCExecutionPlan, out: *mut FFI_ArrowSchema, ) { - let plan = &*((*self_).private_data as *const ExportedExecutionPlan); + debug_assert!(!self_.is_null(), "self pointer is null"); + debug_assert!(!out.is_null(), "out pointer is null"); + let self_ref = &*self_; + debug_assert!(!self_ref.private_data.is_null(), "private_data is null"); + let plan = &*(self_ref.private_data as *const ExportedExecutionPlan); + let schema = plan.schema(); if let Ok(ffi_schema) = FFI_ArrowSchema::try_from(schema.as_ref()) { std::ptr::write(out, ffi_schema); @@ -185,6 +190,7 @@ unsafe extern "C" fn c_exec_plan_get_property_schema( out: *mut FFI_ArrowSchema, err: *mut SedonaCError, ) -> c_int { + debug_assert!(!out.is_null(), "out pointer is null"); // All properties are returned as Utf8 strings (including JSON) use arrow_schema::{DataType, Field}; let field = Field::new("value", DataType::Utf8, false); @@ -209,8 +215,12 @@ unsafe extern "C" fn c_exec_plan_get_property( out: *mut arrow_array::ffi::FFI_ArrowArray, err: *mut SedonaCError, ) -> c_int { - let plan = &*((*self_).private_data as *const ExportedExecutionPlan); - let property_str = CStr::from_ptr(property).to_string_lossy(); + debug_assert!(!self_.is_null(), "self pointer is null"); + debug_assert!(!out.is_null(), "out pointer is null"); + let self_ref = &*self_; + debug_assert!(!self_ref.private_data.is_null(), "private_data is null"); + let plan = &*(self_ref.private_data as *const ExportedExecutionPlan); + let property_str = cstr_from_ptr_or_empty(property); match plan.get_property(&property_str) { Ok(value) => { @@ -238,10 +248,14 @@ unsafe extern "C" fn c_exec_plan_execute( out: *mut FFI_ArrowArrayStream, err: *mut SedonaCError, ) -> c_int { - let plan = &*((*self_).private_data as *const ExportedExecutionPlan); - - // Parse the execute args + debug_assert!(!self_.is_null(), "self pointer is null"); + debug_assert!(!args.is_null(), "args pointer is null"); + debug_assert!(!out.is_null(), "out pointer is null"); + let self_ref = &*self_; let args_ref = &*args; + debug_assert!(!self_ref.private_data.is_null(), "private_data is null"); + let plan = &*(self_ref.private_data as *const ExportedExecutionPlan); + let args_slice = if args_ref.args.is_null() || args_ref.args_len == 0 { &[] } else { @@ -276,11 +290,13 @@ unsafe extern "C" fn c_exec_plan_execute( } unsafe extern "C" fn c_exec_plan_release(self_: *mut SedonaCExecutionPlan) { - if !(*self_).private_data.is_null() { - let _ = Box::from_raw((*self_).private_data as *mut ExportedExecutionPlan); - (*self_).private_data = null_mut(); + debug_assert!(!self_.is_null(), "self pointer is null"); + let self_ref = &mut *self_; + if !self_ref.private_data.is_null() { + let _ = Box::from_raw(self_ref.private_data as *mut ExportedExecutionPlan); + self_ref.private_data = null_mut(); } - (*self_).release = None; + self_ref.release = None; } /// An [ExecutionPlan] that wraps an imported [SedonaCExecutionPlan]. diff --git a/c/sedona-extension/src/expr.rs b/c/sedona-extension/src/expr.rs deleted file mode 100644 index cc25b283ed..0000000000 --- a/c/sedona-extension/src/expr.rs +++ /dev/null @@ -1,346 +0,0 @@ -// 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. - -use std::{ - ffi::{c_char, c_int, c_void, CStr}, - fmt::Debug, - ptr::null_mut, - sync::Arc, -}; - -use arrow_array::{builder::StringBuilder, ffi::FFI_ArrowArray, Array}; -use arrow_schema::SchemaRef; -use datafusion_common::{DataFusionError, Result}; -use datafusion_physical_expr::PhysicalExpr; -use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; -use sedona_expr::spatial_filter::SpatialFilterFactory; -use serde::Deserialize; - -use crate::extension::{SedonaCError, SedonaCExpr}; -use crate::utils::ERRNO_OK; - -/// A wrapper around a DataFusion [PhysicalExpr] with its associated schema. -/// -/// This struct bundles a physical expression with the schema context needed to -/// interpret it, which is required for operations like type resolution -/// and extracting spatial filter information. -pub struct PhysicalExprWithSchema { - expr: Arc, - schema: SchemaRef, -} - -impl Debug for PhysicalExprWithSchema { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("PhysicalExprWithSchema") - .field("expr", &self.expr.to_string()) - .field("schema", &self.schema) - .finish() - } -} - -impl PhysicalExprWithSchema { - /// Create a new PhysicalExprWithSchema - pub fn new(expr: Arc, schema: SchemaRef) -> Self { - Self { expr, schema } - } - - /// Get a reference to the expression - pub fn expr(&self) -> &Arc { - &self.expr - } - - /// Get a reference to the schema - pub fn schema(&self) -> &SchemaRef { - &self.schema - } - - /// Consume self and return the inner expression and schema - pub fn into_parts(self) -> (Arc, SchemaRef) { - (self.expr, self.schema) - } - - /// Extract the bounding box for a spatial filter on the given column - /// - /// Returns the bounding box as a JSON string, or an error if extraction fails. - pub fn filter_bbox(&self, column_name: &str) -> Result { - let factory = SpatialFilterFactory::default(); - let spatial_filter = factory.try_from_expr(&self.expr)?; - let bbox = spatial_filter.filter_bbox(column_name); - serde_json::to_string(&bbox) - .map_err(|e| sedona_internal_datafusion_err!("Failed to serialize bounding box: {}", e)) - } -} - -/// Arguments for the bbox property -#[derive(Debug, Deserialize)] -struct BboxArgs { - column: String, -} - -/// Wrapper around a [SedonaCExpr] that can be used to import an expression -/// from a C implementation. -pub struct ImportedExpr { - inner: SedonaCExpr, -} - -impl Debug for ImportedExpr { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ImportedExpr") - .field("inner", &"") - .finish() - } -} - -impl TryFrom for ImportedExpr { - type Error = DataFusionError; - - fn try_from(value: SedonaCExpr) -> Result { - match (&value.get_property, &value.release) { - (Some(_), Some(_)) => Ok(Self { inner: value }), - _ => Err(sedona_internal_datafusion_err!( - "Can't import released or uninitialized SedonaCExpr" - )), - } - } -} - -impl ImportedExpr { - /// Get a property from this expression - /// - /// # Safety - /// - /// The caller must ensure that the property name and args are valid - /// null-terminated C strings if provided. - pub unsafe fn get_property( - &self, - property: &CStr, - args: Option<&CStr>, - ) -> Result { - let get_property = self - .inner - .get_property - .ok_or_else(|| sedona_internal_datafusion_err!("get_property callback is null"))?; - - let args_ptr = args.map(|a| a.as_ptr()).unwrap_or(std::ptr::null()); - let mut out = FFI_ArrowArray::empty(); - let mut err = SedonaCError::default(); - - let result = get_property(&self.inner, property.as_ptr(), args_ptr, &mut out, &mut err); - - if result != ERRNO_OK { - return sedona_internal_err!("get_property failed: {}", err); - } - - Ok(out) - } -} - -/// Export a [PhysicalExprWithSchema] as a [SedonaCExpr] for use across FFI boundaries. -pub struct ExportedExpr { - inner: Arc, -} - -impl ExportedExpr { - /// Create a new ExportedExpr from a PhysicalExprWithSchema - pub fn new(expr_with_schema: PhysicalExprWithSchema) -> Self { - Self { - inner: Arc::new(expr_with_schema), - } - } - - /// Export this expression as a SedonaCExpr - /// - /// The returned SedonaCExpr takes ownership of the Arc and will release - /// it when the release callback is called. - pub fn export(self) -> SedonaCExpr { - let boxed = Box::new(self.inner); - - SedonaCExpr { - get_property_schema: Some(exported_expr_get_property_schema), - get_property: Some(exported_expr_get_property), - reserved: null_mut(), - release: Some(exported_expr_release), - private_data: Box::into_raw(boxed) as *mut c_void, - } - } -} - -unsafe extern "C" fn exported_expr_get_property_schema( - _self_: *const SedonaCExpr, - _property: *const c_char, - err: *mut SedonaCError, -) -> c_int { - // TODO: Implement property schema retrieval - if !err.is_null() { - *err = SedonaCError::new("get_property_schema not implemented"); - } - libc::EINVAL -} - -unsafe extern "C" fn exported_expr_get_property( - self_: *const SedonaCExpr, - property: *const c_char, - args: *const c_char, - out: *mut FFI_ArrowArray, - err: *mut SedonaCError, -) -> c_int { - if self_.is_null() || (*self_).private_data.is_null() { - if !err.is_null() { - *err = SedonaCError::new("null self pointer"); - } - return libc::EINVAL; - } - - let expr_arc = &*((*self_).private_data as *const Arc); - - if property.is_null() { - if !err.is_null() { - *err = SedonaCError::new("null property pointer"); - } - return libc::EINVAL; - } - - let property_str = match CStr::from_ptr(property).to_str() { - Ok(s) => s, - Err(_) => { - if !err.is_null() { - *err = SedonaCError::new("invalid UTF-8 in property name"); - } - return libc::EINVAL; - } - }; - - match property_str { - "bbox" => { - // Parse args JSON to get column name - if args.is_null() { - if !err.is_null() { - *err = SedonaCError::new("bbox property requires args with column name"); - } - return libc::EINVAL; - } - - let args_str = match CStr::from_ptr(args).to_str() { - Ok(s) => s, - Err(_) => { - if !err.is_null() { - *err = SedonaCError::new("invalid UTF-8 in args"); - } - return libc::EINVAL; - } - }; - - let bbox_args: BboxArgs = match serde_json::from_str(args_str) { - Ok(a) => a, - Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&format!("failed to parse args: {}", e)); - } - return libc::EINVAL; - } - }; - - // Extract bbox using SpatialFilterFactory - let bbox_json = match expr_arc.filter_bbox(&bbox_args.column) { - Ok(json) => json, - Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&format!("failed to extract bbox: {}", e)); - } - return libc::EINVAL; - } - }; - - // Create a length-1 string array with the JSON result - let mut builder = StringBuilder::new(); - builder.append_value(&bbox_json); - let array = builder.finish(); - - // Export the array via FFI - std::ptr::write(out, FFI_ArrowArray::new(&array.to_data())); - ERRNO_OK - } - _ => { - if !err.is_null() { - *err = SedonaCError::new(&format!("unknown property: {}", property_str)); - } - libc::EINVAL - } - } -} - -unsafe extern "C" fn exported_expr_release(self_: *mut SedonaCExpr) { - if self_.is_null() { - return; - } - - let expr = &mut *self_; - if !expr.private_data.is_null() { - let _ = Box::from_raw(expr.private_data as *mut Arc); - expr.private_data = null_mut(); - } - - expr.get_property_schema = None; - expr.get_property = None; - expr.release = None; -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow_schema::{DataType, Field, Schema}; - use datafusion_common::ScalarValue; - use datafusion_physical_expr::expressions::Literal; - - #[test] - fn test_physical_expr_with_schema_new() { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let expr: Arc = Arc::new(Literal::new(ScalarValue::Int32(Some(42)))); - let expr_with_schema = PhysicalExprWithSchema::new(expr.clone(), schema.clone()); - - assert!(Arc::ptr_eq(&expr_with_schema.schema, &schema)); - } - - #[test] - fn test_physical_expr_with_schema_into_parts() { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let expr: Arc = Arc::new(Literal::new(ScalarValue::Int32(Some(42)))); - let expr_with_schema = PhysicalExprWithSchema::new(expr.clone(), schema.clone()); - - let (returned_expr, returned_schema) = expr_with_schema.into_parts(); - assert_eq!(returned_expr.to_string(), expr.to_string()); - assert!(Arc::ptr_eq(&returned_schema, &schema)); - } - - #[test] - fn test_exported_expr_roundtrip() { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let expr: Arc = Arc::new(Literal::new(ScalarValue::Int32(Some(42)))); - let expr_with_schema = PhysicalExprWithSchema::new(expr, schema); - - let exported = ExportedExpr::new(expr_with_schema); - let mut c_expr = exported.export(); - - // Verify release callback works - assert!(c_expr.release.is_some()); - unsafe { - (c_expr.release.unwrap())(&mut c_expr); - } - assert!(c_expr.release.is_none()); - assert!(c_expr.private_data.is_null()); - } -} diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs index 68b82ca87c..de9835521f 100644 --- a/c/sedona-extension/src/lib.rs +++ b/c/sedona-extension/src/lib.rs @@ -16,7 +16,6 @@ // under the License. pub mod execution_plan; -pub mod expr; pub mod extension; pub mod scalar_kernel; pub mod streaming; diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 94e6ad0686..1d2453896a 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -231,7 +231,7 @@ struct SedonaCExpr { /// \brief Extract a serializable property from this expression /// /// This is used to implement PlanProperties and other values that can be - /// easily retreived and serialized. The data type associated with the out + /// easily retrieved and serialized. The data type associated with the out /// array may be retrieved with the get_property_schema callback. int (*get_property)(const struct SedonaCExpr* self, const char* property, const char* args, struct ArrowArray* out, struct SedonaCError* err); @@ -287,7 +287,7 @@ struct SedonaCExecutionPlan { /// \brief Extract a serializable property from this plan /// /// This is used to implement PlanProperties and other values that can be - /// easily retreived and serialized. The data type associated with the out + /// easily retrieved and serialized. The data type associated with the out /// array may be retrieved with the get_property_schema callback. int (*get_property)(const struct SedonaCExecutionPlan* self, const char* property, struct SedonaCExecutionPlanArgs* args, struct ArrowArray* out, @@ -345,7 +345,7 @@ struct SedonaCTableProvider { /// \brief Extract a serializable property from this table provider /// /// This is used to implement PlanProperties and other values that can be - /// easily retreived and serialized. The data type associated with the out + /// easily retrieved and serialized. The data type associated with the out /// array may be retrieved with the get_property_schema callback. int (*get_property)(const struct SedonaCTableProvider* self, const char* property, struct SedonaCExecutionPlanArgs* args, struct ArrowArray* out, diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index f0f371eb22..5935fa5180 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -17,7 +17,7 @@ use std::{ any::Any, - ffi::{c_int, c_void, CStr}, + ffi::{c_int, c_void}, fmt::Debug, ptr::null_mut, sync::Arc, @@ -37,7 +37,7 @@ use crate::execution_plan::{ExportedExecutionPlan, ImportedSedonaCExec}; use crate::extension::{ SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs, SedonaCTableProvider, }; -use crate::utils::{get_table_provider_string_property, ERRNO_OK}; +use crate::utils::{cstr_from_ptr_or_empty, get_table_provider_string_property, ERRNO_OK}; /// A TableProvider wrapper that can be exported across FFI. /// @@ -131,7 +131,12 @@ unsafe extern "C" fn c_table_provider_get_schema( self_: *const SedonaCTableProvider, out: *mut FFI_ArrowSchema, ) { - let provider = &*((*self_).private_data as *const ExportedTableProvider); + debug_assert!(!self_.is_null(), "self pointer is null"); + debug_assert!(!out.is_null(), "out pointer is null"); + let self_ref = &*self_; + debug_assert!(!self_ref.private_data.is_null(), "private_data is null"); + let provider = &*(self_ref.private_data as *const ExportedTableProvider); + let schema = provider.inner.schema(); if let Ok(ffi_schema) = FFI_ArrowSchema::try_from(schema.as_ref()) { std::ptr::write(out, ffi_schema); @@ -144,6 +149,7 @@ unsafe extern "C" fn c_table_provider_get_property_schema( out: *mut FFI_ArrowSchema, err: *mut SedonaCError, ) -> c_int { + debug_assert!(!out.is_null(), "out pointer is null"); // All properties are returned as Utf8 strings use arrow_schema::{DataType, Field}; let field = Field::new("value", DataType::Utf8, false); @@ -168,8 +174,12 @@ unsafe extern "C" fn c_table_provider_get_property( out: *mut FFI_ArrowArray, err: *mut SedonaCError, ) -> c_int { - let provider = &*((*self_).private_data as *const ExportedTableProvider); - let property_str = CStr::from_ptr(property).to_string_lossy(); + debug_assert!(!self_.is_null(), "self pointer is null"); + debug_assert!(!out.is_null(), "out pointer is null"); + let self_ref = &*self_; + debug_assert!(!self_ref.private_data.is_null(), "private_data is null"); + let provider = &*(self_ref.private_data as *const ExportedTableProvider); + let property_str = cstr_from_ptr_or_empty(property); match provider.get_property(&property_str) { Ok(value) => { @@ -197,10 +207,15 @@ unsafe extern "C" fn c_table_provider_scan( out: *mut SedonaCExecutionPlan, err: *mut SedonaCError, ) -> c_int { - let provider = &*((*self_).private_data as *const ExportedTableProvider); + debug_assert!(!self_.is_null(), "self pointer is null"); + debug_assert!(!args.is_null(), "args pointer is null"); + debug_assert!(!out.is_null(), "out pointer is null"); + let self_ref = &*self_; + let args_ref = &*args; + debug_assert!(!self_ref.private_data.is_null(), "private_data is null"); + let provider = &*(self_ref.private_data as *const ExportedTableProvider); // Parse scan args - let args_ref = &*args; let args_slice = if args_ref.args.is_null() || args_ref.args_len == 0 { &[] } else { @@ -242,11 +257,13 @@ unsafe extern "C" fn c_table_provider_scan( } unsafe extern "C" fn c_table_provider_release(self_: *mut SedonaCTableProvider) { - if !(*self_).private_data.is_null() { - let _ = Box::from_raw((*self_).private_data as *mut ExportedTableProvider); - (*self_).private_data = null_mut(); + debug_assert!(!self_.is_null(), "self pointer is null"); + let self_ref = &mut *self_; + if !self_ref.private_data.is_null() { + let _ = Box::from_raw(self_ref.private_data as *mut ExportedTableProvider); + self_ref.private_data = null_mut(); } - (*self_).release = None; + self_ref.release = None; } /// A TableProvider that wraps an imported SedonaCTableProvider. diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index d7a91a8965..04adb06f87 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -17,7 +17,8 @@ //! Utilities for FFI property access and conversion. -use std::ffi::{c_int, CString}; +use std::borrow::Cow; +use std::ffi::{c_char, c_int, CStr, CString}; use std::ptr::null_mut; use arrow_schema::ffi::FFI_ArrowSchema; @@ -34,6 +35,21 @@ use crate::extension::{ /// Success return code for FFI functions. pub const ERRNO_OK: c_int = 0; +/// Safely convert a C string pointer to a Rust string, treating null as empty. +/// +/// # Safety +/// +/// The pointer, if non-null, must point to a valid null-terminated C string. +/// The string must remain valid for the duration of the returned `Cow`. +#[inline] +pub unsafe fn cstr_from_ptr_or_empty(ptr: *const c_char) -> Cow<'static, str> { + if ptr.is_null() { + Cow::Borrowed("") + } else { + CStr::from_ptr(ptr).to_string_lossy() + } +} + /// Get the schema for a property from a [SedonaCExecutionPlan]. /// /// Returns the DataType describing the property's data type. @@ -164,6 +180,12 @@ fn parse_ffi_array_to_bytes( let data = unsafe { arrow_array::ffi::from_ffi_and_data_type(ffi_array, data_type.clone())? }; let array = arrow_array::make_array(data); + if array.len() != 1 || array.null_count() != 0 { + return sedona_internal_err!( + "Expected get_property() to return non-null array of length 1" + ); + } + // Handle different array types match data_type { DataType::Binary => { diff --git a/python/sedonadb/tests/test_dataframe_ffi.py b/python/sedonadb/tests/test_dataframe_ffi.py index 9d0c54eee3..7b230a83aa 100644 --- a/python/sedonadb/tests/test_dataframe_ffi.py +++ b/python/sedonadb/tests/test_dataframe_ffi.py @@ -22,7 +22,7 @@ from sedonadb.testing import skip_if_not_exists -# Test cases: (producer_sql, consumer_sql, description) +# Test cases: (producer_sql, consumer_sql) FFI_TEST_CASES = [ # Simple select all pytest.param( From a8b9c59828d497b4d74dd5dd95fe04c1ff1dcd84 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 30 Jun 2026 09:01:01 -0500 Subject: [PATCH 33/55] fix two more issues --- c/sedona-extension/src/streaming.rs | 4 +++- c/sedona-extension/src/utils.rs | 3 +-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index 2ef03552c7..816fa0ea1b 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -137,7 +137,9 @@ impl StreamingRecordBatchReader { // Create channels for communication // Use bounded channel with size 0 (rendezvous) for backpressure let (request_tx, request_rx) = std::sync::mpsc::sync_channel::<()>(0); - let (response_tx, response_rx) = std::sync::mpsc::sync_channel::(0); + // Use bounded channel with size 1 for response to prevent deadlock if reader + // stops early due to periodic cancellation before receiving the response + let (response_tx, response_rx) = std::sync::mpsc::sync_channel::(1); // Spawn the worker thread let handle = std::thread::spawn(move || { diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index 04adb06f87..fb3da8c70b 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -41,8 +41,7 @@ pub const ERRNO_OK: c_int = 0; /// /// The pointer, if non-null, must point to a valid null-terminated C string. /// The string must remain valid for the duration of the returned `Cow`. -#[inline] -pub unsafe fn cstr_from_ptr_or_empty(ptr: *const c_char) -> Cow<'static, str> { +pub unsafe fn cstr_from_ptr_or_empty<'a>(ptr: *const c_char) -> Cow<'a, str> { if ptr.is_null() { Cow::Borrowed("") } else { From 375aca54289dac6205925dcc904b59414b33f36d Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 16:31:27 -0500 Subject: [PATCH 34/55] simpler cancellation --- c/sedona-extension/src/streaming.rs | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index 816fa0ea1b..04a37fc96d 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -195,6 +195,7 @@ impl StreamingRecordBatchReader { // Check for cancellation if let Some(ref checker) = self.cancel_checker { if checker() { + self.cancelled = true; return Some(Err(ArrowError::ExternalError(Box::new( std::io::Error::new( std::io::ErrorKind::Interrupted, @@ -255,14 +256,6 @@ impl Iterator for StreamingRecordBatchReader { return Some(Ok(batch)); } Some(Err(e)) => { - // Check if this was a cancellation error from periodic checking - if let ArrowError::ExternalError(box_err) = &e { - if let Some(io_err) = box_err.downcast_ref::() { - if io_err.kind() == std::io::ErrorKind::Interrupted { - self.cancelled = true; - } - } - } return Some(Err(e)); } None => return None, From b0d670b0d578f19982d4a41bce7dba9e5cdee406 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 16:33:15 -0500 Subject: [PATCH 35/55] fix release callback --- c/sedona-extension/src/extension.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs index 309c2c2ab1..4e01582123 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -185,7 +185,12 @@ impl Drop for SedonaCError { } } -extern "C" fn sedona_c_noop_release(_self: *mut SedonaCError) {} +extern "C" fn sedona_c_noop_release(self_: *mut SedonaCError) { + // Must set release to NULL per Arrow C Data Interface contract + unsafe { + (*self_).release = None; + } +} pub const UNKNOWN_SEDONA_C_ERROR: SedonaCError = SedonaCError { err: c"Unknown error".as_ptr(), From ee9a22ee1257704d8d4f3b81ee1d0aec427330c8 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 16:41:01 -0500 Subject: [PATCH 36/55] deduplicate property getters --- c/sedona-extension/src/utils.rs | 242 +++++++++++++++----------------- 1 file changed, 114 insertions(+), 128 deletions(-) diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs index fb3da8c70b..ae42851bb1 100644 --- a/c/sedona-extension/src/utils.rs +++ b/c/sedona-extension/src/utils.rs @@ -49,34 +49,21 @@ pub unsafe fn cstr_from_ptr_or_empty<'a>(ptr: *const c_char) -> Cow<'a, str> { } } -/// Get the schema for a property from a [SedonaCExecutionPlan]. -/// -/// Returns the DataType describing the property's data type. -/// If `get_property_schema` is not implemented, defaults to Binary. -fn get_plan_property_data_type(plan: &SedonaCExecutionPlan, property: &str) -> Result { - let Some(get_property_schema) = plan.get_property_schema else { - // Default to Binary if get_property_schema is not implemented - return Ok(DataType::Binary); +/// Get a string property from a [SedonaCTableProvider]. +pub fn get_table_provider_string_property( + provider: &SedonaCTableProvider, + property: &str, +) -> Result { + let Some(get_property) = provider.get_property else { + return sedona_internal_err!("SedonaCTableProvider does not have get_property"); }; - let property_cstr = CString::new(property) - .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; - - let mut ffi_schema = FFI_ArrowSchema::empty(); - let mut err = SedonaCError::default(); - - let code = - unsafe { get_property_schema(plan, property_cstr.as_ptr(), &mut ffi_schema, &mut err) }; - - if code != ERRNO_OK { - return sedona_internal_err!("Failed to get property schema for '{}': {}", property, err); - } - - // Try to convert the FFI schema to a Field - let field = Field::try_from(&ffi_schema).map_err(|e| { - sedona_internal_datafusion_err!("Failed to parse property schema for '{}': {}", property, e) - })?; - Ok(field.data_type().clone()) + call_get_string_property_impl( + property, + "SedonaCTableProvider", + |prop, args, out, err| unsafe { get_property(provider, prop, args, out, err) }, + || get_table_provider_property_data_type(provider, property), + ) } /// Call `get_property` on a [SedonaCExecutionPlan] and deserialize the result. @@ -161,6 +148,98 @@ where parse_ffi_array(ffi_array, &data_type) } +/// Core implementation for getting property data type via FFI. +/// +/// The caller provides a closure that performs the actual FFI call with +/// the correct self pointer type. +fn call_get_property_schema_impl(property: &str, call_ffi: F) -> Result +where + F: FnOnce(*const c_char, *mut FFI_ArrowSchema, *mut SedonaCError) -> c_int, +{ + let property_cstr = CString::new(property) + .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; + + let mut ffi_schema = FFI_ArrowSchema::empty(); + let mut err = SedonaCError::default(); + + let code = call_ffi(property_cstr.as_ptr(), &mut ffi_schema, &mut err); + + if code != ERRNO_OK { + return sedona_internal_err!("Failed to get property schema for '{}': {}", property, err); + } + + // Try to convert the FFI schema to a Field + let field = Field::try_from(&ffi_schema).map_err(|e| { + sedona_internal_datafusion_err!("Failed to parse property schema for '{}': {}", property, e) + })?; + Ok(field.data_type().clone()) +} + +/// Core implementation for getting a string property via FFI. +/// +/// The caller provides closures for the FFI call and data type lookup. +fn call_get_string_property_impl( + property: &str, + type_name: &str, + call_get_property: F, + get_data_type: G, +) -> Result +where + F: FnOnce( + *const c_char, + *mut SedonaCExecutionPlanArgs, + *mut arrow_array::ffi::FFI_ArrowArray, + *mut SedonaCError, + ) -> c_int, + G: FnOnce() -> Result, +{ + let property_cstr = CString::new(property) + .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; + + let mut ffi_args = SedonaCExecutionPlanArgs { + args: std::ptr::null(), + args_len: 0, + exec_plans: std::ptr::null(), + num_exec_plans: 0, + exprs: std::ptr::null(), + num_exprs: 0, + reserved: null_mut(), + }; + + let mut ffi_array = arrow_array::ffi::FFI_ArrowArray::empty(); + let mut err = SedonaCError::default(); + + let code = call_get_property( + property_cstr.as_ptr(), + &mut ffi_args, + &mut ffi_array, + &mut err, + ); + + if code != ERRNO_OK { + return sedona_internal_err!("{} failed to get '{}': {}", type_name, property, err); + } + + let data_type = get_data_type()?; + let bytes = parse_ffi_array_to_bytes(ffi_array, &data_type)?; + String::from_utf8(bytes) + .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) +} + +/// Get the schema for a property from a [SedonaCExecutionPlan]. +/// +/// Returns the DataType describing the property's data type. +/// If `get_property_schema` is not implemented, defaults to Binary. +fn get_plan_property_data_type(plan: &SedonaCExecutionPlan, property: &str) -> Result { + let Some(get_property_schema) = plan.get_property_schema else { + return Ok(DataType::Binary); + }; + + call_get_property_schema_impl(property, |prop, schema, err| unsafe { + get_property_schema(plan, prop, schema, err) + }) +} + /// Parse an FFI array containing JSON and deserialize to the target type. fn parse_ffi_array( ffi_array: arrow_array::ffi::FFI_ArrowArray, @@ -235,42 +314,12 @@ pub fn get_plan_string_property(plan: &SedonaCExecutionPlan, property: &str) -> return sedona_internal_err!("SedonaCExecutionPlan does not have get_property"); }; - let property_cstr = CString::new(property) - .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; - - let mut ffi_args = SedonaCExecutionPlanArgs { - args: std::ptr::null(), - args_len: 0, - exec_plans: std::ptr::null(), - num_exec_plans: 0, - exprs: std::ptr::null(), - num_exprs: 0, - reserved: null_mut(), - }; - - let mut ffi_array = arrow_array::ffi::FFI_ArrowArray::empty(); - let mut err = SedonaCError::default(); - - let code = unsafe { - get_property( - plan, - property_cstr.as_ptr(), - &mut ffi_args, - &mut ffi_array, - &mut err, - ) - }; - - if code != ERRNO_OK { - return sedona_internal_err!("Failed to get '{}': {}", property, err); - } - - // Get the property schema to know how to interpret the array - let data_type = get_plan_property_data_type(plan, property)?; - - let bytes = parse_ffi_array_to_bytes(ffi_array, &data_type)?; - String::from_utf8(bytes) - .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) + call_get_string_property_impl( + property, + "SedonaCExecutionPlan", + |prop, args, out, err| unsafe { get_property(plan, prop, args, out, err) }, + || get_plan_property_data_type(plan, property), + ) } /// Get the schema for a property from a [SedonaCTableProvider]. @@ -282,73 +331,10 @@ fn get_table_provider_property_data_type( property: &str, ) -> Result { let Some(get_property_schema) = provider.get_property_schema else { - // Default to Binary if get_property_schema is not implemented return Ok(DataType::Binary); }; - let property_cstr = CString::new(property) - .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; - - let mut ffi_schema = FFI_ArrowSchema::empty(); - let mut err = SedonaCError::default(); - - let code = - unsafe { get_property_schema(provider, property_cstr.as_ptr(), &mut ffi_schema, &mut err) }; - - if code != ERRNO_OK { - return sedona_internal_err!("Failed to get property schema for '{}': {}", property, err); - } - - // Try to convert the FFI schema to a Field - let field = Field::try_from(&ffi_schema).map_err(|e| { - sedona_internal_datafusion_err!("Failed to parse property schema for '{}': {}", property, e) - })?; - Ok(field.data_type().clone()) -} - -/// Get a string property from a [SedonaCTableProvider]. -pub fn get_table_provider_string_property( - provider: &SedonaCTableProvider, - property: &str, -) -> Result { - let Some(get_property) = provider.get_property else { - return sedona_internal_err!("SedonaCTableProvider does not have get_property"); - }; - - let property_cstr = CString::new(property) - .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?; - - let mut ffi_args = SedonaCExecutionPlanArgs { - args: std::ptr::null(), - args_len: 0, - exec_plans: std::ptr::null(), - num_exec_plans: 0, - exprs: std::ptr::null(), - num_exprs: 0, - reserved: null_mut(), - }; - - let mut ffi_array = arrow_array::ffi::FFI_ArrowArray::empty(); - let mut err = SedonaCError::default(); - - let code = unsafe { - get_property( - provider, - property_cstr.as_ptr(), - &mut ffi_args, - &mut ffi_array, - &mut err, - ) - }; - - if code != ERRNO_OK { - return sedona_internal_err!("Failed to get '{}': {}", property, err); - } - - // Get the property schema to know how to interpret the array - let data_type = get_table_provider_property_data_type(provider, property)?; - - let bytes = parse_ffi_array_to_bytes(ffi_array, &data_type)?; - String::from_utf8(bytes) - .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e)) + call_get_property_schema_impl(property, |prop, schema, err| unsafe { + get_property_schema(provider, prop, schema, err) + }) } From a47cca1a2931d2ee7f5032d371be7a3601afe657 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 16:41:33 -0500 Subject: [PATCH 37/55] remove unused dependency --- Cargo.lock | 1 - c/sedona-extension/Cargo.toml | 1 - 2 files changed, 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d41d70570f..d19387ccd7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5771,7 +5771,6 @@ dependencies = [ "libc", "sedona-common", "sedona-expr", - "sedona-geometry", "sedona-schema", "sedona-testing", "serde", diff --git a/c/sedona-extension/Cargo.toml b/c/sedona-extension/Cargo.toml index f0ca89a2f2..bb6a7acde8 100644 --- a/c/sedona-extension/Cargo.toml +++ b/c/sedona-extension/Cargo.toml @@ -41,7 +41,6 @@ futures = { workspace = true } libc = "0.2.178" sedona-common = { workspace = true } sedona-expr = { workspace = true } -sedona-geometry = { workspace = true } sedona-schema = { workspace = true } sedona-testing = { path = "../../rust/sedona-testing" } serde = { workspace = true } From 6702794b9d47749af9c4b57651cb74179bcc5ae7 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 16:42:53 -0500 Subject: [PATCH 38/55] fix outdated --- c/sedona-extension/src/execution_plan.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 3259a22fa1..ea9152ad0b 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -315,17 +315,17 @@ impl Debug for ImportedSedonaCExec { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { // Try to get debug string from the FFI plan if let Ok(debug_str) = self.get_debug_string() { - f.debug_struct("ImportedTableProviderExec") + f.debug_struct("ImportedSedonaCExec") .field("inner", &debug_str) .finish() } else { - f.debug_struct("ImportedTableProviderExec").finish() + f.debug_struct("ImportedSedonaCExec").finish() } } } impl ImportedSedonaCExec { - /// Create a new ImportedTableProviderExec from a SedonaCExecutionPlan. + /// Create a new ImportedSedonaCExec from a SedonaCExecutionPlan. /// /// This will query the plan for its schema and properties. pub fn try_new(inner: SedonaCExecutionPlan) -> Result { @@ -480,7 +480,7 @@ impl ExecutionPlan for ImportedSedonaCExec { children: Vec>, ) -> Result> { if !children.is_empty() { - return exec_err!("ImportedTableProviderExec does not support children"); + return exec_err!("ImportedSedonaCExec does not support children"); } Ok(self) } From 1281fe3c27d2ccaf9d4f69a55f0b8cb540b9a9c9 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 16:46:01 -0500 Subject: [PATCH 39/55] shared cancel error constructor --- c/sedona-extension/src/streaming.rs | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index 04a37fc96d..323a132049 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -34,6 +34,17 @@ use futures::StreamExt; /// cancellation mechanisms. pub type CancelChecker = Box bool + Send + Sync>; +/// Create a cancellation error for FFI stream operations. +/// +/// This is the single source of truth for cancellation errors, ensuring +/// consistent error format across all cancellation sites. +fn cancellation_arrow_error() -> ArrowError { + ArrowError::ExternalError(Box::new(std::io::Error::new( + std::io::ErrorKind::Interrupted, + "Operation cancelled", + ))) +} + /// Result type for batch fetching from the worker thread. type BatchResult = Option>; @@ -196,12 +207,7 @@ impl StreamingRecordBatchReader { if let Some(ref checker) = self.cancel_checker { if checker() { self.cancelled = true; - return Some(Err(ArrowError::ExternalError(Box::new( - std::io::Error::new( - std::io::ErrorKind::Interrupted, - "Operation cancelled", - ), - )))); + return Some(Err(cancellation_arrow_error())); } } // Continue waiting @@ -241,9 +247,7 @@ impl Iterator for StreamingRecordBatchReader { if let Some(ref checker) = self.cancel_checker { if checker() { self.cancelled = true; - return Some(Err(ArrowError::ExternalError(Box::new( - std::io::Error::new(std::io::ErrorKind::Interrupted, "Operation cancelled"), - )))); + return Some(Err(cancellation_arrow_error())); } } From 53b377a42274a6d3abff0c67a371be4fb6c59c65 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 16:49:04 -0500 Subject: [PATCH 40/55] deduplicate table provider tests --- c/sedona-extension/src/table_provider.rs | 292 +++++++---------------- 1 file changed, 90 insertions(+), 202 deletions(-) diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 5935fa5180..2eb827b2d8 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -471,8 +471,14 @@ mod tests { Ok(ctx) } - #[tokio::test] - async fn test_roundtrip_simple_select() -> Result<()> { + /// Helper to test FFI table provider roundtrip with a SQL query. + /// + /// This handles the common pattern of: + /// 1. Creating a test context with test_data table + /// 2. Exporting the table provider through FFI + /// 3. Importing it and registering in a new context + /// 4. Running the SQL query and asserting results + async fn test_roundtrip_query(sql: &str, expected: &[&str]) -> Result<()> { let ctx = create_test_context().await?; // Get the table provider from the context @@ -492,228 +498,110 @@ mod tests { ctx2.register_table("imported_data", Arc::new(imported))?; // Query and verify - let result = ctx2 - .sql("SELECT id, value_a FROM imported_data ORDER BY id LIMIT 5") - .await? - .collect() - .await?; - - let expected = [ - "+----+---------+", - "| id | value_a |", - "+----+---------+", - "| 1 | 100 |", - "| 2 | 200 |", - "| 3 | 300 |", - "| 4 | 400 |", - "| 5 | 500 |", - "+----+---------+", - ]; - + let result = ctx2.sql(sql).await?.collect().await?; assert_batches_eq!(expected, &result); Ok(()) } #[tokio::test] - async fn test_roundtrip_projection() -> Result<()> { - let ctx = create_test_context().await?; - - // Get the table provider from the context - let table = ctx.table_provider("test_data").await?; - - // Export the table provider - let runtime = tokio::runtime::Handle::current(); - let session = Arc::new(ctx.state()); - let exported = ExportedTableProvider::new(table, session, runtime); - let ffi_provider: SedonaCTableProvider = exported.into(); - - // Import the table provider - let imported = ImportedTableProvider::try_new(ffi_provider)?; - - // Create a new context and register the imported table - let ctx2 = SessionContext::new(); - ctx2.register_table("imported_data", Arc::new(imported))?; - - // Test projection with only specific columns - let result = ctx2 - .sql("SELECT value_b, value_d FROM imported_data ORDER BY value_b LIMIT 3") - .await? - .collect() - .await?; - - let expected = [ - "+---------+---------+", - "| value_b | value_d |", - "+---------+---------+", - "| 1.5 | 1000 |", - "| 3.0 | 2000 |", - "| 4.5 | 3000 |", - "+---------+---------+", - ]; + async fn test_roundtrip_simple_select() -> Result<()> { + test_roundtrip_query( + "SELECT id, value_a FROM imported_data ORDER BY id LIMIT 5", + &[ + "+----+---------+", + "| id | value_a |", + "+----+---------+", + "| 1 | 100 |", + "| 2 | 200 |", + "| 3 | 300 |", + "| 4 | 400 |", + "| 5 | 500 |", + "+----+---------+", + ], + ) + .await + } - assert_batches_eq!(expected, &result); - Ok(()) + #[tokio::test] + async fn test_roundtrip_projection() -> Result<()> { + test_roundtrip_query( + "SELECT value_b, value_d FROM imported_data ORDER BY value_b LIMIT 3", + &[ + "+---------+---------+", + "| value_b | value_d |", + "+---------+---------+", + "| 1.5 | 1000 |", + "| 3.0 | 2000 |", + "| 4.5 | 3000 |", + "+---------+---------+", + ], + ) + .await } #[tokio::test] async fn test_roundtrip_filter() -> Result<()> { - let ctx = create_test_context().await?; - - // Get the table provider from the context - let table = ctx.table_provider("test_data").await?; - - // Export the table provider - let runtime = tokio::runtime::Handle::current(); - let session = Arc::new(ctx.state()); - let exported = ExportedTableProvider::new(table, session, runtime); - let ffi_provider: SedonaCTableProvider = exported.into(); - - // Import the table provider - let imported = ImportedTableProvider::try_new(ffi_provider)?; - - // Create a new context and register the imported table - let ctx2 = SessionContext::new(); - ctx2.register_table("imported_data", Arc::new(imported))?; - - // Test filter (filter is applied on the DataFusion side, not pushed to FFI yet) - let result = ctx2 - .sql("SELECT id, value_a FROM imported_data WHERE id > 20 ORDER BY id LIMIT 5") - .await? - .collect() - .await?; - - let expected = [ - "+----+---------+", - "| id | value_a |", - "+----+---------+", - "| 21 | 2100 |", - "| 22 | 2200 |", - "| 23 | 2300 |", - "| 24 | 2400 |", - "| 25 | 2500 |", - "+----+---------+", - ]; - - assert_batches_eq!(expected, &result); - Ok(()) + test_roundtrip_query( + "SELECT id, value_a FROM imported_data WHERE id > 20 ORDER BY id LIMIT 5", + &[ + "+----+---------+", + "| id | value_a |", + "+----+---------+", + "| 21 | 2100 |", + "| 22 | 2200 |", + "| 23 | 2300 |", + "| 24 | 2400 |", + "| 25 | 2500 |", + "+----+---------+", + ], + ) + .await } #[tokio::test] async fn test_roundtrip_sort() -> Result<()> { - let ctx = create_test_context().await?; - - // Get the table provider from the context - let table = ctx.table_provider("test_data").await?; - - // Export the table provider - let runtime = tokio::runtime::Handle::current(); - let session = Arc::new(ctx.state()); - let exported = ExportedTableProvider::new(table, session, runtime); - let ffi_provider: SedonaCTableProvider = exported.into(); - - // Import the table provider - let imported = ImportedTableProvider::try_new(ffi_provider)?; - - // Create a new context and register the imported table - let ctx2 = SessionContext::new(); - ctx2.register_table("imported_data", Arc::new(imported))?; - - // Test sorting in descending order - let result = ctx2 - .sql("SELECT id, value_c FROM imported_data ORDER BY id DESC LIMIT 5") - .await? - .collect() - .await?; - - let expected = [ - "+----+---------+", - "| id | value_c |", - "+----+---------+", - "| 45 | 90 |", - "| 44 | 88 |", - "| 43 | 86 |", - "| 42 | 84 |", - "| 41 | 82 |", - "+----+---------+", - ]; - - assert_batches_eq!(expected, &result); - Ok(()) + test_roundtrip_query( + "SELECT id, value_c FROM imported_data ORDER BY id DESC LIMIT 5", + &[ + "+----+---------+", + "| id | value_c |", + "+----+---------+", + "| 45 | 90 |", + "| 44 | 88 |", + "| 43 | 86 |", + "| 42 | 84 |", + "| 41 | 82 |", + "+----+---------+", + ], + ) + .await } #[tokio::test] async fn test_roundtrip_limit() -> Result<()> { - let ctx = create_test_context().await?; - - // Get the table provider from the context - let table = ctx.table_provider("test_data").await?; - - // Export the table provider - let runtime = tokio::runtime::Handle::current(); - let session = Arc::new(ctx.state()); - let exported = ExportedTableProvider::new(table, session, runtime); - let ffi_provider: SedonaCTableProvider = exported.into(); - - // Import the table provider - let imported = ImportedTableProvider::try_new(ffi_provider)?; - - // Create a new context and register the imported table - let ctx2 = SessionContext::new(); - ctx2.register_table("imported_data", Arc::new(imported))?; - - // Test limit - let result = ctx2 - .sql("SELECT id FROM imported_data ORDER BY id LIMIT 3") - .await? - .collect() - .await?; - - let expected = [ - "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 3 |", "+----+", - ]; - - assert_batches_eq!(expected, &result); - Ok(()) + test_roundtrip_query( + "SELECT id FROM imported_data ORDER BY id LIMIT 3", + &[ + "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 3 |", "+----+", + ], + ) + .await } #[tokio::test] async fn test_roundtrip_all_columns() -> Result<()> { - let ctx = create_test_context().await?; - - // Get the table provider from the context - let table = ctx.table_provider("test_data").await?; - - // Export the table provider - let runtime = tokio::runtime::Handle::current(); - let session = Arc::new(ctx.state()); - let exported = ExportedTableProvider::new(table, session, runtime); - let ffi_provider: SedonaCTableProvider = exported.into(); - - // Import the table provider - let imported = ImportedTableProvider::try_new(ffi_provider)?; - - // Create a new context and register the imported table - let ctx2 = SessionContext::new(); - ctx2.register_table("imported_data", Arc::new(imported))?; - - // Test selecting all columns - let result = ctx2 - .sql("SELECT * FROM imported_data ORDER BY id LIMIT 2") - .await? - .collect() - .await?; - - let expected = [ - "+----+---------+---------+---------+---------+", - "| id | value_a | value_b | value_c | value_d |", - "+----+---------+---------+---------+---------+", - "| 1 | 100 | 1.5 | 2 | 1000 |", - "| 2 | 200 | 3.0 | 4 | 2000 |", - "+----+---------+---------+---------+---------+", - ]; - - assert_batches_eq!(expected, &result); - Ok(()) + test_roundtrip_query( + "SELECT * FROM imported_data ORDER BY id LIMIT 2", + &[ + "+----+---------+---------+---------+---------+", + "| id | value_a | value_b | value_c | value_d |", + "+----+---------+---------+---------+---------+", + "| 1 | 100 | 1.5 | 2 | 1000 |", + "| 2 | 200 | 3.0 | 4 | 2000 |", + "+----+---------+---------+---------+---------+", + ], + ) + .await } /// A dummy TableProvider with configurable table_type for testing FFI roundtrip. From c5f006afb7c05207a740bb868ea1aeb8b09e5fba Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 16:59:54 -0500 Subject: [PATCH 41/55] allow get_schema to return an error --- c/sedona-extension/src/execution_plan.rs | 22 ++++++++++++++++++---- c/sedona-extension/src/extension.rs | 18 ++++++++++++++---- c/sedona-extension/src/sedona_extension.h | 10 ++++++++-- c/sedona-extension/src/table_provider.rs | 22 ++++++++++++++++++---- 4 files changed, 58 insertions(+), 14 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index ea9152ad0b..330dd1f67e 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -171,7 +171,8 @@ impl From for SedonaCExecutionPlan { unsafe extern "C" fn c_exec_plan_get_schema( self_: *const SedonaCExecutionPlan, out: *mut FFI_ArrowSchema, -) { + err: *mut SedonaCError, +) -> c_int { debug_assert!(!self_.is_null(), "self pointer is null"); debug_assert!(!out.is_null(), "out pointer is null"); let self_ref = &*self_; @@ -179,8 +180,17 @@ unsafe extern "C" fn c_exec_plan_get_schema( let plan = &*(self_ref.private_data as *const ExportedExecutionPlan); let schema = plan.schema(); - if let Ok(ffi_schema) = FFI_ArrowSchema::try_from(schema.as_ref()) { - std::ptr::write(out, ffi_schema); + match FFI_ArrowSchema::try_from(schema.as_ref()) { + Ok(ffi_schema) => { + std::ptr::write(out, ffi_schema); + ERRNO_OK + } + Err(e) => { + if !err.is_null() { + *err = SedonaCError::new(&format!("Failed to convert schema to FFI: {}", e)); + } + libc::EINVAL + } } } @@ -340,7 +350,11 @@ impl ImportedSedonaCExec { }; let mut ffi_schema = FFI_ArrowSchema::empty(); - unsafe { get_schema(&inner, &mut ffi_schema) }; + let mut err = SedonaCError::default(); + let code = unsafe { get_schema(&inner, &mut ffi_schema, &mut err) }; + if code != ERRNO_OK { + return sedona_internal_err!("Failed to get schema: {}", err); + } let schema = Arc::new(Schema::try_from(&ffi_schema)?); // Get plan properties diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs index 4e01582123..07b0bf55b0 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -264,8 +264,13 @@ pub struct SedonaCExecutionPlanArgs { #[derive(Default)] #[repr(C)] pub struct SedonaCExecutionPlan { - pub get_schema: - Option, + pub get_schema: Option< + unsafe extern "C" fn( + self_: *const SedonaCExecutionPlan, + out: *mut FFI_ArrowSchema, + err: *mut SedonaCError, + ) -> c_int, + >, pub get_property_schema: Option< unsafe extern "C" fn( @@ -342,8 +347,13 @@ impl Drop for SedonaCExecutionPlan { #[repr(C)] pub struct SedonaCTableProvider { /// Get the schema of this table provider - pub get_schema: - Option, + pub get_schema: Option< + unsafe extern "C" fn( + self_: *const SedonaCTableProvider, + out: *mut FFI_ArrowSchema, + err: *mut SedonaCError, + ) -> c_int, + >, /// Get the schema of a property pub get_property_schema: Option< diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 1d2453896a..2f1682dcbf 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -277,7 +277,10 @@ struct SedonaCExecutionPlanArgs { /// Instances with a NULL release callback are not valid and must not be used. struct SedonaCExecutionPlan { /// \brief Get the schema associated with the output of this plan - void (*get_schema)(const struct SedonaCExecutionPlan* self, struct ArrowSchema* out); + /// + /// Returns 0 on success, or an errno value on failure. + int (*get_schema)(const struct SedonaCExecutionPlan* self, struct ArrowSchema* out, + struct SedonaCError* err); /// \brief Get the data type of a property int (*get_property_schema)(const struct SedonaCExecutionPlan* self, @@ -335,7 +338,10 @@ struct SedonaCExecutionPlan { /// Instances with a NULL release callback are not valid and must not be used. struct SedonaCTableProvider { /// \brief Get the schema of this table provider - void (*get_schema)(const struct SedonaCTableProvider* self, struct ArrowSchema* out); + /// + /// Returns 0 on success, or an errno value on failure. + int (*get_schema)(const struct SedonaCTableProvider* self, struct ArrowSchema* out, + struct SedonaCError* err); /// \brief Get the data type of a property int (*get_property_schema)(const struct SedonaCTableProvider* self, diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 2eb827b2d8..31a8296955 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -130,7 +130,8 @@ impl From for SedonaCTableProvider { unsafe extern "C" fn c_table_provider_get_schema( self_: *const SedonaCTableProvider, out: *mut FFI_ArrowSchema, -) { + err: *mut SedonaCError, +) -> c_int { debug_assert!(!self_.is_null(), "self pointer is null"); debug_assert!(!out.is_null(), "out pointer is null"); let self_ref = &*self_; @@ -138,8 +139,17 @@ unsafe extern "C" fn c_table_provider_get_schema( let provider = &*(self_ref.private_data as *const ExportedTableProvider); let schema = provider.inner.schema(); - if let Ok(ffi_schema) = FFI_ArrowSchema::try_from(schema.as_ref()) { - std::ptr::write(out, ffi_schema); + match FFI_ArrowSchema::try_from(schema.as_ref()) { + Ok(ffi_schema) => { + std::ptr::write(out, ffi_schema); + ERRNO_OK + } + Err(e) => { + if !err.is_null() { + *err = SedonaCError::new(&format!("Failed to convert schema to FFI: {}", e)); + } + libc::EINVAL + } } } @@ -296,7 +306,11 @@ impl ImportedTableProvider { }; let mut ffi_schema = FFI_ArrowSchema::empty(); - unsafe { get_schema(&inner, &mut ffi_schema) }; + let mut err = SedonaCError::default(); + let code = unsafe { get_schema(&inner, &mut ffi_schema, &mut err) }; + if code != ERRNO_OK { + return sedona_internal_err!("Failed to get schema: {}", err); + } let schema = Arc::new(Schema::try_from(&ffi_schema)?); // Get table type via get_property From 80b87fe4e79aaa42962dc1ecdade31220eb8557c Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 17:01:19 -0500 Subject: [PATCH 42/55] fix the expr signature --- c/sedona-extension/src/extension.rs | 1 + c/sedona-extension/src/sedona_extension.h | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs index 07b0bf55b0..d6098e63a1 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -207,6 +207,7 @@ pub struct SedonaCExpr { unsafe extern "C" fn( self_: *const SedonaCExpr, property: *const c_char, + out: *mut FFI_ArrowSchema, err: *mut SedonaCError, ) -> c_int, >, diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 2f1682dcbf..6760028717 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -226,7 +226,7 @@ struct SedonaCError { struct SedonaCExpr { /// \brief Get the data type of a property int (*get_property_schema)(const struct SedonaCExpr* self, const char* property, - struct SedonaCError* err); + struct ArrowSchema* out, struct SedonaCError* err); /// \brief Extract a serializable property from this expression /// From 1c42831dc63a9ade0aa7720e3a078f7012699e82 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 17:08:04 -0500 Subject: [PATCH 43/55] error setter macro --- c/sedona-extension/src/execution_plan.rs | 21 ++++--------- c/sedona-extension/src/extension.rs | 37 +++++++++++++++++++++++ c/sedona-extension/src/sedona_extension.h | 6 ++++ c/sedona-extension/src/table_provider.rs | 21 ++++--------- 4 files changed, 55 insertions(+), 30 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 330dd1f67e..bf3322345c 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -37,6 +37,7 @@ use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; use serde::{Deserialize, Serialize}; use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs}; +use crate::set_ffi_error; use crate::streaming::{ffi_stream_to_sendable, StreamingRecordBatchReader}; use crate::utils::{cstr_from_ptr_or_empty, get_plan_property, get_plan_string_property, ERRNO_OK}; @@ -186,9 +187,7 @@ unsafe extern "C" fn c_exec_plan_get_schema( ERRNO_OK } Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&format!("Failed to convert schema to FFI: {}", e)); - } + set_ffi_error!(err, "Failed to convert schema to FFI: {}", e); libc::EINVAL } } @@ -210,9 +209,7 @@ unsafe extern "C" fn c_exec_plan_get_property_schema( ERRNO_OK } Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&format!("Failed to convert field to FFI schema: {}", e)); - } + set_ffi_error!(err, "Failed to convert field to FFI schema: {}", e); libc::EINVAL } } @@ -244,9 +241,7 @@ unsafe extern "C" fn c_exec_plan_get_property( ERRNO_OK } Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&e.to_string()); - } + set_ffi_error!(err, "{}", e); libc::EINVAL } } @@ -275,9 +270,7 @@ unsafe extern "C" fn c_exec_plan_execute( let execute_args: ExecuteArgs = match serde_json::from_slice(args_slice) { Ok(a) => a, Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&format!("Failed to parse execute args: {}", e)); - } + set_ffi_error!(err, "Failed to parse execute args: {}", e); return libc::EINVAL; } }; @@ -291,9 +284,7 @@ unsafe extern "C" fn c_exec_plan_execute( ERRNO_OK } Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&e.to_string()); - } + set_ffi_error!(err, "{}", e); libc::EINVAL } } diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs index d6098e63a1..287023f324 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -199,6 +199,43 @@ pub const UNKNOWN_SEDONA_C_ERROR: SedonaCError = SedonaCError { release: Some(sedona_c_noop_release), }; +/// Write an error to an FFI error pointer safely. +/// +/// This uses `ptr::write` instead of assignment to avoid running Drop on +/// potentially uninitialized memory. C consumers may pass uninitialized +/// SedonaCError structs, so we cannot assume the previous value is valid. +/// +/// # Safety +/// +/// The caller must ensure `err` points to valid memory for a SedonaCError, +/// but the memory does not need to be initialized. +#[inline] +pub unsafe fn write_ffi_error(err: *mut SedonaCError, message: &str) { + if !err.is_null() { + std::ptr::write(err, SedonaCError::new(message)); + } +} + +/// Macro to write an error to an FFI error pointer with formatting. +/// +/// This is a convenience wrapper around `write_ffi_error` that supports +/// format strings like `format!()`. +/// +/// # Example +/// +/// ```ignore +/// set_ffi_error!(err, "Failed to parse: {}", e); +/// ``` +#[macro_export] +macro_rules! set_ffi_error { + ($err:expr, $msg:expr) => { + $crate::extension::write_ffi_error($err, $msg) + }; + ($err:expr, $fmt:expr, $($arg:tt)*) => { + $crate::extension::write_ffi_error($err, &format!($fmt, $($arg)*)) + }; +} + /// Raw FFI representation of the SedonaCExpr #[derive(Default)] #[repr(C)] diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 6760028717..843e09fc21 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -204,6 +204,12 @@ struct SedonaCScalarKernel { void* private_data; }; +/// \brief Error information returned by FFI callbacks +/// +/// This structure is written to by FFI callbacks on error. Callers may pass +/// uninitialized memory; the implementation uses ptr::write to avoid reading +/// the previous contents. However, for best practice, callers should +/// zero-initialize: `struct SedonaCError err = {0};` struct SedonaCError { /// \brief A UTF-8 encoded error message /// diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 31a8296955..a1f785708a 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -37,6 +37,7 @@ use crate::execution_plan::{ExportedExecutionPlan, ImportedSedonaCExec}; use crate::extension::{ SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs, SedonaCTableProvider, }; +use crate::set_ffi_error; use crate::utils::{cstr_from_ptr_or_empty, get_table_provider_string_property, ERRNO_OK}; /// A TableProvider wrapper that can be exported across FFI. @@ -145,9 +146,7 @@ unsafe extern "C" fn c_table_provider_get_schema( ERRNO_OK } Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&format!("Failed to convert schema to FFI: {}", e)); - } + set_ffi_error!(err, "Failed to convert schema to FFI: {}", e); libc::EINVAL } } @@ -169,9 +168,7 @@ unsafe extern "C" fn c_table_provider_get_property_schema( ERRNO_OK } Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&format!("Failed to convert field to FFI schema: {}", e)); - } + set_ffi_error!(err, "Failed to convert field to FFI schema: {}", e); libc::EINVAL } } @@ -203,9 +200,7 @@ unsafe extern "C" fn c_table_provider_get_property( ERRNO_OK } Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&e.to_string()); - } + set_ffi_error!(err, "{}", e); libc::EINVAL } } @@ -241,9 +236,7 @@ unsafe extern "C" fn c_table_provider_scan( match serde_json::from_slice(args_slice) { Ok(a) => a, Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&format!("Failed to parse scan args: {}", e)); - } + set_ffi_error!(err, "Failed to parse scan args: {}", e); return libc::EINVAL; } } @@ -258,9 +251,7 @@ unsafe extern "C" fn c_table_provider_scan( ERRNO_OK } Err(e) => { - if !err.is_null() { - *err = SedonaCError::new(&e.to_string()); - } + set_ffi_error!(err, "{}", e); libc::EINVAL } } From 4c65ff5725ecc1987552e9bb147f0e308677ad0f Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 17:17:29 -0500 Subject: [PATCH 44/55] fix session state export --- python/sedonadb/python/sedonadb/dataframe.py | 2 +- python/sedonadb/src/dataframe.rs | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/python/sedonadb/python/sedonadb/dataframe.py b/python/sedonadb/python/sedonadb/dataframe.py index 34730fb50d..308ba2dca3 100644 --- a/python/sedonadb/python/sedonadb/dataframe.py +++ b/python/sedonadb/python/sedonadb/dataframe.py @@ -1291,7 +1291,7 @@ def to_memtable(self) -> "DataFrame": return DataFrame(self._ctx, self._impl.to_memtable(self._ctx._impl)) def __sedonadb_table_provider__(self): - return self._impl.__sedonadb_table_provider__() + return self._impl.__sedonadb_table_provider__(self._ctx._impl) def to_arrow_table(self, schema: Any = None) -> "pa.Table": """Execute and collect results as a PyArrow Table diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs index c2a77f95d2..bb5ee51d76 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -24,7 +24,7 @@ use arrow_schema::{Schema, SchemaRef}; use datafusion::catalog::MemTable; use datafusion::config::ConfigField; use datafusion::logical_expr::SortExpr; -use datafusion::prelude::{DataFrame, SessionContext}; +use datafusion::prelude::DataFrame; use datafusion_common::{Column, DataFusionError, ParamValues}; use datafusion_expr::{ExplainFormat, ExplainOption, Expr, JoinType, LogicalPlanBuilder}; use futures::lock::Mutex; @@ -667,12 +667,12 @@ impl InternalDataFrame { fn __sedonadb_table_provider__<'py>( &self, py: Python<'py>, + ctx: &InternalContext, ) -> Result, PySedonaError> { let provider = self.inner.clone().into_view(); - // Create a session context for FFI - the consuming side will use its own - // session for actual execution, this is just needed for the FFI interface. - let ctx = SessionContext::new(); - let session = Arc::new(ctx.state()); + // Use the actual session state so that object stores, UDFs, and other + // registrations are available when the consumer scans the provider. + let session = Arc::new(ctx.inner.ctx.state()); let exported = sedona_extension::table_provider::ExportedTableProvider::new( provider, session, From 3b1fcdfeeeac2e8fe6f14965912b1141e4ee84e8 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 17:31:36 -0500 Subject: [PATCH 45/55] r impl --- r/sedonadb/R/000-wrappers.R | 5 ++-- r/sedonadb/R/dataframe.R | 15 ++++++++++++ r/sedonadb/src/init.c | 6 ++--- r/sedonadb/src/rust/api.h | 2 +- r/sedonadb/src/rust/src/dataframe.rs | 28 ++++++++++------------ r/sedonadb/src/rust/src/ffi.rs | 25 +++++++++++++------ r/sedonadb/tests/testthat/test-dataframe.R | 2 +- 7 files changed, 54 insertions(+), 29 deletions(-) diff --git a/r/sedonadb/R/000-wrappers.R b/r/sedonadb/R/000-wrappers.R index 95c9058dcd..88d6f513c0 100644 --- a/r/sedonadb/R/000-wrappers.R +++ b/r/sedonadb/R/000-wrappers.R @@ -376,8 +376,9 @@ class(`InternalContext`) <- c( } `InternalDataFrame_to_provider` <- function(self) { - function() { - .Call(savvy_InternalDataFrame_to_provider__impl, `self`) + function(`ctx`) { + `ctx` <- .savvy_extract_ptr(`ctx`, "sedonadb::InternalContext") + .Call(savvy_InternalDataFrame_to_provider__impl, `self`, `ctx`) } } diff --git a/r/sedonadb/R/dataframe.R b/r/sedonadb/R/dataframe.R index 66b1534a17..f9ad5b1355 100644 --- a/r/sedonadb/R/dataframe.R +++ b/r/sedonadb/R/dataframe.R @@ -101,6 +101,21 @@ as_sedonadb_dataframe.datafusion_table_provider <- function( new_sedonadb_dataframe(ctx, df) } +#' @export +as_sedonadb_dataframe.sedonadb_table_provider <- function( + x, + ..., + schema = NULL, + ctx = NULL +) { + if (is.null(ctx)) { + ctx <- ctx() + } + + df <- ctx$data_frame_from_table_provider(x) + new_sedonadb_dataframe(ctx, df) +} + #' Count rows in a DataFrame #' #' @param .data A sedonadb_dataframe or an object that can be coerced to one. diff --git a/r/sedonadb/src/init.c b/r/sedonadb/src/init.c index 400782db70..3eee62091f 100644 --- a/r/sedonadb/src/init.c +++ b/r/sedonadb/src/init.c @@ -251,8 +251,8 @@ SEXP savvy_InternalDataFrame_to_parquet__impl( return handle_result(res); } -SEXP savvy_InternalDataFrame_to_provider__impl(SEXP self__) { - SEXP res = savvy_InternalDataFrame_to_provider__ffi(self__); +SEXP savvy_InternalDataFrame_to_provider__impl(SEXP self__, SEXP c_arg__ctx) { + SEXP res = savvy_InternalDataFrame_to_provider__ffi(self__, c_arg__ctx); return handle_result(res); } @@ -438,7 +438,7 @@ static const R_CallMethodDef CallEntries[] = { {"savvy_InternalDataFrame_to_parquet__impl", (DL_FUNC)&savvy_InternalDataFrame_to_parquet__impl, 8}, {"savvy_InternalDataFrame_to_provider__impl", - (DL_FUNC)&savvy_InternalDataFrame_to_provider__impl, 1}, + (DL_FUNC)&savvy_InternalDataFrame_to_provider__impl, 2}, {"savvy_InternalDataFrame_to_view__impl", (DL_FUNC)&savvy_InternalDataFrame_to_view__impl, 4}, {"savvy_InternalDataFrame_with_params__impl", diff --git a/r/sedonadb/src/rust/api.h b/r/sedonadb/src/rust/api.h index 18cdeb12f2..fb3654abb8 100644 --- a/r/sedonadb/src/rust/api.h +++ b/r/sedonadb/src/rust/api.h @@ -74,7 +74,7 @@ SEXP savvy_InternalDataFrame_to_parquet__ffi( SEXP self__, SEXP c_arg__ctx, SEXP c_arg__path, SEXP c_arg__option_keys, SEXP c_arg__option_values, SEXP c_arg__partition_by, SEXP c_arg__sort_by, SEXP c_arg__single_file_output); -SEXP savvy_InternalDataFrame_to_provider__ffi(SEXP self__); +SEXP savvy_InternalDataFrame_to_provider__ffi(SEXP self__, SEXP c_arg__ctx); SEXP savvy_InternalDataFrame_to_view__ffi(SEXP self__, SEXP c_arg__ctx, SEXP c_arg__table_ref, SEXP c_arg__overwrite); diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index cc150b3f9b..a4093a6214 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -20,17 +20,16 @@ use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{RecordBatchIterator, RecordBatchReader}; use datafusion::catalog::MemTable; use datafusion::config::ConfigField; -use datafusion::prelude::{DataFrame, SessionContext}; +use datafusion::prelude::DataFrame; use datafusion_common::Column; -use datafusion_execution::TaskContextProvider; use datafusion_expr::utils::conjunction; use datafusion_expr::JoinType; use datafusion_expr::{select_expr::SelectExpr, Expr, SortExpr}; -use datafusion_ffi::table_provider::FFI_TableProvider; use savvy::{savvy, savvy_err, sexp, IntoExtPtrSexp, Result}; use sedona::context::{SedonaDataFrame, SedonaWriteOptions}; use sedona::show::{DisplayMode, DisplayTableOptions}; use sedona_extension::streaming::StreamingRecordBatchReader; +use sedona_extension::table_provider::ExportedTableProvider; use sedona_geoparquet::options::TableGeoParquetOptions; use sedona_schema::schema::SedonaSchema; use std::{iter::zip, ptr::swap_nonoverlapping, sync::Arc}; @@ -38,7 +37,7 @@ use tokio::runtime::Runtime; use crate::context::InternalContext; use crate::expression::SedonaDBExprFactory; -use crate::ffi::{import_schema, FFITableProviderR}; +use crate::ffi::{import_schema, SedonaCTableProviderR}; use crate::runtime::wait_for_future_captured_r; #[savvy] @@ -126,22 +125,21 @@ impl InternalDataFrame { Ok(()) } - fn to_provider(&self) -> Result { + fn to_provider(&self, ctx: &InternalContext) -> Result { let provider = self.inner.clone().into_view(); - // Literal true is because the TableProvider that wraps this DataFrame - // can support filters being pushed down. - let ctx = Arc::new(SessionContext::new()) as Arc; - let ffi_provider = FFI_TableProvider::new( + // Use the actual session state so that object stores, UDFs, and other + // registrations are available when the consumer scans the provider. + let session = Arc::new(ctx.inner.ctx.state()); + let exported = ExportedTableProvider::new( provider, - true, - Some(self.runtime.handle().clone()), - &ctx, - None, + session, + self.runtime.handle().clone(), ); + let ffi_provider: sedona_extension::extension::SedonaCTableProvider = exported.into(); - let mut ffi_xptr = FFITableProviderR(ffi_provider).into_external_pointer(); + let mut ffi_xptr = SedonaCTableProviderR(ffi_provider).into_external_pointer(); unsafe { savvy_ffi::Rf_protect(ffi_xptr.0) }; - ffi_xptr.set_class(vec!["datafusion_table_provider"])?; + ffi_xptr.set_class(vec!["sedonadb_table_provider"])?; unsafe { savvy_ffi::Rf_unprotect(1) }; Ok(ffi_xptr) diff --git a/r/sedonadb/src/rust/src/ffi.rs b/r/sedonadb/src/rust/src/ffi.rs index a81d3a60f8..ddb1304c17 100644 --- a/r/sedonadb/src/rust/src/ffi.rs +++ b/r/sedonadb/src/rust/src/ffi.rs @@ -25,8 +25,10 @@ use arrow_array::{ use arrow_schema::{Field, Schema}; use datafusion::catalog::TableProvider; use datafusion_expr::{ScalarUDF, ScalarUDFImpl}; -use datafusion_ffi::{table_provider::FFI_TableProvider, udf::FFI_ScalarUDF}; +use datafusion_ffi::udf::FFI_ScalarUDF; use savvy::{savvy_err, IntoExtPtrSexp}; +use sedona_extension::extension::SedonaCTableProvider; +use sedona_extension::table_provider::ImportedTableProvider; pub fn import_schema(mut xptr: savvy::Sexp) -> savvy::Result { let ffi_schema: &FFI_ArrowSchema = import_xptr(&mut xptr, "nanoarrow_schema")?; @@ -62,10 +64,19 @@ pub fn import_array_stream(mut xptr: savvy::Sexp) -> savvy::Result savvy::Result> { - let ffi_provider: &FFI_TableProvider = - import_xptr(&mut provider_xptr, "datafusion_table_provider")?; - let provider = Arc::::from(ffi_provider); - Ok(provider) + let ffi_provider: &mut SedonaCTableProvider = + import_xptr(&mut provider_xptr, "sedonadb_table_provider")?; + // Move the SedonaCTableProvider out of the external pointer. + // Clear the structure after reading to prevent double-free when R garbage collects. + let ffi_provider = unsafe { + let provider = std::ptr::read(ffi_provider); + // Clear the entire structure to prevent any accidental use + std::ptr::write_bytes(ffi_provider as *mut SedonaCTableProvider, 0, 1); + provider + }; + // try_new validates the release callback + let provider = ImportedTableProvider::try_new(ffi_provider)?; + Ok(Arc::new(provider)) } pub fn import_scalar_udf(mut scalar_udf_xptr: savvy::Sexp) -> savvy::Result { @@ -106,5 +117,5 @@ pub struct FFIScalarUdfR(pub FFI_ScalarUDF); impl IntoExtPtrSexp for FFIScalarUdfR {} #[repr(C)] -pub struct FFITableProviderR(pub FFI_TableProvider); -impl IntoExtPtrSexp for FFITableProviderR {} +pub struct SedonaCTableProviderR(pub SedonaCTableProvider); +impl IntoExtPtrSexp for SedonaCTableProviderR {} diff --git a/r/sedonadb/tests/testthat/test-dataframe.R b/r/sedonadb/tests/testthat/test-dataframe.R index 753dce0d44..78212a7b26 100644 --- a/r/sedonadb/tests/testthat/test-dataframe.R +++ b/r/sedonadb/tests/testthat/test-dataframe.R @@ -55,7 +55,7 @@ test_that("dataframe can be created from nanoarrow objects", { test_that("dataframe can be created from an FFI table provider", { df <- as_sedonadb_dataframe(data.frame(one = 1, two = "two")) - provider <- df$df$to_provider() + provider <- df$df$to_provider(df$ctx) df2 <- as_sedonadb_dataframe(provider) expect_identical( sd_collect(df2), From d977e0fc82bf3d084890b057358163dca7c7d567 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 17:33:56 -0500 Subject: [PATCH 46/55] one last go --- r/sedonadb/src/rust/src/dataframe.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index a4093a6214..8efa234d50 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -130,11 +130,7 @@ impl InternalDataFrame { // Use the actual session state so that object stores, UDFs, and other // registrations are available when the consumer scans the provider. let session = Arc::new(ctx.inner.ctx.state()); - let exported = ExportedTableProvider::new( - provider, - session, - self.runtime.handle().clone(), - ); + let exported = ExportedTableProvider::new(provider, session, self.runtime.handle().clone()); let ffi_provider: sedona_extension::extension::SedonaCTableProvider = exported.into(); let mut ffi_xptr = SedonaCTableProviderR(ffi_provider).into_external_pointer(); From 9e08c11b14b85d1e1516748dcbebaa55150cd65d Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 22:43:47 -0500 Subject: [PATCH 47/55] use an Arc for export/import of streams --- c/sedona-extension/src/execution_plan.rs | 164 +++++++++++++---------- c/sedona-extension/src/streaming.rs | 99 +++++++++----- c/sedona-extension/src/table_provider.rs | 131 ++++++++++-------- python/sedonadb/src/dataframe.rs | 4 +- python/sedonadb/src/reader.rs | 7 +- r/sedonadb/src/rust/src/dataframe.rs | 4 +- rust/sedona-adbc/src/statement.rs | 2 +- 7 files changed, 239 insertions(+), 172 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index bf3322345c..c9d405fb31 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -35,6 +35,7 @@ use datafusion_physical_plan::{ }; use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; use serde::{Deserialize, Serialize}; +use tokio::runtime::Runtime; use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs}; use crate::set_ffi_error; @@ -42,10 +43,13 @@ use crate::streaming::{ffi_stream_to_sendable, StreamingRecordBatchReader}; use crate::utils::{cstr_from_ptr_or_empty, get_plan_property, get_plan_string_property, ERRNO_OK}; /// Wrapper around an [ExecutionPlan] that can be exported across FFI. +/// +/// Holds an `Arc` to ensure the runtime stays alive for the lifetime +/// of the exported plan. pub struct ExportedExecutionPlan { plan: Arc, task_context: Arc, - runtime: tokio::runtime::Handle, + runtime: Arc, } impl Debug for ExportedExecutionPlan { @@ -58,10 +62,13 @@ impl Debug for ExportedExecutionPlan { impl ExportedExecutionPlan { /// Create a new ExportedExecutionPlan from an ExecutionPlan. + /// + /// Takes an `Arc` to ensure the runtime stays alive for the lifetime + /// of the exported plan, preventing "Worker thread terminated" errors. pub fn new( plan: Arc, task_context: Arc, - runtime: tokio::runtime::Handle, + runtime: Arc, ) -> Self { Self { plan, @@ -752,42 +759,53 @@ mod tests { } } + /// Create a test runtime wrapped in Arc. + fn test_runtime() -> Arc { + Arc::new( + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(), + ) + } + /// Helper to set up an imported plan from a DummyExec through FFI roundtrip. - fn setup_imported_plan() -> (ImportedSedonaCExec, Arc) { + /// Returns the runtime to keep it alive for the duration of the test. + fn setup_imported_plan() -> (ImportedSedonaCExec, Arc, Arc) { let dummy = Arc::new(DummyExec::new()); - let runtime = tokio::runtime::Handle::current(); + let runtime = test_runtime(); let task_ctx = Arc::new(TaskContext::default()); - let exported = ExportedExecutionPlan::new(dummy, task_ctx.clone(), runtime); + let exported = ExportedExecutionPlan::new(dummy, task_ctx.clone(), runtime.clone()); let ffi_plan: SedonaCExecutionPlan = exported.into(); let imported = ImportedSedonaCExec::try_new(ffi_plan).unwrap(); - (imported, task_ctx) + (imported, task_ctx, runtime) } fn setup_imported_plan_with( emission_type: EmissionType, boundedness: Boundedness, supports_limit_pushdown: bool, - ) -> (ImportedSedonaCExec, Arc) { + ) -> (ImportedSedonaCExec, Arc, Arc) { let dummy = Arc::new(DummyExec::with_properties( emission_type, boundedness, supports_limit_pushdown, )); - let runtime = tokio::runtime::Handle::current(); + let runtime = test_runtime(); let task_ctx = Arc::new(TaskContext::default()); - let exported = ExportedExecutionPlan::new(dummy, task_ctx.clone(), runtime); + let exported = ExportedExecutionPlan::new(dummy, task_ctx.clone(), runtime.clone()); let ffi_plan: SedonaCExecutionPlan = exported.into(); let imported = ImportedSedonaCExec::try_new(ffi_plan).unwrap(); - (imported, task_ctx) + (imported, task_ctx, runtime) } - #[tokio::test] - async fn test_execution_plan_roundtrip_schema() { - let (imported, _) = setup_imported_plan(); + #[test] + fn test_execution_plan_roundtrip_schema() { + let (imported, _, _runtime) = setup_imported_plan(); // Verify schema matches assert_eq!(imported.schema().fields().len(), 2); @@ -795,21 +813,21 @@ mod tests { assert_eq!(imported.schema().field(1).name(), "value"); } - #[tokio::test] - async fn test_execution_plan_roundtrip_name() { - let (imported, _) = setup_imported_plan(); + #[test] + fn test_execution_plan_roundtrip_name() { + let (imported, _, _runtime) = setup_imported_plan(); assert_eq!(imported.name(), "ImportedSedonaCExec"); } - #[tokio::test] - async fn test_execution_plan_roundtrip_properties() { + #[test] + fn test_execution_plan_roundtrip_properties() { // Test all EmissionType variants for (emission_type, expected_str) in [ (EmissionType::Incremental, "Incremental"), (EmissionType::Final, "Final"), (EmissionType::Both, "Both"), ] { - let (imported, _) = setup_imported_plan_with(emission_type, Boundedness::Bounded, true); + let (imported, _, _runtime) = setup_imported_plan_with(emission_type, Boundedness::Bounded, true); let props = PlanPropertiesArgs::from_plan(&imported); assert_eq!( props.emission_type, expected_str, @@ -834,7 +852,7 @@ mod tests { "UnboundedInfiniteMemory", ), ] { - let (imported, _) = + let (imported, _, _runtime) = setup_imported_plan_with(EmissionType::Incremental, boundedness, true); let props = PlanPropertiesArgs::from_plan(&imported); assert_eq!( @@ -845,16 +863,16 @@ mod tests { } // Test supports_limit_pushdown - let (imported_with, _) = + let (imported_with, _, _runtime) = setup_imported_plan_with(EmissionType::Incremental, Boundedness::Bounded, true); assert!(imported_with.supports_limit_pushdown()); - let (imported_without, _) = + let (imported_without, _, _runtime) = setup_imported_plan_with(EmissionType::Incremental, Boundedness::Bounded, false); assert!(!imported_without.supports_limit_pushdown()); // Test partition count - let (imported, _) = setup_imported_plan(); + let (imported, _, _runtime) = setup_imported_plan(); assert_eq!( imported .properties() @@ -864,9 +882,9 @@ mod tests { ); } - #[tokio::test] - async fn test_execution_plan_roundtrip_display_as() { - let (imported, _) = setup_imported_plan(); + #[test] + fn test_execution_plan_roundtrip_display_as() { + let (imported, _, _runtime) = setup_imported_plan(); // Helper struct to format DisplayAs with a specific format type struct DisplayAsFormat<'a, T: DisplayAs>(&'a T, DisplayFormatType); @@ -894,9 +912,9 @@ mod tests { ); } - #[tokio::test] - async fn test_execution_plan_roundtrip_debug_string() { - let (imported, _) = setup_imported_plan(); + #[test] + fn test_execution_plan_roundtrip_debug_string() { + let (imported, _, _runtime) = setup_imported_plan(); let debug_str = imported.get_debug_string().unwrap(); assert!( @@ -906,48 +924,50 @@ mod tests { ); } - #[tokio::test] - async fn test_execution_plan_roundtrip_execute() { - let (imported, task_ctx) = setup_imported_plan(); - - // Execute partition 0 - let stream = imported.execute(0, task_ctx.clone()).unwrap(); - let batches: Vec = stream - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - let expected = [ - "+----+-------+", - "| id | value |", - "+----+-------+", - "| 1 | 100 |", - "| 2 | 200 |", - "| 3 | 300 |", - "+----+-------+", - ]; - assert_batches_eq!(expected, &batches); - - // Execute partition 1 - needs a fresh import since we consumed the first - let stream2 = imported.execute(1, task_ctx).unwrap(); - let batches2: Vec = stream2 - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - let expected2 = [ - "+----+-------+", - "| id | value |", - "+----+-------+", - "| 11 | 100 |", - "| 12 | 200 |", - "| 13 | 300 |", - "+----+-------+", - ]; - assert_batches_eq!(expected2, &batches2); + #[test] + fn test_execution_plan_roundtrip_execute() { + let (imported, task_ctx, runtime) = setup_imported_plan(); + + runtime.block_on(async { + // Execute partition 0 + let stream = imported.execute(0, task_ctx.clone()).unwrap(); + let batches: Vec = stream + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap(); + + let expected = [ + "+----+-------+", + "| id | value |", + "+----+-------+", + "| 1 | 100 |", + "| 2 | 200 |", + "| 3 | 300 |", + "+----+-------+", + ]; + assert_batches_eq!(expected, &batches); + + // Execute partition 1 - needs a fresh import since we consumed the first + let stream2 = imported.execute(1, task_ctx).unwrap(); + let batches2: Vec = stream2 + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap(); + + let expected2 = [ + "+----+-------+", + "| id | value |", + "+----+-------+", + "| 11 | 100 |", + "| 12 | 200 |", + "| 13 | 300 |", + "+----+-------+", + ]; + assert_batches_eq!(expected2, &batches2); + }); } } diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index 323a132049..f26a8ae686 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -20,12 +20,15 @@ //! This module provides [`StreamingRecordBatchReader`] for exporting DataFusion streams //! over FFI, and [`ffi_stream_to_sendable`] for importing FFI streams back into DataFusion. +use std::sync::Arc; + use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::{ArrowError, SchemaRef}; use datafusion_common::{exec_err, Result}; use datafusion_physical_plan::SendableRecordBatchStream; use futures::StreamExt; +use tokio::runtime::Runtime; /// A cancellation check callback for FFI operations. /// @@ -68,7 +71,9 @@ struct StreamWorker { pub struct StreamingRecordBatchReader { schema: SchemaRef, /// Stream and runtime, wrapped in Option so they can be moved to the worker. - stream_and_runtime: Option<(SendableRecordBatchStream, tokio::runtime::Handle)>, + /// We hold Arc instead of just Handle to keep the runtime alive + /// as long as this reader exists. + stream_and_runtime: Option<(SendableRecordBatchStream, Arc)>, /// Worker thread state, lazily initialized on first fetch. worker: Option, cancel_checker: Option, @@ -79,7 +84,11 @@ pub struct StreamingRecordBatchReader { impl StreamingRecordBatchReader { /// Create a new StreamingRecordBatchReader from a SendableRecordBatchStream. - pub fn new(stream: SendableRecordBatchStream, runtime: tokio::runtime::Handle) -> Self { + /// + /// Takes an `Arc` to ensure the runtime stays alive for the lifetime + /// of the reader. This prevents "Worker thread terminated" errors when the + /// runtime would otherwise be dropped while the stream is still being read. + pub fn new(stream: SendableRecordBatchStream, runtime: Arc) -> Self { Self { schema: stream.schema(), stream_and_runtime: Some((stream, runtime)), @@ -96,9 +105,12 @@ impl StreamingRecordBatchReader { /// The cancellation checker is called before each batch is fetched. If it /// returns `true`, iteration stops with a cancellation error on the next /// call and `None` on subsequent calls. + /// + /// Takes an `Arc` to ensure the runtime stays alive for the lifetime + /// of the reader. pub fn with_cancel_checker( stream: SendableRecordBatchStream, - runtime: tokio::runtime::Handle, + runtime: Arc, cancel_checker: CancelChecker, ) -> Self { Self { @@ -157,7 +169,7 @@ impl StreamingRecordBatchReader { let mut stream = stream; // Process requests until the channel is closed while request_rx.recv().is_ok() { - let result = runtime.block_on(async { + let result = runtime.handle().block_on(async { match stream.next().await { Some(Ok(batch)) => Some(Ok(batch)), Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), @@ -330,6 +342,16 @@ mod tests { use std::sync::Arc; use std::time::Duration; + /// Create a test runtime wrapped in Arc. + fn test_runtime() -> Arc { + Arc::new( + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(), + ) + } + /// Create a slow stream that yields batches with a configurable delay. fn create_slow_stream( num_batches: usize, @@ -363,9 +385,9 @@ mod tests { ) } - #[tokio::test] - async fn test_streaming_reader_basic() { - let runtime = tokio::runtime::Handle::current(); + #[test] + fn test_streaming_reader_basic() { + let runtime = test_runtime(); let (_schema, stream) = create_slow_stream(5, 10); let reader = StreamingRecordBatchReader::new(stream, runtime); @@ -383,9 +405,9 @@ mod tests { } } - #[tokio::test] - async fn test_streaming_reader_cancel() { - let runtime = tokio::runtime::Handle::current(); + #[test] + fn test_streaming_reader_cancel() { + let runtime = test_runtime(); let (_schema, stream) = create_slow_stream(10, 50); // Cancel after reading 3 batches @@ -419,11 +441,11 @@ mod tests { ); } - #[tokio::test] - async fn test_ffi_stream_to_sendable_basic() { + #[test] + fn test_ffi_stream_to_sendable_basic() { use arrow_array::ffi_stream::FFI_ArrowArrayStream; - let runtime = tokio::runtime::Handle::current(); + let runtime = test_runtime(); let (_schema, stream) = create_slow_stream(5, 10); // Export to FFI stream @@ -434,7 +456,7 @@ mod tests { let imported = unsafe { ffi_stream_to_sendable(&mut ffi_stream).unwrap() }; // Collect results - let batches: Vec<_> = imported.collect::>().await; + let batches: Vec<_> = runtime.block_on(imported.collect::>()); assert_eq!(batches.len(), 5); for (i, batch) in batches.iter().enumerate() { @@ -448,11 +470,11 @@ mod tests { } } - #[tokio::test] - async fn test_ffi_stream_to_sendable_cancel() { + #[test] + fn test_ffi_stream_to_sendable_cancel() { use arrow_array::ffi_stream::FFI_ArrowArrayStream; - let runtime = tokio::runtime::Handle::current(); + let runtime = test_runtime(); let (_schema, stream) = create_slow_stream(10, 50); // Export to FFI stream (no cancellation on export side) @@ -470,18 +492,21 @@ mod tests { }; // Collect with cancellation after 3 batches, stop on first error - let mut batches = Vec::new(); - let mut stream = imported; - while let Some(result) = stream.next().await { - let is_err = result.is_err(); - batches.push(result); - if batches.len() == 3 { - cancelled.store(true, Ordering::SeqCst); - } - if is_err { - break; // Stop on first error + let batches = runtime.block_on(async { + let mut batches = Vec::new(); + let mut stream = imported; + while let Some(result) = stream.next().await { + let is_err = result.is_err(); + batches.push(result); + if batches.len() == 3 { + cancelled.store(true, Ordering::SeqCst); + } + if is_err { + break; // Stop on first error + } } - } + batches + }); // Should have 3 successful batches + 1 cancellation error assert_eq!(batches.len(), 4); @@ -529,9 +554,9 @@ mod tests { ) } - #[tokio::test] - async fn test_streaming_reader_skip_empty_batches() { - let runtime = tokio::runtime::Handle::current(); + #[test] + fn test_streaming_reader_skip_empty_batches() { + let runtime = test_runtime(); // Create stream with: 2 rows, 0 rows, 3 rows, 0 rows, 0 rows, 1 row let (_schema, stream) = create_stream_with_empty_batches(vec![2, 0, 3, 0, 0, 1]); @@ -547,9 +572,9 @@ mod tests { assert_eq!(batches[2].as_ref().unwrap().num_rows(), 1); } - #[tokio::test] - async fn test_streaming_reader_no_skip_empty_batches() { - let runtime = tokio::runtime::Handle::current(); + #[test] + fn test_streaming_reader_no_skip_empty_batches() { + let runtime = test_runtime(); // Create stream with: 2 rows, 0 rows, 3 rows let (_schema, stream) = create_stream_with_empty_batches(vec![2, 0, 3]); @@ -564,9 +589,9 @@ mod tests { assert_eq!(batches[2].as_ref().unwrap().num_rows(), 3); } - #[tokio::test] - async fn test_streaming_reader_periodic_check_interval() { - let runtime = tokio::runtime::Handle::current(); + #[test] + fn test_streaming_reader_periodic_check_interval() { + let runtime = test_runtime(); // Create a slow stream where each batch takes 200ms let (_schema, stream) = create_slow_stream(5, 200); diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index a1f785708a..8f79897b57 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -32,6 +32,7 @@ use datafusion_expr::{Expr, TableType}; use datafusion_physical_plan::ExecutionPlan; use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; use serde::{Deserialize, Serialize}; +use tokio::runtime::Runtime; use crate::execution_plan::{ExportedExecutionPlan, ImportedSedonaCExec}; use crate::extension::{ @@ -47,7 +48,9 @@ use crate::utils::{cstr_from_ptr_or_empty, get_table_provider_string_property, E pub struct ExportedTableProvider { inner: Arc, session: Arc, - runtime: tokio::runtime::Handle, + /// We hold Arc instead of just Handle to keep the runtime alive + /// as long as this provider exists. + runtime: Arc, } impl Debug for ExportedTableProvider { @@ -63,10 +66,13 @@ impl ExportedTableProvider { /// /// The session is used during scan operations and must support physical planning /// if the inner TableProvider requires it (e.g., for Views). + /// + /// Takes an `Arc` to ensure the runtime stays alive for the lifetime + /// of the provider, preventing "Worker thread terminated" errors. pub fn new( inner: Arc, session: Arc, - runtime: tokio::runtime::Handle, + runtime: Arc, ) -> Self { Self { inner, @@ -86,7 +92,7 @@ impl ExportedTableProvider { std::thread::spawn(move || { let projection_ref = projection.as_ref(); - runtime.block_on(inner.scan(session.as_ref(), projection_ref, &[], limit)) + runtime.handle().block_on(inner.scan(session.as_ref(), projection_ref, &[], limit)) }) .join() .map_err(|_| { @@ -483,33 +489,35 @@ mod tests { /// 2. Exporting the table provider through FFI /// 3. Importing it and registering in a new context /// 4. Running the SQL query and asserting results - async fn test_roundtrip_query(sql: &str, expected: &[&str]) -> Result<()> { - let ctx = create_test_context().await?; - - // Get the table provider from the context - let table = ctx.table_provider("test_data").await?; - - // Export the table provider - let runtime = tokio::runtime::Handle::current(); - let session = Arc::new(ctx.state()); - let exported = ExportedTableProvider::new(table, session, runtime); - let ffi_provider: SedonaCTableProvider = exported.into(); - - // Import the table provider - let imported = ImportedTableProvider::try_new(ffi_provider)?; - - // Create a new context and register the imported table - let ctx2 = SessionContext::new(); - ctx2.register_table("imported_data", Arc::new(imported))?; - - // Query and verify - let result = ctx2.sql(sql).await?.collect().await?; - assert_batches_eq!(expected, &result); - Ok(()) + fn test_roundtrip_query(sql: &str, expected: &[&str]) -> Result<()> { + let runtime = test_runtime(); + runtime.block_on(async { + let ctx = create_test_context().await?; + + // Get the table provider from the context + let table = ctx.table_provider("test_data").await?; + + // Export the table provider + let session = Arc::new(ctx.state()); + let exported = ExportedTableProvider::new(table, session, runtime.clone()); + let ffi_provider: SedonaCTableProvider = exported.into(); + + // Import the table provider + let imported = ImportedTableProvider::try_new(ffi_provider)?; + + // Create a new context and register the imported table + let ctx2 = SessionContext::new(); + ctx2.register_table("imported_data", Arc::new(imported))?; + + // Query and verify + let result = ctx2.sql(sql).await?.collect().await?; + assert_batches_eq!(expected, &result); + Ok(()) + }) } - #[tokio::test] - async fn test_roundtrip_simple_select() -> Result<()> { + #[test] + fn test_roundtrip_simple_select() { test_roundtrip_query( "SELECT id, value_a FROM imported_data ORDER BY id LIMIT 5", &[ @@ -524,11 +532,11 @@ mod tests { "+----+---------+", ], ) - .await + .unwrap(); } - #[tokio::test] - async fn test_roundtrip_projection() -> Result<()> { + #[test] + fn test_roundtrip_projection() { test_roundtrip_query( "SELECT value_b, value_d FROM imported_data ORDER BY value_b LIMIT 3", &[ @@ -541,11 +549,11 @@ mod tests { "+---------+---------+", ], ) - .await + .unwrap(); } - #[tokio::test] - async fn test_roundtrip_filter() -> Result<()> { + #[test] + fn test_roundtrip_filter() { test_roundtrip_query( "SELECT id, value_a FROM imported_data WHERE id > 20 ORDER BY id LIMIT 5", &[ @@ -560,11 +568,11 @@ mod tests { "+----+---------+", ], ) - .await + .unwrap(); } - #[tokio::test] - async fn test_roundtrip_sort() -> Result<()> { + #[test] + fn test_roundtrip_sort() { test_roundtrip_query( "SELECT id, value_c FROM imported_data ORDER BY id DESC LIMIT 5", &[ @@ -579,22 +587,22 @@ mod tests { "+----+---------+", ], ) - .await + .unwrap(); } - #[tokio::test] - async fn test_roundtrip_limit() -> Result<()> { + #[test] + fn test_roundtrip_limit() { test_roundtrip_query( "SELECT id FROM imported_data ORDER BY id LIMIT 3", &[ "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 3 |", "+----+", ], ) - .await + .unwrap(); } - #[tokio::test] - async fn test_roundtrip_all_columns() -> Result<()> { + #[test] + fn test_roundtrip_all_columns() { test_roundtrip_query( "SELECT * FROM imported_data ORDER BY id LIMIT 2", &[ @@ -606,7 +614,7 @@ mod tests { "+----+---------+---------+---------+---------+", ], ) - .await + .unwrap(); } /// A dummy TableProvider with configurable table_type for testing FFI roundtrip. @@ -651,20 +659,32 @@ mod tests { } } + /// Create a test runtime wrapped in Arc. + fn test_runtime() -> Arc { + Arc::new( + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(), + ) + } + /// Helper to set up an imported table provider from a DummyTableProvider through FFI roundtrip. - fn setup_imported_provider_with(table_type: TableType) -> ImportedTableProvider { + /// Returns the runtime to keep it alive for the duration of the test. + fn setup_imported_provider_with(table_type: TableType) -> (ImportedTableProvider, Arc) { let dummy = Arc::new(DummyTableProvider::with_table_type(table_type)); let ctx = SessionContext::new(); - let runtime = tokio::runtime::Handle::current(); + let runtime = test_runtime(); let session = Arc::new(ctx.state()); - let exported = ExportedTableProvider::new(dummy, session, runtime); + let exported = ExportedTableProvider::new(dummy, session, runtime.clone()); let ffi_provider: SedonaCTableProvider = exported.into(); - ImportedTableProvider::try_new(ffi_provider).expect("Failed to import table provider") + let imported = ImportedTableProvider::try_new(ffi_provider).expect("Failed to import table provider"); + (imported, runtime) } - #[tokio::test] - async fn test_table_provider_roundtrip_schema() -> Result<()> { - let imported = setup_imported_provider_with(TableType::Base); + #[test] + fn test_table_provider_roundtrip_schema() { + let (imported, _runtime) = setup_imported_provider_with(TableType::Base); // Check schema roundtrip let schema = imported.schema(); @@ -673,16 +693,13 @@ mod tests { assert_eq!(schema.field(0).data_type(), &DataType::Int32); assert_eq!(schema.field(1).name(), "value"); assert_eq!(schema.field(1).data_type(), &DataType::Int32); - - Ok(()) } - #[tokio::test] - async fn test_table_provider_roundtrip_table_type() -> Result<()> { + #[test] + fn test_table_provider_roundtrip_table_type() { for table_type in [TableType::Base, TableType::View, TableType::Temporary] { - let imported = setup_imported_provider_with(table_type); + let (imported, _runtime) = setup_imported_provider_with(table_type); assert_eq!(imported.table_type(), table_type); } - Ok(()) } } diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs index bb5ee51d76..81e2f969cd 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -449,7 +449,7 @@ impl InternalDataFrame { simplify: Option, ) -> Result { let stream = wait_for_future(py, &self.runtime, self.inner.clone().execute_stream())??; - let reader = new_py_streaming_reader(stream, self.runtime.handle().clone()); + let reader = new_py_streaming_reader(stream, self.runtime.clone()); let mut reader: Box = Box::new(reader); if simplify.unwrap_or(false) { @@ -676,7 +676,7 @@ impl InternalDataFrame { let exported = sedona_extension::table_provider::ExportedTableProvider::new( provider, session, - self.runtime.handle().clone(), + self.runtime.clone(), ); let ffi_provider: sedona_extension::extension::SedonaCTableProvider = exported.into(); Ok(PyCapsule::new_with_value( diff --git a/python/sedonadb/src/reader.rs b/python/sedonadb/src/reader.rs index fc0b54815b..c1bb7c8217 100644 --- a/python/sedonadb/src/reader.rs +++ b/python/sedonadb/src/reader.rs @@ -14,11 +14,13 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. +use std::sync::Arc; use std::time::Duration; use datafusion::execution::SendableRecordBatchStream; use pyo3::Python; use sedona_extension::streaming::StreamingRecordBatchReader; +use tokio::runtime::Runtime; /// Interval for checking Python signals during batch fetches. const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(2_000); @@ -29,9 +31,12 @@ const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(2_000); /// /// The reader will check Python signals every 2 seconds during batch fetches, /// allowing users to interrupt long-running queries. +/// +/// Takes an `Arc` to ensure the runtime stays alive for the lifetime +/// of the reader, preventing "Worker thread terminated" errors. pub fn new_py_streaming_reader( stream: SendableRecordBatchStream, - runtime: tokio::runtime::Handle, + runtime: Arc, ) -> StreamingRecordBatchReader { // Create a cancel checker that checks Python signals let cancel_checker: Box bool + Send + Sync> = Box::new(|| { diff --git a/r/sedonadb/src/rust/src/dataframe.rs b/r/sedonadb/src/rust/src/dataframe.rs index 8efa234d50..4e643b2172 100644 --- a/r/sedonadb/src/rust/src/dataframe.rs +++ b/r/sedonadb/src/rust/src/dataframe.rs @@ -114,7 +114,7 @@ impl InternalDataFrame { async move { inner.execute_stream().await }, )??; - let reader = StreamingRecordBatchReader::new(stream, self.runtime.handle().clone()) + let reader = StreamingRecordBatchReader::new(stream, self.runtime.clone()) .with_skip_empty_batches(true); let reader: Box = Box::new(reader); @@ -130,7 +130,7 @@ impl InternalDataFrame { // Use the actual session state so that object stores, UDFs, and other // registrations are available when the consumer scans the provider. let session = Arc::new(ctx.inner.ctx.state()); - let exported = ExportedTableProvider::new(provider, session, self.runtime.handle().clone()); + let exported = ExportedTableProvider::new(provider, session, self.runtime.clone()); let ffi_provider: sedona_extension::extension::SedonaCTableProvider = exported.into(); let mut ffi_xptr = SedonaCTableProviderR(ffi_provider).into_external_pointer(); diff --git a/rust/sedona-adbc/src/statement.rs b/rust/sedona-adbc/src/statement.rs index fa672e2c73..3c7592107b 100644 --- a/rust/sedona-adbc/src/statement.rs +++ b/rust/sedona-adbc/src/statement.rs @@ -99,7 +99,7 @@ impl Statement for SedonaStatement { self.runtime.block_on(async { let df = self.ctx.sql(&query).await.map_err(from_datafusion_error)?; let stream = df.execute_stream().await.map_err(from_datafusion_error)?; - let reader = StreamingRecordBatchReader::new(stream, self.runtime.handle().clone()) + let reader = StreamingRecordBatchReader::new(stream, self.runtime.clone()) .with_skip_empty_batches(true); Ok(Box::new(reader) as Box) }) From ceaca402aefbd9d924cbc220c60fad88554da8e3 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 2 Jul 2026 23:14:18 -0500 Subject: [PATCH 48/55] maybe better cancellation --- c/sedona-extension/src/streaming.rs | 62 ++++++++++++++++++++++++++--- 1 file changed, 57 insertions(+), 5 deletions(-) diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index f26a8ae686..81e3907e5f 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -20,6 +20,7 @@ //! This module provides [`StreamingRecordBatchReader`] for exporting DataFusion streams //! over FFI, and [`ffi_stream_to_sendable`] for importing FFI streams back into DataFusion. +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use arrow_array::ffi_stream::FFI_ArrowArrayStream; @@ -57,6 +58,9 @@ struct StreamWorker { request_tx: std::sync::mpsc::SyncSender<()>, /// Channel to receive batch results. response_rx: std::sync::mpsc::Receiver, + /// Cancellation flag shared with worker thread. + /// When set to true, the worker will abort the current fetch and exit. + cancel_flag: Arc, /// Worker thread handle for cleanup. _handle: std::thread::JoinHandle<()>, } @@ -164,21 +168,59 @@ impl StreamingRecordBatchReader { // stops early due to periodic cancellation before receiving the response let (response_tx, response_rx) = std::sync::mpsc::sync_channel::(1); + // Cancellation flag - when set, worker will abort current fetch + let cancel_flag = Arc::new(AtomicBool::new(false)); + let worker_cancel_flag = cancel_flag.clone(); + // Spawn the worker thread let handle = std::thread::spawn(move || { let mut stream = stream; // Process requests until the channel is closed while request_rx.recv().is_ok() { + // Check if cancelled before starting fetch + if worker_cancel_flag.load(Ordering::Relaxed) { + // Send cancellation error and exit + let _ = response_tx.send(Some(Err(cancellation_arrow_error()))); + break; + } + let result = runtime.handle().block_on(async { - match stream.next().await { - Some(Ok(batch)) => Some(Ok(batch)), - Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), - None => None, + // Poll cancellation periodically while waiting for the next batch. + // This allows us to abort the fetch if cancellation is requested. + let cancel_check = async { + loop { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + if worker_cancel_flag.load(Ordering::Relaxed) { + return; + } + } + }; + + tokio::select! { + biased; + // Check cancellation first + _ = cancel_check => { + Some(Err(ArrowError::ExternalError(Box::new(std::io::Error::new( + std::io::ErrorKind::Interrupted, + "Operation cancelled", + ))))) + } + // Then try to get the next batch + batch = stream.next() => { + match batch { + Some(Ok(batch)) => Some(Ok(batch)), + Some(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))), + None => None, + } + } } }); + // If cancelled, exit after sending response + let is_cancelled = worker_cancel_flag.load(Ordering::Relaxed); + // Send the result back; if send fails, the reader was dropped - if response_tx.send(result).is_err() { + if response_tx.send(result).is_err() || is_cancelled { break; } } @@ -187,6 +229,7 @@ impl StreamingRecordBatchReader { self.worker = Some(StreamWorker { request_tx, response_rx, + cancel_flag, _handle: handle, }); } @@ -219,6 +262,15 @@ impl StreamingRecordBatchReader { if let Some(ref checker) = self.cancel_checker { if checker() { self.cancelled = true; + // Signal the worker to abort the current fetch + // This will cause the worker's select! to see the + // cancellation and drop the stream future. + worker.cancel_flag.store(true, Ordering::Relaxed); + // Wait briefly for the worker to send its response + // then return the cancellation error + let _ = worker + .response_rx + .recv_timeout(std::time::Duration::from_millis(200)); return Some(Err(cancellation_arrow_error())); } } From e6d2964849d530257f8c1b7c9399c22751cd68a4 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 6 Jul 2026 13:40:35 -0500 Subject: [PATCH 49/55] tweak exec issues --- Cargo.lock | 12 ++++ Cargo.toml | 1 + c/sedona-extension/Cargo.toml | 1 + c/sedona-extension/src/execution_plan.rs | 3 +- c/sedona-extension/src/streaming.rs | 86 ++++++++++++++---------- c/sedona-extension/src/table_provider.rs | 11 ++- 6 files changed, 75 insertions(+), 39 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d19387ccd7..8857115709 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5776,6 +5776,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "tokio-stream", ] [[package]] @@ -7016,6 +7017,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.18" diff --git a/Cargo.toml b/Cargo.toml index aa3db9a502..5491848e4b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -133,6 +133,7 @@ serde_with = { version = "3" } tempfile = { version = "3"} thiserror = { version = "2" } tokio = { version = "1.52", features = ["macros", "rt", "sync"] } +tokio-stream = "0.1" url = "2.5.7" wkb = "0.9.2" wkt = "0.14.0" diff --git a/c/sedona-extension/Cargo.toml b/c/sedona-extension/Cargo.toml index bb6a7acde8..4800913b39 100644 --- a/c/sedona-extension/Cargo.toml +++ b/c/sedona-extension/Cargo.toml @@ -46,6 +46,7 @@ sedona-testing = { path = "../../rust/sedona-testing" } serde = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } +tokio-stream = { workspace = true } [dev-dependencies] datafusion = { workspace = true, features = ["sql"] } diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index c9d405fb31..c21e9f605c 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -827,7 +827,8 @@ mod tests { (EmissionType::Final, "Final"), (EmissionType::Both, "Both"), ] { - let (imported, _, _runtime) = setup_imported_plan_with(emission_type, Boundedness::Bounded, true); + let (imported, _, _runtime) = + setup_imported_plan_with(emission_type, Boundedness::Bounded, true); let props = PlanPropertiesArgs::from_plan(&imported); assert_eq!( props.emission_type, expected_str, diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index 81e3907e5f..228e7c2e7e 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -53,11 +53,15 @@ fn cancellation_arrow_error() -> ArrowError { type BatchResult = Option>; /// Worker thread state for streaming record batch reads. +/// +/// Uses eager prefetch: the worker fetches batches into a bounded buffer as fast +/// as possible, providing one-batch pipelining. This allows producer fetch and +/// consumer processing to overlap, roughly halving per-batch wall time compared +/// to a zero-capacity rendezvous protocol. struct StreamWorker { - /// Channel to request the next batch (send () to request). - request_tx: std::sync::mpsc::SyncSender<()>, - /// Channel to receive batch results. - response_rx: std::sync::mpsc::Receiver, + /// Channel to receive prefetched batch results. + /// The worker eagerly pushes batches here; bounded capacity provides backpressure. + batch_rx: std::sync::mpsc::Receiver, /// Cancellation flag shared with worker thread. /// When set to true, the worker will abort the current fetch and exit. cancel_flag: Arc, @@ -161,26 +165,24 @@ impl StreamingRecordBatchReader { return; }; - // Create channels for communication - // Use bounded channel with size 0 (rendezvous) for backpressure - let (request_tx, request_rx) = std::sync::mpsc::sync_channel::<()>(0); - // Use bounded channel with size 1 for response to prevent deadlock if reader - // stops early due to periodic cancellation before receiving the response - let (response_tx, response_rx) = std::sync::mpsc::sync_channel::(1); + // Create bounded channel for prefetched batches. + // Capacity 2 allows one-batch pipelining: while consumer processes batch N, + // worker can fetch batch N+1 and have it ready. This roughly halves per-batch + // wall time (max(T_produce, T_consume) instead of T_produce + T_consume). + let (batch_tx, batch_rx) = std::sync::mpsc::sync_channel::(2); // Cancellation flag - when set, worker will abort current fetch let cancel_flag = Arc::new(AtomicBool::new(false)); let worker_cancel_flag = cancel_flag.clone(); - // Spawn the worker thread + // Spawn the worker thread - eagerly fetches batches into the buffer let handle = std::thread::spawn(move || { let mut stream = stream; - // Process requests until the channel is closed - while request_rx.recv().is_ok() { + + loop { // Check if cancelled before starting fetch if worker_cancel_flag.load(Ordering::Relaxed) { - // Send cancellation error and exit - let _ = response_tx.send(Some(Err(cancellation_arrow_error()))); + let _ = batch_tx.send(Some(Err(cancellation_arrow_error()))); break; } @@ -216,19 +218,19 @@ impl StreamingRecordBatchReader { } }); - // If cancelled, exit after sending response + let is_end = result.is_none(); let is_cancelled = worker_cancel_flag.load(Ordering::Relaxed); - // Send the result back; if send fails, the reader was dropped - if response_tx.send(result).is_err() || is_cancelled { + // Send eagerly - blocks if buffer is full (backpressure) + // If send fails, the reader was dropped + if batch_tx.send(result).is_err() || is_end || is_cancelled { break; } } }); self.worker = Some(StreamWorker { - request_tx, - response_rx, + batch_rx, cancel_flag, _handle: handle, }); @@ -243,19 +245,12 @@ impl StreamingRecordBatchReader { ))); }; - // Request the next batch - if worker.request_tx.send(()).is_err() { - return Some(Err(ArrowError::InvalidArgumentError( - "Worker thread terminated".to_string(), - ))); - } - - // Wait for the response, with optional periodic cancellation checking + // Receive from the prefetch buffer, with optional periodic cancellation checking match &self.periodic_check_interval { Some(interval) => { let interval = *interval; loop { - match worker.response_rx.recv_timeout(interval) { + match worker.batch_rx.recv_timeout(interval) { Ok(result) => return result, Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { // Check for cancellation @@ -269,7 +264,7 @@ impl StreamingRecordBatchReader { // Wait briefly for the worker to send its response // then return the cancellation error let _ = worker - .response_rx + .batch_rx .recv_timeout(std::time::Duration::from_millis(200)); return Some(Err(cancellation_arrow_error())); } @@ -285,8 +280,8 @@ impl StreamingRecordBatchReader { } } None => { - // Simple blocking receive - match worker.response_rx.recv() { + // Simple blocking receive from prefetch buffer + match worker.batch_rx.recv() { Ok(result) => result, Err(_) => Some(Err(ArrowError::InvalidArgumentError( "Worker thread terminated".to_string(), @@ -356,7 +351,13 @@ pub unsafe fn ffi_stream_to_sendable( /// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream with cancellation support. /// -/// The cancellation checker is called before each batch is read. If it returns +/// Uses a dedicated OS thread to read from the synchronous FFI stream, sending +/// batches through an async channel. This avoids blocking tokio worker threads, +/// which would otherwise cause deadlocks if the producer side also needs the +/// same runtime via `block_on()` (especially fatal with current-thread +/// runtimes, but problematic with any shared runtime). +/// +/// The cancellation checker is called before each batch is yielded. If it returns /// `true`, the stream yields a cancellation error. /// /// # Safety @@ -368,9 +369,24 @@ pub unsafe fn ffi_stream_to_sendable_with_cancel( cancel_checker: Option, ) -> Result { let reader = arrow_array::ffi_stream::ArrowArrayStreamReader::from_raw(ffi_stream)?; - let schema = reader.schema(); - let stream = futures::stream::iter(reader).map(move |result| { + + // Use an async channel with capacity 2 for one-batch pipelining. + // The dedicated thread reads blocking batches from FFI, the async stream + // receives them without blocking tokio workers. + let (tx, rx) = tokio::sync::mpsc::channel::>(2); + + // Dedicated thread for blocking reads - keeps blocking I/O off tokio workers + std::thread::spawn(move || { + for batch in reader { + // blocking_send blocks this thread if buffer is full (backpressure) + if tx.blocking_send(batch).is_err() { + break; // Receiver dropped, stop reading + } + } + }); + + let stream = tokio_stream::wrappers::ReceiverStream::new(rx).map(move |result| { // Check for cancellation before yielding each batch if let Some(ref checker) = cancel_checker { if checker() { diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 8f79897b57..f96b42af13 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -92,7 +92,9 @@ impl ExportedTableProvider { std::thread::spawn(move || { let projection_ref = projection.as_ref(); - runtime.handle().block_on(inner.scan(session.as_ref(), projection_ref, &[], limit)) + runtime + .handle() + .block_on(inner.scan(session.as_ref(), projection_ref, &[], limit)) }) .join() .map_err(|_| { @@ -671,14 +673,17 @@ mod tests { /// Helper to set up an imported table provider from a DummyTableProvider through FFI roundtrip. /// Returns the runtime to keep it alive for the duration of the test. - fn setup_imported_provider_with(table_type: TableType) -> (ImportedTableProvider, Arc) { + fn setup_imported_provider_with( + table_type: TableType, + ) -> (ImportedTableProvider, Arc) { let dummy = Arc::new(DummyTableProvider::with_table_type(table_type)); let ctx = SessionContext::new(); let runtime = test_runtime(); let session = Arc::new(ctx.state()); let exported = ExportedTableProvider::new(dummy, session, runtime.clone()); let ffi_provider: SedonaCTableProvider = exported.into(); - let imported = ImportedTableProvider::try_new(ffi_provider).expect("Failed to import table provider"); + let imported = + ImportedTableProvider::try_new(ffi_provider).expect("Failed to import table provider"); (imported, runtime) } From e35e663fdcb860775d02a201db78cec1b7bb7f4b Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 6 Jul 2026 13:55:00 -0500 Subject: [PATCH 50/55] pipe cancelation through both sides of the interaction --- c/sedona-extension/src/execution_plan.rs | 24 +++++++++++++++++++++-- c/sedona-extension/src/streaming.rs | 25 ++++-------------------- c/sedona-extension/src/table_provider.rs | 24 ++++++++++++++++++++++- python/sedonadb/src/import_from.rs | 12 ++++++++++++ 4 files changed, 61 insertions(+), 24 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index c21e9f605c..7766a1f8c7 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -39,7 +39,7 @@ use tokio::runtime::Runtime; use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs}; use crate::set_ffi_error; -use crate::streaming::{ffi_stream_to_sendable, StreamingRecordBatchReader}; +use crate::streaming::{ffi_stream_to_sendable, CancelChecker, StreamingRecordBatchReader}; use crate::utils::{cstr_from_ptr_or_empty, get_plan_property, get_plan_string_property, ERRNO_OK}; /// Wrapper around an [ExecutionPlan] that can be exported across FFI. @@ -317,6 +317,8 @@ pub struct ImportedSedonaCExec { properties: PlanProperties, supports_limit_pushdown: bool, name: String, + /// Stored as Arc so it can be cloned for each partition execution. + cancel_checker: Option bool + Send + Sync>>, } impl Debug for ImportedSedonaCExec { @@ -370,9 +372,22 @@ impl ImportedSedonaCExec { properties, supports_limit_pushdown, name, + cancel_checker: None, }) } + /// Set a cancellation checker for this execution plan. + /// + /// The checker is called before each batch is read from the FFI stream. + /// If it returns `true`, the stream yields a cancellation error. + pub fn with_cancel_checker(mut self, cancel_checker: F) -> Self + where + F: Fn() -> bool + Send + Sync + 'static, + { + self.cancel_checker = Some(Arc::new(cancel_checker)); + self + } + fn get_debug_string(&self) -> Result { get_plan_string_property(&self.inner, "debug_string") } @@ -480,7 +495,12 @@ impl ExecutionPlan for ImportedSedonaCExec { } // Convert FFI stream to SendableRecordBatchStream - unsafe { ffi_stream_to_sendable(&mut ffi_stream) } + // Clone the Arc and wrap in a new Box for this execution + let cancel_checker: Option = self.cancel_checker.as_ref().map(|c| { + let c = c.clone(); + Box::new(move || c()) as CancelChecker + }); + unsafe { ffi_stream_to_sendable(&mut ffi_stream, cancel_checker) } } fn children(&self) -> Vec<&Arc> { diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index 228e7c2e7e..3c08ed2f11 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -333,22 +333,6 @@ impl RecordBatchReader for StreamingRecordBatchReader { } } -/// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream. -/// -/// This is the inverse of StreamingRecordBatchReader - it takes an FFI stream -/// and converts it back to a DataFusion stream. Used when importing execution -/// plans from across FFI boundaries. -/// -/// # Safety -/// -/// The caller must ensure that the FFI stream pointer is valid and properly -/// initialized. -pub unsafe fn ffi_stream_to_sendable( - ffi_stream: &mut FFI_ArrowArrayStream, -) -> Result { - ffi_stream_to_sendable_with_cancel(ffi_stream, None) -} - /// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream with cancellation support. /// /// Uses a dedicated OS thread to read from the synchronous FFI stream, sending @@ -364,7 +348,7 @@ pub unsafe fn ffi_stream_to_sendable( /// /// The caller must ensure that the FFI stream pointer is valid and properly /// initialized. -pub unsafe fn ffi_stream_to_sendable_with_cancel( +pub unsafe fn ffi_stream_to_sendable( ffi_stream: &mut FFI_ArrowArrayStream, cancel_checker: Option, ) -> Result { @@ -521,7 +505,7 @@ mod tests { let mut ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); // Import back - let imported = unsafe { ffi_stream_to_sendable(&mut ffi_stream).unwrap() }; + let imported = unsafe { ffi_stream_to_sendable(&mut ffi_stream, None).unwrap() }; // Collect results let batches: Vec<_> = runtime.block_on(imported.collect::>()); @@ -555,9 +539,8 @@ mod tests { let cancel_checker: CancelChecker = Box::new(move || cancelled_clone.load(Ordering::SeqCst)); - let imported = unsafe { - ffi_stream_to_sendable_with_cancel(&mut ffi_stream, Some(cancel_checker)).unwrap() - }; + let imported = + unsafe { ffi_stream_to_sendable(&mut ffi_stream, Some(cancel_checker)).unwrap() }; // Collect with cancellation after 3 batches, stop on first error let batches = runtime.block_on(async { diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index f96b42af13..6e421f3dc3 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -283,6 +283,8 @@ pub struct ImportedTableProvider { inner: SedonaCTableProvider, schema: SchemaRef, table_type: TableType, + /// Stored as Arc so it can be cloned for each scan execution. + cancel_checker: Option bool + Send + Sync>>, } impl Debug for ImportedTableProvider { @@ -319,9 +321,22 @@ impl ImportedTableProvider { inner, schema, table_type, + cancel_checker: None, }) } + /// Set a cancellation checker for this table provider. + /// + /// The checker is called before each batch is read when scanning. + /// If it returns `true`, the stream yields a cancellation error. + pub fn with_cancel_checker(mut self, cancel_checker: F) -> Self + where + F: Fn() -> bool + Send + Sync + 'static, + { + self.cancel_checker = Some(Arc::new(cancel_checker)); + self + } + fn get_table_type(provider: &SedonaCTableProvider) -> Result { let type_str = match get_table_provider_string_property(provider, "table_type") { Ok(s) => s, @@ -392,7 +407,14 @@ impl TableProvider for ImportedTableProvider { return exec_err!("Failed to scan table: {}", err); } - let exec = ImportedSedonaCExec::try_new(ffi_plan)?; + let mut exec = ImportedSedonaCExec::try_new(ffi_plan)?; + + // Pipe through the cancel checker if one is configured + if let Some(ref checker) = self.cancel_checker { + let checker = checker.clone(); + exec = exec.with_cancel_checker(move || checker()); + } + Ok(Arc::new(exec)) } } diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs index aba489be70..bbb404edb3 100644 --- a/python/sedonadb/src/import_from.rs +++ b/python/sedonadb/src/import_from.rs @@ -76,6 +76,18 @@ pub fn import_sedona_ffi_table_provider( }; // try_new validates the release callback let provider = ImportedTableProvider::try_new(ffi_provider)?; + + // Add a Python-aware cancel checker that checks for Ctrl+C signals + let provider = provider.with_cancel_checker(|| { + Python::attach(|py| { + // Run `pass` to process any pending signals, then check for errors + if py.run(cr"pass", None, None).is_err() { + return true; + } + py.check_signals().is_err() + }) + }); + Ok(Arc::new(provider)) } From 5f44619f208f7eeb5ee95f40d0c0ddbad98d8ce9 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 6 Jul 2026 15:59:05 -0500 Subject: [PATCH 51/55] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- c/sedona-extension/src/streaming.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index 3c08ed2f11..ceeb70bf52 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -333,6 +333,15 @@ impl RecordBatchReader for StreamingRecordBatchReader { } } +impl Drop for StreamingRecordBatchReader { + fn drop(&mut self) { + if let Some(worker) = &self.worker { + // Ensure the worker wakes up and exits promptly if it is blocked. + worker.cancel_flag.store(true, Ordering::Relaxed); + } + } +} + /// Convert an FFI ArrowArrayStream into a SendableRecordBatchStream with cancellation support. /// /// Uses a dedicated OS thread to read from the synchronous FFI stream, sending From e2d3a7a3e83b4662a000eeecc545d42bef5139f2 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 6 Jul 2026 16:09:22 -0500 Subject: [PATCH 52/55] periodic cancelation on the import side --- c/sedona-extension/src/execution_plan.rs | 19 +++++++++++++-- c/sedona-extension/src/streaming.rs | 31 +++++++++++++++++++----- c/sedona-extension/src/table_provider.rs | 22 +++++++++++++++-- python/sedonadb/src/import_from.rs | 20 +++++++++------ 4 files changed, 74 insertions(+), 18 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index 7766a1f8c7..df61fbd202 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -21,6 +21,7 @@ use std::{ fmt::{Debug, Display, Formatter}, ptr::null_mut, sync::Arc, + time::Duration, }; use arrow_array::ffi_stream::FFI_ArrowArrayStream; @@ -319,6 +320,8 @@ pub struct ImportedSedonaCExec { name: String, /// Stored as Arc so it can be cloned for each partition execution. cancel_checker: Option bool + Send + Sync>>, + /// Interval for periodic cancellation checking during stream consumption. + check_interval: Option, } impl Debug for ImportedSedonaCExec { @@ -373,12 +376,14 @@ impl ImportedSedonaCExec { supports_limit_pushdown, name, cancel_checker: None, + check_interval: None, }) } /// Set a cancellation checker for this execution plan. /// - /// The checker is called before each batch is read from the FFI stream. + /// The checker is called periodically (controlled by `with_check_interval`) + /// before batches are read from the FFI stream. /// If it returns `true`, the stream yields a cancellation error. pub fn with_cancel_checker(mut self, cancel_checker: F) -> Self where @@ -388,6 +393,16 @@ impl ImportedSedonaCExec { self } + /// Set the interval for periodic cancellation checking. + /// + /// When set, the cancel checker will only be called when this interval + /// has elapsed since the last check, reducing overhead for fast streams. + /// If not set, the checker is called before every batch. + pub fn with_check_interval(mut self, interval: Duration) -> Self { + self.check_interval = Some(interval); + self + } + fn get_debug_string(&self) -> Result { get_plan_string_property(&self.inner, "debug_string") } @@ -500,7 +515,7 @@ impl ExecutionPlan for ImportedSedonaCExec { let c = c.clone(); Box::new(move || c()) as CancelChecker }); - unsafe { ffi_stream_to_sendable(&mut ffi_stream, cancel_checker) } + unsafe { ffi_stream_to_sendable(&mut ffi_stream, cancel_checker, self.check_interval) } } fn children(&self) -> Vec<&Arc> { diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index ceeb70bf52..0edd430d99 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -22,6 +22,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; +use std::time::{Duration, Instant}; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{RecordBatch, RecordBatchReader}; @@ -350,8 +351,9 @@ impl Drop for StreamingRecordBatchReader { /// same runtime via `block_on()` (especially fatal with current-thread /// runtimes, but problematic with any shared runtime). /// -/// The cancellation checker is called before each batch is yielded. If it returns -/// `true`, the stream yields a cancellation error. +/// The cancellation checker is called periodically (at most once per `check_interval`) +/// before yielding batches. If it returns `true`, the stream yields a cancellation error. +/// If `check_interval` is `None`, the checker is called before every batch. /// /// # Safety /// @@ -360,6 +362,7 @@ impl Drop for StreamingRecordBatchReader { pub unsafe fn ffi_stream_to_sendable( ffi_stream: &mut FFI_ArrowArrayStream, cancel_checker: Option, + check_interval: Option, ) -> Result { let reader = arrow_array::ffi_stream::ArrowArrayStreamReader::from_raw(ffi_stream)?; let schema = reader.schema(); @@ -379,10 +382,26 @@ pub unsafe fn ffi_stream_to_sendable( } }); + // Track last check time for periodic cancellation checking + let mut last_check = Instant::now(); + let stream = tokio_stream::wrappers::ReceiverStream::new(rx).map(move |result| { - // Check for cancellation before yielding each batch + // Check for cancellation periodically (or every batch if no interval) if let Some(ref checker) = cancel_checker { - if checker() { + let should_check = match check_interval { + Some(interval) => { + let now = Instant::now(); + if now.duration_since(last_check) >= interval { + last_check = now; + true + } else { + false + } + } + None => true, // No interval means check every batch + }; + + if should_check && checker() { return exec_err!("Operation cancelled"); } } @@ -514,7 +533,7 @@ mod tests { let mut ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); // Import back - let imported = unsafe { ffi_stream_to_sendable(&mut ffi_stream, None).unwrap() }; + let imported = unsafe { ffi_stream_to_sendable(&mut ffi_stream, None, None).unwrap() }; // Collect results let batches: Vec<_> = runtime.block_on(imported.collect::>()); @@ -549,7 +568,7 @@ mod tests { Box::new(move || cancelled_clone.load(Ordering::SeqCst)); let imported = - unsafe { ffi_stream_to_sendable(&mut ffi_stream, Some(cancel_checker)).unwrap() }; + unsafe { ffi_stream_to_sendable(&mut ffi_stream, Some(cancel_checker), None).unwrap() }; // Collect with cancellation after 3 batches, stop on first error let batches = runtime.block_on(async { diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 6e421f3dc3..6bca238c00 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -21,6 +21,7 @@ use std::{ fmt::Debug, ptr::null_mut, sync::Arc, + time::Duration, }; use arrow_array::ffi::FFI_ArrowArray; @@ -285,6 +286,8 @@ pub struct ImportedTableProvider { table_type: TableType, /// Stored as Arc so it can be cloned for each scan execution. cancel_checker: Option bool + Send + Sync>>, + /// Interval for periodic cancellation checking during stream consumption. + check_interval: Option, } impl Debug for ImportedTableProvider { @@ -322,12 +325,14 @@ impl ImportedTableProvider { schema, table_type, cancel_checker: None, + check_interval: None, }) } /// Set a cancellation checker for this table provider. /// - /// The checker is called before each batch is read when scanning. + /// The checker is called periodically (controlled by `with_check_interval`) + /// before batches are read when scanning. /// If it returns `true`, the stream yields a cancellation error. pub fn with_cancel_checker(mut self, cancel_checker: F) -> Self where @@ -337,6 +342,16 @@ impl ImportedTableProvider { self } + /// Set the interval for periodic cancellation checking. + /// + /// When set, the cancel checker will only be called when this interval + /// has elapsed since the last check, reducing overhead for fast streams. + /// If not set, the checker is called before every batch. + pub fn with_check_interval(mut self, interval: Duration) -> Self { + self.check_interval = Some(interval); + self + } + fn get_table_type(provider: &SedonaCTableProvider) -> Result { let type_str = match get_table_provider_string_property(provider, "table_type") { Ok(s) => s, @@ -409,11 +424,14 @@ impl TableProvider for ImportedTableProvider { let mut exec = ImportedSedonaCExec::try_new(ffi_plan)?; - // Pipe through the cancel checker if one is configured + // Pipe through the cancel checker and interval if configured if let Some(ref checker) = self.cancel_checker { let checker = checker.clone(); exec = exec.with_cancel_checker(move || checker()); } + if let Some(interval) = self.check_interval { + exec = exec.with_check_interval(interval); + } Ok(Arc::new(exec)) } diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs index bbb404edb3..068e883ef1 100644 --- a/python/sedonadb/src/import_from.rs +++ b/python/sedonadb/src/import_from.rs @@ -17,6 +17,7 @@ use std::{ ffi::{c_void, CString}, sync::Arc, + time::Duration, }; use arrow_array::{ @@ -78,15 +79,18 @@ pub fn import_sedona_ffi_table_provider( let provider = ImportedTableProvider::try_new(ffi_provider)?; // Add a Python-aware cancel checker that checks for Ctrl+C signals - let provider = provider.with_cancel_checker(|| { - Python::attach(|py| { - // Run `pass` to process any pending signals, then check for errors - if py.run(cr"pass", None, None).is_err() { - return true; - } - py.check_signals().is_err() + // Use a 2 second interval to match the StreamingRecordBatchReader behavior + let provider = provider + .with_cancel_checker(|| { + Python::attach(|py| { + // Run `pass` to process any pending signals, then check for errors + if py.run(cr"pass", None, None).is_err() { + return true; + } + py.check_signals().is_err() + }) }) - }); + .with_check_interval(Duration::from_millis(2_000)); Ok(Arc::new(provider)) } From 7438f1355ea32d53e8769ff214be46e2791d05a9 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 9 Jul 2026 12:12:54 -0500 Subject: [PATCH 53/55] fix cancel checker call frequency --- c/sedona-extension/src/execution_plan.rs | 4 +- c/sedona-extension/src/streaming.rs | 98 ++++++++++++++---------- python/sedonadb/src/reader.rs | 10 ++- 3 files changed, 66 insertions(+), 46 deletions(-) diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs index df61fbd202..73969d69c1 100644 --- a/c/sedona-extension/src/execution_plan.rs +++ b/c/sedona-extension/src/execution_plan.rs @@ -156,6 +156,8 @@ impl ExportedExecutionPlan { } fn execute(&self, partition: usize) -> Result { + // Enter the runtime context so that plan.execute() can spawn tasks + let _guard = self.runtime.enter(); self.plan.execute(partition, self.task_context.clone()) } } @@ -985,7 +987,7 @@ mod tests { ]; assert_batches_eq!(expected, &batches); - // Execute partition 1 - needs a fresh import since we consumed the first + // Execute partition 1 let stream2 = imported.execute(1, task_ctx).unwrap(); let batches2: Vec = stream2 .collect::>() diff --git a/c/sedona-extension/src/streaming.rs b/c/sedona-extension/src/streaming.rs index 0edd430d99..9c1e1343a2 100644 --- a/c/sedona-extension/src/streaming.rs +++ b/c/sedona-extension/src/streaming.rs @@ -77,6 +77,18 @@ struct StreamWorker { /// /// Internally uses a dedicated worker thread to avoid blocking issues when /// called from within a tokio runtime context. +/// +/// # Runtime Requirements +/// +/// For responsive mid-batch cancellation, a **multi-thread runtime** is recommended. +/// On current_thread runtimes, `Handle::block_on` does not drive the time driver, +/// so the cancellation timer cannot fire while waiting for `stream.next()`. In this +/// case, cancellation only takes effect between batch fetches (via reader-side +/// `recv_timeout` and the cancel flag check before each fetch). +/// +/// If your use case requires responsive cancellation during long-running batch +/// fetches, ensure the runtime passed to [`StreamingRecordBatchReader::new`] is +/// created with [`tokio::runtime::Builder::new_multi_thread`]. pub struct StreamingRecordBatchReader { schema: SchemaRef, /// Stream and runtime, wrapped in Option so they can be moved to the worker. @@ -111,9 +123,13 @@ impl StreamingRecordBatchReader { /// Create a new StreamingRecordBatchReader with a cancellation checker. /// - /// The cancellation checker is called before each batch is fetched. If it - /// returns `true`, iteration stops with a cancellation error on the next - /// call and `None` on subsequent calls. + /// The cancellation checker is called periodically at `check_interval` while + /// waiting for batches to be fetched. If it returns `true`, iteration stops + /// with a cancellation error on the next call and `None` on subsequent calls. + /// + /// The `check_interval` controls how frequently the checker is called. This + /// is useful for Python where we need to periodically check for signals + /// (Ctrl+C) during long-running operations. /// /// Takes an `Arc` to ensure the runtime stays alive for the lifetime /// of the reader. @@ -121,6 +137,7 @@ impl StreamingRecordBatchReader { stream: SendableRecordBatchStream, runtime: Arc, cancel_checker: CancelChecker, + check_interval: Duration, ) -> Self { Self { schema: stream.schema(), @@ -129,7 +146,7 @@ impl StreamingRecordBatchReader { cancel_checker: Some(cancel_checker), cancelled: false, skip_empty_batches: false, - periodic_check_interval: None, + periodic_check_interval: Some(check_interval), } } @@ -142,19 +159,6 @@ impl StreamingRecordBatchReader { self } - /// Set a periodic interval for checking cancellation during batch fetches. - /// - /// When set, the cancel checker will be called periodically at this interval - /// even while waiting for a single batch to be fetched. This is useful for - /// Python where we need to periodically check for signals (Ctrl+C) during - /// long-running operations. - /// - /// Without this, the cancel checker is only called between batch fetches. - pub fn with_periodic_check_interval(mut self, interval: std::time::Duration) -> Self { - self.periodic_check_interval = Some(interval); - self - } - /// Ensure the worker thread is running, spawning it if necessary. fn ensure_worker(&mut self) { if self.worker.is_some() { @@ -302,15 +306,6 @@ impl Iterator for StreamingRecordBatchReader { return None; } - // Check for cancellation before fetching the next batch - // (periodic checking during fetch is handled by fetch_next_batch if configured) - if let Some(ref checker) = self.cancel_checker { - if checker() { - self.cancelled = true; - return Some(Err(cancellation_arrow_error())); - } - } - loop { match self.fetch_next_batch() { Some(Ok(batch)) => { @@ -418,7 +413,7 @@ mod tests { use super::*; use arrow_array::Int32Array; use arrow_schema::{DataType, Field, Schema}; - use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -488,25 +483,41 @@ mod tests { #[test] fn test_streaming_reader_cancel() { let runtime = test_runtime(); + // 10 batches, 50ms each = 500ms total let (_schema, stream) = create_slow_stream(10, 50); - // Cancel after reading 3 batches - let counter = Arc::new(AtomicUsize::new(0)); - let counter_clone = counter.clone(); - let cancel_checker: CancelChecker = Box::new(move || { - let count = counter_clone.fetch_add(1, Ordering::SeqCst); - count >= 3 + // Cancel after 175ms - should allow ~3 batches (at 50ms each) to complete + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_clone = cancelled.clone(); + let cancel_checker: CancelChecker = + Box::new(move || cancelled_clone.load(Ordering::SeqCst)); + + let reader = StreamingRecordBatchReader::with_cancel_checker( + stream, + runtime, + cancel_checker, + Duration::from_millis(25), + ); + + // Set cancel flag after 175ms from a separate thread + let cancelled_setter = cancelled.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(175)); + cancelled_setter.store(true, Ordering::SeqCst); }); - let reader = - StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker); let batches: Vec<_> = reader.collect(); - // Should have 3 successful batches + 1 cancellation error - assert_eq!(batches.len(), 4); + // Should have ~3 successful batches + 1 cancellation error + // (timing can vary slightly, so we check for at least 2 and at most 5) + assert!( + batches.len() >= 2 && batches.len() <= 5, + "expected 2-5 batches, got {}", + batches.len() + ); - // First 3 should be Ok - for (i, batch) in batches.iter().enumerate().take(3) { + // All but the last should be Ok + for (i, batch) in batches.iter().take(batches.len() - 1).enumerate() { assert!(batch.is_ok(), "batch {} should be Ok", i); } @@ -680,9 +691,12 @@ mod tests { let cancel_checker: CancelChecker = Box::new(move || cancelled_clone.load(Ordering::SeqCst)); - let reader = - StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker) - .with_periodic_check_interval(Duration::from_millis(50)); + let reader = StreamingRecordBatchReader::with_cancel_checker( + stream, + runtime, + cancel_checker, + Duration::from_millis(50), + ); // Spawn a task to set cancelled after 150ms let cancelled_setter = cancelled.clone(); diff --git a/python/sedonadb/src/reader.rs b/python/sedonadb/src/reader.rs index c1bb7c8217..37b8e2e147 100644 --- a/python/sedonadb/src/reader.rs +++ b/python/sedonadb/src/reader.rs @@ -49,7 +49,11 @@ pub fn new_py_streaming_reader( }) }); - StreamingRecordBatchReader::with_cancel_checker(stream, runtime, cancel_checker) - .with_skip_empty_batches(true) - .with_periodic_check_interval(INTERVAL_CHECK_SIGNALS) + StreamingRecordBatchReader::with_cancel_checker( + stream, + runtime, + cancel_checker, + INTERVAL_CHECK_SIGNALS, + ) + .with_skip_empty_batches(true) } From 3244dec61c7aae7d1275abcd5f35c6413c886df8 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 9 Jul 2026 12:22:56 -0500 Subject: [PATCH 54/55] don't discard panic error --- c/sedona-extension/src/table_provider.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs index 6bca238c00..98c29c21e1 100644 --- a/c/sedona-extension/src/table_provider.rs +++ b/c/sedona-extension/src/table_provider.rs @@ -98,9 +98,7 @@ impl ExportedTableProvider { .block_on(inner.scan(session.as_ref(), projection_ref, &[], limit)) }) .join() - .map_err(|_| { - datafusion_common::DataFusionError::Internal("Scan thread panicked".to_string()) - })? + .map_err(|e| sedona_internal_datafusion_err!("Scan thread panicked {e:?}"))? } fn get_property(&self, property: &str) -> Result { From e4e83a20e3af9b5b97117863ad7534ed74f419a2 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 9 Jul 2026 12:27:54 -0500 Subject: [PATCH 55/55] docs for safety --- c/sedona-extension/src/sedona_extension.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 843e09fc21..128b231e52 100644 --- a/c/sedona-extension/src/sedona_extension.h +++ b/c/sedona-extension/src/sedona_extension.h @@ -281,6 +281,8 @@ struct SedonaCExecutionPlanArgs { /// /// Before using this structure, the release callback MUST be checked. /// Instances with a NULL release callback are not valid and must not be used. +/// +/// Members that accept const* self must be thread safe. struct SedonaCExecutionPlan { /// \brief Get the schema associated with the output of this plan /// @@ -342,6 +344,8 @@ struct SedonaCExecutionPlan { /// /// Before using this structure, the release callback MUST be checked. /// Instances with a NULL release callback are not valid and must not be used. +/// +/// Members that accept const* self must be thread safe. struct SedonaCTableProvider { /// \brief Get the schema of this table provider ///