This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new da4c3c22768f [SPARK-58063][SQL] Test that BIN BY is column-pruned by
the generic ColumnPruning fallback
da4c3c22768f is described below
commit da4c3c22768f443c36cbd530c39971a1e66ce9fd
Author: Nikolina Vraneš <[email protected]>
AuthorDate: Sat Jul 11 17:07:46 2026 +0800
[SPARK-58063][SQL] Test that BIN BY is column-pruned by the generic
ColumnPruning fallback
### 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 [...]
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 `ExprId`s absent from the child, so
subtracting them changes nothing), and both [...]
- `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)
Closes #57157 from vranes/bin-by-column-pruning.
Authored-by: Nikolina Vraneš <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../apache/spark/sql/catalyst/dsl/package.scala | 12 +++++++++
.../catalyst/optimizer/ColumnPruningSuite.scala | 31 +++++++++++++++++++++-
2 files changed, 42 insertions(+), 1 deletion(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
index 3ba538c0e216..b4acde7026b6 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
@@ -524,6 +524,18 @@ package object dsl extends SQLConfHelper {
Generate(generator, unrequiredChildIndex, outer,
alias, outputNames.map(UnresolvedAttribute(_)), logicalPlan)
+ def binBy(
+ rangeStart: Attribute,
+ rangeEnd: Attribute,
+ distributeColumns: Seq[Attribute],
+ scaledDistributeColumns: Seq[Attribute],
+ appendedAttributes: Seq[Attribute],
+ binWidthMicros: Long = 300000000L,
+ originMicros: Long = 0L,
+ timeZoneId: Option[String] = Some("UTC")): LogicalPlan =
+ BinBy(binWidthMicros, rangeStart, rangeEnd, originMicros,
distributeColumns,
+ scaledDistributeColumns, appendedAttributes, logicalPlan, timeZoneId)
+
def insertInto(tableName: String): LogicalPlan =
insertInto(table(tableName))
def insertInto(
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ColumnPruningSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ColumnPruningSuite.scala
index 266c369894ec..e66cb800f773 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ColumnPruningSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ColumnPruningSuite.scala
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.plans.{Inner, PlanTest}
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.rules.RuleExecutor
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{StringType, StructType}
+import org.apache.spark.sql.types.{DoubleType, StringType, StructType,
TimestampType}
class ColumnPruningSuite extends PlanTest {
@@ -482,4 +482,33 @@ class ColumnPruningSuite extends PlanTest {
comparePlans(CustomOptimize.execute(originalQuery.analyze),
correctAnswer.analyze)
}
+
+ test("Column pruning for BinBy") {
+ // Two range inputs, one DISTRIBUTE column, one unused pass-through.
+ val tsStart = AttributeReference("ts_start", TimestampType, nullable =
false)()
+ val tsEnd = AttributeReference("ts_end", TimestampType, nullable = false)()
+ val value = AttributeReference("value", DoubleType, nullable = false)()
+ val label = AttributeReference("label", StringType, nullable = true)()
+ val relation = LocalRelation(tsStart, tsEnd, value, label)
+
+ // Produced attributes: scaled DISTRIBUTE column + three appended columns.
+ val scaledValue = AttributeReference("value", DoubleType, nullable =
true)()
+ val binStart = AttributeReference("bin_start", TimestampType, nullable =
true)()
+ val binEnd = AttributeReference("bin_end", TimestampType, nullable =
true)()
+ val binRatio = AttributeReference("bin_distribute_ratio", DoubleType,
nullable = true)()
+
+ val query = relation
+ .binBy(tsStart, tsEnd, Seq(value), Seq(scaledValue), Seq(binStart,
binEnd, binRatio))
+ .select(binStart)
+
+ val optimized = Optimize.execute(query)
+
+ // `label` is pruned from the child; the range and DISTRIBUTE inputs the
kernel reads stay.
+ val correctAnswer = relation
+ .select(tsStart, tsEnd, value)
+ .binBy(tsStart, tsEnd, Seq(value), Seq(scaledValue), Seq(binStart,
binEnd, binRatio))
+ .select(binStart)
+
+ comparePlans(optimized, correctAnswer, checkAnalysis = false)
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
