anuragmantri commented on code in PR #55518: URL: https://github.com/apache/spark/pull/55518#discussion_r4402822957
########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java: ########## @@ -0,0 +1,111 @@ +/* + * 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}. Rows the command inserts + * 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. Spark matches each declared + * column with the session's resolver, and {@link LogicalWriteInfo#columnUpdateSchema()} uses the + * table's spelling of each column, which can differ in case from the declared reference. + * Implementations must include every column they want to receive: every column reported by + * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or + * routing. Columns that are not declared are absent from the rows the connector receives, so + * the connector must preserve their values itself. Spark may call this method more than once + * for the same operation, and every call must return the same columns. + * <p> + * If {@link RowLevelOperationInfo#updatedColumns()} is empty, the builder should build an + * operation without this mix-in. Review Comment: Superseded by Dongjoon's item 55. An empty declaration now means Spark does not narrow, so builders no longer drop the mix-in. The Javadoc says so. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala: ########## @@ -343,3 +353,63 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { } } } + +object RewriteRowLevelCommand extends SQLConfHelper { + + /** + * Resolves the connector-declared required data attributes for column-update writes against + * the given relation, in declared order and with the relation's spelling. An empty array, + * nested fields, duplicate columns, columns that do not exist in the relation (including a + * reference with no name parts) and metadata columns are rejected with an `AnalysisException`. + */ + private[sql] def resolveRequiredDataAttrs( + relation: LogicalPlan, + operation: SupportsColumnUpdates): Seq[AttributeReference] = { + val refs = operation.requiredDataAttributes.toImmutableArraySeq + if (refs.isEmpty) { Review Comment: Done. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala: ########## @@ -110,8 +111,10 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla // this rule assigns runtime filters to both scan relations (will be shared at runtime) // and must transform the runtime filter condition to use correct expr IDs for each relation // note this only applies to group-based row-level operations (i.e. ReplaceData) + // the map is built from the original table as the write table may be narrower than the Review Comment: Done. Master added the metadata mapping on this line, and the merge keeps it. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala: ########## @@ -1500,6 +1501,21 @@ private class BufferWriter(schema: StructType) extends DataWriter[InternalRow] { buffer.log.append(logEntry) } + // Tag UPDATE/COPY rows distinctly so tests can verify that Spark dispatches + // through writeColumnUpdate (the SupportsColumnUpdates path) rather than write. + override def writeColumnUpdate(metadata: InternalRow, row: InternalRow): Unit = { Review Comment: Done, both options. Mutating the strategy or `DataWithProjectionWritingSparkTask` now fails tests. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/ColumnUpdateTableSharedTests.scala: ########## @@ -0,0 +1,387 @@ +/* + * 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.{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 + } + + // `base` with the given table properties added; a given key replaces the same key in `base` + protected def tablePropsWith( + base: java.util.Map[String, String], + props: (String, String)*): java.util.Map[String, String] = { + (base.asScala ++ props).asJava + } + + // 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'") + + checkLastColumnUpdateSchema("pk", "dep", "salary") + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"), + Row(1, -1, 10, "hr") :: + Row(2, 200, 20, "software") :: + Row(3, -1, 30, "hr") :: Nil) + checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0) + } + + test("column-update: write relation has the declared columns in declared order") { + createAndInitTableWithProps(reqAttrsTableProps("salary,pk")) + + val write = executeAndKeepAnalyzedWrite { + sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'") + } + + assert(!write.query.exists(_.isInstanceOf[Union])) + checkColumnUpdateWrite(write, "salary", "pk") + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"), + Row(1, 1, "hr", 110, 10, 1) :: + Row(2, 2, "hr", 220, 20, 1) :: + Row(3, 3, "software", 300, 30, 1) :: Nil) + } + + test("column-update: the plan writes the declared columns and scans the referenced ones") { + createAndInitTableWithProps(reqAttrsTableProps("pk,salary")) + + // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are read for the + // condition, the assignment and the CHECK constraint + val write = executeAndKeepAnalyzedWrite { + executeAndCheckScan( + s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'", + s"pk INT, dep STRING, salary INT, bonus INT, extra INT$scanMetadataColumns") + } + + val writeNodeName = if (deltaUpdate) "WriteDelta" else "ReplaceData" + val header = write.simpleString(Int.MaxValue) + assert(header.matches(s"""$writeNodeName RelationV2\\[pk#\\d+, salary#\\d+\\] .*"""), header) + 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: undeclared, differently-cased columns in the RHS and WHERE clause are " + + "read with the table's own spelling") { + createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING", + """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" } + |{ "pk": 2, "salary": 200, "bonus": 20, "extra": 5, "dep": "hr" } + |""".stripMargin) + + executeAndCheckScan( + s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND EXTRA > 3", + s"pk INT, salary INT, bonus INT, extra INT, dep STRING$scanMetadataColumns") + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"), + Row(1, 110, 10, 5, "hr") :: Row(2, 200, 20, 5, "hr") :: Nil) + } + + test("column-update: assigned column outside requiredDataAttributes is rejected") { + createAndInitTableWithProps(reqAttrsTableProps("pk,salary")) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep = 'hr'") + }, + condition = "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS", + parameters = Map("connector" -> reqAttrsOperationClass, "missingColumns" -> "[`bonus`]")) + } + + test("column-update: an operation replaced after the rewrite is rejected") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |""".stripMargin) + + val write = executeAndKeepAnalyzedWrite { + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1") + } + val relation = write.table.asInstanceOf[DataSourceV2Relation] + val operationTable = relation.table.asInstanceOf[RowLevelOperationTable] + val t = operationTable.table.asInstanceOf[InMemoryRowLevelOperationTable] + val replacements: Seq[(RowLevelOperation, String)] = if (deltaUpdate) { + Seq( + t.DeltaBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) -> "[pk, salary, dep]", + new t.DeltaBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]") + } else { + Seq( + t.PartitionBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) -> "[pk, salary, dep]", + new t.PartitionBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]") + } + replacements.foreach { case (operation, expectedNames) => + assert(operation != operationTable.operation) + val replaced = write.withNewTable( + relation.copy(table = operationTable.copy(operation = operation)( + operationTable.operationInfo))) + checkError( + exception = intercept[SparkException](V2Writes(replaced)), + condition = "INTERNAL_ERROR", + parameters = Map("message" -> + (s"The write relation of operation ${operation.getClass.getName} has columns " + + s"[pk, dep, salary], but the operation expects $expectedNames."))) + } + } + + test("column-update: 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" } + |{ "pk": 2, "s": { "c1": 3, "c2": 4 }, "dep": "hr" } + |""".stripMargin) + + sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1") + + checkLastUpdatedColumns("s") + checkLastColumnUpdateSchema("pk", "dep", "s") + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"), + Row(1, Row(-1, 2), "hr") :: Row(2, Row(3, 4), "hr") :: Nil) + } + + test("column-update: excluding ColumnPruning keeps the write narrow") { + createAndInitTableWithProps(reqAttrsTableProps("pk,salary")) + + withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName) { + executeAndCheckScan( + s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'", + s"pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT$scanMetadataColumns") + } + + checkLastColumnUpdateSchema("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: reassigning a declared row-ID column keeps the row's other columns") { + createAndInitTableWithProps(reqAttrsTableProps("pk,salary")) + + sql(s"UPDATE $tableNameAsString SET pk = 2, salary = 0 WHERE pk = 1") + + // the copy-on-write fixture finds original rows by `pk`, so it falls back to full-width rows + assert(table.lastWriteInfo.columnUpdateSchema().isPresent == deltaUpdate) + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk, id"), + Row(2, 1, "hr", 0, 10, 1) :: + Row(2, 2, "hr", 200, 20, 1) :: + Row(3, 3, "software", 300, 30, 1) :: Nil) + } + + test("column-update: an UPDATE whose assignments are all identity assignments is narrowed " + + "only if the operation keeps the mix-in") { + Seq( + extraTableProps -> None, + reqAttrsTableProps("pk,salary") -> Some(Seq("pk", "salary")) + ).foreach { case (props, columnUpdateNames) => + withTable(tableNameAsString) { + createAndInitTableWithProps(props) + + sql(s"UPDATE $tableNameAsString SET salary = salary, dep = dep WHERE pk = 1") + + checkLastUpdatedColumns() + columnUpdateNames match { + case Some(names) => checkLastColumnUpdateSchema(names: _*) + case None => + assert(!table.lastWriteInfo.columnUpdateSchema().isPresent, Review Comment: Done. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala: ########## @@ -128,12 +151,61 @@ class InMemoryRowLevelOperationTable private ( copied.replacedPartitions = replacedPartitions copied.lastWriteInfo = lastWriteInfo copied.lastWriteLog = lastWriteLog + copied.lastUpdatedColumns = lastUpdatedColumns copied } override def newRowLevelOperationBuilder( info: RowLevelOperationInfo): RowLevelOperationBuilder = { - if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") { + lastUpdatedColumns = info.updatedColumns() + val splitUpdates = properties.getOrDefault(SPLIT_UPDATES, "false").toBoolean + // the copy-on-write commit finds the original row by `pk` + val updatesPk = info.updatedColumns().exists(_.describe() == "pk") + // The default column-update fixtures drop the mix-in for an UPDATE with no updated columns, + // as SupportsColumnUpdates asks. The fixed-declaration fixtures, the split row-ID fixture and + // all fixtures for DELETE and MERGE keep it. + val updatesNothing = info.command == Command.UPDATE && info.updatedColumns().isEmpty + if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") { + if (updatesNothing) { + () => DeltaBasedOperation(info.command, info.options) + } else if (splitUpdates) { + () => new DeltaBasedColumnUpdateSplitOperation( + info.command, info.updatedColumns().toSeq, info.options) + } else { + () => new DeltaBasedColumnUpdateOperation( + info.command, info.updatedColumns().toSeq, info.options) + } + } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) { + val reqCols = properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim) + if (splitUpdates) { + () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs( + info.command, reqCols, info.options) + } else { + () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command, reqCols, info.options) + } + } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false") == "true") { Review Comment: Done through 55. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala: ########## @@ -109,22 +113,98 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper { WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics) case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) => - val rowSchema = projections.rowProjection.schema + validateWriteRelation(rd) + val (rowSchema, columnUpdateSchema) = + writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd)) val metadataSchema = projections.metadataProjection.map(_.schema) val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap) - val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema) + val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema, + columnUpdateSchema) val write = writeBuilder.build() + validateColumnUpdateWrite(rd, write) val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog) rd.copy(write = Some(write), query = newQuery) case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) => + validateWriteRelation(wd) val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap) - val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections) + val deltaWriteBuilder = + newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd)) val deltaWrite = deltaWriteBuilder.build() + validateColumnUpdateWrite(wd, deltaWrite) val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog) wd.copy(write = Some(deltaWrite), query = newQuery) } + /** + * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records. + * This must match the writes whose write relation RewriteUpdateTable narrows, even when the + * declared columns cover the whole table. + */ + private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = { + write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE + } + + /** + * Checks that the write relation matches the current operation: a write that delivers narrow + * rows must have exactly the declared columns, in declared order, and any other write the full + * table columns. A rule that replaces the operation after the rewrite can break this, and the + * connector would then decode the rows with the wrong layout. + */ + private def validateWriteRelation(rowLevelWrite: RowLevelWrite): Unit = { + val writeNames = rowLevelWrite.table.output.map(_.name) + val expectedNames = rowLevelWrite.operation match { + case op: SupportsColumnUpdates if hasNarrowRows(rowLevelWrite) => + RewriteRowLevelCommand.resolveRequiredDataAttrs(rowLevelWrite.originalTable, op) + .map(_.name) + case _ => + rowLevelWrite.originalTable.output.map(_.name) + } + if (writeNames != expectedNames) { + throw SparkException.internalError( + s"The write relation of operation ${rowLevelWrite.operation.getClass.getName} has " + + s"columns ${writeNames.mkString("[", ", ", "]")}, but the operation expects " + + s"${expectedNames.mkString("[", ", ", "]")}.") + } + } + + /** + * Checks that the distribution and ordering of a write that delivers narrow rows reference only + * columns passed to the write. Other columns stay in the query only until column pruning removes + * them, so they are rejected here whether or not pruning ran. + */ + private def validateColumnUpdateWrite(rowLevelWrite: RowLevelWrite, write: Write): Unit = { + write match { + case w: RequiresDistributionAndOrdering if hasNarrowRows(rowLevelWrite) => + val distribution = w.requiredDistribution match { + case d: ClusteredDistribution => d.clustering.toImmutableArraySeq + case d: OrderedDistribution => d.ordering.toImmutableArraySeq + case _ => Seq.empty Review Comment: Done. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RowLevelWriteColumnPruningSuite.scala: ########## @@ -0,0 +1,278 @@ +/* + * 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.catalyst.optimizer + +import java.util + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.ProjectingInternalRow +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, EqualNullSafe, Expression, If, Literal, Not} +import org.apache.spark.sql.catalyst.plans.{LeftOuter, PlanTest} +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Copy, Keep, Update} +import org.apache.spark.sql.catalyst.rules.RuleExecutor +import org.apache.spark.sql.catalyst.util.{ReplaceDataProjections, WriteDeltaProjections} +import org.apache.spark.sql.catalyst.util.RowDeltaUtils._ +import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability} +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.types.{IntegerType, StructField, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * Tests the [[ColumnPruning]] case that prunes the query columns of a [[RowLevelWrite]] that its + * projections don't read and moves the projections to the ordinals of the pruned query. + */ +class RowLevelWriteColumnPruningSuite extends PlanTest { + + object Optimize extends RuleExecutor[LogicalPlan] { + val batches = Batch("Column Pruning", FixedPoint(10), + ColumnPruning, CollapseProject, RemoveNoopOperators) :: Nil + } + + // t(pk, id, dep, salary, bonus, extra) with the metadata column _partition + private val read = LocalRelation( + $"pk".int.notNull, $"id".int, $"dep".string, $"salary".int, $"bonus".int, $"extra".int, + $"_partition".string.notNull) + + // the write and original table relations; the pruning never inspects them + private val tableRelation: DataSourceV2Relation = { Review Comment: Done. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -228,4 +244,96 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { val expandOutput = generateExpandOutput(attrs, outputs) Expand(outputs, expandOutput, matchedRowsPlan) } + + /** + * Returns the table attributes that these assignments set to a new value, skipping identity + * assignments. + */ + private def collectUpdatedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = { + assignments.collect { + case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key Review Comment: Handled by SPARK-60721 on master, which rejects assignments to metadata columns at analysis. Nothing left to do here. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala: ########## @@ -273,6 +345,440 @@ 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) Review Comment: Done, except moving `pk`. That case fails on the full-width fixture too, because `InMemoryBaseTable.withDeletes` reads the key at stored ordinal 0. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala: ########## @@ -72,10 +72,12 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { protected def buildOperationTable( table: SupportsRowLevelOperations, command: Command, - options: CaseInsensitiveStringMap): RowLevelOperationTable = { - val info = RowLevelOperationInfoImpl(command, options) + options: CaseInsensitiveStringMap, + updatedAttrs: Seq[AttributeReference] = Nil): RowLevelOperationTable = { + val updatedColumns = updatedAttrs.map(attr => FieldReference(Seq(attr.name))) Review Comment: Done. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala: ########## @@ -273,6 +345,440 @@ 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) + + // The case-class equality of DeltaBasedOperation would make this operation equal to a wide + // one, or to another column-update operation, with the same command and options. + override def equals(other: Any): Boolean = other match { + case ref: AnyRef => this eq ref + case _ => false + } + override def hashCode(): Int = System.identityHashCode(this) + override def canEqual(other: Any): Boolean = false + + 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 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]
