This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 39e2f23a3f feat: Range Partitioning FFI (#23520)
39e2f23a3f is described below
commit 39e2f23a3fbeebf4eebc3160d91a29c6a7286be7
Author: Saad Tajwar <[email protected]>
AuthorDate: Wed Jul 22 07:01:29 2026 -0400
feat: Range Partitioning FFI (#23520)
## Which issue does this PR close?
- Closes #22394
## Rationale for this change
Exposing range partition metadata via the FFI for external consumers.
## What changes are included in this PR?
- Added FFI mirror struct for `RangePartitioning` and added new enum
variant for range in `FFI_Partitioning`
- For native -> FFI, added match arm for the new variant, same with FFI
-> native but changed the approach of `From` -> `TryFrom` to utilize the
validation for `RangePartitioning` and modified `plan_properties` to
match
- Added tests
## Are these changes tested?
Yes
## Are there any user-facing changes?
Yes, exposing Range partitioning over FFI. This exposes a new `Range`
variant in the `FFI_Partitioning` enum, which may cause consumers of
this enum to add another arm to match statements to handle the new enum.
New `FFI_RangePartitioning` struct for the `Range` variant.
---------
Co-authored-by: Tim Saucer <[email protected]>
---
datafusion/ffi/src/physical_expr/partitioning.rs | 157 ++++++++++++++++++++---
datafusion/ffi/src/plan_properties.rs | 46 ++++++-
2 files changed, 183 insertions(+), 20 deletions(-)
diff --git a/datafusion/ffi/src/physical_expr/partitioning.rs
b/datafusion/ffi/src/physical_expr/partitioning.rs
index eec437639e..2a9a8528c6 100644
--- a/datafusion/ffi/src/physical_expr/partitioning.rs
+++ b/datafusion/ffi/src/physical_expr/partitioning.rs
@@ -17,20 +17,35 @@
use std::sync::Arc;
-use datafusion_physical_expr::Partitioning;
+use datafusion_common::{DataFusionError, ScalarValue, SplitPoint};
+use datafusion_physical_expr::{
+ LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning,
+};
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
use stabby::vec::Vec as SVec;
+use crate::arrow_wrappers::WrappedArray;
use crate::physical_expr::FFI_PhysicalExpr;
+use crate::physical_expr::sort::FFI_PhysicalSortExpr;
+
+/// A stable struct for sharing [`RangePartitioning`] across FFI boundaries.
+/// See [`RangePartitioning`] for the descriptions of each field.
+#[repr(C)]
+#[derive(Debug)]
+pub struct FFI_RangePartitioning {
+ split_points: SVec<SVec<WrappedArray>>,
+ ordering: SVec<FFI_PhysicalSortExpr>,
+}
/// A stable struct for sharing [`Partitioning`] across FFI boundaries.
-/// See ['Partitioning'] for the meaning of each variant.
+/// See [`Partitioning`] for the meaning of each variant.
#[repr(C)]
#[derive(Debug)]
pub enum FFI_Partitioning {
RoundRobinBatch(usize),
Hash(SVec<FFI_PhysicalExpr>, usize),
UnknownPartitioning(usize),
+ Range(FFI_RangePartitioning),
}
impl From<&Partitioning> for FFI_Partitioning {
@@ -45,49 +60,130 @@ impl From<&Partitioning> for FFI_Partitioning {
.collect();
Self::Hash(exprs, *size)
}
- // FFI does not yet expose range partition metadata.
- // See https://github.com/apache/datafusion/issues/22394
Partitioning::Range(range) => {
- Self::UnknownPartitioning(range.partition_count())
+ // Producer-side conversion should be infallible at ABI
boundary
+ let split_points = range
+ .split_points()
+ .iter()
+ .map(|split_point| {
+ split_point
+ .values()
+ .iter()
+ .map(|value| {
+ WrappedArray::try_from(value).expect(
+ "ScalarValue in RangePartitioning should
convert to WrappedArray",
+ )
+ })
+ .collect()
+ })
+ .collect();
+ let ordering = range
+ .ordering()
+ .iter()
+ .map(FFI_PhysicalSortExpr::from)
+ .collect();
+ Self::Range(FFI_RangePartitioning {
+ split_points,
+ ordering,
+ })
}
Partitioning::UnknownPartitioning(size) =>
Self::UnknownPartitioning(*size),
}
}
}
-impl From<&FFI_Partitioning> for Partitioning {
- fn from(value: &FFI_Partitioning) -> Self {
- match value {
+impl TryFrom<FFI_Partitioning> for Partitioning {
+ type Error = DataFusionError;
+
+ fn try_from(value: FFI_Partitioning) -> Result<Self, Self::Error> {
+ Ok(match value {
FFI_Partitioning::RoundRobinBatch(size) => {
- Partitioning::RoundRobinBatch(*size)
+ Partitioning::RoundRobinBatch(size)
}
FFI_Partitioning::Hash(exprs, size) => {
let exprs = exprs.iter().map(<Arc<dyn
PhysicalExpr>>::from).collect();
- Self::Hash(exprs, *size)
+ Self::Hash(exprs, size)
+ }
+ FFI_Partitioning::Range(range) => {
+ let split_points = range
+ .split_points
+ .into_iter()
+ .map(|split_point| {
+ split_point
+ .into_iter()
+ .map(ScalarValue::try_from)
+ .collect::<Result<Vec<_>, _>>()
+ .map(SplitPoint::new)
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+
+ let ordering =
+
LexOrdering::new(range.ordering.iter().map(PhysicalSortExpr::from))
+ .ok_or_else(|| {
+ DataFusionError::Internal(
+ "FFI Range partitioning ordering must be non-empty"
+ .to_string(),
+ )
+ })?;
+
+ Self::Range(RangePartitioning::try_new(ordering,
split_points)?)
}
FFI_Partitioning::UnknownPartitioning(size) => {
- Self::UnknownPartitioning(*size)
+ Self::UnknownPartitioning(size)
}
- }
+ })
}
}
#[cfg(test)]
mod tests {
- use datafusion_physical_expr::Partitioning;
- use datafusion_physical_expr::expressions::lit;
+ use std::sync::Arc;
+
+ use arrow_schema::SortOptions;
+ use datafusion_common::{Result, ScalarValue, SplitPoint};
+ use datafusion_physical_expr::expressions::{Column, lit};
+ use datafusion_physical_expr::{
+ LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning,
+ };
+ use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+ use stabby::vec::Vec as SVec;
- use crate::physical_expr::partitioning::FFI_Partitioning;
+ use crate::physical_expr::partitioning::{FFI_Partitioning,
FFI_RangePartitioning};
+
+ fn range_partitioning() -> Result<Partitioning> {
+ let a = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
+ let b = Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>;
+ let ordering = LexOrdering::new([
+ PhysicalSortExpr::new(a, SortOptions::default()),
+ PhysicalSortExpr::new(b, SortOptions::new(true, false)),
+ ])
+ .expect("non-empty ordering");
+ let split_points = vec![
+ SplitPoint::new(vec![
+ ScalarValue::Int64(Some(10)),
+ ScalarValue::Utf8(Some("a".to_string())),
+ ]),
+ SplitPoint::new(vec![
+ ScalarValue::Int64(Some(20)),
+ ScalarValue::Utf8(Some("b".to_string())),
+ ]),
+ ];
+ Ok(Partitioning::Range(RangePartitioning::try_new(
+ ordering,
+ split_points,
+ )?))
+ }
#[test]
- fn round_trip_ffi_partitioning() {
+ fn round_trip_ffi_partitioning() -> Result<()> {
for partitioning in [
Partitioning::RoundRobinBatch(10),
Partitioning::Hash(vec![lit(1)], 10),
Partitioning::UnknownPartitioning(10),
+ range_partitioning()?,
] {
let ffi_partitioning: FFI_Partitioning = (&partitioning).into();
- let returned: Partitioning = (&ffi_partitioning).into();
+ let returned: Partitioning = ffi_partitioning.try_into()?;
if let Partitioning::UnknownPartitioning(return_size) = returned {
let Partitioning::UnknownPartitioning(original_size) =
partitioning
@@ -99,5 +195,32 @@ mod tests {
assert_eq!(partitioning, returned);
}
}
+
+ Ok(())
+ }
+
+ #[test]
+ fn round_trip_ffi_range_partitioning_compound_key() -> Result<()> {
+ let partitioning = range_partitioning()?;
+
+ let ffi_partitioning: FFI_Partitioning = (&partitioning).into();
+ let returned: Partitioning = ffi_partitioning.try_into()?;
+ assert_eq!(partitioning, returned);
+
+ Ok(())
+ }
+
+ #[test]
+ fn ffi_range_partitioning_rejects_empty_ordering() {
+ let ffi_partitioning = FFI_Partitioning::Range(FFI_RangePartitioning {
+ split_points: SVec::new(),
+ ordering: SVec::new(),
+ });
+
+ let err = Partitioning::try_from(ffi_partitioning).unwrap_err();
+ assert!(
+ err.to_string().contains("ordering must be non-empty"),
+ "{err}"
+ );
}
}
diff --git a/datafusion/ffi/src/plan_properties.rs
b/datafusion/ffi/src/plan_properties.rs
index b286ee2d7d..09ef26af32 100644
--- a/datafusion/ffi/src/plan_properties.rs
+++ b/datafusion/ffi/src/plan_properties.rs
@@ -20,7 +20,7 @@ use std::sync::Arc;
use arrow::datatypes::SchemaRef;
use datafusion_common::error::{DataFusionError, Result};
-use datafusion_physical_expr::EquivalenceProperties;
+use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
use datafusion_physical_plan::PlanProperties;
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
@@ -172,6 +172,7 @@ impl TryFrom<FFI_PlanProperties> for PlanProperties {
.unwrap_or_default();
let partitioning = unsafe {
(ffi_props.output_partitioning)(&ffi_props) };
+ let partitioning = Partitioning::try_from(partitioning)?;
let eq_properties = if sort_exprs.is_empty() {
EquivalenceProperties::new(Arc::new(schema))
@@ -187,7 +188,7 @@ impl TryFrom<FFI_PlanProperties> for PlanProperties {
Ok(PlanProperties::new(
eq_properties,
- (&partitioning).into(),
+ partitioning,
emission_type,
boundedness,
))
@@ -260,13 +261,15 @@ impl From<FFI_EmissionType> for EmissionType {
#[cfg(test)]
mod tests {
+ use arrow::datatypes::{DataType, Field, Schema};
use datafusion::physical_expr::PhysicalSortExpr;
use datafusion::physical_plan::Partitioning;
+ use datafusion_common::{ScalarValue, SplitPoint};
+ use datafusion_physical_expr::{LexOrdering, RangePartitioning};
use super::*;
fn create_test_props() -> Result<PlanProperties> {
- use arrow::datatypes::{DataType, Field, Schema};
let schema =
Arc::new(Schema::new(vec![Field::new("a", DataType::Float32,
false)]));
@@ -282,6 +285,25 @@ mod tests {
))
}
+ fn create_range_test_props() -> Result<PlanProperties> {
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int64, false)]));
+ 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![
+ SplitPoint::new(vec![ScalarValue::Int64(Some(10))]),
+ SplitPoint::new(vec![ScalarValue::Int64(Some(20))]),
+ ];
+ let range = RangePartitioning::try_new(ordering, split_points)?;
+
+ Ok(PlanProperties::new(
+ EquivalenceProperties::new(schema),
+ Partitioning::Range(range),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ))
+ }
+
#[test]
fn test_round_trip_ffi_plan_properties() -> Result<()> {
let original_props = create_test_props()?;
@@ -314,4 +336,22 @@ mod tests {
Ok(())
}
+
+ #[test]
+ fn test_round_trip_ffi_plan_properties_range_partitioning() -> Result<()> {
+ let original_props = create_range_test_props()?;
+
+ let mut local_props_ptr = FFI_PlanProperties::from(&original_props);
+ local_props_ptr.library_marker_id = crate::mock_foreign_marker_id;
+
+ let foreign_props: PlanProperties = local_props_ptr.try_into()?;
+
+ assert_eq!(
+ format!("{:?}", foreign_props.output_partitioning()),
+ format!("{:?}", original_props.output_partitioning())
+ );
+ assert_eq!(format!("{foreign_props:?}"),
format!("{original_props:?}"));
+
+ Ok(())
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
