diff --git a/datafusion/ffi/src/physical_expr/partitioning.rs b/datafusion/ffi/src/physical_expr/partitioning.rs index eec437639e156..872614d3669d5 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>, + ordering: SVec, +} /// 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, usize), UnknownPartitioning(usize), + Range(FFI_RangePartitioning), } impl From<&Partitioning> for FFI_Partitioning { @@ -45,49 +60,120 @@ 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 for Partitioning { + type Error = DataFusionError; + + fn try_from(value: FFI_Partitioning) -> Result { + Ok(match value { FFI_Partitioning::RoundRobinBatch(size) => { - Partitioning::RoundRobinBatch(*size) + Partitioning::RoundRobinBatch(size) } FFI_Partitioning::Hash(exprs, size) => { let exprs = exprs.iter().map(>::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::, _>>() + .map(SplitPoint::new) + }) + .collect::, _>>()?; + + 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 { + let col_expr = Arc::new(Column::new("a", 0)) as Arc; + let sort_expr = PhysicalSortExpr::new(col_expr, SortOptions::default()); + let ordering = LexOrdering::new([sort_expr]).expect("non-empty ordering"); + let split_points = vec![ + SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), + ]; + 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 +185,50 @@ mod tests { assert_eq!(partitioning, returned); } } + + Ok(()) + } + + #[test] + fn round_trip_ffi_range_partitioning_compound_key() -> Result<()> { + let a = Arc::new(Column::new("a", 0)) as Arc; + let b = Arc::new(Column::new("b", 1)) as Arc; + 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())), + ]), + ]; + let partitioning = + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?); + + 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 b286ee2d7d30c..fbac4972c4d54 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 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 for PlanProperties { Ok(PlanProperties::new( eq_properties, - (&partitioning).into(), + partitioning, emission_type, boundedness, )) @@ -282,6 +283,29 @@ mod tests { )) } + fn create_range_test_props() -> Result { + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{ScalarValue, SplitPoint}; + use datafusion_physical_expr::{LexOrdering, RangePartitioning}; + + 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 +338,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(()) + } }