[SPARK-48037][CORE] Fix SortShuffleWriter lacks shuffle write related metrics resulting in potentially inaccurate data

[cxzl25]

What changes were proposed in this pull request?

This PR aims to fix SortShuffleWriter lacks shuffle write related metrics resulting in potentially inaccurate data.

Why are the changes needed?

When the shuffle writer is SortShuffleWriter, it does not use SQLShuffleWriteMetricsReporter to update metrics, which causes AQE to obtain runtime statistics and the rowCount obtained is 0.

Some optimization rules rely on rowCount statistics, such as EliminateLimits. Because rowCount is 0, it removes the limit operator. At this time, we get data results without limit.

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala

Lines 168 to 172 in 59d5946

	override def runtimeStatistics: Statistics = {
	val dataSize = metrics("dataSize").value
	val rowCount = metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_RECORDS_WRITTEN).value
	Statistics(dataSize, Some(rowCount))
	}

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

Lines 2067 to 2070 in 59d5946

	object EliminateLimits extends Rule[LogicalPlan] {
	private def canEliminate(limitExpr: Expression, child: LogicalPlan): Boolean = {
	limitExpr.foldable && child.maxRows.exists { _ <= limitExpr.eval().asInstanceOf[Int] }
	}

Does this PR introduce any user-facing change?

Yes

How was this patch tested?

Production environment verification.

master metrics

PR metrics

Was this patch authored or co-authored using generative AI tooling?

No

[dongjoon-hyun]

Could you add a test case for this, @cxzl25 ?

[cxzl25]

Sometimes the SQL results are not always ordered, so a parameter has been added to support skip checking.

[dongjoon-hyun]

Could you make a separate PR for this, @cxzl25 ?

[cxzl25]

Sure.

[mridulm]

The changes to core looks fine to me, but I am not very sure of why we are skipping the tests and whether the test is actually testing what we expect.
I will let @dongjoon-hyun review that (along with rest of the PR) better.

[dongjoon-hyun]

Thank you for review, @mridulm .

[cxzl25]

why we are skipping the tests

Limits after group by are not guaranteed to be in order.

[info]   == Results ==
[info]   !== Correct Answer - 1 ==            == Spark Answer - 1 ==
[info]    struct<id:bigint,count(1):bigint>   struct<id:bigint,count(1):bigint>
[info]   ![1,1]                               [0,1] (QueryTest.scala:267)
[info]   org.scalatest.exceptions.TestFailedException:
[info]   at org.scalatest.Assertions.newAssertionFailedException(Assertions.scala:472)
[info]   at org.scalatest.Assertions.newAssertionFailedException$(Assertions.scala:471)
[info]   at org.apache.spark.sql.QueryTest$.newAssertionFailedException(QueryTest.scala:257)
[info]   at org.scalatest.Assertions.fail(Assertions.scala:933)
[info]   at org.scalatest.Assertions.fail$(Assertions.scala:929)
[info]   at org.apache.spark.sql.QueryTest$.fail(QueryTest.scala:257)
[info]   at org.apache.spark.sql.QueryTest$.checkAnswer(QueryTest.scala:267)
[info]   at org.apache.spark.sql.QueryTest.checkAnswer(QueryTest.scala:153)
[info]   at org.apache.spark.sql.execution.adaptive.AdaptiveQueryExecSuite.$anonfun$runAdaptiveAndVerifyResult$1(AdaptiveQueryExecSuite.scala:91)

whether the test is actually testing what we expect

In this added UT check if the final execution plan of AQE contains the limit operator.

./bin/spark-sql --conf spark.driver.memory=6g

set spark.sql.shuffle.partitions=16777217;
create table foo as select id from range(2);
select id, count(*) from foo group by id limit 1;

== Physical Plan ==
AdaptiveSparkPlan (7)
+- == Final Plan ==
   LocalTableScan (1)
+- == Initial Plan ==
   CollectLimit (6)
   +- HashAggregate (5)
      +- Exchange (4)
         +- HashAggregate (3)
            +- Scan hive spark_catalog.default.foo (2)

[dongjoon-hyun]

I merged #46316 . Could you rebase this PR to the master branch, @cxzl25 ?

[dongjoon-hyun]

Also, cc @viirya , too.

[viirya]

Hmm, this only verifies if there is specific operator (i.e., Limit) in the query plan, how is it related to the metrics you want to fix?

[viirya]

Oh, I see. It is due to AQE usage of runtime metrics.

[viirya]

Hmm, the metrics given at getWriter looks like also coming from context.taskMetrics().shuffleWriteMetrics, isn't it looking the same?

[cxzl25]

It may be SQLShuffleWriteMetricsReporter, which may not be the same as context.taskMetrics().shuffleWriteMetrics.

spark/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala

Lines 51 to 56 in 9e12bd7

	writer = manager.getWriter[Any, Any](
	dep.shuffleHandle,
	mapId,
	context,
	createMetricsReporter(context))
	writer.write(inputs.asInstanceOf[Iterator[_ <: Product2[Any, Any]]])

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala

Lines 449 to 455 in 9e12bd7

	def createShuffleWriteProcessor(metrics: Map[String, SQLMetric]): ShuffleWriteProcessor = {
	new ShuffleWriteProcessor {
	override protected def createMetricsReporter(
	context: TaskContext): ShuffleWriteMetricsReporter = {
	new SQLShuffleWriteMetricsReporter(context.taskMetrics().shuffleWriteMetrics, metrics)
	}
	}

[mridulm]

Looks like this missed the 3.4 release by a month @dongjoon-hyun ... might have been a nice addition to it !

[dongjoon-hyun]

Looks like this missed the 3.4 release by a month @dongjoon-hyun ... might have been a nice addition to it !

Ya, I agree and this is still a good addition to 4.0.0-preview.

Cc @cloud-fan as the release manager of 4.0.0-preview

[cloud-fan]

do we have to skip checking the answer? I think the query can be LIMIT 10 so that the result is deterministic.

[cxzl25]

Because AQE is disabled when checking the results, it will have many partitions, and the order is not guaranteed at this time. When AQE is enabled, it will be merged into one partition.

set spark.sql.adaptive.enabled=false;
set spark.sql.shuffle.partitions=1000;
create table foo as select id from range(2);
select id, count(*) from foo group by id limit 1;

output

1	1

[cloud-fan]

AFAIK the checkAnswer util will sort the data before comparison.

[cxzl25]

It does sort, but the final result is sorted locally. Because the result of the limit is uncertain, the local sort has no effect.

spark/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala

Lines 282 to 293 in 08c6bb9

	def getErrorMessageInCheckAnswer(
	df: DataFrame,
	expectedAnswer: Seq[Row],
	checkToRDD: Boolean = true): Option[String] = {
	val isSorted = df.logicalPlan.collect { case s: logical.Sort => s }.nonEmpty
	if (checkToRDD) {
	SQLExecution.withSQLConfPropagated(df.sparkSession) {
	df.rdd.count() // Also attempt to deserialize as an RDD [SPARK-15791]
	}
	}
	val sparkAnswer = try df.collect().toSeq catch {

[cloud-fan]

we can make it certain if the limit is larger than the number of result rows?

[cxzl25]

In this test case, because spark.sql.adaptive.enabled=false, partitions will not be merged.
It has a large number of partitions and tasks, so it requires a large amount of driver memory to execute successfully.

[dongjoon-hyun]

+1, LGTM.

[dongjoon-hyun]

Merged to master for Apache Spark 4.0.0-preview.

Could you make backporting PRs to the release branches, @cxzl25 ?