diff --git a/Cargo.lock b/Cargo.lock index 4468f46f5c..f9585845ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5670,6 +5670,7 @@ dependencies = [ "datafusion", "futures", "sedona", + "sedona-extension", "tokio", ] @@ -5758,13 +5759,24 @@ version = "0.4.0" 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", "sedona-schema", "sedona-testing", + "serde", + "serde_json", + "tokio", + "tokio-stream", ] [[package]] @@ -6382,6 +6394,7 @@ dependencies = [ "sedona-adbc", "sedona-datasource", "sedona-expr", + "sedona-extension", "sedona-gdal", "sedona-geometry", "sedona-geoparquet", @@ -6425,6 +6438,7 @@ dependencies = [ "sedona", "sedona-adbc", "sedona-expr", + "sedona-extension", "sedona-geometry", "sedona-geoparquet", "sedona-proj", @@ -7005,6 +7019,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 e8f8c17c56..4800913b39 100644 --- a/c/sedona-extension/Cargo.toml +++ b/c/sedona-extension/Cargo.toml @@ -30,10 +30,23 @@ 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-expr = { workspace = true } +datafusion-physical-plan = { workspace = true } +futures = { workspace = true } libc = "0.2.178" sedona-common = { workspace = true } sedona-expr = { workspace = true } sedona-schema = { workspace = true } 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 new file mode 100644 index 0000000000..73969d69c1 --- /dev/null +++ b/c/sedona-extension/src/execution_plan.rs @@ -0,0 +1,1011 @@ +// 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, + ffi::{c_int, c_void}, + fmt::{Debug, Display, Formatter}, + ptr::null_mut, + sync::Arc, + time::Duration, +}; + +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::{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 tokio::runtime::Runtime; + +use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs}; +use crate::set_ffi_error; +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. +/// +/// 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: Arc, +} + +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. + /// + /// 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: Arc, + ) -> 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::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)), + "display_default" => { + use std::fmt::Write; + let mut s = String::new(); + 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) + } + "display_tree_render" => { + use std::fmt::Write; + let mut s = String::new(); + 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()) + } + } + _ => exec_err!("Unknown property: {}", property), + } + } + + 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()) + } +} + +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: Some(c_exec_plan_get_property_schema), + 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, + 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_; + 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(); + match FFI_ArrowSchema::try_from(schema.as_ref()) { + Ok(ffi_schema) => { + std::ptr::write(out, ffi_schema); + ERRNO_OK + } + Err(e) => { + set_ffi_error!(err, "Failed to convert schema to FFI: {}", e); + libc::EINVAL + } + } +} + +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 { + 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); + match FFI_ArrowSchema::try_from(&field) { + Ok(ffi_schema) => { + std::ptr::write(out, ffi_schema); + ERRNO_OK + } + Err(e) => { + set_ffi_error!(err, "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, + _args: *mut SedonaCExecutionPlanArgs, + out: *mut arrow_array::ffi::FFI_ArrowArray, + 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_; + 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) => { + // 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); + ERRNO_OK + } + Err(e) => { + set_ffi_error!(err, "{}", e); + 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 { + 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 { + 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) => { + set_ffi_error!(err, "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) => { + set_ffi_error!(err, "{}", e); + libc::EINVAL + } + } +} + +unsafe extern "C" fn c_exec_plan_release(self_: *mut SedonaCExecutionPlan) { + 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_ref.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, + /// 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 { + 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("ImportedSedonaCExec") + .field("inner", &debug_str) + .finish() + } else { + f.debug_struct("ImportedSedonaCExec").finish() + } + } +} + +impl ImportedSedonaCExec { + /// Create a new ImportedSedonaCExec from a SedonaCExecutionPlan. + /// + /// 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"); + }; + + let mut ffi_schema = FFI_ArrowSchema::empty(); + 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 + 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, + cancel_checker: None, + check_interval: None, + }) + } + + /// Set a cancellation checker for this execution plan. + /// + /// 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 + F: Fn() -> bool + Send + Sync + 'static, + { + self.cancel_checker = Some(Arc::new(cancel_checker)); + 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") + } +} + +impl DisplayAs for ImportedSedonaCExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + let property = match t { + DisplayFormatType::Default => "display_default", + DisplayFormatType::Verbose => "display_verbose", + 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, "ImportedSedonaCExec: {}", display_str) + } else { + write!(f, "ImportedSedonaCExec") + } + } +} + +impl ExecutionPlan for ImportedSedonaCExec { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + &self.name + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn partition_statistics(&self, _partition: Option) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + 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 { + // 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 { + get_plan_property::, ()>(&self.inner, "metrics", None) + .ok() + .flatten() + .map(|sm| sm.into_metrics_set()) + } + + fn supports_limit_pushdown(&self) -> bool { + self.supports_limit_pushdown + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + 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 + // 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, self.check_interval) } + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if !children.is_empty() { + return exec_err!("ImportedSedonaCExec does not support children"); + } + Ok(self) + } +} + +/// 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)]), + ))) + } + } + + /// 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. + /// 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 = test_runtime(); + let task_ctx = Arc::new(TaskContext::default()); + + 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, runtime) + } + + fn setup_imported_plan_with( + emission_type: EmissionType, + boundedness: Boundedness, + supports_limit_pushdown: bool, + ) -> (ImportedSedonaCExec, Arc, Arc) { + let dummy = Arc::new(DummyExec::with_properties( + emission_type, + boundedness, + supports_limit_pushdown, + )); + let runtime = test_runtime(); + let task_ctx = Arc::new(TaskContext::default()); + + 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, runtime) + } + + #[test] + fn test_execution_plan_roundtrip_schema() { + let (imported, _, _runtime) = 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"); + } + + #[test] + fn test_execution_plan_roundtrip_name() { + let (imported, _, _runtime) = setup_imported_plan(); + assert_eq!(imported.name(), "ImportedSedonaCExec"); + } + + #[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, _, _runtime) = + 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, _, _runtime) = + 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, _, _runtime) = + setup_imported_plan_with(EmissionType::Incremental, Boundedness::Bounded, true); + assert!(imported_with.supports_limit_pushdown()); + + let (imported_without, _, _runtime) = + setup_imported_plan_with(EmissionType::Incremental, Boundedness::Bounded, false); + assert!(!imported_without.supports_limit_pushdown()); + + // Test partition count + let (imported, _, _runtime) = setup_imported_plan(); + assert_eq!( + imported + .properties() + .output_partitioning() + .partition_count(), + 3 + ); + } + + #[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); + impl std::fmt::Display for DisplayAsFormat<'_, T> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.0.fmt_as(self.1, f) + } + } + + // ImportedSedonaCExec shows itself with the inner plan's display + assert_eq!( + format!("{}", DisplayAsFormat(&imported, DisplayFormatType::Default)), + "ImportedSedonaCExec: DummyExec: default format" + ); + assert_eq!( + format!("{}", DisplayAsFormat(&imported, DisplayFormatType::Verbose)), + "ImportedSedonaCExec: DummyExec: verbose format with schema" + ); + assert_eq!( + format!( + "{}", + DisplayAsFormat(&imported, DisplayFormatType::TreeRender) + ), + "ImportedSedonaCExec: DummyExec: tree render format" + ); + } + + #[test] + fn test_execution_plan_roundtrip_debug_string() { + let (imported, _, _runtime) = setup_imported_plan(); + + let debug_str = imported.get_debug_string().unwrap(); + assert!( + debug_str.contains("DummyExec"), + "debug_string should contain 'DummyExec', got: {}", + debug_str + ); + } + + #[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 + 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/extension.rs b/c/sedona-extension/src/extension.rs index 7d957a57c2..287023f324 100644 --- a/c/sedona-extension/src/extension.rs +++ b/c/sedona-extension/src/extension.rs @@ -21,7 +21,10 @@ use std::{ 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,368 @@ 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 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)) + } + } +} + +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) { + // 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(), + err_len: "Unknown error".len() as u32, + reserved: 0, + 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)] +pub struct SedonaCExpr { + pub get_property_schema: Option< + unsafe extern "C" fn( + self_: *const SedonaCExpr, + property: *const c_char, + out: *mut FFI_ArrowSchema, + 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 +/// +/// 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, + /// 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, +} + +/// Raw FFI representation of the SedonaCExecutionPlan +#[derive(Default)] +#[repr(C)] +pub struct SedonaCExecutionPlan { + 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( + 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(); + } + } +} + +/// 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< + 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< + 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 073ed939bc..de9835521f 100644 --- a/c/sedona-extension/src/lib.rs +++ b/c/sedona-extension/src/lib.rs @@ -15,5 +15,9 @@ // specific language governing permissions and limitations // under the License. +pub mod execution_plan; pub mod extension; pub mod scalar_kernel; +pub mod streaming; +pub mod table_provider; +pub mod utils; diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h index 7191af21f6..128b231e52 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,216 @@ 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 + /// + /// 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 + /// + /// Implementations of this callback must set self->release to NULL. + void (*release)(struct SedonaCError* self); +}; + +struct SedonaCExpr { + /// \brief Get the data type of a property + int (*get_property_schema)(const struct SedonaCExpr* self, const char* property, + struct ArrowSchema* out, struct SedonaCError* err); + + /// \brief Extract a serializable property from this expression + /// + /// This is used to implement PlanProperties and other values that can be + /// 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); + + void* reserved; + + /// \brief Release this instance + /// + /// Implementations of this callback must set self->release to NULL. + void (*release)(struct SedonaCExpr* self); + + /// \brief Opaque implementation-specific data + void* private_data; +}; + +/// Forward declaration of the execution plan +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 { + /// \brief JSON-serialized arguments + const uint8_t* args; + size_t args_len; + + /// \brief Optional array of execution plans + const struct SedonaCExecutionPlan** exec_plans; + size_t num_exec_plans; + + /// \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. +/// +/// Members that accept const* self must be thread safe. +struct SedonaCExecutionPlan { + /// \brief Get the schema associated with the output of this plan + /// + /// 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, + 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 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, + struct SedonaCError* err); + + /// \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); + + /// \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); + + /// \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); + + /// \brief Reserved for future use (must be NULL). + 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; +}; + +/// \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. +/// +/// 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 + /// + /// 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, + 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 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, + struct SedonaCError* err); + + /// \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. + int (*scan)(const struct SedonaCTableProvider* self, + struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, + struct SedonaCError* err); + + /// \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. + /// Returns an execution plan that performs the insert. + int (*insert)(const struct SedonaCTableProvider* self, + struct SedonaCExecutionPlanArgs* args, struct SedonaCExecutionPlan* out, + struct SedonaCError* err); + + /// \brief 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); + + /// \brief 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); + + /// \brief Reserved for future use. Must be NULL. + 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/streaming.rs b/c/sedona-extension/src/streaming.rs new file mode 100644 index 0000000000..9c1e1343a2 --- /dev/null +++ b/c/sedona-extension/src/streaming.rs @@ -0,0 +1,730 @@ +// 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 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}; +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. +/// +/// 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>; + +/// 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>; + +/// 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 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, + /// 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. +/// +/// # 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. + /// 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, + cancelled: bool, + skip_empty_batches: bool, + periodic_check_interval: Option, +} + +impl StreamingRecordBatchReader { + /// Create a new StreamingRecordBatchReader from a SendableRecordBatchStream. + /// + /// 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)), + 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 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. + pub fn with_cancel_checker( + stream: SendableRecordBatchStream, + runtime: Arc, + cancel_checker: CancelChecker, + check_interval: Duration, + ) -> 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: Some(check_interval), + } + } + + /// 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 + } + + /// 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 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 - eagerly fetches batches into the buffer + let handle = std::thread::spawn(move || { + let mut stream = stream; + + loop { + // Check if cancelled before starting fetch + if worker_cancel_flag.load(Ordering::Relaxed) { + let _ = batch_tx.send(Some(Err(cancellation_arrow_error()))); + break; + } + + let result = runtime.handle().block_on(async { + // 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, + } + } + } + }); + + let is_end = result.is_none(); + let is_cancelled = worker_cancel_flag.load(Ordering::Relaxed); + + // 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 { + batch_rx, + cancel_flag, + _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(), + ))); + }; + + // Receive from the prefetch buffer, with optional periodic cancellation checking + match &self.periodic_check_interval { + Some(interval) => { + let interval = *interval; + loop { + match worker.batch_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() { + 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 + .batch_rx + .recv_timeout(std::time::Duration::from_millis(200)); + return Some(Err(cancellation_arrow_error())); + } + } + // Continue waiting + } + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + return Some(Err(ArrowError::InvalidArgumentError( + "Worker thread terminated".to_string(), + ))); + } + } + } + } + None => { + // Simple blocking receive from prefetch buffer + match worker.batch_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; + } + + 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)) => { + return Some(Err(e)); + } + None => return None, + } + } + } +} + +impl RecordBatchReader for StreamingRecordBatchReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +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 +/// 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 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 +/// +/// 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, + cancel_checker: Option, + check_interval: Option, +) -> Result { + let reader = arrow_array::ffi_stream::ArrowArrayStreamReader::from_raw(ffi_stream)?; + let schema = reader.schema(); + + // 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 + } + } + }); + + // 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 periodically (or every batch if no interval) + if let Some(ref checker) = cancel_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"); + } + } + 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, Ordering}; + 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, + 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), + ), + ) + } + + #[test] + fn test_streaming_reader_basic() { + let runtime = test_runtime(); + 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); + } + } + + #[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 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 batches: Vec<_> = reader.collect(); + + // 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() + ); + + // 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); + } + + // 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 + ); + } + + #[test] + fn test_ffi_stream_to_sendable_basic() { + use arrow_array::ffi_stream::FFI_ArrowArrayStream; + + let runtime = test_runtime(); + 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, None, None).unwrap() }; + + // Collect results + let batches: Vec<_> = runtime.block_on(imported.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); + } + } + + #[test] + fn test_ffi_stream_to_sendable_cancel() { + use arrow_array::ffi_stream::FFI_ArrowArrayStream; + + let runtime = test_runtime(); + 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(&mut ffi_stream, Some(cancel_checker), None).unwrap() }; + + // Collect with cancellation after 3 batches, 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); + + // 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), + ), + ) + } + + #[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]); + + 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); + } + + #[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]); + + // 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); + } + + #[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); + + // 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, + 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/table_provider.rs b/c/sedona-extension/src/table_provider.rs new file mode 100644 index 0000000000..98c29c21e1 --- /dev/null +++ b/c/sedona-extension/src/table_provider.rs @@ -0,0 +1,748 @@ +// 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, + ffi::{c_int, c_void}, + fmt::Debug, + ptr::null_mut, + sync::Arc, + time::Duration, +}; + +use arrow_array::ffi::FFI_ArrowArray; +use arrow_schema::{ffi::FFI_ArrowSchema, Schema, SchemaRef}; +use async_trait::async_trait; +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}; +use tokio::runtime::Runtime; + +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. +/// +/// This wraps an inner TableProvider and exposes it via the SedonaCTableProvider +/// FFI interface. +pub struct ExportedTableProvider { + inner: Arc, + session: Arc, + /// We hold Arc instead of just Handle to keep the runtime alive + /// as long as this provider exists. + runtime: Arc, +} + +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. + /// + /// 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: Arc, + ) -> Self { + Self { + inner, + session, + runtime, + } + } + + fn scan( + &self, + projection: Option>, + limit: Option, + ) -> Result> { + let inner = self.inner.clone(); + let session = self.session.clone(); + let runtime = self.runtime.clone(); + + std::thread::spawn(move || { + let projection_ref = projection.as_ref(); + runtime + .handle() + .block_on(inner.scan(session.as_ref(), projection_ref, &[], limit)) + }) + .join() + .map_err(|e| sedona_internal_datafusion_err!("Scan thread panicked {e:?}"))? + } + + 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.to_string()) + } + _ => 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: Some(c_table_provider_get_property_schema), + 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, + 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_; + 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(); + match FFI_ArrowSchema::try_from(schema.as_ref()) { + Ok(ffi_schema) => { + std::ptr::write(out, ffi_schema); + ERRNO_OK + } + Err(e) => { + set_ffi_error!(err, "Failed to convert schema to FFI: {}", e); + libc::EINVAL + } + } +} + +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 { + 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); + match FFI_ArrowSchema::try_from(&field) { + Ok(ffi_schema) => { + std::ptr::write(out, ffi_schema); + ERRNO_OK + } + Err(e) => { + set_ffi_error!(err, "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, + _args: *mut SedonaCExecutionPlanArgs, + out: *mut FFI_ArrowArray, + 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_; + 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) => { + // 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); + ERRNO_OK + } + Err(e) => { + set_ffi_error!(err, "{}", e); + libc::EINVAL + } + } +} + +unsafe extern "C" fn c_table_provider_scan( + self_: *const SedonaCTableProvider, + args: *mut SedonaCExecutionPlanArgs, + out: *mut SedonaCExecutionPlan, + err: *mut SedonaCError, +) -> c_int { + 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_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) => { + set_ffi_error!(err, "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) => { + set_ffi_error!(err, "{}", e); + libc::EINVAL + } + } +} + +unsafe extern "C" fn c_table_provider_release(self_: *mut SedonaCTableProvider) { + 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_ref.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, + /// 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 { + 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 { + // 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"); + }; + + let mut ffi_schema = FFI_ArrowSchema::empty(); + 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 + let table_type = Self::get_table_type(&inner)?; + + Ok(Self { + inner, + schema, + table_type, + cancel_checker: None, + check_interval: None, + }) + } + + /// Set a cancellation checker for this table provider. + /// + /// 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 + F: Fn() -> bool + Send + Sync + 'static, + { + self.cancel_checker = Some(Arc::new(cancel_checker)); + 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, + 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] +impl TableProvider for ImportedTableProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + self.table_type + } + + fn statistics(&self) -> Option { + None + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + limit: Option, + ) -> Result> { + 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 mut exec = ImportedSedonaCExec::try_new(ffi_plan)?; + + // 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)) + } +} + +/// 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::*; + 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 { + 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) + } + + /// 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 + 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(()) + }) + } + + #[test] + fn test_roundtrip_simple_select() { + 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 |", + "+----+---------+", + ], + ) + .unwrap(); + } + + #[test] + fn test_roundtrip_projection() { + 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 |", + "+---------+---------+", + ], + ) + .unwrap(); + } + + #[test] + fn test_roundtrip_filter() { + 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 |", + "+----+---------+", + ], + ) + .unwrap(); + } + + #[test] + fn test_roundtrip_sort() { + 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 |", + "+----+---------+", + ], + ) + .unwrap(); + } + + #[test] + fn test_roundtrip_limit() { + test_roundtrip_query( + "SELECT id FROM imported_data ORDER BY id LIMIT 3", + &[ + "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 3 |", "+----+", + ], + ) + .unwrap(); + } + + #[test] + fn test_roundtrip_all_columns() { + 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 |", + "+----+---------+---------+---------+---------+", + ], + ) + .unwrap(); + } + + /// 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") + } + } + + /// 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. + /// 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 = 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"); + (imported, runtime) + } + + #[test] + fn test_table_provider_roundtrip_schema() { + let (imported, _runtime) = 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); + } + + #[test] + fn test_table_provider_roundtrip_table_type() { + for table_type in [TableType::Base, TableType::View, TableType::Temporary] { + let (imported, _runtime) = setup_imported_provider_with(table_type); + assert_eq!(imported.table_type(), table_type); + } + } +} diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs new file mode 100644 index 0000000000..ae42851bb1 --- /dev/null +++ b/c/sedona-extension/src/utils.rs @@ -0,0 +1,340 @@ +// 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 FFI property access and conversion. + +use std::borrow::Cow; +use std::ffi::{c_char, c_int, CStr, CString}; +use std::ptr::null_mut; + +use arrow_schema::ffi::FFI_ArrowSchema; +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; + +use crate::extension::{ + SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs, SedonaCTableProvider, +}; + +/// 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`. +pub unsafe fn cstr_from_ptr_or_empty<'a>(ptr: *const c_char) -> Cow<'a, str> { + if ptr.is_null() { + Cow::Borrowed("") + } else { + CStr::from_ptr(ptr).to_string_lossy() + } +} + +/// 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"); + }; + + 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. +/// +/// 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); + } + + // 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) +} + +/// 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, + data_type: &DataType, +) -> Result { + 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 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, 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 => { + 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]. +/// +/// 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"); + }; + + 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]. +/// +/// 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 { + return Ok(DataType::Binary); + }; + + call_get_property_schema_impl(property, |prop, schema, err| unsafe { + get_property_schema(provider, prop, schema, err) + }) +} diff --git a/python/sedonadb/Cargo.toml b/python/sedonadb/Cargo.toml index 5306621525..3e82af99f7 100644 --- a/python/sedonadb/Cargo.toml +++ b/python/sedonadb/Cargo.toml @@ -50,6 +50,7 @@ pyo3 = { workspace = true } 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 f65f2b35f4..dc759d5c8d 100644 --- a/python/sedonadb/python/sedonadb/dataframe.py +++ b/python/sedonadb/python/sedonadb/dataframe.py @@ -1384,8 +1384,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__(self._ctx._impl) def to_arrow_table(self, schema: Any = None) -> "pa.Table": """Execute and collect results as a PyArrow Table @@ -1770,7 +1770,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 @@ -1783,7 +1783,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 6269837f3b..81e2f969cd 100644 --- a/python/sedonadb/src/dataframe.rs +++ b/python/sedonadb/src/dataframe.rs @@ -24,11 +24,9 @@ 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_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::*; @@ -44,7 +42,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; @@ -451,7 +449,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.clone()); let mut reader: Box = Box::new(reader); if simplify.unwrap_or(false) { @@ -666,23 +664,25 @@ impl InternalDataFrame { Ok(InternalDataFrame::new(df, self.runtime.clone())) } - fn __datafusion_table_provider__<'py>( + fn __sedonadb_table_provider__<'py>( &self, py: Python<'py>, + ctx: &InternalContext, ) -> Result, PySedonaError> { let provider = self.inner.clone().into_view(); - 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 = sedona_extension::table_provider::ExportedTableProvider::new( provider, - true, - Some(self.runtime.handle().clone()), - &ctx, - None, + session, + self.runtime.clone(), ); + let ffi_provider: sedona_extension::extension::SedonaCTableProvider = exported.into(); Ok(PyCapsule::new_with_value( py, ffi_provider, - c"datafusion_table_provider", + c"sedonadb_table_provider", )?) } } diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs index 52943337e9..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::{ @@ -28,12 +29,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 +47,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 +60,39 @@ 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. + // Clear the structure after reading to prevent double-free when the capsule is dropped. + let ffi_provider = unsafe { + let provider = std::ptr::read(contents); + // 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)?; + + // Add a Python-aware cancel checker that checks for Ctrl+C signals + // 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)) } pub fn import_arrow_array_stream<'py>( diff --git a/python/sedonadb/src/reader.rs b/python/sedonadb/src/reader.rs index 0d8ec2c7da..37b8e2e147 100644 --- a/python/sedonadb/src/reader.rs +++ b/python/sedonadb/src/reader.rs @@ -14,53 +14,46 @@ // 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 pyo3::Python; +use sedona_extension::streaming::StreamingRecordBatchReader; use tokio::runtime::Runtime; -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. +/// +/// 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, -} - -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)) => { - let batch = maybe_batch?; - if batch.num_rows() == 0 { - continue; - } - return Some(Ok(batch)); - } - 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: Arc, +) -> StreamingRecordBatchReader { + // Create a cancel checker that checks Python signals + let cancel_checker: Box bool + Send + Sync> = Box::new(|| { + 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() + }) + }); -impl RecordBatchReader for PySedonaStreamReader { - fn schema(&self) -> SchemaRef { - self.stream.schema() - } + StreamingRecordBatchReader::with_cancel_checker( + stream, + runtime, + cancel_checker, + INTERVAL_CHECK_SIGNALS, + ) + .with_skip_empty_batches(true) } diff --git a/python/sedonadb/src/runtime.rs b/python/sedonadb/src/runtime.rs index 77af227606..38c9bc58c2 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::attach(|py| { - py.run(cr"pass", None, None)?; - py.check_signals() - })?; - } - } - } - }) -} 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 diff --git a/python/sedonadb/tests/test_dataframe_ffi.py b/python/sedonadb/tests/test_dataframe_ffi.py new file mode 100644 index 0000000000..7b230a83aa --- /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) +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) 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/Cargo.toml b/r/sedonadb/src/rust/Cargo.toml index ab7ae42549..a5eb3dea61 100644 --- a/r/sedonadb/src/rust/Cargo.toml +++ b/r/sedonadb/src/rust/Cargo.toml @@ -25,8 +25,8 @@ edition = "2021" crate-type = ["staticlib", "lib"] [dependencies] -arrow-schema = { workspace = true } arrow-array = { workspace = true } +arrow-schema = { workspace = true } datafusion = { workspace = true } datafusion-common = { workspace = true } datafusion-execution = { workspace = true } @@ -37,6 +37,7 @@ savvy-ffi = "*" sedona = { workspace = true } sedona-adbc = { workspace = true } sedona-expr = { workspace = true } +sedona-extension = { workspace = true } sedona-geometry = { workspace = true } sedona-geoparquet = { workspace = true } sedona-proj = { workspace = true } 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 20af9bf5bc..4e643b2172 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::reader::SedonaStreamReader; 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] @@ -115,7 +114,8 @@ impl InternalDataFrame { async move { inner.execute_stream().await }, )??; - let reader = SedonaStreamReader::new(self.runtime.clone(), stream); + let reader = StreamingRecordBatchReader::new(stream, self.runtime.clone()) + .with_skip_empty_batches(true); let reader: Box = Box::new(reader); let mut ffi_stream = FFI_ArrowArrayStream::new(reader); @@ -125,22 +125,17 @@ 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( - provider, - true, - Some(self.runtime.handle().clone()), - &ctx, - None, - ); - - let mut ffi_xptr = FFITableProviderR(ffi_provider).into_external_pointer(); + // 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.clone()); + let ffi_provider: sedona_extension::extension::SedonaCTableProvider = exported.into(); + + 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), 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 aeee806ad6..3c7592107b 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::streaming::StreamingRecordBatchReader; use std::sync::Arc; use tokio::runtime::Runtime; @@ -98,10 +99,9 @@ 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.clone()) + .with_skip_empty_batches(true); + 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 a030ea96d0..0000000000 --- a/rust/sedona/src/reader.rs +++ /dev/null @@ -1,123 +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) => { - let batch = maybe_batch?; - if batch.num_rows() == 0 { - continue; - } - return Some(Ok(batch)); - } - 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()); - } -}