From 00efdbc4bb07927451d0d365709100238792254e Mon Sep 17 00:00:00 2001 From: Hazmi Date: Thu, 3 Sep 2026 16:23:08 +0000 Subject: [PATCH] [VL][Iceberg] Support SQL MERGE INTO --- .../TestGlutenCopyOnWriteDelete.java | 18 +++ .../enhanced/VeloxIcebergSuite.scala | 64 +++++++++++ .../execution/IcebergScanTransformer.scala | 104 ++++++++++-------- .../source/GlutenIcebergSourceUtil.scala | 11 +- .../execution/BasicScanExecTransformer.scala | 12 +- .../execution/BatchScanExecTransformer.scala | 6 +- .../parquet/GlutenParquetRowIndexSuite.scala | 7 +- .../parquet/GlutenParquetRowIndexSuite.scala | 7 +- 8 files changed, 174 insertions(+), 55 deletions(-) diff --git a/backends-velox/src-iceberg-spark34/test/java/org/apache/gluten/extensions/TestGlutenCopyOnWriteDelete.java b/backends-velox/src-iceberg-spark34/test/java/org/apache/gluten/extensions/TestGlutenCopyOnWriteDelete.java index d51d04cf194..a1909b54a99 100644 --- a/backends-velox/src-iceberg-spark34/test/java/org/apache/gluten/extensions/TestGlutenCopyOnWriteDelete.java +++ b/backends-velox/src-iceberg-spark34/test/java/org/apache/gluten/extensions/TestGlutenCopyOnWriteDelete.java @@ -16,12 +16,30 @@ */ package org.apache.gluten.extensions; +import org.apache.gluten.execution.IcebergScanTransformer; + +import org.apache.iceberg.spark.extensions.SparkPlanUtil; import org.apache.iceberg.spark.extensions.TestCopyOnWriteDelete; +import org.apache.spark.sql.execution.SparkPlan; +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec; import org.junit.Test; +import java.util.List; import java.util.concurrent.ExecutionException; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; public class TestGlutenCopyOnWriteDelete extends TestCopyOnWriteDelete { + @Override + protected void assertAllBatchScansVectorized(SparkPlan plan) { + List scans = + SparkPlanUtil.collectLeaves(plan).stream() + .filter(scan -> scan instanceof BatchScanExec || scan instanceof IcebergScanTransformer) + .collect(Collectors.toList()); + assertThat(scans).hasSizeGreaterThan(0).allMatch(SparkPlan::supportsColumnar); + } + @Test public synchronized void testDeleteWithConcurrentTableRefresh() { System.out.println("Run timeout"); diff --git a/backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala b/backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala index 00c6e54b58b..4d99db0de52 100644 --- a/backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala +++ b/backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala @@ -161,6 +161,70 @@ class VeloxIcebergSuite extends IcebergSuite { } } + test("iceberg SQL merge into with native scan and write") { + withTable("iceberg_merge_target", "iceberg_merge_source") { + spark.sql(""" + |create table iceberg_merge_target ( + | id int, + | name string, + | p string + |) using iceberg + |tblproperties ( + | 'format-version' = '2', + | 'write.merge.mode' = 'copy-on-write' + |) + |partitioned by (p) + |""".stripMargin) + spark.sql(""" + |create table iceberg_merge_source ( + | id int, + | name string, + | p string + |) using iceberg + |""".stripMargin) + + spark.sql(""" + |insert into iceberg_merge_target values + | (1, 'old-1', 'p1'), + | (2, 'old-2', 'p1'), + | (3, 'old-3', 'p2') + |""".stripMargin) + spark.sql(""" + |insert into iceberg_merge_source values + | (2, 'new-2', 'p2'), + | (3, 'new-3', 'p1'), + | (4, 'new-4', 'p2') + |""".stripMargin) + + val df = spark.sql(""" + |merge into iceberg_merge_target t + |using iceberg_merge_source s + |on t.id = s.id + |when matched then update set + | name = s.name, + | p = s.p + |when not matched then insert (id, name, p) + | values (s.id, s.name, s.p) + |""".stripMargin) + + val commandPlan = + df.queryExecution.executedPlan.asInstanceOf[CommandResultExec].commandPhysicalPlan + assert(commandPlan.isInstanceOf[VeloxIcebergReplaceDataExec]) + assert( + collect(commandPlan) { + case scan: IcebergScanTransformer => scan + }.size == 2) + + checkAnswer( + spark.sql("select * from iceberg_merge_target order by id"), + Seq( + Row(1, "old-1", "p1"), + Row(2, "new-2", "p2"), + Row(3, "new-3", "p1"), + Row(4, "new-4", "p2"))) + } + } + test("iceberg insert partition table bucket transform") { withTable("iceberg_tb2") { spark.sql(""" diff --git a/gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergScanTransformer.scala b/gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergScanTransformer.scala index 6ac22b956ae..f9f172ae64c 100644 --- a/gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergScanTransformer.scala +++ b/gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergScanTransformer.scala @@ -26,7 +26,7 @@ import org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat import org.apache.spark.Partition import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, DynamicPruningExpression, Expression, Literal} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DynamicPruningExpression, Expression, Literal} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.connector.catalog.Table import org.apache.spark.sql.connector.read.Scan @@ -36,7 +36,7 @@ import org.apache.spark.sql.types.{ArrayType, DataType, StructType} import org.apache.iceberg.{BaseTable, MetadataColumns, SnapshotSummary, TableProperties} import org.apache.iceberg.avro.AvroSchemaUtil -import org.apache.iceberg.spark.source.{GlutenIcebergSourceUtil, SparkTable} +import org.apache.iceberg.spark.source.GlutenIcebergSourceUtil import org.apache.iceberg.spark.source.metrics.NumSplits import org.apache.iceberg.types.{Type, Types} import org.apache.iceberg.types.Type.TypeID @@ -92,18 +92,14 @@ case class IcebergScanTransformer( } if (!BackendsApiManager.getSettings.supportIcebergEqualityDeleteRead()) { - val notSupport = table match { - case t: SparkTable => - t.table() match { - case t: BaseTable => - t.operations() - .current() - .schema() - .columns() - .stream - .anyMatch(c => containsUuidOrFixedType(c.`type`()) || containsMetadataColumn(c)) - case _ => false - } + val notSupport = GlutenIcebergSourceUtil.getTable(scan) match { + case t: BaseTable => + t.operations() + .current() + .schema() + .columns() + .stream + .anyMatch(c => containsUuidOrFixedType(c.`type`()) || containsMetadataColumn(c)) case _ => false } if (notSupport) { @@ -111,7 +107,12 @@ case class IcebergScanTransformer( } // Allow input_file_name() and related metadata functions val allowedMetadataColumns = - IcebergScanTransformer.InputFileRelatedMetadataColumnNames + if (GlutenIcebergSourceUtil.isSparkCopyOnWriteScan(scan)) { + IcebergScanTransformer.InputFileRelatedMetadataColumnNames ++ + IcebergScanTransformer.CopyOnWriteMetadataColumnNames + } else { + IcebergScanTransformer.InputFileRelatedMetadataColumnNames + } val hasUnsupportedMetadata = scan.readSchema().fieldNames.exists { f => MetadataColumns.isMetadataColumn(f) && @@ -120,24 +121,12 @@ case class IcebergScanTransformer( if (hasUnsupportedMetadata) { return ValidationResult.failed("Read unsupported metadata column") } - val containsEqualityDelete = table match { - case t: SparkTable => - t.table() match { - case t: BaseTable => - val snapshot = t - .operations() - .current() - .currentSnapshot() - if (snapshot == null) { - false - } else { - snapshot - .summary() - .getOrDefault(SnapshotSummary.TOTAL_EQ_DELETES_PROP, "0") - .toInt > 0 - } - case _ => false - } + // Row-level operations wrap the Spark table, so use the table held by the scan. + val containsEqualityDelete = GlutenIcebergSourceUtil.getTable(scan) match { + case t: BaseTable => + val snapshot = t.operations().current().currentSnapshot() + snapshot != null && + snapshot.summary().getOrDefault(SnapshotSummary.TOTAL_EQ_DELETES_PROP, "0").toInt > 0 case _ => false } if (containsEqualityDelete) { @@ -150,12 +139,8 @@ case class IcebergScanTransformer( } } - val baseTable = table match { - case t: SparkTable => - t.table() match { - case t: BaseTable => t - case _ => null - } + val baseTable = GlutenIcebergSourceUtil.getTable(scan) match { + case t: BaseTable => t case _ => null } if (baseTable == null) { @@ -206,13 +191,28 @@ case class IcebergScanTransformer( !readSchemaFields.contains(name) } + private lazy val copyOnWriteFilePathMetadataColumns = output.filter { + attr => IcebergScanTransformer.isCopyOnWriteFilePathColumn(attr.name) + } + override def getMetadataColumns(): Seq[AttributeReference] = { - val extraMetadataColumns = inputFileRelatedMetadataColumns.filterNot { - metadataAttr => metadataColumns.exists(_.name.equalsIgnoreCase(metadataAttr.name)) - } + val extraMetadataColumns = + (inputFileRelatedMetadataColumns ++ copyOnWriteFilePathMetadataColumns).filterNot { + metadataAttr => metadataColumns.exists(_.name.equalsIgnoreCase(metadataAttr.name)) + } metadataColumns ++ extraMetadataColumns } + override protected def isRowIndexMetadataColumn(attr: Attribute): Boolean = { + IcebergScanTransformer.isCopyOnWriteRowPositionColumn(attr.name) || + super.isRowIndexMetadataColumn(attr) + } + + override protected def isNativeMetadataColumn(attr: Attribute): Boolean = { + IcebergScanTransformer.isCopyOnWriteFilePathColumn(attr.name) || + super.isNativeMetadataColumn(attr) + } + override lazy val fileFormat: ReadFileFormat = GlutenIcebergSourceUtil.getFileFormat(scan) override def getSplitInfosFromPartitions( @@ -270,12 +270,8 @@ case class IcebergScanTransformer( private[execution] def getKeyGroupPartitioning: Option[Seq[Expression]] = keyGroupedPartitioning private def hasRenamedColumn: Boolean = { - val icebergTable = table match { - case t: SparkTable => - t.table() match { - case t: BaseTable => t - case _ => null - } + val icebergTable = GlutenIcebergSourceUtil.getTable(scan) match { + case t: BaseTable => t case _ => null } if (icebergTable == null) { @@ -351,6 +347,18 @@ object IcebergScanTransformer { private val InputFileRelatedMetadataColumnNames = Set("input_file_name", "input_file_block_start", "input_file_block_length") + private val CopyOnWriteMetadataColumnNames = + Set(MetadataColumns.FILE_PATH.name(), MetadataColumns.ROW_POSITION.name()) + .map(_.toLowerCase(Locale.ROOT)) + + private def isCopyOnWriteFilePathColumn(name: String): Boolean = { + name.equalsIgnoreCase(MetadataColumns.FILE_PATH.name()) + } + + private def isCopyOnWriteRowPositionColumn(name: String): Boolean = { + name.equalsIgnoreCase(MetadataColumns.ROW_POSITION.name()) + } + def apply(batchScan: BatchScanExec): IcebergScanTransformer = { new IcebergScanTransformer( batchScan.output.map(a => a.withName(AvroSchemaUtil.makeCompatibleName(a.name))), diff --git a/gluten-iceberg/src/main/scala/org/apache/iceberg/spark/source/GlutenIcebergSourceUtil.scala b/gluten-iceberg/src/main/scala/org/apache/iceberg/spark/source/GlutenIcebergSourceUtil.scala index 233c06c30cb..0cc8a17ef53 100644 --- a/gluten-iceberg/src/main/scala/org/apache/iceberg/spark/source/GlutenIcebergSourceUtil.scala +++ b/gluten-iceberg/src/main/scala/org/apache/iceberg/spark/source/GlutenIcebergSourceUtil.scala @@ -43,12 +43,15 @@ object GlutenIcebergSourceUtil { private val InputFileBlockLengthCol = "input_file_block_length" def isSupportedScan(sparkScan: Scan): Boolean = sparkScan match { - case _: SparkBatchQueryScan | _: SparkStagedScan => true + case _: SparkBatchQueryScan | _: SparkStagedScan | _: SparkCopyOnWriteScan => true case _ => false } def isSparkStagedScan(sparkScan: Scan): Boolean = sparkScan.isInstanceOf[SparkStagedScan] + def isSparkCopyOnWriteScan(sparkScan: Scan): Boolean = + sparkScan.isInstanceOf[SparkCopyOnWriteScan] + def deleteExists(p: SparkDataSourceRDDPartition): Boolean = { p.inputPartitions.exists { case ip: SparkInputPartition => @@ -149,6 +152,8 @@ object GlutenIcebergSourceUtil { case InputFileNameCol => metadataColumns.put(name, filePath) case InputFileBlockStartCol => metadataColumns.put(name, start.toString) case InputFileBlockLengthCol => metadataColumns.put(name, length.toString) + case filePathCol if filePathCol == MetadataColumns.FILE_PATH.name() => + metadataColumns.put(name, filePath) case _ => } } @@ -206,8 +211,9 @@ object GlutenIcebergSourceUtil { throw new UnsupportedOperationException("Failed to get partition schema from iceberg scan.") } - private def getTable(sparkScan: Scan): Table = sparkScan match { + def getTable(sparkScan: Scan): Table = sparkScan match { case scan: SparkBatchQueryScan => scan.table() + case scan: SparkCopyOnWriteScan => scan.table() case scan: SparkStagedScan => scan.table() case _ => throw new GlutenNotSupportException( @@ -216,6 +222,7 @@ object GlutenIcebergSourceUtil { private def getScanTasks(sparkScan: Scan): List[ScanTask] = sparkScan match { case scan: SparkBatchQueryScan => scan.tasks().asScala.toList + case scan: SparkCopyOnWriteScan => scan.tasks().asScala.toList case scan: SparkStagedScan => scan.taskGroups().asScala.flatMap(_.tasks().asScala).toList case _ => diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala index c765ed333da..40913c39412 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala @@ -159,12 +159,20 @@ trait BasicScanExecTransformer extends LeafTransformSupport with BaseDataSource doNativeValidation(substraitContext, relNode) } + protected def isRowIndexMetadataColumn(attr: Attribute): Boolean = { + BackendsApiManager.getSparkPlanExecApiInstance.isRowIndexMetadataColumn(attr.name) + } + + protected def isNativeMetadataColumn(attr: Attribute): Boolean = { + attr.isMetadataCol || getMetadataColumns().exists(_.exprId == attr.exprId) + } + private def makeColumnTypeNode(attr: Attribute): ColumnTypeNode = { if (getPartitionSchema.exists(_.name.equals(attr.name))) { new ColumnTypeNode(NamedStruct.ColumnType.PARTITION_COL) - } else if (BackendsApiManager.getSparkPlanExecApiInstance.isRowIndexMetadataColumn(attr.name)) { + } else if (isRowIndexMetadataColumn(attr)) { new ColumnTypeNode(NamedStruct.ColumnType.ROWINDEX_COL) - } else if (attr.isMetadataCol) { + } else if (isNativeMetadataColumn(attr)) { new ColumnTypeNode(NamedStruct.ColumnType.METADATA_COL) } else { new ColumnTypeNode(NamedStruct.ColumnType.NORMAL_COL) diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala index 6d2854f30bb..a8807eb6e3d 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala @@ -160,8 +160,12 @@ abstract class BatchScanExecTransformerBase( return ValidationResult.failed(s"Unsupported aggregation push down for $scan.") } + // Keep Spark's schema validation, which also rejects invalid row index column types. + val hasRowIndexColumn = + SparkShimLoader.getSparkShims.findRowIndexColumnIndexInSchema(schema) >= 0 || + output.exists(isRowIndexMetadataColumn) if ( - SparkShimLoader.getSparkShims.findRowIndexColumnIndexInSchema(schema) > 0 && + hasRowIndexColumn && !BackendsApiManager.getSettings.supportNativeRowIndexColumn() ) { return ValidationResult.failed("Unsupported row index column scan in native.") diff --git a/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/GlutenParquetRowIndexSuite.scala b/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/GlutenParquetRowIndexSuite.scala index 570b6d5e0c1..1533246f288 100644 --- a/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/GlutenParquetRowIndexSuite.scala +++ b/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/GlutenParquetRowIndexSuite.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.functions.{col, max, min} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{LongType, StringType} +import org.apache.commons.lang3.exception.ExceptionUtils import org.apache.hadoop.fs.Path import org.apache.parquet.column.ParquetProperties._ import org.apache.parquet.format.converter.ParquetMetadataConverter @@ -353,7 +354,11 @@ class GlutenParquetRowIndexSuite extends ParquetRowIndexSuite with GlutenSQLTest .load(path.getAbsolutePath) val exception = intercept[Exception](dfRead.collect()) - assert(exception.getMessage.contains(ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME)) + // Spark can wrap the validation error in a FAILED_READ_FILE exception. + assert( + ExceptionUtils + .getRootCauseMessage(exception) + .contains(ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME)) } } } diff --git a/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/GlutenParquetRowIndexSuite.scala b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/GlutenParquetRowIndexSuite.scala index 570b6d5e0c1..1533246f288 100644 --- a/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/GlutenParquetRowIndexSuite.scala +++ b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/GlutenParquetRowIndexSuite.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.functions.{col, max, min} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{LongType, StringType} +import org.apache.commons.lang3.exception.ExceptionUtils import org.apache.hadoop.fs.Path import org.apache.parquet.column.ParquetProperties._ import org.apache.parquet.format.converter.ParquetMetadataConverter @@ -353,7 +354,11 @@ class GlutenParquetRowIndexSuite extends ParquetRowIndexSuite with GlutenSQLTest .load(path.getAbsolutePath) val exception = intercept[Exception](dfRead.collect()) - assert(exception.getMessage.contains(ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME)) + // Spark can wrap the validation error in a FAILED_READ_FILE exception. + assert( + ExceptionUtils + .getRootCauseMessage(exception) + .contains(ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME)) } } }