Skip to content

[SPARK-58063][SQL] Test that BIN BY is column-pruned by the generic ColumnPruning fallback#57157

Open
vranes wants to merge 4 commits into
apache:masterfrom
vranes:bin-by-column-pruning
Open

[SPARK-58063][SQL] Test that BIN BY is column-pruned by the generic ColumnPruning fallback#57157
vranes wants to merge 4 commits into
apache:masterfrom
vranes:bin-by-column-pruning

Conversation

@vranes

@vranes vranes commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

This PR adds test coverage asserting that the logical BinBy operator is column-pruned, and a binBy test DSL helper to construct the node.

No production optimizer change: a Project over BinBy is already handled by the generic Project(_, child) fallback in ColumnPruning (Optimizer.scala), which computes required = child.references ++ p.references and rewrites the child via prunedChild. BinBy is not intercepted by GeneratorNestedColumnAliasing (matches only Generate) or NestedColumnAliasing (BinBy is in neither canProjectPushThrough nor canPruneOn), so it falls through to that fallback. Combined with the produced-attributes shape from SPARK-57858 (the scaled DISTRIBUTE columns and appended columns are produced attributes with fresh ExprIds), BinBy.references lists exactly the range and DISTRIBUTE inputs the kernel reads, so the fallback prunes only the unused pass-through columns and never the kernel inputs.

An earlier revision of this PR added a dedicated ColumnPruning arm for Project(_, b: BinBy). Per review, it was found redundant with the generic fallback: after the produced-attributes change the arm's requiredAttrs = (p.references -- b.producedAttributes) ++ b.references prunes the child identically to the fallback's child.references ++ p.references (the subtracted producedAttributes have fresh ExprIds absent from the child, so subtracting them changes nothing), and both guards fire on the same plans. The arm was removed; this PR now only locks in the pruning behavior with a test.

  • ColumnPruningSuite: a Project over BinBy prunes the unused pass-through column while keeping the range and DISTRIBUTE inputs in the child.
  • dsl/package.scala: a binBy test DSL helper (alongside generate) used by the test.

Why are the changes needed?

BIN BY is row-multiplying, so forwarding unused pass-through columns through the operator wastes work proportional to the fan-out. The generic ColumnPruning fallback already prunes them; this PR adds coverage so that behavior is not silently lost.

Does this PR introduce any user-facing change?

No. BIN BY is gated off by default (spark.sql.binByRelationOperator.enabled, SPARK-57440). This is test-only, plus a test DSL helper.

How was this patch tested?

  • ColumnPruningSuite: the new Column pruning for BinBy test.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Anthropic)

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

0 blocking, 1 non-blocking, 0 nits. Correct and result-safe; one non-blocking design suggestion on the pruning mechanism.

Design / architecture (1)

  • Optimizer.scala:1139: unrequiredChildIndex machinery buys only <=16 bytes/row and the Generate analogy doesn't transfer to BinBy — see inline

Verification

Traced the Project(_, BinBy) rewrite: the pruning Project (prunedChild) keeps every kernel-read input (requiredAttrs ⊇ b.references), and the unrequired set reduces to {rangeStart, rangeEnd} -- p.references — DISTRIBUTE columns are gated out, so no kernel input is ever pruned. Result-equivalent for empty / single / many-row input, and the synthesized Project forwards child attributes with unchanged nullability/dataType/exprId (no metadata mismatch). The rewrite is correct; the finding below is a simplification, not a correctness concern.

@vranes vranes force-pushed the bin-by-column-pruning branch from 6672bba to b37456c Compare July 10, 2026 09:45

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-review: 1 addressed, 0 remaining, 1 new (late catch). The prior suggestion (drop the unrequiredChildIndex field / output override) is addressed — switched to the plain prunedChild form. 0 blocking, 1 non-blocking, 0 nits. Correct and result-safe; one non-blocking simplification question on the surviving pruning arm.

Design / architecture (1)

  • Optimizer.scala:1140: the dedicated Project(_, b: BinBy) arm looks redundant with the generic Project(_, child) fallback at :1199BinBy isn't intercepted in between, and both reduce to the same prunedChild(b.child, ...). See inline.

Verification

Traced Project(_, BinBy) routing: no arm between :1140 and :1199 matches it (GeneratorNestedColumnAliasing handles only Generate; NestedColumnAliasing's canProjectPushThrough/canPruneOn both exclude BinBy), so absent the new arm the pattern falls through to the generic fallback. The arm's requiredAttrs = (p.references -- b.producedAttributes) ++ b.references and the fallback's child.references ++ p.references filter b.child.output to the same set — b.producedAttributes carry fresh ExprIds and are never in b.child.output, so subtracting them changes nothing — and the two guards fire on the same plans. The arm is therefore result-equivalent to the fallback (hence a redundancy finding, not a correctness one), and result-safe across empty / single / many-row input with no output-attribute metadata change.


// prune unused pass-through columns from the child of BinBy. The range and DISTRIBUTE inputs
// the kernel reads are kept via `b.references`; only fully-unused pass-throughs are dropped.
case p @ Project(_, b: BinBy)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This arm is correct, but it looks functionally redundant with the generic Project(_, child) fallback at :1199.

Project(_, BinBy) isn't intercepted by any arm between here and :1199: GeneratorNestedColumnAliasing (:1146) matches only Project(_, g: Generate), and NestedColumnAliasing (:1195) matches Project(projectList, child) only when canProjectPushThrough(child) or canPruneOn(child) — and BinBy is in neither list (NestedColumnAliasing.scala:199-219). So without this arm, Project(_, BinBy) falls through to the generic fallback:

case p @ Project(_, child) if !child.isInstanceOf[Project] =>
  val required = child.references ++ p.references
  ...
  val newChildren = child.children.map(c => prunedChild(c, required))

For child = b: BinBy that computes required = b.references ++ p.references and calls the same prunedChild(b.child, required). Compared to this arm's requiredAttrs = (p.references -- b.producedAttributes) ++ b.references: b.producedAttributes (scaled DISTRIBUTE + appended attrs) all carry fresh ExprIds, so none are in b.child.output; subtracting them from p.references can't change which child attributes survive prunedChild's filter. The two reference sets prune the child identically, and the two guards fire on the same plans.

Unlike the Generate arm (which must set unrequiredChildIndex, a field the generic fallback can't populate) or the Expand arm (which rewrites projections), this arm does no irreducible work beyond prunedChild — so it appears eliminable. The new ColumnPruningSuite test would also still pass with the arm deleted, so it doesn't actually distinguish the arm from the fallback.

Is there a reason to keep the explicit arm — a planned future field on BinBy, or a deliberate choice to make the pruning discoverable at this call site? If not, consider dropping it and relying on the generic fallback. Non-blocking either way.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, you're right. Removed the arm. I kept the test (it still passes) and retitled/reframed the PR + JIRA since it's now test-only.

@vranes vranes changed the title [SPARK-58063][SQL] Add column pruning for BIN BY [SPARK-58063][SQL] Test that BIN BY is column-pruned by the generic ColumnPruning fallback Jul 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants