Review comments

Author: szehon-ho

sql/catalyst/src/main/java/org/apache/spark/sql/connector/metric/MergeMetrics.java

@@ -16,148 +16,47 @@
  */
 package org.apache.spark.sql.connector.metric;
 
-public interface MergeMetrics {
-
-  class Builder {
-    private long numTargetRowsCopied = -1;
-    private long numTargetRowsInserted = -1;
-    private long numTargetRowsDeleted = -1;
-    private long numTargetRowsUpdated = -1;
-    private long numTargetRowsMatchedUpdated = -1;
-    private long numTargetRowsMatchedDeleted = -1;
-    private long numTargetRowsNotMatchedBySourceUpdated = -1;
-    private long numTargetRowsNotMatchedBySourceDeleted = -1;
-    private long numSourceRows = -1;
-
-    public Builder numTargetRowsCopied(long numTargetRowsCopied) {
-      this.numTargetRowsCopied = numTargetRowsCopied;
-      return this;
-    }
-
-    public Builder numTargetRowsInserted(long numTargetRowsInserted) {
-      this.numTargetRowsInserted = numTargetRowsInserted;
-      return this;
-    }
-
-    public Builder numTargetRowsDeleted(long numTargetRowsDeleted) {
-      this.numTargetRowsDeleted = numTargetRowsDeleted;
-      return this;
-    }
-
-    public Builder numTargetRowsUpdated(long numTargetRowsUpdated) {
-      this.numTargetRowsUpdated = numTargetRowsUpdated;
-      return this;
-    }
-
-    public Builder numTargetRowsMatchedUpdated(long numTargetRowsMatchedUpdated) {
-      this.numTargetRowsMatchedUpdated = numTargetRowsMatchedUpdated;
-      return this;
-    }
-
-    public Builder numTargetRowsMatchedDeleted(long numTargetRowsMatchedDeleted) {
-      this.numTargetRowsMatchedDeleted = numTargetRowsMatchedDeleted;
-      return this;
-    }
-
-    public Builder numTargetRowsNotMatchedBySourceUpdated(long numTargetRowsNotMatchedBySourceUpdated) {
-      this.numTargetRowsNotMatchedBySourceUpdated = numTargetRowsNotMatchedBySourceUpdated;
-      return this;
-    }
-
-    public Builder numTargetRowsNotMatchedBySourceDeleted(long numTargetRowsNotMatchedBySourceDeleted) {
-      this.numTargetRowsNotMatchedBySourceDeleted = numTargetRowsNotMatchedBySourceDeleted;
-      return this;
-    }
-
-    public MergeMetrics build() {
-      return new MergeMetrics() {
-        @Override
-        public long numTargetRowsCopied() {
-          return numTargetRowsCopied;
-        }
+import java.util.OptionalLong;
 
-        @Override
-        public long numTargetRowsInserted() {
-          return numTargetRowsInserted;
-        }
-
-        @Override
-        public long numTargetRowsDeleted() {
-          return numTargetRowsDeleted;
-        }
-
-        @Override
-        public long numTargetRowsUpdated() {
-          return numTargetRowsUpdated;
-        }
-
-        @Override
-        public long numTargetRowsMatchedUpdated() {
-          return numTargetRowsMatchedUpdated;
-        }
-
-        @Override
-        public long numTargetRowsMatchedDeleted() {
-          return numTargetRowsMatchedDeleted;
-        }
-
-        @Override
-        public long numTargetRowsNotMatchedBySourceUpdated() {
-          return numTargetRowsNotMatchedBySourceUpdated;
-        }
-
-        @Override
-        public long numTargetRowsNotMatchedBySourceDeleted() {
-          return numTargetRowsNotMatchedBySourceDeleted;
-        }
-      };
-    }
-  }
-
-  /**
-   * Returns a new builder for MergeMetrics.
-   */
-  static Builder builder() {
-    return new MergeMetrics.Builder();
-  }
+public interface MergeMetrics {
 
   /**
    * Returns the number of target rows copied unmodified because they did not match any action.
    */
-  long numTargetRowsCopied();
+  OptionalLong numTargetRowsCopied();
 
   /**
    * Returns the number of target rows inserted.
    */
-  long numTargetRowsInserted();
+  OptionalLong numTargetRowsInserted();
 
   /**
    * Returns the number of target rows deleted.
    */
-  long numTargetRowsDeleted();
+  OptionalLong numTargetRowsDeleted();
 
   /**
    * Returns the number of target rows updated.
    */
-  long numTargetRowsUpdated();
+  OptionalLong numTargetRowsUpdated();
 
   /**
    * Returns the number of target rows matched and updated by a matched clause.
    */
-  long numTargetRowsMatchedUpdated();
+  OptionalLong numTargetRowsMatchedUpdated();
 
   /**
    * Returns the number of target rows matched and deleted by a matched clause.
    */
-  long numTargetRowsMatchedDeleted();
+  OptionalLong numTargetRowsMatchedDeleted();
 
   /**
    * Returns the number of target rows updated by a not matched by source clause.
    */
-  long numTargetRowsNotMatchedBySourceUpdated();
+  OptionalLong numTargetRowsNotMatchedBySourceUpdated();
 
   /**
    * Returns the number of target rows deleted by a not matched by source clause.
    */
-  long numTargetRowsNotMatchedBySourceDeleted();
+  OptionalLong numTargetRowsNotMatchedBySourceDeleted();
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/BatchWrite.java

@@ -107,18 +107,11 @@ default void onDataWriterCommit(WriterCommitMessage message) {}
    */
   void abort(WriterCommitMessage[] messages);
 
-  /**
-   * Whether this batch write requests merge execution metrics.
-   */
-  default boolean requestMergeMetrics() {
-    return false;
-  }
-
   /**
    * Similar to {@link #commit(WriterCommitMessage[])}, but providing merge exec metrics to this batch write.
    * @param metrics merge execution metrics
    */
-  default void commitWithMerge(WriterCommitMessage[] messages, MergeMetrics metrics) {
+  default void commitMerge(WriterCommitMessage[] messages, MergeMetrics metrics) {
     commit(messages);
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala

@@ -114,39 +114,36 @@ class InMemoryRowLevelOperationTable(
   }
 
   abstract class RowLevelOperationBatchWrite extends TestBatchWrite {
-    override def requestMergeMetrics(): Boolean = true
 
-    override def commitWithMerge(messages: Array[WriterCommitMessage], metrics: MergeMetrics):
+    override def commitMerge(messages: Array[WriterCommitMessage], metrics: MergeMetrics):
     Unit = {
-      commitProperties += "numTargetRowsCopied" -> metrics.numTargetRowsCopied().toString
-      commitProperties += "numTargetRowsInserted" -> metrics.numTargetRowsInserted().toString
-      commitProperties += "numTargetRowsDeleted" -> metrics.numTargetRowsDeleted().toString
-      commitProperties += "numTargetRowsUpdated" -> metrics.numTargetRowsUpdated().toString
-      commitProperties += "numTargetRowsInserted" -> metrics.numTargetRowsInserted().toString
-      commitProperties += ("numTargetRowsNotMatchedBySourceUpdated"
-        -> metrics.numTargetRowsNotMatchedBySourceUpdated().toString)
-      commitProperties += ("numTargetRowsNotMatchedBySourceDeleted"
-        -> metrics.numTargetRowsNotMatchedBySourceDeleted().toString)
+      commitProperties += "numTargetRowsCopied" -> metrics.numTargetRowsCopied().orElse(-1).toString
+      commitProperties += "numTargetRowsInserted" ->
+        metrics.numTargetRowsInserted().orElse(-1).toString
+      commitProperties += "numTargetRowsDeleted" ->
+        metrics.numTargetRowsDeleted().orElse(-1).toString
+      commitProperties += "numTargetRowsUpdated" ->
+        metrics.numTargetRowsUpdated().orElse(-1).toString
+      commitProperties += "numTargetRowsInserted" ->
+        metrics.numTargetRowsInserted().orElse(-1).toString
       commitProperties += ("numTargetRowsMatchedDeleted"
-        -> metrics.numTargetRowsMatchedDeleted().toString)
+        -> metrics.numTargetRowsMatchedDeleted().orElse(-1).toString)
       commitProperties += ("numTargetRowsMatchedUpdated"
-        -> metrics.numTargetRowsMatchedUpdated().toString)
+        -> metrics.numTargetRowsMatchedUpdated().orElse(-1).toString)
+      commitProperties += ("numTargetRowsNotMatchedBySourceUpdated"
+        -> metrics.numTargetRowsNotMatchedBySourceUpdated().orElse(-1).toString)
+      commitProperties += ("numTargetRowsNotMatchedBySourceDeleted"
+        -> metrics.numTargetRowsNotMatchedBySourceDeleted().orElse(-1).toString)
       commit(messages)
       commits += Commit(Instant.now().toEpochMilli, commitProperties.toMap)
       commitProperties.clear()
     }
-
-    override def commit(messages: Array[WriterCommitMessage]): Unit = {
-      doCommit(messages)
-    }
-
-    def doCommit(messages: Array[WriterCommitMessage]): Unit
   }
 
   private case class PartitionBasedReplaceData(scan: InMemoryBatchScan)
     extends RowLevelOperationBatchWrite {
 
-    override def doCommit(messages: Array[WriterCommitMessage]): Unit = dataMap.synchronized {
+    override def commit(messages: Array[WriterCommitMessage]): Unit = dataMap.synchronized {
       val newData = messages.map(_.asInstanceOf[BufferedRows])
       val readRows = scan.data.flatMap(_.asInstanceOf[BufferedRows].rows)
       val readPartitions = readRows.map(r => getKey(r, schema)).distinct
@@ -203,7 +200,7 @@ class InMemoryRowLevelOperationTable(
       DeltaBufferedRowsWriterFactory
     }
 
-    override def doCommit(messages: Array[WriterCommitMessage]): Unit = {
+    override def commit(messages: Array[WriterCommitMessage]): Unit = {
       val newData = messages.map(_.asInstanceOf[BufferedRows])
       withDeletes(newData)
       withData(newData)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.execution.datasources.v2
 
+import java.util.OptionalLong
+
 import scala.jdk.CollectionConverters._
 
 import org.apache.spark.{SparkEnv, SparkException, TaskContext}
@@ -453,15 +455,11 @@ trait V2TableWriteExec extends V2CommandExec with UnaryExecNode with AdaptiveSpa
         }
       )
 
-      val mergeMetricsOpt = if (batchWrite.requestMergeMetrics()) {
-        Some(getMergeMetrics(query))
-      } else {
-        None
-      }
+      val mergeMetricsOpt = getMergeMetrics(query)
 
       logInfo(log"Data source write support ${MDC(LogKeys.BATCH_WRITE, batchWrite)} is committing.")
       mergeMetricsOpt match {
-        case Some(metrics) => batchWrite.commitWithMerge(messages, metrics)
+        case Some(metrics) => batchWrite.commitMerge(messages, metrics)
         case None => batchWrite.commit(messages)
       }
       logInfo(log"Data source write support ${MDC(LogKeys.BATCH_WRITE, batchWrite)} committed.")
@@ -486,35 +484,29 @@ trait V2TableWriteExec extends V2CommandExec with UnaryExecNode with AdaptiveSpa
     Nil
   }
 
-  private def getMergeMetrics(query: SparkPlan): MergeMetrics = {
-    val mergeMetricsBuilder = new MergeMetrics.Builder()
-    val mergeNode = collectFirst(query) {case c: MergeRowsExec => c}
-    mergeNode.foreach {n => mergeMetricValues(n.metrics, mergeMetricsBuilder)}
-    mergeMetricsBuilder.build()
+  private def getMergeMetrics(query: SparkPlan): Option[MergeMetrics] = {
+    collectFirst(query) { case m: MergeRowsExec => m }.map{ n =>
+      MergeMetricsImpl(
+        numTargetRowsCopied = metric(n.metrics, "numTargetRowsCopied"),
+        numTargetRowsDeleted = metric(n.metrics, "numTargetRowsDeleted"),
+        numTargetRowsUpdated = metric(n.metrics, "numTargetRowsUpdated"),
+        numTargetRowsInserted = metric(n.metrics, "numTargetRowsInserted"),
+        numTargetRowsMatchedDeleted = metric(n.metrics, "numTargetRowsMatchedDeleted"),
+        numTargetRowsMatchedUpdated = metric(n.metrics, "numTargetRowsMatchedUpdated"),
+        numTargetRowsNotMatchedBySourceDeleted =
+          metric(n.metrics, "numTargetRowsNotMatchedBySourceDeleted"),
+        numTargetRowsNotMatchedBySourceUpdated =
+          metric(n.metrics, "numTargetRowsNotMatchedBySourceUpdated")
+      )
+    }
   }
 
-  private def metric(metrics: Map[String, SQLMetric], metric: String): Long = {
+  private def metric(metrics: Map[String, SQLMetric], metric: String): OptionalLong = {
     metrics.get(metric) match {
-      case Some(m) => m.value
-      case None => -1
+      case Some(m) => OptionalLong.of(m.value)
+      case None => OptionalLong.empty()
     }
   }
-
-  private def mergeMetricValues(mergeNodeMetrics: Map[String, SQLMetric],
-                           metricBuilder: MergeMetrics.Builder): Option[MergeMetrics.Builder] = {
-    metricBuilder
-      .numTargetRowsCopied(metric(mergeNodeMetrics, "numTargetRowsCopied"))
-      .numTargetRowsDeleted(metric(mergeNodeMetrics, "numTargetRowsDeleted"))
-      .numTargetRowsUpdated(metric(mergeNodeMetrics, "numTargetRowsUpdated"))
-      .numTargetRowsInserted(metric(mergeNodeMetrics, "numTargetRowsInserted"))
-      .numTargetRowsNotMatchedBySourceDeleted(
-        metric(mergeNodeMetrics, "numTargetRowsNotMatchedBySourceDeleted"))
-      .numTargetRowsNotMatchedBySourceUpdated(
-        metric(mergeNodeMetrics, "numTargetRowsNotMatchedBySourceUpdated"))
-      .numTargetRowsMatchedDeleted(metric(mergeNodeMetrics, "numTargetRowsMatchedDeleted"))
-      .numTargetRowsMatchedUpdated(metric(mergeNodeMetrics, "numTargetRowsMatchedUpdated"))
-    Some(metricBuilder)
-  }
 }
 
 trait WritingSparkTask[W <: DataWriter[InternalRow]] extends Logging with Serializable {
@@ -769,3 +761,13 @@ private[v2] case class DataWritingSparkTaskResult(
  * Sink progress information collected after commit.
  */
 private[sql] case class StreamWriterCommitProgress(numOutputRows: Long)
+
+private case class MergeMetricsImpl(
+    override val numTargetRowsCopied: OptionalLong,
+    override val numTargetRowsDeleted: OptionalLong,
+    override val numTargetRowsUpdated: OptionalLong,
+    override val numTargetRowsInserted: OptionalLong,
+    override val numTargetRowsMatchedUpdated: OptionalLong,
+    override val numTargetRowsMatchedDeleted: OptionalLong,
+    override val numTargetRowsNotMatchedBySourceUpdated: OptionalLong,
+    override val numTargetRowsNotMatchedBySourceDeleted: OptionalLong) extends MergeMetrics