Add metadata evolution handler in v2

spark/src/main/scala/org/apache/spark/sql/delta/actions/actions.scala

@@ -165,7 +165,7 @@ sealed trait Action {
  * Note: Please initialize this class using the companion object's `apply` method, which will
  * assign correct values (`Set()` vs `None`) to [[readerFeatures]] and [[writerFeatures]].
  */
-case class Protocol private (
+case class Protocol (
     minReaderVersion: Int,
     minWriterVersion: Int,
     @JsonInclude(Include.NON_ABSENT) // write to JSON only when the field is not `None`
@@ -1196,8 +1196,7 @@ case class Metadata(
 
   /** Returns the partitionSchema as a [[StructType]] */
   @JsonIgnore
-  lazy val partitionSchema: StructType =
-    new StructType(partitionColumns.map(c => schema(c)).toArray)
+  override lazy val partitionSchema: StructType = super.partitionSchema
 
   /** Partition value keys in the AddFile map. */
   @JsonIgnore

spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaDataSource.scala

@@ -518,11 +518,14 @@ object DeltaDataSource extends DatabricksLogging {
         DeltaSourceMetadataTrackingLog.create(
           spark,
           schemaTrackingLocation,
-          sourceSnapshot,
-          catalogTableOpt,
+          sourceSnapshot.deltaLog.unsafeVolatileTableId,
+          sourceSnapshot.deltaLog.dataPath.toString,
           parameters,
           sourceMetadataPathOpt,
-          mergeConsecutiveSchemaChanges
+          mergeConsecutiveSchemaChanges,
+          currentMetadata =>
+            DeltaSourceMetadataEvolutionSupport.getMergedConsecutiveMetadataChanges(
+              spark, sourceSnapshot.deltaLog, catalogTableOpt, currentMetadata)
         )
       }
   }

spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSourceMetadataEvolutionSupport.scala

@@ -21,12 +21,15 @@ import java.util.Locale
 import scala.collection.mutable
 
 import org.apache.spark.sql.delta._
-import org.apache.spark.sql.delta.actions.{Action, Metadata, Protocol}
+import org.apache.spark.sql.delta.actions.{Action, FileAction, Metadata, Protocol}
 import org.apache.spark.sql.delta.schema.SchemaUtils
 import org.apache.spark.sql.delta.storage.ClosableIterator
 import org.apache.spark.sql.delta.storage.ClosableIterator._
+import org.apache.spark.sql.delta.v2.interop.{AbstractMetadata, AbstractProtocol}
 
+import org.apache.spark.internal.Logging
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.catalog.CatalogTable
 import org.apache.spark.sql.execution.streaming.Offset
 import org.apache.spark.sql.types.StructType
 
@@ -87,26 +90,13 @@ import org.apache.spark.sql.types.StructType
  */
 trait DeltaSourceMetadataEvolutionSupport extends DeltaSourceBase { base: DeltaSource =>
 
-  /**
-   * Whether this DeltaSource is utilizing a schema log entry as its read schema.
-   *
-   * If user explicitly turn on the flag to fall back to using latest schema to read (i.e. the
-   * legacy mode), we will ignore the schema log.
-   */
   protected def trackingMetadataChange: Boolean =
-    !schemaReadOptions.allowUnsafeStreamingReadOnColumnMappingSchemaChanges &&
-      metadataTrackingLog.flatMap(_.getCurrentTrackedMetadata).nonEmpty
+    DeltaSourceMetadataEvolutionSupport.shouldTrackMetadataChange(
+      schemaReadOptions, metadataTrackingLog)
 
-  /**
-   * Whether a schema tracking log is provided (and is empty), so we could initialize eagerly.
-   * This should only be used for the first write to the schema log, after then, schema tracking
-   * should not rely on this state any more.
-   */
   protected def readyToInitializeMetadataTrackingEagerly: Boolean =
-    !schemaReadOptions.allowUnsafeStreamingReadOnColumnMappingSchemaChanges &&
-      metadataTrackingLog.exists { log =>
-        log.getCurrentTrackedMetadata.isEmpty && log.initMetadataLogEagerly
-      }
+    DeltaSourceMetadataEvolutionSupport.shouldInitializeMetadataTrackingEagerly(
+      schemaReadOptions, metadataTrackingLog)
 
 
   /**
@@ -124,44 +114,22 @@ trait DeltaSourceMetadataEvolutionSupport extends DeltaSourceBase { base: DeltaS
     }
   }
 
-  /**
-   * Check the table metadata or protocol changed since the initial read snapshot. We make sure:
-   * 1. The schema is the same, except for internal metadata, AND
-   * 2. The delta related table configurations are strictly equal, AND
-   * 3. The incoming metadata change should not be considered a failure-causing change if we have
-   *    marked the persisted schema and the stream progress is behind that schema version.
-   *    This could happen when we've already merged consecutive schema changes during the analysis
-   *    phase and we are using the merged schema as the read schema. All the schema changes in
-   *    between can be safely ignored because they won't contribute any data.
-   */
   private def hasMetadataOrProtocolChangeComparedToStreamMetadata(
       metadataChangeOpt: Option[Metadata],
       protocolChangeOpt: Option[Protocol],
       newSchemaVersion: Long): Boolean = {
-    if (persistedMetadataAtSourceInit.exists(_.deltaCommitVersion >= newSchemaVersion)) {
-      false
-    } else {
-      protocolChangeOpt.exists(_ != readProtocolAtSourceInit) ||
-      metadataChangeOpt.exists { newMetadata =>
-         hasSchemaChangeComparedToStreamMetadata(newMetadata.schema) ||
-           newMetadata.partitionSchema != readPartitionSchemaAtSourceInit ||
-           newMetadata.configuration.filterKeys(_.startsWith("delta.")).toMap !=
-             readConfigurationsAtSourceInit.filterKeys(_.startsWith("delta.")).toMap
-      }
-    }
+    DeltaSourceMetadataEvolutionSupport.hasMetadataOrProtocolChangeComparedToStreamMetadata(
+      metadataChangeOpt,
+      protocolChangeOpt,
+      newSchemaVersion,
+      persistedMetadataAtSourceInit,
+      readProtocolAtSourceInit,
+      readSchemaAtSourceInit,
+      readPartitionSchemaAtSourceInit,
+      readConfigurationsAtSourceInit,
+      spark)
   }
 
-  /**
-   * Check that the give schema is the same as the schema from the initial read snapshot.
-   */
-  private def hasSchemaChangeComparedToStreamMetadata(newSchema: StructType): Boolean =
-    if (spark.conf.get(DeltaSQLConf.DELTA_STREAMING_IGNORE_INTERNAL_METADATA_FOR_SCHEMA_CHANGE)) {
-      DeltaTableUtils.removeInternalWriterMetadata(spark, newSchema) !=
-        DeltaTableUtils.removeInternalWriterMetadata(spark, readSchemaAtSourceInit)
-    } else {
-      newSchema != readSchemaAtSourceInit
-    }
-
   /**
    * If the current stream metadata is not equal to the metadata change in [[metadataChangeOpt]],
    * return a metadata change barrier [[IndexedFile]].
@@ -437,7 +405,7 @@ trait DeltaSourceMetadataEvolutionSupport extends DeltaSourceBase { base: DeltaS
   }
 }
 
-object DeltaSourceMetadataEvolutionSupport {
+object DeltaSourceMetadataEvolutionSupport extends Logging {
   /** SQL configs that allow unblocking each type of schema changes. */
   private val SQL_CONF_PREFIX = s"${DeltaSQLConf.SQL_CONF_PREFIX}.streaming"
 
@@ -683,6 +651,150 @@ object DeltaSourceMetadataEvolutionSupport {
     spark.sessionState.conf.getConf(
       DeltaSQLConf.DELTA_TYPE_WIDENING_BYPASS_STREAMING_TYPE_CHANGE_CHECK)
 
+  /**
+   * Whether this source should use schema tracking for metadata evolution.
+   * Shared between v1 and v2 connectors.
+   */
+  def shouldTrackMetadataChange(
+      schemaReadOptions: DeltaStreamUtils.SchemaReadOptions,
+      metadataTrackingLog: Option[DeltaSourceMetadataTrackingLog]): Boolean = {
+    !schemaReadOptions.allowUnsafeStreamingReadOnColumnMappingSchemaChanges &&
+      metadataTrackingLog.flatMap(_.getCurrentTrackedMetadata).nonEmpty
+  }
+
+  /**
+   * Whether the tracking log should be initialized eagerly (log is provided but empty).
+   * Shared between v1 and v2 connectors.
+   */
+  def shouldInitializeMetadataTrackingEagerly(
+      schemaReadOptions: DeltaStreamUtils.SchemaReadOptions,
+      metadataTrackingLog: Option[DeltaSourceMetadataTrackingLog]): Boolean = {
+    !schemaReadOptions.allowUnsafeStreamingReadOnColumnMappingSchemaChanges &&
+      metadataTrackingLog.exists { log =>
+        log.getCurrentTrackedMetadata.isEmpty && log.initMetadataLogEagerly
+      }
+  }
+
+  /**
+   * Check the table metadata or protocol changed since the initial read snapshot. We make sure:
+   * 1. The schema is the same, except for internal metadata, AND
+   * 2. The delta related table configurations are strictly equal, AND
+   * 3. The incoming metadata change should not be considered a failure-causing change if we have
+   *    marked the persisted schema and the stream progress is behind that schema version.
+   *    This could happen when we've already merged consecutive schema changes during the analysis
+   *    phase and we are using the merged schema as the read schema. All the schema changes in
+   *    between can be safely ignored because they won't contribute any data.
+   *
+   * @param metadataChangeOpt New metadata action, if any.
+   * @param protocolChangeOpt New protocol action, if any.
+   * @param newSchemaVersion The version of the incoming change.
+   * @param persistedMetadataAtSourceInit The persisted metadata at source init, if any.
+   * @param readProtocolAtSourceInit The protocol at source init.
+   * @param readSchemaAtSourceInit The schema at source init.
+   * @param readPartitionSchemaAtSourceInit The partition schema at source init.
+   * @param readConfigurationsAtSourceInit The table configurations at source init.
+   * @param spark The SparkSession (used for SQL conf checks).
+   */
+  def hasMetadataOrProtocolChangeComparedToStreamMetadata(
+      metadataChangeOpt: Option[AbstractMetadata],
+      protocolChangeOpt: Option[AbstractProtocol],
+      newSchemaVersion: Long,
+      persistedMetadataAtSourceInit: Option[PersistedMetadata],
+      readProtocolAtSourceInit: AbstractProtocol,
+      readSchemaAtSourceInit: StructType,
+      readPartitionSchemaAtSourceInit: StructType,
+      readConfigurationsAtSourceInit: Map[String, String],
+      spark: SparkSession): Boolean = {
+    if (persistedMetadataAtSourceInit.exists(_.deltaCommitVersion >= newSchemaVersion)) {
+      false
+    } else {
+      protocolChangeOpt.exists(p =>
+        p.minReaderVersion != readProtocolAtSourceInit.minReaderVersion ||
+          p.minWriterVersion != readProtocolAtSourceInit.minWriterVersion ||
+          p.readerFeatures != readProtocolAtSourceInit.readerFeatures ||
+          p.writerFeatures != readProtocolAtSourceInit.writerFeatures) ||
+      metadataChangeOpt.exists { newMetadata =>
+        hasSchemaChangeComparedToStreamMetadata(
+          newMetadata.schema, readSchemaAtSourceInit, spark) ||
+          newMetadata.partitionColumns != readPartitionSchemaAtSourceInit ||
+          newMetadata.configuration.filterKeys(_.startsWith("delta.")).toMap !=
+            readConfigurationsAtSourceInit.filterKeys(_.startsWith("delta.")).toMap
+      }
+    }
+  }
+
+  /**
+   * Check that the given schema is the same as the schema from the initial read snapshot.
+   * This is shared between v1 and v2 connectors.
+   */
+  def hasSchemaChangeComparedToStreamMetadata(
+      newSchema: StructType,
+      readSchemaAtSourceInit: StructType,
+      spark: SparkSession): Boolean = {
+    if (spark.conf.get(DeltaSQLConf.DELTA_STREAMING_IGNORE_INTERNAL_METADATA_FOR_SCHEMA_CHANGE)) {
+      DeltaTableUtils.removeInternalWriterMetadata(spark, newSchema) !=
+        DeltaTableUtils.removeInternalWriterMetadata(spark, readSchemaAtSourceInit)
+    } else {
+      newSchema != readSchemaAtSourceInit
+    }
+  }
+
+  /**
+   * Speculate ahead and find the next merged consecutive metadata change if possible.
+   * A metadata change is either:
+   * 1. A [[Metadata]] action change. OR
+   * 2. A [[Protocol]] change.
+   */
+ def getMergedConsecutiveMetadataChanges(
+     spark: SparkSession,
+     deltaLog: DeltaLog,
+     catalogTableOpt: Option[CatalogTable],
+     currentMetadata: PersistedMetadata): Option[PersistedMetadata] = {
+    val currentMetadataVersion = currentMetadata.deltaCommitVersion
+    // We start from the currentSchemaVersion so that we can stop early in case the current
+    // version still has file actions that potentially needs to be processed.
+    val untilMetadataChange =
+      deltaLog.getChangeLogFiles(
+        currentMetadataVersion, catalogTableOpt).map { case (version, fileStatus) =>
+        var metadataAction: Option[Metadata] = None
+        var protocolAction: Option[Protocol] = None
+        var hasFileAction = false
+        DeltaSource.createRewindableActionIterator(spark, deltaLog, fileStatus)
+          .processAndClose { actionsIter =>
+            actionsIter.foreach {
+              case m: Metadata => metadataAction = Some(m)
+              case p: Protocol => protocolAction = Some(p)
+              case _: FileAction => hasFileAction = true
+              case _ =>
+            }
+          }
+        (!hasFileAction && (metadataAction.isDefined || protocolAction.isDefined),
+          version, metadataAction, protocolAction)
+      }.takeWhile(_._1)
+    DeltaSource.iteratorLast(untilMetadataChange.toClosable)
+      .flatMap { case (_, version, metadataOpt, protocolOpt) =>
+        if (version == currentMetadataVersion) {
+          None
+        } else {
+          log.info(s"Looked ahead from version $currentMetadataVersion and " +
+            s"will use metadata at version $version to read Delta stream.")
+          Some(
+            currentMetadata.copy(
+              deltaCommitVersion = version,
+              dataSchemaJson =
+                metadataOpt.map(_.schema.json).getOrElse(currentMetadata.dataSchemaJson),
+              partitionSchemaJson =
+                metadataOpt.map(_.partitionSchema.json)
+                  .getOrElse(currentMetadata.partitionSchemaJson),
+              tableConfigurations = metadataOpt.map(_.configuration)
+                .orElse(currentMetadata.tableConfigurations),
+              protocolJson = protocolOpt.map(_.json).orElse(currentMetadata.protocolJson)
+            )
+          )
+        }
+      }
+  }
+
   // scalastyle:off
   /**
    * Given a non-additive operation type from a previous schema evolution, check we can process