dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3742227035
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -2461,6 +2467,12 @@
],
"sqlState" : "KD009"
},
+ "EMPTY_REQUIRED_DATA_ATTRIBUTES" : {
Review Comment:
nit. The five new error conditions for the same `SupportsColumnUpdates`
contract use three different prefixes (`EMPTY_REQUIRED_DATA_ATTRIBUTES`,
`REQUIRED_DATA_ATTRIBUTES_*`, `SPLIT_UPDATE_*`). Could you consider a common
prefix or a single parent condition with sub-conditions so they group together
and are easier to discover?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand
{
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and
optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs,
scanOnlyDataAttrs)
+ }
+
+
Review Comment:
nit. Redundant empty line.
```suggestion
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand
{
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and
optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs,
scanOnlyDataAttrs)
+ }
+
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable,
metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, assignments, rowIdAttrs)
+ validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs)
+ }
+
+ val narrowDataAttrs = if (supportsColumnUpdate) {
+ computeNarrowReadAttrs(relation, connectorDataAttrs, scanOnlyDataAttrs,
assignments, cond)
+ } else {
+ relation.output
+ }
+
+ val readRelation = if (supportsColumnUpdate) {
+ buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs,
metadataAttrs,
+ rowIdAttrs)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs,
rowIdAttrs)
+ }
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
- val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
- buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ val rowDeltaPlan = if (supportsColumnUpdate) {
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildNarrowDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs,
connectorDataAttrs,
+ scanOnlyDataAttrs)
+ }
} else {
- buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments,
rowIdAttrs)
+ }
}
// build a plan to write the row delta to the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildWriteDeltaProjections(rowDeltaPlan, rowAttrs,
rowIdAttrs, metadataAttrs)
+ val projections = buildWriteDeltaProjections(
+ rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs, connectorDataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections,
groupFilterCond)
}
+ /**
+ * Builds the WriteDelta projection for the column-update path.
+ */
+ private def buildColumnUpdateProjection(
+ plan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute],
+ metadataAttrs: Seq[Attribute],
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference]): LogicalPlan = {
+
+ val assignedValues = assignments.collect {
+ case Assignment(key: Attribute, value) if !isIdentityAssignment(key,
value) =>
+ Alias(value, key.name)()
+ }
+
+ // Connector-required columns whose value isn't being changed by the
UPDATE: pass through the
+ // current value so the connector receives a complete write row. Row-ID
columns are excluded
+ // here even when also connector-required (e.g. a primary key used for
both row lookup and
+ // write payload) they are emitted once, below, via rowIdValues.
+ val assignedKeyIds = assignments.collect {
+ case Assignment(key: AttributeReference, value) if
!isIdentityAssignment(key, value) =>
+ key.exprId
+ }.toSet
+ val rowIdAttrSet = AttributeSet(rowIdAttrs)
+ val connectorPassThroughValues = connectorDataAttrs.filterNot(a =>
+ assignedKeyIds.contains(a.exprId) || rowIdAttrSet.contains(a))
+
+ // scanOnlyDataAttrs are never assigned carry them through the write query
so
Review Comment:
nit. Missing punctuation here and at line 341 (`write payload) they are
emitted once` -> `write payload); they are emitted once`).
```suggestion
// scanOnlyDataAttrs are never assigned; carry them through the write
query so
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -771,6 +777,7 @@ trait WritingSparkTask[W <: DataWriter[InternalRow]]
extends Logging with Serial
case class DataAndMetadataWritingSparkTask(
dataProj: ProjectingInternalRow,
+ updateDataProj: ProjectingInternalRow,
Review Comment:
nit (optional). A nullable `updateDataProj: ProjectingInternalRow`
constructor parameter (here and in `DataWithProjectionWritingSparkTask`) is
consistent with the existing `orNull` usage in the delta tasks, but
`Option[ProjectingInternalRow]` would be more idiomatic for a case class
constructor. Feel free to keep as is if you prefer the symmetry.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -146,6 +203,48 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
Project(operationCol +: updatedValues, plan)
}
+ /**
+ * Variant of `buildReplaceDataUpdateProjection` for the
`SupportsColumnUpdates` narrow-scan
+ * path.
+ * For narrow attributes, looks up assignments by `ExprId` via
`AttributeMap`and passes through
Review Comment:
nit. Missing space.
```suggestion
* For narrow attributes, looks up assignments by `ExprId` via
`AttributeMap` and passes through
```
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here
in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must
override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be
called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like
writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ throw new SparkUnsupportedOperationException(
+ "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED",
+ Map.of("class", getClass().getName()));
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Equivalent to {@link #writeUpdate(Object, Object)} for writers that do
not require metadata.
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here
in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must
override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be
called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like
writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T record) throws IOException {
Review Comment:
`DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED` is defined and thrown here (and
at `:103`), but no test exercises this path. Could you add a test where a
connector mixes in `SupportsColumnUpdates` without overriding `writeUpdate`,
and assert this error condition is raised? That keeps the new error condition
covered and guards the dispatch contract.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -267,13 +268,80 @@ abstract class RowLevelOperationSuiteBase
protected def checkLastWriteInfo(
expectedRowSchema: StructType = new StructType(),
expectedRowIdSchema: Option[StructType] = None,
- expectedMetadataSchema: Option[StructType] = None): Unit = {
+ expectedMetadataSchema: Option[StructType] = None,
+ expectedUpdateSchema: Option[StructType] = None): Unit = {
val info = table.lastWriteInfo
assert(info.schema == expectedRowSchema, "row schema must match")
val actualRowIdSchema = Option(info.rowIdSchema.orElse(null))
assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must
match")
val actualMetadataSchema = Option(info.metadataSchema.orElse(null))
assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema
must match")
+ val actualUpdateSchema = Option(info.updateSchema.orElse(null))
+ assert(actualUpdateSchema == expectedUpdateSchema, "update schema must
match")
+ }
+
+ protected def getUpdateSummary():
org.apache.spark.sql.connector.write.UpdateSummary = {
Review Comment:
nit. Please import `org.apache.spark.sql.connector.write.UpdateSummary`
instead of using the fully-qualified name inline (here and at line 286).
--
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]
