Narrow only the write relation for column updates and address review comments

Rework column-level UPDATE narrowing so that RewriteUpdateTable keeps master's
read relation and plan builders and narrows only the write relation to
requiredDataAttributes(). The rewrite orders the write query so the columns
the write reads come first, and the ColumnPruning case for RowLevelWrite
prunes only the columns after the last one the write reads.

- Rename DataWriter.writeUpdate to writeColumnUpdate and
  LogicalWriteInfo.updateSchema to columnUpdateSchema, and the related error
  to DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED.
- Reject a write whose distribution or ordering references a column the
  column update does not read (COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS),
  and a declared column that does not exist
  (COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE).
- Plan copy-on-write scans from the columns the write query reads for writes
  that deliver narrow rows; other writes read every column as before.
- Build the runtime group filter attribute map from the original table.
- Mark the new APIs @SInCE 4.4.0 and move the MiMa exclude to 4.4.
- Fix Javadoc and error messages, and add tests for the rows and metadata the
  connector receives, cross-column assignments, case-sensitive analysis and
  MERGE updatedColumns().

Author: anuragmantri

common/utils/src/main/resources/error/error-conditions.json

@@ -1260,6 +1260,12 @@
     ],
     "sqlState" : "42000"
   },
+  "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE" : {
+    "message" : [
+      "Connector <connector> mixes in `SupportsColumnUpdates` but declared metadata column(s) <metadataAttributes> in `requiredDataAttributes()`. Declare metadata columns in `requiredMetadataAttributes()` instead."
+    ],
+    "sqlState" : "42000"
+  },
   "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE" : {
     "message" : [
       "Connector <connector> mixes in `SupportsColumnUpdates` but declared nested field(s) <nestedAttributes> in `requiredDataAttributes()`. Declare the root struct column instead of a nested field."
@@ -1274,16 +1280,28 @@
   },
   "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED" : {
     "message" : [
-      "Data source <connector> represents UPDATE as delete + insert but its `requiredDataAttributes()` does not include row ID column(s) <rowIds>. The REINSERT payload is projected down to `requiredDataAttributes()`, so an undeclared row ID column leaves the reinserted row with no identity for the connector to place it by. Include the row ID column(s) in `requiredDataAttributes()`."
+      "Connector <connector> mixes in `SupportsColumnUpdates` and represents UPDATE as delete and insert, but row ID column(s) <rowIds> do not reach the reinserted row, which then has no identity for the connector to place it by. Declare each data row ID column in `requiredDataAttributes()`, and return each metadata row ID column from `requiredMetadataAttributes()` with `MetadataColumn.PRESERVE_ON_REINSERT` set."
     ],
     "sqlState" : "42000"
   },
   "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT" : {
     "message" : [
-      "Cannot reassign row ID column(s) <rowIds> in a column-level UPDATE when data source <connector> represents UPDATE as delete + insert. The REINSERT path has no row-ID channel to reconstruct columns outside `requiredDataAttributes()`, so undeclared columns would become NULL. Either avoid reassigning row ID columns, or include every table column in `requiredDataAttributes()`."
+      "Connector <connector> mixes in `SupportsColumnUpdates` and represents UPDATE as delete and insert, so UPDATE cannot assign row ID column(s) <rowIds>. The reinserted row carries only the columns in `requiredDataAttributes()`, and with a new row ID it cannot be matched to the row whose other columns the connector must preserve. Do not assign row ID columns, or include every table column in `requiredDataAttributes()`."
+    ],
+    "sqlState" : "42000"
+  },
+  "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS" : {
+    "message" : [
+      "Connector <connector> mixes in `SupportsColumnUpdates`, but its write requires a distribution or ordering by column(s) <columns> that the column-level UPDATE does not read. Declare each data column in `requiredDataAttributes()` and each metadata column in `requiredMetadataAttributes()`."
     ],
     "sqlState" : "42000"
   },
+  "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE" : {
+    "message" : [
+      "Connector <connector> mixes in `SupportsColumnUpdates` but declared column(s) <unknownAttributes> in `requiredDataAttributes()` that do not exist in the table."
+    ],
+    "sqlState" : "42703"
+  },
   "COMPARATOR_RETURNS_NULL" : {
     "message" : [
       "The comparator has returned a NULL for a comparison between <firstValue> and <secondValue>.",
@@ -2433,9 +2451,9 @@
     ],
     "sqlState" : "42K03"
   },
-  "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED" : {
+  "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED" : {
     "message" : [
-      "<class> mixes in `SupportsColumnUpdates` but does not implement `writeUpdate`. Connectors that opt into narrow column-level updates must override `writeUpdate` to handle rows in the `LogicalWriteInfo.updateSchema()` layout; the default implementation would forward narrow rows to `write`, which expects the full `LogicalWriteInfo.schema()` layout."
+      "<class> does not override `writeColumnUpdate(record)`. A data writer that receives rows in the `LogicalWriteInfo.columnUpdateSchema()` layout must override `writeColumnUpdate(record)`, or `writeColumnUpdate(metadata, record)` if the operation returns metadata columns from `requiredMetadataAttributes()`."
     ],
     "sqlState" : "0A000"
   },

project/MimaExcludes.scala

@@ -46,7 +46,10 @@ object MimaExcludes {
       "org.apache.spark.ml.regression.DecisionTreeRegressionModel.numLeave"),
     // [SPARK-59154] Remove unused prediction variance helper after inlining its implementation.
     ProblemFilters.exclude[DirectMissingMethodProblem](
-      "org.apache.spark.ml.regression.DecisionTreeRegressionModel.predictVariance")
+      "org.apache.spark.ml.regression.DecisionTreeRegressionModel.predictVariance"),
+    // [SPARK-58111] Write schema narrowing for column-level UPDATE in DSv2
+    ProblemFilters.exclude[ReversedMissingMethodProblem](
+      "org.apache.spark.sql.connector.write.RowLevelOperationInfo.updatedColumns")
   )
 
   // Exclude rules for 4.3.x from 4.2.0 (add 4.3-specific filters below as needed).
@@ -73,10 +76,7 @@ object MimaExcludes {
     // [SPARK-57987] Add desc field to the SQL REST API Node case class
     ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.sql.Node.apply"),
     ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.sql.Node.copy"),
-    ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.status.api.v1.sql.Node$"),
-    // [SPARK-58111] Add scan and write schema narrowing for column-level UPDATEs in DSv2
-    ProblemFilters.exclude[ReversedMissingMethodProblem](
-      "org.apache.spark.sql.connector.write.RowLevelOperationInfo.updatedColumns")
+    ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.status.api.v1.sql.Node$")
   )
 
   // Exclude rules for 4.2.x from 4.1.0

sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java

@@ -85,50 +85,56 @@ default void write(T metadata, T record) throws IOException {
   }
 
   /**
-   * Writes one updated, copied, or reinserted record with metadata.
+   * Writes one updated or copied record with metadata in the column update layout.
    * <p>
-   * Connectors that mix in {@link SupportsColumnUpdates} without also mixing in
-   * {@link SupportsDelta} (i.e. group-based operations) receive records here in the schema
-   * declared by {@link LogicalWriteInfo#updateSchema()}. An operation that mixes in both
-   * instead delivers narrow rows through {@link DeltaWriter#update} / {@link DeltaWriter#reinsert}
-   * and never calls this method. By default, delegates to {@link #writeUpdate(Object)} for
-   * connectors that do not need metadata; implementations that do need it should override this
-   * method directly.
+   * When {@link LogicalWriteInfo#columnUpdateSchema()} is present for a row-level operation that
+   * does not mix in {@link SupportsDelta}, Spark passes updated and copied records to this method
+   * instead of {@link #write(Object, Object)} if the operation returns a non-empty
+   * {@link RowLevelOperation#requiredMetadataAttributes()}, and to
+   * {@link #writeColumnUpdate(Object)} otherwise. The record follows
+   * {@link LogicalWriteInfo#columnUpdateSchema()} and the metadata follows
+   * {@link LogicalWriteInfo#metadataSchema()}. Operations that mix in {@link SupportsDelta} receive
+   * such rows through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert} instead.
+   * <p>
+   * By default, delegates to {@link #writeColumnUpdate(Object)} and drops the metadata.
    * <p>
    * If this method fails (by throwing an exception), {@link #abort()} will be called and this
    * data writer is considered to have been failed.
    *
    * @throws IOException if failure happens during disk/network IO like writing files.
-   * @throws SparkUnsupportedOperationException if the connector mixes in
-   *         {@link SupportsColumnUpdates} but overrides neither {@code writeUpdate} overload.
+   * @throws SparkUnsupportedOperationException if neither this method nor
+   *         {@link #writeColumnUpdate(Object)} is overridden.
    *
-   * @since 4.3.0
+   * @since 4.4.0
    */
-  default void writeUpdate(T metadata, T record) throws IOException {
-    writeUpdate(record);
+  default void writeColumnUpdate(T metadata, T record) throws IOException {
+    writeColumnUpdate(record);
   }
 
   /**
-   * Writes one updated, copied, or reinserted record without metadata.
+   * Writes one updated or copied record without metadata in the column update layout.
+   * <p>
+   * When {@link LogicalWriteInfo#columnUpdateSchema()} is present for a row-level operation that
+   * does not mix in {@link SupportsDelta}, Spark passes updated and copied records to this method
+   * instead of {@link #write(Object)} if the operation returns no
+   * {@link RowLevelOperation#requiredMetadataAttributes()}. The record follows
+   * {@link LogicalWriteInfo#columnUpdateSchema()}.
    * <p>
-   * Connectors that mix in {@link SupportsColumnUpdates} without also mixing in
-   * {@link SupportsDelta} (i.e. group-based operations) receive records here in the schema
-   * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this method,
-   * or {@link #writeUpdate(Object, Object)}, when mixing in {@link SupportsColumnUpdates} without
-   * {@link SupportsDelta}.
+   * A writer for such an operation must override this method, unless the operation returns a
+   * non-empty {@link RowLevelOperation#requiredMetadataAttributes()} and the writer overrides
+   * {@link #writeColumnUpdate(Object, Object)}.
    * <p>
    * If this method fails (by throwing an exception), {@link #abort()} will be called and this
    * data writer is considered to have been failed.
    *
    * @throws IOException if failure happens during disk/network IO like writing files.
-   * @throws SparkUnsupportedOperationException if the connector mixes in
-   *         {@link SupportsColumnUpdates} but overrides neither {@code writeUpdate} overload.
+   * @throws SparkUnsupportedOperationException if this method is not overridden.
    *
-   * @since 4.3.0
+   * @since 4.4.0
    */
-  default void writeUpdate(T record) throws IOException {
+  default void writeColumnUpdate(T record) throws IOException {
     throw new SparkUnsupportedOperationException(
-      "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED",
+      "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
       Map.of("class", getClass().getName()));
   }
 

sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DeltaWriter.java

@@ -41,10 +41,8 @@ public interface DeltaWriter<T> extends DataWriter<T> {
   /**
    * Updates a row.
    * <p>
-   * When {@link LogicalWriteInfo#updateSchema()} is present, the {@code row} follows the narrow
-   * layout it declares rather than the full table schema from {@link LogicalWriteInfo#schema()}.
-   * It is present only for UPDATE on an operation that mixes in {@link SupportsColumnUpdates};
-   * otherwise {@code row} follows the full table schema.
+   * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, the {@code row} follows it;
+   * otherwise it follows {@link LogicalWriteInfo#schema()}.
    *
    * @param metadata values for metadata columns that were projected but are not part of the row ID
    * @param id a row ID to update
@@ -58,10 +56,8 @@ public interface DeltaWriter<T> extends DataWriter<T> {
    * <p>
    * This method handles the insert portion of updated rows split into deletes and inserts.
    * <p>
-   * When {@link LogicalWriteInfo#updateSchema()} is present, the {@code row} follows the narrow
-   * layout it declares rather than the full table schema from {@link LogicalWriteInfo#schema()}.
-   * It is present only for UPDATE on an operation that mixes in {@link SupportsColumnUpdates};
-   * otherwise {@code row} follows the full table schema.
+   * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, the {@code row} follows it;
+   * otherwise it follows {@link LogicalWriteInfo#schema()}.
    *
    * @param metadata values for metadata columns
    * @param row a row to reinsert

sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java

@@ -47,6 +47,10 @@ public interface LogicalWriteInfo {
 
   /**
    * the schema of the input data from Spark to data source.
+   * <p>
+   * When {@link #columnUpdateSchema()} is present, this schema covers only newly inserted rows, as
+   * updated, copied, and reinserted rows follow {@link #columnUpdateSchema()}. It is then empty
+   * when the command inserts no new rows, such as UPDATE.
    */
   StructType schema();
 
@@ -67,12 +71,16 @@ default Optional<StructType> metadataSchema() {
   }
 
   /**
-   * the narrow schema for updates. Present only when the connector mixes in
-   * {@link SupportsColumnUpdates}.
+   * the schema of updated, copied, and reinserted rows from Spark to data source in a
+   * column-level update. Present when the operation mixes in {@link SupportsColumnUpdates} and
+   * Spark delivers these rows with the columns of
+   * {@link SupportsColumnUpdates#requiredDataAttributes()}, which currently happens only for
+   * UPDATE. It covers every table column if every column is declared. When present,
+   * {@link #schema()} covers only newly inserted rows.
    *
-   * @since 4.3.0
+   * @since 4.4.0
    */
-  default Optional<StructType> updateSchema() {
+  default Optional<StructType> columnUpdateSchema() {
     return Optional.empty();
   }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java

@@ -40,13 +40,14 @@ public interface RowLevelOperationInfo {
   Command command();
 
   /**
-   * Returns the columns being updated by this operation. Currently populated only for UPDATE;
-   * DELETE and MERGE report an empty array.
+   * Returns the columns being updated by this operation. Currently only UPDATE populates it;
+   * other commands report an empty array.
    * <p>
-   * Nested struct field updates are reported at root-column granularity
-   * (e.g. {@code SET s.c1 = -1} returns {@code s}).
+   * A column is reported only if it is assigned a new value, so identity assignments such as
+   * {@code SET a = a} are excluded. Nested struct field updates are reported at root-column
+   * granularity (e.g. {@code SET s.c1 = -1} returns {@code s}).
    *
-   * @since 4.3.0
+   * @since 4.4.0
    */
   NamedReference[] updatedColumns();
 }