This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun 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 f521d01350d4 [SPARK-57851][SQL] Shuffle-free single-task execution for
small queries
f521d01350d4 is described below
commit f521d01350d4cee793aec98e2975b81f87c575ba
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Thu Jul 9 13:08:28 2026 -0700
[SPARK-57851][SQL] Shuffle-free single-task execution for small queries
### What changes were proposed in this pull request?
This adds a conservative optimizer rule `MarkSingleTaskExecution` that
marks small single-partition scans, optionally with a shuffle-inducing operator
on top (sort, aggregate, distinct, window, limit/offset, expand) or an
in-memory `LocalRelation`, as candidates for single-task execution. Such a scan
reports a `SinglePartition` output partitioning, allowing `EnsureRequirements`
to elide the shuffle that would otherwise be inserted before the operator on
top.
Details:
- The rule runs as the last optimizer batch (so it sees the final plan
shape) and marks eligible `LogicalRelation`/`LocalRelation` nodes with a
`TreeNodeTag`.
- `FileSourceStrategy`/`SparkStrategies` propagate the mark to
`FileSourceScanExec`/`LocalTableScanExec`.
- `FileSourceScanExec` additionally gates on file count and size thresholds
using the generic `ScanFileListing`, reports `SinglePartition`, and coalesces
its input RDD to a single partition as a correctness backstop when the estimate
does not match the runtime partition count.
- `LocalTableScanExec` reads its data in a single partition and reports
`SinglePartition`.
- `ExpandExec` forwards `SinglePartition` from its child, since Expand only
replicates rows within a partition and never moves rows across partitions.
The feature is controlled by new internal configs under
`spark.sql.optimizer.singleTaskExecution.*` and is disabled by default. Join is
intentionally left out for now and can be added as a follow-up; union is
already covered by the existing `spark.sql.unionOutputPartitioning`.
This is part of the SPIP umbrella
[SPARK-56978](https://issues.apache.org/jira/browse/SPARK-56978) (Faster
queries in local laptop mode), covering the shuffle-free local execution for
small queries category.
### Why are the changes needed?
For small, low-latency queries the fixed cost of a shuffle (scheduling,
serialization, network) dominates the total runtime. When the input is already
a single small partition, the shuffle inserted before a sort/aggregate/window
is unnecessary and can be removed to reduce latency, without affecting
correctness.
### Does this PR introduce _any_ user-facing change?
No. The optimization is behind internal configs
(`spark.sql.optimizer.singleTaskExecution.*`) and is disabled by default.
### How was this patch tested?
New `MarkSingleTaskExecutionSuite` (14 tests) covering:
- the marking decision for the supported plan shapes;
- `SinglePartition` output with no shuffle in the final physical plan;
- empty-scan correctness (a global aggregation over an empty scan still
returns a single row);
- disabled-flag negatives (master flag and per-operator sub-flags);
- ineligibility of unsupported shapes (join) and subquery expressions;
- the leaf-node parallelism override disabling the local-relation case.
`SQLConfSuite` passes as a config-wiring regression check.
### Was this patch authored or co-authored using generative AI tooling?
Yes, using Claude Code.
Closes #56928 from viirya/single-node-execution.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../expressions/aggregate/V2Aggregator.scala | 3 +
.../spark/sql/catalyst/trees/TreePatterns.scala | 1 +
.../org/apache/spark/sql/internal/SQLConf.scala | 106 ++++++++
.../execution/SparkConnectPlanExecution.scala | 2 +-
.../spark/sql/execution/DataSourceScanExec.scala | 62 ++++-
.../apache/spark/sql/execution/ExpandExec.scala | 31 ++-
.../spark/sql/execution/LocalTableScanExec.scala | 31 ++-
.../spark/sql/execution/SparkOptimizer.scala | 7 +-
.../spark/sql/execution/SparkStrategies.scala | 10 +-
.../aggregate/TypedAggregateExpression.scala | 5 +
.../spark/sql/execution/aggregate/udaf.scala | 5 +
.../execution/datasources/FileSourceStrategy.scala | 4 +-
.../datasources/MarkSingleTaskExecution.scala | 177 ++++++++++++
.../org/apache/spark/sql/DataFrameJoinSuite.scala | 2 +-
.../scala/org/apache/spark/sql/SubquerySuite.scala | 2 +-
.../datasources/MarkSingleTaskExecutionSuite.scala | 301 +++++++++++++++++++++
.../scala/org/apache/spark/sql/hive/hiveUDFs.scala | 3 +
17 files changed, 729 insertions(+), 23 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/V2Aggregator.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/V2Aggregator.scala
index 49ba2ec8b904..e92e61245201 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/V2Aggregator.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/V2Aggregator.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.expressions.aggregate
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Expression,
ImplicitCastInputTypes, UnsafeProjection}
+import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern,
USER_DEFINED_AGGREGATION}
import org.apache.spark.sql.connector.catalog.functions.{AggregateFunction =>
V2AggregateFunction}
import org.apache.spark.sql.types.{AbstractDataType, DataType}
import org.apache.spark.util.ArrayImplicits._
@@ -31,6 +32,8 @@ case class V2Aggregator[BUF <: java.io.Serializable, OUT](
inputAggBufferOffset: Int = 0)
extends TypedImperativeAggregate[BUF] with ImplicitCastInputTypes {
+ final override val nodePatterns: Seq[TreePattern] =
Seq(USER_DEFINED_AGGREGATION)
+
private[this] lazy val inputProjection = UnsafeProjection.create(children)
override def nullable: Boolean = aggrFunc.isResultNullable
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala
index 94b4666a88a8..dfb815414dd3 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala
@@ -107,6 +107,7 @@ object TreePattern extends Enumeration {
val TIME_WINDOW: Value = Value
val TIME_ZONE_AWARE_EXPRESSION: Value = Value
val TRUE_OR_FALSE_LITERAL: Value = Value
+ val USER_DEFINED_AGGREGATION: Value = Value
val VARIANT_GET: Value = Value
val WINDOW_EXPRESSION: Value = Value
val WINDOW_TIME: Value = Value
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index ca31519a7069..3315cdc72332 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -7407,6 +7407,112 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val SINGLE_TASK_EXECUTION_ENABLED =
+ buildConf("spark.sql.optimizer.singleTaskExecution.enabled")
+ .doc("When true, eligible query fragments that read a small
single-partition scan can run " +
+ "in a single task, skipping the shuffle that would otherwise be
inserted before an " +
+ "operator such as a sort or aggregation. This avoids the scheduling
overhead of an " +
+ "unnecessary shuffle for small, low-latency queries.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(false)
+
+ val SINGLE_TASK_EXECUTION_AGGREGATION =
+ buildConf("spark.sql.optimizer.singleTaskExecution.aggregation")
+ .internal()
+ .doc("When true, and 'spark.sql.optimizer.singleTaskExecution.enabled'
is also true, " +
+ "enable the single-task optimization for query plans with aggregation
operators.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .fallbackConf(SINGLE_TASK_EXECUTION_ENABLED)
+
+ val SINGLE_TASK_EXECUTION_EXPAND =
+ buildConf("spark.sql.optimizer.singleTaskExecution.expand")
+ .internal()
+ .doc("When true, and 'spark.sql.optimizer.singleTaskExecution.enabled'
is also true, " +
+ "enable the single-task optimization for query plans with expand
operators.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .fallbackConf(SINGLE_TASK_EXECUTION_ENABLED)
+
+ val SINGLE_TASK_EXECUTION_LIMIT_OFFSET =
+ buildConf("spark.sql.optimizer.singleTaskExecution.limitOffset")
+ .internal()
+ .doc("When true, and 'spark.sql.optimizer.singleTaskExecution.enabled'
is also true, " +
+ "enable the single-task optimization for query plans with limit or
offset operators.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .fallbackConf(SINGLE_TASK_EXECUTION_ENABLED)
+
+ val SINGLE_TASK_EXECUTION_SORT =
+ buildConf("spark.sql.optimizer.singleTaskExecution.sort")
+ .internal()
+ .doc("When true, and 'spark.sql.optimizer.singleTaskExecution.enabled'
is also true, " +
+ "enable the single-task optimization for query plans with sort
operators.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .fallbackConf(SINGLE_TASK_EXECUTION_ENABLED)
+
+ val SINGLE_TASK_EXECUTION_WINDOW =
+ buildConf("spark.sql.optimizer.singleTaskExecution.window")
+ .internal()
+ .doc("When true, and 'spark.sql.optimizer.singleTaskExecution.enabled'
is also true, " +
+ "enable the single-task optimization for query plans with window
operators.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .fallbackConf(SINGLE_TASK_EXECUTION_ENABLED)
+
+ val SINGLE_TASK_EXECUTION_MAX_NUM_FILES =
+ buildConf("spark.sql.optimizer.singleTaskExecution.maxNumFiles")
+ .internal()
+ .doc("The maximum number of files that a file scan may have for the
single-task " +
+ "optimization to apply to it.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .createWithDefault(1)
+
+ val SINGLE_TASK_EXECUTION_MIN_NUM_FILES =
+ buildConf("spark.sql.optimizer.singleTaskExecution.minNumFiles")
+ .internal()
+ .doc("The minimum number of files that a file scan may have for the
single-task " +
+ "optimization to apply to it.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .createWithDefault(1)
+
+ val SINGLE_TASK_EXECUTION_MIN_NUM_BYTES =
+ buildConf("spark.sql.optimizer.singleTaskExecution.minNumBytes")
+ .internal()
+ .doc("The minimum total size in bytes that a file scan may have for the
single-task " +
+ "optimization to apply to it.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .longConf
+ .createWithDefault(1)
+
+ val SINGLE_TASK_EXECUTION_LOCAL_TABLE_SCAN_MIN_ROWS =
+ buildConf("spark.sql.optimizer.singleTaskExecution.localTableScan.minRows")
+ .internal()
+ .doc("The minimum number of rows that a local in-memory relation may
have for the " +
+ "single-task optimization to apply to it.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .createWithDefault(1)
+
+ val SINGLE_TASK_EXECUTION_LOCAL_TABLE_SCAN_THRESHOLD =
+
buildConf("spark.sql.optimizer.singleTaskExecution.localTableScan.threshold")
+ .internal()
+ .doc("The maximum number of rows that a local in-memory relation may
have for the " +
+ "single-task optimization to apply to it.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .createWithDefault(1000)
+
val LEGACY_PARSE_QUERY_WITHOUT_EOF =
buildConf("spark.sql.legacy.parseQueryWithoutEof")
.internal()
.doc(
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala
index 5fdfd5d1ccd1..0c4ca9357e84 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala
@@ -213,7 +213,7 @@ private[execution] class
SparkConnectPlanExecution(executeHolder: ExecuteHolder)
}
}
dataframe.queryExecution.executedPlan match {
- case LocalTableScanExec(_, rows, _) =>
+ case LocalTableScanExec(_, rows, _, _) =>
executePlan.eventsManager.postFinished(Some(rows.length))
var offset = 0L
converter(rows.iterator).foreach { case (bytes, count) =>
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala
index be7013188f2f..a727ccf56506 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.{FileSourceOptions,
InternalRow, TableIdent
import org.apache.spark.sql.catalyst.catalog.BucketSpec
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.QueryPlan
-import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning,
Partitioning, UnknownPartitioning}
+import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning,
Partitioning, SinglePartition, UnknownPartitioning}
import org.apache.spark.sql.catalyst.util.{truncatedString, CaseInsensitiveMap}
import org.apache.spark.sql.connector.read.streaming.SparkDataStream
import org.apache.spark.sql.errors.QueryExecutionErrors
@@ -320,6 +320,34 @@ trait FileSourceScanLike extends DataSourceScanExec with
SessionStateHelper {
def requiredSchema: StructType
// Identifier for the table in the metastore.
def tableIdentifier: Option[TableIdentifier]
+ // When true, the `MarkSingleTaskExecution` optimizer rule has marked this
scan's plan shape as a
+ // candidate for single-task execution. The scan is only actually executed
in a single task when
+ // it additionally passes the file count and size thresholds (see
`useSingleTaskExecution`).
+ def markedForSingleTaskExecution: Boolean
+
+ /**
+ * Whether this file scan should run in a single task, reporting a
`SinglePartition` output
+ * partitioning so that a following shuffle can be elided. This is true when
the plan shape was
+ * marked eligible by the optimizer and the statically-selected files fall
within the configured
+ * count and size bounds. Bucketed scans are excluded: they report a
`HashPartitioning` over the
+ * bucket columns, which coalescing to a single partition would invalidate.
It relies on
+ * `selectedPartitions`, so it must not be evaluated before the scan's file
listing is available.
+ */
+ lazy val useSingleTaskExecution: Boolean = {
+ if (!markedForSingleTaskExecution || bucketedScan) {
+ false
+ } else {
+ val sqlConf = getSqlConf(relation.sparkSession)
+ val minNumFiles =
sqlConf.getConf(SQLConf.SINGLE_TASK_EXECUTION_MIN_NUM_FILES)
+ val maxNumFiles =
sqlConf.getConf(SQLConf.SINGLE_TASK_EXECUTION_MAX_NUM_FILES)
+ val minNumBytes =
sqlConf.getConf(SQLConf.SINGLE_TASK_EXECUTION_MIN_NUM_BYTES)
+ val maxPartitionBytes =
sqlConf.getConf(SQLConf.FILES_MAX_PARTITION_BYTES)
+ val numFiles = selectedPartitions.totalNumberOfFiles
+ val numBytes = selectedPartitions.totalFileSize
+ numFiles >= minNumFiles && numFiles <= maxNumFiles &&
+ numBytes >= minNumBytes && numBytes <= maxPartitionBytes
+ }
+ }
lazy val fileConstantMetadataColumns: Seq[AttributeReference] =
output.collect {
@@ -478,6 +506,8 @@ trait FileSourceScanLike extends DataSourceScanExec with
SessionStateHelper {
Nil
}
(partitioning, sortOrder)
+ } else if (useSingleTaskExecution) {
+ (SinglePartition, Nil)
} else {
(UnknownPartitioning(0), Nil)
}
@@ -696,7 +726,8 @@ case class FileSourceScanExec(
override val optionalNumCoalescedBuckets: Option[Int],
override val dataFilters: Seq[Expression],
override val tableIdentifier: Option[TableIdentifier],
- override val disableBucketedScan: Boolean = false)
+ override val disableBucketedScan: Boolean = false,
+ override val markedForSingleTaskExecution: Boolean = false)
extends FileSourceScanLike {
// Note that some vals referring the file-based relation are lazy
intentionally
@@ -744,10 +775,28 @@ case class FileSourceScanExec(
inputRDD :: Nil
}
+ /**
+ * The input RDD, coalesced to a single partition when this scan runs in
single-task mode. This
+ * enforces the `SinglePartition` output partitioning reported by
`outputPartitioning`, which is
+ * estimated from the statically-selected files and may not correspond
exactly to the number of
+ * partitions the input RDD produces after dynamic pruning. Coalescing here
keeps the query
+ * correct in either case.
+ */
+ private[spark] lazy val maybeCoalesceInputRDD: RDD[InternalRow] = {
+ if (useSingleTaskExecution && inputRDD.getNumPartitions > 1) {
+ inputRDD.coalesce(1)
+ } else if (useSingleTaskExecution && inputRDD.getNumPartitions == 0) {
+ // All files were pruned away; produce a single empty partition to match
`SinglePartition`.
+ sparkContext.parallelize[InternalRow](Nil, 1)
+ } else {
+ inputRDD
+ }
+ }
+
protected override def doExecute(): RDD[InternalRow] = {
val numOutputRows = longMetric("numOutputRows")
if (needsUnsafeRowConversion) {
- inputRDD.mapPartitionsWithIndexInternal { (index, iter) =>
+ maybeCoalesceInputRDD.mapPartitionsWithIndexInternal { (index, iter) =>
val toUnsafe = UnsafeProjection.create(schema)
toUnsafe.initialize(index)
iter.map { row =>
@@ -756,7 +805,7 @@ case class FileSourceScanExec(
}
}
} else {
- inputRDD.mapPartitionsInternal { iter =>
+ maybeCoalesceInputRDD.mapPartitionsInternal { iter =>
iter.map { row =>
numOutputRows += 1
row
@@ -768,7 +817,7 @@ case class FileSourceScanExec(
protected override def doExecuteColumnar(): RDD[ColumnarBatch] = {
val numOutputRows = longMetric("numOutputRows")
val scanTime = longMetric("scanTime")
- inputRDD.asInstanceOf[RDD[ColumnarBatch]].mapPartitionsInternal { batches
=>
+
maybeCoalesceInputRDD.asInstanceOf[RDD[ColumnarBatch]].mapPartitionsInternal {
batches =>
new Iterator[ColumnarBatch] {
override def hasNext: Boolean = {
@@ -921,7 +970,8 @@ case class FileSourceScanExec(
optionalNumCoalescedBuckets,
QueryPlan.normalizePredicates(dataFilters, output),
None,
- disableBucketedScan)
+ disableBucketedScan,
+ markedForSingleTaskExecution)
}
override def getStream: Option[SparkDataStream] = stream
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala
index 254772f73208..ba1238564348 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala
@@ -21,7 +21,7 @@ import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.codegen._
-import org.apache.spark.sql.catalyst.plans.physical.{Partitioning,
UnknownPartitioning}
+import org.apache.spark.sql.catalyst.plans.physical.{Partitioning,
SinglePartition, UnknownPartitioning}
import org.apache.spark.sql.execution.metric.SQLMetrics
import org.apache.spark.sql.internal.SQLConf
@@ -36,15 +36,38 @@ import org.apache.spark.sql.internal.SQLConf
case class ExpandExec(
projections: Seq[Seq[Expression]],
output: Seq[Attribute],
- child: SparkPlan)
+ child: SparkPlan,
+ // When true, this Expand is part of a plan marked for single-task
execution by the
+ // `MarkSingleTaskExecution` optimizer rule, and forwards the child's
`SinglePartition`
+ // output partitioning (see `outputPartitioning`).
+ useSingleTask: Boolean = false)
extends UnaryExecNode with CodegenSupport {
override lazy val metrics = Map(
"numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output
rows"))
// The GroupExpressions can output data with arbitrary partitioning, so set
it
- // as UNKNOWN partitioning
- override def outputPartitioning: Partitioning = UnknownPartitioning(0)
+ // as UNKNOWN partitioning. Expand only replicates rows within a partition
and never moves rows
+ // across partitions, so when this Expand is part of a plan marked for
single-task execution
+ // and the child produces a single partition, we can forward the
`SinglePartition` property to
+ // avoid an unneeded shuffle.
+ override def outputPartitioning: Partitioning = {
+ if (useSingleTask && child.outputPartitioning == SinglePartition) {
+ SinglePartition
+ } else {
+ UnknownPartitioning(0)
+ }
+ }
+
+ // Show `useSingleTask` in the string representation only when it is set, so
that plans not
+ // using single-task execution (the default) keep their existing explain
output.
+ override protected def stringArgs: Iterator[Any] = {
+ if (useSingleTask) {
+ super.stringArgs
+ } else {
+ Iterator(projections, output, child)
+ }
+ }
@transient
override lazy val references: AttributeSet =
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/LocalTableScanExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/LocalTableScanExec.scala
index 2d5dbf819959..21c9a43710fa 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/LocalTableScanExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/LocalTableScanExec.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeProjection}
+import org.apache.spark.sql.catalyst.plans.physical.{Partitioning,
SinglePartition, UnknownPartitioning}
import org.apache.spark.sql.connector.read.streaming.SparkDataStream
import org.apache.spark.sql.execution.metric.SQLMetrics
import org.apache.spark.util.ArrayImplicits._
@@ -34,7 +35,10 @@ import org.apache.spark.util.ArrayImplicits._
case class LocalTableScanExec(
output: Seq[Attribute],
@transient rows: Seq[InternalRow],
- @transient stream: Option[SparkDataStream])
+ @transient stream: Option[SparkDataStream],
+ // When true, the relation is scanned in a single partition, so this node
reports a
+ // `SinglePartition` output partitioning. Set by the
`MarkSingleTaskExecution` optimizer rule.
+ useSingleTask: Boolean = false)
extends LeafExecNode
with StreamSourceAwareSparkPlan
with InputRDDCodegen {
@@ -53,14 +57,33 @@ case class LocalTableScanExec(
@transient private lazy val rdd: RDD[InternalRow] = {
if (rows.isEmpty) {
- sparkContext.emptyRDD
+ if (useSingleTask) {
+ // Produce a single empty partition to match the `SinglePartition`
reported by
+ // `outputPartitioning`. `emptyRDD` has zero partitions, and running
e.g. a global
+ // aggregation on a zero-partition RDD with the shuffle elided would
return no rows
+ // instead of the single row expected on empty input.
+ sparkContext.parallelize(Seq.empty[InternalRow], 1)
+ } else {
+ sparkContext.emptyRDD
+ }
} else {
- val numSlices = math.min(
- unsafeRows.length, session.leafNodeDefaultParallelism)
+ val numSlices = if (useSingleTask) {
+ 1
+ } else {
+ math.min(unsafeRows.length, session.leafNodeDefaultParallelism)
+ }
sparkContext.parallelize(unsafeRows.toImmutableArraySeq, numSlices)
}
}
+ override def outputPartitioning: Partitioning = {
+ if (useSingleTask) {
+ SinglePartition
+ } else {
+ UnknownPartitioning(0)
+ }
+ }
+
protected override def doExecute(): RDD[InternalRow] = {
val numOutputRows = longMetric("numOutputRows")
rdd.map { r =>
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
index 1b3b2d3efc72..54158d5bb4a9 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.optimizer._
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.connector.catalog.CatalogManager
-import org.apache.spark.sql.execution.datasources.{PruneFileSourcePartitions,
PushVariantIntoScan, SchemaPruning, V1Writes}
+import org.apache.spark.sql.execution.datasources.{MarkSingleTaskExecution,
PruneFileSourcePartitions, PushVariantIntoScan, SchemaPruning, V1Writes}
import
org.apache.spark.sql.execution.datasources.v2.{GroupBasedRowLevelOperationScanPlanning,
OptimizeMetadataOnlyDeleteFromTable, V2ScanPartitioningAndOrdering,
V2ScanRelationPushDown, V2Writes}
import
org.apache.spark.sql.execution.dynamicpruning.{CleanupDynamicPruningFilters,
PartitionPruning, RowLevelOperationRuntimeGroupFiltering}
import
org.apache.spark.sql.execution.python.{ExtractGroupingPythonUDFFromAggregate,
ExtractPythonUDFFromAggregate, ExtractPythonUDFs, ExtractPythonUDTFs}
@@ -100,7 +100,10 @@ class SparkOptimizer(
ConstantFolding,
EliminateLimits),
Batch("User Provided Optimizers", fixedPoint,
experimentalMethods.extraOptimizations: _*),
- Batch("Replace CTE with Repartition", Once, ReplaceCTERefWithRepartition)))
+ Batch("Replace CTE with Repartition", Once, ReplaceCTERefWithRepartition),
+ // Must run last: it inspects the final plan shape to mark scans that can
run in a single task,
+ // and no subsequent rule should reshape the plan or copy the marked scan
nodes.
+ Batch("MarkSingleTaskExecution", Once, MarkSingleTaskExecution)))
override def nonExcludableRules: Seq[String] = super.nonExcludableRules ++
Seq(
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index 6e761fbe07b2..d89f7a919269 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -1151,7 +1151,9 @@ abstract class SparkStrategies extends
QueryPlanner[SparkPlan] {
case f: logical.TypedFilter =>
execution.FilterExec(f.typedCondition(f.deserializer),
planLater(f.child)) :: Nil
case e @ logical.Expand(_, _, child) =>
- execution.ExpandExec(e.projections, e.output, planLater(child)) :: Nil
+ val useSingleTask = e.getTagValue(
+ datasources.MarkSingleTaskExecution.markTag).getOrElse(false)
+ execution.ExpandExec(e.projections, e.output, planLater(child),
useSingleTask) :: Nil
case logical.Sample(lb, ub, withReplacement, seed, child, sampleMethod)
=>
if (sampleMethod == logical.SampleMethod.System) {
// V2ScanRelationPushDown is non-excludable and always handles
SYSTEM samples
@@ -1161,8 +1163,10 @@ abstract class SparkStrategies extends
QueryPlanner[SparkPlan] {
"TABLESAMPLE SYSTEM node was not properly handled by
V2ScanRelationPushDown.")
}
execution.SampleExec(lb, ub, withReplacement, seed, planLater(child))
:: Nil
- case logical.LocalRelation(output, data, _, stream) =>
- LocalTableScanExec(output, data, stream) :: Nil
+ case r @ logical.LocalRelation(output, data, _, stream) =>
+ val useSingleTask = r.getTagValue(
+ datasources.MarkSingleTaskExecution.markTag).getOrElse(false)
+ LocalTableScanExec(output, data, stream, useSingleTask) :: Nil
case logical.EmptyRelation(l) => EmptyRelationExec(l) :: Nil
case CommandResult(output, _, plan, data) => CommandResultExec(output,
plan, data) :: Nil
// We should match the combination of limit and offset first, to get the
optimal physical
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
index d958790dd09b..df0addad7861 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
@@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateFunction,
DeclarativeAggregate, TypedImperativeAggregate}
import org.apache.spark.sql.catalyst.expressions.codegen.GenerateSafeProjection
import org.apache.spark.sql.catalyst.expressions.objects.Invoke
+import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern,
USER_DEFINED_AGGREGATION}
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.types._
import org.apache.spark.util.Utils
@@ -125,6 +126,8 @@ case class SimpleTypedAggregateExpression(
nullable: Boolean)
extends DeclarativeAggregate with TypedAggregateExpression with
NonSQLExpression {
+ final override val nodePatterns: Seq[TreePattern] =
Seq(USER_DEFINED_AGGREGATION)
+
override lazy val deterministic: Boolean = true
override def children: Seq[Expression] = {
@@ -223,6 +226,8 @@ case class ComplexTypedAggregateExpression(
inputAggBufferOffset: Int = 0)
extends TypedImperativeAggregate[Any] with TypedAggregateExpression with
NonSQLExpression {
+ final override val nodePatterns: Seq[TreePattern] =
Seq(USER_DEFINED_AGGREGATION)
+
override lazy val deterministic: Boolean = true
override def children: Seq[Expression] = {
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
index 492f11607ce6..203ee2d89b7b 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
@@ -25,6 +25,7 @@ import
org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression
import
org.apache.spark.sql.catalyst.expressions.aggregate.{ImperativeAggregate,
TypedImperativeAggregate}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern,
USER_DEFINED_AGGREGATION}
import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer,
UserDefinedAggregateFunction, UserDefinedAggregator}
import org.apache.spark.sql.types._
@@ -358,6 +359,8 @@ case class ScalaUDAF(
with ImplicitCastInputTypes
with UserDefinedExpression {
+ final override val nodePatterns: Seq[TreePattern] =
Seq(USER_DEFINED_AGGREGATION)
+
override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int):
ImperativeAggregate =
copy(mutableAggBufferOffset = newMutableAggBufferOffset)
@@ -500,6 +503,8 @@ case class ScalaAggregator[IN, BUF, OUT](
with ImplicitCastInputTypes
with Logging {
+ final override val nodePatterns: Seq[TreePattern] =
Seq(USER_DEFINED_AGGREGATION)
+
// input and buffer encoders are resolved by ResolveEncodersInScalaAgg
@transient private[this] lazy val inputDeserializer =
inputEncoder.createDeserializer()
@transient private[this] lazy val bufferSerializer =
bufferEncoder.createSerializer()
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala
index b7a1736bc2e9..e2427222d8eb 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala
@@ -337,7 +337,9 @@ object FileSourceStrategy extends Strategy with
PredicateHelper with Logging {
bucketSet,
None,
rebindFileSourceMetadataAttributesInFilters(expandedDataFilters),
- table.map(_.identifier))
+ table.map(_.identifier),
+ markedForSingleTaskExecution =
+ l.getTagValue(MarkSingleTaskExecution.markTag).getOrElse(false))
// extra Project node: wrap flat metadata columns to a metadata struct
val withMetadataProjections = metadataStructOpt.map { metadataStruct =>
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/MarkSingleTaskExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/MarkSingleTaskExecution.scala
new file mode 100644
index 000000000000..c1a5eda5c084
--- /dev/null
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/MarkSingleTaskExecution.scala
@@ -0,0 +1,177 @@
+/*
+ * 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.datasources
+
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * This optimizer rule marks eligible query plans for single-task execution.
The optimization
+ * targets a conservative, specific query shape to ensure predictable and
efficient behavior.
+ *
+ * The rule matches simple query plans with a single small file scan or a
single small in-memory
+ * relation, optionally with a shuffle-inducing operator (sort, aggregation,
window, expand, or
+ * limit/offset) on top. When it detects such a shape, it marks the underlying
scan:
+ *
+ * - a [[LogicalRelation]] or [[LocalRelation]] is marked with the
+ * [[MarkSingleTaskExecution.markTag]] tag, as is any [[Expand]] in the
plan so that the
+ * physical Expand can forward the child's `SinglePartition` output
partitioning.
+ *
+ * The physical scan then reports a `SinglePartition` output partitioning,
which allows
+ * [[org.apache.spark.sql.execution.exchange.EnsureRequirements]] to elide the
shuffle that would
+ * otherwise be inserted before the operator on top. This shuffle is not
required for correctness
+ * of the query, so removing it reduces scheduling overhead for small,
low-latency queries.
+ *
+ * The matching is deliberately strict and conservative to minimize the risk
of unintended
+ * performance regressions. It can be broadened in the future as needed.
+ *
+ * This rule is controlled by [[SQLConf.SINGLE_TASK_EXECUTION_ENABLED]] and
the per-operator
+ * sub-flags in [[SQLConf]].
+ */
+object MarkSingleTaskExecution extends Rule[LogicalPlan] {
+
+ /**
+ * Tag placed on a [[LogicalRelation]] or [[LocalRelation]] that has been
marked eligible for
+ * single-task execution, and on any [[Expand]] in such a plan. The planning
strategies read
+ * this tag to propagate the decision to the physical
+ * [[org.apache.spark.sql.execution.FileSourceScanExec]] /
+ * [[org.apache.spark.sql.execution.LocalTableScanExec]] /
+ * [[org.apache.spark.sql.execution.ExpandExec]].
+ */
+ val markTag: TreeNodeTag[Boolean] =
TreeNodeTag[Boolean]("__single_task_execution")
+
+ /**
+ * Plan patterns that make a query ineligible for the optimization. These
operators either
+ * require shuffles that we cannot safely elide, or run user code whose
behavior we should not
+ * change. User-defined aggregations are excluded defensively: an
optimization that collapses
+ * the partial and final aggregates when no exchange separates them would
skip the user's merge
+ * step, so single-task plans must never be assumed safe for them.
+ */
+ val unsupportedPatterns: Seq[TreePattern] = Seq(
+ EVAL_PYTHON_UDF,
+ EVAL_PYTHON_UDTF,
+ EXISTS_SUBQUERY,
+ FUNCTION_TABLE_RELATION_ARGUMENT_EXPRESSION,
+ LATERAL_SUBQUERY,
+ LIST_SUBQUERY,
+ PYTHON_UDF,
+ SCALAR_SUBQUERY,
+ USER_DEFINED_AGGREGATION)
+
+ /**
+ * The per-operator sub-flags, resolved once per invocation. Each field
indicates whether the
+ * corresponding shuffle-inducing operator is allowed on top of a single
small scan.
+ */
+ private case class EnabledOperators(
+ aggregation: Boolean,
+ expand: Boolean,
+ limitOffset: Boolean,
+ sort: Boolean,
+ window: Boolean)
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ // An explicit leaf-node parallelism override expresses the user's intent
about how many
+ // partitions leaf scans should produce, so do not force scans into a
single partition.
+ if (!conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_ENABLED) ||
+ conf.getConf(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM).isDefined) {
+ return plan
+ }
+ val enabled = EnabledOperators(
+ aggregation = conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_AGGREGATION),
+ expand = conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_EXPAND),
+ limitOffset = conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_LIMIT_OFFSET),
+ sort = conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_SORT),
+ window = conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_WINDOW))
+
+ if (plan.containsAnyPattern(unsupportedPatterns: _*)) {
+ plan
+ } else if (isSupportedShape(plan, enabled)) {
+ // Mark a private clone of the plan rather than the plan itself, so that
this rule never
+ // mutates a node it was handed: a tag set on a shared node would
propagate through
+ // `TreeNode.clone`/`copyTagsFrom` and could leak the marking into
unrelated plans. Note
+ // that marking a per-node copy would not work either: a copy differing
only in tags is
+ // structurally equal to the original, so tree-rebuilding APIs such as
`withNewChildren`
+ // would discard it and keep the original node.
+ val cloned = plan.clone()
+ markSingleTaskExecution(cloned)
+ cloned
+ } else {
+ plan
+ }
+ }
+
+ /**
+ * Returns true if every operator in the plan is one that we support keeping
on top of a single
+ * small scan. Only operators that either do not require a shuffle, or whose
shuffle-inducing
+ * sub-flag is enabled, are allowed. Any other operator makes the plan
ineligible.
+ */
+ private def isSupportedShape(plan: LogicalPlan, enabled: EnabledOperators):
Boolean = plan match {
+ case _: LogicalRelation | _: LocalRelation => true
+ // Operators that never introduce a shuffle by themselves. Note that
`Distinct` and
+ // `SubqueryAlias` need no cases here: they are rewritten away by
non-excludable rules
+ // (`ReplaceDistinctWithAggregate` and `EliminateSubqueryAliases`) long
before this rule runs.
+ case _: Project | _: Filter |
+ _: DeserializeToObject | _: SerializeFromObject =>
+ plan.children.forall(isSupportedShape(_, enabled))
+ // Shuffle-inducing operators, allowed only when the matching sub-flag is
enabled.
+ case _: Aggregate if enabled.aggregation =>
+ plan.children.forall(isSupportedShape(_, enabled))
+ case _: Expand if enabled.expand =>
+ plan.children.forall(isSupportedShape(_, enabled))
+ case (_: GlobalLimit | _: LocalLimit | _: Offset) if enabled.limitOffset =>
+ plan.children.forall(isSupportedShape(_, enabled))
+ case _: Sort if enabled.sort =>
+ plan.children.forall(isSupportedShape(_, enabled))
+ case _: Window if enabled.window =>
+ plan.children.forall(isSupportedShape(_, enabled))
+ case _ => false
+ }
+
+ /**
+ * Sets the mark on each scan in the given (already validated) plan. Marking
mutates the nodes
+ * in place, which is only safe because the caller passes this rule's
private clone of the plan.
+ */
+ private def markSingleTaskExecution(plan: LogicalPlan): Unit = plan match {
+ case lr: LogicalRelation =>
+ lr.setTagValue(markTag, true)
+ case r: LocalRelation =>
+ if (isLocalRelationEligible(r)) {
+ r.setTagValue(markTag, true)
+ }
+ case e: Expand =>
+ // Also mark the Expand itself: the physical `ExpandExec` reads this tag
to forward the
+ // child's `SinglePartition` output partitioning, which it must only do
within a plan
+ // marked for single-task execution.
+ e.setTagValue(markTag, true)
+ e.children.foreach(markSingleTaskExecution)
+ case other =>
+ other.children.foreach(markSingleTaskExecution)
+ }
+
+ /**
+ * A local in-memory relation is eligible when its row count falls within
the configured bounds.
+ */
+ private def isLocalRelationEligible(r: LocalRelation): Boolean = {
+ val minRows =
conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_LOCAL_TABLE_SCAN_MIN_ROWS)
+ val threshold =
conf.getConf(SQLConf.SINGLE_TASK_EXECUTION_LOCAL_TABLE_SCAN_THRESHOLD)
+ r.data.length >= minRows && r.data.length <= threshold
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
index 9733d51a91cb..129d6bb68676 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
@@ -447,7 +447,7 @@ class DataFrameJoinSuite extends SharedSparkSession
}
assert(broadcastExchanges.size == 1)
val tables = broadcastExchanges.head.collect {
- case FileSourceScanExec(_, _, _, _, _, _, _, _,
Some(tableIdent), _) => tableIdent
+ case FileSourceScanExec(_, _, _, _, _, _, _, _,
Some(tableIdent), _, _) => tableIdent
}
assert(tables.size == 1)
assert(tables.head ===
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
index cd3e389d765d..d8d885c9b927 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
@@ -1523,7 +1523,7 @@ class SubquerySuite extends SharedSparkSession
// need to execute the query before we can examine fs.inputRDDs()
assert(stripAQEPlan(df.queryExecution.executedPlan) match {
case WholeStageCodegenExec(ColumnarToRowExec(InputAdapter(
- fs @ FileSourceScanExec(_, _, _, _, partitionFilters, _, _, _, _,
_)))) =>
+ fs @ FileSourceScanExec(_, _, _, _, partitionFilters, _, _, _, _,
_, _)))) =>
partitionFilters.exists(ExecSubqueryExpression.hasSubquery) &&
fs.inputRDDs().forall(
_.asInstanceOf[FileScanRDD].filePartitions.forall(
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/MarkSingleTaskExecutionSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/MarkSingleTaskExecutionSuite.scala
new file mode 100644
index 000000000000..d3f5ac779e54
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/MarkSingleTaskExecutionSuite.scala
@@ -0,0 +1,301 @@
+/*
+ * 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.datasources
+
+import org.apache.spark.sql.{Encoder, Encoders, QueryTest, Row}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.BoundReference
+import org.apache.spark.sql.catalyst.expressions.aggregate.V2Aggregator
+import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation,
LogicalPlan, Sort}
+import org.apache.spark.sql.catalyst.plans.physical.SinglePartition
+import org.apache.spark.sql.catalyst.trees.TreePattern.USER_DEFINED_AGGREGATION
+import org.apache.spark.sql.connector.catalog.functions.{AggregateFunction =>
V2AggregateFunction}
+import org.apache.spark.sql.execution.{FileSourceScanExec, LocalTableScanExec,
SparkPlan}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
AdaptiveSparkPlanHelper}
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
+import org.apache.spark.sql.expressions.Aggregator
+import org.apache.spark.sql.functions.{count, sum, udaf}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{DataType, LongType}
+
+/**
+ * Test suite for the [[MarkSingleTaskExecution]] optimizer rule and its
physical effects. The rule
+ * marks small single-partition scans, optionally under a shuffle-inducing
operator, so that the
+ * scan reports a `SinglePartition` output partitioning and the following
shuffle can be elided.
+ */
+class MarkSingleTaskExecutionSuite extends QueryTest with SharedSparkSession
+ with AdaptiveSparkPlanHelper {
+
+ private val t = "single_task_t"
+ private val t2 = "single_task_t2"
+ private val emptyTable = "single_task_empty"
+
+ private def enabledConfs: Seq[(String, String)] = Seq(
+ SQLConf.SINGLE_TASK_EXECUTION_ENABLED.key -> "true",
+ // Force the optimization to also apply to zero-file / zero-byte scans so
that we can exercise
+ // the empty-scan corner case created by dynamic pruning.
+ SQLConf.SINGLE_TASK_EXECUTION_MIN_NUM_FILES.key -> "0",
+ SQLConf.SINGLE_TASK_EXECUTION_MIN_NUM_BYTES.key -> "0",
+ SQLConf.SINGLE_TASK_EXECUTION_LOCAL_TABLE_SCAN_MIN_ROWS.key -> "0")
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ // A single-file Parquet table with a small amount of data.
+ spark.range(0, 2).selectExpr("id as col", "cast(id as string) as col_str")
+ .repartition(1).write.mode("overwrite").saveAsTable(t)
+ spark.range(0, 3).selectExpr("id % 2 as col")
+ .repartition(1).write.mode("overwrite").saveAsTable(t2)
+ spark.range(0, 0).selectExpr("id as
col").write.mode("overwrite").saveAsTable(emptyTable)
+ }
+
+ override def afterAll(): Unit = {
+ try {
+ sql(s"drop table if exists $t")
+ sql(s"drop table if exists $t2")
+ sql(s"drop table if exists $emptyTable")
+ } finally {
+ super.afterAll()
+ }
+ }
+
+ private def getFinalPhysicalPlan(df: org.apache.spark.sql.DataFrame):
SparkPlan = {
+ df.queryExecution.executedPlan match {
+ case a: AdaptiveSparkPlanExec => a.finalPhysicalPlan
+ case other => other
+ }
+ }
+
+ private def hasShuffle(plan: SparkPlan): Boolean =
+ collectWithSubqueries(plan) { case s: ShuffleExchangeLike => s }.nonEmpty
+
+ private def isMarked(plan: LogicalPlan): Boolean = {
+ val marks = plan.collect {
+ case lr: LogicalRelation =>
lr.getTagValue(MarkSingleTaskExecution.markTag).getOrElse(false)
+ case lr: LocalRelation =>
lr.getTagValue(MarkSingleTaskExecution.markTag).getOrElse(false)
+ }
+ marks.nonEmpty && marks.forall(identity)
+ }
+
+ private def checkMarked(query: String): Unit = withSQLConf(enabledConfs: _*)
{
+ val plan = sql(query).queryExecution.optimizedPlan
+ assert(isMarked(plan), s"expected plan to be marked for single-task
execution:\n$plan")
+ }
+
+ private def checkNotMarked(query: String, confs: Seq[(String, String)] =
enabledConfs): Unit =
+ withSQLConf(confs: _*) {
+ val plan = sql(query).queryExecution.optimizedPlan
+ assert(!isMarked(plan), s"expected plan NOT to be marked:\n$plan")
+ }
+
+ private def checkSinglePartition(
+ query: String,
+ expected: Seq[Row],
+ confs: Seq[(String, String)] = enabledConfs): Unit = withSQLConf(confs:
_*) {
+ val df = sql(query)
+ QueryTest.checkAnswer(df, expected)
+ val plan = getFinalPhysicalPlan(df)
+ assert(!hasShuffle(plan), s"expected no shuffle in:\n$plan")
+ val scans = collect(plan) {
+ case s: FileSourceScanExec => s.outputPartitioning
+ case s: LocalTableScanExec => s.outputPartitioning
+ }
+ assert(scans.nonEmpty, s"expected a scan in:\n$plan")
+ assert(scans.forall(_ == SinglePartition),
+ s"expected all scans to report SinglePartition, got $scans in:\n$plan")
+ }
+
+ test("marks scan + sort") {
+ checkMarked(s"select col from $t order by col")
+ checkMarked(s"select col from (select col from $t where col = 0) order by
col")
+ }
+
+ test("marks scan + aggregation") {
+ checkMarked(s"select count(1) from $t group by col")
+ checkMarked(s"select sum(col) from (select col from $t where col < 42)")
+ }
+
+ test("marks scan + expand (grouping sets)") {
+ checkMarked(s"select col, count(1) from $t group by rollup(col)")
+ }
+
+ test("marks scan + window") {
+ checkMarked(
+ s"select col, row_number() over (partition by col order by col) from $t")
+ }
+
+ test("does not mark when the feature is disabled") {
+ checkNotMarked(
+ s"select col from $t order by col",
+ Seq(SQLConf.SINGLE_TASK_EXECUTION_ENABLED.key -> "false"))
+ }
+
+ test("does not mark when the per-operator flag is disabled") {
+ checkNotMarked(
+ s"select col from $t order by col",
+ enabledConfs :+ (SQLConf.SINGLE_TASK_EXECUTION_SORT.key -> "false"))
+ }
+
+ test("marking does not mutate the input plan's nodes") {
+ // The rule must mark a copy of each eligible node rather than tag the
node it was handed:
+ // an in-place tag on a shared node would propagate through
`TreeNode.clone`/`copyTagsFrom`
+ // and could leak the marking into unrelated queries. Build an eligible
plan shape directly
+ // (rather than through a query, whose optimizer would have already run
this rule) so the
+ // input is guaranteed unmarked, then run the rule on it.
+ withSQLConf(enabledConfs: _*) {
+ val relation = spark.table(t).queryExecution.analyzed.collectFirst {
+ case lr: LogicalRelation => lr
+ }.get
+ val input = Sort(Nil, global = true, relation)
+ val output = MarkSingleTaskExecution(input)
+ // The rule marks a copied relation ...
+ assert(isMarked(output), s"expected output to be marked:\n$output")
+ // ... and leaves the original relation node untagged (it was copied,
not mutated).
+ assert(relation.getTagValue(MarkSingleTaskExecution.markTag).isEmpty,
+ "rule must not tag the input relation node in place")
+ }
+ }
+
+ test("does not mark unsupported plan shapes (join)") {
+ // Join is not a supported operator in this port, so the presence of a
join makes the whole
+ // plan ineligible.
+ checkNotMarked(s"select a.col from $t a join $t b on a.col = b.col order
by a.col")
+ }
+
+ test("does not mark plans with subquery expressions") {
+ checkNotMarked(s"select col from $t where col = (select max(col) from $t2)
order by col")
+ }
+
+ test("does not mark plans with user-defined aggregations") {
+ val strLen = new Aggregator[String, Long, Long] {
+ override def zero: Long = 0L
+ override def reduce(b: Long, a: String): Long = b + a.length
+ override def merge(b1: Long, b2: Long): Long = b1 + b2
+ override def finish(reduction: Long): Long = reduction
+ override def bufferEncoder: Encoder[Long] = Encoders.scalaLong
+ override def outputEncoder: Encoder[Long] = Encoders.scalaLong
+ }
+ // `functions.udaf` produces a `ScalaAggregator` expression.
+ spark.udf.register("test_str_len_agg", udaf(strLen))
+ try {
+ checkNotMarked(s"select test_str_len_agg(col_str) from $t")
+ } finally {
+ spark.sessionState.catalog.dropTempFunction("test_str_len_agg",
ignoreIfNotExists = true)
+ }
+ // A typed Dataset aggregation produces a `TypedAggregateExpression`.
+ withSQLConf(enabledConfs: _*) {
+ import testImplicits._
+ val ds =
spark.table(t).select($"col_str").as[String].select(strLen.toColumn)
+ val optimized = ds.queryExecution.optimizedPlan
+ assert(!isMarked(optimized),
+ s"expected plan with typed aggregation NOT to be marked:\n$optimized")
+ }
+ }
+
+ test("V2Aggregator carries the USER_DEFINED_AGGREGATION pattern") {
+ // ScalaAggregator and the typed aggregate expressions are covered by the
end-to-end test
+ // above; V2Aggregator has no operator-level pattern of its own, so verify
it directly.
+ // HiveUDAFFunction lives in the hive module and is covered there.
+ val v2Func = new V2AggregateFunction[java.lang.Long, java.lang.Long] {
+ override def newAggregationState(): java.lang.Long = 0L
+ override def update(state: java.lang.Long, input: InternalRow):
java.lang.Long =
+ state + input.getLong(0)
+ override def merge(l: java.lang.Long, r: java.lang.Long): java.lang.Long
= l + r
+ override def produceResult(state: java.lang.Long): java.lang.Long = state
+ override def name(): String = "test_v2_sum"
+ override def inputTypes(): Array[DataType] = Array(LongType)
+ override def resultType(): DataType = LongType
+ }
+ val agg = V2Aggregator(v2Func, Seq(BoundReference(0, LongType, nullable =
false)))
+ assert(agg.containsPattern(USER_DEFINED_AGGREGATION))
+ }
+
+ test("output partitioning is SinglePartition, scan + sort") {
+ checkSinglePartition(s"select col from $t order by col", Seq(Row(0),
Row(1)))
+ }
+
+ test("output partitioning is SinglePartition, scan + aggregation with group
by") {
+ checkSinglePartition(
+ s"select count(1) as c from $t2 group by col",
+ Seq(Row(1), Row(2)))
+ }
+
+ test("output partitioning is SinglePartition, scan + aggregation without
group by") {
+ checkSinglePartition(s"select sum(col) from $t", Seq(Row(1)))
+ }
+
+ test("output partitioning is SinglePartition, scan + distinct") {
+ checkSinglePartition(s"select distinct col from $t2", Seq(Row(0), Row(1)))
+ }
+
+ test("output partitioning is SinglePartition, scan + expand") {
+ checkSinglePartition(
+ s"select col, count(1) as c from $t group by rollup(col)",
+ Seq(Row(0, 1), Row(1, 1), Row(null, 2)))
+ }
+
+ test("bucketed scan does not run in a single task") {
+ val bucketed = "single_task_bucketed"
+ withTable(bucketed) {
+ spark.range(0, 2).selectExpr("id as col").write.bucketBy(2,
"col").saveAsTable(bucketed)
+ // Raise the file count bound so that only the bucketing makes the scan
ineligible.
+ val confs = enabledConfs :+
(SQLConf.SINGLE_TASK_EXECUTION_MAX_NUM_FILES.key -> "4")
+ withSQLConf(confs: _*) {
+ val df = sql(s"select col, count(1) as c from $bucketed group by col")
+ checkAnswer(df, Seq(Row(0, 1), Row(1, 1)))
+ val scans = collect(getFinalPhysicalPlan(df)) { case s:
FileSourceScanExec => s }
+ assert(scans.nonEmpty)
+ assert(scans.forall(!_.useSingleTaskExecution),
+ "a bucketed scan must not run in a single task as that would
invalidate its " +
+ "HashPartitioning over the bucket columns")
+ }
+ }
+ }
+
+ test("empty table scan + aggregation is correct and single-partition") {
+ // Without single-task execution eliding the shuffle before the
aggregation, an empty scan
+ // could incorrectly return zero rows instead of a single NULL row for a
global aggregation.
+ checkSinglePartition(s"select sum(col) from $emptyTable", Seq(Row(null)))
+ }
+
+ test("in-memory local relation is scanned in a single partition") {
+ checkSinglePartition(
+ "select col from values (0), (1) as tab(col) order by col",
+ Seq(Row(0), Row(1)))
+ }
+
+ test("empty local relation + global aggregation returns one row") {
+ withSQLConf(enabledConfs: _*) {
+ import testImplicits._
+ val df = Seq.empty[Int].toDF("col").agg(count($"col"), sum($"col"))
+ assert(isMarked(df.queryExecution.optimizedPlan),
+ "expected the empty local relation to be marked for single-task
execution")
+ // A global aggregation over an empty input must still return a single
row.
+ checkAnswer(df, Row(0, null))
+ }
+ }
+
+ test("does not mark when a leaf-node parallelism override is set") {
+ checkNotMarked(
+ "select col from values (0), (1) as tab(col) order by col",
+ enabledConfs :+ (SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "4"))
+ checkNotMarked(
+ s"select col from $t order by col",
+ enabledConfs :+ (SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "4"))
+ }
+}
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
index bf708eecf0c0..129c5e2cc053 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.aggregate._
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext,
CodeGenerator, CodegenFallback, ExprCode}
import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper
+import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern,
USER_DEFINED_AGGREGATION}
import org.apache.spark.sql.hive.HiveShim._
import org.apache.spark.sql.types._
@@ -337,6 +338,8 @@ private[hive] case class HiveUDAFFunction(
with HiveInspectors
with UserDefinedExpression {
+ final override val nodePatterns: Seq[TreePattern] =
Seq(USER_DEFINED_AGGREGATION)
+
override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int):
ImperativeAggregate =
copy(mutableAggBufferOffset = newMutableAggBufferOffset)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
