From ed3fc11e16393dcb488767e5e5213cddd82493b7 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 8 Jun 2026 15:04:39 -0400 Subject: [PATCH] fix(ffi): route partition statistics through statistics_with_args --- datafusion/ffi/src/execution_plan.rs | 115 +++++++++++++----- .../library-user-guide/upgrading/55.0.0.md | 8 ++ 2 files changed, 93 insertions(+), 30 deletions(-) diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 18f92d54ade7e..5501e769efe8c 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -24,7 +24,7 @@ use datafusion_common::{DataFusionError, Result, Statistics}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr_common::metrics::MetricsSet; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs, }; use stabby::string::String as SString; use stabby::vec::Vec as SVec; @@ -209,9 +209,8 @@ unsafe extern "C" fn partition_statistics_fn_wrapper( partition: FFI_Option, ) -> FFI_Result> { let partition: Option = partition.into(); - #[expect(deprecated)] plan.inner() - .partition_statistics(partition) + .statistics_with_args(&StatisticsArgs::new(partition)) .map(|stats| SVec::from(serialize_statistics(stats.as_ref()).as_slice())) .into() } @@ -557,9 +556,9 @@ pub mod tests { self.metrics.clone() } - fn partition_statistics( + fn statistics_with_args( &self, - _partition: Option, + _args: &StatisticsArgs, ) -> Result> { Ok(Arc::new(self.statistics.clone().unwrap_or_else(|| { Statistics::new_unknown(self.props.eq_properties.schema()) @@ -673,28 +672,34 @@ pub mod tests { Ok(()) } - #[test] - #[expect(deprecated)] - fn test_ffi_execution_plan_partition_statistics_round_trip() -> Result<()> { + /// Build an `EmptyExec` carrying `statistics`, then export it across the + /// (mock) FFI boundary and return the resulting foreign plan. + #[cfg(test)] + fn export_empty_exec_over_ffi( + schema: &arrow::datatypes::SchemaRef, + statistics: Option, + ) -> Result> { + let mut plan = EmptyExec::new(Arc::clone(schema)); + if let Some(statistics) = statistics { + plan = plan.with_statistics(statistics); + } + let mut local = FFI_ExecutionPlan::new(Arc::new(plan), None); + local.library_marker_id = crate::mock_foreign_marker_id; + let foreign: Arc = (&local).try_into()?; + Ok(foreign) + } + + /// Schema and a fully-populated `Statistics` (including `ScalarValue`-typed + /// min/max) shared by the FFI statistics round-trip tests. + #[cfg(test)] + fn stats_round_trip_fixture() -> (arrow::datatypes::SchemaRef, Statistics) { use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; let schema = Arc::new(arrow::datatypes::Schema::new(vec![ arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Int32, true), ])); - - // Plans without explicit statistics return Statistics::new_unknown across - // the boundary. - let bare_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let mut bare_local = FFI_ExecutionPlan::new(bare_plan, None); - bare_local.library_marker_id = crate::mock_foreign_marker_id; - let bare_foreign: Arc = (&bare_local).try_into()?; - let bare_stats = bare_foreign.partition_statistics(None)?; - assert_eq!(bare_stats.as_ref(), &Statistics::new_unknown(&schema)); - - // Plans with statistics round-trip them faithfully, including - // ScalarValue-typed min/max. - let original_stats = Statistics { + let statistics = Statistics { num_rows: Precision::Exact(7), total_byte_size: Precision::Inexact(128), column_statistics: vec![ColumnStatistics { @@ -706,18 +711,68 @@ pub mod tests { byte_size: Precision::Exact(28), }], }; - let stats_plan = Arc::new( - EmptyExec::new(Arc::clone(&schema)).with_statistics(original_stats.clone()), + (schema, statistics) + } + + /// Statistics survive an FFI round trip when queried through the + /// **deprecated** `partition_statistics` entry point on the foreign plan. + #[test] + #[expect(deprecated)] + fn test_ffi_execution_plan_partition_statistics_round_trip() -> Result<()> { + let (schema, original_stats) = stats_round_trip_fixture(); + + // A plan without explicit statistics reports new_unknown. + let bare = export_empty_exec_over_ffi(&schema, None)?; + assert_eq!( + bare.partition_statistics(None)?.as_ref(), + &Statistics::new_unknown(&schema) + ); + + // A plan with statistics round-trips them for overall and per-partition queries. + let with_stats = + export_empty_exec_over_ffi(&schema, Some(original_stats.clone()))?; + assert_eq!( + with_stats.partition_statistics(None)?.as_ref(), + &original_stats + ); + assert_eq!( + with_stats.partition_statistics(Some(1))?.as_ref(), + &original_stats ); - let mut stats_local = FFI_ExecutionPlan::new(stats_plan, None); - stats_local.library_marker_id = crate::mock_foreign_marker_id; - let stats_foreign: Arc = (&stats_local).try_into()?; - let observed = stats_foreign.partition_statistics(None)?; - assert_eq!(observed.as_ref(), &original_stats); + Ok(()) + } - let observed_partition = stats_foreign.partition_statistics(Some(1))?; - assert_eq!(observed_partition.as_ref(), &original_stats); + /// Same round trip as + /// [`test_ffi_execution_plan_partition_statistics_round_trip`], but queried + /// through the **new** `statistics_with_args` entry point. + #[test] + fn test_ffi_execution_plan_statistics_with_args_round_trip() -> Result<()> { + let (schema, original_stats) = stats_round_trip_fixture(); + + // A plan without explicit statistics reports new_unknown. + let bare = export_empty_exec_over_ffi(&schema, None)?; + assert_eq!( + bare.statistics_with_args(&StatisticsArgs::new(None))? + .as_ref(), + &Statistics::new_unknown(&schema) + ); + + // A plan with statistics round-trips them for overall and per-partition queries. + let with_stats = + export_empty_exec_over_ffi(&schema, Some(original_stats.clone()))?; + assert_eq!( + with_stats + .statistics_with_args(&StatisticsArgs::new(None))? + .as_ref(), + &original_stats + ); + assert_eq!( + with_stats + .statistics_with_args(&StatisticsArgs::new(Some(1)))? + .as_ref(), + &original_stats + ); Ok(()) } diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index af686571d98f0..ad472df3379d6 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -107,6 +107,14 @@ Existing implementations of `partition_statistics` continue to work unchanged. The default `statistics_with_args` delegates to the deprecated method, so no migration is required until the deprecated method is removed. +> **Warning:** The delegation between the two methods is **one-way**. The default +> `statistics_with_args` calls `partition_statistics`, but the default +> `partition_statistics` does **not** call `statistics_with_args` — it returns +> `Statistics::new_unknown`. So if you migrate a node to override **only** +> `statistics_with_args`, any remaining caller of the deprecated +> `partition_statistics` on that node will silently receive +> `Statistics::new_unknown` instead of your real statistics. + **Who is affected:** - Users who implement custom `ExecutionPlan` nodes (recommended to migrate)