anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4392772308
##########
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:
Done.
##########
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:
Done.
##########
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:
Done.
##########
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:
Done.
##########
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:
Done.
##########
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:
I went with the fallback. The unread-filter branch now honors
`useCatalystRuntimeFiltering`, and its test runs with the flag off and on and
checks the scan type. Sharing the other parts would mean changing master's
`PartitionBasedOperation`, `PartitionBasedReplaceData` and
`newRowLevelScanBuilder`, so I left them as they are.
##########
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:
Done.
##########
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:
Done.
##########
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:
Done.
##########
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:
Done. I also overrode `canEqual`, because with only `equals` and `hashCode`
the wide operation's case-class `equals` still matched the narrow one. The swap
test now builds its replacements from the plan's table and asserts each one
differs from the original operation. Before, `table` returned a different
instance, so the pair never compared equal there.
##########
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:
Done.
##########
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:
Done.
##########
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:
Done.
##########
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:
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]
