kosiew commented on code in PR #24809:
URL: https://github.com/apache/datafusion/pull/24809#discussion_r3947349493
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -4558,6 +4559,166 @@ fn test_replace_order_preserving_variants_with_fetch()
-> Result<()> {
Ok(())
}
+#[test]
+fn preserve_fetch_when_reoptimizing_ordered_merge() -> Result<()> {
+ let schema = schema();
+ let sort_key: LexOrdering =
+ [PhysicalSortExpr::new_default(col("c", &schema)?)].into();
+ let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]);
+ let plan: Arc<dyn ExecutionPlan> =
+ Arc::new(SortPreservingMergeExec::new(sort_key,
input).with_fetch(Some(5)));
+
+ let optimized =
+ EnsureRequirements::new().optimize(plan,
&test_suite_default_config_options())?;
+ let plan = displayable(optimized.as_ref()).indent(true).to_string();
+
+ assert!(
+ plan.contains("SortPreservingMergeExec: [c@2 ASC], fetch=5"),
+ "expected the optimizer to preserve fetch:\n{plan}"
+ );
+
+ Ok(())
+}
+
+#[test]
+fn preserve_fetch_when_reoptimizing_coalesce_partitions() -> Result<()> {
+ let input = parquet_exec_multiple();
+ let plan: Arc<dyn ExecutionPlan> =
+ Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(5)));
+
+ let optimized =
+ EnsureRequirements::new().optimize(plan,
&test_suite_default_config_options())?;
+
+ assert_eq!(optimized.fetch(), Some(5));
+ optimized
+ .downcast_ref::<CoalescePartitionsExec>()
+ .expect("expected CoalescePartitionsExec");
+
+ Ok(())
+}
+
+#[test]
+fn move_fetch_to_replacement_sort() -> Result<()> {
Review Comment:
Could we strengthen this test by executing a small two-partition input and
asserting the resulting TopK values as well? The display assertion confirms
that we constructed a `SortExec` with `fetch=5`, but an execution assertion
would also protect the actual ordering, null handling, tie behavior, and fetch
placement if the implementation changes later. This is non-blocking.
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1722,9 +1756,26 @@ pub fn ensure_distribution(
replace_children_if_necessary(plan, children_plans)?
};
- Ok(Transformed::yes(DistributionContext::new(
- plan, data, children,
- )))
+ let mut optimized_context = DistributionContext::new(plan, data, children);
+
+ // A removed fetch must survive even when this node does not need a new
+ // distribution operator. Otherwise a second optimizer pass can silently
+ // remove the query's LIMIT.
+ if let Some(fetch) = removed_fetch {
Review Comment:
I think there is still a correctness issue when multiple removed
distribution operators have fetches. `removed_fetch` collapses them to the
minimum value, while `fetch_plan` remembers only the outermost fetched
operator. That loses the semantic position of the inner fetch.
For example, consider `CoalescePartitionsExec(fetch=10) ->
SortPreservingMergeExec([c], fetch=5) -> two sorted partitions`. The original
plan gets the global TopK 5 from the ordered merge. After both operators are
removed, we retain `fetch=5` but can restore it as
`CoalescePartitionsExec(fetch=5)` directly over the partitions. That can return
the first five rows in coalesce/input order rather than the global TopK 5.
Could we preserve each fetched operator at its original semantic boundary,
or otherwise replace it with something that is provably equivalent? I think it
would also be useful to add an execution regression for this nested
Coalesce/SPM case, using partition values where concatenation order differs
from global sort order.
--
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]
