feat: write top-level Iceberg unknown columns natively

docs/source/contributor-guide/iceberg-writes.md

@@ -276,7 +276,10 @@ Points where Comet adapts iceberg-rust to match iceberg-java:
   held rows then go through `RowPacer`, so the roll grid does not move.
 - **Field ids and casting.** `decorate_batch_with_field_ids` casts each batch to the
   field-id-annotated Arrow schema derived from the Iceberg schema, with `safe: false`, so a type
-  mismatch fails the task instead of writing NULLs.
+  mismatch fails the task instead of writing NULLs. It first drops top-level `unknown` columns,
+  whose fields `parse_iceberg_schema` prunes. The gate declines a partition field sourced from an
+  `unknown` column: the partition calculator resolves sources by field id in the pruned schema and
+  would fail with `Field not found`.
 - **Deterministic output order.** `FanoutPartitions`, Comet's version of iceberg-rust's
   `FanoutWriter`, closes its files in an order that depends on how the rows arrived and on which
   partitions it closed early, so the fanout path sorts its `DataFile`s by path before returning them

docs/source/user-guide/latest/iceberg.md

@@ -483,8 +483,8 @@ A write is eligible only when ALL of the following hold:
 | resolved `table.locationProvider()`                                                                                                         | Iceberg's built-in `DefaultLocationProvider`                                                                                                                                                                                                                                                                                                                                                                                                                                    |
 | Hadoop S3A settings for an `s3` / `s3a` data location                                                                                       | only `fs.s3a.access.key`, `secret.key`, `session.token`, `endpoint`, `endpoint.region`, and `path.style.access`, including their `fs.s3a.bucket.<data-bucket>.*` forms; any other effective `fs.s3a.*` setting falls back, apart from the ignored settings listed below                                                                                                                                                                                                         |
 | Iceberg `FileIO` S3 settings for an `s3` / `s3a` data location                                                                              | the S3 endpoint, region, static/session credentials, path-style, SSE (`none`, `s3`, `kms`, or `custom`; not `dsse-kms`), assume-role, anonymous/config-chain settings parsed by the pinned iceberg-rust version, plus Comet's credential-provider class and built-in web-identity properties. When a custom provider is configured, its vendor-owned `s3.*` / `client.*` properties are also forwarded; unsupported Iceberg-defined S3 settings still fall back                 |
-| partition spec                                                                                                                              | any, except a `void` field whose source column was dropped beside a live field (see below)                                                                                                                                                                                                                                                                                                                                                                                      |
-| column types                                                                                                                                | any except `uuid` (Spark plans it as a string; no Arrow cast reaches `fixed(16)`) and the v3 types `variant`, `unknown`, `timestamp_ns`, `geometry` and `geography`                                                                                                                                                                                                                                                                                                             |
+| partition spec                                                                                                                              | any, except a `void` field whose source column was dropped beside a live field, or one sourced from an `unknown` column (see below)                                                                                                                                                                                                                                                                                                                                             |
+| column types                                                                                                                                | any except `uuid` (Spark plans it as a string; no Arrow cast reaches `fixed(16)`), the v3 types `variant`, `timestamp_ns`, `geometry` and `geography`, `unknown` inside a nested type, and a table whose columns are all `unknown`                                                                                                                                                                                                                                              |
 
 Within the namespaces that shape data-file bytes — `write.parquet.*` and `parquet.*` —
 everything not listed above must be absent: unvetted `write.parquet.*` keys (e.g.
@@ -553,6 +553,9 @@ cannot write through it either, so the write fails with iceberg-java's own error
 the native writer ([#6141](https://github.com/apache/datafusion-comet/issues/6141)). A spec whose
 fields are all `void` writes unpartitioned and stays eligible.
 
+A partition field sourced from an `unknown` column also falls back: the native writer drops that
+column.
+
 Other `write.*` properties are intentionally not gated because they cannot make the native
 writer produce different data files: distribution and ordering settings shape the Spark plan
 identically on both paths, WAP / branch / snapshot properties act on the JVM committer,

native/core/src/execution/operators/iceberg_write.rs

@@ -58,8 +58,8 @@ use iceberg::arrow::arrow_struct_to_literal;
 use iceberg::io::{FileIO, FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage};
 use iceberg::spec::{
     DataFile, DataFileBuilder, DataFileFormat, Literal, ManifestWriterBuilder, PartitionKey,
-    PartitionSpec, PartitionSpecRef, Schema as IcebergSchema, SchemaRef as IcebergSchemaRef,
-    Struct as IcebergStruct, StructType,
+    PartitionSpec, PartitionSpecRef, PrimitiveType, Schema as IcebergSchema,
+    SchemaRef as IcebergSchemaRef, Struct as IcebergStruct, StructType, Type,
 };
 use iceberg::writer::base_writer::data_file_writer::{DataFileWriter, DataFileWriterBuilder};
 use iceberg::writer::file_writer::location_generator::{
@@ -1488,6 +1488,22 @@ fn parse_iceberg_schema(json: &str) -> DFResult<IcebergSchemaRef> {
     let schema: IcebergSchema = serde_json::from_str(json).map_err(|e| {
         DataFusionError::Internal(format!("Failed to parse iceberg schema JSON: {e}"))
     })?;
+    // iceberg-java writes no column for a top-level `unknown` field.
+    let fields = schema.as_struct().fields();
+    let kept: Vec<_> = fields
+        .iter()
+        .filter(|f| *f.field_type != Type::Primitive(PrimitiveType::Unknown))
+        .cloned()
+        .collect();
+    if kept.len() == fields.len() {
+        return Ok(Arc::new(schema));
+    }
+    let schema = IcebergSchema::builder()
+        .with_schema_id(schema.schema_id())
+        .with_identifier_field_ids(schema.identifier_field_ids())
+        .with_fields(kept)
+        .build()
+        .map_err(|e| DataFusionError::Internal(format!("Failed to prune iceberg schema: {e}")))?;
     Ok(Arc::new(schema))
 }
 
@@ -1597,10 +1613,16 @@ fn decorate_batch_with_field_ids(
     batch: RecordBatch,
     target_schema: &SchemaRef,
 ) -> DFResult<RecordBatch> {
-    if batch.num_columns() != target_schema.fields().len() {
+    // `unknown` columns arrive as Null; `parse_iceberg_schema` dropped their fields.
+    let columns: Vec<&ArrayRef> = batch
+        .columns()
+        .iter()
+        .filter(|c| c.data_type() != &DataType::Null)
+        .collect();
+    if columns.len() != target_schema.fields().len() {
         return Err(DataFusionError::Plan(format!(
             "Iceberg write column count mismatch: arrow batch has {} columns but schema has {}",
-            batch.num_columns(),
+            columns.len(),
             target_schema.fields().len()
         )));
     }
@@ -1609,9 +1631,8 @@ fn decorate_batch_with_field_ids(
         safe: false,
         ..Default::default()
     };
-    let casted: Vec<ArrayRef> = batch
-        .columns()
-        .iter()
+    let casted: Vec<ArrayRef> = columns
+        .into_iter()
         .zip(target_schema.fields().iter())
         .map(|(col, target)| {
             arrow::compute::cast_with_options(col, target.data_type(), &cast_options)
@@ -4584,6 +4605,21 @@ mod tests {
             assert!(format!("{err}").contains("column count mismatch"));
         }
 
+        #[tokio::test]
+        async fn decorate_batch_drops_null_columns() {
+            let schema = iceberg_user_schema();
+            let target = Arc::new(iceberg::arrow::schema_to_arrow_schema(&schema).unwrap());
+            let original = batch(&[1, 2], &["a", "b"]);
+            let mut fields = original.schema().fields().to_vec();
+            fields.insert(1, Arc::new(Field::new("u", DataType::Null, true)));
+            let mut columns = original.columns().to_vec();
+            columns.insert(1, arrow::array::new_null_array(&DataType::Null, 2));
+            let with_null =
+                RecordBatch::try_new(Arc::new(ArrowSchema::new(fields)), columns).unwrap();
+            let decorated = decorate_batch_with_field_ids(with_null, &target).unwrap();
+            assert_eq!(decorated.schema(), target);
+        }
+
         /// A spec whose only partition field is a `void` transform, as iceberg-java's
         /// `UpdatePartitionSpec` leaves a format-version-1 spec after `DROP PARTITION FIELD`.
         /// Built by deserialising the spec JSON rather than through `PartitionSpec::builder`,

spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala

@@ -918,6 +918,26 @@ object IcebergReflection extends Logging {
     }
   }
 
+  /** Names of the partition fields of `spec` sourced from a top-level `unknown` column. */
+  def partitionFieldsWithUnknownSource(spec: Any): Seq[String] = {
+    import scala.jdk.CollectionConverters._
+    val schema = getMethod(spec.getClass, "schema").invoke(spec)
+    val unknownIds = getMethod(schema.getClass, "columns")
+      .invoke(schema)
+      .asInstanceOf[java.util.List[_]]
+      .asScala
+      .filter(c => fieldTypeOf(c).toString == TypeNames.UNKNOWN)
+      .map(fieldIdOf)
+      .toSet
+    getMethod(spec.getClass, "fields")
+      .invoke(spec)
+      .asInstanceOf[java.util.List[_]]
+      .asScala
+      .toSeq
+      .filter(f => unknownIds(getMethod(f.getClass, "sourceId").invoke(f).asInstanceOf[Int]))
+      .map(f => getMethod(f.getClass, "name").invoke(f).asInstanceOf[String])
+  }
+
   /**
    * Gets the partition spec from an Iceberg table.
    */
@@ -1680,17 +1700,24 @@ object IcebergReflection extends Logging {
    * write detection to decline schemas containing types the native writer cannot reproduce.
    * Reflection failures are deliberately not swallowed: the caller's detection wrapper turns them
    * into a fall-back, so a failed walk reads as "cannot verify" rather than "supported".
+   * `nestedOnly` skips the top-level columns.
    */
-  def findFieldWithTypeIds(schema: Any, typeIds: Set[String]): Option[(String, String)] = {
+  def findFieldWithTypeIds(
+      schema: Any,
+      typeIds: Set[String],
+      nestedOnly: Boolean = false): Option[(String, String)] = {
     val queue = new java.util.ArrayDeque[AnyRef]()
     queue.addAll(
       schema.getClass.getMethod("columns").invoke(schema).asInstanceOf[java.util.List[AnyRef]])
+    var topLevelLeft = if (nestedOnly) queue.size else 0
     while (!queue.isEmpty) {
       val field = queue.poll()
       val fieldType = field.getClass.getMethod("type").invoke(field)
       val typeId = fieldType.getClass.getMethod("typeId").invoke(fieldType)
       val typeIdName = typeId.asInstanceOf[Enum[_]].name()
-      if (typeIds.contains(typeIdName)) {
+      val skip = topLevelLeft > 0
+      if (skip) topLevelLeft -= 1
+      if (!skip && typeIds.contains(typeIdName)) {
         val name = field.getClass.getMethod("name").invoke(field).asInstanceOf[String]
         return Some((name, typeIdName))
       }

spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala

@@ -27,6 +27,7 @@ import scala.util.control.NonFatal
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometNativeExec, IcebergWriteExec}
+import org.apache.spark.sql.types.NullType
 
 import org.apache.comet.{CometConf, ConfigEntry}
 import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
@@ -81,11 +82,11 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] {
   }
 
   private val EncryptionPropertyPrefix = "encryption."
-  // `uuid` plus the v3-only types. Iceberg plans `variant` as Spark's VariantType and `unknown`
-  // as NullType, neither of which the native writer handles; Spark cannot plan a write to
-  // `timestamp_ns`, `geometry` or `geography` today, so those are declined in case it learns to.
+  // `uuid` plus the v3-only types. Iceberg plans `variant` as Spark's VariantType, which the
+  // native writer does not handle; Spark cannot plan a write to `timestamp_ns`, `geometry` or
+  // `geography` today, so those are declined in case it learns to. `unknown` is checked below.
   private val UnsupportedWriteTypeIds: Set[String] =
-    Set("UUID", "VARIANT", "UNKNOWN", "TIMESTAMP_NANO", "GEOMETRY", "GEOGRAPHY")
+    Set("UUID", "VARIANT", "TIMESTAMP_NANO", "GEOMETRY", "GEOGRAPHY")
   // Loaded from the native storage factory: `builtin_storage_schemes` in
   // `native/core/src/execution/operators/iceberg_common.rs` is the single point of change and
   // explains why `oss` is read-only and `memory` write-only. Lazy so constructing the serde does
@@ -416,6 +417,7 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] {
     requireNoMetadataColumns,
     requireSupportedColumnTypes,
     requireNoVoidFieldWithDroppedSource,
+    requireNoUnknownPartitionSource,
     requireNoEncryptionPrefix,
     requireNoBloomFilterColumnsEnabled,
     requireRowGroupCheckMinRecordCountAtDefault,
@@ -524,10 +526,18 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] {
       case Some(schema) =>
         IcebergReflection
           .findFieldWithTypeIds(schema, UnsupportedWriteTypeIds)
+          // Only top-level `unknown` fields are pruned, so a nested one would still
+          // get a Parquet column that iceberg-java does not write.
+          .orElse(
+            IcebergReflection.findFieldWithTypeIds(schema, Set("UNKNOWN"), nestedOnly = true))
           .map { case (name, typeId) =>
             s"column $name has Iceberg type ${typeId.toLowerCase(Locale.ROOT)}, " +
               "which the native writer cannot reproduce"
           }
+          // Pruning would leave no columns, which iceberg-java cannot write either.
+          .orElse(if (IcebergReflection.toSparkSchema(schema).forall(_.dataType == NullType)) {
+            Some("every column has Iceberg type unknown")
+          } else None)
     }
 
   // A format-version-1 spec keeps a dropped partition field as a `void` transform, and its source
@@ -536,17 +546,32 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] {
   // decline and let the write fail the way iceberg-java fails it. An all-`void` spec stays
   // eligible, since the native writer writes it unpartitioned.
   // https://github.com/apache/datafusion-comet/issues/6141
-  private val requireNoVoidFieldWithDroppedSource: TriggerRule = ctx =>
+  private val requireNoVoidFieldWithDroppedSource: TriggerRule =
+    requireOutputPartitionSpec { spec =>
+      IcebergReflection.voidFieldsWithDroppedSource(spec).headOption.map { name =>
+        s"partition field $name is a void transform whose source column was dropped, " +
+          "beside a live partition field"
+      }
+    }
+
+  // Top-level `unknown` columns are pruned from the native schema, so a partition field sourced
+  // from one would fail the native partition calculator with `Field not found`.
+  private val requireNoUnknownPartitionSource: TriggerRule =
+    requireOutputPartitionSpec { spec =>
+      IcebergReflection.partitionFieldsWithUnknownSource(spec).headOption.map { name =>
+        s"partition field $name is sourced from a column of Iceberg type unknown"
+      }
+    }
+
+  // Applies `check` to the output partition spec, declining if it cannot be resolved or inspected.
+  private def requireOutputPartitionSpec(check: Any => Option[String]): TriggerRule = ctx =>
     IcebergReflection
       .getOutputSpecIdFromSparkWrite(ctx.sparkWrite)
       .flatMap(IcebergReflection.getPartitionSpecById(ctx.table, _)) match {
-      case None => Some("could not resolve the output partition spec for void field checking")
+      case None => Some("could not resolve the output partition spec")
       case Some(spec) =>
         try {
-          IcebergReflection.voidFieldsWithDroppedSource(spec).headOption.map { name =>
-            s"partition field $name is a void transform whose source column was dropped, " +
-              "beside a live partition field"
-          }
+          check(spec)
         } catch {
           case e: Exception =>
             Some(s"could not inspect the output partition spec: ${e.getMessage}")

spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala

@@ -96,23 +96,29 @@ trait CometIcebergTestBase { this: CometTestBase =>
   /**
    * Adds a column of an Iceberg type Spark DDL cannot declare (e.g. `uuid`, `fixed(N)`) by
    * committing an `UpdateSchema` against the Iceberg table directly. Callers must `REFRESH TABLE`
-   * afterwards so Spark's catalog cache picks up the new schema.
+   * afterwards so Spark's catalog cache picks up the new schema. `parent` nests it in a struct.
    */
   protected def addIcebergColumn(
       icebergTable: AnyRef,
       columnName: String,
-      icebergType: AnyRef): Unit = {
+      icebergType: AnyRef,
+      parent: Option[String] = None): Unit = {
     val update = IcebergReflection
       .loadClass("org.apache.iceberg.Table")
       .getMethod("updateSchema")
       .invoke(icebergTable)
     val updateSchemaClass = IcebergReflection.loadClass("org.apache.iceberg.UpdateSchema")
-    updateSchemaClass
-      .getMethod(
-        "addColumn",
-        classOf[String],
-        IcebergReflection.loadClass("org.apache.iceberg.types.Type"))
-      .invoke(update, columnName, icebergType)
+    val typeClass = IcebergReflection.loadClass("org.apache.iceberg.types.Type")
+    parent match {
+      case None =>
+        updateSchemaClass
+          .getMethod("addColumn", classOf[String], typeClass)
+          .invoke(update, columnName, icebergType)
+      case Some(p) =>
+        updateSchemaClass
+          .getMethod("addColumn", classOf[String], classOf[String], typeClass)
+          .invoke(update, p, columnName, icebergType)
+    }
     updateSchemaClass.getMethod("commit").invoke(update)
   }