lance 14.0.1

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

use std::collections::HashSet;
use std::sync::{
    Arc, Mutex,
    atomic::{AtomicU64, Ordering},
};

use arrow_array::{Array, RecordBatch, UInt8Array, UInt64Array, cast::AsArray, types::UInt64Type};
use arrow_schema::Schema;
use arrow_select;
use datafusion::common::{DataFusionError, Result as DFResult};
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::{
    execution::{SendableRecordBatchStream, TaskContext},
    physical_plan::{
        DisplayAs, ExecutionPlan, PlanProperties,
        execution_plan::{Boundedness, EmissionType},
        stream::RecordBatchStreamAdapter,
    },
};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
use futures::{StreamExt, stream};
use lance_arrow::RecordBatchExt;
use lance_core::{Error, ROW_ADDR, ROW_CREATED_AT_VERSION, ROW_ID};
use lance_table::format::{Fragment, RowDatasetVersionSequence, RowIdMeta};
use lance_table::rowids::version::rechunk_version_sequences;
use lance_table::rowids::{RowIdSequence, rechunk_sequences, write_row_ids};
use roaring::RoaringTreemap;

use crate::dataset::rowids::place_carried_row_lineage;
use crate::dataset::transaction::UpdateMode::RewriteRows;
use crate::dataset::utils::CapturedRowIds;
use crate::dataset::write::merge_insert::inserted_rows::{
    KeyExistenceFilter, KeyExistenceFilterBuilder, extract_key_value_from_batch,
};
use crate::dataset::write::merge_insert::{
    InsertedKeyTracker, MERGE_SOURCE_SENTINEL, SourceDedupeBehavior, canonical_source_schema,
    create_duplicate_row_error, format_key_values_on_columns, resolve_target_bases,
};
use crate::{
    Dataset,
    dataset::{
        transaction::{Operation, Transaction},
        write::{
            TargetBaseInfo, WriteParams, blob_v2_external_base_resolver, cleanup_data_fragments,
            merge_insert::{
                MERGE_ACTION_COLUMN, MergeInsertParams, MergeStats, assign_action::Action,
                exec::MergeInsertMetrics,
            },
            write_fragments_internal,
        },
    },
};

use super::apply_deletions;

/// Consumers own the producer task so cancelling a merge also stops polling
/// its source, including when the producer is waiting on a full output queue.
struct SplitterTask(tokio::task::JoinHandle<()>);

impl Drop for SplitterTask {
    fn drop(&mut self) {
        self.0.abort();
    }
}

/// The serial, inline writer needs an unbounded insert queue because it drains
/// updates first. Spilling uses two concurrent writers and one queued batch per
/// group, so source size does not determine buffered insert data.
#[derive(Clone)]
enum MergeWriteSender {
    Unbounded(tokio::sync::mpsc::UnboundedSender<DFResult<RecordBatch>>),
    Bounded(tokio::sync::mpsc::Sender<DFResult<RecordBatch>>),
}

impl MergeWriteSender {
    fn channel(
        bounded: bool,
    ) -> (
        Self,
        futures::stream::BoxStream<'static, DFResult<RecordBatch>>,
    ) {
        if bounded {
            let (sender, receiver) = tokio::sync::mpsc::channel(1);
            (
                Self::Bounded(sender),
                tokio_stream::wrappers::ReceiverStream::new(receiver).boxed(),
            )
        } else {
            let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
            (
                Self::Unbounded(sender),
                tokio_stream::wrappers::UnboundedReceiverStream::new(receiver).boxed(),
            )
        }
    }

    async fn send(
        &self,
        batch: DFResult<RecordBatch>,
    ) -> Result<(), tokio::sync::mpsc::error::SendError<DFResult<RecordBatch>>> {
        match self {
            Self::Unbounded(sender) => sender.send(batch),
            Self::Bounded(sender) => sender.send(batch).await,
        }
    }
}

/// Shared state for merge insert operations to simplify lock management
struct MergeState {
    /// Row addresses that need to be deleted, due to a row update or delete action
    delete_row_addrs: RoaringTreemap,
    /// Shared collection to capture row ids that need to be updated
    updating_row_ids: Arc<Mutex<CapturedRowIds>>,
    /// Track keys of newly inserted rows (not updates).
    inserted_rows_filter: KeyExistenceFilterBuilder,
    /// Merge operation metrics
    metrics: MergeInsertMetrics,
    /// Whether the dataset uses stable row ids.
    stable_row_ids: bool,
    /// Set to track processed row IDs to detect duplicates
    processed_row_ids: HashSet<u64>,
    /// Set to track non-null keys of rows inserted by FirstSeen mode
    processed_insert_keys: InsertedKeyTracker,
    /// The "on" column names for merge operation
    on_columns: Vec<String>,
    /// How to handle duplicate source rows
    source_dedupe_behavior: SourceDedupeBehavior,
}

impl MergeState {
    fn new(
        metrics: MergeInsertMetrics,
        stable_row_ids: bool,
        on_columns: Vec<String>,
        field_ids: Vec<i32>,
        source_dedupe_behavior: SourceDedupeBehavior,
    ) -> Self {
        Self {
            delete_row_addrs: RoaringTreemap::new(),
            updating_row_ids: Arc::new(Mutex::new(CapturedRowIds::new(stable_row_ids))),
            inserted_rows_filter: KeyExistenceFilterBuilder::new(field_ids),
            metrics,
            stable_row_ids,
            processed_row_ids: HashSet::new(),
            processed_insert_keys: InsertedKeyTracker::default(),
            on_columns,
            source_dedupe_behavior,
        }
    }

    /// Process a single row based on its action, updating internal state
    fn process_row_action(
        &mut self,
        action: Action,
        row_idx: usize,
        row_addr_array: &UInt64Array,
        row_id_array: &UInt64Array,
        batch: &RecordBatch,
    ) -> DFResult<Option<usize>> {
        match action {
            Action::Delete => {
                // Delete action - only delete, don't write back
                if !row_addr_array.is_null(row_idx) {
                    let row_addr = row_addr_array.value(row_idx);
                    let row_id = row_id_array.value(row_idx);

                    // A source with duplicate keys matches the same target row
                    // more than once; apply the same dedupe policy as updates.
                    // (Target-only deletes from `delete_not_matched_by_source`
                    // also reach here but never duplicate, so they never trip
                    // `Fail`.)
                    if !self.processed_row_ids.insert(row_id) {
                        match self.source_dedupe_behavior {
                            SourceDedupeBehavior::Fail => {
                                return Err(create_duplicate_row_error(
                                    batch,
                                    row_idx,
                                    &self.on_columns,
                                ));
                            }
                            SourceDedupeBehavior::FirstSeen => {
                                self.metrics.num_skipped_duplicates.add(1);
                                return Ok(None); // Skip this duplicate row
                            }
                        }
                    }

                    self.delete_row_addrs.insert(row_addr);
                    self.metrics.num_deleted_rows.add(1);
                }
                Ok(None) // Don't keep this row
            }
            Action::UpdateAll => {
                // Update action - delete old row AND insert new data
                if !row_addr_array.is_null(row_idx) {
                    let row_addr = row_addr_array.value(row_idx);
                    let row_id = row_id_array.value(row_idx);

                    // Check for duplicate _rowid in the current merge operation
                    if !self.processed_row_ids.insert(row_id) {
                        match self.source_dedupe_behavior {
                            SourceDedupeBehavior::Fail => {
                                return Err(create_duplicate_row_error(
                                    batch,
                                    row_idx,
                                    &self.on_columns,
                                ));
                            }
                            SourceDedupeBehavior::FirstSeen => {
                                self.metrics.num_skipped_duplicates.add(1);
                                return Ok(None); // Skip this duplicate row
                            }
                        }
                    }

                    self.delete_row_addrs.insert(row_addr);

                    if self.stable_row_ids {
                        self.updating_row_ids.lock().unwrap().capture(&[row_id])?;
                    }
                    // Don't count as actual delete - this is an update
                }

                self.metrics.num_updated_rows.add(1);
                Ok(Some(row_idx)) // Keep this row for writing
            }
            Action::Insert => {
                if self.source_dedupe_behavior == SourceDedupeBehavior::FirstSeen
                    && !self
                        .processed_insert_keys
                        .insert(batch, row_idx, &self.on_columns)?
                {
                    self.metrics.num_skipped_duplicates.add(1);
                    return Ok(None);
                }

                // Capture the key value for conflict detection (only for inserts, not updates)
                if let Some(key_value) =
                    extract_key_value_from_batch(batch, row_idx, &self.on_columns)
                {
                    self.inserted_rows_filter
                        .insert(key_value)
                        .map_err(|e| DataFusionError::External(Box::new(e)))?;
                }
                self.metrics.num_inserted_rows.add(1);
                Ok(Some(row_idx)) // Keep this row for writing
            }
            Action::Nothing => {
                // Do nothing action - keep the row but don't count it
                Ok(None)
            }
            Action::Fail => {
                // Fail action - return an error to fail the operation
                Err(datafusion::error::DataFusionError::Execution(format!(
                    "Merge insert failed: found matching row with key values: {}",
                    format_key_values_on_columns(batch, row_idx, &self.on_columns)
                )))
            }
        }
    }
}

/// Inserts new rows and updates existing rows in the target table.
///
/// This does the actual write.
///
/// This is implemented by moving updated rows to new fragments. This mode
/// is most optimal when updating the full schema.
///
#[derive(Debug)]
pub struct FullSchemaMergeInsertExec {
    input: Arc<dyn ExecutionPlan>,
    dataset: Arc<Dataset>,
    params: MergeInsertParams,
    properties: Arc<PlanProperties>,
    metrics: ExecutionPlanMetricsSet,
    merge_stats: Arc<Mutex<Option<MergeStats>>>,
    transaction: Arc<Mutex<Option<Transaction>>>,
    affected_rows: Arc<Mutex<Option<RoaringTreemap>>>,
    inserted_rows_filter: Arc<Mutex<Option<KeyExistenceFilter>>>,
    source_skipped_duplicates: Arc<AtomicU64>,
    source_blob_columns: Vec<String>,
    update_if_deferred: bool,
    /// Whether the ON columns match the schema's unenforced primary key.
    /// If true, inserted_rows_filter will be included in the transaction for conflict detection.
    is_primary_key: bool,
}

impl FullSchemaMergeInsertExec {
    pub fn try_new(
        input: Arc<dyn ExecutionPlan>,
        dataset: Arc<Dataset>,
        params: MergeInsertParams,
        source_skipped_duplicates: Arc<AtomicU64>,
        source_blob_columns: Vec<String>,
    ) -> DFResult<Self> {
        let empty_schema = Arc::new(arrow_schema::Schema::empty());
        let properties = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(empty_schema),
            Partitioning::UnknownPartitioning(1),
            EmissionType::Final,
            Boundedness::Bounded,
        ));

        // Check if ON columns match the schema's unenforced primary key
        let field_ids: Vec<i32> = params
            .on
            .iter()
            .filter_map(|name| dataset.schema().field(name).map(|f| f.id))
            .collect();
        let pk_field_ids: Vec<i32> = dataset
            .schema()
            .unenforced_primary_key()
            .iter()
            .map(|f| f.id)
            .collect();
        let is_primary_key = !pk_field_ids.is_empty() && field_ids == pk_field_ids;

        Ok(Self {
            input,
            dataset,
            params,
            properties,
            metrics: ExecutionPlanMetricsSet::new(),
            merge_stats: Arc::new(Mutex::new(None)),
            transaction: Arc::new(Mutex::new(None)),
            affected_rows: Arc::new(Mutex::new(None)),
            inserted_rows_filter: Arc::new(Mutex::new(None)),
            source_skipped_duplicates,
            source_blob_columns,
            update_if_deferred: false,
            is_primary_key,
        })
    }

    /// Marks the `UpdateIf` condition as evaluated by `MergeConditionExec`.
    pub fn with_update_if_deferred(mut self) -> Self {
        self.update_if_deferred = true;
        self
    }

    /// Takes the merge statistics if the execution has completed.
    /// Returns `None` if the execution is still in progress or hasn't started.
    pub fn merge_stats(&self) -> Option<MergeStats> {
        self.merge_stats
            .lock()
            .ok()
            .and_then(|mut guard| guard.take())
    }

    /// Takes the transaction if the execution has completed.
    /// Returns `None` if the execution is still in progress or hasn't started.
    pub fn transaction(&self) -> Option<Transaction> {
        self.transaction
            .lock()
            .ok()
            .and_then(|mut guard| guard.take())
    }

    /// Returns the filter for inserted row keys if the execution has completed.
    /// This contains keys of newly inserted rows (not updates) for conflict detection.
    /// Returns `None` if the execution is still in progress or hasn't started.
    pub fn inserted_rows_filter(&self) -> Option<KeyExistenceFilter> {
        self.inserted_rows_filter
            .lock()
            .ok()
            .and_then(|guard| guard.clone())
    }

    /// Takes the affected rows (deleted/updated row addresses) if the execution has completed.
    /// Returns `None` if the execution is still in progress or hasn't started.
    pub fn affected_rows(&self) -> Option<RoaringTreemap> {
        self.affected_rows
            .lock()
            .ok()
            .and_then(|mut guard| guard.take())
    }

    /// Creates a filtered stream that captures row addresses for deletion and returns
    /// a stream with only the source data columns (no _rowaddr or __action columns)
    fn create_filtered_write_stream(
        &self,
        input_stream: SendableRecordBatchStream,
        merge_state: Arc<Mutex<MergeState>>,
    ) -> DFResult<SendableRecordBatchStream> {
        let enable_stable_row_ids = {
            let state = merge_state.lock().map_err(|e| {
                datafusion::error::DataFusionError::Internal(format!(
                    "Failed to lock merge state: {}",
                    e
                ))
            })?;
            state.stable_row_ids
        };

        if enable_stable_row_ids {
            self.create_ordered_update_insert_stream(input_stream, merge_state)
        } else {
            self.create_streaming_write_stream(input_stream, merge_state)
        }
    }

    /// High-performance streaming implementation for non-stable row ID scenarios
    ///
    /// It processes batches one at a time as they arrive from the input stream,
    /// immediately filtering and transforming each batch without buffering.
    fn create_streaming_write_stream(
        &self,
        input_stream: SendableRecordBatchStream,
        merge_state: Arc<Mutex<MergeState>>,
    ) -> DFResult<SendableRecordBatchStream> {
        let (_, rowaddr_idx, rowid_idx, action_idx, data_column_indices, output_schema) =
            self.prepare_stream_schema(input_stream.schema())?;

        let output_schema_clone = output_schema.clone();
        let stream = input_stream.map(move |batch_result| -> DFResult<RecordBatch> {
            let batch = batch_result?;
            let (row_addr_array, row_id_array, action_array) =
                Self::extract_control_arrays(&batch, rowaddr_idx, rowid_idx, action_idx)?;

            // Process each row using the shared state
            let mut keep_rows: Vec<u32> = Vec::with_capacity(batch.num_rows());

            let mut merge_state = merge_state.lock().map_err(|e| {
                datafusion::error::DataFusionError::Internal(format!(
                    "Failed to lock merge state: {}",
                    e
                ))
            })?;

            for row_idx in 0..batch.num_rows() {
                let action_code = action_array.value(row_idx);
                let action = Action::try_from(action_code).map_err(|e| {
                    datafusion::error::DataFusionError::Internal(format!(
                        "Invalid action code {}: {}",
                        action_code, e
                    ))
                })?;

                if merge_state
                    .process_row_action(action, row_idx, row_addr_array, row_id_array, &batch)?
                    .is_some()
                {
                    keep_rows.push(row_idx as u32);
                }
            }

            Self::create_filtered_batch(
                &batch,
                keep_rows,
                &data_column_indices,
                output_schema_clone.clone(),
            )
        });

        Ok(Box::pin(RecordBatchStreamAdapter::new(
            output_schema,
            stream,
        )))
    }

    /// Creates an ordered update-insert stream ensuring updated data before inserted data.
    ///
    /// 1. Separating the input stream into update and insert streams
    /// 2. Using chain operations to guarantee all update batches are processed before any insert batches
    /// 3. Returning the combined ordered stream
    fn create_ordered_update_insert_stream(
        &self,
        input_stream: SendableRecordBatchStream,
        merge_state: Arc<Mutex<MergeState>>,
    ) -> DFResult<SendableRecordBatchStream> {
        let (update_stream, insert_stream, _) =
            self.split_updates_and_inserts(input_stream, merge_state, false)?;

        let output_schema = update_stream.schema();

        // Chain the update and insert streams to ensure order
        let combined_stream = update_stream.chain(insert_stream);

        Ok(Box::pin(RecordBatchStreamAdapter::new(
            output_schema,
            combined_stream,
        )))
    }

    /// Common schema preparation logic
    #[allow(clippy::type_complexity)]
    fn prepare_stream_schema(
        &self,
        input_schema: arrow_schema::SchemaRef,
    ) -> DFResult<(
        arrow_schema::SchemaRef,
        usize,
        usize,
        usize,
        Vec<usize>,
        Arc<Schema>,
    )> {
        // Find column indices
        let (rowaddr_idx, _) = input_schema.column_with_name(ROW_ADDR).ok_or_else(|| {
            datafusion::error::DataFusionError::Internal(
                "Expected _rowaddr column in merge insert input".to_string(),
            )
        })?;

        let (rowid_idx, _) = input_schema.column_with_name(ROW_ID).ok_or_else(|| {
            datafusion::error::DataFusionError::Internal(
                "Expected _rowid column in merge insert input".to_string(),
            )
        })?;

        let (action_idx, _) = input_schema
            .column_with_name(MERGE_ACTION_COLUMN)
            .ok_or_else(|| {
                datafusion::error::DataFusionError::Internal(format!(
                    "Expected {} column in merge insert input",
                    MERGE_ACTION_COLUMN
                ))
            })?;

        // Emit data columns in dataset-schema order, keyed by name. This
        // matters for partial-schema upserts, where `create_plan` adds
        // synthetic unqualified columns (filled from the target side) at
        // the end of the logical projection. Without name-based reordering
        // they would land at the end of the write stream and mismatch the
        // intended writer schema (which is `dataset.schema()`). Using name
        // lookup is also a strictly-safer choice for the full-schema path:
        // it turns an implicit positional assumption into an explicit
        // name-based invariant. The filtered batches are recursively projected
        // to this schema below so nested children follow the same contract.
        let mut name_to_idx: std::collections::HashMap<&str, usize> =
            std::collections::HashMap::with_capacity(input_schema.fields().len());
        for (idx, field) in input_schema.fields().iter().enumerate() {
            let name = field.name();
            // Skip special columns: _rowaddr, _rowid, __action, and the
            // source-presence sentinel.
            if idx == rowaddr_idx
                || idx == rowid_idx
                || idx == action_idx
                || name == ROW_ADDR
                || name == ROW_ID
                || name == ROW_CREATED_AT_VERSION
                || name == MERGE_ACTION_COLUMN
                || name == MERGE_SOURCE_SENTINEL
            {
                continue;
            }
            name_to_idx.insert(name.as_str(), idx);
        }

        let dataset_arrow_schema: arrow_schema::Schema = self.dataset.schema().into();
        let dataset_fields = dataset_arrow_schema.fields();
        let mut data_column_indices: Vec<usize> = Vec::with_capacity(dataset_fields.len());
        for dataset_field in dataset_fields {
            let idx = *name_to_idx
                .get(dataset_field.name().as_str())
                .ok_or_else(|| {
                    datafusion::error::DataFusionError::Internal(format!(
                        "Dataset field {:?} missing from merge insert input schema \
                     — this indicates a logical plan pruning bug",
                        dataset_field.name()
                    ))
                })?;
            data_column_indices.push(idx);
        }

        if data_column_indices.is_empty() {
            return Err(datafusion::error::DataFusionError::Internal(
                "No data columns found in merge insert input".to_string(),
            ));
        }

        let source_data_schema = Schema::new(
            data_column_indices
                .iter()
                .map(|idx| input_schema.field(*idx).clone())
                .collect::<Vec<_>>(),
        );
        let output_schema = Arc::new(
            canonical_source_schema(&source_data_schema, &dataset_arrow_schema)
                .map_err(datafusion::error::DataFusionError::from)?,
        );

        Ok((
            input_schema,
            rowaddr_idx,
            rowid_idx,
            action_idx,
            data_column_indices,
            output_schema,
        ))
    }

    /// Extract control arrays from batch
    fn extract_control_arrays(
        batch: &RecordBatch,
        rowaddr_idx: usize,
        rowid_idx: usize,
        action_idx: usize,
    ) -> DFResult<(&UInt64Array, &UInt64Array, &UInt8Array)> {
        // Get row address, row id and __action arrays
        let row_addr_array = batch
            .column(rowaddr_idx)
            .as_any()
            .downcast_ref::<UInt64Array>()
            .ok_or_else(|| {
                datafusion::error::DataFusionError::Internal(
                    "Expected UInt64Array for _rowaddr column".to_string(),
                )
            })?;

        let row_id_array = batch
            .column(rowid_idx)
            .as_any()
            .downcast_ref::<UInt64Array>()
            .ok_or_else(|| {
                datafusion::error::DataFusionError::Internal(
                    "Expected UInt64Array for _rowid column".to_string(),
                )
            })?;

        let action_array = batch
            .column(action_idx)
            .as_any()
            .downcast_ref::<UInt8Array>()
            .ok_or_else(|| {
                datafusion::error::DataFusionError::Internal(format!(
                    "Expected UInt8Array for {} column",
                    MERGE_ACTION_COLUMN
                ))
            })?;

        Ok((row_addr_array, row_id_array, action_array))
    }

    /// Create filtered batch from selected rows
    fn create_filtered_batch(
        batch: &RecordBatch,
        keep_rows: Vec<u32>,
        data_column_indices: &[usize],
        output_schema: Arc<Schema>,
    ) -> DFResult<RecordBatch> {
        // If no rows to keep, return empty batch
        if keep_rows.is_empty() {
            let empty_columns: Vec<_> = output_schema
                .fields()
                .iter()
                .map(|field| arrow_array::new_empty_array(field.data_type()))
                .collect();
            return RecordBatch::try_new(output_schema, empty_columns)
                .map_err(datafusion::error::DataFusionError::from);
        }

        // Create indices for rows to keep
        let indices = arrow_array::UInt32Array::from(keep_rows);

        // Take only the rows we want to keep
        let filtered_batch = arrow_select::take::take_record_batch(batch, &indices)?;

        // First retain the source field layout, then recursively project it into
        // the dataset layout. The latter is required for nested structs whose
        // children were supplied in a different order.
        let projected = filtered_batch.project(data_column_indices)?;
        projected
            .project_by_schema(output_schema.as_ref())
            .map_err(datafusion::error::DataFusionError::from)
    }

    /// Consume all output groups concurrently, retaining update-first order.
    /// A failed writer closes the splitter and waits for its sibling so every
    /// completed file is available for cleanup before the error is returned.
    async fn write_groups(
        dataset: &Dataset,
        params: &MergeInsertParams,
        streams: Vec<SendableRecordBatchStream>,
        splitter: Option<Arc<SplitterTask>>,
        write_params: &WriteParams,
        source_blob_columns: &[String],
        target_bases: Option<Vec<TargetBaseInfo>>,
    ) -> lance_core::Result<(Vec<Fragment>, usize)> {
        let writes = streams.into_iter().map(|stream| {
            let splitter = splitter.clone();
            let target_bases = target_bases.clone();
            async move {
                let result: lance_core::Result<Vec<Fragment>> = async {
                    // Existing target references can be outside the current
                    // bases. Validate only source-provided blob columns.
                    let stream = if source_blob_columns.is_empty() {
                        stream
                    } else {
                        let resolver = blob_v2_external_base_resolver(
                            Some(dataset),
                            write_params,
                            dataset.schema(),
                        )
                        .await?
                        .ok_or_else(|| {
                            Error::internal("Source has blob v2 columns but target does not")
                        })?;
                        let schema = stream.schema();
                        let indices = source_blob_columns
                            .iter()
                            .map(|name| schema.index_of(name).map_err(Error::from))
                            .collect::<Result<Vec<_>, _>>()?;
                        let validated = stream.then(move |batch_result| {
                            let resolver = resolver.clone();
                            let indices = indices.clone();
                            async move {
                                let batch = batch_result?;
                                let source_blob_batch = batch.project(&indices)?;
                                crate::dataset::blob::validate_external_blob_references(
                                    &resolver,
                                    &source_blob_batch,
                                    &vec![true; batch.num_rows()],
                                )
                                .await
                                .map_err(|error| DataFusionError::External(Box::new(error)))?;
                                Ok(batch)
                            }
                        });
                        Box::pin(RecordBatchStreamAdapter::new(schema, validated))
                            as SendableRecordBatchStream
                    };
                    let (fragments, _) = write_fragments_internal(
                        params.write_version(dataset),
                        Some(dataset),
                        dataset.object_store.clone(),
                        &dataset.base,
                        dataset.schema().clone(),
                        stream,
                        write_params.clone(),
                        target_bases,
                    )
                    .await?;
                    Ok(fragments)
                }
                .await;
                if result.is_err()
                    && let Some(splitter) = splitter
                {
                    // Wake the sibling even if the source never sends another
                    // batch to the group whose writer failed.
                    splitter.0.abort();
                }
                result
            }
        });
        // Do not use try_join_all: dropping a sibling future on error would
        // lose track of files it already completed.
        let results = futures::future::join_all(writes).await;
        let mut fragments = Vec::new();
        let mut update_fragment_count = 0;
        let mut write_error = None;
        for (group, result) in results.into_iter().enumerate() {
            match result {
                Ok(group_fragments) => fragments.extend(group_fragments),
                Err(error) if write_error.is_none() => write_error = Some(error),
                Err(_) => {}
            }
            if group == 0 {
                update_fragment_count = fragments.len();
            }
        }
        if let Some(error) = write_error {
            cleanup_data_fragments(
                &dataset.object_store,
                &dataset.base,
                target_bases.as_deref(),
                &fragments,
            )
            .await;
            return Err(error);
        }
        Ok((fragments, update_fragment_count))
    }

    /// Place only lineage carried by updated rows. With spilling disabled,
    /// fragments may also contain inserts and the final inline sequence is
    /// intentionally partial; commit fills those missing ids and versions.
    async fn place_updated_lineage(
        dataset: &Dataset,
        fragments: &mut [Fragment],
        row_ids: &RowIdSequence,
        created_at: Option<&RowDatasetVersionSequence>,
        spill_budget: Option<usize>,
    ) -> lance_core::Result<()> {
        let sizes = fragments
            .iter()
            .map(|fragment| {
                fragment
                    .physical_rows
                    .map(|rows| rows as u64)
                    .ok_or_else(|| {
                        Error::internal(
                            "Merge insert output fragment is missing its physical row count",
                        )
                    })
            })
            .collect::<lance_core::Result<Vec<_>>>()?;
        let row_ids = rechunk_sequences(
            [row_ids.clone()],
            sizes.iter().copied(),
            spill_budget.is_none(),
        )
        .map_err(|error| {
            Error::internal(format!(
                "Captured row ids do not match merge insert output: {error}"
            ))
        })?;
        let Some(limit) = spill_budget else {
            for (fragment, sequence) in fragments.iter_mut().zip(row_ids) {
                fragment.row_id_meta = Some(RowIdMeta::Inline(write_row_ids(&sequence).into()));
            }
            return Ok(());
        };
        let created_at = created_at.ok_or_else(|| {
            Error::internal("Missing captured created-at versions for merge insert spill")
        })?;
        let created_at =
            rechunk_version_sequences([created_at.clone()], sizes, false).map_err(|error| {
                Error::internal(format!(
                    "Captured created-at versions do not match merge insert output: {error}"
                ))
            })?;
        for ((fragment, row_ids), created_at) in fragments.iter_mut().zip(row_ids).zip(created_at) {
            place_carried_row_lineage(dataset, limit, &row_ids, &created_at)
                .await?
                .apply(fragment);
        }
        Ok(())
    }

    /// Calculate write metrics from new fragments
    fn calculate_write_metrics(new_fragments: &[lance_table::format::Fragment]) -> (usize, usize) {
        let mut total_bytes = 0u64;
        let mut total_files = 0usize;

        for fragment in new_fragments {
            for data_file in &fragment.files {
                if let Some(size) = data_file.file_size_bytes.get() {
                    total_bytes += u64::from(size);
                }
                total_files += 1;
            }
        }

        (total_bytes as usize, total_files)
    }

    fn split_updates_and_inserts(
        &self,
        input_stream: SendableRecordBatchStream,
        merge_state: Arc<Mutex<MergeState>>,
        bounded: bool,
    ) -> DFResult<(
        SendableRecordBatchStream,
        SendableRecordBatchStream,
        Arc<SplitterTask>,
    )> {
        let (_, rowaddr_idx, rowid_idx, action_idx, data_column_indices, output_schema) =
            self.prepare_stream_schema(input_stream.schema())?;

        if input_stream
            .schema()
            .column_with_name(ROW_CREATED_AT_VERSION)
            .is_some()
        {
            let state = merge_state
                .lock()
                .map_err(|e| DataFusionError::Internal(e.to_string()))?;
            let mut captured = state
                .updating_row_ids
                .lock()
                .map_err(|e| DataFusionError::Internal(e.to_string()))?;
            if let CapturedRowIds::SequenceStyle { created_at, .. } = &mut *captured {
                *created_at = Some(RowDatasetVersionSequence::new());
            }
        }

        let (update_tx, update_rx) = MergeWriteSender::channel(bounded);
        let (insert_tx, insert_rx) = MergeWriteSender::channel(bounded);

        let output_schema_clone = output_schema.clone();
        let merge_state_clone = merge_state;

        // The splitter owns both senders, so a panic inside it drops them and
        // the receivers would see end-of-stream — a partial merge committed
        // with no error. Catch the unwind and forward it as a stream error
        // instead, so the writer fails and the transaction is never stored.
        // The clones outlive the unwound future; we never resume it, so the
        // `AssertUnwindSafe` bound is fine.
        let panic_update_tx = update_tx.clone();
        let panic_insert_tx = insert_tx.clone();
        let splitter = Arc::new(SplitterTask(tokio::spawn(async move {
            let result =
                futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(async move {
                    let mut input_stream = input_stream;

                    while let Some(batch_result) = input_stream.next().await {
                        match batch_result {
                            Ok(batch) => {
                                match Self::process_and_split_batch(
                                    &batch,
                                    rowaddr_idx,
                                    rowid_idx,
                                    action_idx,
                                    &data_column_indices,
                                    output_schema_clone.clone(),
                                    merge_state_clone.clone(),
                                ) {
                                    Ok((update_batch_opt, insert_batch_opt)) => {
                                        if let Some(update_batch) = update_batch_opt
                                            && update_tx.send(Ok(update_batch)).await.is_err()
                                        {
                                            break;
                                        }

                                        if let Some(insert_batch) = insert_batch_opt
                                            && insert_tx.send(Ok(insert_batch)).await.is_err()
                                        {
                                            break;
                                        }
                                    }
                                    Err(e) => {
                                        Self::handle_stream_processing_error(
                                            e, &update_tx, &insert_tx,
                                        )
                                        .await;
                                        break;
                                    }
                                }
                            }
                            Err(e) => {
                                Self::handle_stream_processing_error(e, &update_tx, &insert_tx)
                                    .await;
                                break;
                            }
                        }
                    }
                }))
                .await;

            if let Err(panic) = result {
                let panic_msg = if let Some(msg) = panic.downcast_ref::<&str>() {
                    (*msg).to_string()
                } else if let Some(msg) = panic.downcast_ref::<String>() {
                    msg.clone()
                } else {
                    "unknown panic payload".to_string()
                };
                Self::handle_stream_processing_error(
                    datafusion::error::DataFusionError::Execution(format!(
                        "merge insert stream splitter panicked: {panic_msg}"
                    )),
                    &panic_update_tx,
                    &panic_insert_tx,
                )
                .await;
            }
        })));

        // The producer must not outlive both consumers when a merge is cancelled.
        let update_guard = splitter.clone();
        let update_stream = update_rx.map(move |batch| {
            let _guard = &update_guard;
            batch
        });
        let update_stream = Box::pin(RecordBatchStreamAdapter::new(
            output_schema.clone(),
            update_stream,
        ));
        let insert_guard = splitter.clone();
        let insert_stream = insert_rx.map(move |batch| {
            let _guard = &insert_guard;
            batch
        });
        let insert_stream = Box::pin(RecordBatchStreamAdapter::new(output_schema, insert_stream));

        Ok((update_stream, insert_stream, splitter))
    }

    fn process_and_split_batch(
        batch: &RecordBatch,
        rowaddr_idx: usize,
        rowid_idx: usize,
        action_idx: usize,
        data_column_indices: &[usize],
        output_schema: Arc<Schema>,
        merge_state: Arc<Mutex<MergeState>>,
    ) -> DFResult<(Option<RecordBatch>, Option<RecordBatch>)> {
        let (row_addr_array, row_id_array, action_array) =
            Self::extract_control_arrays(batch, rowaddr_idx, rowid_idx, action_idx)?;

        let created_at = batch
            .column_by_name(ROW_CREATED_AT_VERSION)
            .map(|array| {
                array.as_primitive_opt::<UInt64Type>().ok_or_else(|| {
                    DataFusionError::Internal(
                        "Expected UInt64 created-at versions in merge insert input".to_string(),
                    )
                })
            })
            .transpose()?;

        let mut update_indices: Vec<u32> = Vec::new();
        let mut insert_indices: Vec<u32> = Vec::new();

        {
            let mut merge_state = merge_state.lock().map_err(|e| {
                datafusion::error::DataFusionError::Internal(format!(
                    "Failed to lock merge state: {}",
                    e
                ))
            })?;

            for row_idx in 0..batch.num_rows() {
                let action_code = action_array.value(row_idx);
                let action = Action::try_from(action_code).map_err(|e| {
                    datafusion::error::DataFusionError::Internal(format!(
                        "Invalid action code {}: {}",
                        action_code, e
                    ))
                })?;

                if merge_state
                    .process_row_action(action, row_idx, row_addr_array, row_id_array, batch)?
                    .is_some()
                {
                    match action {
                        Action::UpdateAll => {
                            update_indices.push(row_idx as u32);
                            if let Some(created_at) = created_at {
                                if created_at.is_null(row_idx) {
                                    return Err(DataFusionError::Internal(
                                        "Updated row has a null created-at version".to_string(),
                                    ));
                                }
                                merge_state
                                    .updating_row_ids
                                    .lock()
                                    .map_err(|e| DataFusionError::Internal(e.to_string()))?
                                    .capture_created_at(&[created_at.value(row_idx)]);
                            }
                        }
                        Action::Insert => insert_indices.push(row_idx as u32),
                        _ => {}
                    }
                }
            }
        }

        let update_batch = if !update_indices.is_empty() {
            Some(Self::create_filtered_batch(
                batch,
                update_indices,
                data_column_indices,
                output_schema.clone(),
            )?)
        } else {
            None
        };

        let insert_batch = if !insert_indices.is_empty() {
            Some(Self::create_filtered_batch(
                batch,
                insert_indices,
                data_column_indices,
                output_schema,
            )?)
        } else {
            None
        };

        Ok((update_batch, insert_batch))
    }

    async fn handle_stream_processing_error(
        error: datafusion::error::DataFusionError,
        update_tx: &MergeWriteSender,
        insert_tx: &MergeWriteSender,
    ) {
        // Send to first open one. It doesn't matter which one receives it as
        // long as the user gets the error in the end.
        if let Err(tokio::sync::mpsc::error::SendError(error)) = update_tx.send(Err(error)).await {
            let _ = insert_tx.send(error).await;
        }
    }
}

impl DisplayAs for FullSchemaMergeInsertExec {
    fn fmt_as(
        &self,
        t: datafusion::physical_plan::DisplayFormatType,
        f: &mut std::fmt::Formatter,
    ) -> std::fmt::Result {
        match t {
            datafusion::physical_plan::DisplayFormatType::Default
            | datafusion::physical_plan::DisplayFormatType::Verbose => {
                let on_keys = self.params.on.join(", ");
                let when_matched =
                    super::when_matched_label(&self.params.when_matched, self.update_if_deferred);
                let when_not_matched = if self.params.insert_not_matched {
                    "InsertAll"
                } else {
                    "DoNothing"
                };
                let when_not_matched_by_source = match &self.params.delete_not_matched_by_source {
                    crate::dataset::WhenNotMatchedBySource::Keep => "Keep",
                    crate::dataset::WhenNotMatchedBySource::Delete => "Delete",
                    crate::dataset::WhenNotMatchedBySource::DeleteIf(_) => "DeleteIf",
                };

                write!(
                    f,
                    "MergeInsert: on=[{}], when_matched={}, when_not_matched={}, when_not_matched_by_source={}",
                    on_keys, when_matched, when_not_matched, when_not_matched_by_source
                )
            }
            datafusion::physical_plan::DisplayFormatType::TreeRender => {
                write!(f, "MergeInsert[{}]", self.dataset.uri())
            }
        }
    }
}

impl ExecutionPlan for FullSchemaMergeInsertExec {
    fn name(&self) -> &str {
        "FullSchemaMergeInsertExec"
    }

    fn schema(&self) -> arrow_schema::SchemaRef {
        Arc::new(arrow_schema::Schema::empty())
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![&self.input]
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        if children.len() != 1 {
            return Err(datafusion::error::DataFusionError::Internal(
                "FullSchemaMergeInsertExec requires exactly one child".to_string(),
            ));
        }
        Ok(Arc::new(Self {
            input: children[0].clone(),
            dataset: self.dataset.clone(),
            params: self.params.clone(),
            properties: self.properties.clone(),
            metrics: self.metrics.clone(),
            merge_stats: self.merge_stats.clone(),
            transaction: self.transaction.clone(),
            affected_rows: self.affected_rows.clone(),
            inserted_rows_filter: self.inserted_rows_filter.clone(),
            source_skipped_duplicates: self.source_skipped_duplicates.clone(),
            source_blob_columns: self.source_blob_columns.clone(),
            update_if_deferred: self.update_if_deferred,
            is_primary_key: self.is_primary_key,
        }))
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metrics.clone_inner())
    }

    fn properties(&self) -> &Arc<PlanProperties> {
        &self.properties
    }

    fn supports_limit_pushdown(&self) -> bool {
        false
    }

    fn required_input_distribution(&self) -> Vec<datafusion_physical_expr::Distribution> {
        // We require a single partition for the merge operation to ensure all data is processed
        vec![datafusion_physical_expr::Distribution::SinglePartition]
    }

    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
        // We just want one stream.
        vec![false]
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> DFResult<SendableRecordBatchStream> {
        let _baseline_metrics = BaselineMetrics::new(&self.metrics, partition);

        // Input schema structure based on our logical plan:
        // - target._rowaddr: Address of existing rows to update/delete
        // - source.*: Source data columns (variable schema)
        // - __action: Merge action (1=update, 2=insert, 0=delete, etc.)

        // Validate the budget before executing the input or writing any files.
        let spill_budget = self.params.carried_lineage_spill_budget(&self.dataset)?;

        // Execute the input plan to get the merge data stream
        let input_stream = self.input.execute(partition, context)?;
        let has_blob_v2_columns = self
            .dataset
            .schema()
            .fields_pre_order()
            .any(|field| field.is_blob_v2());
        let input_stream = if has_blob_v2_columns {
            let input_schema = input_stream.schema();
            let rewrite_plan = Arc::new(
                crate::dataset::optimize::BlobV2BatchRewritePlan::try_new(
                    self.dataset.schema(),
                    input_schema.as_ref(),
                    true,
                    false,
                )
                .map_err(|error| DataFusionError::External(Box::new(error)))?,
            );
            let output_schema = rewrite_plan.output_schema().clone();
            let dataset = self.dataset.clone();
            let transformed = input_stream.then(move |batch_result| {
                let dataset = dataset.clone();
                let rewrite_plan = rewrite_plan.clone();
                async move {
                    let batch = batch_result?;
                    rewrite_plan
                        .transform_batch(&dataset, batch)
                        .await
                        .map_err(|error| DataFusionError::External(Box::new(error)))
                }
            });
            Box::pin(RecordBatchStreamAdapter::new(output_schema, transformed))
                as SendableRecordBatchStream
        } else {
            input_stream
        };

        // Step 1: Create shared state and streaming processor for row addresses and write data
        // Get field IDs for the ON columns from the dataset schema
        let field_ids: Vec<i32> = self
            .params
            .on
            .iter()
            .filter_map(|name| self.dataset.schema().field(name).map(|f| f.id))
            .collect();
        let merge_state = Arc::new(Mutex::new(MergeState::new(
            MergeInsertMetrics::new(&self.metrics, partition),
            self.dataset.manifest.uses_stable_row_ids(),
            self.params.on.clone(),
            field_ids,
            self.params.source_dedupe_behavior,
        )));
        let (write_data_streams, splitter) = if spill_budget.is_some() {
            // A spilled sequence must cover every row in its fragment. Inserts
            // need commit-assigned ids and creation versions that can change on
            // rebase, so they cannot share a fragment with spilled updates.
            let (updates, inserts, splitter) =
                self.split_updates_and_inserts(input_stream, merge_state.clone(), true)?;
            (vec![updates, inserts], Some(splitter))
        } else {
            (
                vec![self.create_filtered_write_stream(input_stream, merge_state.clone())?],
                None,
            )
        };

        // Use flat_map to handle the async write operation
        let dataset = self.dataset.clone();
        let params = self.params.clone();
        let merge_stats_holder = self.merge_stats.clone();
        let transaction_holder = self.transaction.clone();
        let affected_rows_holder = self.affected_rows.clone();
        let inserted_rows_filter_holder = self.inserted_rows_filter.clone();
        let compacted_sstables = self.params.compacted_sstables.clone();
        let source_skipped_duplicates = self.source_skipped_duplicates.clone();
        let source_blob_columns = self.source_blob_columns.clone();
        let is_primary_key = self.is_primary_key;
        let updating_row_ids = {
            let state = merge_state.lock().unwrap();
            state.updating_row_ids.clone()
        };

        let result_stream = stream::once(async move {
            // Step 2: Write new fragments using the filtered data (inserts + updates)
            let target_bases_info = resolve_target_bases(&dataset, &params).await?;
            // Keep a copy so failures after the write can clean up routed files.
            let cleanup_bases = target_bases_info.clone();
            let write_params = WriteParams {
                allow_external_blob_outside_bases: has_blob_v2_columns,
                ..Default::default()
            };
            let (mut new_fragments, update_fragment_count) = Self::write_groups(
                &dataset,
                &params,
                write_data_streams,
                splitter,
                &write_params,
                &source_blob_columns,
                target_bases_info,
            )
            .await?;
            let lineage_result: lance_core::Result<()> = async {
                let (row_ids, created_at) = {
                    let captured = updating_row_ids.lock().map_err(|e| {
                        Error::internal(format!("Failed to lock captured row lineage: {e}"))
                    })?;
                    (
                        captured.row_id_sequence().cloned(),
                        captured.created_at_sequence().cloned(),
                    )
                };
                if let Some(row_ids) = row_ids {
                    Self::place_updated_lineage(
                        &dataset,
                        &mut new_fragments[..update_fragment_count],
                        &row_ids,
                        created_at.as_ref(),
                        spill_budget,
                    )
                    .await?;
                }
                Ok(())
            }
            .await;
            if let Err(error) = lineage_result {
                cleanup_data_fragments(
                    &dataset.object_store,
                    &dataset.base,
                    cleanup_bases.as_deref(),
                    &new_fragments,
                )
                .await;
                return Err(error.into());
            }

            // Step 2.5: Calculate write metrics from new fragments
            let (total_bytes_written, total_files_written) =
                Self::calculate_write_metrics(&new_fragments);

            // Step 3: Apply deletions to existing fragments
            let merge_state =
                Arc::into_inner(merge_state).expect("MergeState should only have 1 reference now");
            let merge_state =
                Mutex::into_inner(merge_state).expect("MergeState lock should be available");
            let delete_row_addrs_clone = merge_state.delete_row_addrs;
            let inserted_rows_filter = if is_primary_key {
                Some(KeyExistenceFilter::from_bloom_filter(
                    &merge_state.inserted_rows_filter,
                ))
            } else {
                None
            };

            let (updated_fragments, removed_fragment_ids) =
                match apply_deletions(&dataset, &delete_row_addrs_clone).await {
                    Ok(result) => result,
                    Err(e) => {
                        // The new data files are not committed; remove them (including
                        // ones routed to target bases) before surfacing the error.
                        cleanup_data_fragments(
                            &dataset.object_store,
                            &dataset.base,
                            cleanup_bases.as_deref(),
                            &new_fragments,
                        )
                        .await;
                        return Err(e.into());
                    }
                };

            // Step 4: Create the transaction operation
            let operation = Operation::Update {
                removed_fragment_ids,
                updated_fragments,
                new_fragments,
                fields_modified: vec![], // No fields are modified in schema for upsert
                compacted_sstables,
                // Use the full pre-order field list (not just top-level `fields`) so
                // that nested leaf field ids are included. A merge_insert rewrites whole
                // rows, so every field is potentially modified; omitting nested ids would
                // let `register_pure_rewrite_rows_update_frags_in_indices` wrongly extend a
                // nested-field index over the rewritten fragment, silently dropping rows.
                fields_for_preserving_frag_bitmap: dataset
                    .schema()
                    .fields_pre_order()
                    .map(|f| f.id as u32)
                    .collect(),
                update_mode: Some(RewriteRows),
                inserted_rows_filter: inserted_rows_filter.clone(),
                updated_fragment_offsets: None,
            };

            // Step 5: Create and store the transaction
            let transaction = Transaction::new(dataset.manifest.version, operation, None);

            // Step 6: Store transaction, merge stats, and affected rows for later retrieval
            {
                // Update write metrics before converting to stats
                merge_state.metrics.bytes_written.add(total_bytes_written);
                merge_state
                    .metrics
                    .num_files_written
                    .add(total_files_written);

                // Get the final stats from the shared state
                let mut stats = MergeStats::from(&merge_state.metrics);
                stats.num_skipped_duplicates = stats
                    .num_skipped_duplicates
                    .checked_add(source_skipped_duplicates.load(Ordering::Relaxed))
                    .ok_or_else(|| {
                        DataFusionError::Execution(
                            "merge insert skipped duplicate count overflowed u64".to_string(),
                        )
                    })?;

                if let Ok(mut transaction_guard) = transaction_holder.lock() {
                    transaction_guard.replace(transaction);
                }
                if let Ok(mut merge_stats_guard) = merge_stats_holder.lock() {
                    merge_stats_guard.replace(stats);
                }
                if let Ok(mut affected_rows_guard) = affected_rows_holder.lock() {
                    affected_rows_guard.replace(delete_row_addrs_clone);
                }
                if let Ok(mut filter_guard) = inserted_rows_filter_holder.lock() {
                    *filter_guard = inserted_rows_filter;
                }
            };

            // Step 7: Return empty result (write operations don't return data)
            let empty_schema = Arc::new(arrow_schema::Schema::empty());
            let empty_batch = RecordBatch::new_empty(empty_schema);
            Ok(empty_batch)
        });

        let empty_schema = Arc::new(arrow_schema::Schema::empty());
        Ok(Box::pin(RecordBatchStreamAdapter::new(
            empty_schema,
            result_stream,
        )))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_array::{UInt64Array, record_batch};
    use async_trait::async_trait;
    use datafusion::physical_plan::empty::EmptyExec;
    use std::sync::atomic::AtomicUsize;
    use std::time::Duration;
    use tokio::sync::Notify;

    use crate::dataset::progress::WriteFragmentProgress;
    use crate::dataset::rowids::{load_row_id_sequence, read_spilled_versions};
    use crate::dataset::write::InsertBuilder;
    use crate::dataset::write::merge_insert::{MergeInsertBuilder, WhenMatched, WhenNotMatched};
    use lance_table::format::{ROW_CREATED_AT_VERSION_FIELD_ID, RowDatasetVersionMeta};
    use lance_table::rowids::version::write_dataset_versions;
    use rstest::rstest;

    async fn splitter_fixture() -> (FullSchemaMergeInsertExec, Arc<Mutex<MergeState>>) {
        let batch = record_batch!(("key", UInt32, [0]), ("value", UInt32, [0])).unwrap();
        let schema = batch.schema();
        let dataset = Arc::new(
            InsertBuilder::new("memory://")
                .with_params(&WriteParams {
                    enable_stable_row_ids: true,
                    ..Default::default()
                })
                .execute(vec![batch])
                .await
                .unwrap(),
        );
        let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["key".into()])
            .unwrap()
            .when_matched(WhenMatched::UpdateAll)
            .when_not_matched(WhenNotMatched::InsertAll)
            .try_build()
            .unwrap();
        let exec = FullSchemaMergeInsertExec::try_new(
            Arc::new(EmptyExec::new(schema)),
            dataset,
            job.params,
            Arc::new(AtomicU64::new(0)),
            Vec::new(),
        )
        .unwrap();
        let state = Arc::new(Mutex::new(MergeState::new(
            MergeInsertMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
            true,
            vec!["key".into()],
            Vec::new(),
            SourceDedupeBehavior::Fail,
        )));
        (exec, state)
    }

    struct NotifyOnDrop(Arc<Notify>);

    impl Drop for NotifyOnDrop {
        fn drop(&mut self) {
            self.0.notify_one();
        }
    }

    #[tokio::test]
    async fn merge_insert_spill_splitter_backpressure_and_cancellation() {
        let (exec, state) = splitter_fixture().await;
        let batch = record_batch!(
            ("key", UInt32, [100]),
            ("value", UInt32, [100]),
            (ROW_ADDR, UInt64, [None]),
            (ROW_ID, UInt64, [None]),
            (ROW_CREATED_AT_VERSION, UInt64, [None]),
            (MERGE_ACTION_COLUMN, UInt8, [2])
        )
        .unwrap();
        let schema = batch.schema();
        let produced = Arc::new(AtomicUsize::new(0));
        let second = Arc::new(Notify::new());
        let dropped = Arc::new(Notify::new());
        let source = stream::unfold((batch, NotifyOnDrop(dropped.clone())), {
            let produced = produced.clone();
            let second = second.clone();
            move |(batch, guard)| {
                let count = produced.fetch_add(1, Ordering::SeqCst) + 1;
                if count == 2 {
                    second.notify_one();
                }
                async move { (count <= 8).then(|| (Ok(batch.clone()), (batch, guard))) }
            }
        });
        let (updates, mut inserts, splitter) = exec
            .split_updates_and_inserts(
                Box::pin(RecordBatchStreamAdapter::new(schema, source)),
                state,
                true,
            )
            .unwrap();
        tokio::time::timeout(Duration::from_secs(5), second.notified())
            .await
            .unwrap();
        // One batch is queued and the producer is blocked sending the second.
        // An unbounded splitter would consume the remaining ready batches.
        assert_eq!(produced.load(Ordering::SeqCst), 2);
        assert_eq!(inserts.next().await.unwrap().unwrap().num_rows(), 1);
        drop(updates);
        drop(inserts);
        drop(splitter);
        tokio::time::timeout(Duration::from_secs(5), dropped.notified())
            .await
            .unwrap();
    }

    #[derive(Debug, Default)]
    struct FailSecondFile {
        completed: AtomicUsize,
        second_complete: Notify,
    }

    #[async_trait]
    impl WriteFragmentProgress for FailSecondFile {
        async fn begin(&self, _fragment: &Fragment) -> lance_core::Result<()> {
            Ok(())
        }

        async fn complete(&self, _fragment: &Fragment) -> lance_core::Result<()> {
            match self.completed.fetch_add(1, Ordering::SeqCst) {
                // Hold the first writer after it finishes a file, until the
                // sibling has also completed a file and is about to fail.
                0 => self.second_complete.notified().await,
                1 => {
                    self.second_complete.notify_one();
                    return Err(Error::io("injected second-file failure"));
                }
                _ => {}
            }
            Ok(())
        }
    }

    #[tokio::test]
    async fn merge_insert_spill_cleans_both_groups_on_writer_failure() {
        let (exec, state) = splitter_fixture().await;
        let mut before = exec
            .dataset
            .object_store
            .read_dir(exec.dataset.data_dir())
            .await
            .unwrap();
        before.sort();
        let dropped = Arc::new(Notify::new());
        let batches = (0..20_u32)
            .map(|i| {
                record_batch!(
                    ("key", UInt32, [i, 100 + i]),
                    ("value", UInt32, [i, i]),
                    (ROW_ADDR, UInt64, [Some(i as u64), None]),
                    (ROW_ID, UInt64, [Some(i as u64), None]),
                    (ROW_CREATED_AT_VERSION, UInt64, [Some(1), None]),
                    (MERGE_ACTION_COLUMN, UInt8, [1, 2])
                )
                .unwrap()
            })
            .collect::<Vec<_>>();
        let schema = batches[0].schema();
        let source = stream::unfold(
            (batches.into_iter(), NotifyOnDrop(dropped.clone())),
            |(mut batches, guard)| async move {
                batches.next().map(|batch| (Ok(batch), (batches, guard)))
            },
        );
        let (updates, inserts, splitter) = exec
            .split_updates_and_inserts(
                Box::pin(RecordBatchStreamAdapter::new(schema, source)),
                state,
                true,
            )
            .unwrap();
        let progress = Arc::new(FailSecondFile::default());
        let params = WriteParams {
            max_rows_per_file: 1,
            progress: progress.clone(),
            ..Default::default()
        };
        let error = tokio::time::timeout(
            Duration::from_secs(5),
            FullSchemaMergeInsertExec::write_groups(
                &exec.dataset,
                &exec.params,
                vec![updates, inserts],
                Some(splitter),
                &params,
                &[],
                None,
            ),
        )
        .await
        .unwrap()
        .unwrap_err();
        assert!(matches!(error, Error::IO { .. }), "{error}");
        assert!(
            error.to_string().contains("injected second-file failure"),
            "{error}"
        );
        assert!(progress.completed.load(Ordering::SeqCst) >= 2);
        tokio::time::timeout(Duration::from_secs(5), dropped.notified())
            .await
            .unwrap();
        let mut after = exec
            .dataset
            .object_store
            .read_dir(exec.dataset.data_dir())
            .await
            .unwrap();
        after.sort();
        assert_eq!(
            after, before,
            "both writers' uncommitted files must be removed"
        );
    }

    #[rstest]
    #[case::all_spilled(0)]
    #[case::below_boundary(1)]
    #[case::at_boundary(2)]
    #[tokio::test]
    async fn merge_insert_spill_budget_and_fragment_boundaries(#[case] budget_case: u8) {
        let dataset = InsertBuilder::new("memory://")
            .execute(vec![
                arrow_array::record_batch!(("key", UInt32, [0])).unwrap(),
            ])
            .await
            .unwrap();
        let row_ids = RowIdSequence::from(&[10, 3, 99][..]);
        let created_at = RowDatasetVersionSequence::from_versions(&[1, 2, 4]);
        let max_size = write_row_ids(&RowIdSequence::from(&[10, 3][..]))
            .len()
            .max(write_dataset_versions(&RowDatasetVersionSequence::from_versions(&[1, 2])).len());
        let limit = match budget_case {
            0 => 0,
            1 => max_size - 1,
            _ => max_size,
        };
        let mut fragments = vec![Fragment::new(0), Fragment::new(1)];
        fragments[0].physical_rows = Some(2);
        fragments[1].physical_rows = Some(1);
        FullSchemaMergeInsertExec::place_updated_lineage(
            &dataset,
            &mut fragments,
            &row_ids,
            Some(&created_at),
            Some(limit),
        )
        .await
        .unwrap();
        assert_eq!(
            fragments
                .iter()
                .any(|fragment| fragment.has_spilled_row_lineage()),
            budget_case != 2
        );
        for (fragment, (expected_ids, expected_versions)) in fragments
            .iter()
            .zip([(vec![10, 3], vec![1, 2]), (vec![99], vec![4])])
        {
            assert_eq!(
                load_row_id_sequence(&dataset, fragment)
                    .await
                    .unwrap()
                    .iter()
                    .collect::<Vec<_>>(),
                expected_ids
            );
            match &fragment.created_at_version_meta {
                Some(RowDatasetVersionMeta::Column) => assert_eq!(
                    read_spilled_versions(&dataset, fragment, ROW_CREATED_AT_VERSION_FIELD_ID)
                        .await
                        .unwrap()
                        .versions()
                        .collect::<Vec<_>>(),
                    expected_versions
                ),
                Some(meta) => assert_eq!(
                    meta.load_sequence().unwrap().versions().collect::<Vec<_>>(),
                    expected_versions
                ),
                None => assert!(!fragment.has_spilled_row_lineage()),
            }
            assert!(fragment.last_updated_at_version_meta.is_none());
        }
        let too_few = RowDatasetVersionSequence::from_versions(&[1, 2]);
        let error = FullSchemaMergeInsertExec::place_updated_lineage(
            &dataset,
            &mut fragments,
            &row_ids,
            Some(&too_few),
            Some(0),
        )
        .await
        .unwrap_err();
        assert!(matches!(error, Error::Internal { .. }), "{error}");
        assert!(
            error.to_string().contains("Captured created-at versions"),
            "{error}"
        );
    }

    // A panic in the detached splitter task must surface on the output
    // streams as an error. The task owns the channel senders, so an unwind
    // drops them and the receivers cannot distinguish "writer panicked" from
    // end-of-stream — without a guard, a partial merge is silently committed.
    // The poisoned `updating_row_ids` mutex stands in for any panic source.
    #[tokio::test]
    async fn splitter_panic_surfaces_as_stream_error() {
        use arrow_array::Int32Array;
        use datafusion::physical_plan::empty::EmptyExec;
        use futures::TryStreamExt;

        let schema = Arc::new(arrow_schema::Schema::new(vec![
            arrow_schema::Field::new("key", arrow_schema::DataType::UInt64, false),
            arrow_schema::Field::new("value", arrow_schema::DataType::Int32, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt64Array::from(vec![1u64])),
                Arc::new(Int32Array::from(vec![42])),
            ],
        )
        .unwrap();
        let dataset = Dataset::write(
            arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema.clone()),
            "memory://splitter-panic",
            None,
        )
        .await
        .unwrap();

        let job = MergeInsertBuilder::try_new(Arc::new(dataset.clone()), vec!["key".to_string()])
            .unwrap()
            .when_matched(WhenMatched::UpdateAll)
            .when_not_matched(WhenNotMatched::InsertAll)
            .try_build()
            .unwrap();
        let exec = FullSchemaMergeInsertExec::try_new(
            Arc::new(EmptyExec::new(schema.clone())),
            Arc::new(dataset),
            job.params.clone(),
            Arc::new(std::sync::atomic::AtomicU64::new(0)),
            Vec::new(),
        )
        .unwrap();

        let metrics = MergeInsertMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
        let merge_state = Arc::new(std::sync::Mutex::new(MergeState::new(
            metrics,
            true,
            vec!["key".to_string()],
            Vec::new(),
            SourceDedupeBehavior::Fail,
        )));

        // Poison the captured-row-ids mutex so the splitter task panics when
        // the UpdateAll row reaches the capture call.
        {
            let poison_target = merge_state.lock().unwrap().updating_row_ids.clone();
            let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                let _guard = poison_target.lock().unwrap();
                panic!("poison the captured row ids mutex");
            }));
        }

        let input_schema = Arc::new(arrow_schema::Schema::new(vec![
            arrow_schema::Field::new("key", arrow_schema::DataType::UInt64, false),
            arrow_schema::Field::new("value", arrow_schema::DataType::Int32, false),
            arrow_schema::Field::new(ROW_ADDR, arrow_schema::DataType::UInt64, true),
            arrow_schema::Field::new(ROW_ID, arrow_schema::DataType::UInt64, true),
            arrow_schema::Field::new(MERGE_ACTION_COLUMN, arrow_schema::DataType::UInt8, false),
        ]));
        let input_batch = RecordBatch::try_new(
            input_schema.clone(),
            vec![
                Arc::new(UInt64Array::from(vec![1u64])),
                Arc::new(Int32Array::from(vec![42])),
                Arc::new(UInt64Array::from(vec![1u64 << 32])),
                Arc::new(UInt64Array::from(vec![7u64])),
                Arc::new(arrow_array::UInt8Array::from(vec![1u8])), // UpdateAll
            ],
        )
        .unwrap();
        let input_stream = Box::pin(RecordBatchStreamAdapter::new(
            input_schema,
            futures::stream::iter(vec![Ok(input_batch)]),
        ));

        let (update_stream, _insert_stream, _splitter) = exec
            .split_updates_and_inserts(input_stream, merge_state, false)
            .unwrap();

        let result: DFResult<Vec<RecordBatch>> = update_stream.try_collect().await;
        let err = result.expect_err(
            "a splitter panic must surface as a stream error, not a silent end-of-stream",
        );
        assert!(
            err.to_string().to_lowercase().contains("panic"),
            "error should name the panic, got: {err}"
        );
    }

    #[test]
    fn test_merge_state_duplicate_rowid_detection_fail() {
        let metrics = MergeInsertMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
        let mut merge_state = MergeState::new(
            metrics,
            false,
            Vec::new(),
            Vec::new(),
            SourceDedupeBehavior::Fail,
        );

        let row_addr_array = UInt64Array::from(vec![1000, 2000, 3000]);
        let row_id_array = UInt64Array::from(vec![100, 100, 300]); // Duplicate row_id 100

        let result1 = merge_state.process_row_action(
            Action::UpdateAll,
            0,
            &row_addr_array,
            &row_id_array,
            &RecordBatch::new_empty(Arc::new(arrow_schema::Schema::empty())),
        );
        assert!(result1.is_ok(), "First call should succeed");

        let result2 = merge_state.process_row_action(
            Action::UpdateAll,
            1,
            &row_addr_array,
            &row_id_array,
            &RecordBatch::new_empty(Arc::new(arrow_schema::Schema::empty())),
        );
        assert!(
            result2.is_err(),
            "Second call with duplicate _rowid should fail"
        );

        let error_msg = result2.unwrap_err().to_string();
        assert!(
            error_msg.contains("Ambiguous merge insert")
                && error_msg.contains("multiple source rows"),
            "Error message should mention ambiguous merge insert and multiple source rows, got: {}",
            error_msg
        );

        let result3 = merge_state.process_row_action(
            Action::UpdateAll,
            2,
            &row_addr_array,
            &row_id_array,
            &RecordBatch::new_empty(Arc::new(arrow_schema::Schema::empty())),
        );
        assert!(
            result3.is_ok(),
            "Third call with different _rowid should succeed"
        );
    }

    #[test]
    fn test_merge_state_duplicate_rowid_first_seen() {
        let metrics = MergeInsertMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
        let mut merge_state = MergeState::new(
            metrics,
            false,
            Vec::new(),
            Vec::new(),
            SourceDedupeBehavior::FirstSeen,
        );

        let row_addr_array = UInt64Array::from(vec![1000, 2000, 3000]);
        let row_id_array = UInt64Array::from(vec![100, 100, 300]); // Duplicate row_id 100

        let result1 = merge_state.process_row_action(
            Action::UpdateAll,
            0,
            &row_addr_array,
            &row_id_array,
            &RecordBatch::new_empty(Arc::new(arrow_schema::Schema::empty())),
        );
        assert!(result1.is_ok(), "First call should succeed");
        assert_eq!(result1.unwrap(), Some(0), "First row should be kept");

        let result2 = merge_state.process_row_action(
            Action::UpdateAll,
            1,
            &row_addr_array,
            &row_id_array,
            &RecordBatch::new_empty(Arc::new(arrow_schema::Schema::empty())),
        );
        assert!(
            result2.is_ok(),
            "Second call with duplicate _rowid should succeed with FirstSeen"
        );
        assert_eq!(
            result2.unwrap(),
            None,
            "Duplicate row should be skipped (return None)"
        );

        // Verify the metric was incremented
        assert_eq!(
            merge_state.metrics.num_skipped_duplicates.value(),
            1,
            "num_skipped_duplicates should be 1"
        );

        let result3 = merge_state.process_row_action(
            Action::UpdateAll,
            2,
            &row_addr_array,
            &row_id_array,
            &RecordBatch::new_empty(Arc::new(arrow_schema::Schema::empty())),
        );
        assert!(
            result3.is_ok(),
            "Third call with different _rowid should succeed"
        );
        assert_eq!(result3.unwrap(), Some(2), "Third row should be kept");
    }
}