From d6f266582f978f7c0d5fe8fd60af36f479caf988 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 20 Feb 2025 19:53:20 -0500 Subject: [PATCH 1/6] Add unit tests to FFI_ExecutionPlan --- datafusion/ffi/src/execution_plan.rs | 66 +++++++++++++++++++++++++--- 1 file changed, 59 insertions(+), 7 deletions(-) diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 8087acfa33c8..6dd7307e7409 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -219,17 +219,17 @@ impl TryFrom<&FFI_ExecutionPlan> for ForeignExecutionPlan { let properties: PlanProperties = (plan.properties)(plan).try_into()?; let children_rvec = (plan.children)(plan); - let children: Result> = children_rvec + let children = children_rvec .iter() .map(ForeignExecutionPlan::try_from) .map(|child| child.map(|c| Arc::new(c) as Arc)) - .collect(); + .collect::>>()?; Ok(Self { name, plan: plan.clone(), properties, - children: children?, + children, }) } } @@ -281,6 +281,7 @@ impl ExecutionPlan for ForeignExecutionPlan { #[cfg(test)] mod tests { + use arrow::datatypes::{DataType, Field, Schema}; use datafusion::{ physical_plan::{ execution_plan::{Boundedness, EmissionType}, @@ -294,6 +295,7 @@ mod tests { #[derive(Debug)] pub struct EmptyExec { props: PlanProperties, + children: Vec>, } impl EmptyExec { @@ -305,6 +307,7 @@ mod tests { EmissionType::Incremental, Boundedness::Bounded, ), + children: Vec::default(), } } } @@ -333,14 +336,17 @@ mod tests { } fn children(&self) -> Vec<&Arc> { - vec![] + self.children.iter().collect() } fn with_new_children( self: Arc, - _: Vec>, + children: Vec>, ) -> Result> { - unimplemented!() + Ok(Arc::new(EmptyExec { + props: self.props.clone(), + children, + })) } fn execute( @@ -358,7 +364,8 @@ mod tests { #[test] fn test_round_trip_ffi_execution_plan() -> Result<()> { - use arrow::datatypes::{DataType, Field, Schema}; + use std::fmt::Write; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); let ctx = SessionContext::new(); @@ -372,6 +379,51 @@ mod tests { assert!(original_name == foreign_plan.name()); + let display = datafusion::physical_plan::display::DisplayableExecutionPlan::new( + &foreign_plan, + ); + + let mut buf = String::new(); + write!(&mut buf, "{}", display.one_line()).unwrap(); + let buf = buf.trim(); + assert_eq!(buf, "FFI_ExecutionPlan(number_of_children=0)"); + + Ok(()) + } + + #[test] + fn test_ffi_execution_plan_children() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); + let ctx = SessionContext::new(); + + // Version 1: Adding child to the foreign plan + let child_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let child_local = FFI_ExecutionPlan::new(child_plan, ctx.task_ctx(), None); + let child_foreign = Arc::new(ForeignExecutionPlan::try_from(&child_local)?); + + let parent_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let parent_local = FFI_ExecutionPlan::new(parent_plan, ctx.task_ctx(), None); + let parent_foreign = Arc::new(ForeignExecutionPlan::try_from(&parent_local)?); + + assert_eq!(parent_foreign.children().len(), 0); + assert_eq!(child_foreign.children().len(), 0); + + let parent_foreign = parent_foreign.with_new_children(vec![child_foreign])?; + assert_eq!(parent_foreign.children().len(), 1); + + // Version 2: Adding child to the local plan + let child_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let child_local = FFI_ExecutionPlan::new(child_plan, ctx.task_ctx(), None); + let child_foreign = Arc::new(ForeignExecutionPlan::try_from(&child_local)?); + + let parent_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let parent_plan = parent_plan.with_new_children(vec![child_foreign])?; + let parent_local = FFI_ExecutionPlan::new(parent_plan, ctx.task_ctx(), None); + let parent_foreign = Arc::new(ForeignExecutionPlan::try_from(&parent_local)?); + + assert_eq!(parent_foreign.children().len(), 1); + Ok(()) } } From 945d82258db29097b6ad11656a7ac7d3a02acb48 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 20 Feb 2025 20:17:48 -0500 Subject: [PATCH 2/6] Add unit tests for FFI table source --- datafusion/ffi/src/table_source.rs | 40 ++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/datafusion/ffi/src/table_source.rs b/datafusion/ffi/src/table_source.rs index a59836622ee6..418fdf16a564 100644 --- a/datafusion/ffi/src/table_source.rs +++ b/datafusion/ffi/src/table_source.rs @@ -85,3 +85,43 @@ impl From for FFI_TableType { } } } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::error::Result; + + fn round_trip_filter_pushdown(pushdown: TableProviderFilterPushDown) -> Result<()> { + let ffi_pushdown: FFI_TableProviderFilterPushDown = (&pushdown).into(); + let round_trip: TableProviderFilterPushDown = (&ffi_pushdown).into(); + + assert_eq!(pushdown, round_trip); + Ok(()) + } + + #[test] + fn round_trip_all_filter_pushdowns() -> Result<()> { + round_trip_filter_pushdown(TableProviderFilterPushDown::Exact)?; + round_trip_filter_pushdown(TableProviderFilterPushDown::Inexact)?; + round_trip_filter_pushdown(TableProviderFilterPushDown::Unsupported)?; + + Ok(()) + } + + fn round_trip_table_type(table_type: TableType) -> Result<()> { + let ffi_type: FFI_TableType = table_type.into(); + let round_trip_type: TableType = ffi_type.into(); + + assert_eq!(table_type, round_trip_type); + Ok(()) + } + + #[test] + fn test_round_all_trip_table_type() -> Result<()> { + round_trip_table_type(TableType::Base)?; + round_trip_table_type(TableType::Temporary)?; + round_trip_table_type(TableType::View)?; + + Ok(()) + } +} From e265dca9649dfcbbdcda7e409ae5ff459827c249 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 20 Feb 2025 20:24:09 -0500 Subject: [PATCH 3/6] Add round trip tests for volatility --- datafusion/ffi/src/volatility.rs | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/datafusion/ffi/src/volatility.rs b/datafusion/ffi/src/volatility.rs index 8b565b91b76d..0aaf68a174cf 100644 --- a/datafusion/ffi/src/volatility.rs +++ b/datafusion/ffi/src/volatility.rs @@ -46,3 +46,24 @@ impl From<&FFI_Volatility> for Volatility { } } } + +#[cfg(test)] +mod tests { + use datafusion::logical_expr::Volatility; + + use super::FFI_Volatility; + + fn test_round_trip_volatility(volatility: Volatility) { + let ffi_volatility: FFI_Volatility = volatility.into(); + let round_trip: Volatility = (&ffi_volatility).into(); + + assert_eq!(volatility, round_trip); + } + + #[test] + fn test_all_round_trip_volatility() { + test_round_trip_volatility(Volatility::Immutable); + test_round_trip_volatility(Volatility::Stable); + test_round_trip_volatility(Volatility::Volatile); + } +} From bc4718698b5b6ff6d1070030873d9be08e5ad1c0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 20 Feb 2025 20:33:01 -0500 Subject: [PATCH 4/6] Add unit tests for FFI insert op --- datafusion/ffi/src/insert_op.rs | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/datafusion/ffi/src/insert_op.rs b/datafusion/ffi/src/insert_op.rs index e44262377405..8e8693076cc0 100644 --- a/datafusion/ffi/src/insert_op.rs +++ b/datafusion/ffi/src/insert_op.rs @@ -47,3 +47,26 @@ impl From for FFI_InsertOp { } } } + +#[cfg(test)] +mod tests { + use datafusion::logical_expr::dml::InsertOp; + + use super::FFI_InsertOp; + + fn test_round_trip_insert_op(insert_op: InsertOp) { + let ffi_insert_op: FFI_InsertOp = insert_op.into(); + let round_trip: InsertOp = ffi_insert_op.into(); + + assert_eq!(insert_op, round_trip); + } + + /// This test ensures we have not accidentally mapped the FFI + /// enums to the wrong internal enums values. + #[test] + fn test_all_round_trip_insert_ops() { + test_round_trip_insert_op(InsertOp::Append); + test_round_trip_insert_op(InsertOp::Overwrite); + test_round_trip_insert_op(InsertOp::Replace); + } +} From 31bcdeba57d66b793caa98ae4f7123c22a19d5b8 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 21 Feb 2025 07:42:44 -0500 Subject: [PATCH 5/6] Simplify string generation in unit test Co-authored-by: Andrew Lamb --- datafusion/ffi/src/execution_plan.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 6dd7307e7409..8b415fed10d3 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -383,9 +383,7 @@ mod tests { &foreign_plan, ); - let mut buf = String::new(); - write!(&mut buf, "{}", display.one_line()).unwrap(); - let buf = buf.trim(); + let buf = display.one_line().to_string().trim(); assert_eq!(buf, "FFI_ExecutionPlan(number_of_children=0)"); Ok(()) From cd8ce84ff3def6c292408b3c3fc161ed6abc5495 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 21 Feb 2025 08:02:29 -0500 Subject: [PATCH 6/6] Fix drop of borrowed value --- datafusion/ffi/src/execution_plan.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 8b415fed10d3..00602474d621 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -364,8 +364,6 @@ mod tests { #[test] fn test_round_trip_ffi_execution_plan() -> Result<()> { - use std::fmt::Write; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); let ctx = SessionContext::new(); @@ -383,8 +381,8 @@ mod tests { &foreign_plan, ); - let buf = display.one_line().to_string().trim(); - assert_eq!(buf, "FFI_ExecutionPlan(number_of_children=0)"); + let buf = display.one_line().to_string(); + assert_eq!(buf.trim(), "FFI_ExecutionPlan(number_of_children=0)"); Ok(()) }