diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 597a70e01a6e..53ddc53091fb 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -2311,6 +2311,7 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe case _: BatchEvalPython => true case _: ArrowEvalPython => true case _: Expand => true + case _: BinBy => true case _ => false } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala index f457eced4d98..a43be9a1c0a6 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala @@ -1644,4 +1644,50 @@ class FilterPushdownSuite extends PlanTest { .analyze comparePlans(optimizedQueryWithoutStep, correctAnswer) } + + test("push down deterministic predicate through BinBy") { + // Relation: ts_start, ts_end, value (DISTRIBUTE), label (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 output and appended BIN BY 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)() + + def binByOver(child: LogicalPlan): BinBy = BinBy( + binWidthMicros = 300000000L, + rangeStart = tsStart, + rangeEnd = tsEnd, + originMicros = 0L, + distributeColumns = Seq(value), + scaledDistributeColumns = Seq(scaledValue), + appendedAttributes = Seq(binStart, binEnd, binRatio), + child = child, + timeZoneId = Some("UTC")) + + // (a) A predicate on a forwarded pass-through column (label) must push BELOW BinBy. + val queryA = Filter(EqualTo(label, Literal("x")), binByOver(relation)) + val optimizedA = Optimize.execute(queryA) + val expectedA = binByOver(Filter(EqualTo(label, Literal("x")), relation)) + comparePlans(optimizedA, expectedA, checkAnalysis = false) + + // (b) A predicate on a range column (ts_start) must push BELOW BinBy. The range columns are + // consumed by the binning logic yet forwarded unchanged in output, so their value is identical + // on every sub-row and filtering before vs. after binning is equivalent. + val tsLiteral = Literal(0L, TimestampType) + val queryB = Filter(GreaterThan(tsStart, tsLiteral), binByOver(relation)) + val optimizedB = Optimize.execute(queryB) + val expectedB = binByOver(Filter(GreaterThan(tsStart, tsLiteral), relation)) + comparePlans(optimizedB, expectedB, checkAnalysis = false) + + // (c) A predicate on a produced (fresh-ExprId) appended column must stay ABOVE BinBy. + val queryC = Filter(GreaterThan(binRatio, Literal(0.5)), binByOver(relation)) + val optimizedC = Optimize.execute(queryC) + comparePlans(optimizedC, queryC, checkAnalysis = false) + } }