viirya commented on code in PR #27428:
URL: https://github.com/apache/datafusion/pull/27428#discussion_r4574905753
##########
datafusion/ffi/src/execution_plan.rs:
##########
@@ -1005,6 +1163,161 @@ pub mod tests {
Ok(())
}
+ /// A foreign plan reports the same input requirements as the plan it
wraps,
+ /// so optimizer rules on the other side of the boundary see them.
+ #[test]
+ fn test_ffi_execution_plan_input_requirements_round_trip() -> Result<()> {
+ use arrow::compute::SortOptions;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr_common::sort_expr::{
+ LexRequirement, PhysicalSortRequirement,
+ };
+
+ let schema = Arc::new(arrow::datatypes::Schema::new(vec![
+ arrow::datatypes::Field::new("a",
arrow::datatypes::DataType::Int32, false),
+ ]));
+ let a = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
+ let child =
+ || Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc<dyn
ExecutionPlan>;
+ let foreign = |plan: EmptyExec| -> Result<Arc<dyn ExecutionPlan>> {
+ let mut ffi = FFI_ExecutionPlan::new(Arc::new(plan), None);
+ ffi.library_marker_id = crate::mock_foreign_marker_id;
+ (&ffi).try_into()
+ };
+
+ // Defaults survive unchanged.
+ let default_plan = foreign(
+ EmptyExec::new(Arc::clone(&schema)).with_children(vec![child(),
child()]),
+ )?;
+ assert_eq!(
+ format!("{:?}", default_plan.input_distribution_requirements()),
+ format!(
+ "{:?}",
+ InputDistributionRequirements::new(vec![
+ Distribution::UnspecifiedDistribution;
+ 2
+ ])
+ )
+ );
+ assert_eq!(default_plan.required_input_ordering(), vec![None, None]);
+ assert_eq!(default_plan.maintains_input_order(), vec![false, false]);
+ assert_eq!(
+ default_plan.benefits_from_input_partitioning(),
+ vec![true, true]
+ );
+
+ // Non-default requirements, including co-partitioning and soft
+ // ordering, survive unchanged.
+ let ordering = OrderingRequirements::new_soft(
+ LexRequirement::new([PhysicalSortRequirement::new(
+ Arc::clone(&a),
+ Some(SortOptions::new(true, false)),
+ )])
+ .expect("non-empty"),
+ );
+ let reqs = TestInputRequirements {
+ distribution: InputDistributionRequirements::co_partitioned(vec![
+ Distribution::KeyPartitioned(vec![Arc::clone(&a)]),
+ Distribution::KeyPartitioned(vec![Arc::clone(&a)]),
+ ]),
+ ordering: vec![Some(ordering.clone()), None],
+ maintains_input_order: vec![true, false],
+ benefits_from_input_partitioning: vec![false, true],
+ };
+ let plan = foreign(
+ EmptyExec::new(Arc::clone(&schema))
+ .with_children(vec![child(), child()])
+ .with_input_requirements(reqs.clone()),
+ )?;
+ assert!(plan.is::<ForeignExecutionPlan>());
+ assert_eq!(
+ format!("{:?}", plan.input_distribution_requirements()),
+ format!("{:?}", reqs.distribution)
+ );
+ #[expect(deprecated, reason = "the deprecated method must agree too")]
+ let per_child = plan.required_input_distribution();
+ assert_eq!(
+ format!("{per_child:?}"),
+ format!("{:?}", reqs.distribution.clone().into_per_child())
+ );
+ assert_eq!(plan.required_input_ordering(), vec![Some(ordering), None]);
+ assert_eq!(plan.maintains_input_order(), vec![true, false]);
+ assert_eq!(plan.benefits_from_input_partitioning(), vec![false, true]);
+
+ Ok(())
+ }
+
+ /// `EnsureRequirements` running on the consumer side of the boundary sees
+ /// that a foreign `FinalPartitioned` aggregate needs its input
+ /// hash-partitioned, and inserts the repartition. Without the requirement
+ /// crossing FFI it adds only a round-robin repartition, producing a plan
+ /// that groups the same key in several partitions.
+ #[cfg(test)]
+ #[tokio::test]
+ async fn test_ensure_requirements_sees_foreign_plan_requirements() ->
Result<()> {
+ use datafusion::arrow::array::Int32Array;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::arrow::record_batch::RecordBatch;
+ use datafusion::common::tree_node::{Transformed, TreeNode};
+ use datafusion::datasource::MemTable;
+ use datafusion::physical_plan::aggregates::{AggregateExec,
AggregateMode};
+ use datafusion::physical_plan::displayable;
+ use datafusion::physical_plan::repartition::RepartitionExec;
+ use datafusion::prelude::{SessionConfig, SessionContext};
+ use datafusion_physical_optimizer::PhysicalOptimizerRule;
+ use
datafusion_physical_optimizer::ensure_requirements::EnsureRequirements;
+ use
datafusion_physical_optimizer::sanity_checker::check_plan_tree_sanity;
+
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int32, false)]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6]))],
+ )?;
+ let config = SessionConfig::new().with_target_partitions(4);
+ let ctx = SessionContext::new_with_config(config.clone());
+ ctx.register_table("t", Arc::new(MemTable::try_new(schema,
vec![vec![batch]])?))?;
+ let plan = ctx
+ .sql("SELECT a % 2 AS parity, count(*) AS n FROM t GROUP BY 1")
+ .await?
+ .create_physical_plan()
+ .await?;
+
+ // Remove the hash repartition under the final aggregate and send that
+ // aggregate across the (mock) FFI boundary.
+ let plan = plan
+ .transform_up(|node| {
+ let Some(agg) = node.downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(node));
+ };
+ if *agg.mode() != AggregateMode::FinalPartitioned {
+ return Ok(Transformed::no(node));
+ }
+ let repartition = agg
+ .input()
+ .downcast_ref::<RepartitionExec>()
+ .expect("hash repartition under the final aggregate");
+ let stripped = Arc::clone(&node).replace_children(
+ vec![Arc::clone(repartition.input())],
+
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )?;
+ let mut ffi = FFI_ExecutionPlan::new(stripped, None);
+ ffi.library_marker_id = crate::mock_foreign_marker_id;
Review Comment:
Only the plan gets the mock marker. The expressions inside the requirements
are created with the real marker, so they come back as native `Column`s. With a
real library boundary they become `ForeignPhysicalExpr`, whose `eq` is
`ptr::eq`, and every call to these callbacks builds new wrappers. On the host,
a foreign requirement then never matches an existing partitioning or ordering.
The result is still correct, but it can add repartitions or sorts that native
planning wouldn't. Does the datafusion-python `EXPLAIN` show any? For example,
an aggregate over a join on the same key, or a window over sorted input. If it
does, could you open a follow-up issue, or add a test in `ffi_integration` that
crosses a real cdylib?
##########
datafusion/ffi/src/execution_plan.rs:
##########
@@ -462,6 +520,46 @@ impl ExecutionPlan for ForeignExecutionPlan {
self.children.iter().collect()
}
+ fn required_input_distribution(&self) -> Vec<Distribution> {
+ self.input_distribution_requirements().into_per_child()
+ }
+
+ fn input_distribution_requirements(&self) -> InputDistributionRequirements
{
+ let ffi = unsafe {
(self.plan.input_distribution_requirements)(&self.plan) };
+ (&ffi).into()
+ }
+
+ fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+ unsafe { (self.plan.required_input_ordering)(&self.plan) }
+ .iter()
+ .map(|ordering| match ordering {
+ FFI_Option::Some(ordering) =>
OrderingRequirements::try_from(ordering)
+ .inspect_err(|e| {
+ log::error!(
+ "Invalid required input ordering from foreign plan
{}: {e}",
Review Comment:
Returning `None` here drops a hard ordering requirement. As far as I can
tell this can't happen with a well-formed provider, because `LexRequirement`
and `OrderingRequirements` are non-empty by construction. Could a comment say
so, so the fallback doesn't read as a safe default?
##########
datafusion/ffi/src/physical_expr/partitioning.rs:
##########
@@ -128,6 +129,85 @@ impl TryFrom<FFI_Partitioning> for Partitioning {
}
}
+/// A stable enum for sharing a [`Distribution`] across FFI boundaries.
+///
+/// The deprecated [`Distribution::HashPartitioned`] is sent as
+/// `KeyPartitioned`, which the optimizer treats identically.
+#[repr(C)]
+#[derive(Debug)]
+pub enum FFI_Distribution {
+ UnspecifiedDistribution,
+ SinglePartition,
+ KeyPartitioned(SVec<FFI_PhysicalExpr>),
+}
+
+impl From<&Distribution> for FFI_Distribution {
+ #[expect(
+ deprecated,
+ reason = "HashPartitioned is accepted during the KeyPartitioned
migration"
+ )]
+ fn from(value: &Distribution) -> Self {
+ match value {
+ Distribution::UnspecifiedDistribution =>
Self::UnspecifiedDistribution,
+ Distribution::SinglePartition => Self::SinglePartition,
+ Distribution::HashPartitioned(exprs)
+ | Distribution::KeyPartitioned(exprs) => Self::KeyPartitioned(
+ exprs
+ .iter()
+ .map(Arc::clone)
+ .map(FFI_PhysicalExpr::from)
+ .collect(),
+ ),
+ }
+ }
+}
+
+impl From<&FFI_Distribution> for Distribution {
+ fn from(value: &FFI_Distribution) -> Self {
+ match value {
+ FFI_Distribution::UnspecifiedDistribution =>
Self::UnspecifiedDistribution,
+ FFI_Distribution::SinglePartition => Self::SinglePartition,
+ FFI_Distribution::KeyPartitioned(exprs) => Self::KeyPartitioned(
+ exprs.iter().map(<Arc<dyn PhysicalExpr>>::from).collect(),
+ ),
+ }
+ }
+}
+
+/// A stable struct for sharing [`InputDistributionRequirements`] across FFI
+/// boundaries.
+#[repr(C)]
+#[derive(Debug)]
+pub struct FFI_InputDistributionRequirements {
+ /// One distribution per child.
+ distributions: SVec<FFI_Distribution>,
+ /// Whether all children must be co-partitioned.
+ co_partitioned: bool,
Review Comment:
nit: `InputDistributionRequirements` stores co-partitioned child indexes and
is `#[non_exhaustive]`. A short note that the bool relies on all-or-none
construction would flag this spot if subsets become possible.
##########
datafusion/ffi/src/execution_plan.rs:
##########
@@ -690,6 +813,41 @@ pub mod tests {
unimplemented!()
}
+ fn required_input_distribution(&self) -> Vec<Distribution> {
Review Comment:
nit: The test never reaches this override, because
`ForeignExecutionPlan::required_input_distribution` calls its own
`input_distribution_requirements`. Can it be removed?
##########
datafusion/ffi/src/execution_plan.rs:
##########
@@ -462,6 +520,46 @@ impl ExecutionPlan for ForeignExecutionPlan {
self.children.iter().collect()
}
+ fn required_input_distribution(&self) -> Vec<Distribution> {
+ self.input_distribution_requirements().into_per_child()
+ }
+
+ fn input_distribution_requirements(&self) -> InputDistributionRequirements
{
Review Comment:
These are called many times per node during
`EnsureRequirements`/`EnforceSorting`, and each call crosses FFI and rebuilds
the expression wrappers. Requirements are fixed for a plan instance. Could they
be converted once in `TryFrom<FFI_ExecutionPlan>`, the way `properties` is?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
