second attempt

Author: szehon-ho

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

@@ -0,0 +1,17 @@
+package org.apache.spark.sql.connector.metric;
+
+public interface V2ExecMetric {
+  /**
+   * Returns the type of V2 exec metric (ie, firstScan, secondScan, merge, etc).
+   */
+  String metricType();
+  /**
+   * Returns the name of V2 exec metric.
+   */
+  String name();
+
+  /**
+   * Returns the value of V2 exec metric.
+   */
+  String value();
+}

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

@@ -19,6 +19,7 @@
 
 import org.apache.spark.annotation.Evolving;
 import org.apache.spark.sql.connector.metric.CustomTaskMetric;
+import org.apache.spark.sql.connector.metric.V2ExecMetric;
 
 /**
  * An interface that defines how to write the data to data source for batch processing.
@@ -108,16 +109,15 @@ default void onDataWriterCommit(WriterCommitMessage message) {}
 
 
   /**
-   * Whether this batch write requests execution metrics. Returns a row level operation command this batch write
-   * is part of, if requested. Return null if not requested.
+   * Whether this batch write requests execution metrics.
    */
-  default RowLevelOperation.Command requestExecMetrics() {
-    return null;
+  default boolean requestExecMetrics() {
+    return false;
   }
 
   /**
    * Provides an array of query execution metrics to the batch write prior to commit.
    * @param metrics an array of execution metrics
    */
-  default void execMetrics(CustomTaskMetric[] metrics) {}
+  default void execMetrics(V2ExecMetric[] metrics) {}
 }

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

@@ -501,8 +501,6 @@ abstract class InMemoryBaseTable(
       options: CaseInsensitiveStringMap)
     extends BatchScanBaseClass(_data, readSchema, tableSchema) with SupportsRuntimeFiltering {
 
-    var setFilters = Array.empty[Filter]
-
     override def reportDriverMetrics(): Array[CustomTaskMetric] =
       Array(new CustomTaskMetric{
         override def name(): String = "numSplits"
@@ -526,7 +524,6 @@ abstract class InMemoryBaseTable(
     }
 
     override def filter(filters: Array[Filter]): Unit = {
-      this.setFilters = filters
       if (partitioning.length == 1 && partitioning.head.references().length == 1) {
         val ref = partitioning.head.references().head
         filters.foreach {
@@ -598,7 +595,7 @@ abstract class InMemoryBaseTable(
 
   protected abstract class TestBatchWrite extends BatchWrite {
 
-    var commitProperties: mutable.Map[String, String] = mutable.Map.empty[String, String]
+    val commitProperties: mutable.Map[String, String] = mutable.Map.empty[String, String]
 
     override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = {
       BufferedRowsWriterFactory

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

@@ -20,14 +20,12 @@ package org.apache.spark.sql.connector.catalog
 import java.time.Instant
 import java.util
 
-import scala.collection.mutable.ListBuffer
-
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
 import org.apache.spark.sql.connector.catalog.constraints.Constraint
 import org.apache.spark.sql.connector.distributions.{Distribution, Distributions}
 import org.apache.spark.sql.connector.expressions.{FieldReference, LogicalExpressions, NamedReference, SortDirection, SortOrder, Transform}
-import org.apache.spark.sql.connector.metric.CustomTaskMetric
+import org.apache.spark.sql.connector.metric.V2ExecMetric
 import org.apache.spark.sql.connector.read.{Scan, ScanBuilder}
 import org.apache.spark.sql.connector.write.{BatchWrite, DeltaBatchWrite, DeltaWrite, DeltaWriteBuilder, DeltaWriter, DeltaWriterFactory, LogicalWriteInfo, PhysicalWriteInfo, RequiresDistributionAndOrdering, RowLevelOperation, RowLevelOperationBuilder, RowLevelOperationInfo, SupportsDelta, Write, WriteBuilder, WriterCommitMessage}
 import org.apache.spark.sql.connector.write.RowLevelOperation.Command
@@ -50,8 +48,6 @@ class InMemoryRowLevelOperationTable(
     constraints)
   with SupportsRowLevelOperations {
 
-  private val _scans = ListBuffer.empty[Scan]
-
   private final val PARTITION_COLUMN_REF = FieldReference(PartitionKeyColumn.name)
   private final val INDEX_COLUMN_REF = FieldReference(IndexColumn.name)
   private final val SUPPORTS_DELTAS = "supports-deltas"
@@ -74,16 +70,6 @@ class InMemoryRowLevelOperationTable(
     }
   }
 
-  class InMemoryRowLevelOperationScanBuilder(tableSchema: StructType,
-                                             options: CaseInsensitiveStringMap)
-    extends InMemoryScanBuilder(tableSchema, options) {
-    override def build: Scan = {
-      val scan = super.build
-      _scans += scan
-      scan
-    }
-  }
-
   case class PartitionBasedOperation(command: Command) extends RowLevelOperation {
     var configuredScan: InMemoryBatchScan = _
 
@@ -117,7 +103,7 @@ class InMemoryRowLevelOperationTable(
                 SortDirection.ASCENDING.defaultNullOrdering()))
           }
 
-          override def toBatch: BatchWrite = PartitionBasedReplaceData(configuredScan, command)
+          override def toBatch: BatchWrite = PartitionBasedReplaceData(configuredScan)
 
           override def description: String = "InMemoryWrite"
         }
@@ -127,33 +113,14 @@ class InMemoryRowLevelOperationTable(
     override def description(): String = "InMemoryPartitionReplaceOperation"
   }
 
-  abstract class RowLevelOperationBatchWrite(command: Command) extends TestBatchWrite {
-    override def requestExecMetrics(): Command = command
+  abstract class RowLevelOperationBatchWrite extends TestBatchWrite {
+    override def requestExecMetrics(): Boolean = true
 
-    override def execMetrics(metrics: Array[CustomTaskMetric]): Unit = {
-      metrics.foreach(m => commitProperties += (m.name() -> m.value().toString))
+    override def execMetrics(metrics: Array[V2ExecMetric]): Unit = {
+      metrics.foreach(m => commitProperties += s"${m.metricType()}.${m.name()}" -> m.value())
     }
 
     override def commit(messages: Array[WriterCommitMessage]): Unit = {
-      assert(_scans.size <= 2, "Expected at most two scans in row-level operations")
-      assert(_scans.count{ case s: InMemoryBatchScan => s.setFilters.nonEmpty } <= 1,
-        "Expected at most one scan with runtime filters in row-level operations")
-      assert(_scans.count{ case s: InMemoryBatchScan => s.setFilters.isEmpty } <= 1,
-        "Expected at most one scan without runtime filters in row-level operations")
-
-      _scans.foreach{
-        case s: InMemoryBatchScan =>
-          val prefix = if (s.setFilters.isEmpty) {
-            ""
-          } else {
-            "secondScan."
-          }
-          s.reportDriverMetrics().foreach { metric =>
-            commitProperties += (prefix + metric.name() -> metric.value().toString)
-          }
-        case _ =>
-      }
-      _scans.clear()
       doCommit(messages)
       commits += Commit(Instant.now().toEpochMilli, commitProperties.toMap)
       commitProperties.clear()
@@ -162,9 +129,8 @@ class InMemoryRowLevelOperationTable(
     def doCommit(messages: Array[WriterCommitMessage]): Unit
   }
 
-  private case class PartitionBasedReplaceData(scan: InMemoryBatchScan,
-                                               command: RowLevelOperation.Command)
-    extends RowLevelOperationBatchWrite(command) {
+  private case class PartitionBasedReplaceData(scan: InMemoryBatchScan)
+    extends RowLevelOperationBatchWrite {
 
     override def doCommit(messages: Array[WriterCommitMessage]): Unit = dataMap.synchronized {
       val newData = messages.map(_.asInstanceOf[BufferedRows])
@@ -187,7 +153,7 @@ class InMemoryRowLevelOperationTable(
     override def rowId(): Array[NamedReference] = Array(PK_COLUMN_REF)
 
     override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
-      new InMemoryRowLevelOperationScanBuilder(schema, options)
+      new InMemoryScanBuilder(schema, options)
     }
 
     override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
@@ -208,7 +174,7 @@ class InMemoryRowLevelOperationTable(
             )
           }
 
-          override def toBatch: DeltaBatchWrite = TestDeltaBatchWrite(command)
+          override def toBatch: DeltaBatchWrite = TestDeltaBatchWrite
         }
       }
     }
@@ -218,8 +184,8 @@ class InMemoryRowLevelOperationTable(
     }
   }
 
-  private case class TestDeltaBatchWrite(command: Command)
-    extends RowLevelOperationBatchWrite(command) with DeltaBatchWrite{
+  private object TestDeltaBatchWrite extends RowLevelOperationBatchWrite
+    with DeltaBatchWrite{
 
     override def createBatchWriterFactory(info: PhysicalWriteInfo): DeltaWriterFactory = {
       DeltaBufferedRowsWriterFactory

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

@@ -182,7 +182,8 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
         case Some(r) => session.sharedState.cacheManager.uncacheQuery(session, r, cascade = true)
         case None => ()
       }
-      WriteToDataSourceV2Exec(writer, invalidateCacheFunc, planLater(query), customMetrics) :: Nil
+      WriteToDataSourceV2Exec(writer, invalidateCacheFunc, planLater(query), customMetrics,
+        relationOpt.map(_.table)) :: Nil
 
     case c @ CreateTable(ResolvedIdentifier(catalog, ident), columns, partitioning,
         tableSpec: TableSpec, ifNotExists) =>
@@ -275,7 +276,7 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
       }
 
     case AppendData(r: DataSourceV2Relation, query, _, _, Some(write), _) =>
-      AppendDataExec(planLater(query), refreshCache(r), write) :: Nil
+      AppendDataExec(planLater(query), refreshCache(r), write, Some(r.table)) :: Nil
 
     case OverwriteByExpression(r @ DataSourceV2Relation(v1: SupportsWrite, _, _, _, _), _, _,
         _, _, Some(write), analyzedQuery) if v1.supports(TableCapability.V1_BATCH_WRITE) =>
@@ -290,10 +291,10 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
 
     case OverwriteByExpression(
         r: DataSourceV2Relation, _, query, _, _, Some(write), _) =>
-      OverwriteByExpressionExec(planLater(query), refreshCache(r), write) :: Nil
+      OverwriteByExpressionExec(planLater(query), refreshCache(r), write, Some(r.table)) :: Nil
 
     case OverwritePartitionsDynamic(r: DataSourceV2Relation, query, _, _, Some(write)) =>
-      OverwritePartitionsDynamicExec(planLater(query), refreshCache(r), write) :: Nil
+      OverwritePartitionsDynamicExec(planLater(query), refreshCache(r), write, Some(r.table)) :: Nil
 
     case DeleteFromTableWithFilters(r: DataSourceV2Relation, filters) =>
       DeleteFromTableExec(r.table.asDeletable, filters.toArray, refreshCache(r)) :: Nil
@@ -332,15 +333,15 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
           throw SparkException.internalError("Unexpected table relation: " + other)
       }
 
-    case ReplaceData(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections, _,
+    case ReplaceData(o: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections, _,
         Some(write)) =>
       // use the original relation to refresh the cache
-      ReplaceDataExec(planLater(query), refreshCache(r), projections, write) :: Nil
+      ReplaceDataExec(planLater(query), refreshCache(r), projections, write, Some(o.table)) :: Nil
 
-    case WriteDelta(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections,
+    case WriteDelta(o: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections,
         Some(write)) =>
       // use the original relation to refresh the cache
-      WriteDeltaExec(planLater(query), refreshCache(r), projections, write) :: Nil
+      WriteDeltaExec(planLater(query), refreshCache(r), projections, write, Some(o.table)) :: Nil
 
     case MergeRows(isSourceRowPresent, isTargetRowPresent, matchedInstructions,
         notMatchedInstructions, notMatchedBySourceInstructions, checkCardinality, output, child) =>

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

@@ -73,7 +73,6 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
         case _ =>
           throw QueryExecutionErrors.overwriteTableByUnsupportedExpressionError(table)
       }
-
       val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog)
       o.copy(write = Some(write), query = newQuery)