vendor/iceberg-datafusion-arrow59/src/physical_plan/merge.rs
// 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.
//! `SQL MERGE INTO` over an Iceberg table, by **copy-on-write**.
//!
//! DataFusion 55 plans `MERGE INTO` end to end and hands a `TableProvider` the
//! combined target-then-source schema, the `ON` condition and the `WHEN`
//! clauses (`TableProvider::merge_into`, `datafusion-session/src/table.rs`).
//! Everything below that — deciding which rows change, rewriting the data files
//! that hold them, and committing the swap — is the provider's job. This module
//! is that job.
//!
//! # The strategy: copy-on-write, at data-file granularity
//!
//! For every data file live in the current snapshot:
//!
//! 1. read its rows,
//! 2. FULL-join them against the source on the `ON` equality keys,
//! 3. decide, per row, which `WHEN` clause fires,
//! 4. if **no** row of the file is updated or deleted, leave the file alone —
//! it is carried into the new snapshot byte-identical,
//! 5. otherwise write the file's surviving (and updated) rows to a NEW data
//! file, and mark the old one for removal.
//!
//! Source rows that matched no target row anywhere become inserted rows in one
//! more new data file. Then one
//! [`Transaction::overwrite_files`](iceberg::transaction::Transaction::overwrite_files)
//! commit removes the replaced files and adds the new ones — atomically, in one
//! catalog pointer swap.
//!
//! The alternative, merge-on-read, writes an equality-delete file instead of
//! rewriting data (`skade::Table::delete_equality` does exactly that). It is
//! cheaper to write and more expensive to read, and every merge makes the next
//! read more expensive. Copy-on-write pays once, at write time, and leaves a
//! table that reads at full speed. See `.nornir/merge-into-design.md` for the
//! measurement behind choosing it as the default.
//!
//! # The count
//!
//! The single `count` row this plan emits is `updated + deleted + inserted`,
//! and every term is counted off the work that produced the commit — never off
//! the input plan. A merge whose source has four rows of which two match
//! returns 2, not 4.
//!
//! # What it refuses, by name
//!
//! - a non-equality `ON` condition (there is nothing to build a hash join on),
//! - a partitioned table (the partition values would have to be recomputed for
//! every rewritten row; `IcebergWriteExec` gets them from a projection the
//! merge path does not build),
//! - `WHEN NOT MATCHED ... THEN UPDATE/DELETE` and
//! `WHEN NOT MATCHED BY SOURCE THEN INSERT`, which name no row to act on,
//! - a target row matched by more than one source row.
//!
//! It never silently no-ops.
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fmt::{Debug, Formatter};
use std::str::FromStr;
use std::sync::Arc;
use datafusion::arrow::array::{
Array, ArrayRef, BooleanArray, RecordBatch, UInt64Array, new_null_array,
};
use datafusion::arrow::compute::kernels::zip::zip;
use datafusion::arrow::compute::take;
use datafusion::arrow::datatypes::{
DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef,
};
use datafusion::common::{DFSchemaRef, DataFusionError, JoinType, NullEquality};
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::logical_expr::dml::{MergeIntoAction, MergeIntoClause, MergeIntoClauseKind};
use datafusion::logical_expr::{BinaryExpr, Expr, Operator};
use datafusion::physical_expr::expressions::Column as PhysicalColumn;
use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
};
use futures::StreamExt;
use iceberg::arrow::FieldMatchMode;
use iceberg::scan::FileScanTask;
use iceberg::spec::{DataFile, DataFileFormat};
use iceberg::table::Table;
use iceberg::transaction::{ApplyTransactionAction, Transaction};
use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
use iceberg::writer::file_writer::ParquetWriterBuilder;
use iceberg::writer::file_writer::location_generator::{
DefaultFileNameGenerator, DefaultLocationGenerator,
};
use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
use iceberg::{Catalog, Error, ErrorKind, TableIdent};
use uuid::Uuid;
use crate::task_writer::TaskWriter;
use crate::to_datafusion_error;
/// The synthetic non-null ordinal appended to each side of the join so an
/// unmatched row can be told from a matched row whose columns are all NULL.
/// Appended at the END of each side, so the real columns keep the indices the
/// combined merge schema gives them.
const TARGET_ORDINAL: &str = "__iceberg_merge_target_ordinal";
const SOURCE_ORDINAL: &str = "__iceberg_merge_source_ordinal";
/// One `WHEN` clause, with every expression already turned into a
/// [`PhysicalExpr`] over the combined merge schema.
///
/// The compilation happens at PLAN time, inside `merge_into`, using the caller's
/// own `&dyn Session` — that is the only place a user-defined function in a
/// clause can still be resolved.
pub(crate) struct CompiledClause {
/// `Matched`, `NotMatchedByTarget` (both SQL spellings collapse here via
/// [`MergeIntoClauseKind::canonical`]) or `NotMatchedBySource`.
kind: MergeIntoClauseKind,
predicate: Option<Arc<dyn PhysicalExpr>>,
action: CompiledAction,
}
pub(crate) enum CompiledAction {
/// `(target column index, value expression)`.
Update(Vec<(usize, Arc<dyn PhysicalExpr>)>),
/// One entry per target column; `None` means "not named, insert NULL".
Insert(Vec<Option<Arc<dyn PhysicalExpr>>>),
Delete,
}
/// The execution plan behind `MERGE INTO <iceberg table>`.
///
/// Emits exactly one row, `count: UInt64`.
pub(crate) struct IcebergMergeExec {
catalog: Arc<dyn Catalog>,
table_ident: TableIdent,
/// The target table's arrow schema — the first `n_target` columns of the
/// merge schema, and the schema every rewritten data file is written with.
target_schema: ArrowSchemaRef,
/// Target columns followed by source columns, exactly as DataFusion built
/// it. Every compiled expression below is resolved against this.
merge_schema: ArrowSchemaRef,
source: Arc<dyn ExecutionPlan>,
/// `(target column index, source column index)` for each `ON` equality.
equi: Vec<(usize, usize)>,
clauses: Vec<CompiledClause>,
count_schema: ArrowSchemaRef,
plan_properties: Arc<PlanProperties>,
}
impl Debug for IcebergMergeExec {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "IcebergMergeExec: table={}", self.table_ident)
}
}
impl IcebergMergeExec {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
catalog: Arc<dyn Catalog>,
table_ident: TableIdent,
target_schema: ArrowSchemaRef,
merge_schema: ArrowSchemaRef,
source: Arc<dyn ExecutionPlan>,
equi: Vec<(usize, usize)>,
clauses: Vec<CompiledClause>,
) -> Self {
let count_schema = Self::make_count_schema();
let plan_properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(Arc::clone(&count_schema)),
Partitioning::UnknownPartitioning(1),
EmissionType::Final,
Boundedness::Bounded,
));
Self {
catalog,
table_ident,
target_schema,
merge_schema,
source,
equi,
clauses,
count_schema,
plan_properties,
}
}
fn make_count_schema() -> ArrowSchemaRef {
Arc::new(ArrowSchema::new(vec![Field::new(
"count",
DataType::UInt64,
false,
)]))
}
fn make_count_batch(count: u64) -> DFResultBatch {
let count_array = Arc::new(UInt64Array::from(vec![count])) as ArrayRef;
RecordBatch::try_new(Self::make_count_schema(), vec![count_array]).map_err(|e| {
DataFusionError::ArrowError(Box::new(e), Some("Failed to make count batch".into()))
})
}
}
type DFResultBatch = datafusion::common::Result<RecordBatch>;
type DFResult<T> = datafusion::common::Result<T>;
/// Compile the `ON` condition into `(target column index, source column index)`
/// pairs, or refuse by name.
///
/// A hash join needs equalities. Everything else — `t.a > s.a`, `t.a = s.a OR
/// t.b = s.b`, a function of both sides — has no key to build on, and silently
/// degrading to a nested loop over a warehouse table is not a kindness.
pub(crate) fn equi_keys_from_on(
on: &Expr,
merge_schema: &DFSchemaRef,
n_target: usize,
) -> DFResult<Vec<(usize, usize)>> {
let mut conjuncts = Vec::new();
flatten_and(on, &mut conjuncts);
let mut equi = Vec::new();
for c in conjuncts {
let Expr::BinaryExpr(BinaryExpr { left, op, right }) = c else {
return Err(unsupported_on(on));
};
if *op != Operator::Eq {
return Err(unsupported_on(on));
}
let (Expr::Column(lc), Expr::Column(rc)) = (left.as_ref(), right.as_ref()) else {
return Err(unsupported_on(on));
};
let li = merge_schema.index_of_column(lc)?;
let ri = merge_schema.index_of_column(rc)?;
let pair = match (li < n_target, ri < n_target) {
(true, false) => (li, ri - n_target),
(false, true) => (ri, li - n_target),
_ => return Err(unsupported_on(on)),
};
equi.push(pair);
}
if equi.is_empty() {
return Err(unsupported_on(on));
}
Ok(equi)
}
fn flatten_and<'a>(e: &'a Expr, out: &mut Vec<&'a Expr>) {
if let Expr::BinaryExpr(BinaryExpr {
left,
op: Operator::And,
right,
}) = e
{
flatten_and(left, out);
flatten_and(right, out);
} else {
out.push(e);
}
}
fn unsupported_on(on: &Expr) -> DataFusionError {
to_datafusion_error(Error::new(
ErrorKind::FeatureUnsupported,
format!(
"MERGE INTO on an Iceberg table needs an ON condition that is a conjunction of \
`target.col = source.col` equalities; got `{on}`. Refusing rather than falling back \
to a nested-loop scan of the whole table."
),
))
}
impl DisplayAs for IcebergMergeExec {
fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
match t {
DisplayFormatType::Default | DisplayFormatType::TreeRender => {
write!(f, "IcebergMergeExec: table={}", self.table_ident)
}
DisplayFormatType::Verbose => write!(
f,
"IcebergMergeExec: table={}, clauses={}, keys={:?}",
self.table_ident,
self.clauses.len(),
self.equi
),
}
}
}
impl ExecutionPlan for IcebergMergeExec {
fn name(&self) -> &str {
"IcebergMergeExec"
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.plan_properties
}
/// No physical expressions to walk from the framework's side: the clause
/// expressions this node carries are evaluated by hand against batches it
/// builds itself, not against a child's output. Same answer the sibling
/// write/commit nodes give.
fn apply_expressions(
&self,
_f: &mut dyn for<'a> FnMut(
&'a Arc<dyn PhysicalExpr>,
) -> datafusion::error::Result<
datafusion::common::tree_node::TreeNodeRecursion,
>,
) -> datafusion::error::Result<datafusion::common::tree_node::TreeNodeRecursion> {
Ok(datafusion::common::tree_node::TreeNodeRecursion::Continue)
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![&self.source]
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> DFResult<Arc<dyn ExecutionPlan>> {
if children.len() != 1 {
return Err(DataFusionError::Internal(format!(
"IcebergMergeExec expects exactly one child, but provided {}",
children.len()
)));
}
Ok(Arc::new(Self {
catalog: Arc::clone(&self.catalog),
table_ident: self.table_ident.clone(),
target_schema: Arc::clone(&self.target_schema),
merge_schema: Arc::clone(&self.merge_schema),
source: Arc::clone(&children[0]),
equi: self.equi.clone(),
clauses: self.clauses.clone(),
count_schema: Arc::clone(&self.count_schema),
plan_properties: Arc::clone(&self.plan_properties),
}))
}
fn execute(
&self,
partition: usize,
context: Arc<TaskContext>,
) -> DFResult<SendableRecordBatchStream> {
if partition != 0 {
return Err(DataFusionError::Internal(format!(
"IcebergMergeExec only has one partition, but got partition {partition}"
)));
}
let plan = MergePlan {
catalog: Arc::clone(&self.catalog),
table_ident: self.table_ident.clone(),
target_schema: Arc::clone(&self.target_schema),
merge_schema: Arc::clone(&self.merge_schema),
source: Arc::clone(&self.source),
equi: self.equi.clone(),
clauses: self.clauses.clone(),
};
let count_schema = Arc::clone(&self.count_schema);
let stream = futures::stream::once(async move {
let count = plan.run(context).await?;
Self::make_count_batch(count)
})
.boxed();
Ok(Box::pin(RecordBatchStreamAdapter::new(
count_schema,
stream,
)))
}
}
/// The owned half of [`IcebergMergeExec`], moved into the execution future.
struct MergePlan {
catalog: Arc<dyn Catalog>,
table_ident: TableIdent,
target_schema: ArrowSchemaRef,
merge_schema: ArrowSchemaRef,
source: Arc<dyn ExecutionPlan>,
equi: Vec<(usize, usize)>,
clauses: Vec<CompiledClause>,
}
/// What one joined row does to the table.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Outcome {
Keep,
Update,
Delete,
}
impl MergePlan {
async fn run(self, context: Arc<TaskContext>) -> DFResult<u64> {
// Fresh metadata: the plan may have been built against an older load.
let table = self
.catalog
.load_table(&self.table_ident)
.await
.map_err(to_datafusion_error)?;
let n_target = self.target_schema.fields().len();
let source_schema = self.source.schema();
// ── 1. the source, once, in memory ────────────────────────────────
let source_batches = self.collect_source(Arc::clone(&context)).await?;
let source_rows: usize = source_batches.iter().map(|b| b.num_rows()).sum();
let (source_with_ord, source_ord_schema) =
with_ordinal(&source_batches, &source_schema, SOURCE_ORDINAL)?;
// ── 2. every live data file, one at a time ────────────────────────
let file_tasks = plan_files_by_path(&table).await?;
let mut removed_paths: Vec<String> = Vec::new();
let mut added_files: Vec<DataFile> = Vec::new();
let mut matched_source: HashSet<u64> = HashSet::new();
let mut updated: u64 = 0;
let mut deleted: u64 = 0;
for (path, tasks) in file_tasks {
let file_batches = read_file(&table, tasks, &self.target_schema).await?;
let (file_with_ord, file_ord_schema) =
with_ordinal(&file_batches, &self.target_schema, TARGET_ORDINAL)?;
let joined = self
.join(
Arc::clone(&context),
&file_with_ord,
&file_ord_schema,
&source_with_ord,
&source_ord_schema,
)
.await?;
let mut kept: Vec<RecordBatch> = Vec::new();
let mut file_updated: u64 = 0;
let mut file_deleted: u64 = 0;
// Across ALL batches of this file's join: a target row may only be
// matched once.
let mut seen_target: HashSet<u64> = HashSet::new();
for batch in &joined {
let (merge_batch, tgt_ord, src_ord) = split_joined(batch, n_target)?;
let chosen = self.choose_clauses(&merge_batch, &tgt_ord, &src_ord)?;
// Record which source rows found a target row anywhere. Only
// rows with nothing left to match become inserts.
for i in 0..merge_batch.num_rows() {
if !tgt_ord.is_null(i) && !src_ord.is_null(i) {
matched_source.insert(src_ord.value(i));
}
}
let new_cols = self.updated_target_columns(&merge_batch, &chosen, n_target)?;
// Rows of THIS file, in original file order, minus deletions.
let mut order: Vec<(u64, usize)> = Vec::new();
for i in 0..merge_batch.num_rows() {
if tgt_ord.is_null(i) {
continue;
}
let ord = tgt_ord.value(i);
if !seen_target.insert(ord) {
return Err(to_datafusion_error(Error::new(
ErrorKind::DataInvalid,
format!(
"MERGE INTO matched the same target row of {path} with more than \
one source row; the result would depend on which one won. Narrow \
the ON condition or de-duplicate the source."
),
)));
}
let outcome = match chosen[i] {
None => Outcome::Keep,
Some(c) => match self.clauses[c].action {
CompiledAction::Delete => Outcome::Delete,
CompiledAction::Update(_) => Outcome::Update,
// An Insert can only be chosen for a row with no
// target side, which this branch excludes.
CompiledAction::Insert(_) => Outcome::Keep,
},
};
match outcome {
Outcome::Delete => file_deleted += 1,
Outcome::Update => {
file_updated += 1;
order.push((ord, i));
}
Outcome::Keep => order.push((ord, i)),
}
}
order.sort_unstable();
let idx = UInt64Array::from(
order.iter().map(|(_, i)| *i as u64).collect::<Vec<u64>>(),
);
let cols = new_cols
.iter()
.map(|c| take(c.as_ref(), &idx, None))
.collect::<Result<Vec<ArrayRef>, _>>()
.map_err(|e| {
DataFusionError::ArrowError(Box::new(e), Some("merge take".into()))
})?;
if !cols.is_empty() && cols[0].len() > 0 {
kept.push(
RecordBatch::try_new(Arc::clone(&self.target_schema), cols).map_err(
|e| {
DataFusionError::ArrowError(
Box::new(e),
Some("merge rewrite batch".into()),
)
},
)?,
);
}
}
// THE GUARD THAT MAKES A NO-OP MERGE FREE: a file no row of which
// was updated or deleted is not rewritten, and is carried into the
// new snapshot untouched.
if file_updated == 0 && file_deleted == 0 {
continue;
}
updated += file_updated;
deleted += file_deleted;
removed_paths.push(path);
added_files.extend(write_data_files(&table, kept).await?);
}
// ── 3. source rows that matched nothing become inserts ─────────────
let inserted = self
.insert_pass(
&table,
&source_with_ord,
&source_ord_schema,
&matched_source,
n_target,
&mut added_files,
)
.await?;
// ── 4. one overwrite commit, or none at all ───────────────────────
if removed_paths.is_empty() && added_files.is_empty() {
debug_assert_eq!(updated + deleted + inserted, 0);
return Ok(0);
}
let tx = Transaction::new(&table);
let action = tx
.overwrite_files()
.add_data_files(added_files)
.delete_files(removed_paths);
action
.apply(tx)
.map_err(to_datafusion_error)?
.commit(self.catalog.as_ref())
.await
.map_err(to_datafusion_error)?;
let _ = source_rows;
Ok(updated + deleted + inserted)
}
async fn collect_source(&self, context: Arc<TaskContext>) -> DFResult<Vec<RecordBatch>> {
let partitions = self.source.output_partitioning().partition_count();
let mut out = Vec::new();
for p in 0..partitions {
let mut stream = self.source.execute(p, Arc::clone(&context))?;
while let Some(b) = stream.next().await {
let b = b?;
if b.num_rows() > 0 {
out.push(b);
}
}
}
Ok(out)
}
/// FULL hash join of one data file against the source, on the `ON` keys.
async fn join(
&self,
context: Arc<TaskContext>,
left: &[RecordBatch],
left_schema: &ArrowSchemaRef,
right: &[RecordBatch],
right_schema: &ArrowSchemaRef,
) -> DFResult<Vec<RecordBatch>> {
let left_exec =
MemorySourceConfig::try_new_exec(&[left.to_vec()], Arc::clone(left_schema), None)?;
let right_exec =
MemorySourceConfig::try_new_exec(&[right.to_vec()], Arc::clone(right_schema), None)?;
let on: Vec<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> = self
.equi
.iter()
.map(|(t, s)| {
let lf = left_schema.field(*t);
let rf = right_schema.field(*s);
if lf.data_type() != rf.data_type() {
return Err(to_datafusion_error(Error::new(
ErrorKind::FeatureUnsupported,
format!(
"MERGE INTO ON compares `{}` ({}) with `{}` ({}); the join key types \
must match exactly. Cast in the USING subquery.",
lf.name(),
lf.data_type(),
rf.name(),
rf.data_type()
),
)));
}
Ok((
Arc::new(PhysicalColumn::new(lf.name(), *t)) as Arc<dyn PhysicalExpr>,
Arc::new(PhysicalColumn::new(rf.name(), *s)) as Arc<dyn PhysicalExpr>,
))
})
.collect::<DFResult<Vec<_>>>()?;
let join = HashJoinExec::try_new(
left_exec,
right_exec,
on,
None,
&JoinType::Full,
None,
PartitionMode::CollectLeft,
// SQL `=` never matches NULL to NULL, and neither does MERGE's ON.
NullEquality::NullEqualsNothing,
false,
)?;
let mut stream = join.execute(0, context)?;
let mut out = Vec::new();
while let Some(b) = stream.next().await {
let b = b?;
if b.num_rows() > 0 {
out.push(b);
}
}
Ok(out)
}
/// For every joined row, the index of the first clause of the applicable
/// category whose predicate holds. `None` = no clause fires.
fn choose_clauses(
&self,
merge_batch: &RecordBatch,
tgt_ord: &UInt64Array,
src_ord: &UInt64Array,
) -> DFResult<Vec<Option<usize>>> {
let rows = merge_batch.num_rows();
let mut masks: Vec<Option<BooleanArray>> = Vec::with_capacity(self.clauses.len());
for clause in &self.clauses {
masks.push(match &clause.predicate {
None => None,
Some(p) => Some(eval_bool(p, merge_batch)?),
});
}
let mut chosen = vec![None; rows];
for i in 0..rows {
let has_target = !tgt_ord.is_null(i);
let has_source = !src_ord.is_null(i);
let want = if has_target && has_source {
MergeIntoClauseKind::Matched
} else if has_target {
MergeIntoClauseKind::NotMatchedBySource
} else {
MergeIntoClauseKind::NotMatchedByTarget
};
for (c, clause) in self.clauses.iter().enumerate() {
if clause.kind != want {
continue;
}
let fires = match &masks[c] {
None => true,
// A NULL predicate is not true — same rule as WHERE.
Some(m) => !m.is_null(i) && m.value(i),
};
if fires {
chosen[i] = Some(c);
break;
}
}
}
Ok(chosen)
}
/// The target columns of the joined batch after applying every UPDATE
/// assignment that fires. Column `j` is the base value except on rows whose
/// chosen clause assigns `j`.
fn updated_target_columns(
&self,
merge_batch: &RecordBatch,
chosen: &[Option<usize>],
n_target: usize,
) -> DFResult<Vec<ArrayRef>> {
let mut cols: Vec<ArrayRef> = (0..n_target)
.map(|j| Arc::clone(merge_batch.column(j)))
.collect();
for (c, clause) in self.clauses.iter().enumerate() {
let CompiledAction::Update(assignments) = &clause.action else {
continue;
};
let mask = BooleanArray::from(
chosen.iter().map(|x| *x == Some(c)).collect::<Vec<bool>>(),
);
if mask.true_count() == 0 {
continue;
}
for (j, expr) in assignments {
let value = eval_array(expr, merge_batch)?;
let value = cast_to(&value, cols[*j].data_type())?;
cols[*j] = zip(&mask, &value, &cols[*j]).map_err(|e| {
DataFusionError::ArrowError(Box::new(e), Some("merge update zip".into()))
})?;
}
}
Ok(cols)
}
/// Write the rows of `source` that matched no target row anywhere, through
/// the first `WHEN NOT MATCHED [BY TARGET]` clause that fires.
#[allow(clippy::too_many_arguments)]
async fn insert_pass(
&self,
table: &Table,
source_with_ord: &[RecordBatch],
source_ord_schema: &ArrowSchemaRef,
matched_source: &HashSet<u64>,
n_target: usize,
added_files: &mut Vec<DataFile>,
) -> DFResult<u64> {
if !self
.clauses
.iter()
.any(|c| c.kind == MergeIntoClauseKind::NotMatchedByTarget)
{
return Ok(0);
}
let ord_idx = source_ord_schema.fields().len() - 1;
// The target half of an insert-pass row is all NULL by construction, so
// the batch cannot be built against a schema whose target columns are
// REQUIRED. Only the batch is relaxed; the rows finally written go back
// through the table's own schema.
let nullable_merge_schema = Arc::new(ArrowSchema::new(
self.merge_schema
.fields()
.iter()
.map(|f| f.as_ref().clone().with_nullable(true))
.collect::<Vec<Field>>(),
));
let mut batches: Vec<RecordBatch> = Vec::new();
let mut inserted: u64 = 0;
for batch in source_with_ord {
let ord = batch
.column(ord_idx)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("source ordinal is UInt64");
let keep: Vec<u64> = (0..batch.num_rows())
.filter(|i| !matched_source.contains(&ord.value(*i)))
.map(|i| i as u64)
.collect();
if keep.is_empty() {
continue;
}
let idx = UInt64Array::from(keep);
// A merge-schema batch whose target half is all NULL: that is what
// a source row with no target row IS, and it is the schema every
// clause expression was compiled against.
let mut cols: Vec<ArrayRef> = (0..n_target)
.map(|j| new_null_array(self.merge_schema.field(j).data_type(), idx.len()))
.collect();
for s in 0..ord_idx {
cols.push(take(batch.column(s).as_ref(), &idx, None).map_err(|e| {
DataFusionError::ArrowError(Box::new(e), Some("merge insert take".into()))
})?);
}
let merge_batch = RecordBatch::try_new(Arc::clone(&nullable_merge_schema), cols)
.map_err(|e| {
DataFusionError::ArrowError(Box::new(e), Some("merge insert batch".into()))
})?;
let empty_tgt = UInt64Array::new_null(merge_batch.num_rows());
let present_src = UInt64Array::from(
(0..merge_batch.num_rows()).map(|i| i as u64).collect::<Vec<u64>>(),
);
let chosen = self.choose_clauses(&merge_batch, &empty_tgt, &present_src)?;
// Group the rows by which INSERT clause fired; each group becomes
// its own batch of target rows.
let mut per_clause: BTreeMap<usize, Vec<u64>> = BTreeMap::new();
for (i, c) in chosen.iter().enumerate() {
if let Some(c) = c {
per_clause.entry(*c).or_default().push(i as u64);
}
}
for (c, rows) in per_clause {
let CompiledAction::Insert(values) = &self.clauses[c].action else {
continue;
};
let idx = UInt64Array::from(rows);
let mut out: Vec<ArrayRef> = Vec::with_capacity(n_target);
for (j, v) in values.iter().enumerate() {
let col = match v {
None => new_null_array(self.target_schema.field(j).data_type(), idx.len()),
Some(expr) => {
let full = eval_array(expr, &merge_batch)?;
let picked = take(full.as_ref(), &idx, None).map_err(|e| {
DataFusionError::ArrowError(
Box::new(e),
Some("merge insert value take".into()),
)
})?;
cast_to(&picked, self.target_schema.field(j).data_type())?
}
};
out.push(col);
}
inserted += idx.len() as u64;
batches.push(
RecordBatch::try_new(Arc::clone(&self.target_schema), out).map_err(|e| {
DataFusionError::ArrowError(
Box::new(e),
Some("merge inserted batch".into()),
)
})?,
);
}
}
if inserted > 0 {
added_files.extend(write_data_files(table, batches).await?);
}
Ok(inserted)
}
}
impl Clone for CompiledClause {
fn clone(&self) -> Self {
Self {
kind: self.kind,
predicate: self.predicate.clone(),
action: self.action.clone(),
}
}
}
impl Clone for CompiledAction {
fn clone(&self) -> Self {
match self {
Self::Update(v) => Self::Update(v.clone()),
Self::Insert(v) => Self::Insert(v.clone()),
Self::Delete => Self::Delete,
}
}
}
// ─────────────────────────── helpers ────────────────────────────────────────
fn eval_array(expr: &Arc<dyn PhysicalExpr>, batch: &RecordBatch) -> DFResult<ArrayRef> {
expr.evaluate(batch)?.into_array(batch.num_rows())
}
fn eval_bool(expr: &Arc<dyn PhysicalExpr>, batch: &RecordBatch) -> DFResult<BooleanArray> {
let arr = eval_array(expr, batch)?;
arr.as_any()
.downcast_ref::<BooleanArray>()
.cloned()
.ok_or_else(|| {
DataFusionError::Internal(format!(
"MERGE clause predicate did not evaluate to a boolean, got {}",
arr.data_type()
))
})
}
fn cast_to(array: &ArrayRef, to: &DataType) -> DFResult<ArrayRef> {
if array.data_type() == to {
return Ok(Arc::clone(array));
}
datafusion::arrow::compute::cast(array.as_ref(), to).map_err(|e| {
DataFusionError::ArrowError(Box::new(e), Some("merge value cast".into()))
})
}
/// Append a dense non-null ordinal column to every batch, and return the
/// widened schema. The ordinal runs across batches, so it names a row within
/// the whole side.
fn with_ordinal(
batches: &[RecordBatch],
schema: &ArrowSchemaRef,
name: &str,
) -> DFResult<(Vec<RecordBatch>, ArrowSchemaRef)> {
let mut fields: Vec<Field> = schema.fields().iter().map(|f| f.as_ref().clone()).collect();
fields.push(Field::new(name, DataType::UInt64, false));
let wide = Arc::new(ArrowSchema::new(fields));
let mut next: u64 = 0;
let mut out = Vec::with_capacity(batches.len());
for b in batches {
let ords: Vec<u64> = (0..b.num_rows() as u64).map(|i| next + i).collect();
next += b.num_rows() as u64;
let mut cols = b.columns().to_vec();
cols.push(Arc::new(UInt64Array::from(ords)) as ArrayRef);
out.push(
RecordBatch::try_new(Arc::clone(&wide), cols).map_err(|e| {
DataFusionError::ArrowError(Box::new(e), Some("merge ordinal batch".into()))
})?,
);
}
Ok((out, wide))
}
/// Split a join output batch `[target.., target_ord, source.., source_ord]`
/// back into the merge-schema batch the compiled expressions expect, plus the
/// two ordinal arrays.
fn split_joined(
batch: &RecordBatch,
n_target: usize,
) -> DFResult<(RecordBatch, UInt64Array, UInt64Array)> {
let total = batch.num_columns();
let src_ord_idx = total - 1;
let tgt_ord_idx = n_target;
let ord = |i: usize| -> DFResult<UInt64Array> {
batch
.column(i)
.as_any()
.downcast_ref::<UInt64Array>()
.cloned()
.ok_or_else(|| {
DataFusionError::Internal("merge ordinal column is not UInt64".into())
})
};
let tgt_ord = ord(tgt_ord_idx)?;
let src_ord = ord(src_ord_idx)?;
let mut fields = Vec::with_capacity(total - 2);
let mut cols = Vec::with_capacity(total - 2);
for i in 0..total {
if i == tgt_ord_idx || i == src_ord_idx {
continue;
}
fields.push(batch.schema().field(i).clone());
cols.push(Arc::clone(batch.column(i)));
}
let schema = Arc::new(ArrowSchema::new(fields));
let merge_batch = RecordBatch::try_new(schema, cols).map_err(|e| {
DataFusionError::ArrowError(Box::new(e), Some("merge split batch".into()))
})?;
Ok((merge_batch, tgt_ord, src_ord))
}
/// The live data files of the current snapshot, grouped by path.
async fn plan_files_by_path(table: &Table) -> DFResult<BTreeMap<String, Vec<FileScanTask>>> {
if table.metadata().current_snapshot().is_none() {
return Ok(BTreeMap::new());
}
let scan = table.scan().build().map_err(to_datafusion_error)?;
let mut stream = scan.plan_files().await.map_err(to_datafusion_error)?;
let mut out: BTreeMap<String, Vec<FileScanTask>> = BTreeMap::new();
while let Some(t) = stream.next().await {
let t = t.map_err(to_datafusion_error)?;
out.entry(t.data_file_path().to_string()).or_default().push(t);
}
Ok(out)
}
/// Read one data file's live rows, normalised to the table's arrow schema.
///
/// Going through [`ArrowReader`] rather than a raw parquet open is what makes
/// the merge see the table as a reader does: schema evolution is applied, and
/// any equality-delete file already committed against this data file (the
/// merge-on-read half) has removed its rows before the merge sees them.
async fn read_file(
table: &Table,
tasks: Vec<FileScanTask>,
target_schema: &ArrowSchemaRef,
) -> DFResult<Vec<RecordBatch>> {
let reader = table.reader_builder().build();
let stream = futures::stream::iter(tasks.into_iter().map(Ok)).boxed();
let mut batches = reader.read(stream).map_err(to_datafusion_error)?.stream();
let mut out = Vec::new();
while let Some(b) = batches.next().await {
let b = b.map_err(to_datafusion_error)?;
if b.num_rows() == 0 {
continue;
}
out.push(
RecordBatch::try_new(Arc::clone(target_schema), b.columns().to_vec()).map_err(
|e| {
DataFusionError::ArrowError(
Box::new(e),
Some("merge read normalise".into()),
)
},
)?,
);
}
Ok(out)
}
/// Write `batches` as new Iceberg data files, through the same writer stack
/// `IcebergWriteExec` uses for `INSERT`.
async fn write_data_files(table: &Table, batches: Vec<RecordBatch>) -> DFResult<Vec<DataFile>> {
if batches.iter().all(|b| b.num_rows() == 0) {
return Ok(vec![]);
}
// The same writer construction as `IcebergWriteExec::execute` (INSERT):
// table properties parse lazily since iceberg-rust main, the Parquet writer
// honours the table's write.parquet.* properties and its encryption manager.
let table_props = table.metadata().table_properties();
let write_format_default = table_props
.write_format_default()
.map_err(to_datafusion_error)?;
let file_format =
DataFileFormat::from_str(&write_format_default).map_err(to_datafusion_error)?;
if file_format != DataFileFormat::Parquet {
return Err(to_datafusion_error(Error::new(
ErrorKind::FeatureUnsupported,
format!("MERGE INTO cannot rewrite {file_format} data files, only Parquet"),
)));
}
let mut parquet_builder =
ParquetWriterBuilder::from_table_properties(&table_props, table.metadata().current_schema().clone())
.map_err(to_datafusion_error)?
.with_match_mode(FieldMatchMode::Name);
if let Some(encryption_manager) = table.encryption_manager() {
parquet_builder = parquet_builder.with_encryption_manager(encryption_manager.clone());
}
let target_file_size = table_props
.write_target_file_size_bytes()
.map_err(to_datafusion_error)?;
let rolling = RollingFileWriterBuilder::new(
parquet_builder,
target_file_size,
table.file_io().clone(),
DefaultLocationGenerator::new(table.metadata()).map_err(to_datafusion_error)?,
DefaultFileNameGenerator::new(Uuid::now_v7().to_string(), None, file_format),
);
let fanout_enabled = table_props
.write_datafusion_fanout_enabled()
.map_err(to_datafusion_error)?;
let mut writer = TaskWriter::try_new(
DataFileWriterBuilder::new(rolling),
fanout_enabled,
table.metadata().current_schema().clone(),
table.metadata().default_partition_spec().clone(),
)?;
for b in batches {
if b.num_rows() == 0 {
continue;
}
writer.write(b).await?;
}
writer.close().await
}
/// Turn the logical `WHEN` clauses into [`CompiledClause`]s, or refuse by name.
///
/// `compile` is a closure over the caller's `&dyn Session` so a user-defined
/// function inside a clause still resolves — this is the only point at which
/// the session is in scope.
pub(crate) fn compile_clauses<F>(
clauses: Vec<MergeIntoClause>,
target_schema: &ArrowSchemaRef,
mut compile: F,
) -> DFResult<Vec<CompiledClause>>
where
F: FnMut(Expr) -> DFResult<Arc<dyn PhysicalExpr>>,
{
let index_of = |name: &str| -> DFResult<usize> {
target_schema.index_of(name).map_err(|_| {
DataFusionError::Plan(format!("MERGE names target column `{name}`, which does not exist"))
})
};
let mut out = Vec::with_capacity(clauses.len());
for clause in clauses {
let kind = clause.kind.canonical();
let predicate = clause.predicate.map(&mut compile).transpose()?;
let action = match (kind, clause.action) {
(MergeIntoClauseKind::Matched, MergeIntoAction::Update(assignments))
| (MergeIntoClauseKind::NotMatchedBySource, MergeIntoAction::Update(assignments)) => {
let mut compiled = Vec::with_capacity(assignments.len());
for (name, expr) in assignments {
compiled.push((index_of(&name)?, compile(expr)?));
}
CompiledAction::Update(compiled)
}
(MergeIntoClauseKind::Matched, MergeIntoAction::Delete)
| (MergeIntoClauseKind::NotMatchedBySource, MergeIntoAction::Delete) => {
CompiledAction::Delete
}
(MergeIntoClauseKind::NotMatchedByTarget, MergeIntoAction::Insert { columns, values }) => {
let mut by_index: HashMap<usize, Arc<dyn PhysicalExpr>> = HashMap::new();
if columns.is_empty() {
for (j, expr) in values.into_iter().enumerate() {
by_index.insert(j, compile(expr)?);
}
} else {
for (name, expr) in columns.iter().zip(values) {
by_index.insert(index_of(name)?, compile(expr)?);
}
}
CompiledAction::Insert(
(0..target_schema.fields().len())
.map(|j| by_index.remove(&j))
.collect(),
)
}
(kind, action) => {
let action = match action {
MergeIntoAction::Update(_) => "UPDATE",
MergeIntoAction::Insert { .. } => "INSERT",
MergeIntoAction::Delete => "DELETE",
};
let kind = match kind {
MergeIntoClauseKind::Matched => "WHEN MATCHED",
MergeIntoClauseKind::NotMatchedBySource => "WHEN NOT MATCHED BY SOURCE",
_ => "WHEN NOT MATCHED [BY TARGET]",
};
return Err(to_datafusion_error(Error::new(
ErrorKind::FeatureUnsupported,
format!(
"MERGE INTO on an Iceberg table does not support `{kind} THEN {action}`: \
it names no row to act on. Refusing rather than silently doing nothing."
),
)));
}
};
out.push(CompiledClause {
kind,
predicate,
action,
});
}
Ok(out)
}