Skip to content

Commit 7fd04a3

Browse files
authored
fix: rewrite fetch, skip of the Limit node in correct order (#14496)
* fix: rewrite fetch, skip of the Limit node in correct order * style: fix clippy
1 parent d1308f0 commit 7fd04a3

File tree

1 file changed

+46
-18
lines changed
  • datafusion/expr/src/logical_plan

1 file changed

+46
-18
lines changed

datafusion/expr/src/logical_plan/plan.rs

Lines changed: 46 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -959,8 +959,9 @@ impl LogicalPlan {
959959
expr.len()
960960
);
961961
}
962-
let new_skip = skip.as_ref().and_then(|_| expr.pop());
962+
// `LogicalPlan::expressions()` returns in [skip, fetch] order, so we can pop from the end.
963963
let new_fetch = fetch.as_ref().and_then(|_| expr.pop());
964+
let new_skip = skip.as_ref().and_then(|_| expr.pop());
964965
let input = self.only_input(inputs)?;
965966
Ok(LogicalPlan::Limit(Limit {
966967
skip: new_skip.map(Box::new),
@@ -4293,23 +4294,50 @@ digraph {
42934294

42944295
#[test]
42954296
fn test_limit_with_new_children() {
4296-
let limit = LogicalPlan::Limit(Limit {
4297-
skip: None,
4298-
fetch: Some(Box::new(Expr::Literal(
4299-
ScalarValue::new_ten(&DataType::UInt32).unwrap(),
4300-
))),
4301-
input: Arc::new(LogicalPlan::Values(Values {
4302-
schema: Arc::new(DFSchema::empty()),
4303-
values: vec![vec![]],
4304-
})),
4305-
});
4306-
let new_limit = limit
4307-
.with_new_exprs(
4308-
limit.expressions(),
4309-
limit.inputs().into_iter().cloned().collect(),
4310-
)
4311-
.unwrap();
4312-
assert_eq!(limit, new_limit);
4297+
let input = Arc::new(LogicalPlan::Values(Values {
4298+
schema: Arc::new(DFSchema::empty()),
4299+
values: vec![vec![]],
4300+
}));
4301+
let cases = [
4302+
LogicalPlan::Limit(Limit {
4303+
skip: None,
4304+
fetch: None,
4305+
input: Arc::clone(&input),
4306+
}),
4307+
LogicalPlan::Limit(Limit {
4308+
skip: None,
4309+
fetch: Some(Box::new(Expr::Literal(
4310+
ScalarValue::new_ten(&DataType::UInt32).unwrap(),
4311+
))),
4312+
input: Arc::clone(&input),
4313+
}),
4314+
LogicalPlan::Limit(Limit {
4315+
skip: Some(Box::new(Expr::Literal(
4316+
ScalarValue::new_ten(&DataType::UInt32).unwrap(),
4317+
))),
4318+
fetch: None,
4319+
input: Arc::clone(&input),
4320+
}),
4321+
LogicalPlan::Limit(Limit {
4322+
skip: Some(Box::new(Expr::Literal(
4323+
ScalarValue::new_one(&DataType::UInt32).unwrap(),
4324+
))),
4325+
fetch: Some(Box::new(Expr::Literal(
4326+
ScalarValue::new_ten(&DataType::UInt32).unwrap(),
4327+
))),
4328+
input,
4329+
}),
4330+
];
4331+
4332+
for limit in cases {
4333+
let new_limit = limit
4334+
.with_new_exprs(
4335+
limit.expressions(),
4336+
limit.inputs().into_iter().cloned().collect(),
4337+
)
4338+
.unwrap();
4339+
assert_eq!(limit, new_limit);
4340+
}
43134341
}
43144342

43154343
#[test]

0 commit comments

Comments
 (0)