diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 8b15d829aad..77c42a3a3c6 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -421,6 +421,7 @@ jobs: org.apache.spark.sql.comet.ParquetEncryptionITCase org.apache.comet.exec.CometNativeReaderSuite org.apache.comet.CometIcebergNativeSuite + org.apache.comet.CometIcebergStreamingSuite org.apache.comet.CometIcebergEncryptionSuite org.apache.comet.CometIcebergRewriteActionSuite org.apache.comet.CometIcebergWriteActionSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 20046ff5b10..4f0c66cfddc 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -127,6 +127,7 @@ jobs: org.apache.spark.sql.comet.ParquetEncryptionITCase org.apache.comet.exec.CometNativeReaderSuite org.apache.comet.CometIcebergNativeSuite + org.apache.comet.CometIcebergStreamingSuite org.apache.comet.CometIcebergEncryptionSuite org.apache.comet.CometIcebergRewriteActionSuite org.apache.comet.CometIcebergWriteActionSuite diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index cd8ac61bf4d..de038199ad8 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -50,6 +50,113 @@ $SPARK_HOME/bin/spark-shell \ Catalog configuration is standard Iceberg-on-Spark and independent of Comet. The native reader has been tested with Hadoop, Hive, and REST catalogs. The example above uses a Hadoop catalog. For the full catalog configuration reference, see Iceberg's [Spark catalog configuration](https://iceberg.apache.org/docs/latest/spark-configuration/#catalogs). +### Micro-batch streaming reads + +On Spark 4.x, native Iceberg micro-batch source reads can be enabled experimentally with +`spark.comet.scan.icebergNative.streaming.enabled=true`. Native scans and native execution must +also be enabled. Spark 3.x retains its Spark reader because its streaming progress reporter does +not support a replacement source node. + +```python +spark.conf.set("spark.comet.scan.icebergNative.streaming.enabled", "true") +# Allow the foreachBatch DataFrame's RDDScan to enter native batch operators. +spark.conf.set("spark.comet.sparkToColumnar.enabled", "true") + +query = ( + spark.readStream + .option("streaming-max-files-per-micro-batch", "1000") + .table("catalog.db.events") + .writeStream + .option("checkpointLocation", "/checkpoints/iceberg-events") + .trigger(availableNow=True) + .foreachBatch(process_batch) + .start() +) +query.awaitTermination() +``` + +Comet reads the file tasks that Iceberg planned for each batch's start and end offsets. Spark +continues to manage admission limits, checkpoints, and sink commits. Existing +Iceberg scan compatibility checks also apply to streaming. Unsupported sources retain their Spark +reader; continuous processing is not accelerated. + +Trigger support follows the installed Iceberg runtime. Iceberg 1.10 falls back from `AvailableNow` +to a single batch, ignoring its file admission limit. Iceberg 1.11 supports `AvailableNow` with +admission limits. + +Inside `foreachBatch`, use `batch_df.sparkSession` for reference-table reads and temporary views. +Enable `spark.comet.sparkToColumnar.enabled` with `RDDScan` in its supported operator list (the +default) so the callback's DataFrame can feed native operators. These batch queries can use Comet's +existing native joins, aggregations, and Iceberg reference scans. +An Iceberg output table can be appended with `violations.writeTo("catalog.db.violations").append()`. +Spark can retry a callback, so writes still need an idempotency strategy using the batch ID. +Use a new checkpoint when migrating a source from Delta to Iceberg. + +#### Streaming execution and state + +Enable `spark.comet.exec.streaming.enabled=true` to use native operators in Spark 4.x +micro-batches. Supported streaming aggregates perform partial aggregation, merge restored state, +and compute results in Comet. Spark's state store retains its checkpoint format, commit protocol, +watermark tracking, and state eviction. The boundary converts Spark state rows to Arrow batches. +A stateful `foreachBatch` callback must consume the complete batch so Spark can commit every +state partition; returning after only `head()` or another partial action is insufficient. + +Only aggregates with compatible state buffers can use this path. Counts, ordinary numeric sums, +non-decimal averages, and supported min/max expressions are eligible. Aggregates with incompatible +buffers, such as `collect_set`, retain Spark execution. Stream-stream joins, streaming +deduplication, session-window state, and arbitrary user-defined state functions also retain Spark +execution. This option does not provide a native replacement for every Spark stateful operator. + +#### Append streams and change data capture + +The native source accelerates Iceberg's standard `readStream` path. It reads append snapshots +and does not generate change types or update/delete images. Overwrite and delete snapshots fail +by default. Iceberg's `streaming-skip-overwrite-snapshots` and `streaming-skip-delete-snapshots` +options ignore those snapshots; they do not turn the append reader into a change data feed. +See Iceberg's [streaming reads](https://iceberg.apache.org/docs/latest/spark-structured-streaming/#streaming-reads). + +For an integrity pipeline that consumes an append-only event table, the producer must supply +operation types and images. Polaris and Lakekeeper provide the REST catalog without changing +these source semantics. + +### Batch change data capture + +Enable `spark.comet.scan.icebergNative.changelog.enabled=true`, together with the native Iceberg +reader and native execution, to accelerate Iceberg's `.changes` table and +[`create_changelog_view`](https://iceberg.apache.org/docs/latest/spark-procedures/#create_changelog_view). +This experimental path reads the added and removed data-file tasks planned by Iceberg. It uses +the existing native reader and performs carry-over removal, update-image pairing, or net-change +calculation in Rust. The procedure and temporary-view registration remain in Spark. + +```sql +CALL catalog.system.create_changelog_view( + table => 'db.profiles', + changelog_view => 'profile_changes', + options => map('start-snapshot-id', '123', 'end-snapshot-id', '456'), + identifier_columns => array('tenant', 'id'), + compute_updates => true +); +SELECT * FROM profile_changes +WHERE _change_type IN ('INSERT', 'UPDATE_AFTER'); +``` + +Replace the example snapshot IDs with retained snapshots of the source table. The start bound is +exclusive and the end bound inclusive; timestamp bounds follow the installed Iceberg runtime. +The procedure always removes unchanged rows carried over by copy-on-write rewrites. Explicit +`identifier_columns` enable update images by default. With `compute_updates=true`, omitted +identifier columns come from the table schema. `net_changes=true` cancels matching inserts and +deletes across the range; Iceberg rejects combining net changes with update images. + +Results preserve `_change_type`, `_change_ordinal`, and `_commit_snapshot_id`. Binary values and +nested floating-point values retain Iceberg's JVM iterator because their external-row equality +semantics differ from Arrow value equality. Unknown procedure closure layouts also retain the +JVM iterator. The normal native-reader format and type restrictions still apply. + +This is a bounded batch API. It does not supply streaming offsets, checkpoint persistence, or an +exactly-once sink for a polling CDC job. In Iceberg 1.11, changelog planning rejects snapshots with +delete manifests, so merge-on-read changes involving delete files are unsupported before Comet +executes the scan. Snapshot history and removed data files must remain available for the range. + ### Tuning Comet’s native Iceberg reader supports fetching multiple files in parallel to hide I/O latency with the diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 692b0d1bccf..a136725f196 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -349,6 +349,7 @@ fn op_name(op: &OpStruct) -> &'static str { OpStruct::Sample(_) => "Sample", OpStruct::ContribScan(_) => "ContribScan", OpStruct::WindowGroupLimit(_) => "WindowGroupLimit", + OpStruct::IcebergChangelog(_) => "IcebergChangelog", } } diff --git a/native/core/src/execution/operators/iceberg_changelog.rs b/native/core/src/execution/operators/iceberg_changelog.rs new file mode 100644 index 00000000000..a6b57c7bfa3 --- /dev/null +++ b/native/core/src/execution/operators/iceberg_changelog.rs @@ -0,0 +1,737 @@ +// 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. + +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, Float32Array, Float64Array, LargeListArray, ListArray, MapArray, RecordBatch, + RecordBatchOptions, StringArray, StructArray, UInt64Array, +}; +use arrow::datatypes::{DataType, SchemaRef}; +use arrow::row::{RowConverter, Rows, SortField}; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use futures::TryStreamExt; + +/// Native counterpart of Iceberg's sorted changelog iterators. Input ordering and distribution +/// are supplied by create_changelog_view; state is local to one execution, not a streaming checkpoint. +#[derive(Debug)] +pub struct IcebergChangelogExec { + input: Arc, + mode: u32, + metadata: [usize; 3], + identifiers: Vec, + output_indices: Vec, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl IcebergChangelogExec { + pub fn new( + input: Arc, + mode: u32, + metadata: [usize; 3], + identifiers: Vec, + output_indices: Vec, + ) -> Result { + if mode > 2 + || (mode == 1 && identifiers.is_empty()) + || metadata[0] == metadata[1] + || metadata[0] == metadata[2] + || metadata[1] == metadata[2] + { + return Err(DataFusionError::Plan( + "Invalid Iceberg changelog mode or identifiers".into(), + )); + } + let schema = input.schema(); + if identifiers + .iter() + .chain(&output_indices) + .chain(&metadata) + .any(|&i| i >= schema.fields().len()) + { + return Err(DataFusionError::Plan( + "Iceberg changelog column index out of range".into(), + )); + } + let schema = Arc::new(schema.project(&output_indices)?); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + datafusion::physical_plan::Partitioning::UnknownPartitioning( + input.properties().partitioning.partition_count(), + ), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Self { + input, + mode, + metadata, + identifiers, + output_indices, + properties, + metrics: ExecutionPlanMetricsSet::new(), + }) + } +} + +impl DisplayAs for IcebergChangelogExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "IcebergChangelogExec: mode={}", self.mode) + } +} + +impl ExecutionPlan for IcebergChangelogExec { + fn name(&self) -> &str { + "IcebergChangelogExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn apply_expressions( + &self, + _: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new( + Arc::clone(&children[0]), + self.mode, + self.metadata, + self.identifiers.clone(), + self.output_indices.clone(), + )?)) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let processor = Processor::new( + self.input.execute(partition, Arc::clone(&context))?, + self.mode, + self.metadata, + &self.identifiers, + )?; + let size = context.session_config().batch_size(); + let output_indices = self.output_indices.clone(); + let schema = self.schema(); + let output_schema = Arc::clone(&schema); + let baseline = BaselineMetrics::new(&self.metrics, partition); + let stream = futures::stream::try_unfold(processor, move |mut processor| { + let schema = Arc::clone(&output_schema); + let indices = output_indices.clone(); + let baseline = baseline.clone(); + async move { + let mut rows = Vec::with_capacity(size); + while rows.len() < size { + match processor.next().await? { + Some(row) => { + // A sparse result must not pin thousands of full input batches. + let boundary = rows.first().is_some_and(|first: &ChangeRow| { + !Arc::ptr_eq(&first.batch, &row.batch) + }); + rows.push(row); + if boundary { + break; + } + } + None => break, + } + } + if rows.is_empty() { + return Ok(None); + } + let batch = collect_rows(&rows, schema, &indices, processor.source.change_index)?; + baseline.record_output(batch.num_rows()); + Ok(Some((batch, processor))) + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Kind { + Insert, + Delete, + Before, + After, +} +impl Kind { + fn name(self) -> &'static str { + match self { + Self::Insert => "INSERT", + Self::Delete => "DELETE", + Self::Before => "UPDATE_BEFORE", + Self::After => "UPDATE_AFTER", + } + } +} + +struct EncodedBatch { + batch: RecordBatch, + keys: Rows, + identifiers: Option, + kinds: StringArray, +} + +#[derive(Clone)] +struct ChangeRow { + batch: Arc, + index: usize, + kind: Kind, +} +impl ChangeRow { + fn same_record(&self, other: &Self) -> bool { + self.batch.keys.row(self.index) == other.batch.keys.row(other.index) + } + fn same_identifier(&self, other: &Self) -> bool { + self.batch.identifiers.as_ref().unwrap().row(self.index) + == other.batch.identifiers.as_ref().unwrap().row(other.index) + } +} + +// Iceberg compares external Spark values with Objects.equals. Maps compare independently of +// entry order and floating NaNs compare equal, while the signs of zero remain distinct. +fn comparison_values(array: &ArrayRef) -> Result { + Ok(match array.data_type() { + DataType::Float32 => Arc::new( + array + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|v| v.map(|v| if v.is_nan() { f32::NAN } else { v })) + .collect::(), + ), + DataType::Float64 => Arc::new( + array + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|v| v.map(|v| if v.is_nan() { f64::NAN } else { v })) + .collect::(), + ), + DataType::Struct(fields) => { + let values = array.as_any().downcast_ref::().unwrap(); + Arc::new(StructArray::new( + fields.clone(), + values + .columns() + .iter() + .map(comparison_values) + .collect::>>()?, + values.nulls().cloned(), + )) + } + DataType::List(field) => { + let values = array.as_any().downcast_ref::().unwrap(); + Arc::new(ListArray::new( + Arc::clone(field), + values.offsets().clone(), + comparison_values(values.values())?, + values.nulls().cloned(), + )) + } + DataType::LargeList(field) => { + let values = array.as_any().downcast_ref::().unwrap(); + Arc::new(LargeListArray::new( + Arc::clone(field), + values.offsets().clone(), + comparison_values(values.values())?, + values.nulls().cloned(), + )) + } + DataType::Map(field, sorted) => { + let map = array.as_any().downcast_ref::().unwrap(); + let entries = StructArray::new( + map.entries().fields().clone(), + map.entries() + .columns() + .iter() + .map(comparison_values) + .collect::>>()?, + None, + ); + let converter = + RowConverter::new(vec![SortField::new(entries.column(0).data_type().clone())])?; + let keys = converter.convert_columns(&[Arc::clone(entries.column(0))])?; + let mut indices = (0..entries.len()).map(|i| i as u64).collect::>(); + for offsets in map.offsets().windows(2) { + indices[offsets[0] as usize..offsets[1] as usize] + .sort_by(|&a, &b| keys.row(a as usize).cmp(&keys.row(b as usize))); + } + let entries = arrow::compute::take(&entries, &UInt64Array::from(indices), None)?; + Arc::new(MapArray::new( + Arc::clone(field), + map.offsets().clone(), + entries + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + map.nulls().cloned(), + *sorted, + )) + } + _ => Arc::clone(array), + }) +} + +struct Source { + input: SendableRecordBatchStream, + current: Option>, + index: usize, + change_index: usize, + key_indices: Vec, + identifier_indices: Vec, + keys: RowConverter, + identifiers: Option, +} +impl Source { + fn new( + input: SendableRecordBatchStream, + net: bool, + metadata: [usize; 3], + identifiers: &[usize], + ) -> Result { + let schema = input.schema(); + let [change_index, ordinal, snapshot] = metadata; + let key_indices: Vec<_> = (0..schema.fields().len()) + .filter(|&i| i != change_index && (!net || (i != ordinal && i != snapshot))) + .collect(); + let converter = |indices: &[usize]| { + RowConverter::new( + indices + .iter() + .map(|&i| SortField::new(schema.field(i).data_type().clone())) + .collect(), + ) + }; + let keys = converter(&key_indices)?; + let identifiers_converter = if identifiers.is_empty() { + None + } else { + Some(converter(identifiers)?) + }; + Ok(Self { + input, + current: None, + index: 0, + change_index, + key_indices, + identifier_indices: identifiers.to_vec(), + keys, + identifiers: identifiers_converter, + }) + } + async fn next(&mut self) -> Result> { + loop { + if let Some(batch) = &self.current { + if self.index < batch.batch.num_rows() { + let index = self.index; + self.index += 1; + if batch.kinds.is_null(index) { + return Err(DataFusionError::Execution( + "Change type should not be null".into(), + )); + } + let kind = match batch.kinds.value(index) { + "INSERT" => Kind::Insert, + "DELETE" => Kind::Delete, + other => { + return Err(DataFusionError::Execution(format!( + "Unexpected Iceberg change type: {other}" + ))) + } + }; + return Ok(Some(ChangeRow { + batch: Arc::clone(batch), + index, + kind, + })); + } + } + let Some(batch) = self.input.try_next().await? else { + return Ok(None); + }; + if batch.num_rows() == 0 { + continue; + } + let columns = |indices: &[usize]| { + indices + .iter() + .map(|&i| Arc::clone(batch.column(i))) + .collect::>() + }; + let keys = self.keys.convert_columns( + &columns(&self.key_indices) + .iter() + .map(comparison_values) + .collect::>>()?, + )?; + let identifiers = self + .identifiers + .as_ref() + .map(|c| -> Result<_> { + Ok(c.convert_columns( + &columns(&self.identifier_indices) + .iter() + .map(comparison_values) + .collect::>>()?, + )?) + }) + .transpose()?; + let kinds = arrow::compute::cast(batch.column(self.change_index), &DataType::Utf8)?; + let kinds = kinds + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + self.current = Some(Arc::new(EncodedBatch { + batch, + keys, + identifiers, + kinds, + })); + self.index = 0; + } + } +} + +struct Processor { + source: Source, + mode: u32, + lookahead: Option, + repeated: Option<(ChangeRow, usize)>, + update_next: Option, +} +impl Processor { + fn new( + input: SendableRecordBatchStream, + mode: u32, + metadata: [usize; 3], + identifiers: &[usize], + ) -> Result { + Ok(Self { + source: Source::new(input, mode == 2, metadata, identifiers)?, + mode, + lookahead: None, + repeated: None, + update_next: None, + }) + } + async fn remove_carryovers(&mut self) -> Result> { + if let Some((row, count)) = &mut self.repeated { + let row = row.clone(); + *count -= 1; + if *count == 0 { + self.repeated = None; + } + return Ok(Some(row)); + } + loop { + let current = match self.lookahead.take() { + Some(row) => row, + None => match self.source.next().await? { + Some(row) => row, + None => return Ok(None), + }, + }; + if self.mode != 2 && current.kind != Kind::Delete { + return Ok(Some(current)); + } + let mut count = 1usize; + while let Some(next) = self.source.next().await? { + if !current.same_record(&next) { + self.lookahead = Some(next); + break; + } + if current.kind == next.kind { + count += 1; + } else { + count -= 1; + } + if count == 0 { + break; + } + } + if count > 0 { + if count > 1 { + self.repeated = Some((current.clone(), count - 1)); + } + return Ok(Some(current)); + } + } + } + async fn next(&mut self) -> Result> { + if self.mode != 1 { + return self.remove_carryovers().await; + } + let mut current = match self.update_next.take() { + Some(row) => row, + None => match self.remove_carryovers().await? { + Some(row) => row, + None => return Ok(None), + }, + }; + if current.kind == Kind::Delete { + if let Some(mut next) = self.remove_carryovers().await? { + if current.same_identifier(&next) { + if next.kind != Kind::Insert { + return Err(DataFusionError::Execution("Cannot compute updates because there are multiple rows with the same identifier fields. Please make sure the rows are unique.".into())); + } + current.kind = Kind::Before; + next.kind = Kind::After; + } + self.update_next = Some(next); + } + } + Ok(Some(current)) + } +} + +fn collect_rows( + rows: &[ChangeRow], + schema: SchemaRef, + output_indices: &[usize], + change_index: usize, +) -> Result { + let mut batches = Vec::new(); + let mut batch_indices = HashMap::new(); + let indices = rows + .iter() + .map(|row| { + let next_index = batches.len(); + let batch_index = *batch_indices + .entry(Arc::as_ptr(&row.batch) as usize) + .or_insert_with(|| { + batches.push(Arc::clone(&row.batch)); + next_index + }); + (batch_index, row.index) + }) + .collect::>(); + let columns = output_indices + .iter() + .enumerate() + .map(|(output, &input)| -> Result { + if input == change_index { + let values = StringArray::from_iter_values(rows.iter().map(|row| row.kind.name())); + Ok(arrow::compute::cast( + &values, + schema.field(output).data_type(), + )?) + } else { + let arrays = batches + .iter() + .map(|b| b.batch.column(input).as_ref()) + .collect::>(); + Ok(arrow::compute::interleave(&arrays, &indices)?) + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new_with_options( + schema, + columns, + &RecordBatchOptions::new().with_row_count(Some(rows.len())), + )?) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int32Array, Int64Array}; + use arrow::datatypes::{Field, Schema}; + + async fn run(mode: u32, records: &[(&str, &str, i32)]) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + Field::new("_change_type", DataType::Utf8, false), + Field::new("_change_ordinal", DataType::Int32, false), + Field::new("_commit_snapshot_id", DataType::Int64, false), + ])); + // Every row is a separate Arrow batch; pending deletes/updates must survive each boundary. + let batches = records + .iter() + .map(|(value, kind, ordinal)| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec![*value])), + Arc::new(StringArray::from(vec![*kind])), + Arc::new(Int32Array::from(vec![*ordinal])), + Arc::new(Int64Array::from(vec![100 + i64::from(*ordinal)])), + ], + ) + .map_err(DataFusionError::from) + }) + .collect::>(); + let input = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(batches), + )); + let ids = if mode == 1 { vec![0, 3] } else { vec![] }; + let mut processor = Processor::new(input, mode, [2, 3, 4], &ids)?; + let mut rows = Vec::new(); + while let Some(row) = processor.next().await? { + rows.push(row); + } + if rows.is_empty() { + return Ok(vec![]); + } + let batch = collect_rows(&rows, schema, &[0, 1, 2, 3, 4], 2)?; + let value = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let kind = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let ordinal = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + Ok((0..batch.num_rows()) + .map(|i| { + ( + value.value(i).to_owned(), + kind.value(i).to_owned(), + ordinal.value(i), + ) + }) + .collect()) + } + + #[test] + fn comparison_matches_map_equality_and_boxed_float_equality() { + use arrow::array::{Int32Builder, MapBuilder, StringBuilder}; + let mut maps = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + for entries in [[("a", 1), ("b", 2)], [("b", 2), ("a", 1)]] { + for (key, value) in entries { + maps.keys().append_value(key); + maps.values().append_value(value); + } + maps.append(true).unwrap(); + } + let maps: ArrayRef = Arc::new(maps.finish()); + let converter = RowConverter::new(vec![SortField::new(maps.data_type().clone())]).unwrap(); + let keys = converter + .convert_columns(&[comparison_values(&maps).unwrap()]) + .unwrap(); + assert_eq!(keys.row(0), keys.row(1)); + let floats: ArrayRef = Arc::new(Float64Array::from(vec![ + f64::NAN, + f64::from_bits(0x7ff0000000000001), + 0.0, + -0.0, + ])); + let converter = RowConverter::new(vec![SortField::new(DataType::Float64)]).unwrap(); + let keys = converter + .convert_columns(&[comparison_values(&floats).unwrap()]) + .unwrap(); + assert_eq!(keys.row(0), keys.row(1)); + assert_ne!(keys.row(2), keys.row(3)); + } + + #[tokio::test] + async fn carryovers_and_updates_preserve_multiplicity_across_batches() { + let rows = [ + ("a", "DELETE", 1), + ("a", "DELETE", 1), + ("a", "INSERT", 1), + ("b", "INSERT", 1), + ]; + assert_eq!( + run(0, &rows).await.unwrap(), + vec![ + ("a".into(), "DELETE".into(), 1), + ("b".into(), "INSERT".into(), 1) + ] + ); + assert_eq!( + run(1, &rows).await.unwrap(), + vec![ + ("a".into(), "UPDATE_BEFORE".into(), 1), + ("b".into(), "UPDATE_AFTER".into(), 1) + ] + ); + let error = run(1, &[("a", "DELETE", 1), ("b", "DELETE", 1)]) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("multiple rows with the same identifier")); + } + + #[tokio::test] + async fn net_changes_retain_the_surviving_run_snapshot() { + assert!(run(2, &[("a", "INSERT", 0), ("a", "DELETE", 1)]) + .await + .unwrap() + .is_empty()); + assert!(run(0, &[("a", "DELETE", 1), ("a", "INSERT", 1)]) + .await + .unwrap() + .is_empty()); + assert_eq!( + run( + 2, + &[("a", "INSERT", 0), ("a", "DELETE", 1), ("a", "INSERT", 2)] + ) + .await + .unwrap(), + vec![("a".into(), "INSERT".into(), 2)] + ); + assert_eq!( + run( + 2, + &[("a", "INSERT", 0), ("a", "INSERT", 1), ("a", "DELETE", 2)] + ) + .await + .unwrap(), + vec![("a".into(), "INSERT".into(), 0)] + ); + } +} diff --git a/native/core/src/execution/operators/iceberg_scan.rs b/native/core/src/execution/operators/iceberg_scan.rs index 7c3c22ffa60..c64fd191c9b 100644 --- a/native/core/src/execution/operators/iceberg_scan.rs +++ b/native/core/src/execution/operators/iceberg_scan.rs @@ -24,9 +24,9 @@ use std::sync::Arc; use std::task::{Context, Poll}; use arrow::array::{ArrayRef, RecordBatch, RecordBatchOptions}; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{Schema, SchemaRef}; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, Result as DFResult}; +use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue}; use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::expressions::Column; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; @@ -34,9 +34,11 @@ use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, }; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, }; +use datafusion_comet_proto::spark_operator::IcebergChange; use futures::{Stream, StreamExt, TryStreamExt}; use iceberg::arrow::ScanMetrics; use iceberg::io::FileIO; @@ -81,6 +83,8 @@ pub struct IcebergScanExec { catalog_name: String, /// Pre-planned file scan tasks tasks: Vec, + /// Changelog constants aligned with the planned tasks. Absent for ordinary table scans. + changes: Option>, /// Number of data files to read concurrently data_file_concurrency_limit: usize, /// Metrics @@ -108,6 +112,7 @@ impl IcebergScanExec { catalog_properties, catalog_name, tasks, + changes: None, data_file_concurrency_limit, metrics, }) @@ -121,6 +126,20 @@ impl IcebergScanExec { Boundedness::Bounded, )) } + + pub fn with_changes(mut self, changes: Vec) -> DFResult { + if changes.len() != self.tasks.len() + || changes + .iter() + .any(|c| !matches!(c.change_type.as_str(), "INSERT" | "DELETE")) + { + return Err(DataFusionError::Plan( + "Invalid Iceberg changelog tasks".into(), + )); + } + self.changes = Some(changes); + Ok(self) + } } impl ExecutionPlan for IcebergScanExec { @@ -161,7 +180,10 @@ impl ExecutionPlan for IcebergScanExec { _partition: usize, context: Arc, ) -> DFResult { - self.execute_with_tasks(self.tasks.clone(), context) + match &self.changes { + Some(changes) => self.execute_changes(changes, context), + None => self.execute_with_tasks(self.tasks.clone(), context), + } } fn metrics(&self) -> Option { @@ -170,6 +192,93 @@ impl ExecutionPlan for IcebergScanExec { } impl IcebergScanExec { + fn execute_changes( + &self, + changes: &[IcebergChange], + context: Arc, + ) -> DFResult { + // Read only added/removed data files, reusing the ordinary reader and its concurrency. + // A group has constant CDC metadata, so no per-row JVM callback or second table scan is needed. + let mut groups = std::collections::BTreeMap::new(); + for (task, change) in self.tasks.iter().zip(changes) { + groups + .entry(( + change.change_type.clone(), + change.change_ordinal, + change.commit_snapshot_id, + )) + .or_insert_with(Vec::new) + .push(task.clone()); + } + let data_fields = self + .output_schema + .fields() + .iter() + .filter(|f| { + !matches!( + f.name().as_str(), + "_change_type" | "_change_ordinal" | "_commit_snapshot_id" + ) + }) + .cloned() + .collect::>(); + let mut scan = Self::new( + self.metadata_location.clone(), + Arc::new(Schema::new(data_fields)), + self.catalog_properties.clone(), + self.catalog_name.clone(), + vec![], + self.data_file_concurrency_limit, + ) + .map_err(|e| DataFusionError::Execution(e.to_string()))?; + scan.metrics = self.metrics.clone(); + let output_schema = Arc::clone(&self.output_schema); + // Open groups on demand. Each reader already applies the file-concurrency limit; + // overlapping groups would multiply that limit and retain extra FileIO handles. + let streams = groups.into_iter().map( + move |((change_type, ordinal, snapshot), tasks)| -> DFResult<_> { + let input = scan.execute_with_tasks(tasks, Arc::clone(&context))?; + let schema = Arc::clone(&output_schema); + Ok(input.map(move |batch| { + let batch = batch?; + let columns = schema + .fields() + .iter() + .map(|field| { + let constant = match field.name().as_str() { + "_change_type" => { + Some(ScalarValue::Utf8(Some(change_type.clone()))) + } + "_change_ordinal" => Some(ScalarValue::Int32(Some(ordinal))), + "_commit_snapshot_id" => Some(ScalarValue::Int64(Some(snapshot))), + _ => None, + }; + match constant { + Some(value) => Ok(arrow::compute::cast( + &value.to_array_of_size(batch.num_rows())?, + field.data_type(), + )?), + None => Ok(Arc::clone( + batch.column(batch.schema().index_of(field.name())?), + )), + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new_with_options( + Arc::clone(&schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + )?) + })) + }, + ); + let stream = futures::stream::iter(streams).try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.output_schema), + stream, + ))) + } + /// Handles MOR (Merge-On-Read) tables by automatically applying positional and equality /// deletes via iceberg-rust's ArrowReader. fn execute_with_tasks( diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index d09b0b4fb37..2c22778f797 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -34,9 +34,11 @@ mod expand; pub use expand::ExpandExec; mod explode; pub use explode::ExplodeExec; +mod iceberg_changelog; mod iceberg_common; mod iceberg_partition_path; mod iceberg_scan; +pub use iceberg_changelog::IcebergChangelogExec; mod iceberg_write; pub use iceberg_write::IcebergWriteExec; mod parquet_writer; diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 60c206617e3..1c99059077d 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -35,8 +35,8 @@ mod lance_scan; use crate::execution::operators::init_csv_datasource_exec; use crate::execution::operators::AlignedArrowStreamReader; use crate::execution::operators::DynamicFilterJoinExec; -use crate::execution::operators::IcebergScanExec; use crate::execution::operators::IcebergWriteExec; +use crate::execution::operators::{IcebergChangelogExec, IcebergScanExec}; use crate::execution::operators::{PartitionedRankLimitExec, WindowFnKind}; use crate::execution::{ expressions::list_empty_to_null::ListEmptyToNullExpr, @@ -1823,6 +1823,40 @@ impl PhysicalPlanner { Arc::new(SparkPlan::new(spark_plan.plan_id, Arc::new(scan), vec![])), )) } + OpStruct::IcebergChangelog(changelog) => { + assert_eq!(children.len(), 1); + let (scans, shuffle_scans, child) = + self.create_plan(&children[0], inputs, partition_count)?; + let input = Arc::clone(&child.native_plan); + let plan = IcebergChangelogExec::new( + input, + changelog.mode, + [ + changelog.change_type_index as usize, + changelog.change_ordinal_index as usize, + changelog.commit_snapshot_id_index as usize, + ], + changelog + .identifier_indices + .iter() + .map(|&i| i as usize) + .collect(), + changelog + .output_indices + .iter() + .map(|&i| i as usize) + .collect(), + )?; + Ok(( + scans, + shuffle_scans, + Arc::new(SparkPlan::new( + spark_plan.plan_id, + Arc::new(plan), + vec![child], + )), + )) + } OpStruct::IcebergScan(scan) => { // Extract common data and single partition's file tasks // Per-partition injection happens in Scala before sending to native @@ -1855,6 +1889,25 @@ impl PhysicalPlanner { data_file_concurrency_limit, )?; + let iceberg_scan = if scan + .file_scan_tasks + .iter() + .any(|task| task.change.is_some()) + { + let changes = scan + .file_scan_tasks + .iter() + .map(|task| { + task.change.clone().ok_or_else(|| { + GeneralError("Mixed Iceberg data and changelog tasks".into()) + }) + }) + .collect::, _>>()?; + iceberg_scan.with_changes(changes)? + } else { + iceberg_scan + }; + Ok(( vec![], vec![], diff --git a/native/core/src/execution/planner/operator_registry.rs b/native/core/src/execution/planner/operator_registry.rs index face6c3eb5e..d48c36c0b0a 100644 --- a/native/core/src/execution/planner/operator_registry.rs +++ b/native/core/src/execution/planner/operator_registry.rs @@ -162,5 +162,6 @@ fn get_operator_type(spark_operator: &Operator) -> Option { // so the supports-mixed-codegen check skips it. OpStruct::ContribScan(_) => None, OpStruct::WindowGroupLimit(_) => None, // Not yet in OperatorType enum + OpStruct::IcebergChangelog(_) => None, } } diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 819806213df..5b5c6e5a770 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -68,6 +68,7 @@ message Operator { Sample sample = 118; WindowGroupLimit window_group_limit = 119; IcebergWrite iceberg_write = 120; + IcebergChangelog iceberg_changelog = 121; // Extension point for optional, out-of-tree contrib scans (Delta, Lance, ...). The concrete // scan message (e.g. `DeltaScan`) is packed into this envelope on the JVM side and dispatched // by `type_url` on the native side. Using a single permanent field -- rather than a new oneof @@ -421,6 +422,14 @@ message IcebergFileScanTask { optional uint32 partition_data_idx = 20; optional uint32 delete_files_idx = 21; optional uint32 residual_idx = 22; + // Constants for rows from a changed data file, planned by Iceberg's changelog scan. + optional IcebergChange change = 23; +} + +message IcebergChange { + string change_type = 1; + int32 change_ordinal = 2; + int64 commit_snapshot_id = 3; } // Iceberg delete file for MOR tables (positional or equality deletes) @@ -1059,3 +1068,14 @@ enum RankLikeFunction { Rank = 1; DenseRank = 2; } + +// Transform the sorted rows produced by Iceberg create_changelog_view. +message IcebergChangelog { + // 0: remove carryovers, 1: compute update images, 2: net changes. + uint32 mode = 1; + repeated uint32 identifier_indices = 2; + repeated uint32 output_indices = 3; + uint32 change_type_index = 4; + uint32 change_ordinal_index = 5; + uint32 commit_snapshot_id_index = 6; +} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index fc858a75a18..73f6824f343 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -121,6 +121,33 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_ICEBERG_CHANGELOG_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.scan.icebergNative.changelog.enabled") + .category(CATEGORY_SCAN) + .doc("Enable experimental native Iceberg batch changelog scans and change processing.") + .booleanConf + .createWithDefault(false) + + val COMET_ICEBERG_STREAMING_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.scan.icebergNative.streaming.enabled") + .category(CATEGORY_SCAN) + .doc( + "Whether to enable experimental native Iceberg micro-batch source reads on Spark 4.x. " + + "Requires native Iceberg scans and native execution. Spark manages streaming offsets, " + + "checkpoints, and commits. Does not support change data feed or continuous processing.") + .booleanConf + .createWithDefault(false) + + val COMET_STREAMING_EXEC_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.streaming.enabled") + .category(CATEGORY_EXEC) + .doc( + "Enable experimental native operators in Spark 4.x micro-batches, including " + + "streaming aggregates with Spark-compatible state buffers. Spark retains state " + + "persistence, watermarks, checkpoint recovery, and unsupported stateful operators.") + .booleanConf + .createWithDefault(false) + val COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.write.iceberg.splitOperator.enabled") .category(CATEGORY_TESTING) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index ba41a957288..e2bac2cf7c9 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -55,6 +55,9 @@ object IcebergReflection extends Logging { val UNBOUND_PREDICATE = "org.apache.iceberg.expressions.UnboundPredicate" val SPARK_BATCH_QUERY_SCAN = "org.apache.iceberg.spark.source.SparkBatchQueryScan" val SPARK_STAGED_SCAN = "org.apache.iceberg.spark.source.SparkStagedScan" + val SPARK_CHANGELOG_SCAN = "org.apache.iceberg.spark.source.SparkChangelogScan" + val ADDED_ROWS_SCAN_TASK = "org.apache.iceberg.AddedRowsScanTask" + val DELETED_DATA_FILE_SCAN_TASK = "org.apache.iceberg.DeletedDataFileScanTask" val SPARK_SCHEMA_UTIL = "org.apache.iceberg.spark.SparkSchemaUtil" val TABLE = "org.apache.iceberg.Table" val PARTITIONING = "org.apache.iceberg.Partitioning" @@ -81,7 +84,19 @@ object IcebergReflection extends Logging { * instances. */ val ICEBERG_SCAN_CLASSES: Set[String] = - Set(ClassNames.SPARK_BATCH_QUERY_SCAN, ClassNames.SPARK_STAGED_SCAN) + Set( + ClassNames.SPARK_BATCH_QUERY_SCAN, + ClassNames.SPARK_STAGED_SCAN, + ClassNames.SPARK_CHANGELOG_SCAN) + + def isChangelogScan(scan: Any): Boolean = + scan.getClass.getName == ClassNames.SPARK_CHANGELOG_SCAN + + private def changelogScanField(scan: Any, name: String): Option[Any] = { + val field = scan.getClass.getDeclaredField(name) + field.setAccessible(true) + Option(field.get(scan)) + } def isIcebergScanClass(name: String): Boolean = ICEBERG_SCAN_CLASSES.contains(name) @@ -412,6 +427,7 @@ object IcebergReflection extends Logging { * The table() method is protected in SparkScan, requiring reflection to access. */ def getTable(scan: Any): Option[Any] = { + if (isChangelogScan(scan)) return changelogScanField(scan, "table") findMethodInHierarchy(scan.getClass, "table").flatMap { tableMethod => try { Some(tableMethod.invoke(scan)) @@ -437,7 +453,16 @@ object IcebergReflection extends Logging { * and require reflection. */ def getTasks(scan: Any): Option[java.util.List[_]] = - if (isStagedScan(scan)) tasksFromTaskGroups(scan) else tasksFromTasksAccessor(scan) + if (isStagedScan(scan) || isChangelogScan(scan)) tasksFromTaskGroups(scan) + else tasksFromTasksAccessor(scan) + + /** Tasks already bounded by the streaming source's start/end offsets, or a batch partition. */ + def tasksFromInputPartition(partition: AnyRef): java.util.Collection[_] = { + val taskGroup = getDeclaredMethod(partition.getClass, "taskGroup").invoke(partition) + getMethod(taskGroup.getClass, "tasks") + .invoke(taskGroup) + .asInstanceOf[java.util.Collection[_]] + } private def tasksFromTasksAccessor(scan: Any): Option[java.util.List[_]] = findMethodInHierarchy(scan.getClass, "tasks") match { @@ -492,7 +517,9 @@ object IcebergReflection extends Logging { * method we know isn't there. */ def getFilterExpressions(scan: Any): Option[java.util.List[_]] = - if (isStagedScan(scan)) { + if (isChangelogScan(scan)) { + changelogScanField(scan, "filters").map(_.asInstanceOf[java.util.List[_]]) + } else if (isStagedScan(scan)) { Some(java.util.Collections.emptyList[AnyRef]()) } else { // Iceberg 1.11 renamed SparkScan.filterExpressions() to filters(); 1.8-1.10 use the old name. @@ -709,10 +736,10 @@ object IcebergReflection extends Logging { * * Returns an empty sequence when the task has no partition spec. */ - def partitionSourceFieldIds(task: Any, fileScanTaskClass: Class[_]): Seq[Int] = { + def partitionSourceFieldIds(task: Any): Seq[Int] = { val spec = try { - getMethod(fileScanTaskClass, "spec").invoke(task) + getMethod(loadClass(ClassNames.CONTENT_SCAN_TASK), "spec").invoke(task) } catch { case _: Exception => null } @@ -882,7 +909,14 @@ object IcebergReflection extends Logging { * if reflection fails (callers must handle appropriately based on context) */ def getDeleteFilesFromTask(task: Any, fileScanTaskClass: Class[_]): java.util.List[_] = { - val deletesMethod = getMethod(fileScanTaskClass, "deletes") + val deletesMethod = if (fileScanTaskClass.isInstance(task)) { + getMethod(fileScanTaskClass, "deletes") + } else if (loadClass(ClassNames.ADDED_ROWS_SCAN_TASK).isInstance(task)) { + getMethod(loadClass(ClassNames.ADDED_ROWS_SCAN_TASK), "deletes") + } else { + // Deliberately rejects other changelog task kinds, including row-level delete tasks. + getMethod(loadClass(ClassNames.DELETED_DATA_FILE_SCAN_TASK), "existingDeletes") + } val deletes = deletesMethod.invoke(task).asInstanceOf[java.util.List[_]] if (deletes == null) new java.util.ArrayList[Any]() else deletes } @@ -953,6 +987,10 @@ object IcebergReflection extends Logging { * The expected Iceberg Schema, or None if reflection fails */ def getExpectedSchema(scan: Any): Option[Any] = { + if (isChangelogScan(scan)) { + return try changelogScanField(scan, "projection") + catch { case _: NoSuchFieldException => changelogScanField(scan, "expectedSchema") } + } // Iceberg 1.11 renamed SparkScan.expectedSchema() to projection() (the projected read // schema); 1.8-1.10 still expose expectedSchema(). Try the new name first, then fall back. findMethodInHierarchy(scan.getClass, "projection") diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index f8ec2e6ab29..eb8557a8b87 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -75,6 +75,7 @@ object CometExecRule { val nativeExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = Map[Class[_ <: SparkPlan], CometOperatorSerde[_]]( classOf[ProjectExec] -> CometProjectExec, + classOf[IcebergChangelogExec] -> CometIcebergChangelogExec, classOf[FilterExec] -> CometFilterExec, classOf[LocalLimitExec] -> CometLocalLimitExec, classOf[GlobalLimitExec] -> CometGlobalLimitExec, @@ -686,9 +687,10 @@ case class CometExecRule(session: SparkSession) // We shouldn't transform Spark query plan if Comet is not loaded. if (!isCometLoaded(conf)) return plan - // Comet does not support structured streaming. Fall back to Spark for any plan that - // belongs to a streaming query (detected via StreamSourceAwareSparkPlan.getStream). - if (ShimCometStreaming.isStreamingPlan(plan)) return plan + if (ShimCometStreaming.isStreamingPlan(plan) && + !ShimCometStreaming.nativeExecutionEnabled(conf, plan)) { + return plan + } if (!CometConf.COMET_EXEC_ENABLED.get(conf)) { // Comet exec is disabled, but for Spark shuffle, we still can use Comet columnar shuffle @@ -698,7 +700,7 @@ case class CometExecRule(session: SparkSession) plan } } else { - val normalizedPlan = normalizePlan(plan) + val normalizedPlan = normalizePlan(CometIcebergChangelogExec.rewrite(plan)) val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { normalizedPlan.transformUp { case p => @@ -1027,6 +1029,12 @@ case class CometExecRule(session: SparkSession) val fallbackReasons = new ListBuffer[String]() if (CometSparkToColumnarExec.isSchemaSupported(op.schema, fallbackReasons)) { op match { + // Restore/save retain Spark's versioned state store and watermark/commit protocol. + // Convert their rows back to Arrow so the aggregate merging restored state runs natively. + case state + if ShimCometStreaming.isStateBoundary(state) && + ShimCometStreaming.nativeExecutionEnabled(conf, state) => + true // Convert Spark DS v1 scan to Arrow format case scan: FileSourceScanExec => scan.relation.fileFormat match { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index f7e0fe5cb98..6ff167335c4 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -35,10 +35,10 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, DynamicPruningExpre import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.{sideBySide, ArrayBasedMapData, GenericArrayData, MetadataColumnHelper} import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns.getExistenceDefaultValues -import org.apache.spark.sql.comet.{CometBatchScanExec, CometScanExec} +import org.apache.spark.sql.comet.{CometBatchScanExec, CometIcebergNativeScanExec, CometScanExec} import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec} import org.apache.spark.sql.execution.datasources.HadoopFsRelation -import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanExecBase} import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -77,10 +77,16 @@ case class CometScanRule(session: SparkSession) private def _apply(plan: SparkPlan): SparkPlan = { if (!isCometLoaded(conf)) return plan - // Comet does not support structured streaming. The parallel guard in - // CometExecRule only stops operator wrapping, so without this check we - // would still rewrite scans to CometScanExec in a streaming plan. - if (ShimCometStreaming.isStreamingPlan(plan)) return plan + // Only the Iceberg micro-batch source can be replaced in a streaming plan. Spark retains + // the streaming operators, offset planning, checkpointing, and sink commit protocol. + if (ShimCometStreaming.isStreamingPlan(plan)) { + return if (COMET_NATIVE_SCAN_ENABLED.get(conf) && COMET_ICEBERG_STREAMING_ENABLED.get( + conf)) { + ShimCometStreaming.transformIcebergScans(plan, scan => transformV2Scan(scan)) + } else { + plan + } + } def isSupportedScanNode(plan: SparkPlan): Boolean = plan match { case _: FileSourceScanExec => true @@ -357,13 +363,17 @@ case class CometScanRule(session: SparkSession) Some(CometScanExec(scanExec, session)) } - private def transformV2Scan(scanExec: BatchScanExec): SparkPlan = { + private def transformV2Scan(scanExec: DataSourceV2ScanExecBase): SparkPlan = { + val batchScan = scanExec match { + case batch: BatchScanExec => Some(batch) + case _ => None + } // Give any optional, out-of-tree scan contrib (e.g. Lance) first crack at this V2 scan. On a // default build no contrib is registered, so this returns None and we proceed with Comet's // built-in V2 handling below. A registered contrib either claims the scan or declines via its // own `withFallbackReason` fallback message. - CometScanContrib.tryTransformV2(scanExec) match { + batchScan.flatMap(CometScanContrib.tryTransformV2) match { case Some(handled) => return handled case None => // proceed with vanilla logic } @@ -372,7 +382,7 @@ case class CometScanRule(session: SparkSession) // a contrib's own table could legitimately end with. Running the check here -- after the // contrib hook has declined -- means a contrib that owns such a table still gets to claim it, // while the fallback for a genuine Iceberg metadata table is unchanged. - if (isIcebergMetadataTable(scanExec)) { + if (batchScan.exists(isIcebergMetadataTable)) { return withFallbackReason(scanExec, "Iceberg Metadata tables are not supported") } @@ -383,7 +393,7 @@ case class CometScanRule(session: SparkSession) // both regress that per-path support and decline a contrib's scan before it was offered. scanExec.scan match { - case scan: CSVScan if COMET_CSV_V2_NATIVE_ENABLED.get() => + case scan: CSVScan if batchScan.isDefined && COMET_CSV_V2_NATIVE_ENABLED.get() => if (scanExec.output.exists(_.isMetadataCol)) { return withFallbackReason( scanExec, @@ -424,7 +434,7 @@ case class CometScanRule(session: SparkSession) && !isInferSchemaEnabled && isSingleCharacterDelimiter) { CometBatchScanExec( scanExec.clone().asInstanceOf[BatchScanExec], - runtimeFilters = scanExec.runtimeFilters) + runtimeFilters = batchScan.get.runtimeFilters) } else { withFallbackReasons(scanExec, fallbackReasons.toSet) } @@ -433,6 +443,7 @@ case class CometScanRule(session: SparkSession) // RewriteDataFiles (and similar maintenance actions) where the planner has already // staged FileScanTasks via ScanTaskSetManager. case _ if IcebergReflection.isIcebergScanClass(scanExec.scan.getClass.getName) => + val runtimeFilters = batchScan.map(_.runtimeFilters).getOrElse(Seq.empty) val fallbackReasons = new ListBuffer[String]() // Native Iceberg scan requires both configs to be enabled @@ -448,9 +459,15 @@ case class CometScanRule(session: SparkSession) return withFallbackReasons(scanExec, fallbackReasons.toSet) } + val changelog = IcebergReflection.isChangelogScan(scanExec.scan) + if (changelog && !CometConf.COMET_ICEBERG_CHANGELOG_ENABLED.get()) { + return withFallbackReason(scanExec, "Native Iceberg changelog scans are disabled") + } + // Check for unsupported metadata columns in Iceberg scans val unsupportedMetadataCols = scanExec.output.filter(_.isMetadataCol).filterNot { attr => - CometIcebergNativeScan.MetadataFieldIds.keySet.contains(attr.name) + CometIcebergNativeScan.MetadataFieldIds.keySet.contains(attr.name) || + (changelog && CometIcebergNativeScan.ChangeFieldIds.contains(attr.name)) } if (unsupportedMetadataCols.nonEmpty) { fallbackReasons += "Unsupported Iceberg metadata columns: " + @@ -489,7 +506,7 @@ case class CometScanRule(session: SparkSession) // handed to `extract` below so the reflective accessor runs once per scan. val (icebergTasks, taskValidation) = try { - IcebergReflection.getTasks(scanExec.scan) match { + ShimCometStreaming.icebergTasks(scanExec) match { case Some(tasks) => (tasks, CometScanRule.validateIcebergFileScanTasks(tasks, s3CompliantSchemes)) case None => @@ -988,7 +1005,7 @@ case class CometScanRule(session: SparkSession) // Check that all DPP subqueries use InSubqueryExec which we know how to handle. // Future Spark versions might introduce new subquery types we haven't tested. val dppSubqueriesSupported = { - val unsupportedSubqueries = scanExec.runtimeFilters.collect { + val unsupportedSubqueries = runtimeFilters.collect { case DynamicPruningExpression(e) if !e.isInstanceOf[InSubqueryExec] => e.getClass.getSimpleName } @@ -997,7 +1014,7 @@ case class CometScanRule(session: SparkSession) // as a preparatory refactor for future features (Null Safe Equality DPP, multiple // equality predicates). Currently indices always has one element, but future Spark // versions might use multiple indices. - val multiIndexDpp = scanExec.runtimeFilters.exists { + val multiIndexDpp = runtimeFilters.exists { case DynamicPruningExpression(e: InSubqueryExec) => e.plan match { case sab: SubqueryAdaptiveBroadcastExec => @@ -1028,10 +1045,26 @@ case class CometScanRule(session: SparkSession) metadataSchemeSupported && partitionTypesSupported && unifiedPartitionTypeSupported && complexTypePredicatesSupported && transformFunctionsSupported && deleteFileTypesSupported && dppSubqueriesSupported) { - CometBatchScanExec( - scanExec.clone().asInstanceOf[BatchScanExec], - runtimeFilters = scanExec.runtimeFilters, - nativeIcebergScanMetadata = Some(metadata)) + scanExec match { + case batch: BatchScanExec => + CometBatchScanExec( + batch.clone().asInstanceOf[BatchScanExec], + runtimeFilters = runtimeFilters, + nativeIcebergScanMetadata = Some(metadata)) + case _ => + val nativeOp = CometIcebergNativeScan.placeholder( + scanExec, + metadata, + org.apache.comet.serde.OperatorOuterClass.Operator + .newBuilder() + .setPlanId(scanExec.id)) + CometIcebergNativeScanExec( + nativeOp, + scanExec, + session, + metadata.metadataLocation, + metadata).convertBlock() + } } else { withFallbackReasons(scanExec, fallbackReasons.toSet) } @@ -1300,7 +1333,6 @@ object CometScanRule extends Logging { val formatMethod = IcebergReflection.getMethod(contentFileClass, "format") val pathMethod = IcebergReflection.getMethod(contentFileClass, "path") val residualMethod = IcebergReflection.getMethod(contentScanTaskClass, "residual") - val deletesMethod = IcebergReflection.getMethod(fileScanTaskClass, "deletes") val termMethod = IcebergReflection.getMethod(unboundPredicateClass, "term") var allParquet = true @@ -1349,8 +1381,13 @@ object CometScanRule extends Logging { inspectLocation(pathMethod.invoke(dataFile).toString) + if (!fileScanTaskClass.isInstance(task)) { + val deletes = IcebergReflection.getDeleteFilesFromTask(task, fileScanTaskClass) + require(deletes.isEmpty, "Native changelog scans do not support delete files") + } + // Residual transform check (short-circuit if already found unsupported) - if (nonIdentityTransform.isEmpty && fileScanTaskClass.isInstance(task)) { + if (nonIdentityTransform.isEmpty) { try { val residual = residualMethod.invoke(task) if (unboundPredicateClass.isInstance(residual)) { @@ -1371,7 +1408,7 @@ object CometScanRule extends Logging { // Collect delete files and check their schemes if (fileScanTaskClass.isInstance(task)) { try { - val deletes = deletesMethod.invoke(task).asInstanceOf[java.util.List[_]] + val deletes = IcebergReflection.getDeleteFilesFromTask(task, fileScanTaskClass) deleteFiles.addAll(deletes) deletes.asScala.foreach { deleteFile => diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index ac2538288e2..c4b1b7b03f8 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -34,7 +34,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeExec} import org.apache.spark.sql.comet.shims.ShimDataSourceRDDPartition -import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceRDD, DataSourceRDDPartition} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceRDD, DataSourceRDDPartition, DataSourceV2ScanExecBase} import org.apache.spark.sql.types._ import com.google.protobuf.ByteString @@ -45,6 +45,7 @@ import org.apache.comet.objectstore.NativeConfig import org.apache.comet.serde.{CometOperatorSerde, OperatorOuterClass} import org.apache.comet.serde.OperatorOuterClass.{Operator, SparkStructField} import org.apache.comet.serde.QueryPlanSerde.serializeDataType +import org.apache.comet.shims.ShimCometStreaming object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] with Logging { @@ -100,6 +101,11 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit "_spec_id" -> (Int.MaxValue - 4), "_partition" -> (Int.MaxValue - 5)) + val ChangeFieldIds: Map[String, Int] = Map( + "_change_type" -> (Int.MaxValue - 104), + "_change_ordinal" -> (Int.MaxValue - 105), + "_commit_snapshot_id" -> (Int.MaxValue - 106)) + /** * Wraps an Iceberg partition value (a typed primitive) in a PartitionValue. The value encoding * is shared with predicate literals via [[icebergLiteralToProto]]. @@ -409,13 +415,12 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit private def serializePartitionData( task: Any, contentScanTaskClass: Class[_], - fileScanTaskClass: Class[_], taskBuilder: OperatorOuterClass.IcebergFileScanTask.Builder, commonBuilder: OperatorOuterClass.IcebergScanCommon.Builder, partitionSpecToPoolIndex: mutable.HashMap[String, Int], partitionDataToPoolIndex: mutable.HashMap[String, Int]): Unit = { try { - val specMethod = IcebergReflection.getMethod(fileScanTaskClass, "spec") + val specMethod = IcebergReflection.getMethod(contentScanTaskClass, "spec") val spec = specMethod.invoke(task) if (spec != null) { @@ -876,6 +881,13 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit "Metadata should have been extracted in CometScanRule.") } + Some(placeholder(scan.wrapped, metadata, builder)) + } + + def placeholder( + scan: DataSourceV2ScanExecBase, + metadata: CometIcebergNativeScanMetadata, + builder: Operator.Builder): Operator = { val icebergScanBuilder = OperatorOuterClass.IcebergScan.newBuilder() val commonBuilder = OperatorOuterClass.IcebergScanCommon.newBuilder() @@ -884,13 +896,13 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // required_schema, pools) is set by serializePartitions() at execution time, so setting it // here would be wasted work. commonBuilder.setMetadataLocation(metadata.metadataLocation) - commonBuilder.setScanHashCode(scan.scan.hashCode()) + commonBuilder.setScanHashCode(ShimCometStreaming.icebergScanHash(scan)) icebergScanBuilder.setCommon(commonBuilder.build()) // partition field intentionally empty - will be populated at execution time builder.clearChildren() - Some(builder.setIcebergScan(icebergScanBuilder).build()) + builder.setIcebergScan(icebergScanBuilder).build() } /** @@ -913,7 +925,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * DeleteFileIndex does. * * @param scanExec - * The BatchScanExec whose inputRDD contains the DPP-filtered partitions + * The batch or micro-batch scan whose inputRDD contains the planned partitions * @param output * The output attributes for the scan * @param metadata @@ -922,7 +934,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * Tuple of (commonBytes, perPartitionBytes) for native execution */ def serializePartitions( - scanExec: BatchScanExec, + scanExec: DataSourceV2ScanExecBase, output: Seq[Attribute], metadata: CometIcebergNativeScanMetadata): (Array[Byte], Array[Array[Byte]]) = { @@ -967,7 +979,10 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit lazy val hasHistoricalColumns = { val tableSchemaFieldIds = fieldIdMapping(metadata.tableSchema.asInstanceOf[AnyRef]).values.toSet - metadata.globalFieldIdMapping.values.exists(id => !tableSchemaFieldIds.contains(id)) + metadata.globalFieldIdMapping + .filterNot { case (name, _) => ChangeFieldIds.contains(name) } + .values + .exists(id => !tableSchemaFieldIds.contains(id)) } // Columns whose Iceberg type iceberg-rust cannot use for page-index pruning; residual // predicates over them are dropped (see icebergExprToProto). Computed once from the full table @@ -1033,16 +1048,8 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit .inputPartitions(partition.asInstanceOf[DataSourceRDDPartition]) inputPartitions.foreach { inputPartition => - val inputPartClass = inputPartition.getClass - { - val taskGroupMethod = - IcebergReflection.getDeclaredMethod(inputPartClass, "taskGroup") - val taskGroup = taskGroupMethod.invoke(inputPartition) - - val tasksMethod = IcebergReflection.getMethod(taskGroup.getClass, "tasks") - val tasksCollection = - tasksMethod.invoke(taskGroup).asInstanceOf[java.util.Collection[_]] + val tasksCollection = IcebergReflection.tasksFromInputPartition(inputPartition) tasksCollection.asScala.foreach { task => totalTasks += 1 @@ -1079,7 +1086,31 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // verbatim for iceberg-rust to decode. Unencrypted files leave the field unset. keyMetadataBytes(keyMetadataMethod, dataFile).foreach(taskBuilder.setKeyMetadata) - val taskSchema = taskSchemaMethod.invoke(task) + val taskSchema = if (fileScanTaskClass.isInstance(task)) { + taskSchemaMethod.invoke(task) + } else { + val changeClass = + IcebergReflection.loadClass("org.apache.iceberg.ChangelogScanTask") + val changeType = + IcebergReflection.getMethod(changeClass, "operation").invoke(task).toString + require( + Set("INSERT", "DELETE").contains(changeType), + s"Unsupported changelog operation: $changeType") + taskBuilder.setChange( + OperatorOuterClass.IcebergChange + .newBuilder() + .setChangeType(changeType) + .setChangeOrdinal( + IcebergReflection + .getMethod(changeClass, "changeOrdinal") + .invoke(task) + .asInstanceOf[Int]) + .setCommitSnapshotId(IcebergReflection + .getMethod(changeClass, "commitSnapshotId") + .invoke(task) + .asInstanceOf[Long])) + metadata.tableSchema.asInstanceOf[AnyRef] + } val deletes = IcebergReflection.getDeleteFilesFromTask(task, fileScanTaskClass) @@ -1123,7 +1154,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit .schemaWithRequiredFields( baseSchema, metadata.table, - IcebergReflection.partitionSourceFieldIds(task, fileScanTaskClass)) + IcebergReflection.partitionSourceFieldIds(task)) .asInstanceOf[AnyRef] val schemaIdx = schemaToPoolIndex.getOrElseUpdate( @@ -1137,18 +1168,22 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val nameToFieldId = fieldIdMapping(schema) - val projectFieldIds = output.map { attr => - nameToFieldId - .get(attr.name) - .orElse(metadata.globalFieldIdMapping.get(attr.name)) - .orElse(CometIcebergNativeScan.MetadataFieldIds.get(attr.name)) - .getOrElse { - throw new IllegalStateException( - s"Column '${attr.name}' not found in task schema, global schema, " + - "or metadata field IDs. This indicates a bug in CometScanRule " + - "validation -- all output columns should be resolvable.") - } - } + val projectFieldIds = output + .filterNot { attr => + taskBuilder.hasChange && ChangeFieldIds.contains(attr.name) + } + .map { attr => + nameToFieldId + .get(attr.name) + .orElse(metadata.globalFieldIdMapping.get(attr.name)) + .orElse(CometIcebergNativeScan.MetadataFieldIds.get(attr.name)) + .getOrElse { + throw new IllegalStateException( + s"Column '${attr.name}' not found in task schema, global schema, " + + "or metadata field IDs. This indicates a bug in CometScanRule " + + "validation -- all output columns should be resolvable.") + } + } val projectFieldIdsIdx = projectFieldIdsToPoolIndex.getOrElseUpdate( projectFieldIds, { @@ -1216,7 +1251,6 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit serializePartitionData( task, contentScanTaskClass, - fileScanTaskClass, taskBuilder, commonBuilder, partitionSpecToPoolIndex, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergChangelogExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergChangelogExec.scala new file mode 100644 index 00000000000..cb569de3900 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergChangelogExec.scala @@ -0,0 +1,188 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import java.lang.invoke.SerializedLambda + +import scala.util.control.NonFatal + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet} +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.execution.{DeserializeToObjectExec, MapPartitionsExec, SerializeFromObjectExec} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.types._ + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.serde.CometOperatorSerde +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** Preserves the original Iceberg iterator as the fallback for its native replacement. */ +case class IcebergChangelogExec( + original: SerializeFromObjectExec, + mode: Int, + identifierIndices: Seq[Int], + outputIndices: Seq[Int], + child: SparkPlan) + extends UnaryExecNode { + override def output: Seq[Attribute] = original.output + override def producedAttributes: AttributeSet = outputSet + override def outputPartitioning: Partitioning = + UnknownPartitioning(child.outputPartitioning.numPartitions) + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) + override protected def doExecute(): RDD[InternalRow] = { + val mapper = original.child.asInstanceOf[MapPartitionsExec] + val decoder = mapper.child.asInstanceOf[DeserializeToObjectExec] + original.copy(child = mapper.copy(child = decoder.copy(child = child))).execute() + } +} + +object CometIcebergChangelogExec extends CometOperatorSerde[IcebergChangelogExec] { + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_ICEBERG_CHANGELOG_ENABLED) + override def requiresNativeChildren: Boolean = true + + // Match the procedure's serialized Java lambda, never an arbitrary mapPartitions function. + // Unknown Spark/Iceberg closure layouts retain the original Spark iterator. + private def lambda(function: AnyRef): Option[SerializedLambda] = { + val method = function.getClass.getDeclaredMethod("writeReplace") + method.setAccessible(true) + method.invoke(function) match { + case value: SerializedLambda => Some(value) + case _ => None + } + } + + private def comparable(dt: DataType, nested: Boolean = false): Boolean = dt match { + // Iceberg's external-row iterator uses reference equality for binary values in some + // containers. Preserve that JVM behavior until a common value-equality contract exists. + case BinaryType => false + case FloatType | DoubleType if nested => false + case StructType(fields) => fields.forall(f => comparable(f.dataType, nested = true)) + case ArrayType(element, _) => comparable(element, nested = true) + case MapType(key, value, _) => + comparable(key, nested = true) && comparable(value, nested = true) + case _ => true + } + + def rewrite(plan: SparkPlan): SparkPlan = { + if (!CometConf.COMET_ICEBERG_CHANGELOG_ENABLED.get() || + !CometConf.COMET_ICEBERG_NATIVE_ENABLED.get()) { + return plan + } + plan.transformDown { case original: SerializeFromObjectExec => + val replacement = + try { + original.child match { + case mapper: MapPartitionsExec => + mapper.child match { + case decoder: DeserializeToObjectExec => + for { + wrapper <- lambda(mapper.func.asInstanceOf[AnyRef]) + if wrapper.getImplClass == "org/apache/spark/sql/internal/ToScalaUDF$" + if wrapper.getCapturedArgCount == 1 + function <- lambda(wrapper.getCapturedArg(0)) + if function.getImplClass == + "org/apache/iceberg/spark/procedures/CreateChangelogViewProcedure" + arguments = (0 until function.getCapturedArgCount).map( + function.getCapturedArg) + schema <- arguments.collectFirst { case s: StructType => s } + if schema.fields.forall(f => comparable(f.dataType)) || { + withFallbackReason( + original, + "Iceberg changelog binary or nested floating comparison " + + "uses the JVM iterator") + false + } + if schema.fields.map(f => f.name -> f.dataType).toSeq == + decoder.child.output.map(a => a.name -> a.dataType) + method = function.getImplMethodName + mode <- + if (method.startsWith("lambda$applyChangelogIterator$")) { + Some(1) + } else if (method.startsWith("lambda$applyCarryoverRemoveIterator$")) { + arguments.collectFirst { case b: java.lang.Boolean => if (b) 2 else 0 } + } else { + None + } + identifiers = + if (mode == 1) { + arguments.collectFirst { case a: Array[String] => a.toSeq }.get + } else { + Seq.empty[String] + } + indices = original.output.map(a => schema.fieldIndex(a.name)) + if original.output.zip(indices).forall { case (a, i) => + a.dataType == schema(i).dataType + } + } yield IcebergChangelogExec( + original, + mode, + identifiers.map(schema.fieldIndex), + indices, + decoder.child) + case _ => None + } + case _ => None + } + } catch { case NonFatal(_) => None } + replacement.getOrElse(original) + } + } + + override def convert( + op: IcebergChangelogExec, + builder: Operator.Builder, + children: Operator*): Option[Operator] = { + if (children.isEmpty) { + withFallbackReason(op, "Native Iceberg changelog requires a native input") + return None + } + val changelog = OperatorOuterClass.IcebergChangelog.newBuilder().setMode(op.mode) + val names = op.child.output.map(_.name) + changelog.setChangeTypeIndex(names.indexOf("_change_type")) + changelog.setChangeOrdinalIndex(names.indexOf("_change_ordinal")) + changelog.setCommitSnapshotIdIndex(names.indexOf("_commit_snapshot_id")) + op.identifierIndices.foreach(changelog.addIdentifierIndices) + op.outputIndices.foreach(changelog.addOutputIndices) + Some(builder.setIcebergChangelog(changelog).build()) + } + + override def createExec(nativeOp: Operator, op: IcebergChangelogExec): CometNativeExec = + CometIcebergChangelogExec(nativeOp, op, op.output, op.child, SerializedPlan(None)) +} + +case class CometIcebergChangelogExec( + override val nativeOp: Operator, + override val originalPlan: SparkPlan, + override val output: Seq[Attribute], + child: SparkPlan, + override val serializedPlanOpt: SerializedPlan) + extends CometUnaryExec { + override def producedAttributes: AttributeSet = outputSet + override def outputPartitioning: Partitioning = + UnknownPartitioning(child.outputPartitioning.numPartitions) + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala index 3182a9daeb4..5a89162bc65 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala @@ -27,7 +27,9 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, SortOrder} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} -import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.comet.shims.ShimStreamSourceAwareSparkPlan +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanExecBase} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.AccumulatorV2 @@ -37,14 +39,18 @@ import com.google.common.base.Objects import org.apache.comet.iceberg.CometIcebergNativeScanMetadata import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.operator.CometIcebergNativeScan +import org.apache.comet.shims.ShimCometStreaming /** * Native Iceberg scan operator that delegates file reading to iceberg-rust. * - * Replaces Spark's Iceberg BatchScanExec to bypass the DataSource V2 API and enable native - * execution. Iceberg's catalog and planning run in Spark to produce FileScanTasks, which are - * serialized to protobuf for the native side to execute using iceberg-rust's FileIO and - * ArrowReader. This provides better performance than reading through Spark's abstraction layers. + * Replaces Spark's Iceberg batch or micro-batch scan to enable native file reading. Micro-batch + * scans retain their Spark stream identity and use its offset-bounded input partitions. + * + * Bypasses the DataSource V2 reader API for native execution. Iceberg's catalog and planning run + * in Spark to produce FileScanTasks, which are serialized to protobuf for the native side to + * execute using iceberg-rust's FileIO and ArrowReader. This provides better performance than + * reading through Spark's abstraction layers. * * Supports Dynamic Partition Pruning (DPP) via top-level `runtimeFilters` (mirroring Spark's * `BatchScanExec.runtimeFilters`). Because the field is a constructor parameter, Spark's standard @@ -57,12 +63,15 @@ case class CometIcebergNativeScanExec( override val nativeOp: Operator, override val output: Seq[Attribute], runtimeFilters: Seq[Expression], - @transient override val originalPlan: BatchScanExec, + @transient override val originalPlan: DataSourceV2ScanExecBase, override val serializedPlanOpt: SerializedPlan, metadataLocation: String, scanHashCode: Int, @transient nativeIcebergScanMetadata: CometIcebergNativeScanMetadata) - extends CometLeafExec { + extends CometLeafExec + with ShimStreamSourceAwareSparkPlan { + + override protected def streamSourcePlan: SparkPlan = originalPlan override val supportsColumnar: Boolean = true @@ -99,12 +108,11 @@ case class CometIcebergNativeScanExec( // would re-translate the original (unresolved) InSubqueryExec and throw "no subquery // result". This makes the top-level runtimeFilters the single source of truth at // serialization time. - val effectiveOriginalPlan = - if (originalPlan.runtimeFilters != runtimeFilters) { - originalPlan.copy(runtimeFilters = runtimeFilters) - } else { - originalPlan - } + val effectiveOriginalPlan = originalPlan match { + case batch: BatchScanExec if batch.runtimeFilters != runtimeFilters => + batch.copy(runtimeFilters = runtimeFilters) + case _ => originalPlan + } CometIcebergNativeScan.serializePartitions( effectiveOriginalPlan, output, @@ -213,7 +221,15 @@ case class CometIcebergNativeScanExec( // Add num_splits as a runtime metric (incremented on the native side during execution) val numSplitsMetric = SQLMetrics.createMetric(sparkContext, "number of file splits processed") - baseMetrics ++ icebergPlanningMetrics + ("num_splits" -> numSplitsMetric) + // Spark's streaming progress reporter looks up this name on the source-aware scan. Share + // the native counter so repeated foreachBatch actions have Spark's usual input-row accounting. + val sourceMetrics = + if (originalPlan != null && ShimCometStreaming.isStreamingPlan(originalPlan)) { + Map("numOutputRows" -> baseMetrics("output_rows")) + } else { + Map.empty[String, SQLMetric] + } + baseMetrics ++ icebergPlanningMetrics ++ sourceMetrics + ("num_splits" -> numSplitsMetric) } /** Executes using CometExecRDD - planning data is computed lazily on first access. */ @@ -343,7 +359,7 @@ object CometIcebergNativeScanExec { /** Creates a CometIcebergNativeScanExec with deferred partition serialization. */ def apply( nativeOp: Operator, - scanExec: BatchScanExec, + scanExec: DataSourceV2ScanExecBase, session: SparkSession, metadataLocation: String, nativeIcebergScanMetadata: CometIcebergNativeScanMetadata): CometIcebergNativeScanExec = { @@ -351,13 +367,16 @@ object CometIcebergNativeScanExec { val exec = CometIcebergNativeScanExec( nativeOp, scanExec.output, - scanExec.runtimeFilters, + scanExec match { + case batch: BatchScanExec => batch.runtimeFilters + case _ => Nil + }, scanExec, SerializedPlan(None), metadataLocation, // Capture Iceberg's scan hash now, while the transient scan is still available; it is // needed for equality after canonicalization nulls originalPlan (see #4774). - scanExec.scan.hashCode(), + ShimCometStreaming.icebergScanHash(scanExec), nativeIcebergScanMetadata) scanExec.logicalLink.foreach(exec.setLogicalLink) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 0e160faf47a..20d22d06552 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -31,7 +31,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, Expression, ExpressionSet, Generator, NamedExpression, SortOrder} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, Mode, Partial, PartialMerge, Percentile} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Count, Final, Mode, Partial, PartialMerge, Percentile} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.physical._ @@ -1640,6 +1640,16 @@ trait CometBaseAggregate { val modes = aggregate.aggregateExpressions.map(_.mode).distinct val modeSet = modes.toSet val hasPartialMerge = modeSet.contains(PartialMerge) + // Streaming buffers are persisted by Spark and may be restored by a later Spark-only run. + // An all-native aggregate chain does not make an incompatible checkpoint format safe. + // Count has a compatible buffer; its mixed batch exclusion concerns AQE/subquery rewrites, + // neither of which applies to streaming aggregates. + if (aggregate.isStreaming && + !QueryPlanSerde.allAggsSupportMixedExecution( + aggregate.aggregateExpressions.filterNot(_.aggregateFunction.isInstanceOf[Count]))) { + withFallbackReason(aggregate, "Streaming aggregate requires Spark-compatible state buffers") + return None + } // In distinct aggregates there can be a combination of modes. // We support {Partial, PartialMerge} mix; other combinations are rejected. val multiMode = modes.size > 1 && modeSet != Set(Partial, PartialMerge) diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometStreaming.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometStreaming.scala index f80eab95ee7..ae1f4061af3 100644 --- a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometStreaming.scala +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometStreaming.scala @@ -20,12 +20,31 @@ package org.apache.comet.shims import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanExecBase +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.iceberg.IcebergReflection // StreamSourceAwareSparkPlan does not exist in Spark 3.x, so fall back to // walking the physical tree and checking each node's logical link. Inspecting // only the root's logical link would silently miss streaming plans whenever a // rule produced a fresh root node without copying the link over. object ShimCometStreaming { + def nativeExecutionEnabled(conf: SQLConf, plan: SparkPlan): Boolean = false + + def isStateBoundary(plan: SparkPlan): Boolean = false + def isStreamingPlan(plan: SparkPlan): Boolean = plan.exists(_.logicalLink.exists(_.isStreaming)) + + // Spark 3.x progress reporting requires a MicroBatchScanExec node and cannot associate a + // replacement source with its stream. Keep the existing Spark execution path on these versions. + def transformIcebergScans( + plan: SparkPlan, + transform: DataSourceV2ScanExecBase => SparkPlan): SparkPlan = plan + + def icebergTasks(scan: DataSourceV2ScanExecBase): Option[java.util.List[_]] = + IcebergReflection.getTasks(scan.scan) + + def icebergScanHash(scan: DataSourceV2ScanExecBase): Int = scan.scan.hashCode() } diff --git a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala index 31156779dec..059d0b00f0a 100644 --- a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala +++ b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala @@ -19,4 +19,8 @@ package org.apache.spark.sql.comet.shims -trait ShimStreamSourceAwareSparkPlan {} +import org.apache.spark.sql.execution.SparkPlan + +trait ShimStreamSourceAwareSparkPlan { + protected def streamSourcePlan: SparkPlan = null +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometStreaming.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometStreaming.scala index ed4cb37e2d0..3f70b16632d 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometStreaming.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometStreaming.scala @@ -19,11 +19,67 @@ package org.apache.comet.shims +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream import org.apache.spark.sql.execution.{SparkPlan, StreamSourceAwareSparkPlan} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2ScanExecBase, MicroBatchScanExec} +import org.apache.spark.sql.execution.streaming.Source +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf +import org.apache.comet.iceberg.IcebergReflection object ShimCometStreaming { - def isStreamingPlan(plan: SparkPlan): Boolean = plan.exists { - case p: StreamSourceAwareSparkPlan => p.getStream.isDefined - case _ => false + def nativeExecutionEnabled(conf: SQLConf, plan: SparkPlan): Boolean = + CometConf.COMET_STREAMING_EXEC_ENABLED.get(conf) && !plan.exists { + case source: StreamSourceAwareSparkPlan => + source.getStream.exists { + case _: Source | _: MicroBatchStream => false + case _ => true + } + case _ => false + } + + def isStateBoundary(plan: SparkPlan): Boolean = { + // These classes moved from execution.streaming to operators.stateful in Spark 4.1. + Set("StateStoreRestoreExec", "StateStoreSaveExec", "EventTimeWatermarkExec") + .contains(plan.getClass.getSimpleName) + } + + def isStreamingPlan(plan: SparkPlan): Boolean = plan.exists { p => + // No-data batches can have a streaming LocalRelation with no source attached. + p.logicalLink.exists(_.isStreaming) || (p match { + case source: StreamSourceAwareSparkPlan => source.getStream.isDefined + case _ => false + }) + } + + def transformIcebergScans( + plan: SparkPlan, + transform: DataSourceV2ScanExecBase => SparkPlan): SparkPlan = plan.transformUp { + // A vendor CDF source can expose different task semantics. Only claim the Apache Iceberg + // append-snapshot source whose offset-bounded FileScanTasks the native reader understands. + case scan: MicroBatchScanExec + if IcebergReflection.isIcebergScanClass(scan.scan.getClass.getName) && + scan.stream.getClass.getName == "org.apache.iceberg.spark.source.SparkMicroBatchStream" => + transform(scan) + } + + def icebergTasks(scan: DataSourceV2ScanExecBase): Option[java.util.List[_]] = scan match { + case microBatch: MicroBatchScanExec => + // inputPartitions is a lazy val on MicroBatchScanExec. Serialization uses the same + // materialized partitions via inputRDD; never call scan.tasks() or scan.toBatch here. + Some(microBatch.inputPartitions.flatMap { partition => + IcebergReflection.tasksFromInputPartition(partition).asScala + }.asJava) + case _ => IcebergReflection.getTasks(scan.scan) + } + + def icebergScanHash(scan: DataSourceV2ScanExecBase): Int = scan match { + case microBatch: MicroBatchScanExec => + (microBatch.stream, microBatch.start.json(), microBatch.end.json(), scan.scan.hashCode()) + .hashCode() + case _ => scan.scan.hashCode() } } diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala index 93552fc00fd..7e2f39f19c8 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala @@ -20,8 +20,13 @@ package org.apache.spark.sql.comet.shims import org.apache.spark.sql.connector.read.streaming.SparkDataStream -import org.apache.spark.sql.execution.StreamSourceAwareSparkPlan +import org.apache.spark.sql.execution.{SparkPlan, StreamSourceAwareSparkPlan} trait ShimStreamSourceAwareSparkPlan extends StreamSourceAwareSparkPlan { - override def getStream: Option[SparkDataStream] = None + protected def streamSourcePlan: SparkPlan = null + + override def getStream: Option[SparkDataStream] = streamSourcePlan match { + case source: StreamSourceAwareSparkPlan => source.getStream + case _ => None + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometOperatorSerdeBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometOperatorSerdeBenchmark.scala index 2f3904c5df6..0e97dc445ef 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometOperatorSerdeBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometOperatorSerdeBenchmark.scala @@ -26,6 +26,7 @@ import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.comet.{CometBatchScanExec, CometIcebergNativeScanExec} import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.comet.CometConf import org.apache.comet.rules.CometScanRule @@ -100,7 +101,7 @@ object CometOperatorSerdeBenchmark extends CometBenchmarkBase { private def reconstructBatchScanExec( nativeScan: CometIcebergNativeScanExec): CometBatchScanExec = { CometBatchScanExec( - wrapped = nativeScan.originalPlan, + wrapped = nativeScan.originalPlan.asInstanceOf[BatchScanExec], runtimeFilters = Seq.empty, nativeIcebergScanMetadata = Some(nativeScan.nativeIcebergScanMetadata)) } diff --git a/spark/src/test/spark-4.x/org/apache/comet/CometIcebergStreamingSuite.scala b/spark/src/test/spark-4.x/org/apache/comet/CometIcebergStreamingSuite.scala new file mode 100644 index 00000000000..920b8c1e570 --- /dev/null +++ b/spark/src/test/spark-4.x/org/apache/comet/CometIcebergStreamingSuite.scala @@ -0,0 +1,830 @@ +/* + * 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. + */ + +package org.apache.comet + +import java.io.{File, FileInputStream} +import java.util.Properties + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ +import scala.util.Using + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame, QueryTest, Row} +import org.apache.spark.sql.comet.{CometHashAggregateExec, CometIcebergChangelogExec, CometIcebergNativeScanExec} +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.datasources.v2.MicroBatchScanExec +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.streaming.{StreamingQuery, StreamingQueryException, Trigger} + +class CometIcebergStreamingSuite + extends CometTestBase + with CometIcebergTestBase + with AdaptiveSparkPlanHelper { + + override protected def sparkConf: SparkConf = + super.sparkConf.set( + "spark.sql.extensions", + "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions") + + private def withCatalog(f: (File, String, String) => Unit): Unit = { + assume(icebergAvailable, "Compatible Iceberg runtime required") + withTempIcebergDir { dir => + val catalog = "stream_cat_" + dir.getName.replace("-", "_") + val namespace = "ns_" + dir.getName.replace("-", "_") + // Run the same assertions against a disposable REST catalog during integration testing. + val properties = sys.env.get("COMET_TEST_ICEBERG_CATALOG_PROPERTIES") match { + case Some(path) => + val values = new Properties() + Using.resource(new FileInputStream(path))(values.load) + values.stringPropertyNames().asScala.toSeq.map(k => k -> values.getProperty(k)) + case None => + Seq("type" -> "hadoop", "warehouse" -> new File(dir, "warehouse").toString) + } + val configs = properties.map { case (key, value) => + s"spark.sql.catalog.$catalog.$key" -> value + } ++ Seq( + s"spark.sql.catalog.$catalog" -> "org.apache.iceberg.spark.SparkCatalog", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_CHANGELOG_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_STREAMING_ENABLED.key -> "true") + withSQLConf(configs: _*) { + sql(s"CREATE NAMESPACE $catalog.$namespace") + try f(dir, catalog, namespace) + finally sql(s"DROP NAMESPACE $catalog.$namespace") + } + } + } + + // StreamingQueryWrapper moved packages in Spark 4.1. The public accessors are unchanged. + private def executedPlan(query: StreamingQuery): SparkPlan = { + val execution = query.getClass.getMethod("streamingQuery").invoke(query) + execution.getClass + .getMethod("lastExecution") + .invoke(execution) + .asInstanceOf[QueryExecution] + .executedPlan + } + + private def runAvailable(source: DataFrame, checkpoint: File, outputMode: String = "append")( + process: (DataFrame, Long) => Unit): StreamingQuery = { + val query = source.writeStream + .outputMode(outputMode) + .option("checkpointLocation", checkpoint.toString) + .trigger(Trigger.AvailableNow()) + .foreachBatch(process) + .start() + try { + assert(query.awaitTermination(60000), "AvailableNow query did not terminate") + query + } finally { + if (query.isActive) query.stop() + } + } + + test( + "offset-bounded native batches preserve admission limits, progress, and checkpoint restart") { + withCatalog { (dir, catalog, namespace) => + val table = s"$catalog.$namespace.events" + withTable(table) { + sql(s"CREATE TABLE $table (id BIGINT) USING iceberg") + (0 until 3).foreach { batch => + spark.range(batch * 2L, batch * 2L + 2).coalesce(1).writeTo(table).append() + } + val checkpoint = new File(dir, "checkpoint") + def source = spark.readStream + .option("streaming-max-files-per-micro-batch", "1") + .table(table) + // Iceberg 1.10 falls back from AvailableNow to a single batch. 1.11 supports the + // trigger and its admission limit. Preserve the source runtime's own batch boundaries. + val expected = ArrayBuffer.empty[(Long, Seq[Long])] + val baseline = withSQLConf(CometConf.COMET_ICEBERG_STREAMING_ENABLED.key -> "false") { + runAvailable(source, new File(dir, "baseline")) { (batch, id) => + expected += id -> batch.collect().toSeq.map(_.getLong(0)).sorted + } + } + val batches = ArrayBuffer.empty[(Long, Seq[Long])] + val query = runAvailable(source, checkpoint) { (batch, id) => + batches += id -> batch.collect().toSeq.map(_.getLong(0)).sorted + } + assert(batches == expected) + assert(batches.flatMap(_._2).sorted == (0L until 6L)) + assert( + query.recentProgress.map(_.numInputRows).toSeq == + baseline.recentProgress.map(_.numInputRows).toSeq) + if (icebergVersionAtLeast(1, 11)) { + assert(batches.map(_._2) == Seq(Seq(0L, 1L), Seq(2L, 3L), Seq(4L, 5L))) + } + val native = executedPlan(query).collect { case s: CometIcebergNativeScanExec => s } + assert(native.size == 1, executedPlan(query).toString) + assert(native.head.getStream.isDefined) + assert(native.head.metrics("numOutputRows") eq native.head.metrics("output_rows")) + + spark.range(6, 8).coalesce(1).writeTo(table).append() + val resumed = ArrayBuffer.empty[(Long, Seq[Long])] + val restart = runAvailable(source, checkpoint) { (batch, id) => + resumed += id -> batch.collect().toSeq.map(_.getLong(0)).sorted + } + assert(resumed.toSeq == Seq((batches.last._1 + 1) -> Seq(6L, 7L))) + assert(restart.recentProgress.map(_.numInputRows).sum == 2) + assert(executedPlan(restart).exists(_.isInstanceOf[CometIcebergNativeScanExec])) + } + } + } + + test("streaming scan opt-in and native execution gates retain Spark source and progress") { + withCatalog { (dir, catalog, namespace) => + val table = s"$catalog.$namespace.fallback_events" + withTable(table) { + sql(s"CREATE TABLE $table (id INT) USING iceberg") + sql(s"INSERT INTO $table VALUES (1), (2)") + Seq( + CometConf.COMET_ICEBERG_STREAMING_ENABLED.key, + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key, + CometConf.COMET_NATIVE_SCAN_ENABLED.key, + CometConf.COMET_EXEC_ENABLED.key).zipWithIndex.foreach { case (config, index) => + withSQLConf(config -> "false") { + val rows = ArrayBuffer.empty[Row] + val query = + runAvailable(spark.readStream.table(table), new File(dir, s"fallback-$index")) { + (batch, _) => + rows ++= batch.collect() + } + assert(rows.toSet == Set(Row(1), Row(2))) + assert(query.recentProgress.map(_.numInputRows).sum == 2) + assert(executedPlan(query).exists(_.isInstanceOf[MicroBatchScanExec])) + assert(!executedPlan(query).exists(_.isInstanceOf[CometIcebergNativeScanExec])) + } + } + } + } + } + + test("failed foreachBatch is replayed without replaying an earlier committed batch") { + withCatalog { (dir, catalog, namespace) => + val table = s"$catalog.$namespace.retry_events" + withTable(table) { + sql(s"CREATE TABLE $table (id BIGINT) USING iceberg") + spark.range(0, 1).coalesce(1).writeTo(table).append() + def source = spark.readStream + .option("streaming-max-files-per-micro-batch", "1") + .table(table) + val checkpoint = new File(dir, "retry") + val committed = ArrayBuffer.empty[Long] + runAvailable(source, checkpoint) { (batch, id) => + assert(id == 0) + committed ++= batch.collect().map(_.getLong(0)) + } + assert(committed.toSeq == Seq(0L)) + spark.range(1, 2).coalesce(1).writeTo(table).append() + val error = intercept[StreamingQueryException] { + runAvailable(source, checkpoint) { (batch, id) => + assert(id == 1) + assert(batch.collect().map(_.getLong(0)).toSeq == Seq(1L)) + throw new IllegalStateException("injected callback failure") + } + } + assert(error.getMessage.contains("injected callback failure")) + assert(committed.toSeq == Seq(0L)) + val retried = ArrayBuffer.empty[Long] + val query = runAvailable(source, checkpoint) { (batch, id) => + val rows = batch.collect().map(_.getLong(0)).toSeq + assert(rows == Seq(id)) + retried ++= rows + } + assert(retried.toSeq == Seq(1L)) + assert(executedPlan(query).exists(_.isInstanceOf[CometIcebergNativeScanExec])) + } + } + } + + test("unsupported Iceberg files and other streaming sources retain their Spark readers") { + withCatalog { (dir, catalog, namespace) => + val table = s"$catalog.$namespace.avro_events" + withTable(table) { + sql(s"""CREATE TABLE $table (id INT) USING iceberg + TBLPROPERTIES ('write.format.default' = 'avro')""") + sql(s"INSERT INTO $table VALUES (1), (2)") + val rows = ArrayBuffer.empty[Row] + val query = runAvailable(spark.readStream.table(table), new File(dir, "avro")) { + (batch, _) => rows ++= batch.collect() + } + assert(rows.toSet == Set(Row(1), Row(2))) + assert(executedPlan(query).exists(_.isInstanceOf[MicroBatchScanExec])) + assert(!executedPlan(query).exists(_.isInstanceOf[CometIcebergNativeScanExec])) + } + val query = runAvailable( + spark.readStream.format("rate-micro-batch").option("rowsPerBatch", 2).load(), + new File(dir, "other-source")) { (batch, _) => + assert(batch.count() == 2) + } + assert(executedPlan(query).exists(_.isInstanceOf[MicroBatchScanExec])) + assert(!executedPlan(query).exists(_.isInstanceOf[CometIcebergNativeScanExec])) + } + } + + test("append streaming rejects overwrite and delete snapshots instead of losing changes") { + withCatalog { (dir, catalog, namespace) => + Seq("overwrite", "delete").foreach { operation => + val table = s"$catalog.$namespace.mutations_$operation" + withTable(table) { + sql(s"""CREATE TABLE $table (id INT, target STRING) USING iceberg + PARTITIONED BY (id) TBLPROPERTIES ('format-version' = '2', + 'write.update.mode' = 'copy-on-write')""") + sql(s"INSERT INTO $table VALUES (1, 'old'), (2, 'removed')") + val checkpoint = new File(dir, operation) + def source = spark.readStream.table(table) + val initial = runAvailable(source, checkpoint) { (batch, _) => + assert(batch.collect().toSet == Set(Row(1, "old"), Row(2, "removed"))) + } + assert(executedPlan(initial).exists(_.isInstanceOf[CometIcebergNativeScanExec])) + if (operation == "overwrite") { + sql(s"UPDATE $table SET target = 'new' WHERE id = 1") + } else { + sql(s"DELETE FROM $table WHERE id = 2") + } + val snapshot = sql(s"SELECT operation FROM $table.snapshots ORDER BY committed_at DESC") + assert(snapshot.head().getString(0) == operation) + val error = intercept[StreamingQueryException] { + runAvailable(source, checkpoint) { (batch, _) => batch.collect() } + } + assert(error.getMessage.toLowerCase.contains(operation), error.getMessage) + } + } + } + } + + test("Iceberg changelog computes actual mutation images within snapshot bounds") { + withCatalog { (_, catalog, namespace) => + val table = s"$catalog.$namespace.changelog_mappings" + withTable(table) { + sql(s"""CREATE TABLE $table (tenant INT, id INT, target STRING) USING iceberg + TBLPROPERTIES ('format-version' = '2', 'write.update.mode' = 'copy-on-write', + 'write.delete.mode' = 'copy-on-write')""") + sql(s"INSERT INTO $table VALUES (1, 1, 'old'), (1, 2, 'stable'), (1, 3, 'removed')") + def latestSnapshot: Long = + sql(s"SELECT snapshot_id FROM $table.snapshots ORDER BY committed_at DESC") + .head() + .getLong(0) + val start = latestSnapshot + sql(s"UPDATE $table SET target = 'new' WHERE id = 1") + val update = latestSnapshot + sql(s"DELETE FROM $table WHERE id = 3") + val delete = latestSnapshot + sql(s"INSERT INTO $table VALUES (1, 4, 'added')") + val end = latestSnapshot + // Prove that an explicitly bounded read excludes a later commit. + sql(s"INSERT INTO $table VALUES (1, 5, 'too-late')") + withTempView("actual_changes") { + sql(s"""CALL $catalog.system.create_changelog_view( + table => '$namespace.changelog_mappings', changelog_view => 'actual_changes', + options => map('start-snapshot-id', '$start', 'end-snapshot-id', '$end'), + identifier_columns => array('tenant', 'id'), compute_updates => true)""") + val changes = spark + .table("actual_changes") + .select("id", "target", "_change_type", "_commit_snapshot_id") + val expected = Seq( + Row(1, "old", "UPDATE_BEFORE", update), + Row(1, "new", "UPDATE_AFTER", update), + Row(3, "removed", "DELETE", delete), + Row(4, "added", "INSERT", end)) + checkAnswer(changes, expected) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + checkAnswer( + spark + .table("actual_changes") + .select("id", "target", "_change_type", "_commit_snapshot_id"), + expected) + } + // The standard Iceberg procedure remains a batch API, with native data processing. + assert(!changes.isStreaming) + assert( + collect(changes.queryExecution.executedPlan) { + case scan: CometIcebergNativeScanExec => scan + }.nonEmpty, + changes.queryExecution.executedPlan.toString) + assert( + collect(changes.queryExecution.executedPlan) { case node: CometIcebergChangelogExec => + node + }.nonEmpty, + changes.queryExecution.executedPlan.toString) + } + } + } + } + + private def checkNativeChangelog(query: => DataFrame, processed: Boolean): Unit = { + val expected = withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + query.collect().toSeq + } + val native = query + // The procedure sorts within partitions; it promises no global output order. + QueryTest.sameRows(expected, native.collect().toSeq, false).foreach(message => fail(message)) + val plan = native.queryExecution.executedPlan + assert(collect(plan) { case s: CometIcebergNativeScanExec => s }.nonEmpty, plan.toString) + if (processed) { + assert(collect(plan) { case c: CometIcebergChangelogExec => c }.nonEmpty, plan.toString) + } + } + + test("native raw changelog, carryovers and net changes preserve bounds and duplicates") { + withCatalog { (_, catalog, namespace) => + for (version <- Seq(1, 2, 3) if version < 3 || icebergVersionAtLeast(1, 9)) { + val table = s"$catalog.$namespace.cdc_v$version" + withTable(table) { + sql(s"""CREATE TABLE $table (id INT, value STRING) USING iceberg + TBLPROPERTIES ('format-version' = '$version', + 'write.update.mode' = 'copy-on-write', 'write.delete.mode' = 'copy-on-write')""") + sql(s"""INSERT INTO $table SELECT /*+ COALESCE(1) */ * FROM VALUES + (1, 'old'), (2, NULL), (2, NULL)""") + assert(sql(s"SELECT * FROM $table.data_files").count() == 1) + def snapshot = sql(s"""SELECT snapshot_id, committed_at FROM $table.snapshots + ORDER BY committed_at DESC""").head() + val start = snapshot + sql(s"UPDATE $table SET value = 'middle' WHERE id = 1") + sql(s"UPDATE $table SET value = 'final' WHERE id = 1") + sql(s"INSERT INTO $table VALUES (3, 'temporary')") + sql(s"DELETE FROM $table WHERE id = 3") + val end = snapshot + sql(s"INSERT INTO $table VALUES (4, 'outside')") + val bounds = s"""map('start-snapshot-id', '${start.getLong(0)}', + 'end-snapshot-id', '${end.getLong(0)}')""" + def raw = spark.read + .option("start-snapshot-id", start.getLong(0)) + .option("end-snapshot-id", end.getLong(0)) + .table(s"$table.changes") + checkNativeChangelog(raw, processed = false) + // Iceberg 1.11's JVM changelog reader misprojects noncontiguous metadata fields. + // Derive this projection from complete JVM rows instead of exercising that bug. + val rawRows = withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + raw.collect().toSeq + } + checkAnswer( + raw.select("_change_type", "_commit_snapshot_id"), + rawRows.map(r => Row(r.getString(2), r.getLong(4)))) + checkNativeChangelog(raw.groupBy("_change_type").count(), processed = false) + checkNativeChangelog(raw.filter("id = 1"), processed = false) + assert(raw.filter("id = 2").count() > 0, "Fixture must contain actual carryovers") + withTempView("cdc_default", "cdc_net", "cdc_time") { + sql(s"""CALL $catalog.system.create_changelog_view( + table => '$namespace.cdc_v$version', changelog_view => 'cdc_default', + options => $bounds)""") + checkNativeChangelog(spark.table("cdc_default"), processed = true) + checkNativeChangelog( + spark.table("cdc_default").filter("id = 1 AND _change_type = 'DELETE'"), + processed = true) + checkNativeChangelog( + spark.table("cdc_default").select("_change_type"), + processed = true) + assert(spark.table("cdc_default").filter("id = 2 OR id = 4").count() == 0) + sql(s"""CALL $catalog.system.create_changelog_view( + table => '$namespace.cdc_v$version', changelog_view => 'cdc_net', + options => $bounds, net_changes => true)""") + checkNativeChangelog(spark.table("cdc_net"), processed = true) + checkAnswer( + spark.table("cdc_net").select("id", "value", "_change_type").orderBy("value"), + Seq(Row(1, "final", "INSERT"), Row(1, "old", "DELETE"))) + sql(s"""CALL $catalog.system.create_changelog_view( + table => '$namespace.cdc_v$version', changelog_view => 'cdc_time', + options => map('start-timestamp', '${start.getTimestamp(1).getTime}', + 'end-timestamp', '${end.getTimestamp(1).getTime}'), net_changes => true)""") + checkNativeChangelog(spark.table("cdc_time"), processed = true) + checkAnswer(spark.table("cdc_time"), spark.table("cdc_net")) + } + } + } + } + } + + test("native update images use table identifiers and preserve complex values") { + withCatalog { (_, catalog, namespace) => + val table = s"$catalog.$namespace.cdc_complex" + withTable(table) { + sql(s"""CREATE TABLE $table (id INT NOT NULL, value STRING, tags MAP, + details STRUCT>, score DOUBLE) USING iceberg + TBLPROPERTIES ('format-version' = '2', 'write.update.mode' = 'copy-on-write')""") + sql(s"ALTER TABLE $table SET IDENTIFIER FIELDS id") + sql(s"""INSERT INTO $table SELECT /*+ COALESCE(1) */ * FROM VALUES + (1, 'stable', map('b', 2, 'a', 1), named_struct('items', array(1, NULL)), double('NaN')), + (2, 'old', NULL, NULL, -0.0D)""") + assert(sql(s"SELECT * FROM $table.data_files").count() == 1) + val start = sql(s"SELECT snapshot_id FROM $table.snapshots").head().getLong(0) + sql(s"UPDATE $table SET value = 'new' WHERE id = 2") + withTempView("cdc_identifiers") { + // Explicit compute_updates uses the table's identifier fields when columns are omitted. + sql(s"""CALL $catalog.system.create_changelog_view( + table => '$namespace.cdc_complex', changelog_view => 'cdc_identifiers', + options => map('start-snapshot-id', '$start'), compute_updates => true)""") + checkNativeChangelog(spark.table("cdc_identifiers"), processed = true) + checkAnswer( + spark.table("cdc_identifiers").select("id", "value", "_change_type"), + Seq(Row(2, "old", "UPDATE_BEFORE"), Row(2, "new", "UPDATE_AFTER"))) + // Explicit columns imply compute_updates=true when that option is omitted. + sql(s"""CALL $catalog.system.create_changelog_view( + table => '$namespace.cdc_complex', changelog_view => 'cdc_identifiers', + options => map('start-snapshot-id', '$start'), identifier_columns => array('id'))""") + checkNativeChangelog(spark.table("cdc_identifiers"), processed = true) + checkNativeChangelog( + spark.table("cdc_identifiers").selectExpr("id + 10 AS id", "upper(value) AS value"), + processed = true) + withSQLConf(CometConf.COMET_ICEBERG_CHANGELOG_ENABLED.key -> "false") { + val fallback = spark.table("cdc_identifiers") + assert(fallback.count() == 2) + assert(collect(fallback.queryExecution.executedPlan) { + case c: CometIcebergChangelogExec => c + }.isEmpty) + } + } + } + } + } + + test("native changelog rejects ambiguous update identifiers and keeps procedure validation") { + withCatalog { (_, catalog, namespace) => + val table = s"$catalog.$namespace.cdc_duplicates" + withTable(table) { + sql(s"""CREATE TABLE $table (id INT, value STRING) USING iceberg + TBLPROPERTIES ('write.update.mode' = 'copy-on-write')""") + sql(s"INSERT INTO $table VALUES (1, 'a'), (1, 'b')") + val start = sql(s"SELECT snapshot_id FROM $table.snapshots").head().getLong(0) + sql(s"UPDATE $table SET value = concat(value, '_new')") + withTempView("cdc_ambiguous") { + def causeContains(error: Throwable, message: String): Boolean = + Iterator.iterate(error)(_.getCause).takeWhile(_ != null).exists { cause => + Option(cause.getMessage).exists(_.contains(message)) + } + Seq("true", "false").foreach { enabled => + withSQLConf(CometConf.COMET_ENABLED.key -> enabled) { + sql(s"""CALL $catalog.system.create_changelog_view( + table => '$namespace.cdc_duplicates', changelog_view => 'cdc_ambiguous', + options => map('start-snapshot-id', '$start'), identifier_columns => array('id'))""") + val error = intercept[Exception](spark.table("cdc_ambiguous").collect()) + assert( + causeContains(error, "multiple rows with the same identifier"), + error.toString) + } + } + val error = intercept[Exception] { + sql(s"""CALL $catalog.system.create_changelog_view( + table => '$namespace.cdc_duplicates', identifier_columns => array('id'), + net_changes => true)""") + } + assert(causeContains(error, "Not support net changes with update images")) + } + } + } + } + + test("vanilla Iceberg batch changelog rejects merge-on-read delete files") { + withCatalog { (_, catalog, namespace) => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val table = s"$catalog.$namespace.mor_changes" + withTable(table) { + sql(s"""CREATE TABLE $table (id INT, target STRING) USING iceberg + TBLPROPERTIES ('format-version' = '2', 'write.update.mode' = 'merge-on-read')""") + sql(s"INSERT INTO $table VALUES (1, 'old'), (2, 'stable')") + sql(s"UPDATE $table SET target = 'new' WHERE id = 1") + assert(sql(s"SELECT * FROM $table.delete_files").count() > 0) + val error = intercept[Exception] { + spark.table(s"$table.changes").collect() + } + assert( + Iterator.iterate[Throwable](error)(_.getCause).takeWhile(_ != null).exists { cause => + Option(cause.getMessage).exists( + _.contains("Delete files are currently not supported in changelog scans")) + }, + error.toString) + } + } + } + } + + test("append event micro-batches filter postimages and detect historical mapping conflicts") { + // The application supplies change events; the append source does not derive them from updates. + withCatalog { (dir, catalog, namespace) => + val mappings = s"$catalog.$namespace.mappings" + val events = s"$catalog.$namespace.mapping_events" + withTable(mappings, events) { + sql(s"""CREATE TABLE $mappings ( + tenant INT, before_key STRING, target STRING, observed TIMESTAMP) USING iceberg""") + sql(s"""INSERT INTO $mappings VALUES + (1, 'a', 'old', TIMESTAMP '2020-01-01'), + (1, 'a', 'new', TIMESTAMP '2020-01-02'), + (1, 'b', 'unchanged', TIMESTAMP '2020-01-01'), + (1, 'c', 'inserted', TIMESTAMP '2020-01-01'), + (2, 'b', 'different-tenant', TIMESTAMP '2020-01-01')""") + val records = sql("""SELECT * FROM VALUES + (1, 'a', 'old', 'update_preimage', TIMESTAMP '2020-01-02'), + (1, 'a', 'new', 'update_postimage', TIMESTAMP '2020-01-02'), + (1, 'b', 'discarded-preimage', 'update_preimage', TIMESTAMP '2020-01-02'), + (1, 'b', 'unchanged', 'update_postimage', TIMESTAMP '2020-01-02'), + (1, 'c', 'discarded-delete', 'delete', TIMESTAMP '2020-01-02'), + (1, 'c', 'inserted', 'insert', TIMESTAMP '2020-01-02') + AS changes(tenant, before_key, target, _change_type, observed)""") + records.coalesce(1).writeTo(events).using("iceberg").create() + val query = + runAvailable(spark.readStream.table(events), new File(dir, "mapping-events")) { + (batch, _) => + val batchSession = batch.sparkSession + batch.createOrReplaceTempView("mapping_events_batch") + try { + val violations = batchSession.sql(s""" + WITH current_mappings AS ( + SELECT tenant, before_key, target, observed FROM mapping_events_batch + WHERE _change_type IN ('insert', 'update_postimage') + ), historical_mappings AS ( + SELECT r.* FROM $mappings r + JOIN (SELECT DISTINCT tenant, before_key FROM current_mappings) k + ON r.tenant = k.tenant AND r.before_key = k.before_key + WHERE r.observed < (SELECT MIN(observed) FROM current_mappings) + ), all_mappings AS ( + SELECT * FROM current_mappings UNION ALL SELECT * FROM historical_mappings + ) + SELECT tenant, before_key, COUNT(DISTINCT target), SORT_ARRAY(COLLECT_SET(target)) + FROM all_mappings GROUP BY tenant, before_key HAVING COUNT(DISTINCT target) > 1 + """) + assert(violations.collect().toSeq == Seq(Row(1, "a", 2L, Seq("new", "old")))) + val plan = violations.queryExecution.executedPlan + assert( + collect(plan) { case s: CometIcebergNativeScanExec => s }.nonEmpty, + plan.toString) + assert( + collect(plan) { + case p if p.nodeName.contains("Comet") && p.nodeName.contains("Join") => p + }.nonEmpty, + plan.toString) + } finally { + batchSession.catalog.dropTempView("mapping_events_batch") + } + } + assert(executedPlan(query).exists(_.isInstanceOf[CometIcebergNativeScanExec])) + } + } + } + + test( + "foreachBatch integrity checks read references, repeat actions, and append Iceberg results") { + withCatalog { (dir, catalog, namespace) => + val events = s"$catalog.$namespace.integrity_events" + val refs = s"$catalog.$namespace.integrity_references" + val result = s"$catalog.$namespace.integrity_violations" + withTable(events, refs, result) { + sql(s"""CREATE TABLE $events ( + tenant INT, entity STRING, session STRUCT, observed TIMESTAMP) + USING iceberg PARTITIONED BY (tenant)""") + sql(s"CREATE TABLE $refs (tenant INT, entity STRING, session_id STRING) USING iceberg") + sql(s"""CREATE TABLE $result ( + kind STRING, tenant INT, entity STRING, events BIGINT, hard BOOLEAN) USING iceberg""") + sql( + s"INSERT INTO $refs VALUES (1, 'known', 'known-session'), (2, 'missing', 'missing-session')") + sql(s"""INSERT INTO $events VALUES + (1, 'known', named_struct('id', 'known-session'), TIMESTAMP '2020-01-01 00:00:00'), + (1, 'known', named_struct('id', 'missing-session'), TIMESTAMP '2020-01-01 00:00:00'), + (1, 'missing', named_struct('id', 'missing-session'), TIMESTAMP '2020-01-01 00:00:00'), + (1, 'missing', named_struct('id', 'missing-session'), TIMESTAMP '2020-01-02 00:00:00'), + (1, 'recent', named_struct('id', 'recent-session'), TIMESTAMP '2100-01-01 00:00:00')""") + val query = runAvailable(spark.readStream.table(events), new File(dir, "integrity")) { + (batch, _) => + assert(!batch.isEmpty) + val batchSession = batch.sparkSession + import org.apache.spark.sql.functions._ + val minTime = batch.select(min("observed")).head().getTimestamp(0) + assert(minTime.toString == "2020-01-01 00:00:00.0") + val tenants = batch.select("tenant").distinct().collect().map(_.getInt(0)) + batch.createOrReplaceTempView("stream_events") + val references = + batchSession.table(refs).filter(col("tenant").isin(tenants.toSeq: _*)) + references.createOrReplaceTempView("stream_references") + try { + // Check entity and nested session references independently, as separate callbacks + // would. A known entity with a missing session must fail only the session check. + Seq(("entity", "entity", "entity"), ("session", "session.id", "session_id")) + .foreach { case (kind, sourceKey, referenceKey) => + val violations = batchSession.sql(s""" + SELECT '$kind' AS kind, e.tenant, e.$sourceKey AS entity, COUNT(*) AS events, + MIN(e.observed) < current_timestamp() - INTERVAL 15 MINUTES AS hard + FROM stream_events e LEFT ANTI JOIN stream_references r + ON e.tenant = r.tenant AND e.$sourceKey = r.$referenceKey + GROUP BY e.tenant, e.$sourceKey""") + assert(violations.count() == 2) + assert(violations.filter("hard").count() == 1) + assert(violations.filter("NOT hard").count() == 1) + val expected = if (kind == "entity") { + Set(Row(kind, 1, "missing", 2L, true), Row(kind, 1, "recent", 1L, false)) + } else { + Set( + Row(kind, 1, "missing-session", 3L, true), + Row(kind, 1, "recent-session", 1L, false)) + } + assert(violations.collect().toSet == expected) + val plan = violations.queryExecution.executedPlan + assert( + collect(plan) { case s: CometIcebergNativeScanExec => s }.nonEmpty, + plan.toString) + assert( + collect(plan) { + case p if p.nodeName.contains("Comet") && p.nodeName.contains("Join") => p + }.nonEmpty, + plan.toString) + violations.writeTo(result).append() + } + } finally { + batchSession.catalog.dropTempView("stream_events") + batchSession.catalog.dropTempView("stream_references") + } + } + assert(executedPlan(query).exists(_.isInstanceOf[CometIcebergNativeScanExec])) + assert(query.recentProgress.map(_.numInputRows).sum >= 5) + // foreachBatch committed through its own session/catalog instance. + spark.catalog.refreshTable(result) + checkAnswer( + spark.table(result), + Seq( + Row("entity", 1, "missing", 2L, true), + Row("entity", 1, "recent", 1L, false), + Row("session", 1, "missing-session", 3L, true), + Row("session", 1, "recent-session", 1L, false))) + } + } + } + + test("native streaming aggregate merges persisted state and restarts with Spark checkpoints") { + withCatalog { (dir, catalog, namespace) => + val table = s"$catalog.$namespace.state_events" + withTable(table) { + sql(s"CREATE TABLE $table (tenant INT, value BIGINT) USING iceberg") + sql(s"INSERT INTO $table VALUES (1, 10), (1, 20), (2, NULL)") + val checkpoint = new File(dir, "state-checkpoint") + def source = spark.readStream + .table(table) + .groupBy("tenant") + .agg(count(lit(1)).as("n"), sum("value").as("total")) + def run(native: Boolean, expected: Set[Row]): StreamingQuery = { + withSQLConf(CometConf.COMET_STREAMING_EXEC_ENABLED.key -> native.toString) { + val rows = ArrayBuffer.empty[Row] + val query = runAvailable(source, checkpoint, "complete") { (batch, _) => + rows.clear() + rows ++= batch.collect() + } + assert(rows.toSet == expected) + val plan = executedPlan(query) + val aggregates = plan.collect { case a: CometHashAggregateExec => a } + if (native) { + assert( + aggregates.exists(_.aggregateExpressions.exists( + _.mode == org.apache.spark.sql.catalyst.expressions.aggregate.PartialMerge)), + plan.toString) + assert(plan.exists(_.getClass.getSimpleName == "StateStoreRestoreExec")) + assert(query.lastProgress.stateOperators.map(_.numRowsTotal).sum == expected.size) + } else assert(aggregates.isEmpty, plan.toString) + query + } + } + run(false, Set(Row(1, 2L, 30L), Row(2, 1L, null))) + sql(s"INSERT INTO $table VALUES (1, 5), (2, 7), (3, 9)") + run(true, Set(Row(1, 3L, 35L), Row(2, 2L, 7L), Row(3, 1L, 9L))) + sql(s"INSERT INTO $table VALUES (3, 1)") + run(false, Set(Row(1, 3L, 35L), Row(2, 2L, 7L), Row(3, 2L, 10L))) + } + } + } + + test("native state merge replays a failed batch without applying it twice") { + withCatalog { (dir, catalog, namespace) => + val table = s"$catalog.$namespace.state_retry" + withTable(table) { + sql(s"CREATE TABLE $table (id INT, value BIGINT) USING iceberg") + sql(s"INSERT INTO $table VALUES (1, 10)") + val checkpoint = new File(dir, "state-retry") + def source = spark.readStream.table(table).groupBy("id").agg(sum("value")) + withSQLConf(CometConf.COMET_STREAMING_EXEC_ENABLED.key -> "true") { + runAvailable(source, checkpoint, "update") { (batch, _) => + assert(batch.collect().toSeq == Seq(Row(1, 10L))) + } + sql(s"INSERT INTO $table VALUES (1, 5)") + val failure = intercept[StreamingQueryException] { + runAvailable(source, checkpoint, "update") { (batch, _) => + assert(batch.collect().toSeq == Seq(Row(1, 15L))) + throw new IllegalStateException("state callback failure") + } + } + assert(failure.getMessage.contains("state callback failure")) + val replay = runAvailable(source, checkpoint, "update") { (batch, _) => + assert(batch.collect().toSeq == Seq(Row(1, 15L))) + } + assert(executedPlan(replay).exists(_.isInstanceOf[CometHashAggregateExec])) + } + } + } + } + + test("native window state preserves late-event filtering and watermark eviction") { + withCatalog { (dir, catalog, namespace) => + val table = s"$catalog.$namespace.window_events" + withTable(table) { + sql(s"CREATE TABLE $table (tenant INT, time TIMESTAMP) USING iceberg") + def source = spark.readStream + .table(table) + .withWatermark("time", "10 seconds") + .groupBy(window(col("time"), "10 seconds"), col("tenant")) + .count() + val actual = ArrayBuffer.empty[Row] + val expected = ArrayBuffer.empty[Row] + for ((second, index) <- Seq(1, 25, 2, 45).zipWithIndex) { + sql(s"INSERT INTO $table VALUES (1, TIMESTAMP '2026-01-01 00:00:$second')") + for (native <- Seq(false, true)) { + withSQLConf(CometConf.COMET_STREAMING_EXEC_ENABLED.key -> native.toString) { + val query = runAvailable(source, new File(dir, s"window-$native")) { (batch, _) => + val rows = batch.collect() + if (native) actual ++= rows else expected ++= rows + } + if (native) { + assert( + executedPlan(query).exists(_.isInstanceOf[CometHashAggregateExec]), + executedPlan(query).toString) + } + } + } + assert(actual.toSeq == expected.toSeq, s"window mismatch after input $index") + } + // Iceberg 1.10 falls back to one batch per trigger, so it does not run the extra + // no-data batch that advances eviction before the late event arrives. Both engines + // must follow that runtime's schedule, as checked after every input above. + assert(actual.size == (if (icebergVersionAtLeast(1, 11)) 2 else 1)) + assert(actual.forall(_.getLong(2) == (if (icebergVersionAtLeast(1, 11)) 1L else 2L))) + } + } + } + + test("streaming aggregates with incompatible checkpoint buffers retain Spark state") { + withCatalog { (dir, catalog, namespace) => + val table = s"$catalog.$namespace.collected_state" + withTable(table) { + sql(s"CREATE TABLE $table (id INT, value STRING) USING iceberg") + sql(s"INSERT INTO $table VALUES (1, 'a')") + def source = spark.readStream.table(table).groupBy("id").agg(collect_set("value")) + val checkpoint = new File(dir, "collected-state") + for (native <- Seq(true, false)) { + if (!native) sql(s"INSERT INTO $table VALUES (1, 'b')") + withSQLConf(CometConf.COMET_STREAMING_EXEC_ENABLED.key -> native.toString) { + val query = runAvailable(source, checkpoint, "complete") { (batch, _) => + assert( + batch.collect().head.getSeq[String](1).toSet == + (if (native) Set("a") else Set("a", "b"))) + } + assert(!executedPlan(query).exists(_.isInstanceOf[CometHashAggregateExec])) + } + } + } + } + } + + test("streaming execution keeps continuous sources and disabled no-data batches in Spark") { + val attr = org.apache.spark.sql.catalyst.expressions + .AttributeReference("id", org.apache.spark.sql.types.IntegerType)() + val empty = org.apache.spark.sql.execution.LocalTableScanExec(Seq(attr), Seq.empty, None) + empty.setLogicalLink( + org.apache.spark.sql.catalyst.plans.logical.LocalRelation(Seq(attr), isStreaming = true)) + withSQLConf( + CometConf.COMET_STREAMING_EXEC_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + assert(org.apache.comet.rules.CometExecRule(spark).apply(empty) eq empty) + } + val continuousClass = classOf[org.apache.spark.sql.connector.read.streaming.ContinuousStream] + val continuous = java.lang.reflect.Proxy + .newProxyInstance( + continuousClass.getClassLoader, + Array(continuousClass), + (_, _, _) => + throw new AssertionError("Continuous source must not be executed by the rule")) + .asInstanceOf[org.apache.spark.sql.connector.read.streaming.ContinuousStream] + val scan = org.apache.spark.sql.execution.RDDScanExec( + Seq(attr), + spark.sparkContext.emptyRDD[org.apache.spark.sql.catalyst.InternalRow], + "continuous", + stream = Some(continuous)) + withSQLConf( + CometConf.COMET_STREAMING_EXEC_ENABLED.key -> "true", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true") { + assert(org.apache.comet.rules.CometExecRule(spark).apply(scan) eq scan) + } + } +}