You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@iceberg.apache.org by "amogh-jahagirdar (via GitHub)" <gi...@apache.org> on 2025/04/07 02:57:36 UTC
[PR] Spark 3.5 row lineage [iceberg]
amogh-jahagirdar opened a new pull request, #12736:
URL: https://github.com/apache/iceberg/pull/12736
(no comment)
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060856557
##########
core/src/main/java/org/apache/iceberg/MetadataColumns.java:
##########
@@ -152,4 +154,19 @@ public static boolean isMetadataColumn(int id) {
public static boolean nonMetadataColumn(String name) {
return !isMetadataColumn(name);
}
+
+ public static Schema schemaWithRowLineage(Table table) {
+ Preconditions.checkArgument(null != table, "Invalid table: null");
+ Preconditions.checkArgument(
+ !(table instanceof BaseMetadataTable),
+ "Cannot produce row lineage for metadata table: %s",
+ table);
+ Schema rowLineageSchema =
+ new Schema(
+ MetadataColumns.metadataColumn(table, MetadataColumns.ROW_ID.name()).asOptional(),
+ MetadataColumns.metadataColumn(
+ table, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name())
+ .asOptional());
+ return TypeUtil.join(table.schema(), rowLineageSchema);
Review Comment:
Looks like this is constructing `Schema` instances just to satisfy the types expected by `TypeUtil.join`. Is there an alternative method? Should there be?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060969325
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,75 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m.targetTable) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ 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
+ }
+
+ // Updates target rows without a matching source row.
+ // For NOT MATCHED BY TARGET (INSERT), null literals are assigned
+ // during alignment in ResolveRowLevelCommandAssignments.
Review Comment:
Okay, so if `ResolveRowLevelCommandAssignments` has to run again after this rule, why not just add the `null` insert cases in this rule?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059092298
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata attribute
+ // This works around the logic in
+ // ExposesMetadataColumns, used later in metadata attribute resolution,
+ // which prevents surfacing other metadata columns when a single metadata column is in the output
+ val rowLineageAsDataColumns = rowLineageAttributes.map(removeMetadataColumnAttribute)
+
+ val tableWithLineage = r.copy(output =
+ r.output ++ rowLineageAsDataColumns)
+
+ mergeIntoTable.copy(
+ targetTable = tableWithLineage,
+ matchedActions = matchedAssignmentsForLineage,
+ notMatchedBySourceActions = notMatchedBySourceActionsForLineage)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
Review Comment:
I think I prefer the solution of checking the table version. It's a table property injected by `SparkTable` so we don't need to go too deep. Then we could always have `_row_id` as a metadata column, even if for v2 tables it's always null.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059375632
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,436 @@
+/*
+ * 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.TableUtil.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 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.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> UNPARTITIONED_GENERATOR = record -> null;
+ 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, "data", Types.IntegerType.get()),
+ MetadataColumns.ROW_ID,
+ MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER));
+
+ static final List<Record> INITIAL_RECORDS =
+ ImmutableList.of(
+ createRecord(0, 0L, 1L),
+ createRecord(1, 1L, 1L),
+ createRecord(2, 2L, 1L),
+ createRecord(3, 3L, 1L),
+ createRecord(4, 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);
+ assumeThat(vectorized).isFalse();
Review Comment:
I've changed this test class to assert against the row IDs/last udpated sequence via SQL, which will go through inheritance (as opposed to reading from files). I went back and forth on this, and concluded that testing with inheritance ends up actually being simpler.
The original goal of testing without inheritance was to have independent tests that just verify our Spark changes are preserving or nulling out when appropriate, but building and reading via the low level reader has its own complications imo for someone wanting to verify this test does what it should.
For example, with reading the files directly, we have to specify the idToConstant mapping to have the baseRowId/seq number otherwise the reader will return nulls for those. After specifying the mapping, the seq number passed in will override instead of returning nulls. We could work around this when building our list of expected rows for specific records, but then someone reading the test has to gauge that and additionally verify that logic is correct, and we're not getting lucky somehow.
Reading the raw data file records is made more complicated by the fact that expectations need to filter out deleted or updated rows and treat those separately.
So with all that said, it ended up being simpler to just assert the row ID/seq number including inheritance. Since it's no guarantee what row Ids new rows would be assigned we just need to assert that the new row has a higher row ID than the prev max. I think that's OK because at this abstraction the actual row ID assigned for new rows is inconsequential so long as we're not somehow writing a duplicate (which the assertion mentioned before would catch).
cc @rdblue
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2045142385
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala:
##########
@@ -148,9 +151,29 @@ case class ExtendedDataSourceV2Strategy(spark: SparkSession) extends Strategy wi
case UnsetViewProperties(ResolvedV2View(catalog, ident), propertyKeys, ifExists) =>
AlterV2ViewUnsetPropertiesExec(catalog, ident, propertyKeys, ifExists) :: Nil
+ case ReplaceData(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, _, Some(write)) =>
+ ReplaceDataExec(planLater(query), refreshCache(r), write) :: Nil
+
+ case WriteDelta(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections,
+ Some(write)) =>
+ WriteDeltaExec(planLater(query), refreshCache(r), projections, write) :: Nil
+
case _ => Nil
}
+ /**
+ * In case of a table enabled with row lineage, the Dsv2 relation output for the write will have
+ * row lineage in the outputs. However, the cached Dsv2 relation won't have these fields
+ * so refresh after a write should be performed with an updated plan with row lineage outputs removed.
+ * @param r
+ */
+ private def refreshCache(r: DataSourceV2Relation)(): Unit = {
+ val planToCacheBy = r.copy(output = r.output.filterNot(
+ attr => attr.name == MetadataColumns.ROW_ID.name() ||
+ attr.name == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()))
+ spark.sharedState.cacheManager.recacheByPlan(spark, planToCacheBy)
+ }
+
Review Comment:
And even if we go with the current approach, probably also want to call `reacacheByPlan` with the full relation just in case that's also cached.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060979914
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,490 @@
+/*
+ * 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.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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+
+ createAndInitTable("id INT, data STRING", null);
Review Comment:
Is this creating a partitioned table?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5: Update MERGE and UPDATE for row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2061132912
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,492 @@
+/*
+ * 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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(1L, 101, "b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ 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, s._file FROM %s"
Review Comment:
It's not, I was was sanity checking my query by projecting it but it's not needed. Removed!
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,492 @@
+/*
+ * 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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(1L, 101, "b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ 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, s._file FROM %s"
Review Comment:
It's not, I was was sanity checking my query locally by projecting it but it's not needed. Removed!
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5: Update MERGE and UPDATE for row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue merged PR #12736:
URL: https://github.com/apache/iceberg/pull/12736
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053127517
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTableForRowLineage.scala:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.UpdateTable
+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@UpdateTable(_, _, _) if shouldUpdatePlan(updateTable) =>
+ updatePlanWithRowLineage(updateTable)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
+ // on the output of operation which indicates the rule already ran
+ private def shouldUpdatePlan(updateTable: UpdateTable): Boolean = {
+ val rowLineageAttrs = findRowLineageAttributes(updateTable.metadataOutput)
+ val allLineageAttrsPresent = rowLineageAttrs.nonEmpty && rowLineageAttrs.forall(updateTable.metadataOutput.contains)
+ val rowIdAbsentFromOutput = !updateTable.output.exists(_.name == ROW_ID_ATTRIBUTE_NAME)
+
+ allLineageAttrsPresent && rowIdAbsentFromOutput
+ }
+
+
+ private def updatePlanWithRowLineage(updateTable: UpdateTable): LogicalPlan = {
+ EliminateSubqueryAliases(updateTable.table) match {
+ case r @ DataSourceV2Relation(_: SupportsRowLevelOperations, _, _, _, _) =>
+ val rowLineageAttributes = findRowLineageAttributes(updateTable.metadataOutput)
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val lineageAssignments = updateTable.assignments ++
+ Seq(Assignment(lastUpdatedSequence, Literal(null)))
Review Comment:
Why does this not add a row ID assignment?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#issuecomment-2818857924
Forgot to link the test I had for that case @RussellSpitzer: https://github.com/apache/iceberg/pull/12736/files#diff-88253eaa3307bd26810d9499757f4bc789dda80547321f513ea04ac06d8e59ebR142
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059370149
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/SparkRowLevelOperationsTestBase.java:
##########
@@ -177,7 +177,19 @@ public static Object[][] parameters() {
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,
Review Comment:
Currently, in our DML tests for the parquet case we randomize between vectorized or not. I've changed this for V3 at least to run both the vectorized and non-vectorized case.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5: Update MERGE and UPDATE for row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2061122782
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,492 @@
+/*
+ * 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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(1L, 101, "b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ 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, s._file FROM %s"
Review Comment:
Nit: `_file` isn't needed, right?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#issuecomment-2817463379
Thanks @RussellSpitzer for taking a look!
>One thing I want to make sure we test is "WHEN NOT MATCHED" with no matching clauses. Spark has some logic which transforms these from row level operations into Appends and I am a little worried some of the logic here for manipulating the output and input table will not work in that case.
https://github.com/apache/spark/blob/597cafc05d9a4ecd5a111d1dc9c92fb37c77ce3f/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala#L47-L75
https://github.com/apache/iceberg/issues/12653
I think we may be fine there because we essentially do not care about writing row_lineage or last_updated sequence values in this case, so writing null or not writing are both valid.
Exactly, the implementation as it is currently does not have an explicit rule case for projecting row IDs or sequence values in the Append case because that is a case where we know we would only be adding new rows where all the row lDs or sequence values will be null. Which means they don't need to explicitly be persisted as per the spec.
This is what I was trying to say in this comment https://github.com/apache/iceberg/pull/12736/files#diff-cc38222cfec9d6a4142d6c102e3654a23ea9ec9cc1d9226ba19fbf12ae7c0329R120-R124 but let me know if I could make it more clear. Essentially when building the SparkWrite, it could either be for an an
1. overwrite files (CoW)
2. an append
3. overwrite dynamic partitions/by filter.
So in both 2,3 we know it will only be addition of new rows or an addition + removal of existing rows. In these cases the row lineage fields will be null for the entire write. 1 is the only case where we need to make sure the write schema has row lineage.
Let me know if I'm missing anything! I used to have an explicit handling for the merge cases which translated into appends but I realized it's not really needed so I removed it.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053123798
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata attribute
+ // This works around the logic in
+ // ExposesMetadataColumns, used later in metadata attribute resolution,
+ // which prevents surfacing other metadata columns when a single metadata column is in the output
Review Comment:
Looks like this is formatted strangely.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053127840
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RemoveRowLineageOutputFromOriginalTable.scala:
##########
@@ -0,0 +1,56 @@
+/*
+ * 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.ReplaceData
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+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.
Review Comment:
Is this comment wrapped correctly?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5: Update MERGE and UPDATE for row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#issuecomment-2836046881
I'm going to merge this since it's been ready for a few days. We can still follow up if there is additional feedback.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2031189720
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,266 @@
+/*
+ * 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.Alias
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.plans.logical.AppendData
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows.Keep
+import org.apache.spark.sql.catalyst.plans.logical.Project
+import org.apache.spark.sql.catalyst.plans.logical.ReplaceData
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.types.Metadata
+
+object RewriteMergeIntoTableForRowLineage extends RewriteOperationForRowLineage {
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ RewriteMergeIntoTable.apply(plan) resolveOperators {
+ case writeDelta@WriteDelta(_, _, _, _, _, _) if updatePlanForRowLineage(writeDelta) => {
+ updateDeltaPlan(writeDelta)
+ }
+
+ case replaceData@ReplaceData(_, _, _, _, _, _) if updatePlanForRowLineage(replaceData) => {
+ updateReplaceDataPlan(replaceData)
+ }
+
+ case append@AppendData(table, _, _, _, _, _) if updatePlanForRowLineage(append) =>
Review Comment:
For the append case, users may not being go through Spark extensions so we may need to handle appends at execution time by creating a new internal row with 2 additional fields for the row lineage columns
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059375632
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,436 @@
+/*
+ * 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.TableUtil.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 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.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> UNPARTITIONED_GENERATOR = record -> null;
+ 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, "data", Types.IntegerType.get()),
+ MetadataColumns.ROW_ID,
+ MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER));
+
+ static final List<Record> INITIAL_RECORDS =
+ ImmutableList.of(
+ createRecord(0, 0L, 1L),
+ createRecord(1, 1L, 1L),
+ createRecord(2, 2L, 1L),
+ createRecord(3, 3L, 1L),
+ createRecord(4, 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);
+ assumeThat(vectorized).isFalse();
Review Comment:
I've changed this test class to assert against the row IDs/last udpated sequence via SQL, which will go through inheritance (as opposed to reading from files). I went back and forth on this, and concluded that testing with inheritance ends up actually being simpler.
The original goal of testing without inheritance was to have independent tests that just verify our Spark changes are preserving or nulling out when appropriate, but building and reading via the low level reader has its own complications imo for someone wanting to verify this test does what it should.
For example, with reading the files directly, we have to specify the idToConstant mapping to have the baseRowId/seq number otherwise the reader will return nulls for those. After specifying the mapping, the seq number passed in will override instead of returning nulls. We could work around this when building our list of expected rows for specific records, but then someone reading the test has to gauge that and additionally verify that logic is correct, and that we're not getting lucky somehow.
Reading the raw data file records is made more complicated by the fact that expectations need to filter out deleted or updated rows and treat those separately.
So with all that said, it ended up being simpler to just assert the row ID/seq number including inheritance. Since it's no guarantee what row Ids new rows would be assigned we just need to assert that the new row has a higher row ID than the prev max. I think that's OK because at this abstraction the actual row ID assigned for new rows is inconsequential so long as we're not somehow writing a duplicate (which the assertion mentioned before would catch).
The implication of this is this test class cannot run for Spark readers where inheritance is not implemented. I also think that's OK because as we add the implementations we'd tighten up these assumptions.
cc @rdblue
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060870333
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java:
##########
@@ -258,23 +256,16 @@ public Set<TableCapability> capabilities() {
@Override
public MetadataColumn[] metadataColumns() {
DataType sparkPartitionType = SparkSchemaUtil.convert(Partitioning.partitionType(table()));
- ImmutableList.Builder<SparkMetadataColumn> metadataColumns = ImmutableList.builder();
- metadataColumns.add(
- new SparkMetadataColumn(MetadataColumns.SPEC_ID.name(), DataTypes.IntegerType, false),
- 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));
-
- if (TableUtil.formatVersion(table()) >= 3) {
- metadataColumns.add(
- new SparkMetadataColumn(MetadataColumns.ROW_ID.name(), DataTypes.LongType, true));
- metadataColumns.add(
- new SparkMetadataColumn(
- MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name(), DataTypes.LongType, true));
- }
-
- return metadataColumns.build().toArray(SparkMetadataColumn[]::new);
+ return new MetadataColumn[] {
+ new SparkMetadataColumn(MetadataColumns.SPEC_ID.name(), DataTypes.IntegerType, false),
+ 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.ROW_ID.name(), DataTypes.LongType, true),
+ new SparkMetadataColumn(
+ MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name(), DataTypes.LongType, true)
+ };
Review Comment:
Changed this to always allow projecting row lineage columns regardless of format version. For the unsupported format versions, it'll return null
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkTable.java:
##########
@@ -258,23 +256,16 @@ public Set<TableCapability> capabilities() {
@Override
public MetadataColumn[] metadataColumns() {
DataType sparkPartitionType = SparkSchemaUtil.convert(Partitioning.partitionType(table()));
- ImmutableList.Builder<SparkMetadataColumn> metadataColumns = ImmutableList.builder();
- metadataColumns.add(
- new SparkMetadataColumn(MetadataColumns.SPEC_ID.name(), DataTypes.IntegerType, false),
- 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));
-
- if (TableUtil.formatVersion(table()) >= 3) {
- metadataColumns.add(
- new SparkMetadataColumn(MetadataColumns.ROW_ID.name(), DataTypes.LongType, true));
- metadataColumns.add(
- new SparkMetadataColumn(
- MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name(), DataTypes.LongType, true));
- }
-
- return metadataColumns.build().toArray(SparkMetadataColumn[]::new);
+ return new MetadataColumn[] {
+ new SparkMetadataColumn(MetadataColumns.SPEC_ID.name(), DataTypes.IntegerType, false),
+ 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.ROW_ID.name(), DataTypes.LongType, true),
+ new SparkMetadataColumn(
+ MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name(), DataTypes.LongType, true)
+ };
Review Comment:
Changed this to always allow projecting row lineage columns regardless of format version. For the unsupported format versions, it'll return null for the row IDs
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060887711
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaWriteBuilder.java:
##########
@@ -104,6 +109,16 @@ private void validateMetadataSchema() {
new Schema(
MetadataColumns.SPEC_ID,
MetadataColumns.metadataColumn(table, MetadataColumns.PARTITION_COLUMN_NAME));
+ if (TableUtil.supportsRowLineage(table)) {
+ Schema rowLineageSchema =
+ new Schema(
+ MetadataColumns.metadataColumn(table, MetadataColumns.ROW_ID.name()).asOptional(),
+ MetadataColumns.metadataColumn(
+ table, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name())
+ .asOptional());
+ expectedMetadataSchema = TypeUtil.join(expectedMetadataSchema, rowLineageSchema);
Review Comment:
schemaWithRowLineage?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060972410
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,490 @@
+/*
+ * 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.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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
Review Comment:
This is a little weaker than it should be. It would be better to use `Iterables.getOnlyElement` or to assert that the number of rows in `allRows` accounts for the total number of expected rows. That way this doesn't leave open the possibility that there is an unexpected added row.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053119414
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata attribute
+ // This works around the logic in
+ // ExposesMetadataColumns, used later in metadata attribute resolution,
+ // which prevents surfacing other metadata columns when a single metadata column is in the output
+ val rowLineageAsDataColumns = rowLineageAttributes.map(removeMetadataColumnAttribute)
+
+ val tableWithLineage = r.copy(output =
+ r.output ++ rowLineageAsDataColumns)
+
+ mergeIntoTable.copy(
+ targetTable = tableWithLineage,
+ matchedActions = matchedAssignmentsForLineage,
+ notMatchedBySourceActions = notMatchedBySourceActionsForLineage)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
Review Comment:
By "present in the target table" you mean "in the relation's metadataOutput" right?
I'm not sure I agree. Don't we always want to have `_row_id` available but read `null` if the table isn't v3? That would mean that we don't need to produce different `metadataOutput` based on the version of the table so you always have a consistent set of metadata columns, just with different behavior.
That said, if `metadataOutput` always has `_row_id` then you'd need to look at the table in the relation and see if it is a v3 table, probably using the `format-version` table property.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053163626
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata attribute
+ // This works around the logic in
+ // ExposesMetadataColumns, used later in metadata attribute resolution,
+ // which prevents surfacing other metadata columns when a single metadata column is in the output
+ val rowLineageAsDataColumns = rowLineageAttributes.map(removeMetadataColumnAttribute)
+
+ val tableWithLineage = r.copy(output =
+ r.output ++ rowLineageAsDataColumns)
+
+ mergeIntoTable.copy(
+ targetTable = tableWithLineage,
+ matchedActions = matchedAssignmentsForLineage,
+ notMatchedBySourceActions = notMatchedBySourceActionsForLineage)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
Review Comment:
>Don't we always want to have _row_id available but read null if the table isn't v3?
Hm I looked at it as we only want to have _row_id available for v3+ since the existence of that column in the projection could be potentially misleading, though arguably since it would just be null means that projecting it is fine since that implicitly indicates there is no row lineage.
>That said, if metadataOutput always has _row_id then you'd need to look at the table in the relation and see if it is a v3 table, probably using the format-version table property.
Yeah, we'd have to extract the Iceberg table from the DSV2 table. I don't think that's too bad, I just was concerned that merely indicating _row_id in older tables was problematic
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2050020503
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLineagePropagation.java:
##########
@@ -0,0 +1,489 @@
+/*
+ * 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.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.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.iceberg.AppendFiles;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.MetadataColumns;
+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.data.GenericAppenderFactory;
+import org.apache.iceberg.data.GenericRecord;
+import org.apache.iceberg.data.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.types.TypeUtil;
+import org.apache.iceberg.types.Types;
+import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
+import org.apache.spark.sql.catalyst.parser.ParseException;
+import org.apache.spark.sql.execution.SparkPlan;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.TestTemplate;
+
+public abstract class TestRowLineagePropagation extends SparkRowLevelOperationsTestBase {
+
+ static Schema rowLineageSchema =
+ new Schema(MetadataColumns.ROW_ID, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER);
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0),
+ row(2),
+ row(3),
+ row(4),
+ row(1234), // updated
+ row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", "PARTITIONED BY (data)", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<StructLike> partitions =
+ IntStream.range(0, 5)
+ .mapToObj(
+ partitionValue ->
+ createPartitionStructure(
+ table.spec().partitionType(), ImmutableMap.of("data", partitionValue)))
+ .collect(Collectors.toList());
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5, partitions);
+
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0),
+ row(2),
+ row(3),
+ row(4),
+ row(1234), // updated
+ row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ private StructLike createPartitionStructure(
+ Types.StructType partitionType, Map<String, Object> partitionValues) {
+ Record record = GenericRecord.create(partitionType);
+ return record.copy(partitionValues);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyNonMatched()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0), row(1), row(2), row(3), row(4), row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(generatedRecords);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.add(insertedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyMatched()
+ throws IOException, NoSuchTableException, ParseException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 ",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0), row(2), row(3), row(4), row(1234) // updated
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ expectedRecords.add(updatedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeMatchedDelete() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ SparkPlan plan =
+ executeAndKeepPlan(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN DELETE",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySource()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN UPDATE set data = 5678",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ // Every record should preserve its row ID and null out its sequence number.
+ // Every record other than id 1 should be 5678 and
+ for (Record record : generatedRecords) {
+ Record newRecord = record.copy();
+ if (record.get(0, Integer.class) == 1) {
+ newRecord.set(0, 1234);
+ newRecord.set(2, null);
+ } else {
+ newRecord.set(0, 5678);
+ newRecord.set(2, null);
+ }
+
+ expectedRecords.add(newRecord);
+ }
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ Record recordWithData1 =
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class).equals(1))
+ .findFirst()
+ .get();
+ List<Record> expectedRecords = Lists.newArrayList();
+ Record record = recordWithData1.copy();
+ record.set(0, 1234);
+ record.set(2, null);
+ expectedRecords.add(record);
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testUpdate() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+
+ sql("UPDATE %s AS t set data = 5678 WHERE data = 1", commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 5678, 1L, null);
+ expectedRecords.add(updatedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ protected Record createRecord(Schema schema, int id, Long rowId, Long lastUpdatedSequenceNumber) {
+ Record record = GenericRecord.create(schema);
+ record.set(0, id);
+ record.set(1, rowId);
+ record.set(2, lastUpdatedSequenceNumber);
+ return record;
+ }
+
+ protected List<Record> recordsSortedByRowId(Table table) throws IOException {
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> records = Lists.newArrayList();
+ table.refresh();
+ Snapshot snapshot = table.currentSnapshot();
+ if (branch != null) {
+ snapshot = table.snapshot(branch);
+ }
+
+ try (CloseableIterable<Record> writtenRecords =
+ IcebergGenerics.read(table).project(schema).useSnapshot(snapshot.snapshotId()).build()) {
+ for (Record record : writtenRecords) {
+ records.add(record);
+ }
+ }
+
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ protected List<Record> sortedRecordsByRowId(List<Record> records) {
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ protected List<Record> generateRecordsAndAppendToTable(
+ Table table, int numRecords, List<StructLike> partitions) throws IOException {
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> records = Lists.newArrayList();
+ Map<Integer, List<Record>> partitionToRecords = Maps.newHashMap();
+ Map<Integer, StructLike> partitionToPartitions = Maps.newHashMap();
+ for (int i = 0; i < numRecords; i++) {
+ records.add(createRecord(schema, i, (long) i, (long) i * 2));
+ int partitionKey = partitions.get(i) != null ? i : -1;
+ List<Record> existingRecords =
+ partitionToRecords.getOrDefault(partitionKey, Lists.newArrayList());
+ existingRecords.add(records.get(i));
+ partitionToRecords.put(partitionKey, existingRecords);
+ }
+
+ for (StructLike partition : partitions) {
+ if (partition != null) {
+ partitionToPartitions.put(partition.get(0, Integer.class), partition);
+ }
+ }
+
+ AppendFiles append = table.newAppend();
+ for (Map.Entry<Integer, List<Record>> entry : partitionToRecords.entrySet()) {
+ OutputFile file = Files.localOutput(temp.resolve(UUID.randomUUID().toString()).toFile());
Review Comment:
Can't build a FanoutWriter without providing a FileWriterFactory implementation which the spark one isn't exposed here, but we can use PartitionMap to simplify a few things here
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2045584975
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala:
##########
@@ -148,9 +151,29 @@ case class ExtendedDataSourceV2Strategy(spark: SparkSession) extends Strategy wi
case UnsetViewProperties(ResolvedV2View(catalog, ident), propertyKeys, ifExists) =>
AlterV2ViewUnsetPropertiesExec(catalog, ident, propertyKeys, ifExists) :: Nil
+ case ReplaceData(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, _, Some(write)) =>
+ ReplaceDataExec(planLater(query), refreshCache(r), write) :: Nil
+
+ case WriteDelta(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections,
+ Some(write)) =>
+ WriteDeltaExec(planLater(query), refreshCache(r), projections, write) :: Nil
+
case _ => Nil
}
+ /**
+ * In case of a table enabled with row lineage, the Dsv2 relation output for the write will have
+ * row lineage in the outputs. However, the cached Dsv2 relation won't have these fields
+ * so refresh after a write should be performed with an updated plan with row lineage outputs removed.
+ * @param r
+ */
+ private def refreshCache(r: DataSourceV2Relation)(): Unit = {
+ val planToCacheBy = r.copy(output = r.output.filterNot(
+ attr => attr.name == MetadataColumns.ROW_ID.name() ||
+ attr.name == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()))
+ spark.sharedState.cacheManager.recacheByPlan(spark, planToCacheBy)
+ }
+
Review Comment:
I went ahead and added the optimizer rule instead
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059536198
##########
spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkMetadataColumns.java:
##########
@@ -317,28 +316,6 @@ public void testConflictingColumns() {
sql("SELECT _spec_id, _partition, _renamed_spec_id FROM %s", TABLE_NAME));
}
- @TestTemplate
- public void testRowLineageColumnsResolvedInV3OrHigher() {
Review Comment:
I'm removing this test now, it's not really useful anymore since any format version can project _row_id, and in the unsupported versions it'll be null.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053123443
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
Review Comment:
This seems suspicious. Shouldn't this be `InsertAction` for not matched?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053235855
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
Review Comment:
Let me give these better names
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053133869
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.TableUtil.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.Comparator;
+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.Files;
+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.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.Lists;
+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.TestTemplate;
+
+public abstract class TestRowLevelOperationsWithLineage extends SparkRowLevelOperationsTestBase {
+ static final Function<Record, StructLike> UNPARTITIONED_GENERATOR = record -> null;
+ // Use the first field in the row as identity partition
+ static final Function<Record, StructLike> IDENTITY_PARTITIONED_GENERATOR =
+ record -> TestHelpers.Row.of(record.get(0, Integer.class));
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(
Review Comment:
`numRecords` is always 5 and there are only 2 partition functions. I think it may be better to just have constants up at the top, rather than generating the records. I still have some trouble reading these tests because I have to remember that the IDs are assigned 0-4.
Also, a constant sequence number would be helpful to make this match reality a bit more closely.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2031189720
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,266 @@
+/*
+ * 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.Alias
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.plans.logical.AppendData
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows.Keep
+import org.apache.spark.sql.catalyst.plans.logical.Project
+import org.apache.spark.sql.catalyst.plans.logical.ReplaceData
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.types.Metadata
+
+object RewriteMergeIntoTableForRowLineage extends RewriteOperationForRowLineage {
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ RewriteMergeIntoTable.apply(plan) resolveOperators {
+ case writeDelta@WriteDelta(_, _, _, _, _, _) if updatePlanForRowLineage(writeDelta) => {
+ updateDeltaPlan(writeDelta)
+ }
+
+ case replaceData@ReplaceData(_, _, _, _, _, _) if updatePlanForRowLineage(replaceData) => {
+ updateReplaceDataPlan(replaceData)
+ }
+
+ case append@AppendData(table, _, _, _, _, _) if updatePlanForRowLineage(append) =>
Review Comment:
For the append case, users may not being go through Spark extensions (they may just be inserting/df.append'ing some records) so we may need to handle appends at execution time by creating a new internal row with 2 additional fields for the row lineage columns, where the fields will be nulled out.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#issuecomment-2826336284
It looks like tests are failing after I rebased off the recent inheritance changes, looking into it...
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059370149
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/SparkRowLevelOperationsTestBase.java:
##########
@@ -177,7 +177,19 @@ public static Object[][] parameters() {
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,
Review Comment:
Currently, in our DML tests for the parquet case we randomize between vectorized or not. I've changed this for V3 at least to run both the vectorized and non-vectorized case. I think we should get away from the random alternation here, I don't think it adds too much time and we can also remove some other cases from V2 down the line. It's better for reproducibility when trying to debug and isolate issues.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053119851
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata attribute
+ // This works around the logic in
+ // ExposesMetadataColumns, used later in metadata attribute resolution,
+ // which prevents surfacing other metadata columns when a single metadata column is in the output
+ val rowLineageAsDataColumns = rowLineageAttributes.map(removeMetadataColumnAttribute)
+
+ val tableWithLineage = r.copy(output =
+ r.output ++ rowLineageAsDataColumns)
+
+ mergeIntoTable.copy(
+ targetTable = tableWithLineage,
+ matchedActions = matchedAssignmentsForLineage,
+ notMatchedBySourceActions = notMatchedBySourceActionsForLineage)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
+ // on the output of operation which indicates the rule already ran
+ private def shouldUpdatePlan(mergeIntoTable: MergeIntoTable): Boolean = {
+ val metadataOutput = mergeIntoTable.targetTable.metadataOutput
+ val rowLineageAttrs = findRowLineageAttributes(metadataOutput)
+ val allLineageAttrsPresent = rowLineageAttrs.nonEmpty && rowLineageAttrs.forall(metadataOutput.contains)
Review Comment:
Is the `nonEmpty` helpful? If it is empty then the first `forall` will fail and short circuit.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053135137
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.TableUtil.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.Comparator;
+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.Files;
+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.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.Lists;
+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.TestTemplate;
+
+public abstract class TestRowLevelOperationsWithLineage extends SparkRowLevelOperationsTestBase {
+ static final Function<Record, StructLike> UNPARTITIONED_GENERATOR = record -> null;
+ // Use the first field in the row as identity partition
+ static final Function<Record, StructLike> IDENTITY_PARTITIONED_GENERATOR =
Review Comment:
I usually like to have test cases that are a bit more realistic. You _could_ partition by the ID but it's uncommon and I would argue it's better to have groups of records like real partitions. Single-record partitions leaves this open to bugs when there are multiple records per data file. For instance, what if there is a record reuse bug that only affects these metadata columns?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053165632
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
Review Comment:
Correct, this would only be a `RemoveAction`
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053145064
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
Review Comment:
So this is the "not matched by source", so this will only be an update or a removeAction (this is the only SQL syntax docs I could find, that outline the grammar and possible options) https://docs.databricks.com/aws/en/sql/language-manual/delta-merge-into#syntax
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060855402
##########
core/src/main/java/org/apache/iceberg/MetadataColumns.java:
##########
@@ -152,4 +154,19 @@ public static boolean isMetadataColumn(int id) {
public static boolean nonMetadataColumn(String name) {
return !isMetadataColumn(name);
}
+
+ public static Schema schemaWithRowLineage(Table table) {
Review Comment:
I think that accepting `Table` is too broad. This should just need the table schema. There's no need to check for the partition column here so you don't need to call `metadataColumn` and pass `table`.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2061028456
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,490 @@
+/*
+ * 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.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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+
+ createAndInitTable("id INT, data STRING", 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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), 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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "b", 1L, 1L, ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ private List<Object[]> rowsWithLineageAndFilePos() {
+ return sql(
+ "SELECT id, data, _row_id, _last_updated_sequence_number, _file, _pos FROM %s ORDER BY _row_id",
Review Comment:
I meant to comment on this but I thought with the JOIN + updating the assertions to include ANY matching + ordering the seq number first, it was becoming a bit harder to parse what was going on in these tests. I'll take another pass at it.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#issuecomment-2827851100
Ok, we randomize vectorized execution in our tests https://github.com/apache/iceberg/blob/main/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/SparkRowLevelOperationsTestBase.java#L180 . That'd explain it since that path hasn't been updated yet.
I also feel like we should make these tests deterministic, I'd rather have some additional cases and have CI take a bit longer
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2048073422
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java:
##########
@@ -117,7 +119,10 @@ public WriteBuilder overwrite(Filter[] filters) {
@Override
public Write build() {
// Validate
- Schema writeSchema = validateOrMergeWriteSchema(table, dsSchema, writeConf);
+ Schema writeSchema =
+ validateOrMergeWriteSchema(
+ table, dsSchema, writeConf, TableUtil.formatVersion(table) >= 3 && overwriteFiles);
Review Comment:
Add a comment for this, but in the context of building a regular `SparkWrite` (so appends, overwriteFIles, overwriteDynamicPartitions) the only one case where we need to include row lineage explicitly in the output schema is the `overwriteFiles` case for CoW.
In every other case, both fields will be explicitly null for all the fields produced, and as per the spec that means we do not need to write them out. cc @rdblue
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java:
##########
@@ -117,7 +119,10 @@ public WriteBuilder overwrite(Filter[] filters) {
@Override
public Write build() {
// Validate
- Schema writeSchema = validateOrMergeWriteSchema(table, dsSchema, writeConf);
+ Schema writeSchema =
+ validateOrMergeWriteSchema(
+ table, dsSchema, writeConf, TableUtil.formatVersion(table) >= 3 && overwriteFiles);
Review Comment:
I'll Add a comment for this, but in the context of building a regular `SparkWrite` (so appends, overwriteFIles, overwriteDynamicPartitions) the only one case where we need to include row lineage explicitly in the output schema is the `overwriteFiles` case for CoW.
In every other case, both fields will be explicitly null for all the fields produced, and as per the spec that means we do not need to write them out. cc @rdblue
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2048093293
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkPositionDeltaOperation.java:
##########
@@ -98,6 +99,12 @@ public DeltaWriteBuilder newWriteBuilder(LogicalWriteInfo info) {
public NamedReference[] requiredMetadataAttributes() {
NamedReference specId = Expressions.column(MetadataColumns.SPEC_ID.name());
NamedReference partition = Expressions.column(MetadataColumns.PARTITION_COLUMN_NAME);
+ if (TableUtil.formatVersion(table) >= 3) {
Review Comment:
instead of >= 3, maybe we want to move this to table util as well TableUtil.supportsRowLineage. It avoids the "magic format" number that's littered throughout now
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2048073422
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java:
##########
@@ -117,7 +119,10 @@ public WriteBuilder overwrite(Filter[] filters) {
@Override
public Write build() {
// Validate
- Schema writeSchema = validateOrMergeWriteSchema(table, dsSchema, writeConf);
+ Schema writeSchema =
+ validateOrMergeWriteSchema(
+ table, dsSchema, writeConf, TableUtil.formatVersion(table) >= 3 && overwriteFiles);
Review Comment:
I'll Add a comment for this, but in the context of building a regular `SparkWrite`, the only one case where we need to include row lineage explicitly in the output schema is the `overwriteFiles` case for CoW.
In every other case for `SparkWrite`, both fields will be explicitly null for all the fields produced, and as per the spec that means we do not need to write them out. cc @rdblue
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053145064
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
Review Comment:
So this is the "not matched by source", so this will only be an update or a removeAction https://docs.databricks.com/aws/en/sql/language-manual/delta-merge-into#syntax
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053126840
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
Review Comment:
I think it would be better to return a tuple with the row ID and last updated seq attrs instead of a list.
That has a couple of nice consequences. Then you wouldn't need to do the `filter` in the next couple lines. But the more important effect is that it would be more reliable because the code below that adds the attrs to the table output doesn't check the order of the attributes. It just assumes that they're in the same order as the assignments added to the actions. I don't like relying on the incoming order.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053122266
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
Review Comment:
Nit: I think it would be easier to understand if the inner `actions` were called `assignments`, since the action is `UpdateAction`.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2061122165
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,492 @@
+/*
+ * 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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(1L, 101, "b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ 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, s._file 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 all data columns and the last updated sequence at the end
Review Comment:
Nit: I don't think the comment is still correct, but it is no longer relevant
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059384245
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.TableUtil.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.Comparator;
+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.Files;
+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.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.Lists;
+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.TestTemplate;
+
+public abstract class TestRowLevelOperationsWithLineage extends SparkRowLevelOperationsTestBase {
+ static final Function<Record, StructLike> UNPARTITIONED_GENERATOR = record -> null;
+ // Use the first field in the row as identity partition
+ static final Function<Record, StructLike> IDENTITY_PARTITIONED_GENERATOR =
Review Comment:
Good point, I've updated to use a bucket(2) partition on the data column. That'll force a grouping of records per data file so we can catch any record reuse issues.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059090462
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
Review Comment:
Got it. I probably should have seen that and it makes sense that we would need to add this. That raises a couple things:
1. For inserts, where do we add the `null` assignment? Is that automatic somehow, like the alignment rule runs again and adds the assignment?
2. For this, can you add a comment to clarify that this is actually an update for unmatched rows in the target table? Just making it a bit more obvious may help people like me reading in the future.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059363985
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,454 @@
+/*
+ * 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.TableUtil.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 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.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> UNPARTITIONED_GENERATOR = record -> null;
+ 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, "data", Types.IntegerType.get()),
+ MetadataColumns.ROW_ID,
+ MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER));
+
+ static final List<Record> INITIAL_RECORDS =
+ ImmutableList.of(
+ createRecord(0, 0L, 1L),
+ createRecord(1, 1L, 1L),
+ createRecord(2, 2L, 1L),
+ createRecord(3, 3L, 1L),
+ createRecord(4, 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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ appendRecords(table, partitionRecords(INITIAL_RECORDS, table.spec(), UNPARTITIONED_GENERATOR));
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 5678}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT *",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ List<Object[]> rowsWhereIdShouldBePreserved =
+ sql(
+ "SELECT data, _row_id, _last_updated_sequence_number FROM %s WHERE data != 5678 ORDER BY _row_id",
+ selectTarget());
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(0, 0L, 1L),
+ row(1234, 1L, updateSequenceNumber),
+ row(2, 2L, 1L),
+ row(3, 3L, 1L),
+ row(4, 4L, 1L)),
+ rowsWhereIdShouldBePreserved);
+
+ // New row with data 5678 should have any row ID higher than the max in the previous commit
+ // version (4), and have the latest seq number
+ long previousHighestRowId = 4;
+ List<Object[]> newRows =
+ sql(
+ "SELECT data, _row_id, _last_updated_sequence_number FROM %s WHERE data = 5678 ORDER BY _row_id",
+ selectTarget());
+ assertAddedRowLineage(newRows, 5678, previousHighestRowId, updateSequenceNumber);
Review Comment:
There may be a simpler way to express this without the filtering in the query...we can get all the records sorted by row ID. Since we know the new record must be at the end we can sublist every element except the last, and assert that. Then take the new record and assert the same stuff as 135 and below.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060879933
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteOperation.java:
##########
@@ -95,13 +98,18 @@ public WriteBuilder newWriteBuilder(LogicalWriteInfo info) {
@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};
+ if (TableUtil.supportsRowLineage(table)) {
+ metadataAttributes.add(Expressions.column(MetadataColumns.ROW_POSITION.name()));
+ metadataAttributes.add(Expressions.column(MetadataColumns.ROW_ID.name()));
+ metadataAttributes.add(
+ Expressions.column(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()));
+ } else if (command == DELETE || command == UPDATE) {
+ metadataAttributes.add(Expressions.column(MetadataColumns.ROW_POSITION.name()));
}
+
+ return metadataAttributes.toArray(NamedReference[]::new);
Review Comment:
Checking if I can simplify this a bit more
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2061029326
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,490 @@
+/*
+ * 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.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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
Review Comment:
Agreed, I changed to use Iterables.getOnlyElement.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2044986176
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteOperationForRowLineage.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.analysis
+
+import org.apache.spark.sql.catalyst.expressions.Attribute
+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.plans.logical.ReplaceData
+import org.apache.spark.sql.catalyst.plans.logical.V2WriteCommand
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.catalyst.util.METADATA_COL_ATTR_KEY
+import org.apache.spark.sql.catalyst.util.WriteDeltaProjections
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+
+trait RewriteOperationForRowLineage extends RewriteRowLevelCommand {
+
+ protected val ROW_ID_ATTRIBUTE_NAME = "_row_id"
+ protected val LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME =
+ "_last_updated_sequence_number"
Review Comment:
Use the existing MetadataColumn constants
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2047516346
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLineagePropagation.java:
##########
@@ -0,0 +1,489 @@
+/*
+ * 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.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.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.iceberg.AppendFiles;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.MetadataColumns;
+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.data.GenericAppenderFactory;
+import org.apache.iceberg.data.GenericRecord;
+import org.apache.iceberg.data.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.types.TypeUtil;
+import org.apache.iceberg.types.Types;
+import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
+import org.apache.spark.sql.catalyst.parser.ParseException;
+import org.apache.spark.sql.execution.SparkPlan;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.TestTemplate;
+
+public abstract class TestRowLineagePropagation extends SparkRowLevelOperationsTestBase {
+
+ static Schema rowLineageSchema =
+ new Schema(MetadataColumns.ROW_ID, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER);
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0),
+ row(2),
+ row(3),
+ row(4),
+ row(1234), // updated
+ row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", "PARTITIONED BY (data)", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<StructLike> partitions =
+ IntStream.range(0, 5)
+ .mapToObj(
+ partitionValue ->
+ createPartitionStructure(
+ table.spec().partitionType(), ImmutableMap.of("data", partitionValue)))
+ .collect(Collectors.toList());
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5, partitions);
+
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0),
+ row(2),
+ row(3),
+ row(4),
+ row(1234), // updated
+ row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ private StructLike createPartitionStructure(
+ Types.StructType partitionType, Map<String, Object> partitionValues) {
+ Record record = GenericRecord.create(partitionType);
+ return record.copy(partitionValues);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyNonMatched()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0), row(1), row(2), row(3), row(4), row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(generatedRecords);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.add(insertedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyMatched()
+ throws IOException, NoSuchTableException, ParseException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 ",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0), row(2), row(3), row(4), row(1234) // updated
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ expectedRecords.add(updatedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeMatchedDelete() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ SparkPlan plan =
+ executeAndKeepPlan(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN DELETE",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySource()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN UPDATE set data = 5678",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ // Every record should preserve its row ID and null out its sequence number.
+ // Every record other than id 1 should be 5678 and
+ for (Record record : generatedRecords) {
+ Record newRecord = record.copy();
+ if (record.get(0, Integer.class) == 1) {
+ newRecord.set(0, 1234);
+ newRecord.set(2, null);
+ } else {
+ newRecord.set(0, 5678);
+ newRecord.set(2, null);
+ }
+
+ expectedRecords.add(newRecord);
+ }
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ Record recordWithData1 =
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class).equals(1))
+ .findFirst()
+ .get();
+ List<Record> expectedRecords = Lists.newArrayList();
+ Record record = recordWithData1.copy();
+ record.set(0, 1234);
+ record.set(2, null);
+ expectedRecords.add(record);
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testUpdate() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+
+ sql("UPDATE %s AS t set data = 5678 WHERE data = 1", commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 5678, 1L, null);
+ expectedRecords.add(updatedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ protected Record createRecord(Schema schema, int id, Long rowId, Long lastUpdatedSequenceNumber) {
+ Record record = GenericRecord.create(schema);
+ record.set(0, id);
+ record.set(1, rowId);
+ record.set(2, lastUpdatedSequenceNumber);
+ return record;
+ }
+
+ protected List<Record> recordsSortedByRowId(Table table) throws IOException {
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> records = Lists.newArrayList();
+ table.refresh();
+ Snapshot snapshot = table.currentSnapshot();
+ if (branch != null) {
+ snapshot = table.snapshot(branch);
+ }
+
+ try (CloseableIterable<Record> writtenRecords =
+ IcebergGenerics.read(table).project(schema).useSnapshot(snapshot.snapshotId()).build()) {
+ for (Record record : writtenRecords) {
+ records.add(record);
+ }
+ }
+
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ protected List<Record> sortedRecordsByRowId(List<Record> records) {
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ protected List<Record> generateRecordsAndAppendToTable(
+ Table table, int numRecords, List<StructLike> partitions) throws IOException {
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> records = Lists.newArrayList();
+ Map<Integer, List<Record>> partitionToRecords = Maps.newHashMap();
+ Map<Integer, StructLike> partitionToPartitions = Maps.newHashMap();
+ for (int i = 0; i < numRecords; i++) {
+ records.add(createRecord(schema, i, (long) i, (long) i * 2));
+ int partitionKey = partitions.get(i) != null ? i : -1;
+ List<Record> existingRecords =
+ partitionToRecords.getOrDefault(partitionKey, Lists.newArrayList());
+ existingRecords.add(records.get(i));
+ partitionToRecords.put(partitionKey, existingRecords);
+ }
+
+ for (StructLike partition : partitions) {
+ if (partition != null) {
+ partitionToPartitions.put(partition.get(0, Integer.class), partition);
+ }
+ }
+
+ AppendFiles append = table.newAppend();
+ for (Map.Entry<Integer, List<Record>> entry : partitionToRecords.entrySet()) {
+ OutputFile file = Files.localOutput(temp.resolve(UUID.randomUUID().toString()).toFile());
Review Comment:
can replace all this with `FanoutDataWriter`
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2047516895
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLineagePropagation.java:
##########
@@ -0,0 +1,489 @@
+/*
+ * 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.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.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.iceberg.AppendFiles;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.MetadataColumns;
+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.data.GenericAppenderFactory;
+import org.apache.iceberg.data.GenericRecord;
+import org.apache.iceberg.data.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.types.TypeUtil;
+import org.apache.iceberg.types.Types;
+import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
+import org.apache.spark.sql.catalyst.parser.ParseException;
+import org.apache.spark.sql.execution.SparkPlan;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.TestTemplate;
+
+public abstract class TestRowLineagePropagation extends SparkRowLevelOperationsTestBase {
+
+ static Schema rowLineageSchema =
+ new Schema(MetadataColumns.ROW_ID, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER);
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0),
+ row(2),
+ row(3),
+ row(4),
+ row(1234), // updated
+ row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", "PARTITIONED BY (data)", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<StructLike> partitions =
+ IntStream.range(0, 5)
+ .mapToObj(
+ partitionValue ->
+ createPartitionStructure(
+ table.spec().partitionType(), ImmutableMap.of("data", partitionValue)))
+ .collect(Collectors.toList());
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5, partitions);
+
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0),
+ row(2),
+ row(3),
+ row(4),
+ row(1234), // updated
+ row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ private StructLike createPartitionStructure(
+ Types.StructType partitionType, Map<String, Object> partitionValues) {
+ Record record = GenericRecord.create(partitionType);
+ return record.copy(partitionValues);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyNonMatched()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0), row(1), row(2), row(3), row(4), row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(generatedRecords);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.add(insertedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyMatched()
+ throws IOException, NoSuchTableException, ParseException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 ",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0), row(2), row(3), row(4), row(1234) // updated
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ expectedRecords.add(updatedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeMatchedDelete() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ SparkPlan plan =
+ executeAndKeepPlan(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN DELETE",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySource()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN UPDATE set data = 5678",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ // Every record should preserve its row ID and null out its sequence number.
+ // Every record other than id 1 should be 5678 and
+ for (Record record : generatedRecords) {
+ Record newRecord = record.copy();
+ if (record.get(0, Integer.class) == 1) {
+ newRecord.set(0, 1234);
+ newRecord.set(2, null);
+ } else {
+ newRecord.set(0, 5678);
+ newRecord.set(2, null);
+ }
+
+ expectedRecords.add(newRecord);
+ }
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ Record recordWithData1 =
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class).equals(1))
+ .findFirst()
+ .get();
+ List<Record> expectedRecords = Lists.newArrayList();
+ Record record = recordWithData1.copy();
+ record.set(0, 1234);
+ record.set(2, null);
+ expectedRecords.add(record);
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testUpdate() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+
+ sql("UPDATE %s AS t set data = 5678 WHERE data = 1", commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 5678, 1L, null);
+ expectedRecords.add(updatedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ protected Record createRecord(Schema schema, int id, Long rowId, Long lastUpdatedSequenceNumber) {
+ Record record = GenericRecord.create(schema);
+ record.set(0, id);
+ record.set(1, rowId);
+ record.set(2, lastUpdatedSequenceNumber);
+ return record;
+ }
+
+ protected List<Record> recordsSortedByRowId(Table table) throws IOException {
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> records = Lists.newArrayList();
+ table.refresh();
+ Snapshot snapshot = table.currentSnapshot();
+ if (branch != null) {
+ snapshot = table.snapshot(branch);
+ }
+
+ try (CloseableIterable<Record> writtenRecords =
+ IcebergGenerics.read(table).project(schema).useSnapshot(snapshot.snapshotId()).build()) {
+ for (Record record : writtenRecords) {
+ records.add(record);
+ }
+ }
+
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ protected List<Record> sortedRecordsByRowId(List<Record> records) {
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ protected List<Record> generateRecordsAndAppendToTable(
Review Comment:
Separate the generation of records from the append
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "RussellSpitzer (via GitHub)" <gi...@apache.org>.
RussellSpitzer commented on PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#issuecomment-2816868801
This looks good to me although I only had time to skim it today. I'll try to take a full pass later. One thing I want to make sure we test is "WHEN NOT MATCHED" with no matching clauses. Spark has some logic which transforms these from row level operations into Appends and I am a little worried some of the logic here for manipulating the output and input table will not work in that case.
https://github.com/apache/iceberg/issues/12653
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060974460
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,490 @@
+/*
+ * 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.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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+
+ createAndInitTable("id INT, data STRING", 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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), 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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "b", 1L, 1L, ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ private List<Object[]> rowsWithLineageAndFilePos() {
+ return sql(
+ "SELECT id, data, _row_id, _last_updated_sequence_number, _file, _pos FROM %s ORDER BY _row_id",
+ 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 all data columns and the last updated sequence at the end
+ // Actual should all data columns, followed by row ID, last updated sequence, file, and then pos
+ private void assertAddedRowLineage(Object[] expected, Object[] actual) {
+ // validate all the data columns
+ for (int pos = 0; pos < expected.length - 1; pos++) {
+ assertThat(actual[pos]).isEqualTo(expected[pos]);
+ }
+
+ // Determine the expected row ID from the firstRowId of the file and the pos
+ int actualRowIdPos = actual.length - 4;
+ int filePos = actual.length - 2;
+ int rowPositionPos = actual.length - 1;
+ String file = actual[filePos].toString();
+ long fileFirstRowId =
+ (Long)
+ sql("SELECT first_row_id FROM %s.files WHERE file_path = '%s'", selectTarget(), file)
+ .get(0)[0];
+ long rowPosition = (long) actual[rowPositionPos];
+ long expectedFirstRowId = fileFirstRowId + rowPosition;
+ assertThat(actual[actualRowIdPos]).isEqualTo(expectedFirstRowId);
+
+ int expectedLastSequencePos = expected.length - 1;
+ int actualLastSequencePos = 3;
+ assertThat(actual[actualLastSequencePos]).isEqualTo(expected[expectedLastSequencePos]);
Review Comment:
Why not check this in the assertion on line 470? Doesn't it just require selecting `_last_updated_sequence_number` after the content? And, isn't that where this is anyway because `actualLastSequencePos` is hard-coded to 3?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060973380
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,490 @@
+/*
+ * 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.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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+
+ createAndInitTable("id INT, data STRING", 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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), 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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "b", 1L, 1L, ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ private List<Object[]> rowsWithLineageAndFilePos() {
+ return sql(
+ "SELECT id, data, _row_id, _last_updated_sequence_number, _file, _pos FROM %s ORDER BY _row_id",
+ 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 all data columns and the last updated sequence at the end
+ // Actual should all data columns, followed by row ID, last updated sequence, file, and then pos
+ private void assertAddedRowLineage(Object[] expected, Object[] actual) {
+ // validate all the data columns
+ for (int pos = 0; pos < expected.length - 1; pos++) {
+ assertThat(actual[pos]).isEqualTo(expected[pos]);
+ }
+
+ // Determine the expected row ID from the firstRowId of the file and the pos
+ int actualRowIdPos = actual.length - 4;
+ int filePos = actual.length - 2;
+ int rowPositionPos = actual.length - 1;
+ String file = actual[filePos].toString();
+ long fileFirstRowId =
+ (Long)
+ sql("SELECT first_row_id FROM %s.files WHERE file_path = '%s'", selectTarget(), file)
+ .get(0)[0];
Review Comment:
This is okay, but running a SQL query for every row would get slow if this were scaled up in the future. Probably better to run a join to get `first_row_id`?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5: Update MERGE and UPDATE for row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2061132781
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,492 @@
+/*
+ * 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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(updateSnapshot.sequenceNumber(), 101, "updated_b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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, ANY),
+ row(1L, 101, "b", 1L, ANY, ANY, ANY),
+ row(1L, 102, "c", 2L, ANY, ANY, ANY),
+ row(1L, 103, "d", 3L, ANY, ANY, ANY),
+ row(1L, 104, "e", 4L, ANY, 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ 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, s._file 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 all data columns and the last updated sequence at the end
Review Comment:
Fixed!
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059375632
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,436 @@
+/*
+ * 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.TableUtil.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 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.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> UNPARTITIONED_GENERATOR = record -> null;
+ 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, "data", Types.IntegerType.get()),
+ MetadataColumns.ROW_ID,
+ MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER));
+
+ static final List<Record> INITIAL_RECORDS =
+ ImmutableList.of(
+ createRecord(0, 0L, 1L),
+ createRecord(1, 1L, 1L),
+ createRecord(2, 2L, 1L),
+ createRecord(3, 3L, 1L),
+ createRecord(4, 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);
+ assumeThat(vectorized).isFalse();
Review Comment:
I've changed this test class to assert against the row IDs/last udpated sequence via SQL, which will go through inheritance (as opposed to reading from files). I went back and forth on this, and concluded that testing with inheritance ends up actually being simpler.
The original goal of testing without inheritance was to have independent tests that just verify our Spark changes are preserving or nulling out when appropriate, but building and reading via the low level reader has its own complications imo for someone wanting to verify this test does what it should.
For example, with reading the files directly, we have to specify the idToConstant mapping to have the baseRowId/seq number otherwise the reader will return nulls for those. After specifying the mapping, the seq number passed in will override instead of returning nulls. We could work around this when building our list of expected rows for specific records, but then someone reading the test has to gauge that and additionally verify that logic is correct, and we're not getting lucky somehow.
Reading the raw data file records is made more complicated by the fact that expectations need to filter out deleted or updated rows and treat those separately.
So with all that said, it ended up being simpler to just assert the row ID/seq number including inheritance. Since it's no guarantee what row Ids new rows would be assigned we just need to assert that the new row has a higher row ID than the prev max. I think that's OK because at this abstraction the actual row ID assigned for new rows is inconsequential so long as we're not somehow writing a duplicate (which the assertion mentioned before would catch).
The implication of this is this test class cannot run for Spark readers where inheritance is not implemented. I also think that's OK because as we add the implementations we'd tighten up these assumptions.
cc @rdblue
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "nastra (via GitHub)" <gi...@apache.org>.
nastra commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2036234436
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLineagePropogation.java:
##########
@@ -0,0 +1,489 @@
+/*
+ * 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.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.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.iceberg.AppendFiles;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.MetadataColumns;
+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.data.GenericAppenderFactory;
+import org.apache.iceberg.data.GenericRecord;
+import org.apache.iceberg.data.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.types.TypeUtil;
+import org.apache.iceberg.types.Types;
+import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
+import org.apache.spark.sql.catalyst.parser.ParseException;
+import org.apache.spark.sql.execution.SparkPlan;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.TestTemplate;
+
+public abstract class TestRowLineagePropogation extends SparkRowLevelOperationsTestBase {
+
+ static Schema rowLineageSchema =
+ new Schema(MetadataColumns.ROW_ID, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER);
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0),
+ row(2),
+ row(3),
+ row(4),
+ row(1234), // updated
+ row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", "PARTITIONED BY (data)", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<StructLike> partitions =
+ IntStream.range(0, 5)
+ .mapToObj(
+ partitionValue ->
+ createPartitionStructure(
+ table.spec().partitionType(), ImmutableMap.of("data", partitionValue)))
+ .collect(Collectors.toList());
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5, partitions);
+
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0),
+ row(2),
+ row(3),
+ row(4),
+ row(1234), // updated
+ row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ private StructLike createPartitionStructure(
+ Types.StructType partitionType, Map<String, Object> partitionValues) {
+ Record record = GenericRecord.create(partitionType);
+ return record.copy(partitionValues);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyNonMatched()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT (data) VALUES (5678)",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0), row(1), row(2), row(3), row(4), row(5678) // inserted
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(generatedRecords);
+ Record insertedRecord = createRecord(schema, 5678, null, null);
+ expectedRecords.add(insertedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyMatched()
+ throws IOException, NoSuchTableException, ParseException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 ",
+ commitTarget());
+
+ ImmutableList<Object[]> expectedRows =
+ ImmutableList.of(
+ row(0), row(2), row(3), row(4), row(1234) // updated
+ );
+
+ assertEquals(
+ "Output should match", expectedRows, sql("SELECT * FROM %s ORDER BY data", selectTarget()));
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 1234, 1L, null);
+ expectedRecords.add(updatedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeMatchedDelete() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ SparkPlan plan =
+ executeAndKeepPlan(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN DELETE",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySource()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN UPDATE set data = 5678",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ // Every record should preserve its row ID and null out its sequence number.
+ // Every record other than id 1 should be 5678 and
+ for (Record record : generatedRecords) {
+ Record newRecord = record.copy();
+ if (record.get(0, Integer.class) == 1) {
+ newRecord.set(0, 1234);
+ newRecord.set(2, null);
+ } else {
+ newRecord.set(0, 5678);
+ newRecord.set(2, null);
+ }
+
+ expectedRecords.add(newRecord);
+ }
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ Record recordWithData1 =
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class).equals(1))
+ .findFirst()
+ .get();
+ List<Record> expectedRecords = Lists.newArrayList();
+ Record record = recordWithData1.copy();
+ record.set(0, 1234);
+ record.set(2, null);
+ expectedRecords.add(record);
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ @TestTemplate
+ public void testUpdate() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> generatedRecords = generateRecordsAndAppendToTable(table, 5);
+
+ sql("UPDATE %s AS t set data = 5678 WHERE data = 1", commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(
+ generatedRecords.stream()
+ .filter(record -> record.get(0, Integer.class) != 1)
+ .collect(Collectors.toList()));
+ Record updatedRecord = createRecord(schema, 5678, 1L, null);
+ expectedRecords.add(updatedRecord);
+ table.refresh();
+
+ assertRecordsAreEqual(sortedRecordsByRowId(expectedRecords), recordsSortedByRowId(table));
+ }
+
+ protected Record createRecord(Schema schema, int id, Long rowId, Long lastUpdatedSequenceNumber) {
+ Record record = GenericRecord.create(schema);
+ record.set(0, id);
+ record.set(1, rowId);
+ record.set(2, lastUpdatedSequenceNumber);
+ return record;
+ }
+
+ protected List<Record> recordsSortedByRowId(Table table) throws IOException {
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> records = Lists.newArrayList();
+ table.refresh();
+ Snapshot snapshot = table.currentSnapshot();
+ if (branch != null) {
+ snapshot = table.snapshot(branch);
+ }
+
+ try (CloseableIterable<Record> writtenRecords =
+ IcebergGenerics.read(table).project(schema).useSnapshot(snapshot.snapshotId()).build()) {
+ for (Record record : writtenRecords) {
+ records.add(record);
+ }
+ }
+
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ protected List<Record> sortedRecordsByRowId(List<Record> records) {
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ protected List<Record> generateRecordsAndAppendToTable(
+ Table table, int numRecords, List<StructLike> partitions) throws IOException {
+ Schema schema = TypeUtil.join(table.schema(), rowLineageSchema);
+ List<Record> records = Lists.newArrayList();
+ Map<Integer, List<Record>> partitionToRecords = Maps.newHashMap();
+ Map<Integer, StructLike> partitionToPartitions = Maps.newHashMap();
+ for (int i = 0; i < numRecords; i++) {
+ records.add(createRecord(schema, i, (long) i, (long) i * 2));
+ int partitionKey = partitions.get(i) != null ? i : -1;
+ List<Record> existingRecords =
+ partitionToRecords.getOrDefault(partitionKey, Lists.newArrayList());
+ existingRecords.add(records.get(i));
+ partitionToRecords.put(partitionKey, existingRecords);
+ }
+
+ for (StructLike partition : partitions) {
+ if (partition != null) {
+ partitionToPartitions.put(partition.get(0, Integer.class), partition);
+ }
+ }
+
+ AppendFiles append = table.newAppend();
+ for (Map.Entry<Integer, List<Record>> entry : partitionToRecords.entrySet()) {
+ OutputFile file = Files.localOutput(temp.resolve(UUID.randomUUID().toString()).toFile());
+ StructLike partition =
+ entry.getKey() == -1 ? null : partitionToPartitions.get(entry.getKey());
+ DataWriter<Record> writer =
+ new GenericAppenderFactory(schema, table.spec())
+ .newDataWriter(EncryptionUtil.plainAsEncryptedOutput(file), fileFormat, partition);
+ List<Record> recordsForPartition = partitionToRecords.get(entry.getKey());
+ writer.write(recordsForPartition);
+ writer.close();
+ append =
+ append
+ .appendFile(writer.toDataFile())
+ .toBranch(branch != null ? branch : SnapshotRef.MAIN_BRANCH);
+ }
+
+ append.commit();
+
+ return records;
+ }
+
+ protected List<Record> generateRecordsAndAppendToTable(Table table, int numRecords)
+ throws IOException {
+ List<StructLike> partitions = Lists.newArrayList();
+ for (int i = 0; i < numRecords; i++) {
+ partitions.add(null);
+ }
+
+ return generateRecordsAndAppendToTable(table, numRecords, partitions);
+ }
+
+ protected void assertRecordsAreEqual(List<Record> expected, List<Record> actual) {
+ assertThat(actual.size()).isEqualTo(expected.size());
Review Comment:
```suggestion
assertThat(actual).hasSameSizeAs(expected);
```
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "nastra (via GitHub)" <gi...@apache.org>.
nastra commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2036232707
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java:
##########
@@ -327,7 +327,11 @@ public void pruneColumns(StructType requestedSchema) {
StructType requestedProjection =
new StructType(
Stream.of(requestedSchema.fields())
- .filter(field -> MetadataColumns.nonMetadataColumn(field.name()))
+ .filter(
+ field ->
+ MetadataColumns.nonMetadataColumn(field.name())
+ || field.name().equals("_row_id")
Review Comment:
would be good to use `MetadataColumns.ROW_ID.name()` here
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060975885
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,490 @@
+/*
+ * 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.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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+
+ createAndInitTable("id INT, data STRING", 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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), 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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "b", 1L, 1L, ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), 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 testMergeWhenNotMatchedBySource()
+ 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 BY SOURCE THEN UPDATE set data = 'not_matched_by_source'",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "not_matched_by_source", 0L, updateSequenceNumber),
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "not_matched_by_source", 2L, updateSequenceNumber),
+ row(103, "not_matched_by_source", 3L, updateSequenceNumber),
+ row(104, "not_matched_by_source", 4L, updateSequenceNumber)),
+ rowsWithLineage());
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ 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\": 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 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ long updateSequenceNumber = latestSnapshot(table).sequenceNumber();
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(101, "updated_b", 1L, updateSequenceNumber),
+ row(102, "updated_c", 2L, updateSequenceNumber)),
+ 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());
+ }
+
+ private List<Object[]> rowsWithLineageAndFilePos() {
+ return sql(
+ "SELECT id, data, _row_id, _last_updated_sequence_number, _file, _pos FROM %s ORDER BY _row_id",
Review Comment:
This would probably be easier to use in the assertion for new rows if `_last_updated_sequence_number` came first.
Also, why not join on `_file` here to the metadata table to recover `first_row_id` instead of running another select?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060895675
##########
core/src/main/java/org/apache/iceberg/MetadataColumns.java:
##########
@@ -152,4 +154,19 @@ public static boolean isMetadataColumn(int id) {
public static boolean nonMetadataColumn(String name) {
return !isMetadataColumn(name);
}
+
+ public static Schema schemaWithRowLineage(Table table) {
+ Preconditions.checkArgument(null != table, "Invalid table: null");
+ Preconditions.checkArgument(
+ !(table instanceof BaseMetadataTable),
+ "Cannot produce row lineage for metadata table: %s",
+ table);
+ Schema rowLineageSchema =
+ new Schema(
+ MetadataColumns.metadataColumn(table, MetadataColumns.ROW_ID.name()).asOptional(),
+ MetadataColumns.metadataColumn(
+ table, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name())
+ .asOptional());
+ return TypeUtil.join(table.schema(), rowLineageSchema);
Review Comment:
I think using TypeUtil.join as it is today is right. My mental model is the row lineage fields are a separate schema themselves that we combine with normal schemas when we need to.
Alternatively we could have an API that joins vararg fields? (`join(Types.NestedField... fields)`) that would avoid clients creating the schema
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053112107
##########
core/src/main/java/org/apache/iceberg/TableUtil.java:
##########
@@ -60,4 +61,28 @@ public static String metadataFileLocation(Table table) {
"%s does not have a metadata file location", table.getClass().getSimpleName()));
}
}
+
+ public static boolean supportsRowLineage(Table table) {
+ Preconditions.checkArgument(null != table, "Invalid table: null");
+ if (table instanceof BaseMetadataTable) {
+ return false;
+ }
+
+ return formatVersion(table) >= TableMetadata.MIN_FORMAT_VERSION_ROW_LINEAGE;
+ }
+
+ public static Schema schemaWithRowLineage(Table table) {
Review Comment:
Should this be in `TableUtil`? It's a good util method, but it seems to fit better in `MetadataColumns` to me. This class just has work-arounds for things that aren't exposed through the `Table` API like version.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2056312051
##########
core/src/main/java/org/apache/iceberg/TableUtil.java:
##########
@@ -60,4 +61,28 @@ public static String metadataFileLocation(Table table) {
"%s does not have a metadata file location", table.getClass().getSimpleName()));
}
}
+
+ public static boolean supportsRowLineage(Table table) {
+ Preconditions.checkArgument(null != table, "Invalid table: null");
+ if (table instanceof BaseMetadataTable) {
+ return false;
+ }
+
+ return formatVersion(table) >= TableMetadata.MIN_FORMAT_VERSION_ROW_LINEAGE;
+ }
+
+ public static Schema schemaWithRowLineage(Table table) {
Review Comment:
I agree TableUtil probably isn't the rigt place to keep this method MetadataColumns.withRowLineageSchema(Table table) is better.
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTableForRowLineage.scala:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.UpdateTable
+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@UpdateTable(_, _, _) if shouldUpdatePlan(updateTable) =>
+ updatePlanWithRowLineage(updateTable)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
+ // on the output of operation which indicates the rule already ran
+ private def shouldUpdatePlan(updateTable: UpdateTable): Boolean = {
+ val rowLineageAttrs = findRowLineageAttributes(updateTable.metadataOutput)
+ val allLineageAttrsPresent = rowLineageAttrs.nonEmpty && rowLineageAttrs.forall(updateTable.metadataOutput.contains)
+ val rowIdAbsentFromOutput = !updateTable.output.exists(_.name == ROW_ID_ATTRIBUTE_NAME)
+
+ allLineageAttrsPresent && rowIdAbsentFromOutput
+ }
+
+
+ private def updatePlanWithRowLineage(updateTable: UpdateTable): LogicalPlan = {
+ EliminateSubqueryAliases(updateTable.table) match {
+ case r @ DataSourceV2Relation(_: SupportsRowLevelOperations, _, _, _, _) =>
+ val rowLineageAttributes = findRowLineageAttributes(updateTable.metadataOutput)
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val lineageAssignments = updateTable.assignments ++
+ Seq(Assignment(lastUpdatedSequence, Literal(null)))
Review Comment:
I originally only explicitly assigned the last updated sequence number because the row ID would just be carried over as part of the projection, but it's probably safer to just be very explicit here (just like the merge rule is).
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.TableUtil.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.Comparator;
+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.Files;
+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.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.Lists;
+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.TestTemplate;
+
+public abstract class TestRowLevelOperationsWithLineage extends SparkRowLevelOperationsTestBase {
+ static final Function<Record, StructLike> UNPARTITIONED_GENERATOR = record -> null;
+ // Use the first field in the row as identity partition
+ static final Function<Record, StructLike> IDENTITY_PARTITIONED_GENERATOR =
+ record -> TestHelpers.Row.of(record.get(0, Integer.class));
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(
Review Comment:
>numRecords is always 5 and there are only 2 partition functions. I think it may be better to just have constants up at the top, rather than generating the records. I still have some trouble reading these tests because I have to remember that the IDs are assigned 0-4.
For this, would it help if we pass in an explicit range of values (since it's always 5, this range is constant) but I think the fact taht it's a range may help a reader be able to explicitly see that the ids are always 0-4.....
>Also, a constant sequence number would be helpful to make this match reality a bit more closely.
Sounds good, I'll update to make the seq number constant. I originally made them different so that it more easily distinguishes the values but arguably that doesn't matter for these tests. What we care about is if they're nulled out or preserved (the actual value is inconsequential)
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata attribute
+ // This works around the logic in
+ // ExposesMetadataColumns, used later in metadata attribute resolution,
+ // which prevents surfacing other metadata columns when a single metadata column is in the output
+ val rowLineageAsDataColumns = rowLineageAttributes.map(removeMetadataColumnAttribute)
+
+ val tableWithLineage = r.copy(output =
+ r.output ++ rowLineageAsDataColumns)
+
+ mergeIntoTable.copy(
+ targetTable = tableWithLineage,
+ matchedActions = matchedAssignmentsForLineage,
+ notMatchedBySourceActions = notMatchedBySourceActionsForLineage)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
+ // on the output of operation which indicates the rule already ran
+ private def shouldUpdatePlan(mergeIntoTable: MergeIntoTable): Boolean = {
+ val metadataOutput = mergeIntoTable.targetTable.metadataOutput
+ val rowLineageAttrs = findRowLineageAttributes(metadataOutput)
+ val allLineageAttrsPresent = rowLineageAttrs.nonEmpty && rowLineageAttrs.forall(metadataOutput.contains)
Review Comment:
Hm if `rowLineageAttrs` is empty, the `forall` returns true which we don't want since that means the lineage attributes are not present (in tables not supporting row lineage). Maybe with the suggestion above to always include _row_id this would get reworked anyways
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059093631
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata attribute
+ // This works around the logic in
+ // ExposesMetadataColumns, used later in metadata attribute resolution,
+ // which prevents surfacing other metadata columns when a single metadata column is in the output
+ val rowLineageAsDataColumns = rowLineageAttributes.map(removeMetadataColumnAttribute)
+
+ val tableWithLineage = r.copy(output =
+ r.output ++ rowLineageAsDataColumns)
+
+ mergeIntoTable.copy(
+ targetTable = tableWithLineage,
+ matchedActions = matchedAssignmentsForLineage,
+ notMatchedBySourceActions = notMatchedBySourceActionsForLineage)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
+ // on the output of operation which indicates the rule already ran
+ private def shouldUpdatePlan(mergeIntoTable: MergeIntoTable): Boolean = {
+ val metadataOutput = mergeIntoTable.targetTable.metadataOutput
+ val rowLineageAttrs = findRowLineageAttributes(metadataOutput)
+ val allLineageAttrsPresent = rowLineageAttrs.nonEmpty && rowLineageAttrs.forall(metadataOutput.contains)
Review Comment:
Yeah, I think if `findRowLineageAttributes` returns a tuple of `(rowIdAttr, lastUpdatedAttr)` then we don't need this. You'd just check if both are defined.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2061003382
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,490 @@
+/*
+ * 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.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);
+ assumeThat(vectorized).isFalse();
+ }
+
+ @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[2] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ ImmutableList.of(
+ row(100, "a", 0L, 1L, ANY, ANY),
+ row(101, "updated_b", 1L, updateSnapshot.sequenceNumber(), ANY, ANY),
+ row(102, "c", 2L, 1L, ANY, ANY),
+ row(103, "d", 3L, 1L, ANY, ANY),
+ row(104, "e", 4L, 1L, ANY, ANY)),
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ allRows.stream().filter(row -> (long) row[2] >= updateSnapshotFirstRowId).findFirst().get();
+ assertAddedRowLineage(row(200, "f", updateSnapshot.sequenceNumber()), newRow);
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+
+ createAndInitTable("id INT, data STRING", null);
Review Comment:
Ah copy paste error on my recent update :( good catch!
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2030390082
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java:
##########
@@ -327,7 +327,11 @@ public void pruneColumns(StructType requestedSchema) {
StructType requestedProjection =
new StructType(
Stream.of(requestedSchema.fields())
- .filter(field -> MetadataColumns.nonMetadataColumn(field.name()))
+ .filter(
+ field ->
+ MetadataColumns.nonMetadataColumn(field.name())
+ || field.name().equals("_row_id")
+ || field.name().equals("_last_updated_sequence_number"))
Review Comment:
Fix this so it checks metadata columns with this name
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2040339981
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,266 @@
+/*
+ * 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.Alias
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.plans.logical.AppendData
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows.Keep
+import org.apache.spark.sql.catalyst.plans.logical.Project
+import org.apache.spark.sql.catalyst.plans.logical.ReplaceData
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.types.Metadata
+
+object RewriteMergeIntoTableForRowLineage extends RewriteOperationForRowLineage {
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ RewriteMergeIntoTable.apply(plan) resolveOperators {
+ case writeDelta@WriteDelta(_, _, _, _, _, _) if updatePlanForRowLineage(writeDelta) => {
+ updateDeltaPlan(writeDelta)
+ }
+
+ case replaceData@ReplaceData(_, _, _, _, _, _) if updatePlanForRowLineage(replaceData) => {
+ updateReplaceDataPlan(replaceData)
+ }
+
+ case append@AppendData(table, _, _, _, _, _) if updatePlanForRowLineage(append) =>
Review Comment:
Actually after looking into it, I was overcomplicating it. We don't really need this. In any append case we don't even need an updated plan because we don't need to append explicit nulls.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2048072021
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java:
##########
@@ -192,11 +200,23 @@ private static Schema validateOrMergeWriteSchema(
// if the validation passed, update the table schema
update.commit();
} else {
- writeSchema = SparkSchemaUtil.convert(table.schema(), dsSchema, caseSensitive);
+ Schema schema = table.schema();
+ if (writeIncludesRowLineage) {
+ schema = TypeUtil.join(schema, rowLineageSchema(table));
+ }
+
+ writeSchema = SparkSchemaUtil.convert(schema, dsSchema, caseSensitive);
TypeUtil.validateWriteSchema(
table.schema(), writeSchema, writeConf.checkNullability(), writeConf.checkOrdering());
}
return writeSchema;
}
+
+ private static Schema rowLineageSchema(Table table) {
+ return new Schema(
+ MetadataColumns.metadataColumn(table, MetadataColumns.ROW_ID.name()).asOptional(),
+ MetadataColumns.metadataColumn(table, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name())
+ .asOptional());
Review Comment:
I'm doing this in a few places, maybe move this to TableUtil
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2048073422
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java:
##########
@@ -117,7 +119,10 @@ public WriteBuilder overwrite(Filter[] filters) {
@Override
public Write build() {
// Validate
- Schema writeSchema = validateOrMergeWriteSchema(table, dsSchema, writeConf);
+ Schema writeSchema =
+ validateOrMergeWriteSchema(
+ table, dsSchema, writeConf, TableUtil.formatVersion(table) >= 3 && overwriteFiles);
Review Comment:
I'll Add a comment for this, but in the context of building a regular `SparkWrite`, the only one case where we need to include row lineage explicitly in the output schema is the `overwriteFiles` case for CoW.
In every other case for `SparkWrite`, both fields will be explicitly null for all the fields produced, and as per the spec that means we do not need to write them out. This simplifies a bit because otherwise we'd need to explicitly handle Append cases in the rules to update the DSV2 output but that's not strictly neccessary. cc @rdblue
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2045124799
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala:
##########
@@ -148,9 +151,29 @@ case class ExtendedDataSourceV2Strategy(spark: SparkSession) extends Strategy wi
case UnsetViewProperties(ResolvedV2View(catalog, ident), propertyKeys, ifExists) =>
AlterV2ViewUnsetPropertiesExec(catalog, ident, propertyKeys, ifExists) :: Nil
+ case ReplaceData(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, _, Some(write)) =>
+ ReplaceDataExec(planLater(query), refreshCache(r), write) :: Nil
+
+ case WriteDelta(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections,
+ Some(write)) =>
+ WriteDeltaExec(planLater(query), refreshCache(r), projections, write) :: Nil
+
case _ => Nil
}
+ /**
+ * In case of a table enabled with row lineage, the Dsv2 relation output for the write will have
+ * row lineage in the outputs. However, the cached Dsv2 relation won't have these fields
+ * so refresh after a write should be performed with an updated plan with row lineage outputs removed.
+ * @param r
+ */
+ private def refreshCache(r: DataSourceV2Relation)(): Unit = {
+ val planToCacheBy = r.copy(output = r.output.filterNot(
+ attr => attr.name == MetadataColumns.ROW_ID.name() ||
+ attr.name == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()))
+ spark.sharedState.cacheManager.recacheByPlan(spark, planToCacheBy)
+ }
+
Review Comment:
The only uncertainty is on adding new optimizer rules, but again adding a single optimizer rule for this case seems a lot less intrusive than manipulating the physical execution strategy
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala:
##########
@@ -148,9 +151,29 @@ case class ExtendedDataSourceV2Strategy(spark: SparkSession) extends Strategy wi
case UnsetViewProperties(ResolvedV2View(catalog, ident), propertyKeys, ifExists) =>
AlterV2ViewUnsetPropertiesExec(catalog, ident, propertyKeys, ifExists) :: Nil
+ case ReplaceData(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, _, Some(write)) =>
+ ReplaceDataExec(planLater(query), refreshCache(r), write) :: Nil
+
+ case WriteDelta(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections,
+ Some(write)) =>
+ WriteDeltaExec(planLater(query), refreshCache(r), projections, write) :: Nil
+
case _ => Nil
}
+ /**
+ * In case of a table enabled with row lineage, the Dsv2 relation output for the write will have
+ * row lineage in the outputs. However, the cached Dsv2 relation won't have these fields
+ * so refresh after a write should be performed with an updated plan with row lineage outputs removed.
+ * @param r
+ */
+ private def refreshCache(r: DataSourceV2Relation)(): Unit = {
+ val planToCacheBy = r.copy(output = r.output.filterNot(
+ attr => attr.name == MetadataColumns.ROW_ID.name() ||
+ attr.name == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()))
+ spark.sharedState.cacheManager.recacheByPlan(spark, planToCacheBy)
+ }
+
Review Comment:
The only uncertainty is on adding new optimizer rules, but again adding a single optimizer rule for this case seems a lot less intrusive than manipulating the physical execution strategy like what's being done here.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2050866505
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLineagePropagation.java:
##########
@@ -0,0 +1,441 @@
+/*
+ * 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.TableUtil.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.Comparator;
+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.Files;
+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.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.Lists;
+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.TestTemplate;
+
+public abstract class TestRowLineagePropagation extends SparkRowLevelOperationsTestBase {
Review Comment:
`Propagation` is probably a wrong word to use here.. change it to `RowLevelOperationsWithLineage`?
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLineagePropagation.java:
##########
@@ -0,0 +1,441 @@
+/*
+ * 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.TableUtil.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.Comparator;
+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.Files;
+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.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.Lists;
+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.TestTemplate;
+
+public abstract class TestRowLineagePropagation extends SparkRowLevelOperationsTestBase {
Review Comment:
`Propagation` is probably a wrong word to use here.. change it to `TestRowLevelOperationsWithLineage`?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2051292495
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/iceberg/spark/extensions/RemoveRowLineageOutputFromOriginalTable.scala:
##########
@@ -0,0 +1,56 @@
+/*
+ * 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 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.ReplaceData
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+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] {
Review Comment:
This should be moved to the optimizer. Also, tests for regular deletes
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2040528348
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/SparkParquetWriters.java:
##########
@@ -557,14 +559,39 @@ public Map.Entry<K, V> next() {
private static class InternalRowWriter extends ParquetValueWriters.StructWriter<InternalRow> {
private final DataType[] types;
+ private final Integer[] rowLineageAttributeIndices = {null, null};
private InternalRowWriter(List<ParquetValueWriter<?>> writers, List<DataType> types) {
super(writers);
+ int writerIndex = 0;
+ for (ParquetValueWriter<?> writer : writers) {
+ if (writer.columns().size() == 1 && writer.columns().get(0) instanceof ColumnWriter) {
+ ColumnWriter<?> columnWriter = (ColumnWriter<?>) writer.columns().get(0);
+ if (columnWriter.desc().getPrimitiveType().getId().intValue()
+ == MetadataColumns.ROW_ID.fieldId()) {
+ rowLineageAttributeIndices[0] = writerIndex;
+ }
+ if (columnWriter.desc().getPrimitiveType().getId().intValue()
+ == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.fieldId()) {
+ rowLineageAttributeIndices[1] = writerIndex;
+ }
+ writerIndex++;
Review Comment:
Not needed, I was trying something.....
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2045115139
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala:
##########
@@ -148,9 +151,29 @@ case class ExtendedDataSourceV2Strategy(spark: SparkSession) extends Strategy wi
case UnsetViewProperties(ResolvedV2View(catalog, ident), propertyKeys, ifExists) =>
AlterV2ViewUnsetPropertiesExec(catalog, ident, propertyKeys, ifExists) :: Nil
+ case ReplaceData(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, _, Some(write)) =>
+ ReplaceDataExec(planLater(query), refreshCache(r), write) :: Nil
+
+ case WriteDelta(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections,
+ Some(write)) =>
+ WriteDeltaExec(planLater(query), refreshCache(r), projections, write) :: Nil
+
case _ => Nil
}
+ /**
+ * In case of a table enabled with row lineage, the Dsv2 relation output for the write will have
+ * row lineage in the outputs. However, the cached Dsv2 relation won't have these fields
+ * so refresh after a write should be performed with an updated plan with row lineage outputs removed.
+ * @param r
+ */
+ private def refreshCache(r: DataSourceV2Relation)(): Unit = {
+ val planToCacheBy = r.copy(output = r.output.filterNot(
+ attr => attr.name == MetadataColumns.ROW_ID.name() ||
+ attr.name == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()))
+ spark.sharedState.cacheManager.recacheByPlan(spark, planToCacheBy)
+ }
+
Review Comment:
Alternatively, I could add an optimizer rule which takes the existing ReplaceData/WriteDelta where the original table has the row lineage output, and then remove it. Then adding to the extended DSV2 strategy isn't needed. I tried this out and it also seems to work. It's a bit more code but not much more, and arguably it's a lot less risky than manipulating the logic to recache by a certain plan (we just would do what Spark does, which seems the most robust)
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2045142385
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala:
##########
@@ -148,9 +151,29 @@ case class ExtendedDataSourceV2Strategy(spark: SparkSession) extends Strategy wi
case UnsetViewProperties(ResolvedV2View(catalog, ident), propertyKeys, ifExists) =>
AlterV2ViewUnsetPropertiesExec(catalog, ident, propertyKeys, ifExists) :: Nil
+ case ReplaceData(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, _, Some(write)) =>
+ ReplaceDataExec(planLater(query), refreshCache(r), write) :: Nil
+
+ case WriteDelta(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections,
+ Some(write)) =>
+ WriteDeltaExec(planLater(query), refreshCache(r), projections, write) :: Nil
+
case _ => Nil
}
+ /**
+ * In case of a table enabled with row lineage, the Dsv2 relation output for the write will have
+ * row lineage in the outputs. However, the cached Dsv2 relation won't have these fields
+ * so refresh after a write should be performed with an updated plan with row lineage outputs removed.
+ * @param r
+ */
+ private def refreshCache(r: DataSourceV2Relation)(): Unit = {
+ val planToCacheBy = r.copy(output = r.output.filterNot(
+ attr => attr.name == MetadataColumns.ROW_ID.name() ||
+ attr.name == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()))
+ spark.sharedState.cacheManager.recacheByPlan(spark, planToCacheBy)
+ }
+
Review Comment:
And even if we go with the current approach, probably also want to call `reacacheByPlan` with the full relation just in case that's also cached. This level of complexity probably means this isn't the right way
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2042588893
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestMerge.java:
##########
@@ -1337,6 +1337,7 @@ public void testMergeWithBucketTransform() {
row(1, "emp-id-1"), // updated
row(2, "emp-id-2") // new
);
+
Review Comment:
Unnecessary whitespace change.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2050842658
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java:
##########
@@ -327,7 +327,11 @@ public void pruneColumns(StructType requestedSchema) {
StructType requestedProjection =
new StructType(
Stream.of(requestedSchema.fields())
- .filter(field -> MetadataColumns.nonMetadataColumn(field.name()))
+ .filter(
+ field ->
+ MetadataColumns.nonMetadataColumn(field.name())
+ || field.name().equals("_row_id")
Review Comment:
Actually this code path is no longer needed, removed it entirely!
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053121924
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
Review Comment:
This is only delete right?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2053163626
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ // Treat row lineage columns as data columns by removing the metadata attribute
+ // This works around the logic in
+ // ExposesMetadataColumns, used later in metadata attribute resolution,
+ // which prevents surfacing other metadata columns when a single metadata column is in the output
+ val rowLineageAsDataColumns = rowLineageAttributes.map(removeMetadataColumnAttribute)
+
+ val tableWithLineage = r.copy(output =
+ r.output ++ rowLineageAsDataColumns)
+
+ mergeIntoTable.copy(
+ targetTable = tableWithLineage,
+ matchedActions = matchedAssignmentsForLineage,
+ notMatchedBySourceActions = notMatchedBySourceActionsForLineage)
+ }
+ }
+
+ // The plan should only be updated if row lineage metadata attributes are present
+ // in the target table AND lineage attributes are not already
Review Comment:
>Don't we always want to have _row_id available but read null if the table isn't v3?
Hm I looked at it as we only want to have _row_id available for v3+ since the existence of that column in the projection could be potentially misleading, though arguably since it would just be null means that projecting it is fine since that implicitly indicates there is no row lineage.
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2048073422
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkWriteBuilder.java:
##########
@@ -117,7 +119,10 @@ public WriteBuilder overwrite(Filter[] filters) {
@Override
public Write build() {
// Validate
- Schema writeSchema = validateOrMergeWriteSchema(table, dsSchema, writeConf);
+ Schema writeSchema =
+ validateOrMergeWriteSchema(
+ table, dsSchema, writeConf, TableUtil.formatVersion(table) >= 3 && overwriteFiles);
Review Comment:
I'll Add a comment for this, but in the context of building a regular `SparkWrite`, the only one case where we need to include row lineage explicitly in the output schema is the `overwriteFiles` case for CoW.
In every other case, both fields will be explicitly null for all the fields produced, and as per the spec that means we do not need to write them out. cc @rdblue
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2050867791
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLineagePropagation.java:
##########
@@ -0,0 +1,441 @@
+/*
+ * 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.TableUtil.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.Comparator;
+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.Files;
+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.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.Lists;
+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.TestTemplate;
+
+public abstract class TestRowLineagePropagation extends SparkRowLevelOperationsTestBase {
+ static final Function<Record, StructLike> UNPARTITIONED_GENERATOR = record -> null;
+ // Use the first field in the row as identity partition
+ static final Function<Record, StructLike> IDENTITY_PARTITIONED_GENERATOR =
+ record -> TestHelpers.Row.of(record.get(0, Integer.class));
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(
+ schemaWithRowLineage(table), table.spec(), 5, UNPARTITIONED_GENERATOR);
+ appendRecords(table, recordsByPartition);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 5678}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT *",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(filterByDataNotEqualTo(allRecords(recordsByPartition), 1));
+ Record updatedRecord = createRecordWithRowLineage(schemaWithRowLineage(table), 1234, 1L, null);
+ Record insertedRecord =
+ createRecordWithRowLineage(schemaWithRowLineage(table), 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ assertRecordsAreEqual(sortByRowId(expectedRecords), sortByRowId(actualRecords(table)));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithBothMatchedAndNonMatchedPartitioned()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", "PARTITIONED BY (data)", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = schemaWithRowLineage(table);
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(schema, table.spec(), 5, IDENTITY_PARTITIONED_GENERATOR);
+ appendRecords(table, recordsByPartition);
+
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 5678}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT *",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(filterByDataNotEqualTo(allRecords(recordsByPartition), 1));
+ Record updatedRecord = createRecordWithRowLineage(schema, 1234, 1L, null);
+ Record insertedRecord = createRecordWithRowLineage(schema, 5678, null, null);
+ expectedRecords.addAll(ImmutableList.of(insertedRecord, updatedRecord));
+
+ assertRecordsAreEqual(sortByRowId(expectedRecords), sortByRowId(actualRecords(table)));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyNonMatched()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = schemaWithRowLineage(table);
+
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(schema, table.spec(), 5, UNPARTITIONED_GENERATOR);
+ appendRecords(table, recordsByPartition);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 5678}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data "
+ + "WHEN NOT MATCHED THEN "
+ + "INSERT *",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(allRecords(recordsByPartition));
+ Record insertedRecord = createRecordWithRowLineage(schema, 5678, null, null);
+ expectedRecords.add(insertedRecord);
+
+ assertRecordsAreEqual(sortByRowId(expectedRecords), sortByRowId(actualRecords(table)));
+ }
+
+ @TestTemplate
+ public void testMergeIntoWithOnlyMatched()
+ throws IOException, NoSuchTableException, ParseException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = schemaWithRowLineage(table);
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(
+ schemaWithRowLineage(table), table.spec(), 5, UNPARTITIONED_GENERATOR);
+
+ appendRecords(table, recordsByPartition);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = 1234 ",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(filterByDataNotEqualTo(allRecords(recordsByPartition), 1));
+ Record updatedRecord = createRecordWithRowLineage(schema, 1234, 1L, null);
+ expectedRecords.add(updatedRecord);
+
+ assertRecordsAreEqual(sortByRowId(expectedRecords), sortByRowId(actualRecords(table)));
+ }
+
+ @TestTemplate
+ public void testMergeMatchedDelete() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = schemaWithRowLineage(table);
+
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(schema, table.spec(), 5, UNPARTITIONED_GENERATOR);
+ appendRecords(table, recordsByPartition);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 2}\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.data == s.data AND t.data = 1 "
+ + "WHEN MATCHED THEN DELETE",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(filterByDataNotEqualTo(allRecords(recordsByPartition), 1));
+ assertRecordsAreEqual(sortByRowId(expectedRecords), sortByRowId(actualRecords(table)));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySource()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(
+ schemaWithRowLineage(table), table.spec(), 5, UNPARTITIONED_GENERATOR);
+ appendRecords(table, recordsByPartition);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN UPDATE set data = 5678",
+ commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ // Every record should preserve its row ID and null out its sequence number.
+ // Every record other than id 1 should be 5678 and
+ for (Record record : allRecords(recordsByPartition)) {
+ Record newRecord = record.copy();
+ if (record.get(0, Integer.class) == 1) {
+ newRecord.set(0, 1234);
+ newRecord.set(2, null);
+ } else {
+ newRecord.set(0, 5678);
+ newRecord.set(2, null);
+ }
+
+ expectedRecords.add(newRecord);
+ }
+
+ assertRecordsAreEqual(sortByRowId(expectedRecords), sortByRowId(actualRecords(table)));
+ }
+
+ @TestTemplate
+ public void testMergeWhenNotMatchedBySourceDelete()
+ throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(
+ schemaWithRowLineage(table), table.spec(), 5, UNPARTITIONED_GENERATOR);
+ appendRecords(table, recordsByPartition);
+ createOrReplaceView("source", "data INT", "{ \"data\": 1}\n" + "{ \"data\": 42}\n");
+
+ sql(
+ "MERGE INTO %s AS t USING source AS s ON t.data == s.data AND t.data = 1"
+ + " WHEN MATCHED THEN UPDATE set data = 1234 "
+ + "WHEN NOT MATCHED BY SOURCE THEN DELETE",
+ commitTarget());
+
+ Record recordWithData1 =
+ allRecords(recordsByPartition).stream()
+ .filter(record -> record.get(0, Integer.class).equals(1))
+ .findFirst()
+ .get();
+ List<Record> expectedRecords = Lists.newArrayList();
+ Record record = recordWithData1.copy();
+ record.set(0, 1234);
+ record.set(2, null);
+ expectedRecords.add(record);
+
+ assertRecordsAreEqual(sortByRowId(expectedRecords), sortByRowId(actualRecords(table)));
+ }
+
+ @TestTemplate
+ public void testUpdate() throws NoSuchTableException, ParseException, IOException {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+ // write the data file with <data, _row_id, _last_updated_sequence_number>
+ // merge, update data to data + 1 on odd values of data. _row_id should be preserved for those
+ // cases, _last_updated_sequence_number should be null
+ // even values of data should be preserved as is
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ Schema schema = schemaWithRowLineage(table);
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(schema, table.spec(), 5, UNPARTITIONED_GENERATOR);
+ appendRecords(table, recordsByPartition);
+
+ sql("UPDATE %s AS t set data = 5678 WHERE data = 1", commitTarget());
+
+ List<Record> expectedRecords = Lists.newArrayList();
+ expectedRecords.addAll(filterByDataNotEqualTo(allRecords(recordsByPartition), 1));
+ Record updatedRecord = createRecordWithRowLineage(schema, 5678, 1L, null);
+ expectedRecords.add(updatedRecord);
+
+ assertRecordsAreEqual(sortByRowId(expectedRecords), sortByRowId(actualRecords(table)));
+ }
+
+ protected Record createRecordWithRowLineage(
+ Schema schema, int id, Long rowId, Long lastUpdatedSequenceNumber) {
+ Record record = GenericRecord.create(schema);
+ record.set(0, id);
+ record.set(1, rowId);
+ record.set(2, lastUpdatedSequenceNumber);
+ return record;
+ }
+
+ protected List<Record> actualRecords(Table table) throws IOException {
+ List<Record> records = Lists.newArrayList();
+ table.refresh();
+ Snapshot snapshot = branch != null ? table.snapshot(branch) : table.currentSnapshot();
+ try (CloseableIterable<Record> recordsOnDisk =
+ IcebergGenerics.read(table)
+ .project(schemaWithRowLineage(table))
+ .useSnapshot(snapshot.snapshotId())
+ .build()) {
+ recordsOnDisk.forEach(records::add);
+ }
+
+ return records;
+ }
+
+ protected List<Record> sortByRowId(List<Record> records) {
+ records.sort(
+ Comparator.comparing(
+ record -> record.get(1, Long.class), Comparator.nullsFirst(Comparator.naturalOrder())));
+ return records;
+ }
+
+ /**
+ * Generates numRecords records with data values as 0 to numRecords - 1, with row lineage, each
+ * record is partitioned based on the provided partition generator. The row ID and last updated
+ * sequence number assigned is the same as the data value for that row.
+ *
+ * @return a partitioned map
+ */
+ protected PartitionMap<List<Record>> createRecordsWithRowLineage(
+ Schema schema,
+ PartitionSpec spec,
+ int numRecords,
+ Function<Record, StructLike> partitionGenerator) {
+ PartitionMap<List<Record>> recordsByPartition =
+ PartitionMap.create(Map.of(spec.specId(), spec));
+ for (int i = 0; i < numRecords; i++) {
+ long rowId = i;
+ long lastUpdatedSequenceNumber = i;
+ Record record = createRecordWithRowLineage(schema, i, rowId, lastUpdatedSequenceNumber);
+ 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 List<Record> filterByDataNotEqualTo(List<Record> records, int data) {
+ return records.stream()
+ .filter(record -> record.get(0, Integer.class) != data)
+ .collect(Collectors.toList());
+ }
+
+ 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), 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 void assertRecordsAreEqual(List<Record> expected, List<Record> actual) {
+ assertThat(actual.size()).isEqualTo(expected.size());
+ for (int i = 0; i < expected.size(); i++) {
+ Record expectedRecord = expected.get(i);
+ Record actualRecord = actual.get(i);
+ assertThat(expectedRecord.get(0)).isEqualTo(actualRecord.get(0));
+ assertThat(expectedRecord.get(1)).isEqualTo(actualRecord.get(1));
+ assertThat(expectedRecord.get(2)).isEqualTo(actualRecord.get(2));
Review Comment:
in MoR the records returned by the generics code includes the file position for pos deletes so an exact equality faills since there's an additional field...that's why this compares the data, row id, sequence values directly. I'll clean this up so it's more obvious to a reader
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2056321860
##########
spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.TableUtil.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.Comparator;
+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.Files;
+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.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.Lists;
+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.TestTemplate;
+
+public abstract class TestRowLevelOperationsWithLineage extends SparkRowLevelOperationsTestBase {
+ static final Function<Record, StructLike> UNPARTITIONED_GENERATOR = record -> null;
+ // Use the first field in the row as identity partition
+ static final Function<Record, StructLike> IDENTITY_PARTITIONED_GENERATOR =
+ record -> TestHelpers.Row.of(record.get(0, Integer.class));
+
+ @BeforeAll
+ public static void setupSparkConf() {
+ spark.conf().set("spark.sql.shuffle.partitions", "4");
+ }
+
+ @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 {
+ assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
+
+ createAndInitTable("data INT", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+ PartitionMap<List<Record>> recordsByPartition =
+ createRecordsWithRowLineage(
Review Comment:
>I think it may be better to just have constants up at the top, rather than generating the records
Oh I see what you mean, at least currently, the records themselves are constants. Good point!
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060873913
##########
core/src/main/java/org/apache/iceberg/MetadataColumns.java:
##########
@@ -152,4 +154,19 @@ public static boolean isMetadataColumn(int id) {
public static boolean nonMetadataColumn(String name) {
return !isMetadataColumn(name);
}
+
+ public static Schema schemaWithRowLineage(Table table) {
Review Comment:
I was just looking at this again and totally agree. Users should just pass in a schema, and we can directly use the MetadataColumns.ROW_ID internally
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060881490
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteOperation.java:
##########
@@ -95,13 +98,18 @@ public WriteBuilder newWriteBuilder(LogicalWriteInfo info) {
@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};
+ if (TableUtil.supportsRowLineage(table)) {
+ metadataAttributes.add(Expressions.column(MetadataColumns.ROW_POSITION.name()));
+ metadataAttributes.add(Expressions.column(MetadataColumns.ROW_ID.name()));
+ metadataAttributes.add(
+ Expressions.column(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name()));
+ } else if (command == DELETE || command == UPDATE) {
+ metadataAttributes.add(Expressions.column(MetadataColumns.ROW_POSITION.name()));
}
+
+ return metadataAttributes.toArray(NamedReference[]::new);
Review Comment:
If it's a delete or update add the row position fist. Then if it's row lineage supported table add those columns. That's easier to parse I think compared to the else if structure now
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "rdblue (via GitHub)" <gi...@apache.org>.
rdblue commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2060952937
##########
core/src/main/java/org/apache/iceberg/MetadataColumns.java:
##########
@@ -152,4 +154,19 @@ public static boolean isMetadataColumn(int id) {
public static boolean nonMetadataColumn(String name) {
return !isMetadataColumn(name);
}
+
+ public static Schema schemaWithRowLineage(Table table) {
+ Preconditions.checkArgument(null != table, "Invalid table: null");
+ Preconditions.checkArgument(
+ !(table instanceof BaseMetadataTable),
+ "Cannot produce row lineage for metadata table: %s",
+ table);
+ Schema rowLineageSchema =
+ new Schema(
+ MetadataColumns.metadataColumn(table, MetadataColumns.ROW_ID.name()).asOptional(),
+ MetadataColumns.metadataColumn(
+ table, MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name())
+ .asOptional());
+ return TypeUtil.join(table.schema(), rowLineageSchema);
Review Comment:
The update looks good. We may want to have a few different forms, like `join(Schema, NestedField...)` but this is much simpler now that the `Table` has been removed. Thanks!
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2030390082
##########
spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java:
##########
@@ -327,7 +327,11 @@ public void pruneColumns(StructType requestedSchema) {
StructType requestedProjection =
new StructType(
Stream.of(requestedSchema.fields())
- .filter(field -> MetadataColumns.nonMetadataColumn(field.name()))
+ .filter(
+ field ->
+ MetadataColumns.nonMetadataColumn(field.name())
+ || field.name().equals("_row_id")
+ || field.name().equals("_last_updated_sequence_number"))
Review Comment:
Fix this so it checks metadata columns with this name
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Spark 3.5 row lineage [iceberg]
Posted by "amogh-jahagirdar (via GitHub)" <gi...@apache.org>.
amogh-jahagirdar commented on code in PR #12736:
URL: https://github.com/apache/iceberg/pull/12736#discussion_r2059551800
##########
spark/v3.5/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTableForRowLineage.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * 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.MergeIntoTable
+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 @ MergeIntoTable(_, _, _, matchedActions, _, notMatchedBySourceActions)
+ if m.resolved && m.rewritable && m.aligned &&
+ (matchedActions.nonEmpty || notMatchedBySourceActions.nonEmpty) &&
+ shouldUpdatePlan(m) =>
+ updateMergeIntoForRowLineage(m)
+ }
+ }
+
+ protected def updateMergeIntoForRowLineage(mergeIntoTable: MergeIntoTable): LogicalPlan = {
+ EliminateSubqueryAliases(mergeIntoTable.targetTable) match {
+ case r: DataSourceV2Relation =>
+ val matchedActions = mergeIntoTable.matchedActions
+ val notMatchedBySourceActions = mergeIntoTable.notMatchedBySourceActions
+ val rowLineageAttributes = findRowLineageAttributes(r.metadataOutput)
+ val rowId = rowLineageAttributes.filter(
+ attr => attr.name == ROW_ID_ATTRIBUTE_NAME).head
+ val lastUpdatedSequence = rowLineageAttributes.filter(
+ attr => attr.name == LAST_UPDATED_SEQUENCE_NUMBER_ATTRIBUTE_NAME).head
+
+ val matchedAssignmentsForLineage = matchedActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
+ Assignment(lastUpdatedSequence, Literal(null))))
+
+ case p => p
+ }
+
+ val notMatchedBySourceActionsForLineage = notMatchedBySourceActions.map {
+ case UpdateAction(cond, actions) =>
+ UpdateAction(cond, actions ++ Seq(Assignment(rowId, rowId),
Review Comment:
1. Yeah this is automatic and exactly in the alignment rule. Though to be clear this isn't relying on the alignment rule running again. The [alignment rule ](https://github.com/apache/spark/blame/84b9848f9a69c32faa9baf623a90bbb2f4939d04/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala#L414) runs before extended resolution rules.
Specifically, at https://github.com/apache/spark/blob/84b9848f9a69c32faa9baf623a90bbb2f4939d04/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AssignmentUtils.scala#L109 we'll get a null literal, and at then assign it to the row lineage fields.
2. Yes definitely agree, I'll add a comment (for the insert case as well)!
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org