anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3732493322
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -226,4 +389,98 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Variant of `buildDeletesAndInserts` for the `SupportsColumnUpdates`
narrow-scan path.
+ * This variant realigns the assignments to one value per surviving rowAttr
padding unassigned
+ * rowAttrs with identity, so the reinsert output arity matches the delete
output arity in
+ * the resulting Expand.
+ */
+ private def buildNarrowDeletesAndInserts(
+ matchedRowsPlan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute]): Expand = {
+
+ val (metadataAttrs, rowAttrs) = matchedRowsPlan.output.partition { attr =>
+ MetadataAttribute.isValid(attr.metadata)
+ }
+ val assignmentMap = AttributeMap(assignments.collect {
+ case a @ Assignment(key: Attribute, _) => key -> a
+ })
+ val reinsertAssignments = rowAttrs.map { attr =>
+ assignmentMap.get(attr) match {
+ case Some(a) => a
+ case None => Assignment(attr, attr)
+ }
+ }
+ val deleteOutput = deltaDeleteOutput(rowAttrs, rowIdAttrs, metadataAttrs)
+ val insertOutput = deltaReinsertOutput(reinsertAssignments, metadataAttrs)
+ val outputs = Seq(deleteOutput, insertOutput)
+ val operationTypeAttr = AttributeReference(OPERATION_COLUMN, IntegerType,
nullable = false)()
+ val attrs = operationTypeAttr +: matchedRowsPlan.output
+ val expandOutput = generateExpandOutput(attrs, outputs)
+ Expand(outputs, expandOutput, matchedRowsPlan)
+ }
+
+ /**
+ * Resolves the connector's `requiredDataAttributes()` if the operation opts
into column
+ * updates. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation): Seq[AttributeReference] = operation match
{
+ case scu: SupportsColumnUpdates => resolveRequiredDataAttrs(relation, scu)
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan
for a column-update
+ * write: connector-declared attrs, unioned with any table columns
referenced by non-identity
+ * assignment RHS expressions, the operation condition, and the table's
partition expressions.
+ * Partition-column refs are always kept so downstream rules
(V2ScanPartitioningAndOrdering,
+ * GroupBasedRowLevelOperationScanPlanning) can resolve the table's
partitioning expressions
+ * against the scan output.
+ */
+ private def computeNarrowReadAttrs(
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ assignments: Seq[Assignment],
+ cond: Expression): Seq[AttributeReference] = {
+ val relationSet = relation.outputSet
+ val nonIdentityRhsRefs = assignments.iterator
+ .filterNot(a => a.key.isInstanceOf[Attribute] &&
+ isIdentityAssignment(a.key.asInstanceOf[Attribute], a.value))
+ .flatMap(_.value.references.toSeq)
+ .toSeq
+ val extraRefs = (cond.references.toSeq ++ nonIdentityRhsRefs)
+ .collect { case a: AttributeReference => a }
+ .filter(relationSet.contains)
+ val partitionRefNames = relation.table.partitioning().toImmutableArraySeq
Review Comment:
I have implemented Option 1 in
https://github.com/apache/spark/pull/55518/commits/cef886e16b7fb400b90841caa8975355f4c7ff42
please let me know what you think.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can
implement this interface to
+ * receive a narrow row containing only the columns declared via {@link
#requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ *
+ * @since 4.3.0
Review Comment:
I was hoping to maybe still land this 4.3 as I requested
[here](https://lists.apache.org/thread/0d1b7gg3bd7r6ml1m0fs4z0tcry23ost) since
this PR has been around for a while. I don't know if we have missed the train
already. I will let you and other reviewers decide. Let me know if we should
move on to Spark 4.4. I will update this.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,541 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector
+
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.connector.catalog.{CatalogV2Util, TableInfo}
+import
org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity,
reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField,
StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that return true from
+ *
[[org.apache.spark.sql.connector.write.RowLevelOperation#supportsColumnUpdates]].
Review Comment:
Updated the comments. Thanks.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +82,42 @@ 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}; the
default delegates
+ * to {@link #write(Object, Object)} so existing connectors are unaffected.
+ * <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.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ write(metadata, record);
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Equivalent to {@link #writeUpdate(Object, Object)} for writers that do
not require metadata.
+ * The default delegates to {@link #write(Object)}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be
called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like
writing files.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T record) throws IOException {
Review Comment:
Done. Implementation must override this or we will throw.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -418,7 +432,29 @@ case class ReplaceData(
// validates row projection output is compatible with table attributes
private def rowAttrsResolved: Boolean = {
val inRowAttrs =
DataTypeUtils.toAttributes(projections.rowProjection.schema)
- table.skipSchemaResolution || areCompatible(inRowAttrs, table.output)
+ val inUpdateAttrs = projections.updateRowProjection match {
+ case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
+ case None => Nil
+ }
+ // `rowProjection` (INSERT-tagged rows) validates against `table.output`
-- the full table
+ // shape. `updateRowProjection` (UPDATE/COPY-tagged rows) is narrow for
column-update
+ // connectors, so it validates against `projectedDataAttrs` (the
connector-declared narrow
+ // set) instead. When the connector does not mix in
`SupportsColumnUpdates`,
+ // `updateRowProjection` is absent and `updateResolved` is trivially true.
+ val insertResolved = table.skipSchemaResolution || inRowAttrs.isEmpty ||
Review Comment:
Done, updated all the sites you mentioned.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuite.scala:
##########
@@ -45,11 +45,13 @@ class DeltaBasedUpdateTableSuite extends
DeltaBasedUpdateTableSuiteBase {
sql(s"SELECT * FROM $tableNameAsString"),
Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // info.schema() reflects the row projection: `id` is non-nullable because
the assignment
+ // supplies a non-null literal, matching master's projection-derived write
schema semantics.
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -226,4 +389,98 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Variant of `buildDeletesAndInserts` for the `SupportsColumnUpdates`
narrow-scan path.
+ * This variant realigns the assignments to one value per surviving rowAttr
padding unassigned
+ * rowAttrs with identity, so the reinsert output arity matches the delete
output arity in
+ * the resulting Expand.
+ */
+ private def buildNarrowDeletesAndInserts(
+ matchedRowsPlan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute]): Expand = {
+
+ val (metadataAttrs, rowAttrs) = matchedRowsPlan.output.partition { attr =>
+ MetadataAttribute.isValid(attr.metadata)
+ }
+ val assignmentMap = AttributeMap(assignments.collect {
+ case a @ Assignment(key: Attribute, _) => key -> a
+ })
+ val reinsertAssignments = rowAttrs.map { attr =>
+ assignmentMap.get(attr) match {
+ case Some(a) => a
+ case None => Assignment(attr, attr)
+ }
+ }
+ val deleteOutput = deltaDeleteOutput(rowAttrs, rowIdAttrs, metadataAttrs)
+ val insertOutput = deltaReinsertOutput(reinsertAssignments, metadataAttrs)
Review Comment:
Thanks. I'm going with the smaller option (a). Rejecting the combination is
contained enough for this PR and preserves narrow writes for every non-row-ID
split-update case.
I added test coverage for this in `DeltaBasedColumnUpdateTableSuite`
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can
implement this interface to
+ * receive a narrow row containing only the columns declared via {@link
#requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
Review Comment:
Added the lines on all mentioned API docs.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -267,13 +268,97 @@ 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 = {
+ catalog.loadTable(ident).asInstanceOf[InMemoryTable]
+ .commits.last.writeSummary.get
+ .asInstanceOf[org.apache.spark.sql.connector.write.UpdateSummary]
+ }
+
+ /**
+ * Asserts the last UPDATE's write summary metrics. `deltaUpdate` controls
the expected
+ * COPY count: MoR connectors emit only deltas (no COPY rows) so
`numCopiedRows` is forced
+ * to 0; CoW connectors emit COPY rows for unchanged rows in matched groups.
+ */
+ protected def checkUpdateMetrics(
+ numUpdatedRows: Long,
+ numCopiedRows: Long,
+ deltaUpdate: Boolean = false): Unit = {
+ val summary = getUpdateSummary()
+ assert(summary.numUpdatedRows() === numUpdatedRows,
+ s"Expected numUpdatedRows=$numUpdatedRows, got
${summary.numUpdatedRows()}")
+ val expectedCopied = if (deltaUpdate) 0L else numCopiedRows
+ assert(summary.numCopiedRows() === expectedCopied,
+ s"Expected numCopiedRows=$expectedCopied, got
${summary.numCopiedRows()}")
+ }
+
+ /**
+ * Asserts that the column names in RowLevelOperationInfo.updatedColumns()
received by the
+ * last operation match exactly the expected set. Order is ignored.
+ */
+ protected def checkLastUpdatedColumns(expectedNames: String*): Unit = {
+ val actual = table.lastUpdatedColumns.map(_.describe()).toSet
+ val expected = expectedNames.toSet
+ assert(actual == expected,
+ s"updatedColumns mismatch: expected ${expected.mkString("[", ", ", "]")}
" +
+ s"but got ${actual.mkString("[", ", ", "]")}")
+ }
+
+ /**
+ * Asserts the set of top-level column names present in the last connector
scan schema.
+ * Metadata columns are filtered out so tests can focus on data-column
narrowing without
+ * having to enumerate `_partition` / `index` on every assertion.
+ */
+ protected def checkLastScanDataColumns(expectedNames: String*): Unit = {
Review Comment:
Verified there are no callers and removed this method.
--
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]
