From a426574ccbaf110faec01cbcb7620876e128df32 Mon Sep 17 00:00:00 2001 From: Amogh Ramesh Date: Sun, 26 Jul 2026 18:45:19 +0530 Subject: [PATCH 1/4] fix(ffi): preserve TableProvider DML overrides --- datafusion/ffi/src/table_provider.rs | 440 ++++++++++++++++++++++++--- 1 file changed, 404 insertions(+), 36 deletions(-) diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 5a4b2fa27256f..5de5e7ab77f46 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -36,6 +36,7 @@ use datafusion_proto::logical_plan::{ use datafusion_proto::protobuf::LogicalExprList; use prost::Message; +use stabby::string::String as SString; use stabby::vec::Vec as SVec; use tokio::runtime::Handle; @@ -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, + ) -> FfiFuture>, + + update: unsafe extern "C" fn( + provider: &Self, + session: FFI_SessionRef, + assignments: SVec, + filters_serialized: SVec, + ) -> FfiFuture>, + + truncate: unsafe extern "C" fn( + provider: &Self, + session: FFI_SessionRef, + ) -> FfiFuture>, } unsafe impl Send for FFI_TableProvider {} unsafe impl Sync for FFI_TableProvider {} +#[repr(C)] +#[derive(Debug)] +struct FFI_TableProviderUpdateAssignment { + column: SString, + expr_serialized: SVec, +} + struct ProviderPrivateData { provider: Arc, runtime: Option, @@ -202,21 +228,45 @@ unsafe extern "C" fn table_type_fn_wrapper( provider.inner().table_type().into() } -fn supports_filters_pushdown_internal( - provider: &Arc, - filters_serialized: &[u8], +fn parse_serialized_exprs( + exprs_serialized: &[u8], task_ctx: &Arc, codec: &dyn LogicalExtensionCodec, -) -> Result> { - let filters = match filters_serialized.is_empty() { - true => vec![], +) -> Result> { + match exprs_serialized.is_empty() { + true => Ok(vec![]), false => { - let proto_filters = LogicalExprList::decode(filters_serialized) + let proto_exprs = LogicalExprList::decode(exprs_serialized) .map_err(|e| DataFusionError::Plan(e.to_string()))?; - parse_exprs(proto_filters.expr.iter(), task_ctx.as_ref(), codec)? + Ok(parse_exprs( + proto_exprs.expr.iter(), + task_ctx.as_ref(), + codec, + )?) } - }; + } +} + +fn serialize_expr_list<'a>( + exprs: impl IntoIterator, + codec: &dyn LogicalExtensionCodec, +) -> Result> { + Ok(LogicalExprList { + expr: serialize_exprs(exprs, codec)?, + } + .encode_to_vec() + .into_iter() + .collect()) +} + +fn supports_filters_pushdown_internal( + provider: &Arc, + filters_serialized: &[u8], + task_ctx: &Arc, + codec: &dyn LogicalExtensionCodec, +) -> Result> { + let filters = parse_serialized_exprs(filters_serialized, task_ctx, codec)?; let filters_borrowed: Vec<&Expr> = filters.iter().collect(); let results: SVec<_> = provider @@ -271,19 +321,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> = projections.into_option().map(|p| p.into_iter().collect()); @@ -336,6 +378,133 @@ 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, +) -> FfiFuture> { + let task_ctx: Result, DataFusionError> = + (&provider.logical_codec.task_ctx_provider).try_into(); + let runtime = provider.runtime().clone(); + let logical_codec: Arc = (&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 + Send + Sync), 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.clone())) + } + .into_ffi() +} + +unsafe extern "C" fn update_fn_wrapper( + provider: &FFI_TableProvider, + session: FFI_SessionRef, + assignments: SVec, + filters_serialized: SVec, +) -> FfiFuture> { + let task_ctx: Result, DataFusionError> = + (&provider.logical_codec.task_ctx_provider).try_into(); + let runtime = provider.runtime().clone(); + let logical_codec: Arc = (&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 + Send + Sync), 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 mut exprs = parse_serialized_exprs( + &assignment.expr_serialized, + &task_ctx, + logical_codec.as_ref(), + )?; + let expr = match exprs.len() { + 1 => exprs.remove(0), + _ => { + return Err(DataFusionError::Plan( + "Expected exactly one expression for update assignment" + .to_string(), + )); + } + }; + Ok((assignment.column.to_string(), expr)) + }) + .collect::>>() + ); + 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.clone())) + } + .into_ffi() +} + +unsafe extern "C" fn truncate_fn_wrapper( + provider: &FFI_TableProvider, + session: FFI_SessionRef, +) -> FfiFuture> { + 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 + Send + Sync), 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.clone())) + } + .into_ffi() +} + unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_TableProvider) { unsafe { debug_assert!(!provider.private_data.is_null()); @@ -368,6 +537,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 +601,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 +674,7 @@ impl TableProvider for ForeignTableProvider { .into(); let codec: Arc = (&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 +711,8 @@ impl TableProvider for ForeignTableProvider { let codec: Arc = (&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 +744,79 @@ impl TableProvider for ForeignTableProvider { Ok(plan) } + + async fn delete_from( + &self, + session: &dyn Session, + filters: Vec, + ) -> Result> { + let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + let codec: Arc = (&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; + + >::try_from(&df_result!(maybe_plan)?)? + }; + + Ok(plan) + } + + async fn update( + &self, + session: &dyn Session, + assignments: Vec<(String, Expr)>, + filters: Vec, + ) -> Result> { + let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + let codec: Arc = (&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::>>()? + .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; + + >::try_from(&df_result!(maybe_plan)?)? + }; + + Ok(plan) + } + + async fn truncate(&self, session: &dyn Session) -> Result> { + 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; + + >::try_from(&df_result!(maybe_plan)?)? + }; + + Ok(plan) + } } #[cfg(test)] mod tests { - use arrow::datatypes::Schema; + use std::sync::Mutex; + + use arrow::datatypes::{DataType, Field, Schema}; use datafusion::prelude::{SessionContext, col, lit}; use datafusion_execution::TaskContextProvider; @@ -671,11 +906,144 @@ mod tests { Ok(()) } + #[derive(Debug, Default)] + struct DmlCalls { + delete_filters: Option>, + update_assignments: Option>, + update_filters: Option>, + truncated: bool, + } + + #[derive(Debug)] + struct DmlTableProvider { + calls: Arc>, + schema: SchemaRef, + } + + fn dml_count_plan() -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new( + "count", + DataType::UInt64, + false, + )])); + Arc::new(crate::execution_plan::tests::EmptyExec::new(schema)) + } + + #[async_trait] + impl TableProvider for DmlTableProvider { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _session: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Err(DataFusionError::Internal( + "DmlTableProvider scan should not be called".to_string(), + )) + } + + async fn delete_from( + &self, + _state: &dyn Session, + filters: Vec, + ) -> Result> { + self.calls.lock().unwrap().delete_filters = Some(filters); + Ok(dml_count_plan()) + } + + async fn update( + &self, + _state: &dyn Session, + assignments: Vec<(String, Expr)>, + filters: Vec, + ) -> Result> { + let mut calls = self.calls.lock().unwrap(); + calls.update_assignments = Some(assignments); + calls.update_filters = Some(filters); + Ok(dml_count_plan()) + } + + async fn truncate(&self, _state: &dyn Session) -> Result> { + self.calls.lock().unwrap().truncated = true; + Ok(dml_count_plan()) + } + } + + #[tokio::test] + async fn test_round_trip_ffi_table_provider_dml() -> Result<()> { + let calls = Arc::new(Mutex::new(DmlCalls::default())); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Float64, true), + ])); + let provider = Arc::new(DmlTableProvider { + calls: Arc::clone(&calls), + schema, + }); + let ctx = Arc::new(SessionContext::new()); + let task_ctx_provider = Arc::clone(&ctx) as Arc; + 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 = (&ffi_provider).into(); + assert!( + foreign_table_provider + .downcast_ref::() + .is_some() + ); + + let state = ctx.state(); + let delete_filter = col("a").gt(lit(10_i64)); + let delete_plan = foreign_table_provider + .delete_from(&state, vec![delete_filter.clone()]) + .await?; + assert_eq!(delete_plan.schema().field(0).name(), "count"); + assert_eq!( + calls.lock().unwrap().delete_filters.clone(), + Some(vec![delete_filter]) + ); + + let update_expr = lit(42_f64); + let update_filter = col("a").eq(lit(7_i64)); + let update_plan = foreign_table_provider + .update( + &state, + vec![("b".to_string(), update_expr.clone())], + vec![update_filter.clone()], + ) + .await?; + assert_eq!(update_plan.schema().field(0).name(), "count"); + assert_eq!( + calls.lock().unwrap().update_assignments.clone(), + Some(vec![("b".to_string(), update_expr)]) + ); + assert_eq!( + calls.lock().unwrap().update_filters.clone(), + Some(vec![update_filter]) + ); + + let truncate_plan = foreign_table_provider.truncate(&state).await?; + assert_eq!(truncate_plan.schema().field(0).name(), "count"); + assert!(calls.lock().unwrap().truncated); + + 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; From 55688992c955212d3870c7dfe208a62c95d3009c Mon Sep 17 00:00:00 2001 From: Amogh Ramesh Date: Sun, 26 Jul 2026 18:45:24 +0530 Subject: [PATCH 2/4] test(ffi): cover TableProvider DML across libraries --- datafusion/ffi/src/tests/mod.rs | 45 +++++++++++++++++++++++-- datafusion/ffi/tests/ffi_integration.rs | 30 +++++++++++++++++ 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index d372dcf9177e6..8f52e96e96d06 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -26,8 +26,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_plan::ExecutionPlan; use sync_provider::create_sync_table_provider; use udf_udaf_udwf::{ @@ -226,6 +226,38 @@ impl TableProvider for TableWithStats { ) -> Result> { self.inner.scan(session, projection, filters, limit).await } + + async fn delete_from( + &self, + _state: &dyn Session, + filters: Vec, + ) -> Result> { + if filters != vec![col("a").gt(lit(10_i32))] { + return exec_err!("Unexpected DELETE filters"); + } + + Ok(dml_count_plan()) + } + + async fn update( + &self, + _state: &dyn Session, + assignments: Vec<(String, Expr)>, + filters: Vec, + ) -> Result> { + if assignments != vec![("b".to_string(), lit(42_f64))] { + return exec_err!("Unexpected UPDATE assignments"); + } + if filters != vec![col("a").eq(lit(7_i32))] { + return exec_err!("Unexpected UPDATE filters"); + } + + Ok(dml_count_plan()) + } + + async fn truncate(&self, _state: &dyn Session) -> Result> { + Ok(dml_count_plan()) + } } pub(crate) extern "C" fn create_table_with_statistics( @@ -241,6 +273,15 @@ pub(crate) extern "C" fn create_table_with_statistics( FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec) } +fn dml_count_plan() -> Arc { + 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/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs index 86f953e262ead..42a8a0ff51443 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,35 @@ 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 = (&ffi_provider).into(); + let state = ctx.state(); + + let delete_plan = foreign + .delete_from(&state, vec![col("a").gt(lit(10_i32))]) + .await?; + assert_eq!(delete_plan.schema().field(0).name(), "count"); + + let update_plan = foreign + .update( + &state, + vec![("b".to_string(), lit(42_f64))], + vec![col("a").eq(lit(7_i32))], + ) + .await?; + assert_eq!(update_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()?; From 28a8e203ca15594941b361c414c6f8034035756f Mon Sep 17 00:00:00 2001 From: Amogh Ramesh Date: Sun, 26 Jul 2026 20:00:41 +0530 Subject: [PATCH 3/4] test(ffi): cover DML validation errors --- datafusion/ffi/src/table_provider.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 5de5e7ab77f46..671dc6e95001a 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -1038,6 +1038,24 @@ mod tests { assert_eq!(truncate_plan.schema().field(0).name(), "count"); assert!(calls.lock().unwrap().truncated); + 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") + ); + Ok(()) } From 64923e2ec8e48045fe5bd88899d9f62949dc2350 Mon Sep 17 00:00:00 2001 From: Amogh Ramesh Date: Sun, 26 Jul 2026 21:01:42 +0530 Subject: [PATCH 4/4] fix(ffi): tighten DML forwarding --- datafusion/ffi/src/table_provider.rs | 29 ++++++++++++------------- datafusion/ffi/src/tests/mod.rs | 11 +++++++--- datafusion/ffi/tests/ffi_integration.rs | 12 +++++++--- 3 files changed, 31 insertions(+), 21 deletions(-) diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 671dc6e95001a..65ca69afb5670 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -410,7 +410,7 @@ unsafe extern "C" fn delete_from_fn_wrapper( let plan = sresult_return!(internal_provider.delete_from(session, filters).await); - FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime.clone())) + FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime)) } .into_ffi() } @@ -474,7 +474,7 @@ unsafe extern "C" fn update_fn_wrapper( .await ); - FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime.clone())) + FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime)) } .into_ffi() } @@ -500,7 +500,7 @@ unsafe extern "C" fn truncate_fn_wrapper( let plan = sresult_return!(internal_provider.truncate(session).await); - FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime.clone())) + FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime)) } .into_ffi() } @@ -1005,33 +1005,32 @@ mod tests { ); let state = ctx.state(); - let delete_filter = col("a").gt(lit(10_i64)); + 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, vec![delete_filter.clone()]) + .delete_from(&state, delete_filters.clone()) .await?; assert_eq!(delete_plan.schema().field(0).name(), "count"); assert_eq!( calls.lock().unwrap().delete_filters.clone(), - Some(vec![delete_filter]) + Some(delete_filters) ); - let update_expr = lit(42_f64); - let update_filter = col("a").eq(lit(7_i64)); + 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, - vec![("b".to_string(), update_expr.clone())], - vec![update_filter.clone()], - ) + .update(&state, update_assignments.clone(), update_filters.clone()) .await?; assert_eq!(update_plan.schema().field(0).name(), "count"); assert_eq!( calls.lock().unwrap().update_assignments.clone(), - Some(vec![("b".to_string(), update_expr)]) + Some(update_assignments) ); assert_eq!( calls.lock().unwrap().update_filters.clone(), - Some(vec![update_filter]) + Some(update_filters) ); let truncate_plan = foreign_table_provider.truncate(&state).await?; diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 8f52e96e96d06..881835e568087 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -232,7 +232,7 @@ impl TableProvider for TableWithStats { _state: &dyn Session, filters: Vec, ) -> Result> { - if filters != vec![col("a").gt(lit(10_i32))] { + if filters != vec![col("a").gt(lit(10_i32)), col("b").lt(lit(2.5_f64))] { return exec_err!("Unexpected DELETE filters"); } @@ -245,10 +245,15 @@ impl TableProvider for TableWithStats { assignments: Vec<(String, Expr)>, filters: Vec, ) -> Result> { - if assignments != vec![("b".to_string(), lit(42_f64))] { + if assignments + != vec![ + ("b".to_string(), lit(42_f64)), + ("a".to_string(), lit(7_i32)), + ] + { return exec_err!("Unexpected UPDATE assignments"); } - if filters != vec![col("a").eq(lit(7_i32))] { + if filters != vec![col("a").eq(lit(7_i32)), col("b").gt(lit(1.5_f64))] { return exec_err!("Unexpected UPDATE filters"); } diff --git a/datafusion/ffi/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs index 42a8a0ff51443..10129d9e756a4 100644 --- a/datafusion/ffi/tests/ffi_integration.rs +++ b/datafusion/ffi/tests/ffi_integration.rs @@ -106,15 +106,21 @@ mod tests { let state = ctx.state(); let delete_plan = foreign - .delete_from(&state, vec![col("a").gt(lit(10_i32))]) + .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 update_plan = foreign .update( &state, - vec![("b".to_string(), lit(42_f64))], - vec![col("a").eq(lit(7_i32))], + vec![ + ("b".to_string(), lit(42_f64)), + ("a".to_string(), lit(7_i32)), + ], + vec![col("a").eq(lit(7_i32)), col("b").gt(lit(1.5_f64))], ) .await?; assert_eq!(update_plan.schema().field(0).name(), "count");