jordepic commented on code in PR #5361: URL: https://github.com/apache/datafusion-comet/pull/5361#discussion_r3789753797
########## native/core/src/execution/operators/iceberg_write.rs: ########## @@ -0,0 +1,1161 @@ +// 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, BooleanArray, 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, MetricBuilder, MetricsSet, Time, +}; +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, PartitionKey, 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> { + // Time spent inside the iceberg-rust writer stack (write + close), excluding time spent + // waiting on the upstream input stream. Surfaced on the JVM exec's SQL metrics by name. + let write_time = MetricBuilder::new(&self.metrics).subset_time("write_time", partition); + 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, + write_time, + ) + .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>, + write_time: Time, +) -> 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})" + ))) + } + }; + + let clustered_splitter = match &writer { + InnerWriter::Clustered(_) => Some(ClusteredBatchSplitter::try_new( + Arc::clone(&partition_spec), + Arc::clone(&iceberg_schema), + )?), + _ => None, + }; + let fanout_splitter = match &writer { + InnerWriter::Fanout(_) => Some( + RecordBatchPartitionSplitter::try_new_with_computed_values( + Arc::clone(&iceberg_schema), + Arc::clone(&partition_spec), + ) + .map_err(iceberg_err)?, + ), + _ => None, + }; + + // 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)?; + let _timer = write_time.timer(); + writer + .write( + decorated, + fanout_splitter.as_ref(), + clustered_splitter.as_ref(), + ) + .await?; + } + let _timer = write_time.timer(); + 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, + fanout_splitter: Option<&RecordBatchPartitionSplitter>, + clustered_splitter: Option<&ClusteredBatchSplitter>, + ) -> DFResult<()> { + use iceberg::writer::partitioning::PartitioningWriter; + match self { + InnerWriter::Unpartitioned(w) => w.write(batch).await.map_err(iceberg_err), + InnerWriter::Fanout(w) => { + let parts = fanout_splitter + .expect("fanout splitter must be Some for fanout 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 parts = clustered_splitter + .expect("clustered splitter must be Some for clustered writes") + .split(&batch)?; + 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}") +} + +/// Splits each batch into contiguous runs of equal partition value, in batch order. +/// +/// `RecordBatchPartitionSplitter::split` computes the partition transforms internally and groups +/// rows through a HashMap, which emits parts in unspecified order -- and `ClusteredWriter` +/// hard-errors when a closed partition is revisited, so the clustered path needs the batch's own +/// (partition-clustered) order back. Splitting on run boundaries preserves that order by +/// construction and computes the transforms exactly once. Input that is not actually clustered +/// yields multiple runs with the same key and surfaces the same `ClusteredWriter` error the +/// splitter path would have produced. +struct ClusteredBatchSplitter { + calculator: PartitionValueCalculator, + partition_type: StructType, + partition_spec: PartitionSpecRef, + schema: IcebergSchemaRef, +} + +impl ClusteredBatchSplitter { + fn try_new(partition_spec: PartitionSpecRef, schema: IcebergSchemaRef) -> DFResult<Self> { + Ok(Self { + calculator: PartitionValueCalculator::try_new(&partition_spec, &schema) + .map_err(iceberg_err)?, + partition_type: partition_spec + .partition_type(&schema) + .map_err(iceberg_err)?, + partition_spec, + schema, + }) + } + + fn split(&self, batch: &RecordBatch) -> DFResult<Vec<(PartitionKey, RecordBatch)>> { + let partition_array = self.calculator.calculate(batch).map_err(iceberg_err)?; + let literals = + arrow_struct_to_literal(&partition_array, &self.partition_type).map_err(iceberg_err)?; + let mut runs: Vec<(IcebergStruct, usize, usize)> = Vec::new(); + for (row, literal) in literals.into_iter().enumerate() { + let value = match literal { + Some(Literal::Struct(value)) => value, + other => { + return Err(DataFusionError::Internal(format!( + "partition value is not a struct literal: {other:?}" + ))) + } + }; + match runs.last_mut() { + Some((current, _, len)) if *current == value => *len += 1, + _ => runs.push((value, row, 1)), + } + } + runs.into_iter() + .map(|(value, start, len)| { + let key = PartitionKey::new( + self.partition_spec.as_ref().clone(), + Arc::clone(&self.schema), + value, + ); + Ok((key, materialize_run(batch, start, len)?)) + }) + .collect() + } +} + +/// Copy `len` rows starting at `start` out of `batch`. A zero-copy `RecordBatch::slice` would be +/// cheaper, but the parquet writer's NaN-count visitor reads list/map children via +/// `list_array.values()`, which ignores a slice's offset window -- sliced list-of-float columns +/// would over-count NaNs. `filter_record_batch` materialises compacted children (exactly what +/// `RecordBatchPartitionSplitter` produces), keeping those counts correct. +fn materialize_run(batch: &RecordBatch, start: usize, len: usize) -> DFResult<RecordBatch> { + let mut mask = vec![false; batch.num_rows()]; Review Comment: Done, exactly as you sketched: one sequential `UInt32Array` of `0..num_rows` per batch, and each run gathers through `take_record_batch(batch, indices.slice(start, len))` — O(run length) per run, O(batch rows) per batch total, and `take` compacts list/map children the same way `filter` did, so the NaN-visitor constraint holds. Single-run batches (the common case) short-circuit to a zero-copy `batch.clone()` — arrays keep their original zero offsets there, so the slice hazard doesn't apply. On the `PartitionSpec` clone: I couldn't remove it — `PartitionKey::new` takes its spec by value, and `copy_with_data` clones internally too, so one spec clone per run is the floor with the current iceberg-rust API. Left a comment saying so; it's now the only per-run cost that isn't O(run length), and specs are a handful of fields. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala: ########## @@ -354,10 +354,22 @@ object CometMetricNode { /** * Creates a [[CometMetricNode]] from a [[CometPlan]]. + * + * The `metrics` access is guarded against the one way it can fail: forcing an unmaterialised + * `metrics` lazy val on a node whose `@transient session` is `null` NPEs inside + * `SQLMetrics.createMetric(sparkContext, ...)`. That happens for a JVM-side exec constructed + * off the planning thread (e.g. `AQEShuffleReadExec` built by AQE's stage-finalisation rules + * when `SparkSession.getActiveSession` was `None`). It must NOT be pre-checked via `session`: + * this method also runs inside task closures on executors, where `session` is always `null` + * after deserialisation but `metrics` was materialised on the driver and shipped with the plan. + * A node whose metrics are unreachable contributes an empty map, but its subtree is still + * walked so every other operator keeps reporting. */ def fromCometPlan(cometPlan: SparkPlan): CometMetricNode = { - val children = cometPlan.children.map(fromCometPlan) - CometMetricNode(cometPlan.metrics, children) + val nodeMetrics = + try cometPlan.metrics + catch { case _: NullPointerException => Map.empty[String, SQLMetric] } Review Comment: Narrowed as suggested: `CometPlan` nodes read `metrics` unguarded, so a genuine NPE inside a Comet operator's metric construction fails loudly; only foreign nodes get the catch. Doc comment updated with the reasoning (Comet operators reach here on the driver with a live session or on executors with `metrics` already materialised, so neither case needs it). `CometTaskMetricsSuite` still passes on all four profiles. ########## spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala: ########## @@ -569,6 +569,675 @@ 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") { + assumeNativeAcceleration() + withIcebergCatalog { _ => + // Three nested kinds in one schema so the recursive PARQUET_FIELD_ID_META_KEY decoration + // is exercised end-to-end. Reading back via Iceberg's reader is the proof point: Iceberg + // matches columns by field id, not by name, so a row that round-trips means every nested + // field id made it into the Parquet metadata. + spark.sql(s""" + CREATE TABLE $catalog.$ns.native_complex ( + id INT, + addr STRUCT<city: STRING, zip: INT>, + tags ARRAY<STRING>, + attrs MAP<STRING, INT> + ) USING iceberg + """) + val snapshot = withNativeEnabled { + captureWrite("native_complex") { + spark.sql(""" + INSERT INTO cat.db.native_complex VALUES + (1, named_struct('city', 'NYC', 'zip', 10001), array('a', 'b'), map('k1', 1, 'k2', 2)), + (2, named_struct('city', 'SF', 'zip', 94016), array('c'), map('k3', 3)) + """) + } + } + assert(snapshot.snapshotDelta == 1L) + val nativeExecs = snapshot.plans.flatMap { p => + collectWithSubqueries(p) { case e: CometIcebergWriteExec => e } + } + assert(nativeExecs.nonEmpty, "expected native write exec in captured plans") + val rows = spark + .sql(s"SELECT id, addr.city, addr.zip, tags, attrs FROM $catalog.$ns.native_complex" + + " ORDER BY id") + .collect() + assert(rows.length == 2) + assert(rows(0).getInt(0) == 1) + assert(rows(0).getString(1) == "NYC") + assert(rows(0).getInt(2) == 10001) + assert(rows(1).getString(1) == "SF") + assert(rows(1).getInt(2) == 94016) + } + } + + // Java sources float/double manifest metrics from writer-tracked state (FloatFieldMetrics): + // NaN counted separately, bounds computed over non-NaN values, bounds dropped when every + // value is NaN. The native path must reproduce those decisions via the JVM-side metrics + // rebuild, so write identical data through both paths and pin the aggregated per-column + // manifest metrics against each other. + test("native acceleration: NaN float/double manifest metrics match the JVM writer") { + assumeNativeAcceleration() + withIcebergCatalog { _ => + Seq("nan_native", "nan_jvm").foreach { t => + spark.sql(s""" + CREATE TABLE $catalog.$ns.$t ( + id INT, + f FLOAT, + d DOUBLE, + all_nan FLOAT + ) USING iceberg + """) + } + // `f`'s zero is deliberately -0.0: parquet-rs normalises zero lower bounds to -0.0 while + // Java preserves the sign it saw (an accepted divergence), so -0.0 is the one zero where + // both paths agree bit-for-bit and Row equality below stays exact. + def insert(t: String): Unit = { + spark.sql(s""" + INSERT INTO $catalog.$ns.$t VALUES + (1, CAST('NaN' AS FLOAT), 1.5D, CAST('NaN' AS FLOAT)), + (2, 2.5, CAST('NaN' AS DOUBLE), CAST('NaN' AS FLOAT)), + (3, -0.0, CAST('NaN' AS DOUBLE), CAST('NaN' AS FLOAT)), + (4, NULL, 0.25D, NULL) + """) + } + + val snapshot = withNativeEnabled { captureWrite("nan_native")(insert("nan_native")) } + assert(snapshot.snapshotDelta == 1L) + val nativeExecs = snapshot.plans.flatMap { p => + collectWithSubqueries(p) { case e: CometIcebergWriteExec => e } + } + assert(nativeExecs.nonEmpty, "expected the NaN write to engage the native path") + insert("nan_jvm") + + // Aggregate across data files so the assertion is robust to how the insert splits into + // tasks; both paths run the identical upstream plan, so the file split is the same. + def columnMetrics(t: String): Seq[Row] = { + spark + .sql(s""" + SELECT + sum(readable_metrics.f.nan_value_count), + min(readable_metrics.f.lower_bound), + max(readable_metrics.f.upper_bound), + sum(readable_metrics.d.nan_value_count), + min(readable_metrics.d.lower_bound), + max(readable_metrics.d.upper_bound), + sum(readable_metrics.all_nan.nan_value_count), + min(readable_metrics.all_nan.lower_bound), + max(readable_metrics.all_nan.upper_bound) + FROM $catalog.$ns.$t.data_files + """) + .collect() + .toSeq + } + + val native = columnMetrics("nan_native") + val jvm = columnMetrics("nan_jvm") + assert( + native == jvm, + s"native manifest metrics ${native.mkString} != JVM manifest metrics ${jvm.mkString}") + + val row = native.head + assert(row.getLong(0) == 1L, "f has exactly one NaN") + assert(row.getFloat(1) == -0.0f && row.getFloat(2) == 2.5f, "f bounds skip NaN") + assert(row.getLong(3) == 2L, "d has exactly two NaNs") + assert(row.getLong(6) == 3L, "all_nan counts every non-null value as NaN") + assert(row.isNullAt(7) && row.isNullAt(8), "all-NaN column has no bounds") + } + } + + test("native acceleration: CTAS runs its inner append through the native writer") { + assumeNativeAcceleration() + assume(isSpark35Plus, "CTAS re-plans its inner append only on Spark 3.5+") + withIcebergCatalog { _ => + // A brand-new table has no metadata file yet, so this also pins the empty + // metadata-location path in the proto builder. + val snapshot = withNativeEnabled { + captureWrite("ctas_native") { + spark.sql(s""" + CREATE TABLE $catalog.$ns.ctas_native USING iceberg AS + SELECT * FROM VALUES (1, 'us', 1.0), (2, 'eu', 2.0) AS t(id, region, amount) + """) + } + } + assert(snapshot.snapshotDelta == 1L) + val nativeExecs = snapshot.plans.flatMap { p => + collectWithSubqueries(p) { case e: CometIcebergWriteExec => e } + } + assert(nativeExecs.nonEmpty, "expected CTAS's inner append to engage the native writer") + assertRows("ctas_native", Seq(1, 2)) + } + } + + test("native acceleration: RTAS replaces table contents through the native writer") { + assumeNativeAcceleration() + assume(isSpark35Plus, "RTAS re-plans its inner append only on Spark 3.5+") + withIcebergCatalog { warehouseDir => + createTable(warehouseDir, "rtas_native", partitionSpec = "") + coalesceInsert("rtas_native", Seq((1, "old", 1.0))) + val snapshot = withNativeEnabled { + captureWrite("rtas_native") { + spark.sql(s""" + REPLACE TABLE $catalog.$ns.rtas_native USING iceberg AS + SELECT * FROM VALUES (10, 'new', 10.0), (11, 'new', 11.0) AS t(id, region, amount) + """) + } + } + val nativeExecs = snapshot.plans.flatMap { p => + collectWithSubqueries(p) { case e: CometIcebergWriteExec => e } + } + assert(nativeExecs.nonEmpty, "expected RTAS's inner append to engage the native writer") + assertRows("rtas_native", Seq(10, 11)) + } + } + + test("native acceleration: fanout writer handles unsorted partitioned input") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable( + warehouseDir, + "native_fanout", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'write.spark.fanout.enabled'='true'")) + // Fanout writes skip Spark's partition-local sort, so partition values arrive + // interleaved -- the mode the clustered writer would reject. + assertNativeWriteEngages("native_fanout", Seq(1, 2, 3, 4)) { + spark.sql( + "INSERT INTO cat.db.native_fanout VALUES " + + "(1, 'us-east', 1.0), (2, 'eu', 2.0), (3, 'us-east', 3.0), (4, 'eu', 4.0)") + } + val byRegion = spark + .sql("SELECT region, count(*) FROM cat.db.native_fanout GROUP BY region ORDER BY region") + .collect() + .map(r => r.getString(0) -> r.getLong(1)) + .toSeq + assert(byRegion == Seq("eu" -> 2L, "us-east" -> 2L), s"got $byRegion") + val files = spark + .sql("SELECT count(*) FROM cat.db.native_fanout.data_files") + .collect() + .head + .getLong(0) + assert(files >= 2L, s"expected at least one data file per partition, got $files") + } + } + + test("native acceleration: target-file-size rolls one task across multiple files") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable( + warehouseDir, + "native_roll", + partitionSpec = "", + properties = Some("'write.target-file-size-bytes'='1'")) + val values = (1 to 300).map(i => s"($i, 'r', $i.0)").mkString(", ") + // One upstream slice + small Comet batches: the rolling writer checks the target size + // per batch, so three batches against a 1-byte target must roll into multiple files. + withSQLConf( + "spark.sql.leafNodeDefaultParallelism" -> "1", + CometConf.COMET_BATCH_SIZE.key -> "100") { + assertNativeWriteEngages("native_roll", 1 to 300) { + spark.sql(s"INSERT INTO cat.db.native_roll VALUES $values") + } + } + val fileRows = spark + .sql("SELECT record_count FROM cat.db.native_roll.data_files") + .collect() + .map(_.getLong(0)) + assert(fileRows.length >= 2, s"expected a multi-file roll, got ${fileRows.length} file(s)") + assert(fileRows.sum == 300L, s"rows across rolled files must sum to 300, got $fileRows") + } + } + + test("native acceleration: empty append commits exactly once with zero data files") { Review Comment: Added `native acceleration: empty append to a partitioned table commits with zero data files`, covering both partitioned writers: an identity-partitioned table on the default clustered path and a fanout twin (`write.spark.fanout.enabled=true`), each fed an `INSERT INTO ... SELECT` matching nothing. Both commit exactly once with zero data files through `CometIcebergWriteExec`, with the splitter never invoked. I also read both writers' `close()` at the pinned rev — `ClusteredWriter` takes the (absent) current writer and `FanoutWriter` iterates an empty map, so both return empty output rather than erroring — and the test now pins that. Green on all four profiles. -- 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]
