From 725f18b6753a8875da101291e10554becdbf5723 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Tue, 25 Aug 2026 17:08:34 +0800 Subject: [PATCH 01/22] Add GPU support for Iceberg v3 row lineage writes Signed-off-by: Chong Gao --- .../spark/source/GpuSparkScanAccess.java | 7 +- .../spark/source/GpuSparkWriteAccess.java | 7 + .../rapids/iceberg/IcebergProviderBase.scala | 15 + .../source/GpuSparkPositionDeltaWrite.scala | 4 +- .../iceberg/spark/source/GpuSparkScan.scala | 6 +- .../spark/source/GpuSparkStagedScan.scala | 47 +++ .../iceberg/spark/source/GpuSparkWrite.scala | 8 +- .../apache/iceberg/spark/source/write.scala | 10 +- .../main/python/iceberg/iceberg_ctas_test.py | 37 +- .../src/main/python/iceberg/iceberg_test.py | 333 +++++++++++++++++- .../rapids/iceberg/IcebergProvider.scala | 1 + .../datasources/v2/GpuMergeRowsExec.scala | 29 +- 12 files changed, 489 insertions(+), 15 deletions(-) create mode 100644 iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkStagedScan.scala diff --git a/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkScanAccess.java b/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkScanAccess.java index 26d6e46ccde..9719c783cf8 100644 --- a/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkScanAccess.java +++ b/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkScanAccess.java @@ -47,7 +47,8 @@ private GpuSparkScanAccess() { } public static boolean supports(Scan scan) { - return scan instanceof SparkBatchQueryScan || scan instanceof SparkCopyOnWriteScan; + return scan instanceof SparkBatchQueryScan || scan instanceof SparkCopyOnWriteScan + || scan instanceof SparkStagedScan; } public static boolean isBatchQueryScan(Scan scan) { @@ -58,6 +59,10 @@ public static boolean isCopyOnWriteScan(Scan scan) { return scan instanceof SparkCopyOnWriteScan; } + public static boolean isStagedScan(Scan scan) { + return scan instanceof SparkStagedScan; + } + public static boolean isMetadataScan(Scan scan) { return sparkScan(scan).table() instanceof BaseMetadataTable; } diff --git a/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkWriteAccess.java b/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkWriteAccess.java index ed044b48250..007d3e8b998 100644 --- a/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkWriteAccess.java +++ b/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkWriteAccess.java @@ -17,14 +17,17 @@ package org.apache.iceberg.spark.source; import java.lang.reflect.Field; +import java.util.List; import java.util.Map; +import org.apache.iceberg.ContentFile; import org.apache.iceberg.DataFile; import org.apache.iceberg.FileFormat; import org.apache.iceberg.Schema; import org.apache.iceberg.Table; import org.apache.iceberg.deletes.DeleteGranularity; import org.apache.iceberg.io.DeleteWriteResult; +import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.WriteResult; import org.apache.spark.api.java.JavaSparkContext; import org.apache.spark.sql.connector.write.RowLevelOperation.Command; @@ -53,6 +56,10 @@ public static String sparkWriteClassName() { return SparkWrite.class.getName(); } + public static void deleteTaskFiles(FileIO io, List> files) { + SparkCleanupUtil.deleteTaskFiles(io, files); + } + public static Table table(Write write) { return readField(sparkWrite(write), "table", Table.class); } diff --git a/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala b/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala index 2c53190a658..fd3de6c2a45 100644 --- a/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala +++ b/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala @@ -43,6 +43,8 @@ abstract class IcebergProviderBase extends IcebergProvider { IcebergProvider.cpuBatchQueryScanClassName) val cpuCopyOnWriteScanClass = ShimReflectionUtils.loadClass( IcebergProvider.cpuCopyOnWriteScanClassName) + val cpuStagedScanClass = ShimReflectionUtils.loadClass( + IcebergProvider.cpuStagedScanClassName) Seq( new ScanRule[Scan]( @@ -75,6 +77,19 @@ abstract class IcebergProviderBase extends IcebergProvider { "Iceberg copy on write scan", ClassTag(cpuCopyOnWriteScanClass) ), + new ScanRule[Scan]( + (a, conf, p, r) => new ScanMeta[Scan](a, conf, p, r) { + private lazy val convertedScan: Try[GpuSparkScan] = GpuSparkScan.tryConvert(a, this.conf) + + override def tagSelfForGpu(): Unit = { + GpuSparkScan.tagForGpu(this, convertedScan) + } + + override def convertToGpu(): GpuScan = convertedScan.get + }, + "Iceberg staged scan", + ClassTag(cpuStagedScanClass) + ), ).map(r => (r.getClassFor.asSubclass(classOf[Scan]), r)).toMap } diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala index fd1b24a880a..2c827981603 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala @@ -548,7 +548,7 @@ trait GpuDeleteAndDataDeltaWriter extends GpuDeltaWriter { val files = mutable.ListBuffer[ContentFile[_]]() files ++= result.dataFiles().map(_.asInstanceOf[ContentFile[_]]) files ++= result.deleteFiles().map(_.asInstanceOf[ContentFile[_]]) - SparkCleanupUtil.deleteTaskFiles(io, files.asJava) + GpuSparkWriteAccess.deleteTaskFiles(io, files.asJava) } override def close(): Unit = { @@ -651,7 +651,7 @@ class GpuDeleteOnlyDeltaWriter( override def abort(): Unit = { close() val result = delegate.result() - SparkCleanupUtil.deleteTaskFiles(io, result.deleteFiles()) + GpuSparkWriteAccess.deleteTaskFiles(io, result.deleteFiles()) } override def close(): Unit = { diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkScan.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkScan.scala index 1da2dd40a94..680adf9f05d 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkScan.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkScan.scala @@ -83,10 +83,12 @@ object GpuSparkScan { new GpuSparkBatchQueryScan(cpuScan, rapidsConf, false) } else if (GpuSparkScanAccess.isCopyOnWriteScan(cpuScan)) { ShimUtils.newCopyOnWriteScan(cpuScan, rapidsConf, false) + } else if (GpuSparkScanAccess.isStagedScan(cpuScan)) { + new GpuSparkStagedScan(cpuScan, rapidsConf, false) } else { throw new IllegalArgumentException( - s"Currently iceberg support only supports batch query scan and copy-on-write scan, " + - s"but got ${cpuScan.getClass.getName}") + s"Currently Iceberg support only supports batch query, copy-on-write, and staged " + + s"scans, but got ${cpuScan.getClass.getName}") } } } diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkStagedScan.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkStagedScan.scala new file mode 100644 index 00000000000..04f92bf01c5 --- /dev/null +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkStagedScan.scala @@ -0,0 +1,47 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed 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.source + +import scala.collection.JavaConverters._ + +import com.nvidia.spark.rapids.{GpuScan, RapidsConf} +import org.apache.iceberg.ScanTaskGroup +import org.apache.iceberg.types.Types + +import org.apache.spark.sql.connector.read.Scan + +/** GPU scan for file groups staged by Iceberg's rewrite_data_files action. */ +class GpuSparkStagedScan( + override val cpuScan: Scan, + override val rapidsConf: RapidsConf, + override val queryUsesInputFile: Boolean) + extends GpuSparkScan(cpuScan, rapidsConf, queryUsesInputFile) { + + override def groupingKeyType(): Types.StructType = + GpuSparkScanAccess.groupingKeyType(cpuScan) + + override def taskGroups(): Seq[_ <: ScanTaskGroup[_]] = + GpuSparkScanAccess.taskGroups(cpuScan).asScala.toSeq + + override def withInputFile(): GpuScan = + new GpuSparkStagedScan(cpuScan, rapidsConf, true) + + override def toString: String = + s"GpuSparkStagedScan(table=${GpuSparkScanAccess.table(cpuScan)}, " + + s"type=${GpuSparkScanAccess.expectedSchema(cpuScan).asStruct()}, " + + s"queryUseInputFile=$queryUsesInputFile)" +} diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala index f2659d5ddea..6c08571cb74 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala @@ -64,7 +64,8 @@ class GpuSparkWrite(cpu: Write) extends GpuWrite with RequiresDistributionAndOrd // Iceberg's SparkWrite returns different implementations based on write mode: // - BatchAppend for append operations // - DynamicOverwrite for dynamic partition overwrite - // - BatchRewrite for copy-on-write operations (DELETE) + // - CopyOnWriteOperation for row-level copy-on-write operations + // - RewriteFiles for rewrite_data_files // Since these are private classes, we check the class name to determine which GPU version // to use val cpuBatch = cpu.toBatch @@ -75,6 +76,7 @@ class GpuSparkWrite(cpu: Write) extends GpuWrite with RequiresDistributionAndOrd case "DynamicOverwrite" => new GpuDynamicOverwrite(this, cpuBatch) case "OverwriteByFilter" => new GpuOverwriteByFilter(this, cpuBatch) case "CopyOnWriteOperation" => new GpuCopyOnWriteOperation(this, cpuBatch) + case "RewriteFiles" => new GpuRewriteFiles(this, cpuBatch) case _ => throw new UnsupportedOperationException( s"Unsupported Iceberg batch write type: $cpuBatchClassName") @@ -446,7 +448,7 @@ class GpuUnpartitionedDataWriter( close() val result = delegate.result() - SparkCleanupUtil.deleteTaskFiles(io, result.dataFiles()) + GpuSparkWriteAccess.deleteTaskFiles(io, result.dataFiles()) } override def close(): Unit = { @@ -494,7 +496,7 @@ class GpuPartitionedDataWriter( close() val result = delegate.result() - SparkCleanupUtil.deleteTaskFiles(io, result.dataFiles()) + GpuSparkWriteAccess.deleteTaskFiles(io, result.dataFiles()) } override def close(): Unit = { diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala index 4861d846e73..a03f015d108 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala @@ -74,6 +74,14 @@ class GpuCopyOnWriteOperation(write: GpuSparkWrite, cpuBatchWrite: BatchWrite) } } +/** GPU version of the batch write used by the rewrite_data_files procedure. */ +class GpuRewriteFiles(write: GpuSparkWrite, cpuBatchWrite: BatchWrite) + extends GpuBaseBatchWrite(write, cpuBatchWrite) { + override def commit(messages: Array[WriterCommitMessage]): Unit = { + cpuBatchWrite.commit(messages) + } +} + /** * GPU version of position delta batch write for merge-on-read DELETE operations. * This wraps the CPU PositionDeltaBatchWrite to handle position delete files. @@ -100,4 +108,4 @@ class GpuPositionDeltaBatchWrite(write: GpuSparkPositionDeltaWrite, override def createBatchWriterFactory(info: PhysicalWriteInfo): DeltaWriterFactory = { write.createDeltaWriterFactory } -} \ No newline at end of file +} diff --git a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py index fae07bb9999..ad44b7bc323 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py @@ -27,7 +27,9 @@ get_full_table_name, iceberg_write_enabled_conf, iceberg_unsupported_mark, _build_tblprops, ctas_partition_transforms, supports_iceberg_v3, - ICEBERG_V3_UNSUPPORTED_REASON) + ICEBERG_V3_UNSUPPORTED_REASON, + supports_iceberg_row_lineage_inheritance, + ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) from marks import iceberg, ignore_order, allow_non_gpu, allow_non_gpu_conditional, datagen_overrides from spark_session import with_gpu_session, with_cpu_session, is_spark_400_or_later @@ -133,6 +135,39 @@ def run_ctas(spark): conf=iceberg_write_enabled_conf) +@iceberg +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +@ignore_order(local=True) +def test_ctas_v3_row_lineage(spark_tmp_table_factory): + table_name = get_full_table_name(spark_tmp_table_factory) + conf = copy_and_update(iceberg_write_enabled_conf, { + "spark.rapids.sql.format.iceberg.v3.enabled": "true" + }) + + with_gpu_session( + lambda spark: _execute_ctas( + spark, + table_name, + spark_tmp_table_factory, + lambda sp: sp.range(3), + {"format-version": "3"}, + ret=False), + conf=conf) + + rows = with_cpu_session( + lambda spark: spark.sql( + f"SELECT id, _row_id, _last_updated_sequence_number FROM {table_name} " + "ORDER BY id").collect()) + assert [(row["id"], row["_row_id"], row["_last_updated_sequence_number"]) + for row in rows] == [ + (0, 0, 1), + (1, 1, 1), + (2, 2, 1) + ] + + @iceberg @pytest.mark.skipif(not supports_iceberg_v3, reason=ICEBERG_V3_UNSUPPORTED_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Requires a local Hadoop catalog") diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index d3610d71225..674bce2a064 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -22,7 +22,7 @@ from conftest import is_iceberg_remote_catalog, is_iceberg_rest_catalog from data_gen import * from iceberg import get_full_table_name, iceberg_unsupported_mark, _build_tblprops, \ - _BASE_TBLPROPS_SQL, create_iceberg_table, supports_iceberg_v3, \ + _BASE_TBLPROPS_SQL, create_iceberg_table, iceberg_write_enabled_conf, supports_iceberg_v3, \ ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, \ ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, row_lineage_df from marks import allow_non_gpu, iceberg, ignore_order @@ -50,6 +50,14 @@ rapids_reader_types = ['PERFILE', 'MULTITHREADED', 'COALESCING'] _NO_FANOUT = _BASE_TBLPROPS_SQL +_ROW_LINEAGE_WRITE_CONF = { + **iceberg_write_enabled_conf, + "spark.rapids.sql.format.iceberg.v3.enabled": "true" +} + + +def _with_gpu_lineage_write(write_func): + return with_gpu_session(write_func, conf=_ROW_LINEAGE_WRITE_CONF) pytestmark = iceberg_unsupported_mark @@ -469,6 +477,234 @@ def setup_iceberg_table(spark): conf=read_conf) +@iceberg +@ignore_order(local=True) +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +@pytest.mark.parametrize("reader_type", rapids_reader_types) +def test_iceberg_v3_row_lineage_append(spark_tmp_table_factory, reader_type): + full_table = get_full_table_name(spark_tmp_table_factory) + + def setup_iceberg_table(spark): + spark.sql(f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG " + f"TBLPROPERTIES ('format-version' = '2')") + spark.sql(f"INSERT INTO {full_table} VALUES (1), (2)") + v2_snapshot_id = spark.sql( + f"SELECT snapshot_id FROM {full_table}.snapshots ORDER BY committed_at DESC") \ + .head()[0] + spark.sql( + f"ALTER TABLE {full_table} SET TBLPROPERTIES (" + "'format-version' = '3', " + "'write.parquet.row-group-size-bytes' = '4096', " + "'read.split.target-size' = '4096', " + "'read.split.open-file-cost' = '0')") + + legacy = spark.sql( + f"SELECT id, _row_id, _last_updated_sequence_number FROM {full_table} " + f"VERSION AS OF {v2_snapshot_id}").collect() + assert len(legacy) == 2 + assert all(row["_row_id"] is None for row in legacy) + assert all(row["_last_updated_sequence_number"] is None for row in legacy) + + _with_gpu_lineage_write( + lambda gpu: gpu.range(3, 1503).coalesce(1).writeTo(full_table).append()) + + current = { + row.id: row for row in spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + assert len(current) == 1502 + assert current[3]["_pos"] == 0 + assert current[3]["_row_id"] == 0 + assert current[1502]["_row_id"] == 1499 + assert current[3]["_last_updated_sequence_number"] == 2 + assert current[1502]["_last_updated_sequence_number"] == 2 + assert {current[1]["_row_id"], current[2]["_row_id"]} == {1500, 1501} + assert current[1]["_last_updated_sequence_number"] == 1 + assert current[2]["_last_updated_sequence_number"] == 1 + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": reader_type + }) + + +@iceberg +@ignore_order(local=True) +@allow_non_gpu("BatchScanExec") +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +@pytest.mark.parametrize("reader_type", rapids_reader_types) +def test_iceberg_v3_row_lineage_update(spark_tmp_table_factory, reader_type): + full_table = get_full_table_name(spark_tmp_table_factory) + + def setup_iceberg_table(spark): + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3', 'write.update.mode' = 'copy-on-write')") + _with_gpu_lineage_write( + lambda gpu: gpu.range(0, 2).selectExpr("id", "CAST(0 AS BIGINT) AS v") + .coalesce(1).writeTo(full_table).append()) + + before = { + row.id: row for row in spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + assert before[1]["_row_id"] == 1 + assert before[1]["_last_updated_sequence_number"] == 1 + + _with_gpu_lineage_write( + lambda gpu: gpu.sql(f"UPDATE {full_table} SET v = v + 1 WHERE id = 1").collect()) + after = { + row.id: row for row in spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + data_sequence_numbers = [ + row.sequence_number for row in + spark.sql( + f"SELECT sequence_number FROM {full_table}.entries WHERE status != 2").collect() + ] + + assert after[1]["v"] == 1 + assert after[1]["_row_id"] == before[1]["_row_id"] == 1 + assert after[1]["_last_updated_sequence_number"] == 2 + assert after[0]["_row_id"] == before[0]["_row_id"] + assert after[0]["_last_updated_sequence_number"] == \ + before[0]["_last_updated_sequence_number"] == 1 + assert data_sequence_numbers == [2] + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, v, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": reader_type + }) + + +@iceberg +@ignore_order(local=True) +@allow_non_gpu("BatchScanExec", "DeleteFromTableExec") +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +@pytest.mark.parametrize("reader_type", rapids_reader_types) +def test_iceberg_v3_row_lineage_delete_leading_rows(spark_tmp_table_factory, reader_type): + full_table = get_full_table_name(spark_tmp_table_factory) + + def setup_iceberg_table(spark): + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3', 'write.delete.mode' = 'copy-on-write')") + _with_gpu_lineage_write( + lambda gpu: gpu.range(0, 1).writeTo(full_table).append()) + _with_gpu_lineage_write( + lambda gpu: gpu.sql(f"DELETE FROM {full_table} WHERE id = 0").collect()) + _with_gpu_lineage_write( + lambda gpu: gpu.range(1, 4).coalesce(1).writeTo(full_table).append()) + + before = { + row.id: row for row in spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + assert before[3]["_pos"] == 2 + assert before[3]["_row_id"] == 3 + assert before[3]["_last_updated_sequence_number"] == 3 + + _with_gpu_lineage_write( + lambda gpu: gpu.sql(f"DELETE FROM {full_table} WHERE id < 3").collect()) + after = spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") \ + .collect() + data_sequence_numbers = [ + row.sequence_number for row in + spark.sql( + f"SELECT sequence_number FROM {full_table}.entries WHERE status != 2").collect() + ] + + assert len(after) == 1 + assert after[0]["id"] == 3 + assert after[0]["_pos"] == 0 + assert after[0]["_row_id"] == 3 + assert after[0]["_last_updated_sequence_number"] == 3 + assert data_sequence_numbers == [4] + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": reader_type + }) + + +@iceberg +@ignore_order(local=True) +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +@pytest.mark.parametrize("reader_type", rapids_reader_types) +def test_iceberg_v3_row_lineage_merge_update_insert( + spark_tmp_table_factory, reader_type): + full_table = get_full_table_name(spark_tmp_table_factory) + source_view = spark_tmp_table_factory.get() + + def setup_iceberg_table(spark): + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3', 'write.merge.mode' = 'copy-on-write')") + _with_gpu_lineage_write( + lambda gpu: gpu.range(0, 3).selectExpr("id", "CAST(0 AS BIGINT) AS v") + .coalesce(1).writeTo(full_table).append()) + + def merge(gpu): + gpu.range(1, 4, 2).selectExpr("id", "id * 10 AS v") \ + .createOrReplaceTempView(source_view) + gpu.sql( + f"MERGE INTO {full_table} t USING {source_view} s ON t.id = s.id " + "WHEN MATCHED THEN UPDATE SET v = s.v " + "WHEN NOT MATCHED THEN INSERT (id, v) VALUES (s.id, s.v)").collect() + + _with_gpu_lineage_write(merge) + rows = { + row.id: row for row in spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + + assert rows[0]["_row_id"] == 0 + assert rows[0]["_last_updated_sequence_number"] == 1 + assert rows[1]["v"] == 10 + assert rows[1]["_row_id"] == 1 + assert rows[1]["_last_updated_sequence_number"] == 2 + assert rows[2]["_row_id"] == 2 + assert rows[2]["_last_updated_sequence_number"] == 1 + assert rows[3]["v"] == 30 + assert rows[3]["_row_id"] == 6 + assert rows[3]["_last_updated_sequence_number"] == 2 + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": reader_type + }) + + @iceberg @ignore_order(local=True) @pytest.mark.skipif( @@ -502,6 +738,101 @@ def setup_iceberg_table(spark): }) +@iceberg +@ignore_order(local=True) +@allow_non_gpu("BatchScanExec", "CallExec") +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +def test_iceberg_v3_row_lineage_gpu_rewrite_data_files(spark_tmp_table_factory): + full_table = get_full_table_name(spark_tmp_table_factory) + + def setup_iceberg_table(spark): + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3')") + _with_gpu_lineage_write( + lambda gpu: gpu.range(0, 2).writeTo(full_table).append()) + _with_gpu_lineage_write( + lambda gpu: gpu.range(2, 4).writeTo(full_table).append()) + before = { + row.id: row for row in spark.sql( + f"SELECT id, _file, _pos, _row_id, _last_updated_sequence_number " + f"FROM {full_table}").collect() + } + assert len({row["_file"] for row in before.values()}) > 1 + + _with_gpu_lineage_write( + lambda gpu: gpu.sql( + f"CALL spark_catalog.system.rewrite_data_files(table => '{full_table}', " + "options => map('min-input-files', '2'))").collect()) + after = { + row.id: row for row in spark.sql( + f"SELECT id, _file, _pos, _row_id, _last_updated_sequence_number " + f"FROM {full_table}").collect() + } + + assert len({row["_file"] for row in after.values()}) == 1 + assert any(after[row_id]["_file"] != before[row_id]["_file"] for row_id in before) + for row_id in before: + assert after[row_id]["_row_id"] == before[row_id]["_row_id"] + assert after[row_id]["_last_updated_sequence_number"] == \ + before[row_id]["_last_updated_sequence_number"] + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": "COALESCING" + }) + + +@iceberg +@ignore_order(local=True) +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +def test_iceberg_v3_row_lineage_insert_overwrite(spark_tmp_table_factory): + full_table = get_full_table_name(spark_tmp_table_factory) + source_view = spark_tmp_table_factory.get() + + def setup_iceberg_table(spark): + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3')") + _with_gpu_lineage_write( + lambda gpu: gpu.range(0, 2).selectExpr("id", "CAST(0 AS BIGINT) AS v") + .writeTo(full_table).append()) + + def overwrite(gpu): + gpu.range(10, 12).selectExpr("id", "CAST(1 AS BIGINT) AS v") \ + .createOrReplaceTempView(source_view) + gpu.sql(f"INSERT OVERWRITE {full_table} SELECT * FROM {source_view}").collect() + + _with_gpu_lineage_write(overwrite) + rows = { + row.id: row for row in spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + + assert set(rows) == {10, 11} + assert {rows[10]["_row_id"], rows[11]["_row_id"]} == {2, 3} + assert rows[10]["_last_updated_sequence_number"] == 2 + assert rows[11]["_last_updated_sequence_number"] == 2 + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": "COALESCING" + }) + + @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering # Compression codec to test and whether the codec is supported by cudf diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala index c636c601391..07b26d39f88 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala @@ -51,6 +51,7 @@ trait IcebergProbe { object IcebergProvider { val cpuBatchQueryScanClassName: String = "org.apache.iceberg.spark.source.SparkBatchQueryScan" val cpuCopyOnWriteScanClassName: String = "org.apache.iceberg.spark.source.SparkCopyOnWriteScan" + val cpuStagedScanClassName: String = "org.apache.iceberg.spark.source.SparkStagedScan" private lazy val probe: IcebergProbe = ShimLoaderTemp.newIcebergProbe() diff --git a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala index d1c4953e307..8ee9c002ca0 100644 --- a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala +++ b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala @@ -171,10 +171,18 @@ case class GpuMergeRowsExec( val boundMatchedBySourceInsts = GpuBindReferences.bindGpuReferences( notMatchedBySourceInstructions, child.output, allMetrics) .asInstanceOf[Seq[GpuInstruction]] + val instructionOutputs = (boundMatchedInsts ++ boundNotMatchedInsts ++ + boundMatchedBySourceInsts).flatMap(_.outputs) + val outputDataTypes = if (instructionOutputs.nonEmpty) { + instructionOutputs.maxBy(_.length).map(_.dataType).toArray + } else { + GpuColumnVector.extractTypes(schema) + } child.executeColumnar().mapPartitions { iter => new GpuMergeBatchIterator( dataTypes, + outputDataTypes, iter, boundTargetRowPresent, boundSourceRowPresent, @@ -325,8 +333,20 @@ object GpuMergeRowsExec { condition.columnarEval(batch) } - def applyOutputs(batch: ColumnarBatch): Seq[ColumnarBatch] = { - outputs.map(output => GpuProjectExec.project(batch, output)) + def applyOutputs( + batch: ColumnarBatch, + outputDataTypes: Array[DataType]): Seq[ColumnarBatch] = { + outputs.map { output => + require(output.length <= outputDataTypes.length, + s"Merge output has ${output.length} columns, expected at most ${outputDataTypes.length}") + // Spark permits merge actions to omit trailing columns from their InternalRow. Iceberg v3 + // uses this for unchanged rows while update and insert actions append row-lineage fields. + // cuDF tables require identical schemas for concatenation, so materialize the omitted + // trailing fields as correctly typed null columns. + val paddedOutput = output ++ outputDataTypes.drop(output.length) + .map(GpuLiteral(null, _)) + GpuProjectExec.project(batch, paddedOutput) + } } override def nullable: Boolean = false @@ -368,6 +388,7 @@ object GpuMergeRowsExec { * Similar to Spark's MergeRowIterator but operates on batches instead of rows. * * @param inputDataTypes Spark data types of input iterator. + * @param outputDataTypes Spark data types of the merge output. * @param inputIter Iterator of input columnar batches * @param isTargetRowPresent Bound GPU expression to check if target row is present * @param isSourceRowPresent Bound GPU expression to check if source row is present @@ -381,6 +402,7 @@ object GpuMergeRowsExec { */ class GpuMergeBatchIterator( inputDataTypes: Array[DataType], + outputDataTypes: Array[DataType], inputIter: Iterator[ColumnarBatch], isTargetRowPresent: GpuExpression, isSourceRowPresent: GpuExpression, @@ -518,7 +540,7 @@ class GpuMergeBatchIterator( if (writeSummaryEnabled) { attemptMetrics.record(instructionExec, filtered.numRows(), sourcePresent) } - outputs ++= instructionExec.applyOutputs(filtered) + outputs ++= instructionExec.applyOutputs(filtered, outputDataTypes) .map(SpillableColumnarBatch .apply(_, SpillPriorities.ACTIVE_ON_DECK_PRIORITY)) } @@ -529,4 +551,3 @@ class GpuMergeBatchIterator( sourcePresent, attemptMetrics) } } - From 396d97e2a4099b6197312dc3d98ed29fb8bc4402 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Wed, 2 Sep 2026 15:19:59 +0800 Subject: [PATCH 02/22] Verify GPU CTAS execution for Iceberg row lineage Signed-off-by: Chong Gao --- .../main/python/iceberg/iceberg_ctas_test.py | 37 +++++++++++++------ 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py index ad44b7bc323..e575d206cbd 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py @@ -19,7 +19,7 @@ from asserts import (assert_equal_with_local_sort, assert_gpu_and_cpu_are_equal_collect, assert_gpu_fallback_collect) -from conftest import is_iceberg_remote_catalog +from conftest import is_iceberg_remote_catalog, spark_jvm from data_gen import gen_df, copy_and_update, RepeatSeqGen from iceberg import (create_iceberg_table, iceberg_base_table_cols, @@ -30,7 +30,8 @@ ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -from marks import iceberg, ignore_order, allow_non_gpu, allow_non_gpu_conditional, datagen_overrides +from marks import (iceberg, ignore_order, allow_non_gpu, allow_non_gpu_conditional, + datagen_overrides) from spark_session import with_gpu_session, with_cpu_session, is_spark_400_or_later pytestmark = [ @@ -146,15 +147,29 @@ def test_ctas_v3_row_lineage(spark_tmp_table_factory): "spark.rapids.sql.format.iceberg.v3.enabled": "true" }) - with_gpu_session( - lambda spark: _execute_ctas( - spark, - table_name, - spark_tmp_table_factory, - lambda sp: sp.range(3), - {"format-version": "3"}, - ret=False), - conf=conf) + callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback + callback.startCapture() + try: + with_gpu_session( + lambda spark: _execute_ctas( + spark, + table_name, + spark_tmp_table_factory, + lambda sp: sp.range(3), + {"format-version": "3"}, + ret=False), + conf=conf) + captured_plans = callback.getResultsWithTimeout(10000) + assert any( + callback.contains(plan, "GpuAtomicCreateTableAsSelectExec") + for plan in captured_plans + ), "GpuAtomicCreateTableAsSelectExec is not found in the captured CTAS plans" + assert not any( + callback.didFallBack(plan, "AtomicCreateTableAsSelectExec") + for plan in captured_plans + ), "Captured CTAS plan contains CPU AtomicCreateTableAsSelectExec" + finally: + callback.endCapture() rows = with_cpu_session( lambda spark: spark.sql( From a730620453977b3829e4088a0dfeaa59f1a874b6 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Sat, 5 Sep 2026 09:40:53 +0800 Subject: [PATCH 03/22] Fix row lineage test coverage Update the Spark 4.1 retry suite for the merge output schema. Keep Iceberg DML tests in operation-specific suites and require staged scans to execute on GPU. Signed-off-by: Chong Gao --- .../python/iceberg/iceberg_delete_test.py | 66 ++++++ .../main/python/iceberg/iceberg_merge_test.py | 59 ++++++ .../src/main/python/iceberg/iceberg_test.py | 192 ++---------------- .../python/iceberg/iceberg_update_test.py | 63 ++++++ .../v2/GpuMergeBatchIteratorRetrySuite.scala | 5 +- 5 files changed, 208 insertions(+), 177 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py index ce680ba15c4..3bc5d6e5a6f 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py @@ -33,6 +33,9 @@ # Configuration for copy-on-write DELETE operations iceberg_delete_cow_enabled_conf = copy_and_update(iceberg_write_enabled_conf, {}) +iceberg_delete_v3_enabled_conf = copy_and_update( + iceberg_delete_cow_enabled_conf, + {"spark.rapids.sql.format.iceberg.v3.enabled": "true"}) # Configuration for merge-on-read DELETE operations iceberg_delete_mor_enabled_conf = copy_and_update(iceberg_write_enabled_conf, {}) @@ -200,6 +203,69 @@ def setup_iceberg_table(spark): }) +@iceberg +@ignore_order(local=True) +@allow_non_gpu("BatchScanExec", "DeleteFromTableExec") +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +@pytest.mark.parametrize("reader_type", rapids_reader_types) +def test_iceberg_v3_row_lineage_gpu_delete_leading_rows( + spark_tmp_table_factory, reader_type): + full_table = get_full_table_name(spark_tmp_table_factory) + + def setup_iceberg_table(spark): + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3', 'write.delete.mode' = 'copy-on-write')") + with_gpu_session( + lambda gpu: gpu.range(0, 1).writeTo(full_table).append(), + conf=iceberg_delete_v3_enabled_conf) + with_gpu_session( + lambda gpu: gpu.sql(f"DELETE FROM {full_table} WHERE id = 0").collect(), + conf=iceberg_delete_v3_enabled_conf) + with_gpu_session( + lambda gpu: gpu.range(1, 4).coalesce(1).writeTo(full_table).append(), + conf=iceberg_delete_v3_enabled_conf) + + before = { + row.id: row for row in spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + assert before[3]["_pos"] == 2 + assert before[3]["_row_id"] == 3 + assert before[3]["_last_updated_sequence_number"] == 3 + + with_gpu_session( + lambda gpu: gpu.sql(f"DELETE FROM {full_table} WHERE id < 3").collect(), + conf=iceberg_delete_v3_enabled_conf) + after = spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") \ + .collect() + data_sequence_numbers = [ + row.sequence_number for row in + spark.sql( + f"SELECT sequence_number FROM {full_table}.entries WHERE status != 2").collect() + ] + + assert len(after) == 1 + assert after[0]["id"] == 3 + assert after[0]["_pos"] == 0 + assert after[0]["_row_id"] == 3 + assert after[0]["_last_updated_sequence_number"] == 3 + assert data_sequence_numbers == [4] + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": reader_type + }) + + def _do_test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col_sql, delete_mode, table_properties=None): """Helper function for partitioned table DELETE tests.""" do_delete_test( diff --git a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py index 0d26fd341aa..88a936e8239 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py @@ -32,6 +32,9 @@ # Base configuration for Iceberg MERGE tests iceberg_merge_enabled_conf = copy_and_update(iceberg_write_enabled_conf, {}) +iceberg_merge_v3_enabled_conf = copy_and_update( + iceberg_merge_enabled_conf, + {"spark.rapids.sql.format.iceberg.v3.enabled": "true"}) def create_iceberg_table_with_merge_data( table_name: str, @@ -264,6 +267,62 @@ def setup_iceberg_table(spark): }) +@iceberg +@ignore_order(local=True) +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +@pytest.mark.parametrize("reader_type", rapids_reader_types) +def test_iceberg_v3_row_lineage_gpu_merge_update_insert( + spark_tmp_table_factory, reader_type): + full_table = get_full_table_name(spark_tmp_table_factory) + source_view = spark_tmp_table_factory.get() + + def setup_iceberg_table(spark): + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3', 'write.merge.mode' = 'copy-on-write')") + with_gpu_session( + lambda gpu: gpu.range(0, 3).selectExpr("id", "CAST(0 AS BIGINT) AS v") + .coalesce(1).writeTo(full_table).append(), + conf=iceberg_merge_v3_enabled_conf) + + def merge(gpu): + gpu.range(1, 4, 2).selectExpr("id", "id * 10 AS v") \ + .createOrReplaceTempView(source_view) + gpu.sql( + f"MERGE INTO {full_table} t USING {source_view} s ON t.id = s.id " + "WHEN MATCHED THEN UPDATE SET v = s.v " + "WHEN NOT MATCHED THEN INSERT (id, v) VALUES (s.id, s.v)").collect() + + with_gpu_session(merge, conf=iceberg_merge_v3_enabled_conf) + rows = { + row.id: row for row in spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + + assert rows[0]["_row_id"] == 0 + assert rows[0]["_last_updated_sequence_number"] == 1 + assert rows[1]["v"] == 10 + assert rows[1]["_row_id"] == 1 + assert rows[1]["_last_updated_sequence_number"] == 2 + assert rows[2]["_row_id"] == 2 + assert rows[2]["_last_updated_sequence_number"] == 1 + assert rows[3]["v"] == 30 + assert rows[3]["_row_id"] == 6 + assert rows[3]["_last_updated_sequence_number"] == 2 + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": reader_type + }) + + @allow_non_gpu("MergeRows$Keep", "MergeRows$Discard", "MergeRows$Split") @iceberg @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index 674bce2a064..bd9217f29e0 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -535,176 +535,6 @@ def setup_iceberg_table(spark): }) -@iceberg -@ignore_order(local=True) -@allow_non_gpu("BatchScanExec") -@pytest.mark.skipif( - not supports_iceberg_row_lineage_inheritance, - reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -@pytest.mark.parametrize("reader_type", rapids_reader_types) -def test_iceberg_v3_row_lineage_update(spark_tmp_table_factory, reader_type): - full_table = get_full_table_name(spark_tmp_table_factory) - - def setup_iceberg_table(spark): - spark.sql( - f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " - "TBLPROPERTIES ('format-version' = '3', 'write.update.mode' = 'copy-on-write')") - _with_gpu_lineage_write( - lambda gpu: gpu.range(0, 2).selectExpr("id", "CAST(0 AS BIGINT) AS v") - .coalesce(1).writeTo(full_table).append()) - - before = { - row.id: row for row in spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - assert before[1]["_row_id"] == 1 - assert before[1]["_last_updated_sequence_number"] == 1 - - _with_gpu_lineage_write( - lambda gpu: gpu.sql(f"UPDATE {full_table} SET v = v + 1 WHERE id = 1").collect()) - after = { - row.id: row for row in spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - data_sequence_numbers = [ - row.sequence_number for row in - spark.sql( - f"SELECT sequence_number FROM {full_table}.entries WHERE status != 2").collect() - ] - - assert after[1]["v"] == 1 - assert after[1]["_row_id"] == before[1]["_row_id"] == 1 - assert after[1]["_last_updated_sequence_number"] == 2 - assert after[0]["_row_id"] == before[0]["_row_id"] - assert after[0]["_last_updated_sequence_number"] == \ - before[0]["_last_updated_sequence_number"] == 1 - assert data_sequence_numbers == [2] - - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, v, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": reader_type - }) - - -@iceberg -@ignore_order(local=True) -@allow_non_gpu("BatchScanExec", "DeleteFromTableExec") -@pytest.mark.skipif( - not supports_iceberg_row_lineage_inheritance, - reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -@pytest.mark.parametrize("reader_type", rapids_reader_types) -def test_iceberg_v3_row_lineage_delete_leading_rows(spark_tmp_table_factory, reader_type): - full_table = get_full_table_name(spark_tmp_table_factory) - - def setup_iceberg_table(spark): - spark.sql( - f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG " - "TBLPROPERTIES ('format-version' = '3', 'write.delete.mode' = 'copy-on-write')") - _with_gpu_lineage_write( - lambda gpu: gpu.range(0, 1).writeTo(full_table).append()) - _with_gpu_lineage_write( - lambda gpu: gpu.sql(f"DELETE FROM {full_table} WHERE id = 0").collect()) - _with_gpu_lineage_write( - lambda gpu: gpu.range(1, 4).coalesce(1).writeTo(full_table).append()) - - before = { - row.id: row for row in spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - assert before[3]["_pos"] == 2 - assert before[3]["_row_id"] == 3 - assert before[3]["_last_updated_sequence_number"] == 3 - - _with_gpu_lineage_write( - lambda gpu: gpu.sql(f"DELETE FROM {full_table} WHERE id < 3").collect()) - after = spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") \ - .collect() - data_sequence_numbers = [ - row.sequence_number for row in - spark.sql( - f"SELECT sequence_number FROM {full_table}.entries WHERE status != 2").collect() - ] - - assert len(after) == 1 - assert after[0]["id"] == 3 - assert after[0]["_pos"] == 0 - assert after[0]["_row_id"] == 3 - assert after[0]["_last_updated_sequence_number"] == 3 - assert data_sequence_numbers == [4] - - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": reader_type - }) - - -@iceberg -@ignore_order(local=True) -@pytest.mark.skipif( - not supports_iceberg_row_lineage_inheritance, - reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -@pytest.mark.parametrize("reader_type", rapids_reader_types) -def test_iceberg_v3_row_lineage_merge_update_insert( - spark_tmp_table_factory, reader_type): - full_table = get_full_table_name(spark_tmp_table_factory) - source_view = spark_tmp_table_factory.get() - - def setup_iceberg_table(spark): - spark.sql( - f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " - "TBLPROPERTIES ('format-version' = '3', 'write.merge.mode' = 'copy-on-write')") - _with_gpu_lineage_write( - lambda gpu: gpu.range(0, 3).selectExpr("id", "CAST(0 AS BIGINT) AS v") - .coalesce(1).writeTo(full_table).append()) - - def merge(gpu): - gpu.range(1, 4, 2).selectExpr("id", "id * 10 AS v") \ - .createOrReplaceTempView(source_view) - gpu.sql( - f"MERGE INTO {full_table} t USING {source_view} s ON t.id = s.id " - "WHEN MATCHED THEN UPDATE SET v = s.v " - "WHEN NOT MATCHED THEN INSERT (id, v) VALUES (s.id, s.v)").collect() - - _with_gpu_lineage_write(merge) - rows = { - row.id: row for row in spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - - assert rows[0]["_row_id"] == 0 - assert rows[0]["_last_updated_sequence_number"] == 1 - assert rows[1]["v"] == 10 - assert rows[1]["_row_id"] == 1 - assert rows[1]["_last_updated_sequence_number"] == 2 - assert rows[2]["_row_id"] == 2 - assert rows[2]["_last_updated_sequence_number"] == 1 - assert rows[3]["v"] == 30 - assert rows[3]["_row_id"] == 6 - assert rows[3]["_last_updated_sequence_number"] == 2 - - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": reader_type - }) - - @iceberg @ignore_order(local=True) @pytest.mark.skipif( @@ -740,7 +570,7 @@ def setup_iceberg_table(spark): @iceberg @ignore_order(local=True) -@allow_non_gpu("BatchScanExec", "CallExec") +@allow_non_gpu("CallExec") @pytest.mark.skipif( not supports_iceberg_row_lineage_inheritance, reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) @@ -762,10 +592,22 @@ def setup_iceberg_table(spark): } assert len({row["_file"] for row in before.values()}) > 1 - _with_gpu_lineage_write( - lambda gpu: gpu.sql( - f"CALL spark_catalog.system.rewrite_data_files(table => '{full_table}', " - "options => map('min-input-files', '2'))").collect()) + callback = spark._sc._jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback + callback.startCapture() + try: + _with_gpu_lineage_write( + lambda gpu: gpu.sql( + f"CALL spark_catalog.system.rewrite_data_files(table => '{full_table}', " + "options => map('min-input-files', '2'))").collect()) + captured_plans = callback.getResultsWithTimeout(10000) + assert any( + callback.contains(plan, "GpuBatchScanExec") for plan in captured_plans + ), "GpuBatchScanExec is not found in the captured rewrite plans" + assert not any( + callback.didFallBack(plan, "BatchScanExec") for plan in captured_plans + ), "Captured rewrite plan contains CPU BatchScanExec" + finally: + callback.endCapture() after = { row.id: row for row in spark.sql( f"SELECT id, _file, _pos, _row_id, _last_updated_sequence_number " diff --git a/integration_tests/src/main/python/iceberg/iceberg_update_test.py b/integration_tests/src/main/python/iceberg/iceberg_update_test.py index c30f1384783..3dc7732e1c4 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_update_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_update_test.py @@ -32,6 +32,9 @@ # Configuration for copy-on-write UPDATE operations iceberg_update_cow_enabled_conf = copy_and_update(iceberg_write_enabled_conf, {}) +iceberg_update_v3_enabled_conf = copy_and_update( + iceberg_update_cow_enabled_conf, + {"spark.rapids.sql.format.iceberg.v3.enabled": "true"}) # Fixed seed for reproducible test data. Iceberg's update test plan will be different with different data and filter. UPDATE_TEST_SEED = 42 @@ -188,6 +191,66 @@ def setup_iceberg_table(spark): }) +@iceberg +@ignore_order(local=True) +@allow_non_gpu("BatchScanExec") +@pytest.mark.skipif( + not supports_iceberg_row_lineage_inheritance, + reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) +@pytest.mark.parametrize("reader_type", rapids_reader_types) +def test_iceberg_v3_row_lineage_gpu_update(spark_tmp_table_factory, reader_type): + full_table = get_full_table_name(spark_tmp_table_factory) + + def setup_iceberg_table(spark): + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3', 'write.update.mode' = 'copy-on-write')") + with_gpu_session( + lambda gpu: gpu.range(0, 2).selectExpr("id", "CAST(0 AS BIGINT) AS v") + .coalesce(1).writeTo(full_table).append(), + conf=iceberg_update_v3_enabled_conf) + + before = { + row.id: row for row in spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + assert before[1]["_row_id"] == 1 + assert before[1]["_last_updated_sequence_number"] == 1 + + with_gpu_session( + lambda gpu: gpu.sql( + f"UPDATE {full_table} SET v = v + 1 WHERE id = 1").collect(), + conf=iceberg_update_v3_enabled_conf) + after = { + row.id: row for row in spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") + .collect() + } + data_sequence_numbers = [ + row.sequence_number for row in + spark.sql( + f"SELECT sequence_number FROM {full_table}.entries WHERE status != 2").collect() + ] + + assert after[1]["v"] == 1 + assert after[1]["_row_id"] == before[1]["_row_id"] == 1 + assert after[1]["_last_updated_sequence_number"] == 2 + assert after[0]["_row_id"] == before[0]["_row_id"] + assert after[0]["_last_updated_sequence_number"] == \ + before[0]["_last_updated_sequence_number"] == 1 + assert data_sequence_numbers == [2] + + with_cpu_session(setup_iceberg_table) + assert_gpu_and_cpu_are_equal_collect( + lambda spark: spark.sql( + f"SELECT id, v, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), + conf={ + "spark.rapids.sql.format.iceberg.v3.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": reader_type + }) + + @iceberg @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) diff --git a/tests/src/test/spark411/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeBatchIteratorRetrySuite.scala b/tests/src/test/spark411/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeBatchIteratorRetrySuite.scala index 144c9253e21..ff3fa24c4dc 100644 --- a/tests/src/test/spark411/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeBatchIteratorRetrySuite.scala +++ b/tests/src/test/spark411/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeBatchIteratorRetrySuite.scala @@ -73,8 +73,9 @@ class GpuMergeBatchIteratorRetrySuite extends RmmSparkRetrySuiteBase { Seq(GpuBoundReference(0, IntegerType, nullable = true)(ExprId(0), "id")), ACTION_INSERT) val it = new GpuMergeBatchIterator( - Array(IntegerType), - Seq(buildBatch()).iterator, + inputDataTypes = Array(IntegerType), + outputDataTypes = Array(IntegerType), + inputIter = Seq(buildBatch()).iterator, isTargetRowPresent = GpuLiteral.create(false, BooleanType), isSourceRowPresent = GpuLiteral.create(true, BooleanType), matchedInstructionExecs = Nil, From 8b6b4515081aec3f374045b500680e1c5c41a407 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Mon, 7 Sep 2026 16:41:43 +0800 Subject: [PATCH 04/22] Simplify Iceberg row lineage CTAS test Signed-off-by: Chong Gao --- .../main/python/iceberg/iceberg_ctas_test.py | 80 ++++++++----------- 1 file changed, 34 insertions(+), 46 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py index e575d206cbd..859922e919f 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py @@ -17,7 +17,8 @@ import pytest from pyspark.sql.types import ArrayType, BinaryType -from asserts import (assert_equal_with_local_sort, assert_gpu_and_cpu_are_equal_collect, +from asserts import (assert_cpu_and_gpu_are_equal_collect_with_capture, + assert_equal_with_local_sort, assert_gpu_and_cpu_are_equal_collect, assert_gpu_fallback_collect) from conftest import is_iceberg_remote_catalog, spark_jvm from data_gen import gen_df, copy_and_update, RepeatSeqGen @@ -82,7 +83,9 @@ def _assert_gpu_equals_cpu_ctas(spark_tmp_table_factory, df_gen: Callable, table_prop: Dict[str, str], partition_col_sql: Optional[str] = None, - conf: Optional[Dict[str, str]] = None): + conf: Optional[Dict[str, str]] = None, + read_func: Optional[Callable] = None, + gpu_plan_assertion: Optional[Callable] = None): if conf is None: conf = iceberg_write_enabled_conf @@ -90,17 +93,23 @@ def _assert_gpu_equals_cpu_ctas(spark_tmp_table_factory, gpu_table = f"{base_name}_gpu" cpu_table = f"{base_name}_cpu" - def run_gpu_ctas(spark): - _execute_ctas(spark, gpu_table, spark_tmp_table_factory, - df_gen, table_prop, partition_col_sql, ret=True) + def run_ctas(spark): + gpu_enabled = str(spark.conf.get("spark.rapids.sql.enabled", "false")).lower() == "true" + target_table = gpu_table if gpu_enabled else cpu_table + return _execute_ctas( + spark, target_table, spark_tmp_table_factory, + df_gen, table_prop, partition_col_sql, ret=True) - with_gpu_session(run_gpu_ctas, conf=conf) - with_cpu_session(lambda spark: _execute_ctas(spark, cpu_table, spark_tmp_table_factory, - df_gen, table_prop, partition_col_sql, False), - conf=conf) + assert_cpu_and_gpu_are_equal_collect_with_capture( + run_ctas, + conf=conf, + gpu_plan_assertion=gpu_plan_assertion) - cpu_data = with_cpu_session(lambda spark: spark.table(cpu_table).collect()) - gpu_data = with_cpu_session(lambda spark: spark.table(gpu_table).collect()) + def read_table(spark, table_name): + return spark.table(table_name) if read_func is None else read_func(spark, table_name) + + cpu_data = with_cpu_session(lambda spark: read_table(spark, cpu_table).collect(), conf=conf) + gpu_data = with_cpu_session(lambda spark: read_table(spark, gpu_table).collect(), conf=conf) assert_equal_with_local_sort(cpu_data, gpu_data) @@ -142,45 +151,24 @@ def run_ctas(spark): reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) @ignore_order(local=True) def test_ctas_v3_row_lineage(spark_tmp_table_factory): - table_name = get_full_table_name(spark_tmp_table_factory) conf = copy_and_update(iceberg_write_enabled_conf, { "spark.rapids.sql.format.iceberg.v3.enabled": "true" }) - callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback - callback.startCapture() - try: - with_gpu_session( - lambda spark: _execute_ctas( - spark, - table_name, - spark_tmp_table_factory, - lambda sp: sp.range(3), - {"format-version": "3"}, - ret=False), - conf=conf) - captured_plans = callback.getResultsWithTimeout(10000) - assert any( - callback.contains(plan, "GpuAtomicCreateTableAsSelectExec") - for plan in captured_plans - ), "GpuAtomicCreateTableAsSelectExec is not found in the captured CTAS plans" - assert not any( - callback.didFallBack(plan, "AtomicCreateTableAsSelectExec") - for plan in captured_plans - ), "Captured CTAS plan contains CPU AtomicCreateTableAsSelectExec" - finally: - callback.endCapture() - - rows = with_cpu_session( - lambda spark: spark.sql( - f"SELECT id, _row_id, _last_updated_sequence_number FROM {table_name} " - "ORDER BY id").collect()) - assert [(row["id"], row["_row_id"], row["_last_updated_sequence_number"]) - for row in rows] == [ - (0, 0, 1), - (1, 1, 1), - (2, 2, 1) - ] + def assert_gpu_ctas(plan): + callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback + ctas_plan = callback.extractExecutedPlan(plan) + callback.assertContains(ctas_plan, "GpuAtomicCreateTableAsSelectExec") + callback.assertNotContain(ctas_plan, "AtomicCreateTableAsSelectExec") + + _assert_gpu_equals_cpu_ctas( + spark_tmp_table_factory, + lambda spark: spark.range(3), + {"format-version": "3"}, + conf=conf, + read_func=lambda spark, table: spark.sql( + f"SELECT id, _row_id, _last_updated_sequence_number FROM {table}"), + gpu_plan_assertion=assert_gpu_ctas) @iceberg From 25794adb7a97a43a6da0890adbd7bec74e31ff4e Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Mon, 7 Sep 2026 17:11:38 +0800 Subject: [PATCH 05/22] Simplify Iceberg row lineage delete test Reuse the DELETE comparison helper with a custom metadata read so CPU and GPU results validate row positions and lineage fields. Signed-off-by: Chong Gao --- .../python/iceberg/iceberg_delete_test.py | 83 +++++-------------- 1 file changed, 22 insertions(+), 61 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py index 3bc5d6e5a6f..ed7f88bf097 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py @@ -78,7 +78,8 @@ def insert_data(spark): def do_delete_test(spark_tmp_table_factory, delete_sql_func, data_gen_func=None, partition_col_sql=None, table_properties=None, - delete_mode='copy-on-write'): + delete_mode='copy-on-write', conf=iceberg_delete_cow_enabled_conf, + read_func=None): """ Helper function to test DELETE operations by comparing CPU and GPU results. @@ -89,6 +90,8 @@ def do_delete_test(spark_tmp_table_factory, delete_sql_func, data_gen_func=None, partition_col_sql: SQL for partitioning clause table_properties: Additional table properties delete_mode: 'copy-on-write' or 'merge-on-read' + conf: Spark configuration used for DELETE and result reads + read_func: Optional function that takes (spark, table_name) and returns a DataFrame """ base_table_name = get_full_table_name(spark_tmp_table_factory) cpu_table_name = f"{base_table_name}_cpu" @@ -104,17 +107,21 @@ def do_delete_test(spark_tmp_table_factory, delete_sql_func, data_gen_func=None, def do_gpu_delete(spark): delete_sql_func(spark, gpu_table_name) - with_gpu_session(do_gpu_delete, conf=iceberg_delete_cow_enabled_conf) + with_gpu_session(do_gpu_delete, conf=conf) # Execute DELETE on CPU def do_cpu_delete(spark): delete_sql_func(spark, cpu_table_name) - with_cpu_session(do_cpu_delete) + with_cpu_session(do_cpu_delete, conf=conf) # Compare results - cpu_data = with_cpu_session(lambda spark: spark.table(cpu_table_name).collect()) - gpu_data = with_cpu_session(lambda spark: spark.table(gpu_table_name).collect()) + if read_func is None: + read_func = lambda spark, table_name: spark.table(table_name) + cpu_data = with_cpu_session( + lambda spark: read_func(spark, cpu_table_name).collect(), conf=conf) + gpu_data = with_cpu_session( + lambda spark: read_func(spark, gpu_table_name).collect(), conf=conf) assert_equal_with_local_sort(cpu_data, gpu_data) @@ -205,65 +212,19 @@ def setup_iceberg_table(spark): @iceberg @ignore_order(local=True) -@allow_non_gpu("BatchScanExec", "DeleteFromTableExec") +@allow_non_gpu("BatchScanExec") @pytest.mark.skipif( not supports_iceberg_row_lineage_inheritance, reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -@pytest.mark.parametrize("reader_type", rapids_reader_types) -def test_iceberg_v3_row_lineage_gpu_delete_leading_rows( - spark_tmp_table_factory, reader_type): - full_table = get_full_table_name(spark_tmp_table_factory) - - def setup_iceberg_table(spark): - spark.sql( - f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG " - "TBLPROPERTIES ('format-version' = '3', 'write.delete.mode' = 'copy-on-write')") - with_gpu_session( - lambda gpu: gpu.range(0, 1).writeTo(full_table).append(), - conf=iceberg_delete_v3_enabled_conf) - with_gpu_session( - lambda gpu: gpu.sql(f"DELETE FROM {full_table} WHERE id = 0").collect(), - conf=iceberg_delete_v3_enabled_conf) - with_gpu_session( - lambda gpu: gpu.range(1, 4).coalesce(1).writeTo(full_table).append(), - conf=iceberg_delete_v3_enabled_conf) - - before = { - row.id: row for row in spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - assert before[3]["_pos"] == 2 - assert before[3]["_row_id"] == 3 - assert before[3]["_last_updated_sequence_number"] == 3 - - with_gpu_session( - lambda gpu: gpu.sql(f"DELETE FROM {full_table} WHERE id < 3").collect(), - conf=iceberg_delete_v3_enabled_conf) - after = spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") \ - .collect() - data_sequence_numbers = [ - row.sequence_number for row in - spark.sql( - f"SELECT sequence_number FROM {full_table}.entries WHERE status != 2").collect() - ] - - assert len(after) == 1 - assert after[0]["id"] == 3 - assert after[0]["_pos"] == 0 - assert after[0]["_row_id"] == 3 - assert after[0]["_last_updated_sequence_number"] == 3 - assert data_sequence_numbers == [4] - - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": reader_type - }) +def test_iceberg_v3_row_lineage_gpu_delete_leading_rows(spark_tmp_table_factory): + do_delete_test( + spark_tmp_table_factory, + lambda spark, table: spark.sql(f"DELETE FROM {table} WHERE id < 3"), + data_gen_func=lambda spark: spark.range(1, 4).coalesce(1), + table_properties={"format-version": "3"}, + conf=iceberg_delete_v3_enabled_conf, + read_func=lambda spark, table: spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {table}")) def _do_test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col_sql, delete_mode, table_properties=None): From 4a8b7d8b156d60a18ec68ec37e80fef6a49fc6e0 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Mon, 7 Sep 2026 17:23:51 +0800 Subject: [PATCH 06/22] Simplify Iceberg row lineage merge test Compare CPU and GPU MERGE results through a shared helper instead of asserting literal row-lineage values. Signed-off-by: Chong Gao --- .../main/python/iceberg/iceberg_merge_test.py | 111 ++++++++---------- 1 file changed, 49 insertions(+), 62 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py index 88a936e8239..bc6fa462806 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py @@ -98,6 +98,21 @@ def insert_data(spark): with_cpu_session(insert_data) +def _assert_gpu_and_cpu_merge_writes_are_equal( + cpu_table_name, gpu_table_name, merge_func, read_func, conf): + def run_merge(spark, table_name): + merge_func(spark, table_name) + + with_cpu_session(lambda spark: run_merge(spark, cpu_table_name), conf=conf) + with_gpu_session(lambda spark: run_merge(spark, gpu_table_name), conf=conf) + + cpu_data = with_cpu_session( + lambda spark: read_func(spark, cpu_table_name).collect(), conf=conf) + gpu_data = with_cpu_session( + lambda spark: read_func(spark, gpu_table_name).collect(), conf=conf) + assert_equal_with_local_sort(cpu_data, gpu_data) + + def do_merge_test( spark_tmp_table_factory, merge_sql_func, @@ -134,22 +149,12 @@ def do_merge_test( create_iceberg_table_with_merge_data(source_table, partition_col_sql, table_properties, ensure_distinct_key=True, seed=42, merge_mode=merge_mode, iceberg_base_table_cols=iceberg_base_table_cols, iceberg_gens_list=iceberg_gens_list) - # Execute MERGE on GPU - def do_gpu_merge(spark): - merge_sql_func(spark, gpu_target_table, source_table) - - with_gpu_session(do_gpu_merge, conf=iceberg_merge_enabled_conf) - - # Execute MERGE on CPU - def do_cpu_merge(spark): - merge_sql_func(spark, cpu_target_table, source_table) - - with_cpu_session(do_cpu_merge) - - # Compare results - cpu_data = with_cpu_session(lambda spark: spark.table(cpu_target_table).collect()) - gpu_data = with_cpu_session(lambda spark: spark.table(gpu_target_table).collect()) - assert_equal_with_local_sort(cpu_data, gpu_data) + _assert_gpu_and_cpu_merge_writes_are_equal( + cpu_target_table, + gpu_target_table, + lambda spark, target: merge_sql_func(spark, target, source_table), + lambda spark, target: spark.table(target), + iceberg_merge_enabled_conf) def _do_test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mode, table_properties=None): @@ -272,55 +277,37 @@ def setup_iceberg_table(spark): @pytest.mark.skipif( not supports_iceberg_row_lineage_inheritance, reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -@pytest.mark.parametrize("reader_type", rapids_reader_types) -def test_iceberg_v3_row_lineage_gpu_merge_update_insert( - spark_tmp_table_factory, reader_type): - full_table = get_full_table_name(spark_tmp_table_factory) +def test_iceberg_v3_row_lineage_gpu_merge_update_insert(spark_tmp_table_factory): + base_table = get_full_table_name(spark_tmp_table_factory) + cpu_table = f"{base_table}_cpu" + gpu_table = f"{base_table}_gpu" source_view = spark_tmp_table_factory.get() - def setup_iceberg_table(spark): + def setup_iceberg_tables(spark): + for table in [cpu_table, gpu_table]: + spark.sql( + f"CREATE TABLE {table} (id BIGINT, v BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3', " + "'write.merge.mode' = 'copy-on-write')") + spark.range(0, 3).selectExpr("id", "CAST(0 AS BIGINT) AS v") \ + .coalesce(1).writeTo(table).append() + + def merge(spark, table): + spark.range(1, 4, 2).selectExpr("id", "id * 10 AS v") \ + .createOrReplaceTempView(source_view) spark.sql( - f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " - "TBLPROPERTIES ('format-version' = '3', 'write.merge.mode' = 'copy-on-write')") - with_gpu_session( - lambda gpu: gpu.range(0, 3).selectExpr("id", "CAST(0 AS BIGINT) AS v") - .coalesce(1).writeTo(full_table).append(), - conf=iceberg_merge_v3_enabled_conf) - - def merge(gpu): - gpu.range(1, 4, 2).selectExpr("id", "id * 10 AS v") \ - .createOrReplaceTempView(source_view) - gpu.sql( - f"MERGE INTO {full_table} t USING {source_view} s ON t.id = s.id " - "WHEN MATCHED THEN UPDATE SET v = s.v " - "WHEN NOT MATCHED THEN INSERT (id, v) VALUES (s.id, s.v)").collect() - - with_gpu_session(merge, conf=iceberg_merge_v3_enabled_conf) - rows = { - row.id: row for row in spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - - assert rows[0]["_row_id"] == 0 - assert rows[0]["_last_updated_sequence_number"] == 1 - assert rows[1]["v"] == 10 - assert rows[1]["_row_id"] == 1 - assert rows[1]["_last_updated_sequence_number"] == 2 - assert rows[2]["_row_id"] == 2 - assert rows[2]["_last_updated_sequence_number"] == 1 - assert rows[3]["v"] == 30 - assert rows[3]["_row_id"] == 6 - assert rows[3]["_last_updated_sequence_number"] == 2 - - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": reader_type - }) + f"MERGE INTO {table} t USING {source_view} s ON t.id = s.id " + "WHEN MATCHED THEN UPDATE SET v = s.v " + "WHEN NOT MATCHED THEN INSERT (id, v) VALUES (s.id, s.v)").collect() + + with_cpu_session(setup_iceberg_tables) + _assert_gpu_and_cpu_merge_writes_are_equal( + cpu_table, + gpu_table, + merge, + lambda spark, table: spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {table}"), + iceberg_merge_v3_enabled_conf) @allow_non_gpu("MergeRows$Keep", "MergeRows$Discard", "MergeRows$Split") From 890a9f1a09f4208cb6a3b5457769f45aa41cf3df Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Mon, 7 Sep 2026 17:31:51 +0800 Subject: [PATCH 07/22] Simplify Iceberg row lineage write tests Compare CPU and GPU append, rewrite, and overwrite results through a shared helper instead of asserting literal metadata values. Signed-off-by: Chong Gao --- .../src/main/python/iceberg/iceberg_test.py | 196 ++++++------------ 1 file changed, 68 insertions(+), 128 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index bd9217f29e0..3d3283566b6 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -56,8 +56,28 @@ } -def _with_gpu_lineage_write(write_func): - return with_gpu_session(write_func, conf=_ROW_LINEAGE_WRITE_CONF) +def _assert_gpu_and_cpu_lineage_writes_are_equal( + spark_tmp_table_factory, setup_func, write_func, read_func): + base_table = get_full_table_name(spark_tmp_table_factory) + cpu_table = f"{base_table}_cpu" + gpu_table = f"{base_table}_gpu" + + def setup_tables(spark): + setup_func(spark, cpu_table) + setup_func(spark, gpu_table) + + def run_write(spark, table): + write_func(spark, table) + + with_cpu_session(setup_tables) + with_cpu_session(lambda spark: run_write(spark, cpu_table), conf=_ROW_LINEAGE_WRITE_CONF) + with_gpu_session(lambda spark: run_write(spark, gpu_table), conf=_ROW_LINEAGE_WRITE_CONF) + + cpu_data = with_cpu_session( + lambda spark: read_func(spark, cpu_table).collect(), conf=_ROW_LINEAGE_WRITE_CONF) + gpu_data = with_cpu_session( + lambda spark: read_func(spark, gpu_table).collect(), conf=_ROW_LINEAGE_WRITE_CONF) + assert_equal_with_local_sort(cpu_data, gpu_data) pytestmark = iceberg_unsupported_mark @@ -482,57 +502,24 @@ def setup_iceberg_table(spark): @pytest.mark.skipif( not supports_iceberg_row_lineage_inheritance, reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -@pytest.mark.parametrize("reader_type", rapids_reader_types) -def test_iceberg_v3_row_lineage_append(spark_tmp_table_factory, reader_type): - full_table = get_full_table_name(spark_tmp_table_factory) - - def setup_iceberg_table(spark): - spark.sql(f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG " +def test_iceberg_v3_row_lineage_append(spark_tmp_table_factory): + def setup_iceberg_table(spark, table): + spark.sql(f"CREATE TABLE {table} (id BIGINT) USING ICEBERG " f"TBLPROPERTIES ('format-version' = '2')") - spark.sql(f"INSERT INTO {full_table} VALUES (1), (2)") - v2_snapshot_id = spark.sql( - f"SELECT snapshot_id FROM {full_table}.snapshots ORDER BY committed_at DESC") \ - .head()[0] + spark.sql(f"INSERT INTO {table} VALUES (1), (2)") spark.sql( - f"ALTER TABLE {full_table} SET TBLPROPERTIES (" + f"ALTER TABLE {table} SET TBLPROPERTIES (" "'format-version' = '3', " "'write.parquet.row-group-size-bytes' = '4096', " "'read.split.target-size' = '4096', " "'read.split.open-file-cost' = '0')") - legacy = spark.sql( - f"SELECT id, _row_id, _last_updated_sequence_number FROM {full_table} " - f"VERSION AS OF {v2_snapshot_id}").collect() - assert len(legacy) == 2 - assert all(row["_row_id"] is None for row in legacy) - assert all(row["_last_updated_sequence_number"] is None for row in legacy) - - _with_gpu_lineage_write( - lambda gpu: gpu.range(3, 1503).coalesce(1).writeTo(full_table).append()) - - current = { - row.id: row for row in spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - assert len(current) == 1502 - assert current[3]["_pos"] == 0 - assert current[3]["_row_id"] == 0 - assert current[1502]["_row_id"] == 1499 - assert current[3]["_last_updated_sequence_number"] == 2 - assert current[1502]["_last_updated_sequence_number"] == 2 - assert {current[1]["_row_id"], current[2]["_row_id"]} == {1500, 1501} - assert current[1]["_last_updated_sequence_number"] == 1 - assert current[2]["_last_updated_sequence_number"] == 1 - - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": reader_type - }) + _assert_gpu_and_cpu_lineage_writes_are_equal( + spark_tmp_table_factory, + setup_iceberg_table, + lambda spark, table: spark.range(3, 1503).coalesce(1).writeTo(table).append(), + lambda spark, table: spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {table}")) @iceberg @@ -575,60 +562,29 @@ def setup_iceberg_table(spark): not supports_iceberg_row_lineage_inheritance, reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) def test_iceberg_v3_row_lineage_gpu_rewrite_data_files(spark_tmp_table_factory): - full_table = get_full_table_name(spark_tmp_table_factory) - - def setup_iceberg_table(spark): + def setup_iceberg_table(spark, table): spark.sql( - f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG " + f"CREATE TABLE {table} (id BIGINT) USING ICEBERG " "TBLPROPERTIES ('format-version' = '3')") - _with_gpu_lineage_write( - lambda gpu: gpu.range(0, 2).writeTo(full_table).append()) - _with_gpu_lineage_write( - lambda gpu: gpu.range(2, 4).writeTo(full_table).append()) - before = { - row.id: row for row in spark.sql( - f"SELECT id, _file, _pos, _row_id, _last_updated_sequence_number " - f"FROM {full_table}").collect() - } - assert len({row["_file"] for row in before.values()}) > 1 + spark.range(0, 2).writeTo(table).append() + spark.range(2, 4).writeTo(table).append() - callback = spark._sc._jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback - callback.startCapture() - try: - _with_gpu_lineage_write( - lambda gpu: gpu.sql( - f"CALL spark_catalog.system.rewrite_data_files(table => '{full_table}', " - "options => map('min-input-files', '2'))").collect()) - captured_plans = callback.getResultsWithTimeout(10000) - assert any( - callback.contains(plan, "GpuBatchScanExec") for plan in captured_plans - ), "GpuBatchScanExec is not found in the captured rewrite plans" - assert not any( - callback.didFallBack(plan, "BatchScanExec") for plan in captured_plans - ), "Captured rewrite plan contains CPU BatchScanExec" - finally: - callback.endCapture() - after = { - row.id: row for row in spark.sql( - f"SELECT id, _file, _pos, _row_id, _last_updated_sequence_number " - f"FROM {full_table}").collect() - } - - assert len({row["_file"] for row in after.values()}) == 1 - assert any(after[row_id]["_file"] != before[row_id]["_file"] for row_id in before) - for row_id in before: - assert after[row_id]["_row_id"] == before[row_id]["_row_id"] - assert after[row_id]["_last_updated_sequence_number"] == \ - before[row_id]["_last_updated_sequence_number"] + def rewrite_data_files(spark, table): + spark.sql( + f"CALL spark_catalog.system.rewrite_data_files(table => '{table}', " + "options => map('min-input-files', '2'))").collect() - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": "COALESCING" - }) + def read_data_and_file_count(spark, table): + return spark.sql( + f"SELECT d.id, d._row_id, d._last_updated_sequence_number, f.data_file_count " + f"FROM {table} d CROSS JOIN (" + f"SELECT count(*) AS data_file_count FROM {table}.data_files) f") + + _assert_gpu_and_cpu_lineage_writes_are_equal( + spark_tmp_table_factory, + setup_iceberg_table, + rewrite_data_files, + read_data_and_file_count) @iceberg @@ -637,42 +593,26 @@ def setup_iceberg_table(spark): not supports_iceberg_row_lineage_inheritance, reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) def test_iceberg_v3_row_lineage_insert_overwrite(spark_tmp_table_factory): - full_table = get_full_table_name(spark_tmp_table_factory) source_view = spark_tmp_table_factory.get() - def setup_iceberg_table(spark): + def setup_iceberg_table(spark, table): spark.sql( - f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " + f"CREATE TABLE {table} (id BIGINT, v BIGINT) USING ICEBERG " "TBLPROPERTIES ('format-version' = '3')") - _with_gpu_lineage_write( - lambda gpu: gpu.range(0, 2).selectExpr("id", "CAST(0 AS BIGINT) AS v") - .writeTo(full_table).append()) - - def overwrite(gpu): - gpu.range(10, 12).selectExpr("id", "CAST(1 AS BIGINT) AS v") \ - .createOrReplaceTempView(source_view) - gpu.sql(f"INSERT OVERWRITE {full_table} SELECT * FROM {source_view}").collect() - - _with_gpu_lineage_write(overwrite) - rows = { - row.id: row for row in spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - - assert set(rows) == {10, 11} - assert {rows[10]["_row_id"], rows[11]["_row_id"]} == {2, 3} - assert rows[10]["_last_updated_sequence_number"] == 2 - assert rows[11]["_last_updated_sequence_number"] == 2 - - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": "COALESCING" - }) + spark.range(0, 2).selectExpr("id", "CAST(0 AS BIGINT) AS v") \ + .writeTo(table).append() + + def overwrite(spark, table): + spark.range(10, 12).selectExpr("id", "CAST(1 AS BIGINT) AS v") \ + .createOrReplaceTempView(source_view) + spark.sql(f"INSERT OVERWRITE {table} SELECT * FROM {source_view}").collect() + + _assert_gpu_and_cpu_lineage_writes_are_equal( + spark_tmp_table_factory, + setup_iceberg_table, + overwrite, + lambda spark, table: spark.sql( + f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {table}")) @iceberg From 0f0151d5b21a66218afbfd6a656bcd3eb1e33251 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Mon, 7 Sep 2026 17:40:05 +0800 Subject: [PATCH 08/22] Simplify Iceberg row lineage update test Reuse the UPDATE comparison helper with a custom lineage read instead of asserting literal metadata values. Signed-off-by: Chong Gao --- .../python/iceberg/iceberg_update_test.py | 79 ++++++------------- 1 file changed, 22 insertions(+), 57 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_update_test.py b/integration_tests/src/main/python/iceberg/iceberg_update_test.py index 3dc7732e1c4..ef48536443f 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_update_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_update_test.py @@ -79,7 +79,8 @@ def insert_data(spark): def do_update_test(spark_tmp_table_factory, update_sql_func, data_gen_func=None, partition_col_sql=None, table_properties=None, - update_mode='copy-on-write'): + update_mode='copy-on-write', conf=iceberg_update_cow_enabled_conf, + read_func=None): """ Helper function to test UPDATE operations by comparing CPU and GPU results. @@ -90,6 +91,8 @@ def do_update_test(spark_tmp_table_factory, update_sql_func, data_gen_func=None, partition_col_sql: SQL for partitioning clause table_properties: Additional table properties update_mode: Update mode - 'copy-on-write' or 'merge-on-read' + conf: Spark configuration used for UPDATE and result reads + read_func: Optional function that takes (spark, table_name) and returns a DataFrame """ base_table_name = get_full_table_name(spark_tmp_table_factory) cpu_table_name = f"{base_table_name}_cpu" @@ -105,17 +108,21 @@ def do_update_test(spark_tmp_table_factory, update_sql_func, data_gen_func=None, def do_gpu_update(spark): update_sql_func(spark, gpu_table_name) - with_gpu_session(do_gpu_update, conf=iceberg_update_cow_enabled_conf) + with_gpu_session(do_gpu_update, conf=conf) # Execute UPDATE on CPU def do_cpu_update(spark): update_sql_func(spark, cpu_table_name) - with_cpu_session(do_cpu_update) + with_cpu_session(do_cpu_update, conf=conf) # Compare results - cpu_data = with_cpu_session(lambda spark: spark.table(cpu_table_name).collect()) - gpu_data = with_cpu_session(lambda spark: spark.table(gpu_table_name).collect()) + if read_func is None: + read_func = lambda spark, table_name: spark.table(table_name) + cpu_data = with_cpu_session( + lambda spark: read_func(spark, cpu_table_name).collect(), conf=conf) + gpu_data = with_cpu_session( + lambda spark: read_func(spark, gpu_table_name).collect(), conf=conf) assert_equal_with_local_sort(cpu_data, gpu_data) @@ -197,58 +204,16 @@ def setup_iceberg_table(spark): @pytest.mark.skipif( not supports_iceberg_row_lineage_inheritance, reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -@pytest.mark.parametrize("reader_type", rapids_reader_types) -def test_iceberg_v3_row_lineage_gpu_update(spark_tmp_table_factory, reader_type): - full_table = get_full_table_name(spark_tmp_table_factory) - - def setup_iceberg_table(spark): - spark.sql( - f"CREATE TABLE {full_table} (id BIGINT, v BIGINT) USING ICEBERG " - "TBLPROPERTIES ('format-version' = '3', 'write.update.mode' = 'copy-on-write')") - with_gpu_session( - lambda gpu: gpu.range(0, 2).selectExpr("id", "CAST(0 AS BIGINT) AS v") - .coalesce(1).writeTo(full_table).append(), - conf=iceberg_update_v3_enabled_conf) - - before = { - row.id: row for row in spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - assert before[1]["_row_id"] == 1 - assert before[1]["_last_updated_sequence_number"] == 1 - - with_gpu_session( - lambda gpu: gpu.sql( - f"UPDATE {full_table} SET v = v + 1 WHERE id = 1").collect(), - conf=iceberg_update_v3_enabled_conf) - after = { - row.id: row for row in spark.sql( - f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {full_table}") - .collect() - } - data_sequence_numbers = [ - row.sequence_number for row in - spark.sql( - f"SELECT sequence_number FROM {full_table}.entries WHERE status != 2").collect() - ] - - assert after[1]["v"] == 1 - assert after[1]["_row_id"] == before[1]["_row_id"] == 1 - assert after[1]["_last_updated_sequence_number"] == 2 - assert after[0]["_row_id"] == before[0]["_row_id"] - assert after[0]["_last_updated_sequence_number"] == \ - before[0]["_last_updated_sequence_number"] == 1 - assert data_sequence_numbers == [2] - - with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect( - lambda spark: spark.sql( - f"SELECT id, v, _pos, _row_id, _last_updated_sequence_number FROM {full_table}"), - conf={ - "spark.rapids.sql.format.iceberg.v3.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": reader_type - }) +def test_iceberg_v3_row_lineage_gpu_update(spark_tmp_table_factory): + do_update_test( + spark_tmp_table_factory, + lambda spark, table: spark.sql(f"UPDATE {table} SET v = v + 1 WHERE id = 1"), + data_gen_func=lambda spark: spark.range(0, 2).selectExpr( + "id", "CAST(0 AS BIGINT) AS v").coalesce(1), + table_properties={"format-version": "3"}, + conf=iceberg_update_v3_enabled_conf, + read_func=lambda spark, table: spark.sql( + f"SELECT id, v, _pos, _row_id, _last_updated_sequence_number FROM {table}")) @iceberg From 683c7eef5ab1674fead0c9efd48af85e49ad1e80 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Tue, 8 Sep 2026 13:53:05 +0800 Subject: [PATCH 09/22] Remove Iceberg rewrite data files support Signed-off-by: Chong Gao --- .../spark/source/GpuSparkScanAccess.java | 7 +-- .../rapids/iceberg/IcebergProviderBase.scala | 16 ------- .../iceberg/spark/source/GpuSparkScan.scala | 6 +-- .../spark/source/GpuSparkStagedScan.scala | 47 ------------------- .../iceberg/spark/source/GpuSparkWrite.scala | 2 - .../apache/iceberg/spark/source/write.scala | 8 ---- .../src/main/python/iceberg/iceberg_test.py | 32 ------------- .../rapids/iceberg/IcebergProvider.scala | 2 - 8 files changed, 3 insertions(+), 117 deletions(-) delete mode 100644 iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkStagedScan.scala diff --git a/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkScanAccess.java b/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkScanAccess.java index 9719c783cf8..26d6e46ccde 100644 --- a/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkScanAccess.java +++ b/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkScanAccess.java @@ -47,8 +47,7 @@ private GpuSparkScanAccess() { } public static boolean supports(Scan scan) { - return scan instanceof SparkBatchQueryScan || scan instanceof SparkCopyOnWriteScan - || scan instanceof SparkStagedScan; + return scan instanceof SparkBatchQueryScan || scan instanceof SparkCopyOnWriteScan; } public static boolean isBatchQueryScan(Scan scan) { @@ -59,10 +58,6 @@ public static boolean isCopyOnWriteScan(Scan scan) { return scan instanceof SparkCopyOnWriteScan; } - public static boolean isStagedScan(Scan scan) { - return scan instanceof SparkStagedScan; - } - public static boolean isMetadataScan(Scan scan) { return sparkScan(scan).table() instanceof BaseMetadataTable; } diff --git a/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala b/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala index fd3de6c2a45..b7f2296e564 100644 --- a/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala +++ b/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala @@ -43,9 +43,6 @@ abstract class IcebergProviderBase extends IcebergProvider { IcebergProvider.cpuBatchQueryScanClassName) val cpuCopyOnWriteScanClass = ShimReflectionUtils.loadClass( IcebergProvider.cpuCopyOnWriteScanClassName) - val cpuStagedScanClass = ShimReflectionUtils.loadClass( - IcebergProvider.cpuStagedScanClassName) - Seq( new ScanRule[Scan]( (a, conf, p, r) => new ScanMeta[Scan](a, conf, p, r) { @@ -77,19 +74,6 @@ abstract class IcebergProviderBase extends IcebergProvider { "Iceberg copy on write scan", ClassTag(cpuCopyOnWriteScanClass) ), - new ScanRule[Scan]( - (a, conf, p, r) => new ScanMeta[Scan](a, conf, p, r) { - private lazy val convertedScan: Try[GpuSparkScan] = GpuSparkScan.tryConvert(a, this.conf) - - override def tagSelfForGpu(): Unit = { - GpuSparkScan.tagForGpu(this, convertedScan) - } - - override def convertToGpu(): GpuScan = convertedScan.get - }, - "Iceberg staged scan", - ClassTag(cpuStagedScanClass) - ), ).map(r => (r.getClassFor.asSubclass(classOf[Scan]), r)).toMap } diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkScan.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkScan.scala index 680adf9f05d..1da2dd40a94 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkScan.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkScan.scala @@ -83,12 +83,10 @@ object GpuSparkScan { new GpuSparkBatchQueryScan(cpuScan, rapidsConf, false) } else if (GpuSparkScanAccess.isCopyOnWriteScan(cpuScan)) { ShimUtils.newCopyOnWriteScan(cpuScan, rapidsConf, false) - } else if (GpuSparkScanAccess.isStagedScan(cpuScan)) { - new GpuSparkStagedScan(cpuScan, rapidsConf, false) } else { throw new IllegalArgumentException( - s"Currently Iceberg support only supports batch query, copy-on-write, and staged " + - s"scans, but got ${cpuScan.getClass.getName}") + s"Currently iceberg support only supports batch query scan and copy-on-write scan, " + + s"but got ${cpuScan.getClass.getName}") } } } diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkStagedScan.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkStagedScan.scala deleted file mode 100644 index 04f92bf01c5..00000000000 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkStagedScan.scala +++ /dev/null @@ -1,47 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. - * - * Licensed 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.source - -import scala.collection.JavaConverters._ - -import com.nvidia.spark.rapids.{GpuScan, RapidsConf} -import org.apache.iceberg.ScanTaskGroup -import org.apache.iceberg.types.Types - -import org.apache.spark.sql.connector.read.Scan - -/** GPU scan for file groups staged by Iceberg's rewrite_data_files action. */ -class GpuSparkStagedScan( - override val cpuScan: Scan, - override val rapidsConf: RapidsConf, - override val queryUsesInputFile: Boolean) - extends GpuSparkScan(cpuScan, rapidsConf, queryUsesInputFile) { - - override def groupingKeyType(): Types.StructType = - GpuSparkScanAccess.groupingKeyType(cpuScan) - - override def taskGroups(): Seq[_ <: ScanTaskGroup[_]] = - GpuSparkScanAccess.taskGroups(cpuScan).asScala.toSeq - - override def withInputFile(): GpuScan = - new GpuSparkStagedScan(cpuScan, rapidsConf, true) - - override def toString: String = - s"GpuSparkStagedScan(table=${GpuSparkScanAccess.table(cpuScan)}, " + - s"type=${GpuSparkScanAccess.expectedSchema(cpuScan).asStruct()}, " + - s"queryUseInputFile=$queryUsesInputFile)" -} diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala index 6c08571cb74..20bf98cde03 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala @@ -65,7 +65,6 @@ class GpuSparkWrite(cpu: Write) extends GpuWrite with RequiresDistributionAndOrd // - BatchAppend for append operations // - DynamicOverwrite for dynamic partition overwrite // - CopyOnWriteOperation for row-level copy-on-write operations - // - RewriteFiles for rewrite_data_files // Since these are private classes, we check the class name to determine which GPU version // to use val cpuBatch = cpu.toBatch @@ -76,7 +75,6 @@ class GpuSparkWrite(cpu: Write) extends GpuWrite with RequiresDistributionAndOrd case "DynamicOverwrite" => new GpuDynamicOverwrite(this, cpuBatch) case "OverwriteByFilter" => new GpuOverwriteByFilter(this, cpuBatch) case "CopyOnWriteOperation" => new GpuCopyOnWriteOperation(this, cpuBatch) - case "RewriteFiles" => new GpuRewriteFiles(this, cpuBatch) case _ => throw new UnsupportedOperationException( s"Unsupported Iceberg batch write type: $cpuBatchClassName") diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala index a03f015d108..80c6426ff21 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala @@ -74,14 +74,6 @@ class GpuCopyOnWriteOperation(write: GpuSparkWrite, cpuBatchWrite: BatchWrite) } } -/** GPU version of the batch write used by the rewrite_data_files procedure. */ -class GpuRewriteFiles(write: GpuSparkWrite, cpuBatchWrite: BatchWrite) - extends GpuBaseBatchWrite(write, cpuBatchWrite) { - override def commit(messages: Array[WriterCommitMessage]): Unit = { - cpuBatchWrite.commit(messages) - } -} - /** * GPU version of position delta batch write for merge-on-read DELETE operations. * This wraps the CPU PositionDeltaBatchWrite to handle position delete files. diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index 3d3283566b6..9cefb96eb62 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -555,38 +555,6 @@ def setup_iceberg_table(spark): }) -@iceberg -@ignore_order(local=True) -@allow_non_gpu("CallExec") -@pytest.mark.skipif( - not supports_iceberg_row_lineage_inheritance, - reason=ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) -def test_iceberg_v3_row_lineage_gpu_rewrite_data_files(spark_tmp_table_factory): - def setup_iceberg_table(spark, table): - spark.sql( - f"CREATE TABLE {table} (id BIGINT) USING ICEBERG " - "TBLPROPERTIES ('format-version' = '3')") - spark.range(0, 2).writeTo(table).append() - spark.range(2, 4).writeTo(table).append() - - def rewrite_data_files(spark, table): - spark.sql( - f"CALL spark_catalog.system.rewrite_data_files(table => '{table}', " - "options => map('min-input-files', '2'))").collect() - - def read_data_and_file_count(spark, table): - return spark.sql( - f"SELECT d.id, d._row_id, d._last_updated_sequence_number, f.data_file_count " - f"FROM {table} d CROSS JOIN (" - f"SELECT count(*) AS data_file_count FROM {table}.data_files) f") - - _assert_gpu_and_cpu_lineage_writes_are_equal( - spark_tmp_table_factory, - setup_iceberg_table, - rewrite_data_files, - read_data_and_file_count) - - @iceberg @ignore_order(local=True) @pytest.mark.skipif( diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala index 07b26d39f88..495d41402de 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala @@ -51,8 +51,6 @@ trait IcebergProbe { object IcebergProvider { val cpuBatchQueryScanClassName: String = "org.apache.iceberg.spark.source.SparkBatchQueryScan" val cpuCopyOnWriteScanClassName: String = "org.apache.iceberg.spark.source.SparkCopyOnWriteScan" - val cpuStagedScanClassName: String = "org.apache.iceberg.spark.source.SparkStagedScan" - private lazy val probe: IcebergProbe = ShimLoaderTemp.newIcebergProbe() From a84b7fc64d2e8775696ebea085114668f3504b01 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Tue, 8 Sep 2026 14:07:43 +0800 Subject: [PATCH 10/22] Use datagen for Iceberg row lineage tests Signed-off-by: Chong Gao --- .../main/python/iceberg/iceberg_ctas_test.py | 5 +++-- .../main/python/iceberg/iceberg_delete_test.py | 2 +- .../main/python/iceberg/iceberg_merge_test.py | 10 ++++++---- .../src/main/python/iceberg/iceberg_test.py | 17 +++++++++++------ .../main/python/iceberg/iceberg_update_test.py | 3 +-- 5 files changed, 22 insertions(+), 15 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py index 859922e919f..d615db9b30f 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py @@ -30,7 +30,8 @@ ctas_partition_transforms, supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, - ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON) + ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, + row_lineage_df) from marks import (iceberg, ignore_order, allow_non_gpu, allow_non_gpu_conditional, datagen_overrides) from spark_session import with_gpu_session, with_cpu_session, is_spark_400_or_later @@ -163,7 +164,7 @@ def assert_gpu_ctas(plan): _assert_gpu_equals_cpu_ctas( spark_tmp_table_factory, - lambda spark: spark.range(3), + lambda spark: row_lineage_df(spark), {"format-version": "3"}, conf=conf, read_func=lambda spark, table: spark.sql( diff --git a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py index ed7f88bf097..5cdd2f38553 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py @@ -220,7 +220,7 @@ def test_iceberg_v3_row_lineage_gpu_delete_leading_rows(spark_tmp_table_factory) do_delete_test( spark_tmp_table_factory, lambda spark, table: spark.sql(f"DELETE FROM {table} WHERE id < 3"), - data_gen_func=lambda spark: spark.range(1, 4).coalesce(1), + data_gen_func=lambda spark: row_lineage_df(spark, start=1), table_properties={"format-version": "3"}, conf=iceberg_delete_v3_enabled_conf, read_func=lambda spark, table: spark.sql( diff --git a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py index bc6fa462806..8bad7709b25 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py @@ -289,12 +289,14 @@ def setup_iceberg_tables(spark): f"CREATE TABLE {table} (id BIGINT, v BIGINT) USING ICEBERG " "TBLPROPERTIES ('format-version' = '3', " "'write.merge.mode' = 'copy-on-write')") - spark.range(0, 3).selectExpr("id", "CAST(0 AS BIGINT) AS v") \ - .coalesce(1).writeTo(table).append() + row_lineage_df(spark, with_value=True).writeTo(table).append() def merge(spark, table): - spark.range(1, 4, 2).selectExpr("id", "id * 10 AS v") \ - .createOrReplaceTempView(source_view) + row_lineage_df( + spark, + start=DEFAULT_DATA_GEN_LENGTH // 2, + with_value=True, + value_start=DEFAULT_DATA_GEN_LENGTH).createOrReplaceTempView(source_view) spark.sql( f"MERGE INTO {table} t USING {source_view} s ON t.id = s.id " "WHEN MATCHED THEN UPDATE SET v = s.v " diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index 9cefb96eb62..c1f010bb928 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -506,7 +506,7 @@ def test_iceberg_v3_row_lineage_append(spark_tmp_table_factory): def setup_iceberg_table(spark, table): spark.sql(f"CREATE TABLE {table} (id BIGINT) USING ICEBERG " f"TBLPROPERTIES ('format-version' = '2')") - spark.sql(f"INSERT INTO {table} VALUES (1), (2)") + row_lineage_df(spark, start=1).writeTo(table).append() spark.sql( f"ALTER TABLE {table} SET TBLPROPERTIES (" "'format-version' = '3', " @@ -514,10 +514,14 @@ def setup_iceberg_table(spark, table): "'read.split.target-size' = '4096', " "'read.split.open-file-cost' = '0')") + def append_data(spark, table): + row_lineage_df( + spark, start=DEFAULT_DATA_GEN_LENGTH + 1).writeTo(table).append() + _assert_gpu_and_cpu_lineage_writes_are_equal( spark_tmp_table_factory, setup_iceberg_table, - lambda spark, table: spark.range(3, 1503).coalesce(1).writeTo(table).append(), + append_data, lambda spark, table: spark.sql( f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {table}")) @@ -567,12 +571,13 @@ def setup_iceberg_table(spark, table): spark.sql( f"CREATE TABLE {table} (id BIGINT, v BIGINT) USING ICEBERG " "TBLPROPERTIES ('format-version' = '3')") - spark.range(0, 2).selectExpr("id", "CAST(0 AS BIGINT) AS v") \ - .writeTo(table).append() + row_lineage_df(spark, with_value=True).writeTo(table).append() def overwrite(spark, table): - spark.range(10, 12).selectExpr("id", "CAST(1 AS BIGINT) AS v") \ - .createOrReplaceTempView(source_view) + row_lineage_df( + spark, + start=DEFAULT_DATA_GEN_LENGTH, + with_value=True).createOrReplaceTempView(source_view) spark.sql(f"INSERT OVERWRITE {table} SELECT * FROM {source_view}").collect() _assert_gpu_and_cpu_lineage_writes_are_equal( diff --git a/integration_tests/src/main/python/iceberg/iceberg_update_test.py b/integration_tests/src/main/python/iceberg/iceberg_update_test.py index ef48536443f..f1128918396 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_update_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_update_test.py @@ -208,8 +208,7 @@ def test_iceberg_v3_row_lineage_gpu_update(spark_tmp_table_factory): do_update_test( spark_tmp_table_factory, lambda spark, table: spark.sql(f"UPDATE {table} SET v = v + 1 WHERE id = 1"), - data_gen_func=lambda spark: spark.range(0, 2).selectExpr( - "id", "CAST(0 AS BIGINT) AS v").coalesce(1), + data_gen_func=lambda spark: row_lineage_df(spark, with_value=True), table_properties={"format-version": "3"}, conf=iceberg_update_v3_enabled_conf, read_func=lambda spark, table: spark.sql( From e41202c3d962e99116549c995648ed2b4dbdff6a Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Tue, 8 Sep 2026 15:37:50 +0800 Subject: [PATCH 11/22] Make Iceberg row lineage writes deterministic Signed-off-by: Chong Gao --- .../main/python/iceberg/iceberg_delete_test.py | 18 ++++++++++++------ .../main/python/iceberg/iceberg_merge_test.py | 5 ++++- .../main/python/iceberg/iceberg_update_test.py | 18 ++++++++++++------ 3 files changed, 28 insertions(+), 13 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py index 5cdd2f38553..8fcd97de6c5 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py @@ -50,7 +50,8 @@ def create_iceberg_table_with_data(table_name: str, partition_col_sql=None, data_gen_func=None, table_properties=None, - delete_mode='copy-on-write'): + delete_mode='copy-on-write', + write_order=None): """Helper function to create and populate an Iceberg table for DELETE tests.""" # Default to copy-on-write mode, but allow override for merge-on-read tests base_props = { @@ -71,6 +72,8 @@ def create_iceberg_table_with_data(table_name: str, # Insert data def insert_data(spark): + if write_order: + spark.sql(f"ALTER TABLE {table_name} WRITE ORDERED BY {write_order}").collect() df = data_gen_func(spark) df.writeTo(table_name).append() @@ -79,7 +82,7 @@ def insert_data(spark): def do_delete_test(spark_tmp_table_factory, delete_sql_func, data_gen_func=None, partition_col_sql=None, table_properties=None, delete_mode='copy-on-write', conf=iceberg_delete_cow_enabled_conf, - read_func=None): + read_func=None, write_order=None): """ Helper function to test DELETE operations by comparing CPU and GPU results. @@ -92,6 +95,7 @@ def do_delete_test(spark_tmp_table_factory, delete_sql_func, data_gen_func=None, delete_mode: 'copy-on-write' or 'merge-on-read' conf: Spark configuration used for DELETE and result reads read_func: Optional function that takes (spark, table_name) and returns a DataFrame + write_order: Optional deterministic Iceberg write order """ base_table_name = get_full_table_name(spark_tmp_table_factory) cpu_table_name = f"{base_table_name}_cpu" @@ -99,9 +103,9 @@ def do_delete_test(spark_tmp_table_factory, delete_sql_func, data_gen_func=None, # Create identical tables for CPU and GPU create_iceberg_table_with_data(cpu_table_name, partition_col_sql, - data_gen_func, table_properties, delete_mode) + data_gen_func, table_properties, delete_mode, write_order) create_iceberg_table_with_data(gpu_table_name, partition_col_sql, - data_gen_func, table_properties, delete_mode) + data_gen_func, table_properties, delete_mode, write_order) # Execute DELETE on GPU def do_gpu_delete(spark): @@ -222,9 +226,11 @@ def test_iceberg_v3_row_lineage_gpu_delete_leading_rows(spark_tmp_table_factory) lambda spark, table: spark.sql(f"DELETE FROM {table} WHERE id < 3"), data_gen_func=lambda spark: row_lineage_df(spark, start=1), table_properties={"format-version": "3"}, - conf=iceberg_delete_v3_enabled_conf, + conf=copy_and_update( + iceberg_delete_v3_enabled_conf, {"spark.sql.shuffle.partitions": "1"}), read_func=lambda spark, table: spark.sql( - f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {table}")) + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {table}"), + write_order="id") def _do_test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col_sql, delete_mode, table_properties=None): diff --git a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py index 8bad7709b25..64e1396047c 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py @@ -289,6 +289,7 @@ def setup_iceberg_tables(spark): f"CREATE TABLE {table} (id BIGINT, v BIGINT) USING ICEBERG " "TBLPROPERTIES ('format-version' = '3', " "'write.merge.mode' = 'copy-on-write')") + spark.sql(f"ALTER TABLE {table} WRITE ORDERED BY id").collect() row_lineage_df(spark, with_value=True).writeTo(table).append() def merge(spark, table): @@ -303,13 +304,15 @@ def merge(spark, table): "WHEN NOT MATCHED THEN INSERT (id, v) VALUES (s.id, s.v)").collect() with_cpu_session(setup_iceberg_tables) + conf = copy_and_update( + iceberg_merge_v3_enabled_conf, {"spark.sql.shuffle.partitions": "1"}) _assert_gpu_and_cpu_merge_writes_are_equal( cpu_table, gpu_table, merge, lambda spark, table: spark.sql( f"SELECT id, v, _row_id, _last_updated_sequence_number FROM {table}"), - iceberg_merge_v3_enabled_conf) + conf) @allow_non_gpu("MergeRows$Keep", "MergeRows$Discard", "MergeRows$Split") diff --git a/integration_tests/src/main/python/iceberg/iceberg_update_test.py b/integration_tests/src/main/python/iceberg/iceberg_update_test.py index f1128918396..881859bfad1 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_update_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_update_test.py @@ -44,7 +44,8 @@ def create_iceberg_table_with_data(table_name: str, partition_col_sql=None, data_gen_func=None, table_properties=None, - update_mode='copy-on-write'): + update_mode='copy-on-write', + write_order=None): """Helper function to create and populate an Iceberg table for UPDATE tests. Args: @@ -72,6 +73,8 @@ def create_iceberg_table_with_data(table_name: str, # Insert data def insert_data(spark): + if write_order: + spark.sql(f"ALTER TABLE {table_name} WRITE ORDERED BY {write_order}").collect() df = data_gen_func(spark) df.writeTo(table_name).append() @@ -80,7 +83,7 @@ def insert_data(spark): def do_update_test(spark_tmp_table_factory, update_sql_func, data_gen_func=None, partition_col_sql=None, table_properties=None, update_mode='copy-on-write', conf=iceberg_update_cow_enabled_conf, - read_func=None): + read_func=None, write_order=None): """ Helper function to test UPDATE operations by comparing CPU and GPU results. @@ -93,6 +96,7 @@ def do_update_test(spark_tmp_table_factory, update_sql_func, data_gen_func=None, update_mode: Update mode - 'copy-on-write' or 'merge-on-read' conf: Spark configuration used for UPDATE and result reads read_func: Optional function that takes (spark, table_name) and returns a DataFrame + write_order: Optional deterministic Iceberg write order """ base_table_name = get_full_table_name(spark_tmp_table_factory) cpu_table_name = f"{base_table_name}_cpu" @@ -100,9 +104,9 @@ def do_update_test(spark_tmp_table_factory, update_sql_func, data_gen_func=None, # Create identical tables for CPU and GPU create_iceberg_table_with_data(cpu_table_name, partition_col_sql, - data_gen_func, table_properties, update_mode) + data_gen_func, table_properties, update_mode, write_order) create_iceberg_table_with_data(gpu_table_name, partition_col_sql, - data_gen_func, table_properties, update_mode) + data_gen_func, table_properties, update_mode, write_order) # Execute UPDATE on GPU def do_gpu_update(spark): @@ -210,9 +214,11 @@ def test_iceberg_v3_row_lineage_gpu_update(spark_tmp_table_factory): lambda spark, table: spark.sql(f"UPDATE {table} SET v = v + 1 WHERE id = 1"), data_gen_func=lambda spark: row_lineage_df(spark, with_value=True), table_properties={"format-version": "3"}, - conf=iceberg_update_v3_enabled_conf, + conf=copy_and_update( + iceberg_update_v3_enabled_conf, {"spark.sql.shuffle.partitions": "1"}), read_func=lambda spark, table: spark.sql( - f"SELECT id, v, _pos, _row_id, _last_updated_sequence_number FROM {table}")) + f"SELECT id, v, _pos, _row_id, _last_updated_sequence_number FROM {table}"), + write_order="id") @iceberg From 7e4a47f39ac1eb9f58f89ac0dbccfc10d44907f2 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Tue, 8 Sep 2026 16:00:04 +0800 Subject: [PATCH 12/22] Fix Iceberg row lineage replace data writes Signed-off-by: Chong Gao --- .../iceberg/spark/source/GpuSparkWrite.scala | 52 +++++++++++++++++-- .../com/nvidia/spark/rapids/GpuWrite.scala | 11 +++- .../datasources/v2/GpuReplaceDataExec.scala | 16 +++++- 3 files changed, 71 insertions(+), 8 deletions(-) diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala index 20bf98cde03..a4e9014f811 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala @@ -22,7 +22,7 @@ import scala.collection.JavaConverters._ import scala.util.{Failure, Success} import com.nvidia.spark.rapids._ -import com.nvidia.spark.rapids.Arm.closeOnExcept +import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} import com.nvidia.spark.rapids.RapidsPluginImplicits.AutoCloseableSeq import com.nvidia.spark.rapids.SpillPriorities.ACTIVE_ON_DECK_PRIORITY import com.nvidia.spark.rapids.fileio.iceberg.IcebergFileIO @@ -48,7 +48,7 @@ import org.apache.spark.sql.execution.datasources.v2.{AtomicCreateTableAsSelectE import org.apache.spark.sql.rapids.GpuWriteJobStatsTracker import org.apache.spark.sql.rapids.shims.SparkSessionUtils import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.util.SerializableConfiguration @@ -412,13 +412,55 @@ class GpuWriterFactory(val tableBroadcast: Broadcast[Table], } } +private trait GpuDataWriterWithRowLineage extends GpuDataWriterWithMetadata { + protected def dataSparkType: StructType + + def write(record: ColumnarBatch): Unit + + override def writeWithMetadata( + metadata: ColumnarBatch, + metadataSchema: StructType, + record: ColumnarBatch): Unit = { + val missingColumnCount = dataSparkType.length - record.numCols() + if (missingColumnCount == 0) { + write(record) + } else { + require(metadata.numRows() == record.numRows(), + s"Metadata row count ${metadata.numRows()} does not match record row count " + + s"${record.numRows()}") + require(missingColumnCount == GpuDataWriterWithRowLineage.lineageColumnNames.length, + s"Expected ${GpuDataWriterWithRowLineage.lineageColumnNames.length} row lineage " + + s"columns but record is missing $missingColumnCount columns") + + val lineageColumns = closeOnExcept(new Array[ColumnVector](missingColumnCount)) { columns => + GpuDataWriterWithRowLineage.lineageColumnNames.zipWithIndex.foreach { + case (name, index) => + val ordinal = metadataSchema.fieldIndex(name) + columns(index) = metadata.column(ordinal).asInstanceOf[GpuColumnVector].incRefCount() + } + columns + } + + withResource(new ColumnarBatch(lineageColumns, metadata.numRows())) { lineage => + write(GpuColumnVector.combineColumns(record, lineage)) + } + } + } +} + +private object GpuDataWriterWithRowLineage { + val lineageColumnNames: Seq[String] = Seq("_row_id", "_last_updated_sequence_number") +} + class GpuUnpartitionedDataWriter( val fileWriterFactory: GpuSparkFileWriterFactory, val fileFactory: OutputFileFactory, val io: FileIO, val spec: PartitionSpec, val targetFileSize: Long) - extends DataWriter[ColumnarBatch] { + extends DataWriter[ColumnarBatch] with GpuDataWriterWithRowLineage { + override protected def dataSparkType: StructType = fileWriterFactory.dataSparkType + private val delegate = new GpuRollingDataWriter( fileWriterFactory, fileFactory, @@ -460,10 +502,10 @@ class GpuPartitionedDataWriter( val io: FileIO, val spec: PartitionSpec, val dataSchema: Schema, - val dataSparkType: StructType, + override val dataSparkType: StructType, val targetFileSize: Long, val fanoutEnabled: Boolean, -) extends DataWriter[ColumnarBatch] { +) extends DataWriter[ColumnarBatch] with GpuDataWriterWithRowLineage { private val delegate: PartitioningWriter[SpillableColumnarBatch, DataWriteResult] = if (fanoutEnabled) { diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala index aa53567c4e0..66fc11cec11 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,11 +17,20 @@ package com.nvidia.spark.rapids import org.apache.spark.sql.connector.write.Write +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch trait GpuWrite extends Write { var metrics: Map[String, GpuMetric] = Map.empty } +trait GpuDataWriterWithMetadata { + def writeWithMetadata( + metadata: ColumnarBatch, + metadataSchema: StructType, + record: ColumnarBatch): Unit +} + // Allows use of GpuWrite from Java code abstract class GpuWriteWrapper extends GpuWrite { } diff --git a/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala b/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala index 89ab7235f86..75d065ac9cb 100644 --- a/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala +++ b/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala @@ -29,7 +29,7 @@ spark-rapids-shim-json-lines ***/ package org.apache.spark.sql.execution.datasources.v2 import com.nvidia.spark.rapids.Arm.withResource -import com.nvidia.spark.rapids.GpuWrite +import com.nvidia.spark.rapids.{GpuDataWriterWithMetadata, GpuWrite} import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.GpuProjectingColumnarBatch @@ -66,11 +66,23 @@ case class GpuReplaceDataWritingSparkTask( extends GpuWritingSparkTask[DataWriter[ColumnarBatch]] { private lazy val rowProjection = GpuProjectingColumnarBatch(projs.rowProjection) + private lazy val metadataProjection = projs.metadataProjection.map(GpuProjectingColumnarBatch(_)) + override protected def write( writer: DataWriter[ColumnarBatch], batch: ColumnarBatch): Unit = { withResource(rowProjection.project(batch)) { projected => - writer.write(projected) + metadataProjection match { + case Some(projection) => + withResource(projection.project(batch)) { metadata => + writer match { + case metadataWriter: GpuDataWriterWithMetadata => + metadataWriter.writeWithMetadata(metadata, projection.schema, projected) + case _ => writer.write(metadata, projected) + } + } + case None => writer.write(projected) + } } } } From aaf88ac31a56374c9dc2350f778e65155152e092 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Tue, 8 Sep 2026 16:08:08 +0800 Subject: [PATCH 13/22] Fix Scala import order Signed-off-by: Chong Gao --- .../sql/execution/datasources/v2/GpuReplaceDataExec.scala | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala b/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala index 75d065ac9cb..7f2a0fb10d0 100644 --- a/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala +++ b/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala @@ -28,8 +28,8 @@ spark-rapids-shim-json-lines ***/ package org.apache.spark.sql.execution.datasources.v2 -import com.nvidia.spark.rapids.Arm.withResource import com.nvidia.spark.rapids.{GpuDataWriterWithMetadata, GpuWrite} +import com.nvidia.spark.rapids.Arm.withResource import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.GpuProjectingColumnarBatch @@ -76,6 +76,8 @@ case class GpuReplaceDataWritingSparkTask( case Some(projection) => withResource(projection.project(batch)) { metadata => writer match { + // Newer Iceberg versions use this bridge to append row-lineage columns, while + // older Iceberg versions continue through Spark's metadata-aware writer path. case metadataWriter: GpuDataWriterWithMetadata => metadataWriter.writeWithMetadata(metadata, projection.schema, projected) case _ => writer.write(metadata, projected) From e48289eb18928e71663143fc62482c96a88653c4 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Tue, 8 Sep 2026 16:23:14 +0800 Subject: [PATCH 14/22] Fix Iceberg row lineage writer visibility Signed-off-by: Chong Gao --- .../scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala index a4e9014f811..ad871a8a24c 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala @@ -412,7 +412,7 @@ class GpuWriterFactory(val tableBroadcast: Broadcast[Table], } } -private trait GpuDataWriterWithRowLineage extends GpuDataWriterWithMetadata { +trait GpuDataWriterWithRowLineage extends GpuDataWriterWithMetadata { protected def dataSparkType: StructType def write(record: ColumnarBatch): Unit @@ -448,7 +448,7 @@ private trait GpuDataWriterWithRowLineage extends GpuDataWriterWithMetadata { } } -private object GpuDataWriterWithRowLineage { +object GpuDataWriterWithRowLineage { val lineageColumnNames: Seq[String] = Seq("_row_id", "_last_updated_sequence_number") } From 1d89ff1a4083ea49645ba36c18d34edd6382d09a Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Wed, 9 Sep 2026 17:55:43 +0800 Subject: [PATCH 15/22] Align GPU data writers with Spark API Signed-off-by: Chong Gao --- .../iceberg/spark/source/GpuSparkWrite.scala | 30 +++++++++----- .../com/nvidia/spark/rapids/GpuWrite.scala | 14 +++---- .../nvidia/spark/rapids/GpuDataWriter.scala | 41 +++++++++++++++++++ .../v2/WriteToDataSourceV2Exec.scala | 9 +++- .../nvidia/spark/rapids/GpuDataWriter.scala | 38 +++++++++++++++++ .../datasources/v2/GpuReplaceDataExec.scala | 25 ++++++----- 6 files changed, 128 insertions(+), 29 deletions(-) create mode 100644 sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/GpuDataWriter.scala create mode 100644 sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/GpuDataWriter.scala diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala index ad871a8a24c..5454741eb71 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala @@ -374,11 +374,18 @@ class GpuWriterFactory(val tableBroadcast: Broadcast[Table], val outputWriterFactory: ColumnarOutputWriterFactory, val statsTracker: GpuWriteJobStatsTracker, val hadoopConf: SerializableConfiguration -) extends DataWriterFactory { +) extends GpuDataWriterFactory { private lazy val fileIO: IcebergFileIO = new IcebergFileIO(tableBroadcast.value.io()) override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = { + createWriter(partitionId, taskId, null) + } + + override def createWriter( + partitionId: Int, + taskId: Long, + metadataSchema: StructType): DataWriter[InternalRow] = { val table = tableBroadcast.value val spec = table.specs().get(outputSpecId) val io = table.io() @@ -402,24 +409,25 @@ class GpuWriterFactory(val tableBroadcast: Broadcast[Table], fileIO) if (spec.isUnpartitioned) { - new GpuUnpartitionedDataWriter(writerFactory, outputFileFactory, io, spec, targetFileSize) + new GpuUnpartitionedDataWriter( + writerFactory, outputFileFactory, io, spec, targetFileSize, metadataSchema) .asInstanceOf[DataWriter[InternalRow]] } else { new GpuPartitionedDataWriter(writerFactory, outputFileFactory, io, spec, writeSchema, - dsSchema, targetFileSize, useFanout) + dsSchema, targetFileSize, useFanout, metadataSchema) .asInstanceOf[DataWriter[InternalRow]] } } } -trait GpuDataWriterWithRowLineage extends GpuDataWriterWithMetadata { +trait GpuDataWriterWithRowLineage extends GpuDataWriter { protected def dataSparkType: StructType + protected def metadataSchema: StructType - def write(record: ColumnarBatch): Unit + override def write(record: ColumnarBatch): Unit - override def writeWithMetadata( + override def write( metadata: ColumnarBatch, - metadataSchema: StructType, record: ColumnarBatch): Unit = { val missingColumnCount = dataSparkType.length - record.numCols() if (missingColumnCount == 0) { @@ -457,8 +465,9 @@ class GpuUnpartitionedDataWriter( val fileFactory: OutputFileFactory, val io: FileIO, val spec: PartitionSpec, - val targetFileSize: Long) - extends DataWriter[ColumnarBatch] with GpuDataWriterWithRowLineage { + val targetFileSize: Long, + override protected val metadataSchema: StructType) + extends GpuDataWriterWithRowLineage { override protected def dataSparkType: StructType = fileWriterFactory.dataSparkType private val delegate = new GpuRollingDataWriter( @@ -505,7 +514,8 @@ class GpuPartitionedDataWriter( override val dataSparkType: StructType, val targetFileSize: Long, val fanoutEnabled: Boolean, -) extends DataWriter[ColumnarBatch] with GpuDataWriterWithRowLineage { + override protected val metadataSchema: StructType, +) extends GpuDataWriterWithRowLineage { private val delegate: PartitioningWriter[SpillableColumnarBatch, DataWriteResult] = if (fanoutEnabled) { diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala index 66fc11cec11..355ba4bdfc3 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala @@ -16,19 +16,19 @@ package com.nvidia.spark.rapids -import org.apache.spark.sql.connector.write.Write +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.write.{DataWriter, DataWriterFactory, Write} import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.vectorized.ColumnarBatch trait GpuWrite extends Write { var metrics: Map[String, GpuMetric] = Map.empty } -trait GpuDataWriterWithMetadata { - def writeWithMetadata( - metadata: ColumnarBatch, - metadataSchema: StructType, - record: ColumnarBatch): Unit +trait GpuDataWriterFactory extends DataWriterFactory { + def createWriter( + partitionId: Int, + taskId: Long, + metadataSchema: StructType): DataWriter[InternalRow] } // Allows use of GpuWrite from Java code diff --git a/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/GpuDataWriter.scala b/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/GpuDataWriter.scala new file mode 100644 index 00000000000..1bcde7e2b57 --- /dev/null +++ b/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/GpuDataWriter.scala @@ -0,0 +1,41 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed 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. + */ +/*** spark-rapids-shim-json-lines +{"spark": "350"} +{"spark": "351"} +{"spark": "352"} +{"spark": "353"} +{"spark": "354"} +{"spark": "355"} +{"spark": "356"} +{"spark": "357"} +{"spark": "358"} +{"spark": "359"} +spark-rapids-shim-json-lines ***/ + +package com.nvidia.spark.rapids + +import org.apache.spark.sql.connector.write.DataWriter +import org.apache.spark.sql.vectorized.ColumnarBatch + +trait GpuDataWriter extends DataWriter[ColumnarBatch] { + override def write(record: ColumnarBatch): Unit + + def write(metadata: ColumnarBatch, record: ColumnarBatch): Unit = { + throw new UnsupportedOperationException( + "Writing records with metadata is not supported before Spark 4.0") + } +} diff --git a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala index 9fe2e7b830b..18ae40b9db0 100644 --- a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala +++ b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala @@ -358,6 +358,13 @@ trait GpuWritingSparkTask[W <: DataWriter[ColumnarBatch]] extends Logging with S protected def write(writer: W, row: ColumnarBatch): Unit + protected def createWriter( + writerFactory: DataWriterFactory, + partitionId: Int, + taskId: Long): W = { + writerFactory.createWriter(partitionId, taskId).asInstanceOf[W] + } + def run( writerFactory: DataWriterFactory, context: TaskContext, @@ -369,7 +376,7 @@ trait GpuWritingSparkTask[W <: DataWriter[ColumnarBatch]] extends Logging with S val partId = context.partitionId() val taskId = context.taskAttemptId() val attemptId = context.attemptNumber() - val dataWriter = writerFactory.createWriter(partId, taskId).asInstanceOf[W] + val dataWriter = createWriter(writerFactory, partId, taskId) var count = 0L // write the data and commit this writer. diff --git a/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/GpuDataWriter.scala b/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/GpuDataWriter.scala new file mode 100644 index 00000000000..06da0e2349d --- /dev/null +++ b/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/GpuDataWriter.scala @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed 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. + */ +/*** spark-rapids-shim-json-lines +{"spark": "400"} +{"spark": "401"} +{"spark": "402"} +{"spark": "403"} +{"spark": "404"} +{"spark": "411"} +{"spark": "412"} +{"spark": "413"} +{"spark": "420"} +{"spark": "500"} +spark-rapids-shim-json-lines ***/ + +package com.nvidia.spark.rapids + +import org.apache.spark.sql.connector.write.DataWriter +import org.apache.spark.sql.vectorized.ColumnarBatch + +trait GpuDataWriter extends DataWriter[ColumnarBatch] { + override def write(record: ColumnarBatch): Unit + + override def write(metadata: ColumnarBatch, record: ColumnarBatch): Unit +} diff --git a/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala b/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala index 7f2a0fb10d0..ed970ea96af 100644 --- a/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala +++ b/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala @@ -28,13 +28,13 @@ spark-rapids-shim-json-lines ***/ package org.apache.spark.sql.execution.datasources.v2 -import com.nvidia.spark.rapids.{GpuDataWriterWithMetadata, GpuWrite} +import com.nvidia.spark.rapids.{GpuDataWriter, GpuDataWriterFactory, GpuWrite} import com.nvidia.spark.rapids.Arm.withResource import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.GpuProjectingColumnarBatch import org.apache.spark.sql.catalyst.util.ReplaceDataProjections -import org.apache.spark.sql.connector.write.DataWriter +import org.apache.spark.sql.connector.write.DataWriterFactory import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.vectorized.ColumnarBatch @@ -63,25 +63,28 @@ case class GpuReplaceDataExec( case class GpuReplaceDataWritingSparkTask( projs: ReplaceDataProjections) - extends GpuWritingSparkTask[DataWriter[ColumnarBatch]] { + extends GpuWritingSparkTask[GpuDataWriter] { private lazy val rowProjection = GpuProjectingColumnarBatch(projs.rowProjection) private lazy val metadataProjection = projs.metadataProjection.map(GpuProjectingColumnarBatch(_)) + override protected def createWriter( + writerFactory: DataWriterFactory, + partitionId: Int, + taskId: Long): GpuDataWriter = { + writerFactory.asInstanceOf[GpuDataWriterFactory] + .createWriter(partitionId, taskId, metadataProjection.map(_.schema).orNull) + .asInstanceOf[GpuDataWriter] + } + override protected def write( - writer: DataWriter[ColumnarBatch], + writer: GpuDataWriter, batch: ColumnarBatch): Unit = { withResource(rowProjection.project(batch)) { projected => metadataProjection match { case Some(projection) => withResource(projection.project(batch)) { metadata => - writer match { - // Newer Iceberg versions use this bridge to append row-lineage columns, while - // older Iceberg versions continue through Spark's metadata-aware writer path. - case metadataWriter: GpuDataWriterWithMetadata => - metadataWriter.writeWithMetadata(metadata, projection.schema, projected) - case _ => writer.write(metadata, projected) - } + writer.write(metadata, projected) } case None => writer.write(projected) } From 3b5e4e8266055a6669a20609898d4154a011c898 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Thu, 10 Sep 2026 09:52:27 +0800 Subject: [PATCH 16/22] Remove duplicate Iceberg cleanup bridge Signed-off-by: Chong Gao --- .../org/apache/iceberg/spark/source/GpuSparkWriteAccess.java | 4 ---- 1 file changed, 4 deletions(-) diff --git a/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkWriteAccess.java b/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkWriteAccess.java index 50925b82507..2ae5d96e244 100644 --- a/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkWriteAccess.java +++ b/iceberg-common/src/main/java/org/apache/iceberg/spark/source/GpuSparkWriteAccess.java @@ -56,10 +56,6 @@ public static String sparkWriteClassName() { return SparkWrite.class.getName(); } - public static void deleteTaskFiles(FileIO io, List> files) { - SparkCleanupUtil.deleteTaskFiles(io, files); - } - public static Table table(Write write) { return readField(sparkWrite(write), "table", Table.class); } From ce95927c19be9abaee27646ba592e076e4ab1443 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Thu, 10 Sep 2026 09:56:44 +0800 Subject: [PATCH 17/22] Revert unrelated Iceberg formatting changes Signed-off-by: Chong Gao --- .../com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala | 1 + .../src/main/scala/org/apache/iceberg/spark/source/write.scala | 2 +- .../scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala | 1 + 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala b/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala index b7f2296e564..2c53190a658 100644 --- a/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala +++ b/iceberg/common/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderBase.scala @@ -43,6 +43,7 @@ abstract class IcebergProviderBase extends IcebergProvider { IcebergProvider.cpuBatchQueryScanClassName) val cpuCopyOnWriteScanClass = ShimReflectionUtils.loadClass( IcebergProvider.cpuCopyOnWriteScanClassName) + Seq( new ScanRule[Scan]( (a, conf, p, r) => new ScanMeta[Scan](a, conf, p, r) { diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala index 80c6426ff21..4861d846e73 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/write.scala @@ -100,4 +100,4 @@ class GpuPositionDeltaBatchWrite(write: GpuSparkPositionDeltaWrite, override def createBatchWriterFactory(info: PhysicalWriteInfo): DeltaWriterFactory = { write.createDeltaWriterFactory } -} +} \ No newline at end of file diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala index 495d41402de..c636c601391 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProvider.scala @@ -51,6 +51,7 @@ trait IcebergProbe { object IcebergProvider { val cpuBatchQueryScanClassName: String = "org.apache.iceberg.spark.source.SparkBatchQueryScan" val cpuCopyOnWriteScanClassName: String = "org.apache.iceberg.spark.source.SparkCopyOnWriteScan" + private lazy val probe: IcebergProbe = ShimLoaderTemp.newIcebergProbe() From e5b657194d15a855a8f25e83a607ffc5dfc78317 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Thu, 10 Sep 2026 13:49:39 +0800 Subject: [PATCH 18/22] Cache Iceberg row lineage metadata ordinals Signed-off-by: Chong Gao --- .../org/apache/iceberg/spark/source/GpuSparkWrite.scala | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala index 5454741eb71..935b03f9078 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala @@ -424,6 +424,9 @@ trait GpuDataWriterWithRowLineage extends GpuDataWriter { protected def dataSparkType: StructType protected def metadataSchema: StructType + private lazy val lineageColumnOrdinals = + GpuDataWriterWithRowLineage.lineageColumnNames.map(metadataSchema.fieldIndex) + override def write(record: ColumnarBatch): Unit override def write( @@ -441,9 +444,8 @@ trait GpuDataWriterWithRowLineage extends GpuDataWriter { s"columns but record is missing $missingColumnCount columns") val lineageColumns = closeOnExcept(new Array[ColumnVector](missingColumnCount)) { columns => - GpuDataWriterWithRowLineage.lineageColumnNames.zipWithIndex.foreach { - case (name, index) => - val ordinal = metadataSchema.fieldIndex(name) + lineageColumnOrdinals.zipWithIndex.foreach { + case (ordinal, index) => columns(index) = metadata.column(ordinal).asInstanceOf[GpuColumnVector].incRefCount() } columns From 28cb8cebcbffff5241ff22f5e99c2b1574d85a40 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Thu, 10 Sep 2026 14:03:48 +0800 Subject: [PATCH 19/22] Remove redundant GPU merge output padding Signed-off-by: Chong Gao --- .../datasources/v2/GpuMergeRowsExec.scala | 28 ++----------------- .../v2/GpuMergeBatchIteratorRetrySuite.scala | 5 ++-- 2 files changed, 5 insertions(+), 28 deletions(-) diff --git a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala index 8ee9c002ca0..e219f8df0aa 100644 --- a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala +++ b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala @@ -171,18 +171,10 @@ case class GpuMergeRowsExec( val boundMatchedBySourceInsts = GpuBindReferences.bindGpuReferences( notMatchedBySourceInstructions, child.output, allMetrics) .asInstanceOf[Seq[GpuInstruction]] - val instructionOutputs = (boundMatchedInsts ++ boundNotMatchedInsts ++ - boundMatchedBySourceInsts).flatMap(_.outputs) - val outputDataTypes = if (instructionOutputs.nonEmpty) { - instructionOutputs.maxBy(_.length).map(_.dataType).toArray - } else { - GpuColumnVector.extractTypes(schema) - } child.executeColumnar().mapPartitions { iter => new GpuMergeBatchIterator( dataTypes, - outputDataTypes, iter, boundTargetRowPresent, boundSourceRowPresent, @@ -333,20 +325,8 @@ object GpuMergeRowsExec { condition.columnarEval(batch) } - def applyOutputs( - batch: ColumnarBatch, - outputDataTypes: Array[DataType]): Seq[ColumnarBatch] = { - outputs.map { output => - require(output.length <= outputDataTypes.length, - s"Merge output has ${output.length} columns, expected at most ${outputDataTypes.length}") - // Spark permits merge actions to omit trailing columns from their InternalRow. Iceberg v3 - // uses this for unchanged rows while update and insert actions append row-lineage fields. - // cuDF tables require identical schemas for concatenation, so materialize the omitted - // trailing fields as correctly typed null columns. - val paddedOutput = output ++ outputDataTypes.drop(output.length) - .map(GpuLiteral(null, _)) - GpuProjectExec.project(batch, paddedOutput) - } + def applyOutputs(batch: ColumnarBatch): Seq[ColumnarBatch] = { + outputs.map(output => GpuProjectExec.project(batch, output)) } override def nullable: Boolean = false @@ -388,7 +368,6 @@ object GpuMergeRowsExec { * Similar to Spark's MergeRowIterator but operates on batches instead of rows. * * @param inputDataTypes Spark data types of input iterator. - * @param outputDataTypes Spark data types of the merge output. * @param inputIter Iterator of input columnar batches * @param isTargetRowPresent Bound GPU expression to check if target row is present * @param isSourceRowPresent Bound GPU expression to check if source row is present @@ -402,7 +381,6 @@ object GpuMergeRowsExec { */ class GpuMergeBatchIterator( inputDataTypes: Array[DataType], - outputDataTypes: Array[DataType], inputIter: Iterator[ColumnarBatch], isTargetRowPresent: GpuExpression, isSourceRowPresent: GpuExpression, @@ -540,7 +518,7 @@ class GpuMergeBatchIterator( if (writeSummaryEnabled) { attemptMetrics.record(instructionExec, filtered.numRows(), sourcePresent) } - outputs ++= instructionExec.applyOutputs(filtered, outputDataTypes) + outputs ++= instructionExec.applyOutputs(filtered) .map(SpillableColumnarBatch .apply(_, SpillPriorities.ACTIVE_ON_DECK_PRIORITY)) } diff --git a/tests/src/test/spark411/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeBatchIteratorRetrySuite.scala b/tests/src/test/spark411/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeBatchIteratorRetrySuite.scala index ff3fa24c4dc..144c9253e21 100644 --- a/tests/src/test/spark411/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeBatchIteratorRetrySuite.scala +++ b/tests/src/test/spark411/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeBatchIteratorRetrySuite.scala @@ -73,9 +73,8 @@ class GpuMergeBatchIteratorRetrySuite extends RmmSparkRetrySuiteBase { Seq(GpuBoundReference(0, IntegerType, nullable = true)(ExprId(0), "id")), ACTION_INSERT) val it = new GpuMergeBatchIterator( - inputDataTypes = Array(IntegerType), - outputDataTypes = Array(IntegerType), - inputIter = Seq(buildBatch()).iterator, + Array(IntegerType), + Seq(buildBatch()).iterator, isTargetRowPresent = GpuLiteral.create(false, BooleanType), isSourceRowPresent = GpuLiteral.create(true, BooleanType), matchedInstructionExecs = Nil, From 680c184944fda5370106bb50f58b90f5e54717c6 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Thu, 10 Sep 2026 14:44:17 +0800 Subject: [PATCH 20/22] Verify Iceberg next row ID after GPU writes Signed-off-by: Chong Gao --- .../src/main/python/iceberg/iceberg_test.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index 97602fe8d2a..6e8ee33d439 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -69,6 +69,11 @@ def setup_tables(spark): def run_write(spark, table): write_func(spark, table) + def next_row_id(spark, table): + iceberg_table = spark._jvm.org.apache.iceberg.spark.Spark3Util.loadIcebergTable( + spark._jsparkSession, table) + return iceberg_table.operations().current().nextRowId() + with_cpu_session(setup_tables) with_cpu_session(lambda spark: run_write(spark, cpu_table), conf=_ROW_LINEAGE_WRITE_CONF) with_gpu_session(lambda spark: run_write(spark, gpu_table), conf=_ROW_LINEAGE_WRITE_CONF) @@ -78,6 +83,11 @@ def run_write(spark, table): gpu_data = with_cpu_session( lambda spark: read_func(spark, gpu_table).collect(), conf=_ROW_LINEAGE_WRITE_CONF) assert_equal_with_local_sort(cpu_data, gpu_data) + cpu_next_row_id = with_cpu_session( + lambda spark: next_row_id(spark, cpu_table), conf=_ROW_LINEAGE_WRITE_CONF) + gpu_next_row_id = with_cpu_session( + lambda spark: next_row_id(spark, gpu_table), conf=_ROW_LINEAGE_WRITE_CONF) + assert cpu_next_row_id == gpu_next_row_id pytestmark = iceberg_unsupported_mark From 17f975d1059e05e37b1c3240180873b374261af8 Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Fri, 11 Sep 2026 09:41:32 +0800 Subject: [PATCH 21/22] Preserve row lineage and partition order in GPU Iceberg writes Keep Spark 4 INSERT and REINSERT rows in partition order with an aligned reinsert mask. Restore missing lineage from reinsert metadata and let new rows inherit lineage. Align mixed MERGE action outputs before concatenation, and cover row order, metadata and ownership in unit tests. Addresses review: https://github.com/NVIDIA/cudf-spark/pull/15866#discussion_r3975139421 Signed-off-by: Chong Gao --- .../source/GpuSparkPositionDeltaWrite.scala | 46 ++++- .../iceberg/spark/source/GpuSparkWrite.scala | 101 +++++++++-- .../spark/rapids/GpuDeltaBatchWriter.java | 37 ++++ .../rapids/shims/DeltaInsertFilter.scala | 2 + .../datasources/v2/GpuMergeRowsExec.scala | 27 ++- .../v2/WriteToDataSourceV2Exec.scala | 130 +++++++++++--- .../rapids/shims/DeltaInsertFilter.scala | 15 +- .../source/GpuRowLineageWriterSuite.scala | 146 ++++++++++++++++ .../v2/GpuDeltaWritingSparkTaskSuite.scala | 162 ++++++++++++++++++ 9 files changed, 600 insertions(+), 66 deletions(-) create mode 100644 sql-plugin-api/src/main/java/com/nvidia/spark/rapids/GpuDeltaBatchWriter.java create mode 100644 tests/src/test/spark350/scala/org/apache/iceberg/spark/source/GpuRowLineageWriterSuite.scala create mode 100644 tests/src/test/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuDeltaWritingSparkTaskSuite.scala diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala index 2c827981603..d876829be91 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala @@ -252,7 +252,7 @@ class GpuPositionDeltaWriterFactory( } -trait GpuDeltaWriter extends DeltaWriter[ColumnarBatch] { +trait GpuIcebergDeltaWriter extends DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter { def context: GpuWriteContext @@ -470,7 +470,7 @@ class GpuBasePositionDeltaWriter( * Base trait for delta writers that handle both deletes and data writes. * This is the GPU equivalent of Java's DeleteAndDataDeltaWriter. */ -trait GpuDeleteAndDataDeltaWriter extends GpuDeltaWriter { +trait GpuDeleteAndDataDeltaWriter extends GpuIcebergDeltaWriter { protected val table: Table protected val delegate: GpuBasePositionDeltaWriter protected val io: FileIO @@ -492,6 +492,29 @@ trait GpuDeleteAndDataDeltaWriter extends GpuDeltaWriter { private var closed: Boolean = false + protected def insertData(row: ColumnarBatch): Unit + + override def insert(row: ColumnarBatch): Unit = reinsert(null, row) + + override def reinsert(metadata: ColumnarBatch, row: ColumnarBatch): Unit = { + val physicalRow = withResource(Seq(metadata, row)) { _ => + GpuDataWriterWithRowLineage.appendLineage( + row, metadata, context.dataSparkType, context.metadataSparkType) + } + insertData(physicalRow) + } + + override def insertAndReinsert( + metadata: ColumnarBatch, + row: ColumnarBatch, + reinsertMask: CudfColumnVector): Unit = { + val physicalRow = withResource(Seq(metadata, row, reinsertMask)) { _ => + GpuDataWriterWithRowLineage.appendLineage( + row, metadata, context.dataSparkType, context.metadataSparkType, reinsertMask) + } + insertData(physicalRow) + } + override def delete(metadata: ColumnarBatch, rowId: ColumnarBatch): Unit = { require(metadata != null, "Metadata batch must be non null") @@ -569,7 +592,7 @@ class GpuDeleteOnlyDeltaWriter( table: Table, writerFactory: GpuSparkFileWriterFactory, deleteFileFactory: OutputFileFactory, - override val context: GpuWriteContext) extends GpuDeltaWriter { + override val context: GpuWriteContext) extends GpuIcebergDeltaWriter { private val io: FileIO = table.io() private val specs: mutable.Map[Integer, PartitionSpec] = table.specs().asScala @@ -642,6 +665,19 @@ class GpuDeleteOnlyDeltaWriter( throw new UnsupportedOperationException("Delete-only writer does not support inserts") } + override def insertAndReinsert( + metadata: ColumnarBatch, + row: ColumnarBatch, + reinsertMask: CudfColumnVector): Unit = { + withResource(Seq(metadata, row, reinsertMask)) { _ => + throw new UnsupportedOperationException("Delete-only writer does not support inserts") + } + } + + override def reinsert(metadata: ColumnarBatch, row: ColumnarBatch): Unit = { + throw new UnsupportedOperationException("Delete-only writer does not support reinserts") + } + override def commit(): WriterCommitMessage = { close() val result = delegate.result() @@ -704,7 +740,7 @@ class GpuUnpartitionedDeltaWriter( delegate.writeDelete(batch, spec, partition) } - override def insert(row: ColumnarBatch): Unit = { + override protected def insertData(row: ColumnarBatch): Unit = { val spillBatch = closeOnExcept(row) { _ => SpillableColumnarBatch(row, ACTIVE_ON_DECK_PRIORITY) } @@ -756,7 +792,7 @@ class GpuPartitionedDeltaWriter( delegate.writeDelete(batch, spec, partition) } - override def insert(row: ColumnarBatch): Unit = { + override protected def insertData(row: ColumnarBatch): Unit = { // Partition the data and write each partition dataPartitioner.partition(row) .safeConsume { part => diff --git a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala index 935b03f9078..d1207ceae34 100644 --- a/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala +++ b/iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala @@ -21,6 +21,7 @@ import java.util.Locale import scala.collection.JavaConverters._ import scala.util.{Failure, Success} +import ai.rapids.cudf.{ColumnVector => CudfColumnVector} import com.nvidia.spark.rapids._ import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} import com.nvidia.spark.rapids.RapidsPluginImplicits.AutoCloseableSeq @@ -47,7 +48,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.datasources.v2.{AtomicCreateTableAsSelectExec, AtomicReplaceTableAsSelectExec} import org.apache.spark.sql.rapids.GpuWriteJobStatsTracker import org.apache.spark.sql.rapids.shims.SparkSessionUtils -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{LongType, StructType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.util.SerializableConfiguration @@ -424,44 +425,106 @@ trait GpuDataWriterWithRowLineage extends GpuDataWriter { protected def dataSparkType: StructType protected def metadataSchema: StructType - private lazy val lineageColumnOrdinals = - GpuDataWriterWithRowLineage.lineageColumnNames.map(metadataSchema.fieldIndex) - override def write(record: ColumnarBatch): Unit override def write( metadata: ColumnarBatch, record: ColumnarBatch): Unit = { + write(GpuDataWriterWithRowLineage.appendLineage( + record, metadata, dataSparkType, metadataSchema)) + } +} + +object GpuDataWriterWithRowLineage { + val lineageColumnNames: Seq[String] = Seq("_row_id", "_last_updated_sequence_number") + + /** + * Returns an owned physical row batch matching dataSparkType without consuming record, metadata, + * or reinsertMask. Records that already contain all physical columns keep their existing values. + * Otherwise, appends the missing _row_id and _last_updated_sequence_number columns at the end. + * + * Without metadata, both appended columns are null. With metadata, lineage columns are looked up + * by name in metadataSchema. If reinsertMask is absent, their values are copied for every row. + * With a mask, only REINSERT rows (true) copy metadata; INSERT rows (false) receive nulls even + * when their metadata contains values. Metadata nulls are preserved for Iceberg's lineage + * inheritance mechanism; this method does not assign row IDs or sequence numbers. + * + * For example, a batch containing a REINSERT followed by an INSERT: + * {{{ + * record: + * id amount + * 1 200 + * 2 300 + * + * metadata: + * _row_id _last_updated_sequence_number + * 101 null + * 999 8 + * + * reinsertMask: [true, false] + * + * result (columns in dataSparkType order): + * id amount _row_id _last_updated_sequence_number + * 1 200 101 null + * 2 300 null null + * }}} + * The REINSERT preserves row ID 101; the INSERT ignores metadata values 999 and 8 so its lineage + * can be inherited. Input row order is unchanged. + */ + def appendLineage( + record: ColumnarBatch, + metadata: ColumnarBatch, + dataSparkType: StructType, + metadataSchema: StructType, + reinsertMask: CudfColumnVector = null): ColumnarBatch = { + if (reinsertMask != null) { + require(reinsertMask.getRowCount == record.numRows(), + "Reinsert mask row count does not match record row count") + } val missingColumnCount = dataSparkType.length - record.numCols() if (missingColumnCount == 0) { - write(record) + GpuColumnVector.combineColumns(record) } else { - require(metadata.numRows() == record.numRows(), - s"Metadata row count ${metadata.numRows()} does not match record row count " + - s"${record.numRows()}") - require(missingColumnCount == GpuDataWriterWithRowLineage.lineageColumnNames.length, - s"Expected ${GpuDataWriterWithRowLineage.lineageColumnNames.length} row lineage " + + require(missingColumnCount == lineageColumnNames.length, + s"Expected ${lineageColumnNames.length} row lineage " + s"columns but record is missing $missingColumnCount columns") + require(dataSparkType.takeRight(missingColumnCount).map(_.name).toSeq == lineageColumnNames, + "Expected row lineage columns at the end of the write schema") + if (metadata != null) { + require(metadata.numRows() == record.numRows(), + s"Metadata row count ${metadata.numRows()} does not match record row count " + + s"${record.numRows()}") + } val lineageColumns = closeOnExcept(new Array[ColumnVector](missingColumnCount)) { columns => - lineageColumnOrdinals.zipWithIndex.foreach { - case (ordinal, index) => - columns(index) = metadata.column(ordinal).asInstanceOf[GpuColumnVector].incRefCount() + lineageColumnNames.zipWithIndex.foreach { case (name, index) => + columns(index) = if (metadata == null) { + // Newly inserted rows inherit both lineage values from the Iceberg commit. + GpuColumnVector.fromNull(record.numRows(), LongType) + } else { + val column = metadata.column(metadataSchema.fieldIndex(name)) + .asInstanceOf[GpuColumnVector] + if (reinsertMask == null) { + column.incRefCount() + } else { + // INSERT rows inherit lineage even when their metadata projection has values. + withResource(GpuScalar.from(null, LongType)) { nullValue => + GpuColumnVector.from( + reinsertMask.ifElse(column.getBase, nullValue), LongType) + } + } + } } columns } - withResource(new ColumnarBatch(lineageColumns, metadata.numRows())) { lineage => - write(GpuColumnVector.combineColumns(record, lineage)) + withResource(new ColumnarBatch(lineageColumns, record.numRows())) { lineage => + GpuColumnVector.combineColumns(record, lineage) } } } } -object GpuDataWriterWithRowLineage { - val lineageColumnNames: Seq[String] = Seq("_row_id", "_last_updated_sequence_number") -} - class GpuUnpartitionedDataWriter( val fileWriterFactory: GpuSparkFileWriterFactory, val fileFactory: OutputFileFactory, diff --git a/sql-plugin-api/src/main/java/com/nvidia/spark/rapids/GpuDeltaBatchWriter.java b/sql-plugin-api/src/main/java/com/nvidia/spark/rapids/GpuDeltaBatchWriter.java new file mode 100644 index 00000000000..658014f4be0 --- /dev/null +++ b/sql-plugin-api/src/main/java/com/nvidia/spark/rapids/GpuDeltaBatchWriter.java @@ -0,0 +1,37 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed 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 com.nvidia.spark.rapids; + +import ai.rapids.cudf.ColumnVector; + +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** GPU batch operations independent of Spark's version-specific delta writer interface. */ +public interface GpuDeltaBatchWriter { + /** Reinserts rows with optional metadata, consuming both input batches. */ + void reinsert(ColumnarBatch metadata, ColumnarBatch rows); + + /** + * Writes INSERT and REINSERT rows in their input order, consuming all three inputs. + * Metadata applies only to rows whose reinsertMask value is true. Keeping both operations + * together preserves partition ordering required by clustered data writers. + */ + void insertAndReinsert( + ColumnarBatch metadata, + ColumnarBatch rows, + ColumnVector reinsertMask); +} diff --git a/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/DeltaInsertFilter.scala b/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/DeltaInsertFilter.scala index 792713989db..5b0868fc135 100644 --- a/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/DeltaInsertFilter.scala +++ b/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/DeltaInsertFilter.scala @@ -35,6 +35,8 @@ import org.apache.spark.sql.catalyst.util.RowDeltaUtils.INSERT_OPERATION import org.apache.spark.sql.vectorized.ColumnarBatch object DeltaInsertFilter { + val reinsertOperation: Option[Int] = None + def filterInsertRows(batch: ColumnarBatch): CudfColumnVector = { withResource(CudfScalar.fromInt(INSERT_OPERATION)) { s => batch.column(0).asInstanceOf[GpuColumnVector].getBase.equalTo(s) diff --git a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala index e219f8df0aa..ca82f17979a 100644 --- a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala +++ b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/GpuMergeRowsExec.scala @@ -325,8 +325,19 @@ object GpuMergeRowsExec { condition.columnarEval(batch) } - def applyOutputs(batch: ColumnarBatch): Seq[ColumnarBatch] = { - outputs.map(output => GpuProjectExec.project(batch, output)) + def applyOutputs( + batch: ColumnarBatch, + outputDataTypes: Array[DataType]): Seq[ColumnarBatch] = { + outputs.map { output => + // Spark InternalRow merge actions can have different widths, as with Iceberg row lineage. + // cuDF concatenation requires equal schemas, so materialize omitted trailing fields. + val paddedOutput = if (output.length == outputDataTypes.length) { + output + } else { + output ++ outputDataTypes.drop(output.length).map(GpuLiteral(null, _)) + } + GpuProjectExec.project(batch, paddedOutput) + } } override def nullable: Boolean = false @@ -394,6 +405,16 @@ class GpuMergeBatchIterator( import GpuMergeRowsExec.MergeRowMetrics + private val outputDataTypes: Array[DataType] = { + val instructionOutputs = (matchedInstructionExecs ++ notMatchedInstructionExecs ++ + notMatchedBySourceInstructionExecs).flatMap(_.outputs) + if (instructionOutputs.isEmpty) { + Array.empty[DataType] + } else { + instructionOutputs.maxBy(_.length).map(_.dataType).toArray + } + } + // Skip staging/publish work before Spark 4.1 where WriteSummary is unused. private val writeSummaryEnabled: Boolean = GpuMergeRowMetricsShims.writeSummaryEnabled // Reused across batches/attempts; reset() clears counts before each retry attempt. @@ -518,7 +539,7 @@ class GpuMergeBatchIterator( if (writeSummaryEnabled) { attemptMetrics.record(instructionExec, filtered.numRows(), sourcePresent) } - outputs ++= instructionExec.applyOutputs(filtered) + outputs ++= instructionExec.applyOutputs(filtered, outputDataTypes) .map(SpillableColumnarBatch .apply(_, SpillPriorities.ACTIVE_ON_DECK_PRIORITY)) } diff --git a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala index 18ae40b9db0..127e1b16c7d 100644 --- a/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala +++ b/sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala @@ -42,8 +42,8 @@ package org.apache.spark.sql.execution.datasources.v2 import scala.util.control.NonFatal -import ai.rapids.cudf.{ColumnVector => CudfColumnVector, Scalar => CudfScalar} -import com.nvidia.spark.rapids.{GpuColumnarToRowExec, GpuColumnVector, GpuDeltaWrite, GpuExec, GpuMetric, GpuWrite} +import ai.rapids.cudf.{ColumnVector => CudfColumnVector, Scalar => CudfScalar, Table => CudfTable} +import com.nvidia.spark.rapids.{GpuColumnarToRowExec, GpuColumnVector, GpuDeltaBatchWriter, GpuDeltaWrite, GpuExec, GpuMetric, GpuWrite} import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} import com.nvidia.spark.rapids.RmmRapidsRetryIterator.withRetryNoSplit import com.nvidia.spark.rapids.shims.DeltaInsertFilter @@ -60,7 +60,7 @@ import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{ExplainMode, QueryExecution, SparkPlan, SparkPlanInfo} import org.apache.spark.sql.execution.{SQLExecution, UnaryExecNode} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec -import org.apache.spark.sql.execution.datasources.v2.GpuDelteWritingSparkTask.filterByOperation +import org.apache.spark.sql.execution.datasources.v2.GpuDelteWritingSparkTask.{filterByOperation, writeInserts} import org.apache.spark.sql.execution.metric.{CustomMetrics, SQLMetric, SQLMetrics} import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.rapids.execution.TrampolineUtil @@ -440,7 +440,8 @@ object GpuDataWritingSparkTask extends GpuWritingSparkTask[DataWriter[ColumnarBa * Applies projections to extract row data and metadata before writing. */ case class GpuDeltaWritingSparkTask( - projs: WriteDeltaProjections) extends GpuWritingSparkTask[DeltaWriter[ColumnarBatch]] { + projs: WriteDeltaProjections) + extends GpuWritingSparkTask[DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter] { private lazy val rowProjection = projs.rowProjection .map(GpuProjectingColumnarBatch(_)) @@ -452,7 +453,9 @@ case class GpuDeltaWritingSparkTask( private lazy val rowIdProjection = GpuProjectingColumnarBatch(projs.rowIdProjection) private lazy val rowIdDataTypes = rowIdProjection.schema.fields.map(_.dataType) - override protected def write(writer: DeltaWriter[ColumnarBatch], batch: ColumnarBatch): Unit = { + override protected def write( + writer: DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter, + batch: ColumnarBatch): Unit = { withRetryNoSplit(batch) { _ => val deleteFilter = filterByOperation(batch, DELETE_OPERATION) withResource(deleteFilter) { _ => @@ -486,17 +489,7 @@ case class GpuDeltaWritingSparkTask( } } - val insertFilter = DeltaInsertFilter.filterInsertRows(batch) - withResource(insertFilter) { _ => - withResource(rowProjection.project(batch)) { rows => - val filteredRows = GpuColumnVector.filter(rows, rowDataTypes, insertFilter) - if (filteredRows.numRows() > 0) { - writer.insert(filteredRows) - } else { - filteredRows.close() - } - } - } + writeInserts(writer, batch, rowProjection, None) } } } @@ -507,7 +500,8 @@ case class GpuDeltaWritingSparkTask( * Applies both row and metadata projections before writing. */ case class GpuDeltaWithMetadataWritingSparkTask( - projs: WriteDeltaProjections) extends GpuWritingSparkTask[DeltaWriter[ColumnarBatch]] { + projs: WriteDeltaProjections) + extends GpuWritingSparkTask[DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter] { private lazy val rowProjection = projs.rowProjection .map(GpuProjectingColumnarBatch(_)) @@ -526,7 +520,9 @@ case class GpuDeltaWithMetadataWritingSparkTask( .map(_.schema.fields.map(f => f.dataType)) .orNull - override protected def write(writer: DeltaWriter[ColumnarBatch], batch: ColumnarBatch): Unit = { + override protected def write( + writer: DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter, + batch: ColumnarBatch): Unit = { withRetryNoSplit(batch) { _ => if (metadataProjection != null) { val deleteFilter = filterByOperation(batch, DELETE_OPERATION) @@ -577,23 +573,101 @@ case class GpuDeltaWithMetadataWritingSparkTask( } if (rowProjection != null) { - val insertFilter = DeltaInsertFilter.filterInsertRows(batch) - withResource(insertFilter) { _ => - withResource(rowProjection.project(batch)) { rows => - val filterRows = GpuColumnVector.filter(rows, rowDataTypes, insertFilter) - if (filterRows.numRows() > 0) { - writer.insert(filterRows) - } else { - filterRows.close() + writeInserts(writer, batch, rowProjection, Option(metadataProjection)) + } + } + } +} + +object GpuDelteWritingSparkTask { + /** + * Writes INSERT and REINSERT rows from a mixed-operation batch in their original order. + * The operation is in the first column. A combined INSERT-or-REINSERT filter selects both the + * projected data rows and optional metadata; applying that same filter to the REINSERT flags + * produces a reinsertMask aligned with the selected rows. The writer receives all three inputs + * in one insertAndReinsert call. Metadata values are preserved here for the writer to interpret. + * + * For example, on Spark 4, with other columns omitted: + * {{{ + * batch: + * operation partition id metadata._row_id + * DELETE A 9 109 + * REINSERT A 1 101 + * INSERT A 2 999 + * REINSERT B 3 103 + * INSERT B 4 888 + * + * insertFilter = [false, false, true, false, true] + * reinsertFilter = [false, true, false, true, false] + * dataFilter = [false, true, true, true, true] + * + * rows: + * partition id + * A 1 + * A 2 + * B 3 + * B 4 + * + * metadata._row_id = [101, 999, 103, 888] + * reinsertMask = [true, false, true, false] + * writer.insertAndReinsert(metadata, rows, reinsertMask) + * }}} + * Partition order remains A, A, B, B. Writing all REINSERTs before all INSERTs would instead + * produce A, B, A, B, causing clustered writers to revisit a closed partition. For Iceberg, the + * downstream appendLineage uses the mask to produce row IDs [101, null, 103, null]; this method + * does not clear the INSERT metadata itself. + * + * Spark 3 has no REINSERT operation; only INSERT rows are selected and passed to writer.insert. + * No writer call is made when no rows are selected. This method borrows batch and transfers + * ownership of the selected rows, metadata and reinsertMask to the writer. + */ + private[v2] def writeInserts( + writer: DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter, + batch: ColumnarBatch, + rowProjection: GpuProjectingColumnarBatch, + metadataProjection: Option[GpuProjectingColumnarBatch]): Unit = { + val reinsertFilter = DeltaInsertFilter.reinsertOperation + .map(filterByOperation(batch, _)).orNull + withResource(reinsertFilter) { _ => + val dataFilter = withResource(DeltaInsertFilter.filterInsertRows(batch)) { insertFilter => + if (reinsertFilter == null) { + insertFilter.incRefCount() + } else { + insertFilter.or(reinsertFilter) + } + } + withResource(dataFilter) { _ => + val rows = withResource(rowProjection.project(batch)) { projected => + GpuColumnVector.filter(projected, rowProjection.schema.map(_.dataType).toArray, + dataFilter) + } + if (rows.numRows() == 0) { + withResource(rows) { _ => () } + } else if (reinsertFilter == null) { + writer.insert(rows) + } else { + // Keep INSERT and REINSERT rows together so partition ordering is not disturbed. + val metadata = closeOnExcept(rows) { _ => + metadataProjection.map { projection => + withResource(projection.project(batch)) { projected => + GpuColumnVector.filter(projected, projection.schema.map(_.dataType).toArray, + dataFilter) + } + }.orNull + } + val reinsertMask = closeOnExcept(Seq(metadata, rows)) { _ => + withResource(new CudfTable(reinsertFilter)) { masks => + withResource(masks.filter(dataFilter)) { filtered => + filtered.getColumn(0).incRefCount() + } } } + writer.insertAndReinsert(metadata, rows, reinsertMask) } } } } -} -object GpuDelteWritingSparkTask { private[v2] def filterByOperation(batch: ColumnarBatch, op: Int): CudfColumnVector = { withResource(CudfScalar.fromInt(op)) { cudfOp => batch.column(0).asInstanceOf[GpuColumnVector].getBase.equalTo(cudfOp) diff --git a/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/DeltaInsertFilter.scala b/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/DeltaInsertFilter.scala index 6ca958a6f05..a080d83419c 100644 --- a/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/DeltaInsertFilter.scala +++ b/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/DeltaInsertFilter.scala @@ -36,18 +36,11 @@ import org.apache.spark.sql.catalyst.util.RowDeltaUtils.{INSERT_OPERATION, REINS import org.apache.spark.sql.vectorized.ColumnarBatch object DeltaInsertFilter { + val reinsertOperation: Option[Int] = Some(REINSERT_OPERATION) + def filterInsertRows(batch: ColumnarBatch): CudfColumnVector = { - val opCol = batch.column(0).asInstanceOf[GpuColumnVector].getBase - val insertMatch = withResource(CudfScalar.fromInt(INSERT_OPERATION)) { s => - opCol.equalTo(s) - } - withResource(insertMatch) { _ => - val reinsertMatch = withResource(CudfScalar.fromInt(REINSERT_OPERATION)) { s => - opCol.equalTo(s) - } - withResource(reinsertMatch) { _ => - insertMatch.or(reinsertMatch) - } + withResource(CudfScalar.fromInt(INSERT_OPERATION)) { s => + batch.column(0).asInstanceOf[GpuColumnVector].getBase.equalTo(s) } } } diff --git a/tests/src/test/spark350/scala/org/apache/iceberg/spark/source/GpuRowLineageWriterSuite.scala b/tests/src/test/spark350/scala/org/apache/iceberg/spark/source/GpuRowLineageWriterSuite.scala new file mode 100644 index 00000000000..4e26ef063ae --- /dev/null +++ b/tests/src/test/spark350/scala/org/apache/iceberg/spark/source/GpuRowLineageWriterSuite.scala @@ -0,0 +1,146 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed 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. + */ + +/*** spark-rapids-shim-json-lines +{"spark": "350"} +{"spark": "351"} +{"spark": "352"} +{"spark": "353"} +{"spark": "354"} +{"spark": "355"} +{"spark": "356"} +{"spark": "357"} +{"spark": "358"} +{"spark": "359"} +spark-rapids-shim-json-lines ***/ +package org.apache.iceberg.spark.source + +import ai.rapids.cudf.{ColumnVector => CudfColumnVector, Table} +import com.nvidia.spark.rapids.{GpuColumnVector, RmmSparkRetrySuiteBase} +import com.nvidia.spark.rapids.Arm.withResource + +import org.apache.spark.sql.types.{DataType, LongType, StructType} +import org.apache.spark.sql.vectorized.ColumnarBatch + +class GpuRowLineageWriterSuite extends RmmSparkRetrySuiteBase { + private val writeSchema = new StructType().add("id", LongType) + .add("_row_id", LongType).add("_last_updated_sequence_number", LongType) + private val metadataSchema = new StructType().add("_last_updated_sequence_number", LongType) + .add("_spec_id", LongType).add("_row_id", LongType) + + private def values(batch: ColumnarBatch, ordinal: Int): Seq[Option[Long]] = { + withResource(batch.column(ordinal).asInstanceOf[GpuColumnVector].copyToHost()) { column => + (0 until batch.numRows()).map { row => + if (column.isNullAt(row)) None else Some(column.getLong(row)) + } + } + } + + private def dataBatch(): ColumnarBatch = { + withResource(new Table.TestBuilder().column(Long.box(1L), 2L).build()) { table => + GpuColumnVector.from(table, Array[DataType](LongType)) + } + } + + test("reinsert preserves row IDs and sequence nulls using metadata column names") { + withResource(dataBatch()) { record => + withResource(new Table.TestBuilder() + .column(Long.box(7L), null.asInstanceOf[java.lang.Long]) + .column(Long.box(0L), 0L).column(Long.box(101L), 102L).build()) { table => + withResource(GpuColumnVector.from(table, Array.fill[DataType](3)(LongType))) { metadata => + withResource(GpuDataWriterWithRowLineage.appendLineage( + record, metadata, writeSchema, metadataSchema)) { physical => + assert(physical.numCols() == 3) + assert(values(physical, 0) == Seq(Some(1L), Some(2L))) + assert(values(physical, 1) == Seq(Some(101L), Some(102L))) + assert(values(physical, 2) == Seq(Some(7L), None)) + } + // Appending lineage borrows its inputs; closing the result must not release them. + assert(values(metadata, 2) == Seq(Some(101L), Some(102L))) + } + } + assert(values(record, 0) == Seq(Some(1L), Some(2L))) + } + } + + test("mixed inserts and reinserts preserve order and inherit lineage only for inserts") { + withResource(dataBatch()) { record => + withResource(new Table.TestBuilder() + .column(Long.box(7L), 8L).column(Long.box(0L), 0L) + .column(Long.box(101L), 102L).build()) { table => + withResource(GpuColumnVector.from(table, Array.fill[DataType](3)(LongType))) { metadata => + withResource(CudfColumnVector.fromBooleans(true, false)) { reinsertMask => + withResource(GpuDataWriterWithRowLineage.appendLineage( + record, metadata, writeSchema, metadataSchema, reinsertMask)) { physical => + assert(values(physical, 0) == Seq(Some(1L), Some(2L))) + assert(values(physical, 1) == Seq(Some(101L), None)) + assert(values(physical, 2) == Seq(Some(7L), None)) + } + } + assert(values(metadata, 2) == Seq(Some(101L), Some(102L))) + assert(values(metadata, 0) == Seq(Some(7L), Some(8L))) + } + } + assert(values(record, 0) == Seq(Some(1L), Some(2L))) + } + } + + test("insert appends inheritable null lineage") { + withResource(dataBatch()) { record => + withResource(GpuDataWriterWithRowLineage.appendLineage( + record, null, writeSchema, null)) { physical => + assert(values(physical, 0) == Seq(Some(1L), Some(2L))) + assert(values(physical, 1) == Seq(None, None)) + assert(values(physical, 2) == Seq(None, None)) + } + } + } + + test("complete Spark 3 rows and v2 rows do not acquire extra columns") { + withResource(new Table.TestBuilder() + .column(Long.box(1L), 2L) + .column(Long.box(101L), 102L) + .column(Long.box(7L), 8L).build()) { table => + withResource(GpuColumnVector.from(table, Array.fill[DataType](3)(LongType))) { record => + withResource(GpuDataWriterWithRowLineage.appendLineage( + record, null, writeSchema, null)) { physical => + assert(physical.numCols() == 3) + assert(values(physical, 1) == Seq(Some(101L), Some(102L))) + } + } + } + withResource(dataBatch()) { record => + withResource(GpuDataWriterWithRowLineage.appendLineage( + record, null, new StructType().add("id", LongType), null)) { physical => + assert(physical.numCols() == 1) + assert(values(physical, 0) == Seq(Some(1L), Some(2L))) + } + } + } + + test("mismatched metadata row counts fail without consuming the record") { + withResource(dataBatch()) { record => + withResource(new ColumnarBatch(Array.empty, 1)) { metadata => + val error = intercept[IllegalArgumentException] { + GpuDataWriterWithRowLineage.appendLineage( + record, metadata, writeSchema, metadataSchema) + } + assert(error.getMessage.contains("Metadata row count")) + } + assert(values(record, 0) == Seq(Some(1L), Some(2L))) + } + } +} diff --git a/tests/src/test/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuDeltaWritingSparkTaskSuite.scala b/tests/src/test/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuDeltaWritingSparkTaskSuite.scala new file mode 100644 index 00000000000..dfc8495cd24 --- /dev/null +++ b/tests/src/test/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuDeltaWritingSparkTaskSuite.scala @@ -0,0 +1,162 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed 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. + */ +/*** spark-rapids-shim-json-lines +{"spark": "400"} +{"spark": "401"} +{"spark": "402"} +{"spark": "403"} +{"spark": "404"} +{"spark": "411"} +{"spark": "412"} +{"spark": "413"} +{"spark": "420"} +{"spark": "500"} +spark-rapids-shim-json-lines ***/ + +package org.apache.spark.sql.execution.datasources.v2 + +import ai.rapids.cudf.{ColumnVector => CudfColumnVector, Table} +import com.nvidia.spark.rapids.{GpuColumnVector, GpuDeltaBatchWriter, RmmSparkRetrySuiteBase} +import com.nvidia.spark.rapids.Arm.withResource + +import org.apache.spark.sql.catalyst.ProjectingInternalRow +import org.apache.spark.sql.catalyst.util.RowDeltaUtils.{DELETE_OPERATION, INSERT_OPERATION, REINSERT_OPERATION, UPDATE_OPERATION} +import org.apache.spark.sql.catalyst.util.WriteDeltaProjections +import org.apache.spark.sql.connector.write.{DeltaWriter, WriterCommitMessage} +import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StructType} +import org.apache.spark.sql.vectorized.ColumnarBatch + +class GpuDeltaWritingSparkTaskSuite extends RmmSparkRetrySuiteBase { + private val rowSchema = new StructType().add("value", LongType) + private val metadataSchema = new StructType().add("_row_id", LongType) + .add("_last_updated_sequence_number", LongType) + + private def longValues(batch: ColumnarBatch, ordinal: Int): Seq[Option[Long]] = { + withResource(batch.column(ordinal).asInstanceOf[GpuColumnVector].copyToHost()) { column => + (0 until batch.numRows()).map { row => + if (column.isNullAt(row)) None else Some(column.getLong(row)) + } + } + } + + private class RecordingWriter extends DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter { + var dataOrder = Seq.empty[Option[Long]] + var inserted = Seq.empty[Option[Long]] + var reinserted = Seq.empty[Option[Long]] + var lineage = Seq.empty[Seq[Option[Long]]] + var deletedRows = 0 + var updatedRows = 0 + + override def insert(row: ColumnarBatch): Unit = withResource(row) { _ => + dataOrder ++= longValues(row, 0) + inserted ++= longValues(row, 0) + } + + override def reinsert(metadata: ColumnarBatch, row: ColumnarBatch): Unit = { + withResource(Seq(metadata, row)) { _ => + dataOrder ++= longValues(row, 0) + reinserted ++= longValues(row, 0) + if (metadata != null) { + lineage = Seq(longValues(metadata, 0), longValues(metadata, 1)) + } + } + } + + override def insertAndReinsert( + metadata: ColumnarBatch, + row: ColumnarBatch, + reinsertMask: CudfColumnVector): Unit = { + withResource(Seq(metadata, row, reinsertMask)) { _ => + val flags = withResource(reinsertMask.copyToHost()) { host => + (0 until row.numRows()).map(index => host.getBoolean(index)) + } + val values = longValues(row, 0) + dataOrder ++= values + inserted ++= values.zip(flags).collect { case (value, false) => value } + reinserted ++= values.zip(flags).collect { case (value, true) => value } + if (metadata != null) { + lineage = Seq(0, 1).map { ordinal => + longValues(metadata, ordinal).zip(flags).collect { case (value, true) => value } + } + } + } + } + + override def delete(metadata: ColumnarBatch, rowId: ColumnarBatch): Unit = { + withResource(Seq(metadata, rowId)) { _ => deletedRows += rowId.numRows() } + } + + override def update( + metadata: ColumnarBatch, + rowId: ColumnarBatch, + row: ColumnarBatch): Unit = { + withResource(Seq(metadata, rowId, row)) { _ => updatedRows += row.numRows() } + } + + override def commit(): WriterCommitMessage = null + override def abort(): Unit = () + override def close(): Unit = () + } + + private class MetadataTask(projections: WriteDeltaProjections) + extends GpuDeltaWithMetadataWritingSparkTask(projections) { + def writeBatch( + writer: DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter, + batch: ColumnarBatch): Unit = write(writer, batch) + } + + private class PlainTask(projections: WriteDeltaProjections) + extends GpuDeltaWritingSparkTask(projections) { + def writeBatch( + writer: DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter, + batch: ColumnarBatch): Unit = write(writer, batch) + } + + Seq(false, true).foreach { withMetadata => + test(s"preserve insert and reinsert order with metadata=$withMetadata") { + val projections = WriteDeltaProjections( + Some(ProjectingInternalRow(rowSchema, Seq(1))), + ProjectingInternalRow(rowSchema, Seq(1)), + if (withMetadata) Some(ProjectingInternalRow(metadataSchema, Seq(2, 3))) else None) + val writer = new RecordingWriter + val batch = withResource(new Table.TestBuilder() + .column(Int.box(INSERT_OPERATION), REINSERT_OPERATION, DELETE_OPERATION, + INSERT_OPERATION, REINSERT_OPERATION, UPDATE_OPERATION) + .column(Long.box(10L), 20L, 30L, 40L, 50L, 60L) + .column(Long.box(110L), 220L, 330L, 440L, 550L, 660L) + .column(Long.box(5L), null.asInstanceOf[java.lang.Long], 7L, 8L, 9L, 11L) + .build()) { table => + GpuColumnVector.from(table, Array[DataType](IntegerType, LongType, LongType, LongType)) + } + // Writing tasks own and close the input batch. + if (withMetadata) { + new MetadataTask(projections).writeBatch(writer, batch) + } else { + new PlainTask(projections).writeBatch(writer, batch) + } + assert(writer.dataOrder == Seq(Some(10L), Some(20L), Some(40L), Some(50L))) + assert(writer.inserted == Seq(Some(10L), Some(40L))) + assert(writer.reinserted == Seq(Some(20L), Some(50L))) + assert(writer.deletedRows == 1) + assert(writer.updatedRows == 1) + if (withMetadata) { + assert(writer.lineage == Seq(Seq(Some(220L), Some(550L)), Seq(None, Some(9L)))) + } else { + assert(writer.lineage.isEmpty) + } + } + } +} From 71a95da74f753912d7f0a825361a62b63be4928a Mon Sep 17 00:00:00 2001 From: Chong Gao Date: Fri, 11 Sep 2026 09:41:58 +0800 Subject: [PATCH 22/22] Exercise Iceberg v3 across existing test matrices Extend existing DML, read, SQL UI and view tests to format v3 while preserving parameter interactions, IDs and marks. Retain tests with specific delete-file-format contracts and verify the current v3 MOR fallback, including the insert-only MERGE rewrite to GPU append. Addresses review: https://github.com/NVIDIA/cudf-spark/pull/15866#discussion_r3975097345 Signed-off-by: Chong Gao --- .../src/main/python/iceberg/__init__.py | 75 ++++- .../python/iceberg/iceberg_append_test.py | 98 +++--- .../main/python/iceberg/iceberg_ctas_test.py | 84 ++--- .../python/iceberg/iceberg_delete_test.py | 162 ++++++---- .../main/python/iceberg/iceberg_merge_test.py | 206 ++++++++---- .../iceberg/iceberg_overwrite_dynamic_test.py | 75 +++-- .../iceberg/iceberg_overwrite_static_test.py | 95 +++--- .../main/python/iceberg/iceberg_rtas_test.py | 91 +++--- .../src/main/python/iceberg/iceberg_test.py | 295 +++++++++++------- .../python/iceberg/iceberg_update_test.py | 190 +++++++---- .../main/python/iceberg/iceberg_view_test.py | 12 +- .../iceberg/iceberg_write_sql_ui_test.py | 17 +- 12 files changed, 895 insertions(+), 505 deletions(-) diff --git a/integration_tests/src/main/python/iceberg/__init__.py b/integration_tests/src/main/python/iceberg/__init__.py index 62c19708c37..2603d06a0d4 100644 --- a/integration_tests/src/main/python/iceberg/__init__.py +++ b/integration_tests/src/main/python/iceberg/__init__.py @@ -25,7 +25,7 @@ from conftest import is_iceberg_rest_catalog, spark_jvm from data_gen import * -from spark_session import is_iceberg_supported_spark, with_cpu_session +from spark_session import is_iceberg_supported_spark, with_cpu_session, with_gpu_session iceberg_unsupported_mark = pytest.mark.skipif( not is_iceberg_supported_spark(), @@ -43,6 +43,73 @@ ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON = \ "Iceberg row lineage inheritance requires iceberg 1.10.0 or later" +# Keep format versions crossed with every existing semantic test parameter. +_v3_unsupported_mark = pytest.mark.skipif( + not supports_iceberg_v3, reason=ICEBERG_V3_UNSUPPORTED_REASON) +_v3_mor_fallback_mark = pytest.mark.allow_non_gpu_conditional( + True, "WriteDeltaExec", "MergeRowsExec", "BatchScanExec", "ColumnarToRowExec", + "ShuffleExchangeExec", "SortExec", "ProjectExec") +# Row-level COW planning can retain the CPU scan used for file pruning. +_v3_cow_scan_mark = pytest.mark.allow_non_gpu_conditional(True, "BatchScanExec") +iceberg_format_versions = [ + pytest.param("2", id="v2"), + pytest.param("3", marks=_v3_unsupported_mark, id="v3")] +iceberg_read_format_versions = [pytest.param("1", id="v1"), *iceberg_format_versions] +iceberg_read_enabled_conf = {"spark.rapids.sql.format.iceberg.v3.enabled": "true"} +iceberg_mor_format_versions = [ + pytest.param("2", id="v2"), + pytest.param("3", marks=[_v3_unsupported_mark, _v3_mor_fallback_mark], id="v3")] +iceberg_cow_format_versions = [ + pytest.param("2", id="v2"), + pytest.param("3", marks=[_v3_unsupported_mark, _v3_cow_scan_mark], id="v3")] + + +def with_iceberg_format_versions(parameters): + """Cross existing parameter rows with v2/v3, preserving collection-time marks.""" + result = [] + for parameter in parameters: + if hasattr(parameter, "values"): + values, marks, case_id = parameter.values, list(parameter.marks), parameter.id + else: + values = parameter if isinstance(parameter, tuple) else (parameter,) + marks, case_id = [], None + for version in ("2", "3"): + version_marks = list(marks) + if version == "3": + version_marks.append(_v3_unsupported_mark) + if "merge-on-read" in values: + version_marks.append(_v3_mor_fallback_mark) + elif "copy-on-write" in values: + version_marks.append(_v3_cow_scan_mark) + result.append(pytest.param( + version, *values, marks=version_marks, + id=f"v{version}-{case_id}" if case_id is not None else None)) + return result + + +def with_iceberg_dml_session(func, format_version, mode, conf): + """Require the current Puffin fallback specifically for v3 MOR writes.""" + if format_version != "3" or mode != "merge-on-read": + return with_gpu_session(func, conf=conf) + callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback + callback.startCapture() + try: + result = with_gpu_session(func, conf=conf) + plans = callback.getResultsWithTimeout(10000) + # Spark rewrites insert-only MERGE to append even on a MOR table. + assert any(callback.didFallBack(plan, "WriteDeltaExec") or + callback.contains(plan, "GpuAppendDataExec") for plan in plans), \ + "Expected v3 MOR fallback or GPU append for insert-only MERGE:\n" + \ + "\n".join(str(plan) for plan in plans) + return result + finally: + callback.endCapture() + + +def iceberg_table_properties_sql(format_version): + props = _build_tblprops({'format-version': format_version}) + return "TBLPROPERTIES (" + ", ".join(f"'{k}' = '{v}'" for k, v in props.items()) + ")" + # iceberg supported types iceberg_table_gen = MappingProxyType({ '_c0': byte_gen, '_c1': short_gen, '_c2': int_gen, @@ -204,6 +271,7 @@ def row_lineage_df( "spark.sql.parquet.int96RebaseModeInWrite": "CORRECTED", "spark.rapids.sql.format.iceberg.enabled": "true", "spark.rapids.sql.format.iceberg.write.enabled": "true", + "spark.rapids.sql.format.iceberg.v3.enabled": "true", # WriteDeltaExec is disabled by default as it's experimental, but we need it enabled # for merge-on-read (MOR) DML operations (UPDATE/DELETE/MERGE with write.*.mode='merge-on-read') "spark.rapids.sql.exec.WriteDeltaExec": "true", @@ -387,9 +455,12 @@ def assert_iceberg_files_use_codec(spark: SparkSession, table_name: str, expecte def create_iceberg_table(table_name: str, partition_col_sql: Optional[str] = None, table_prop: Optional[Dict[str, str]] = None, - df_gen: Optional[Callable[[SparkSession], DataFrame]] = None) -> str: + df_gen: Optional[Callable[[SparkSession], DataFrame]] = None, + format_version: Optional[str] = None) -> str: if table_prop is None: table_prop = {'format-version':'1'} + if format_version is not None: + table_prop = {**table_prop, 'format-version': format_version} table_prop = _build_tblprops(table_prop) if df_gen is None: diff --git a/integration_tests/src/main/python/iceberg/iceberg_append_test.py b/integration_tests/src/main/python/iceberg/iceberg_append_test.py index e352ae39d27..4c0716713d0 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_append_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_append_test.py @@ -19,12 +19,11 @@ assert_gpu_fallback_write_sql from conftest import is_iceberg_remote_catalog from data_gen import gen_df, copy_and_update -from iceberg import create_iceberg_table, \ - iceberg_base_table_cols, iceberg_gens_list, get_full_table_name, \ - iceberg_full_gens_list, \ - iceberg_write_enabled_conf, iceberg_unsupported_mark, _build_tblprops, \ - full_coverage_partition_transforms, assert_iceberg_files_use_codec, \ - supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON +from iceberg import ( + iceberg_format_versions, create_iceberg_table, iceberg_base_table_cols, iceberg_gens_list, + get_full_table_name, iceberg_full_gens_list, iceberg_write_enabled_conf, + iceberg_unsupported_mark, _build_tblprops, full_coverage_partition_transforms, + assert_iceberg_files_use_codec, supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON) from marks import iceberg, ignore_order, allow_non_gpu, datagen_overrides from spark_session import with_gpu_session, with_cpu_session @@ -58,8 +57,9 @@ def insert_data(spark, table_name: str): @iceberg @ignore_order(local=True) -def test_insert_into_unpartitioned_table(spark_tmp_table_factory): - table_prop = {"format-version": "2"} +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_unpartitioned_table(format_version, spark_tmp_table_factory): + table_prop = {"format-version": format_version} do_test_insert_into_table_sql( spark_tmp_table_factory, @@ -91,7 +91,7 @@ def insert_data(spark, table_name): lambda spark, table_name: spark.sql(f"SELECT * FROM {table_name}"), base_table_name, ["AppendDataExec"], - conf=iceberg_write_enabled_conf) + conf=copy_and_update(iceberg_write_enabled_conf, {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) @iceberg @@ -99,14 +99,15 @@ def insert_data(spark, table_name): @allow_non_gpu('AppendDataExec') @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("partition_table", [True, False], ids=lambda x: f"partition_table={x}") -def test_insert_into_unpartitioned_table_values(spark_tmp_table_factory, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_unpartitioned_table_values(format_version, spark_tmp_table_factory, partition_table): base_table_name = get_full_table_name(spark_tmp_table_factory) cpu_table_name = f"{base_table_name}_cpu" gpu_table_name = f"{base_table_name}_gpu" def create_table(spark, table_name: str): - props = _build_tblprops({"format-version": "2"}) + props = _build_tblprops({"format-version": format_version}) props_sql = ", ".join(f"'{k}' = '{v}'" for k, v in props.items()) sql = f"CREATE TABLE {table_name} (id int, name string) USING ICEBERG " if partition_table: @@ -139,14 +140,15 @@ def insert_data(spark, table_name: str): @allow_non_gpu('LocalTableScanExec', 'ShuffleExchangeExec') @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("partition_table", [True, False], ids=lambda x: f"partition_table={x}") -def test_insert_into_table_values_aqe(spark_tmp_table_factory, partition_table): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_table_values_aqe(format_version, spark_tmp_table_factory, partition_table): """Regression test for GPU V2 writes with AQE and a CPU VALUES input plan.""" base_table_name = get_full_table_name(spark_tmp_table_factory) cpu_table_name = f"{base_table_name}_cpu" gpu_table_name = f"{base_table_name}_gpu" def create_table(spark, table_name: str): - props = _build_tblprops({"format-version": "2"}) + props = _build_tblprops({"format-version": format_version}) props_sql = ", ".join(f"'{k}' = '{v}'" for k, v in props.items()) sql = f"CREATE TABLE {table_name} (id int, name string) USING ICEBERG " if partition_table: @@ -176,8 +178,9 @@ def insert_data(spark, table_name: str): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_into_unpartitioned_table_all_cols(spark_tmp_table_factory): - table_prop = {"format-version": "2"} +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_unpartitioned_table_all_cols(format_version, spark_tmp_table_factory): + table_prop = {"format-version": format_version} cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] gen_list = list(zip(cols, iceberg_full_gens_list)) @@ -208,10 +211,10 @@ def insert_data(spark, table_name: str): def _do_test_insert_into_partitioned_table(spark_tmp_table_factory, partition_col_sql, - table_prop=None): + table_prop=None, format_version="2"): """Helper function for partitioned table insert tests.""" if table_prop is None: - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} def create_table_and_set_write_order(table_name: str): create_iceberg_table( @@ -233,9 +236,13 @@ def create_table_and_set_write_order(table_name: str): @pytest.mark.parametrize("partition_col_sql", [ pytest.param("year(_c9)", id="year(timestamp_col)"), ]) -def test_insert_into_partitioned_table(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_partitioned_table(format_version, spark_tmp_table_factory, partition_col_sql): """Basic partition test - runs for all catalogs including remote.""" - _do_test_insert_into_partitioned_table(spark_tmp_table_factory, partition_col_sql) + _do_test_insert_into_partitioned_table( + spark_tmp_table_factory, + partition_col_sql, + format_version=format_version) @iceberg @@ -243,20 +250,25 @@ def test_insert_into_partitioned_table(spark_tmp_table_factory, partition_col_sq @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("partition_col_sql", full_coverage_partition_transforms) -def test_insert_into_partitioned_table_full_coverage(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_partitioned_table_full_coverage(format_version, spark_tmp_table_factory, partition_col_sql): """Partition-transform coverage anchor: this is the single test that exercises the partition writer against every transform in full_coverage_partition_transforms. Every other DML op's _full_coverage test picks only a few distinct transforms from this list (via ctas_/rtas_/overwrite_*_/delete_/update_/merge_-prefixed constants in iceberg/__init__.py), so the partition writer's coverage of all 26 transforms lives here. Skipped for remote catalogs.""" - _do_test_insert_into_partitioned_table(spark_tmp_table_factory, partition_col_sql) + _do_test_insert_into_partitioned_table( + spark_tmp_table_factory, + partition_col_sql, + format_version=format_version) @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_into_partitioned_table_all_cols(spark_tmp_table_factory): - table_prop = {"format-version": "2"} +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_partitioned_table_all_cols(format_version, spark_tmp_table_factory): + table_prop = {"format-version": format_version} cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] gen_list = list(zip(cols, iceberg_full_gens_list)) @@ -301,9 +313,10 @@ def insert_data(spark, table_name: str): @allow_non_gpu('AppendDataExec', 'ShuffleExchangeExec', 'ProjectExec') @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("file_format", ["orc", "avro"], ids=lambda x: f"file_format={x}") -def test_insert_into_table_unsupported_file_format_fallback( +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_table_unsupported_file_format_fallback(format_version, spark_tmp_table_factory, file_format): - table_prop = {"format-version": "2", + table_prop = {"format-version": format_version, "write.format.default": file_format} def insert_data(spark, table_name: str): @@ -328,10 +341,11 @@ def insert_data(spark, table_name: str): pytest.param("truncate(3, contact.email)", id="truncate_nested_struct_field"), pytest.param("contact.email", id="identity_nested_struct_field"), ], ) -def test_insert_into_table_nested_partition_source_fallback( +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_table_nested_partition_source_fallback(format_version, spark_tmp_table_factory, partition_col_sql): table_name = get_full_table_name(spark_tmp_table_factory) - table_prop = _build_tblprops({"format-version": "2"}) + table_prop = _build_tblprops({"format-version": format_version}) props_sql = ", ".join(f"'{k}' = '{v}'" for k, v in table_prop.items()) def create_table(spark): @@ -362,9 +376,10 @@ def insert_data(spark): @pytest.mark.parametrize("conf_key", ["spark.rapids.sql.format.iceberg.enabled", "spark.rapids.sql.format.iceberg.write.enabled"], ids=lambda x: f"{x}=False") -def test_insert_into_iceberg_table_fallback_when_conf_disabled( +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_iceberg_table_fallback_when_conf_disabled(format_version, spark_tmp_table_factory, conf_key): - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} def insert_data(spark, table_name: str): df = gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) @@ -387,11 +402,12 @@ def insert_data(spark, table_name: str): pytest.param(None, id="unpartitioned"), pytest.param("year(_c9)", id="year_partition"), ]) -def test_insert_into_aqe(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_aqe(format_version, spark_tmp_table_factory, partition_col_sql): """ Test INSERT INTO with AQE enabled. """ - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} # Configuration with AQE enabled conf = copy_and_update(iceberg_write_enabled_conf, { @@ -429,7 +445,8 @@ def insert_data(spark, table_name: str): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_after_drop_partition_field(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_after_drop_partition_field(format_version, spark_tmp_table_factory): """Test INSERT on table after dropping a partition field (void transform). When a partition field is dropped, Iceberg creates a 'void transform' - @@ -440,7 +457,7 @@ def test_insert_after_drop_partition_field(spark_tmp_table_factory): cpu_table_name = f"{base_table_name}_cpu" gpu_table_name = f"{base_table_name}_gpu" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} # Use two partition columns so after dropping one, we still have at least one partition_col_sql = "bucket(8, _c2), bucket(8, _c3)" @@ -486,11 +503,12 @@ def insert_data(spark, table_name): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_into_partitioned_table_fanout_enabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_partitioned_table_fanout_enabled(format_version, spark_tmp_table_factory): # Use bucket(2, ...) to keep partition count low and avoid OOM from Iceberg's FanoutDataWriter. _do_test_insert_into_partitioned_table( spark_tmp_table_factory, "bucket(2, _c9)", - table_prop={"format-version": "2", "write.spark.fanout.enabled": "true"}) + table_prop={"format-version": format_version, "write.spark.fanout.enabled": "true"}, format_version=format_version) # Regression for https://github.com/NVIDIA/spark-rapids/issues/14905 — the GPU writer @@ -510,11 +528,12 @@ def test_insert_into_partitioned_table_fanout_enabled(spark_tmp_table_factory): (None, "zstd"), # No override: Iceberg's default (zstd) must win on GPU too. ("zstd", "zstd"), ("uncompressed", "uncompressed")]) -def test_insert_into_table_honors_iceberg_compression_codec( +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_table_honors_iceberg_compression_codec(format_version, spark_tmp_table_factory, table_codec, expected_codec): table_name = get_full_table_name(spark_tmp_table_factory) - extra_props = {"format-version": "2"} + extra_props = {"format-version": format_version} if table_codec is not None: extra_props["write.parquet.compression-codec"] = table_codec @@ -551,11 +570,12 @@ def create_table(spark): @allow_non_gpu('AppendDataExec', 'ShuffleExchangeExec', 'ProjectExec') @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("codec", ["gzip", "lz4"]) -def test_insert_into_table_falls_back_on_unsupported_codec(spark_tmp_table_factory, codec): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_into_table_falls_back_on_unsupported_codec(format_version, spark_tmp_table_factory, codec): table_name = get_full_table_name(spark_tmp_table_factory) create_iceberg_table( table_name, - table_prop={"format-version": "2", "write.parquet.compression-codec": codec}) + table_prop={"format-version": format_version, "write.parquet.compression-codec": codec}) def insert_data(spark): df = gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) diff --git a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py index d615db9b30f..e66893a7d1b 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py @@ -22,16 +22,12 @@ assert_gpu_fallback_collect) from conftest import is_iceberg_remote_catalog, spark_jvm from data_gen import gen_df, copy_and_update, RepeatSeqGen -from iceberg import (create_iceberg_table, - iceberg_base_table_cols, - iceberg_gens_list, iceberg_full_gens_list, - get_full_table_name, iceberg_write_enabled_conf, - iceberg_unsupported_mark, _build_tblprops, - ctas_partition_transforms, supports_iceberg_v3, - ICEBERG_V3_UNSUPPORTED_REASON, - supports_iceberg_row_lineage_inheritance, - ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, - row_lineage_df) +from iceberg import ( + iceberg_format_versions, create_iceberg_table, iceberg_base_table_cols, iceberg_gens_list, + iceberg_full_gens_list, get_full_table_name, iceberg_write_enabled_conf, + iceberg_unsupported_mark, _build_tblprops, ctas_partition_transforms, supports_iceberg_v3, + ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, + ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, row_lineage_df) from marks import (iceberg, ignore_order, allow_non_gpu, allow_non_gpu_conditional, datagen_overrides) from spark_session import with_gpu_session, with_cpu_session, is_spark_400_or_later @@ -116,9 +112,10 @@ def read_table(spark, table_name): @iceberg @ignore_order(local=True) -def test_ctas_unpartitioned_table(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_unpartitioned_table(format_version, spark_tmp_table_factory): table_prop = { - "format-version": "2" + "format-version": format_version } df_gen = lambda spark: gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) @@ -143,7 +140,7 @@ def run_ctas(spark): assert_gpu_fallback_collect( run_ctas, "AtomicCreateTableAsSelectExec", - conf=iceberg_write_enabled_conf) + conf=copy_and_update(iceberg_write_enabled_conf, {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) @iceberg @@ -156,7 +153,7 @@ def test_ctas_v3_row_lineage(spark_tmp_table_factory): "spark.rapids.sql.format.iceberg.v3.enabled": "true" }) - def assert_gpu_ctas(plan): + def assert_gpu_ctas(_cpu_plan, plan): callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback ctas_plan = callback.extractExecutedPlan(plan) callback.assertContains(ctas_plan, "GpuAtomicCreateTableAsSelectExec") @@ -206,14 +203,17 @@ def run_ctas(spark): lambda sp: gen_df(sp, list(zip(iceberg_base_table_cols, iceberg_gens_list))), table_prop) - assert_gpu_and_cpu_are_equal_collect(run_ctas, conf=conf) + assert_gpu_and_cpu_are_equal_collect( + run_ctas, + conf=copy_and_update(conf, {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) -def _do_test_ctas_partitioned_table(spark_tmp_table_factory, partition_col_sql, table_prop=None): +def _do_test_ctas_partitioned_table( + spark_tmp_table_factory, partition_col_sql, table_prop=None, format_version="2"): """Helper function for partitioned table CTAS tests.""" if table_prop is None: table_prop = { - "format-version": "2" + "format-version": format_version } df_gen = lambda spark: gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) @@ -230,9 +230,10 @@ def _do_test_ctas_partitioned_table(spark_tmp_table_factory, partition_col_sql, @pytest.mark.parametrize("partition_col_sql", [ pytest.param("year(_c9)", id="year(timestamp_col)"), ]) -def test_ctas_partitioned_table(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_partitioned_table(format_version, spark_tmp_table_factory, partition_col_sql): """Basic partition test - runs for all catalogs including remote.""" - _do_test_ctas_partitioned_table(spark_tmp_table_factory, partition_col_sql) + _do_test_ctas_partitioned_table(spark_tmp_table_factory, partition_col_sql, format_version=format_version) @iceberg @@ -241,11 +242,12 @@ def test_ctas_partitioned_table(spark_tmp_table_factory, partition_col_sql): @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("partition_col_sql", ctas_partition_transforms) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_ctas_partitioned_table_full_coverage(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_partitioned_table_full_coverage(format_version, spark_tmp_table_factory, partition_col_sql): """Sanity-check CTAS against a few partition transforms distinct from those picked by other DML ops. The 26-transform partition-writer coverage anchor lives in iceberg_append_test.py::test_insert_into_partitioned_table_full_coverage.""" - _do_test_ctas_partitioned_table(spark_tmp_table_factory, partition_col_sql) + _do_test_ctas_partitioned_table(spark_tmp_table_factory, partition_col_sql, format_version=format_version) @iceberg @@ -253,10 +255,11 @@ def test_ctas_partitioned_table_full_coverage(spark_tmp_table_factory, partition @allow_non_gpu('AtomicCreateTableAsSelectExec', 'AppendDataExec') @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("file_format", ["orc", "avro"], ids=lambda x: f"file_format={x}") -def test_ctas_unsupported_file_format_fallback(spark_tmp_table_factory, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_unsupported_file_format_fallback(format_version, spark_tmp_table_factory, file_format): table_prop = { - "format-version": "2", + "format-version": format_version, "write.format.default": file_format } @@ -280,10 +283,11 @@ def run_ctas(spark): @pytest.mark.parametrize("conf_key", ["spark.rapids.sql.format.iceberg.enabled", "spark.rapids.sql.format.iceberg.write.enabled"], ids=lambda x: f"{x}=False") -def test_ctas_fallback_when_conf_disabled(spark_tmp_table_factory, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_fallback_when_conf_disabled(format_version, spark_tmp_table_factory, conf_key): table_prop = { - "format-version": "2" + "format-version": format_version } def run_ctas(spark): @@ -304,9 +308,10 @@ def run_ctas(spark): @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("gen_list", _BINARY_CTAS_GEN_LISTS, ids=["binary", "array_binary"]) -def test_ctas_unpartitioned_table_binary_types(spark_tmp_table_factory, gen_list): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_unpartitioned_table_binary_types(format_version, spark_tmp_table_factory, gen_list): table_prop = { - "format-version": "2" + "format-version": format_version } df_gen = lambda spark: gen_df(spark, gen_list, length=32) @@ -318,9 +323,10 @@ def test_ctas_unpartitioned_table_binary_types(spark_tmp_table_factory, gen_list @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_ctas_unpartitioned_table_all_cols(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_unpartitioned_table_all_cols(format_version, spark_tmp_table_factory): table_prop = { - "format-version": "2" + "format-version": format_version } cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] @@ -334,9 +340,10 @@ def test_ctas_unpartitioned_table_all_cols(spark_tmp_table_factory): @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_ctas_partitioned_table_all_cols(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_partitioned_table_all_cols(format_version, spark_tmp_table_factory): table_prop = { - "format-version": "2" + "format-version": format_version } cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] @@ -355,9 +362,10 @@ def test_ctas_partitioned_table_all_cols(spark_tmp_table_factory): @pytest.mark.parametrize("partition_table", [True, False], ids=lambda x: f"partition_table={x}") @allow_non_gpu('AtomicCreateTableAsSelectExec', 'AppendDataExec', 'ShuffleExchangeExec', 'SortExec', 'ProjectExec') @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_ctas_from_values(spark_tmp_table_factory, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_from_values(format_version, spark_tmp_table_factory, partition_table): - table_prop = _build_tblprops({"format-version": "2"}) + table_prop = _build_tblprops({"format-version": format_version}) base_name = get_full_table_name(spark_tmp_table_factory) gpu_table = f"{base_name}_gpu" @@ -391,7 +399,8 @@ def execute_ctas_from_values(spark, target_table: str): pytest.param("year(_c9)", id="triple_datetime_transforms"), ]) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_ctas_aqe(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_aqe(format_version, spark_tmp_table_factory, partition_col_sql): """ Test CTAS with multiple partition transforms on the same column with AQE enabled. @@ -404,7 +413,7 @@ def test_ctas_aqe(spark_tmp_table_factory, partition_col_sql): - GpuShuffleCoalesceExec ends up as a child of GpuRowToColumnarExec """ table_prop = { - "format-version": "2", + "format-version": format_version, } df_gen = lambda spark: gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) @@ -426,9 +435,10 @@ def test_ctas_aqe(spark_tmp_table_factory, partition_col_sql): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_ctas_partitioned_table_fanout_enabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_ctas_partitioned_table_fanout_enabled(format_version, spark_tmp_table_factory): # Use bucket(2, ...) to keep partition count low and avoid OOM from Iceberg's FanoutDataWriter. _do_test_ctas_partitioned_table( spark_tmp_table_factory, "bucket(2, _c9)", - table_prop={"format-version": "2", "write.spark.fanout.enabled": "true"}) + table_prop={"format-version": format_version, "write.spark.fanout.enabled": "true"}, format_version=format_version) diff --git a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py index 8fcd97de6c5..797490a5aed 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_delete_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_delete_test.py @@ -18,14 +18,14 @@ assert_gpu_fallback_write_sql from conftest import is_iceberg_remote_catalog from data_gen import * -from iceberg import (create_iceberg_table, get_full_table_name, iceberg_write_enabled_conf, - iceberg_base_table_cols, iceberg_gens_list, iceberg_nested_write_gens_list, - iceberg_unsupported_mark, delete_partition_transforms_distributed, - _build_tblprops, assert_iceberg_files_use_codec, - supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON, - supports_iceberg_row_lineage_inheritance, - ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, - row_lineage_df, rapids_reader_types) +from iceberg import ( + iceberg_cow_format_versions, iceberg_mor_format_versions, with_iceberg_dml_session, + with_iceberg_format_versions, create_iceberg_table, get_full_table_name, + iceberg_write_enabled_conf, iceberg_base_table_cols, iceberg_gens_list, + iceberg_nested_write_gens_list, iceberg_unsupported_mark, + delete_partition_transforms_distributed, _build_tblprops, assert_iceberg_files_use_codec, + supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, + ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, row_lineage_df, rapids_reader_types) from marks import allow_non_gpu, allow_non_gpu_conditional, iceberg, ignore_order, datagen_overrides from spark_session import is_spark_35x, is_spark_400_or_later, with_cpu_session, with_gpu_session @@ -51,11 +51,11 @@ def create_iceberg_table_with_data(table_name: str, data_gen_func=None, table_properties=None, delete_mode='copy-on-write', - write_order=None): + write_order=None, format_version="2"): """Helper function to create and populate an Iceberg table for DELETE tests.""" # Default to copy-on-write mode, but allow override for merge-on-read tests base_props = { - 'format-version': '2', + 'format-version': format_version, 'write.delete.mode': delete_mode } if table_properties: @@ -82,7 +82,7 @@ def insert_data(spark): def do_delete_test(spark_tmp_table_factory, delete_sql_func, data_gen_func=None, partition_col_sql=None, table_properties=None, delete_mode='copy-on-write', conf=iceberg_delete_cow_enabled_conf, - read_func=None, write_order=None): + read_func=None, write_order=None, format_version="2"): """ Helper function to test DELETE operations by comparing CPU and GPU results. @@ -103,15 +103,15 @@ def do_delete_test(spark_tmp_table_factory, delete_sql_func, data_gen_func=None, # Create identical tables for CPU and GPU create_iceberg_table_with_data(cpu_table_name, partition_col_sql, - data_gen_func, table_properties, delete_mode, write_order) + data_gen_func, table_properties, delete_mode, write_order, format_version=format_version) create_iceberg_table_with_data(gpu_table_name, partition_col_sql, - data_gen_func, table_properties, delete_mode, write_order) + data_gen_func, table_properties, delete_mode, write_order, format_version=format_version) # Execute DELETE on GPU def do_gpu_delete(spark): delete_sql_func(spark, gpu_table_name) - with_gpu_session(do_gpu_delete, conf=conf) + with_iceberg_dml_session(do_gpu_delete, format_version, delete_mode, conf=conf) # Execute DELETE on CPU def do_cpu_delete(spark): @@ -133,14 +133,17 @@ def do_cpu_delete(spark): @iceberg @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) -@pytest.mark.parametrize('delete_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,delete_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_delete_unpartitioned_table(spark_tmp_table_factory, delete_mode): +def test_iceberg_delete_unpartitioned_table(format_version, spark_tmp_table_factory, delete_mode): """Test DELETE on unpartitioned table with both copy-on-write and merge-on-read modes""" do_delete_test( spark_tmp_table_factory, lambda spark, table: spark.sql(f"DELETE FROM {table} WHERE _c2 % 3 = 0"), - delete_mode=delete_mode + delete_mode=delete_mode, + format_version=format_version ) @@ -181,7 +184,9 @@ def delete_data(spark, table_name): lambda spark, table_name: spark.sql(f"SELECT * FROM {table_name}"), base_table_name, [fallback_exec], - conf=iceberg_delete_cow_enabled_conf) + conf=copy_and_update( + iceberg_delete_cow_enabled_conf, + {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) @iceberg @@ -233,14 +238,16 @@ def test_iceberg_v3_row_lineage_gpu_delete_leading_rows(spark_tmp_table_factory) write_order="id") -def _do_test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col_sql, delete_mode, table_properties=None): +def _do_test_iceberg_delete_partitioned_table( + spark_tmp_table_factory, partition_col_sql, delete_mode, table_properties=None, format_version="2"): """Helper function for partitioned table DELETE tests.""" do_delete_test( spark_tmp_table_factory, lambda spark, table: spark.sql(f"DELETE FROM {table} WHERE _c2 % 3 = 0"), partition_col_sql=partition_col_sql, table_properties=table_properties, - delete_mode=delete_mode + delete_mode=delete_mode, + format_version=format_version ) @@ -252,11 +259,17 @@ def _do_test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition @pytest.mark.parametrize("partition_col_sql", [ pytest.param("year(_c9)", id="year(timestamp_col)"), ]) -@pytest.mark.parametrize('delete_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,delete_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col_sql, delete_mode): +def test_iceberg_delete_partitioned_table(format_version, spark_tmp_table_factory, partition_col_sql, delete_mode): """Basic partition test - runs for all catalogs including remote.""" - _do_test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col_sql, delete_mode) + _do_test_iceberg_delete_partitioned_table( + spark_tmp_table_factory, + partition_col_sql, + delete_mode, + format_version=format_version) # This requires reading of _partition field, which is a struct @@ -265,29 +278,39 @@ def test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize("partition_col_sql,delete_mode", delete_partition_transforms_distributed) +@pytest.mark.parametrize( + 'format_version,partition_col_sql,delete_mode', + with_iceberg_format_versions(delete_partition_transforms_distributed)) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_delete_partitioned_table_full_coverage(spark_tmp_table_factory, partition_col_sql, delete_mode): +def test_iceberg_delete_partitioned_table_full_coverage( + format_version, spark_tmp_table_factory, partition_col_sql, delete_mode): """Sanity-check DELETE across the two write modes against partition transforms distinct from those picked by other DML ops. The 26-transform partition-writer coverage anchor lives in iceberg_append_test.py::test_insert_into_partitioned_table_full_coverage.""" - _do_test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col_sql, delete_mode) + _do_test_iceberg_delete_partitioned_table( + spark_tmp_table_factory, + partition_col_sql, + delete_mode, + format_version=format_version) @iceberg @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('delete_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,delete_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_delete_with_complex_predicate(spark_tmp_table_factory, delete_mode): +def test_iceberg_delete_with_complex_predicate(format_version, spark_tmp_table_factory, delete_mode): """Test DELETE with complex predicate""" do_delete_test( spark_tmp_table_factory, lambda spark, table: spark.sql( f"DELETE FROM {table} WHERE _c2 > 100 AND _c3 < 50 OR _c1 IS NULL" ), - delete_mode=delete_mode + delete_mode=delete_mode, + format_version=format_version ) @@ -296,20 +319,21 @@ def test_iceberg_delete_with_complex_predicate(spark_tmp_table_factory, delete_m @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('delete_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,delete_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) +])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_delete_fallback_write_disabled(spark_tmp_table_factory, delete_mode, fallback_exec): +def test_iceberg_delete_fallback_write_disabled( + format_version, spark_tmp_table_factory, delete_mode, fallback_exec): """Test DELETE falls back when Iceberg write is disabled (both modes)""" base_table_name = get_full_table_name(spark_tmp_table_factory) # Phase 1: Initialize tables with data (separate for CPU and GPU) cpu_table_name = f'{base_table_name}_cpu' gpu_table_name = f'{base_table_name}_gpu' - create_iceberg_table_with_data(cpu_table_name, delete_mode=delete_mode) - create_iceberg_table_with_data(gpu_table_name, delete_mode=delete_mode) + create_iceberg_table_with_data(cpu_table_name, delete_mode=delete_mode, format_version=format_version) + create_iceberg_table_with_data(gpu_table_name, delete_mode=delete_mode, format_version=format_version) # Phase 2: DELETE operation (to be tested with fallback) def write_func(spark, table_name): @@ -334,13 +358,14 @@ def read_func(spark, table_name): @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('delete_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,delete_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) +])) @pytest.mark.parametrize("file_format", ["orc", "avro"], ids=lambda x: f"file_format={x}") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_delete_fallback_unsupported_file_format(spark_tmp_table_factory, delete_mode, fallback_exec, file_format): +def test_iceberg_delete_fallback_unsupported_file_format( + format_version, spark_tmp_table_factory, delete_mode, fallback_exec, file_format): """Test DELETE falls back with unsupported file formats (ORC, Avro) for both modes""" base_table_name = get_full_table_name(spark_tmp_table_factory) @@ -351,7 +376,7 @@ def data_gen(spark): def init_table(table_name): # Step 1: Create table with parquet as default write format table_props = { - 'format-version': '2', + 'format-version': format_version, 'write.delete.mode': delete_mode, 'write.format.default': 'parquet' } @@ -400,9 +425,11 @@ def read_func(spark, table_name): @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('delete_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,delete_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_delete_nested_types(spark_tmp_table_factory, delete_mode): +def test_iceberg_delete_nested_types(format_version, spark_tmp_table_factory, delete_mode): """Test DELETE with supported nested types.""" cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_nested_write_gens_list)] data_gen_func = lambda spark: gen_df(spark, list(zip(cols, iceberg_nested_write_gens_list))) @@ -411,7 +438,8 @@ def test_iceberg_delete_nested_types(spark_tmp_table_factory, delete_mode): spark_tmp_table_factory, lambda spark, table: spark.sql(f"DELETE FROM {table} WHERE _c0 % 3 = 0"), data_gen_func=data_gen_func, - delete_mode=delete_mode + delete_mode=delete_mode, + format_version=format_version ) @@ -420,19 +448,20 @@ def test_iceberg_delete_nested_types(spark_tmp_table_factory, delete_mode): @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('delete_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,delete_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) -def test_iceberg_delete_fallback_iceberg_disabled(spark_tmp_table_factory, delete_mode, fallback_exec): +])) +def test_iceberg_delete_fallback_iceberg_disabled( + format_version, spark_tmp_table_factory, delete_mode, fallback_exec): """Test DELETE falls back when Iceberg is completely disabled (both modes)""" base_table_name = get_full_table_name(spark_tmp_table_factory) # Phase 1: Initialize tables with data (separate for CPU and GPU) cpu_table_name = f'{base_table_name}_cpu' gpu_table_name = f'{base_table_name}_gpu' - create_iceberg_table_with_data(cpu_table_name, delete_mode=delete_mode) - create_iceberg_table_with_data(gpu_table_name, delete_mode=delete_mode) + create_iceberg_table_with_data(cpu_table_name, delete_mode=delete_mode, format_version=format_version) + create_iceberg_table_with_data(gpu_table_name, delete_mode=delete_mode, format_version=format_version) # Phase 2: DELETE operation (to be tested with fallback) def write_func(spark, table_name): @@ -457,7 +486,8 @@ def read_func(spark, table_name): @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_iceberg_delete_mor_fallback_writedelta_disabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_mor_format_versions) +def test_iceberg_delete_mor_fallback_writedelta_disabled(format_version, spark_tmp_table_factory): """Test merge-on-read DELETE falls back when WriteDeltaExec is disabled This test verifies that when WriteDeltaExec is explicitly disabled (it's disabled by default @@ -468,8 +498,8 @@ def test_iceberg_delete_mor_fallback_writedelta_disabled(spark_tmp_table_factory # Phase 1: Initialize tables with data (separate for CPU and GPU) cpu_table_name = f'{base_table_name}_cpu' gpu_table_name = f'{base_table_name}_gpu' - create_iceberg_table_with_data(cpu_table_name, delete_mode='merge-on-read') - create_iceberg_table_with_data(gpu_table_name, delete_mode='merge-on-read') + create_iceberg_table_with_data(cpu_table_name, delete_mode='merge-on-read', format_version=format_version) + create_iceberg_table_with_data(gpu_table_name, delete_mode='merge-on-read', format_version=format_version) # Phase 2: DELETE operation (to be tested with fallback) def write_func(spark, table_name): @@ -494,18 +524,20 @@ def read_func(spark, table_name): @iceberg @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) -@pytest.mark.parametrize('update_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,update_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @pytest.mark.parametrize("partition_col_sql", [ pytest.param(None, id="unpartitioned"), pytest.param("year(_c9)", id="year_partition"), ]) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_delete_aqe(spark_tmp_table_factory, update_mode, partition_col_sql): +def test_delete_aqe(format_version, spark_tmp_table_factory, update_mode, partition_col_sql): """ Test DELETE with AQE enabled. """ table_prop = { - 'format-version': '2', + 'format-version': format_version, 'write.delete.mode': update_mode } @@ -529,7 +561,11 @@ def initialize_table(table_name): def delete_from_table(spark, table_name): spark.sql(f"DELETE FROM {table_name} WHERE _c2 % 3 = 0") - with_gpu_session(lambda spark: delete_from_table(spark, gpu_table), conf=conf) + with_iceberg_dml_session( + lambda spark: delete_from_table(spark, gpu_table), + format_version, + update_mode, + conf=conf) with_cpu_session(lambda spark: delete_from_table(spark, cpu_table), conf=conf) cpu_data = with_cpu_session(lambda spark: spark.table(cpu_table).collect()) @@ -541,9 +577,11 @@ def delete_from_table(spark, table_name): @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.datagen_overrides(seed=DELETE_TEST_SEED, reason=DELETE_TEST_SEED_OVERRIDE_REASON) -@pytest.mark.parametrize('delete_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,delete_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_delete_after_drop_partition_field(spark_tmp_table_factory, delete_mode): +def test_iceberg_delete_after_drop_partition_field(format_version, spark_tmp_table_factory, delete_mode): """Test DELETE on table after dropping a partition field (void transform). When a partition field is dropped, Iceberg creates a 'void transform' - @@ -559,9 +597,9 @@ def test_iceberg_delete_after_drop_partition_field(spark_tmp_table_factory, dele # Create partitioned tables with data create_iceberg_table_with_data(cpu_table_name, partition_col_sql=partition_col_sql, - delete_mode=delete_mode) + delete_mode=delete_mode, format_version=format_version) create_iceberg_table_with_data(gpu_table_name, partition_col_sql=partition_col_sql, - delete_mode=delete_mode) + delete_mode=delete_mode, format_version=format_version) # Drop one partition field on both tables (creates void transform) def drop_partition_field(spark, table_name): @@ -574,8 +612,9 @@ def drop_partition_field(spark, table_name): def do_delete(spark, table_name): spark.sql(f"DELETE FROM {table_name} WHERE _c2 % 3 = 0") - with_gpu_session(lambda spark: do_delete(spark, gpu_table_name), - conf=iceberg_delete_cow_enabled_conf) + with_iceberg_dml_session( + lambda spark: do_delete(spark, gpu_table_name), + format_version, delete_mode, conf=iceberg_delete_cow_enabled_conf) with_cpu_session(lambda spark: do_delete(spark, cpu_table_name)) # Compare results @@ -588,13 +627,14 @@ def do_delete(spark, table_name): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_iceberg_delete_partitioned_table_fanout_enabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_cow_format_versions) +def test_iceberg_delete_partitioned_table_fanout_enabled(format_version, spark_tmp_table_factory): # Use bucket(2, ...) to keep partition count low and avoid OOM from Iceberg's FanoutDataWriter. _do_test_iceberg_delete_partitioned_table( spark_tmp_table_factory, partition_col_sql="bucket(2, _c9)", delete_mode='copy-on-write', - table_properties={"write.spark.fanout.enabled": "true"}) + table_properties={"write.spark.fanout.enabled": "true"}, format_version=format_version) # Regression for the GpuSparkPositionDeltaWrite branch of issue #14905. diff --git a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py index 64e1396047c..09113e6c9b4 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_merge_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_merge_test.py @@ -18,13 +18,13 @@ assert_gpu_fallback_write_sql from conftest import is_iceberg_remote_catalog from data_gen import * -from iceberg import (create_iceberg_table, get_full_table_name, iceberg_write_enabled_conf, - iceberg_base_table_cols, iceberg_gens_list, iceberg_nested_write_gens_list, - iceberg_unsupported_mark, merge_partition_transforms_distributed, - supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON, - supports_iceberg_row_lineage_inheritance, - ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, - row_lineage_df, rapids_reader_types) +from iceberg import ( + iceberg_cow_format_versions, iceberg_mor_format_versions, with_iceberg_dml_session, + with_iceberg_format_versions, create_iceberg_table, get_full_table_name, + iceberg_write_enabled_conf, iceberg_base_table_cols, iceberg_gens_list, + iceberg_nested_write_gens_list, iceberg_unsupported_mark, merge_partition_transforms_distributed, + supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, + ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, row_lineage_df, rapids_reader_types) from marks import allow_non_gpu, allow_non_gpu_conditional, iceberg, ignore_order, datagen_overrides from spark_session import is_spark_400_or_later, with_gpu_session, with_cpu_session @@ -44,7 +44,7 @@ def create_iceberg_table_with_merge_data( seed=None, merge_mode='copy-on-write', iceberg_base_table_cols=iceberg_base_table_cols, - iceberg_gens_list=iceberg_gens_list): + iceberg_gens_list=iceberg_gens_list, format_version="2"): """ Helper function to create and populate an Iceberg table for MERGE tests. @@ -58,7 +58,7 @@ def create_iceberg_table_with_merge_data( merge_mode: Merge mode - 'copy-on-write' or 'merge-on-read' """ base_props = { - 'format-version': '2', + 'format-version': format_version, 'write.merge.mode': merge_mode, # See https://github.com/NVIDIA/spark-rapids/issues/13698 'read.parquet.vectorization.enabled': 'false' @@ -99,12 +99,14 @@ def insert_data(spark): def _assert_gpu_and_cpu_merge_writes_are_equal( - cpu_table_name, gpu_table_name, merge_func, read_func, conf): + cpu_table_name, gpu_table_name, merge_func, read_func, conf, + format_version="2", merge_mode="copy-on-write"): def run_merge(spark, table_name): merge_func(spark, table_name) with_cpu_session(lambda spark: run_merge(spark, cpu_table_name), conf=conf) - with_gpu_session(lambda spark: run_merge(spark, gpu_table_name), conf=conf) + with_iceberg_dml_session( + lambda spark: run_merge(spark, gpu_table_name), format_version, merge_mode, conf=conf) cpu_data = with_cpu_session( lambda spark: read_func(spark, cpu_table_name).collect(), conf=conf) @@ -120,7 +122,7 @@ def do_merge_test( table_properties=None, merge_mode='copy-on-write', iceberg_base_table_cols=iceberg_base_table_cols, - iceberg_gens_list=iceberg_gens_list): + iceberg_gens_list=iceberg_gens_list, format_version="2"): """ Helper function to test MERGE operations by comparing CPU and GPU results. @@ -140,24 +142,47 @@ def do_merge_test( source_table = f"{base_table_name}_source" # Create identical target tables for CPU and GPU (using runtime seed) - create_iceberg_table_with_merge_data(cpu_target_table, partition_col_sql, table_properties, merge_mode=merge_mode, iceberg_base_table_cols=iceberg_base_table_cols, iceberg_gens_list=iceberg_gens_list) - create_iceberg_table_with_merge_data(gpu_target_table, partition_col_sql, table_properties, merge_mode=merge_mode, iceberg_base_table_cols=iceberg_base_table_cols, iceberg_gens_list=iceberg_gens_list) + create_iceberg_table_with_merge_data( + cpu_target_table, + partition_col_sql, + table_properties, + merge_mode=merge_mode, + iceberg_base_table_cols=iceberg_base_table_cols, + iceberg_gens_list=iceberg_gens_list, + format_version=format_version) + create_iceberg_table_with_merge_data( + gpu_target_table, + partition_col_sql, + table_properties, + merge_mode=merge_mode, + iceberg_base_table_cols=iceberg_base_table_cols, + iceberg_gens_list=iceberg_gens_list, + format_version=format_version) # Create source table with different seed and distinct keys to satisfy MERGE cardinality constraint # (each target row matches at most one source row) # Using a fixed different seed ensures source data differs from target data - create_iceberg_table_with_merge_data(source_table, partition_col_sql, table_properties, - ensure_distinct_key=True, seed=42, merge_mode=merge_mode, iceberg_base_table_cols=iceberg_base_table_cols, iceberg_gens_list=iceberg_gens_list) + create_iceberg_table_with_merge_data( + source_table, + partition_col_sql, + table_properties, + ensure_distinct_key=True, + seed=42, + merge_mode=merge_mode, + iceberg_base_table_cols=iceberg_base_table_cols, + iceberg_gens_list=iceberg_gens_list, + format_version=format_version) _assert_gpu_and_cpu_merge_writes_are_equal( cpu_target_table, gpu_target_table, lambda spark, target: merge_sql_func(spark, target, source_table), lambda spark, target: spark.table(target), - iceberg_merge_enabled_conf) + iceberg_merge_enabled_conf, format_version=format_version, merge_mode=merge_mode) -def _do_test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mode, table_properties=None): +def _do_test_iceberg_merge( + spark_tmp_table_factory, partition_col_sql, merge_mode, table_properties=None, format_version="2"): """Helper function for MERGE tests.""" merge_sql = """ MERGE INTO {target} t USING {source} s ON t._c0 = s._c0 @@ -169,7 +194,8 @@ def _do_test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mod lambda spark, target, source: spark.sql(merge_sql.format(target=target, source=source)), partition_col_sql=partition_col_sql, table_properties=table_properties, - merge_mode=merge_mode + merge_mode=merge_mode, + format_version=format_version ) @@ -177,15 +203,17 @@ def _do_test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mod @iceberg @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) -@pytest.mark.parametrize('merge_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,merge_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @pytest.mark.parametrize('partition_col_sql', [ None, pytest.param("year(_c9)", id="year(timestamp_col)"), ]) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mode): +def test_iceberg_merge(format_version, spark_tmp_table_factory, partition_col_sql, merge_mode): """Basic partition test - runs for all catalogs including remote.""" - _do_test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mode) + _do_test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mode, format_version=format_version) @allow_non_gpu( @@ -232,7 +260,7 @@ def merge_data(spark, target_table): lambda spark, table_name: spark.sql(f"SELECT * FROM {table_name}"), target_base_name, [fallback_exec], - conf=iceberg_merge_enabled_conf) + conf=copy_and_update(iceberg_merge_enabled_conf, {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) @iceberg @@ -320,21 +348,25 @@ def merge(spark, table): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize("partition_col_sql,merge_mode", merge_partition_transforms_distributed) +@pytest.mark.parametrize( + 'format_version,partition_col_sql,merge_mode', + with_iceberg_format_versions(merge_partition_transforms_distributed)) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_merge_full_coverage(spark_tmp_table_factory, partition_col_sql, merge_mode): +def test_iceberg_merge_full_coverage(format_version, spark_tmp_table_factory, partition_col_sql, merge_mode): """Sanity-check MERGE across the two write modes against partition transforms distinct from those picked by other DML ops. The 26-transform partition-writer coverage anchor lives in iceberg_append_test.py::test_insert_into_partitioned_table_full_coverage.""" - _do_test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mode) + _do_test_iceberg_merge(spark_tmp_table_factory, partition_col_sql, merge_mode, format_version=format_version) @allow_non_gpu("MergeRows$Keep", "MergeRows$Discard", "MergeRows$Split") @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('merge_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,merge_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @pytest.mark.parametrize('partition_col_sql', [ pytest.param(None, id="unpartitioned"), pytest.param("year(_c9)", id="year(timestamp_col)"), @@ -384,20 +416,22 @@ def test_iceberg_merge_full_coverage(spark_tmp_table_factory, partition_col_sql, id="conditional_not_matched_by_source"), ]) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_merge_additional_patterns(spark_tmp_table_factory, partition_col_sql, merge_sql, merge_mode): +def test_iceberg_merge_additional_patterns( + format_version, spark_tmp_table_factory, partition_col_sql, merge_sql, merge_mode): """Test additional MERGE patterns (conditional updates, deletes, not matched by source) on Iceberg tables.""" do_merge_test( spark_tmp_table_factory, lambda spark, target, source: spark.sql(merge_sql.format(target=target, source=source)), partition_col_sql=partition_col_sql, - merge_mode=merge_mode + merge_mode=merge_mode, + format_version=format_version ) @allow_non_gpu("MergeRows$Keep", "MergeRows$Discard", "MergeRows$Split") @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('merge_mode', ['copy-on-write']) +@pytest.mark.parametrize('format_version,merge_mode', with_iceberg_format_versions(['copy-on-write'])) @pytest.mark.parametrize('partition_col_sql', [pytest.param("year(_c9)", id="year(timestamp_col)")]) @pytest.mark.parametrize('merge_sql', [ pytest.param( @@ -410,24 +444,26 @@ def test_iceberg_merge_additional_patterns(spark_tmp_table_factory, partition_co id="multiple_matched_clauses"), ]) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_merge_additional_patterns_bug(spark_tmp_table_factory, partition_col_sql, merge_sql, merge_mode): +def test_iceberg_merge_additional_patterns_bug( + format_version, spark_tmp_table_factory, partition_col_sql, merge_sql, merge_mode): """Test additional MERGE patterns (conditional updates, deletes, not matched by source) on Iceberg tables.""" do_merge_test( spark_tmp_table_factory, lambda spark, target, source: spark.sql(merge_sql.format(target=target, source=source)), partition_col_sql=partition_col_sql, - merge_mode=merge_mode + merge_mode=merge_mode, + format_version=format_version ) @allow_non_gpu("ReplaceDataExec", "WriteDeltaExec", "MergeRowsExec", "BatchScanExec", "ColumnarToRowExec", "ShuffleExchangeExec", "SortExec", "ProjectExec") @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('merge_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,merge_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) -def test_iceberg_merge_fallback_write_disabled(spark_tmp_table_factory, merge_mode, fallback_exec): +])) +def test_iceberg_merge_fallback_write_disabled(format_version, spark_tmp_table_factory, merge_mode, fallback_exec): """Test MERGE falls back when Iceberg write is disabled""" base_table_name = get_full_table_name(spark_tmp_table_factory) @@ -436,10 +472,15 @@ def test_iceberg_merge_fallback_write_disabled(spark_tmp_table_factory, merge_mo gpu_target_table = f'{base_table_name}_target_gpu' source_table = f'{base_table_name}_source' - create_iceberg_table_with_merge_data(cpu_target_table, merge_mode=merge_mode) - create_iceberg_table_with_merge_data(gpu_target_table, merge_mode=merge_mode) + create_iceberg_table_with_merge_data(cpu_target_table, merge_mode=merge_mode, format_version=format_version) + create_iceberg_table_with_merge_data(gpu_target_table, merge_mode=merge_mode, format_version=format_version) # Source table needs distinct keys for MERGE cardinality constraint, with different seed - create_iceberg_table_with_merge_data(source_table, ensure_distinct_key=True, seed=42, merge_mode=merge_mode) + create_iceberg_table_with_merge_data( + source_table, + ensure_distinct_key=True, + seed=42, + merge_mode=merge_mode, + format_version=format_version) # Phase 2: MERGE operation (to be tested with fallback) def write_func(spark, target_table_name): @@ -470,12 +511,13 @@ def read_func(spark, table_name): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('merge_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,merge_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) +])) @pytest.mark.parametrize("file_format", ["orc", "avro"], ids=lambda x: f"file_format={x}") -def test_iceberg_merge_fallback_unsupported_file_format(spark_tmp_table_factory, file_format, merge_mode, fallback_exec): +def test_iceberg_merge_fallback_unsupported_file_format( + format_version, spark_tmp_table_factory, file_format, merge_mode, fallback_exec): """Test MERGE falls back with unsupported file formats (ORC, Avro)""" base_table_name = get_full_table_name(spark_tmp_table_factory) @@ -486,7 +528,7 @@ def data_gen(spark): def init_table(table_name, ensure_distinct_key=False): # Create table with parquet, insert data, then change format table_props = { - 'format-version': '2', + 'format-version': format_version, 'write.merge.mode': merge_mode, 'write.format.default': 'parquet' } @@ -559,9 +601,11 @@ def read_func(spark, table_name): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('merge_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,merge_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_merge_nested_types(spark_tmp_table_factory, merge_mode): +def test_iceberg_merge_nested_types(format_version, spark_tmp_table_factory, merge_mode): """Test MERGE on tables containing supported nested types.""" nested_cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_nested_write_gens_list)] merge_sql = """ @@ -574,7 +618,8 @@ def test_iceberg_merge_nested_types(spark_tmp_table_factory, merge_mode): lambda spark, target, source: spark.sql(merge_sql.format(target=target, source=source)), merge_mode=merge_mode, iceberg_base_table_cols=nested_cols, - iceberg_gens_list=iceberg_nested_write_gens_list + iceberg_gens_list=iceberg_nested_write_gens_list, + format_version=format_version ) @@ -582,11 +627,12 @@ def test_iceberg_merge_nested_types(spark_tmp_table_factory, merge_mode): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('merge_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,merge_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) -def test_iceberg_merge_fallback_iceberg_disabled(spark_tmp_table_factory, merge_mode, fallback_exec): +])) +def test_iceberg_merge_fallback_iceberg_disabled( + format_version, spark_tmp_table_factory, merge_mode, fallback_exec): """Test MERGE falls back when Iceberg is completely disabled""" base_table_name = get_full_table_name(spark_tmp_table_factory) @@ -594,10 +640,15 @@ def test_iceberg_merge_fallback_iceberg_disabled(spark_tmp_table_factory, merge_ gpu_target_table = f'{base_table_name}_target_gpu' source_table = f'{base_table_name}_source' - create_iceberg_table_with_merge_data(cpu_target_table, merge_mode=merge_mode) - create_iceberg_table_with_merge_data(gpu_target_table, merge_mode=merge_mode) + create_iceberg_table_with_merge_data(cpu_target_table, merge_mode=merge_mode, format_version=format_version) + create_iceberg_table_with_merge_data(gpu_target_table, merge_mode=merge_mode, format_version=format_version) # Source table needs distinct keys for MERGE cardinality constraint, with different seed - create_iceberg_table_with_merge_data(source_table, ensure_distinct_key=True, seed=42, merge_mode=merge_mode) + create_iceberg_table_with_merge_data( + source_table, + ensure_distinct_key=True, + seed=42, + merge_mode=merge_mode, + format_version=format_version) def write_func(spark, target_table_name): spark.sql(f""" @@ -625,7 +676,8 @@ def read_func(spark, table_name): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_iceberg_merge_mor_fallback_writedelta_disabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_mor_format_versions) +def test_iceberg_merge_mor_fallback_writedelta_disabled(format_version, spark_tmp_table_factory): """Test merge-on-read MERGE falls back when WriteDeltaExec is disabled This test verifies that when WriteDeltaExec is explicitly disabled (it's disabled by default @@ -637,10 +689,21 @@ def test_iceberg_merge_mor_fallback_writedelta_disabled(spark_tmp_table_factory) gpu_target_table = f'{base_table_name}_target_gpu' source_table = f'{base_table_name}_source' - create_iceberg_table_with_merge_data(cpu_target_table, merge_mode='merge-on-read') - create_iceberg_table_with_merge_data(gpu_target_table, merge_mode='merge-on-read') + create_iceberg_table_with_merge_data( + cpu_target_table, + merge_mode='merge-on-read', + format_version=format_version) + create_iceberg_table_with_merge_data( + gpu_target_table, + merge_mode='merge-on-read', + format_version=format_version) # Source table needs distinct keys for MERGE cardinality constraint, with different seed - create_iceberg_table_with_merge_data(source_table, ensure_distinct_key=True, seed=42, merge_mode='merge-on-read') + create_iceberg_table_with_merge_data( + source_table, + ensure_distinct_key=True, + seed=42, + merge_mode='merge-on-read', + format_version=format_version) def write_func(spark, target_table_name): spark.sql(f""" @@ -673,12 +736,13 @@ def read_func(spark, table_name): pytest.param("year(_c9)", id="year_partition"), ]) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_merge_aqe(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_cow_format_versions) +def test_merge_aqe(format_version, spark_tmp_table_factory, partition_col_sql): """ Test MERGE INTO with AQE enabled. """ table_prop = { - 'format-version': '2', + 'format-version': format_version, } # Configuration with AQE enabled @@ -709,7 +773,11 @@ def merge_table(spark, target_table): WHEN NOT MATCHED THEN INSERT * """) - with_gpu_session(lambda spark: merge_table(spark, gpu_target), conf=conf) + with_iceberg_dml_session( + lambda spark: merge_table(spark, gpu_target), + format_version, + "copy-on-write", + conf=conf) with_cpu_session(lambda spark: merge_table(spark, cpu_target), conf=conf) cpu_data = with_cpu_session(lambda spark: spark.table(cpu_target).collect()) @@ -720,9 +788,11 @@ def merge_table(spark, target_table): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('merge_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,merge_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_merge_after_drop_partition_field(spark_tmp_table_factory, merge_mode): +def test_iceberg_merge_after_drop_partition_field(format_version, spark_tmp_table_factory, merge_mode): """Test MERGE on table after dropping a partition field (void transform). When a partition field is dropped, Iceberg creates a 'void transform' - @@ -739,12 +809,12 @@ def test_iceberg_merge_after_drop_partition_field(spark_tmp_table_factory, merge # Create partitioned target tables with data - use same seed for both to ensure same data create_iceberg_table_with_merge_data(cpu_target_table, partition_col_sql=partition_col_sql, - merge_mode=merge_mode, seed=42) + merge_mode=merge_mode, seed=42, format_version=format_version) create_iceberg_table_with_merge_data(gpu_target_table, partition_col_sql=partition_col_sql, - merge_mode=merge_mode, seed=42) + merge_mode=merge_mode, seed=42, format_version=format_version) # Source table needs distinct keys for MERGE cardinality constraint, with different seed create_iceberg_table_with_merge_data(source_table, partition_col_sql=partition_col_sql, - ensure_distinct_key=True, seed=43, merge_mode=merge_mode) + ensure_distinct_key=True, seed=43, merge_mode=merge_mode, format_version=format_version) # Drop one partition field on target tables and source table (creates void transform) def drop_partition_field(spark, table_name): @@ -764,8 +834,9 @@ def do_merge(spark, target_table): WHEN NOT MATCHED THEN INSERT * """) - with_gpu_session(lambda spark: do_merge(spark, gpu_target_table), - conf=iceberg_merge_enabled_conf) + with_iceberg_dml_session( + lambda spark: do_merge(spark, gpu_target_table), + format_version, merge_mode, conf=iceberg_merge_enabled_conf) with_cpu_session(lambda spark: do_merge(spark, cpu_target_table)) # Compare results @@ -779,10 +850,11 @@ def do_merge(spark, target_table): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_iceberg_merge_partitioned_fanout_enabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_cow_format_versions) +def test_iceberg_merge_partitioned_fanout_enabled(format_version, spark_tmp_table_factory): # Use bucket(2, ...) to keep partition count low and avoid OOM from Iceberg's FanoutDataWriter. _do_test_iceberg_merge( spark_tmp_table_factory, "bucket(2, _c9)", merge_mode='copy-on-write', - table_properties={"write.spark.fanout.enabled": "true"}) + table_properties={"write.spark.fanout.enabled": "true"}, format_version=format_version) diff --git a/integration_tests/src/main/python/iceberg/iceberg_overwrite_dynamic_test.py b/integration_tests/src/main/python/iceberg/iceberg_overwrite_dynamic_test.py index 384974514f8..9835189463d 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_overwrite_dynamic_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_overwrite_dynamic_test.py @@ -18,12 +18,11 @@ from asserts import assert_equal_with_local_sort, assert_gpu_fallback_collect from conftest import is_iceberg_remote_catalog from data_gen import gen_df, copy_and_update -from iceberg import create_iceberg_table, \ - iceberg_base_table_cols, iceberg_gens_list, \ - get_full_table_name, iceberg_full_gens_list, iceberg_nested_write_gens_list, \ - iceberg_write_enabled_conf, iceberg_unsupported_mark, \ - overwrite_dynamic_partition_transforms, supports_iceberg_v3, \ - ICEBERG_V3_UNSUPPORTED_REASON +from iceberg import ( + iceberg_format_versions, create_iceberg_table, iceberg_base_table_cols, iceberg_gens_list, + get_full_table_name, iceberg_full_gens_list, iceberg_nested_write_gens_list, + iceberg_write_enabled_conf, iceberg_unsupported_mark, overwrite_dynamic_partition_transforms, + supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON) from marks import iceberg, ignore_order, allow_non_gpu, allow_non_gpu_conditional, datagen_overrides from spark_session import with_gpu_session, with_cpu_session, is_spark_400_or_later @@ -84,9 +83,10 @@ def overwrite_data(spark, table_name: str): @iceberg @ignore_order(local=True) -def test_insert_overwrite_dynamic_unpartitioned_table(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_unpartitioned_table(format_version, spark_tmp_table_factory): """Test dynamic overwrite on unpartitioned tables - should run on GPU.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} do_test_insert_overwrite_dynamic( spark_tmp_table_factory, @@ -115,13 +115,14 @@ def insert_data(spark, seed): assert_gpu_fallback_collect( lambda spark: insert_data(spark, None), "OverwritePartitionsDynamicExec", - conf=dynamic_overwrite_conf) + conf=copy_and_update(dynamic_overwrite_conf, {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) -def _do_test_insert_overwrite_dynamic_partitioned(spark_tmp_table_factory, partition_col_sql, table_prop=None): +def _do_test_insert_overwrite_dynamic_partitioned( + spark_tmp_table_factory, partition_col_sql, table_prop=None, format_version="2"): """Helper function for partitioned table dynamic overwrite tests.""" if table_prop is None: - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} def create_table_with_partition(table_name: str): create_iceberg_table( @@ -140,9 +141,13 @@ def create_table_with_partition(table_name: str): @pytest.mark.parametrize("partition_col_sql", [ pytest.param("year(_c9)", id="year(timestamp_col)"), ]) -def test_insert_overwrite_dynamic_bucket_partitioned(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_bucket_partitioned(format_version, spark_tmp_table_factory, partition_col_sql): """Basic partition test - runs for all catalogs including remote.""" - _do_test_insert_overwrite_dynamic_partitioned(spark_tmp_table_factory, partition_col_sql) + _do_test_insert_overwrite_dynamic_partitioned( + spark_tmp_table_factory, + partition_col_sql, + format_version=format_version) @iceberg @@ -151,20 +156,26 @@ def test_insert_overwrite_dynamic_bucket_partitioned(spark_tmp_table_factory, pa @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("partition_col_sql", overwrite_dynamic_partition_transforms) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_insert_overwrite_dynamic_bucket_partitioned_full_coverage(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_bucket_partitioned_full_coverage( + format_version, spark_tmp_table_factory, partition_col_sql): """Sanity-check dynamic INSERT OVERWRITE against a few partition transforms distinct from those picked by other DML ops. The 26-transform partition-writer coverage anchor lives in iceberg_append_test.py::test_insert_into_partitioned_table_full_coverage.""" - _do_test_insert_overwrite_dynamic_partitioned(spark_tmp_table_factory, partition_col_sql) + _do_test_insert_overwrite_dynamic_partitioned( + spark_tmp_table_factory, + partition_col_sql, + format_version=format_version) @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_dynamic_nested_types(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_nested_types(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE with dynamic mode on Iceberg-native nested types on GPU.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_nested_write_gens_list)] gen_list = list(zip(cols, iceberg_nested_write_gens_list)) @@ -209,9 +220,10 @@ def overwrite_data(spark, table_name): @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_insert_overwrite_dynamic_all_cols(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_all_cols(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE with dynamic mode on all Iceberg write types on GPU.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] gen_list = list(zip(cols, iceberg_full_gens_list)) @@ -258,10 +270,11 @@ def overwrite_data(spark, table_name): @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("file_format", ["orc", "avro"], ids=lambda x: f"file_format={x}") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_insert_overwrite_dynamic_unsupported_file_format_fallback( +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_unsupported_file_format_fallback(format_version, spark_tmp_table_factory, file_format): """Test that INSERT OVERWRITE falls back to CPU with unsupported file formats.""" - table_prop = {"format-version": "2", + table_prop = {"format-version": format_version, "write.format.default": file_format} table_name = get_full_table_name(spark_tmp_table_factory) @@ -300,10 +313,11 @@ def overwrite_data(spark, table_name: str): @pytest.mark.parametrize("conf_key", ["spark.rapids.sql.format.iceberg.enabled", "spark.rapids.sql.format.iceberg.write.enabled"], ids=lambda x: f"{x}=False") -def test_insert_overwrite_dynamic_fallback_when_conf_disabled( +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_fallback_when_conf_disabled(format_version, spark_tmp_table_factory, conf_key): """Test that INSERT OVERWRITE falls back to CPU when Iceberg write is disabled.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} table_name = get_full_table_name(spark_tmp_table_factory) @@ -340,12 +354,13 @@ def overwrite_data(spark, table_name: str): @pytest.mark.parametrize("partition_col_sql", [ pytest.param("year(_c9)", id="year_partition"), ]) -def test_overwrite_dynamic_aqe(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_overwrite_dynamic_aqe(format_version, spark_tmp_table_factory, partition_col_sql): """ Test INSERT OVERWRITE (dynamic partitions) with AQE enabled. """ table_prop = { - 'format-version': '2', + 'format-version': format_version, } # Configuration with AQE enabled @@ -385,7 +400,8 @@ def overwrite_dynamic(spark, table_name): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_dynamic_after_drop_partition_field(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_after_drop_partition_field(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE (dynamic mode) on table after dropping a partition field (void transform). When a partition field is dropped, Iceberg creates a 'void transform' - @@ -397,7 +413,7 @@ def test_insert_overwrite_dynamic_after_drop_partition_field(spark_tmp_table_fac cpu_table_name = f"{base_table_name}_cpu" gpu_table_name = f"{base_table_name}_gpu" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} # Use two partition columns so after dropping one, we still have at least one partition_col_sql = "bucket(8, _c2), bucket(8, _c3)" @@ -450,9 +466,10 @@ def overwrite_data(spark, table_name): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_dynamic_partitioned_fanout_enabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_dynamic_partitioned_fanout_enabled(format_version, spark_tmp_table_factory): # Use bucket(2, ...) to keep partition count low and avoid OOM from Iceberg's FanoutDataWriter. _do_test_insert_overwrite_dynamic_partitioned( spark_tmp_table_factory, "bucket(2, _c9)", - table_prop={"format-version": "2", "write.spark.fanout.enabled": "true"}) + table_prop={"format-version": format_version, "write.spark.fanout.enabled": "true"}, format_version=format_version) diff --git a/integration_tests/src/main/python/iceberg/iceberg_overwrite_static_test.py b/integration_tests/src/main/python/iceberg/iceberg_overwrite_static_test.py index 417f4c040fd..68500c6bb2b 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_overwrite_static_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_overwrite_static_test.py @@ -20,13 +20,13 @@ assert_gpu_fallback_collect from conftest import is_iceberg_remote_catalog from data_gen import DEFAULT_DATA_GEN_LENGTH, StringGen, copy_and_update, gen_df -from iceberg import create_iceberg_table, \ - iceberg_base_table_cols, iceberg_gens_list, \ - get_full_table_name, iceberg_full_gens_list, iceberg_nested_write_gens_list, \ - iceberg_write_enabled_conf, iceberg_unsupported_mark, _build_tblprops, \ - overwrite_static_partition_transforms, supports_iceberg_v3, \ - ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, \ - ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, row_lineage_df +from iceberg import ( + iceberg_format_versions, create_iceberg_table, iceberg_base_table_cols, iceberg_gens_list, + get_full_table_name, iceberg_full_gens_list, iceberg_nested_write_gens_list, + iceberg_write_enabled_conf, iceberg_unsupported_mark, _build_tblprops, + overwrite_static_partition_transforms, supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON, + supports_iceberg_row_lineage_inheritance, ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, + row_lineage_df) from marks import iceberg, ignore_order, allow_non_gpu, datagen_overrides from spark_session import with_gpu_session, with_cpu_session @@ -85,9 +85,10 @@ def overwrite_data(spark, table_name: str): @iceberg @ignore_order(local=True) -def test_insert_overwrite_unpartitioned_table(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_unpartitioned_table(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE on unpartitioned Iceberg tables.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} do_test_insert_overwrite_table_sql( spark_tmp_table_factory, @@ -115,7 +116,9 @@ def insert_data(spark, seed): assert_gpu_fallback_collect( lambda spark: insert_data(spark, None), "OverwriteByExpressionExec", - conf=iceberg_static_overwrite_conf) + conf=copy_and_update( + iceberg_static_overwrite_conf, + {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) @iceberg @@ -154,14 +157,15 @@ def setup_iceberg_table(spark): @ignore_order(local=True) @allow_non_gpu('OverwriteByExpressionExec', 'AppendDataExec') @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_unpartitioned_table_values(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_unpartitioned_table_values(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE on unpartitioned tables with VALUES syntax.""" base_table_name = get_full_table_name(spark_tmp_table_factory) cpu_table_name = f"{base_table_name}_cpu" gpu_table_name = f"{base_table_name}_gpu" def create_table(spark, table_name: str): - props = _build_tblprops({"format-version": "2"}) + props = _build_tblprops({"format-version": format_version}) props_sql = ", ".join(f"'{k}' = '{v}'" for k, v in props.items()) spark.sql(f"CREATE TABLE {table_name} (id int, name string) USING ICEBERG " f"TBLPROPERTIES ({props_sql})") @@ -192,10 +196,11 @@ def overwrite_data(spark, table_name: str): assert_equal_with_local_sort(cpu_data, gpu_data) -def _do_test_insert_overwrite_partitioned_table(spark_tmp_table_factory, partition_col_sql, table_prop=None): +def _do_test_insert_overwrite_partitioned_table( + spark_tmp_table_factory, partition_col_sql, table_prop=None, format_version="2"): """Helper function for partitioned table INSERT OVERWRITE tests.""" if table_prop is None: - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} def create_table_and_set_write_order(table_name: str): create_iceberg_table( @@ -217,9 +222,13 @@ def create_table_and_set_write_order(table_name: str): @pytest.mark.parametrize("partition_col_sql", [ pytest.param("year(_c9)", id="year(timestamp_col)"), ]) -def test_insert_overwrite_partitioned_table(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_partitioned_table(format_version, spark_tmp_table_factory, partition_col_sql): """Basic partition test - runs for all catalogs including remote.""" - _do_test_insert_overwrite_partitioned_table(spark_tmp_table_factory, partition_col_sql) + _do_test_insert_overwrite_partitioned_table( + spark_tmp_table_factory, + partition_col_sql, + format_version=format_version) @iceberg @@ -227,19 +236,25 @@ def test_insert_overwrite_partitioned_table(spark_tmp_table_factory, partition_c @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("partition_col_sql", overwrite_static_partition_transforms) -def test_insert_overwrite_partitioned_table_full_coverage(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_partitioned_table_full_coverage( + format_version, spark_tmp_table_factory, partition_col_sql): """Sanity-check INSERT OVERWRITE against a few partition transforms distinct from those picked by other DML ops. The 26-transform partition-writer coverage anchor lives in iceberg_append_test.py::test_insert_into_partitioned_table_full_coverage.""" - _do_test_insert_overwrite_partitioned_table(spark_tmp_table_factory, partition_col_sql) + _do_test_insert_overwrite_partitioned_table( + spark_tmp_table_factory, + partition_col_sql, + format_version=format_version) @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_unpartitioned_table_nested_types(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_unpartitioned_table_nested_types(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE with Iceberg-native nested types on GPU.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_nested_write_gens_list)] gen_list = list(zip(cols, iceberg_nested_write_gens_list)) @@ -283,9 +298,10 @@ def overwrite_data(spark, table_name): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_unpartitioned_table_all_cols(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_unpartitioned_table_all_cols(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE on unpartitioned table with all Iceberg write types on GPU.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] gen_list = list(zip(cols, iceberg_full_gens_list)) @@ -329,9 +345,10 @@ def overwrite_data(spark, table_name): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_partitioned_table_nested_types(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_partitioned_table_nested_types(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE on partitioned table with Iceberg-native nested types on GPU.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_nested_write_gens_list)] gen_list = list(zip(cols, iceberg_nested_write_gens_list)) partition_col_sql = "bucket(16, _c0), bucket(16, _c1)" @@ -378,9 +395,10 @@ def overwrite_data(spark, table_name): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_partitioned_table_all_cols(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_partitioned_table_all_cols(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE on partitioned table with all Iceberg write types on GPU.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] gen_list = list(zip(cols, iceberg_full_gens_list)) partition_col_sql = "bucket(16, _c2), bucket(16, _c3)" @@ -429,10 +447,11 @@ def overwrite_data(spark, table_name): @allow_non_gpu('OverwriteByExpressionExec', 'ShuffleExchangeExec', 'SortExec', 'ProjectExec') @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("file_format", ["orc", "avro"], ids=lambda x: f"file_format={x}") -def test_insert_overwrite_table_unsupported_file_format_fallback( +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_table_unsupported_file_format_fallback(format_version, spark_tmp_table_factory, file_format): """Test that unsupported file formats fall back to CPU.""" - table_prop = {"format-version": "2", + table_prop = {"format-version": format_version, "write.format.default": file_format} def insert_initial_data(spark, table_name: str): @@ -467,10 +486,11 @@ def overwrite_data(spark, table_name: str): @pytest.mark.parametrize("conf_key", ["spark.rapids.sql.format.iceberg.enabled", "spark.rapids.sql.format.iceberg.write.enabled"], ids=lambda x: f"{x}=False") -def test_insert_overwrite_iceberg_table_fallback_when_conf_disabled( +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_iceberg_table_fallback_when_conf_disabled(format_version, spark_tmp_table_factory, conf_key): """Test that overwrite falls back to CPU when Iceberg write is disabled.""" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} def insert_initial_data(spark, table_name: str): df = gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list)), seed=INITIAL_INSERT_SEED) @@ -501,7 +521,8 @@ def overwrite_data(spark, table_name: str): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_static_after_drop_partition_field(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_static_after_drop_partition_field(format_version, spark_tmp_table_factory): """Test INSERT OVERWRITE (static mode) on table after dropping a partition field (void transform). When a partition field is dropped, Iceberg creates a 'void transform' - @@ -512,7 +533,7 @@ def test_insert_overwrite_static_after_drop_partition_field(spark_tmp_table_fact cpu_table_name = f"{base_table_name}_cpu" gpu_table_name = f"{base_table_name}_gpu" - table_prop = {"format-version": "2"} + table_prop = {"format-version": format_version} # Use two partition columns so after dropping one, we still have at least one partition_col_sql = "bucket(8, _c2), bucket(8, _c3)" @@ -557,7 +578,8 @@ def overwrite_data(spark, table_name): @ignore_order(local=True) @allow_non_gpu('ShuffleExchangeExec') @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_static_df_api_truncate_string(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_static_df_api_truncate_string(format_version, spark_tmp_table_factory): """Test static overwrite via DataFrame writeTo().overwrite() API with truncate(5, string_col) partitioning. Verifies GPU writes produce Parquet files with correct Iceberg field IDs so that file-level statistics are available for overwrite validation. @@ -568,7 +590,7 @@ def test_insert_overwrite_static_df_api_truncate_string(spark_tmp_table_factory) partition_col_sql = f"truncate({truncate_width}, _c6)" partition_filter = f"_c6 >= '{prefix}10' AND _c6 < '{prefix}20'" - table_prop = _build_tblprops({"format-version": "2", + table_prop = _build_tblprops({"format-version": format_version, "write.format.default": "parquet"}) conf = copy_and_update(iceberg_static_overwrite_conf, { @@ -615,9 +637,10 @@ def overwrite_data(spark, table_name): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_insert_overwrite_partitioned_table_fanout_enabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_insert_overwrite_partitioned_table_fanout_enabled(format_version, spark_tmp_table_factory): # Use bucket(2, ...) to keep partition count low and avoid OOM from Iceberg's FanoutDataWriter. _do_test_insert_overwrite_partitioned_table( spark_tmp_table_factory, "bucket(2, _c9)", - table_prop={"format-version": "2", "write.spark.fanout.enabled": "true"}) + table_prop={"format-version": format_version, "write.spark.fanout.enabled": "true"}, format_version=format_version) diff --git a/integration_tests/src/main/python/iceberg/iceberg_rtas_test.py b/integration_tests/src/main/python/iceberg/iceberg_rtas_test.py index dd3af2518c1..71ef441f155 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_rtas_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_rtas_test.py @@ -19,14 +19,11 @@ from asserts import assert_equal_with_local_sort, assert_gpu_fallback_collect from conftest import is_iceberg_remote_catalog from data_gen import gen_df, copy_and_update -from iceberg import (create_iceberg_table, - iceberg_base_table_cols, - iceberg_gens_list, iceberg_full_gens_list, - iceberg_nested_write_gens_list, - get_full_table_name, iceberg_write_enabled_conf, - iceberg_unsupported_mark, _build_tblprops, - rtas_partition_transforms, supports_iceberg_v3, - ICEBERG_V3_UNSUPPORTED_REASON) +from iceberg import ( + iceberg_format_versions, create_iceberg_table, iceberg_base_table_cols, iceberg_gens_list, + iceberg_full_gens_list, iceberg_nested_write_gens_list, get_full_table_name, + iceberg_write_enabled_conf, iceberg_unsupported_mark, _build_tblprops, rtas_partition_transforms, + supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON) from marks import iceberg, ignore_order, allow_non_gpu, allow_non_gpu_conditional, datagen_overrides from spark_session import with_gpu_session, with_cpu_session, is_spark_400_or_later @@ -100,9 +97,10 @@ def run_gpu_rtas(spark): @iceberg @ignore_order(local=True) -def test_rtas_unpartitioned_table(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_unpartitioned_table(format_version, spark_tmp_table_factory): table_prop = { - "format-version": "2" + "format-version": format_version } df_gen = lambda spark: gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) @@ -132,14 +130,15 @@ def run_rtas(spark): assert_gpu_fallback_collect( run_rtas, "AtomicReplaceTableAsSelectExec", - conf=iceberg_write_enabled_conf) + conf=copy_and_update(iceberg_write_enabled_conf, {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) -def _do_test_rtas_partitioned_table(spark_tmp_table_factory, partition_col_sql, table_prop=None): +def _do_test_rtas_partitioned_table( + spark_tmp_table_factory, partition_col_sql, table_prop=None, format_version="2"): """Helper function for partitioned table RTAS tests.""" if table_prop is None: table_prop = { - "format-version": "2" + "format-version": format_version } df_gen = lambda spark: gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) @@ -156,9 +155,10 @@ def _do_test_rtas_partitioned_table(spark_tmp_table_factory, partition_col_sql, @pytest.mark.parametrize("partition_col_sql", [ pytest.param("year(_c9)", id="year(timestamp_col)"), ]) -def test_rtas_partitioned_table(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_partitioned_table(format_version, spark_tmp_table_factory, partition_col_sql): """Basic partition test - runs for all catalogs including remote.""" - _do_test_rtas_partitioned_table(spark_tmp_table_factory, partition_col_sql) + _do_test_rtas_partitioned_table(spark_tmp_table_factory, partition_col_sql, format_version=format_version) @iceberg @@ -167,20 +167,22 @@ def test_rtas_partitioned_table(spark_tmp_table_factory, partition_col_sql): @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("partition_col_sql", rtas_partition_transforms) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_rtas_partitioned_table_full_coverage(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_partitioned_table_full_coverage(format_version, spark_tmp_table_factory, partition_col_sql): """Sanity-check RTAS against a few partition transforms distinct from those picked by other DML ops. The 26-transform partition-writer coverage anchor lives in iceberg_append_test.py::test_insert_into_partitioned_table_full_coverage.""" - _do_test_rtas_partitioned_table(spark_tmp_table_factory, partition_col_sql) + _do_test_rtas_partitioned_table(spark_tmp_table_factory, partition_col_sql, format_version=format_version) @iceberg @ignore_order(local=True) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_create_or_replace_table(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_create_or_replace_table(format_version, spark_tmp_table_factory): """Test CREATE OR REPLACE TABLE AS SELECT when table doesn't exist""" table_prop = { - "format-version": "2" + "format-version": format_version } base_name = get_full_table_name(spark_tmp_table_factory) @@ -208,10 +210,11 @@ def test_create_or_replace_table(spark_tmp_table_factory): @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("file_format", ["orc", "avro"], ids=lambda x: f"file_format={x}") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_rtas_unsupported_file_format_fallback(spark_tmp_table_factory, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_unsupported_file_format_fallback(format_version, spark_tmp_table_factory, file_format): table_prop = { - "format-version": "2", + "format-version": format_version, "write.format.default": file_format } @@ -241,10 +244,11 @@ def run_rtas(spark): "spark.rapids.sql.format.iceberg.write.enabled"], ids=lambda x: f"{x}=False") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_rtas_fallback_when_conf_disabled(spark_tmp_table_factory, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_fallback_when_conf_disabled(format_version, spark_tmp_table_factory, conf_key): table_prop = { - "format-version": "2" + "format-version": format_version } def run_rtas(spark): @@ -269,9 +273,10 @@ def run_rtas(spark): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_rtas_unpartitioned_table_nested_types(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_unpartitioned_table_nested_types(format_version, spark_tmp_table_factory): table_prop = { - "format-version": "2" + "format-version": format_version } cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_nested_write_gens_list)] @@ -285,10 +290,11 @@ def test_rtas_unpartitioned_table_nested_types(spark_tmp_table_factory): @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_rtas_unpartitioned_table_all_cols(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_unpartitioned_table_all_cols(format_version, spark_tmp_table_factory): """Test RTAS on unpartitioned table with all Iceberg write types on GPU.""" table_prop = { - "format-version": "2" + "format-version": format_version } cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] @@ -301,9 +307,10 @@ def test_rtas_unpartitioned_table_all_cols(spark_tmp_table_factory): @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_rtas_partitioned_table_nested_types(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_partitioned_table_nested_types(format_version, spark_tmp_table_factory): table_prop = { - "format-version": "2" + "format-version": format_version } cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_nested_write_gens_list)] @@ -320,10 +327,11 @@ def test_rtas_partitioned_table_nested_types(spark_tmp_table_factory): @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_rtas_partitioned_table_all_cols(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_partitioned_table_all_cols(format_version, spark_tmp_table_factory): """Test RTAS on partitioned table with all Iceberg write types on GPU.""" table_prop = { - "format-version": "2" + "format-version": format_version } cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_full_gens_list)] @@ -342,10 +350,11 @@ def test_rtas_partitioned_table_all_cols(spark_tmp_table_factory): @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.parametrize("partition_table", [True, False], ids=lambda x: f"partition_table={x}") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_rtas_from_values(spark_tmp_table_factory, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_from_values(format_version, spark_tmp_table_factory, partition_table): table_prop = { - "format-version": "2" + "format-version": format_version } base_name = get_full_table_name(spark_tmp_table_factory) @@ -356,7 +365,11 @@ def execute_rtas_from_values(spark, target_table: str): # Create initial table initial_df_gen = lambda sp: gen_df(sp, [("id", iceberg_gens_list[0]), ("name", iceberg_gens_list[1])]) partition_col_sql = "bucket(8, id)" if partition_table else None - create_iceberg_table(target_table, partition_col_sql, table_prop, initial_df_gen) + create_iceberg_table( + target_table, + partition_col_sql, + table_prop, + initial_df_gen) # Execute RTAS partition_clause = "" if not partition_table else "PARTITIONED BY (bucket(8, id)) " @@ -383,12 +396,13 @@ def execute_rtas_from_values(spark, target_table: str): pytest.param(None, id="unpartitioned"), pytest.param("year(_c9)", id="year_partition"), ]) -def test_rtas_aqe(spark_tmp_table_factory, partition_col_sql): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_aqe(format_version, spark_tmp_table_factory, partition_col_sql): """ Test REPLACE TABLE AS SELECT with AQE enabled. """ table_prop = { - "format-version": "2", + "format-version": format_version, } df_gen = lambda spark: gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) @@ -411,9 +425,10 @@ def test_rtas_aqe(spark_tmp_table_factory, partition_col_sql): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_rtas_partitioned_table_fanout_enabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_rtas_partitioned_table_fanout_enabled(format_version, spark_tmp_table_factory): # Use bucket(2, ...) to keep partition count low and avoid OOM from Iceberg's FanoutDataWriter. _do_test_rtas_partitioned_table( spark_tmp_table_factory, "bucket(2, _c9)", - table_prop={"format-version": "2", "write.spark.fanout.enabled": "true"}) + table_prop={"format-version": format_version, "write.spark.fanout.enabled": "true"}, format_version=format_version) diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index 07b8970c605..fd82c1fc164 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -21,10 +21,12 @@ assert_gpu_and_cpu_row_counts_equal, assert_gpu_fallback_collect, assert_spark_exception from conftest import is_iceberg_remote_catalog, is_iceberg_rest_catalog from data_gen import * -from iceberg import get_full_table_name, iceberg_unsupported_mark, _build_tblprops, \ - _BASE_TBLPROPS_SQL, create_iceberg_table, iceberg_write_enabled_conf, supports_iceberg_v3, \ - ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, \ - ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, row_lineage_df +from iceberg import ( + iceberg_format_versions, iceberg_read_enabled_conf, iceberg_read_format_versions, + iceberg_table_properties_sql, get_full_table_name, iceberg_unsupported_mark, _build_tblprops, + _BASE_TBLPROPS_SQL, create_iceberg_table, iceberg_write_enabled_conf, supports_iceberg_v3, + ICEBERG_V3_UNSUPPORTED_REASON, supports_iceberg_row_lineage_inheritance, + ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, row_lineage_df) from marks import allow_non_gpu, iceberg, ignore_order from spark_session import is_databricks_runtime, is_spark_35x, is_spark_400_or_later, \ is_spark_40x, is_spark_41x, spark_version, with_cpu_session, with_gpu_session @@ -157,10 +159,12 @@ def _assert_partial_clustering_spj_plan(_cpu_plan, plan): or (is_spark_41x() and _is_spark_patch_at_least(spark_version(), 2)) ), reason="Requires Spark's partial-clustering correctness fix and GPU Iceberg scan support") -def test_iceberg_spj_partial_clustering_distinct(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_spj_partial_clustering_distinct(format_version, spark_tmp_table_factory): left_table = get_full_table_name(spark_tmp_table_factory) right_table = get_full_table_name(spark_tmp_table_factory) table_props = _build_tblprops({ + "format-version": format_version, # Keep separate INSERTs as separate scan splits so that id=1 is partially clustered. "read.split.target-size": "1", "read.split.open-file-cost": "1", @@ -205,7 +209,7 @@ def distinct_after_spj(spark): # Comparing the results also exercises the replicated and padded scan partitions. assert_cpu_and_gpu_are_equal_collect_with_capture( distinct_after_spj, - conf=conf, + conf={**iceberg_read_enabled_conf, **conf}, require_non_empty=True, gpu_plan_assertion=_assert_partial_clustering_spj_plan) @@ -218,11 +222,13 @@ def distinct_after_spj(spark): @pytest.mark.parametrize("partition_filter", [True, False], ids=["filtered", "unfiltered"]) @pytest.mark.parametrize("partially_clustered", [True, False], ids=["partially_clustered", "clustered"]) -def test_iceberg_spj_partition_filter(spark_tmp_table_factory, partition_filter, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_spj_partition_filter(format_version, spark_tmp_table_factory, partition_filter, partially_clustered): left_table = get_full_table_name(spark_tmp_table_factory) right_table = get_full_table_name(spark_tmp_table_factory) table_props = _build_tblprops({ + "format-version": format_version, # Keep separate INSERTs as separate scan splits so that id=1 is partially clustered. "read.split.target-size": "1", "read.split.open-file-cost": "1", @@ -287,7 +293,7 @@ def assert_plan(_cpu_plan, plan): # what forces the patch-level gate on test_iceberg_spj_partial_clustering_distinct. assert_cpu_and_gpu_are_equal_collect_with_capture( join_after_spj, - conf=conf, + conf={**iceberg_read_enabled_conf, **conf}, require_non_empty=True, gpu_plan_assertion=assert_plan) @@ -319,7 +325,8 @@ def assert_plan(_cpu_plan, plan): ids=["reduced", "control"]) @pytest.mark.parametrize("key_ddl, key_value_sql, left_transform, right_transform", _spj_reducible_transforms) -def test_iceberg_spj_reducible_transforms(spark_tmp_table_factory, key_ddl, key_value_sql, +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_spj_reducible_transforms(format_version, spark_tmp_table_factory, key_ddl, key_value_sql, left_transform, right_transform, allow_compatible_transforms): left_table = get_full_table_name(spark_tmp_table_factory) @@ -328,10 +335,10 @@ def test_iceberg_spj_reducible_transforms(spark_tmp_table_factory, key_ddl, key_ def setup_iceberg_tables(spark): spark.sql( f"CREATE TABLE {left_table} ({key_ddl}, price DOUBLE) USING ICEBERG " - f"PARTITIONED BY ({left_transform}) {_NO_FANOUT}") + f"PARTITIONED BY ({left_transform}) {iceberg_table_properties_sql(format_version)}") spark.sql( f"CREATE TABLE {right_table} ({key_ddl}, value STRING) USING ICEBERG " - f"PARTITIONED BY ({right_transform}) {_NO_FANOUT}") + f"PARTITIONED BY ({right_transform}) {iceberg_table_properties_sql(format_version)}") spark.sql( f"INSERT INTO {left_table} SELECT {key_value_sql}, CAST(id AS DOUBLE) " f"FROM range({_SPJ_REDUCIBLE_ROWS})") @@ -366,7 +373,7 @@ def join_on_reducible_transforms(spark): # is the control: it must stay correct whether or not the reduced grouping works. assert_cpu_and_gpu_are_equal_collect_with_capture( join_on_reducible_transforms, - conf=conf, + conf={**iceberg_read_enabled_conf, **conf}, require_non_empty=True, gpu_plan_assertion=lambda _cpu_plan, plan: _assert_spj_join_shape( plan, allow_compatible_transforms)) @@ -375,24 +382,28 @@ def join_on_reducible_transforms(spark): @allow_non_gpu("BatchScanExec") @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering -def test_iceberg_fallback_not_unsafe_row(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_fallback_not_unsafe_row(format_version, spark_tmp_table_factory): full_table = get_full_table_name(spark_tmp_table_factory) def setup_iceberg_table(spark): - spark.sql(f"CREATE TABLE {full_table} (id BIGINT, data STRING) USING ICEBERG {_NO_FANOUT}") + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT, data STRING) USING ICEBERG {iceberg_table_properties_sql(format_version)}") spark.sql(f"INSERT INTO {full_table} VALUES (1, 'a'), (2, 'b'), (3, 'c')") with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql(f"SELECT COUNT(DISTINCT id) from {full_table}"), - conf={"spark.rapids.sql.format.iceberg.enabled": "false"} + conf={**iceberg_read_enabled_conf, "spark.rapids.sql.format.iceberg.enabled": "false"} ) @iceberg -def test_iceberg_scan_from_background_thread(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_scan_from_background_thread(format_version, spark_tmp_table_factory): full_table = get_full_table_name(spark_tmp_table_factory) def setup_iceberg_table(spark): - spark.sql(f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG {_NO_FANOUT}") + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT) USING ICEBERG {iceberg_table_properties_sql(format_version)}") spark.sql(f"INSERT INTO {full_table} VALUES (1), (2), (3)") with_cpu_session(setup_iceberg_table) @@ -409,26 +420,28 @@ def run_query(): callback = spark._sc._jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback callback.assertContains(df._jdf, "GpuBatchScanExec") - with_gpu_session(scan_iceberg_table) + with_gpu_session(scan_iceberg_table, conf=iceberg_read_enabled_conf) @iceberg @ignore_order(local=True) @pytest.mark.skipif(is_databricks_runtime(), reason="AQE+DPP not supported until Spark 3.2.0+ and AQE+DPP not supported on Databricks") @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_aqe_dpp(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_aqe_dpp(format_version, spark_tmp_table_factory, reader_type): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = two_col_df(spark, int_gen, int_gen) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} (a INT, b INT) USING ICEBERG PARTITIONED BY (a) {_NO_FANOUT}") + spark.sql( + f"CREATE TABLE {full_table} (a INT, b INT) USING ICEBERG PARTITIONED BY (a) {iceberg_table_properties_sql(format_version)}") spark.sql(f"INSERT INTO {full_table} SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql(f"SELECT * from {full_table} as X JOIN {full_table} as Y ON X.a = Y.a " f"WHERE Y.a > 0"), - conf={"spark.sql.adaptive.enabled": "true", + conf={**iceberg_read_enabled_conf, "spark.sql.adaptive.enabled": "true", "spark.rapids.sql.format.parquet.reader.type": reader_type, "spark.sql.optimizer.dynamicPartitionPruning.enabled": "true"}) @@ -436,66 +449,76 @@ def setup_iceberg_table(spark): @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize("data_gens", iceberg_gens_list, ids=idfn) @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_parquet_read_round_trip_select_one(spark_tmp_table_factory, data_gens, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_parquet_read_round_trip_select_one( + format_version, spark_tmp_table_factory, data_gens, reader_type): gen_list = [('_c' + str(i), gen) for i, gen in enumerate(data_gens)] full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = gen_df(spark, gen_list) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {_NO_FANOUT} AS SELECT * FROM {tmpview}") + spark.sql( + f"CREATE TABLE {full_table} USING ICEBERG {iceberg_table_properties_sql(format_version)} AS SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) # explicitly only select 1 column to make sure we test that path in the schema parsing code assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql(f"SELECT _c0 FROM {full_table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize("data_gens", iceberg_primitive_gens_list, ids=idfn) @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_parquet_read_round_trip(spark_tmp_table_factory, data_gens, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_parquet_read_round_trip(format_version, spark_tmp_table_factory, data_gens, reader_type): gen_list = [('_c' + str(i), gen) for i, gen in enumerate(data_gens)] full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = gen_df(spark, gen_list) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {_NO_FANOUT} AS SELECT * FROM {tmpview}") + spark.sql( + f"CREATE TABLE {full_table} USING ICEBERG {iceberg_table_properties_sql(format_version)} AS SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql(f"SELECT * FROM {full_table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize("data_gens", iceberg_gens_list, ids=idfn) @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_parquet_read_round_trip_all_types(spark_tmp_table_factory, data_gens, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_parquet_read_round_trip_all_types( + format_version, spark_tmp_table_factory, data_gens, reader_type): gen_list = [('_c' + str(i), gen) for i, gen in enumerate(data_gens)] full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = gen_df(spark, gen_list) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {_NO_FANOUT} AS SELECT * FROM {tmpview}") + spark.sql( + f"CREATE TABLE {full_table} USING ICEBERG {iceberg_table_properties_sql(format_version)} AS SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql(f"SELECT * FROM {full_table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @pytest.mark.parametrize("data_gens", [[long_gen]], ids=idfn) @pytest.mark.parametrize("iceberg_format", ["orc", "avro"], ids=idfn) @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_unsupported_formats(spark_tmp_table_factory, data_gens, iceberg_format, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_unsupported_formats( + format_version, spark_tmp_table_factory, data_gens, iceberg_format, reader_type): gen_list = [('_c' + str(i), gen) for i, gen in enumerate(data_gens)] full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = gen_df(spark, gen_list) df.createOrReplaceTempView(tmpview) - props = _build_tblprops({'write.format.default': iceberg_format}) + props = _build_tblprops({"format-version": format_version, 'write.format.default': iceberg_format}) props_sql = ", ".join(f"'{k}' = '{v}'" for k, v in props.items()) spark.sql(f"CREATE TABLE {full_table} USING ICEBERG " f"TBLPROPERTIES({props_sql}) " @@ -504,7 +527,7 @@ def setup_iceberg_table(spark): assert_spark_exception( lambda : with_gpu_session( lambda spark : spark.sql(f"SELECT * FROM {full_table}").collect(), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}), + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}), "UnsupportedOperationException") @iceberg @@ -512,16 +535,18 @@ def setup_iceberg_table(spark): @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize("disable_conf", ["spark.rapids.sql.format.iceberg.enabled", "spark.rapids.sql.format.iceberg.read.enabled"], ids=idfn) -def test_iceberg_read_fallback(spark_tmp_table_factory, disable_conf): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_fallback(format_version, spark_tmp_table_factory, disable_conf): full_table = get_full_table_name(spark_tmp_table_factory) def setup_iceberg_table(spark): - spark.sql(f"CREATE TABLE {full_table} (id BIGINT, data STRING) USING ICEBERG {_NO_FANOUT}") + spark.sql( + f"CREATE TABLE {full_table} (id BIGINT, data STRING) USING ICEBERG {iceberg_table_properties_sql(format_version)}") spark.sql(f"INSERT INTO {full_table} VALUES (1, 'a'), (2, 'b'), (3, 'c')") with_cpu_session(setup_iceberg_table) assert_gpu_fallback_collect( lambda spark : spark.sql(f"SELECT * FROM {full_table}"), "BatchScanExec", - conf = {disable_conf : "false"}) + conf = {**iceberg_read_enabled_conf, disable_conf : "false"}) @iceberg @@ -692,14 +717,15 @@ def overwrite(spark, table): pytest.param(("lz4", "Unsupported Parquet compression type")), ("zstd", None)], ids=idfn) @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_read_parquet_compression_codec(spark_tmp_table_factory, codec_info, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_parquet_compression_codec(format_version, spark_tmp_table_factory, codec_info, reader_type): codec, error_msg = codec_info full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) - props = _build_tblprops({'write.parquet.compression-codec': codec}) + props = _build_tblprops({"format-version": format_version, 'write.parquet.compression-codec': codec}) props_sql = ", ".join(f"'{k}' = '{v}'" for k, v in props.items()) spark.sql(f"CREATE TABLE {full_table} (id BIGINT, data BIGINT) USING ICEBERG " f"TBLPROPERTIES({props_sql})") @@ -709,90 +735,97 @@ def setup_iceberg_table(spark): read_conf = {'spark.rapids.sql.format.parquet.reader.type': reader_type} if error_msg: assert_spark_exception( - lambda : with_gpu_session(lambda spark : spark.sql(query).collect(), conf=read_conf), + lambda : with_gpu_session(lambda spark : spark.sql(query).collect(), conf={**iceberg_read_enabled_conf, **read_conf}), error_msg) else: - assert_gpu_and_cpu_are_equal_collect(lambda spark : spark.sql(query), conf=read_conf) + assert_gpu_and_cpu_are_equal_collect(lambda spark : spark.sql(query), conf={**iceberg_read_enabled_conf, **read_conf}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize("key_gen", [int_gen, long_gen, string_gen, boolean_gen, date_gen, timestamp_gen, decimal_gen_64bit], ids=idfn) @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_read_partition_key(spark_tmp_table_factory, key_gen, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_partition_key(format_version, spark_tmp_table_factory, key_gen, reader_type): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = two_col_df(spark, key_gen, long_gen).orderBy("a") df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG PARTITIONED BY (a) {_NO_FANOUT} " + \ + spark.sql(f"CREATE TABLE {full_table} USING ICEBERG PARTITIONED BY (a) {iceberg_table_properties_sql( + format_version)} " + \ f"AS SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql(f"SELECT a FROM {full_table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_input_meta(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_input_meta(format_version, spark_tmp_table_factory, reader_type): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen).orderBy("a") df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG PARTITIONED BY (a) {_NO_FANOUT} " + \ + spark.sql(f"CREATE TABLE {full_table} USING ICEBERG PARTITIONED BY (a) {iceberg_table_properties_sql( + format_version)} " + \ f"AS SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql( "SELECT a, input_file_name(), input_file_block_start(), input_file_block_length() " + \ f"FROM {full_table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_disorder_read_schema(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_disorder_read_schema(format_version, spark_tmp_table_factory, reader_type): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = three_col_df(spark, long_gen, string_gen, float_gen) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {_NO_FANOUT} " + \ + spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {iceberg_table_properties_sql(format_version)} " + \ f"AS SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql(f"SELECT b,c,a FROM {full_table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering -def test_iceberg_read_appended_table(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_appended_table(format_version, spark_tmp_table_factory): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {_NO_FANOUT} " + \ + spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {iceberg_table_properties_sql(format_version)} " + \ f"AS SELECT * FROM {tmpview}") df = binary_op_df(spark, long_gen, seed=1) df.createOrReplaceTempView(tmpview) spark.sql(f"INSERT INTO {full_table} " + \ f"SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) - assert_gpu_and_cpu_are_equal_collect(lambda spark : spark.sql(f"SELECT * FROM {full_table}")) + assert_gpu_and_cpu_are_equal_collect(lambda spark : spark.sql(f"SELECT * FROM {full_table}"), conf=iceberg_read_enabled_conf) @iceberg # Some metadata files have types that are not supported on the GPU yet (e.g.: BinaryType) @allow_non_gpu("BatchScanExec", "ProjectExec") @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering -def test_iceberg_read_metadata_fallback(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_metadata_fallback(format_version, spark_tmp_table_factory): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {_NO_FANOUT} " + \ + spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {iceberg_table_properties_sql(format_version)} " + \ f"AS SELECT * FROM {tmpview}") df = binary_op_df(spark, long_gen, seed=1) df.createOrReplaceTempView(tmpview) @@ -804,18 +837,19 @@ def setup_iceberg_table(spark): # SQL does not have syntax to read table metadata assert_gpu_fallback_collect( lambda spark : spark.read.format("iceberg").load(f"{full_table}.{subtable}"), - "BatchScanExec") + "BatchScanExec", conf=iceberg_read_enabled_conf) @iceberg # Some metadata files have types that are not supported on the GPU yet (e.g.: BinaryType) @allow_non_gpu("BatchScanExec", "ProjectExec") -def test_iceberg_read_metadata_count(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_metadata_count(format_version, spark_tmp_table_factory): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {_NO_FANOUT} " + \ + spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {iceberg_table_properties_sql(format_version)} " + \ f"AS SELECT * FROM {tmpview}") df = binary_op_df(spark, long_gen, seed=1) df.createOrReplaceTempView(tmpview) @@ -826,18 +860,19 @@ def setup_iceberg_table(spark): "manifests", "partitions", "snapshots"]: # SQL does not have syntax to read table metadata assert_gpu_and_cpu_row_counts_equal( - lambda spark : spark.read.format("iceberg").load(f"{full_table}.{subtable}")) + lambda spark : spark.read.format("iceberg").load(f"{full_table}.{subtable}"), conf=iceberg_read_enabled_conf) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_read_timetravel(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_timetravel(format_version, spark_tmp_table_factory, reader_type): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_snapshots(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {_NO_FANOUT} " + \ + spark.sql(f"CREATE TABLE {full_table} USING ICEBERG {iceberg_table_properties_sql(format_version)} " + \ f"AS SELECT * FROM {tmpview}".format(tmpview)) df = binary_op_df(spark, long_gen, seed=1) df.createOrReplaceTempView(tmpview) @@ -849,19 +884,20 @@ def setup_snapshots(spark): assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.read.option("versionAsOf", first_snapshot_id) \ .format("iceberg").load("{}".format(full_table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_incremental_read(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_incremental_read(format_version, spark_tmp_table_factory, reader_type): full_table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_snapshots(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) spark.sql("CREATE TABLE {} USING ICEBERG ".format(full_table) + \ - _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) df = binary_op_df(spark, long_gen, seed=1) df.createOrReplaceTempView(tmpview) spark.sql("INSERT INTO {} ".format(full_table) + \ @@ -879,19 +915,20 @@ def setup_snapshots(spark): .option("start-snapshot-id", start_snapshot) \ .option("end-snapshot-id", end_snapshot) \ .format("iceberg").load("{}".format(full_table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_reorder_columns(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_reorder_columns(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) spark.sql("CREATE TABLE {} USING ICEBERG ".format(table) + \ - _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) spark.sql("ALTER TABLE {} ALTER COLUMN b FIRST".format(table)) df = binary_op_df(spark, long_gen, seed=1) df.createOrReplaceTempView(tmpview) @@ -900,19 +937,20 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_rename_column(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_rename_column(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) spark.sql("CREATE TABLE {} USING ICEBERG ".format(table) + \ - _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) spark.sql("ALTER TABLE {} RENAME COLUMN a TO c".format(table)) df = binary_op_df(spark, long_gen, seed=1) df.createOrReplaceTempView(tmpview) @@ -921,19 +959,20 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_column_names_swapped(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_column_names_swapped(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) spark.sql("CREATE TABLE {} USING ICEBERG ".format(table) + \ - _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) spark.sql("ALTER TABLE {} RENAME COLUMN a TO c".format(table)) spark.sql("ALTER TABLE {} RENAME COLUMN b TO a".format(table)) spark.sql("ALTER TABLE {} RENAME COLUMN c TO b".format(table)) @@ -944,19 +983,20 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_alter_column_type(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_alter_column_type(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = three_col_df(spark, int_gen, float_gen, DecimalGen(precision=7, scale=3)) df.createOrReplaceTempView(tmpview) spark.sql("CREATE TABLE {} USING ICEBERG ".format(table) + \ - _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) spark.sql("ALTER TABLE {} ALTER COLUMN a TYPE BIGINT".format(table)) spark.sql("ALTER TABLE {} ALTER COLUMN b TYPE DOUBLE".format(table)) spark.sql("ALTER TABLE {} ALTER COLUMN c TYPE DECIMAL(17, 3)".format(table)) @@ -967,19 +1007,20 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_add_column(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_add_column(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) spark.sql("CREATE TABLE {} USING ICEBERG ".format(table) + \ - _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) spark.sql("ALTER TABLE {} ADD COLUMNS (c DOUBLE)".format(table)) df = three_col_df(spark, long_gen, long_gen, double_gen) df.createOrReplaceTempView(tmpview) @@ -988,19 +1029,20 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_remove_column(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_remove_column(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) spark.sql("CREATE TABLE {} USING ICEBERG ".format(table) + \ - _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) spark.sql("ALTER TABLE {} DROP COLUMN a".format(table)) df = unary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) @@ -1009,19 +1051,20 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_add_partition_field(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_add_partition_field(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, int_gen) df.createOrReplaceTempView(tmpview) spark.sql("CREATE TABLE {} USING ICEBERG ".format(table) + \ - _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) spark.sql("ALTER TABLE {} ADD PARTITION FIELD b".format(table)) df = binary_op_df(spark, int_gen) df.createOrReplaceTempView(tmpview) @@ -1030,18 +1073,20 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_drop_partition_field(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_drop_partition_field(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, int_gen) df.createOrReplaceTempView(tmpview) - spark.sql("CREATE TABLE {} (a INT, b INT) USING ICEBERG PARTITIONED BY (b) ".format(table) + _NO_FANOUT) + spark.sql( + "CREATE TABLE {} (a INT, b INT) USING ICEBERG PARTITIONED BY (b) ".format(table) + iceberg_table_properties_sql(format_version)) spark.sql("INSERT INTO {} SELECT * FROM {} ORDER BY b".format(table, tmpview)) spark.sql("ALTER TABLE {} DROP PARTITION FIELD b".format(table)) df = binary_op_df(spark, int_gen) @@ -1051,7 +1096,7 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @@ -1073,28 +1118,31 @@ def setup_iceberg_table(spark): @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_parquet_read_with_input_file(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_parquet_read_with_input_file(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmpview = spark_tmp_table_factory.get() def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) - spark.sql("CREATE TABLE {} USING ICEBERG ".format(table) + _NO_FANOUT + " AS SELECT * FROM {}".format(tmpview)) + spark.sql( + "CREATE TABLE {} USING ICEBERG ".format(table) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmpview)) with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark : spark.sql("SELECT *, input_file_name() FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) @pytest.mark.skipif(not is_iceberg_remote_catalog(), reason="Filecache is only meaningful with remote storage, skipping for local Hadoop filesystem") -def test_iceberg_read_with_filecache(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_with_filecache(format_version, spark_tmp_table_factory, reader_type): """Create a table on CPU, read it twice on GPU with file cache enabled, and verify both reads match the CPU result.""" filecache_enabled = with_gpu_session( - lambda spark: spark.conf.get("spark.rapids.filecache.enabled", "false")) + lambda spark: spark.conf.get("spark.rapids.filecache.enabled", "false"), conf=iceberg_read_enabled_conf) assert filecache_enabled == "true", \ "spark.rapids.filecache.enabled must be set to true to run this test" table = get_full_table_name(spark_tmp_table_factory) @@ -1102,7 +1150,8 @@ def test_iceberg_read_with_filecache(spark_tmp_table_factory, reader_type): def setup_iceberg_table(spark): df = binary_op_df(spark, long_gen) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {table} USING ICEBERG {_NO_FANOUT} AS SELECT * FROM {tmpview}") + spark.sql( + f"CREATE TABLE {table} USING ICEBERG {iceberg_table_properties_sql(format_version)} AS SELECT * FROM {tmpview}") with_cpu_session(setup_iceberg_table) query = f"SELECT * FROM {table}" cpu_result = with_cpu_session(lambda spark: spark.sql(query).collect()) @@ -1111,15 +1160,16 @@ def setup_iceberg_table(spark): filecache_conf = { 'spark.rapids.sql.format.parquet.reader.type': reader_type, } - gpu_result_1 = with_gpu_session(lambda spark: spark.sql(query).collect(), conf=filecache_conf) - gpu_result_2 = with_gpu_session(lambda spark: spark.sql(query).collect(), conf=filecache_conf) + gpu_result_1 = with_gpu_session(lambda spark: spark.sql(query).collect(), conf={**iceberg_read_enabled_conf, **filecache_conf}) + gpu_result_2 = with_gpu_session(lambda spark: spark.sql(query).collect(), conf={**iceberg_read_enabled_conf, **filecache_conf}) assert_equal_with_local_sort(cpu_result, gpu_result_1) assert_equal_with_local_sort(cpu_result, gpu_result_2) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_parquet_read_from_url_encoded_path(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_parquet_read_from_url_encoded_path(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmp_view = spark_tmp_table_factory.get() partition_gen = StringGen(pattern="(.|\n){1,10}", nullable=False)\ @@ -1130,18 +1180,20 @@ def test_iceberg_parquet_read_from_url_encoded_path(spark_tmp_table_factory, rea def setup_iceberg_table(spark): df = two_col_df(spark, long_gen, partition_gen).sortWithinPartitions('b') df.createOrReplaceTempView(tmp_view) - spark.sql("CREATE TABLE {} USING ICEBERG PARTITIONED BY (b) ".format(table) + _NO_FANOUT + " AS SELECT * FROM {}".format(tmp_view)) + spark.sql( + "CREATE TABLE {} USING ICEBERG PARTITIONED BY (b) ".format(table) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmp_view)) with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark: spark.sql("SELECT * FROM {}".format(table)), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) @pytest.mark.skipif(not is_iceberg_rest_catalog(), reason="S3 path handling is exercised only with the REST catalog") -def test_iceberg_parquet_read_from_uri_invalid_s3_path(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_parquet_read_from_uri_invalid_s3_path(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) tmp_view = spark_tmp_table_factory.get() partition_gen = StringGen(pattern="(.|\n){1,10}", nullable=False)\ @@ -1151,21 +1203,23 @@ def setup_iceberg_table(spark): df = two_col_df(spark, long_gen, partition_gen).sortWithinPartitions('b') df.createOrReplaceTempView(tmp_view) spark.sql("CREATE TABLE {} USING ICEBERG PARTITIONED BY (b) ".format(table) + - _NO_FANOUT + " AS SELECT * FROM {}".format(tmp_view)) + iceberg_table_properties_sql(format_version) + " AS SELECT * FROM {}".format(tmp_view)) with_cpu_session(setup_iceberg_table) assert with_gpu_session( lambda spark: spark.sparkContext.getConf().get( - 'spark.rapids.perfio.s3.enabled', 'false') == 'true'), \ + 'spark.rapids.perfio.s3.enabled', 'false') == 'true', conf=iceberg_read_enabled_conf), \ "PerfIO S3 must be enabled at Spark startup for REST catalog tests" assert_gpu_and_cpu_are_equal_collect( lambda spark: spark.sql(f"SELECT * FROM {table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_read_metadata_columns_with_partition_evolution(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_metadata_columns_with_partition_evolution( + format_version, spark_tmp_table_factory, reader_type): """ Test reading Iceberg metadata columns (_file, _pos, _spec_id, _partition) with partition evolution. """ @@ -1175,7 +1229,8 @@ def setup_iceberg_table(spark): # Create table partitioned by a df = three_col_df(spark, long_gen, int_gen, string_gen) df.createOrReplaceTempView(tmpview) - spark.sql(f"CREATE TABLE {table} (a BIGINT, b INT, c STRING) USING ICEBERG PARTITIONED BY (a) {_NO_FANOUT}") + spark.sql( + f"CREATE TABLE {table} (a BIGINT, b INT, c STRING) USING ICEBERG PARTITIONED BY (a) {iceberg_table_properties_sql(format_version)}") spark.sql(f"INSERT INTO {table} SELECT * FROM {tmpview}") # Evolve partition: add b as partition field @@ -1199,13 +1254,14 @@ def setup_iceberg_table(spark): # Test reading all metadata columns along with data columns assert_gpu_and_cpu_are_equal_collect( lambda spark: spark.sql(f"SELECT a, b, c, _file, _pos, _spec_id, _partition FROM {table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_read_pos_with_split_file(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_iceberg_read_pos_with_split_file(format_version, spark_tmp_table_factory, reader_type): # Writes a single Parquet data file containing many row groups, then forces # Iceberg's planner to split that file across multiple scan tasks at row-group # byte boundaries via a tiny row-group size, a tiny split target, and a zero @@ -1218,7 +1274,7 @@ def test_iceberg_read_pos_with_split_file(spark_tmp_table_factory, reader_type): # holds in whichever reader is chosen. table = get_full_table_name(spark_tmp_table_factory) def setup_iceberg_table(spark): - spark.sql(f"CREATE TABLE {table} (id BIGINT) USING ICEBERG {_NO_FANOUT}") + spark.sql(f"CREATE TABLE {table} (id BIGINT) USING ICEBERG {iceberg_table_properties_sql(format_version)}") spark.sql( f"ALTER TABLE {table} SET TBLPROPERTIES (" "'write.parquet.row-group-size-bytes' = '4096', " @@ -1228,7 +1284,7 @@ def setup_iceberg_table(spark): with_cpu_session(setup_iceberg_table) assert_gpu_and_cpu_are_equal_collect( lambda spark: spark.sql(f"SELECT id, _pos FROM {table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @@ -1265,7 +1321,8 @@ def setup_iceberg_table(spark): @iceberg @ignore_order(local=True) @pytest.mark.parametrize('reader_type', rapids_reader_types) -def test_iceberg_small_file_combine_with_schema_evolution(spark_tmp_table_factory, reader_type): +@pytest.mark.parametrize("format_version", iceberg_read_format_versions) +def test_iceberg_small_file_combine_with_schema_evolution(format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) schema_evolution_gens_v1 = [('a', long_gen), ('b', int_gen)] schema_evolution_gens_v2 = schema_evolution_gens_v1 + [('c', string_gen)] @@ -1273,7 +1330,7 @@ def test_iceberg_small_file_combine_with_schema_evolution(spark_tmp_table_factor create_iceberg_table( table, partition_col_sql='bucket(2, a)', - df_gen=lambda spark: gen_df(spark, schema_evolution_gens_v1)) + df_gen=lambda spark: gen_df(spark, schema_evolution_gens_v1), format_version=format_version) def setup_iceberg_table(spark): for seed_offset in range(4): @@ -1303,21 +1360,22 @@ def setup_iceberg_table(spark): assert_gpu_and_cpu_are_equal_collect( lambda spark: spark.sql(f"SELECT a, b, c FROM {table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) @pytest.mark.parametrize('reader_type', rapids_reader_types) +@pytest.mark.parametrize("format_version", iceberg_read_format_versions) def test_iceberg_small_file_combine_with_partition_spec_evolution( - spark_tmp_table_factory, reader_type): + format_version, spark_tmp_table_factory, reader_type): table = get_full_table_name(spark_tmp_table_factory) partition_evolution_gens = [('a', long_gen), ('b', int_gen), ('c', string_gen)] base_seed = get_datagen_seed() create_iceberg_table( table, partition_col_sql='bucket(10, a)', - df_gen=lambda spark: gen_df(spark, partition_evolution_gens)) + df_gen=lambda spark: gen_df(spark, partition_evolution_gens), format_version=format_version) def setup_iceberg_table(spark): for seed_offset in range(4): @@ -1356,21 +1414,22 @@ def setup_iceberg_table(spark): assert_gpu_and_cpu_are_equal_collect( lambda spark: spark.sql(f"SELECT a, b, c, _spec_id, _partition FROM {table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) @iceberg @ignore_order(local=True) @pytest.mark.parametrize('reader_type', rapids_reader_types) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason = "S3tables catalog is managed") +@pytest.mark.parametrize("format_version", iceberg_read_format_versions) def test_iceberg_small_file_combine_with_add_files_identity_partition( - spark_tmp_table_factory, reader_type): + format_version, spark_tmp_table_factory, reader_type): target_table = get_full_table_name(spark_tmp_table_factory) source_table = get_full_table_name(spark_tmp_table_factory) create_iceberg_table( target_table, partition_col_sql='a', - df_gen=lambda spark: spark.createDataFrame([], 'a long, b string')) + df_gen=lambda spark: spark.createDataFrame([], 'a long, b string'), format_version=format_version) def setup_imported_table(spark): spark.sql( @@ -1412,4 +1471,4 @@ def setup_imported_table(spark): # Imported partitioned Parquet files materialize `a` from the path, not the file payload. assert_gpu_and_cpu_are_equal_collect( lambda spark: spark.sql(f"SELECT a, b FROM {target_table}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) diff --git a/integration_tests/src/main/python/iceberg/iceberg_update_test.py b/integration_tests/src/main/python/iceberg/iceberg_update_test.py index 881859bfad1..eda18d09085 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_update_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_update_test.py @@ -18,13 +18,14 @@ assert_gpu_fallback_write_sql from conftest import is_iceberg_remote_catalog from data_gen import * -from iceberg import (create_iceberg_table, get_full_table_name, iceberg_write_enabled_conf, - iceberg_base_table_cols, iceberg_gens_list, iceberg_nested_write_gens_list, - iceberg_unsupported_mark, update_partition_transforms_distributed, - supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON, - supports_iceberg_row_lineage_inheritance, - ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, row_lineage_df, - rapids_reader_types) +from iceberg import ( + iceberg_cow_format_versions, iceberg_mor_format_versions, with_iceberg_dml_session, + with_iceberg_format_versions, create_iceberg_table, get_full_table_name, + iceberg_write_enabled_conf, iceberg_base_table_cols, iceberg_gens_list, + iceberg_nested_write_gens_list, iceberg_unsupported_mark, + update_partition_transforms_distributed, supports_iceberg_v3, ICEBERG_V3_UNSUPPORTED_REASON, + supports_iceberg_row_lineage_inheritance, ICEBERG_ROW_LINEAGE_INHERITANCE_UNSUPPORTED_REASON, + row_lineage_df, rapids_reader_types) from marks import allow_non_gpu, allow_non_gpu_conditional, disable_ansi_mode, iceberg, ignore_order, datagen_overrides from spark_session import is_spark_400_or_later, with_cpu_session, with_gpu_session @@ -45,7 +46,7 @@ def create_iceberg_table_with_data(table_name: str, data_gen_func=None, table_properties=None, update_mode='copy-on-write', - write_order=None): + write_order=None, format_version="2"): """Helper function to create and populate an Iceberg table for UPDATE tests. Args: @@ -56,7 +57,7 @@ def create_iceberg_table_with_data(table_name: str, update_mode: Update mode - 'copy-on-write' or 'merge-on-read' """ base_props = { - 'format-version': '2', + 'format-version': format_version, 'write.update.mode': update_mode } if table_properties: @@ -83,7 +84,7 @@ def insert_data(spark): def do_update_test(spark_tmp_table_factory, update_sql_func, data_gen_func=None, partition_col_sql=None, table_properties=None, update_mode='copy-on-write', conf=iceberg_update_cow_enabled_conf, - read_func=None, write_order=None): + read_func=None, write_order=None, format_version="2"): """ Helper function to test UPDATE operations by comparing CPU and GPU results. @@ -104,15 +105,15 @@ def do_update_test(spark_tmp_table_factory, update_sql_func, data_gen_func=None, # Create identical tables for CPU and GPU create_iceberg_table_with_data(cpu_table_name, partition_col_sql, - data_gen_func, table_properties, update_mode, write_order) + data_gen_func, table_properties, update_mode, write_order, format_version=format_version) create_iceberg_table_with_data(gpu_table_name, partition_col_sql, - data_gen_func, table_properties, update_mode, write_order) + data_gen_func, table_properties, update_mode, write_order, format_version=format_version) # Execute UPDATE on GPU def do_gpu_update(spark): update_sql_func(spark, gpu_table_name) - with_gpu_session(do_gpu_update, conf=conf) + with_iceberg_dml_session(do_gpu_update, format_version, update_mode, conf=conf) # Execute UPDATE on CPU def do_cpu_update(spark): @@ -133,14 +134,17 @@ def do_cpu_update(spark): @iceberg @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) -@pytest.mark.parametrize('update_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,update_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_unpartitioned_table_single_column(spark_tmp_table_factory, update_mode): +def test_iceberg_update_unpartitioned_table_single_column(format_version, spark_tmp_table_factory, update_mode): """Test UPDATE on unpartitioned table with single column update""" do_update_test( spark_tmp_table_factory, lambda spark, table: spark.sql(f"UPDATE {table} SET _c2 = _c2 + 100 WHERE _c2 % 3 = 0"), - update_mode=update_mode + update_mode=update_mode, + format_version=format_version ) @@ -173,7 +177,9 @@ def update_data(spark, table_name): lambda spark, table_name: spark.sql(f"SELECT * FROM {table_name}"), base_table_name, [fallback_exec], - conf=iceberg_update_cow_enabled_conf) + conf=copy_and_update( + iceberg_update_cow_enabled_conf, + {"spark.rapids.sql.format.iceberg.v3.enabled": "false"})) @iceberg @@ -225,24 +231,29 @@ def test_iceberg_v3_row_lineage_gpu_update(spark_tmp_table_factory): @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('update_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,update_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_unpartitioned_table_multiple_columns(spark_tmp_table_factory, update_mode): +def test_iceberg_update_unpartitioned_table_multiple_columns(format_version, spark_tmp_table_factory, update_mode): """Test UPDATE on unpartitioned table with multiple column updates""" do_update_test( spark_tmp_table_factory, lambda spark, table: spark.sql(f"UPDATE {table} SET _c2 = _c2 + 100, _c6 = 'updated' WHERE _c2 % 3 = 0"), - update_mode=update_mode + update_mode=update_mode, + format_version=format_version ) -def _do_test_iceberg_update_partitioned_table_single_column(spark_tmp_table_factory, update_mode, partition_col_sql, table_properties=None): +def _do_test_iceberg_update_partitioned_table_single_column( + spark_tmp_table_factory, update_mode, partition_col_sql, table_properties=None, format_version="2"): """Helper function for partitioned table UPDATE tests.""" do_update_test( spark_tmp_table_factory, lambda spark, table: spark.sql(f"UPDATE {table} SET _c2 = _c2 + 100 WHERE _c2 % 3 = 0"), partition_col_sql=partition_col_sql, table_properties=table_properties, - update_mode=update_mode + update_mode=update_mode, + format_version=format_version ) @@ -250,14 +261,21 @@ def _do_test_iceberg_update_partitioned_table_single_column(spark_tmp_table_fact @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) -@pytest.mark.parametrize('update_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,update_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @pytest.mark.parametrize("partition_col_sql", [ pytest.param("year(_c9)", id="year(timestamp_col)"), ]) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_partitioned_table_single_column(spark_tmp_table_factory, update_mode, partition_col_sql): +def test_iceberg_update_partitioned_table_single_column( + format_version, spark_tmp_table_factory, update_mode, partition_col_sql): """Basic partition test - runs for all catalogs including remote.""" - _do_test_iceberg_update_partitioned_table_single_column(spark_tmp_table_factory, update_mode, partition_col_sql) + _do_test_iceberg_update_partitioned_table_single_column( + spark_tmp_table_factory, + update_mode, + partition_col_sql, + format_version=format_version) @iceberg @@ -265,28 +283,38 @@ def test_iceberg_update_partitioned_table_single_column(spark_tmp_table_factory, @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize("partition_col_sql,update_mode", update_partition_transforms_distributed) +@pytest.mark.parametrize( + 'format_version,partition_col_sql,update_mode', + with_iceberg_format_versions(update_partition_transforms_distributed)) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_partitioned_table_single_column_full_coverage(spark_tmp_table_factory, update_mode, partition_col_sql): +def test_iceberg_update_partitioned_table_single_column_full_coverage( + format_version, spark_tmp_table_factory, update_mode, partition_col_sql): """Sanity-check UPDATE across the two write modes against partition transforms distinct from those picked by other DML ops. The 26-transform partition-writer coverage anchor lives in iceberg_append_test.py::test_insert_into_partitioned_table_full_coverage.""" - _do_test_iceberg_update_partitioned_table_single_column(spark_tmp_table_factory, update_mode, partition_col_sql) + _do_test_iceberg_update_partitioned_table_single_column( + spark_tmp_table_factory, + update_mode, + partition_col_sql, + format_version=format_version) @iceberg @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('update_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,update_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_partitioned_table_multiple_columns(spark_tmp_table_factory, update_mode): +def test_iceberg_update_partitioned_table_multiple_columns(format_version, spark_tmp_table_factory, update_mode): """Test UPDATE on bucket-partitioned table with multiple column updates""" do_update_test( spark_tmp_table_factory, lambda spark, table: spark.sql(f"UPDATE {table} SET _c2 = _c2 + 100, _c6 = 'updated' WHERE _c2 % 3 = 0"), partition_col_sql="year(_c8)", - update_mode=update_mode + update_mode=update_mode, + format_version=format_version ) @@ -295,7 +323,8 @@ def test_iceberg_update_partitioned_table_multiple_columns(spark_tmp_table_facto @disable_ansi_mode @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_mor_then_select_count(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_mor_format_versions) +def test_iceberg_update_mor_then_select_count(format_version, spark_tmp_table_factory): """Test UPDATE with merge-on-read mode, then select count with the same update filter. This test verifies that after a merge-on-read UPDATE operation, subsequent COUNT(*) @@ -307,8 +336,16 @@ def test_iceberg_update_mor_then_select_count(spark_tmp_table_factory): # Phase 1: Initialize tables with data (separate for CPU and GPU) cpu_table_name = f'{base_table_name}_cpu' gpu_table_name = f'{base_table_name}_gpu' - create_iceberg_table_with_data(cpu_table_name, update_mode='merge-on-read', partition_col_sql="hour(_c9)") - create_iceberg_table_with_data(gpu_table_name, update_mode='merge-on-read', partition_col_sql="hour(_c9)") + create_iceberg_table_with_data( + cpu_table_name, + update_mode='merge-on-read', + partition_col_sql="hour(_c9)", + format_version=format_version) + create_iceberg_table_with_data( + gpu_table_name, + update_mode='merge-on-read', + partition_col_sql="hour(_c9)", + format_version=format_version) # Phase 2: Execute UPDATE on both CPU and GPU tables def _do_update(spark, table_name): @@ -318,7 +355,9 @@ def _do_update(spark, table_name): with_cpu_session(lambda spark: _do_update(spark, cpu_table_name)) # UPDATE on GPU - with_gpu_session(lambda spark: _do_update(spark, gpu_table_name), conf=iceberg_update_cow_enabled_conf) + with_iceberg_dml_session( + lambda spark: _do_update(spark, gpu_table_name), + format_version, "merge-on-read", conf=iceberg_update_cow_enabled_conf) # Phase 3: Query COUNT(*) with the same filter and compare results def _query_count(spark, table_name): @@ -328,7 +367,8 @@ def _query_count(spark, table_name): cpu_count = with_cpu_session(lambda spark: _query_count(spark, cpu_table_name)) # Query count on GPU - gpu_count = with_gpu_session(lambda spark: _query_count(spark, gpu_table_name)) + gpu_count = with_gpu_session( + lambda spark: _query_count(spark, gpu_table_name), conf=iceberg_update_cow_enabled_conf) # Phase 4: Compare CPU and GPU counts assert cpu_count == gpu_count, f"Count mismatch: CPU={cpu_count}, GPU={gpu_count}" @@ -339,20 +379,21 @@ def _query_count(spark, table_name): @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('update_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,update_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) +])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_fallback_write_disabled(spark_tmp_table_factory, update_mode, fallback_exec): +def test_iceberg_update_fallback_write_disabled( + format_version, spark_tmp_table_factory, update_mode, fallback_exec): """Test UPDATE falls back when Iceberg write is disabled""" base_table_name = get_full_table_name(spark_tmp_table_factory) # Phase 1: Initialize tables with data (separate for CPU and GPU) cpu_table_name = f'{base_table_name}_cpu' gpu_table_name = f'{base_table_name}_gpu' - create_iceberg_table_with_data(cpu_table_name, update_mode=update_mode) - create_iceberg_table_with_data(gpu_table_name, update_mode=update_mode) + create_iceberg_table_with_data(cpu_table_name, update_mode=update_mode, format_version=format_version) + create_iceberg_table_with_data(gpu_table_name, update_mode=update_mode, format_version=format_version) # Phase 2: UPDATE operation (to be tested with fallback) def write_func(spark, table_name): @@ -377,13 +418,14 @@ def read_func(spark, table_name): @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('update_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,update_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) +])) @pytest.mark.parametrize("file_format", ["orc", "avro"], ids=lambda x: f"file_format={x}") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_fallback_unsupported_file_format(spark_tmp_table_factory, file_format, update_mode, fallback_exec): +def test_iceberg_update_fallback_unsupported_file_format( + format_version, spark_tmp_table_factory, file_format, update_mode, fallback_exec): """Test UPDATE falls back with unsupported file formats (ORC, Avro) This test creates a table with parquet format, inserts data, then changes the @@ -399,7 +441,7 @@ def data_gen(spark): def init_table(table_name): # Step 1: Create table with parquet as default write format table_props = { - 'format-version': '2', + 'format-version': format_version, 'write.update.mode': update_mode, 'write.format.default': 'parquet' } @@ -449,8 +491,10 @@ def read_func(spark, table_name): @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -@pytest.mark.parametrize('update_mode', ['copy-on-write', 'merge-on-read']) -def test_iceberg_update_nested_types(spark_tmp_table_factory, update_mode): +@pytest.mark.parametrize( + 'format_version,update_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) +def test_iceberg_update_nested_types(format_version, spark_tmp_table_factory, update_mode): """Test UPDATE with supported nested types.""" cols = [f"_c{idx}" for idx, _ in enumerate(iceberg_nested_write_gens_list)] data_gen_func = lambda spark: gen_df(spark, list(zip(cols, iceberg_nested_write_gens_list))) @@ -459,7 +503,8 @@ def test_iceberg_update_nested_types(spark_tmp_table_factory, update_mode): spark_tmp_table_factory, lambda spark, table: spark.sql(f"UPDATE {table} SET _c0 = _c0 + 100 WHERE _c0 % 3 = 0"), data_gen_func=data_gen_func, - update_mode=update_mode + update_mode=update_mode, + format_version=format_version ) @allow_non_gpu("ReplaceDataExec", "WriteDeltaExec", "BatchScanExec", "ColumnarToRowExec") @@ -467,20 +512,21 @@ def test_iceberg_update_nested_types(spark_tmp_table_factory, update_mode): @ignore_order(local=True) @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -@pytest.mark.parametrize('update_mode,fallback_exec', [ +@pytest.mark.parametrize('format_version,update_mode,fallback_exec', with_iceberg_format_versions([ pytest.param('copy-on-write', 'ReplaceDataExec', id='cow'), pytest.param('merge-on-read', 'WriteDeltaExec', id='mor') -]) +])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_fallback_iceberg_disabled(spark_tmp_table_factory, update_mode, fallback_exec): +def test_iceberg_update_fallback_iceberg_disabled( + format_version, spark_tmp_table_factory, update_mode, fallback_exec): """Test UPDATE falls back when Iceberg is completely disabled""" base_table_name = get_full_table_name(spark_tmp_table_factory) # Phase 1: Initialize tables with data (separate for CPU and GPU) cpu_table_name = f'{base_table_name}_cpu' gpu_table_name = f'{base_table_name}_gpu' - create_iceberg_table_with_data(cpu_table_name, update_mode=update_mode) - create_iceberg_table_with_data(gpu_table_name, update_mode=update_mode) + create_iceberg_table_with_data(cpu_table_name, update_mode=update_mode, format_version=format_version) + create_iceberg_table_with_data(gpu_table_name, update_mode=update_mode, format_version=format_version) # Phase 2: UPDATE operation (to be tested with fallback) def write_func(spark, table_name): @@ -506,7 +552,8 @@ def read_func(spark, table_name): @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_mor_fallback_writedelta_disabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_mor_format_versions) +def test_iceberg_update_mor_fallback_writedelta_disabled(format_version, spark_tmp_table_factory): """Test merge-on-read UPDATE falls back when WriteDeltaExec is disabled This test verifies that when WriteDeltaExec is explicitly disabled (it's disabled by default @@ -517,8 +564,8 @@ def test_iceberg_update_mor_fallback_writedelta_disabled(spark_tmp_table_factory # Phase 1: Initialize tables with data (separate for CPU and GPU) cpu_table_name = f'{base_table_name}_cpu' gpu_table_name = f'{base_table_name}_gpu' - create_iceberg_table_with_data(cpu_table_name, update_mode='merge-on-read') - create_iceberg_table_with_data(gpu_table_name, update_mode='merge-on-read') + create_iceberg_table_with_data(cpu_table_name, update_mode='merge-on-read', format_version=format_version) + create_iceberg_table_with_data(gpu_table_name, update_mode='merge-on-read', format_version=format_version) # Phase 2: UPDATE operation (to be tested with fallback) def write_func(spark, table_name): @@ -542,18 +589,20 @@ def read_func(spark, table_name): @iceberg @ignore_order(local=True) -@pytest.mark.parametrize('update_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,update_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @pytest.mark.parametrize("partition_col_sql", [ pytest.param(None, id="unpartitioned"), pytest.param("year(_c9)", id="year_partition"), ]) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_update_aqe(spark_tmp_table_factory, update_mode, partition_col_sql): +def test_update_aqe(format_version, spark_tmp_table_factory, update_mode, partition_col_sql): """ Test UPDATE with AQE enabled. """ table_prop = { - 'format-version': '2', + 'format-version': format_version, 'write.update.mode': update_mode } @@ -577,7 +626,8 @@ def initialize_table(table_name): def update_table(spark, table_name): spark.sql(f"UPDATE {table_name} SET _c2 = _c2 + 1 WHERE _c0 > 0") - with_gpu_session(lambda spark: update_table(spark, gpu_table), conf=conf) + with_iceberg_dml_session( + lambda spark: update_table(spark, gpu_table), format_version, update_mode, conf=conf) with_cpu_session(lambda spark: update_table(spark, cpu_table), conf=conf) cpu_data = with_cpu_session(lambda spark: spark.table(cpu_table).collect()) @@ -589,9 +639,11 @@ def update_table(spark, table_name): @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") @pytest.mark.datagen_overrides(seed=UPDATE_TEST_SEED, reason=UPDATE_TEST_SEED_OVERRIDE_REASON) -@pytest.mark.parametrize('update_mode', ['copy-on-write', 'merge-on-read']) +@pytest.mark.parametrize( + 'format_version,update_mode', + with_iceberg_format_versions(['copy-on-write', 'merge-on-read'])) @allow_non_gpu_conditional(is_spark_400_or_later(), "EmptyRelationExec") -def test_iceberg_update_after_drop_partition_field(spark_tmp_table_factory, update_mode): +def test_iceberg_update_after_drop_partition_field(format_version, spark_tmp_table_factory, update_mode): """Test UPDATE on table after dropping a partition field (void transform). When a partition field is dropped, Iceberg creates a 'void transform' - @@ -607,9 +659,9 @@ def test_iceberg_update_after_drop_partition_field(spark_tmp_table_factory, upda # Create partitioned tables with data create_iceberg_table_with_data(cpu_table_name, partition_col_sql=partition_col_sql, - update_mode=update_mode) + update_mode=update_mode, format_version=format_version) create_iceberg_table_with_data(gpu_table_name, partition_col_sql=partition_col_sql, - update_mode=update_mode) + update_mode=update_mode, format_version=format_version) # Drop one partition field on both tables (creates void transform) def drop_partition_field(spark, table_name): @@ -622,8 +674,9 @@ def drop_partition_field(spark, table_name): def do_update(spark, table_name): spark.sql(f"UPDATE {table_name} SET _c2 = _c2 + 100 WHERE _c2 % 3 = 0") - with_gpu_session(lambda spark: do_update(spark, gpu_table_name), - conf=iceberg_update_cow_enabled_conf) + with_iceberg_dml_session( + lambda spark: do_update(spark, gpu_table_name), + format_version, update_mode, conf=iceberg_update_cow_enabled_conf) with_cpu_session(lambda spark: do_update(spark, cpu_table_name)) # Compare results @@ -637,10 +690,11 @@ def do_update(spark, table_name): @datagen_overrides(seed=0, reason='https://github.com/NVIDIA/spark-rapids-jni/issues/4016') @ignore_order(local=True) @pytest.mark.skipif(is_iceberg_remote_catalog(), reason="Skip for remote catalog to reduce test time") -def test_iceberg_update_partitioned_table_fanout_enabled(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_cow_format_versions) +def test_iceberg_update_partitioned_table_fanout_enabled(format_version, spark_tmp_table_factory): # Use bucket(2, ...) to keep partition count low and avoid OOM from Iceberg's FanoutDataWriter. _do_test_iceberg_update_partitioned_table_single_column( spark_tmp_table_factory, update_mode='copy-on-write', partition_col_sql="bucket(2, _c9)", - table_properties={"write.spark.fanout.enabled": "true"}) + table_properties={"write.spark.fanout.enabled": "true"}, format_version=format_version) diff --git a/integration_tests/src/main/python/iceberg/iceberg_view_test.py b/integration_tests/src/main/python/iceberg/iceberg_view_test.py index a7d78cecc5d..90c2af36cd9 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_view_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_view_test.py @@ -18,7 +18,10 @@ from asserts import assert_gpu_and_cpu_are_equal_collect from conftest import is_iceberg_rest_catalog from data_gen import * -from iceberg import get_full_table_name, rapids_reader_types, create_iceberg_table, iceberg_base_table_cols, iceberg_gens_list, iceberg_unsupported_mark +from iceberg import ( + iceberg_format_versions, iceberg_read_enabled_conf, iceberg_read_format_versions, + iceberg_table_properties_sql, get_full_table_name, rapids_reader_types, create_iceberg_table, + iceberg_base_table_cols, iceberg_gens_list, iceberg_unsupported_mark) from marks import iceberg, ignore_order from spark_session import with_cpu_session @@ -40,7 +43,8 @@ pytest.param("SELECT _c0, _c2, _c6 FROM {table_name}", id="projection"), pytest.param("SELECT _c7, COUNT(*) as cnt, SUM(_c2) as sum_c2 FROM {table_name} GROUP BY _c7", id="aggregation"), ]) -def test_iceberg_view(spark_tmp_table_factory, reader_type, view_sql): +@pytest.mark.parametrize("format_version", iceberg_read_format_versions) +def test_iceberg_view(format_version, spark_tmp_table_factory, reader_type, view_sql): """Test reading from an Iceberg view.""" table_name = get_full_table_name(spark_tmp_table_factory) @@ -48,7 +52,7 @@ def test_iceberg_view(spark_tmp_table_factory, reader_type, view_sql): view_name = "iceberg_view_" + view_uuid # Create an Iceberg table, and insert data into it - create_iceberg_table(table_name) + create_iceberg_table(table_name, format_version=format_version) def insert_data(spark): df = gen_df(spark, list(zip(iceberg_base_table_cols, iceberg_gens_list))) df.writeTo(table_name).append() @@ -61,4 +65,4 @@ def setup_iceberg_view(spark): assert_gpu_and_cpu_are_equal_collect( lambda spark: spark.sql(f"SELECT * FROM {view_name}"), - conf={'spark.rapids.sql.format.parquet.reader.type': reader_type}) + conf={**iceberg_read_enabled_conf, 'spark.rapids.sql.format.parquet.reader.type': reader_type}) diff --git a/integration_tests/src/main/python/iceberg/iceberg_write_sql_ui_test.py b/integration_tests/src/main/python/iceberg/iceberg_write_sql_ui_test.py index 760de7a8aa6..cdf5f32bd5e 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_write_sql_ui_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_write_sql_ui_test.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from iceberg import get_full_table_name, iceberg_write_enabled_conf, \ - iceberg_unsupported_mark, _BASE_TBLPROPS_SQL +from iceberg import ( + iceberg_format_versions, iceberg_table_properties_sql, get_full_table_name, + iceberg_write_enabled_conf, iceberg_unsupported_mark, _BASE_TBLPROPS_SQL) +import pytest + from marks import allow_non_gpu, iceberg from spark_session import with_gpu_session @@ -34,7 +37,8 @@ def _is_write_node(name): # write execution we actually assert on. @allow_non_gpu('CreateTableExec') @iceberg -def test_v2_write_sql_ui_shows_gpu_child_operators(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_v2_write_sql_ui_shows_gpu_child_operators(format_version, spark_tmp_table_factory): """Regression test: the SQL UI / History Server must show the GPU child operators under a DataSource V2 table write (GpuV2TableWriteExec), not just the write node. GpuV2TableWriteExec executes a columnar copy of the query's @@ -61,7 +65,7 @@ def max_execution_id(): return mx spark.sql(f"CREATE TABLE {table_name} (grp BIGINT, cnt BIGINT) " - f"USING ICEBERG {_BASE_TBLPROPS_SQL}") + f"USING ICEBERG {iceberg_table_properties_sql(format_version)}") # The integration-test Spark session is shared across tests, so the SQL status # store accumulates executions from earlier tests -- including CPU V2 writes # (e.g. a VALUES-based AppendData over a LocalTableScan / Scan ExistingRDD). We @@ -111,7 +115,8 @@ def max_execution_id(): @allow_non_gpu('CreateTableExec') @iceberg -def test_v2_write_sql_ui_gpu_child_operator_metrics_are_visible(spark_tmp_table_factory): +@pytest.mark.parametrize("format_version", iceberg_format_versions) +def test_v2_write_sql_ui_gpu_child_operator_metrics_are_visible(format_version, spark_tmp_table_factory): """Regression test: the SQL UI / History Server must show metric values (op times) on the GPU child operators of a DataSource V2 table write, not just their names. GpuV2TableWriteExec executes a columnar copy of the query's @@ -138,7 +143,7 @@ def max_execution_id(): return mx spark.sql(f"CREATE TABLE {table_name} (grp BIGINT, cnt BIGINT) " - f"USING ICEBERG {_BASE_TBLPROPS_SQL}") + f"USING ICEBERG {iceberg_table_properties_sql(format_version)}") # See test_v2_write_sql_ui_shows_gpu_child_operators: drain the listener bus # and scope to executions created by our own INSERT (the IT Spark session is # shared, so the status store accumulates executions from earlier tests).