dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4354576399
##########
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:
**26. `withTable` does not say that its one result must replace the write
relation and every read of the old table.** Each call builds a new operation
and a new `RowLevelOperationTable`, while scan planning pairs the reads with
the write by reference: `GroupBasedRowLevelOperation` collects the reads with
`r.table eq table` (`patterns.scala:480-483`), and
`GroupBasedRowLevelOperationScanPlanning` replaces only reads of that instance
(`GroupBasedRowLevelOperationScanPlanning.scala:93`).
`V2Writes.validateWriteRelation` compares only column names, so it does not
notice a mismatch.
If a rule swaps only the write relation, as the new test does before running
`V2Writes` alone (`GroupBasedColumnUpdateTableSuite.scala:551-555`), or calls
`withTable` once per relation, the write finds no read relation, so its
copy-on-write scan is never configured and runtime group filtering is skipped.
The in-memory fixture then fails on `scan.data` at commit
(`InMemoryRowLevelOperationTable.scala:728`). In Iceberg,
`CopyOnWriteOperation.overwrittenFiles()` returns an empty set when the scan is
null ([SparkWrite.java
L450-452](https://github.com/apache/iceberg/blob/9303e4c89705f139290ae8c1266fb4485785e279/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkWrite.java#L450-L452)),
so the commit adds the rewritten files without removing the old ones, and the
rows are duplicated silently. Iceberg's `ResolveBranch` builds one instance and
uses it for both the write and the query, which is the right shape. This hazard
is not new with `withTable`, but after item 9 the
PR description points rule authors to this helper.
Could the Scaladoc say that a rule should call `withTable` once and use the
result for the write relation and every read of the old table?
##########
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:
**27.** Minor: this default is new in 22332b3 (the previous commit, 9b6655a,
had `useWriteColumnUpdate: Boolean,`), no caller uses it, and `false` is the
wrong value for every narrow write. The only construction passes the flag by
name (`DataSourceV2Strategy.scala:594`), and `withTransaction` and
`withNewChildInternal` copy it. A construction in the pre-PR six-argument shape
still compiles and sends narrow UPDATE and COPY rows to
`DataWriter.write(...)`, whose layout is the empty `schema()`, and nothing
fails, because the flag cannot be derived from `write`.
`DataAndMetadataWritingSparkTask` and `DataWithProjectionWritingSparkTask`
already require it. Could the default be dropped?
##########
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:
**29. Item 5's fallback covers only `column-update-cow`, so the req-attrs
copy-on-write flavor still commits wrong rows when an UPDATE assigns `pk`.**
L185-186 build a full-width operation when `updatedColumns()` contains `pk`, as
`SupportsColumnUpdates.java:62-67` now asks, but this branch always builds the
narrow operation, and its commit looks the base row up by the new `pk` value
with `find` (L747). With `reqAttrsTableProps("pk,salary")` on the shared rows
(1, 1, hr, 100, 10, 1), (2, 2, hr, 200, 20, 1) and (3, 3, software, 300, 30,
1), `UPDATE t SET pk = 2, salary = 0 WHERE pk = 1` is accepted because `pk` is
declared. The narrow row [2, 0] is then overlaid on the existing `pk = 2` row,
so the commit writes (2, 2, hr, 0, 20, 1) instead of (2, 1, hr, 0, 10, 1). With
duplicate `pk` values, the overlay uses whichever row comes first, while the
delta commit now fails unless exactly one row matches. No test assigns `pk` on
this flavor. Could this branch use the same fallback (or fa
il clearly), and could the copy-on-write commit also require exactly one match?
##########
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:
**30.** Minor: items 14 and 23 made the split commit locate rows by
`rowId()`, but this writer still drops a row ID that is not an `INT` and treats
an `INT` row ID as `pk`. With `no-metadata`, a delete arrives as `delete(null,
rowIdProjection)` (`WriteToDataSourceV2Exec.scala:890`). So on the suite's own
`column-update-split-row-id=_partition` + `no-metadata` table
(`DeltaBasedColumnUpdateTableSuite.scala:865-868`, which today runs only a
rejected UPDATE), `DELETE FROM t WHERE pk = 1` logs `(DELETE, null, null,
null)`, and the null-ID branch (L572-574) calls
`metadataSchema.fieldIndex("_partition")` on an empty schema (L559), which
fails with `FIELD_NOT_FOUND`. A MERGE with `WHEN MATCHED THEN DELETE` fails the
same way. With `column-update-split-row-id=id` on `pk INT NOT NULL, id INT NOT
NULL, dep STRING` and rows (1, 2, hr) and (2, 1, hr), `DELETE ... WHERE pk = 1`
adds 2 to `buffer.deletes`, and `InMemoryBaseTable.scala:416` removes `pk = 2`
instead. Nothing fails today, but the
fixture accepts both configurations. Could the writer keep a copy of the row
ID beside the logged entry (leaving the asserted `id` slot as it is) and
resolve such DELETEs through `currentRowByRowId`, or fail clearly for row-ID
setups it does not support?
##########
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:
**39.** Nit: this overload re-inlines the column conversion (L117), the
default `dep` partitioning (L127) and the `TableInfo` build of
`createTable(columns, transforms)` (L142-149); only
`.withProperties(tableProps)` differs, and about 30 call sites create their
tables through it. `UpdateTableSuiteBase.scala:162-166` likewise rebuilds what
`createTable(columns, Array.empty[Transform])` does. Round 4
([r4074793101](https://github.com/apache/spark/pull/55518#discussion_r4074793101))
suggested a props parameter on the base helper. Could `createTable(columns,
transforms)` take `tableProps: java.util.Map[String, String] =
extraTableProps`, with both `createAndInitTable` overloads and
`UpdateTableSuiteBase` going through it?
##########
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:
**28. `references` walks the `List` behind `query.output` once per read
ordinal, which is quadratic in the query width.** `query.output` is a `List` in
the rewrite shapes (`Project.output` is `projectList.map`, and the relation
output comes from `StructType.map`), so `output.length` (L371) and
`output(ordinal)` (L376) are linear and run once per ordinal. A full-row write
reads every column, so one evaluation takes about 1.5 * W^2 list steps for a
W-column query, e.g. 37.5M at 5,000 columns. This now applies to every DSv2
DELETE, UPDATE and MERGE, not only to column updates: `CheckAnalysis` forces it
through `missingInput`, and the new `ColumnPruning` guard forces it again on
every new `ReplaceData` or `WriteDelta` copy. Master used the cached
`query.outputSet`, and 9b6655a's `query.output.take(lastReadOrdinal + 1)` was
linear. `Optimizer.scala:1541` repeats `output(i)` inside
`output.indices.filter`.
Also, L371 checks only the upper bound, so a `-1` ordinal from a third-party
builder escapes as a raw `IndexOutOfBoundsException` instead of this internal
error.
Could this use `val output = query.output.toIndexedSeq` with
`!output.isDefinedAt(ordinal)`, and could the `ColumnPruning` case use
`w.query.output.zipWithIndex.collect { case (a, i) if w.references.contains(a)
=> i }`?
##########
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:
**36.** Minor: item 18's `ColumnUpdateTableSharedTests` took the tests
listed there, but eight more UPDATE tests are still copied between this suite
and `GroupBasedColumnUpdateTableSuite` (this suite's line / the group suite's
line):
- "write relation has the declared columns in declared order" (L883 / L245)
- "scan reads only the referenced columns" (L915 / L307), which differ only
by the trait's `scanMetadataColumns`
- "analyzed plan shows the declared columns as the write target" (L958 /
L500), which differ only by `write.nodeName`
- "an operation replaced after the rewrite is rejected" (L931 / L511), both
new in 22332b3
- "undeclared, differently-cased columns in the RHS and WHERE clause ..."
(L476 / L598)
- the column update schema test (L653 / L47) and "data correctness" (L665 /
L59)
- the assignment outside `requiredDataAttributes` (L1024 / L469)
That is about 110 lines on each side, every fix has to be made twice, and
`DeltaBasedNoMetadataColumnUpdateTableSuite` runs none of them. This suite also
hard-codes `java.util.Map.of("column-update-req-attrs", ...)` at L884, L916 and
L959 and in `createAndInitTableWithReqAttrs` (L289-295) instead of using its
own `reqAttrsTableProps`. Could these move into the trait, with hooks for the
replacement operations and `reqAttrsOperationClass`?
##########
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:
**40.** Nit: this sits between `createTable(columns)` (L126-129) and the
overload it delegates to (L142-149), so the `createTable` overloads, which were
contiguous before this PR, are split. AGENTS.md asks to "put the new member
with the code it belongs with" (as in item 19). Could it move next to
`extraTableProps` (L94-96) and `noMetadata` (L101-102)?
##########
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:
**34.** Nit: a declaration with no name parts, i.e. a `NamedReference` whose
`fieldNames()` is empty, passes this `length != 1` filter, so it is reported as
a nested field, and since `toSQLId(Seq())` is empty, the message reads
"declared nested field(s) [] in `requiredDataAttributes()`. Declare the root
struct column instead of a nested field." This came with item 17's `toSQLId`
change; the same PR prints `<empty>` for this case in
`undeclaredWriteRequirementColumnsError` (`QueryCompilationErrors.scala:4725`)
and `V2ExpressionUtils.unresolvableColumns`. Could empty declarations be
rejected separately, e.g. as an unknown column shown as `<empty>`, before L390?
##########
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:
**32.** Minor: 22332b3 moved these six `updatedColumns()` tests from
`DeltaBasedUpdateTableSuiteBase` here, so they now run in all eight UPDATE
suites, i.e. 48 table setups and UPDATE jobs instead of 18. The value they
check is computed from the assignments before the operation type matters
(`RewriteUpdateTable.scala:47-48`), so it does not depend on the table kind,
and "no WHERE clause still reports assigned columns" (L96) checks nothing
beyond L56, because the condition is not involved. Could they become one
table-driven test, e.g. `Seq("id = -1" -> Seq("id"), "id = -1, dep = 'eng'" ->
Seq("id", "dep"), ...)`, or stay in one suite?
##########
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:
**38.** Nit: this is `SchemaUtils.foldName` (`SchemaUtils.scala:551-553`),
whose Scaladoc says that `AttributeSeq` looks attributes up by this key, which
is what the duplicate check relies on. Could this be `refs.map(ref =>
SchemaUtils.foldName(ref.fieldNames.head, conf.caseSensitiveAnalysis))`,
dropping the `java.util.Locale` import?
##########
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:
**35.** Nit: the PR description's compatibility paragraph says that
`RowLevelWrite` has a new abstract `withPrunedQuery`, but this
`queryProjections` is abstract too and drives the new `references`, so an
extension that adds only `withPrunedQuery` still fails to compile. Since the
description becomes the commit message, could it name both?
##########
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:
**31.** Minor: this checks `attribute` and `method`, but not
`relationOutput`, which is the pruned scan output this test is about, nor
`scanClass`, so a change in which other columns the narrow scan reads would go
unnoticed. Similarly, the three CHECK tests
(`ColumnUpdateTableSharedTests.scala:165-169` and `:184-188`, and
`DeltaBasedColumnUpdateTableSuite.scala:1094-1098`) check only the
`CHECK_CONSTRAINT_VIOLATION` condition, so they cannot tell which constraint or
value failed. `DataSourceV2CatalystRuntimeFilterSuite` and
`CheckConstraintSuite` use `checkError` with every parameter for these
conditions. Could these use `checkError`, or `checkErrorMatchPVals` for the
scan class?
##########
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:
**37.** Nit: after item 25, `column-update` and `column-update-req-attrs`
with `split-updates` (L163-179) build the same operations with the same
arguments as `column-update-split` and `column-update-split-req-attrs` here, so
these keys are now synonyms, and the req-attrs plus `split-updates` branch is
never exercised. As a result, "column-update: split-updates sends a narrow
UPDATE as delete and reinsert" (`DeltaBasedColumnUpdateTableSuite.scala:718`)
and "column-update split: delete and reinsert carry the row ID and preserved
metadata" (`:735`) run the same DDL, data, UPDATE and write-log assertions.
Likewise, `column-update-cow-no-write-column-update` is still a dispatch key
(L195-198), although item 18 made it a writer flag. Could the old
`column-update-split` and `column-update-split-req-attrs` keys go (keeping
`column-update-split-row-id`), with the split tests spelled as `column-update`
plus `split-updates`, and one of the two tests dropped?
##########
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:
**33.** Nit: this brings back a hand-written `lastUpdatedColumns` assertion,
which round 4
([r4074793101](https://github.com/apache/spark/pull/55518#discussion_r4074793101))
asked to replace with `checkLastUpdatedColumns`. L549 resets the field to
`null` precisely to catch `withTable` not calling the builder, and in that case
this line fails with a `NullPointerException`, while the helper fails with "no
row-level operation was built" on the same live table. Could this be
`checkLastUpdatedColumns("salary")`?
--
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]
