jordepic commented on code in PR #5361:
URL: https://github.com/apache/datafusion-comet/pull/5361#discussion_r3787171078
##########
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:
Added `native acceleration: wide primitive types keep JVM-parity values and
manifest metrics`: `BIGINT, BOOLEAN, FLOAT, DECIMAL(9,2), DECIMAL(38,10), DATE,
TIMESTAMP, TIMESTAMP_NTZ, BINARY` written under
`spark.sql.session.timeZone=America/New_York`, comparing both the full row sets
and the aggregated `readable_metrics` bounds/counts against a JVM-written twin.
On top of the twin parity there's an absolute check: the `timestamptz` value
must read back as the same zoned instant it was written as and the ntz value
must pass through untouched, so a shift that happened to hit both paths equally
would still fail. The `fixed(N)` test from the other thread rides the same
helper shape. Green on all four Spark/Iceberg 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(),
+ })
+ }
+
+ 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:
Fixed — the clustered path no longer uses `RecordBatchPartitionSplitter` at
all. A new `ClusteredBatchSplitter` runs `PartitionValueCalculator::calculate`
once per batch, walks the literals for contiguous-run boundaries (preserving
batch order by construction, so the sort and the first-occurrence map are gone
too), and emits one part per run. Input that isn't actually clustered produces
multiple runs with the same key and surfaces the same `ClusteredWriter` error
as before.
One deliberate non-optimisation, recorded in a comment on `materialize_run`:
the runs are materialised with `filter_record_batch` (what the splitter
produced before) rather than zero-copy `RecordBatch::slice`. iceberg-rust's
NaN-count visitor reads list/map children via `list_array.values()`, which
ignores a slice's offset window, so sliced list-of-float columns would
over-count NaNs in the manifest. The transform double-compute is gone; the copy
the splitter was already paying stays.
--
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]
