Add create table as select support for iceberg. (#13595)

Author: liurenjie1024

iceberg/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderImpl.scala

@@ -19,15 +19,17 @@ package com.nvidia.spark.rapids.iceberg
 import scala.reflect.ClassTag
 import scala.util.{Failure, Success, Try}
 
-import com.nvidia.spark.rapids.{AppendDataExecMeta, FileFormatChecks, GpuExec, GpuExpression, GpuRowToColumnarExec, GpuScan, IcebergFormatType, RapidsConf, ReadFileOp, ScanMeta, ScanRule, ShimReflectionUtils, StaticInvokeMeta, TargetSize, WriteFileOp}
+import com.nvidia.spark.rapids.{AppendDataExecMeta, AtomicCreateTableAsSelectExecMeta, FileFormatChecks, GpuExec, GpuExpression, GpuRowToColumnarExec, GpuScan, IcebergFormatType, RapidsConf, ReadFileOp, ScanMeta, ScanRule, ShimReflectionUtils, StaticInvokeMeta, TargetSize, WriteFileOp}
 import org.apache.iceberg.spark.functions.{BucketFunction, GpuBucketExpression}
 import org.apache.iceberg.spark.source.{GpuSparkBatchQueryScan, GpuSparkWrite}
+import org.apache.iceberg.spark.supportsCatalog
 
 import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
 import org.apache.spark.sql.connector.read.Scan
 import org.apache.spark.sql.connector.write.Write
 import org.apache.spark.sql.execution.SparkPlan
-import org.apache.spark.sql.execution.datasources.v2.{AppendDataExec, GpuAppendDataExec}
+import org.apache.spark.sql.execution.datasources.v2.{AppendDataExec, AtomicCreateTableAsSelectExec, GpuAppendDataExec}
+import org.apache.spark.sql.execution.datasources.v2.rapids.GpuAtomicCreateTableAsSelectExec
 
 class IcebergProviderImpl extends IcebergProvider {
   override def getScans: Map[Class[_ <: Scan], ScanRule[_ <: Scan]] = {
@@ -98,6 +100,41 @@ class IcebergProviderImpl extends IcebergProvider {
     GpuSparkWrite.supports(write)
   }
 
+  override def isSupportedCatalog(catalogClass: Class[_]): Boolean = {
+    supportsCatalog(catalogClass)
+  }
+
+  override def tagForGpu(
+      cpuExec: AtomicCreateTableAsSelectExec,
+      meta: AtomicCreateTableAsSelectExecMeta): Unit = {
+    if (!meta.conf.isIcebergEnabled) {
+      meta.willNotWorkOnGpu("Iceberg input and output has been disabled. To enable set " +
+        s"${RapidsConf.ENABLE_ICEBERG.key} to true")
+    }
+
+    if (!meta.conf.isIcebergWriteEnabled) {
+      meta.willNotWorkOnGpu("Iceberg output has been disabled. To enable set " +
+        s"${RapidsConf.ENABLE_ICEBERG_WRITE.key} to true")
+    }
+
+    FileFormatChecks.tag(meta, cpuExec.query.schema, IcebergFormatType, WriteFileOp)
+
+    GpuSparkWrite.tagForGpuCtas(cpuExec, meta)
+  }
+
+  override def convertToGpu(
+      cpuExec: AtomicCreateTableAsSelectExec,
+      meta: AtomicCreateTableAsSelectExecMeta): GpuExec = {
+    GpuAtomicCreateTableAsSelectExec(
+      cpuExec.catalog,
+      cpuExec.ident,
+      cpuExec.partitioning,
+      cpuExec.query,
+      cpuExec.tableSpec,
+      cpuExec.writeOptions,
+      cpuExec.ifNotExists)
+  }
+
   override def tagForGpu(cpuExec: AppendDataExec, meta: AppendDataExecMeta): Unit = {
     if (!meta.conf.isIcebergEnabled) {
       meta.willNotWorkOnGpu("Iceberg input and output has been disabled. To enable set " +

iceberg/src/main/scala/org/apache/iceberg/spark/package.scala

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.iceberg
+
+package object spark {
+  def supportsCatalog(catalogClass: Class[_]): Boolean = {
+    classOf[BaseCatalog].isAssignableFrom(catalogClass)
+  }
+}

iceberg/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala

@@ -26,8 +26,9 @@ import com.nvidia.spark.rapids.fileio.iceberg.IcebergFileIO
 import com.nvidia.spark.rapids.iceberg.GpuIcebergPartitioner
 import org.apache.hadoop.mapreduce.Job
 import org.apache.hadoop.shaded.org.apache.commons.lang3.reflect.{FieldUtils, MethodUtils}
-import org.apache.iceberg.{DataFile, FileFormat, PartitionSpec, Schema, SerializableTable, SnapshotUpdate, Table}
+import org.apache.iceberg.{DataFile, FileFormat, PartitionSpec, Schema, SerializableTable, SnapshotUpdate, Table, TableProperties}
 import org.apache.iceberg.io.{DataWriteResult, FileIO, GpuClusteredDataWriter, GpuFanoutDataWriter, GpuRollingDataWriter, OutputFileFactory, PartitioningWriter}
+import org.apache.iceberg.spark.{Spark3Util, SparkSchemaUtil}
 import org.apache.iceberg.spark.functions.{GpuFieldTransform, GpuTransform}
 import org.apache.iceberg.spark.source.SparkWrite.TaskCommit
 
@@ -39,6 +40,7 @@ import org.apache.spark.sql.connector.distributions.Distribution
 import org.apache.spark.sql.connector.expressions.SortOrder
 import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, RequiresDistributionAndOrdering, Write, WriterCommitMessage}
 import org.apache.spark.sql.connector.write.streaming.StreamingWrite
+import org.apache.spark.sql.execution.datasources.v2.AtomicCreateTableAsSelectExec
 import org.apache.spark.sql.rapids.GpuWriteJobStatsTracker
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.ColumnarBatch
@@ -141,28 +143,30 @@ object GpuSparkWrite {
     classOf[SparkWrite].isAssignableFrom(cpuClass)
   }
 
-  def tagForGpu(cpuWrite: Write, meta: SparkPlanMeta[_]): Unit = {
-    if (!supports(cpuWrite.getClass)) {
-      meta.willNotWorkOnGpu(s"GpuSparkWrite only supports ${classOf[SparkWrite].getName}, " +
-        s"but got: ${cpuWrite.getClass.getName}")
-      return
-    }
-
-    val dataFileFormat: FileFormat = FieldUtils.readField(cpuWrite, "format", true)
-      .asInstanceOf[FileFormat]
-
+  /**
+   * Tag for GPU support for Iceberg write operations.
+   * This method checks:
+   * 1. File format is supported (only Parquet)
+   * 2. Partition transforms are supported
+   *
+   * @param dataFileFormat The file format for the write operation
+   * @param partitionSpec The partition specification
+   * @param dsSchema The DataFrame schema
+   * @param writeSchema The write schema
+   * @param meta The metadata for tagging
+   */
+  private def tagForGpuWrite(
+      dataFileFormat: FileFormat,
+      partitionSpec: PartitionSpec,
+      dsSchema: StructType,
+      writeSchema: Schema,
+      meta: SparkPlanMeta[_]): Unit = {
+    // Check file format support
     if (!dataFileFormat.equals(FileFormat.PARQUET)) {
       meta.willNotWorkOnGpu(s"GpuSparkWrite only supports Parquet, but got: $dataFileFormat")
     }
 
-    val table: Table = FieldUtils.readField(cpuWrite, "table", true).asInstanceOf[Table]
-    val partitionSpec = table.spec()
-
-    val dsSchema = FieldUtils.readField(cpuWrite, "dsSchema", true)
-      .asInstanceOf[StructType]
-    val writeSchema = FieldUtils.readField(cpuWrite, "writeSchema", true)
-      .asInstanceOf[Schema]
-
+    // Check partition transform support
     if (partitionSpec.isPartitioned) {
       for (partitionField <- partitionSpec.fields().asScala) {
         val transform = partitionField.transform()
@@ -180,6 +184,55 @@ object GpuSparkWrite {
     }
   }
 
+  def tagForGpu(cpuWrite: Write, meta: SparkPlanMeta[_]): Unit = {
+    if (!supports(cpuWrite.getClass)) {
+      meta.willNotWorkOnGpu(s"GpuSparkWrite only supports ${classOf[SparkWrite].getName}, " +
+        s"but got: ${cpuWrite.getClass.getName}")
+      return
+    }
+
+    val dataFileFormat: FileFormat = FieldUtils.readField(cpuWrite, "format", true)
+      .asInstanceOf[FileFormat]
+
+    val table: Table = FieldUtils.readField(cpuWrite, "table", true).asInstanceOf[Table]
+    val partitionSpec = table.spec()
+
+    val dsSchema = FieldUtils.readField(cpuWrite, "dsSchema", true)
+      .asInstanceOf[StructType]
+    val writeSchema = FieldUtils.readField(cpuWrite, "writeSchema", true)
+      .asInstanceOf[Schema]
+
+    tagForGpuWrite(dataFileFormat, partitionSpec, dsSchema, writeSchema, meta)
+  }
+
+  /**
+   * Tag for GPU support for Iceberg CTAS operations.
+   * This method checks file format and partitioning support for CREATE TABLE AS SELECT.
+   *
+   * @param cpuExec The CPU AtomicCreateTableAsSelectExec
+   * @param meta The metadata for tagging
+   */
+  def tagForGpuCtas(
+      cpuExec: AtomicCreateTableAsSelectExec,
+      meta: SparkPlanMeta[_]): Unit = {
+    val properties: Map[String, String] = Spark3Util
+      .rebuildCreateProperties(cpuExec.tableSpec.properties.asJava)
+      .asScala.toMap
+    val fileFormatStr = properties.getOrElse(TableProperties.DEFAULT_FILE_FORMAT,
+      TableProperties.DEFAULT_FILE_FORMAT_DEFAULT)
+
+    val fileFormat = FileFormat.fromString(fileFormatStr)
+
+    // Convert Spark schema to Iceberg schema
+    val querySchema = cpuExec.query.schema
+    val icebergSchema = SparkSchemaUtil.convert(querySchema)
+
+    // Convert Spark connector transforms to Iceberg PartitionSpec
+    val partitionSpec = Spark3Util.toPartitionSpec(icebergSchema, cpuExec.partitioning.toArray)
+
+    // Reuse tagForGpuWrite for validation
+    tagForGpuWrite(fileFormat, partitionSpec, querySchema, icebergSchema, meta)
+  }
 
   def convert(cpuWrite: Write): GpuSparkWrite = {
     new GpuSparkWrite(cpuWrite.asInstanceOf[SparkWrite])