This is an automated email from the ASF dual-hosted git repository.
amoghj pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/main by this push:
new 56d81e72ff Spark 3.4: Backport UPDATE/MERGE logic for row lineage
(#13344)
56d81e72ff is described below
commit 56d81e72ff138d026fc6c79575d280a6a9b0bf48
Author: Drew Gallardo <[email protected]>
AuthorDate: Fri Jun 20 08:51:38 2025 +0700
Spark 3.4: Backport UPDATE/MERGE logic for row lineage (#13344)
---
.../extensions/IcebergSparkSessionExtensions.scala | 6 +
.../RewriteMergeIntoTableForRowLineage.scala | 71 +++
.../analysis/RewriteOperationForRowLineage.scala | 87 ++++
.../analysis/RewriteUpdateTableForRowLineage.scala | 53 +++
.../RemoveRowLineageOutputFromOriginalTable.scala | 54 +++
.../SparkRowLevelOperationsTestBase.java | 14 +-
.../extensions/TestCopyOnWriteWithLineage.java | 35 ++
.../extensions/TestMergeOnReadWithLineage.java | 35 ++
.../TestRowLevelOperationsWithLineage.java | 510 +++++++++++++++++++++
.../spark/source/SparkCopyOnWriteOperation.java | 20 +-
.../iceberg/spark/source/SparkCopyOnWriteScan.java | 40 ++
.../spark/source/SparkPositionDeltaOperation.java | 17 +-
.../source/SparkPositionDeltaWriteBuilder.java | 13 +-
.../apache/iceberg/spark/source/SparkTable.java | 5 +-
.../iceberg/spark/source/SparkWriteBuilder.java | 24 +-
15 files changed, 967 insertions(+), 17 deletions(-)
diff --git
a/spark/v3.4/spark-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala
index af6a64911e..350ed6988f 100644
---
a/spark/v3.4/spark-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala
+++
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala
@@ -30,9 +30,12 @@ import
org.apache.spark.sql.catalyst.analysis.ResolveMergeIntoTableReferences
import org.apache.spark.sql.catalyst.analysis.ResolveProcedures
import org.apache.spark.sql.catalyst.analysis.ResolveViews
import org.apache.spark.sql.catalyst.analysis.RewriteMergeIntoTable
+import
org.apache.spark.sql.catalyst.analysis.RewriteMergeIntoTableForRowLineage
import org.apache.spark.sql.catalyst.analysis.RewriteUpdateTable
+import org.apache.spark.sql.catalyst.analysis.RewriteUpdateTableForRowLineage
import
org.apache.spark.sql.catalyst.optimizer.ExtendedReplaceNullWithFalseInPredicate
import
org.apache.spark.sql.catalyst.optimizer.ExtendedSimplifyConditionalsInPredicate
+import
org.apache.spark.sql.catalyst.optimizer.RemoveRowLineageOutputFromOriginalTable
import org.apache.spark.sql.catalyst.optimizer.ReplaceStaticInvoke
import
org.apache.spark.sql.catalyst.parser.extensions.IcebergSparkSqlExtensionsParser
import
org.apache.spark.sql.execution.datasources.v2.ExtendedDataSourceV2Strategy
@@ -54,6 +57,8 @@ class IcebergSparkSessionExtensions extends
(SparkSessionExtensions => Unit) {
extensions.injectResolutionRule { _ => CheckMergeIntoTableConditions }
extensions.injectResolutionRule { _ => ProcedureArgumentCoercion }
extensions.injectResolutionRule { _ => AlignRowLevelCommandAssignments }
+ extensions.injectResolutionRule { _ => RewriteUpdateTableForRowLineage}
+ extensions.injectResolutionRule { _ => RewriteMergeIntoTableForRowLineage}
extensions.injectResolutionRule { _ => RewriteUpdateTable }
extensions.injectResolutionRule { _ => RewriteMergeIntoTable }
extensions.injectCheckRule { _ => CheckViews }
@@ -64,6 +69,7 @@ class IcebergSparkSessionExtensions extends
(SparkSessionExtensions => Unit) {
extensions.injectOptimizerRule { _ =>
ExtendedSimplifyConditionalsInPredicate }
extensions.injectOptimizerRule { _ =>
ExtendedReplaceNullWithFalseInPredicate }
extensions.injectOptimizerRule { _ => ReplaceStaticInvoke }
+ extensions.injectOptimizerRule { _ =>
RemoveRowLineageOutputFromOriginalTable}
// pre-CBO rules run only once and the order of the rules is important
// - dynamic filters should be added before replacing commands with
rewrite plans
// - scans must be planned before building writes
diff --git
a/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala
new file mode 100644
index 0000000000..cebb360a29
--- /dev/null
+++
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.catalyst.analysis
+
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.plans.logical.Assignment
+import org.apache.spark.sql.catalyst.plans.logical.InsertAction
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.MergeIntoIcebergTable
+import org.apache.spark.sql.catalyst.plans.logical.UpdateAction
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+
+object RewriteMergeIntoTableForRowLineage extends
RewriteOperationForRowLineage {
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ plan.resolveOperators {
+ case m @ MergeIntoIcebergTable(_, _, _, matchedActions, _, _)
+ if m.resolved && m.aligned &&
+ matchedActions.nonEmpty &&
+ shouldUpdatePlan(m.targetTable) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable:
MergeIntoIcebergTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedActions = mergeIntoTable.notMatchedActions
+ val (rowId, lastUpdatedSequenceNumber) =
findRowLineageAttributes(r.metadataOutput).get
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, assignments) =>
+ UpdateAction(cond, assignments ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequenceNumber, Literal(null))))
+
+ case deleteAction => deleteAction
+ }
+
+ val notMatchedActionsForLineage = notMatchedActions.map {
+ case InsertAction(cond, assignments) =>
+ InsertAction(cond, assignments ++ Seq(
+ Assignment(rowId, Literal(null)),
+ Assignment(lastUpdatedSequenceNumber, Literal(null))))
+ }
+
+ val tableWithLineage = r.copy(output = r.output ++ Seq(rowId,
lastUpdatedSequenceNumber))
+ mergeIntoTable.copy(
+ targetTable = tableWithLineage,
+ matchedActions = matchedAssignmentsForLineage,
+ notMatchedActions = notMatchedActionsForLineage)
+ }
+ }
+}
diff --git
a/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteOperationForRowLineage.scala
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteOperationForRowLineage.scala
new file mode 100644
index 0000000000..611e62b75e
--- /dev/null
+++
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteOperationForRowLineage.scala
@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.catalyst.analysis
+
+import org.apache.iceberg.MetadataColumns
+import org.apache.iceberg.TableUtil
+import org.apache.iceberg.spark.source.SparkTable
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.util.METADATA_COL_ATTR_KEY
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.types.MetadataBuilder
+
+trait RewriteOperationForRowLineage extends RewriteRowLevelIcebergCommand {
+
+ protected val ROW_ID_ATTRIBUTE_NAME = MetadataColumns.ROW_ID.name()
+ protected val LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME =
MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()
+
+ // The plan should only be updated if the underlying Iceberg table supports
row lineage AND
+ // lineage attributes are not already on the output of operation which
indicates the rule already ran
+ protected def shouldUpdatePlan(table: LogicalPlan): Boolean = {
+ val supportsRowLineage = EliminateSubqueryAliases(table) match {
+ case r: DataSourceV2Relation =>
+ r.table match {
+ case sparkTable: SparkTable =>
+ TableUtil.supportsRowLineage(sparkTable.table())
+ }
+ case _ => false
+ }
+
+ val rowIdAbsentFromOutput = !table.output.exists(_.name ==
ROW_ID_ATTRIBUTE_NAME)
+
+ supportsRowLineage && rowIdAbsentFromOutput
+ }
+
+ protected def findRowLineageAttributes(
+ expressions: Seq[Expression]
+ ): Option[(AttributeReference, AttributeReference)] = {
+ val rowIdAttr = expressions.collectFirst {
+ case attr: AttributeReference
+ if isMetadataColumn(attr) && attr.name == ROW_ID_ATTRIBUTE_NAME =>
attr
+ }
+
+ val lastUpdatedAttr = expressions.collectFirst {
+ case attr: AttributeReference
+ if isMetadataColumn(attr) && attr.name ==
LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME => attr
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata
attribute
+ // This works around the logic in ExposesMetadataColumns,
+ // which prevents surfacing other metadata columns when a single metadata
column is in the output
+ (rowIdAttr, lastUpdatedAttr) match {
+ case (Some(rowId), Some(lastUpdated)) =>
+ Some((removeMetadataColumnAttribute(rowId),
removeMetadataColumnAttribute(lastUpdated)))
+ case _ => None
+ }
+ }
+
+ protected def removeMetadataColumnAttribute(attr: AttributeReference):
AttributeReference = {
+ attr.withMetadata(
+ new MetadataBuilder()
+ .withMetadata(attr.metadata)
+ .remove(METADATA_COL_ATTR_KEY).build())
+ }
+
+ private def isMetadataColumn(attributeReference: AttributeReference):
Boolean = {
+ attributeReference.metadata.contains(METADATA_COL_ATTR_KEY)
+ }
+}
diff --git
a/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTableForRowLineage.scala
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTableForRowLineage.scala
new file mode 100644
index 0000000000..00513d89fc
--- /dev/null
+++
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTableForRowLineage.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.catalyst.analysis
+
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.plans.logical.Assignment
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.UpdateIcebergTable
+import org.apache.spark.sql.connector.catalog.SupportsRowLevelOperations
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+
+object RewriteUpdateTableForRowLineage extends RewriteOperationForRowLineage {
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ plan resolveOperators {
+ case updateTable@UpdateIcebergTable(_, _, _, _) if
shouldUpdatePlan(updateTable.table) =>
+ updatePlanWithRowLineage(updateTable)
+ }
+ }
+
+
+ private def updatePlanWithRowLineage(updateTable: UpdateIcebergTable):
LogicalPlan = {
+ EliminateSubqueryAliases(updateTable.table) match {
+ case r @ DataSourceV2Relation(_: SupportsRowLevelOperations, _, _, _, _)
=>
+ val lineageAttributes = findRowLineageAttributes(r.metadataOutput).get
+ val (rowId, lastUpdatedSequence) =
(removeMetadataColumnAttribute(lineageAttributes._1),
+ removeMetadataColumnAttribute(lineageAttributes._2))
+
+ val lineageAssignments = updateTable.assignments ++
+ Seq(Assignment(lastUpdatedSequence, Literal(null)),
Assignment(rowId, rowId))
+
+ val tableWithLineage = r.copy(output = r.output ++ Seq(rowId,
lastUpdatedSequence))
+ updateTable.copy(table = tableWithLineage, assignments =
lineageAssignments)
+ }
+ }
+}
diff --git
a/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RemoveRowLineageOutputFromOriginalTable.scala
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RemoveRowLineageOutputFromOriginalTable.scala
new file mode 100644
index 0000000000..a3e9ad372b
--- /dev/null
+++
b/spark/v3.4/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RemoveRowLineageOutputFromOriginalTable.scala
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.iceberg.MetadataColumns
+import org.apache.spark.sql.catalyst.analysis.NamedRelation
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.ReplaceIcebergData
+import org.apache.spark.sql.catalyst.plans.logical.WriteIcebergDelta
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+
+/**
+ * RemoveRowLineageOutputFromOriginalTable removes row lineage outputs from
Dsv2 write's
+ * originalTable so that downstream behaviors like relation caching just work,
without having to
+ * modify physical planning strategies.
+ */
+object RemoveRowLineageOutputFromOriginalTable extends Rule[LogicalPlan] {
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ plan.transform {
+ case writeDelta@WriteIcebergDelta(_, _, originalTable, _, _) =>
+ writeDelta.copy(originalTable = removeRowLineageOutput(originalTable))
+ case replaceData@ReplaceIcebergData(_, _, originalTable, _) =>
+ replaceData.copy(originalTable = removeRowLineageOutput(originalTable))
+ }
+ }
+
+ private def removeRowLineageOutput(table: NamedRelation):
DataSourceV2Relation = {
+ table match {
+ case dsv2Relation@DataSourceV2Relation(_, _, _, _, _) =>
+ dsv2Relation.copy(output = dsv2Relation.output.filterNot(
+ attr => attr.name == MetadataColumns.ROW_ID.name() ||
+ attr.name == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()))
+ }
+ }
+}
+
diff --git
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/SparkRowLevelOperationsTestBase.java
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/SparkRowLevelOperationsTestBase.java
index b5d6415763..893f9931cf 100644
---
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/SparkRowLevelOperationsTestBase.java
+++
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/SparkRowLevelOperationsTestBase.java
@@ -177,7 +177,19 @@ public abstract class SparkRowLevelOperationsTestBase
extends ExtensionsTestBase
SparkCatalog.class.getName(),
ImmutableMap.of("type", "hadoop"),
FileFormat.PARQUET,
- RANDOM.nextBoolean(),
+ true,
+ WRITE_DISTRIBUTION_MODE_HASH,
+ true,
+ null,
+ LOCAL,
+ 3
+ },
+ {
+ "testhadoop",
+ SparkCatalog.class.getName(),
+ ImmutableMap.of("type", "hadoop"),
+ FileFormat.PARQUET,
+ false,
WRITE_DISTRIBUTION_MODE_HASH,
true,
null,
diff --git
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestCopyOnWriteWithLineage.java
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestCopyOnWriteWithLineage.java
new file mode 100644
index 0000000000..5b82f55ead
--- /dev/null
+++
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestCopyOnWriteWithLineage.java
@@ -0,0 +1,35 @@
+/*
+ * 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.iceberg.spark.extensions;
+
+import java.util.Map;
+import org.apache.iceberg.RowLevelOperationMode;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+
+public class TestCopyOnWriteWithLineage extends
TestRowLevelOperationsWithLineage {
+
+ @Override
+ protected Map<String, String> extraTableProperties() {
+ return ImmutableMap.of(
+ TableProperties.MERGE_MODE,
RowLevelOperationMode.COPY_ON_WRITE.modeName(),
+ TableProperties.UPDATE_MODE,
RowLevelOperationMode.COPY_ON_WRITE.modeName(),
+ TableProperties.DELETE_MODE,
RowLevelOperationMode.COPY_ON_WRITE.modeName());
+ }
+}
diff --git
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestMergeOnReadWithLineage.java
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestMergeOnReadWithLineage.java
new file mode 100644
index 0000000000..3f038276af
--- /dev/null
+++
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestMergeOnReadWithLineage.java
@@ -0,0 +1,35 @@
+/*
+ * 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.iceberg.spark.extensions;
+
+import java.util.Map;
+import org.apache.iceberg.RowLevelOperationMode;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+
+public class TestMergeOnReadWithLineage extends
TestRowLevelOperationsWithLineage {
+
+ @Override
+ protected Map<String, String> extraTableProperties() {
+ return ImmutableMap.of(
+ TableProperties.MERGE_MODE,
RowLevelOperationMode.MERGE_ON_READ.modeName(),
+ TableProperties.UPDATE_MODE,
RowLevelOperationMode.MERGE_ON_READ.modeName(),
+ TableProperties.DELETE_MODE,
RowLevelOperationMode.MERGE_ON_READ.modeName());
+ }
+}
diff --git
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java
new file mode 100644
index 0000000000..008e97480e
--- /dev/null
+++
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java
@@ -0,0 +1,510 @@
+/*
+ * 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.iceberg.spark.extensions;
+
+import static org.apache.iceberg.MetadataColumns.schemaWithRowLineage;
+import static org.apache.iceberg.spark.Spark3Util.loadIcebergTable;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assumptions.assumeThat;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import org.apache.iceberg.AppendFiles;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.MetadataColumns;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.SnapshotRef;
+import org.apache.iceberg.StructLike;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TestHelpers;
+import org.apache.iceberg.data.GenericAppenderFactory;
+import org.apache.iceberg.data.GenericRecord;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.DataWriter;
+import org.apache.iceberg.io.OutputFile;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Iterables;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.spark.functions.BucketFunction;
+import org.apache.iceberg.types.Types;
+import org.apache.iceberg.util.Pair;
+import org.apache.iceberg.util.PartitionMap;
+import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
+import org.apache.spark.sql.catalyst.parser.ParseException;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+
+public abstract class TestRowLevelOperationsWithLineage extends
SparkRowLevelOperationsTestBase {
+ static final Function<StructLike, StructLike> BUCKET_PARTITION_GENERATOR =
+ record ->
+ TestHelpers.Row.of(BucketFunction.BucketInt.invoke(2, record.get(0,
Integer.class)));
+
+ static final Schema SCHEMA =
+ new Schema(
+ ImmutableList.of(
+ Types.NestedField.required(1, "id", Types.IntegerType.get()),
+ Types.NestedField.required(2, "data", Types.StringType.get()),
+ MetadataColumns.ROW_ID,
+ MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER));
+
+ static final List<Record> INITIAL_RECORDS =
+ ImmutableList.of(
+ createRecord(SCHEMA, 100, "a", 0L, 1L),
+ createRecord(SCHEMA, 101, "b", 1L, 1L),
+ createRecord(SCHEMA, 102, "c", 2L, 1L),
+ createRecord(SCHEMA, 103, "d", 3L, 1L),
+ createRecord(SCHEMA, 104, "e", 4L, 1L));
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @BeforeEach
+ public void beforeEach() {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // ToDo: Remove these as row lineage inheritance gets implemented in the
other readers
+ assumeThat(fileFormat).isEqualTo(FileFormat.PARQUET);
+ }
+
+ @AfterEach
+ public void removeTables() {
+ sql("DROP TABLE IF EXISTS %s", tableName);
+ sql("DROP TABLE IF EXISTS source");
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatched()
+ throws NoSuchTableException, ParseException, IOException {
+ createAndInitTable("id INT, data STRING", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ appendUnpartitionedRecords(table, INITIAL_RECORDS);
+ createOrReplaceView(
+ "source",
+ "id int, data string",
+ "{ \"id\": 101, \"data\": \"updated_b\" }\n " + "{ \"id\": 200,
\"data\": \"f\" }\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.id == s.id "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = s.data "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT *",
+ commitTarget());
+
+ Snapshot updateSnapshot = latestSnapshot(table);
+ long updateSnapshotFirstRowId = updateSnapshot.firstRowId();
+ List<Object[]> allRows = rowsWithLineageAndFilePos();
+ List<Object[]> carriedOverAndUpdatedRows =
+ allRows.stream()
+ .filter(row -> (long) row[3] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ // Project sequence numbers first for easier comparison on the added row
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(1L, 100, "a", 0L, ANY, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY,
ANY),
+ row(1L, 102, "c", 2L, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ Iterables.getOnlyElement(
+ allRows.stream()
+ .filter(row -> (long) row[3] >= updateSnapshotFirstRowId)
+ .collect(Collectors.toList()));
+ assertAddedRowLineage(row(updateSnapshot.sequenceNumber(), 200, "f"),
newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+ createAndInitTable("id INT, data STRING", "PARTITIONED BY (bucket(2,
id))", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ appendRecords(
+ table, partitionRecords(INITIAL_RECORDS, table.spec(),
BUCKET_PARTITION_GENERATOR));
+ createOrReplaceView(
+ "source",
+ "id int, data string",
+ "{ \"id\": 101, \"data\": \"updated_b\" }\n " + "{ \"id\": 200,
\"data\": \"f\" }\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.id == s.id "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = s.data "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT *",
+ commitTarget());
+
+ Snapshot updateSnapshot = latestSnapshot(table);
+ long updateSnapshotFirstRowId = updateSnapshot.firstRowId();
+ List<Object[]> allRows = rowsWithLineageAndFilePos();
+
+ List<Object[]> carriedOverAndUpdatedRows =
+ allRows.stream()
+ .filter(row -> (long) row[3] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ // Project sequence numbers first for easier comparison on the added row
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(1L, 100, "a", 0L, ANY, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY,
ANY),
+ row(1L, 102, "c", 2L, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ Iterables.getOnlyElement(
+ allRows.stream()
+ .filter(row -> (long) row[3] >= updateSnapshotFirstRowId)
+ .collect(Collectors.toList()));
+ assertAddedRowLineage(row(updateSnapshot.sequenceNumber(), 200, "f"),
newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyNonMatched()
+ throws NoSuchTableException, ParseException, IOException {
+ createAndInitTable("id INT, data string", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ appendUnpartitionedRecords(table, INITIAL_RECORDS);
+ createOrReplaceView(
+ "source",
+ "id INT, data STRING",
+ "{ \"id\": 101, \"data\": \"updated_b\" }\n " + "{ \"id\": 200,
\"data\": \"f\" }\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.id == s.id "
+ + "WHEN NOT MATCHED THEN "
+ + "INSERT *",
+ commitTarget());
+
+ Snapshot updateSnapshot = latestSnapshot(table);
+ long updateSnapshotFirstRowId = updateSnapshot.firstRowId();
+
+ List<Object[]> allRows = rowsWithLineageAndFilePos();
+ List<Object[]> carriedOverAndUpdatedRows =
+ allRows.stream()
+ .filter(row -> (long) row[3] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ // Project sequence numbers first for easier comparison on the added row
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(1L, 100, "a", 0L, ANY, ANY),
+ row(1L, 101, "b", 1L, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ Iterables.getOnlyElement(
+ allRows.stream()
+ .filter(row -> (long) row[3] >= updateSnapshotFirstRowId)
+ .collect(Collectors.toList()));
+ assertAddedRowLineage(row(updateSnapshot.sequenceNumber(), 200, "f"),
newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyMatched()
+ throws IOException, NoSuchTableException, ParseException {
+ createAndInitTable("id INT, data STRING", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ appendUnpartitionedRecords(table, INITIAL_RECORDS);
+ createOrReplaceView(
+ "source",
+ "id INT, data string",
+ "{ \"id\": 101, \"data\": \"updated_b\" }\n "
+ + "{ \"id\": 102, \"data\": \"updated_c\" }\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.id == s.id "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = s.data ",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber),
+ row(103, "d", 3L, 1L),
+ row(104, "e", 4L, 1L)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeMatchedDelete() throws NoSuchTableException,
ParseException, IOException {
+ createAndInitTable("id INT, data STRING", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ appendUnpartitionedRecords(table, INITIAL_RECORDS);
+ createOrReplaceView(
+ "source",
+ "id INT, data string",
+ "{ \"id\": 101, \"data\": \"delete_101\" }\n "
+ + "{ \"id\": 102, \"data\": \"delete_102\" }\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s " + "ON t.id == s.id " + "WHEN
MATCHED THEN DELETE",
+ commitTarget());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(row(100, "a", 0L, 1L), row(103, "d", 3L, 1L),
row(104, "e", 4L, 1L)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testUpdate() throws NoSuchTableException, ParseException,
IOException {
+ createAndInitTable("id INT, data STRING", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ appendUnpartitionedRecords(table, INITIAL_RECORDS);
+
+ sql("UPDATE %s AS t set data = 'updated_b' WHERE id = 101",
commitTarget());
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "c", 2L, 1L),
+ row(103, "d", 3L, 1L),
+ row(104, "e", 4L, 1L)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testDelete() throws NoSuchTableException, ParseException,
IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("id int, data STRING", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ appendUnpartitionedRecords(table, INITIAL_RECORDS);
+
+ sql("DELETE FROM %s WHERE id = 101", commitTarget());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L),
+ row(102, "c", 2L, 1L),
+ row(103, "d", 3L, 1L),
+ row(104, "e", 4L, 1L)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWithManyRecords() throws NoSuchTableException,
ParseException, IOException {
+ createAndInitTable("id INT, data STRING", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+
+ int numRecords = 25000;
+ int startingId = 100;
+
+ List<Record> initialRecords = Lists.newArrayList();
+ int rowId = 0;
+ for (int id = 100; id < startingId + numRecords; id++) {
+ initialRecords.add(createRecord(SCHEMA, id, "data_" + id, rowId++, 1L));
+ }
+
+ appendUnpartitionedRecords(table, initialRecords);
+ createOrReplaceView(
+ "source",
+ "id int, data string",
+ "{ \"id\": 101, \"data\": \"updated_data_101\" }\n "
+ + "{ \"id\": 26000, \"data\": \"data_26000\" }\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.id == s.id "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = s.data "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT *",
+ commitTarget());
+
+ Snapshot updateSnapshot = latestSnapshot(table);
+ long updateSnapshotFirstRowId = updateSnapshot.firstRowId();
+ List<Object[]> allRows = rowsWithLineageAndFilePos();
+ List<Object[]> carriedOverAndUpdatedRows =
+ allRows.stream()
+ .filter(row -> (long) row[3] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ int newlyInsertedId = 26000;
+ int updatedId = 101;
+ List<Object[]> expectedCarriedOverAndUpdatedRows =
+ ImmutableList.<Object[]>builder()
+ .add(row(1L, 100, "data_100", 0L, ANY, ANY))
+ .add(row(updateSnapshot.sequenceNumber(), updatedId,
"updated_data_101", 1L, ANY, ANY))
+ .addAll(
+ // Every record with higher ids than the updated excluding the
new row should be a
+ // carry over
+ initialRecords.stream()
+ .filter(
+ initialRecord -> {
+ int id = initialRecord.get(0, Integer.class);
+ return id > updatedId && id != newlyInsertedId;
+ })
+ .map(this::recordToExpectedRow)
+ .collect(Collectors.toList()))
+ .build();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ expectedCarriedOverAndUpdatedRows,
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ Iterables.getOnlyElement(
+ allRows.stream()
+ .filter(row -> (long) row[3] >= updateSnapshotFirstRowId)
+ .collect(Collectors.toList()));
+ assertAddedRowLineage(row(updateSnapshot.sequenceNumber(), 26000,
"data_26000"), newRow);
+ }
+
+ private Object[] recordToExpectedRow(Record record) {
+ int id = record.get(0, Integer.class);
+ String data = record.get(1, String.class);
+ long rowId = record.get(2, Long.class);
+ long lastUpdated = record.get(3, Long.class);
+ return row(lastUpdated, id, data, rowId, ANY, ANY);
+ }
+
+ private List<Object[]> rowsWithLineageAndFilePos() {
+ return sql(
+ "SELECT s._last_updated_sequence_number, s.id, s.data, s._row_id,
files.first_row_id, s._pos FROM %s"
+ + " AS s JOIN %s.files AS files ON files.file_path = s._file ORDER
BY s._row_id",
+ selectTarget(), selectTarget());
+ }
+
+ private List<Object[]> rowsWithLineage() {
+ return sql(
+ "SELECT id, data, _row_id, _last_updated_sequence_number FROM %s ORDER
BY _row_id",
+ selectTarget());
+ }
+
+ /**
+ * Partitions the provided records based on the spec and partition function
+ *
+ * @return a partitioned map
+ */
+ protected PartitionMap<List<Record>> partitionRecords(
+ List<Record> records,
+ PartitionSpec spec,
+ Function<StructLike, StructLike> partitionGenerator) {
+ PartitionMap<List<Record>> recordsByPartition =
+ PartitionMap.create(Map.of(spec.specId(), spec));
+ for (Record record : records) {
+ StructLike partition = partitionGenerator != null ?
partitionGenerator.apply(record) : null;
+ List<Record> recordsForPartition = recordsByPartition.get(spec.specId(),
partition);
+ if (recordsForPartition == null) {
+ recordsForPartition = Lists.newArrayList();
+ }
+
+ recordsForPartition.add(record);
+ recordsByPartition.put(spec.specId(), partition, recordsForPartition);
+ }
+
+ return recordsByPartition;
+ }
+
+ protected void appendUnpartitionedRecords(Table table, List<Record> records)
throws IOException {
+ appendRecords(table, partitionRecords(records, table.spec(), record ->
null));
+ }
+
+ // Append unpartitioned records?
+ protected void appendRecords(Table table, PartitionMap<List<Record>>
partitionedRecords)
+ throws IOException {
+ AppendFiles append = table.newAppend();
+
+ for (Map.Entry<Pair<Integer, StructLike>, List<Record>> entry :
partitionedRecords.entrySet()) {
+ OutputFile file =
Files.localOutput(temp.resolve(UUID.randomUUID().toString()).toFile());
+ DataWriter<Record> writer =
+ new GenericAppenderFactory(schemaWithRowLineage(table.schema()),
table.spec())
+ .newDataWriter(
+ EncryptionUtil.plainAsEncryptedOutput(file), fileFormat,
entry.getKey().second());
+ List<Record> recordsForPartition = entry.getValue();
+ writer.write(recordsForPartition);
+ writer.close();
+ append =
+ append
+ .appendFile(writer.toDataFile())
+ .toBranch(branch != null ? branch : SnapshotRef.MAIN_BRANCH);
+ }
+
+ append.commit();
+ }
+
+ protected static Record createRecord(
+ Schema schema, int id, String data, long rowId, long
lastUpdatedSequenceNumber) {
+ Record record = GenericRecord.create(schema);
+ record.set(0, id);
+ record.set(1, data);
+ record.set(2, rowId);
+ record.set(3, lastUpdatedSequenceNumber);
+ return record;
+ }
+
+ private Snapshot latestSnapshot(Table table) {
+ return branch != null ? table.snapshot(branch) : table.currentSnapshot();
+ }
+
+ // Expected should have last updated sequence number followed by data columns
+ // Actual should have the contents of expected followed by the file first
row ID and position
+ private void assertAddedRowLineage(Object[] expected, Object[] actual) {
+ // validate the sequence number and all the data columns
+ for (int pos = 0; pos < expected.length; pos++) {
+ assertThat(actual[pos]).isEqualTo(expected[pos]);
+ }
+
+ int rowIdPos = expected.length;
+ int firstRowIdPos = rowIdPos + 1;
+ int positionPos = firstRowIdPos + 1;
+ long expectedRowId = (Long) actual[firstRowIdPos] + (Long)
actual[positionPos];
+ assertThat(actual[rowIdPos]).isEqualTo(expectedRowId);
+ }
+}
diff --git
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteOperation.java
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteOperation.java
index 4fca05345a..dd3d785b3b 100644
---
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteOperation.java
+++
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteOperation.java
@@ -21,9 +21,12 @@ package org.apache.iceberg.spark.source;
import static
org.apache.spark.sql.connector.write.RowLevelOperation.Command.DELETE;
import static
org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE;
+import java.util.List;
import org.apache.iceberg.IsolationLevel;
import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.Table;
+import org.apache.iceberg.TableUtil;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.connector.expressions.Expressions;
import org.apache.spark.sql.connector.expressions.NamedReference;
@@ -95,13 +98,18 @@ class SparkCopyOnWriteOperation implements
RowLevelOperation {
@Override
public NamedReference[] requiredMetadataAttributes() {
- NamedReference file = Expressions.column(MetadataColumns.FILE_PATH.name());
- NamedReference pos =
Expressions.column(MetadataColumns.ROW_POSITION.name());
-
+ List<NamedReference> metadataAttributes = Lists.newArrayList();
+
metadataAttributes.add(Expressions.column(MetadataColumns.FILE_PATH.name()));
if (command == DELETE || command == UPDATE) {
- return new NamedReference[] {file, pos};
- } else {
- return new NamedReference[] {file};
+
metadataAttributes.add(Expressions.column(MetadataColumns.ROW_POSITION.name()));
+ }
+
+ if (TableUtil.supportsRowLineage(table)) {
+
metadataAttributes.add(Expressions.column(MetadataColumns.ROW_ID.name()));
+ metadataAttributes.add(
+
Expressions.column(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()));
}
+
+ return metadataAttributes.toArray(NamedReference[]::new);
}
}
diff --git
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteScan.java
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteScan.java
index 7a6025b073..733f35fb2a 100644
---
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteScan.java
+++
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteScan.java
@@ -35,6 +35,7 @@ import org.apache.iceberg.metrics.ScanReport;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.Sets;
import org.apache.iceberg.spark.SparkReadConf;
+import org.apache.iceberg.spark.SparkSchemaUtil;
import org.apache.iceberg.util.SnapshotUtil;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.connector.expressions.Expressions;
@@ -43,6 +44,10 @@ import org.apache.spark.sql.connector.read.Statistics;
import org.apache.spark.sql.connector.read.SupportsRuntimeFiltering;
import org.apache.spark.sql.sources.Filter;
import org.apache.spark.sql.sources.In;
+import org.apache.spark.sql.types.Metadata;
+import org.apache.spark.sql.types.MetadataBuilder;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -53,6 +58,7 @@ class SparkCopyOnWriteScan extends
SparkPartitioningAwareScan<FileScanTask>
private final Snapshot snapshot;
private Set<String> filteredLocations = null;
+ private StructType readSchema;
SparkCopyOnWriteScan(
SparkSession spark,
@@ -82,6 +88,15 @@ class SparkCopyOnWriteScan extends
SparkPartitioningAwareScan<FileScanTask>
}
}
+ @Override
+ public StructType readSchema() {
+ if (readSchema == null) {
+ this.readSchema =
rowLineageAsDataCols(SparkSchemaUtil.convert(expectedSchema()));
+ }
+
+ return readSchema;
+ }
+
Long snapshotId() {
return snapshot != null ? snapshot.snapshotId() : null;
}
@@ -187,4 +202,29 @@ class SparkCopyOnWriteScan extends
SparkPartitioningAwareScan<FileScanTask>
Snapshot currentSnapshot = SnapshotUtil.latestSnapshot(table(), branch());
return currentSnapshot != null ? currentSnapshot.snapshotId() : null;
}
+
+ // Indicate to Spark to treat the row id and sequence number as data columns
since some optimizer
+ // rules for DELETE will not output the row lineage columns otherwise
+ private StructType rowLineageAsDataCols(StructType schema) {
+ StructField[] fields = new StructField[schema.fields().length];
+ for (int i = 0; i < schema.fields().length; i++) {
+ StructField field = schema.fields()[i];
+ if (isRowLineageField(field)) {
+ Metadata updatedMetadata =
+ new
MetadataBuilder().withMetadata(field.metadata()).remove("__metadata_col").build();
+ fields[i] = field.copy(field.name(), field.dataType(),
field.nullable(), updatedMetadata);
+ } else {
+ fields[i] = field;
+ }
+ }
+
+ return new StructType(fields);
+ }
+
+ private boolean isRowLineageField(StructField field) {
+ boolean hasLineageFieldName =
+ field.name().equals(MetadataColumns.ROW_ID.name())
+ ||
field.name().equals(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name());
+ return hasLineageFieldName && field.metadata().contains("__metadata_col");
+ }
}
diff --git
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaOperation.java
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaOperation.java
index 8acd87d3cb..7d71142364 100644
---
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaOperation.java
+++
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaOperation.java
@@ -18,9 +18,12 @@
*/
package org.apache.iceberg.spark.source;
+import java.util.List;
import org.apache.iceberg.IsolationLevel;
import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.Table;
+import org.apache.iceberg.TableUtil;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.connector.expressions.Expressions;
import org.apache.spark.sql.connector.expressions.NamedReference;
@@ -96,9 +99,17 @@ class SparkPositionDeltaOperation implements
RowLevelOperation, SupportsDelta {
@Override
public NamedReference[] requiredMetadataAttributes() {
- NamedReference specId = Expressions.column(MetadataColumns.SPEC_ID.name());
- NamedReference partition =
Expressions.column(MetadataColumns.PARTITION_COLUMN_NAME);
- return new NamedReference[] {specId, partition};
+ List<NamedReference> metadataAttributes = Lists.newArrayList();
+ metadataAttributes.add(Expressions.column(MetadataColumns.SPEC_ID.name()));
+
metadataAttributes.add(Expressions.column(MetadataColumns.PARTITION_COLUMN_NAME));
+
+ if (TableUtil.supportsRowLineage(table)) {
+
metadataAttributes.add(Expressions.column(MetadataColumns.ROW_ID.name()));
+ metadataAttributes.add(
+
Expressions.column(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()));
+ }
+
+ return metadataAttributes.toArray(new NamedReference[0]);
}
@Override
diff --git
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaWriteBuilder.java
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaWriteBuilder.java
index c58935206b..70cd735e1d 100644
---
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaWriteBuilder.java
+++
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaWriteBuilder.java
@@ -22,6 +22,7 @@ import org.apache.iceberg.IsolationLevel;
import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.Schema;
import org.apache.iceberg.Table;
+import org.apache.iceberg.TableUtil;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.spark.SparkSchemaUtil;
import org.apache.iceberg.spark.SparkUtil;
@@ -85,8 +86,12 @@ class SparkPositionDeltaWriteBuilder implements
DeltaWriteBuilder {
if (info.schema() == null || info.schema().isEmpty()) {
return null;
} else {
- Schema dataSchema = SparkSchemaUtil.convert(table.schema(),
info.schema());
- validateSchema("data", table.schema(), dataSchema);
+ Schema writeSchema =
+ TableUtil.supportsRowLineage(table)
+ ? MetadataColumns.schemaWithRowLineage(table.schema())
+ : table.schema();
+ Schema dataSchema = SparkSchemaUtil.convert(writeSchema, info.schema());
+ validateSchema("data", writeSchema, dataSchema);
return dataSchema;
}
}
@@ -104,6 +109,10 @@ class SparkPositionDeltaWriteBuilder implements
DeltaWriteBuilder {
new Schema(
MetadataColumns.SPEC_ID,
MetadataColumns.metadataColumn(table,
MetadataColumns.PARTITION_COLUMN_NAME));
+ if (TableUtil.supportsRowLineage(table)) {
+ expectedMetadataSchema =
MetadataColumns.schemaWithRowLineage(expectedMetadataSchema);
+ }
+
StructType metadataSparkType = info.metadataSchema().get();
Schema metadataSchema = SparkSchemaUtil.convert(expectedMetadataSchema,
metadataSparkType);
validateSchema("metadata", expectedMetadataSchema, metadataSchema);
diff --git
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java
index bbc7434138..e4c0789971 100644
---
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java
+++
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java
@@ -261,7 +261,10 @@ public class SparkTable
new SparkMetadataColumn(MetadataColumns.PARTITION_COLUMN_NAME,
sparkPartitionType, true),
new SparkMetadataColumn(MetadataColumns.FILE_PATH.name(),
DataTypes.StringType, false),
new SparkMetadataColumn(MetadataColumns.ROW_POSITION.name(),
DataTypes.LongType, false),
- new SparkMetadataColumn(MetadataColumns.IS_DELETED.name(),
DataTypes.BooleanType, false)
+ new SparkMetadataColumn(MetadataColumns.IS_DELETED.name(),
DataTypes.BooleanType, false),
+ new SparkMetadataColumn(MetadataColumns.ROW_ID.name(),
DataTypes.LongType, true),
+ new SparkMetadataColumn(
+ MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name(),
DataTypes.LongType, true)
};
}
diff --git
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java
index 602b692d73..ece0a5e6ff 100644
---
a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java
+++
b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java
@@ -19,8 +19,10 @@
package org.apache.iceberg.spark.source;
import org.apache.iceberg.IsolationLevel;
+import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.Schema;
import org.apache.iceberg.Table;
+import org.apache.iceberg.TableUtil;
import org.apache.iceberg.UpdateSchema;
import org.apache.iceberg.expressions.Expression;
import org.apache.iceberg.expressions.Expressions;
@@ -116,8 +118,14 @@ class SparkWriteBuilder implements WriteBuilder,
SupportsDynamicOverwrite, Suppo
@Override
public Write build() {
- // Validate
- Schema writeSchema = validateOrMergeWriteSchema(table, dsSchema,
writeConf);
+ // The write schema should only include row lineage in the output if it's
an overwrite
+ // operation.
+ // In any other case, only null row IDs and sequence numbers would be
produced which
+ // means the row lineage columns can be excluded from the output files
+ boolean writeIncludesRowLineage = TableUtil.supportsRowLineage(table) &&
overwriteFiles;
+ Schema writeSchema =
+ validateOrMergeWriteSchema(table, dsSchema, writeConf,
writeIncludesRowLineage);
+
SparkUtil.validatePartitionTransforms(table.spec());
// Get application id
@@ -170,7 +178,7 @@ class SparkWriteBuilder implements WriteBuilder,
SupportsDynamicOverwrite, Suppo
}
private static Schema validateOrMergeWriteSchema(
- Table table, StructType dsSchema, SparkWriteConf writeConf) {
+ Table table, StructType dsSchema, SparkWriteConf writeConf, boolean
writeIncludesRowLineage) {
Schema writeSchema;
boolean caseSensitive = writeConf.caseSensitive();
if (writeConf.mergeSchema()) {
@@ -182,6 +190,10 @@ class SparkWriteBuilder implements WriteBuilder,
SupportsDynamicOverwrite, Suppo
UpdateSchema update =
table.updateSchema().caseSensitive(caseSensitive).unionByNameWith(newSchema);
Schema mergedSchema = update.apply();
+ if (writeIncludesRowLineage) {
+ mergedSchema =
+ TypeUtil.join(mergedSchema,
MetadataColumns.schemaWithRowLineage(table.schema()));
+ }
// reconvert the dsSchema without assignment to use the ids assigned by
UpdateSchema
writeSchema = SparkSchemaUtil.convert(mergedSchema, dsSchema,
caseSensitive);
@@ -192,7 +204,11 @@ class SparkWriteBuilder implements WriteBuilder,
SupportsDynamicOverwrite, Suppo
// if the validation passed, update the table schema
update.commit();
} else {
- writeSchema = SparkSchemaUtil.convert(table.schema(), dsSchema,
caseSensitive);
+ Schema schema =
+ writeIncludesRowLineage
+ ? MetadataColumns.schemaWithRowLineage(table.schema())
+ : table.schema();
+ writeSchema = SparkSchemaUtil.convert(schema, dsSchema, caseSensitive);
TypeUtil.validateWriteSchema(
table.schema(), writeSchema, writeConf.checkNullability(),
writeConf.checkOrdering());
}
