This is an automated email from the ASF dual-hosted git repository.
zhouyuan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 9f6dcb599d [GLUTEN-12273][CORE] Support Iceberg staged scans (#12962)
9f6dcb599d is described below
commit 9f6dcb599d3ebb8ef9f55450163aae4f025300e9
Author: inf <[email protected]>
AuthorDate: Mon Sep 14 13:41:55 2026 +0000
[GLUTEN-12273][CORE] Support Iceberg staged scans (#12962)
---
.../gluten/execution/IcebergScanTransformer.scala | 2 +-
.../spark/source/GlutenIcebergSourceUtil.scala | 150 ++++++++++-----------
.../org/apache/gluten/execution/IcebergSuite.scala | 70 ++++++++++
3 files changed, 146 insertions(+), 76 deletions(-)
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 6735c92d85..6ac22b956a 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
@@ -363,7 +363,7 @@ object IcebergScanTransformer {
}
def supportsBatchScan(scan: Scan): Boolean = {
- scan.getClass == GlutenIcebergSourceUtil.getClassOfSparkBatchQueryScan
+ GlutenIcebergSourceUtil.isSupportedScan(scan)
}
private def containsUuidOrFixedType(dataType: Type): Boolean = {
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 db1bb024af..233c06c30c 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
@@ -31,7 +31,7 @@ import org.apache.spark.sql.types.StructType
import org.apache.iceberg._
import org.apache.iceberg.spark.SparkSchemaUtil
-import java.lang.{Class, Long => JLong}
+import java.lang.{Long => JLong}
import java.util.{ArrayList => JArrayList, HashMap => JHashMap, List => JList,
Map => JMap}
import java.util.Locale
@@ -42,10 +42,13 @@ object GlutenIcebergSourceUtil {
private val InputFileBlockStartCol = "input_file_block_start"
private val InputFileBlockLengthCol = "input_file_block_length"
- def getClassOfSparkBatchQueryScan(): Class[SparkBatchQueryScan] = {
- classOf[SparkBatchQueryScan]
+ def isSupportedScan(sparkScan: Scan): Boolean = sparkScan match {
+ case _: SparkBatchQueryScan | _: SparkStagedScan => true
+ case _ => false
}
+ def isSparkStagedScan(sparkScan: Scan): Boolean =
sparkScan.isInstanceOf[SparkStagedScan]
+
def deleteExists(p: SparkDataSourceRDDPartition): Boolean = {
p.inputPartitions.exists {
case ip: SparkInputPartition =>
@@ -114,32 +117,22 @@ object GlutenIcebergSourceUtil {
def getFieldIds(sparkScan: Scan): JHashMap[String, Integer] = {
val fieldIds = new JHashMap[String, Integer]()
- sparkScan match {
- case scan: SparkBatchQueryScan =>
- scan.table().schema().columns().asScala.foreach {
- field => fieldIds.put(field.name(), field.fieldId())
- }
- case _ =>
- throw new GlutenNotSupportException("Only support iceberg
SparkBatchQueryScan.")
+ getTable(sparkScan).schema().columns().asScala.foreach {
+ field => fieldIds.put(field.name(), field.fieldId())
}
fieldIds
}
def getInitialDefaults(sparkScan: Scan): JHashMap[String, String] = {
val initialDefaults = new JHashMap[String, String]()
- sparkScan match {
- case scan: SparkBatchQueryScan =>
- scan.table().schema().columns().asScala.foreach {
- field =>
- val defaultValue = IcebergDefaultValueUtil.getInitialDefault(field)
- if (defaultValue != null) {
- initialDefaults.put(
- field.name(),
- TypeUtil.getPartitionValueString(field.`type`(), defaultValue))
- }
+ getTable(sparkScan).schema().columns().asScala.foreach {
+ field =>
+ val defaultValue = IcebergDefaultValueUtil.getInitialDefault(field)
+ if (defaultValue != null) {
+ initialDefaults.put(
+ field.name(),
+ TypeUtil.getPartitionValueString(field.`type`(), defaultValue))
}
- case _ =>
- throw new GlutenNotSupportException("Only support iceberg
SparkBatchQueryScan.")
}
initialDefaults
}
@@ -162,65 +155,72 @@ object GlutenIcebergSourceUtil {
metadataColumns
}
- def getFileFormat(sparkScan: Scan): ReadFileFormat = sparkScan match {
- case scan: SparkBatchQueryScan =>
- val tasks = scan.tasks().asScala
- asFileScanTask(tasks.toList).foreach {
- task =>
- task.file().format() match {
- case FileFormat.PARQUET => return ReadFileFormat.ParquetReadFormat
- case FileFormat.ORC => return ReadFileFormat.OrcReadFormat
- case _ =>
- }
- }
- throw new GlutenNotSupportException("Iceberg Only support parquet and
orc file format.")
- case _ =>
- throw new GlutenNotSupportException("Only support iceberg
SparkBatchQueryScan.")
+ def getFileFormat(sparkScan: Scan): ReadFileFormat = {
+ asFileScanTask(getScanTasks(sparkScan)).foreach {
+ task =>
+ task.file().format() match {
+ case FileFormat.PARQUET => return ReadFileFormat.ParquetReadFormat
+ case FileFormat.ORC => return ReadFileFormat.OrcReadFormat
+ case _ =>
+ }
+ }
+ throw new GlutenNotSupportException("Iceberg Only support parquet and orc
file format.")
}
- def getReadPartitionSchema(sparkScan: Scan): StructType = sparkScan match {
- case scan: SparkBatchQueryScan =>
- val tasks = scan.tasks().asScala
- asFileScanTask(tasks.toList).foreach {
- task =>
- val spec = task.spec()
- if (spec.isPartitioned) {
- val readFields = scan.readSchema().fields.map(_.name).toSet
- // Iceberg will generate some non-table fields as partition
fields, such as x_bucket,
- // which will not appear in readFields, they also cannot be
filtered.
- val tableFields =
spec.schema().columns().asScala.map(_.name()).toSet
- val voidTransformFields = scan
- .table()
- .spec()
+ def getReadPartitionSchema(sparkScan: Scan): StructType = {
+ asFileScanTask(getScanTasks(sparkScan)).foreach {
+ task =>
+ val spec = task.spec()
+ if (spec.isPartitioned) {
+ val readFields = sparkScan.readSchema().fields.map(_.name).toSet
+ // Iceberg will generate some non-table fields as partition fields,
such as x_bucket,
+ // which will not appear in readFields, they also cannot be filtered.
+ val tableFields = spec.schema().columns().asScala.map(_.name()).toSet
+ val voidTransformFields = getTable(sparkScan)
+ .spec()
+ .fields()
+ .asScala
+ .filter(
+ f => {
+ f.transform().isVoid
+ })
+ .map(_.name())
+ .toSet
+ val partitionFields =
+ spec
+ .partitionType()
.fields()
.asScala
- .filter(
- f => {
- f.transform().isVoid
- })
- .map(_.name())
- .toSet
- val partitionFields =
- spec
- .partitionType()
- .fields()
- .asScala
- .filter(f => !tableFields.contains(f.name) ||
readFields.contains(f.name()))
- .filter(f => !voidTransformFields.contains(f.name()))
- partitionFields.foreach {
- field =>
TypeUtil.validatePartitionColumnType(field.`type`().typeId())
- }
-
- val icebergSchema = new Schema(partitionFields.toList.asJava)
- return SparkSchemaUtil.convert(icebergSchema)
- } else {
- return new StructType()
+ .filter(f => !tableFields.contains(f.name) ||
readFields.contains(f.name()))
+ .filter(f => !voidTransformFields.contains(f.name()))
+ partitionFields.foreach {
+ field =>
TypeUtil.validatePartitionColumnType(field.`type`().typeId())
}
- }
- throw new UnsupportedOperationException(
- "Failed to get partition schema from iceberg SparkBatchQueryScan.")
+
+ val icebergSchema = new Schema(partitionFields.toList.asJava)
+ return SparkSchemaUtil.convert(icebergSchema)
+ } else {
+ return new StructType()
+ }
+ }
+ throw new UnsupportedOperationException("Failed to get partition schema
from iceberg scan.")
+ }
+
+ private def getTable(sparkScan: Scan): Table = sparkScan match {
+ case scan: SparkBatchQueryScan => scan.table()
+ case scan: SparkStagedScan => scan.table()
+ case _ =>
+ throw new GlutenNotSupportException(
+ s"Unsupported Iceberg scan: ${sparkScan.getClass.getName}.")
+ }
+
+ private def getScanTasks(sparkScan: Scan): List[ScanTask] = sparkScan match {
+ case scan: SparkBatchQueryScan => scan.tasks().asScala.toList
+ case scan: SparkStagedScan =>
+ scan.taskGroups().asScala.flatMap(_.tasks().asScala).toList
case _ =>
- throw new UnsupportedOperationException("Only support iceberg
SparkBatchQueryScan.")
+ throw new GlutenNotSupportException(
+ s"Unsupported Iceberg scan: ${sparkScan.getClass.getName}.")
}
private def asFileScanTask(tasks: List[ScanTask]): List[FileScanTask] = {
diff --git
a/gluten-iceberg/src/test/scala/org/apache/gluten/execution/IcebergSuite.scala
b/gluten-iceberg/src/test/scala/org/apache/gluten/execution/IcebergSuite.scala
index 0ee193ce0f..5a194fc9a9 100644
---
a/gluten-iceberg/src/test/scala/org/apache/gluten/execution/IcebergSuite.scala
+++
b/gluten-iceberg/src/test/scala/org/apache/gluten/execution/IcebergSuite.scala
@@ -20,7 +20,13 @@ import org.apache.gluten.config.GlutenIcebergConfig
import org.apache.spark.SparkConf
import org.apache.spark.sql.Row
+import org.apache.spark.sql.execution.QueryExecution
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.util.QueryExecutionListener
+
+import org.apache.iceberg.spark.source.GlutenIcebergSourceUtil
+
+import java.util.concurrent.{CountDownLatch, TimeUnit}
abstract class IcebergSuite extends WholeStageTransformerSuite {
protected val rootPath: String = getClass.getResource("/").getPath
@@ -60,6 +66,70 @@ abstract class IcebergSuite extends
WholeStageTransformerSuite {
}
}
+ test("rewrite_data_files uses an iceberg staged scan transformer") {
+ val tableName = "iceberg_rewrite_tb"
+ withTable(tableName) {
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ spark.sql(s"CREATE TABLE $tableName (id INT, data STRING) USING
iceberg")
+ (1 to 5).foreach {
+ id => spark.sql(s"INSERT INTO $tableName VALUES ($id, 'value-$id')")
+ }
+
+ def dataFileCount: Long =
+ spark.table(s"spark_catalog.default.$tableName.files").count()
+
+ assert(dataFileCount == 5)
+ val stagedScanSeen = new CountDownLatch(1)
+ val listener = new QueryExecutionListener {
+ override def onSuccess(
+ funcName: String,
+ qe: QueryExecution,
+ durationNs: Long): Unit = {
+ if (
+ qe.executedPlan.exists {
+ case scan: IcebergScanTransformer =>
+ GlutenIcebergSourceUtil.isSparkStagedScan(scan.scan)
+ case _ => false
+ }
+ ) {
+ stagedScanSeen.countDown()
+ }
+ }
+
+ override def onFailure(
+ funcName: String,
+ qe: QueryExecution,
+ exception: Exception): Unit = {}
+ }
+
+ try {
+ spark.listenerManager.register(listener)
+ val result = spark
+ .sql(s"""
+ |CALL spark_catalog.system.rewrite_data_files(
+ | table => 'default.$tableName',
+ | options => map('min-input-files', '2'))
+ |""".stripMargin)
+ .collect()
+
+ assert(result.length == 1)
+ assert(result.head.getInt(0) == 5)
+ assert(result.head.getInt(1) == 1)
+ assert(
+ stagedScanSeen.await(10, TimeUnit.SECONDS),
+ "Rewrite read did not use IcebergScanTransformer with
SparkStagedScan")
+ } finally {
+ spark.listenerManager.unregister(listener)
+ }
+
+ assert(dataFileCount == 1)
+ checkAnswer(
+ spark.sql(s"SELECT * FROM $tableName ORDER BY id"),
+ (1 to 5).map(id => Row(id, s"value-$id")))
+ }
+ }
+ }
+
test("iceberg input_file_name") {
withTable("iceberg_input_file_tb") {
spark.sql("""
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
