You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2024/08/28 14:20:33 UTC
(datafusion-comet) branch main updated: chore: Improve CometExchange metrics (#873)
This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new f7f0bb1e chore: Improve CometExchange metrics (#873)
f7f0bb1e is described below
commit f7f0bb1ed68367b8d3e1c88010c1f943f480ea11
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Wed Aug 28 07:20:29 2024 -0700
chore: Improve CometExchange metrics (#873)
* init
* chore: Improve CometExchange metrics
* fix format
* Add data size
---
.../src/execution/datafusion/shuffle_writer.rs | 59 +++++++++++++++++++---
.../spark/sql/comet/CometCollectLimitExec.scala | 3 --
.../sql/comet/CometTakeOrderedAndProjectExec.scala | 6 +--
.../apache/spark/sql/comet/CometWindowExec.scala | 13 +----
.../shuffle/CometShuffleExchangeExec.scala | 15 +++---
.../comet/exec/CometNativeShuffleSuite.scala | 9 ++++
6 files changed, 69 insertions(+), 36 deletions(-)
diff --git a/native/core/src/execution/datafusion/shuffle_writer.rs b/native/core/src/execution/datafusion/shuffle_writer.rs
index 751484db..9668359f 100644
--- a/native/core/src/execution/datafusion/shuffle_writer.rs
+++ b/native/core/src/execution/datafusion/shuffle_writer.rs
@@ -32,6 +32,7 @@ use arrow::{datatypes::*, ipc::writer::StreamWriter};
use async_trait::async_trait;
use bytes::Buf;
use crc32fast::Hasher;
+use datafusion::physical_plan::metrics::Time;
use datafusion::{
arrow::{
array::*,
@@ -239,14 +240,19 @@ impl PartitionBuffer {
}
/// Appends all rows of given batch into active array builders.
- fn append_batch(&mut self, batch: &RecordBatch) -> Result<isize> {
+ fn append_batch(&mut self, batch: &RecordBatch, time_metric: &Time) -> Result<isize> {
let columns = batch.columns();
let indices = (0..batch.num_rows()).collect::<Vec<usize>>();
- self.append_rows(columns, &indices)
+ self.append_rows(columns, &indices, time_metric)
}
/// Appends rows of specified indices from columns into active array builders.
- fn append_rows(&mut self, columns: &[ArrayRef], indices: &[usize]) -> Result<isize> {
+ fn append_rows(
+ &mut self,
+ columns: &[ArrayRef],
+ indices: &[usize],
+ time_metric: &Time,
+ ) -> Result<isize> {
let mut mem_diff = 0;
let mut start = 0;
@@ -263,7 +269,10 @@ impl PartitionBuffer {
});
self.num_active_rows += end - start;
if self.num_active_rows >= self.batch_size {
+ let mut timer = time_metric.timer();
mem_diff += self.flush()?;
+ timer.stop();
+
mem_diff += self.init_active_if_necessary()?;
}
start = end;
@@ -592,6 +601,9 @@ struct ShuffleRepartitionerMetrics {
/// total spilled bytes during the execution of the operator
spilled_bytes: Count,
+
+ /// The original size of spilled data. Different to `spilled_bytes` because of compression.
+ data_size: Count,
}
impl ShuffleRepartitionerMetrics {
@@ -600,6 +612,7 @@ impl ShuffleRepartitionerMetrics {
baseline: BaselineMetrics::new(metrics, partition),
spill_count: MetricBuilder::new(metrics).spill_count(partition),
spilled_bytes: MetricBuilder::new(metrics).spilled_bytes(partition),
+ data_size: MetricBuilder::new(metrics).counter("data_size", partition),
}
}
}
@@ -683,7 +696,10 @@ impl ShuffleRepartitioner {
));
}
- let _timer = self.metrics.baseline.elapsed_compute().timer();
+ // Update data size metric
+ self.metrics.data_size.add(input.get_array_memory_size());
+
+ let time_metric = self.metrics.baseline.elapsed_compute();
// NOTE: in shuffle writer exec, the output_rows metrics represents the
// number of rows those are written to output data file.
@@ -754,8 +770,11 @@ impl ShuffleRepartitioner {
// If the range of indices is not big enough, just appending the rows into
// active array builders instead of directly adding them as a record batch.
- mem_diff +=
- output.append_rows(input.columns(), &shuffled_partition_ids[start..end])?;
+ mem_diff += output.append_rows(
+ input.columns(),
+ &shuffled_partition_ids[start..end],
+ time_metric,
+ )?;
}
if mem_diff > 0 {
@@ -782,7 +801,7 @@ impl ShuffleRepartitioner {
);
let output = &mut buffered_partitions[0];
- output.append_batch(&input)?;
+ output.append_batch(&input, time_metric)?;
}
other => {
// this should be unreachable as long as the validation logic
@@ -798,7 +817,6 @@ impl ShuffleRepartitioner {
/// Writes buffered shuffled record batches into Arrow IPC bytes.
async fn shuffle_write(&mut self) -> Result<SendableRecordBatchStream> {
- let _timer = self.metrics.baseline.elapsed_compute().timer();
let num_output_partitions = self.num_output_partitions;
let mut buffered_partitions = self.buffered_partitions.lock().await;
let mut output_batches: Vec<Vec<u8>> = vec![vec![]; num_output_partitions];
@@ -823,14 +841,21 @@ impl ShuffleRepartitioner {
.map_err(|e| DataFusionError::Execution(format!("shuffle write error: {:?}", e)))?;
for i in 0..num_output_partitions {
+ let mut timer = self.metrics.baseline.elapsed_compute().timer();
+
offsets[i] = output_data.stream_position()?;
output_data.write_all(&output_batches[i])?;
+
+ timer.stop();
+
output_batches[i].clear();
// append partition in each spills
for spill in &output_spills {
let length = spill.offsets[i + 1] - spill.offsets[i];
if length > 0 {
+ let mut timer = self.metrics.baseline.elapsed_compute().timer();
+
let mut spill_file =
BufReader::new(File::open(spill.file.path()).map_err(|e| {
DataFusionError::Execution(format!("shuffle write error: {:?}", e))
@@ -839,15 +864,22 @@ impl ShuffleRepartitioner {
std::io::copy(&mut spill_file.take(length), &mut output_data).map_err(|e| {
DataFusionError::Execution(format!("shuffle write error: {:?}", e))
})?;
+
+ timer.stop();
}
}
}
+ let mut timer = self.metrics.baseline.elapsed_compute().timer();
output_data.flush()?;
+ timer.stop();
// add one extra offset at last to ease partition length computation
offsets[num_output_partitions] = output_data
.stream_position()
.map_err(|e| DataFusionError::Execution(format!("shuffle write error: {:?}", e)))?;
+
+ let mut timer = self.metrics.baseline.elapsed_compute().timer();
+
let mut output_index =
BufWriter::new(File::create(index_file).map_err(|e| {
DataFusionError::Execution(format!("shuffle write error: {:?}", e))
@@ -859,6 +891,8 @@ impl ShuffleRepartitioner {
}
output_index.flush()?;
+ timer.stop();
+
let used = self.reservation.size();
self.reservation.shrink(used);
@@ -878,6 +912,10 @@ impl ShuffleRepartitioner {
self.metrics.spill_count.value()
}
+ fn data_size(&self) -> usize {
+ self.metrics.data_size.value()
+ }
+
async fn spill(&self) -> Result<usize> {
log::debug!(
"ShuffleRepartitioner spilling shuffle data of {} to disk while inserting ({} time(s) so far)",
@@ -891,6 +929,8 @@ impl ShuffleRepartitioner {
return Ok(0);
}
+ let mut timer = self.metrics.baseline.elapsed_compute().timer();
+
let spillfile = self
.runtime
.disk_manager
@@ -902,6 +942,8 @@ impl ShuffleRepartitioner {
)
.await?;
+ timer.stop();
+
let mut spills = self.spills.lock().await;
let used = self.reservation.size();
self.metrics.spill_count.add(1);
@@ -955,6 +997,7 @@ impl Debug for ShuffleRepartitioner {
.field("memory_used", &self.used())
.field("spilled_bytes", &self.spilled_bytes())
.field("spilled_count", &self.spill_count())
+ .field("data_size", &self.data_size())
.finish()
}
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala
index 3effcca4..8ea0b176 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala
@@ -25,7 +25,6 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, SinglePartition}
import org.apache.spark.sql.comet.execution.shuffle.{CometShuffledBatchRDD, CometShuffleExchangeExec}
-import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec.{METRIC_NATIVE_TIME_DESCRIPTION, METRIC_NATIVE_TIME_NAME}
import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode, UnsafeRowSerializer}
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter}
import org.apache.spark.sql.vectorized.ColumnarBatch
@@ -56,8 +55,6 @@ case class CometCollectLimitExec(
SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext)
override lazy val metrics: Map[String, SQLMetric] = Map(
"dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
- METRIC_NATIVE_TIME_NAME ->
- SQLMetrics.createNanoTimingMetric(sparkContext, METRIC_NATIVE_TIME_DESCRIPTION),
"numPartitions" -> SQLMetrics.createMetric(
sparkContext,
"number of partitions")) ++ readMetrics ++ writeMetrics
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
index f808a2b1..a8e2a98d 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
@@ -25,10 +25,8 @@ import org.apache.spark.serializer.Serializer
import org.apache.spark.sql.catalyst.expressions.{Attribute, NamedExpression, SortOrder}
import org.apache.spark.sql.catalyst.util.truncatedString
import org.apache.spark.sql.comet.execution.shuffle.{CometShuffledBatchRDD, CometShuffleExchangeExec}
-import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec.{METRIC_NATIVE_TIME_DESCRIPTION, METRIC_NATIVE_TIME_NAME}
import org.apache.spark.sql.execution.{SparkPlan, TakeOrderedAndProjectExec, UnaryExecNode, UnsafeRowSerializer}
-import org.apache.spark.sql.execution.metric.{SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter}
-import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter}
import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.comet.serde.QueryPlanSerde.exprToProto
@@ -57,8 +55,6 @@ case class CometTakeOrderedAndProjectExec(
SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext)
override lazy val metrics: Map[String, SQLMetric] = Map(
"dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
- METRIC_NATIVE_TIME_NAME ->
- SQLMetrics.createNanoTimingMetric(sparkContext, METRIC_NATIVE_TIME_DESCRIPTION),
"numPartitions" -> SQLMetrics.createMetric(
sparkContext,
"number of partitions")) ++ readMetrics ++ writeMetrics
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala
index bb187170..93b51617 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala
@@ -21,9 +21,8 @@ package org.apache.spark.sql.comet
import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, NamedExpression, SortOrder}
import org.apache.spark.sql.catalyst.plans.physical.Partitioning
-import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec.{METRIC_NATIVE_TIME_DESCRIPTION, METRIC_NATIVE_TIME_NAME}
import org.apache.spark.sql.execution.SparkPlan
-import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
import com.google.common.base.Objects
@@ -49,17 +48,9 @@ case class CometWindowExec(
override def nodeName: String = "CometWindowExec"
- private lazy val writeMetrics =
- SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext)
- private lazy val readMetrics =
- SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext)
override lazy val metrics: Map[String, SQLMetric] = Map(
"dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
- METRIC_NATIVE_TIME_NAME ->
- SQLMetrics.createNanoTimingMetric(sparkContext, METRIC_NATIVE_TIME_DESCRIPTION),
- "numPartitions" -> SQLMetrics.createMetric(
- sparkContext,
- "number of partitions")) ++ readMetrics ++ writeMetrics
+ "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of partitions"))
override def outputOrdering: Seq[SortOrder] = child.outputOrdering
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index cec64d7a..028ba243 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -39,7 +39,6 @@ import org.apache.spark.sql.catalyst.expressions.codegen.LazilyGeneratedOrdering
import org.apache.spark.sql.catalyst.plans.logical.Statistics
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.comet.{CometExec, CometMetricNode, CometPlan}
-import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec.{METRIC_NATIVE_TIME_DESCRIPTION, METRIC_NATIVE_TIME_NAME}
import org.apache.spark.sql.comet.shims.ShimCometShuffleWriteProcessor
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeLike, ShuffleOrigin}
@@ -78,8 +77,6 @@ case class CometShuffleExchangeExec(
SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext)
override lazy val metrics: Map[String, SQLMetric] = Map(
"dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
- METRIC_NATIVE_TIME_NAME ->
- SQLMetrics.createNanoTimingMetric(sparkContext, METRIC_NATIVE_TIME_DESCRIPTION),
"numPartitions" -> SQLMetrics.createMetric(
sparkContext,
"number of partitions")) ++ readMetrics ++ writeMetrics
@@ -221,9 +218,6 @@ case class CometShuffleExchangeExec(
}
object CometShuffleExchangeExec extends ShimCometShuffleExchangeExec {
- val METRIC_NATIVE_TIME_NAME = "shuffleNativeTotalTime"
- val METRIC_NATIVE_TIME_DESCRIPTION = "shuffle native code time"
-
def prepareShuffleDependency(
rdd: RDD[ColumnarBatch],
outputAttributes: Seq[Attribute],
@@ -468,6 +462,7 @@ class CometShuffleWriteProcessor(
mapId: Long,
mapIndex: Int,
context: TaskContext): MapStatus = {
+ val metricsReporter = createMetricsReporter(context)
val shuffleBlockResolver =
SparkEnv.get.shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver]
val dataFile = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
@@ -483,11 +478,13 @@ class CometShuffleWriteProcessor(
// Maps native metrics to SQL metrics
val nativeSQLMetrics = Map(
"output_rows" -> metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_RECORDS_WRITTEN),
- "elapsed_compute" -> metrics("shuffleNativeTotalTime"))
+ "data_size" -> metrics("dataSize"),
+ "elapsed_compute" -> metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_WRITE_TIME))
val nativeMetrics = CometMetricNode(nativeSQLMetrics)
// Getting rid of the fake partitionId
val newInputs = inputs.asInstanceOf[Iterator[_ <: Product2[Any, Any]]].map(_._2)
+
val cometIter = CometExec.getCometIterator(
Seq(newInputs.asInstanceOf[Iterator[ColumnarBatch]]),
nativePlan,
@@ -512,8 +509,8 @@ class CometShuffleWriteProcessor(
})
.toArray
- // Update Spark metrics from native metrics
- metrics("dataSize") += Files.size(tempDataFilePath)
+ // Total written bytes at native
+ metricsReporter.incBytesWritten(Files.size(tempDataFilePath))
// commit
shuffleBlockResolver.writeMetadataFileAndCommit(
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
index aec122be..329688fc 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
@@ -169,6 +169,15 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
assert(metrics.contains("shuffleRecordsWritten"))
assert(metrics("shuffleRecordsWritten").value == 5L)
+
+ assert(metrics.contains("shuffleBytesWritten"))
+ assert(metrics("shuffleBytesWritten").value > 0)
+
+ assert(metrics.contains("dataSize"))
+ assert(metrics("dataSize").value > 0L)
+
+ assert(metrics.contains("shuffleWriteTime"))
+ assert(metrics("shuffleWriteTime").value > 0L)
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org