jayzhan211 commented on code in PR #27428:
URL: https://github.com/apache/datafusion/pull/27428#discussion_r4576340448
##########
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:
Agree with @viirya, and it's worse than extra repartitions: a correct
foreign subtree now fails host planning. I returned a plan from the
integration-test cdylib from a `TableProvider::scan`: its top node requires
`KeyPartitioned([a])` and its child is already `Hash([a], 4)` (same with a hard
sort on `a` over `SortExec(a)`). `SELECT * FROM t` plans on the base of this
PR, but here it fails with `SanityCheckPlan: the plan produced by physical
optimizer rule 'EnsureRequirements' does not satisfy the order and distribution
requirements`, because each requirement call builds new `ForeignPhysicalExpr`
wrappers, which only equal themselves.
A structural `eq` callback on `FFI_PhysicalExpr` fixes both shapes, with no
extra repartition or sort, and the FFI unit and integration suites pass with
it. Since this PR already changes the ABI, could it go in here, with the test
below and a line about `FFI_PhysicalExpr` in the 57 upgrade note?
```diff
--- a/datafusion/ffi/src/physical_expr/mod.rs
+++ b/datafusion/ffi/src/physical_expr/mod.rs
@@ pub struct FFI_PhysicalExpr {
pub hash: unsafe extern "C" fn(&Self) -> u64,
+ /// Structural equality with `other`. Returns `false` when `other` was
+ /// created by a different library.
+ pub eq: unsafe extern "C" fn(&Self, other: &Self) -> bool,
+
@@
+unsafe extern "C" fn eq_fn_wrapper(
+ expr: &FFI_PhysicalExpr,
+ other: &FFI_PhysicalExpr,
+) -> bool {
+ // `other.inner()` is only valid when this library created `other`.
+ (other.library_marker_id)() == crate::get_library_marker_id()
+ && expr.inner() == other.inner()
+}
@@ (both constructors: `clone_fn_wrapper` and `From<Arc<dyn PhysicalExpr>>`)
hash: hash_fn_wrapper,
+ eq: eq_fn_wrapper,
@@ impl PartialEq for ForeignPhysicalExpr {
- // FFI_PhysicalExpr cannot be compared, so identity equality is the
best we can do.
- std::ptr::eq(self, other)
+ std::ptr::eq(self, other) || unsafe { (self.expr.eq)(&self.expr,
&other.expr) }
@@ fn ffi_physical_expr_hash() {
- assert_ne!(&foreign_1, &foreign_2);
+ assert_eq!(&foreign_1, &foreign_2);
```
<details>
<summary>Cross-library test: fails on this PR with SanityCheckPlan, passes
with the fix</summary>
`datafusion/ffi/src/tests/mod.rs` (producer, exported as its own symbol like
`datafusion_ffi_test_create_exec_with_byte_metrics`):
```rs
/// A plan whose top node requires its input hash-partitioned (or sorted) on
/// `a`, and whose child already provides it.
#[unsafe(no_mangle)]
pub extern "C" fn datafusion_ffi_test_create_exec_with_requirements(
hard_ordering: bool,
) -> FFI_ExecutionPlan {
use crate::execution_plan::tests::TestInputRequirements;
use datafusion_physical_expr::Distribution;
use datafusion_physical_expr_common::sort_expr::{
LexRequirement, OrderingRequirements, PhysicalSortRequirement,
};
use datafusion_physical_plan::InputDistributionRequirements;
use datafusion_physical_plan::repartition::RepartitionExec;
use datafusion_physical_plan::sorts::sort::SortExec;
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32,
false)]));
let a = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
let leaf = Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc<dyn
ExecutionPlan>;
let (child, reqs): (Arc<dyn ExecutionPlan>, _) = if hard_ordering {
let ord =
LexOrdering::new([PhysicalSortExpr::new_default(Arc::clone(&a))]).unwrap();
let req = OrderingRequirements::new(
LexRequirement::new([PhysicalSortRequirement::new(Arc::clone(&a), None)])
.unwrap(),
);
(
Arc::new(SortExec::new(ord,
leaf).with_preserve_partitioning(true)),
TestInputRequirements {
distribution: InputDistributionRequirements::new(vec![
Distribution::UnspecifiedDistribution,
]),
ordering: vec![Some(req)],
maintains_input_order: vec![true],
benefits_from_input_partitioning: vec![false],
},
)
} else {
(
Arc::new(
RepartitionExec::try_new(
leaf,
Partitioning::Hash(vec![Arc::clone(&a)], 4),
)
.unwrap(),
),
TestInputRequirements {
distribution: InputDistributionRequirements::new(vec![
Distribution::KeyPartitioned(vec![Arc::clone(&a)]),
]),
ordering: vec![None],
maintains_input_order: vec![false],
benefits_from_input_partitioning: vec![true],
},
)
};
let plan = EmptyExec::new(schema)
.with_children(vec![child])
.with_input_requirements(reqs);
FFI_ExecutionPlan::new(Arc::new(plan), None)
}
```
`datafusion/ffi/src/tests/utils.rs` (loader):
```rs
pub fn get_requirements_exec(
hard_ordering: bool,
) -> Result<crate::execution_plan::FFI_ExecutionPlan> {
let lib_path = find_library()?;
let lib = unsafe {
libloading::Library::new(&lib_path)
.map_err(|e| DataFusionError::External(Box::new(e)))?
};
let create_exec: libloading::Symbol<
extern "C" fn(bool) -> crate::execution_plan::FFI_ExecutionPlan,
> = unsafe {
lib.get(b"datafusion_ffi_test_create_exec_with_requirements")
.map_err(|e| DataFusionError::External(Box::new(e)))?
};
let plan = create_exec(hard_ordering);
#[expect(clippy::mem_forget)]
std::mem::forget(lib);
Ok(plan)
}
```
`datafusion/ffi/integration-tests/tests/cases/execution_plan.rs` (consumer;
needs `async-trait` in the crate's dev-dependencies):
```rs
mod foreign_requirements {
use std::sync::Arc;
use async_trait::async_trait;
use datafusion::catalog::{Session, TableProvider};
use datafusion::logical_expr::{Expr, TableType};
use datafusion::prelude::SessionContext;
use datafusion_common::Result;
use datafusion_ffi::tests::utils::get_requirements_exec;
use datafusion_physical_plan::{ExecutionPlan, displayable};
/// A host-side table whose scan returns a plan built by another library.
#[derive(Debug)]
struct ForeignScanTable(Arc<dyn ExecutionPlan>);
#[async_trait]
impl TableProvider for ForeignScanTable {
fn schema(&self) -> arrow_schema::SchemaRef {
self.0.schema()
}
fn table_type(&self) -> TableType {
TableType::Base
}
async fn scan(
&self,
_: &dyn Session,
_: Option<&[usize]>,
_: &[Expr],
_: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
Ok(Arc::clone(&self.0))
}
}
/// A correct foreign subtree (its top node's requirement is met by its
own
/// child) must plan on the host without errors or extra
repartitions/sorts.
#[tokio::test]
async fn foreign_requirements_met_by_foreign_child() -> Result<()> {
for hard_ordering in [false, true] {
let ffi = get_requirements_exec(hard_ordering)?;
let foreign: Arc<dyn ExecutionPlan> = (&ffi).try_into()?;
let ctx = SessionContext::new();
ctx.register_table("t", Arc::new(ForeignScanTable(foreign)))?;
let plan = ctx
.sql("SELECT * FROM t")
.await?
.create_physical_plan()
.await?;
let display =
displayable(plan.as_ref()).indent(false).to_string();
assert!(
!display.contains("\n RepartitionExec: partitioning=Hash")
&& !display.contains("\n SortExec"),
"redundant repartition or sort added:\n{display}"
);
}
Ok(())
}
}
```
</details>
--
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]
