dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4372866939
##########
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,
Review Comment:
**42.** Minor: after item 27, this flag is the only record of whether UPDATE
and COPY rows follow `columnUpdateSchema()`, and it cannot be derived from
`write` or `projections`. Plugins that replace `ReplaceDataExec` rebuild their
node from the other fields, so they compile unchanged and drop it:
- NVIDIA/cudf-spark builds `GpuReplaceDataExec(childPlan,
cpuExec.refreshCache, cpuExec.projections, gpuWrite)`
([ReplaceDataExecShim.scala
L41](https://github.com/NVIDIA/cudf-spark/blob/a97083ad4c47cd35d894c7b5d91a5899fce748b5/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/ReplaceDataExecShim.scala#L41),
whose shim list includes 420 and 500), and its task calls
`writer.write(metadata, projected)` or `writer.write(projected)`
([GpuReplaceDataExec.scala
L87-L89](https://github.com/NVIDIA/cudf-spark/blob/a97083ad4c47cd35d894c7b5d91a5899fce748b5/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala#L87-L89)).
- Gluten builds `VeloxIcebergReplaceDataExec(original.query,
original.refreshCache, original.write)` ([VeloxIcebergReplaceDataExec.scala
L31-L37](https://github.com/apache/gluten/blob/ffb8af3fedfc5c9c74244c83a4c33adb87caccd8/backends-velox/src-iceberg/main/scala/org/apache/gluten/execution/VeloxIcebergReplaceDataExec.scala#L31-L37)).
Nothing is affected today, but once a connector they offload, such as
Iceberg copy-on-write, mixes in `SupportsColumnUpdates`, narrow UPDATE and COPY
rows would reach `write(...)` instead of `writeColumnUpdate(...)`, without any
error. Could the Scaladoc of this parameter say that a node replacing this one
must send UPDATE and COPY rows to `writeColumnUpdate` when it is set, or not
replace it, and could the compatibility paragraph of the PR description say the
same?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data
column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name =>
buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's
values (or the
+ // existence default for a column added after the base row was written),
overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by
the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]):
Array[NamedReference] = {
+ val base = Seq("pk",
"dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r =>
baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends
only the declared
+ // columns in the row projection instead of the full row schema. It declares
`pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever
columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
Review Comment:
**46.** Nit: this extends the case class `DeltaBasedOperation` (L275), so
this class and its subclasses (L507, L515, L597, L608) inherit `equals`,
`hashCode` and `copy` on `(command, options)`. For the same table, `new
DeltaBasedColumnUpdateOperation(UPDATE)` equals `DeltaBasedOperation(UPDATE,
CaseInsensitiveStringMap.empty())`, which is exactly the pair at
`ColumnUpdateTableSharedTests.scala:157-158`, and a split and a non-split
operation, or two with different `requiredDataAttributes()`, are equal too.
`RowLevelOperationTable` and the relation then compare equal as well, so a
future test that swaps one of these operations in with `plan.transform` would
keep the old node (`fastEquals`), and a `Set` or `Map` of them, unlike the
`Seq` at `ColumnUpdateTableSharedTests.scala:155`, would collapse them. No
current test is affected. Could this class override `equals` and `hashCode`
with reference equality?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,105 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can
implement this interface to
+ * receive a narrow row containing only the columns declared via {@link
#requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated,
copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in
{@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}.
Inserted rows always follow
+ * {@link LogicalWriteInfo#schema()} and arrive through {@link
DataWriter#write(Object)} and
+ * {@link DeltaWriter#insert}. When {@link
LogicalWriteInfo#columnUpdateSchema()} is absent, every
+ * row follows {@link LogicalWriteInfo#schema()}, as for any other operation.
Connectors must
+ * decide which layout to expect by whether {@link
LogicalWriteInfo#columnUpdateSchema()} is
+ * present, not by the command. A builder that wants full-width rows for some
command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this
mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link
org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that
the scan
+ * does not read columns the command neither references nor declares.
Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only
top-level columns
+ * are pruned: a struct column the command references through a nested field
is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level
operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order.
Implementations must include
Review Comment:
**43.** Minor: this says that the declared columns become
`columnUpdateSchema()` "in declared order", but not that they take the table's
spelling. Declarations resolve case-insensitively, and Spark maps each one back
to the relation's own attribute (`RewriteRowLevelCommand.scala:409-413`), as
"requiredDataAttributes resolves case-insensitively without adopting the
declared spelling" checks (`DeltaBasedColumnUpdateTableSuite.scala:260-274`:
declaring `PK,SALARY,dep` gives `pk, salary, dep`). `updatedColumns()` uses the
table's spelling as well (`RewriteRowLevelCommand.scala:77`). So a connector
that declares `FieldReference("PK")` for a column `pk` and follows L62-67 with
`Arrays.asList(info.updatedColumns()).contains(FieldReference("PK"))` gets
`false` for `UPDATE t SET pk = 2`, because `FieldReference` equality is
case-sensitive. It then builds the narrow operation and looks rows up by the
new `pk`, as in item 5, and its `columnUpdateSchema().get().fieldIndex("PK")`
throws as wel
l. Could this paragraph and `RowLevelOperationInfo.updatedColumns()` say that
the names follow the table's spelling?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DeltaWriter.java:
##########
@@ -52,6 +55,9 @@ public interface DeltaWriter<T> extends DataWriter<T> {
* Reinserts a row with metadata.
* <p>
* This method handles the insert portion of updated rows split into deletes
and inserts.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, the {@code
row} follows it;
+ * otherwise it follows {@link LogicalWriteInfo#schema()}.
Review Comment:
**41. The 4.0 default `reinsert(metadata, row)` still forwards to
`insert(row)` (L68-70), so a split-update `SupportsColumnUpdates` writer that
keeps it gets narrow rows, without their metadata, in `insert()`.** For a
narrow UPDATE on an operation that represents updates as deletes and inserts,
`WriteToDataSourceV2Exec.scala:903` and `:957` pass the narrow row projection
to `reinsert`, so this default hands a `columnUpdateSchema()` row to
`insert()`. However, `SupportsColumnUpdates.java:32-34` says that inserted rows
always follow `schema()` and arrive through `DeltaWriter#insert`, and
`schema()` is empty for UPDATE (`LogicalWriteInfo.java:51-53`). The default
also drops `metadata`, which is the only identity of the reinserted row when
the row ID is a metadata column preserved on reinsert, a setup that
`SupportsColumnUpdates.java:78-81` and `validateRowIdDeclared`
(`RewriteUpdateTable.scala:323-337`) accept.
With the setup of "metadata row ID preserved on reinsert need not be
declared" (`DeltaBasedColumnUpdateTableSuite.scala:486-504`: `_partition` as
the row ID, `columnUpdateSchema()` = `[salary]`) and a writer that implements
`delete`, `update` and `insert` but keeps this default, `UPDATE t SET salary =
-1 WHERE dep = 'hr'` deletes the matched row and then calls `insert([-1])`.
Depending on the writer, the row is lost or written without `pk` and `dep`, and
a writer that decodes `insert` rows by the table layout fails with an
`IndexOutOfBoundsException` from `ProjectingInternalRow`. The group-based path
is guarded against the same mistake, because `DataWriter.writeColumnUpdate(T)`
throws by default (`DataWriter.java:138-142`), but nothing guards this one, and
no test reaches it, since the only split writer, `DeltaBufferWriter`, overrides
`reinsert` (`InMemoryRowLevelOperationTable.scala:876-881`). Because this
default has been public since 4.0.0, changing it would affect existing wri
ters. Could this Javadoc say that a writer for an operation that mixes in
`SupportsColumnUpdates` and represents updates as deletes and inserts must
override `reinsert`, and could `SupportsColumnUpdates.java:32-34` be qualified
accordingly?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,362 @@
+/*
+ * 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 java.util.regex.Pattern
+
+import org.apache.spark.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.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}
+
+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: runtime group filtering reads only the
referenced columns") {
+ 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")
+
+ checkLastWriteInfo(expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD, DEP_FIELD, StructField("salary", IntegerType, nullable =
true)))))
+ 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: connector missing writeColumnUpdate
override 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-cow", "true",
"column-update-cow-no-write-column-update", "true"))
+
+ 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"))
+ }
+
+ test("column-update ReplaceData: the subquery plan narrows the write
relation and the scan") {
+ withTempView("deps") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ 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")
+ }
+
+ 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
+ createAndInitTableWithProps(reqAttrsTableProps("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: unread partition column skips runtime group
filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithProps(reqAttrsTableProps("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")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT
* FROM matched_ids)")
+ },
+ condition =
"DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE",
+ sqlState = Some("KD000"),
+ parameters = Map(
+ "attribute" -> Pattern.quote("`dep`"),
+ "method" -> Pattern.quote("filterAttributes()"),
+ "scanClass" -> (Pattern.quote(
+
classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperation].getName)
+
+ """\$\$anon\$\d+\$\$anon\$\d+"""),
+ "relationOutput" -> Pattern.quote("\"STRUCT<pk: INT NOT NULL, id:
INT, salary: INT, " +
+ "_partition: STRING NOT NULL COMMENT 'Partition key used to store
the row', " +
+ "index: INT NOT NULL COMMENT 'Metadata column used to conflict
with a data " +
+ "column'>\"")),
+ matchPVals = true)
+ }
+ }
+
+ 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: UPDATE without column updates, DELETE and
MERGE read every " +
+ "column") {
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING,
_partition STRING, index INT"
+ Seq(fullWidthTableProps, extraTableProps).foreach { props =>
+ withTable(tableNameAsString) {
+ 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,
+ props)
+
+ if (props == fullWidthTableProps) {
+ 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: 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)
+
+ checkLastUpdatedColumns("salary")
+ assert(rebuilt.operationInfo eq operationTable.operationInfo)
+ val rebuiltQuery = write.query.transform {
+ case r: DataSourceV2Relation if r.table eq operationTable =>
r.copy(table = rebuilt)
+ }
+ assert(rebuiltQuery.exists {
+ case r: DataSourceV2Relation => r.table eq rebuilt
+ case _ => false
+ })
+ V2Writes(write.withNewTable(relation.copy(table =
rebuilt)).withNewQuery(rebuiltQuery))
+ }
+
+ test("column-update ReplaceData: column update schema nullability covers
UPDATE and COPY " +
+ "rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT NOT NULL, bonus INT, dep
STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": null, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow-req-attrs", "pk,salary,bonus"))
Review Comment:
**52.** Nit: this map is `reqAttrsTableProps("pk,salary,bonus")` (L41-43),
which L158, L182 and L206 use. Item 36 asked the delta suite to use its own
hook instead of such literals, and this one is left in this suite, so it would
not follow an override of the hook. Could it be
`reqAttrsTableProps("pk,salary,bonus")`?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data
column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name =>
buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's
values (or the
+ // existence default for a column added after the base row was written),
overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by
the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]):
Array[NamedReference] = {
+ val base = Seq("pk",
"dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r =>
baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends
only the declared
+ // columns in the row projection instead of the full row schema. It declares
`pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever
columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] =
declaredColumns(updatedCols)
+
+ private val clusterExprs: Seq[Expression] = {
+ val refs = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(_.split(",").map(name => FieldReference(name.trim)).toSeq)
+ .getOrElse(Seq(PARTITION_COLUMN_REF))
+ if (properties.getOrDefault(COLUMN_UPDATE_CLUSTER_BY_NESTED_TRANSFORM,
"false") == "true") {
+ refs.map(ref => LogicalExpressions.apply("f",
LogicalExpressions.apply("g", ref)))
+ } else {
+ refs
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ val writeSchema = info.columnUpdateSchema().orElse(info.schema())
+ val batchWrite = new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ physicalInfo: PhysicalWriteInfo): DeltaWriterFactory = {
+ newWriterFactory(info, writeSchema)
+ }
+
+ override protected def doCommit(messages: Array[WriterCommitMessage]):
Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ commitColumnUpdate(info, writeSchema, newData)
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ if (noMetadata && !properties.containsKey(COLUMN_UPDATE_CLUSTER_BY))
{
+ new DeltaWrite {
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ } else {
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(clusterExprs.toArray)
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ clusterExprs.map { expr =>
+ LogicalExpressions.sort(
+ expr,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()): SortOrder
+ }.toArray
+ }
+
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ }
+ }
+ }
+
+ protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(writeSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the
full
+ // table schema by overlaying them on the base row found by pk.
+ protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3,
writeSchema).asInstanceOf[InternalRow]
+ val baseRow = Some(currentRowByRowId(Seq("pk" -> pk)))
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via
writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
Review Comment:
**51.** Nit: the dispatcher always passes `info.options` and the updated
columns (L164-201), so this `options` default and those at L518, L600, L612 and
L697, as well as `updatedCols = Nil` at L611, are never used, and neither is
`pkRowId = true` of `DeltaBufferWriter` (L851), whose only construction passes
it (L811). The SPARK-57681 options tests (`UpdateTableSuiteBase.scala:1274`,
`:1292` and `:1349`) run only with each suite's `extraTableProps`, so if a
req-attrs, split or row-ID branch stopped passing `info.options`, it would
still compile, silently use empty options, and no test would notice. Could
these defaults go?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/ColumnUpdateTableSharedTests.scala:
##########
@@ -0,0 +1,355 @@
+/*
+ * 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}
+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.{InMemoryBaseTable,
InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.connector.write.{RowLevelOperation,
RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+// Column-update tests that hold for both delta-based and group-based
operations; each suite
+// supplies its column-update table properties through `extraTableProps`.
+trait ColumnUpdateTableSharedTests extends UpdateTableSuiteBase {
+
+ // the properties of a column-update table whose connector declares exactly
`reqAttrs`
+ protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String,
String]
+
+ // the properties of a table of the same kind whose operation does not mix
in column updates
+ protected def fullWidthTableProps: java.util.Map[String, String]
+
+ // the metadata columns the scan of a column-update table reads, in DDL form
+ protected def scanMetadataColumns: String = ", _partition STRING, index INT"
+
+ // the class name of the operation that `reqAttrsTableProps` builds
+ protected def reqAttrsOperationClass: String = if (deltaUpdate) {
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+ } else {
+
classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with CHECK (extra > 0), partitioned
by dep
+ protected def createAndInitTableWithProps(props: java.util.Map[String,
String]): Unit = {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT, bonus
INT, extra INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100, "bonus": 10, "extra":
1 }
+ |{ "pk": 2, "id": 2, "dep": "hr", "salary": 200, "bonus": 20, "extra":
1 }
+ |{ "pk": 3, "id": 3, "dep": "software", "salary": 300, "bonus": 30,
"extra": 1 }
+ |""".stripMargin,
+ props)
+ sql(s"ALTER TABLE $tableNameAsString ADD CONSTRAINT positive_extra CHECK
(extra > 0)")
+ }
+
+ test("column-update: undeclared columns are not written and keep their
values") {
+ 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'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
Review Comment:
**45.** Minor: this reads `columnUpdateSchema.get` without a guard, as do
L186, L201 and L348 here, `DeltaBasedColumnUpdateTableSuite.scala:147`, `:269`,
`:301` and `:498`, and `checkLastColumnUpdateWriteLog`
(`RowLevelOperationSuiteBase.scala:424`), which seven tests use. If a
regression stops narrowing, they all fail with
`java.util.NoSuchElementException: No value present`, without the expected
names or the actual write info. Like `checkLastUpdatedColumns`
(`RowLevelOperationSuiteBase.scala:409-416`, item 33), could a
`checkLastColumnUpdateSchema(expectedNames: String*)` helper first assert that
the schema is present, with a message, and could these sites and
`checkColumnUpdateWrite` use it?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala:
##########
@@ -92,6 +95,38 @@ object GroupBasedRowLevelOperationScanPlanning extends
Rule[LogicalPlan] with Pr
}
}
+ /**
+ * Returns the columns of `relation` that the projections and filters
directly above each read
+ * of the table use, as the rest of the query sees only their output. Only
whole columns are
+ * pruned, as the query is not rewritten for a pruned nested schema.
+ *
+ * An UPDATE with a subquery in its condition reads the table twice, once
per Union child, and
+ * the analyzer gives the second read new expr IDs. Both reads share one
scan, so the columns
+ * are matched to `relation` by name.
+ */
+ private def readAttrs(
+ plan: LogicalPlan,
+ table: RowLevelOperationTable,
+ relation: DataSourceV2Relation): Seq[AttributeReference] = {
+ val readNames = tableReads(plan, table).flatMap { case (projects, filters,
read) =>
+ val references = AttributeSet((projects ++
filters).flatMap(_.references))
+ read.output.filter(references.contains).map(_.name)
+ }.toSet
+ relation.output.filter(attr => readNames.contains(attr.name))
+ }
+
+ private def tableReads(
Review Comment:
**48.** Nit: `table` is always `relation.table` (L52 only casts it), and
this method only returns a tuple that `readAttrs` takes apart at L111; the
`NamedExpression` (L21) and `RowLevelOperationTable` (L30) imports exist only
for these signatures. Could `readAttrs` take only the relation and use a local
recursive function? For example:
```scala
private def readAttrs(
plan: LogicalPlan,
relation: DataSourceV2Relation): Seq[AttributeReference] = {
def readNames(node: LogicalPlan): Seq[String] = node match {
case PhysicalOperation(projects, filters, r: DataSourceV2Relation)
if r.table eq relation.table =>
val references = AttributeSet((projects ++
filters).flatMap(_.references))
r.output.filter(references.contains).map(_.name)
case _ =>
node.children.flatMap(readNames)
}
val names = readNames(plan).toSet
relation.output.filter(attr => names.contains(attr.name))
}
```
L70 would then call `readAttrs(rd.query, relation)`. It has to keep stopping
at the first match, because `plan.collect` would also match the `Filter` under
a `Project`, with `projects = r.output`.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data
column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name =>
buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's
values (or the
+ // existence default for a column added after the base row was written),
overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by
the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]):
Array[NamedReference] = {
+ val base = Seq("pk",
"dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r =>
baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends
only the declared
+ // columns in the row projection instead of the full row schema. It declares
`pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever
columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] =
declaredColumns(updatedCols)
+
+ private val clusterExprs: Seq[Expression] = {
+ val refs = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(_.split(",").map(name => FieldReference(name.trim)).toSeq)
+ .getOrElse(Seq(PARTITION_COLUMN_REF))
+ if (properties.getOrDefault(COLUMN_UPDATE_CLUSTER_BY_NESTED_TRANSFORM,
"false") == "true") {
+ refs.map(ref => LogicalExpressions.apply("f",
LogicalExpressions.apply("g", ref)))
+ } else {
+ refs
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ val writeSchema = info.columnUpdateSchema().orElse(info.schema())
+ val batchWrite = new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ physicalInfo: PhysicalWriteInfo): DeltaWriterFactory = {
+ newWriterFactory(info, writeSchema)
+ }
+
+ override protected def doCommit(messages: Array[WriterCommitMessage]):
Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ commitColumnUpdate(info, writeSchema, newData)
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ if (noMetadata && !properties.containsKey(COLUMN_UPDATE_CLUSTER_BY))
{
+ new DeltaWrite {
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ } else {
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(clusterExprs.toArray)
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ clusterExprs.map { expr =>
+ LogicalExpressions.sort(
+ expr,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()): SortOrder
+ }.toArray
+ }
+
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ }
+ }
+ }
+
+ protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(writeSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the
full
+ // table schema by overlaying them on the base row found by pk.
+ protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3,
writeSchema).asInstanceOf[InternalRow]
+ val baseRow = Some(currentRowByRowId(Seq("pk" -> pk)))
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via
writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] =
reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ val pkRowId = rowId().map(_.describe()).toSeq == Seq("pk") &&
+ info.rowIdSchema().get().fields.map(_.dataType).toSeq ==
Seq(IntegerType)
+ new DeltaBufferedRowsWriterFactory(writeSchema, pkRowId)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the
connector-declared
+ // columns (requiredDataAttributes order). The original row is located by
the row ID, taking a
+ // data row-ID column from the narrow row and a metadata row-ID column
from the preserved
+ // metadata, and the full row is rebuilt by overlaying the narrow row onto
it. A DELETE logged
+ // without a `pk` row ID is resolved the same way from its metadata, so
such a row ID must
+ // consist of metadata columns.
+ override protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val metadataSchema = info.metadataSchema().orElse(new StructType())
+ val coversTable = schema.fieldNames.forall(writeFieldIdx.contains)
+ val deleteOpName = UTF8String.fromString(Delete.toString)
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ def rowIdValues(narrowRow: Option[InternalRow], metadata: InternalRow):
Seq[(String, Any)] = {
+ rowId().toSeq.map(_.describe()).map { name =>
+ (writeFieldIdx.get(name), narrowRow) match {
+ case (Some(i), Some(row)) => name -> row.get(i,
writeSchema(i).dataType)
+ case _ =>
+ if (metadata == null ||
!metadataSchema.fieldNames.contains(name)) {
+ throw new IllegalStateException(s"The split fixture can only
locate a row by pk " +
+ s"or by metadata row-ID columns, but row-ID column $name is
in neither")
+ }
+ val i = metadataSchema.fieldIndex(name)
+ name -> metadata.get(i, metadataSchema(i).dataType)
+ }
+ }
+ }
+
+ def metadataOf(logRow: InternalRow): InternalRow =
+ if (logRow.isNullAt(2)) null else logRow.getStruct(2,
metadataSchema.length)
+
+ val expandedData = newData.map { buf =>
+ val expanded = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == deleteOpName && logRow.isNullAt(1)) {
+ val values = rowIdValues(None, metadataOf(logRow))
+ buf.deletes += currentRowByRowId(values).getInt(0)
+ } else if (opName == reinsertOpName) {
+ val narrowRow = logRow.get(3,
writeSchema).asInstanceOf[InternalRow]
+ val baseRow = if (coversTable) {
+ None
+ } else {
+ Some(currentRowByRowId(rowIdValues(Some(narrowRow),
metadataOf(logRow))))
+ }
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ expanded.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via
writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ expanded.rows.append(insertRow.copy())
+ }
+ }
+ expanded
+ }
+
+ withDeletes(newData)
+ withData(expandedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ // Test-only: a split-update fixture whose requiredDataAttributes() is the
given column list.
+ class DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateSplitOperation(command, options = options)
{
+ override def requiredDataAttributes(): Array[NamedReference] =
reqCols.map(FieldReference(_))
+ }
+
+ // Test-only: a split-update fixture whose row ID is the given column and
whose
+ // requiredDataAttributes() lists only the updated columns, to exercise
which undeclared
+ // row-ID columns still reach the REINSERT row.
+ class DeltaBasedColumnUpdateSplitOperationWithRowId(
+ command: Command,
+ rowIdCol: String,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateSplitOperation(command, updatedCols,
options) {
+ override def rowId(): Array[NamedReference] =
Array(FieldReference(rowIdCol))
+ override def requiredDataAttributes(): Array[NamedReference] =
updatedCols.toArray
+ }
+
+ class PartitionBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ override val options: CaseInsensitiveStringMap =
CaseInsensitiveStringMap.empty())
+ extends RowLevelOperation with SupportsColumnUpdates with
RowLevelOperationWithOptions {
+ var configuredScan: BatchScanBaseClass = _
+
+ override def command(): Command = command
+
+ override def requiredDataAttributes(): Array[NamedReference] =
declaredColumns(updatedCols)
+
+ override def requiredMetadataAttributes(): Array[NamedReference] = {
+ if (noMetadata) {
+ Array.empty
+ } else {
+ Array(PARTITION_COLUMN_REF, INDEX_COLUMN_REF)
+ }
+ }
+
+ private def clusterColumnRef: NamedReference =
+ if (noMetadata) FieldReference("dep") else PARTITION_COLUMN_REF
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap):
ScanBuilder = {
+ val onBuild: BatchScanBaseClass => Unit = { scan =>
+ configuredScan = scan
+ }
+ if (properties.getOrDefault(COLUMN_UPDATE_COW_UNREAD_FILTER_ATTRS,
"false") == "true") {
Review Comment:
**47.** Nit: `PartitionBasedColumnUpdateOperation` and
`PartitionBasedNarrowReplaceData` copy parts of `PartitionBasedOperation` and
`PartitionBasedReplaceData` instead of sharing them:
`requiredMetadataAttributes()` (L629-635 and L213-219), the default scan
builder (L663 and L222-224), the commit bookkeeping (L722-726 and L265-269) and
the `build()` override here (L656-660 and L780-784). This branch has already
drifted: it always builds a plain `InMemoryBatchScan`, so it skips the
`useCatalystRuntimeFiltering` switch of `newRowLevelScanBuilder` (L773-777),
and a table with both `use-catalyst-runtime-filtering` and
`column-update-cow-unread-filter-attrs` would silently plan the plain scan. No
test combines them today. Could the shared parts go through helpers, or could
this branch at least honor `useCatalystRuntimeFiltering`?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,623 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkRuntimeException,
SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.analysis.RewriteRowLevelCommand
+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.expressions.{FieldReference,
NamedReference}
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, 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
Review Comment:
**44. No test checks what the connector receives for a narrow delta UPDATE
with metadata columns, or for a narrow copy-on-write UPDATE with a subquery
condition.** The tests that assert the metadata, row IDs and rows in the write
log on these paths, "update handles metadata columns correctly" and "update
with subquery handles metadata columns correctly", live in
`DeltaBasedUpdateTableSuite.scala:35` and `:60` and in
`GroupBasedUpdateTableSuite.scala:31` and `:53`, which these suites do not
extend. The narrow write-log checks cover only delete and reinsert (L401,
L499), the no-metadata delta writer
(`DeltaBasedNoMetadataColumnUpdateTableSuite.scala:54` and `:69`) and
copy-on-write without a subquery (`GroupBasedColumnUpdateTableSuite.scala:96`,
`:113` and `:131`). The subquery test there
(`GroupBasedColumnUpdateTableSuite.scala:156-177`) checks the scan, the write
relation and the result, and L85-102 here checks only the schemas. So nothing
asserts the `update(metadata, id, row)` e
ntries of `DeltaWithMetadataWritingSparkTask` (`_partition` kept, `index`
nulled, the row ID and the narrow row), or the `writeColumnUpdate(metadata,
record)` entries of the `Union` plan. The fixture commits do not read metadata
on these paths (`InMemoryRowLevelOperationTable.scala:486-491` and `:736-745`),
so a metadata projection remapped to the wrong ordinal in either plan shape, or
an `index` that is no longer nulled, would pass every suite. The PR description
also says that the suites cover "the rows, row IDs and metadata the connector
receives". Could each path get a `checkLastColumnUpdateWriteLog(...)`
assertion, e.g. by porting the two metadata tests?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -246,18 +290,55 @@ abstract class RowLevelOperationSuiteBase
assert(scans.isEmpty, "plan must not contain scan relations")
}
+ // the schemas of the scans of the target table, directly or through a
row-level operation, in
+ // the plan of the query
+ private def executeAndCollectScanSchemas(query: String): Seq[StructType] = {
+ val executedPlan = executeAndKeepPlan {
+ sql(query)
+ }
+ val targetName = table.name
+ def scansTarget(scan: BatchScanExec): Boolean = scan.table match {
Review Comment:
**50.** Nit: both cases compute `t.name == targetName`, because
`RowLevelOperationTable.name` already returns `table.name`
(`RowLevelOperationTable.scala:47`), so this can be `scan.table.name ==
targetName`. Also, `executeAndCheckScanExcludes` (L334-342) has a single
caller, `DeltaBasedColumnUpdateTableSuite.scala:254`, where the exact check
fits: that scan should read exactly `pk INT, _partition STRING`, which the
inherited "update with schema pruning"
(`DeltaBasedUpdateTableSuiteBase.scala:79-90`) also expects in that suite.
Could L254 use `executeAndCheckScan(..., "pk INT, _partition STRING")`, which
would also catch a scan that starts reading `id` or `index`, and could the
Excludes helper go?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,623 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkRuntimeException,
SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.analysis.RewriteRowLevelCommand
+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.expressions.{FieldReference,
NamedReference}
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, 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 splitOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitWithRowIdOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ private def splitReqAttrsTableProps(reqAttrs: String): java.util.Map[String,
String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs, "split-updates",
"true")
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData, reqAttrsTableProps(reqAttrs))
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String):
Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update", "true", "split-updates", "true"))
+ }
+
+ 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: 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 and the column update schema follow the
non-identity " +
+ "assignments") {
+ Seq(
+ ("id = -1", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("id = -1, dep = 'eng'", Seq("id", "dep"), Seq("pk", "dep", "id"),
Row(1, -1, "eng")),
+ ("id = -1, dep = dep", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1,
"hr")),
+ ("id = id, dep = 'eng'", Seq("dep"), Seq("pk", "dep"), Row(1, 10,
"eng")),
+ ("id = id, dep = dep", Seq.empty, Seq("pk", "dep"), Row(1, 10, "hr")),
+ ("id = pk", Seq("id"), Seq("pk", "dep", "id"), Row(1, 1, "hr"))
+ ).foreach { case (assignments, updatedColumns, columnUpdateNames,
updatedRow) =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 10, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns(updatedColumns: _*)
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
columnUpdateNames)
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), updatedRow ::
Nil)
+ }
+ }
+ }
+
+ test("column-update split: row-level DELETE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin)
+ 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") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update: undeclared write requirement columns are rejected") {
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
ColumnPruning.ruleName
+ val splitReqAttrsOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+ Seq(
+ ("pk,salary", "dep", Map.empty[String, String], reqAttrsOperationClass,
"[`dep`]"),
+ ("pk,salary", "dep", Map("split-updates" -> "true"),
splitReqAttrsOperationClass,
+ "[`dep`]"),
+ ("pk,salary", "dep", Map("column-update-cluster-by-nested-transform" ->
"true"),
+ reqAttrsOperationClass, "[`dep`]"),
+ ("pk,salary", "id,dep", Map.empty[String, String],
reqAttrsOperationClass,
+ "[`dep`, `id`]"),
+ ("pk,salary,s", "s.missing", Map.empty[String, String],
reqAttrsOperationClass,
+ "[`s`.`missing` (FIELD_NOT_FOUND)]")
+ ).foreach { case (reqAttrs, clusterBy, props, operationClass, columns) =>
+ withTable(tableNameAsString) {
+ val tableProps = props ++ Map(
+ "column-update-req-attrs" -> reqAttrs, "column-update-cluster-by" ->
clusterBy)
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, s STRUCT<c1:
INT>, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 1, "s": { "c1": 1 }, "dep": "hr" }
+ |""".stripMargin,
+ tableProps.asJava)
+
+ 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" -> operationClass, "columns" ->
columns))
+ }
+ }
+ }
+ }
+ }
+
+ 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: distribution by an undeclared row ID column is
allowed") {
+ Seq("salary = -1", "salary = -1, pk = pk").foreach { assignments =>
+ withTable(tableNameAsString) {
+ 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 $assignments 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: an undeclared, unreferenced partition column is not
read") {
+ 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
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: 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: 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"))
+
+ // `id = id` assigns nothing, so only the non-empty check can reject the
declaration
+ Seq("id = -1", "id = id").foreach { assignment =>
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET $assignment WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ parameters = Map("connector" -> reqAttrsOperationClass))
+ }
+ }
+
+ test("column-update: invalid requiredDataAttributes are rejected") {
+ Seq(
+ ("pk,s.c1,dep", "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ "nestedAttributes" -> "[`s`.`c1`]"),
+ ("pk,pk,salary", "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`pk`]"),
+ ("salary,PK,SALARY,pk",
"COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`salary`, `PK`]"),
+ ("pk,_partition,salary",
"COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ "metadataAttributes" -> "[`_partition`]"),
+ ("pk,nonexistent_col,salary",
"COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ "unknownAttributes" -> "[`nonexistent_col`]"),
+ ("pk,salary",
"COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ "missingColumns" -> "[`a.b`]")
+ ).foreach { case (reqAttrs, condition, columns) =>
+ withTable(tableNameAsString) {
+ createAndInitTableWithReqAttrs(reqAttrs,
+ "pk INT NOT NULL, salary INT, s STRUCT<c1: INT, c2: INT>, `a.b` INT,
dep STRING",
+ """{ "pk": 1, "salary": 100, "s": { "c1": 1, "c2": 2 }, "a.b": 1,
"dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1, `a.b` = -1 WHERE
pk = 1")
+ },
+ condition = condition,
+ parameters = Map("connector" -> reqAttrsOperationClass, columns))
+ }
+ }
+ }
+
+ test("column-update: a declared column with no name parts is reported as
<empty>") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+ val t = table
+ val operation = new t.DeltaBasedColumnUpdateOperation(UPDATE) {
+ override def requiredDataAttributes(): Array[NamedReference] =
+ Array(FieldReference(Seq.empty[String]), FieldReference("pk"),
+ FieldReference(Seq.empty[String]))
+ }
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ RewriteRowLevelCommand.resolveRequiredDataAttrs(
+ spark.table(tableNameAsString).queryExecution.analyzed, operation)
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> operation.getClass.getName, "unknownAttributes" ->
"[<empty>, <empty>]"))
+ }
+
+ test("column-update split: column update schema has the declared columns") {
Review Comment:
**49.** Nit: this test repeats "column-update: column update schema contains
a single assigned column" (L85-102): the same DDL, the same `UPDATE ... SET id
= -1 WHERE pk = 1` and the same eight-line `checkLastWriteInfo(...)` (L383-390
and L94-101), with only the table properties and one data row differing. The
two configurations take different rewrite paths, so both are worth running, but
the expected write info is maintained twice. Could they become one test that
loops over the two property sets, like L129-151, with the split properties
shared with `createAndInitTableSplit` (L80-83)?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala:
##########
@@ -56,6 +56,56 @@ object V2ExpressionUtils extends SQLConfHelper with Logging {
plan.resolve(ref.fieldNames.toImmutableArraySeq, conf.resolver)
}
+ /**
+ * Returns the names of the columns referenced by the given expressions that
cannot be resolved
+ * against the plan, at any depth, including inside transform arguments.
+ */
+ private[sql] def unresolvableColumns(
+ exprs: Seq[V2Expression],
+ relation: LogicalPlan): Seq[String] = {
+ unresolvableReferences(exprs, relation).map {
+ case (Seq(), _) => "<empty>"
+ case (parts, None) => parts.quoted
+ case (parts, Some(condition)) => s"${parts.quoted} ($condition)"
+ }.distinct
Review Comment:
**54.** Nit: this `.distinct` became a no-op when `unresolvableReferences`
got its own `.distinct` (L106): distinct pairs map to distinct strings, since
`quoted` is injective, an entry with a condition ends with `)` while a quoted
name cannot, and `<empty>` is not a quoted name. Could it go?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -321,25 +804,63 @@ class InMemoryRowLevelOperationTable private (
}
}
-private class DeltaBufferedRowsWriterFactory(schema: StructType) extends
DeltaWriterFactory {
+private class DeltaBufferedRowsWriterFactory(
+ schema: StructType,
+ pkRowId: Boolean = true) extends DeltaWriterFactory {
override def createWriter(partitionId: Int, taskId: Long):
DeltaWriter[InternalRow] = {
- new DeltaBufferWriter(schema)
+ new DeltaBufferWriter(schema, pkRowId)
+ }
+}
+
+// Test-only: a writer factory for a connector that never overrides
writeColumnUpdate(), to
Review Comment:
**53.** Nit: these four copy-on-write writer classes (L815-847) sit between
`DeltaBufferedRowsWriterFactory` (L807-813) and the `DeltaBufferWriter` it
creates (L811, L849-894), which were adjacent at the merge-base, so the
factory's `pkRowId` (L809) is explained only at the writer (L849-850), 40 lines
and four unrelated classes away. As in item 19, AGENTS.md asks to "put the new
member with the code it belongs with". Could they move below
`DeltaBufferWriter`?
--
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]
