sunchao commented on code in PR #57437: URL: https://github.com/apache/spark/pull/57437#discussion_r3641014955
########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReusableBroadcastValueProjection.scala: ########## @@ -0,0 +1,126 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.optimizer + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.trees.TreePattern.{CTE, DYNAMIC_PRUNING_SUBQUERY} +import org.apache.spark.sql.types.{DateType, StringType, TimestampNTZType, TimestampType} + +/** + * Finds an existing broadcast whose stored rows provide a safe superset of a pruning domain. + */ +private[sql] object ReusableBroadcastValueProjection extends PredicateHelper { + + private def isSafeValueExpression(expression: Expression): Boolean = { Review Comment: Fixed in `5c57963`. Plain default-collation `StringType` attributes are now supported; the existing distinct-value byte limit still bounds variable-width output, and non-binary collations remain conservatively excluded. `DateFormatClass` now accepts any non-null, resolved string-literal pattern instead of hardcoding `yyyy-MM-dd`. The new `reuse ancestor broadcast string values with a residual join predicate` regression uses the `products.category` example, verifies both the exact query result and the projected `books`, `toys`, `clothing` superset, and runs for V1/V2 scans with adaptive execution both enabled and disabled. The reordered-column regression separately exercises the `yyyyMMdd` format. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala: ########## @@ -46,47 +46,86 @@ case class PlanDynamicPruningFilters(sparkSession: SparkSession) extends Rule[Sp HashedRelationBroadcastMode(packedKeys) } + private def reusableBroadcast( Review Comment: Fixed in `5c57963`. The non-adaptive direct-reuse path now preserves the existing `candidatePlan.sameResult(sparkPlan)` behavior, including mismatched hash modes and null-aware joins. Exact hash-mode matching and the null-aware restriction apply only to the new projected-broadcast path. I also added `preserve direct broadcast pruning for existing hash and null-aware modes`, which explicitly verifies both pre-existing cases with the new optimization disabled. This keeps the feature from changing released direct-reuse behavior. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/ProjectedBroadcastValueSubqueryExec.scala: ########## @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution + +import java.util.concurrent.{Future => JFuture} + +import scala.concurrent.duration.Duration + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Cast, Expression, NamedExpression, UnsafeRow} +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.sql.execution.BroadcastValueResult.{Available, Unavailable} +import org.apache.spark.sql.execution.joins.HashedRelation +import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.util.ThreadUtils + +/** Collects pruning values from the full rows of an already required hash broadcast. */ +case class ProjectedBroadcastValueSubqueryExec( + name: String, + valueExpression: Expression, + child: SparkPlan) extends BaseSubqueryExec with UnaryExecNode { + + override def output: Seq[Attribute] = { + val outputName = valueExpression match { + case named: NamedExpression => named.name + case Cast(named: NamedExpression, _, _, _) => named.name + case _ => "key" + } + Seq(AttributeReference(outputName, valueExpression.dataType, valueExpression.nullable)()) + } + + override lazy val metrics = Map( + "numInputRows" -> SQLMetrics.createMetric(sparkContext, "number of input rows"), + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), + "dataSize" -> SQLMetrics.createMetric(sparkContext, "data size (bytes)"), + "projectionDisabled" -> SQLMetrics.createMetric(sparkContext, "projection disabled"), + "projectionErrors" -> SQLMetrics.createMetric(sparkContext, "projection errors"), + "collectTime" -> SQLMetrics.createMetric(sparkContext, "time to collect (ms)")) + + override def doCanonicalize(): SparkPlan = { + ProjectedBroadcastValueSubqueryExec( + "dpp", + QueryPlan.normalizeExpressions(valueExpression, child.output), + child.canonicalized) + } + + @transient + private lazy val relationFuture: JFuture[BroadcastValueResult[Array[InternalRow]]] = { + val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + SQLExecution.withThreadLocalCaptured[BroadcastValueResult[Array[InternalRow]]]( + session, SubqueryBroadcastExec.executionContext) { + SQLExecution.withExecutionId(session, executionId) { + val beforeCollect = System.nanoTime() + val broadcast = child.executeBroadcast[HashedRelation]() + val limits = BroadcastValueProjectionLimits( + maxInputRows = conf.dynamicPartitionPruningBroadcastProjectionMaxRows.toLong, + maxOutputBytes = conf.dynamicPartitionPruningBroadcastProjectionMaxBytes, + maxSourceBytes = conf.dynamicPartitionPruningBroadcastProjectionMaxSourceBytes) + val result = BroadcastValueProjector.collectExactValueDomain( + broadcast, + child, + valueExpression, + limits, + onInputRow = () => longMetric("numInputRows") += 1, + onError = () => longMetric("projectionErrors") += 1) + + longMetric("collectTime") += (System.nanoTime() - beforeCollect) / 1000000 + result match { + case Available(rows) => + longMetric("numOutputRows") += rows.length + longMetric("dataSize") += + rows.iterator.map(_.asInstanceOf[UnsafeRow].getSizeInBytes.toLong).sum + case Unavailable => + longMetric("projectionDisabled") += 1 + } + SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq) + result + } + } + } + + override protected def doPrepare(): Unit = { + relationFuture + } + + override protected def doExecute(): RDD[InternalRow] = { + throw QueryExecutionErrors.executeCodePathUnsupportedError( + "ProjectedBroadcastValueSubqueryExec") + } + + override def executeCollect(): Array[InternalRow] = executeCollectResult() match { + case Available(rows) => rows + case Unavailable => Array.empty Review Comment: Fixed in `5c57963`. `executeCollect()` now throws `SparkException.internalError` when the projected domain is unavailable rather than returning an empty array. Dynamic partition pruning explicitly consumes the typed `executeCollectResult()`/`resultOf` path and still fails open. The regression exercises the row, projected-byte, and source-byte limits and asserts that direct `executeCollect()` raises `INTERNAL_ERROR` in each case. An unavailable domain therefore cannot silently become an empty pruning domain. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/subquery.scala: ########## @@ -127,23 +127,33 @@ case class InSubqueryExec( override def nullable: Boolean = child.nullable override def toString: String = s"$child IN ${plan.name}" - override def withNewPlan(plan: BaseSubqueryExec): InSubqueryExec = copy(plan = plan) + override def withNewPlan(plan: BaseSubqueryExec): InSubqueryExec = + copy(plan = plan, resultBroadcast = null, result = null) final override def nodePatternsInternal(): Seq[TreePattern] = Seq(IN_SUBQUERY_EXEC) def updateResult(): Unit = { - val rows = plan.executeCollect() - result = if (plan.output.length > 1) { + val (rows, unavailable) = ProjectedBroadcastValueSubqueryExec.resultOf(plan) match { + case Some(BroadcastValueResult.Available(values)) => (values, false) + case Some(BroadcastValueResult.Unavailable) => (Array.empty[InternalRow], true) + case None => (plan.executeCollect(), false) + } + result = if (unavailable) { Review Comment: Fixed in `5c57963`; `updateResult()` now asserts `isDynamicPruning` before installing the unavailable-result marker. The limit regressions also copy the same expression with `isDynamicPruning = false` and verify that `updateResult()` raises `AssertionError`. Ordinary SQL `IN` therefore cannot accidentally use dynamic pruning's fail-open semantics. -- 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]
