From 4125e4291a4e6c1795fff8e2ca76d6c1f447ddfb Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 6 Oct 2026 09:43:20 -0700 Subject: [PATCH 1/9] perf: read only the nested fields Spark keeps in the native Iceberg scan The native Iceberg scan sent iceberg-rust the full table schema as the FileScanTask schema, so iceberg-rust read and decoded every leaf of a projected struct, list, or map column even when Spark's nested schema pruning kept only some of them. Pass the pruned scan schema instead, without metadata columns, and still add partition sources and equality-delete keys. Memoize those additions so tasks share one schema in the pool, and keep the full schema when a needed column is nested inside a pruned struct. Tasks with deletes read with the pruned schema too, so a VERSION AS OF read of a column dropped after the snapshot no longer fails natively when the table has deletes. The pruned schema is built at planning time, where a reflection failure falls back to Spark, and findFieldObject now propagates reflection failures. A column added to a task schema resolves from the current table schema before the schema history, so it keeps its current name and type. Add spark.comet.scan.icebergNative.nestedSchemaPruning.enabled to turn it off. --- .../scala/org/apache/comet/CometConf.scala | 10 + .../comet/iceberg/IcebergReflection.scala | 74 +++-- .../apache/comet/rules/CometScanRule.scala | 12 +- .../operator/CometIcebergNativeScan.scala | 81 +++--- .../operator/CometIcebergNativeWrite.scala | 4 +- .../iceberg/nested_schema_pruning.sql | 135 ++++++++++ .../comet/CometIcebergNativeSuite.scala | 253 +++++++++++++++++- .../benchmark/CometIcebergReadBenchmark.scala | 109 +++++++- 8 files changed, 615 insertions(+), 63 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 248772fb6c4..a55fa6648e6 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -152,6 +152,16 @@ object CometConf extends ShimCometConf { .checkValue(v => v > 0, "Data file concurrency limit must be positive") .createWithDefault(1) + val COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.scan.icebergNative.nestedSchemaPruning.enabled") + .category(CATEGORY_SCAN) + .doc( + "Whether the native Iceberg scan reads only the nested fields of a struct, list, or map " + + "column that Spark's nested schema pruning keeps. When false, it reads every nested " + + "field of a projected column and drops the unused ones after decoding.") + .booleanConf + .createWithDefault(true) + val COMET_ICEBERG_IO_TIMEOUT: ConfigEntry[Long] = conf("spark.comet.iceberg.ioTimeout") .category(CATEGORY_SCAN) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index 96626631393..d4e368970d9 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -707,20 +707,20 @@ object IcebergReflection extends Logging { } } - /** Returns the `Types.NestedField` for `fieldId` in `schema`, or None. */ - def findFieldObject(schema: Any, fieldId: Int): Option[Any] = { - try { - val findFieldMethod = getMethod(schema.getClass, "findField", classOf[Int]) - Option(findFieldMethod.invoke(schema, fieldId.asInstanceOf[AnyRef])) - } catch { - case _: Exception => None - } - } + /** + * Returns the `Types.NestedField` for `fieldId` at any depth of `schema`, or None when `schema` + * has no such field. A reflection failure propagates rather than reading as a missing field, + * because the callers run at serialization time, when the scan can no longer fall back. + */ + def findFieldObject(schema: Any, fieldId: Int): Option[Any] = + Option( + getMethod(schema.getClass, "findField", classOf[Int]) + .invoke(schema, fieldId.asInstanceOf[AnyRef])) /** * Returns a schema equal to `baseSchema` but guaranteed to contain `requiredFieldIds`. Any id - * not already present is resolved from the table's schema history (`table.schemas()`) and - * appended. + * not already present is resolved from the table's current schema, or from its schema history + * (`table.schemas()`) for a column since dropped, and appended. * * Two callers need this: an equality delete may be keyed on a column since dropped from the * current schema (mirroring Iceberg-Java's `DeleteFilter.fileProjection`), and a partition @@ -741,12 +741,14 @@ object IcebergReflection extends Logging { logDebug( s"Native Iceberg scan schema is missing field id(s) ${missingIds.mkString(",")}; " + "resolving them from table schema history") - val history = getAllSchemas(table) + // table.schemas() lists the oldest schema first, so a column the table still has is looked + // up in the current schema before it, to keep its current name and type. An older name + // could clash with a column that took it over. + val schemas = getMethod(table.getClass, "schema").invoke(table) +: getAllSchemas(table) val resolvedFields = missingIds.map { id => - history.iterator - .flatMap(s => findFieldObject(s, id)) - .toSeq - .headOption + schemas.iterator + .map(findFieldObject(_, id)) + .collectFirst { case Some(field) => field } .getOrElse(throw new IllegalStateException( s"Cannot resolve field id $id in table schema history")) } @@ -763,6 +765,40 @@ object IcebergReflection extends Logging { } } + /** The spec reserves field ids above `Integer.MAX_VALUE - 200` for metadata columns. */ + val MaxDataFieldId: Int = Int.MaxValue - 200 + + /** + * Returns `schema` without its top-level metadata columns (`_file`, `_pos`, `_partition` and + * the other reserved ids), or `schema` itself when it has none. A scan's expected schema + * includes the metadata columns the query selects, which a table schema never has. + */ + def withoutMetadataColumns(schema: Any): Any = { + import scala.jdk.CollectionConverters._ + val columns = + getMethod(schema.getClass, "columns").invoke(schema).asInstanceOf[java.util.List[_]] + val dataColumns = columns.asScala.filter(fieldIdOf(_) <= MaxDataFieldId) + if (dataColumns.size == columns.size) { + schema + } else { + schema.getClass + .getConstructor(classOf[java.util.List[_]]) + .newInstance(new java.util.ArrayList[Any](dataColumns.asJava)) + .asInstanceOf[AnyRef] + } + } + + /** Whether `fieldId` is a field of `schema` nested below its top level. */ + def isNestedField(schema: Any, fieldId: Int): Boolean = { + import scala.jdk.CollectionConverters._ + findFieldObject(schema, fieldId).isDefined && + !getMethod(schema.getClass, "columns") + .invoke(schema) + .asInstanceOf[java.util.List[_]] + .asScala + .exists(fieldIdOf(_) == fieldId) + } + /** * The source column field ids referenced by a task's partition spec. * @@ -2344,6 +2380,10 @@ object IcebergReflection extends Logging { * List of FileScanTask objects from Iceberg planning * @param scanSchema * The expectedSchema from the SparkScan (for schema evolution / VERSION AS OF) + * @param prunedScanSchema + * scanSchema without its metadata columns, which iceberg-rust resolves by reserved field id + * rather than from the task schema. Holds only the nested fields Spark's schema pruning keeps, + * so a task reading with it skips the others. * @param globalFieldIdMapping * Mapping from column names to Iceberg field IDs (built from scanSchema) * @param catalogProperties @@ -2358,6 +2398,7 @@ case class CometIcebergNativeScanMetadata( nameMapping: Option[String], @transient tasks: java.util.List[_], scanSchema: Any, + prunedScanSchema: Any, tableSchema: Any, globalFieldIdMapping: Map[String, Int], catalogProperties: Map[String, String], @@ -2418,6 +2459,7 @@ object CometIcebergNativeScanMetadata extends Logging { nameMapping = nameMapping, tasks = tasks, scanSchema = scanSchema, + prunedScanSchema = withoutMetadataColumns(scanSchema), tableSchema = tableSchema, globalFieldIdMapping = globalFieldIdMapping, catalogProperties = catalogProperties, diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index c7c9a31909d..acd149a966d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -774,12 +774,12 @@ case class CometScanRule(session: SparkSession) // fields to the table's by field id (apache/iceberg-rust#2617). A file written before a // nested field was added (to a struct, or to a struct inside a list or map) fails the // native scan with "Incorrect number of arrays for StructArray fields", and a nested field - // renamed since the file was written can read back as NULL. The native read asks for a - // projected column's full nested type even when Spark prunes it, taken from the current - // table schema, or from the scan schema when VERSION AS OF reads a dropped column (see - // CometIcebergNativeScan). A FileScanTask does not record which schema wrote its file, so - // fall back when any schema in the table's history lacks a nested field of a projected - // column or names it differently. + // renamed since the file was written can read back as NULL. The native read usually asks + // only for the nested fields Spark's schema pruning keeps, but some tasks still read with + // the full table schema (see CometIcebergNativeScan), so this check walks both schemas and + // falls back even when the changed field is pruned away. A FileScanTask does not record + // which schema wrote its file, so fall back when any schema in the table's history lacks + // a nested field of a projected column or names it differently. val nestedFieldsSupported = try { val schemas = Seq(metadata.tableSchema, metadata.scanSchema) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index 335f02ff034..ca6941ce16c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -1026,6 +1026,32 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val pageIndexUnsupportedColumns = IcebergReflection.pageIndexUnsupportedColumns(metadata.tableSchema) + // iceberg-rust reads every leaf of a projected column that the task schema holds. Reading with + // the pruned scan schema skips the nested fields Spark's schema pruning removed, which the + // table schema would read and decode only to drop later. + val nestedSchemaPruningEnabled = CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.get() + // Partition sources and equality-delete keys that the task schema lacks are appended at its + // top level. A nested one that the query pruned away would land outside its struct, so a task + // that needs one reads with the full schema instead. + val canUsePrunedSchemaCache = mutable.HashMap[Seq[Int], Boolean]() + def canUsePrunedSchema(requiredFieldIds: Seq[Int]): Boolean = + nestedSchemaPruningEnabled && canUsePrunedSchemaCache.getOrElseUpdate( + requiredFieldIds, + requiredFieldIds.forall { id => + IcebergReflection.findFieldObject(metadata.prunedScanSchema, id).isDefined || + !IcebergReflection.isNestedField(metadata.tableSchema, id) + }) + // schemaWithRequiredFields builds a new Schema whenever it appends a column. Memoize it by + // base schema (Iceberg's Schema compares by identity) and ids, so tasks that need the same + // columns share one schema object, and with it one schema pool entry. + val requiredFieldsCache = mutable.HashMap[(AnyRef, Seq[Int]), AnyRef]() + def withRequiredFields(baseSchema: AnyRef, requiredFieldIds: Seq[Int]): AnyRef = + requiredFieldsCache.getOrElseUpdate( + (baseSchema, requiredFieldIds), + IcebergReflection + .schemaWithRequiredFields(baseSchema, metadata.table, requiredFieldIds) + .asInstanceOf[AnyRef]) + val perPartitionBuilders = mutable.ArrayBuffer[OperatorOuterClass.IcebergScan]() var totalTasks = 0 @@ -1130,49 +1156,38 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // verbatim for iceberg-rust to decode. Unencrypted files leave the field unset. keyMetadataBytes(keyMetadataMethod, dataFile).foreach(taskBuilder.setKeyMetadata) - val taskSchema = taskSchemaMethod.invoke(task) - val deletes = IcebergReflection.getDeleteFilesFromTask(task, fileScanTaskClass) val hasDeletes = !deletes.isEmpty - val baseSchema: AnyRef = - if (hasDeletes) { - // An equality delete may be keyed on a column dropped from the current schema - // (schema evolution). iceberg-rust must read that column to apply the delete, - // so union the equality-delete field ids into the task schema, resolving any - // dropped ones from the table's schema history (mirrors Iceberg-Java's - // DeleteFilter.fileProjection). - val equalityFieldIds = deletes.asScala.flatMap { df => - requiredEqualityFieldIds(deleteFileClass, df).asScala.map(_.intValue()) - }.toSeq - if (equalityFieldIds.nonEmpty) { - IcebergReflection - .schemaWithRequiredFields(taskSchema, metadata.table, equalityFieldIds) - .asInstanceOf[AnyRef] - } else { - taskSchema - } - } else { - if (hasHistoricalColumns) { - metadata.scanSchema.asInstanceOf[AnyRef] - } else { - metadata.tableSchema.asInstanceOf[AnyRef] - } - } + // An equality delete may be keyed on a column dropped from the current schema + // (schema evolution). iceberg-rust must read that column to apply the delete, + // so union the equality-delete field ids into the task schema, resolving any + // dropped ones from the table's schema history (mirrors Iceberg-Java's + // DeleteFilter.fileProjection). + val equalityFieldIds = deletes.asScala.flatMap { df => + requiredEqualityFieldIds(deleteFileClass, df).asScala.map(_.intValue()) + }.toSeq // iceberg-rust validates a FileScanTask by resolving its partition spec against // the task schema, so a task carrying a partition spec needs that spec's source // columns present even when the query projects them out (e.g. selecting only // _spec_id / _partition). Union them in. project_field_ids still drives the read, // so these columns are not materialized into the output. - val schema: AnyRef = - IcebergReflection - .schemaWithRequiredFields( - baseSchema, - metadata.table, - IcebergReflection.partitionSourceFieldIds(task, fileScanTaskClass)) - .asInstanceOf[AnyRef] + val requiredFieldIds = equalityFieldIds ++ + IcebergReflection.partitionSourceFieldIds(task, fileScanTaskClass) + + val baseSchema: AnyRef = + if (canUsePrunedSchema(requiredFieldIds)) { + metadata.prunedScanSchema.asInstanceOf[AnyRef] + } else if (hasDeletes) { + taskSchemaMethod.invoke(task) + } else if (hasHistoricalColumns) { + metadata.scanSchema.asInstanceOf[AnyRef] + } else { + metadata.tableSchema.asInstanceOf[AnyRef] + } + val schema: AnyRef = withRequiredFields(baseSchema, requiredFieldIds) val schemaIdx = schemaToPoolIndex.getOrElseUpdate( schema, { diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala index 34c07ff8dd5..31278f7325e 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala @@ -95,8 +95,6 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { // scheme reads its bucket from the URL host (`requireSupportedStorageScheme`). private val LocalStorageSchemes: Set[String] = Set("file", "memory") private val MaxSupportedFormatVersion = 3 - // The Iceberg spec reserves field ids above `Integer.MAX_VALUE - 200` for metadata columns. - private val MaxDataFieldId = Int.MaxValue - 200 private val ParquetWritePropertyPrefix = "write.parquet." private val ParquetMrPropertyPrefix = "parquet." private val CometS3CredentialProviderClassProperty = @@ -489,7 +487,7 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { case None => Some("could not resolve the write schema's field ids") case Some(fields) => fields.collectFirst { - case (name, id) if id > MaxDataFieldId => + case (name, id) if id > IcebergReflection.MaxDataFieldId => s"write schema includes metadata column $name, which iceberg-java fills with row " + "lineage and the native writer does not write" } diff --git a/spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql b/spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql new file mode 100644 index 00000000000..b9a9589da29 --- /dev/null +++ b/spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql @@ -0,0 +1,135 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you 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. + +-- The native Iceberg scan reads only the nested fields that Spark's nested schema pruning keeps. +-- These queries check the pruned reads against Spark. CometIcebergNativeSuite checks that they +-- scan fewer bytes, and covers deletes, time travel, and files without field ids. + +-- No Iceberg spark-runtime is published for Spark 4.2 yet, and the 4.0 runtime the build reuses +-- is binary-incompatible with it. See https://github.com/apache/datafusion-comet/issues/4969. +-- MaxSparkVersion: 4.1 + +-- Config: spark.sql.catalog.test_cat=org.apache.iceberg.spark.SparkCatalog +-- Config: spark.sql.catalog.test_cat.type=hadoop +-- Config: spark.sql.catalog.test_cat.warehouse=/tmp/comet-iceberg-sql-test +-- Config: spark.comet.scan.icebergNative.enabled=true + +statement +DROP TABLE IF EXISTS test_cat.db.nested_pruning + +statement +CREATE TABLE test_cat.db.nested_pruning ( + id INT, + s STRUCT>, + items ARRAY>, + m MAP> +) USING iceberg + +-- NULL structs, lists, maps, list elements, map values, and fields, so that the pruned read has +-- to rebuild the validity of every level from the leaves it keeps. +statement +INSERT INTO test_cat.db.nested_pruning +SELECT + CAST(id AS INT), + IF(id % 7 = 0, NULL, named_struct( + 'a', IF(id % 5 = 0, NULL, CAST(id AS INT)), + 'pad', repeat('p', 20), + 'inner', IF(id % 3 = 0, NULL, named_struct('b', CAST(id * 2 AS INT), 'pad', repeat('q', 20))))), + IF(id % 7 = 0, NULL, array( + named_struct('x', CAST(id AS INT), 'pad', repeat('r', 20)), + IF(id % 5 = 0, NULL, named_struct('x', CAST(-id AS INT), 'pad', repeat('s', 20))))), + IF(id % 7 = 0, NULL, map('k', IF(id % 3 = 0, NULL, + named_struct('v', CAST(id AS INT), 'pad', repeat('t', 20))))) +FROM range(100) + +query +SELECT id, s.a FROM test_cat.db.nested_pruning ORDER BY id + +-- `s IS NULL` reads the validity of `s` pruned to `s.inner.b`. +query +SELECT id, s.inner.b, s IS NULL FROM test_cat.db.nested_pruning ORDER BY id + +query +SELECT id, items.x FROM test_cat.db.nested_pruning ORDER BY id + +query +SELECT id, m['k'].v FROM test_cat.db.nested_pruning ORDER BY id + +query +SELECT id FROM test_cat.db.nested_pruning WHERE s.inner.b > 100 ORDER BY id + +-- The whole struct, so nothing is pruned. +query +SELECT id, s FROM test_cat.db.nested_pruning ORDER BY id + +statement +DROP TABLE test_cat.db.nested_pruning + +-- A partition source the query does not project is appended to the task schema at its top level. +-- `s.region` is nested, so a query that prunes it away reads with the full table schema instead, +-- where an appended `region` would clash with the top-level one. +statement +DROP TABLE IF EXISTS test_cat.db.nested_pruning_nested_source + +statement +CREATE TABLE test_cat.db.nested_pruning_nested_source ( + id INT, region STRING, s STRUCT +) USING iceberg PARTITIONED BY (s.region) + +statement +INSERT INTO test_cat.db.nested_pruning_nested_source +SELECT CAST(id AS INT), 'top', named_struct( + 'region', IF(id % 2 = 0, 'east', 'west'), 'a', CAST(id AS INT), 'pad', repeat('p', 20)) +FROM range(100) + +query +SELECT id, region, s.a FROM test_cat.db.nested_pruning_nested_source ORDER BY id + +query +SELECT id, s.region FROM test_cat.db.nested_pruning_nested_source ORDER BY id + +statement +DROP TABLE test_cat.db.nested_pruning_nested_source + +-- The partition source `p` is renamed to `q`, and a new `p` is added. A query that does not +-- project `q` appends it to the task schema under its current name, apart from the new `p`. A +-- bucket partition, because Iceberg rejects a new column named like an identity partition field. +statement +DROP TABLE IF EXISTS test_cat.db.nested_pruning_renamed_source + +statement +CREATE TABLE test_cat.db.nested_pruning_renamed_source ( + id INT, p STRING, s STRUCT +) USING iceberg PARTITIONED BY (bucket(4, p)) + +statement +INSERT INTO test_cat.db.nested_pruning_renamed_source +SELECT CAST(id AS INT), CAST(id AS STRING), named_struct( + 'a', CAST(id AS INT), 'pad', repeat('p', 20)) +FROM range(100) + +statement +ALTER TABLE test_cat.db.nested_pruning_renamed_source RENAME COLUMN p TO q + +statement +ALTER TABLE test_cat.db.nested_pruning_renamed_source ADD COLUMN p STRING + +query +SELECT id, p, s.a FROM test_cat.db.nested_pruning_renamed_source ORDER BY id + +statement +DROP TABLE test_cat.db.nested_pruning_renamed_source diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index f271f7d6ec2..5a5979403c8 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -1801,7 +1801,8 @@ class CometIcebergNativeSuite checkNestedFieldEvolutionFallback( s"SELECT id FROM $table WHERE s IS NULL ORDER BY id", "s.b (added)") - // The native read asks for the column's full nested type even when Spark prunes it. + // The native read asks only for s.a here, but the check also covers the column's full + // type in the table schema, so it still falls back. checkNestedFieldEvolutionFallback( s"SELECT id, s.a FROM $table ORDER BY id", "s.b (added)") @@ -2284,6 +2285,256 @@ class CometIcebergNativeSuite } } + // Spark's nested schema pruning reaches the native scan through the scan schema, so only the + // nested fields a query uses are read and the wide `pad` strings beside them are skipped. + // sql-tests/iceberg/nested_schema_pruning.sql checks the results. One data file holds every + // row, so each `pad` column chunk is larger than iceberg-rust's 1 MiB read coalescing, which + // would otherwise merge the reads of the kept chunks across the skipped ones. + test("nested schema pruning reads only the nested fields the query uses") { + assume(icebergAvailable, "Iceberg not available in classpath") + + withTempIcebergDir { warehouseDir => + withSQLConf( + "spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.test_cat.type" -> "hadoop", + "spark.sql.catalog.test_cat.warehouse" -> warehouseDir.getAbsolutePath, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + + val table = "test_cat.db.nested_pruning" + spark.sql(s""" + CREATE TABLE $table ( + id INT, + s STRUCT, + items ARRAY>, + m MAP> + ) USING iceberg + """) + spark.sql(s""" + INSERT INTO $table + SELECT + CAST(id AS INT), + named_struct('a', CAST(id AS INT), 'pad', pad), + array(named_struct('x', CAST(id AS INT), 'pad', pad)), + map('k', named_struct('v', CAST(id AS INT), 'pad', pad)) + FROM ( + SELECT id, concat_ws('', transform(array('a', 'b', 'c', 'd'), + salt -> sha2(concat(CAST(id AS STRING), salt), 256))) AS pad + FROM range(0, 20000, 1, 1)) + """) + + def bytesScanned(pruneNestedFields: Boolean): Long = { + var bytes = 0L + withSQLConf( + CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.key -> + pruneNestedFields.toString) { + val df = spark.sql(s"SELECT sum(s.a), count(items.x), count(m['k'].v) FROM $table") + df.collect() + val scans = collectIcebergNativeScans(df.queryExecution.executedPlan) + assert(scans.length == 1, s"expected one native scan, got ${scans.length}") + bytes = scans.head.metrics("bytes_scanned").value + } + bytes + } + val prunedBytes = bytesScanned(pruneNestedFields = true) + val fullBytes = bytesScanned(pruneNestedFields = false) + assert( + prunedBytes * 4 < fullBytes, + s"pruned read should skip the pad fields: pruned=$prunedBytes, full=$fullBytes") + + spark.sql(s"DROP TABLE $table") + } + } + } + + // A pruned task schema still needs the columns iceberg-rust uses beyond the projection: the + // partition source and the equality-delete key when the query projects neither. Tasks with + // deletes read with the pruned schema too, and the tasks of every partition share the one schema + // that appending the partition source builds. + test("nested schema pruning with deletes, partitions, and time travel") { + assume(icebergAvailable, "Iceberg not available in classpath") + + withTempIcebergDir { warehouseDir => + withSQLConf( + "spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.test_cat.type" -> "hadoop", + "spark.sql.catalog.test_cat.warehouse" -> warehouseDir.getAbsolutePath, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + + // Also checks that no task, tasks with deletes included, reads the pruned `pad` fields. + // With `oneSchema`, every task needs the same columns, so they share one task schema. + def checkPrunedNativeScan(query: String, oneSchema: Boolean = false): Unit = { + val (_, cometPlan) = checkSparkAnswer(query) + val scans = collectIcebergNativeScans(cometPlan) + assert(scans.length == 1, s"expected one native scan, got ${scans.length}\n$cometPlan") + // Planning commonData leaks manifest streams on Iceberg versions before 1.8.0. + if (icebergVersionAtLeast(1, 8)) { + val schemas = OperatorOuterClass.IcebergScanCommon + .parseFrom(scans.head.commonData) + .getSchemaPoolList + .asScala + assert( + schemas.forall(!_.contains("\"pad\"")), + s"$query: pruned field in a task schema:\n${schemas.mkString("\n")}") + assert( + !oneSchema || schemas.length == 1, + s"$query: expected one task schema, got:\n${schemas.mkString("\n")}") + } + } + def latestSnapshotId(table: String): Long = spark + .sql(s"SELECT snapshot_id FROM $table.snapshots ORDER BY committed_at DESC LIMIT 1") + .collect()(0) + .getLong(0) + + val morProperties = """ + TBLPROPERTIES ( + 'format-version' = '2', + 'write.delete.mode' = 'merge-on-read', + 'write.update.mode' = 'merge-on-read', + 'write.merge.mode' = 'merge-on-read') + """ + val rows = """ + SELECT CAST(id AS INT) AS id, IF(id % 2 = 0, 'even', 'odd') AS p, + named_struct('a', CAST(id AS INT), 'pad', repeat('x', 100)) AS s + FROM range(200) + """ + + val mor = "test_cat.db.nested_pruning_mor" + spark.sql(s""" + CREATE TABLE $mor (id INT, c STRING, s STRUCT) + USING iceberg $morProperties + """) + spark.sql(s"INSERT INTO $mor SELECT id, p, s FROM ($rows)") + val snapshotBeforeDeletes = latestSnapshotId(mor) + spark.sql(s"DELETE FROM $mor WHERE id % 10 = 0") + val snapshotWithDeletes = latestSnapshotId(mor) + commitEqualityDelete("test_cat", "db", "nested_pruning_mor", "id", 7, warehouseDir) + spark.sql(s"ALTER TABLE $mor DROP COLUMN c") + checkPrunedNativeScan(s"SELECT id, s.a FROM $mor ORDER BY id") + // The equality-delete key `id` is not projected. Iceberg gives the delete only to the data + // file whose `id` range holds 7, so only that file's task appends `id`. + checkPrunedNativeScan(s"SELECT s.a FROM $mor ORDER BY s.a") + checkPrunedNativeScan( + s"SELECT id, s.a FROM $mor VERSION AS OF $snapshotBeforeDeletes ORDER BY id") + // `c` was dropped after this snapshot, so the current table schema lacks it, and these + // tasks with deletes must read with the scan schema. + checkPrunedNativeScan( + s"SELECT id, c, s.a FROM $mor VERSION AS OF $snapshotWithDeletes ORDER BY id") + + val partitioned = "test_cat.db.nested_pruning_partitioned" + spark.sql(s""" + CREATE TABLE $partitioned (id INT, p STRING, s STRUCT) + USING iceberg PARTITIONED BY (p) $morProperties + """) + spark.sql(s"INSERT INTO $partitioned $rows") + spark.sql(s"DELETE FROM $partitioned WHERE id % 10 = 0") + // The partition source `p` is not projected. Every task appends it, and they share the + // memoized result. + checkPrunedNativeScan(s"SELECT id, s.a FROM $partitioned ORDER BY id", oneSchema = true) + checkPrunedNativeScan(s"SELECT p, count(s.a) FROM $partitioned GROUP BY p ORDER BY p") + + Seq(mor, partitioned).foreach(t => spark.sql(s"DROP TABLE $t")) + } + } + } + + // Data files written without Iceberg field ids and read with no name mapping: iceberg-rust then + // projects whole top-level columns by position, and a by-name struct cast narrows each one to + // the pruned type. Spark's own reader returns NULL nested fields for such a table, so the pruned + // read is checked against the native read of every nested field instead. + test("nested schema pruning on data files without field ids") { + assume(icebergAvailable, "Iceberg not available in classpath") + + withTempIcebergDir { warehouseDir => + withSQLConf( + "spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.test_cat.type" -> "hadoop", + "spark.sql.catalog.test_cat.warehouse" -> warehouseDir.getAbsolutePath, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + + val dataPath = s"${warehouseDir.getAbsolutePath}/nested_no_ids" + spark + .sql(""" + SELECT + CAST(id AS INT) AS id, + named_struct('a', CAST(id AS INT), 'pad', repeat('x', 50), 'b', CAST(-id AS INT)) + AS s, + array(named_struct('x', CAST(id AS INT), 'pad', repeat('y', 50))) AS items + FROM range(100) + """) + .coalesce(1) + .write + .parquet(dataPath) + spark.sql("CREATE NAMESPACE IF NOT EXISTS test_cat.db") + val table = "test_cat.db.nested_no_ids" + spark.sql(s""" + CREATE TABLE $table ( + id INT, + s STRUCT, + items ARRAY> + ) USING iceberg + """) + try { + val tableUtilClass = Class.forName("org.apache.iceberg.spark.SparkTableUtil") + val icebergTable = spark.sessionState.catalogManager + .catalog("test_cat") + .asInstanceOf[org.apache.iceberg.spark.SparkCatalog] + .loadTable( + org.apache.spark.sql.connector.catalog.Identifier.of(Array("db"), "nested_no_ids")) + .asInstanceOf[org.apache.iceberg.spark.source.SparkTable] + .table() + spark.sql(s"CREATE TABLE parquet_no_ids USING parquet LOCATION '$dataPath'") + tableUtilClass + .getMethod( + "importSparkTable", + classOf[org.apache.spark.sql.SparkSession], + classOf[org.apache.spark.sql.catalyst.TableIdentifier], + classOf[org.apache.iceberg.Table], + classOf[String]) + .invoke( + null, + spark, + new org.apache.spark.sql.catalyst.TableIdentifier("parquet_no_ids"), + icebergTable, + s"${warehouseDir.getAbsolutePath}/staging") + spark.sql( + s"ALTER TABLE $table UNSET TBLPROPERTIES IF EXISTS ('schema.name-mapping.default')") + + def nativeRows(query: String, pruneNestedFields: Boolean): Seq[Row] = { + var rows = Seq.empty[Row] + withSQLConf( + CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.key -> + pruneNestedFields.toString) { + val df = spark.sql(query) + rows = df.collect().toSeq + assertSingleNativeScan(df.queryExecution.executedPlan) + } + rows + } + // `s.b` is the struct's last field, so a positional cast would read `s.a` instead. + Seq( + s"SELECT id, s.b FROM $table ORDER BY id", + s"SELECT id, items.x FROM $table ORDER BY id").foreach { query => + val full = nativeRows(query, pruneNestedFields = false) + assert(full.length == 100 && full.forall(!_.isNullAt(1)), s"$query: $full") + assert(nativeRows(query, pruneNestedFields = true) == full, query) + } + + spark.sql(s"DROP TABLE $table") + spark.sql("DROP TABLE parquet_no_ids") + } catch { + case _: ClassNotFoundException => + cancel("SparkTableUtil not available") + } + } + } + } + test("UUID type - native Iceberg UUID column (reproduces type mismatch)") { assume(icebergAvailable, "Iceberg not available in classpath") diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergReadBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergReadBenchmark.scala index 1a171a5adac..ab5449bcc39 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergReadBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergReadBenchmark.scala @@ -19,7 +19,10 @@ package org.apache.spark.sql.benchmark +import java.io.File + import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.comet.CometIcebergNativeScanExec import org.apache.spark.sql.types._ import org.apache.comet.CometConf @@ -28,7 +31,8 @@ import org.apache.comet.CometConf * Benchmark to measure Comet Iceberg read performance. To run this benchmark: * `SPARK_GENERATE_BENCHMARK_FILES=1 make * benchmark-org.apache.spark.sql.benchmark.CometIcebergReadBenchmark` Results will be written to - * "spark/benchmarks/CometIcebergReadBenchmark-**results.txt". + * "spark/benchmarks/CometIcebergReadBenchmark-**results.txt". Pass `--nested-only` to run only + * the nested column cases. */ object CometIcebergReadBenchmark extends CometBenchmarkBase { @@ -75,10 +79,107 @@ object CometIcebergReadBenchmark extends CometBenchmarkBase { } } + /** + * Reads a TPC-H style nested `partsupp` table, whose `partsupp_data` column is an + * `array` dominated by a text field, as in the nested TPC-H schema. Spark prunes the + * struct to the fields a query uses. The pruned Comet case passes that pruned schema to + * iceberg-rust, and the other reads every nested field and drops the unused ones after + * decoding. + */ + def icebergNestedProjectionBenchmark(values: Int): Unit = { + withTempPath { dir => + configureIcebergHadoopCatalog(new File(dir, "iceberg-warehouse")) + val table = s"$defaultIcebergCatalog.db.partsupp_nested" + spark.sql(s""" + CREATE TABLE $table ( + ps_partkey BIGINT, + ps_suppkey BIGINT, + partsupp_data ARRAY> + ) USING iceberg + TBLPROPERTIES ('format-version' = '2', 'write.parquet.compression-codec' = 'zstd') + """) + // TPC-H style comments: 16 words from a small vocabulary, about 128 characters each. + val words = ("furiously carefully quickly slyly blithely fluffily ironic final regular " + + "express special pending unusual bold even silent idle ruthless packages requests " + + "accounts deposits foxes ideas theodolites pinto beans instructions dependencies " + + "excuses platelets asymptotes courts dolphins sleep wake haggle nag use boost") + .split(" ") + .map(w => s"'$w'") + .mkString(", ") + spark.sql(s""" + INSERT INTO $table + SELECT id DIV 4, id % 4 + 1, array(named_struct( + 'ps_availqty', CAST(pmod(xxhash64(id, 1), 9999) + 1 AS INT), + 'ps_supplycost', CAST(pmod(xxhash64(id, 2), 99901) / 100 + 1 AS DECIMAL(15, 2)), + 'ps_comment', concat_ws(' ', transform(sequence(1, 16), + i -> element_at(array($words), CAST(pmod(xxhash64(id, i), 40) + 1 AS INT)))))) + FROM range($values) + """) + + val cometIceberg = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") + def pruneNestedFields(enabled: Boolean): (String, String) = + CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.key -> enabled.toString + + def explodeField(aggregate: String, field: String): String = + s"SELECT $aggregate FROM $table LATERAL VIEW explode(partsupp_data.$field) e AS v" + + Seq( + "one small field, as TPC-H q20 reads partsupp" -> explodeField("sum(v)", "ps_availqty"), + "the text field, which pruning cannot skip" -> + explodeField("sum(length(v))", "ps_comment")) + .foreach { case (name, query) => + val benchmark = + new Benchmark(s"Nested Iceberg column, $name", values.toLong, output = output) + + benchmark.addCase("SQL Iceberg - Spark") { _ => + spark.sql(query).noop() + } + benchmark.addCase("SQL Iceberg - Comet, every nested field") { _ => + withSQLConf(cometIceberg :+ pruneNestedFields(false): _*) { + spark.sql(query).noop() + } + } + benchmark.addCase("SQL Iceberg - Comet, pruned nested fields") { _ => + withSQLConf(cometIceberg :+ pruneNestedFields(true): _*) { + spark.sql(query).noop() + } + } + + val bytesScanned = Seq(false, true).map { enabled => + var bytes = 0L + withSQLConf(cometIceberg :+ pruneNestedFields(enabled): _*) { + val df = spark.sql(query) + df.collect() + bytes = collect(df.queryExecution.executedPlan) { + case scan: CometIcebergNativeScanExec => scan.metrics("bytes_scanned").value + }.sum + } + bytes / 1024.0 / 1024.0 + } + warn( + benchmark, + f"Comet bytes scanned: every nested field ${bytesScanned.head}%.1f MiB, " + + f"pruned nested fields ${bytesScanned.last}%.1f MiB") + benchmark.run() + } + + spark.sql(s"DROP TABLE $table") + } + } + override def runCometBenchmark(mainArgs: Array[String]): Unit = { - runBenchmarkWithTable("SQL Single Numeric Iceberg Column Scan", 1024 * 1024 * 128) { v => - Seq(BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, DoubleType) - .foreach(icebergScanBenchmark(v, _)) + if (!mainArgs.contains("--nested-only")) { + runBenchmarkWithTable("SQL Single Numeric Iceberg Column Scan", 1024 * 1024 * 128) { v => + Seq(BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, DoubleType) + .foreach(icebergScanBenchmark(v, _)) + } + } + runBenchmark("SQL Nested Iceberg Column Scan") { + icebergNestedProjectionBenchmark(1024 * 1024 * 4) } } } From 335e7e0cafb80d43d3f2d808345187eebc20461e Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 9 Oct 2026 13:10:30 -0700 Subject: [PATCH 2/9] feedback --- .../comet/iceberg/IcebergReflection.scala | 34 +++++++++------- .../comet/CometIcebergNativeSuite.scala | 40 ++++++++++++++----- 2 files changed, 51 insertions(+), 23 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index d4e368970d9..8ef3f04d7d7 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -719,8 +719,9 @@ object IcebergReflection extends Logging { /** * Returns a schema equal to `baseSchema` but guaranteed to contain `requiredFieldIds`. Any id - * not already present is resolved from the table's current schema, or from its schema history - * (`table.schemas()`) for a column since dropped, and appended. + * not already present is appended, resolved from the table's current schema, or from its schema + * history (`table.schemas()`) when the column has since been dropped or its current name is + * taken. * * Two callers need this: an equality delete may be keyed on a column since dropped from the * current schema (mirroring Iceberg-Java's `DeleteFilter.fileProjection`), and a partition @@ -730,6 +731,7 @@ object IcebergReflection extends Logging { * CometScanRule is responsible for falling back before we get here. */ def schemaWithRequiredFields(baseSchema: Any, table: Any, requiredFieldIds: Seq[Int]): Any = { + import scala.jdk.CollectionConverters._ // `findFieldObject` searches recursively, so a source column already present as a nested field // (e.g. `s.region`) is not re-appended at the top level, which would create a duplicate field // id. Empty `requiredFieldIds` (the common non-partitioned, no-delete task) does no lookups. @@ -741,23 +743,27 @@ object IcebergReflection extends Logging { logDebug( s"Native Iceberg scan schema is missing field id(s) ${missingIds.mkString(",")}; " + "resolving them from table schema history") - // table.schemas() lists the oldest schema first, so a column the table still has is looked - // up in the current schema before it, to keep its current name and type. An older name - // could clash with a column that took it over. - val schemas = getMethod(table.getClass, "schema").invoke(table) +: getAllSchemas(table) - val resolvedFields = missingIds.map { id => - schemas.iterator - .map(findFieldObject(_, id)) - .collectFirst { case Some(field) => field } - .getOrElse(throw new IllegalStateException( - s"Cannot resolve field id $id in table schema history")) - } val existing = getMethod(baseSchema.getClass, "columns") .invoke(baseSchema) .asInstanceOf[java.util.List[_]] val newColumns = new java.util.ArrayList[Any](existing) - resolvedFields.foreach(newColumns.add) + // A field is appended under a name it has had that the task schema does not use yet. The + // current schema comes first, to keep a live column's current name and type, and then + // table.schemas(), oldest first. A VERSION AS OF scan schema carries the snapshot's names, + // so a current name can clash with a column the snapshot still has. + val schemas = getMethod(table.getClass, "schema").invoke(table) +: getAllSchemas(table) + val names = scala.collection.mutable.Set(existing.asScala.map(fieldName).toSeq: _*) + missingIds.foreach { id => + val (field, name) = schemas.iterator + .flatMap(findFieldObject(_, id)) + .map(f => (f, fieldName(f))) + .find { case (_, n) => !names.contains(n) } + .getOrElse(throw new IllegalStateException( + s"Cannot resolve field id $id in table schema history under an unused name")) + names += name + newColumns.add(field) + } baseSchema.getClass .getConstructor(classOf[java.util.List[_]]) .newInstance(newColumns) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 5a5979403c8..d090448433c 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -2351,7 +2351,8 @@ class CometIcebergNativeSuite // A pruned task schema still needs the columns iceberg-rust uses beyond the projection: the // partition source and the equality-delete key when the query projects neither. Tasks with // deletes read with the pruned schema too, and the tasks of every partition share the one schema - // that appending the partition source builds. + // that appending the partition source builds. VERSION AS OF reads of columns dropped, renamed, + // or swapped after the snapshot use the snapshot's names. test("nested schema pruning with deletes, partitions, and time travel") { assume(icebergAvailable, "Iceberg not available in classpath") @@ -2404,37 +2405,58 @@ class CometIcebergNativeSuite val mor = "test_cat.db.nested_pruning_mor" spark.sql(s""" - CREATE TABLE $mor (id INT, c STRING, s STRUCT) + CREATE TABLE $mor ( + id INT, c STRING, r STRING, x INT, y INT, s STRUCT) USING iceberg $morProperties """) - spark.sql(s"INSERT INTO $mor SELECT id, p, s FROM ($rows)") + spark.sql(s""" + INSERT INTO $mor + SELECT id, p, concat('r', CAST(id AS STRING)), id, -id, s FROM ($rows) + """) val snapshotBeforeDeletes = latestSnapshotId(mor) spark.sql(s"DELETE FROM $mor WHERE id % 10 = 0") val snapshotWithDeletes = latestSnapshotId(mor) commitEqualityDelete("test_cat", "db", "nested_pruning_mor", "id", 7, warehouseDir) - spark.sql(s"ALTER TABLE $mor DROP COLUMN c") + // Drop `c`, rename `r`, and swap the names of `x` and `y`. + Seq( + "DROP COLUMN c", + "RENAME COLUMN r TO r2", + "RENAME COLUMN x TO tmp", + "RENAME COLUMN y TO x", + "RENAME COLUMN tmp TO y").foreach(change => spark.sql(s"ALTER TABLE $mor $change")) checkPrunedNativeScan(s"SELECT id, s.a FROM $mor ORDER BY id") // The equality-delete key `id` is not projected. Iceberg gives the delete only to the data // file whose `id` range holds 7, so only that file's task appends `id`. checkPrunedNativeScan(s"SELECT s.a FROM $mor ORDER BY s.a") checkPrunedNativeScan( s"SELECT id, s.a FROM $mor VERSION AS OF $snapshotBeforeDeletes ORDER BY id") - // `c` was dropped after this snapshot, so the current table schema lacks it, and these - // tasks with deletes must read with the scan schema. + // The current table schema lacks `c` and names `r`, `x`, and `y` differently, so these + // tasks with deletes must read with the scan schema, which keeps the snapshot's names. checkPrunedNativeScan( - s"SELECT id, c, s.a FROM $mor VERSION AS OF $snapshotWithDeletes ORDER BY id") + s"SELECT id, c, r, x, y, s.a FROM $mor VERSION AS OF $snapshotWithDeletes ORDER BY id") val partitioned = "test_cat.db.nested_pruning_partitioned" spark.sql(s""" - CREATE TABLE $partitioned (id INT, p STRING, s STRUCT) + CREATE TABLE $partitioned (id INT, p STRING, c STRING, s STRUCT) USING iceberg PARTITIONED BY (p) $morProperties """) - spark.sql(s"INSERT INTO $partitioned $rows") + spark.sql(s""" + INSERT INTO $partitioned + SELECT id, p, concat('c', CAST(id AS STRING)), s FROM ($rows) + """) spark.sql(s"DELETE FROM $partitioned WHERE id % 10 = 0") + val partitionedSnapshot = latestSnapshotId(partitioned) // The partition source `p` is not projected. Every task appends it, and they share the // memoized result. checkPrunedNativeScan(s"SELECT id, s.a FROM $partitioned ORDER BY id", oneSchema = true) checkPrunedNativeScan(s"SELECT p, count(s.a) FROM $partitioned GROUP BY p ORDER BY p") + // Once `p` takes the dropped `c`'s name, the partition source's current name clashes with + // the snapshot's `c`, so it is appended under its older name. + spark.sql(s"ALTER TABLE $partitioned DROP COLUMN c") + spark.sql(s"ALTER TABLE $partitioned RENAME COLUMN p TO c") + checkPrunedNativeScan( + s"SELECT id, c, s.a FROM $partitioned VERSION AS OF $partitionedSnapshot ORDER BY id", + oneSchema = true) Seq(mor, partitioned).foreach(t => spark.sql(s"DROP TABLE $t")) } From d88698c4425563e5a498e4bfeffd86cadc3a8a2b Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 10 Oct 2026 11:29:11 -0700 Subject: [PATCH 3/9] Resolve appended task-schema columns from one schema Take every missing partition source and equality-delete key from one schema, the current one or the newest in the history that has them all under names the task schema does not use. Picking a name per field clashed when two partition sources were renamed into each other's names. Tasks that need no extra columns skip the memo lookups, and tasks without deletes skip the equality-key work. --- .../comet/iceberg/IcebergReflection.scala | 46 ++++++++++--------- .../operator/CometIcebergNativeScan.scala | 41 +++++++++++------ .../comet/CometIcebergNativeSuite.scala | 19 +++++++- 3 files changed, 69 insertions(+), 37 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index 8ef3f04d7d7..f6dd0998287 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -718,10 +718,10 @@ object IcebergReflection extends Logging { .invoke(schema, fieldId.asInstanceOf[AnyRef])) /** - * Returns a schema equal to `baseSchema` but guaranteed to contain `requiredFieldIds`. Any id - * not already present is appended, resolved from the table's current schema, or from its schema - * history (`table.schemas()`) when the column has since been dropped or its current name is - * taken. + * Returns a schema equal to `baseSchema` but guaranteed to contain `requiredFieldIds`. The ids + * not already present are appended as they are in one schema: the table's current schema, or + * the newest schema in its history (`table.schemas()`) that has them all under names + * `baseSchema` does not use. * * Two callers need this: an equality delete may be keyed on a column since dropped from the * current schema (mirroring Iceberg-Java's `DeleteFilter.fileProjection`), and a partition @@ -747,23 +747,27 @@ object IcebergReflection extends Logging { getMethod(baseSchema.getClass, "columns") .invoke(baseSchema) .asInstanceOf[java.util.List[_]] - val newColumns = new java.util.ArrayList[Any](existing) - // A field is appended under a name it has had that the task schema does not use yet. The - // current schema comes first, to keep a live column's current name and type, and then - // table.schemas(), oldest first. A VERSION AS OF scan schema carries the snapshot's names, - // so a current name can clash with a column the snapshot still has. - val schemas = getMethod(table.getClass, "schema").invoke(table) +: getAllSchemas(table) - val names = scala.collection.mutable.Set(existing.asScala.map(fieldName).toSeq: _*) - missingIds.foreach { id => - val (field, name) = schemas.iterator - .flatMap(findFieldObject(_, id)) - .map(f => (f, fieldName(f))) - .find { case (_, n) => !names.contains(n) } - .getOrElse(throw new IllegalStateException( - s"Cannot resolve field id $id in table schema history under an unused name")) - names += name - newColumns.add(field) - } + val existingNames = existing.asScala.map(fieldName).toSet + // Taking every field from one schema keeps their names apart, which a name picked per field + // cannot promise once columns have been renamed into each other's names. A VERSION AS OF + // scan schema carries the snapshot's names, so a schema also must not name a field like a + // column `baseSchema` has. The current schema comes first, to keep current names, and then + // the others newest first, so that a promoted column keeps its widest type. + val schemas = getMethod(table.getClass, "schema").invoke(table) +: + getAllSchemas(table).reverse + val resolvedFields = schemas.iterator + .map(schema => missingIds.flatMap(findFieldObject(schema, _))) + .find { fields => + val names = fields.map(fieldName) + fields.length == missingIds.length && names.distinct.length == names.length && + !names.exists(existingNames.contains) + } + .getOrElse(throw new IllegalStateException( + s"Cannot resolve field ids ${missingIds.mkString(",")} from one schema in the table " + + "schema history under names the task schema does not use")) + val newColumns = new java.util.ArrayList[Any](existing.size + resolvedFields.length) + newColumns.addAll(existing) + resolvedFields.foreach(newColumns.add) baseSchema.getClass .getConstructor(classOf[java.util.List[_]]) .newInstance(newColumns) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index ca6941ce16c..adc0b64fc5b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -1035,22 +1035,28 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // that needs one reads with the full schema instead. val canUsePrunedSchemaCache = mutable.HashMap[Seq[Int], Boolean]() def canUsePrunedSchema(requiredFieldIds: Seq[Int]): Boolean = - nestedSchemaPruningEnabled && canUsePrunedSchemaCache.getOrElseUpdate( - requiredFieldIds, - requiredFieldIds.forall { id => - IcebergReflection.findFieldObject(metadata.prunedScanSchema, id).isDefined || - !IcebergReflection.isNestedField(metadata.tableSchema, id) - }) + nestedSchemaPruningEnabled && (requiredFieldIds.isEmpty || + canUsePrunedSchemaCache.getOrElseUpdate( + requiredFieldIds, + requiredFieldIds.forall { id => + IcebergReflection.findFieldObject(metadata.prunedScanSchema, id).isDefined || + !IcebergReflection.isNestedField(metadata.tableSchema, id) + })) // schemaWithRequiredFields builds a new Schema whenever it appends a column. Memoize it by // base schema (Iceberg's Schema compares by identity) and ids, so tasks that need the same - // columns share one schema object, and with it one schema pool entry. + // columns share one schema object, and with it one schema pool entry. Most tasks need no extra + // columns, so they skip both lookups. val requiredFieldsCache = mutable.HashMap[(AnyRef, Seq[Int]), AnyRef]() def withRequiredFields(baseSchema: AnyRef, requiredFieldIds: Seq[Int]): AnyRef = - requiredFieldsCache.getOrElseUpdate( - (baseSchema, requiredFieldIds), - IcebergReflection - .schemaWithRequiredFields(baseSchema, metadata.table, requiredFieldIds) - .asInstanceOf[AnyRef]) + if (requiredFieldIds.isEmpty) { + baseSchema + } else { + requiredFieldsCache.getOrElseUpdate( + (baseSchema, requiredFieldIds), + IcebergReflection + .schemaWithRequiredFields(baseSchema, metadata.table, requiredFieldIds) + .asInstanceOf[AnyRef]) + } val perPartitionBuilders = mutable.ArrayBuffer[OperatorOuterClass.IcebergScan]() @@ -1165,9 +1171,14 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // so union the equality-delete field ids into the task schema, resolving any // dropped ones from the table's schema history (mirrors Iceberg-Java's // DeleteFilter.fileProjection). - val equalityFieldIds = deletes.asScala.flatMap { df => - requiredEqualityFieldIds(deleteFileClass, df).asScala.map(_.intValue()) - }.toSeq + val equalityFieldIds = + if (hasDeletes) { + deletes.asScala.flatMap { df => + requiredEqualityFieldIds(deleteFileClass, df).asScala.map(_.intValue()) + }.toSeq + } else { + Nil + } // iceberg-rust validates a FileScanTask by resolving its partition spec against // the task schema, so a task carrying a partition spec needs that spec's source diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index d090448433c..8ebff034097 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -2458,7 +2458,24 @@ class CometIcebergNativeSuite s"SELECT id, c, s.a FROM $partitioned VERSION AS OF $partitionedSnapshot ORDER BY id", oneSchema = true) - Seq(mor, partitioned).foreach(t => spark.sql(s"DROP TABLE $t")) + // Two partition sources renamed into each other's names: `q` takes the dropped `c`'s name + // and `p` takes `q`'s. Names picked one field at a time clash here, so both are appended + // as one older schema names them. + val twoSources = "test_cat.db.nested_pruning_two_sources" + spark.sql(s""" + CREATE TABLE $twoSources (id INT, p INT, q INT, c INT, s STRUCT) + USING iceberg PARTITIONED BY (bucket(4, p), bucket(4, q)) + """) + spark.sql(s"INSERT INTO $twoSources SELECT id, id % 3, id % 5, id * 10, s FROM ($rows)") + val twoSourcesSnapshot = latestSnapshotId(twoSources) + Seq("DROP COLUMN c", "RENAME COLUMN q TO c", "RENAME COLUMN p TO q").foreach { change => + spark.sql(s"ALTER TABLE $twoSources $change") + } + checkPrunedNativeScan( + s"SELECT id, c, s.a FROM $twoSources VERSION AS OF $twoSourcesSnapshot ORDER BY id", + oneSchema = true) + + Seq(mor, partitioned, twoSources).foreach(t => spark.sql(s"DROP TABLE $t")) } } } From 6213d0576ec743ea5ef567788469b095d1371a52 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 12 Oct 2026 11:34:37 -0700 Subject: [PATCH 4/9] Build one task schema per scan for native Iceberg reads iceberg-rust caches an equality-delete predicate per delete file, with the column names of the task that loads it first, and binds it by name in every task of a reader. Task schemas built per task could name a key differently across tasks, which failed the scan or, after two columns swapped names, deleted the wrong rows. Build one pruned and one full task schema per scan during planning. The names come from the newest schema whose top-level columns match the scan schema, and an equality key dropped before it takes one name from the schema history. A scan whose keys cannot be named that way falls back to Spark. Serialization now builds the task schema and the projection once per scan. --- .../comet/iceberg/IcebergReflection.scala | 197 ++++++++++-------- .../apache/comet/rules/CometScanRule.scala | 74 +++++-- .../operator/CometIcebergNativeScan.scala | 153 +++----------- .../comet/CometIcebergNativeSuite.scala | 160 +++++++++++--- 4 files changed, 323 insertions(+), 261 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index f6dd0998287..24af10ccf15 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -709,8 +709,7 @@ object IcebergReflection extends Logging { /** * Returns the `Types.NestedField` for `fieldId` at any depth of `schema`, or None when `schema` - * has no such field. A reflection failure propagates rather than reading as a missing field, - * because the callers run at serialization time, when the scan can no longer fall back. + * has no such field. A reflection failure propagates rather than reading as a missing field. */ def findFieldObject(schema: Any, fieldId: Int): Option[Any] = Option( @@ -718,62 +717,92 @@ object IcebergReflection extends Logging { .invoke(schema, fieldId.asInstanceOf[AnyRef])) /** - * Returns a schema equal to `baseSchema` but guaranteed to contain `requiredFieldIds`. The ids - * not already present are appended as they are in one schema: the table's current schema, or - * the newest schema in its history (`table.schemas()`) that has them all under names - * `baseSchema` does not use. + * The task schemas of a native Iceberg scan, built once for all of its tasks during planning: + * the pruned scan schema, and the full schema read when pruning is off or cannot apply. Both + * hold the partition sources and equality-delete keys (`requiredFieldIds`) the tasks need + * beyond the projection: iceberg-rust resolves a task's partition spec against its schema, and + * reads the keys to apply equality deletes. The pruned one is None when a required column is + * nested inside a struct the query pruned away, as appending it at the top level would put it + * outside its struct. The full one is None, so that the scan falls back, when no name fits. * - * Two callers need this: an equality delete may be keyed on a column since dropped from the - * current schema (mirroring Iceberg-Java's `DeleteFilter.fileProjection`), and a partition - * spec's source columns must be present for iceberg-rust to resolve `partition_type` when the - * query projects them out. Called at serialization time, so it throws on failure (a required id - * that cannot be resolved, or any reflection error) rather than silently degrading. - * CometScanRule is responsible for falling back before we get here. - */ - def schemaWithRequiredFields(baseSchema: Any, table: Any, requiredFieldIds: Seq[Int]): Any = { + * Every task reads with the same schema because iceberg-rust builds an equality-delete + * predicate once per delete file, with the column names of the task that loads it first, and + * binds it by name in every other task of the reader. The names come from the full schema, the + * newest schema, current first, that names every column of the scan schema the same way. That + * is the current schema unless a VERSION AS OF snapshot named a column differently. A key that + * schema lacks, dropped before it, takes the newest name that no other column has. + */ + def nativeTaskSchemas( + table: Any, + tableSchema: Any, + scanSchema: Any, + requiredFieldIds: Set[Int]): (Option[Any], Option[Any]) = { + val prunedScanSchema = withoutMetadataColumns(scanSchema) + val scanColumns = columnNamesById(prunedScanSchema).toSet + val schemas = tableSchema +: getAllSchemas(table).reverse + val full = schemas.iterator + .map(schema => schema -> columnNamesById(schema)) + .find { case (_, columns) => scanColumns.subsetOf(columns.toSet) } + full match { + case None => (None, None) + case Some((fullSchema, fullColumns)) => + val fullNames = fullColumns.values.toSet + val ids = requiredFieldIds.toSeq.sorted + val inFull = ids.flatMap(id => findFieldObject(fullSchema, id).map(id -> _)).toMap + // Keys dropped before `fullSchema`, each under a name that no other column has. + val dropped = + ids.filterNot(inFull.contains).foldLeft(Option(Map.empty[Int, Any])) { (resolved, id) => + resolved.flatMap { fields => + val taken = fullNames ++ fields.values.map(fieldName) + schemas.iterator + .flatMap(findFieldObject(_, id)) + .find(field => !taken.contains(fieldName(field))) + .map(field => fields + (id -> field)) + } + } + dropped match { + case None => (None, None) + case Some(droppedFields) => + val missing = ids.filter(findFieldObject(prunedScanSchema, _).isEmpty) + val prunedTaskSchema = + if (missing.forall(id => fullColumns.contains(id) || !inFull.contains(id))) { + val fields = missing.map(id => inFull.getOrElse(id, droppedFields(id))) + Some(withColumns(prunedScanSchema, fields)) + } else { + None + } + val fullFields = droppedFields.toSeq.sortBy(_._1).map(_._2) + (prunedTaskSchema, Some(withColumns(fullSchema, fullFields))) + } + } + } + + /** The top-level columns of `schema`, as field id to name. */ + private def columnNamesById(schema: Any): Map[Int, String] = { import scala.jdk.CollectionConverters._ - // `findFieldObject` searches recursively, so a source column already present as a nested field - // (e.g. `s.region`) is not re-appended at the top level, which would create a duplicate field - // id. Empty `requiredFieldIds` (the common non-partitioned, no-delete task) does no lookups. - val missingIds = - requiredFieldIds.distinct.filter(id => findFieldObject(baseSchema, id).isEmpty) - if (missingIds.isEmpty) { - baseSchema + getMethod(schema.getClass, "columns") + .invoke(schema) + .asInstanceOf[java.util.List[_]] + .asScala + .map(column => fieldIdOf(column) -> fieldName(column)) + .toMap + } + + /** `schema` with `fields` appended at its top level, or `schema` itself when there are none. */ + private def withColumns(schema: Any, fields: Seq[Any]): Any = + if (fields.isEmpty) { + schema } else { - logDebug( - s"Native Iceberg scan schema is missing field id(s) ${missingIds.mkString(",")}; " + - "resolving them from table schema history") - val existing = - getMethod(baseSchema.getClass, "columns") - .invoke(baseSchema) - .asInstanceOf[java.util.List[_]] - val existingNames = existing.asScala.map(fieldName).toSet - // Taking every field from one schema keeps their names apart, which a name picked per field - // cannot promise once columns have been renamed into each other's names. A VERSION AS OF - // scan schema carries the snapshot's names, so a schema also must not name a field like a - // column `baseSchema` has. The current schema comes first, to keep current names, and then - // the others newest first, so that a promoted column keeps its widest type. - val schemas = getMethod(table.getClass, "schema").invoke(table) +: - getAllSchemas(table).reverse - val resolvedFields = schemas.iterator - .map(schema => missingIds.flatMap(findFieldObject(schema, _))) - .find { fields => - val names = fields.map(fieldName) - fields.length == missingIds.length && names.distinct.length == names.length && - !names.exists(existingNames.contains) - } - .getOrElse(throw new IllegalStateException( - s"Cannot resolve field ids ${missingIds.mkString(",")} from one schema in the table " + - "schema history under names the task schema does not use")) - val newColumns = new java.util.ArrayList[Any](existing.size + resolvedFields.length) - newColumns.addAll(existing) - resolvedFields.foreach(newColumns.add) - baseSchema.getClass + val columns = + getMethod(schema.getClass, "columns").invoke(schema).asInstanceOf[java.util.List[_]] + val newColumns = new java.util.ArrayList[Any](columns.size + fields.length) + newColumns.addAll(columns) + fields.foreach(newColumns.add) + schema.getClass .getConstructor(classOf[java.util.List[_]]) .newInstance(newColumns) .asInstanceOf[AnyRef] } - } /** The spec reserves field ids above `Integer.MAX_VALUE - 200` for metadata columns. */ val MaxDataFieldId: Int = Int.MaxValue - 200 @@ -798,44 +827,16 @@ object IcebergReflection extends Logging { } } - /** Whether `fieldId` is a field of `schema` nested below its top level. */ - def isNestedField(schema: Any, fieldId: Int): Boolean = { - import scala.jdk.CollectionConverters._ - findFieldObject(schema, fieldId).isDefined && - !getMethod(schema.getClass, "columns") - .invoke(schema) - .asInstanceOf[java.util.List[_]] - .asScala - .exists(fieldIdOf(_) == fieldId) - } - /** - * The source column field ids referenced by a task's partition spec. + * The `sourceId`s of a partition spec's fields, in spec order. Empty for an unpartitioned spec. * * iceberg-rust validates a `FileScanTask` by resolving its partition spec against the task - * schema (`partition_type(schema)`), so a task carrying a partition spec needs those source - * columns present in the schema even when the query projects them out (for example selecting - * only `_spec_id` or `_partition`). These ids are unioned into the native task schema via - * [[schemaWithRequiredFields]]. project_field_ids still drives the read, so the columns are not - * materialized into the scan output. All ids resolve from the table's schema history, including - * a source column later dropped by partition evolution. - * - * Returns an empty sequence when the task has no partition spec. - */ - def partitionSourceFieldIds(task: Any, fileScanTaskClass: Class[_]): Seq[Int] = { - val spec = - try { - getMethod(fileScanTaskClass, "spec").invoke(task) - } catch { - case _: Exception => null - } - if (spec == null) Seq.empty else partitionFieldSourceIds(spec) - } - - /** - * The `sourceId`s of a partition spec's fields, in spec order. Empty for an unpartitioned spec. + * schema (`partition_type(schema)`), so the source columns must be in the task schema even when + * the query projects them out (for example selecting only `_spec_id` or `_partition`). See + * [[nativeTaskSchemas]]. project_field_ids still drives the read, so the columns are not + * materialized into the scan output. */ - private def partitionFieldSourceIds(spec: Any): Seq[Int] = { + def partitionFieldSourceIds(spec: Any): Seq[Int] = { import scala.jdk.CollectionConverters._ // Exclude fields whose source column was dropped (unknown result type): their source is absent // from the current schema and the serialized spec/values omit them too (see @@ -2390,10 +2391,14 @@ object IcebergReflection extends Logging { * List of FileScanTask objects from Iceberg planning * @param scanSchema * The expectedSchema from the SparkScan (for schema evolution / VERSION AS OF) - * @param prunedScanSchema - * scanSchema without its metadata columns, which iceberg-rust resolves by reserved field id - * rather than from the task schema. Holds only the nested fields Spark's schema pruning keeps, - * so a task reading with it skips the others. + * @param prunedTaskSchema + * The schema every task reads with when nested schema pruning is on: scanSchema without its + * metadata columns, which iceberg-rust resolves by reserved field id, plus the partition + * sources and equality-delete keys the tasks need. None when one of those is nested inside a + * struct the query pruned away. See [[IcebergReflection.nativeTaskSchemas]]. + * @param fullTaskSchema + * The schema every task reads with otherwise: the scan's full table schema plus the + * equality-delete keys it lacks. None when no name fits every task, which falls back to Spark. * @param globalFieldIdMapping * Mapping from column names to Iceberg field IDs (built from scanSchema) * @param catalogProperties @@ -2408,7 +2413,8 @@ case class CometIcebergNativeScanMetadata( nameMapping: Option[String], @transient tasks: java.util.List[_], scanSchema: Any, - prunedScanSchema: Any, + prunedTaskSchema: Option[Any], + fullTaskSchema: Option[Any], tableSchema: Any, globalFieldIdMapping: Map[String, Int], catalogProperties: Map[String, String], @@ -2433,6 +2439,9 @@ object CometIcebergNativeScanMetadata extends Logging { * The scan's FileScanTasks (already extracted). Passed in rather than re-read via * [[IcebergReflection.getTasks]], which for a staged scan rebuilds a flattened list of every * task on each call. + * @param requiredFieldIds + * The partition sources and equality-delete keys of those tasks, which every task schema + * holds * @return * Some(metadata) if all reflection succeeds, None to trigger fallback */ @@ -2440,7 +2449,8 @@ object CometIcebergNativeScanMetadata extends Logging { scan: Any, metadataLocation: String, catalogProperties: Map[String, String], - tasks: java.util.List[_]): Option[CometIcebergNativeScanMetadata] = { + tasks: java.util.List[_], + requiredFieldIds: Set[Int]): Option[CometIcebergNativeScanMetadata] = { import org.apache.comet.iceberg.IcebergReflection._ for { @@ -2459,6 +2469,8 @@ object CometIcebergNativeScanMetadata extends Logging { } val globalFieldIdMapping = buildFieldIdMapping(scanSchema) + val (prunedTaskSchema, fullTaskSchema) = + nativeTaskSchemas(table, tableSchema, scanSchema, requiredFieldIds) // File format is always PARQUET, // validated in CometScanRule.validateIcebergFileScanTasks() @@ -2469,7 +2481,8 @@ object CometIcebergNativeScanMetadata extends Logging { nameMapping = nameMapping, tasks = tasks, scanSchema = scanSchema, - prunedScanSchema = withoutMetadataColumns(scanSchema), + prunedTaskSchema = prunedTaskSchema, + fullTaskSchema = fullTaskSchema, tableSchema = tableSchema, globalFieldIdMapping = globalFieldIdMapping, catalogProperties = catalogProperties, diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index acd149a966d..ed23b09179f 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -618,8 +618,12 @@ case class CometScanRule(session: SparkSession) .get(COMET_S3_COMPLIANT_SCHEMES_KEY) .map(COMET_S3_COMPLIANT_SCHEMES_KEY -> _) - val result = CometIcebergNativeScanMetadata - .extract(scanExec.scan, effectiveLocation, catalogProperties, icebergTasks) + val result = CometIcebergNativeScanMetadata.extract( + scanExec.scan, + effectiveLocation, + catalogProperties, + icebergTasks, + taskValidation.requiredFieldIds) result } catch { @@ -735,15 +739,16 @@ case class CometScanRule(session: SparkSession) val resolvedProjectedFieldIds = projectedFieldIds.flatten.toSet val hasUnresolvedProjectedFieldIds = projectedFieldIds.exists(_.isEmpty) - // The whole Iceberg table schema is serialized to native, but iceberg-rust can represent + // The full task schema can be serialized to native, but iceberg-rust can represent // Variant in that schema as long as no projected field contains one. Check projected // roots strictly, and allow Variant only under entirely unprojected roots. Other // unsupported types still fail closed everywhere. An empty data projection is also strict // because iceberg-rust currently interprets an empty field-id list as a request for every // column. + val taskSchema = metadata.fullTaskSchema.getOrElse(metadata.tableSchema) val schemaTypesSupported = try { - val fullSchema = IcebergReflection.toSparkSchema(metadata.tableSchema) + val fullSchema = IcebergReflection.toSparkSchema(taskSchema) // DataTypeSupport recursively dispatches back to this override for struct fields, // array elements, and map entries, so Variant is allowed at any nesting depth only // when its entire top-level Iceberg field is unprojected. @@ -754,7 +759,7 @@ case class CometScanRule(session: SparkSession) reasons: ListBuffer[String]): Boolean = isVariantType(dt) || super.isTypeSupported(dt, name, reasons) } - val tableFieldIds = IcebergReflection.buildFieldIdMapping(metadata.tableSchema) + val tableFieldIds = IcebergReflection.buildFieldIdMapping(taskSchema) fullSchema.fields.forall { field => val isProjected = projectedDataColumns.isEmpty || @@ -775,11 +780,11 @@ case class CometScanRule(session: SparkSession) // nested field was added (to a struct, or to a struct inside a list or map) fails the // native scan with "Incorrect number of arrays for StructArray fields", and a nested field // renamed since the file was written can read back as NULL. The native read usually asks - // only for the nested fields Spark's schema pruning keeps, but some tasks still read with - // the full table schema (see CometIcebergNativeScan), so this check walks both schemas and - // falls back even when the changed field is pruned away. A FileScanTask does not record - // which schema wrote its file, so fall back when any schema in the table's history lacks - // a nested field of a projected column or names it differently. + // only for the nested fields Spark's schema pruning keeps, but a scan can still read with + // the full schema (see IcebergReflection.nativeTaskSchemas), so this check walks both + // schemas and falls back even when the changed field is pruned away. A FileScanTask does + // not record which schema wrote its file, so fall back when any schema in the table's + // history lacks a nested field of a projected column or names it differently. val nestedFieldsSupported = try { val schemas = Seq(metadata.tableSchema, metadata.scanSchema) @@ -1039,12 +1044,21 @@ case class CometScanRule(session: SparkSession) } } + // Every task needs each partition source and equality-delete key under one name (see + // IcebergReflection.nativeTaskSchemas). + val taskSchemaNamesSupported = metadata.fullTaskSchema.isDefined + if (!taskSchemaNamesSupported) { + fallbackReasons += "Comet's native Iceberg reader cannot name a partition source or " + + "equality-delete key the same way in every task: no schema in the table's history " + + "names it apart from the columns the scan reads" + } + if (schemaSupported && fileIOCompatible && formatVersionSupported && defaultValuesSupported && schemaTypesSupported && encryptionKeyLengthSupported && taskValidation.allParquet && allSupportedFilesystems && allLocationsOpenable && metadataSchemeSupported && partitionTypesSupported && unifiedPartitionTypeSupported && transformFunctionsSupported && deleteFileTypesSupported && dppSubqueriesSupported && - nestedFieldsSupported) { + nestedFieldsSupported && taskSchemaNamesSupported) { CometBatchScanExec( scanExec.clone().asInstanceOf[BatchScanExec], runtimeFilters = scanExec.runtimeFilters, @@ -1367,6 +1381,8 @@ object CometScanRule extends Logging { IcebergReflection.loadClass(IcebergReflection.ClassNames.FILE_SCAN_TASK) val unboundPredicateClass = IcebergReflection.loadClass(IcebergReflection.ClassNames.UNBOUND_PREDICATE) + val deleteFileClass = + IcebergReflection.loadClass(IcebergReflection.ClassNames.DELETE_FILE) // Resolve all method lookups outside the loop val fileMethod = IcebergReflection.getMethod(contentScanTaskClass, "file") @@ -1374,11 +1390,18 @@ object CometScanRule extends Logging { val pathMethod = IcebergReflection.getMethod(contentFileClass, "path") val residualMethod = IcebergReflection.getMethod(contentScanTaskClass, "residual") val deletesMethod = IcebergReflection.getMethod(fileScanTaskClass, "deletes") + val specMethod = IcebergReflection.getMethod(fileScanTaskClass, "spec") val termMethod = IcebergReflection.getMethod(unboundPredicateClass, "term") var allParquet = true val unsupportedSchemes = mutable.Set[String]() var nonIdentityTransform: Option[String] = None + // Many tasks share a partition spec, and a delete file can apply to many data files, so each + // distinct instance is inspected once. + def identitySet(): java.util.Set[Any] = + java.util.Collections.newSetFromMap(new java.util.IdentityHashMap[Any, java.lang.Boolean]()) + val specs = identitySet() + val seenDeleteFiles = identitySet() val deleteFiles = new java.util.ArrayList[Any]() // Buckets the native FileIO must read (data + delete files), for the single-config check in // CometScanRule; only S3-family locations contribute (see NativeConfig.bucketForUri). @@ -1439,13 +1462,17 @@ object CometScanRule extends Logging { // Collect delete files and check their schemes if (fileScanTaskClass.isInstance(task)) { + // Unlike the checks around it, a failure here propagates: the spec decides which columns + // every task schema needs, and a schema that misses one fails the native scan. + Option(specMethod.invoke(task)).foreach(specs.add) try { val deletes = deletesMethod.invoke(task).asInstanceOf[java.util.List[_]] - deleteFiles.addAll(deletes) - deletes.asScala.foreach { deleteFile => - IcebergReflection.extractFileLocation(contentFileClass, deleteFile).foreach { - deletePath => inspectLocation(deletePath) + if (seenDeleteFiles.add(deleteFile)) { + deleteFiles.add(deleteFile) + IcebergReflection.extractFileLocation(contentFileClass, deleteFile).foreach { + deletePath => inspectLocation(deletePath) + } } } } catch { @@ -1454,13 +1481,25 @@ object CometScanRule extends Logging { } } + // The columns every native task schema holds beyond the projection: partition sources and + // equality-delete keys (see IcebergReflection.nativeTaskSchemas). + val requiredFieldIds = + specs.asScala.flatMap(IcebergReflection.partitionFieldSourceIds).toSet ++ + deleteFiles.asScala.flatMap { deleteFile => + IcebergReflection + .getEqualityFieldIds(deleteFileClass, deleteFile) + .asScala + .map(_.asInstanceOf[java.lang.Integer].intValue()) + } + IcebergTaskValidationResult( allParquet, unsupportedSchemes.toSet, nonIdentityTransform, deleteFiles, dataFileBuckets.toSet, - hostlessLocation) + hostlessLocation, + requiredFieldIds) } } @@ -1473,4 +1512,5 @@ case class IcebergTaskValidationResult( nonIdentityTransform: Option[String], deleteFiles: java.util.List[_], dataFileBuckets: Set[String], - hostlessLocation: Option[String]) + hostlessLocation: Option[String], + requiredFieldIds: Set[Int]) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index adc0b64fc5b..ebd7fa4a9f6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -980,10 +980,8 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val commonBuilder = OperatorOuterClass.IcebergScanCommon.newBuilder() // Deduplication structures - map unique values to pool indices - val schemaToPoolIndex = mutable.HashMap[AnyRef, Int]() val partitionSpecToPoolIndex = mutable.HashMap[String, Int]() val nameMappingToPoolIndex = mutable.HashMap[String, Int]() - val projectFieldIdsToPoolIndex = mutable.HashMap[Seq[Int], Int]() val partitionDataToPoolIndex = mutable.HashMap[String, Int]() // Individual delete files are interned into a flat pool; deleteFilesToPoolIndex then dedups // the per-task sets as lists of indices into it, so a delete file that applies to many data @@ -1006,20 +1004,6 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit idx }) val residualToPoolIndex = mutable.HashMap[OperatorOuterClass.IcebergPredicate, Int]() - // Field-id mappings are read out of an Iceberg schema by reflection, one lookup per column, so - // memoize them. Keyed like schemaToPoolIndex above: a task schema that Iceberg materializes - // fresh per task misses, but the table/scan schema shared by every task hits. - val fieldIdMappingCache = mutable.HashMap[AnyRef, Map[String, Int]]() - def fieldIdMapping(schema: AnyRef): Map[String, Int] = - fieldIdMappingCache.getOrElseUpdate(schema, IcebergReflection.buildFieldIdMapping(schema)) - // Whether the scan schema references field ids the current table schema no longer has (a - // dropped column read through VERSION AS OF). Loop-invariant, and lazy so a scan whose tasks - // all carry deletes never walks the table schema at all. - lazy val hasHistoricalColumns = { - val tableSchemaFieldIds = - fieldIdMapping(metadata.tableSchema.asInstanceOf[AnyRef]).values.toSet - metadata.globalFieldIdMapping.values.exists(id => !tableSchemaFieldIds.contains(id)) - } // Columns whose Iceberg type iceberg-rust cannot use for page-index pruning; residual // predicates over them are dropped (see icebergExprToProto). Computed once from the full table // schema so a filter column projected out of the scan output is still recognized. @@ -1027,36 +1011,33 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit IcebergReflection.pageIndexUnsupportedColumns(metadata.tableSchema) // iceberg-rust reads every leaf of a projected column that the task schema holds. Reading with - // the pruned scan schema skips the nested fields Spark's schema pruning removed, which the - // table schema would read and decode only to drop later. - val nestedSchemaPruningEnabled = CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.get() - // Partition sources and equality-delete keys that the task schema lacks are appended at its - // top level. A nested one that the query pruned away would land outside its struct, so a task - // that needs one reads with the full schema instead. - val canUsePrunedSchemaCache = mutable.HashMap[Seq[Int], Boolean]() - def canUsePrunedSchema(requiredFieldIds: Seq[Int]): Boolean = - nestedSchemaPruningEnabled && (requiredFieldIds.isEmpty || - canUsePrunedSchemaCache.getOrElseUpdate( - requiredFieldIds, - requiredFieldIds.forall { id => - IcebergReflection.findFieldObject(metadata.prunedScanSchema, id).isDefined || - !IcebergReflection.isNestedField(metadata.tableSchema, id) - })) - // schemaWithRequiredFields builds a new Schema whenever it appends a column. Memoize it by - // base schema (Iceberg's Schema compares by identity) and ids, so tasks that need the same - // columns share one schema object, and with it one schema pool entry. Most tasks need no extra - // columns, so they skip both lookups. - val requiredFieldsCache = mutable.HashMap[(AnyRef, Seq[Int]), AnyRef]() - def withRequiredFields(baseSchema: AnyRef, requiredFieldIds: Seq[Int]): AnyRef = - if (requiredFieldIds.isEmpty) { - baseSchema - } else { - requiredFieldsCache.getOrElseUpdate( - (baseSchema, requiredFieldIds), - IcebergReflection - .schemaWithRequiredFields(baseSchema, metadata.table, requiredFieldIds) - .asInstanceOf[AnyRef]) - } + // the pruned task schema skips the nested fields Spark's schema pruning removed, which the + // full schema would read and decode only to drop later. Every task reads with the same schema + // (see IcebergReflection.nativeTaskSchemas), so its pool entry and the projected field ids are + // the same for every task too. + val pruningEnabled = CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.get() + val taskSchema: AnyRef = metadata.prunedTaskSchema + .filter(_ => pruningEnabled) + .orElse(metadata.fullTaskSchema) + .getOrElse(throw new IllegalStateException( + "No native Iceberg task schema. CometScanRule should have fallen back to Spark.")) + .asInstanceOf[AnyRef] + val nameToFieldId = IcebergReflection.buildFieldIdMapping(taskSchema) + val projectFieldIds = output.map { attr => + nameToFieldId + .get(attr.name) + .orElse(metadata.globalFieldIdMapping.get(attr.name)) + .orElse(CometIcebergNativeScan.MetadataFieldIds.get(attr.name)) + .getOrElse { + throw new IllegalStateException( + s"Column '${attr.name}' not found in task schema, global schema, " + + "or metadata field IDs. This indicates a bug in CometScanRule " + + "validation -- all output columns should be resolvable.") + } + } + val projectFieldIdsList = OperatorOuterClass.ProjectFieldIdList.newBuilder() + projectFieldIds.foreach(id => projectFieldIdsList.addFieldIds(id)) + commonBuilder.addProjectFieldIdsPool(projectFieldIdsList.build()) val perPartitionBuilders = mutable.ArrayBuffer[OperatorOuterClass.IcebergScan]() @@ -1101,9 +1082,9 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // keyMetadata() is declared on ContentFile (present across all supported Iceberg versions). // For encrypted tables it returns the plaintext StandardKeyMetadata blob; null otherwise. val keyMetadataMethod = IcebergReflection.getMethod(contentFileClass, "keyMetadata") - val taskSchemaMethod = IcebergReflection.getMethod(fileScanTaskClass, "schema") val toJsonMethod = IcebergReflection.getMethod(schemaParserClass, "toJson", schemaClass) + commonBuilder.addSchemaPool(toJsonMethod.invoke(null, taskSchema).asInstanceOf[String]) // Access inputRDD - safe now, DPP is resolved scanExec.inputRDD match { @@ -1162,77 +1143,9 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // verbatim for iceberg-rust to decode. Unencrypted files leave the field unset. keyMetadataBytes(keyMetadataMethod, dataFile).foreach(taskBuilder.setKeyMetadata) - val deletes = - IcebergReflection.getDeleteFilesFromTask(task, fileScanTaskClass) - val hasDeletes = !deletes.isEmpty - - // An equality delete may be keyed on a column dropped from the current schema - // (schema evolution). iceberg-rust must read that column to apply the delete, - // so union the equality-delete field ids into the task schema, resolving any - // dropped ones from the table's schema history (mirrors Iceberg-Java's - // DeleteFilter.fileProjection). - val equalityFieldIds = - if (hasDeletes) { - deletes.asScala.flatMap { df => - requiredEqualityFieldIds(deleteFileClass, df).asScala.map(_.intValue()) - }.toSeq - } else { - Nil - } - - // iceberg-rust validates a FileScanTask by resolving its partition spec against - // the task schema, so a task carrying a partition spec needs that spec's source - // columns present even when the query projects them out (e.g. selecting only - // _spec_id / _partition). Union them in. project_field_ids still drives the read, - // so these columns are not materialized into the output. - val requiredFieldIds = equalityFieldIds ++ - IcebergReflection.partitionSourceFieldIds(task, fileScanTaskClass) - - val baseSchema: AnyRef = - if (canUsePrunedSchema(requiredFieldIds)) { - metadata.prunedScanSchema.asInstanceOf[AnyRef] - } else if (hasDeletes) { - taskSchemaMethod.invoke(task) - } else if (hasHistoricalColumns) { - metadata.scanSchema.asInstanceOf[AnyRef] - } else { - metadata.tableSchema.asInstanceOf[AnyRef] - } - val schema: AnyRef = withRequiredFields(baseSchema, requiredFieldIds) - - val schemaIdx = schemaToPoolIndex.getOrElseUpdate( - schema, { - val idx = schemaToPoolIndex.size - val schemaJson = toJsonMethod.invoke(null, schema).asInstanceOf[String] - commonBuilder.addSchemaPool(schemaJson) - idx - }) - taskBuilder.setSchemaIdx(schemaIdx) - - val nameToFieldId = fieldIdMapping(schema) - - val projectFieldIds = output.map { attr => - nameToFieldId - .get(attr.name) - .orElse(metadata.globalFieldIdMapping.get(attr.name)) - .orElse(CometIcebergNativeScan.MetadataFieldIds.get(attr.name)) - .getOrElse { - throw new IllegalStateException( - s"Column '${attr.name}' not found in task schema, global schema, " + - "or metadata field IDs. This indicates a bug in CometScanRule " + - "validation -- all output columns should be resolvable.") - } - } - - val projectFieldIdsIdx = projectFieldIdsToPoolIndex.getOrElseUpdate( - projectFieldIds, { - val idx = projectFieldIdsToPoolIndex.size - val listBuilder = OperatorOuterClass.ProjectFieldIdList.newBuilder() - projectFieldIds.foreach(id => listBuilder.addFieldIds(id)) - commonBuilder.addProjectFieldIdsPool(listBuilder.build()) - idx - }) - taskBuilder.setProjectFieldIdsIdx(projectFieldIdsIdx) + // One task schema and one projection for the whole scan, both at pool index 0. + taskBuilder.setSchemaIdx(0) + taskBuilder.setProjectFieldIdsIdx(0) val deleteFilesList = extractDeleteFilesList( @@ -1325,10 +1238,10 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // Log deduplication summary val allPoolSizes = Seq( - schemaToPoolIndex.size, + commonBuilder.getSchemaPoolCount, partitionSpecToPoolIndex.size, nameMappingToPoolIndex.size, - projectFieldIdsToPoolIndex.size, + commonBuilder.getProjectFieldIdsPoolCount, partitionDataToPoolIndex.size, deleteFileToPoolIndex.size, deleteFilesToPoolIndex.size, diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 8ebff034097..cf84ac928a3 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -1099,6 +1099,16 @@ class CometIcebergNativeSuite tableName: String, column: String, value: Any, + warehouseDir: File): Unit = + commitEqualityDeletes(catalogName, namespace, tableName, column, Seq(value), warehouseDir) + + /** Like [[commitEqualityDelete]], with one delete file holding every value in `values`. */ + private def commitEqualityDeletes( + catalogName: String, + namespace: String, + tableName: String, + column: String, + values: Seq[Any], warehouseDir: File): Unit = { val catalog = spark.sessionState.catalogManager .catalog(catalogName) @@ -1110,9 +1120,12 @@ class CometIcebergNativeSuite .table() val deleteRowSchema = table.schema().select(column) - val record: org.apache.iceberg.data.Record = - org.apache.iceberg.data.GenericRecord.create(deleteRowSchema) - record.setField(column, value) + val records = values.map { value => + val record: org.apache.iceberg.data.Record = + org.apache.iceberg.data.GenericRecord.create(deleteRowSchema) + record.setField(column, value) + record + } val out = table .io() .newOutputFile(new File(warehouseDir, s"eq-delete-${System.nanoTime()}.parquet").toString) @@ -1120,11 +1133,16 @@ class CometIcebergNativeSuite table, out, null, - java.util.Collections.singletonList(record), + records.asJava, deleteRowSchema) table.newRowDelta().addDeletes(deleteFile).commit() } + private def latestSnapshotId(table: String): Long = spark + .sql(s"SELECT snapshot_id FROM $table.snapshots ORDER BY committed_at DESC LIMIT 1") + .collect()(0) + .getLong(0) + // Regression for the native scan applying an equality delete keyed on a column that has since // been dropped (schema evolution). The delete's field id is absent from the current schema, so // the serde must recover it from the table's schema history to build the task schema; before the @@ -2350,9 +2368,9 @@ class CometIcebergNativeSuite // A pruned task schema still needs the columns iceberg-rust uses beyond the projection: the // partition source and the equality-delete key when the query projects neither. Tasks with - // deletes read with the pruned schema too, and the tasks of every partition share the one schema - // that appending the partition source builds. VERSION AS OF reads of columns dropped, renamed, - // or swapped after the snapshot use the snapshot's names. + // deletes read with the pruned schema too, and every task of a scan reads with the same one. + // VERSION AS OF reads of columns dropped, renamed, or swapped after the snapshot use the + // snapshot's names. test("nested schema pruning with deletes, partitions, and time travel") { assume(icebergAvailable, "Iceberg not available in classpath") @@ -2365,9 +2383,9 @@ class CometIcebergNativeSuite CometConf.COMET_EXEC_ENABLED.key -> "true", CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { - // Also checks that no task, tasks with deletes included, reads the pruned `pad` fields. - // With `oneSchema`, every task needs the same columns, so they share one task schema. - def checkPrunedNativeScan(query: String, oneSchema: Boolean = false): Unit = { + // Also checks that every task, tasks with deletes included, reads with one schema that + // leaves out the pruned `pad` fields. + def checkPrunedNativeScan(query: String): Unit = { val (_, cometPlan) = checkSparkAnswer(query) val scans = collectIcebergNativeScans(cometPlan) assert(scans.length == 1, s"expected one native scan, got ${scans.length}\n$cometPlan") @@ -2378,17 +2396,10 @@ class CometIcebergNativeSuite .getSchemaPoolList .asScala assert( - schemas.forall(!_.contains("\"pad\"")), - s"$query: pruned field in a task schema:\n${schemas.mkString("\n")}") - assert( - !oneSchema || schemas.length == 1, - s"$query: expected one task schema, got:\n${schemas.mkString("\n")}") + schemas.length == 1 && !schemas.head.contains("\"pad\""), + s"$query: expected one task schema without pad, got:\n${schemas.mkString("\n")}") } } - def latestSnapshotId(table: String): Long = spark - .sql(s"SELECT snapshot_id FROM $table.snapshots ORDER BY committed_at DESC LIMIT 1") - .collect()(0) - .getLong(0) val morProperties = """ TBLPROPERTIES ( @@ -2425,8 +2436,7 @@ class CometIcebergNativeSuite "RENAME COLUMN y TO x", "RENAME COLUMN tmp TO y").foreach(change => spark.sql(s"ALTER TABLE $mor $change")) checkPrunedNativeScan(s"SELECT id, s.a FROM $mor ORDER BY id") - // The equality-delete key `id` is not projected. Iceberg gives the delete only to the data - // file whose `id` range holds 7, so only that file's task appends `id`. + // The equality-delete key `id` is not projected, so the task schema appends it. checkPrunedNativeScan(s"SELECT s.a FROM $mor ORDER BY s.a") checkPrunedNativeScan( s"SELECT id, s.a FROM $mor VERSION AS OF $snapshotBeforeDeletes ORDER BY id") @@ -2446,21 +2456,18 @@ class CometIcebergNativeSuite """) spark.sql(s"DELETE FROM $partitioned WHERE id % 10 = 0") val partitionedSnapshot = latestSnapshotId(partitioned) - // The partition source `p` is not projected. Every task appends it, and they share the - // memoized result. - checkPrunedNativeScan(s"SELECT id, s.a FROM $partitioned ORDER BY id", oneSchema = true) + // The partition source `p` is not projected, so the task schema appends it. + checkPrunedNativeScan(s"SELECT id, s.a FROM $partitioned ORDER BY id") checkPrunedNativeScan(s"SELECT p, count(s.a) FROM $partitioned GROUP BY p ORDER BY p") - // Once `p` takes the dropped `c`'s name, the partition source's current name clashes with - // the snapshot's `c`, so it is appended under its older name. + // Once `p` takes the dropped `c`'s name, the task schema keeps the snapshot's names: `c` + // for the dropped column, and `p` for the partition source. spark.sql(s"ALTER TABLE $partitioned DROP COLUMN c") spark.sql(s"ALTER TABLE $partitioned RENAME COLUMN p TO c") checkPrunedNativeScan( - s"SELECT id, c, s.a FROM $partitioned VERSION AS OF $partitionedSnapshot ORDER BY id", - oneSchema = true) + s"SELECT id, c, s.a FROM $partitioned VERSION AS OF $partitionedSnapshot ORDER BY id") // Two partition sources renamed into each other's names: `q` takes the dropped `c`'s name - // and `p` takes `q`'s. Names picked one field at a time clash here, so both are appended - // as one older schema names them. + // and `p` takes `q`'s. The task schema names both as the snapshot does, `p` and `q`. val twoSources = "test_cat.db.nested_pruning_two_sources" spark.sql(s""" CREATE TABLE $twoSources (id INT, p INT, q INT, c INT, s STRUCT) @@ -2472,14 +2479,103 @@ class CometIcebergNativeSuite spark.sql(s"ALTER TABLE $twoSources $change") } checkPrunedNativeScan( - s"SELECT id, c, s.a FROM $twoSources VERSION AS OF $twoSourcesSnapshot ORDER BY id", - oneSchema = true) + s"SELECT id, c, s.a FROM $twoSources VERSION AS OF $twoSourcesSnapshot ORDER BY id") Seq(mor, partitioned, twoSources).foreach(t => spark.sql(s"DROP TABLE $t")) } } } + // iceberg-rust builds an equality-delete predicate once per delete file, with the names of the + // task that loads it first, and binds it by name in every task that shares the reader, so every + // task must give a key the same name. Adaptive split sizing is off, so the files of each table + // land in one task group and share a reader. + test("nested schema pruning gives each equality-delete key one name in every task") { + assume(icebergAvailable, "Iceberg not available in classpath") + // Equality deletes keyed on a since-dropped column need Iceberg 1.11's schema-history lookup. + assume( + icebergVersionAtLeast(1, 11), + "Equality delete on a dropped column requires Iceberg 1.11+ schema-history resolution") + + withTempIcebergDir { warehouseDir => + withSQLConf( + "spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.test_cat.type" -> "hadoop", + "spark.sql.catalog.test_cat.warehouse" -> warehouseDir.getAbsolutePath, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + + val properties = """ + TBLPROPERTIES ('format-version' = '2', 'read.split.adaptive-size.enabled' = 'false') + """ + + // A delete keyed on `y` covers only the second file. After `y` is dropped and `x0` is + // renamed to `x`, a delete keyed on `x` covers both, so one task needs `x` and the other + // needs `x` and `y`. + val renamed = "test_cat.db.key_names_renamed" + spark.sql(s"CREATE TABLE $renamed (id INT, x0 INT, y INT) USING iceberg $properties") + spark.sql(s"INSERT INTO $renamed VALUES (1, 10, 100), (2, 20, 100)") + spark.sql(s"INSERT INTO $renamed VALUES (3, 25, 200), (4, 30, 300)") + commitEqualityDelete("test_cat", "db", "key_names_renamed", "y", 200, warehouseDir) + spark.sql(s"ALTER TABLE $renamed DROP COLUMN y") + spark.sql(s"ALTER TABLE $renamed RENAME COLUMN x0 TO x") + commitEqualityDeletes( + "test_cat", + "db", + "key_names_renamed", + "x", + Seq(10, 25), + warehouseDir) + checkIcebergNativeScan(s"SELECT id FROM $renamed ORDER BY id") + + // Keys from schemas that never held both: `p` is dropped before `q` is added, and a delete + // keyed on each applies to the one data file. + val generations = "test_cat.db.key_names_generations" + spark.sql(s"CREATE TABLE $generations (id INT, p INT) USING iceberg $properties") + spark.sql(s"INSERT INTO $generations VALUES (1, 10), (2, 20), (3, 30)") + commitEqualityDelete("test_cat", "db", "key_names_generations", "p", 10, warehouseDir) + spark.sql(s"ALTER TABLE $generations DROP COLUMN p") + spark.sql(s"ALTER TABLE $generations ADD COLUMN q INT") + commitEqualityDelete("test_cat", "db", "key_names_generations", "q", 99, warehouseDir) + checkIcebergNativeScan(s"SELECT id FROM $generations ORDER BY id") + + // The first spec's files need the nested partition source `s.region`, which the query + // prunes away, so this VERSION AS OF read uses the full schema. A global delete keyed on + // `k` covers files of both specs, and `k` and `j` swap names after the snapshot, so the + // full schema must carry the snapshot's names, as the pruned one does. + val swapped = "test_cat.db.key_names_swapped" + spark.sql(s""" + CREATE TABLE $swapped (id INT, k INT, j INT, s STRUCT) + USING iceberg PARTITIONED BY (s.region) $properties + """) + spark.sql(s""" + INSERT INTO $swapped VALUES + (1, 10, 100, named_struct('region', 'east', 'a', 1)), + (2, 20, 200, named_struct('region', 'west', 'a', 2)) + """) + loadIcebergTable(spark, "test_cat", "db", "key_names_swapped") + .asInstanceOf[org.apache.iceberg.Table] + .updateSpec() + .removeField("s.region") + .commit() + spark.sql(s""" + INSERT INTO $swapped VALUES + (3, 10, 300, named_struct('region', 'east', 'a', 3)), + (4, 30, 400, named_struct('region', 'west', 'a', 4)) + """) + commitEqualityDelete("test_cat", "db", "key_names_swapped", "k", 10, warehouseDir) + val snapshot = latestSnapshotId(swapped) + Seq("RENAME COLUMN k TO tmp", "RENAME COLUMN j TO k", "RENAME COLUMN tmp TO j").foreach { + change => spark.sql(s"ALTER TABLE $swapped $change") + } + checkIcebergNativeScan(s"SELECT id, k FROM $swapped VERSION AS OF $snapshot ORDER BY id") + + Seq(renamed, generations, swapped).foreach(t => spark.sql(s"DROP TABLE $t")) + } + } + } + // Data files written without Iceberg field ids and read with no name mapping: iceberg-rust then // projects whole top-level columns by position, and a by-name struct cast narrows each one to // the pruned type. Spark's own reader returns NULL nested fields for such a table, so the pruned From 8566834663235f9cb47d2f5af3578ed46f57028b Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 13 Oct 2026 08:50:43 -0700 Subject: [PATCH 5/9] Check the full task schema for nested evolution The full task schema of a VERSION AS OF read can be older than the table's schema and hold a nested field that an earlier data file lacks, which the native scan cannot read. Walk it in the nested-evolution check too, so such a scan falls back to Spark. --- .../apache/comet/rules/CometScanRule.scala | 10 +++++--- .../comet/CometIcebergNativeSuite.scala | 25 +++++++++++++++++-- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index ed23b09179f..d88f240b873 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -781,13 +781,15 @@ case class CometScanRule(session: SparkSession) // native scan with "Incorrect number of arrays for StructArray fields", and a nested field // renamed since the file was written can read back as NULL. The native read usually asks // only for the nested fields Spark's schema pruning keeps, but a scan can still read with - // the full schema (see IcebergReflection.nativeTaskSchemas), so this check walks both - // schemas and falls back even when the changed field is pruned away. A FileScanTask does - // not record which schema wrote its file, so fall back when any schema in the table's + // its full task schema (see IcebergReflection.nativeTaskSchemas), which for VERSION AS OF + // can be an older schema than the table's. So this check walks the table, scan, and full + // task schemas, and falls back even when the changed field is pruned away. A FileScanTask + // does not record which schema wrote its file, so fall back when any schema in the table's // history lacks a nested field of a projected column or names it differently. val nestedFieldsSupported = try { - val schemas = Seq(metadata.tableSchema, metadata.scanSchema) + val schemas = + Seq(metadata.tableSchema, metadata.scanSchema) ++ metadata.fullTaskSchema val fieldIds = if (hasUnresolvedProjectedFieldIds) { schemas.flatMap(IcebergReflection.buildFieldIdMapping(_).values).toSet diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index cf84ac928a3..0c044b15cb4 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -1827,7 +1827,8 @@ class CometIcebergNativeSuite checkNestedFieldEvolutionFallback( s"SELECT l.id, r.id FROM $table l JOIN $table r ON l.s = r.s ORDER BY l.id", "s.b (added)") - // The old snapshot's files lack s.b too, and the native read uses the current schema. + // The old snapshot's files lack s.b too, and the scan's full schema is the current one, + // which has it. checkNestedFieldEvolutionFallback( s"SELECT id, s FROM $table VERSION AS OF $snapshotBeforeAdd ORDER BY id", "s.b (added)") @@ -1835,7 +1836,27 @@ class CometIcebergNativeSuite // A read that does not project the evolved column stays native. checkIcebergNativeScan(s"SELECT id FROM $table ORDER BY id") - spark.sql(s"DROP TABLE $table") + // A VERSION AS OF read whose full schema is older than the table's: the pruned-away + // nested partition source `s.k` makes the scan read its full schema, and the dropped `c` + // makes that the snapshot's, which still has the s.b that the first file predates. + val historical = "test_cat.db.evo_struct_historical" + spark.sql(s""" + CREATE TABLE $historical (id INT, c INT, s STRUCT) + USING iceberg PARTITIONED BY (s.k) + """) + spark.sql(s"INSERT INTO $historical VALUES (1, 100, named_struct('k', 'east', 'a', 10))") + spark.sql(s"ALTER TABLE $historical ADD COLUMN s.b INT") + spark.sql(s""" + INSERT INTO $historical VALUES (2, 200, named_struct('k', 'east', 'a', 20, 'b', 99)) + """) + val snapshotWithB = latestSnapshotId(historical) + spark.sql(s"ALTER TABLE $historical DROP COLUMN c") + spark.sql(s"ALTER TABLE $historical DROP COLUMN s.b") + checkNestedFieldEvolutionFallback( + s"SELECT id, c, s.a FROM $historical VERSION AS OF $snapshotWithB ORDER BY id", + "s.b (added)") + + Seq(table, historical).foreach(t => spark.sql(s"DROP TABLE $t")) } } } From 5395e50ba30a857b4101449a4ae4136cf5fa02b9 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 13 Oct 2026 13:49:33 -0700 Subject: [PATCH 6/9] Check the task schema the scan reads for nested evolution The nested-evolution check now walks the schema the scan reads, the pruned task schema while pruning applies and the full one otherwise, through one method that serialization uses too. A pruned VERSION AS OF read no longer falls back because of a nested field only the full schema holds. The Iceberg user guide lists the new naming fallback. The SQL fixture adds empty lists and maps, and the pruned-read tests also check the answers with pruning off. --- docs/source/user-guide/latest/iceberg.md | 3 ++ .../comet/iceberg/IcebergReflection.scala | 14 ++++++- .../apache/comet/rules/CometScanRule.scala | 16 ++++---- .../operator/CometIcebergNativeScan.scala | 10 ++--- .../iceberg/nested_schema_pruning.sql | 11 ++--- .../comet/CometIcebergNativeSuite.scala | 41 ++++++++++++++++--- 6 files changed, 69 insertions(+), 26 deletions(-) diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 4c3abe00216..551aea71c59 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -196,6 +196,9 @@ The following scenarios will fall back to the JVM Iceberg reader: or renamed. The native reader cannot yet match such a field to data files written before the change. The check uses the table's schema history, so the fallback stays after those files are rewritten +- Scans where no schema in the table's history gives a partition source or equality-delete key a + name that the scan's other columns don't use (every task of a native scan must give such a + column the same name) Writes are not covered by this list. By default Iceberg writes use Spark's own writer; see [Iceberg Writes](iceberg-writes.md) for the experimental native writer and when it applies. diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index 24af10ccf15..80e892f7b01 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -28,6 +28,7 @@ import org.apache.hadoop.conf.Configuration import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession +import org.apache.comet.CometConf import org.apache.comet.util.ClassLoaders /** @@ -2419,7 +2420,18 @@ case class CometIcebergNativeScanMetadata( globalFieldIdMapping: Map[String, Int], catalogProperties: Map[String, String], catalogName: Option[String], - fileFormat: String) + fileFormat: String) { + + /** + * The schema every task of the scan reads with: the pruned task schema while nested schema + * pruning is on and can apply, the full one otherwise. CometScanRule checks this schema and + * CometIcebergNativeScan serializes it, so the two cannot disagree. None falls back to Spark. + */ + def taskSchema: Option[Any] = + prunedTaskSchema + .filter(_ => CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.get()) + .orElse(fullTaskSchema) +} object CometIcebergNativeScanMetadata extends Logging { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index d88f240b873..460b505a1de 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -779,17 +779,15 @@ case class CometScanRule(session: SparkSession) // fields to the table's by field id (apache/iceberg-rust#2617). A file written before a // nested field was added (to a struct, or to a struct inside a list or map) fails the // native scan with "Incorrect number of arrays for StructArray fields", and a nested field - // renamed since the file was written can read back as NULL. The native read usually asks - // only for the nested fields Spark's schema pruning keeps, but a scan can still read with - // its full task schema (see IcebergReflection.nativeTaskSchemas), which for VERSION AS OF - // can be an older schema than the table's. So this check walks the table, scan, and full - // task schemas, and falls back even when the changed field is pruned away. A FileScanTask - // does not record which schema wrote its file, so fall back when any schema in the table's - // history lacks a nested field of a projected column or names it differently. + // renamed since the file was written can read back as NULL. This check walks the schema + // the scan reads (see CometIcebergNativeScanMetadata.taskSchema), which for VERSION AS OF + // can be older than the table's, and also the table and scan schemas, so it can fall back + // when the changed field is pruned away. A FileScanTask does not record which schema wrote + // its file, so fall back when any schema in the table's history lacks a nested field of a + // projected column or names it differently. val nestedFieldsSupported = try { - val schemas = - Seq(metadata.tableSchema, metadata.scanSchema) ++ metadata.fullTaskSchema + val schemas = Seq(metadata.tableSchema, metadata.scanSchema) ++ metadata.taskSchema val fieldIds = if (hasUnresolvedProjectedFieldIds) { schemas.flatMap(IcebergReflection.buildFieldIdMapping(_).values).toSet diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index ebd7fa4a9f6..6fafa65c756 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -1015,12 +1015,10 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // full schema would read and decode only to drop later. Every task reads with the same schema // (see IcebergReflection.nativeTaskSchemas), so its pool entry and the projected field ids are // the same for every task too. - val pruningEnabled = CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.get() - val taskSchema: AnyRef = metadata.prunedTaskSchema - .filter(_ => pruningEnabled) - .orElse(metadata.fullTaskSchema) - .getOrElse(throw new IllegalStateException( - "No native Iceberg task schema. CometScanRule should have fallen back to Spark.")) + val taskSchema: AnyRef = metadata.taskSchema + .getOrElse( + throw new IllegalStateException( + "No native Iceberg task schema. CometScanRule should have fallen back to Spark.")) .asInstanceOf[AnyRef] val nameToFieldId = IcebergReflection.buildFieldIdMapping(taskSchema) val projectFieldIds = output.map { attr => diff --git a/spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql b/spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql index b9a9589da29..0488b6b26a4 100644 --- a/spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql +++ b/spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql @@ -40,7 +40,8 @@ CREATE TABLE test_cat.db.nested_pruning ( ) USING iceberg -- NULL structs, lists, maps, list elements, map values, and fields, so that the pruned read has --- to rebuild the validity of every level from the leaves it keeps. +-- to rebuild the validity of every level from the leaves it keeps. Empty lists and maps are valid +-- with no child rows, which is another case. statement INSERT INTO test_cat.db.nested_pruning SELECT @@ -49,11 +50,11 @@ SELECT 'a', IF(id % 5 = 0, NULL, CAST(id AS INT)), 'pad', repeat('p', 20), 'inner', IF(id % 3 = 0, NULL, named_struct('b', CAST(id * 2 AS INT), 'pad', repeat('q', 20))))), - IF(id % 7 = 0, NULL, array( + CASE WHEN id % 7 = 0 THEN NULL WHEN id % 11 = 0 THEN array() ELSE array( named_struct('x', CAST(id AS INT), 'pad', repeat('r', 20)), - IF(id % 5 = 0, NULL, named_struct('x', CAST(-id AS INT), 'pad', repeat('s', 20))))), - IF(id % 7 = 0, NULL, map('k', IF(id % 3 = 0, NULL, - named_struct('v', CAST(id AS INT), 'pad', repeat('t', 20))))) + IF(id % 5 = 0, NULL, named_struct('x', CAST(-id AS INT), 'pad', repeat('s', 20)))) END, + CASE WHEN id % 7 = 0 THEN NULL WHEN id % 11 = 0 THEN map() ELSE map('k', IF(id % 3 = 0, NULL, + named_struct('v', CAST(id AS INT), 'pad', repeat('t', 20)))) END FROM range(100) query diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 0c044b15cb4..45ee015c3d6 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -85,6 +85,14 @@ class CometIcebergNativeSuite assertSingleNativeScan(cometPlan) } + /** + * [[checkIcebergNativeScan]] with nested schema pruning off, so the scan reads its full schema. + */ + private def checkFullSchemaNativeScan(query: String): Unit = + withSQLConf(CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.key -> "false") { + checkIcebergNativeScan(query) + } + /** * Asserts a query stays native even when its DPP filter falls back (SMJ / non-Comet BHJ / reuse * disabled), where only the runtime filter, not the scan, reverts to Spark. @@ -1856,7 +1864,24 @@ class CometIcebergNativeSuite s"SELECT id, c, s.a FROM $historical VERSION AS OF $snapshotWithB ORDER BY id", "s.b (added)") - Seq(table, historical).foreach(t => spark.sql(s"DROP TABLE $t")) + // Unpartitioned, the same read takes the pruned schema, which leaves s.b out, so it stays + // native. With pruning off it reads the full schema and falls back. + val unpartitioned = "test_cat.db.evo_struct_historical_unpartitioned" + spark.sql(s"CREATE TABLE $unpartitioned (id INT, c INT, s STRUCT) USING iceberg") + spark.sql(s"INSERT INTO $unpartitioned VALUES (1, 100, named_struct('a', 10))") + spark.sql(s"ALTER TABLE $unpartitioned ADD COLUMN s.b INT") + spark.sql(s"INSERT INTO $unpartitioned VALUES (2, 200, named_struct('a', 20, 'b', 99))") + val unpartitionedSnapshot = latestSnapshotId(unpartitioned) + spark.sql(s"ALTER TABLE $unpartitioned DROP COLUMN c") + spark.sql(s"ALTER TABLE $unpartitioned DROP COLUMN s.b") + val unpartitionedQuery = + s"SELECT id, c, s.a FROM $unpartitioned VERSION AS OF $unpartitionedSnapshot ORDER BY id" + checkIcebergNativeScan(unpartitionedQuery) + withSQLConf(CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.key -> "false") { + checkNestedFieldEvolutionFallback(unpartitionedQuery, "s.b (added)") + } + + Seq(table, historical, unpartitioned).foreach(t => spark.sql(s"DROP TABLE $t")) } } } @@ -2405,7 +2430,8 @@ class CometIcebergNativeSuite CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { // Also checks that every task, tasks with deletes included, reads with one schema that - // leaves out the pruned `pad` fields. + // leaves out the pruned `pad` fields, and that the full schema read with pruning off gives + // the same answer. def checkPrunedNativeScan(query: String): Unit = { val (_, cometPlan) = checkSparkAnswer(query) val scans = collectIcebergNativeScans(cometPlan) @@ -2420,6 +2446,7 @@ class CometIcebergNativeSuite schemas.length == 1 && !schemas.head.contains("\"pad\""), s"$query: expected one task schema without pad, got:\n${schemas.mkString("\n")}") } + checkFullSchemaNativeScan(query) } val morProperties = """ @@ -2530,6 +2557,10 @@ class CometIcebergNativeSuite val properties = """ TBLPROPERTIES ('format-version' = '2', 'read.split.adaptive-size.enabled' = 'false') """ + def checkPrunedAndFullReads(query: String): Unit = { + checkIcebergNativeScan(query) + checkFullSchemaNativeScan(query) + } // A delete keyed on `y` covers only the second file. After `y` is dropped and `x0` is // renamed to `x`, a delete keyed on `x` covers both, so one task needs `x` and the other @@ -2548,7 +2579,7 @@ class CometIcebergNativeSuite "x", Seq(10, 25), warehouseDir) - checkIcebergNativeScan(s"SELECT id FROM $renamed ORDER BY id") + checkPrunedAndFullReads(s"SELECT id FROM $renamed ORDER BY id") // Keys from schemas that never held both: `p` is dropped before `q` is added, and a delete // keyed on each applies to the one data file. @@ -2559,7 +2590,7 @@ class CometIcebergNativeSuite spark.sql(s"ALTER TABLE $generations DROP COLUMN p") spark.sql(s"ALTER TABLE $generations ADD COLUMN q INT") commitEqualityDelete("test_cat", "db", "key_names_generations", "q", 99, warehouseDir) - checkIcebergNativeScan(s"SELECT id FROM $generations ORDER BY id") + checkPrunedAndFullReads(s"SELECT id FROM $generations ORDER BY id") // The first spec's files need the nested partition source `s.region`, which the query // prunes away, so this VERSION AS OF read uses the full schema. A global delete keyed on @@ -2590,7 +2621,7 @@ class CometIcebergNativeSuite Seq("RENAME COLUMN k TO tmp", "RENAME COLUMN j TO k", "RENAME COLUMN tmp TO j").foreach { change => spark.sql(s"ALTER TABLE $swapped $change") } - checkIcebergNativeScan(s"SELECT id, k FROM $swapped VERSION AS OF $snapshot ORDER BY id") + checkPrunedAndFullReads(s"SELECT id, k FROM $swapped VERSION AS OF $snapshot ORDER BY id") Seq(renamed, generations, swapped).foreach(t => spark.sql(s"DROP TABLE $t")) } From 86eaa60d37c2b527f04b95f3208add0ecda50203 Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 14 Oct 2026 11:02:21 -0700 Subject: [PATCH 7/9] Match field types when picking the full task schema The full task schema was the newest schema whose top-level columns match the scan schema by id and name. For a VERSION AS OF read across a type promotion, that could be a newer schema with the promoted type, which the batch adaptation after iceberg-rust cannot narrow back, so the native scan failed. Pick the newest schema in which every field the scan reads, at any depth, has the same name and type instead. The time travel test reads a promoted int, decimal, and nested field, on partitioned and unpartitioned tables, with pruning on and off. --- .../comet/iceberg/IcebergReflection.scala | 44 +++++++++++++++---- .../comet/CometIcebergNativeSuite.scala | 34 ++++++++++++++ 2 files changed, 69 insertions(+), 9 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index 80e892f7b01..651a7fcfc38 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -728,10 +728,11 @@ object IcebergReflection extends Logging { * * Every task reads with the same schema because iceberg-rust builds an equality-delete * predicate once per delete file, with the column names of the task that loads it first, and - * binds it by name in every other task of the reader. The names come from the full schema, the - * newest schema, current first, that names every column of the scan schema the same way. That - * is the current schema unless a VERSION AS OF snapshot named a column differently. A key that - * schema lacks, dropped before it, takes the newest name that no other column has. + * binds it by name in every other task of the reader. The full schema is the newest schema, + * current first, that has every field the scan reads, at any depth, with the same name and + * type, as the batch adaptation after iceberg-rust cannot narrow a promoted type back. That is + * the current schema unless a VERSION AS OF snapshot named or typed a field differently. A key + * that schema lacks, dropped before it, takes the newest name that no other column has. */ def nativeTaskSchemas( table: Any, @@ -739,14 +740,19 @@ object IcebergReflection extends Logging { scanSchema: Any, requiredFieldIds: Set[Int]): (Option[Any], Option[Any]) = { val prunedScanSchema = withoutMetadataColumns(scanSchema) - val scanColumns = columnNamesById(prunedScanSchema).toSet + val scanFields = fieldsAtAnyDepth(prunedScanSchema) val schemas = tableSchema +: getAllSchemas(table).reverse - val full = schemas.iterator - .map(schema => schema -> columnNamesById(schema)) - .find { case (_, columns) => scanColumns.subsetOf(columns.toSet) } + val full = schemas.find { schema => + scanFields.forall { case (id, name, leafType) => + findFieldObject(schema, id).exists { field => + fieldName(field) == name && leafType.forall(_ == fieldTypeOf(field)) + } + } + } full match { case None => (None, None) - case Some((fullSchema, fullColumns)) => + case Some(fullSchema) => + val fullColumns = columnNamesById(fullSchema) val fullNames = fullColumns.values.toSet val ids = requiredFieldIds.toSeq.sorted val inFull = ids.flatMap(id => findFieldObject(fullSchema, id).map(id -> _)).toMap @@ -789,6 +795,24 @@ object IcebergReflection extends Logging { .toMap } + /** + * The fields of `schema` at any depth, as id, name, and the type of a field with no children. + */ + private def fieldsAtAnyDepth(schema: Any): Seq[(Int, String, Option[Any])] = { + import scala.jdk.CollectionConverters._ + def walk(field: Any): Seq[(Int, String, Option[Any])] = { + val children = childFields(field) + val leafType = if (children.isEmpty) Some(fieldTypeOf(field)) else None + (fieldIdOf(field), fieldName(field), leafType) +: children.flatMap(walk) + } + getMethod(schema.getClass, "columns") + .invoke(schema) + .asInstanceOf[java.util.List[_]] + .asScala + .toSeq + .flatMap(walk) + } + /** `schema` with `fields` appended at its top level, or `schema` itself when there are none. */ private def withColumns(schema: Any, fields: Seq[Any]): Any = if (fields.isEmpty) { @@ -1384,6 +1408,8 @@ object IcebergReflection extends Logging { private def fieldIdOf(field: Any): Int = getMethod(field.getClass, "fieldId").invoke(field).asInstanceOf[Int] + private def fieldTypeOf(field: Any): Any = getMethod(field.getClass, "type").invoke(field) + /** * Converts an Iceberg `Schema` to the Spark `StructType` it reads as, via * `SparkSchemaUtil.convert`. Comet serializes the whole table/scan schema to native (not just diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 45ee015c3d6..57cb658c162 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -2529,6 +2529,40 @@ class CometIcebergNativeSuite checkPrunedNativeScan( s"SELECT id, c, s.a FROM $twoSources VERSION AS OF $twoSourcesSnapshot ORDER BY id") + // VERSION AS OF reads across type promotions. The batch adaptation after iceberg-rust + // cannot narrow a promoted type back, so the full schema must have the snapshot's types, + // not just be the newest schema that still has the dropped `c`. Each query reads one + // promoted column, so another promotion cannot steer the choice. The pruned schema cannot + // hold the nested partition source `s.k`, so the partitioned table reads the full schema + // with pruning on too. + Seq("promoted" -> "", "promoted_partitioned" -> "PARTITIONED BY (s.k)").foreach { + case (name, partitioning) => + val promoted = s"test_cat.db.nested_pruning_$name" + spark.sql(s""" + CREATE TABLE $promoted ( + id INT, x INT, d DECIMAL(9, 2), c INT, s STRUCT) + USING iceberg $partitioning + """) + spark.sql(s""" + INSERT INTO $promoted VALUES + (1, 10, 1.25, 100, named_struct('k', 'east', 'a', 1)), + (2, NULL, NULL, NULL, named_struct('k', 'west', 'a', CAST(NULL AS INT))) + """) + val snapshot = latestSnapshotId(promoted) + Seq( + "ALTER COLUMN x TYPE BIGINT", + "ALTER COLUMN d TYPE DECIMAL(18, 2)", + "ALTER COLUMN s.a TYPE BIGINT", + "DROP COLUMN c").foreach(change => spark.sql(s"ALTER TABLE $promoted $change")) + Seq("x", "d", "s.a").foreach { column => + val query = + s"SELECT id, $column, c FROM $promoted VERSION AS OF $snapshot ORDER BY id" + checkIcebergNativeScan(query) + checkFullSchemaNativeScan(query) + } + spark.sql(s"DROP TABLE $promoted") + } + Seq(mor, partitioned, twoSources).foreach(t => spark.sql(s"DROP TABLE $t")) } } From 170e5eb48d2b0a1fc22934eae9498a4acf4e1f66 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 15 Oct 2026 08:41:46 -0700 Subject: [PATCH 8/9] Fix the native Iceberg task schema at planning `taskSchema` read the nested schema pruning config every time it was called: during planning, where CometScanRule checks the schema, and again when the scan ran and serialized it. A query planned with pruning on and run with pruning off, such as a cached DataFrame, read the full schema, which the nested-evolution check never walked. Make it a val, built with the metadata during planning. The struct fallback test plans a VERSION AS OF read with pruning on and runs it with pruning off. --- .../org/apache/comet/iceberg/IcebergReflection.scala | 8 +++++--- .../scala/org/apache/comet/CometIcebergNativeSuite.scala | 7 +++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index 651a7fcfc38..d9f24b1874c 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -2450,10 +2450,12 @@ case class CometIcebergNativeScanMetadata( /** * The schema every task of the scan reads with: the pruned task schema while nested schema - * pruning is on and can apply, the full one otherwise. CometScanRule checks this schema and - * CometIcebergNativeScan serializes it, so the two cannot disagree. None falls back to Spark. + * pruning is on and can apply, the full one otherwise. It is fixed when `extract` builds the + * metadata during planning, so the schema CometScanRule checks is the one + * CometIcebergNativeScan serializes at execution, even if the config changes in between. None + * falls back to Spark. */ - def taskSchema: Option[Any] = + val taskSchema: Option[Any] = prunedTaskSchema .filter(_ => CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.get()) .orElse(fullTaskSchema) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 57cb658c162..1d2b3e0f1a0 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -1880,6 +1880,13 @@ class CometIcebergNativeSuite withSQLConf(CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.key -> "false") { checkNestedFieldEvolutionFallback(unpartitionedQuery, "s.b (added)") } + // The task schema is chosen during planning, so a query planned with pruning on still + // reads the pruned schema when it runs with pruning off. + val planned = spark.sql(unpartitionedQuery) + assertSingleNativeScan(planned.queryExecution.executedPlan) + withSQLConf(CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.key -> "false") { + checkAnswer(planned, Seq(Row(1, 100, 10), Row(2, 200, 20))) + } Seq(table, historical, unpartitioned).foreach(t => spark.sql(s"DROP TABLE $t")) } From 7d8138a0551c26d00366578aafa65387731bbcae Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 15 Oct 2026 12:18:50 -0700 Subject: [PATCH 9/9] Simplify the native Iceberg task schema code and tests Build the scan schema's data columns and the projected field ids in one pass each, and name the type check's schema `fullTaskSchema` so it reads apart from `taskSchema`, the schema the scan reads. The promotion reads reuse `checkPrunedAndFullReads`, and a time travel query that could not catch a wrong schema choice is gone. --- .../comet/iceberg/IcebergReflection.scala | 7 +++-- .../apache/comet/rules/CometScanRule.scala | 6 ++-- .../operator/CometIcebergNativeScan.scala | 28 +++++++++---------- .../comet/CometIcebergNativeSuite.scala | 19 ++++++------- 4 files changed, 30 insertions(+), 30 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index d9f24b1874c..e832c55e169 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -841,13 +841,16 @@ object IcebergReflection extends Logging { import scala.jdk.CollectionConverters._ val columns = getMethod(schema.getClass, "columns").invoke(schema).asInstanceOf[java.util.List[_]] - val dataColumns = columns.asScala.filter(fieldIdOf(_) <= MaxDataFieldId) + val dataColumns = new java.util.ArrayList[Any](columns.size) + columns.asScala.iterator + .filter(fieldIdOf(_) <= MaxDataFieldId) + .foreach(column => dataColumns.add(column)) if (dataColumns.size == columns.size) { schema } else { schema.getClass .getConstructor(classOf[java.util.List[_]]) - .newInstance(new java.util.ArrayList[Any](dataColumns.asJava)) + .newInstance(dataColumns) .asInstanceOf[AnyRef] } } diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index 460b505a1de..0f2cced757c 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -745,10 +745,10 @@ case class CometScanRule(session: SparkSession) // unsupported types still fail closed everywhere. An empty data projection is also strict // because iceberg-rust currently interprets an empty field-id list as a request for every // column. - val taskSchema = metadata.fullTaskSchema.getOrElse(metadata.tableSchema) + val fullTaskSchema = metadata.fullTaskSchema.getOrElse(metadata.tableSchema) val schemaTypesSupported = try { - val fullSchema = IcebergReflection.toSparkSchema(taskSchema) + val fullSchema = IcebergReflection.toSparkSchema(fullTaskSchema) // DataTypeSupport recursively dispatches back to this override for struct fields, // array elements, and map entries, so Variant is allowed at any nesting depth only // when its entire top-level Iceberg field is unprojected. @@ -759,7 +759,7 @@ case class CometScanRule(session: SparkSession) reasons: ListBuffer[String]): Boolean = isVariantType(dt) || super.isTypeSupported(dt, name, reasons) } - val tableFieldIds = IcebergReflection.buildFieldIdMapping(taskSchema) + val tableFieldIds = IcebergReflection.buildFieldIdMapping(fullTaskSchema) fullSchema.fields.forall { field => val isProjected = projectedDataColumns.isEmpty || diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index 6fafa65c756..d1d5db3787a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -1021,21 +1021,21 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit "No native Iceberg task schema. CometScanRule should have fallen back to Spark.")) .asInstanceOf[AnyRef] val nameToFieldId = IcebergReflection.buildFieldIdMapping(taskSchema) - val projectFieldIds = output.map { attr => - nameToFieldId - .get(attr.name) - .orElse(metadata.globalFieldIdMapping.get(attr.name)) - .orElse(CometIcebergNativeScan.MetadataFieldIds.get(attr.name)) - .getOrElse { - throw new IllegalStateException( - s"Column '${attr.name}' not found in task schema, global schema, " + - "or metadata field IDs. This indicates a bug in CometScanRule " + - "validation -- all output columns should be resolvable.") - } + val projectFieldIds = OperatorOuterClass.ProjectFieldIdList.newBuilder() + output.foreach { attr => + projectFieldIds.addFieldIds( + nameToFieldId + .get(attr.name) + .orElse(metadata.globalFieldIdMapping.get(attr.name)) + .orElse(CometIcebergNativeScan.MetadataFieldIds.get(attr.name)) + .getOrElse { + throw new IllegalStateException( + s"Column '${attr.name}' not found in task schema, global schema, " + + "or metadata field IDs. This indicates a bug in CometScanRule " + + "validation -- all output columns should be resolvable.") + }) } - val projectFieldIdsList = OperatorOuterClass.ProjectFieldIdList.newBuilder() - projectFieldIds.foreach(id => projectFieldIdsList.addFieldIds(id)) - commonBuilder.addProjectFieldIdsPool(projectFieldIdsList.build()) + commonBuilder.addProjectFieldIdsPool(projectFieldIds.build()) val perPartitionBuilders = mutable.ArrayBuffer[OperatorOuterClass.IcebergScan]() diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 1d2b3e0f1a0..0e10bfa5d8f 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -93,6 +93,12 @@ class CometIcebergNativeSuite checkIcebergNativeScan(query) } + /** [[checkIcebergNativeScan]] with nested schema pruning on, as by default, and then off. */ + private def checkPrunedAndFullReads(query: String): Unit = { + checkIcebergNativeScan(query) + checkFullSchemaNativeScan(query) + } + /** * Asserts a query stays native even when its DPP filter falls back (SMJ / non-Comet BHJ / reuse * disabled), where only the runtime filter, not the scan, reverts to Spark. @@ -2479,7 +2485,6 @@ class CometIcebergNativeSuite INSERT INTO $mor SELECT id, p, concat('r', CAST(id AS STRING)), id, -id, s FROM ($rows) """) - val snapshotBeforeDeletes = latestSnapshotId(mor) spark.sql(s"DELETE FROM $mor WHERE id % 10 = 0") val snapshotWithDeletes = latestSnapshotId(mor) commitEqualityDelete("test_cat", "db", "nested_pruning_mor", "id", 7, warehouseDir) @@ -2493,8 +2498,6 @@ class CometIcebergNativeSuite checkPrunedNativeScan(s"SELECT id, s.a FROM $mor ORDER BY id") // The equality-delete key `id` is not projected, so the task schema appends it. checkPrunedNativeScan(s"SELECT s.a FROM $mor ORDER BY s.a") - checkPrunedNativeScan( - s"SELECT id, s.a FROM $mor VERSION AS OF $snapshotBeforeDeletes ORDER BY id") // The current table schema lacks `c` and names `r`, `x`, and `y` differently, so these // tasks with deletes must read with the scan schema, which keeps the snapshot's names. checkPrunedNativeScan( @@ -2562,10 +2565,8 @@ class CometIcebergNativeSuite "ALTER COLUMN s.a TYPE BIGINT", "DROP COLUMN c").foreach(change => spark.sql(s"ALTER TABLE $promoted $change")) Seq("x", "d", "s.a").foreach { column => - val query = - s"SELECT id, $column, c FROM $promoted VERSION AS OF $snapshot ORDER BY id" - checkIcebergNativeScan(query) - checkFullSchemaNativeScan(query) + checkPrunedAndFullReads( + s"SELECT id, $column, c FROM $promoted VERSION AS OF $snapshot ORDER BY id") } spark.sql(s"DROP TABLE $promoted") } @@ -2598,10 +2599,6 @@ class CometIcebergNativeSuite val properties = """ TBLPROPERTIES ('format-version' = '2', 'read.split.adaptive-size.enabled' = 'false') """ - def checkPrunedAndFullReads(query: String): Unit = { - checkIcebergNativeScan(query) - checkFullSchemaNativeScan(query) - } // A delete keyed on `y` covers only the second file. After `y` is dropped and `x0` is // renamed to `x`, a delete keyed on `x` covers both, so one task needs `x` and the other