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-22157-c8b784a01f5d0bcbe0dac806730fb61afc0be8ef in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 1a416dafc9be2619337a16b7a8169d2883dba5a3 Author: Liam Feehery <[email protected]> AuthorDate: Tue May 19 09:49:48 2026 -0400 Expose `ExecutionPlan` statistics across the FFI boundary (#22157) ## Which issue does this PR close? - Closes #22152 ## Rationale for this change `ExecutionPlan::partition_statistics` and `TableProvider::statistics` are not currently transported across the DataFusion FFI boundary, so foreign plans and providers always report `Statistics::new_unknown` / `None`. This blocks optimizer rules that depend on statistics (e.g. join reordering, partition pruning) from working with out-of-process plugins, which defeats the point of exposing those hooks to plugin authors. `Statistics` contains `Precision<ScalarValue>` for column min/max/sum. `ScalarValue` is a large enum that's impractical to mirror in `#[repr(C)]`, so I reuse the existing `datafusion_proto_common::Statistics` prost encoding — the same pattern this crate already uses for filter expressions. ## What changes are included in this PR? - New `datafusion_ffi::statistics` module with `[de]serialize_statistics` helpers wrapping the`datafusion_proto_common::Statistics` round-trip. - New `partition_statistics` field on `FFI_ExecutionPlan` and corresponding `ExecutionPlan::partition_statistics` impl on `ForeignExecutionPlan` - New `statistics` field on `FFI_TableProvider` and corresponding `TableProvider::statistics` impl on `ForeignTableProvider`. Since the trait returns `Option<Statistics>`, the implementation cannot propagate decode errors, it logs a `log::warn!` and triggers a `debug_assert!`. This PR is expected to be merged after #22136 so it includes those changes. ## Are these changes tested? Yes: - Unit tests in `statistics.rs` cover three round-trip cases: `Statistics::new_unknown`, fully-exact statistics with `ScalarValue::Int32`/`Int64`/`Utf8` min/max/sum, and mixed `Precision::Exact`/`Inexact`/`Absent` values. - A new round-trip integration test in `execution_plan.rs` exercises `ForeignExecutionPlan::partition_statistics` with both `None` and `Some(idx)` partitions, against a plan with no statistics (returns `Statistics::new_unknown`) and a plan with concrete statistics. - A new round-trip integration test in `table_provider.rs` uses a thin `TableWithStats` wrapper over `MemTable` to verify both the `None` path and the concrete `Statistics` path through `ForeignTableProvider::statistics`. ## Are there any user-facing changes? This is a breaking ABI change for the `datafusion-ffi` crate: - `FFI_ExecutionPlan` gains a `partition_statistics` field. - `FFI_TableProvider` gains a `statistics` field. Plugins compiled against earlier versions of `datafusion-ffi` will need to be recompiled. There are no breaking changes to the Rust trait surface or to `Statistics` itself; downstream `ExecutionPlan` / `TableProvider` implementations require no changes. --- datafusion/ffi/src/execution_plan.rs | 96 +++++++++++++++++++- datafusion/ffi/src/lib.rs | 1 + datafusion/ffi/src/physical_expr/metrics.rs | 14 ++- datafusion/ffi/src/statistics.rs | 124 ++++++++++++++++++++++++++ datafusion/ffi/src/table_provider.rs | 131 ++++++++++++++++++++++++++++ datafusion/ffi/src/tests/mod.rs | 95 ++++++++++++++++++++ datafusion/ffi/tests/ffi_execution_plan.rs | 34 ++++++++ datafusion/ffi/tests/ffi_integration.rs | 15 ++++ 8 files changed, 500 insertions(+), 10 deletions(-) diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 1a8c9767fb..ddad605081 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{DataFusionError, Result}; +use datafusion_common::{DataFusionError, Result, Statistics}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr_common::metrics::MetricsSet; use datafusion_physical_plan::{ @@ -36,6 +36,7 @@ use crate::execution::FFI_TaskContext; use crate::physical_expr::metrics::FFI_MetricsSet; use crate::plan_properties::FFI_PlanProperties; use crate::record_batch_stream::FFI_RecordBatchStream; +use crate::statistics::{deserialize_statistics, serialize_statistics}; use crate::util::{FFI_Option, FFI_Result}; use crate::{df_result, sresult, sresult_return}; @@ -74,6 +75,15 @@ pub struct FFI_ExecutionPlan { /// underlying [`ExecutionPlan::metrics`] returned `None`. pub metrics: unsafe extern "C" fn(plan: &Self) -> FFI_Option<FFI_MetricsSet>, + /// Snapshot partition statistics. `partition == None` corresponds to + /// statistics over all partitions; `Some(idx)` corresponds to a specific + /// partition. The returned bytes are a prost-encoded + /// `datafusion_proto_common::Statistics`. + pub partition_statistics: unsafe extern "C" fn( + plan: &Self, + partition: FFI_Option<usize>, + ) -> FFI_Result<SVec<u8>>, + /// Used to create a clone on the provider of the execution plan. This should /// only need to be called by the receiver of the plan. pub clone: unsafe extern "C" fn(plan: &Self) -> Self, @@ -195,6 +205,17 @@ unsafe extern "C" fn metrics_fn_wrapper( .into() } +unsafe extern "C" fn partition_statistics_fn_wrapper( + plan: &FFI_ExecutionPlan, + partition: FFI_Option<usize>, +) -> FFI_Result<SVec<u8>> { + let partition: Option<usize> = partition.into(); + plan.inner() + .partition_statistics(partition) + .map(|stats| SVec::from(serialize_statistics(stats.as_ref()).as_slice())) + .into() +} + unsafe extern "C" fn release_fn_wrapper(plan: &mut FFI_ExecutionPlan) { unsafe { debug_assert!(!plan.private_data.is_null()); @@ -287,6 +308,7 @@ impl FFI_ExecutionPlan { execute: execute_fn_wrapper, repartitioned: repartitioned_fn_wrapper, metrics: metrics_fn_wrapper, + partition_statistics: partition_statistics_fn_wrapper, clone: clone_fn_wrapper, release: release_fn_wrapper, private_data: Box::into_raw(private_data) as *mut c_void, @@ -454,6 +476,13 @@ impl ExecutionPlan for ForeignExecutionPlan { unsafe { (self.plan.metrics)(&self.plan) }.into(); ffi.map(MetricsSet::from) } + + fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> { + let bytes = df_result!(unsafe { + (self.plan.partition_statistics)(&self.plan, partition.into()) + })?; + Ok(Arc::new(deserialize_statistics(bytes.as_slice())?)) + } } #[cfg(any(test, feature = "integration-tests"))] @@ -468,6 +497,7 @@ pub mod tests { props: Arc<PlanProperties>, children: Vec<Arc<dyn ExecutionPlan>>, metrics: Option<MetricsSet>, + statistics: Option<Statistics>, } impl EmptyExec { @@ -481,6 +511,7 @@ pub mod tests { )), children: Vec::default(), metrics: None, + statistics: None, } } @@ -488,6 +519,11 @@ pub mod tests { self.metrics = Some(metrics); self } + + pub fn with_statistics(mut self, statistics: Statistics) -> Self { + self.statistics = Some(statistics); + self + } } impl DisplayAs for EmptyExec { @@ -521,6 +557,7 @@ pub mod tests { props: Arc::clone(&self.props), children, metrics: self.metrics.clone(), + statistics: self.statistics.clone(), })) } @@ -536,6 +573,15 @@ pub mod tests { self.metrics.clone() } + fn partition_statistics( + &self, + _partition: Option<usize>, + ) -> Result<Arc<Statistics>> { + Ok(Arc::new(self.statistics.clone().unwrap_or_else(|| { + Statistics::new_unknown(self.props.eq_properties.schema()) + }))) + } + fn apply_expressions( &self, f: &mut dyn FnMut( @@ -659,6 +705,54 @@ pub mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_partition_statistics_round_trip() -> Result<()> { + 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<dyn ExecutionPlan> = (&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 { + num_rows: Precision::Exact(7), + total_byte_size: Precision::Inexact(128), + column_statistics: vec![ColumnStatistics { + null_count: Precision::Exact(1), + max_value: Precision::Exact(ScalarValue::Int32(Some(10))), + min_value: Precision::Exact(ScalarValue::Int32(Some(-3))), + sum_value: Precision::Absent, + distinct_count: Precision::Inexact(6), + byte_size: Precision::Exact(28), + }], + }; + let stats_plan = Arc::new( + EmptyExec::new(Arc::clone(&schema)).with_statistics(original_stats.clone()), + ); + let mut stats_local = FFI_ExecutionPlan::new(stats_plan, None); + stats_local.library_marker_id = crate::mock_foreign_marker_id; + let stats_foreign: Arc<dyn ExecutionPlan> = (&stats_local).try_into()?; + + let observed = stats_foreign.partition_statistics(None)?; + assert_eq!(observed.as_ref(), &original_stats); + + let observed_partition = stats_foreign.partition_statistics(Some(1))?; + assert_eq!(observed_partition.as_ref(), &original_stats); + + Ok(()) + } + #[test] fn test_ffi_execution_plan_local_bypass() { let schema = Arc::new(arrow::datatypes::Schema::new(vec![ diff --git a/datafusion/ffi/src/lib.rs b/datafusion/ffi/src/lib.rs index de3caf8c17..4df6c4b570 100644 --- a/datafusion/ffi/src/lib.rs +++ b/datafusion/ffi/src/lib.rs @@ -41,6 +41,7 @@ pub mod proto; pub mod record_batch_stream; pub mod schema_provider; pub mod session; +pub mod statistics; pub mod table_provider; pub mod table_provider_factory; pub mod table_source; diff --git a/datafusion/ffi/src/physical_expr/metrics.rs b/datafusion/ffi/src/physical_expr/metrics.rs index 6c29bd0ea6..ebef728e05 100644 --- a/datafusion/ffi/src/physical_expr/metrics.rs +++ b/datafusion/ffi/src/physical_expr/metrics.rs @@ -59,7 +59,7 @@ pub struct FFI_MetricsSet { pub struct FFI_Metric { pub value: FFI_MetricValue, pub labels: SVec<FFI_Label>, - pub partition: FFI_Option<u64>, + pub partition: FFI_Option<usize>, pub metric_type: FFI_MetricType, pub metric_category: FFI_Option<FFI_MetricCategory>, } @@ -203,7 +203,7 @@ impl From<&Metric> for FFI_Metric { Self { value: FFI_MetricValue::from(m.value()), labels: m.labels().iter().map(FFI_Label::from).collect(), - partition: m.partition().map(|p| p as u64).into(), + partition: m.partition().into(), metric_type: m.metric_type().into(), metric_category: m.metric_category().map(FFI_MetricCategory::from).into(), } @@ -213,14 +213,10 @@ impl From<&Metric> for FFI_Metric { impl From<FFI_Metric> for Metric { fn from(m: FFI_Metric) -> Self { let labels: Vec<Label> = m.labels.into_iter().map(Label::from).collect(); - let partition: Option<u64> = m.partition.into(); + let partition: Option<usize> = m.partition.into(); let category: Option<FFI_MetricCategory> = m.metric_category.into(); - let mut metric = Metric::new_with_labels( - m.value.into(), - partition.map(|p| p as usize), - labels, - ) - .with_type(m.metric_type.into()); + let mut metric = Metric::new_with_labels(m.value.into(), partition, labels) + .with_type(m.metric_type.into()); if let Some(c) = category { metric = metric.with_category(c.into()); } diff --git a/datafusion/ffi/src/statistics.rs b/datafusion/ffi/src/statistics.rs new file mode 100644 index 0000000000..019a8b2f22 --- /dev/null +++ b/datafusion/ffi/src/statistics.rs @@ -0,0 +1,124 @@ +// 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. + +//! Helpers for moving [`Statistics`] across the FFI boundary as prost-encoded +//! `datafusion_proto_common::Statistics` bytes. +//! +//! [`Statistics`] contains [`Precision<ScalarValue>`] for column min/max/sum, +//! and `ScalarValue` is a large enum that's impractical to mirror in +//! `#[repr(C)]`. The proto round-trip already exists in `datafusion-proto-common` +//! and is the same pattern used to ship filter expressions across the FFI +//! boundary, so we reuse it here. +//! +//! [`Precision<ScalarValue>`]: datafusion_common::stats::Precision + +use datafusion_common::{DataFusionError, Result, Statistics}; +use prost::Message; + +/// Serialize [`Statistics`] to prost-encoded +/// `datafusion_proto_common::Statistics` bytes. +pub(crate) fn serialize_statistics(stats: &Statistics) -> Vec<u8> { + datafusion_proto_common::Statistics::from(stats).encode_to_vec() +} + +/// Decode prost-encoded `datafusion_proto_common::Statistics` bytes back into +/// [`Statistics`]. +pub(crate) fn deserialize_statistics(bytes: &[u8]) -> Result<Statistics> { + let proto = datafusion_proto_common::Statistics::decode(bytes).map_err(|e| { + DataFusionError::Plan(format!("failed to decode Statistics: {e}")) + })?; + Statistics::try_from(&proto) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::ScalarValue; + use datafusion_common::stats::Precision; + use datafusion_common::{ColumnStatistics, Statistics}; + + use super::*; + + #[test] + fn round_trip_unknown_statistics() { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let original = Statistics::new_unknown(&Arc::new(schema)); + + let bytes = serialize_statistics(&original); + let observed = deserialize_statistics(&bytes).expect("decode"); + + assert_eq!(observed, original); + } + + #[test] + fn round_trip_exact_statistics_with_scalar_values() { + let original = Statistics { + num_rows: Precision::Exact(100), + total_byte_size: Precision::Exact(4096), + column_statistics: vec![ + ColumnStatistics { + null_count: Precision::Exact(2), + max_value: Precision::Exact(ScalarValue::Int32(Some(50))), + min_value: Precision::Exact(ScalarValue::Int32(Some(-10))), + sum_value: Precision::Exact(ScalarValue::Int64(Some(1234))), + distinct_count: Precision::Exact(40), + byte_size: Precision::Exact(800), + }, + ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Exact(ScalarValue::Utf8(Some( + "zebra".to_string(), + ))), + min_value: Precision::Exact(ScalarValue::Utf8(Some( + "ant".to_string(), + ))), + sum_value: Precision::Absent, + distinct_count: Precision::Inexact(95), + byte_size: Precision::Inexact(2048), + }, + ], + }; + + let bytes = serialize_statistics(&original); + let observed = deserialize_statistics(&bytes).expect("decode"); + + assert_eq!(observed, original); + } + + #[test] + fn round_trip_mixed_precision() { + let original = Statistics { + num_rows: Precision::Inexact(42), + total_byte_size: Precision::Absent, + column_statistics: vec![ColumnStatistics { + null_count: Precision::Absent, + max_value: Precision::Inexact(ScalarValue::Float64(Some(1.5))), + min_value: Precision::Absent, + sum_value: Precision::Inexact(ScalarValue::Float64(Some(63.0))), + distinct_count: Precision::Absent, + byte_size: Precision::Absent, + }], + }; + + let bytes = serialize_statistics(&original); + let observed = deserialize_statistics(&bytes).expect("decode"); + + assert_eq!(observed, original); + } +} diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index b6c077526c..5a4b2fa272 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -22,6 +22,7 @@ use arrow::datatypes::SchemaRef; use async_ffi::{FfiFuture, FutureExt}; use async_trait::async_trait; use datafusion_catalog::{Session, TableProvider}; +use datafusion_common::Statistics; use datafusion_common::error::{DataFusionError, Result}; use datafusion_execution::TaskContext; use datafusion_expr::dml::InsertOp; @@ -44,6 +45,7 @@ use crate::arrow_wrappers::WrappedSchema; use crate::execution::FFI_TaskContextProvider; use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use crate::session::{FFI_SessionRef, ForeignSession}; +use crate::statistics::{deserialize_statistics, serialize_statistics}; use crate::table_source::{FFI_TableProviderFilterPushDown, FFI_TableType}; use crate::util::{FFI_Option, FFI_Result}; use crate::{df_result, sresult_return}; @@ -133,6 +135,11 @@ pub struct FFI_TableProvider { insert_op: FFI_InsertOp, ) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>>, + /// Snapshot the provider's table-level statistics. [`FFI_Option::None`] + /// corresponds to [`TableProvider::statistics`] returning `None`; + /// `Some(bytes)` is a prost-encoded `datafusion_proto_common::Statistics`. + pub statistics: unsafe extern "C" fn(provider: &Self) -> FFI_Option<SVec<u8>>, + pub logical_codec: FFI_LogicalExtensionCodec, /// Used to create a clone on the provider of the execution plan. This should @@ -179,6 +186,16 @@ unsafe extern "C" fn schema_fn_wrapper(provider: &FFI_TableProvider) -> WrappedS provider.inner().schema().into() } +unsafe extern "C" fn statistics_fn_wrapper( + provider: &FFI_TableProvider, +) -> FFI_Option<SVec<u8>> { + let serialized: Option<SVec<u8>> = provider + .inner() + .statistics() + .map(|s| SVec::from(&*serialize_statistics(&s))); + serialized.into() +} + unsafe extern "C" fn table_type_fn_wrapper( provider: &FFI_TableProvider, ) -> FFI_TableType { @@ -344,6 +361,7 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_TableProvider) -> FFI_Table table_type: table_type_fn_wrapper, supports_filters_pushdown: provider.supports_filters_pushdown, insert_into: provider.insert_into, + statistics: statistics_fn_wrapper, logical_codec: provider.logical_codec.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -404,6 +422,7 @@ impl FFI_TableProvider { false => None, }, insert_into: insert_into_fn_wrapper, + statistics: statistics_fn_wrapper, logical_codec, clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -451,6 +470,21 @@ impl TableProvider for ForeignTableProvider { unsafe { (self.0.table_type)(&self.0).into() } } + fn statistics(&self) -> Option<Statistics> { + let ffi_opt = unsafe { (self.0.statistics)(&self.0) }; + let bytes: Option<SVec<u8>> = ffi_opt.into(); + let bytes = bytes?; + match deserialize_statistics(bytes.as_slice()) { + Ok(stats) => Some(stats), + Err(e) => { + log::warn!("Failed to deserialize FFI statistics: {e}"); + // Fires in debug builds to surface encoding bugs early; callers see None. + debug_assert!(false, "Failed to deserialize FFI statistics: {e}"); + None + } + } + } + async fn scan( &self, session: &dyn Session, @@ -772,4 +806,101 @@ mod tests { Ok(()) } + + #[test] + fn test_ffi_table_provider_statistics_round_trip() -> Result<()> { + use arrow::datatypes::{DataType, Field}; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::MemTable; + 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( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + + 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); + + // Provider without statistics should cross the boundary as None. + let no_stats_inner = Arc::new(MemTable::try_new( + Arc::clone(&schema), + vec![vec![batch.clone()]], + )?); + let no_stats_provider = Arc::new(TableWithStats { + inner: no_stats_inner, + stats: None, + }); + let mut ffi_provider = FFI_TableProvider::new( + no_stats_provider, + true, + None, + task_ctx_provider.clone(), + None, + ); + ffi_provider.library_marker_id = crate::mock_foreign_marker_id; + let foreign: Arc<dyn TableProvider> = (&ffi_provider).into(); + assert!(foreign.statistics().is_none()); + + // Provider with statistics should round-trip faithfully. + let original_stats = Statistics { + num_rows: Precision::Exact(3), + total_byte_size: Precision::Inexact(12), + column_statistics: vec![ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Exact(ScalarValue::Int32(Some(3))), + min_value: Precision::Exact(ScalarValue::Int32(Some(1))), + sum_value: Precision::Exact(ScalarValue::Int64(Some(6))), + distinct_count: Precision::Exact(3), + byte_size: Precision::Exact(12), + }], + }; + 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 mut ffi_provider = + FFI_TableProvider::new(stats_provider, true, None, task_ctx_provider, None); + ffi_provider.library_marker_id = crate::mock_foreign_marker_id; + let foreign: Arc<dyn TableProvider> = (&ffi_provider).into(); + assert_eq!(foreign.statistics().as_ref(), Some(&original_stats)); + + Ok(()) + } } diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 41fdd2699a..62e62d8235 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -20,8 +20,16 @@ use std::sync::Arc; use arrow::array::RecordBatch; use arrow_schema::{DataType, Field, Schema}; use async_provider::create_async_table_provider; +use async_trait::async_trait; use catalog::create_catalog_provider; +use datafusion_catalog::MemTable; +use datafusion_catalog::{Session, TableProvider}; use datafusion_common::record_batch; +use datafusion_common::stats::Precision; +use datafusion_common::{ColumnStatistics, Statistics}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{Expr, TableType}; +use datafusion_physical_plan::ExecutionPlan; use sync_provider::create_sync_table_provider; use udf_udaf_udwf::{ create_ffi_abs_func, create_ffi_random_func, create_ffi_rank_func, @@ -98,6 +106,11 @@ pub struct ForeignLibraryModule { pub create_empty_exec: extern "C" fn() -> FFI_ExecutionPlan, + pub create_exec_with_statistics: extern "C" fn() -> FFI_ExecutionPlan, + + pub create_table_with_statistics: + extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider, + pub create_physical_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, pub version: extern "C" fn() -> u64, @@ -145,6 +158,86 @@ pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan { FFI_ExecutionPlan::new(plan, None) } +/// Returns canonical statistics used by both the producer and consumer sides of +/// the integration tests so round-trips can be asserted without hard-coding +/// the values in two places. +pub fn make_test_statistics() -> Statistics { + Statistics { + num_rows: Precision::Exact(42), + total_byte_size: Precision::Exact(672), + column_statistics: vec![ + ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Exact(ScalarValue::Int32(Some(100))), + min_value: Precision::Exact(ScalarValue::Int32(Some(-10))), + sum_value: Precision::Exact(ScalarValue::Int64(Some(1890))), + distinct_count: Precision::Inexact(40), + byte_size: Precision::Exact(168), + }, + ColumnStatistics { + null_count: Precision::Exact(1), + max_value: Precision::Exact(ScalarValue::Float64(Some(99.5))), + min_value: Precision::Exact(ScalarValue::Float64(Some(-1.5))), + sum_value: Precision::Absent, + distinct_count: Precision::Absent, + byte_size: Precision::Exact(328), + }, + ], + } +} + +pub(crate) extern "C" fn create_exec_with_statistics() -> FFI_ExecutionPlan { + let schema = create_test_schema(); + let plan = Arc::new(EmptyExec::new(schema).with_statistics(make_test_statistics())); + FFI_ExecutionPlan::new(plan, None) +} + +/// Thin wrapper that attaches a fixed [`Statistics`] snapshot to any inner +/// [`TableProvider`] without changing its scan behaviour. +#[derive(Debug)] +struct TableWithStats { + inner: Arc<dyn TableProvider>, + stats: Statistics, +} + +#[async_trait] +impl TableProvider for TableWithStats { + fn schema(&self) -> arrow_schema::SchemaRef { + self.inner.schema() + } + + fn table_type(&self) -> TableType { + self.inner.table_type() + } + + fn statistics(&self) -> Option<Statistics> { + Some(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 + } +} + +pub(crate) extern "C" fn create_table_with_statistics( + codec: FFI_LogicalExtensionCodec, +) -> FFI_TableProvider { + let schema = create_test_schema(); + let batch = create_record_batch(1, 5); + let inner = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let provider = Arc::new(TableWithStats { + inner, + stats: make_test_statistics(), + }); + FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec) +} + /// This defines the entry point for using the module. #[unsafe(no_mangle)] pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { @@ -162,6 +255,8 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_rank_udwf: create_ffi_rank_func, create_extension_options: config::create_extension_options, create_empty_exec, + create_exec_with_statistics, + create_table_with_statistics, create_physical_optimizer_rule: physical_optimizer::create_physical_optimizer_rule, version: super::version, diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index 66a5ce4b0a..bd84e064de 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -28,6 +28,40 @@ mod tests { use datafusion_physical_plan::ExecutionPlan; use std::sync::Arc; + #[test] + fn test_ffi_execution_plan_partition_statistics_cross_library() + -> Result<(), DataFusionError> { + let module = get_module()?; + + // Producer: plan with no explicit statistics → expects Statistics::new_unknown. + let bare = (module.create_empty_exec)(); + let bare: Arc<dyn ExecutionPlan> = (&bare).try_into()?; + assert!(bare.is::<ForeignExecutionPlan>()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); + let bare_stats = bare.partition_statistics(None)?; + assert_eq!( + bare_stats.as_ref(), + &datafusion_common::Statistics::new_unknown(&schema), + ); + + // Producer: plan with known statistics — round-trip through the cdylib boundary. + let expected = datafusion_ffi::tests::make_test_statistics(); + let with_stats = (module.create_exec_with_statistics)(); + let with_stats: Arc<dyn ExecutionPlan> = (&with_stats).try_into()?; + assert!(with_stats.is::<ForeignExecutionPlan>()); + + // Both None (all-partition aggregate) and Some(idx) must return the + // same statistics because EmptyExec ignores the partition argument. + let observed_all = with_stats.partition_statistics(None)?; + assert_eq!(observed_all.as_ref(), &expected); + + let observed_part = with_stats.partition_statistics(Some(0))?; + assert_eq!(observed_part.as_ref(), &expected); + + Ok(()) + } + #[test] fn test_ffi_execution_plan_new_sets_runtimes_on_children() -> Result<(), DataFusionError> { diff --git a/datafusion/ffi/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs index 4186bafc83..6a6b6b3100 100644 --- a/datafusion/ffi/tests/ffi_integration.rs +++ b/datafusion/ffi/tests/ffi_integration.rs @@ -71,6 +71,21 @@ mod tests { test_table_provider(true).await } + #[test] + fn test_ffi_table_provider_statistics_cross_library() -> Result<()> { + let module = get_module()?; + let (_, codec) = super::utils::ctx_and_codec(); + + let expected = datafusion_ffi::tests::make_test_statistics(); + + let ffi_provider = (module.create_table_with_statistics)(codec); + let foreign: Arc<dyn TableProvider> = (&ffi_provider).into(); + + assert_eq!(foreign.statistics().as_ref(), Some(&expected)); + + Ok(()) + } + #[tokio::test] async fn test_table_provider_factory() -> Result<()> { let table_provider_module = get_module()?; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
