jordepic commented on code in PR #5361:
URL: https://github.com/apache/datafusion-comet/pull/5361#discussion_r3787172717
##########
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:
Done — `stampSortOrderId` (and its `setAccessible` field write) is gone. The
driver resolves the write's `SortOrder` from `Table.sortOrders()` (falling back
to `SortOrder.unsorted()` for id 0, which isn't always in the map), ships it in
the task closure (`SortOrder` is `Serializable`), and
`rebuildDataFilesWithJavaMetrics` applies it through the public
`DataFiles.Builder.withSortOrder` it was already building through. Verified
`withSortOrder(SortOrder)` is identical on 1.5.2 / 1.8.1 / 1.10.0 / 1.11.0. The
existing sort_order_id JVM-parity test still passes on all four profiles.
##########
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:
Wired up. The native operator now reports `write_time` — time inside the
iceberg-rust writer stack (`write` + `close`), excluding waiting on upstream —
which flows through the root `CometMetricNode` by name into a new nano-timing
SQL metric ("time in native Iceberg writer") on `CometIcebergWriteExec`. Rows /
bytes / files-rolled deliberately stay JVM-derived: they're already surfaced as
`numFiles` / `numOutputRows` / `numOutputBytes` from the decoded manifest (the
committed values, which is what the stock Spark UI row shows), so native
counters with the same meaning would just double-report.
##########
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:
Narrowed to the metrics access itself. `fromCometPlan` now recurses through
every node exactly as upstream does; the only remaining change is that a node
whose `session` is `null` (the actual NPE trigger — `metrics` is a lazy val
that calls `SQLMetrics.createMetric(sparkContext, ...)`) contributes an empty
metric map while its subtree is still walked. Any node with a live session —
Comet or not — reports exactly as before, so no existing operator's metrics can
regress, and the null-session node's children are no longer pruned either (they
were under the old guard).
##########
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:
It owns its store. opendal's memory service creates a fresh `MemoryCore {
data: Arc<Mutex<BTreeMap>> }` inside `Builder::build()` — there's no
process-global registry — and `load_file_io` builds a new `FileIO` per call, so
the manifest bytes are freed when the function returns and the
`FileIO`/operator drops. Added a sentence to the comment at the call site so
the next reader doesn't have to re-derive it.
--
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]
