anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4364254863
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -380,14 +380,17 @@ case class ReplaceDataExec(
write: Write,
rowLevelCommand: RowLevelOperation.Command,
tableName: String,
+ useWriteColumnUpdate: Boolean = false,
Review Comment:
Done in d8a4641: `useWriteColumnUpdate` has no default again. The only
construction, in `DataSourceV2Strategy`, passes it by name.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -356,6 +357,32 @@ trait RowLevelWrite extends V2WriteCommand with
SupportsSubquery {
def condition: Expression
def originalTable: NamedRelation
+ /** The projections that read query columns by ordinal. */
+ protected def queryProjections: Seq[ProjectingInternalRow]
+
+ /**
+ * The write reads the operation column at ordinal 0 and other query columns
only through
+ * projections, by ordinal. Any other query column can be pruned with
[[withPrunedQuery]].
+ */
+ override lazy val references: AttributeSet = {
+ val output = query.output
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -321,23 +810,61 @@ class InMemoryRowLevelOperationTable private (
}
}
-private class DeltaBufferedRowsWriterFactory(schema: StructType) extends
DeltaWriterFactory {
+private class DeltaBufferedRowsWriterFactory(
+ schema: StructType,
+ intRowId: Boolean = true) extends DeltaWriterFactory {
override def createWriter(partitionId: Int, taskId: Long):
DeltaWriter[InternalRow] = {
- new DeltaBufferWriter(schema)
+ new DeltaBufferWriter(schema, intRowId)
+ }
+}
+
+// Test-only: a writer factory for a connector that never overrides
writeColumnUpdate(), to
+// exercise the DataWriter#writeColumnUpdate default that throws
+// DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED.
+// Declared top-level (no reference to the enclosing table) so it stays
task-serializable.
+private class NoWriteColumnUpdateOverrideWriterFactory extends
DataWriterFactory {
+ override def createWriter(partitionId: Int, taskId: Long):
DataWriter[InternalRow] = {
+ new NoWriteColumnUpdateOverrideWriter
}
}
-private class DeltaBufferWriter(schema: StructType) extends
BufferWriter(schema)
- with DeltaWriter[InternalRow] {
+private class NoWriteColumnUpdateOverrideWriter extends
DataWriter[InternalRow] {
+ override def write(record: InternalRow): Unit = {}
+ override def commit(): WriterCommitMessage = new BufferedRows(Seq.empty,
StructType(Nil))
+ override def abort(): Unit = {}
+ override def close(): Unit = {}
+}
+
+// Test-only: a writer that overrides writeColumnUpdate(record) but not
+// writeColumnUpdate(metadata, record), so rows with metadata reach it through
the default.
+private class RecordOnlyColumnUpdateWriterFactory(schema: StructType) extends
DataWriterFactory {
+ override def createWriter(partitionId: Int, taskId: Long):
DataWriter[InternalRow] = {
+ new RecordOnlyColumnUpdateWriter(schema)
+ }
+}
+
+private class RecordOnlyColumnUpdateWriter(schema: StructType) extends
DataWriter[InternalRow] {
+ private val delegate = new BufferWriter(schema)
+ override def write(record: InternalRow): Unit = delegate.write(record)
+ override def writeColumnUpdate(record: InternalRow): Unit =
delegate.writeColumnUpdate(record)
+ override def commit(): WriterCommitMessage = delegate.commit()
+ override def abort(): Unit = delegate.abort()
+ override def close(): Unit = delegate.close()
+}
+
+// Logs the row ID of a DELETE as an integer when `intRowId` is set and as
null otherwise; the
+// commit then resolves such a DELETE from its metadata.
+private class DeltaBufferWriter(schema: StructType, intRowId: Boolean = true)
+ extends BufferWriter(schema) with DeltaWriter[InternalRow] {
private final val DELETE = UTF8String.fromString(Delete.toString)
private final val UPDATE = UTF8String.fromString(Update.toString)
private final val REINSERT = UTF8String.fromString(Reinsert.toString)
private final val INSERT = UTF8String.fromString(Insert.toString)
override def delete(meta: InternalRow, id: InternalRow): Unit = {
- val pk = id.getInt(0)
- buffer.deletes += pk
+ val pk: Any = if (intRowId) id.getInt(0) else null
Review Comment:
Done in d8a4641
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -356,6 +357,32 @@ trait RowLevelWrite extends V2WriteCommand with
SupportsSubquery {
def condition: Expression
def originalTable: NamedRelation
+ /** The projections that read query columns by ordinal. */
+ protected def queryProjections: Seq[ProjectingInternalRow]
Review Comment:
Done: the compatibility paragraph now says "`RowLevelWrite` has new abstract
`queryProjections` and `withPrunedQuery`".
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -128,12 +153,61 @@ class InMemoryRowLevelOperationTable private (
copied.replacedPartitions = replacedPartitions
copied.lastWriteInfo = lastWriteInfo
copied.lastWriteLog = lastWriteLog
+ copied.lastUpdatedColumns = lastUpdatedColumns
copied
}
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
- if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
+ lastUpdatedColumns = info.updatedColumns()
+ val splitUpdates = properties.getOrDefault(SPLIT_UPDATES,
"false").toBoolean
+ if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") {
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) {
+ val reqCols =
properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim)
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command,
reqCols, info.options)
+ }
+ } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false")
== "true") {
+ // Test-only: returns an empty requiredDataAttributes() so we can verify
Spark rejects it.
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(
+ info.command, Array.empty, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_COW, "false") == "true") {
+ if (info.updatedColumns().exists(_.describe() == "pk")) {
+ () => PartitionBasedOperation(info.command, info.options)
+ } else {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_COW_REQ_ATTRS)) {
+ val reqCols =
properties.get(COLUMN_UPDATE_COW_REQ_ATTRS).split(",").map(_.trim)
+ () => new PartitionBasedColumnUpdateOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else if (
+ properties.getOrDefault(COLUMN_UPDATE_COW_NO_WRITE_COLUMN_UPDATE,
"false") == "true") {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_SPLIT, "false") ==
"true") {
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -343,3 +365,60 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Resolves the connector-declared required data attributes for
column-update writes against
+ * the given relation, in declared order and with the relation's spelling.
An empty declaration,
+ * nested fields, duplicate columns, columns that do not exist in the
relation and metadata
+ * columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ if (refs.isEmpty) {
+ throw
QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.fieldNames.toSeq)
Review Comment:
Done in d8a4641.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,637 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector
+
+import org.apache.spark.{SparkException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTable
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.connector.write.RowLevelOperationTable
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] =
{
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String):
java.util.Map[String, String] = {
+ java.util.Map.of("column-update-cow-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] =
java.util.Map.of()
+
+ test("column-update ReplaceData: column update schema has only the declared
columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update ReplaceData: data correctness -- bonus preserved, salary
updated") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ // CoW: pk=1 and pk=3 match (2 updates); the 'hr' partition has no other
rows so no copies.
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: subquery WHERE condition data correctness")
{
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ withTempView("target_deps") {
+ import testImplicits._
+ val subqueryDF = Seq("hr").toDF()
+ subqueryDF.createOrReplaceTempView("target_deps")
+
+ sql(
+ s"""UPDATE $tableNameAsString
+ |SET salary = -1
+ |WHERE dep IN (SELECT * FROM target_deps)
+ |""".stripMargin)
+ }
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: runtime group filtering data correctness") {
+ Seq(true, false).foreach { dppEnabled =>
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key ->
dppEnabled.toString,
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
+ withTable(tableNameAsString) {
+ withTempView("matched_pk") {
+ createAndInitTable("pk INT, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ import testImplicits._
+ val matchedPkDF = Seq(Some(1), None).toDF()
+ matchedPkDF.createOrReplaceTempView("matched_pk")
+
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = -1 " +
+ s"WHERE pk IN (SELECT * FROM matched_pk)",
+ primaryScanSchema =
+ "pk INT, salary INT, dep STRING, _partition STRING, index
INT",
+ groupFilterScanSchema = Some("pk INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, 300, 30, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update ReplaceData: updated and copied rows reach
writeColumnUpdate with " +
+ "metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", null), data =
Row(1, "hr", -1)),
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", 1), data =
Row(2, "hr", 200)))
+ }
+
+ test("column-update ReplaceData: without metadata, rows reach
writeColumnUpdate(record)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Row(3, 300, 30,
"software") :: Nil)
+ }
+
+ test("column-update ReplaceData: a writer overriding only
writeColumnUpdate(record) " +
+ "receives rows with metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true",
"column-update-cow-record-only-writer", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: without metadata, column update schema is
the declared " +
+ "layout") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedColumnUpdateSchema = Some(StructType(Seq(PK_FIELD, DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)))))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: connector missing writeColumnUpdate
override is rejected") {
+ createAndInitTableReplaceDataNoWriteColumnUpdate("pk INT NOT NULL, salary
INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
+ parameters = Map(
+ "class" ->
"org.apache.spark.sql.connector.catalog.NoWriteColumnUpdateOverrideWriter"))
+ }
+
+ private def createAndInitTableReplaceDataNoWriteColumnUpdate(
+ schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-cow-no-write-column-update", "true"))
+ }
+
+ //
---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to
the
+ // declared columns, on both copy-on-write plan shapes.
+ //
---------------------------------------------------------------------------
+
+ private val reqAttrsOperationClass =
+
classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private def createAndInitTableWithReqAttrs(reqAttrs: String): Unit = {
+ createAndInitTableWithProps(reqAttrsTableProps(reqAttrs))
+ }
+
+ test("column-update ReplaceData: write relation has the declared columns in
declared order") {
+ createAndInitTableWithReqAttrs("salary,pk")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep =
'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery write relation has the declared
columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)")
+ }
+
+ assert(write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: runtime group filtering on an undeclared
condition column") {
+ withTempView("matched_ids") {
+ // `dep` is declared because the scan filters groups at runtime by the
partition column
+ createAndInitTableWithReqAttrs("pk,dep,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `id` is neither declared nor in the write relation; the group filter
must still map it
+ // onto the second read of the table that the subquery rewrite adds
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, dep STRING, salary INT, bonus
INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = Some("id INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+
+ test("column-update ReplaceData: scan reads only the referenced columns") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are
read for the
+ // condition, the assignment and the CHECK constraint
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep =
'hr'",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition
STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery scan reads only the referenced
columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition
STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: unread partition column skips runtime group
filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `dep` is neither declared nor referenced, so the scan doesn't read it
and doesn't
+ // report it from filterAttributes(); every group is read and replaced
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, salary INT, bonus INT, extra INT,
" +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = None)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr", "software"))
+ }
+ }
+
+ test("column-update ReplaceData: runtime filter attribute the scan doesn't
read is rejected") {
+ withTempView("matched_ids") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100 }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow-req-attrs", "pk,salary",
+ "column-update-cow-unread-filter-attrs", "true"))
+ import testImplicits._
+ Seq(1).toDF().createOrReplaceTempView("matched_ids")
+
+ val ex = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT *
FROM matched_ids)")
+ }
+ assert(ex.getCondition ==
"DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE")
Review Comment:
Done in d8a4641.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,637 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector
+
+import org.apache.spark.{SparkException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTable
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.connector.write.RowLevelOperationTable
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] =
{
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String):
java.util.Map[String, String] = {
+ java.util.Map.of("column-update-cow-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] =
java.util.Map.of()
+
+ test("column-update ReplaceData: column update schema has only the declared
columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update ReplaceData: data correctness -- bonus preserved, salary
updated") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ // CoW: pk=1 and pk=3 match (2 updates); the 'hr' partition has no other
rows so no copies.
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: subquery WHERE condition data correctness")
{
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ withTempView("target_deps") {
+ import testImplicits._
+ val subqueryDF = Seq("hr").toDF()
+ subqueryDF.createOrReplaceTempView("target_deps")
+
+ sql(
+ s"""UPDATE $tableNameAsString
+ |SET salary = -1
+ |WHERE dep IN (SELECT * FROM target_deps)
+ |""".stripMargin)
+ }
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: runtime group filtering data correctness") {
+ Seq(true, false).foreach { dppEnabled =>
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key ->
dppEnabled.toString,
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
+ withTable(tableNameAsString) {
+ withTempView("matched_pk") {
+ createAndInitTable("pk INT, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ import testImplicits._
+ val matchedPkDF = Seq(Some(1), None).toDF()
+ matchedPkDF.createOrReplaceTempView("matched_pk")
+
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = -1 " +
+ s"WHERE pk IN (SELECT * FROM matched_pk)",
+ primaryScanSchema =
+ "pk INT, salary INT, dep STRING, _partition STRING, index
INT",
+ groupFilterScanSchema = Some("pk INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, 300, 30, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update ReplaceData: updated and copied rows reach
writeColumnUpdate with " +
+ "metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", null), data =
Row(1, "hr", -1)),
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", 1), data =
Row(2, "hr", 200)))
+ }
+
+ test("column-update ReplaceData: without metadata, rows reach
writeColumnUpdate(record)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Row(3, 300, 30,
"software") :: Nil)
+ }
+
+ test("column-update ReplaceData: a writer overriding only
writeColumnUpdate(record) " +
+ "receives rows with metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true",
"column-update-cow-record-only-writer", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: without metadata, column update schema is
the declared " +
+ "layout") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedColumnUpdateSchema = Some(StructType(Seq(PK_FIELD, DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)))))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: connector missing writeColumnUpdate
override is rejected") {
+ createAndInitTableReplaceDataNoWriteColumnUpdate("pk INT NOT NULL, salary
INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
+ parameters = Map(
+ "class" ->
"org.apache.spark.sql.connector.catalog.NoWriteColumnUpdateOverrideWriter"))
+ }
+
+ private def createAndInitTableReplaceDataNoWriteColumnUpdate(
+ schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-cow-no-write-column-update", "true"))
+ }
+
+ //
---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to
the
+ // declared columns, on both copy-on-write plan shapes.
+ //
---------------------------------------------------------------------------
+
+ private val reqAttrsOperationClass =
+
classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private def createAndInitTableWithReqAttrs(reqAttrs: String): Unit = {
+ createAndInitTableWithProps(reqAttrsTableProps(reqAttrs))
+ }
+
+ test("column-update ReplaceData: write relation has the declared columns in
declared order") {
+ createAndInitTableWithReqAttrs("salary,pk")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep =
'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery write relation has the declared
columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)")
+ }
+
+ assert(write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: runtime group filtering on an undeclared
condition column") {
+ withTempView("matched_ids") {
+ // `dep` is declared because the scan filters groups at runtime by the
partition column
+ createAndInitTableWithReqAttrs("pk,dep,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `id` is neither declared nor in the write relation; the group filter
must still map it
+ // onto the second read of the table that the subquery rewrite adds
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, dep STRING, salary INT, bonus
INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = Some("id INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+
+ test("column-update ReplaceData: scan reads only the referenced columns") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are
read for the
+ // condition, the assignment and the CHECK constraint
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep =
'hr'",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition
STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery scan reads only the referenced
columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition
STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: unread partition column skips runtime group
filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `dep` is neither declared nor referenced, so the scan doesn't read it
and doesn't
+ // report it from filterAttributes(); every group is read and replaced
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, salary INT, bonus INT, extra INT,
" +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = None)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr", "software"))
+ }
+ }
+
+ test("column-update ReplaceData: runtime filter attribute the scan doesn't
read is rejected") {
+ withTempView("matched_ids") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100 }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow-req-attrs", "pk,salary",
+ "column-update-cow-unread-filter-attrs", "true"))
+ import testImplicits._
+ Seq(1).toDF().createOrReplaceTempView("matched_ids")
+
+ val ex = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT *
FROM matched_ids)")
+ }
+ assert(ex.getCondition ==
"DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE")
+ assert(ex.getMessageParameters.get("attribute") == "`dep`")
+ assert(ex.getMessageParameters.get("method") == "filterAttributes()")
+ }
+ }
+
+ test("column-update ReplaceData: a connector without column updates reads
every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of[String, String]())
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = 0",
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT")
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 0, "hr") ::
Nil)
+ }
+
+ test("column-update ReplaceData: undeclared distribution column is
rejected") {
+ // without metadata, the write clusters by `dep`
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow-req-attrs", "pk,salary",
"no-metadata", "true"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
ColumnPruning.ruleName
+ for {
+ cond <- Seq("pk = 1", "dep = 'hr'")
+ confs <- Seq(Nil, Seq(excludeColumnPruning))
+ } withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE $cond")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" ->
"[`dep`]"))
+ }
+ }
+
+ test("column-update ReplaceData: other operations read every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of())
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING,
_partition STRING, index INT"
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1",
fullScanSchema)
+
+ executeAndCheckScan(s"DELETE FROM $tableNameAsString WHERE pk > 1",
fullScanSchema)
+
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150)).toDF("pk", "salary").createOrReplaceTempView("source")
+ executeAndCheckScan(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |""".stripMargin,
+ fullScanSchema)
+ }
+
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 150, 10,
"hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: DELETE and MERGE on a column-update table
read every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin)
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING,
_partition STRING, index INT"
+
+ executeAndCheckScan(s"DELETE FROM $tableNameAsString WHERE pk > 1",
fullScanSchema)
+
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150)).toDF("pk", "salary").createOrReplaceTempView("source")
+ executeAndCheckScan(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |""".stripMargin,
+ fullScanSchema)
+ }
+
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 150, 10,
"hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: assigned column outside
requiredDataAttributes is rejected") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep = 'hr'")
+ },
+ condition =
"COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "missingColumns" -> "[`bonus`]"))
+ }
+
+ test("column-update ReplaceData: assigned column outside
requiredDataAttributes is rejected " +
+ "on the subquery path") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep IN (SELECT
dep FROM deps)")
+ },
+ condition =
"COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "missingColumns" -> "[`bonus`]"))
+ }
+ }
+
+ test("column-update ReplaceData: analyzed plan shows the declared columns as
the write target") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep =
'hr'")
+ }
+
+ val header = write.simpleString(Int.MaxValue)
+ assert(header.matches("""ReplaceData RelationV2\[pk#\d+, salary#\d+\]
.*"""), header)
+ }
+
+ test("column-update ReplaceData: an operation replaced after the rewrite is
rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = table
+ Seq(
+ t.PartitionBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) ->
"[pk, salary, dep]",
+ new t.PartitionBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]"
+ ).foreach { case (operation, expectedNames) =>
+ val replaced = write.withNewTable(
+ relation.copy(table = operationTable.copy(operation = operation)(
+ operationTable.operationInfo)))
+ checkError(
+ exception = intercept[SparkException](V2Writes(replaced)),
+ condition = "INTERNAL_ERROR",
+ parameters = Map("message" ->
+ (s"The write relation of operation ${operation.getClass.getName} has
columns " +
+ s"[pk, dep, salary], but the operation expects $expectedNames.")))
+ }
+ }
+
+ test("column-update ReplaceData: an operation rebuilt with withTable keeps
the updated columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = table
+ t.lastUpdatedColumns = null
+
+ val rebuilt = operationTable.withTable(t)
+
+ assert(t.lastUpdatedColumns.map(_.describe()).toSeq == Seq("salary"))
Review Comment:
Done in d8a4641: it is `checkLastUpdatedColumns("salary")` now.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala:
##########
@@ -53,6 +53,68 @@ abstract class UpdateTableSuiteBase extends
RowLevelOperationSuiteBase {
s"Expected numCopiedRows=$expectedCopied, got
${summary.numCopiedRows()}")
}
+ test("SPARK-58111: updatedColumns: single non-identity assignment") {
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala:
##########
@@ -58,4 +63,13 @@ private[sql] case class RowLevelOperationTable(
override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
operation.newWriteBuilder(info)
}
+
+ /**
+ * Builds the operation for another table from the same info, for a rule
that redirects the
+ * write after the rewrite, such as to a branch of the table.
+ */
+ def withTable(newTable: Table with SupportsRowLevelOperations):
RowLevelOperationTable = {
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -128,12 +153,61 @@ class InMemoryRowLevelOperationTable private (
copied.replacedPartitions = replacedPartitions
copied.lastWriteInfo = lastWriteInfo
copied.lastWriteLog = lastWriteLog
+ copied.lastUpdatedColumns = lastUpdatedColumns
copied
}
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
- if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
+ lastUpdatedColumns = info.updatedColumns()
+ val splitUpdates = properties.getOrDefault(SPLIT_UPDATES,
"false").toBoolean
+ if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") {
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) {
+ val reqCols =
properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim)
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command,
reqCols, info.options)
+ }
+ } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false")
== "true") {
+ // Test-only: returns an empty requiredDataAttributes() so we can verify
Spark rejects it.
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(
+ info.command, Array.empty, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_COW, "false") == "true") {
+ if (info.updatedColumns().exists(_.describe() == "pk")) {
+ () => PartitionBasedOperation(info.command, info.options)
+ } else {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_COW_REQ_ATTRS)) {
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -343,3 +365,60 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Resolves the connector-declared required data attributes for
column-update writes against
+ * the given relation, in declared order and with the relation's spelling.
An empty declaration,
+ * nested fields, duplicate columns, columns that do not exist in the
relation and metadata
+ * columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ if (refs.isEmpty) {
+ throw
QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.fieldNames.toSeq)
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames = refs.map { ref =>
Review Comment:
Done in d8a4641: it is `SchemaUtils.foldName(ref.fieldNames.head,
conf.caseSensitiveAnalysis)` now, and the `Locale` import is gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -141,6 +153,19 @@ abstract class RowLevelOperationSuiteBase
append(schemaString, jsonData)
}
+ protected def createAndInitTable(
Review Comment:
Done in d8a4641.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,1114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector
+
+import org.apache.spark.{SparkException, SparkRuntimeException,
SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import
org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable,
InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.write.{LogicalWriteInfo,
RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField,
StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the
declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] =
{
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String):
java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = {
+ java.util.Map.of("supports-deltas", "true")
+ }
+
+ override protected def scanMetadataColumns: String = ", _partition STRING"
+
+ private val reqAttrsOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private val splitOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitReqAttrsOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+
+ private val splitWithRowIdOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ test("column-update: column update schema contains a single assigned
column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema contains multiple assigned
columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk
= 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false),
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema is pk and dep for a full identity
update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty; pk and dep
remain because the
+ // connector declares them as base columns.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Array(PK_FIELD,
DEP_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column
narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1,
3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ 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"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: column update schema excludes identity assignments in a
mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk
= 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: CHAR and VARCHAR self-assignments are reported as
updated") {
+ createTable("pk INT NOT NULL, id INT, c CHAR(3), v VARCHAR(10), dep
STRING")
+ sql(s"INSERT INTO $tableNameAsString VALUES (1, 1, 'a', 'x', 'hr')")
+
+ sql(s"UPDATE $tableNameAsString SET id = id, c = c, v = v WHERE pk = 1")
+
+ checkLastUpdatedColumns("c", "v")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 1, "a ", "x", "hr") :: Nil)
+ }
+
+ test("column-update: a LogicalWriteInfo that does not implement
columnUpdateSchema throws") {
+ val info = new LogicalWriteInfo {
+ override def options(): CaseInsensitiveStringMap =
CaseInsensitiveStringMap.empty()
+ override def queryId(): String = "query"
+ override def schema(): StructType = new StructType()
+ }
+
+ checkError(
+ exception =
intercept[SparkUnsupportedOperationException](info.columnUpdateSchema()),
+ condition = "DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED",
+ parameters = Map("class" -> info.getClass.getName))
+ }
+
+ test("column-update: updatedColumns is empty for DELETE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ checkLastUpdatedColumns()
+ }
+
+ test("column-update split: row-level DELETE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split", "true"))
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "hr"))
+ }
+
+ test("column-update split: row-level DELETE by a metadata row ID falls back
to full-width " +
+ "rows") {
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition"))
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "software"))
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split", "true"))
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ 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"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- mixed identity and real
assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk
= 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") ::
Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-req-attrs", reqAttrs))
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep
STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: undeclared distribution column is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
ColumnPruning.ruleName
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" ->
"[`dep`]"))
+ }
+ }
+ }
+
+ test("column-update split: undeclared distribution column is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-split-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> splitReqAttrsOperationClass, "columns"
-> "[`dep`]"))
+ }
+
+ test("column-update: undeclared write requirement columns are listed once
each, sorted") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 1, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "id,dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" ->
"[`dep`, `id`]"))
+ }
+
+ test("column-update: a write requirement on a missing nested field reports
why") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1 }, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,s",
+ "column-update-cluster-by", "s.missing"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "columns" -> "[`s`.`missing` (FIELD_NOT_FOUND)]"))
+ }
+
+ test("column-update: an empty write requirement column name is reported as
<empty>") {
+ checkError(
+ exception =
QueryCompilationErrors.undeclaredWriteRequirementColumnsError(
+ reqAttrsOperationClass, Seq(Seq.empty[String] ->
None)).asInstanceOf[AnalysisException],
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" ->
"[<empty>]"))
+ }
+
+ test("column-update: undeclared column inside nested transforms is
rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep",
+ "column-update-cluster-by-nested-transform", "true"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
ColumnPruning.ruleName
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" ->
"[`dep`]"))
+ }
+ }
+ }
+
+ test("column-update: distribution by an undeclared row ID column is
allowed") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: distribution by an undeclared row ID column assigned to
itself is " +
+ "allowed") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1, pk = pk WHERE pk = 1")
+
+ checkLastUpdatedColumns("salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: partition source column not in requiredDataAttributes")
{
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep
STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScanExcludes(s"UPDATE $tableNameAsString SET id = -1 WHERE
pk = 1", "dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on
the row-ID " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("PK,salary,dep", "pk INT NOT NULL, salary
INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: undeclared, differently-cased columns in the RHS and
WHERE clause " +
+ "are read with the table's own spelling") {
+ createAndInitTableWithReqAttrs("pk,dep,salary",
+ "pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND
EXTRA > 3",
+ "pk INT, salary INT, bonus INT, extra INT, dep STRING, _partition
STRING")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 110, 10, 5, "hr") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is
unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on
an assigned " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("pk,SALARY,dep", "pk INT NOT NULL, salary
INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ
only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "SALARY", "dep"))
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200,
"hr") :: Nil)
+ }
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException")
{
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin,
+ java.util.Map.of("column-update-empty-req-attrs", "true"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ parameters = Map("connector" -> reqAttrsOperationClass))
+ }
+
+ test("column-update: nested requiredDataAttributes throws
AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,s.c1,dep",
+ "pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass,
"nestedAttributes" -> "[`s`.`c1`]"))
+ }
+
+ test("column-update: duplicate requiredDataAttributes throws
AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,pk,salary", "pk INT NOT NULL, salary
INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass,
"duplicateAttributes" -> "[`pk`]"))
+ }
+
+ test("column-update: case-only duplicate requiredDataAttributes throws
AnalysisException") {
+ createAndInitTableWithReqAttrs("PK,pk,salary", "pk INT NOT NULL, salary
INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass,
"duplicateAttributes" -> "[`PK`]"))
+ }
+
+ test("column-update: duplicate requiredDataAttributes are reported in
declared order") {
+ createAndInitTableWithReqAttrs("salary,pk,SALARY,PK", "pk INT NOT NULL,
salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass, "duplicateAttributes" ->
"[`salary`, `pk`]"))
+ }
+
+ test("column-update: metadata column in requiredDataAttributes throws
AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,_partition,dep", "pk INT NOT NULL,
salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "metadataAttributes" -> "[`_partition`]"))
+ }
+
+ test("column-update split: metadata column in requiredDataAttributes throws
" +
+ "AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-req-attrs", "pk,_partition,dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> splitReqAttrsOperationClass,
+ "metadataAttributes" -> "[`_partition`]"))
+ }
+
+ test("column-update: requiredDataAttributes rejects a column that does not
exist") {
+ createAndInitTableWithReqAttrs("pk,nonexistent_col,id", "pk INT NOT NULL,
id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "unknownAttributes" -> "[`nonexistent_col`]"))
+ }
+
+ test("column-update: column update schema excludes undeclared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update: data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String):
Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split", "true"))
+ }
+
+ test("column-update split: column update schema has the declared columns") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ test("column-update: split-updates sends a narrow UPDATE as delete and
reinsert") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update", "true", "split-updates", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ deleteWriteLogEntry(id = 1, metadata = Row("hr", null)),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(1, "hr",
-1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "hr") :: Nil)
+ }
+
+ test("column-update split: delete and reinsert carry the row ID and
preserved metadata") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ deleteWriteLogEntry(id = 1, metadata = Row("hr", null)),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(1, "hr",
-1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected")
{
+ // `extra` is neither declared, updated, nor referenced, so
requiredDataAttributes()
+ // ([pk, dep, salary]) doesn't cover every table column -- the write stays
genuinely narrow
+ // and reassigning the row ID must be rejected.
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra
STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE
dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ parameters = Map("connector" -> splitOperationClass, "rowIds" ->
"[`pk`]"))
+ }
+
+ test("column-update split: row-ID reassignment is allowed when
requiredDataAttributes " +
+ "covers every table column") {
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep =
'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment is allowed when
requiredDataAttributes " +
+ "covers every data column of a relation with metadata columns") {
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ // `_partition` and the coerced `salary > 0.5` make the analyzer add the
metadata columns to
+ // the target relation in a later iteration. Only analyzes the plan.
+ val parsed = spark.sessionState.sqlParser.parsePlan(
+ s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 " +
+ s"WHERE dep = 'hr' AND _partition = 'hr' AND salary > 0.5")
+ val write = spark.sessionState.executePlan(parsed).analyzed.collectFirst {
+ case w: WriteDelta => w
+ }.getOrElse(fail("couldn't find WriteDelta in analyzed plan"))
+
+ assert(write.table.output.map(_.name) == Seq("pk", "salary", "dep"))
+ }
+
+ private def createAndInitTableSplitWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split-req-attrs", reqAttrs))
+ }
+
+ test("column-update split: undeclared row-ID column on narrow write is
rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "pk"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds"
-> "[`pk`]"))
+ }
+
+ test("column-update split: metadata row ID preserved on reinsert need not be
declared") {
+ // `_partition` is a required metadata attribute with
PRESERVE_ON_REINSERT, so the REINSERT
+ // row carries it through the metadata projection.
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
Seq("salary"))
+ checkLastColumnUpdateWriteLog(
+ WriteLogEntry(operation = Delete, metadata = Some(Row("hr", null))),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(-1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update split: metadata row ID not preserved on reinsert must be
declared") {
+ // `index` is a required metadata attribute, but its REINSERT value is
null.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "index"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds"
-> "[`index`]"))
+ }
+
+ test("column-update split: metadata row ID outside
requiredMetadataAttributes must be " +
+ "declared") {
+ // Without required metadata attributes, nothing projects `_partition`
into the REINSERT row.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition",
"no-metadata", "true"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds"
-> "[`_partition`]"))
+ }
+
+ //
---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to
the
+ // declared columns, for in-place and delete+insert deltas.
+ //
---------------------------------------------------------------------------
+
+ test("column-update: write relation has the declared columns in declared
order") {
Review Comment:
Done in d8a4641. All eight tests moved into `ColumnUpdateTableSharedTests`,
so the no-metadata suite runs them too. The trait takes the table properties
through `reqAttrsTableProps` and `fullWidthTableProps`, and the scan's metadata
columns through `scanMetadataColumns`. It picks the write node name, the
replacement operations and the req-attrs operation class from `deltaUpdate`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -127,6 +128,17 @@ abstract class RowLevelOperationSuiteBase
createTable(columns, transforms)
}
+ // `extraTableProps` as the TBLPROPERTIES clause of a CREATE TABLE statement
+ protected def extraTablePropertiesClause: String = {
Review Comment:
Done in d8a4641: it sits next to `extraTableProps` and `noMetadata` now, as
a pure move, and the `createTable` overloads are contiguous again.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
