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 0525904f72c..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,8 +21,9 @@ 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 +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 @@ -47,8 +48,8 @@ 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.vectorized.ColumnarBatch +import org.apache.spark.sql.types.{LongType, StructType} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.util.SerializableConfiguration @@ -64,7 +65,7 @@ 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 // Since these are private classes, we check the class name to determine which GPU version // to use val cpuBatch = cpu.toBatch @@ -374,11 +375,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,23 +410,131 @@ 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 GpuDataWriter { + protected def dataSparkType: StructType + protected def metadataSchema: StructType + + 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) { + GpuColumnVector.combineColumns(record) + } else { + 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 => + 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, record.numRows())) { lineage => + GpuColumnVector.combineColumns(record, lineage) + } + } + } +} + class GpuUnpartitionedDataWriter( val fileWriterFactory: GpuSparkFileWriterFactory, val fileFactory: OutputFileFactory, val io: FileIO, val spec: PartitionSpec, - val targetFileSize: Long) - extends DataWriter[ColumnarBatch] { + val targetFileSize: Long, + override protected val metadataSchema: StructType) + extends GpuDataWriterWithRowLineage { + override protected def dataSparkType: StructType = fileWriterFactory.dataSparkType + private val delegate = new GpuRollingDataWriter( fileWriterFactory, fileFactory, @@ -460,10 +576,11 @@ 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] { + override protected val metadataSchema: StructType, +) extends GpuDataWriterWithRowLineage { private val delegate: PartitioningWriter[SpillableColumnarBatch, DataWriteResult] = if (fanoutEnabled) { 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 fae07bb9999..e66893a7d1b 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_ctas_test.py @@ -17,18 +17,19 @@ 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 +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) -from marks import iceberg, ignore_order, allow_non_gpu, allow_non_gpu_conditional, datagen_overrides +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 pytestmark = [ @@ -79,7 +80,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 @@ -87,25 +90,32 @@ 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) @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))) @@ -130,7 +140,33 @@ 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 +@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): + conf = copy_and_update(iceberg_write_enabled_conf, { + "spark.rapids.sql.format.iceberg.v3.enabled": "true" + }) + + 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") + callback.assertNotContain(ctas_plan, "AtomicCreateTableAsSelectExec") + + _assert_gpu_equals_cpu_ctas( + spark_tmp_table_factory, + lambda spark: row_lineage_df(spark), + {"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 @@ -167,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))) @@ -191,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 @@ -202,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 @@ -214,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 } @@ -241,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): @@ -265,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) @@ -279,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)] @@ -295,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)] @@ -316,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" @@ -352,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. @@ -365,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))) @@ -387,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 ce680ba15c4..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 @@ -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, {}) @@ -47,11 +50,12 @@ 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, 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: @@ -68,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() @@ -75,7 +81,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, write_order=None, format_version="2"): """ Helper function to test DELETE operations by comparing CPU and GPU results. @@ -86,6 +93,9 @@ 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 + 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" @@ -93,25 +103,29 @@ 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, format_version=format_version) 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, 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=iceberg_delete_cow_enabled_conf) + with_iceberg_dml_session(do_gpu_delete, format_version, delete_mode, 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) @@ -119,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 ) @@ -167,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 @@ -200,14 +219,35 @@ def setup_iceberg_table(spark): }) -def _do_test_iceberg_delete_partitioned_table(spark_tmp_table_factory, partition_col_sql, delete_mode, table_properties=None): +@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) +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: row_lineage_df(spark, start=1), + table_properties={"format-version": "3"}, + 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}"), + write_order="id") + + +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 ) @@ -219,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 @@ -232,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 ) @@ -263,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): @@ -301,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) @@ -318,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' } @@ -367,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))) @@ -378,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 ) @@ -387,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): @@ -424,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 @@ -435,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): @@ -461,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 } @@ -496,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()) @@ -508,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' - @@ -526,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): @@ -541,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 @@ -555,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 0d26fd341aa..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 @@ -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, @@ -41,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. @@ -55,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' @@ -95,6 +98,23 @@ 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, + 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_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) + 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, @@ -102,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. @@ -122,34 +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) - - # 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) + 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) - # 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, 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 @@ -161,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 ) @@ -169,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( @@ -224,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 @@ -264,26 +300,73 @@ 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) +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_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.sql(f"ALTER TABLE {table} WRITE ORDERED BY id").collect() + row_lineage_df(spark, with_value=True).writeTo(table).append() + + def merge(spark, table): + 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 " + "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}"), + conf) + + @allow_non_gpu("MergeRows$Keep", "MergeRows$Discard", "MergeRows$Split") @iceberg @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)"), @@ -333,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( @@ -359,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) @@ -385,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): @@ -419,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) @@ -435,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' } @@ -508,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 = """ @@ -523,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 ) @@ -531,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) @@ -543,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""" @@ -574,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 @@ -586,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""" @@ -622,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 @@ -658,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()) @@ -669,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' - @@ -688,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): @@ -713,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 @@ -728,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 9c5ad714302..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, 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 @@ -50,6 +52,44 @@ 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 _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) + + 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) + + 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) + 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 @@ -119,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", @@ -167,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) @@ -180,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", @@ -249,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) @@ -281,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) @@ -290,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})") @@ -328,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)) @@ -337,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) @@ -371,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"}) @@ -398,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}) " @@ -466,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 @@ -474,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 @@ -551,6 +614,35 @@ 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) +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')") + row_lineage_df(spark, start=1).writeTo(table).append() + spark.sql( + 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')") + + 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, + append_data, + lambda spark, table: spark.sql( + f"SELECT id, _pos, _row_id, _last_updated_sequence_number FROM {table}")) + + @iceberg @ignore_order(local=True) @pytest.mark.skipif( @@ -584,6 +676,35 @@ 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) +def test_iceberg_v3_row_lineage_insert_overwrite(spark_tmp_table_factory): + source_view = spark_tmp_table_factory.get() + + def setup_iceberg_table(spark, table): + spark.sql( + f"CREATE TABLE {table} (id BIGINT, v BIGINT) USING ICEBERG " + "TBLPROPERTIES ('format-version' = '3')") + row_lineage_df(spark, with_value=True).writeTo(table).append() + + def overwrite(spark, table): + 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( + 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 @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 @@ -596,14 +717,15 @@ def setup_iceberg_table(spark): 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})") @@ -613,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) @@ -708,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) @@ -730,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) @@ -753,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) + \ @@ -783,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) @@ -804,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) @@ -825,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)) @@ -848,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)) @@ -871,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) @@ -892,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) @@ -913,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) @@ -934,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) @@ -955,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 @@ -977,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) @@ -1006,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()) @@ -1015,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)\ @@ -1034,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)\ @@ -1055,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. """ @@ -1079,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 @@ -1103,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 @@ -1122,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', " @@ -1132,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 @@ -1169,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)] @@ -1177,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): @@ -1207,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): @@ -1260,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( @@ -1316,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 c30f1384783..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 @@ -32,6 +33,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 @@ -41,7 +45,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, format_version="2"): """Helper function to create and populate an Iceberg table for UPDATE tests. Args: @@ -52,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: @@ -69,6 +74,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() @@ -76,7 +83,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, write_order=None, format_version="2"): """ Helper function to test UPDATE operations by comparing CPU and GPU results. @@ -87,6 +95,9 @@ 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 + 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" @@ -94,39 +105,46 @@ 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, format_version=format_version) 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, 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=iceberg_update_cow_enabled_conf) + with_iceberg_dml_session(do_gpu_update, format_version, update_mode, 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) @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 ) @@ -159,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 @@ -188,28 +208,52 @@ 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) +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: row_lineage_df(spark, with_value=True), + table_properties={"format-version": "3"}, + 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}"), + write_order="id") + + @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_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 ) @@ -217,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 @@ -232,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 ) @@ -262,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(*) @@ -274,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): @@ -285,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): @@ -295,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}" @@ -306,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): @@ -344,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 @@ -366,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' } @@ -416,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))) @@ -426,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") @@ -434,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): @@ -473,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 @@ -484,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): @@ -509,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 } @@ -544,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()) @@ -556,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' - @@ -574,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): @@ -589,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 @@ -604,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). 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/scala/com/nvidia/spark/rapids/GpuWrite.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuWrite.scala index aa53567c4e0..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 @@ -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. @@ -16,12 +16,21 @@ 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 trait GpuWrite extends Write { var metrics: Map[String, GpuMetric] = Map.empty } +trait GpuDataWriterFactory extends DataWriterFactory { + def createWriter( + partitionId: Int, + taskId: Long, + metadataSchema: StructType): DataWriter[InternalRow] +} + // Allows use of GpuWrite from Java code abstract class GpuWriteWrapper extends GpuWrite { } 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/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 d1c4953e307..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)) } @@ -529,4 +550,3 @@ class GpuMergeBatchIterator( sourcePresent, attemptMetrics) } } - 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..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 @@ -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. @@ -433,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(_)) @@ -445,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) { _ => @@ -479,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) } } } @@ -500,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(_)) @@ -519,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) @@ -570,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/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/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/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..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.{GpuDataWriter, GpuDataWriterFactory, GpuWrite} import com.nvidia.spark.rapids.Arm.withResource -import com.nvidia.spark.rapids.GpuWrite 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,14 +63,31 @@ 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 => - writer.write(projected) + metadataProjection match { + case Some(projection) => + withResource(projection.project(batch)) { metadata => + writer.write(metadata, projected) + } + case None => writer.write(projected) + } } } } 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) + } + } + } +}