Skip to content

feat: [iceberg] support iceberg metadata columns _file, _partition #4752

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
parthchandra wants to merge 11 commits into
base: main
Choose a base branch
from
Draft

feat: [iceberg] support iceberg metadata columns _file, _partition #4752

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
9d5ec72
feat: [iceberg] support iceberg metadata columns _file, _partition
parthchandra Jun 15, 2026
d5265be
update
parthchandra Jul 26, 2026
b5f2f92
fix
parthchandra Jul 27, 2026
9832468
fix
parthchandra Jul 27, 2026
5dab501
fix fix
parthchandra Jul 27, 2026
f68c0fb
point to pr branch
parthchandra Jul 27, 2026
aea9cf7
Fix 4.1 ci failure
parthchandra Jul 28, 2026
d21dc1e
skip 4.2 ci
parthchandra Jul 28, 2026
d5d4e78
Add cross spec compatibility check
parthchandra Jul 29, 2026
3bef00e
update to latest
parthchandra Jul 29, 2026
a6a8d3b
style
parthchandra Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 48 additions & 34 deletions dev/diffs/iceberg/1.11.0.diff
View file
Open in desktop

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion native/Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions native/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,9 @@ iceberg = { git = "https://github.com/apache/iceberg-rust", rev = "4532da8" }
iceberg-storage-opendal = { git = "https://github.com/apache/iceberg-rust", rev = "4532da8", features = ["opendal-memory", "opendal-fs", "opendal-s3", "opendal-gcs", "opendal-oss", "opendal-azdls"] }
reqsign-core = "3"

[patch."https://github.com/apache/iceberg-rust"]
iceberg = { git = "https://github.com/parthchandra/iceberg-rust", branch = "metadata-columns" }

[profile.release]
debug = true
overflow-checks = false
Expand Down
1 change: 1 addition & 0 deletions native/core/src/execution/operators/iceberg_scan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -660,6 +660,7 @@ mod tests {
partition: None,
partition_spec: None,
name_mapping: None,
unified_partition_type: None,
case_sensitive: false,
key_metadata: None,
}
Expand Down
61 changes: 61 additions & 0 deletions native/core/src/execution/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3813,6 +3813,41 @@ fn parse_file_scan_tasks_from_common(
})
.collect::<Result<Vec<_>, _>>()?;

// Compute unified partition type from the partition_type_pool.
// Each entry is a StructType JSON with the resolved field types for one partition spec.
// Merge all specs into a single unified type (dedup by field_id).
//
// NOTE: The type information here comes from Iceberg Java's PartitionSpec.partitionType()
// (via Scala reflection). This is the same type computation as iceberg-rust's
// Transform::result_type() -- both implement the Iceberg spec's transform result rules.
// When iceberg-rust's own scan planning is used (not Comet's proto path), it computes
// this via compute_unified_partition_type(specs, schema) instead.
let unified_partition_type = {
let mut seen_field_ids = std::collections::HashSet::new();
let mut struct_fields: Vec<iceberg::spec::NestedFieldRef> = Vec::new();

for type_json in &proto_common.partition_type_pool {
match serde_json::from_str::<iceberg::spec::StructType>(type_json) {
Ok(struct_type) => {
for field in struct_type.fields() {
if seen_field_ids.insert(field.id) {
struct_fields.push(Arc::clone(field));
}
}
}
Err(e) => {
return Err(ExecutionError::GeneralError(format!(
"Failed to deserialize partition type JSON from pool: {e}"
)));
}
}
}

iceberg::spec::StructType::new(struct_fields)
};

let unified_partition_type_arc = Arc::new(unified_partition_type);

let results: Result<Vec<_>, _> = proto_tasks
.iter()
.map(|proto_task| {
Expand Down Expand Up @@ -3917,6 +3952,14 @@ fn parse_file_scan_tasks_from_common(
.field_ids
.clone();

let unified_partition_type_for_task = if project_field_ids
.contains(&iceberg::metadata_columns::RESERVED_FIELD_ID_PARTITION)
{
Some(Arc::clone(&unified_partition_type_arc))
} else {
None
};

Ok(iceberg::scan::FileScanTask {
file_size_in_bytes: proto_task.file_size_in_bytes,
data_file_path: proto_task.data_file_path.clone(),
Expand All @@ -3931,6 +3974,7 @@ fn parse_file_scan_tasks_from_common(
partition,
partition_spec,
name_mapping,
unified_partition_type: unified_partition_type_for_task,
case_sensitive: false,
// Plaintext StandardKeyMetadata forwarded verbatim from the JVM; decoded by
// iceberg-rust with no KMS unwrap. None for unencrypted data files.
Expand Down Expand Up @@ -5432,4 +5476,21 @@ mod tests {
}
});
}

#[test]
fn test_metadata_field_id_constants_match_iceberg_rust() {
// These constants are duplicated in Scala (CometIcebergNativeScan.MetadataFieldIds)
// as Int.MaxValue - 1 and Int.MaxValue - 5. This test ensures the Rust constants
// haven't drifted, which would cause a silent mismatch with the Scala side.
assert_eq!(
iceberg::metadata_columns::RESERVED_FIELD_ID_FILE,
i32::MAX - 1,
"RESERVED_FIELD_ID_FILE must be i32::MAX - 1 to match Scala MetadataFieldIds"
);
assert_eq!(
iceberg::metadata_columns::RESERVED_FIELD_ID_PARTITION,
i32::MAX - 5,
"RESERVED_FIELD_ID_PARTITION must be i32::MAX - 5 to match Scala MetadataFieldIds"
);
}
}
36 changes: 36 additions & 0 deletions spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@ object IcebergReflection extends Logging {
val SPARK_BATCH_QUERY_SCAN = "org.apache.iceberg.spark.source.SparkBatchQueryScan"
val SPARK_STAGED_SCAN = "org.apache.iceberg.spark.source.SparkStagedScan"
val SPARK_SCHEMA_UTIL = "org.apache.iceberg.spark.SparkSchemaUtil"
val TABLE = "org.apache.iceberg.Table"
val PARTITIONING = "org.apache.iceberg.Partitioning"
}

/**
Expand Down Expand Up @@ -461,6 +463,40 @@ object IcebergReflection extends Logging {
}
}

/**
* Validates that the table's unified partition type can be computed -- the merge of every
* historical partition spec, which is what the `_partition` metadata column projects.
*
* Iceberg Java's `Partitioning.partitionType(table)` runs the same cross-spec compatibility
* check that iceberg-rust does natively: a V1 table does not guarantee partition field ids are
* unique across specs, so two specs can bind the same id to incompatible source/transform
* pairs, which cannot be merged into one struct field. iceberg-rust returns a DataInvalid error
* in that case, but only at scan time -- too late for Comet to fall back. Calling the Java
* check here, at plan time, lets `CometScanRule` fall back to Spark instead of failing inside
* the native reader.
*
* Returns None when the unified type is computable, or Some(reason) when it is not -- either
* the specs conflict or the reflection call itself failed. Both mean Comet cannot safely serve
* `_partition`, so both map to a fallback.
*/
def validateUnifiedPartitionType(table: Any): Option[String] = {
try {
val tableClass = loadClass(ClassNames.TABLE)
val partitioningClass = loadClass(ClassNames.PARTITIONING)
partitioningClass
.getMethod("partitionType", tableClass)
.invoke(null, table.asInstanceOf[AnyRef])
None
} catch {
// A conflict surfaces as the ValidationException thrown by partitionType(), wrapped by
// reflection in InvocationTargetException; unwrap it for a meaningful reason.
case e: java.lang.reflect.InvocationTargetException =>
Some(Option(e.getCause).getOrElse(e).getMessage)
case e: Exception =>
Some(e.getMessage)
}
}

/**
* Gets the table metadata from an Iceberg table.
*
Expand Down
89 changes: 80 additions & 9 deletions spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ case class CometScanRule(session: SparkSession)
case scan if !CometConf.COMET_NATIVE_SCAN_ENABLED.get(conf) =>
withFallbackReason(scan, "Comet Scan is not enabled")

case scan if metadataCols(scan).nonEmpty =>
case scan: FileSourceScanExec if metadataCols(scan).nonEmpty =>

@hsiang-c hsiang-c Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

@hsiang-c hsiang-c Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can move it under the data source V1 branch?

withFallbackReason(
scan,
s"Metadata column(s) ${metadataCols(scan).mkString(", ")} is not supported")
Expand Down Expand Up @@ -302,6 +302,11 @@ case class CometScanRule(session: SparkSession)

scanExec.scan match {
case scan: CSVScan if COMET_CSV_V2_NATIVE_ENABLED.get() =>
if (scanExec.output.exists(_.isMetadataCol)) {
return withFallbackReason(
scanExec,
"Metadata columns are not supported for CSV V2 scans")
}
val fallbackReasons = new ListBuffer[String]()
val schemaSupported =
CometBatchScanExec.isSchemaSupported(scan.readDataSchema, fallbackReasons)
Expand Down Expand Up @@ -361,9 +366,30 @@ case class CometScanRule(session: SparkSession)
return withFallbackReasons(scanExec, fallbackReasons.toSet)
}

// Check for unsupported metadata columns in Iceberg scans
val unsupportedMetadataCols = scanExec.output.filter(_.isMetadataCol).filterNot { attr =>
CometIcebergNativeScan.MetadataFieldIds.keySet.contains(attr.name)
}
if (unsupportedMetadataCols.nonEmpty) {
fallbackReasons += "Unsupported Iceberg metadata columns: " +
unsupportedMetadataCols.map(_.name).mkString(", ")
return withFallbackReasons(scanExec, fallbackReasons.toSet)
}

val typeChecker = CometScanTypeChecker()
// Filter out metadata columns from schema check -- their types are handled
// by iceberg-rust directly (e.g., _partition can be an empty struct for
// unpartitioned tables which the general type checker rejects).
val metadataColNames = scanExec.output.filter(_.isMetadataCol).map(_.name).toSet
val dataSchema = if (metadataColNames.nonEmpty) {
val filtered =
scanExec.scan.readSchema().filter(f => !metadataColNames.contains(f.name))
new org.apache.spark.sql.types.StructType(filtered.toArray)
} else {
scanExec.scan.readSchema()
}
val schemaSupported =
typeChecker.isSchemaSupported(scanExec.scan.readSchema(), fallbackReasons)
typeChecker.isSchemaSupported(dataSchema, fallbackReasons)

if (!schemaSupported) {
fallbackReasons += "Comet extension is not enabled for " +
Expand Down Expand Up @@ -451,15 +477,17 @@ case class CometScanRule(session: SparkSession)
}

// Now perform all validation using the pre-extracted metadata
// Check if table uses a FileIO implementation compatible with iceberg-rust

// Check if table uses a FileIO implementation compatible with iceberg-rust.
// Comet's native reader uses object_store (Rust) for I/O, bypassing Iceberg Java's
// FileIO entirely. Only allow known-compatible implementations whose underlying
// storage object_store can reach via standard URL schemes.
val fileIOCompatible = IcebergReflection.getFileIO(metadata.table) match {
case Some(fileIO)
if fileIO.getClass.getName == "org.apache.iceberg.inmemory.InMemoryFileIO" =>
fallbackReasons += "InMemoryFileIO is not supported by Comet's native reader"
false
case Some(_) =>
case Some(fileIO) if CometScanRule.isCompatibleFileIO(fileIO.getClass.getName) =>
true
case Some(fileIO) =>
fallbackReasons += s"FileIO ${fileIO.getClass.getName} is not supported by " +
"Comet's native reader (object_store bypasses Iceberg Java FileIO)"
false
case None =>
fallbackReasons += "Could not check FileIO compatibility"
false
Expand Down Expand Up @@ -603,6 +631,25 @@ case class CometScanRule(session: SparkSession)
false
}

// The `_partition` metadata column projects the table's unified partition type -- the
// merge of all historical specs. iceberg-rust computes that merge at scan time and errors
// (DataInvalid) when specs bind one field id to incompatible source/transform pairs (only
// possible in V1 tables, which do not keep partition field ids unique across specs). That
// error would fail the native scan with no chance to fall back, so when `_partition` is
// projected we run Iceberg Java's equivalent check here and fall back if it cannot merge.
val unifiedPartitionTypeSupported =
if (scanExec.output.exists(a => a.isMetadataCol && a.name == "_partition")) {
IcebergReflection.validateUnifiedPartitionType(metadata.table) match {
case Some(reason) =>
fallbackReasons += "Iceberg table has partition specs whose unified partition " +
s"type cannot be computed for the _partition metadata column: $reason"
false
case None => true
}
} else {
true
}

// Get filter expressions for complex predicates check
val filterExpressionsOpt = IcebergReflection.getFilterExpressions(scanExec.scan)

Expand Down Expand Up @@ -788,6 +835,7 @@ case class CometScanRule(session: SparkSession)
if (schemaSupported && fileIOCompatible && formatVersionSupported &&
defaultValuesSupported && schemaTypesSupported && encryptionKeyLengthSupported &&
taskValidation.allParquet && allSupportedFilesystems && partitionTypesSupported &&
unifiedPartitionTypeSupported &&
complexTypePredicatesSupported && transformFunctionsSupported &&
deleteFileTypesSupported && dppSubqueriesSupported) {
CometBatchScanExec(
Expand Down Expand Up @@ -882,6 +930,29 @@ case class CometScanTypeChecker() extends DataTypeSupport with CometTypeShim {

object CometScanRule extends Logging {

// Iceberg FileIO implementations whose underlying storage object_store can reach.
// Custom/test FileIO classes (e.g. CustomFileIO in TestSparkExecutorCache) are not compatible
// because Comet's native reader bypasses Java FileIO entirely.
private val CompatibleFileIOClasses: Set[String] = Set(
"org.apache.iceberg.hadoop.HadoopFileIO",
"org.apache.iceberg.aws.s3.S3FileIO",
"org.apache.iceberg.gcp.gcs.GCSFileIO",
"org.apache.iceberg.io.ResolvingFileIO",
"org.apache.iceberg.spark.SparkFileIO",
"org.apache.iceberg.azure.adlsv2.ADLSFileIO",
"org.apache.iceberg.CachingFileIO")

// Prefix of the EncryptingFileIO family. An encrypted table's io() is not the bare
// EncryptingFileIO but a nested variant chosen from the wrapped delegate's capabilities
// (e.g. EncryptingFileIO$WithSupportsPrefixOperations when the delegate is HadoopFileIO), so
// an exact class-name match misses it. Comet forwards each file's key_metadata to iceberg-rust
// and reads the ciphertext via object_store, so any EncryptingFileIO variant is compatible.
private val EncryptingFileIOPrefix = "org.apache.iceberg.encryption.EncryptingFileIO"

/** True if `className` is a FileIO whose backing storage Comet's native reader can reach. */
def isCompatibleFileIO(className: String): Boolean =
CompatibleFileIOClasses.contains(className) || className.startsWith(EncryptingFileIOPrefix)

// Per-scheme memo of `NativeBase.isObjectStoreSchemeSupported`. The answer depends only on the
// URL scheme, so we cache by scheme and never re-cross the JNI boundary for a repeated scheme.
private val schemeSupportCache =
Expand Down
Loading
Loading