feat: allow scaling RangePartitioning

datafusion/ffi/src/execution_plan.rs

@@ -600,6 +600,12 @@ pub mod tests {
             self
         }
 
+        pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self {
+            self.props =
+                Arc::new(self.props.as_ref().clone().with_partitioning(partitioning));
+            self
+        }
+
         pub fn with_expressions(
             mut self,
             expressions: Vec<Arc<dyn PhysicalExpr>>,

datafusion/ffi/src/physical_expr/partitioning.rs

@@ -33,8 +33,9 @@ use crate::physical_expr::sort::FFI_PhysicalSortExpr;
 #[repr(C)]
 #[derive(Debug)]
 pub struct FFI_RangePartitioning {
-    split_points: SVec<SVec<WrappedArray>>,
+    samples: SVec<SVec<WrappedArray>>,
     ordering: SVec<FFI_PhysicalSortExpr>,
+    partition_count: usize,
 }
 
 /// A stable struct for sharing [`Partitioning`] across FFI boundaries.
@@ -62,8 +63,8 @@ impl From<&Partitioning> for FFI_Partitioning {
             }
             Partitioning::Range(range) => {
                 // Producer-side conversion should be infallible at ABI boundary
-                let split_points = range
-                    .split_points()
+                let samples = range
+                    .samples()
                     .iter()
                     .map(|split_point| {
                         split_point
@@ -83,8 +84,9 @@ impl From<&Partitioning> for FFI_Partitioning {
                     .map(FFI_PhysicalSortExpr::from)
                     .collect();
                 Self::Range(FFI_RangePartitioning {
-                    split_points,
+                    samples,
                     ordering,
+                    partition_count: range.partition_count(),
                 })
             }
             Partitioning::UnknownPartitioning(size) => Self::UnknownPartitioning(*size),
@@ -105,8 +107,8 @@ impl TryFrom<FFI_Partitioning> for Partitioning {
                 Self::Hash(exprs, size)
             }
             FFI_Partitioning::Range(range) => {
-                let split_points = range
-                    .split_points
+                let samples = range
+                    .samples
                     .into_iter()
                     .map(|split_point| {
                         split_point
@@ -126,7 +128,11 @@ impl TryFrom<FFI_Partitioning> for Partitioning {
                         )
                     })?;
 
-                Self::Range(RangePartitioning::try_new(ordering, split_points)?)
+                Self::Range(RangePartitioning::try_new_with_samples(
+                    ordering,
+                    samples,
+                    range.partition_count,
+                )?)
             }
             FFI_Partitioning::UnknownPartitioning(size) => {
                 Self::UnknownPartitioning(size)
@@ -174,13 +180,28 @@ mod tests {
         )?))
     }
 
+    fn sampled_range_partitioning() -> Result<Partitioning> {
+        let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new(
+            Column::new("a", 0),
+        ))])
+        .expect("non-empty ordering");
+        let samples = [10, 20, 30, 40, 50]
+            .into_iter()
+            .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))]))
+            .collect();
+        Ok(Partitioning::Range(
+            RangePartitioning::try_new_with_samples(ordering, samples, 3)?,
+        ))
+    }
+
     #[test]
     fn round_trip_ffi_partitioning() -> Result<()> {
         for partitioning in [
             Partitioning::RoundRobinBatch(10),
             Partitioning::Hash(vec![lit(1)], 10),
             Partitioning::UnknownPartitioning(10),
             range_partitioning()?,
+            sampled_range_partitioning()?,
         ] {
             let ffi_partitioning: FFI_Partitioning = (&partitioning).into();
             let returned: Partitioning = ffi_partitioning.try_into()?;
@@ -210,11 +231,33 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn round_trip_ffi_sampled_range_partitioning() -> Result<()> {
+        let partitioning = sampled_range_partitioning()?;
+
+        let ffi_partitioning: FFI_Partitioning = (&partitioning).into();
+        let returned: Partitioning = ffi_partitioning.try_into()?;
+        let Partitioning::Range(returned) = returned else {
+            panic!("expected range partitioning");
+        };
+        let Partitioning::Range(original) = partitioning else {
+            panic!("expected range partitioning");
+        };
+
+        assert_eq!(returned, original);
+        assert_eq!(returned.samples(), original.samples());
+        assert_eq!(returned.max_partition_count(), 6);
+        assert_eq!(returned.partition_count(), 3);
+
+        Ok(())
+    }
+
     #[test]
     fn ffi_range_partitioning_rejects_empty_ordering() {
         let ffi_partitioning = FFI_Partitioning::Range(FFI_RangePartitioning {
-            split_points: SVec::new(),
+            samples: SVec::new(),
             ordering: SVec::new(),
+            partition_count: 1,
         });
 
         let err = Partitioning::try_from(ffi_partitioning).unwrap_err();

datafusion/ffi/src/plan_properties.rs

@@ -290,11 +290,14 @@ mod tests {
         let col = datafusion::physical_plan::expressions::col("a", &schema)?;
         let ordering = LexOrdering::new([PhysicalSortExpr::new_default(col)])
             .expect("non-empty ordering");
-        let split_points = vec![
+        let samples = vec![
             SplitPoint::new(vec![ScalarValue::Int64(Some(10))]),
             SplitPoint::new(vec![ScalarValue::Int64(Some(20))]),
+            SplitPoint::new(vec![ScalarValue::Int64(Some(30))]),
+            SplitPoint::new(vec![ScalarValue::Int64(Some(40))]),
+            SplitPoint::new(vec![ScalarValue::Int64(Some(50))]),
         ];
-        let range = RangePartitioning::try_new(ordering, split_points)?;
+        let range = RangePartitioning::try_new_with_samples(ordering, samples, 3)?;
 
         Ok(PlanProperties::new(
             EquivalenceProperties::new(schema),
@@ -314,7 +317,6 @@ mod tests {
         let foreign_props: PlanProperties = local_props_ptr.try_into()?;
 
         assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}"));
-
         Ok(())
     }
 
@@ -351,6 +353,10 @@ mod tests {
             format!("{:?}", original_props.output_partitioning())
         );
         assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}"));
+        let Partitioning::Range(range) = foreign_props.output_partitioning() else {
+            panic!("expected range partitioning");
+        };
+        assert_eq!(range.max_partition_count(), 6);
 
         Ok(())
     }

datafusion/ffi/src/tests/mod.rs

@@ -27,10 +27,13 @@ use datafusion_catalog::MemTable;
 use datafusion_catalog::{Session, TableProvider};
 use datafusion_common::stats::Precision;
 use datafusion_common::{ColumnStatistics, Statistics};
-use datafusion_common::{Result, ScalarValue, exec_err};
+use datafusion_common::{Result, ScalarValue, SplitPoint, exec_err};
 use datafusion_expr::{Expr, TableType, col, lit};
-use datafusion_physical_expr::PhysicalExpr;
-use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr::{
+    LexOrdering, PhysicalExpr, PhysicalSortExpr, RangePartitioning,
+};
+use datafusion_physical_plan::{ExecutionPlan, Partitioning};
 use sync_provider::create_sync_table_provider;
 use udf_udaf_udwf::{
     create_ffi_abs_func, create_ffi_first_value_func, create_ffi_random_func,
@@ -228,7 +231,22 @@ pub fn make_test_statistics() -> Statistics {
 
 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()));
+    let ordering =
+        LexOrdering::new([PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))])
+            .expect("non-empty ordering");
+    let samples = [10, 20, 30, 40, 50]
+        .into_iter()
+        .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))]))
+        .collect();
+    let partitioning = Partitioning::Range(
+        RangePartitioning::try_new_with_samples(ordering, samples, 3)
+            .expect("valid sampled range partitioning"),
+    );
+    let plan = Arc::new(
+        EmptyExec::new(schema)
+            .with_statistics(make_test_statistics())
+            .with_partitioning(partitioning),
+    );
     FFI_ExecutionPlan::new(plan, None)
 }
 

datafusion/ffi/tests/ffi_execution_plan.rs

@@ -29,7 +29,7 @@ mod tests {
     use datafusion_ffi::tests::utils::get_module;
     use datafusion_physical_plan::execution_plan::InvariantLevel;
     use datafusion_physical_plan::{
-        ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions,
+        ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions,
     };
     use std::sync::Arc;
 
@@ -68,6 +68,30 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_ffi_range_partitioning_cross_library() -> Result<(), DataFusionError> {
+        let module = get_module()?;
+        let plan = (module.create_exec_with_statistics)();
+        let plan: Arc<dyn ExecutionPlan> = (&plan).try_into()?;
+        let Partitioning::Range(range) = plan.properties().output_partitioning() else {
+            panic!("expected range partitioning");
+        };
+
+        assert_eq!(range.partition_count(), 3);
+        assert_eq!(range.max_partition_count(), 6);
+        assert_eq!(range.samples().len(), 5);
+        assert_eq!(
+            range
+                .split_points()
+                .iter()
+                .map(|point| point.to_string())
+                .collect::<Vec<_>>(),
+            vec!["(20)", "(40)"]
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn test_ffi_execution_plan_expressions_cross_library() -> Result<(), DataFusionError>
     {