Skip to content

Commit 8098706

Browse files
xudong963xudong963
committed
fix: preserve fetched distribution operator boundaries
1 parent c90e58b commit 8098706

5 files changed

Lines changed: 243 additions & 191 deletions

File tree

‎datafusion/core/tests/physical_optimizer/enforce_distribution.rs‎

Lines changed: 179 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ use datafusion_physical_plan::filter::FilterExec;
7272
use datafusion_physical_plan::joins::utils::JoinOn;
7373
use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
7474
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
75+
use datafusion_physical_plan::sorts::sort::SortExec;
7576
use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
7677
use datafusion_physical_plan::union::UnionExec;
7778
use datafusion_physical_plan::{
@@ -4543,11 +4544,12 @@ fn test_replace_order_preserving_variants_with_fetch() -> Result<()> {
45434544
// Apply the function
45444545
let result = replace_order_preserving_variants(dist_context)?;
45454546

4546-
// Verify the plan was transformed to CoalescePartitionsExec
4547+
// A fetched ordered merge must still select the TopK rows.
4548+
let result = check_integrity(result)?;
45474549
result
45484550
.plan
4549-
.downcast_ref::<CoalescePartitionsExec>()
4550-
.expect("Expected CoalescePartitionsExec");
4551+
.downcast_ref::<SortExec>()
4552+
.expect("Expected a TopK SortExec");
45514553

45524554
// Verify fetch was preserved
45534555
assert_eq!(
@@ -4597,31 +4599,183 @@ fn preserve_fetch_when_reoptimizing_coalesce_partitions() -> Result<()> {
45974599
Ok(())
45984600
}
45994601

4600-
#[test]
4601-
fn move_fetch_to_replacement_sort() -> Result<()> {
4602-
let schema = schema();
4603-
let sort_key: LexOrdering =
4604-
[PhysicalSortExpr::new_default(col("c", &schema)?)].into();
4605-
let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]);
4606-
let merge: Arc<dyn ExecutionPlan> = Arc::new(
4607-
SortPreservingMergeExec::new(sort_key.clone(), input).with_fetch(Some(5)),
4608-
);
4609-
let plan = sort_required_exec_with_req(merge, sort_key);
4602+
#[tokio::test]
4603+
async fn move_fetch_to_replacement_sort() -> Result<()> {
4604+
for (options, partitions, expected) in [
4605+
(
4606+
SortOptions::default(),
4607+
[
4608+
vec![None, Some(1), Some(1), Some(6)],
4609+
vec![None, Some(1), Some(2), Some(7)],
4610+
],
4611+
vec![None, None, Some(1), Some(1), Some(1)],
4612+
),
4613+
(
4614+
SortOptions {
4615+
descending: true,
4616+
nulls_first: false,
4617+
},
4618+
[vec![Some(7), Some(1), None], vec![Some(6), Some(1), None]],
4619+
vec![Some(7), Some(6), Some(1), Some(1), None],
4620+
),
4621+
] {
4622+
let (input, sort_key) = sorted_memory_input(partitions, options)?;
4623+
let merge: Arc<dyn ExecutionPlan> = Arc::new(
4624+
SortPreservingMergeExec::new(sort_key.clone(), input).with_fetch(Some(5)),
4625+
);
4626+
assert_eq!(fetch_test_values(Arc::clone(&merge)).await?, expected);
4627+
let plan = sort_required_exec_with_req(merge, sort_key);
4628+
let optimized = ensure_distribution_helper(plan, 10, false)?;
4629+
let replacement = Arc::clone(optimized.children()[0]);
4630+
let sort = replacement
4631+
.downcast_ref::<SortExec>()
4632+
.expect("expected a replacement sort");
4633+
assert_eq!(sort.fetch(), Some(5));
4634+
assert_eq!(fetch_test_values(replacement).await?, expected);
4635+
}
4636+
Ok(())
4637+
}
46104638

4611-
let optimized = ensure_distribution_helper(plan, 10, false)?;
4612-
let plan = displayable(optimized.as_ref()).indent(true).to_string();
4639+
#[tokio::test]
4640+
async fn preserve_fetch_in_nested_distribution_operators() -> Result<()> {
4641+
for outer_fetch in [0, 3, 10] {
4642+
let (input, sort_key) = sorted_memory_input(
4643+
[0, 1].map(|start| (start..10).step_by(2).map(Some).collect()),
4644+
SortOptions::default(),
4645+
)?;
4646+
let merge: Arc<dyn ExecutionPlan> =
4647+
Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5)));
4648+
let plan: Arc<dyn ExecutionPlan> =
4649+
Arc::new(CoalescePartitionsExec::new(merge).with_fetch(Some(outer_fetch)));
4650+
let expected = (0..outer_fetch.min(5))
4651+
.map(|value| Some(value as i64))
4652+
.collect::<Vec<_>>();
4653+
assert_reoptimized_fetch_values(plan, &expected).await?;
4654+
}
4655+
Ok(())
4656+
}
46134657

4614-
assert!(
4615-
plan.contains(
4616-
"SortExec: TopK(fetch=5), expr=[c@2 ASC], preserve_partitioning=[false]"
4617-
),
4618-
"expected the replacement sort to preserve fetch:\n{plan}"
4619-
);
4620-
assert!(
4621-
!plan.contains("CoalescePartitionsExec: fetch=5"),
4622-
"fetch below the replacement sort would change TopK results:\n{plan}"
4623-
);
4658+
#[tokio::test]
4659+
async fn preserve_topk_when_parent_changes_ordering() -> Result<()> {
4660+
let (input, sort_key) = sorted_memory_input(
4661+
[0, 1].map(|start| (start..10).step_by(2).map(Some).collect()),
4662+
SortOptions::default(),
4663+
)?;
4664+
let descending = [PhysicalSortExpr::new(
4665+
col("c", &input.schema())?,
4666+
SortOptions {
4667+
descending: true,
4668+
nulls_first: false,
4669+
},
4670+
)]
4671+
.into();
4672+
let merge: Arc<dyn ExecutionPlan> =
4673+
Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5)));
4674+
let plan: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(descending, merge));
4675+
assert_reoptimized_fetch_values(plan, &[Some(4), Some(3), Some(2), Some(1), Some(0)])
4676+
.await
4677+
}
4678+
4679+
#[tokio::test]
4680+
async fn preserve_fetch_when_parallelizing_sort_above_filter() -> Result<()> {
4681+
let (input, sort_key) = sorted_memory_input(
4682+
[
4683+
vec![Some(-4), Some(-2), Some(2), Some(4), Some(6)],
4684+
vec![Some(-3), Some(-1), Some(3), Some(5), Some(7)],
4685+
],
4686+
SortOptions::default(),
4687+
)?;
4688+
let predicate = Arc::new(BinaryExpr::new(
4689+
col("c", &input.schema())?,
4690+
Operator::Gt,
4691+
lit(0_i64),
4692+
));
4693+
let coalesce: Arc<dyn ExecutionPlan> =
4694+
Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(5)));
4695+
let filter: Arc<dyn ExecutionPlan> =
4696+
Arc::new(FilterExec::try_new(predicate, coalesce)?);
4697+
let mut plan: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(sort_key, filter));
4698+
let mut config = test_suite_default_config_options();
4699+
config.optimizer.enable_round_robin_repartition = false;
4700+
config.optimizer.repartition_sorts = true;
4701+
for iteration in 0..3 {
4702+
if iteration > 0 {
4703+
plan = EnsureRequirements::new().optimize(plan, &config)?;
4704+
}
4705+
// Either input batch can arrive first. Both contain three positive
4706+
// rows, so keeping the limit below the filter always returns three.
4707+
assert_eq!(
4708+
fetch_test_values(Arc::clone(&plan)).await?.len(),
4709+
3,
4710+
"iteration {iteration}:\n{}",
4711+
displayable(plan.as_ref()).indent(true)
4712+
);
4713+
}
4714+
Ok(())
4715+
}
46244716

4717+
fn sorted_memory_input(
4718+
partitions: [Vec<Option<i64>>; 2],
4719+
options: SortOptions,
4720+
) -> Result<(Arc<dyn ExecutionPlan>, LexOrdering)> {
4721+
let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int64, true)]));
4722+
let order: LexOrdering = [PhysicalSortExpr::new(col("c", &schema)?, options)].into();
4723+
let partitions = partitions
4724+
.into_iter()
4725+
.map(|values| {
4726+
RecordBatch::try_new(
4727+
Arc::clone(&schema),
4728+
vec![Arc::new(Int64Array::from(values))],
4729+
)
4730+
.map(|batch| vec![batch])
4731+
})
4732+
.collect::<std::result::Result<Vec<_>, _>>()?;
4733+
let source = MemorySourceConfig::try_new(&partitions, schema, None)?
4734+
.try_with_sort_information(vec![order.clone()])?;
4735+
Ok((DataSourceExec::from_data_source(source), order))
4736+
}
4737+
4738+
async fn fetch_test_values(plan: Arc<dyn ExecutionPlan>) -> Result<Vec<Option<i64>>> {
4739+
let batches = collect(plan, SessionContext::new().task_ctx()).await?;
4740+
Ok(batches
4741+
.iter()
4742+
.flat_map(|batch| {
4743+
batch
4744+
.column(0)
4745+
.as_any()
4746+
.downcast_ref::<Int64Array>()
4747+
.unwrap()
4748+
.iter()
4749+
})
4750+
.collect())
4751+
}
4752+
4753+
async fn assert_reoptimized_fetch_values(
4754+
plan: Arc<dyn ExecutionPlan>,
4755+
expected: &[Option<i64>],
4756+
) -> Result<()> {
4757+
for repartition_sorts in [false, true] {
4758+
let mut optimized = Arc::clone(&plan);
4759+
let mut config = test_suite_default_config_options();
4760+
config.optimizer.enable_round_robin_repartition = false;
4761+
config.optimizer.repartition_sorts = repartition_sorts;
4762+
for iteration in 0..3 {
4763+
if iteration > 0 {
4764+
let distribution =
4765+
DistributionContext::new_default(Arc::clone(&optimized))
4766+
.transform_up(|context| ensure_distribution(context, &config))?
4767+
.data;
4768+
check_integrity(distribution)?;
4769+
optimized = EnsureRequirements::new().optimize(optimized, &config)?;
4770+
}
4771+
assert_eq!(
4772+
fetch_test_values(Arc::clone(&optimized)).await?,
4773+
expected,
4774+
"iteration {iteration}, repartition_sorts={repartition_sorts}:\n{}",
4775+
displayable(optimized.as_ref()).indent(true)
4776+
);
4777+
}
4778+
}
46254779
Ok(())
46264780
}
46274781

‎datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs‎

Lines changed: 16 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
3131
use arrow::record_batch::RecordBatch;
3232
use insta::{allow_duplicates, assert_snapshot};
3333
use datafusion_common::tree_node::{TransformedResult, TreeNode};
34-
use datafusion_common::{assert_contains, NullEquality, Result};
34+
use datafusion_common::{NullEquality, Result};
3535
use datafusion_common::config::ConfigOptions;
3636
use datafusion_datasource::source::DataSourceExec;
3737
use datafusion_execution::TaskContext;
@@ -1199,43 +1199,22 @@ fn test_plan_with_order_preserving_variants_preserves_fetch() -> Result<()> {
11991199
.with_fetch(Some(10))
12001200
.unwrap();
12011201

1202-
// Test sort's fetch is greater than coalesce fetch, return error because it's not reasonable
1203-
let requirements = OrderPreservationContext::new(
1204-
coalesced.clone(),
1205-
false,
1206-
vec![OrderPreservationContext::new(
1207-
parquet_exec.clone(),
1208-
false,
1209-
vec![],
1210-
)],
1211-
);
1212-
let res = plan_with_order_preserving_variants(requirements, false, true, Some(15));
1213-
assert_contains!(
1214-
res.unwrap_err().to_string(),
1215-
"CoalescePartitionsExec fetch [10] should be greater than or equal to SortExec fetch [15]"
1216-
);
1217-
1218-
// Test sort is without fetch, expected to get the fetch value from the coalesced
1219-
let requirements = OrderPreservationContext::new(
1220-
coalesced.clone(),
1221-
false,
1222-
vec![OrderPreservationContext::new(
1223-
parquet_exec.clone(),
1202+
// Keep the coalesce's row selection independently of an ancestor's fetch.
1203+
for sort_fetch in [Some(15), None, Some(5)] {
1204+
let requirements = OrderPreservationContext::new(
1205+
Arc::clone(&coalesced),
12241206
false,
1225-
vec![],
1226-
)],
1227-
);
1228-
let res = plan_with_order_preserving_variants(requirements, false, true, None)?;
1229-
assert_eq!(res.plan.fetch(), Some(10),);
1230-
1231-
// Test sort's fetch is less than coalesces fetch, expected to get the fetch value from the sort
1232-
let requirements = OrderPreservationContext::new(
1233-
coalesced,
1234-
false,
1235-
vec![OrderPreservationContext::new(parquet_exec, false, vec![])],
1236-
);
1237-
let res = plan_with_order_preserving_variants(requirements, false, true, Some(5))?;
1238-
assert_eq!(res.plan.fetch(), Some(5),);
1207+
vec![OrderPreservationContext::new(
1208+
parquet_exec.clone(),
1209+
false,
1210+
vec![],
1211+
)],
1212+
);
1213+
let res =
1214+
plan_with_order_preserving_variants(requirements, false, true, sort_fetch)?;
1215+
assert!(Arc::ptr_eq(&res.plan, &coalesced));
1216+
assert_eq!(res.plan.fetch(), Some(10));
1217+
}
12391218
Ok(())
12401219
}
12411220

0 commit comments

Comments
 (0)