xudong963 commented on code in PR #26660:
URL: https://github.com/apache/datafusion/pull/26660#discussion_r4419519142
##########
datafusion/physical-optimizer/src/limit_pushdown.rs:
##########
@@ -129,29 +106,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,
Review Comment:
for the `Step::Pass`, This contradicts
[LimitBehavior::CopyToChildren](http://127.0.0.1:9000/Users/xudong/opensource/all-contexts/_df_pr_26660_current/datafusion/physical-plan/src/execution_plan.rs:2106)
and can produce too many rows for a single-output multi-child custom operator.
##########
datafusion/physical-plan/src/execution_plan.rs:
##########
@@ -2064,6 +2088,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:
Merge carries a fetch that “must equal” `ExecutionPlan::fetch()`.
The consumers already disagree about which declaration they trust:
-
[apply_bound()](http://127.0.0.1:9000/Users/xudong/opensource/all-contexts/_df_pr_26660_current/datafusion/physical-optimizer/src/limit_pushdown.rs:243)
reads plan.fetch().
- Sort and window pushdown match the Merge payload.
##########
datafusion/physical-optimizer/src/limit_pushdown.rs:
##########
@@ -129,29 +106,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),
Review Comment:
ditto
--
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]
