Schema evolution

Author: andreaschat-db

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

@@ -75,6 +75,10 @@ abstract class InMemoryBaseTable(
 
   override def columns(): Array[Column] = tableColumns
 
+  private[catalog] def updateColumns(newColumns: Array[Column]): Unit = {
+    tableColumns = newColumns
+  }
+
   override def version(): String = tableVersion.toString
 
   def setVersion(version: String): Unit = {

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

@@ -71,9 +71,9 @@ class TxnTable(val delegate: InMemoryRowLevelOperationTable)
     delegate.properties,
     delegate.constraints) {
 
-  // TODO: Revise schema evolution.
-  alterTableWithData(delegate.data, schema)
+  withData(delegate.data)
 
+  // Keep initial version to detect any changes during the transaction.
   private val initialVersion: String = version()
 
   // A tracker of filters used in each scan.
@@ -86,8 +86,8 @@ class TxnTable(val delegate: InMemoryRowLevelOperationTable)
   def commit(): Unit = {
     if (version() != initialVersion) {
       delegate.dataMap.clear()
-      // TODO: Revise schema evolution.
-      delegate.alterTableWithData(data, delegate.schema)
+      delegate.alterTableWithData(data, schema)
+      delegate.updateColumns(columns()) // Evolve schema if needed.
       delegate.replacedPartitions = replacedPartitions
       delegate.lastWriteInfo = lastWriteInfo
       delegate.lastWriteLog = lastWriteLog
@@ -124,6 +124,10 @@ class TxnTableCatalog(delegate: InMemoryRowLevelOperationTableCatalog) extends T
   }
 
   override def alterTable(ident: Identifier, changes: TableChange*): Table = {
+    // TODO: This evicts the staged TxnTable, losing any in-flight DML changes. The correct
+    // approach is to apply only the schema change to the existing TxnTable so that the ongoing
+    // DML can observe the new schema and reconcile at commit time. Concurrent DDL + DML is not
+    // supported in this test catalog for now.
     val newDelegateTable = delegate.alterTable(ident, changes: _*)
     tables.remove(ident) // Load again.
     newDelegateTable

sql/core/src/test/scala/org/apache/spark/sql/connector/AppendDataTransactionSuite.scala

@@ -399,4 +399,102 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase {
         Row(1, 100, "hr"),
         Row(2, 200, "software")))
   }
+
+  test("SQL INSERT WITH SCHEMA EVOLUTION adds new column with transactional checks") {
+    createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+      """{ "pk": 1, "salary": 100, "dep": "hr" }
+        |{ "pk": 2, "salary": 200, "dep": "software" }
+        |""".stripMargin)
+
+    sql(
+      s"""CREATE TABLE $sourceNameAsString
+         |(pk INT NOT NULL, salary INT, dep STRING, active BOOLEAN)""".stripMargin)
+    sql(s"INSERT INTO $sourceNameAsString VALUES (3, 300, 'hr', true), (4, 400, 'software', false)")
+
+    val (txn, txnTables) = executeTransaction {
+      sql(s"INSERT WITH SCHEMA EVOLUTION INTO $tableNameAsString SELECT * FROM $sourceNameAsString")
+    }
+
+    assert(txn.currentState === Committed)
+    assert(txn.isClosed)
+
+    // the new column must be visible in the committed delegate's schema
+    assert(table.schema.fieldNames.toSeq === Seq("pk", "salary", "dep", "active"))
+
+    checkAnswer(
+      sql(s"SELECT * FROM $tableNameAsString"),
+      Seq(
+        Row(1, 100, "hr", null),      // pre-existing rows: active is null
+        Row(2, 200, "software", null),
+        Row(3, 300, "hr", true),      // inserted with active
+        Row(4, 400, "software", false)))
+  }
+
+  for (isDynamic <- Seq(false, true))
+  test(s"SQL INSERT OVERWRITE WITH SCHEMA EVOLUTION adds new column with transactional checks " +
+      s"isDynamic: $isDynamic") {
+    createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+      """{ "pk": 1, "salary": 100, "dep": "hr" }
+        |{ "pk": 2, "salary": 200, "dep": "software" }
+        |{ "pk": 3, "salary": 300, "dep": "hr" }
+        |""".stripMargin)
+
+    sql(
+      s"""CREATE TABLE $sourceNameAsString
+         |(pk INT NOT NULL, salary INT, dep STRING, active BOOLEAN)""".stripMargin)
+    sql(s"INSERT INTO $sourceNameAsString VALUES (11, 999, 'hr', true), (12, 888, 'hr', false)")
+
+    val insertOverwrite = if (isDynamic) {
+      s"""INSERT WITH SCHEMA EVOLUTION OVERWRITE TABLE $tableNameAsString
+         |SELECT * FROM $sourceNameAsString
+         |""".stripMargin
+    } else {
+      s"""INSERT WITH SCHEMA EVOLUTION OVERWRITE TABLE $tableNameAsString
+         |PARTITION (dep = 'hr')
+         |SELECT pk, salary, active FROM $sourceNameAsString
+         |""".stripMargin
+    }
+
+    val confValue = if (isDynamic) PartitionOverwriteMode.DYNAMIC else PartitionOverwriteMode.STATIC
+    val (txn, _) = withSQLConf(SQLConf.PARTITION_OVERWRITE_MODE.key -> confValue.toString) {
+      executeTransaction { sql(insertOverwrite) }
+    }
+
+    assert(txn.currentState === Committed)
+    assert(txn.isClosed)
+
+    // the new column must be visible in the committed delegate's schema
+    assert(table.schema.fieldNames.contains("active"))
+
+    checkAnswer(
+      sql(s"SELECT * FROM $tableNameAsString"),
+      Seq(
+        Row(2, 200, "software", null), // unchanged (different partition)
+        Row(11, 999, "hr", true),      // overwrote hr partition
+        Row(12, 888, "hr", false)))
+  }
+
+  test("SQL INSERT WITH SCHEMA EVOLUTION analysis failure aborts transaction") {
+    createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+      """{ "pk": 1, "salary": 100, "dep": "hr" }
+        |{ "pk": 2, "salary": 200, "dep": "software" }
+        |""".stripMargin)
+
+    sql(
+      s"""CREATE TABLE $sourceNameAsString
+         |(pk INT NOT NULL, salary INT, dep STRING, active BOOLEAN)""".stripMargin)
+
+    val e = intercept[AnalysisException] {
+      sql(
+        s"""INSERT WITH SCHEMA EVOLUTION INTO $tableNameAsString
+           |SELECT nonexistent_col FROM $sourceNameAsString
+           |""".stripMargin)
+    }
+
+    assert(e.getMessage.contains("nonexistent_col"))
+    assert(catalog.lastTransaction.currentState === Aborted)
+    assert(catalog.lastTransaction.isClosed)
+    // schema must be unchanged after the aborted transaction
+    assert(table.schema.fieldNames.toSeq === Seq("pk", "salary", "dep"))
+  }
 }