adriangb commented on code in PR #26660:
URL: https://github.com/apache/datafusion/pull/26660#discussion_r4366396375
##########
datafusion/physical-plan/src/execution_plan.rs:
##########
@@ -777,13 +777,37 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send +
Sync {
/// Returns `true` if a limit can be safely pushed down through this
/// `ExecutionPlan` node.
///
- /// If this method returns `true`, and the query plan contains a limit at
- /// the output of this node, DataFusion will push the limit to the input
- /// of this node.
+ /// DataFusion's own operators still return the same values as before, so
+ /// a wrapper that forwards this to its input keeps working.
+ #[deprecated(
+ since = "56.0.0",
+ note = "implement `limit_behavior` instead; DataFusion no longer calls
this"
Review Comment:
The note says "DataFusion no longer calls this", but the default
`limit_behavior()` below still calls it. A user that implements only this
method still changes the optimizer's behavior. Please reword, for example:
```suggestion
note = "implement `limit_behavior` instead; the default
`limit_behavior` reads this"
```
##########
datafusion/physical-plan/src/execution_plan.rs:
##########
@@ -777,13 +777,37 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send +
Sync {
/// Returns `true` if a limit can be safely pushed down through this
/// `ExecutionPlan` node.
///
- /// If this method returns `true`, and the query plan contains a limit at
- /// the output of this node, DataFusion will push the limit to the input
- /// of this node.
+ /// DataFusion's own operators still return the same values as before, so
+ /// a wrapper that forwards this to its input keeps working.
+ #[deprecated(
+ since = "56.0.0",
+ note = "implement `limit_behavior` instead; DataFusion no longer calls
this"
+ )]
fn supports_limit_pushdown(&self) -> bool {
false
}
+ /// What a row limit (`LIMIT` / `OFFSET`) on this node's output means for
+ /// its input. The optimizer rules that move limits read only this; see
+ /// [`LimitBehavior`].
+ ///
+ /// The default keeps the behavior of nodes that still implement the
+ /// deprecated [`Self::supports_limit_pushdown`]: [`LimitBehavior::Stops`]
+ /// if it returns `false`, otherwise [`LimitBehavior::MoveToInput`] when
+ /// [`Self::cardinality_effect`] is [`CardinalityEffect::Equal`] and
+ /// [`LimitBehavior::CopyToChildren`] when it is not.
+ fn limit_behavior(&self) -> LimitBehavior {
+ #[expect(deprecated)]
+ let supports_limit_pushdown = self.supports_limit_pushdown();
+ if !supports_limit_pushdown {
+ return LimitBehavior::Stops;
+ }
+ match self.cardinality_effect() {
+ CardinalityEffect::Equal => LimitBehavior::MoveToInput,
+ _ => LimitBehavior::CopyToChildren,
Review Comment:
This default changes `LimitPushPastWindows` for more external operators than
the PR description says. The default `cardinality_effect()` is `Unknown`, so a
wrapper like this one now gets `CopyToChildren`:
```rust
impl ExecutionPlan for MyWrapper {
fn supports_limit_pushdown(&self) -> bool {
self.input.supports_limit_pushdown() // forwards; no
cardinality_effect override
}
}
```
| | before | after |
|---|---|---|
| window limit at `MyWrapper` | reset | pushed through |
Can you name this case in the upgrade guide?
##########
datafusion/physical-plan/src/execution_plan.rs:
##########
@@ -2075,6 +2099,53 @@ pub fn get_plan_string(plan: &Arc<dyn ExecutionPlan>) ->
Vec<String> {
actual.iter().map(|elem| (*elem).to_string()).collect()
}
+/// What a row limit on an operator's output means for the operator's input.
+/// Returned by [`ExecutionPlan::limit_behavior`].
+///
+/// A limit can be *copied* below a node as an early-stop hint while it also
+/// stays above the node, or *moved* below a node, leaving nothing above it.
+/// Copying is safe whenever the first `n` output rows need at most the first
+/// `n` rows of each child. Moving also needs the node to output exactly the
+/// rows of its input.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LimitBehavior {
+ /// A limit above this node must stay above it. The default.
+ Stops,
+ /// This node applies `OFFSET skip LIMIT fetch` to its single input
+ /// (`GlobalLimitExec`, `LocalLimitExec`).
+ Limit { skip: usize, fetch: Option<usize> },
+ /// A limit may be copied to each child, but may not move: this node can
+ /// output more rows than any one child (`UnionExec`).
+ CopyToChildren,
+ /// This node outputs exactly the rows of its input: a limit may move to
+ /// its input (`ProjectionExec`).
+ MoveToInput,
+ /// This node merges its input partitions, so it sees every row and can
+ /// enforce a limit with its own fetch (`CoalescePartitionsExec`,
+ /// `SortPreservingMergeExec`). Without a `fetch` of its own it outputs
+ /// exactly the rows of its input, like [`Self::MoveToInput`]. With one,
+ /// it already selects rows, so a limit above it may not move below it.
+ /// `fetch` must equal the node's [`ExecutionPlan::fetch`].
+ Merge { fetch: Option<usize> },
Review Comment:
Only this doc comment keeps `fetch` equal to `ExecutionPlan::fetch()`. A
wrapper that forwards `limit_behavior()` from an inner merge, but not
`fetch()`, breaks it silently. Could the rules read `plan.fetch()` instead, so
`Merge` has no field? If you keep the field, a `debug_assert_eq!(fetch,
plan.fetch())` in `limit_pushdown.rs` would catch the mismatch.
##########
datafusion/physical-optimizer/src/limit_pushdown.rs:
##########
@@ -142,29 +119,225 @@ impl PhysicalOptimizerRule for LimitPushdown {
}
}
-struct LimitInfo {
- input: Arc<dyn ExecutionPlan>,
+/// The row bound a node receives from its ancestors.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Above {
+ /// No bound.
+ Nothing,
+ /// `OFFSET skip LIMIT fetch` that no ancestor enforces yet: it must be
+ /// placed at or above the first node it can't pass.
+ Open(RowBound),
+ /// A bound an ancestor already enforces, with its `skip` applied. A node
+ /// below may stop after `fetch` rows; nothing else is required of it. A
+ /// limit node further down is still folded in and placed again.
+ Enforced(Option<usize>),
+}
+
+impl Above {
+ fn open(bound: RowBound) -> Self {
+ if bound == RowBound::default() {
+ Above::Nothing
+ } else {
+ Above::Open(bound)
+ }
+ }
+
+ fn bound(self) -> RowBound {
+ match self {
+ Above::Nothing => RowBound::default(),
+ Above::Open(bound) => bound,
+ Above::Enforced(fetch) => RowBound::fetch(fetch),
+ }
+ }
+
+ /// `self` above a node that enforces its own `fetch` (a TopK, a scan with
+ /// a row limit, ...). The node enforces the combined bound unless rows
+ /// still have to be skipped above it.
+ fn above_fetching_node(self, fetch: usize) -> Self {
+ let bound = self.bound().over(RowBound::fetch(Some(fetch)));
+ if bound.skip > 0 {
+ Above::Open(bound)
+ } else {
+ Above::Enforced(bound.fetch)
+ }
+ }
+}
+
+/// What a node does with a bound from its parent.
+enum Step {
+ /// A limit node: its bound is folded in and the node removed.
+ Limit(RowBound),
+ /// Passes the bound to every child.
+ Pass,
+ /// Combines partitions, so it sees every row and can enforce the whole
+ /// bound itself.
+ Merge,
+ /// The bound can't go below this node.
+ Stop,
+}
+
+fn step(plan: &Arc<dyn ExecutionPlan>) -> Step {
+ match plan.limit_behavior() {
+ LimitBehavior::Limit { skip, fetch } => Step::Limit(RowBound { skip,
fetch }),
+ LimitBehavior::CopyToChildren | LimitBehavior::MoveToInput =>
Step::Pass,
+ LimitBehavior::Merge { .. } => Step::Merge,
+ // LimitPushPastWindows handles windows.
+ LimitBehavior::Stops | LimitBehavior::Window(_) => Step::Stop,
+ }
+}
+
+/// Pushes the bounds of all limit nodes in `plan` down as far as they pass,
+/// and writes each one once, at the first node that can enforce it:
+///
+/// - a merge, which takes it as its own fetch;
+/// - otherwise the first node the bound can't pass, which takes it as its
+/// fetch or gets a limit node on top.
+///
+/// Below the place a bound is enforced, nodes that can stop early are given
+/// it as a fetch too. Operators can't skip rows, so a bound with an `OFFSET`
+/// always keeps a [`GlobalLimitExec`] where it is enforced.
+///
+/// [`GlobalLimitExec`]: datafusion_physical_plan::limit::GlobalLimitExec
+fn push_limits(
+ mut plan: Arc<dyn ExecutionPlan>,
+ mut above: Above,
+ stats_ctx: &StatisticsContext,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ while let Some((input, folded)) = fold_limit(&plan, above, stats_ctx)? {
+ (plan, above) = (input, folded);
+ }
+ let (plan, below) = apply_bound(plan, above);
+
+ let children = plan.children();
+ let mut changed = false;
+ let new_children = children
+ .into_iter()
+ .map(|child| {
+ let new_child = push_limits(Arc::clone(child), below, stats_ctx)?;
+ changed |= !Arc::ptr_eq(child, &new_child);
+ Ok(new_child)
+ })
+ .collect::<Result<_>>()?;
+ if changed {
+ replace_children_if_necessary(plan, new_children)
+ } else {
+ Ok(plan)
+ }
+}
+
+/// If `plan` is a limit node, returns its input and `above` combined with the
+/// limit's bound. Otherwise returns `None`.
+fn fold_limit(
+ plan: &Arc<dyn ExecutionPlan>,
+ above: Above,
+ stats_ctx: &StatisticsContext,
+) -> Result<Option<(Arc<dyn ExecutionPlan>, Above)>> {
+ let Step::Limit(limit) = step(plan) else {
+ return Ok(None);
+ };
+ let input = Arc::clone(plan.children()[0]);
+ let bound = above.bound().over(limit);
+ // Exact statistics may prove the input already has few enough rows. The
+ // bound is kept so a nested limit can still be combined with it.
+ let above = if limit_satisfied_by_input(&input, bound, stats_ctx)? {
+ Above::Enforced(bound.fetch)
+ } else {
+ Above::open(bound)
+ };
+ Ok(Some((input, above)))
+}
+
+/// Applies `above` to `plan`, which is not a limit node. Returns the new plan
+/// and the bound its children receive.
+fn apply_bound(
+ plan: Arc<dyn ExecutionPlan>,
+ above: Above,
+) -> (Arc<dyn ExecutionPlan>, Above) {
+ let above = match plan.fetch() {
+ Some(fetch) => above.above_fetching_node(fetch),
+ None => above,
+ };
+ match (step(&plan), above) {
+ (_, Above::Nothing | Above::Enforced(None)) => (plan, Above::Nothing),
+ // Only an OFFSET: it is placed here, operators can't skip.
+ (_, Above::Open(bound)) if bound.fetch.is_none() => {
+ (limit_node(plan, bound), Above::Nothing)
+ }
+ (Step::Limit(_), _) => unreachable!("limit nodes are folded before
apply_bound"),
+ (Step::Pass, above) => (plan, above),
+ (Step::Merge, Above::Open(bound)) => match
plan.with_fetch(bound.input_rows()) {
+ Some(_) => (place(plan, bound),
Above::Enforced(bound.input_rows())),
+ // Not reachable in-tree: both merges absorb a fetch. Unlike the
+ // other branches, the bound below does not add `skip`.
Review Comment:
This arm drops `skip` from the bound below. An external merge that returns
`Merge { fetch: None }` and has no `with_fetch` gives wrong results:
```text
GlobalLimitExec: skip=5, fetch=3
ExtMergeExec -- limit_behavior() = Merge { fetch: None },
with_fetch() = None
SortExec -- becomes fetch=3
```
- Each partition gives at most 3 rows.
- The merge must give 8 rows (5 skipped + 3 kept). It can give fewer.
- Result: short or empty output.
The old code did not have this problem, because it found merges only by
downcast to the two in-tree types.
```suggestion
None => (limit_node(plan, bound),
Above::Enforced(bound.input_rows())),
```
##########
datafusion/physical-optimizer/src/limit_pushdown.rs:
##########
@@ -176,172 +349,43 @@ pub fn pushdown_limit_helper(
)
}
-/// Like [`pushdown_limit_helper`], but computes statistics with `stats_ctx`,
-/// so its cache can be shared with other calls and other rules.
+/// Like [`pushdown_limit_helper`], but computes statistics with `stats_ctx`.
+#[deprecated(
+ since = "56.0.0",
+ note = "LimitPushdown no longer works one node at a time; run the
`LimitPushdown` rule"
+)]
+#[expect(deprecated)]
pub fn pushdown_limit_helper_with_stats(
Review Comment:
`pushdown_limit_helper_with_stats` is not in any 55.x tag (55.0.0 to
55.2.0). It came from https://github.com/apache/datafusion/pull/26094 in this
cycle. Can we remove it instead of deprecating it? Then the upgrade guide entry
is not necessary.
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:
##########
@@ -549,7 +522,10 @@ fn pushdown_requirement_to_children(
Ok(None)
}
} else if plan.fetch().is_some()
- && plan.supports_limit_pushdown()
+ && !matches!(
+ plan.limit_behavior(),
+ LimitBehavior::Stops | LimitBehavior::Window(_)
Review Comment:
`Stops | Window(_)` is also matched in `limit_pushdown.rs:185`. A new
variant needs an edit in both places. Nit: a method on `LimitBehavior` (for
example `blocks_limit()`) would keep the two rules in sync.
--
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]
