dongjoon-hyun commented on code in PR #57355:
URL: https://github.com/apache/spark/pull/57355#discussion_r4081857954
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala:
##########
@@ -360,6 +370,7 @@ object ShuffleExchangeExec {
newPartitioning: Partitioning,
serializer: Serializer,
writeMetrics: Map[String, SQLMetric],
+ customWriteMetrics: Map[String, SQLMetric],
Review Comment:
This is still binary-incompatible. In `branch-4.3` (4.3.0 is released) and
`master`, the signature is `(rdd, outputAttributes, newPartitioning,
serializer, writeMetrics: Map, pipelined: Boolean = false)`. Since Scala
compiles an existing 5-arg call site into a call to the 6-arg JVM method plus
`prepareShuffleDependency$default$6`, downstream jars compiled against 4.3 will
hit `NoSuchMethodError` because `(..., Map, boolean)` no longer exists. The new
5-arg overload does not help for this. In addition, a caller passing
`pipelined` positionally, `(..., writeMetrics, true)`, fails at source level
too.
Could you keep the existing `(..., writeMetrics: Map[String, SQLMetric],
pipelined: Boolean = false)` signature as-is (delegating to the new method),
and make the new method with `customWriteMetrics` have no default arguments?
##########
core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputMetricsReporter.java:
##########
@@ -0,0 +1,38 @@
+/*
+ * 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.shuffle.api;
+
+import org.apache.spark.annotation.Private;
+import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric;
+
+/**
+ * :: Private ::
+ * A map output writer that reports custom shuffle metrics collected while
writing a map task.
+ *
+ * @since 4.3.0
Review Comment:
`branch-4.3` is already `4.3.1-SNAPSHOT`. This should be `4.4.0`
(`branch-4.x`). The same applies to `CustomShuffleMetric` and
`CustomShuffleTaskMetric`. Also, the newly added
`ShuffleDriverComponents.supportedCustomMetrics()` has no `@since` tag.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala:
##########
@@ -0,0 +1,91 @@
+/*
+ * 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.metric
+
+import org.apache.spark.SparkContext
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.METRIC_NAME
+import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric,
CustomShuffleTaskMetric}
+
+object CustomShuffleMetrics extends Logging {
+
+ private[spark] val SHUFFLE_CUSTOM = "shuffleCustom"
+
+ /**
+ * Given a custom shuffle metric, builds and returns a metric type for a
`CustomShuffleMetric`
+ * class.
+ */
+ def buildMetricTypeName(metric: CustomShuffleMetric): String = {
+ s"${SHUFFLE_CUSTOM}_${metric.getClass.getName}"
+ }
+
+ /**
+ * Given a custom shuffle metric type name, this method parses it and
returns the corresponding
+ * `CustomShuffleMetric` class name.
+ */
+ def parseMetricType(metricType: String): Option[String] = {
+ if (metricType.startsWith(s"${SHUFFLE_CUSTOM}_")) {
+ Some(metricType.drop(SHUFFLE_CUSTOM.length + 1))
+ } else {
+ None
+ }
+ }
+
+ /**
+ * Creates collision-filtered custom shuffle metric [[SQLMetric]]s for an
operator that already
+ * exposes the given Spark-owned metric names. Plugin metrics whose names
collide with a
+ * Spark-owned name are dropped so they can't shadow the real accumulators
(e.g.
+ * `shuffleRecordsWritten`, which feeds AQE). Operators pass the reserved
names they own; the
+ * declared metrics come from the shuffle driver components.
+ */
+ def createFilteredMetrics(
+ sc: SparkContext,
+ reservedNames: Set[String]): Map[String, SQLMetric] = {
+ val (reserved, allowed) =
sc.shuffleDriverComponents.supportedCustomMetrics()
Review Comment:
This plugin call is not guarded. Since it runs while evaluating the
operator's lazy `metrics`, if a plugin throws here (or returns `null`, which
causes an NPE at `.partition`), every query containing a shuffle fails. Could
we catch `NonFatal`, log a warning, and return `Map.empty`?
In addition, duplicated metric names are silently collapsed by `.toMap` in
`createMetrics`. It would be good to log a warning for that case too.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala:
##########
@@ -0,0 +1,89 @@
+/*
+ * 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.metric
+
+import org.apache.spark.SharedSparkContext
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric,
CustomShuffleTaskMetric}
+
+class CustomShuffleMetricsSuite extends SparkFunSuite with SharedSparkContext {
+
+ test("buildMetricTypeName and parseMetricType round-trip the declaring class
name") {
+ val metric = shuffleMetric("s3BytesUploaded", "s3 bytes uploaded")
+ val tag = CustomShuffleMetrics.buildMetricTypeName(metric)
+ assert(tag ===
s"${CustomShuffleMetrics.SHUFFLE_CUSTOM}_${metric.getClass.getName}")
+
assert(CustomShuffleMetrics.parseMetricType(tag).contains(metric.getClass.getName))
+ // A built-in SQLMetric type is not a custom shuffle metric.
+ assert(CustomShuffleMetrics.parseMetricType(MetricUtilsSumTag).isEmpty)
+ }
+
+ test("createMetrics tags each SQLMetric so the UI delegates to the plugin's
aggregation") {
+ val metrics = CustomShuffleMetrics.createMetrics(sc, Array(
+ shuffleMetric("s3BytesUploaded", "s3 bytes uploaded"),
+ shuffleMetric("s3BlockUploads", "s3 block uploads")))
+
+ assert(metrics.keySet === Set("s3BytesUploaded", "s3BlockUploads"))
+ metrics.values.foreach { m =>
+ assert(CustomShuffleMetrics.parseMetricType(m.metricType).isDefined)
+ }
+ }
+
+ test("updateMetrics folds reported values into the matching SQLMetrics by
name") {
+ val metrics = CustomShuffleMetrics.createMetrics(sc, Array(
+ shuffleMetric("s3BytesUploaded", "s3 bytes uploaded"),
+ shuffleMetric("s3BlockUploads", "s3 block uploads")))
+
+ CustomShuffleMetrics.updateMetrics(
+ Array(taskMetric("s3BytesUploaded", 1024L), taskMetric("s3BlockUploads",
3L)), metrics)
+
+ assert(metrics("s3BytesUploaded").value === 1024L)
+ assert(metrics("s3BlockUploads").value === 3L)
+ }
+
+ test("updateMetrics ignores reported values with no matching declaration") {
+ val metrics = CustomShuffleMetrics.createMetrics(sc, Array(
+ shuffleMetric("s3BytesUploaded", "s3 bytes uploaded")))
+
+ CustomShuffleMetrics.updateMetrics(
+ Array(taskMetric("s3BytesUploaded", 512L), taskMetric("undeclared",
99L)), metrics)
+
+ assert(metrics("s3BytesUploaded").value === 512L)
+ assert(!metrics.contains("undeclared"))
+ }
+
+ private val MetricUtilsSumTag = "sum"
+
+ private def shuffleMetric(name: String, description: String):
CustomShuffleMetric =
+ new TestShuffleMetric(name, description)
+
+ private def taskMetric(name: String, value: Long): CustomShuffleTaskMetric =
{
+ val metricName = name
+ val metricValue = value
+ new CustomShuffleTaskMetric {
+ override def name(): String = metricName
+ override def value(): Long = metricValue
+ }
+ }
+}
+
+private class TestShuffleMetric(metricName: String, metricDescription: String)
Review Comment:
`CustomShuffleMetric` javadoc requires a 0-arg constructor because
`SQLAppStatusListener` re-instantiates the class via reflection, but this test
metric has a 2-arg constructor, so it would render as `N/A` in the real UI
path. It would be better to follow the contract in tests.
Since the aggregator is also cached by class name, a single generic class
cannot serve multiple metrics; effectively one class per metric is required. It
would be helpful to mention this explicitly in the javadoc.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala:
##########
@@ -38,14 +38,34 @@ trait LimitExec extends UnaryExecNode {
def limit: Int
}
+/**
+ * Shuffle read/write metric plumbing for operators that shuffle via
+ * [[ShuffleExchangeExec.prepareShuffleDependency]]. Publishes the built-in
read/write metrics plus
+ * collision-filtered plugin metrics (`customWriteMetrics`) through `metrics`.
Operators with
+ * additional Spark-owned metrics reserve those names against plugin
collisions via
+ * [[extraReservedMetrics]].
+ */
+trait ShuffleMetricsSupport { self: SparkPlan =>
Review Comment:
Since this trait is shared by `ShuffleExchangeExec` (in the `exchange`
package) and the limit operators, `limit.scala` looks like an unnatural place.
Could you move it into a separate file?
--
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]
