Support schema tracking log at analysis level

Author: PorridgeSwim

spark-unified/src/main/scala/io/delta/internal/ApplyV2Streaming.scala

@@ -21,7 +21,7 @@ import scala.jdk.OptionConverters._
 
 import io.delta.spark.internal.v2.catalog.SparkTable
 import io.delta.spark.internal.v2.utils.ScalaUtils
-import org.apache.spark.sql.delta.DeltaV2Mode
+import org.apache.spark.sql.delta.{DeltaOptions, DeltaV2Mode}
 import org.apache.spark.sql.delta.sources.DeltaSourceUtils
 
 import org.apache.spark.sql.SparkSession
@@ -71,6 +71,16 @@ class ApplyV2Streaming(
     deltaV2Mode.isStreamingReadsEnabled(s.dataSource.catalogTable.toJava)
   }
 
+  /** True when schema-tracking is set on extraOptions but not yet on the SparkTable. */
+  private def needsSchemaTrackingRebuild(
+      table: SparkTable, extraOptions: CaseInsensitiveStringMap): Boolean = {
+    val tableOptions = new CaseInsensitiveStringMap(table.getOptions)
+    (extraOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION) ||
+        extraOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION_ALIAS)) &&
+      !tableOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION) &&
+      !tableOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION_ALIAS)
+  }
+
   override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperators {
     case s: StreamingRelation if shouldApplyV2Streaming(s) =>
       // catalogTable is guaranteed to be defined because shouldApplyV2Streaming checks it
@@ -102,5 +112,18 @@ class ApplyV2Streaming(
         identifier = Some(ident),
         // Keep this None to force the V2 path; we don't want to fall back to V1 here.
         v1Relation = None)
+
+    // TODO(#5319): remove this rule after Spark supports directly create table reflect trackingLog
+    case s @ StreamingRelationV2(_, _, table: SparkTable, extraOptions, _, _, _, _)
+        if needsSchemaTrackingRebuild(table, extraOptions) =>
+      val merged = new java.util.HashMap[String, String]()
+      merged.putAll(table.getOptions)
+      merged.putAll(extraOptions.asCaseSensitiveMap())
+      val rebuilt = if (table.getCatalogTable.isPresent) {
+        new SparkTable(table.getIdentifier, table.getCatalogTable.get, merged)
+      } else {
+        new SparkTable(table.getIdentifier, table.getTablePath.toString, merged)
+      }
+      s.copy(table = rebuilt, output = toAttributes(rebuilt.schema))
   }
 }

spark-unified/src/test/scala/io/delta/internal/ApplyV2StreamingSuite.scala

@@ -21,20 +21,25 @@ import java.util.{HashMap => JHashMap}
 
 import scala.jdk.CollectionConverters._
 
+import io.delta.kernel.internal.SnapshotImpl
 import io.delta.spark.internal.v2.catalog.SparkTable
+import io.delta.spark.internal.v2.snapshot.PathBasedSnapshotManager
 import io.delta.storage.commit.uccommitcoordinator.UCCommitCoordinatorClient
+import org.apache.spark.sql.delta.DeltaLog
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.catalog.CatalogTable
 import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTableType}
 import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2
+import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
+import org.apache.spark.sql.delta.DeltaOptions
 import org.apache.spark.sql.delta.Relocated.StreamingRelation
-import org.apache.spark.sql.delta.sources.DeltaSQLConf
+import org.apache.spark.sql.delta.sources.{DeltaSourceMetadataTrackingLog, DeltaSQLConf, PersistedMetadata}
 import org.apache.spark.sql.delta.test.DeltaSQLCommandTest
 import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
 import org.apache.spark.sql.execution.datasources.DataSource
 import org.apache.spark.sql.connector.catalog.Identifier
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{StringType, StructType}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
 class ApplyV2StreamingSuite extends DeltaSQLCommandTest {
@@ -140,4 +145,223 @@ class ApplyV2StreamingSuite extends DeltaSQLCommandTest {
       }
     }
   }
+
+  // ---------------------------------------------------------------------------
+  // Rebuild StreamingRelationV2 if provided schema tracking log provided
+  // ---------------------------------------------------------------------------
+
+  /** The data-schema seeded into the tracking log by [[seedSchemaLogWithExtraColumn]]. */
+  private val seededFieldNames: Seq[String] = Seq("id", "extra")
+
+  private def buildStreamingRelationV2(
+      table: SparkTable, extraOptions: Map[String, String]): StreamingRelationV2 = {
+    StreamingRelationV2(
+      source = None,
+      sourceName = "delta",
+      table = table,
+      extraOptions = new CaseInsensitiveStringMap(extraOptions.asJava),
+      output = toAttributes(table.schema),
+      catalog = None,
+      identifier = Some(table.getIdentifier),
+      v1Relation = None)
+  }
+
+  /**
+   * Pre-seed the schema-tracking log at `schemaLogPath` with a 2-column schema
+   * (`id LONG, extra STRING`) that differs from the underlying snapshot's 1-column schema
+   */
+  private def seedSchemaLogWithExtraColumn(tablePath: String, schemaLogPath: String): Unit = {
+    val deltaLog = DeltaLog.forTable(spark, tablePath)
+    val snapshotManager =
+      new PathBasedSnapshotManager(tablePath, deltaLog.newDeltaHadoopConf())
+    val tableId =
+      snapshotManager.loadLatestSnapshot.asInstanceOf[SnapshotImpl].getMetadata.getId
+    val trackingLog = DeltaSourceMetadataTrackingLog.create(
+      spark, schemaLogPath, tableId, tablePath, parameters = Map.empty[String, String])
+    val customSchemaJson =
+      """{"type":"struct","fields":[
+        |{"name":"id","type":"long","nullable":true,"metadata":{}},
+        |{"name":"extra","type":"string","nullable":true,"metadata":{}}]}""".stripMargin
+    val emptyPartitionJson = """{"type":"struct","fields":[]}"""
+    val seededEntry = PersistedMetadata(
+      tableId,
+      deltaCommitVersion = 0L,
+      dataSchemaJson = customSchemaJson,
+      partitionSchemaJson = emptyPartitionJson,
+      sourceMetadataPath = tablePath + "/_delta_log/_streaming_metadata")
+    trackingLog.writeNewMetadata(seededEntry, replaceCurrent = false)
+  }
+
+  /** Asserts the table's schema matches the entry written by [[seedSchemaLogWithExtraColumn]]. */
+  private def assertSchemaMatchesSeededLogEntry(table: SparkTable): Unit = {
+    assert(table.schema.fieldNames.toSeq == seededFieldNames)
+    assert(table.schema.fields(1).dataType == StringType)
+  }
+
+  /**
+   * Build a catalog-backed SparkTable rooted at `tableLocationUri`. Mirrors the common production
+   * path through DeltaCatalog and is the default for tests that do not specifically distinguish
+   * between path-based and catalog-based construction.
+   */
+  private def buildCatalogBasedSparkTable(
+      tableLocationUri: URI, options: JHashMap[String, String]): SparkTable = {
+    val catalogTable = createCatalogTable(tableLocationUri, ucManaged = false)
+    val identifier = Identifier.of(
+      catalogTable.identifier.database.toArray, catalogTable.identifier.table)
+    new SparkTable(identifier, catalogTable, options)
+  }
+
+  test("schema-tracking rebuild: path-based SparkTable picks up the persisted schema") {
+    withTempDir { tableDir =>
+      withTempDir { schemaLogDir =>
+        val tablePath = tableDir.getCanonicalPath
+        createDeltaTable(tablePath) // snapshot schema: id BIGINT
+        val schemaLogPath = schemaLogDir.getCanonicalPath
+        seedSchemaLogWithExtraColumn(tablePath, schemaLogPath)
+
+        val identifier = Identifier.of(Array("default"), "tbl")
+        val table = new SparkTable(identifier, tablePath)
+        assert(!table.getOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION))
+
+        val plan = buildStreamingRelationV2(
+          table, Map(DeltaOptions.SCHEMA_TRACKING_LOCATION -> schemaLogPath))
+        val result = applyRule(plan).asInstanceOf[StreamingRelationV2]
+        val rebuiltTable = result.table.asInstanceOf[SparkTable]
+
+        assert(rebuiltTable ne table, "rebuild should produce a new SparkTable")
+        assert(rebuiltTable.getOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION))
+        assert(rebuiltTable.getOptions.get(DeltaOptions.SCHEMA_TRACKING_LOCATION) ==
+          schemaLogPath)
+        assert(!rebuiltTable.getCatalogTable.isPresent,
+          "path branch should not have catalogTable")
+        // Rebuilt schema is driven by the persisted entry, not the snapshot.
+        assertSchemaMatchesSeededLogEntry(rebuiltTable)
+        // And the rule's output is re-derived from that rebuilt schema.
+        assert(result.output.map(_.name) == seededFieldNames)
+
+        // Idempotent: re-applying the rule does not rebuild a second time.
+        val reappliedResult = applyRule(result).asInstanceOf[StreamingRelationV2]
+        assert(reappliedResult.table eq rebuiltTable, "re-applying rule should not rebuild")
+      }
+    }
+  }
+
+  test("schema-tracking rebuild: catalog-based SparkTable picks up the persisted schema and " +
+      "keeps its CatalogTable") {
+    withTempDir { tableDir =>
+      withTempDir { schemaLogDir =>
+        val tablePath = tableDir.getCanonicalPath
+        createDeltaTable(tablePath)
+        val schemaLogPath = schemaLogDir.getCanonicalPath
+        seedSchemaLogWithExtraColumn(tablePath, schemaLogPath)
+
+        val table = buildCatalogBasedSparkTable(tableDir.toURI, new JHashMap[String, String]())
+        assert(!table.getOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION))
+        assert(table.getCatalogTable.isPresent)
+
+        val plan = buildStreamingRelationV2(
+          table, Map(DeltaOptions.SCHEMA_TRACKING_LOCATION -> schemaLogPath))
+        val result = applyRule(plan).asInstanceOf[StreamingRelationV2]
+        val rebuiltTable = result.table.asInstanceOf[SparkTable]
+
+        assert(rebuiltTable.getOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION))
+        assert(rebuiltTable.getCatalogTable.isPresent,
+          "catalog branch should keep CatalogTable")
+        assertSchemaMatchesSeededLogEntry(rebuiltTable)
+      }
+    }
+  }
+
+  test("schema-tracking rebuild: triggered by SCHEMA_TRACKING_LOCATION_ALIAS option key") {
+    withTempDir { tableDir =>
+      withTempDir { schemaLogDir =>
+        val tablePath = tableDir.getCanonicalPath
+        createDeltaTable(tablePath)
+        val schemaLogPath = schemaLogDir.getCanonicalPath
+        seedSchemaLogWithExtraColumn(tablePath, schemaLogPath)
+
+        val table = buildCatalogBasedSparkTable(tableDir.toURI, new JHashMap[String, String]())
+
+        val plan = buildStreamingRelationV2(
+          table, Map(DeltaOptions.SCHEMA_TRACKING_LOCATION_ALIAS -> schemaLogPath))
+        val result = applyRule(plan).asInstanceOf[StreamingRelationV2]
+        val rebuiltTable = result.table.asInstanceOf[SparkTable]
+
+        assert(rebuiltTable.getOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION_ALIAS))
+        assert(rebuiltTable.getOptions.get(DeltaOptions.SCHEMA_TRACKING_LOCATION_ALIAS) ==
+          schemaLogPath)
+        assertSchemaMatchesSeededLogEntry(rebuiltTable)
+      }
+    }
+  }
+
+  test("schema-tracking rebuild: skipped when extraOptions has no schema-tracking option") {
+    withTempDir { tableDir =>
+      val tablePath = tableDir.getCanonicalPath
+      createDeltaTable(tablePath)
+      val table = buildCatalogBasedSparkTable(tableDir.toURI, new JHashMap[String, String]())
+
+      val plan = buildStreamingRelationV2(table, Map.empty)
+      val result = applyRule(plan)
+      assert(result eq plan, "no rebuild expected when schema-tracking option not present")
+    }
+  }
+
+  test("schema-tracking rebuild: skipped when SparkTable already carries the " +
+      "schema-tracking option") {
+    withTempDir { tableDir =>
+      withTempDir { schemaLogDir =>
+        val tablePath = tableDir.getCanonicalPath
+        createDeltaTable(tablePath)
+        val schemaLogPath = schemaLogDir.getCanonicalPath
+        val tableOptions = new JHashMap[String, String]()
+        tableOptions.put(DeltaOptions.SCHEMA_TRACKING_LOCATION, schemaLogPath)
+        val table = buildCatalogBasedSparkTable(tableDir.toURI, tableOptions)
+        assert(table.getOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION))
+
+        val plan = buildStreamingRelationV2(
+          table, Map(DeltaOptions.SCHEMA_TRACKING_LOCATION -> schemaLogPath))
+        val result = applyRule(plan)
+        assert(result eq plan, "no rebuild expected when table already carries the option")
+      }
+    }
+  }
+
+  test("schema-tracking via V1 StreamingRelation: option propagates through V1 -> V2 conversion") {
+    // Counterpart to the V2 rebuild tests above: those start from StreamingRelationV2 and exercise
+    // the rebuild branch. This test starts from a V1 StreamingRelation carrying the schema-tracking
+    // option in dataSource.options, and verifies the V1 -> V2 conversion branch hands the option to
+    // the new SparkTable so its schema is driven by the persisted log entry.
+    withTempDir { tableDir =>
+      withTempDir { schemaLogDir =>
+        val tablePath = tableDir.getCanonicalPath
+        createDeltaTable(tablePath)
+        val schemaLogPath = schemaLogDir.getCanonicalPath
+        seedSchemaLogWithExtraColumn(tablePath, schemaLogPath)
+
+        val catalogTable = createCatalogTable(tableDir.toURI, ucManaged = false)
+        val dataSource = DataSource(
+          sparkSession = spark,
+          userSpecifiedSchema = None,
+          className = "delta",
+          options = Map(
+            "path" -> tablePath,
+            DeltaOptions.SCHEMA_TRACKING_LOCATION -> schemaLogPath),
+          catalogTable = Some(catalogTable))
+        val plan = StreamingRelation(dataSource)
+
+        // STRICT mode forces V1 -> V2 conversion in ApplyV2Streaming.
+        withSQLConf(DeltaSQLConf.V2_ENABLE_MODE.key -> "STRICT") {
+          val result = applyRule(plan).asInstanceOf[StreamingRelationV2]
+          val convertedTable = result.table.asInstanceOf[SparkTable]
+          assert(convertedTable.getOptions.containsKey(DeltaOptions.SCHEMA_TRACKING_LOCATION))
+          assert(convertedTable.getOptions.get(DeltaOptions.SCHEMA_TRACKING_LOCATION) ==
+            schemaLogPath)
+          // Schema is driven by the seeded log entry, not the underlying snapshot.
+          assertSchemaMatchesSeededLogEntry(convertedTable)
+          assert(result.output.map(_.name) == seededFieldNames)
+        }
+      }
+    }
+  }
 }

spark-unified/src/test/scala/org/apache/spark/sql/delta/test/DeltaV2SourceSchemaEvolutionSuite.scala

@@ -71,12 +71,8 @@ trait DeltaV2SourceSchemaEvolutionSuiteBase extends V2ForceTest {
   }
 
   // TODO(#5319): Move tests to shouldPassTests as V2 schema tracking log support is implemented.
-  override protected def shouldPassTests: Set[String] = Set.empty[String]
-
-  // All tests from StreamingSchemaEvolutionSuiteBase.
-  // Override in CDC suites to add CDC-specific tests.
-  override protected def shouldFailTests: Set[String] = Set(
-    // ========== Schema location validation ==========
+  override protected def shouldPassTests: Set[String] = Set(
+    // ========== Schema log unit test ==========
     "schema location not under checkpoint",
     "schema location same as checkpoint",
     "schema location using a different file system",
@@ -91,11 +87,18 @@ trait DeltaV2SourceSchemaEvolutionSuiteBase extends V2ForceTest {
     "schema / checkpoint location unit tests - " +
       "schema location and checkpoint location are the same but with explicit file scheme",
     "schema / checkpoint location unit tests - special characters in schema location",
+    "concurrent schema log modification should be detected",
+    "schema log replace current",
+    "backward-compat: latest version can read back older JSON",
+    "forward-compat: older version can read back newer JSON",
 
     // ========== Schema log core ==========
-    "multiple delta source sharing same schema log is blocked",
+    "multiple delta source sharing same schema log is blocked"
+  )
+
+  override protected def shouldFailTests: Set[String] = Set(
+    // ========== Schema log core ==========
     "schema log is applied",
-    "concurrent schema log modification should be detected",
     "schema log initialization with additive schema changes",
     "detect incompatible schema change while streaming",
     "detect incompatible schema change during first getBatch",
@@ -123,12 +126,7 @@ trait DeltaV2SourceSchemaEvolutionSuiteBase extends V2ForceTest {
     "unblock with sql conf",
     "schema tracking interacting with unsafe escape flag",
     "streaming with a column mapping upgrade",
-    "partition evolution",
-    "schema log replace current",
-
-    // ========== Backward/forward compatibility ==========
-    "backward-compat: latest version can read back older JSON",
-    "forward-compat: older version can read back newer JSON"
+    "partition evolution"
   )
 }
 
@@ -143,25 +141,24 @@ class DeltaV2SourceSchemaEvolutionIdColumnMappingSuite
     with DeltaV2SourceSchemaEvolutionSuiteBase
 
 // CDC suites
+// TODO(#5319): Support CDC non-additive schema evolution
+trait DeltaV2SourceSchemaEvolutionCDCSuiteBase extends DeltaV2SourceSchemaEvolutionSuiteBase {
+  self: StreamingSchemaEvolutionSuiteBase =>
 
-class DeltaV2SourceSchemaEvolutionCDCNameColumnMappingSuite
-  extends DeltaSourceSchemaEvolutionCDCNameColumnMappingSuite
-    with DeltaV2SourceSchemaEvolutionSuiteBase {
+  override protected def shouldPassTests: Set[String] = Set.empty[String]
 
-  override protected def shouldFailTests: Set[String] = super.shouldFailTests ++ Set(
-    // Additional tests from CDCStreamingSchemaEvolutionSuiteBase
-    "CDC streaming with schema evolution",
-    "protocol and configuration evolution"
-  )
+  override protected def shouldFailTests: Set[String] =
+    super.shouldPassTests ++ super.shouldFailTests ++ Set(
+      // Additional tests from CDCStreamingSchemaEvolutionSuiteBase
+      "CDC streaming with schema evolution",
+      "protocol and configuration evolution"
+    )
 }
 
+class DeltaV2SourceSchemaEvolutionCDCNameColumnMappingSuite
+  extends DeltaSourceSchemaEvolutionCDCNameColumnMappingSuite
+    with DeltaV2SourceSchemaEvolutionCDCSuiteBase
+
 class DeltaV2SourceSchemaEvolutionCDCIdColumnMappingSuite
   extends DeltaSourceSchemaEvolutionCDCIdColumnMappingSuite
-    with DeltaV2SourceSchemaEvolutionSuiteBase {
-
-  override protected def shouldFailTests: Set[String] = super.shouldFailTests ++ Set(
-    // Additional tests from CDCStreamingSchemaEvolutionSuiteBase
-    "CDC streaming with schema evolution",
-    "protocol and configuration evolution"
-  )
-}
+    with DeltaV2SourceSchemaEvolutionCDCSuiteBase