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 f2d11a668b16 [SPARK-57313][SQL] Fix SampleExec numOutputRows metric
when whole-stage codegen is disabled
f2d11a668b16 is described below
commit f2d11a668b165a396f449efd6c64713989e66b75
Author: Eric Yang <[email protected]>
AuthorDate: Wed Jun 10 17:44:00 2026 -0700
[SPARK-57313][SQL] Fix SampleExec numOutputRows metric when whole-stage
codegen is disabled
### What changes were proposed in this pull request?
Update the metric `number of output rows` for non-WSCG code path
### Why are the changes needed?
`number of output rows` metric is not updated/incremented when WSCG is off
- it can be useful to show the numOutputRows and be consistent with the WSCG
path.
<img width="370" height="339" alt="image"
src="https://github.com/user-attachments/assets/3531d904-16c6-456c-a9a0-6b667b2449f8";
/>
### Does this PR introduce _any_ user-facing change?
Yes. `number of output rows` metric gets correct value now when WSCG is off
### How was this patch tested?
New test case added.
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #56363 from jiwen624/SPARK-57313.
Authored-by: Eric Yang <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../spark/sql/execution/basicPhysicalOperators.scala | 9 ++++++++-
.../spark/sql/execution/metric/SQLMetricsSuite.scala | 17 +++++++++++++++++
2 files changed, 25 insertions(+), 1 deletion(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
index 81c048022c79..88c74ab7adc4 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
@@ -492,7 +492,8 @@ case class SampleExec(
"numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output
rows"))
protected override def doExecute(): RDD[InternalRow] = {
- if (withReplacement) {
+ val numOutputRows = longMetric("numOutputRows")
+ val sampled = if (withReplacement) {
// Disable gap sampling since the gap sampling method buffers two rows
internally,
// requiring us to copy the row, which is more expensive than the random
number generator.
new PartitionwiseSampledRDD[InternalRow, InternalRow](
@@ -503,6 +504,12 @@ case class SampleExec(
} else {
child.execute().randomSampleWithRange(lowerBound, upperBound,
resolvedSeed)
}
+ sampled.mapPartitionsInternal { iter =>
+ iter.map { row =>
+ numOutputRows += 1
+ row
+ }
+ }
}
// Mark this as empty. This plan doesn't need to evaluate any inputs and can
defer the evaluation
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index 6f06a7348080..8d6f0a5da5e1 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -125,6 +125,23 @@ class SQLMetricsSuite extends SharedSparkSession with
SQLMetricsTestUtils
}
}
+ test("SPARK-57313: Sample numOutputRows metric") {
+ Seq(false, true).foreach { withReplacement =>
+ Seq("false", "true").foreach { enableWholeStage =>
+ withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key ->
enableWholeStage) {
+ val df = spark.range(0, 1000, 1, 1)
+ .sample(withReplacement = withReplacement, fraction = 0.5, seed =
1)
+ val expectedRows = df.collect().length
+ sparkContext.listenerBus.waitUntilEmpty()
+ val sample = df.queryExecution.executedPlan.collect {
+ case s: SampleExec => s
+ }
+ assert(sample.size == 1)
+ assert(sample.head.metrics("numOutputRows").value == expectedRows)
+ }
+ }
+ }
+ }
test("Recursive CTEs metrics") {
withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> "") {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
