anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3635902078
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala:
##########
@@ -139,17 +138,13 @@ class
RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
tableAttrs: Seq[Attribute],
scanAttrs: Seq[Attribute]): AttributeMap[Attribute] = {
- val attrMapping = tableAttrs.map { tableAttr =>
+ // The scan may be narrowed to exclude columns not needed by the connector.
+ // Attributes absent from the scan are skipped here; the caller must ensure
+ // that any attribute referenced in the condition is present in the scan.
+ val attrMapping = tableAttrs.flatMap { tableAttr =>
scanAttrs
.find(scanAttr => conf.resolver(scanAttr.name, tableAttr.name))
.map(scanAttr => tableAttr -> scanAttr)
- .getOrElse {
- throw new AnalysisException(
- errorClass = "_LEGACY_ERROR_TEMP_3075",
- messageParameters = Map(
- "tableAttr" -> tableAttr.toString,
- "scanAttrs" -> scanAttrs.mkString(",")))
- }
}
Review Comment:
Makes sense. Similar to other changes for column updates paths, I created a
conditional method `buildNarrowTableToScanAttrMap()` which is called only
during column updates and throws when any of the condition references are
missing. My rationale is that the runtime filtering applies to only the filters
so it is sufficient remap the filters only. Let me know if this understanding
is incorrect.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level
operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted
rows, in declared
+ * order. Implementations must include every column they want to receive
(typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any
columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link
RowLevelOperationInfo#updatedColumns()}) is
Review Comment:
Done.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,13 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation.
Review Comment:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,540 @@
+/*
+ * 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.sql.Row
+import org.apache.spark.sql.connector.catalog.{CatalogV2Util, TableInfo}
+import
org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity,
reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField,
StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that return true from
+ *
[[org.apache.spark.sql.connector.write.RowLevelOperation#supportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the row projection
+ * (LogicalWriteInfo.schema()) to contain only the assigned/changed columns
rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] =
{
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ test("column-update: rowSchema contains only the 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")
+
+ // info.schema() is empty: UPDATE-only column-update writes have no
INSERT-shaped rows,
+ // so nothing flows through the writer's write() path. The narrow
update-row layout is
+ // carried by info.updateSchema().
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema 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))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false),
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema is empty 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 -- the
connector still declares
+ // pk for row lookup, so the narrow update schema is just [pk].
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Array(PK_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))),
+ expectedUpdateSchema = 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))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: rowSchema 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))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: cross-column assignment is not treated as identity") {
+ 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 dep = dep, id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: nested struct field update narrows to the root struct
column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep
STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"expected [s] in updatedColumns (root struct) but got: $updatedNames")
+
+ // info.updateSchema() carries the narrow row layout; info.schema() is the
full table.
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("s"),
+ s"s must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("dep"),
+ s"dep must not be in update schema: $updateSchema")
+ }
+
+ test("column-update: nested field identity update reports root struct as
updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep
STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"nested identity is reported as an update at root-column granularity:
$updatedNames")
+
+ // Data correctness: the struct is rewritten but with equal values, so
rows are unchanged.
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns contains non-identity assigned columns")
{
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'eng' WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id", "dep"),
+ s"expected [id, dep] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns excludes identity assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = dep WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id"),
+ s"expected only [id] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns is empty 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")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"expected empty updatedColumns but got:
${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: updatedColumns is empty for DELETE (Javadoc contract)")
{
+ 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'")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"DELETE must pass empty updatedColumns but got:
${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ 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)
+ // 1 row matched WHERE pk = 1 -> 1 update. MoR emits no COPY rows.
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate =
true)
+ }
+
+ 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, deltaUpdate =
true)
+ }
+
+ 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, deltaUpdate =
true)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-req-attrs", reqAttrs)
+ val columns =
CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ 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: empty requiredDataAttributes throws AnalysisException")
{
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-empty-req-attrs", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(
+ StructType.fromDDL("pk INT NOT NULL, id INT, dep STRING"))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin)
+
+ // A connector that mixes in SupportsColumnUpdates but returns an empty
+ // requiredDataAttributes() violates the mix-in contract -- analysis must
reject the
+ // operation rather than silently falling through to the wide write path.
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getCondition == "EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ s"expected EMPTY_REQUIRED_DATA_ATTRIBUTES but got: ${ex.getCondition}")
+ }
+
+ test("column-update: requiredDataAttributes throws AnalysisException for
invalid column") {
+ createAndInitTableWithReqAttrs("nonexistent_col", "pk INT NOT NULL, id
INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getMessage.contains("nonexistent_col"),
+ s"Expected error about unresolvable column but got: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableFromInfo(schemaString: String, jsonData:
String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-from-info", "true")
+ val columns =
CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update from-info: write schema is updatedColumns + pk
pass-through") {
+ createAndInitTableFromInfo("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")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"salary must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"pk must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("id"),
+ s"id must not be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("dep"),
+ s"dep must not be in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: pk already in updatedColumns is not
duplicated") {
+ createAndInitTableFromInfo("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'")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ val pkCount = updateSchema.fieldNames.count(_ == "pk")
+ assert(pkCount == 1, s"pk must appear exactly once in update schema:
$updateSchema")
+ }
+
+ test("column-update from-info: data correctness") {
+ createAndInitTableFromInfo("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 = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split", "true")
+ val columns =
CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: write schema is narrow (assigned + pk
pass-through)") {
+ 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")
+
+ // For representUpdateAsDeleteAndInsert connectors, reinsert-tagged rows
still route through
+ // writeUpdate() (info.updateSchema() = narrow), not write().
info.schema() stays empty.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_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)
+ }
+
+ //
---------------------------------------------------------------------------
+ // Scan-narrowing tests: verify that when the connector opts into column
updates,
+ // the physical scan reads only the columns actually needed by the operation.
+ //
+ // Assertions use `checkLastScanExcludes` (proves narrowing) and
`checkLastScanIncludes`
+ // (proves transparent widening) rather than exact-set matches --
ColumnPruning is
+ // free to further tighten the scan beyond our analysis-time narrow set
(e.g. when an
+ // assignment RHS is a literal, the assigned column doesn't need to be read).
+ //
---------------------------------------------------------------------------
+
+ test("column-update: scan excludes columns outside connector-declared + cond
+ RHS refs") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [pk] only; RHS is a
literal.
+ // `dep` is neither declared nor referenced -- it must not appear in the
scan.
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastScanExcludes("dep")
+ }
+
+ test("column-update: scan transparently widens for cond-referenced columns")
{
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [dep].
+ // `dep` must appear even though the connector didn't declare it.
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkLastScanIncludes("dep")
+ }
+
+ test("column-update: scan transparently widens for assignment RHS
references") {
+ 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)
+
+ // requiredDataAttributes = [pk, salary]; RHS `salary + bonus` references
`bonus`.
+ // `bonus` must appear in the scan even though it's not declared as
required.
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastScanIncludes("bonus", "salary")
+ checkLastScanExcludes("dep")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update: analysis fails when assignment key is outside
requiredDataAttributes") {
+ // Connector declares only [pk] but the user assigns to `id`. We enforce
+ // updatedColumns ⊆ requiredDataAttributes at analysis time (root-column
granularity).
Review Comment:
Yes, this was auto generated. I rewrote it as per your suggestion.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -210,11 +247,25 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildReplaceDataProjections(
plan: LogicalPlan,
rowAttrs: Seq[Attribute],
- metadataAttrs: Seq[Attribute]): ReplaceDataProjections = {
+ metadataAttrs: Seq[Attribute],
+ updateRowAttrs: Seq[Attribute] = Nil): ReplaceDataProjections = {
val outputs = extractOutputs(plan)
- val outputsWithRow = filterOutputs(outputs, OPERATIONS_WITH_ROW)
- val rowProjection = newLazyProjection(plan, outputsWithRow, rowAttrs)
+ val rowProjection = if (updateRowAttrs.nonEmpty) {
+ val outputsForInsert = filterOutputs(outputs,
+ OPERATIONS_WITH_ROW -- Set(UPDATE_OPERATION, COPY_OPERATION))
+ newLazyProjection(plan, outputsForInsert, rowAttrs)
Review Comment:
Thanks for catching this. I mirrored the guard in
`buildWriteDeltaProjections`
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -83,11 +84,20 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
table: RowLevelOperationTable,
metadataAttrs: Seq[AttributeReference],
rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
-
Review Comment:
Done
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -41,7 +44,11 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
EliminateSubqueryAliases(aliasedTable) match {
case r @ ExtractV2Table(tbl: SupportsRowLevelOperations) =>
checkNoGeneratedColumns(r, UPDATE)
- val table = buildOperationTable(tbl, UPDATE, r.options)
+ val updatedCols = assignments.collect {
+ case Assignment(key: AttributeReference, value) if
!isIdentityAssignment(key, value) =>
+ FieldReference(key.name)
Review Comment:
Thanks for the pointer. Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -223,19 +274,39 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
None
}
- ReplaceDataProjections(rowProjection, metadataProjection)
+ ReplaceDataProjections(rowProjection, updateRowProjection,
metadataProjection)
}
protected def buildWriteDeltaProjections(
plan: LogicalPlan,
rowAttrs: Seq[Attribute],
rowIdAttrs: Seq[Attribute],
- metadataAttrs: Seq[Attribute]): WriteDeltaProjections = {
+ metadataAttrs: Seq[Attribute],
+ updateRowAttrs: Seq[Attribute] = Nil): WriteDeltaProjections = {
val outputs = extractOutputs(plan)
- val rowProjection = if (rowAttrs.nonEmpty) {
+ // Always produce Some(rowProjection) even for empty rowAttrs
(identity-only column updates).
+ // When updateRowAttrs is non-empty, the row projection covers
INSERT-shaped rows only and a
+ // separate updateRowProjection handles UPDATE/COPY/REINSERT-shaped rows.
+ val rowProjection = if (updateRowAttrs.nonEmpty) {
+ val outputsForInsert = filterOutputs(outputs,
+ OPERATIONS_WITH_ROW -- Set(UPDATE_OPERATION, COPY_OPERATION,
REINSERT_OPERATION))
+ if (outputsForInsert.isEmpty) {
+ Some(ProjectingInternalRow(StructType(Nil), Nil))
+ } else {
+ Some(newLazyProjection(plan, outputsForInsert, rowAttrs))
+ }
+ } else if (rowAttrs.nonEmpty) {
val outputsWithRow = filterOutputs(outputs, OPERATIONS_WITH_ROW)
Some(newLazyProjection(plan, outputsWithRow, rowAttrs))
+ } else {
+ Some(ProjectingInternalRow(StructType(Nil), Nil))
Review Comment:
Thanks for the explanation. I moved it back to `None` and updated the
comment.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -244,6 +281,327 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends
only the
+ // declared + assigned columns in the row projection instead of the full row
schema. The base
+ // class composes its required-attrs set as `pk` (the row-lookup key) plus
whatever columns
+ // Spark reports as being assigned via
`RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedOperation(command, CaseInsensitiveStringMap.empty())
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val pk: NamedReference = FieldReference("pk")
+ val updatedNames = updatedCols.map(_.describe()).toSet
+ if (updatedNames.contains("pk")) updatedCols.toArray
+ else (pk +: updatedCols).toArray
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a
stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ // Use the narrow update schema for UPDATE/COPY/REINSERT
rows when present.
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, rows contain only the assigned
columns
+ // (narrow schema from LogicalWriteInfo). We expand each row
to the full table
+ // schema by overlaying write-schema columns on the base row
found by pk.
+ override protected def doCommit(messages:
Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx =
writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ 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 = dataMap.values.iterator.flatten
+ .flatMap(_.rows)
+ .find(r => r.getInt(schema.fieldIndex("pk")) == pk)
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ fullRow.update(i, base.get(i,
schema.fields(i).dataType))
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i)
=>
+ writeFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j,
field.dataType))
+ }
+ }
+ 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
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit
= {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(command: Command, reqCols:
Array[String])
+ extends DeltaBasedColumnUpdateOperation(command) {
+ override def requiredDataAttributes(): Array[NamedReference] =
reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateOperationFromInfo(
+ command: Command,
+ updatedCols: Seq[NamedReference])
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols) {
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a
stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+ override def requiredDistribution(): Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ override def requiredOrdering(): Array[SortOrder] =
Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row
has only the
+ // connector-declared columns (requiredDataAttributes order).
+ // pk is the first field in the write schema (declared before
updatedCols).
+ // Reconstruct the full row by overlaying the narrow row onto
the original.
+ override protected def doCommit(messages:
Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx =
writeSchema.fieldNames.zipWithIndex.toMap
+ val reinsertOpName =
UTF8String.fromString(Reinsert.toString)
+ val pkIdx = writeFieldIdx("pk")
+
+ val expandedData = newData.map { buf =>
+ val expanded = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == reinsertOpName) {
+ val narrowRow = logRow.get(3,
writeSchema).asInstanceOf[InternalRow]
+ val pk = narrowRow.getInt(pkIdx)
+ val baseRow = dataMap.values.iterator.flatten
+ .flatMap(_.rows)
+ .find(r => r.getInt(schema.fieldIndex("pk")) == pk)
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ fullRow.update(i, base.get(i,
schema.fields(i).dataType))
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i)
=>
+ writeFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j,
field.dataType))
+ }
+ }
+ expanded.rows.append(fullRow)
+ }
+ }
+ expanded
+ }
+
+ withDeletes(newData)
+ withData(expandedData)
+ lastWriteLog = newData.flatMap(buffer =>
buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit
= {}
+ }
+ }
+ }
+ }
+ }
+
+ class PartitionBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends RowLevelOperation with SupportsColumnUpdates {
+ var configuredScan: InMemoryBatchScan = _
+
+ override def command(): Command = command
+
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r =>
baseNames.contains(r.describe()))).toArray
+ }
+
+ override def requiredMetadataAttributes(): Array[NamedReference] =
+ Array(PARTITION_COLUMN_REF, INDEX_COLUMN_REF)
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap):
ScanBuilder = {
+ new InMemoryScanBuilder(schema, options) {
+ override def build(): Scan = {
+ val scan = super.build()
+
InMemoryRowLevelOperationTable.recordLastScanSchema(scan.readSchema())
+ configuredScan = scan.asInstanceOf[InMemoryBatchScan]
+ scan
+ }
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
+ lastWriteInfo = info
+ new WriteBuilder {
+ override def build(): Write = new Write with
RequiresDistributionAndOrdering {
+ override def requiredDistribution: Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+
+ override def requiredOrdering: Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+
+ override def toBatch: BatchWrite = {
+ val narrowSchema = if (info.updateSchema().isPresent) {
+ info.updateSchema().get()
+ } else {
+ info.schema()
+ }
+ PartitionBasedNarrowReplaceData(configuredScan, narrowSchema,
info.schema())
Review Comment:
Done
--
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]
