alamb commented on code in PR #20048:
URL: https://github.com/apache/datafusion/pull/20048#discussion_r2742819576
##########
datafusion/core/tests/physical_optimizer/limit_pushdown.rs:
##########
@@ -343,3 +354,104 @@ fn merges_local_limit_with_global_limit() -> Result<()> {
Ok(())
}
+
+#[test]
+fn preserves_nested_global_limit() -> Result<()> {
+ // If there are multiple limits in an execution plan, they all need to be
+ // preserved in the optimized plan.
+ //
+ // Plan structure:
+ // GlobalLimitExec: skip=1, fetch=1
+ // NestedLoopJoinExec (Left)
+ // EmptyExec (left side)
+ // GlobalLimitExec: skip=2, fetch=1
+ // NestedLoopJoinExec (Right)
+ // EmptyExec (left side)
+ // EmptyExec (right side)
+ let schema = create_schema();
+
+ // Build inner join: NestedLoopJoin(Empty, Empty)
+ let inner_left = empty_exec(Arc::clone(&schema));
+ let inner_right = empty_exec(Arc::clone(&schema));
+ let inner_join = nested_loop_join_exec(inner_left, inner_right,
JoinType::Right)?;
+
+ // Add inner limit: GlobalLimitExec: skip=2, fetch=1
+ let inner_limit = global_limit_exec(inner_join, 2, Some(1));
+
+ // Build outer join: NestedLoopJoin(Empty, GlobalLimit)
+ let outer_left = empty_exec(Arc::clone(&schema));
+ let outer_join = nested_loop_join_exec(outer_left, inner_limit,
JoinType::Left)?;
+
+ // Add outer limit: GlobalLimitExec: skip=1, fetch=1
+ let outer_limit = global_limit_exec(outer_join, 1, Some(1));
+
+ let initial = get_plan_string(&outer_limit);
+ let expected_initial = [
+ "GlobalLimitExec: skip=1, fetch=1",
+ " NestedLoopJoinExec: join_type=Left",
+ " EmptyExec",
+ " GlobalLimitExec: skip=2, fetch=1",
+ " NestedLoopJoinExec: join_type=Right",
+ " EmptyExec",
+ " EmptyExec",
+ ];
+ assert_eq!(initial, expected_initial);
+
+ let after_optimize =
+ LimitPushdown::new().optimize(outer_limit, &ConfigOptions::new())?;
+ let expected = [
+ "GlobalLimitExec: skip=1, fetch=1",
+ " NestedLoopJoinExec: join_type=Left",
+ " EmptyExec",
+ " GlobalLimitExec: skip=2, fetch=1",
+ " NestedLoopJoinExec: join_type=Right",
+ " EmptyExec",
+ " EmptyExec",
+ ];
+ assert_eq!(get_plan_string(&after_optimize), expected);
+
+ Ok(())
+}
+
+#[test]
+fn preserves_skip_before_sort() -> Result<()> {
+ // If there's a limit with skip before a node that (1) supports fetch but
+ // (2) does not support limit pushdown, that limit should not be removed.
+ //
+ // Plan structure:
+ // GlobalLimitExec: skip=1, fetch=None
+ // SortExec: TopK(fetch=4)
+ // EmptyExec
+ let schema = create_schema();
+
+ let empty = empty_exec(Arc::clone(&schema));
+
+ let ordering = [PhysicalSortExpr {
+ expr: col("c1", &schema)?,
+ options: SortOptions::default(),
+ }];
+ let sort = sort_exec(ordering.into(), empty)
+ .with_fetch(Some(4))
+ .unwrap();
+
+ let outer_limit = global_limit_exec(sort, 1, None);
+
+ let initial = get_plan_string(&outer_limit);
+ let expected_initial = [
+ "GlobalLimitExec: skip=1, fetch=None",
+ " SortExec: TopK(fetch=4), expr=[c1@0 ASC],
preserve_partitioning=[false]",
+ " EmptyExec",
+ ];
+ assert_eq!(initial, expected_initial);
+
+ let after_optimize =
+ LimitPushdown::new().optimize(outer_limit, &ConfigOptions::new())?;
+ let expected = [
+ "GlobalLimitExec: skip=1, fetch=3",
Review Comment:
Without the code change in the PR, the actual code looks like
```
"SortExec: TopK(fetch=4), expr=[c1@0 ASC],
preserve_partitioning=[false]"
" EmptyExec"
```
Note the offset (aka the `skip`) was dropped
--
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]
