andygrove commented on code in PR #5361:
URL: https://github.com/apache/datafusion-comet/pull/5361#discussion_r3786613036
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -569,6 +569,569 @@ class CometIcebergWriteActionSuite
}
}
+ // --- Round-trip parity vs Spark default path
---------------------------------------------------
+
+ // --- Native acceleration
--------------------------------------------------------------------
+
+ test("native acceleration: AppendData INSERT FROM SELECT") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "native_source", partitionSpec = "")
+ createTable(warehouseDir, "native_target", partitionSpec = "")
+ spark.sql(
+ "INSERT INTO cat.db.native_source VALUES " +
+ "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)")
+ assertNativeWriteEngages("native_target", Seq(1, 2, 3)) {
+ spark.sql(
+ "INSERT INTO cat.db.native_target SELECT id, region, amount FROM
cat.db.native_source")
+ }
+ }
+ }
+
+ test("native acceleration: AppendData unpartitioned VALUES") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "native_append_values", partitionSpec = "")
+ assertNativeWriteEngages("native_append_values", Seq(1, 2, 3)) {
+ spark.sql(
+ "INSERT INTO cat.db.native_append_values VALUES " +
+ "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)")
+ }
+ }
+ }
+
+ // What Iceberg's Spark writer stamps for `sort_order_id` on appended files
changed across
+ // releases: through 1.10 `SparkWrite$WriterFactory` never wires the table
sort order (files
+ // get 0 even on a sorted table); 1.11 added
`SparkWriteConf.outputSortOrderId` and stamps the
+ // resolved order id. The native path reflects the resolver when present and
defaults to 0
+ // otherwise, so pin parity against the JVM writer on the same runtime
instead of a literal.
+ test("native acceleration: appended files carry the same sort_order_id as
the JVM writer") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ // WRITE ORDERED BY (provided by IcebergSparkSessionExtensions, enabled
in this suite)
+ // bumps the table's sort order id to a non-default value (1).
+ Seq("sorted_native", "sorted_jvm").foreach { t =>
+ createTable(warehouseDir, t, partitionSpec = "")
+ spark.sql(s"ALTER TABLE cat.db.$t WRITE ORDERED BY id")
+ }
+ val insert = (t: String) =>
+ spark.sql(
+ s"INSERT INTO cat.db.$t VALUES " +
+ "(3, 'eu', 30.7), (1, 'us-east', 10.5), (2, 'us-west', 20.3)")
+ assertNativeWriteEngages("sorted_native", Seq(1, 2,
3))(insert("sorted_native"))
+ insert("sorted_jvm")
+
+ def sortOrderIds(t: String): Set[Int] = spark
+ .sql(s"SELECT DISTINCT sort_order_id FROM cat.db.$t.data_files")
+ .collect()
+ .map(_.getInt(0))
+ .toSet
+ val nativeIds = sortOrderIds("sorted_native")
+ val jvmIds = sortOrderIds("sorted_jvm")
+ assert(nativeIds == jvmIds, s"native sort_order_ids $nativeIds != JVM
$jvmIds")
+ assert(nativeIds.size == 1, s"expected one distinct sort_order_id, got
$nativeIds")
+ }
+ }
+
+ test("native acceleration: AppendData partitioned by identity") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "native_append_part", partitionSpec =
"PARTITIONED BY (region)")
+ assertNativeWriteEngages("native_append_part", Seq(1, 2, 3)) {
+ spark.sql(
+ "INSERT INTO cat.db.native_append_part VALUES " +
+ "(1, 'us-east', 10.5), (2, 'us-east', 20.3), (3, 'eu', 30.7)")
+ }
+ }
+ }
+
+ test("native acceleration: OverwriteByExpression (INSERT OVERWRITE STATIC)")
{
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "native_overwrite_static", partitionSpec = "")
+ spark.sql(
+ "INSERT INTO cat.db.native_overwrite_static VALUES " +
+ "(1, 'old', 1.0), (2, 'old', 2.0), (3, 'old', 3.0)")
+ assertNativeWriteEngages("native_overwrite_static", Seq(10, 11)) {
+ withSQLConf("spark.sql.sources.partitionOverwriteMode" -> "STATIC") {
+ spark.sql(
+ "INSERT OVERWRITE cat.db.native_overwrite_static VALUES " +
+ "(10, 'new', 100.0), (11, 'new', 110.0)")
+ }
+ }
+ }
+ }
+
+ test("native acceleration: OverwritePartitionsDynamic") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "native_overwrite_dyn", partitionSpec =
"PARTITIONED BY (region)")
+ spark.sql(
+ "INSERT INTO cat.db.native_overwrite_dyn VALUES " +
+ "(1, 'us-east', 1.0), (2, 'us-west', 2.0), (3, 'eu', 3.0)")
+ assertNativeWriteEngages("native_overwrite_dyn", Seq(2, 3, 10)) {
+ withSQLConf("spark.sql.sources.partitionOverwriteMode" -> "DYNAMIC") {
+ spark.sql("INSERT OVERWRITE cat.db.native_overwrite_dyn VALUES (10,
'us-east', 100.0)")
+ }
+ }
+ }
+ }
+
+ test("native acceleration: ReplaceData (CoW DELETE)") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "native_cow_delete",
+ partitionSpec = "",
+ properties = Some("'write.delete.mode'='copy-on-write'"))
+ // Seed via the JVM path so the assertion isolates native engagement to
the DELETE.
+ withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key ->
"false") {
+ coalesceInsert(
+ "native_cow_delete",
+ Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0), (4,
"us-east", 40.0)))
+ }
+ assertNativeWriteEngages("native_cow_delete", Seq(1, 3, 4)) {
+ spark.sql("DELETE FROM cat.db.native_cow_delete WHERE id = 2")
+ }
+ }
+ }
+
+ test("native acceleration: ReplaceData (CoW UPDATE)") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "native_cow_update",
+ partitionSpec = "",
+ properties = Some("'write.update.mode'='copy-on-write'"))
+ withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key ->
"false") {
+ coalesceInsert(
+ "native_cow_update",
+ Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0)))
+ }
+ // Engage natively; expected rows are the ids 1..3 (UPDATE keeps
cardinality).
+ assertNativeWriteEngages("native_cow_update", Seq(1, 2, 3)) {
+ spark.sql("UPDATE cat.db.native_cow_update SET amount = amount * 2
WHERE id = 2")
+ }
+ // Spot-check the UPDATE actually rewrote row 2 (cardinality unchanged +
value flipped).
+ val r =
+ spark.sql("SELECT id, amount FROM cat.db.native_cow_update WHERE id =
2").collect()
+ assert(r.length == 1 && r(0).getDouble(1) == 40.0, s"got ${r.toSeq}")
+ }
+ }
+
+ test("native acceleration: ReplaceData (CoW MERGE) falls back (MergeRowsExec
not Comet)") {
+ // TODO(comet-merge-rows): native MERGE engagement requires a Comet
equivalent of Iceberg's
+ // `MergeRowsExec` (the per-row dispatch operator that assigns
__row_operation codes from
+ // MATCHED/NOT MATCHED clauses). Without it, `MergeRowsExec` stays JVM,
the upstream chain
+ // breaks Comet-native partway, and `requiresNativeChildren=true` declines
the
+ // `IcebergWriteExec -> CometIcebergWriteExec` conversion. Until that
lands, MERGE
+ // falls back to the JVM two-op path -- this test pins that contract so a
future MERGE-row-exec
+ // addition surfaces clearly (the test will start failing and need to flip
back to
+ // `assertNativeWriteEngages`).
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "native_cow_merge",
+ partitionSpec = "",
+ properties = Some("'write.merge.mode'='copy-on-write'"))
+ withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key ->
"false") {
+ coalesceInsert("native_cow_merge", Seq((1, "us-east", 10.0), (2,
"us-west", 20.0)))
+ }
+ assertNativeWriteDoesNotEngage("native_cow_merge", Seq(1, 2, 3)) {
+ spark.sql("""
+ |MERGE INTO cat.db.native_cow_merge t
+ |USING (SELECT 2 AS id, 'us-west' AS region, 200.0 AS amount UNION
ALL
+ | SELECT 3 AS id, 'eu' AS region, 30.0 AS amount) s
+ |ON t.id = s.id
+ |WHEN MATCHED THEN UPDATE SET t.amount = s.amount
+ |WHEN NOT MATCHED THEN INSERT (id, region, amount) VALUES (s.id,
s.region, s.amount)
+ |""".stripMargin)
+ }
+ }
+ }
+
+ test("native acceleration: complex types (struct, array, map) round-trip
with field IDs") {
Review Comment:
The native-acceleration tests all write `(id INT, region STRING, amount
DOUBLE)`, and this complex-types test adds STRING and INT leaves, so the type
surface actually exercised through the native writer is fairly narrow.
Would you mind adding a round-trip parity test over the wider primitive set:
`DATE`, `TIMESTAMP`, `TIMESTAMP_NTZ`, `DECIMAL` at a couple of precisions,
`BINARY`, `BOOLEAN`, `BIGINT`, `FLOAT`? The one I would most like pinned is
timestamps under a non-UTC `spark.sql.session.timeZone`, with both `timestamp`
and `timestamptz` columns. That is exactly where a silent value shift would
hide, and none of the current tests would notice it. Your
`manifestMetricsParity` helper looks like it would adapt nicely to compare row
values as well as metrics.
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala:
##########
@@ -163,12 +159,12 @@ object CometIcebergNativeWrite extends
CometOperatorSerde[IcebergWriteExec] with
"custom location provider unsupported"),
requireFormatVersionAtMostTwo,
requireNoEncryptionPrefix,
- requireSupportedMetricsModes,
requireNoBloomFilterColumnsEnabled,
requireRowGroupCheckMinRecordCountAtDefault,
requireRowGroupCheckMaxRecordCountAtDefault,
requireParquetPageVersionDefault,
requireShredVariantsDisabled,
+ requireParseableCompressionLevel,
Review Comment:
I think a table with a `uuid` column would be judged eligible here and then
fail the task rather than falling back.
Iceberg's `TypeToSparkType` maps `uuid` to `StringType`, so Comet hands the
native writer a Utf8 column, while `schema_to_arrow_schema` makes the target
`FixedSizeBinary(16)`. Since `decorate_batch_with_field_ids` casts with `safe:
false`, a 36-character UUID string has nowhere to go. The gate is entirely
table-property based with no schema-type check, so detection would return
`Compatible` first and the failure would surface at execution time.
Have you tried a table with a `uuid` column? If it does fail, a trigger rule
that declines on column types the native writer cannot reproduce (plus a
detection test pinning it) would keep this a fallback rather than a task error.
`fixed(N)` should be fine via `Binary -> FixedSizeBinary` as long as every
value is exactly N bytes, but it would be good to have that pinned too.
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -0,0 +1,1111 @@
+// 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.
+
+//! Native Iceberg write operator using iceberg-rust.
+//!
+//! Drains the upstream Arrow stream through iceberg-rust's writer stack
+//! (`ParquetWriterBuilder` -> `RollingFileWriterBuilder` ->
`DataFileWriterBuilder`
+//! -> `Unpartitioned`/`Fanout`/`Clustered`Writer) and emits a single-row,
single-column
+//! Arrow batch carrying the `Vec<DataFile>` produced for the task, packed as
an Iceberg V2
+//! data manifest via iceberg-rust's `ManifestWriter` against an in-memory
`FileIO`. The JVM
+//! decodes the bytes with `ManifestFiles.read(...)` to recover the
`DataFile`s for commit.
+
+use std::fmt;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, BinaryArray, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef};
+use datafusion::error::{DataFusionError, Result as DFResult};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
Partitioning,
+ PlanProperties, SendableRecordBatchStream,
+};
+use futures::TryStreamExt;
+use iceberg::arrow::{
+ arrow_struct_to_literal, PartitionValueCalculator,
RecordBatchPartitionSplitter,
+};
+use iceberg::spec::{
+ DataFile, DataFileFormat, Literal, ManifestWriterBuilder, PartitionSpec,
PartitionSpecRef,
+ Schema as IcebergSchema, SchemaRef as IcebergSchemaRef, Struct as
IcebergStruct, StructType,
+};
+use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
+use iceberg::writer::file_writer::location_generator::{
+ DefaultFileNameGenerator, DefaultLocationGenerator,
+};
+use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
+use iceberg::writer::file_writer::ParquetWriterBuilder;
+use iceberg::writer::partitioning::clustered_writer::ClusteredWriter;
+use iceberg::writer::partitioning::fanout_writer::FanoutWriter;
+use iceberg::writer::partitioning::unpartitioned_writer::UnpartitionedWriter;
+#[cfg(test)]
+use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
+use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
+use parquet::file::properties::{EnabledStatistics, WriterProperties};
+
+use datafusion_comet_proto::spark_operator::{
+ CompressionCodec as ProtoCompressionCodec, IcebergParquetWriteSettings,
IcebergWrite,
+ IcebergWriteCommon, IcebergWriterMode as ProtoIcebergWriterMode,
+};
+
+use crate::cloud::s3::credential_bridge::AccessMode;
+use crate::execution::operators::iceberg_common::load_file_io;
+
+/// Builder chain instantiated once per task and handed to the partitioning
wrapper.
+type IcebergDataFileWriterBuilder =
+ DataFileWriterBuilder<ParquetWriterBuilder, DefaultLocationGenerator,
DefaultFileNameGenerator>;
+
+/// Native Iceberg write operator. Owns the parsed Iceberg schema/spec and the
parquet writer
+/// properties; at task execution it builds the iceberg-rust writer stack,
drains the upstream
+/// Arrow stream into it, and emits a single Avro-encoded `Vec<DataFile>` row.
+pub struct IcebergWriteExec {
+ input: Arc<dyn ExecutionPlan>,
+ common: Arc<IcebergWriteCommon>,
+ iceberg_schema: IcebergSchemaRef,
+ partition_spec: PartitionSpecRef,
+ writer_mode: ProtoIcebergWriterMode,
+ writer_properties: Arc<WriterProperties>,
+ partition_id: Option<i32>,
+ task_attempt_id: Option<i64>,
+ output_schema: SchemaRef,
+ plan_properties: Arc<PlanProperties>,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl IcebergWriteExec {
+ pub fn try_new(input: Arc<dyn ExecutionPlan>, proto: IcebergWrite) ->
DFResult<Self> {
+ let IcebergWrite {
+ common,
+ partition_id,
+ task_attempt_id,
+ } = proto;
+ let common = common.ok_or_else(|| {
+ DataFusionError::Internal("IcebergWrite missing common
payload".into())
+ })?;
+ let settings = common.parquet_settings.as_ref().ok_or_else(|| {
+ DataFusionError::Internal("IcebergWriteCommon missing
parquet_settings".into())
+ })?;
+ let writer_properties = build_writer_properties(settings)?;
+ let iceberg_schema =
parse_iceberg_schema(&common.iceberg_schema_json)?;
+ let partition_spec =
parse_partition_spec(&common.partition_spec_json)?;
+ let writer_mode =
ProtoIcebergWriterMode::try_from(common.writer_mode).map_err(|_| {
+ DataFusionError::Internal(format!(
+ "Unknown IcebergWriterMode proto value: {}",
+ common.writer_mode
+ ))
+ })?;
+ let output_schema = build_output_schema();
+ let plan_properties = Self::compute_properties(&input,
Arc::clone(&output_schema));
+ Ok(Self {
+ input,
+ common: Arc::new(common),
+ iceberg_schema,
+ partition_spec,
+ writer_mode,
+ writer_properties: Arc::new(writer_properties),
+ partition_id,
+ task_attempt_id,
+ output_schema,
+ plan_properties,
+ metrics: ExecutionPlanMetricsSet::new(),
+ })
+ }
+
+ fn compute_properties(
+ input: &Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ ) -> Arc<PlanProperties> {
+ Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(schema),
+
Partitioning::UnknownPartitioning(input.output_partitioning().partition_count()),
+ EmissionType::Final,
+ Boundedness::Bounded,
+ ))
+ }
+}
+
+impl ExecutionPlan for IcebergWriteExec {
+ fn name(&self) -> &str {
+ "IcebergWriteExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.plan_properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.input]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> DFResult<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return Err(DataFusionError::Internal(
+ "IcebergWriteExec requires exactly one child".into(),
+ ));
+ }
+ Ok(Arc::new(Self {
+ input: children.pop().unwrap(),
+ common: Arc::clone(&self.common),
+ iceberg_schema: Arc::clone(&self.iceberg_schema),
+ partition_spec: Arc::clone(&self.partition_spec),
+ writer_mode: self.writer_mode,
+ writer_properties: Arc::clone(&self.writer_properties),
+ partition_id: self.partition_id,
+ task_attempt_id: self.task_attempt_id,
+ output_schema: Arc::clone(&self.output_schema),
+ plan_properties: Arc::clone(&self.plan_properties),
+ metrics: self.metrics.clone(),
+ }))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> DFResult<SendableRecordBatchStream> {
+ let input_stream = self.input.execute(partition, context)?;
+ let common = Arc::clone(&self.common);
+ let iceberg_schema = Arc::clone(&self.iceberg_schema);
+ let partition_spec = Arc::clone(&self.partition_spec);
+ let writer_mode = self.writer_mode;
+ let writer_properties = Arc::clone(&self.writer_properties);
+ let partition_id = self.partition_id;
+ let task_attempt_id = self.task_attempt_id;
+ let output_schema = Arc::clone(&self.output_schema);
+
+ let task = async move {
+ let data_files = run_write_task(
+ input_stream,
+ Arc::clone(&common),
+ Arc::clone(&iceberg_schema),
+ Arc::clone(&partition_spec),
+ writer_mode,
+ writer_properties.as_ref().clone(),
+ partition_id,
+ task_attempt_id,
+ )
+ .await?;
+ let manifest_bytes = encode_data_files_as_manifest(
+ data_files,
+ iceberg_schema,
+ partition_spec,
+ partition_id,
+ task_attempt_id,
+ &common.operation_id,
+ )
+ .await?;
+ let batch = build_output_batch(manifest_bytes, &output_schema)?;
+ Ok::<_, DataFusionError>(futures::stream::iter(vec![Ok(batch)]))
+ };
+
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&self.output_schema),
+ futures::stream::once(task).try_flatten(),
+ )))
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+}
+
+impl fmt::Debug for IcebergWriteExec {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("IcebergWriteExec")
+ .field("metadata_location", &self.common.metadata_location)
+ .field("data_location", &self.common.data_location)
+ .field("operation_id", &self.common.operation_id)
+ .field("writer_mode", &self.writer_mode)
+ .field("partition_id", &self.partition_id)
+ .field("task_attempt_id", &self.task_attempt_id)
+ .finish()
+ }
+}
+
+impl DisplayAs for IcebergWriteExec {
+ fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) ->
fmt::Result {
+ write!(
+ f,
+ "IcebergWriteExec: metadata_location={}, data_location={},
operation_id={}",
+ self.common.metadata_location, self.common.data_location,
self.common.operation_id
+ )
+ }
+}
+
+/// One-shot per-task write coroutine. Builds the iceberg-rust writer stack,
decorates each input
+/// batch with `PARQUET_FIELD_ID_META_KEY` metadata so iceberg-rust can match
Arrow columns to
+/// Iceberg field IDs, and routes through
`UnpartitionedWriter`/`FanoutWriter`/`ClusteredWriter`
+/// depending on `writer_mode`.
+#[allow(clippy::too_many_arguments)]
+async fn run_write_task(
+ mut input: SendableRecordBatchStream,
+ common: Arc<IcebergWriteCommon>,
+ iceberg_schema: IcebergSchemaRef,
+ partition_spec: PartitionSpecRef,
+ writer_mode: ProtoIcebergWriterMode,
+ writer_properties: WriterProperties,
+ partition_id: Option<i32>,
+ task_attempt_id: Option<i64>,
+) -> DFResult<Vec<DataFile>> {
+ // The JVM exec wrapper stamps both ids per task; a missing id means the
plan template was
+ // executed directly, and defaulting would make every task collide on the
same file names.
+ let partition_id = partition_id.ok_or_else(|| {
+ DataFusionError::Internal("IcebergWrite executed without a
partition_id".into())
+ })?;
+ let task_attempt_id = task_attempt_id.ok_or_else(|| {
+ DataFusionError::Internal("IcebergWrite executed without a
task_attempt_id".into())
+ })?;
+ let catalog_properties = common
+ .catalog_properties
+ .iter()
+ .map(|(k, v)| (k.clone(), v.clone()))
+ .collect();
+ let file_io = load_file_io(
+ &catalog_properties,
+ &common.data_location,
+ &common.catalog_name,
+ AccessMode::Write,
+ )?;
+
+ let location_generator =
+
DefaultLocationGenerator::with_data_location(common.data_location.clone());
+ let file_name_generator = DefaultFileNameGenerator::new(
+ file_name_prefix(partition_id, task_attempt_id, &common.operation_id),
+ None,
+ DataFileFormat::Parquet,
+ );
+ let parquet_builder = ParquetWriterBuilder::new(writer_properties,
Arc::clone(&iceberg_schema));
+ let rolling_builder = RollingFileWriterBuilder::new(
+ parquet_builder,
+ common.target_file_size_bytes as usize,
+ file_io,
+ location_generator,
+ file_name_generator,
+ );
+ let data_file_builder = DataFileWriterBuilder::new(rolling_builder);
+
+ let unpartitioned = partition_spec.is_unpartitioned();
+ let mut writer = match (unpartitioned, writer_mode) {
+ (true, ProtoIcebergWriterMode::IcebergWriterUnpartitioned) => {
+
InnerWriter::Unpartitioned(UnpartitionedWriter::new(data_file_builder))
+ }
+ (false, ProtoIcebergWriterMode::IcebergWriterFanout) => {
+ InnerWriter::Fanout(FanoutWriter::new(data_file_builder))
+ }
+ (false, ProtoIcebergWriterMode::IcebergWriterClustered) => {
+ InnerWriter::Clustered(ClusteredWriter::new(data_file_builder))
+ }
+ (actual, mode) => {
+ return Err(DataFusionError::Internal(format!(
+ "IcebergWrite writer_mode {mode:?} is inconsistent with the
partition spec \
+ (unpartitioned={actual})"
+ )))
+ }
+ };
+
+ // `RecordBatchPartitionSplitter::split` groups rows through a HashMap and
emits the parts in
+ // unspecified order. `ClusteredWriter` hard-errors when a closed
partition is revisited, so
+ // the parts of every batch must be written in the batch's own
(partition-sorted) order --
+ // recover it by computing each row's partition value and recording first
occurrences.
+ let clustered_order = match &writer {
+ InnerWriter::Clustered(_) => Some((
+ PartitionValueCalculator::try_new(&partition_spec, &iceberg_schema)
+ .map_err(iceberg_err)?,
+ partition_spec
+ .partition_type(&iceberg_schema)
+ .map_err(iceberg_err)?,
+ )),
+ _ => None,
+ };
+
+ let splitter = if unpartitioned {
+ None
+ } else {
+ Some(
+ RecordBatchPartitionSplitter::try_new_with_computed_values(
+ Arc::clone(&iceberg_schema),
+ Arc::clone(&partition_spec),
+ )
+ .map_err(iceberg_err)?,
+ )
+ };
+
+ // Build the field-id-decorated target schema once per task; every batch
is cast against it.
+ let target_schema =
+
Arc::new(iceberg::arrow::schema_to_arrow_schema(&iceberg_schema).map_err(iceberg_err)?);
+ while let Some(batch) = input.try_next().await? {
+ let decorated = decorate_batch_with_field_ids(batch, &target_schema)?;
+ writer
+ .write(decorated, splitter.as_ref(), clustered_order.as_ref())
+ .await?;
+ }
+ writer.close().await
+}
+
+/// Enum-based dispatch over the three iceberg-rust partitioning writers. Each
variant takes the
+/// same builder chain so we can keep the type fixed.
+enum InnerWriter {
+ Unpartitioned(UnpartitionedWriter<IcebergDataFileWriterBuilder>),
+ Fanout(FanoutWriter<IcebergDataFileWriterBuilder>),
+ Clustered(ClusteredWriter<IcebergDataFileWriterBuilder>),
+}
+
+impl InnerWriter {
+ async fn write(
+ &mut self,
+ batch: RecordBatch,
+ splitter: Option<&RecordBatchPartitionSplitter>,
+ clustered_order: Option<&(PartitionValueCalculator, StructType)>,
+ ) -> DFResult<()> {
+ use iceberg::writer::partitioning::PartitioningWriter;
+ match self {
+ InnerWriter::Unpartitioned(w) =>
w.write(batch).await.map_err(iceberg_err),
+ InnerWriter::Fanout(w) => {
+ let parts = splitter
+ .expect("partition splitter must be Some for partitioned
writes")
+ .split(&batch)
+ .map_err(iceberg_err)?;
+ for (key, part) in parts {
+ w.write(key, part).await.map_err(iceberg_err)?;
+ }
+ Ok(())
+ }
+ InnerWriter::Clustered(w) => {
+ let mut parts = splitter
+ .expect("partition splitter must be Some for partitioned
writes")
+ .split(&batch)
+ .map_err(iceberg_err)?;
+ let (calculator, partition_type) = clustered_order
+ .expect("clustered order helper must be Some for clustered
writes");
+ let order = partition_first_occurrence_order(&batch,
calculator, partition_type)?;
+ parts.sort_by_key(|(key, _)|
order.get(key.data()).copied().unwrap_or(usize::MAX));
Review Comment:
This computes the partition values a second time.
`partition_first_occurrence_order` runs `PartitionValueCalculator::calculate`
over the whole batch, but `RecordBatchPartitionSplitter::split` (two lines up)
has already evaluated the same partition transforms internally. Clustered is
the default for partitioned tables, so this doubles the partition-transform
cost on the common path.
Is there a way to recover the ordering without the second pass? The input is
partition-sorted by construction, so anything that gives each part's first row
index would do. Alternatively, a single `calculate` whose result feeds both the
split and the ordering, if iceberg-rust exposes that shape.
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala:
##########
@@ -354,10 +354,21 @@ object CometMetricNode {
/**
* Creates a [[CometMetricNode]] from a [[CometPlan]].
+ *
+ * Stops walking at non-Comet nodes: a JVM-side `AQEShuffleReadExec` (or any
other Spark exec
+ * constructed off the planning thread) captures `SparkPlan.session` eagerly
as a `@transient
+ * final val`, which can be `null` if `SparkSession.getActiveSession`
returned `None` at the
+ * moment AQE's stage-finalisation rules built it. Forcing such a node's
`metrics` lazy val NPEs
+ * in `SQLMetrics.createMetric(sparkContext, ...)`. We don't own those
metrics anyway -- the
+ * native side only sources updates against operators it actually planned,
and JVM-side
+ * AQE-stage nodes belong to a different stage that's already been
materialised independently.
*/
- def fromCometPlan(cometPlan: SparkPlan): CometMetricNode = {
- val children = cometPlan.children.map(fromCometPlan)
- CometMetricNode(cometPlan.metrics, children)
+ def fromCometPlan(cometPlan: SparkPlan): CometMetricNode = cometPlan match {
Review Comment:
The AQE reasoning in the doc comment makes sense, but this changes behaviour
for every caller, not just the Iceberg write: `CometNativeScanExec`,
`CometShuffleExchangeExec`, `CometNativeWriteExec` and `operators.scala` all
come through here, and a whole subtree now silently reports `Map.empty` where
it previously would have contributed metrics.
I went looking and could not find a non-`CometPlan` node living inside a
native block today, so I think it is safe in practice. Could you either narrow
the guard to the node type that actually triggers this (or guard the `metrics`
access itself), or add a test pinning that no existing operator's metrics
regress? A silent metric drop is the kind of thing that goes unnoticed for a
long time.
##########
spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala:
##########
@@ -1124,6 +1280,464 @@ object IcebergReflection extends Logging {
logError(s"Iceberg reflection failure: Failed to get data location:
${e.getMessage}", e)
None
}
+
+ /**
+ * Returns the top-level column names of an Iceberg `Schema`, in declared
order. Used by the
+ * native write serde to project Spark 4.x `ReplaceData` row streams (which
carry an
+ * `__row_operation` column plus optional file-metadata columns) down to
just the data columns
+ * the native iceberg-rust writer expects.
+ */
+ def getSchemaFieldNames(schema: Any): Option[Seq[String]] =
+ try {
+ val cols = schema.getClass
+ .getMethod("columns")
+ .invoke(schema)
+ .asInstanceOf[java.util.List[_]]
+ val names = new scala.collection.mutable.ArrayBuffer[String](cols.size())
+ val it = cols.iterator()
+ while (it.hasNext) {
+ val col = it.next().asInstanceOf[AnyRef]
+ names +=
col.getClass.getMethod("name").invoke(col).asInstanceOf[String]
+ }
+ Some(names.toSeq)
+ } catch {
+ case e: Exception =>
+ logError(s"Iceberg reflection failure: Schema.columns():
${e.getMessage}")
+ None
+ }
+
+ /**
+ * Sum `recordCount` and `fileSizeInBytes` across `dataFiles` for SQL-metric
reporting. The
+ * concrete `DataFile` impl (`BaseFile`) is package-private in Iceberg, so
look the accessors up
+ * on the public `DataFile` interface instead; virtual dispatch still hits
the concrete
+ * implementation at invoke time.
+ */
+ def sumDataFileMetrics(dataFiles: java.util.List[_]): (Long, Long) = {
+ if (dataFiles.isEmpty) return (0L, 0L)
+ val dataFileClass = loadClass(ClassNames.DATA_FILE)
+ val recordCountMethod = dataFileClass.getMethod("recordCount")
+ val fileSizeMethod = dataFileClass.getMethod("fileSizeInBytes")
+ var rows = 0L
+ var bytes = 0L
+ val it = dataFiles.iterator()
+ while (it.hasNext) {
+ val df = it.next().asInstanceOf[AnyRef]
+ rows +=
recordCountMethod.invoke(df).asInstanceOf[java.lang.Long].longValue()
+ bytes +=
fileSizeMethod.invoke(df).asInstanceOf[java.lang.Long].longValue()
+ }
+ (rows, bytes)
+ }
+
+ /**
+ * Stamp `sortOrderId` on every `DataFile` in `dataFiles`. iceberg-rust's
writer doesn't expose
+ * the field (it's `pub(crate)` on `DataFile`), so the manifest comes back
with sort_order_id
+ * unset. iceberg-java's `BaseFile` declares a private `sortOrderId:
Integer` field that the
+ * normal `SparkWrite` path populates from the table's `outputSortOrderId`;
we mirror that here
+ * by writing the same value via reflection before handing files to the
committer.
+ */
+ def stampSortOrderId(dataFiles: java.util.List[_], sortOrderId: Int): Unit =
{
Review Comment:
This is a step beyond the rest of the reflection in this file, which reads
fields and calls package-private constructors but does not mutate iceberg-java
state. An uncaught `NoSuchFieldException` here would also be a task failure
rather than a fallback, unlike `newDataManifestFile` which soft-fails on
exactly that.
`DataFiles.Builder.withSortOrder(SortOrder)` is public, and
`rebuildDataFilesWithJavaMetrics` already builds through that builder
immediately afterwards. Could the sort order be applied there instead, by
resolving the `SortOrder` on the driver (it is `Serializable`) and shipping it
into the task closure alongside `metricsConfig`? That would drop the
private-field write entirely. If you would rather keep the current shape, could
it at least soft-fail on `NoSuchFieldException` for consistency?
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -0,0 +1,1111 @@
+// 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.
+
+//! Native Iceberg write operator using iceberg-rust.
+//!
+//! Drains the upstream Arrow stream through iceberg-rust's writer stack
+//! (`ParquetWriterBuilder` -> `RollingFileWriterBuilder` ->
`DataFileWriterBuilder`
+//! -> `Unpartitioned`/`Fanout`/`Clustered`Writer) and emits a single-row,
single-column
+//! Arrow batch carrying the `Vec<DataFile>` produced for the task, packed as
an Iceberg V2
+//! data manifest via iceberg-rust's `ManifestWriter` against an in-memory
`FileIO`. The JVM
+//! decodes the bytes with `ManifestFiles.read(...)` to recover the
`DataFile`s for commit.
+
+use std::fmt;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, BinaryArray, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef};
+use datafusion::error::{DataFusionError, Result as DFResult};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
Partitioning,
+ PlanProperties, SendableRecordBatchStream,
+};
+use futures::TryStreamExt;
+use iceberg::arrow::{
+ arrow_struct_to_literal, PartitionValueCalculator,
RecordBatchPartitionSplitter,
+};
+use iceberg::spec::{
+ DataFile, DataFileFormat, Literal, ManifestWriterBuilder, PartitionSpec,
PartitionSpecRef,
+ Schema as IcebergSchema, SchemaRef as IcebergSchemaRef, Struct as
IcebergStruct, StructType,
+};
+use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
+use iceberg::writer::file_writer::location_generator::{
+ DefaultFileNameGenerator, DefaultLocationGenerator,
+};
+use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
+use iceberg::writer::file_writer::ParquetWriterBuilder;
+use iceberg::writer::partitioning::clustered_writer::ClusteredWriter;
+use iceberg::writer::partitioning::fanout_writer::FanoutWriter;
+use iceberg::writer::partitioning::unpartitioned_writer::UnpartitionedWriter;
+#[cfg(test)]
+use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
+use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
+use parquet::file::properties::{EnabledStatistics, WriterProperties};
+
+use datafusion_comet_proto::spark_operator::{
+ CompressionCodec as ProtoCompressionCodec, IcebergParquetWriteSettings,
IcebergWrite,
+ IcebergWriteCommon, IcebergWriterMode as ProtoIcebergWriterMode,
+};
+
+use crate::cloud::s3::credential_bridge::AccessMode;
+use crate::execution::operators::iceberg_common::load_file_io;
+
+/// Builder chain instantiated once per task and handed to the partitioning
wrapper.
+type IcebergDataFileWriterBuilder =
+ DataFileWriterBuilder<ParquetWriterBuilder, DefaultLocationGenerator,
DefaultFileNameGenerator>;
+
+/// Native Iceberg write operator. Owns the parsed Iceberg schema/spec and the
parquet writer
+/// properties; at task execution it builds the iceberg-rust writer stack,
drains the upstream
+/// Arrow stream into it, and emits a single Avro-encoded `Vec<DataFile>` row.
+pub struct IcebergWriteExec {
+ input: Arc<dyn ExecutionPlan>,
+ common: Arc<IcebergWriteCommon>,
+ iceberg_schema: IcebergSchemaRef,
+ partition_spec: PartitionSpecRef,
+ writer_mode: ProtoIcebergWriterMode,
+ writer_properties: Arc<WriterProperties>,
+ partition_id: Option<i32>,
+ task_attempt_id: Option<i64>,
+ output_schema: SchemaRef,
+ plan_properties: Arc<PlanProperties>,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl IcebergWriteExec {
+ pub fn try_new(input: Arc<dyn ExecutionPlan>, proto: IcebergWrite) ->
DFResult<Self> {
+ let IcebergWrite {
+ common,
+ partition_id,
+ task_attempt_id,
+ } = proto;
+ let common = common.ok_or_else(|| {
+ DataFusionError::Internal("IcebergWrite missing common
payload".into())
+ })?;
+ let settings = common.parquet_settings.as_ref().ok_or_else(|| {
+ DataFusionError::Internal("IcebergWriteCommon missing
parquet_settings".into())
+ })?;
+ let writer_properties = build_writer_properties(settings)?;
+ let iceberg_schema =
parse_iceberg_schema(&common.iceberg_schema_json)?;
+ let partition_spec =
parse_partition_spec(&common.partition_spec_json)?;
+ let writer_mode =
ProtoIcebergWriterMode::try_from(common.writer_mode).map_err(|_| {
+ DataFusionError::Internal(format!(
+ "Unknown IcebergWriterMode proto value: {}",
+ common.writer_mode
+ ))
+ })?;
+ let output_schema = build_output_schema();
+ let plan_properties = Self::compute_properties(&input,
Arc::clone(&output_schema));
+ Ok(Self {
+ input,
+ common: Arc::new(common),
+ iceberg_schema,
+ partition_spec,
+ writer_mode,
+ writer_properties: Arc::new(writer_properties),
+ partition_id,
+ task_attempt_id,
+ output_schema,
+ plan_properties,
+ metrics: ExecutionPlanMetricsSet::new(),
+ })
+ }
+
+ fn compute_properties(
+ input: &Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ ) -> Arc<PlanProperties> {
+ Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(schema),
+
Partitioning::UnknownPartitioning(input.output_partitioning().partition_count()),
+ EmissionType::Final,
+ Boundedness::Bounded,
+ ))
+ }
+}
+
+impl ExecutionPlan for IcebergWriteExec {
+ fn name(&self) -> &str {
+ "IcebergWriteExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.output_schema)
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.plan_properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.input]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> DFResult<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return Err(DataFusionError::Internal(
+ "IcebergWriteExec requires exactly one child".into(),
+ ));
+ }
+ Ok(Arc::new(Self {
+ input: children.pop().unwrap(),
+ common: Arc::clone(&self.common),
+ iceberg_schema: Arc::clone(&self.iceberg_schema),
+ partition_spec: Arc::clone(&self.partition_spec),
+ writer_mode: self.writer_mode,
+ writer_properties: Arc::clone(&self.writer_properties),
+ partition_id: self.partition_id,
+ task_attempt_id: self.task_attempt_id,
+ output_schema: Arc::clone(&self.output_schema),
+ plan_properties: Arc::clone(&self.plan_properties),
+ metrics: self.metrics.clone(),
+ }))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> DFResult<SendableRecordBatchStream> {
+ let input_stream = self.input.execute(partition, context)?;
+ let common = Arc::clone(&self.common);
+ let iceberg_schema = Arc::clone(&self.iceberg_schema);
+ let partition_spec = Arc::clone(&self.partition_spec);
+ let writer_mode = self.writer_mode;
+ let writer_properties = Arc::clone(&self.writer_properties);
+ let partition_id = self.partition_id;
+ let task_attempt_id = self.task_attempt_id;
+ let output_schema = Arc::clone(&self.output_schema);
+
+ let task = async move {
+ let data_files = run_write_task(
+ input_stream,
+ Arc::clone(&common),
+ Arc::clone(&iceberg_schema),
+ Arc::clone(&partition_spec),
+ writer_mode,
+ writer_properties.as_ref().clone(),
+ partition_id,
+ task_attempt_id,
+ )
+ .await?;
+ let manifest_bytes = encode_data_files_as_manifest(
+ data_files,
+ iceberg_schema,
+ partition_spec,
+ partition_id,
+ task_attempt_id,
+ &common.operation_id,
+ )
+ .await?;
+ let batch = build_output_batch(manifest_bytes, &output_schema)?;
+ Ok::<_, DataFusionError>(futures::stream::iter(vec![Ok(batch)]))
+ };
+
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&self.output_schema),
+ futures::stream::once(task).try_flatten(),
+ )))
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+}
+
+impl fmt::Debug for IcebergWriteExec {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("IcebergWriteExec")
+ .field("metadata_location", &self.common.metadata_location)
+ .field("data_location", &self.common.data_location)
+ .field("operation_id", &self.common.operation_id)
+ .field("writer_mode", &self.writer_mode)
+ .field("partition_id", &self.partition_id)
+ .field("task_attempt_id", &self.task_attempt_id)
+ .finish()
+ }
+}
+
+impl DisplayAs for IcebergWriteExec {
+ fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) ->
fmt::Result {
+ write!(
+ f,
+ "IcebergWriteExec: metadata_location={}, data_location={},
operation_id={}",
+ self.common.metadata_location, self.common.data_location,
self.common.operation_id
+ )
+ }
+}
+
+/// One-shot per-task write coroutine. Builds the iceberg-rust writer stack,
decorates each input
+/// batch with `PARQUET_FIELD_ID_META_KEY` metadata so iceberg-rust can match
Arrow columns to
+/// Iceberg field IDs, and routes through
`UnpartitionedWriter`/`FanoutWriter`/`ClusteredWriter`
+/// depending on `writer_mode`.
+#[allow(clippy::too_many_arguments)]
+async fn run_write_task(
+ mut input: SendableRecordBatchStream,
+ common: Arc<IcebergWriteCommon>,
+ iceberg_schema: IcebergSchemaRef,
+ partition_spec: PartitionSpecRef,
+ writer_mode: ProtoIcebergWriterMode,
+ writer_properties: WriterProperties,
+ partition_id: Option<i32>,
+ task_attempt_id: Option<i64>,
+) -> DFResult<Vec<DataFile>> {
+ // The JVM exec wrapper stamps both ids per task; a missing id means the
plan template was
+ // executed directly, and defaulting would make every task collide on the
same file names.
+ let partition_id = partition_id.ok_or_else(|| {
+ DataFusionError::Internal("IcebergWrite executed without a
partition_id".into())
+ })?;
+ let task_attempt_id = task_attempt_id.ok_or_else(|| {
+ DataFusionError::Internal("IcebergWrite executed without a
task_attempt_id".into())
+ })?;
+ let catalog_properties = common
+ .catalog_properties
+ .iter()
+ .map(|(k, v)| (k.clone(), v.clone()))
+ .collect();
+ let file_io = load_file_io(
+ &catalog_properties,
+ &common.data_location,
+ &common.catalog_name,
+ AccessMode::Write,
+ )?;
+
+ let location_generator =
+
DefaultLocationGenerator::with_data_location(common.data_location.clone());
+ let file_name_generator = DefaultFileNameGenerator::new(
+ file_name_prefix(partition_id, task_attempt_id, &common.operation_id),
+ None,
+ DataFileFormat::Parquet,
+ );
+ let parquet_builder = ParquetWriterBuilder::new(writer_properties,
Arc::clone(&iceberg_schema));
+ let rolling_builder = RollingFileWriterBuilder::new(
+ parquet_builder,
+ common.target_file_size_bytes as usize,
+ file_io,
+ location_generator,
+ file_name_generator,
+ );
+ let data_file_builder = DataFileWriterBuilder::new(rolling_builder);
+
+ let unpartitioned = partition_spec.is_unpartitioned();
+ let mut writer = match (unpartitioned, writer_mode) {
+ (true, ProtoIcebergWriterMode::IcebergWriterUnpartitioned) => {
+
InnerWriter::Unpartitioned(UnpartitionedWriter::new(data_file_builder))
+ }
+ (false, ProtoIcebergWriterMode::IcebergWriterFanout) => {
+ InnerWriter::Fanout(FanoutWriter::new(data_file_builder))
+ }
+ (false, ProtoIcebergWriterMode::IcebergWriterClustered) => {
+ InnerWriter::Clustered(ClusteredWriter::new(data_file_builder))
+ }
+ (actual, mode) => {
+ return Err(DataFusionError::Internal(format!(
+ "IcebergWrite writer_mode {mode:?} is inconsistent with the
partition spec \
+ (unpartitioned={actual})"
+ )))
+ }
+ };
+
+ // `RecordBatchPartitionSplitter::split` groups rows through a HashMap and
emits the parts in
+ // unspecified order. `ClusteredWriter` hard-errors when a closed
partition is revisited, so
+ // the parts of every batch must be written in the batch's own
(partition-sorted) order --
+ // recover it by computing each row's partition value and recording first
occurrences.
+ let clustered_order = match &writer {
+ InnerWriter::Clustered(_) => Some((
+ PartitionValueCalculator::try_new(&partition_spec, &iceberg_schema)
+ .map_err(iceberg_err)?,
+ partition_spec
+ .partition_type(&iceberg_schema)
+ .map_err(iceberg_err)?,
+ )),
+ _ => None,
+ };
+
+ let splitter = if unpartitioned {
+ None
+ } else {
+ Some(
+ RecordBatchPartitionSplitter::try_new_with_computed_values(
+ Arc::clone(&iceberg_schema),
+ Arc::clone(&partition_spec),
+ )
+ .map_err(iceberg_err)?,
+ )
+ };
+
+ // Build the field-id-decorated target schema once per task; every batch
is cast against it.
+ let target_schema =
+
Arc::new(iceberg::arrow::schema_to_arrow_schema(&iceberg_schema).map_err(iceberg_err)?);
+ while let Some(batch) = input.try_next().await? {
+ let decorated = decorate_batch_with_field_ids(batch, &target_schema)?;
+ writer
+ .write(decorated, splitter.as_ref(), clustered_order.as_ref())
+ .await?;
+ }
+ writer.close().await
+}
+
+/// Enum-based dispatch over the three iceberg-rust partitioning writers. Each
variant takes the
+/// same builder chain so we can keep the type fixed.
+enum InnerWriter {
+ Unpartitioned(UnpartitionedWriter<IcebergDataFileWriterBuilder>),
+ Fanout(FanoutWriter<IcebergDataFileWriterBuilder>),
+ Clustered(ClusteredWriter<IcebergDataFileWriterBuilder>),
+}
+
+impl InnerWriter {
+ async fn write(
+ &mut self,
+ batch: RecordBatch,
+ splitter: Option<&RecordBatchPartitionSplitter>,
+ clustered_order: Option<&(PartitionValueCalculator, StructType)>,
+ ) -> DFResult<()> {
+ use iceberg::writer::partitioning::PartitioningWriter;
+ match self {
+ InnerWriter::Unpartitioned(w) =>
w.write(batch).await.map_err(iceberg_err),
+ InnerWriter::Fanout(w) => {
+ let parts = splitter
+ .expect("partition splitter must be Some for partitioned
writes")
+ .split(&batch)
+ .map_err(iceberg_err)?;
+ for (key, part) in parts {
+ w.write(key, part).await.map_err(iceberg_err)?;
+ }
+ Ok(())
+ }
+ InnerWriter::Clustered(w) => {
+ let mut parts = splitter
+ .expect("partition splitter must be Some for partitioned
writes")
+ .split(&batch)
+ .map_err(iceberg_err)?;
+ let (calculator, partition_type) = clustered_order
+ .expect("clustered order helper must be Some for clustered
writes");
+ let order = partition_first_occurrence_order(&batch,
calculator, partition_type)?;
+ parts.sort_by_key(|(key, _)|
order.get(key.data()).copied().unwrap_or(usize::MAX));
+ for (key, part) in parts {
+ w.write(key, part).await.map_err(iceberg_err)?;
+ }
+ Ok(())
+ }
+ }
+ }
+
+ async fn close(self) -> DFResult<Vec<DataFile>> {
+ use iceberg::writer::partitioning::PartitioningWriter;
+ match self {
+ InnerWriter::Unpartitioned(w) =>
w.close().await.map_err(iceberg_err),
+ InnerWriter::Fanout(w) => w.close().await.map_err(iceberg_err),
+ InnerWriter::Clustered(w) => w.close().await.map_err(iceberg_err),
+ }
+ }
+}
+
+// --- helpers -------------------------------------------------------------
+
+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}"))
+ })?;
+ Ok(Arc::new(schema))
+}
+
+fn parse_partition_spec(json: &str) -> DFResult<PartitionSpecRef> {
+ let spec: PartitionSpec = serde_json::from_str(json).map_err(|e| {
+ DataFusionError::Internal(format!("Failed to parse partition spec
JSON: {e}"))
+ })?;
+ Ok(Arc::new(spec))
+}
+
+fn iceberg_err(e: iceberg::Error) -> DataFusionError {
+ DataFusionError::External(Box::new(e))
+}
+
+fn build_output_schema() -> SchemaRef {
+ Arc::new(ArrowSchema::new(vec![Field::new(
+ "iceberg_manifest",
+ DataType::Binary,
+ false,
+ )]))
+}
+
+/// Align an input batch with the field-id-decorated target schema by casting
each column. The
+/// caller is responsible for building `target_schema` once per task via
+/// `iceberg::arrow::schema_to_arrow_schema` — it carries
`PARQUET_FIELD_ID_META_KEY` on every
+/// nested field, and `arrow::compute::cast` rebuilds the column structure to
match while
+/// reusing data buffers. This is the same conformance step the iceberg-rust
DataFusion
+/// integration gets for free from DataFusion's `INSERT INTO` planner.
+fn decorate_batch_with_field_ids(
+ batch: RecordBatch,
+ target_schema: &SchemaRef,
+) -> DFResult<RecordBatch> {
+ if batch.num_columns() != target_schema.fields().len() {
+ return Err(DataFusionError::Plan(format!(
+ "Iceberg write column count mismatch: arrow batch has {} columns
but schema has {}",
+ batch.num_columns(),
+ target_schema.fields().len()
+ )));
+ }
+ // safe:false so a lossy type divergence fails the task instead of writing
silent NULLs.
+ let cast_options = arrow::compute::CastOptions {
+ safe: false,
+ ..Default::default()
+ };
+ let casted: Vec<ArrayRef> = batch
+ .columns()
+ .iter()
+ .zip(target_schema.fields().iter())
+ .map(|(col, target)| {
+ arrow::compute::cast_with_options(col, target.data_type(),
&cast_options)
+ })
+ .collect::<Result<_, _>>()
+ .map_err(DataFusionError::from)?;
+ RecordBatch::try_new(Arc::clone(target_schema),
casted).map_err(DataFusionError::from)
+}
+
+fn file_name_prefix(partition_id: i32, task_attempt_id: i64, operation_id:
&str) -> String {
+ format!("{partition_id:05}-{task_attempt_id:05}-{operation_id}")
+}
+
+/// First-occurrence rank of each distinct partition value in `batch`, used to
restore the
+/// batch's own partition order after `RecordBatchPartitionSplitter::split`'s
HashMap grouping.
+fn partition_first_occurrence_order(
+ batch: &RecordBatch,
+ calculator: &PartitionValueCalculator,
+ partition_type: &StructType,
+) -> DFResult<std::collections::HashMap<IcebergStruct, usize>> {
+ let partition_array = calculator.calculate(batch).map_err(iceberg_err)?;
+ let literals =
+ arrow_struct_to_literal(&partition_array,
partition_type).map_err(iceberg_err)?;
+ let mut order = std::collections::HashMap::new();
+ for literal in literals {
+ match literal {
+ Some(Literal::Struct(value)) => {
+ let rank = order.len();
+ order.entry(value).or_insert(rank);
+ }
+ other => {
+ return Err(DataFusionError::Internal(format!(
+ "partition value is not a struct literal: {other:?}"
+ )))
+ }
+ }
+ }
+ Ok(order)
+}
+
+/// Serialise the produced data files as an in-memory Iceberg V2 data
manifest, then read the
+/// manifest bytes back out. The JVM side decodes these bytes with
`ManifestFiles.read(...)` to
+/// recover the `DataFile`s.
+///
+/// The manifest entries carry a placeholder `snapshot_id` (`None` ->
`UNASSIGNED_SNAPSHOT_ID =
+/// -1`) and a placeholder `sequence_number` of `0`. Neither is meaningful
here: the JVM ignores
+/// the entry-level fields and only consumes the embedded `DataFile`s, which
the driver later
+/// re-stamps with the real snapshot id during `BatchWrite.commit`.
+async fn encode_data_files_as_manifest(
+ data_files: Vec<DataFile>,
+ iceberg_schema: IcebergSchemaRef,
+ partition_spec: PartitionSpecRef,
+ partition_id: Option<i32>,
+ task_attempt_id: Option<i64>,
+ operation_id: &str,
+) -> DFResult<Vec<u8>> {
+ // The manifest is assembled entirely in-process via the `memory` scheme,
so the credential
+ // dispatch key / access mode are inert here.
+ let memory_io = load_file_io(
+ &std::collections::HashMap::new(),
+ "memory:///",
Review Comment:
Small thing to confirm: does this `memory://` `FileIO` own its backing
store, so it is dropped when the function returns, or is opendal's memory
service process-global? If the store is shared, every task's manifest bytes
would stay resident for the lifetime of the executor, since nothing deletes the
entry after the `read()` below.
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -0,0 +1,1111 @@
+// 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.
+
+//! Native Iceberg write operator using iceberg-rust.
+//!
+//! Drains the upstream Arrow stream through iceberg-rust's writer stack
+//! (`ParquetWriterBuilder` -> `RollingFileWriterBuilder` ->
`DataFileWriterBuilder`
+//! -> `Unpartitioned`/`Fanout`/`Clustered`Writer) and emits a single-row,
single-column
+//! Arrow batch carrying the `Vec<DataFile>` produced for the task, packed as
an Iceberg V2
+//! data manifest via iceberg-rust's `ManifestWriter` against an in-memory
`FileIO`. The JVM
+//! decodes the bytes with `ManifestFiles.read(...)` to recover the
`DataFile`s for commit.
+
+use std::fmt;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, BinaryArray, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef};
+use datafusion::error::{DataFusionError, Result as DFResult};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
Partitioning,
+ PlanProperties, SendableRecordBatchStream,
+};
+use futures::TryStreamExt;
+use iceberg::arrow::{
+ arrow_struct_to_literal, PartitionValueCalculator,
RecordBatchPartitionSplitter,
+};
+use iceberg::spec::{
+ DataFile, DataFileFormat, Literal, ManifestWriterBuilder, PartitionSpec,
PartitionSpecRef,
+ Schema as IcebergSchema, SchemaRef as IcebergSchemaRef, Struct as
IcebergStruct, StructType,
+};
+use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
+use iceberg::writer::file_writer::location_generator::{
+ DefaultFileNameGenerator, DefaultLocationGenerator,
+};
+use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
+use iceberg::writer::file_writer::ParquetWriterBuilder;
+use iceberg::writer::partitioning::clustered_writer::ClusteredWriter;
+use iceberg::writer::partitioning::fanout_writer::FanoutWriter;
+use iceberg::writer::partitioning::unpartitioned_writer::UnpartitionedWriter;
+#[cfg(test)]
+use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
+use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
+use parquet::file::properties::{EnabledStatistics, WriterProperties};
+
+use datafusion_comet_proto::spark_operator::{
+ CompressionCodec as ProtoCompressionCodec, IcebergParquetWriteSettings,
IcebergWrite,
+ IcebergWriteCommon, IcebergWriterMode as ProtoIcebergWriterMode,
+};
+
+use crate::cloud::s3::credential_bridge::AccessMode;
+use crate::execution::operators::iceberg_common::load_file_io;
+
+/// Builder chain instantiated once per task and handed to the partitioning
wrapper.
+type IcebergDataFileWriterBuilder =
+ DataFileWriterBuilder<ParquetWriterBuilder, DefaultLocationGenerator,
DefaultFileNameGenerator>;
+
+/// Native Iceberg write operator. Owns the parsed Iceberg schema/spec and the
parquet writer
+/// properties; at task execution it builds the iceberg-rust writer stack,
drains the upstream
+/// Arrow stream into it, and emits a single Avro-encoded `Vec<DataFile>` row.
+pub struct IcebergWriteExec {
+ input: Arc<dyn ExecutionPlan>,
+ common: Arc<IcebergWriteCommon>,
+ iceberg_schema: IcebergSchemaRef,
+ partition_spec: PartitionSpecRef,
+ writer_mode: ProtoIcebergWriterMode,
+ writer_properties: Arc<WriterProperties>,
+ partition_id: Option<i32>,
+ task_attempt_id: Option<i64>,
+ output_schema: SchemaRef,
+ plan_properties: Arc<PlanProperties>,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl IcebergWriteExec {
+ pub fn try_new(input: Arc<dyn ExecutionPlan>, proto: IcebergWrite) ->
DFResult<Self> {
+ let IcebergWrite {
+ common,
+ partition_id,
+ task_attempt_id,
+ } = proto;
+ let common = common.ok_or_else(|| {
+ DataFusionError::Internal("IcebergWrite missing common
payload".into())
+ })?;
+ let settings = common.parquet_settings.as_ref().ok_or_else(|| {
+ DataFusionError::Internal("IcebergWriteCommon missing
parquet_settings".into())
+ })?;
+ let writer_properties = build_writer_properties(settings)?;
+ let iceberg_schema =
parse_iceberg_schema(&common.iceberg_schema_json)?;
+ let partition_spec =
parse_partition_spec(&common.partition_spec_json)?;
+ let writer_mode =
ProtoIcebergWriterMode::try_from(common.writer_mode).map_err(|_| {
+ DataFusionError::Internal(format!(
+ "Unknown IcebergWriterMode proto value: {}",
+ common.writer_mode
+ ))
+ })?;
+ let output_schema = build_output_schema();
+ let plan_properties = Self::compute_properties(&input,
Arc::clone(&output_schema));
+ Ok(Self {
+ input,
+ common: Arc::new(common),
+ iceberg_schema,
+ partition_spec,
+ writer_mode,
+ writer_properties: Arc::new(writer_properties),
+ partition_id,
+ task_attempt_id,
+ output_schema,
+ plan_properties,
+ metrics: ExecutionPlanMetricsSet::new(),
Review Comment:
I do not see anything writing to this `ExecutionPlanMetricsSet`, so
`metrics()` always returns an empty set, and `CometIcebergWriteExec` passes
`CometMetricNode(metrics, Nil)` with no child nodes either. That leaves no
native-side visibility at all for the writer.
Some counters would be valuable here: time spent inside the iceberg-rust
writer stack, rows and bytes handed to parquet-rs, files rolled. If you would
rather defer that, dropping the field for now would be clearer than keeping a
metrics set that never reports anything.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
