Expose ExecutionPlan statistics across the FFI boundary

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<u64>,
+    ) -> 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<u64>,
+) -> FFI_Result<SVec<u8>> {
+    let partition: Option<usize> = Option::<u64>::from(partition).map(|p| p as usize);
+    plan.inner()
+        .partition_statistics(partition)
+        .map(|stats| serialize_statistics(stats.as_ref()).into_iter().collect())
+        .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,10 +476,24 @@ 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.map(|p| p as u64).into(),
+            )
+        })?;
+        Ok(Arc::new(deserialize_statistics(bytes.as_slice())?))
+    }
 }
 
 #[cfg(any(test, feature = "integration-tests"))]
 pub mod tests {
+    #[cfg(test)]
+    use datafusion_common::stats::Precision;
+    #[cfg(test)]
+    use datafusion_common::{ColumnStatistics, ScalarValue};
     use datafusion_physical_plan::Partitioning;
     use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 
@@ -468,6 +504,7 @@ pub mod tests {
         props: Arc<PlanProperties>,
         children: Vec<Arc<dyn ExecutionPlan>>,
         metrics: Option<MetricsSet>,
+        statistics: Option<Statistics>,
     }
 
     impl EmptyExec {
@@ -481,13 +518,19 @@ pub mod tests {
                 )),
                 children: Vec::default(),
                 metrics: None,
+                statistics: None,
             }
         }
 
         pub fn with_metrics(mut self, metrics: MetricsSet) -> Self {
             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 +564,7 @@ pub mod tests {
                 props: Arc::clone(&self.props),
                 children,
                 metrics: self.metrics.clone(),
+                statistics: self.statistics.clone(),
             }))
         }
 
@@ -536,6 +580,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 +712,51 @@ pub mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_ffi_execution_plan_partition_statistics_round_trip() -> Result<()> {
+        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![

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;

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);
+    }
+}