added tests

Author: saadtajwar

datafusion/ffi/src/physical_expr/partitioning.rs

@@ -17,14 +17,19 @@
 
 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, physical_expr::{FFI_PhysicalExpr, sort::FFI_PhysicalSortExpr}};
+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.
+/// See [`RangePartitioning`] for the descriptions of each field.
 #[repr(C)]
 #[derive(Debug)]
 pub struct FFI_RangePartitioning {
@@ -33,7 +38,7 @@ pub struct FFI_RangePartitioning {
 }
 
 /// 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 {
@@ -56,49 +61,119 @@ impl From<&Partitioning> for FFI_Partitioning {
                 Self::Hash(exprs, *size)
             }
             Partitioning::Range(range) => {
-                let split_points = range.split_points().iter().map(|s| s.values().iter().map(|v| v.try_into().expect("Fail") ).collect() ).collect();
-                let ordering = range.ordering().iter().map(|e| FFI_PhysicalSortExpr::from(e)).collect();
-                Self::Range(FFI_RangePartitioning { split_points, ordering})
+                // 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::Range()
             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 col_expr = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
+        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
@@ -110,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<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())),
+            ]),
+        ];
+        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}"
+        );
     }
 }

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,
         ))
@@ -282,6 +283,29 @@ mod tests {
         ))
     }
 
+    fn create_range_test_props() -> Result<PlanProperties> {
+        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(())
+    }
 }