andygrove commented on code in PR #5361:
URL: https://github.com/apache/datafusion-comet/pull/5361#discussion_r3789709537
##########
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:
This is much better than pruning the subtree, and the note about why the
`session` pre-check does not work is worth having. My remaining worry is that
the catch is wider than the cause you diagnosed. It will also swallow an NPE
thrown from inside any operator's own `metrics` definition, so a genuine bug in
a Comet operator's metric construction turns into that operator silently
reporting nothing instead of failing loudly.
Since the failure you are guarding is specifically a foreign JVM node built
off the planning thread, would gating on that narrow the blast radius enough?
Something like only swallowing when the node is not a `CometPlan`. Comet
operators reach here either on the driver with a live session or on executors
with `metrics` already materialised, so neither case should need the guard, and
a real NPE inside a Comet operator's metrics would still surface.
##########
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:
The run-boundary split is a much better shape than what it replaced, and the
comment on why a plain `slice` is unsafe here is a good catch. But I think this
reintroduces the cost problem from the other direction.
`materialize_run` allocates a `Vec<bool>` the length of the whole batch,
bit-packs it into a `BooleanArray`, and then `filter_record_batch` has to scan
that full-length predicate to build its filter strategy. That is O(batch rows)
per run rather than O(run length), so a batch with R runs costs O(R x
batch_rows) across every column. For a low-cardinality partition that is
nothing, but a clustered write on something like `bucket(1000, id)` or an
hourly partition over a wide range can easily produce hundreds of runs in one
8192-row batch, and then this is a good deal more work than the two transform
evaluations I originally asked you to remove.
Would building one sequential `UInt32Array` of `0..num_rows` per batch and
then `take`-ing `indices.slice(start, len)` per run work here? `take` on a list
or map column gathers the referenced children into fresh compacted arrays, same
as `filter` does, so the NaN-visitor concern in your comment should still be
satisfied, and each run becomes O(run length) with the total back to O(batch
rows).
Separately, `PartitionKey::new(self.partition_spec.as_ref().clone(), ...)`
clones the whole `PartitionSpec` once per run. Small next to the above, but it
scales the same way, and it would be nice not to have both.
##########
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:
This pins the empty append on an unpartitioned table, so it exercises
`UnpartitionedWriter::close()` with nothing written. The partitioned equivalent
goes through `ClusteredWriter::close()` or `FanoutWriter::close()` having never
seen a batch, and that path is not covered anywhere. An `INSERT INTO ...
SELECT` whose filter matches nothing is an entirely ordinary thing to do
against a partitioned table, so if either of those writers errors or panics
with no current writer it would be a hard failure on a common case.
Could you add a partitioned variant here? It should be a couple of lines
given the helper already exists, and it would also cover the
`ClusteredBatchSplitter` never being invoked at all.
--
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]
