You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/13 10:05:22 UTC
(datafusion) branch gh-readonly-queue/main/pr-23906-186f96fba6b023c8e7310314bb891124e915abc7 created (now ab12f5e4be)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23906-186f96fba6b023c8e7310314bb891124e915abc7
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at ab12f5e4be fix(ffi): preserve TableProvider DML overrides (#23906)
This branch includes the following new commits:
new ab12f5e4be fix(ffi): preserve TableProvider DML overrides (#23906)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix(ffi): preserve TableProvider DML overrides (#23906)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23906-186f96fba6b023c8e7310314bb891124e915abc7
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit ab12f5e4be594e7287602dd0d60381d971a9389a
Author: Amogh Ramesh <ra...@gmail.com>
AuthorDate: Thu Aug 13 10:04:58 2026 +0000
fix(ffi): preserve TableProvider DML overrides (#23906)
## Which issue does this PR close?
Part of #22328.
## Rationale for this change
`ForeignTableProvider` inherits the default DML methods, so producer
overrides are lost across the FFI boundary and writable tables appear
read-only.
## What changes are included in this PR?
- Forward `delete_from`, `update`, and `truncate` through
`FFI_TableProvider`.
- Preserve ordered filters and assignments across the boundary.
- Move expression-list byte encoding and decoding into reusable
`datafusion-proto` helpers.
- Cover forced-foreign and dynamic-library paths, including empty
filters and malformed update payloads.
## Are these changes tested?
- `cargo test -p datafusion-ffi --features integration-tests`
- `cargo clippy --all-targets --all-features -- -D warnings`
## Are there any user-facing changes?
The FFI ABI changes. Foreign libraries must rebuild against the new
DataFusion version.
---
datafusion/ffi/src/table_provider.rs | 530 +++++++++++++++++++++++++-----
datafusion/ffi/src/tests/mod.rs | 68 +++-
datafusion/ffi/src/udtf.rs | 41 +--
datafusion/ffi/tests/ffi_integration.rs | 42 +++
datafusion/proto/src/bytes/mod.rs | 33 +-
datafusion/proto/tests/cases/serialize.rs | 61 +++-
6 files changed, 666 insertions(+), 109 deletions(-)
diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs
index ee9377bff0..f1a79added 100644
--- a/datafusion/ffi/src/table_provider.rs
+++ b/datafusion/ffi/src/table_provider.rs
@@ -28,14 +28,15 @@ use datafusion_execution::TaskContext;
use datafusion_expr::dml::InsertOp;
use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType};
use datafusion_physical_plan::ExecutionPlan;
-use datafusion_proto::logical_plan::from_proto::parse_exprs;
-use datafusion_proto::logical_plan::to_proto::serialize_exprs;
+use datafusion_proto::bytes::{
+ logical_exprs_from_bytes_with_extension_codec,
+ logical_exprs_to_bytes_with_extension_codec,
+};
use datafusion_proto::logical_plan::{
DefaultLogicalExtensionCodec, LogicalExtensionCodec,
};
-use datafusion_proto::protobuf::LogicalExprList;
-use prost::Message;
+use stabby::string::String as SString;
use stabby::vec::Vec as SVec;
use tokio::runtime::Handle;
@@ -103,7 +104,7 @@ pub struct FFI_TableProvider {
/// * `session` - session
/// * `projections` - if specified, only a subset of the columns are returned
/// * `filters_serialized` - filters to apply to the scan, which are a
- /// [`LogicalExprList`] protobuf message serialized into bytes to pass
+ /// [`LogicalExprList`][datafusion_proto::protobuf::LogicalExprList] protobuf message serialized into bytes to pass
/// across the FFI boundary.
/// * `limit` - if specified, limit the number of rows returned
scan: unsafe extern "C" fn(
@@ -118,7 +119,7 @@ pub struct FFI_TableProvider {
table_type: unsafe extern "C" fn(provider: &Self) -> FFI_TableType,
/// Based upon the input filters, identify which are supported. The filters
- /// are a [`LogicalExprList`] protobuf message serialized into bytes to pass
+ /// are a [`LogicalExprList`][datafusion_proto::protobuf::LogicalExprList] protobuf message serialized into bytes to pass
/// across the FFI boundary.
supports_filters_pushdown: Option<
unsafe extern "C" fn(
@@ -160,11 +161,36 @@ pub struct FFI_TableProvider {
/// the foreign interface. See [`crate::get_library_marker_id`] and
/// the crate's `README.md` for more information.
pub library_marker_id: extern "C" fn() -> usize,
+
+ delete_from: unsafe extern "C" fn(
+ provider: &Self,
+ session: FFI_SessionRef,
+ filters_serialized: SVec<u8>,
+ ) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>>,
+
+ update: unsafe extern "C" fn(
+ provider: &Self,
+ session: FFI_SessionRef,
+ assignments: SVec<FFI_TableProviderUpdateAssignment>,
+ filters_serialized: SVec<u8>,
+ ) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>>,
+
+ truncate: unsafe extern "C" fn(
+ provider: &Self,
+ session: FFI_SessionRef,
+ ) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>>,
}
unsafe impl Send for FFI_TableProvider {}
unsafe impl Sync for FFI_TableProvider {}
+#[repr(C)]
+#[derive(Debug)]
+struct FFI_TableProviderUpdateAssignment {
+ column: SString,
+ expr_serialized: SVec<u8>,
+}
+
struct ProviderPrivateData {
provider: Arc<dyn TableProvider>,
runtime: Option<Handle>,
@@ -202,21 +228,33 @@ unsafe extern "C" fn table_type_fn_wrapper(
provider.inner().table_type().into()
}
+fn parse_serialized_exprs(
+ exprs_serialized: &[u8],
+ task_ctx: &Arc<TaskContext>,
+ codec: &dyn LogicalExtensionCodec,
+) -> Result<Vec<Expr>> {
+ logical_exprs_from_bytes_with_extension_codec(
+ exprs_serialized,
+ task_ctx.as_ref(),
+ codec,
+ )
+}
+
+fn serialize_expr_list<'a>(
+ exprs: impl IntoIterator<Item = &'a Expr>,
+ codec: &dyn LogicalExtensionCodec,
+) -> Result<SVec<u8>> {
+ let bytes = logical_exprs_to_bytes_with_extension_codec(exprs, codec)?;
+ Ok(SVec::from(bytes.as_ref()))
+}
+
fn supports_filters_pushdown_internal(
provider: &Arc<dyn TableProvider>,
filters_serialized: &[u8],
task_ctx: &Arc<TaskContext>,
codec: &dyn LogicalExtensionCodec,
) -> Result<SVec<FFI_TableProviderFilterPushDown>> {
- let filters = match filters_serialized.is_empty() {
- true => vec![],
- false => {
- let proto_filters = LogicalExprList::decode(filters_serialized)
- .map_err(|e| DataFusionError::Plan(e.to_string()))?;
-
- parse_exprs(proto_filters.expr.iter(), task_ctx.as_ref(), codec)?
- }
- };
+ let filters = parse_serialized_exprs(filters_serialized, task_ctx, codec)?;
let filters_borrowed: Vec<&Expr> = filters.iter().collect();
let results: SVec<_> = provider
@@ -271,19 +309,11 @@ unsafe extern "C" fn scan_fn_wrapper(
);
let task_ctx = sresult_return!(task_ctx);
- let filters = match filters_serialized.is_empty() {
- true => vec![],
- false => {
- let proto_filters =
- sresult_return!(LogicalExprList::decode(filters_serialized.as_ref()));
-
- sresult_return!(parse_exprs(
- proto_filters.expr.iter(),
- task_ctx.as_ref(),
- logical_codec.as_ref(),
- ))
- }
- };
+ let filters = sresult_return!(parse_serialized_exprs(
+ &filters_serialized,
+ &task_ctx,
+ logical_codec.as_ref(),
+ ));
let projections: Option<Vec<usize>> =
projections.into_option().map(|p| p.into_iter().collect());
@@ -336,6 +366,135 @@ unsafe extern "C" fn insert_into_fn_wrapper(
.into_ffi()
}
+unsafe extern "C" fn delete_from_fn_wrapper(
+ provider: &FFI_TableProvider,
+ session: FFI_SessionRef,
+ filters_serialized: SVec<u8>,
+) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>> {
+ let task_ctx: Result<Arc<TaskContext>, DataFusionError> =
+ (&provider.logical_codec.task_ctx_provider).try_into();
+ let runtime = provider.runtime().clone();
+ let logical_codec: Arc<dyn LogicalExtensionCodec> = (&provider.logical_codec).into();
+ let internal_provider = Arc::clone(provider.inner());
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&dyn Session, DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session = Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let task_ctx = sresult_return!(task_ctx);
+ let filters = sresult_return!(parse_serialized_exprs(
+ &filters_serialized,
+ &task_ctx,
+ logical_codec.as_ref(),
+ ));
+
+ let plan = sresult_return!(internal_provider.delete_from(session, filters).await);
+
+ FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime))
+ }
+ .into_ffi()
+}
+
+unsafe extern "C" fn update_fn_wrapper(
+ provider: &FFI_TableProvider,
+ session: FFI_SessionRef,
+ assignments: SVec<FFI_TableProviderUpdateAssignment>,
+ filters_serialized: SVec<u8>,
+) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>> {
+ let task_ctx: Result<Arc<TaskContext>, DataFusionError> =
+ (&provider.logical_codec.task_ctx_provider).try_into();
+ let runtime = provider.runtime().clone();
+ let logical_codec: Arc<dyn LogicalExtensionCodec> = (&provider.logical_codec).into();
+ let internal_provider = Arc::clone(provider.inner());
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&dyn Session, DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session = Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let task_ctx = sresult_return!(task_ctx);
+ let assignments = sresult_return!(
+ assignments
+ .into_iter()
+ .map(|assignment| {
+ let column = assignment.column.to_string();
+ let mut exprs = parse_serialized_exprs(
+ &assignment.expr_serialized,
+ &task_ctx,
+ logical_codec.as_ref(),
+ )?;
+ let expression_count = exprs.len();
+ let expr = match expression_count {
+ 1 => exprs.remove(0),
+ _ => {
+ return Err(DataFusionError::Plan(format!(
+ "Expected exactly one expression for update assignment to column \
+ '{column}', got {expression_count}"
+ )));
+ }
+ };
+ Ok((column, expr))
+ })
+ .collect::<Result<Vec<_>>>()
+ );
+ let filters = sresult_return!(parse_serialized_exprs(
+ &filters_serialized,
+ &task_ctx,
+ logical_codec.as_ref(),
+ ));
+
+ let plan = sresult_return!(
+ internal_provider
+ .update(session, assignments, filters)
+ .await
+ );
+
+ FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime))
+ }
+ .into_ffi()
+}
+
+unsafe extern "C" fn truncate_fn_wrapper(
+ provider: &FFI_TableProvider,
+ session: FFI_SessionRef,
+) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>> {
+ let runtime = provider.runtime().clone();
+ let internal_provider = Arc::clone(provider.inner());
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&dyn Session, DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session = Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let plan = sresult_return!(internal_provider.truncate(session).await);
+
+ FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime))
+ }
+ .into_ffi()
+}
+
unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_TableProvider) {
unsafe {
debug_assert!(!provider.private_data.is_null());
@@ -368,6 +527,9 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_TableProvider) -> FFI_Table
version: super::version,
private_data,
library_marker_id: crate::get_library_marker_id,
+ delete_from: provider.delete_from,
+ update: provider.update,
+ truncate: provider.truncate,
}
}
@@ -429,6 +591,9 @@ impl FFI_TableProvider {
version: super::version,
private_data: Box::into_raw(private_data) as *mut c_void,
library_marker_id: crate::get_library_marker_id,
+ delete_from: delete_from_fn_wrapper,
+ update: update_fn_wrapper,
+ truncate: truncate_fn_wrapper,
}
}
}
@@ -499,10 +664,7 @@ impl TableProvider for ForeignTableProvider {
.into();
let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
- let filter_list = LogicalExprList {
- expr: serialize_exprs(filters, codec.as_ref())?,
- };
- let filters_serialized = filter_list.encode_to_vec().into_iter().collect();
+ let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?;
let plan = unsafe {
let maybe_plan = (self.0.scan)(
@@ -539,13 +701,8 @@ impl TableProvider for ForeignTableProvider {
let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
- let expr_list = LogicalExprList {
- expr: serialize_exprs(
- filters.iter().map(|f| f.to_owned()),
- codec.as_ref(),
- )?,
- };
- let serialized_filters = expr_list.encode_to_vec();
+ let serialized_filters =
+ serialize_expr_list(filters.iter().copied(), codec.as_ref())?;
let pushdowns = df_result!(pushdown_fn(
&self.0,
@@ -577,11 +734,79 @@ impl TableProvider for ForeignTableProvider {
Ok(plan)
}
+
+ async fn delete_from(
+ &self,
+ session: &dyn Session,
+ filters: Vec<Expr>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone());
+ let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
+ let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?;
+
+ let plan = unsafe {
+ let maybe_plan =
+ (self.0.delete_from)(&self.0, session, filters_serialized).await;
+
+ <Arc<dyn ExecutionPlan>>::try_from(&df_result!(maybe_plan)?)?
+ };
+
+ Ok(plan)
+ }
+
+ async fn update(
+ &self,
+ session: &dyn Session,
+ assignments: Vec<(String, Expr)>,
+ filters: Vec<Expr>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone());
+ let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
+
+ let assignments: SVec<_> = assignments
+ .iter()
+ .map(|(column, expr)| {
+ Ok(FFI_TableProviderUpdateAssignment {
+ column: SString::from(column.as_str()),
+ expr_serialized: serialize_expr_list(
+ std::iter::once(expr),
+ codec.as_ref(),
+ )?,
+ })
+ })
+ .collect::<Result<Vec<_>>>()?
+ .into_iter()
+ .collect();
+ let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?;
+
+ let plan = unsafe {
+ let maybe_plan =
+ (self.0.update)(&self.0, session, assignments, filters_serialized).await;
+
+ <Arc<dyn ExecutionPlan>>::try_from(&df_result!(maybe_plan)?)?
+ };
+
+ Ok(plan)
+ }
+
+ async fn truncate(&self, session: &dyn Session) -> Result<Arc<dyn ExecutionPlan>> {
+ let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone());
+
+ let plan = unsafe {
+ let maybe_plan = (self.0.truncate)(&self.0, session).await;
+
+ <Arc<dyn ExecutionPlan>>::try_from(&df_result!(maybe_plan)?)?
+ };
+
+ Ok(plan)
+ }
}
#[cfg(test)]
mod tests {
- use arrow::datatypes::Schema;
+ use std::sync::atomic::{AtomicUsize, Ordering};
+
+ use arrow::datatypes::{DataType, Field, Schema};
use datafusion::prelude::{SessionContext, col, lit};
use datafusion_execution::TaskContextProvider;
@@ -613,6 +838,105 @@ mod tests {
)?))
}
+ #[derive(Debug)]
+ struct TableWithStats {
+ inner: Arc<dyn TableProvider>,
+ stats: Option<Statistics>,
+ delete_calls: AtomicUsize,
+ update_calls: AtomicUsize,
+ }
+
+ impl TableWithStats {
+ fn new(inner: Arc<dyn TableProvider>, stats: Option<Statistics>) -> Self {
+ Self {
+ inner,
+ stats,
+ delete_calls: AtomicUsize::new(0),
+ update_calls: AtomicUsize::new(0),
+ }
+ }
+ }
+
+ #[async_trait]
+ impl TableProvider for TableWithStats {
+ fn schema(&self) -> SchemaRef {
+ self.inner.schema()
+ }
+
+ fn table_type(&self) -> TableType {
+ self.inner.table_type()
+ }
+
+ fn statistics(&self) -> Option<Statistics> {
+ self.stats.clone()
+ }
+
+ async fn scan(
+ &self,
+ session: &dyn Session,
+ projection: Option<&Vec<usize>>,
+ filters: &[Expr],
+ limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.inner.scan(session, projection, filters, limit).await
+ }
+
+ async fn delete_from(
+ &self,
+ _state: &dyn Session,
+ filters: Vec<Expr>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let call = self.delete_calls.fetch_add(1, Ordering::Relaxed);
+ let valid = match call {
+ 0 => filters == vec![col("a").gt(lit(10_i64)), col("b").lt(lit(2.5_f64))],
+ 1 => filters.is_empty(),
+ _ => false,
+ };
+
+ if !valid {
+ return Err(DataFusionError::Internal(format!(
+ "Unexpected DELETE filters for call {call}"
+ )));
+ }
+ Ok(dml_count_plan())
+ }
+
+ async fn update(
+ &self,
+ _state: &dyn Session,
+ assignments: Vec<(String, Expr)>,
+ filters: Vec<Expr>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if assignments
+ != vec![
+ ("b".to_string(), lit(42_f64)),
+ ("a".to_string(), lit(7_i64)),
+ ]
+ {
+ return Err(DataFusionError::Internal(
+ "Unexpected UPDATE assignments".to_string(),
+ ));
+ }
+ let call = self.update_calls.fetch_add(1, Ordering::Relaxed);
+ let valid = match call {
+ 0 => filters == vec![col("a").eq(lit(7_i64)), col("b").gt(lit(1.5_f64))],
+ 1 => filters.is_empty(),
+ _ => false,
+ };
+
+ if !valid {
+ return Err(DataFusionError::Internal(format!(
+ "Unexpected UPDATE filters for call {call}"
+ )));
+ }
+ Ok(dml_count_plan())
+ }
+
+ async fn truncate(&self, _state: &dyn Session) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(dml_count_plan())
+ }
+ }
+
#[tokio::test]
async fn test_round_trip_ffi_table_provider_scan() -> Result<()> {
let provider = create_test_table_provider()?;
@@ -671,11 +995,91 @@ mod tests {
Ok(())
}
+ fn dml_count_plan() -> Arc<dyn ExecutionPlan> {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "count",
+ DataType::UInt64,
+ false,
+ )]));
+ Arc::new(crate::execution_plan::tests::EmptyExec::new(schema))
+ }
+
+ #[tokio::test]
+ async fn test_round_trip_ffi_table_provider_dml() -> Result<()> {
+ use datafusion::datasource::MemTable;
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Float64, true),
+ ]));
+ let inner = Arc::new(MemTable::try_new(schema, vec![vec![]])?);
+ let provider = Arc::new(TableWithStats::new(inner, None));
+ let ctx = Arc::new(SessionContext::new());
+ let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
+ let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider);
+
+ let mut ffi_provider =
+ FFI_TableProvider::new(provider, true, None, task_ctx_provider, None);
+ ffi_provider.library_marker_id = crate::mock_foreign_marker_id;
+
+ let foreign_table_provider: Arc<dyn TableProvider> = (&ffi_provider).into();
+ assert!(
+ foreign_table_provider
+ .downcast_ref::<ForeignTableProvider>()
+ .is_some()
+ );
+
+ let state = ctx.state();
+ let delete_filters = vec![col("a").gt(lit(10_i64)), col("b").lt(lit(2.5_f64))];
+
+ let delete_plan = foreign_table_provider
+ .delete_from(&state, delete_filters)
+ .await?;
+ assert_eq!(delete_plan.schema().field(0).name(), "count");
+
+ let delete_all_plan = foreign_table_provider.delete_from(&state, vec![]).await?;
+ assert_eq!(delete_all_plan.schema().field(0).name(), "count");
+
+ let update_assignments = vec![
+ ("b".to_string(), lit(42_f64)),
+ ("a".to_string(), lit(7_i64)),
+ ];
+ let update_filters = vec![col("a").eq(lit(7_i64)), col("b").gt(lit(1.5_f64))];
+
+ let update_plan = foreign_table_provider
+ .update(&state, update_assignments.clone(), update_filters)
+ .await?;
+ assert_eq!(update_plan.schema().field(0).name(), "count");
+
+ let update_all_plan = foreign_table_provider
+ .update(&state, update_assignments, vec![])
+ .await?;
+ assert_eq!(update_all_plan.schema().field(0).name(), "count");
+
+ let truncate_plan = foreign_table_provider.truncate(&state).await?;
+ assert_eq!(truncate_plan.schema().field(0).name(), "count");
+
+ let session =
+ FFI_SessionRef::new(&state, None, ffi_provider.logical_codec.clone());
+ let assignments = [FFI_TableProviderUpdateAssignment {
+ column: SString::from("b"),
+ expr_serialized: SVec::new(),
+ }]
+ .into_iter()
+ .collect();
+ let result = unsafe {
+ (ffi_provider.update)(&ffi_provider, session, assignments, SVec::new()).await
+ };
+ assert!(df_result!(result).unwrap_err().to_string().contains(
+ "Expected exactly one expression for update assignment to column 'b', got 0",
+ ));
+
+ Ok(())
+ }
+
#[tokio::test]
async fn test_aggregation() -> Result<()> {
- use arrow::datatypes::Field;
use datafusion::arrow::array::Float32Array;
- use datafusion::arrow::datatypes::DataType;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::assert_batches_eq;
use datafusion::datasource::MemTable;
@@ -816,35 +1220,6 @@ mod tests {
use datafusion_common::stats::Precision;
use datafusion_common::{ColumnStatistics, ScalarValue};
- // A thin wrapper that lets us inject statistics onto any TableProvider.
- #[derive(Debug)]
- struct TableWithStats {
- inner: Arc<dyn TableProvider>,
- stats: Option<Statistics>,
- }
-
- #[async_trait]
- impl TableProvider for TableWithStats {
- fn schema(&self) -> SchemaRef {
- self.inner.schema()
- }
- fn table_type(&self) -> TableType {
- self.inner.table_type()
- }
- fn statistics(&self) -> Option<Statistics> {
- self.stats.clone()
- }
- async fn scan(
- &self,
- session: &dyn Session,
- projection: Option<&Vec<usize>>,
- filters: &[Expr],
- limit: Option<usize>,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- self.inner.scan(session, projection, filters, limit).await
- }
- }
-
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
let batch = RecordBatch::try_new(
@@ -861,10 +1236,7 @@ mod tests {
Arc::clone(&schema),
vec![vec![batch.clone()]],
)?);
- let no_stats_provider = Arc::new(TableWithStats {
- inner: no_stats_inner,
- stats: None,
- });
+ let no_stats_provider = Arc::new(TableWithStats::new(no_stats_inner, None));
let mut ffi_provider = FFI_TableProvider::new(
no_stats_provider,
true,
@@ -891,10 +1263,10 @@ mod tests {
};
let stats_inner =
Arc::new(MemTable::try_new(Arc::clone(&schema), vec![vec![batch]])?);
- let stats_provider = Arc::new(TableWithStats {
- inner: stats_inner,
- stats: Some(original_stats.clone()),
- });
+ let stats_provider = Arc::new(TableWithStats::new(
+ stats_inner,
+ Some(original_stats.clone()),
+ ));
let mut ffi_provider =
FFI_TableProvider::new(stats_provider, true, None, task_ctx_provider, None);
ffi_provider.library_marker_id = crate::mock_foreign_marker_id;
diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs
index 74310d9c28..ce922de664 100644
--- a/datafusion/ffi/src/tests/mod.rs
+++ b/datafusion/ffi/src/tests/mod.rs
@@ -16,6 +16,7 @@
// under the License.
use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
use arrow::array::{RecordBatch, record_batch};
use arrow_schema::{DataType, Field, Schema};
@@ -26,8 +27,8 @@ use datafusion_catalog::MemTable;
use datafusion_catalog::{Session, TableProvider};
use datafusion_common::stats::Precision;
use datafusion_common::{ColumnStatistics, Statistics};
-use datafusion_common::{Result, ScalarValue};
-use datafusion_expr::{Expr, TableType};
+use datafusion_common::{Result, ScalarValue, exec_err};
+use datafusion_expr::{Expr, TableType, col, lit};
use datafusion_physical_expr::PhysicalExpr;
use datafusion_physical_plan::ExecutionPlan;
use sync_provider::create_sync_table_provider;
@@ -237,6 +238,8 @@ pub(crate) extern "C" fn create_exec_with_statistics() -> FFI_ExecutionPlan {
struct TableWithStats {
inner: Arc<dyn TableProvider>,
stats: Statistics,
+ delete_calls: AtomicUsize,
+ update_calls: AtomicUsize,
}
#[async_trait]
@@ -262,6 +265,56 @@ impl TableProvider for TableWithStats {
) -> Result<Arc<dyn ExecutionPlan>> {
self.inner.scan(session, projection, filters, limit).await
}
+
+ async fn delete_from(
+ &self,
+ _state: &dyn Session,
+ filters: Vec<Expr>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let call = self.delete_calls.fetch_add(1, Ordering::Relaxed);
+ let valid = match call {
+ 0 => filters == vec![col("a").gt(lit(10_i32)), col("b").lt(lit(2.5_f64))],
+ 1 => filters.is_empty(),
+ _ => false,
+ };
+ if !valid {
+ return exec_err!("Unexpected DELETE filters for call {call}");
+ }
+ Ok(dml_count_plan())
+ }
+
+ async fn update(
+ &self,
+ _state: &dyn Session,
+ assignments: Vec<(String, Expr)>,
+ filters: Vec<Expr>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if assignments
+ != vec![
+ ("b".to_string(), lit(42_f64)),
+ ("a".to_string(), lit(7_i32)),
+ ]
+ {
+ return exec_err!("Unexpected UPDATE assignments");
+ }
+
+ let call = self.update_calls.fetch_add(1, Ordering::Relaxed);
+ let valid = match call {
+ 0 => filters == vec![col("a").eq(lit(7_i32)), col("b").gt(lit(1.5_f64))],
+ 1 => filters.is_empty(),
+ _ => false,
+ };
+
+ if !valid {
+ return exec_err!("Unexpected UPDATE filters for call {call}");
+ }
+
+ Ok(dml_count_plan())
+ }
+
+ async fn truncate(&self, _state: &dyn Session) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(dml_count_plan())
+ }
}
pub(crate) extern "C" fn create_table_with_statistics(
@@ -273,10 +326,21 @@ pub(crate) extern "C" fn create_table_with_statistics(
let provider = Arc::new(TableWithStats {
inner,
stats: make_test_statistics(),
+ delete_calls: AtomicUsize::new(0),
+ update_calls: AtomicUsize::new(0),
});
FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec)
}
+fn dml_count_plan() -> Arc<dyn ExecutionPlan> {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "count",
+ DataType::UInt64,
+ false,
+ )]));
+ Arc::new(EmptyExec::new(schema))
+}
+
/// This defines the entry point for using the module.
#[unsafe(no_mangle)]
pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule {
diff --git a/datafusion/ffi/src/udtf.rs b/datafusion/ffi/src/udtf.rs
index fa28519d58..eb1744b7b1 100644
--- a/datafusion/ffi/src/udtf.rs
+++ b/datafusion/ffi/src/udtf.rs
@@ -23,14 +23,14 @@ use datafusion_catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
use datafusion_common::DataFusionError;
use datafusion_common::error::Result;
use datafusion_execution::TaskContext;
-use datafusion_proto::logical_plan::from_proto::parse_exprs;
-use datafusion_proto::logical_plan::to_proto::serialize_exprs;
+use datafusion_proto::bytes::{
+ logical_exprs_from_bytes_with_extension_codec,
+ logical_exprs_to_bytes_with_extension_codec,
+};
use datafusion_proto::logical_plan::{
DefaultLogicalExtensionCodec, LogicalExtensionCodec,
};
-use datafusion_proto::protobuf::LogicalExprList;
use datafusion_session::Session;
-use prost::Message;
use stabby::vec::Vec as SVec;
use tokio::runtime::Handle;
@@ -113,12 +113,10 @@ unsafe extern "C" fn call_fn_wrapper(
sresult_return!((&udtf.logical_codec.task_ctx_provider).try_into());
let codec: Arc<dyn LogicalExtensionCodec> = (&udtf.logical_codec).into();
- let proto_filters = sresult_return!(LogicalExprList::decode(args.as_ref()));
-
- let args = sresult_return!(parse_exprs(
- proto_filters.expr.iter(),
+ let args = sresult_return!(logical_exprs_from_bytes_with_extension_codec(
+ args.as_ref(),
ctx.as_ref(),
- codec.as_ref()
+ codec.as_ref(),
));
#[expect(deprecated)]
@@ -143,12 +141,10 @@ unsafe extern "C" fn call_with_args_wrapper(
sresult_return!((&udtf.logical_codec.task_ctx_provider).try_into());
let codec: Arc<dyn LogicalExtensionCodec> = (&udtf.logical_codec).into();
- let proto_filters = sresult_return!(LogicalExprList::decode(args.as_ref()));
-
- let args = sresult_return!(parse_exprs(
- proto_filters.expr.iter(),
+ let args = sresult_return!(logical_exprs_from_bytes_with_extension_codec(
+ args.as_ref(),
ctx.as_ref(),
- codec.as_ref()
+ codec.as_ref(),
));
let mut foreign_session = None;
@@ -280,13 +276,12 @@ impl TableFunctionImpl for ForeignTableFunction {
self.0.logical_codec.clone(),
);
let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
- let expr_list = LogicalExprList {
- expr: serialize_exprs(args.exprs(), codec.as_ref())?,
- };
- let filters_serialized = expr_list.encode_to_vec().into_iter().collect();
+ let bytes =
+ logical_exprs_to_bytes_with_extension_codec(args.exprs(), codec.as_ref())?;
+ let args_serialized = SVec::from(bytes.as_ref());
let table_provider =
- unsafe { (self.0.call_with_args)(&self.0, filters_serialized, session) };
+ unsafe { (self.0.call_with_args)(&self.0, args_serialized, session) };
let table_provider = df_result!(table_provider)?;
let table_provider: Arc<dyn TableProvider> = (&table_provider).into();
@@ -296,13 +291,11 @@ impl TableFunctionImpl for ForeignTableFunction {
fn call(&self, args: &[datafusion_expr::Expr]) -> Result<Arc<dyn TableProvider>> {
let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
- let expr_list = LogicalExprList {
- expr: serialize_exprs(args, codec.as_ref())?,
- };
- let filters_serialized = expr_list.encode_to_vec().into_iter().collect();
+ let bytes = logical_exprs_to_bytes_with_extension_codec(args, codec.as_ref())?;
+ let args_serialized = SVec::from(bytes.as_ref());
#[expect(deprecated)]
- let table_provider = unsafe { (self.0.call)(&self.0, filters_serialized) };
+ let table_provider = unsafe { (self.0.call)(&self.0, args_serialized) };
let table_provider = df_result!(table_provider)?;
let table_provider: Arc<dyn TableProvider> = (&table_provider).into();
diff --git a/datafusion/ffi/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs
index 86f953e262..cba334e8fa 100644
--- a/datafusion/ffi/tests/ffi_integration.rs
+++ b/datafusion/ffi/tests/ffi_integration.rs
@@ -27,6 +27,7 @@ mod tests {
use arrow::datatypes::Schema;
use datafusion::catalog::{TableProvider, TableProviderFactory};
use datafusion::error::Result;
+ use datafusion::prelude::{col, lit};
use datafusion_common::TableReference;
use datafusion_common::ToDFSchema;
use datafusion_expr::CreateExternalTable;
@@ -95,6 +96,47 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_ffi_table_provider_dml_cross_library() -> Result<()> {
+ let module = get_module()?;
+ let (ctx, codec) = super::utils::ctx_and_codec();
+
+ let ffi_provider = (module.create_table_with_statistics)(codec);
+ let foreign: Arc<dyn TableProvider> = (&ffi_provider).into();
+ let state = ctx.state();
+
+ let delete_plan = foreign
+ .delete_from(
+ &state,
+ vec![col("a").gt(lit(10_i32)), col("b").lt(lit(2.5_f64))],
+ )
+ .await?;
+ assert_eq!(delete_plan.schema().field(0).name(), "count");
+ let delete_all_plan = foreign.delete_from(&state, vec![]).await?;
+ assert_eq!(delete_all_plan.schema().field(0).name(), "count");
+
+ let update_assignments = vec![
+ ("b".to_string(), lit(42_f64)),
+ ("a".to_string(), lit(7_i32)),
+ ];
+ let update_plan = foreign
+ .update(
+ &state,
+ update_assignments.clone(),
+ vec![col("a").eq(lit(7_i32)), col("b").gt(lit(1.5_f64))],
+ )
+ .await?;
+ assert_eq!(update_plan.schema().field(0).name(), "count");
+
+ let update_all_plan = foreign.update(&state, update_assignments, vec![]).await?;
+ assert_eq!(update_all_plan.schema().field(0).name(), "count");
+
+ let truncate_plan = foreign.truncate(&state).await?;
+ assert_eq!(truncate_plan.schema().field(0).name(), "count");
+
+ Ok(())
+ }
+
#[tokio::test]
async fn test_table_provider_factory() -> Result<()> {
let table_provider_module = get_module()?;
diff --git a/datafusion/proto/src/bytes/mod.rs b/datafusion/proto/src/bytes/mod.rs
index 388e373c3f..fd745041eb 100644
--- a/datafusion/proto/src/bytes/mod.rs
+++ b/datafusion/proto/src/bytes/mod.rs
@@ -16,7 +16,7 @@
// under the License.
//! Serialization / Deserialization to Bytes
-use crate::logical_plan::to_proto::serialize_expr;
+use crate::logical_plan::to_proto::{serialize_expr, serialize_exprs};
use crate::logical_plan::{
self, AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec,
};
@@ -107,6 +107,37 @@ impl Serializeable for Expr {
}
}
+/// Serialize logical expressions as protobuf bytes using the provided
+/// extension codec.
+pub fn logical_exprs_to_bytes_with_extension_codec<'a>(
+ exprs: impl IntoIterator<Item = &'a Expr>,
+ extension_codec: &dyn LogicalExtensionCodec,
+) -> Result<Bytes> {
+ Ok(protobuf::LogicalExprList {
+ expr: serialize_exprs(exprs, extension_codec)?,
+ }
+ .encode_to_vec()
+ .into())
+}
+
+/// Deserialize logical expressions from protobuf bytes using the provided
+/// extension codec.
+pub fn logical_exprs_from_bytes_with_extension_codec(
+ bytes: &[u8],
+ ctx: &TaskContext,
+ extension_codec: &dyn LogicalExtensionCodec,
+) -> Result<Vec<Expr>> {
+ let protobuf = protobuf::LogicalExprList::decode(bytes).map_err(|e| {
+ plan_datafusion_err!("Error decoding expressions as protobuf: {e}")
+ })?;
+
+ Ok(logical_plan::from_proto::parse_exprs(
+ protobuf.expr.iter(),
+ ctx,
+ extension_codec,
+ )?)
+}
+
/// Serialize a LogicalPlan as bytes
pub fn logical_plan_to_bytes(plan: &LogicalPlan) -> Result<Bytes> {
let extension_codec = DefaultLogicalExtensionCodec {};
diff --git a/datafusion/proto/tests/cases/serialize.rs b/datafusion/proto/tests/cases/serialize.rs
index a0a917e323..46ab8511f7 100644
--- a/datafusion/proto/tests/cases/serialize.rs
+++ b/datafusion/proto/tests/cases/serialize.rs
@@ -25,14 +25,18 @@ use datafusion::prelude::SessionContext;
use datafusion_common::ScalarValue;
use datafusion_expr::expr::{HigherOrderFunction, LambdaVariable, Placeholder};
use datafusion_expr::{ColumnarValue, HigherOrderUDF, col, create_udf, lambda, lit};
-use datafusion_expr::{Expr, Volatility};
+use datafusion_expr::{Expr, ScalarUDF, Volatility};
use datafusion_functions::string;
-use datafusion_proto::bytes::Serializeable;
+use datafusion_proto::bytes::{
+ Serializeable, logical_exprs_from_bytes_with_extension_codec,
+ logical_exprs_to_bytes_with_extension_codec,
+};
use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
use datafusion_proto::logical_plan::from_proto::parse_expr;
use datafusion_proto::logical_plan::to_proto::serialize_expr;
-use crate::cases::MyHigherOrderUDF;
+use crate::cases::roundtrip_logical_plan::UDFExtensionCodec;
+use crate::cases::{MyHigherOrderUDF, MyRegexUdf};
#[test]
#[should_panic(
@@ -42,6 +46,57 @@ fn bad_decode() {
Expr::from_bytes(b"Leet").unwrap();
}
+#[test]
+fn logical_exprs_roundtrip() {
+ let ctx = SessionContext::new();
+ let codec = DefaultLogicalExtensionCodec {};
+ let exprs = vec![col("a").gt(lit(10_i32)), col("b").lt(lit(2.5_f64))];
+
+ let bytes = logical_exprs_to_bytes_with_extension_codec(&exprs, &codec).unwrap();
+
+ let decoded = logical_exprs_from_bytes_with_extension_codec(
+ &bytes,
+ ctx.task_ctx().as_ref(),
+ &codec,
+ )
+ .unwrap();
+
+ assert_eq!(decoded, exprs);
+
+ let bytes =
+ logical_exprs_to_bytes_with_extension_codec(std::iter::empty::<&Expr>(), &codec)
+ .unwrap();
+
+ assert!(bytes.is_empty());
+
+ let decoded = logical_exprs_from_bytes_with_extension_codec(
+ &bytes,
+ ctx.task_ctx().as_ref(),
+ &codec,
+ )
+ .unwrap();
+
+ assert!(decoded.is_empty());
+}
+
+#[test]
+fn logical_exprs_roundtrip_with_extension_codec() {
+ let ctx = SessionContext::new();
+ let codec = UDFExtensionCodec;
+ let udf = ScalarUDF::from(MyRegexUdf::new(".*".to_owned()));
+ let exprs = vec![udf.call(vec![lit("foo")])];
+
+ let bytes = logical_exprs_to_bytes_with_extension_codec(&exprs, &codec).unwrap();
+
+ let decoded = logical_exprs_from_bytes_with_extension_codec(
+ &bytes,
+ ctx.task_ctx().as_ref(),
+ &codec,
+ )
+ .unwrap();
+ assert_eq!(format!("{exprs:?}"), format!("{decoded:?}"));
+}
+
#[test]
#[cfg(feature = "json")]
fn plan_to_json() {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org