spiral-sdk 0.1.11

//! Merging source rows into a table by key, as one version.
//!
//! The join runs in the host: the source is held in memory, keyed by its `on` columns in
//! Arrow's row format, and the table's keys stream past it. Matched rows become patches or
//! deletions and unmatched source rows become appended rows, all in one [`TableWrite`].

use std::collections::HashMap;
use std::sync::Arc;

use arrow_array::cast::AsArray as _;
use arrow_array::types::{
    Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
};
use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, RecordBatch, UInt32Array};
use arrow_row::{RowConverter, Rows, SortField};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use arrow_select::concat::concat_batches;
use arrow_select::take::take;
use futures::{StreamExt as _, TryStreamExt as _, stream};
use tokio::sync::mpsc;
use tokio::task::{JoinHandle, spawn};
use vortex_array::expr::{
    Expression as VortexExpression, and, and_collect, col, gt_eq, lit, lt_eq,
};
use vortex_array::scalar::Scalar;

use crate::scandal::RecordBatchStream;
use crate::scandal::error::ScandalError;
use crate::scandal::types::FragmentReceipt;
use crate::scandal_backend::{RowEdit, reader_stream};
use crate::table_write::ROW_ID_COLUMN;
use crate::{Error, Expression, IntoArrowStream, ScanOptions, Table, TableWrite};

/// What a merge does with each kind of row.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MergeActions {
    /// What happens to a table row whose key a source row has.
    pub matched: Matched,
    /// The columns a match changes; every source column but the keys when absent.
    pub columns: Option<Vec<String>>,
    /// What happens to a source row whose key no table row has.
    pub not_matched: NotMatched,
    /// What happens to a table row whose key no source row has.
    pub not_matched_by_source: BySource,
}

impl Default for MergeActions {
    /// An upsert: patch matched rows, insert unmatched source rows, keep the other table rows.
    fn default() -> Self {
        Self {
            matched: Matched::Patch,
            columns: None,
            not_matched: NotMatched::Insert,
            not_matched_by_source: BySource::Keep,
        }
    }
}

/// What a merge does with a table row whose key a source row has.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Matched {
    /// Record the source row's values beside the data.
    Patch,
    /// Delete the table row.
    Delete,
    /// Leave the table row as it is.
    Ignore,
}

/// What a merge does with a source row whose key no table row has.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum NotMatched {
    /// Append the source row, which must carry the table's full schema.
    Insert,
    /// Leave the source row out.
    Ignore,
}

/// What a merge does with a table row whose key no source row has.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BySource {
    /// Leave the table row as it is.
    Keep,
    /// Delete the table row.
    Delete,
}

/// Rows a merge changed, and the version with the changes.
#[derive(Clone, Debug)]
pub struct MergedRows {
    /// The version with the changes; the original table when nothing changed.
    pub table: Table,
    /// Source rows appended.
    pub inserted: u64,
    /// Table rows whose values a match changed.
    pub updated: u64,
    /// Table rows deleted.
    pub deleted: u64,
}

impl Table {
    /// Write the next version of this table with `source` merged into it by the `on` columns,
    /// as `actions` say: one version with every patch, deletion and appended row.
    ///
    /// The source is held in memory, so a merge suits batches of upserts; join larger sources in
    /// DataFusion. Duplicate source keys are an error, and null keys match nothing. One source
    /// row may match several table rows and updates all of them. A matched row whose values
    /// are unchanged is not written. Unless unmatched table rows are deleted, the table scan is
    /// pruned by the bounds of the source's integer and string keys. When nothing changes,
    /// nothing is written and the result holds this table.
    pub async fn merge_rows(
        &self,
        source: impl IntoArrowStream,
        on: &[&str],
        actions: MergeActions,
    ) -> Result<MergedRows, Error> {
        let reader = source.into_arrow_stream()?;
        let source_schema = reader.schema();
        let batches = reader_stream(reader).try_collect::<Vec<_>>().await?;
        let source = concat_batches(&source_schema, &batches).map_err(Error::arrow)?;
        let schema = self.schema().await?;
        let merge = Merge::new(&schema, source, on, &actions)?;
        merge.write(self).await
    }
}

/// A merge's source rows, keyed by their `on` columns, and the actions that apply to them.
struct Merge {
    source: RecordBatch,
    on: Vec<String>,
    keys: RowConverter,
    /// The source row of each non-null key, by the key's row format.
    index: HashMap<Box<[u8]>, usize>,
    matched: Matched,
    not_matched: NotMatched,
    by_source: BySource,
    /// The columns a match updates, with the table's fields, and their source values in the
    /// row format, which compares them with the table's values.
    updates: Option<Updates>,
    /// The table's schema, which inserted rows take.
    table_schema: SchemaRef,
}

/// The columns a patch of matched rows writes.
struct Updates {
    /// The row ID column and the updated columns, with the table's fields.
    schema: SchemaRef,
    converter: RowConverter,
    values: Rows,
}

/// What a merge writes for one batch of table rows.
#[derive(Debug, Default)]
struct BatchEdits {
    /// Row IDs and new values of matched rows whose values change.
    updates: Option<RecordBatch>,
    /// Row IDs of the rows to delete.
    deletions: Option<RecordBatch>,
}

impl Merge {
    /// Check `source` against the table's `schema` and the `actions`, and key its rows.
    fn new(
        schema: &Schema,
        source: RecordBatch,
        on: &[&str],
        actions: &MergeActions,
    ) -> Result<Self, Error> {
        if on.is_empty() {
            return Err(Error::invalid(
                "a merge matches rows on at least one column",
            ));
        }
        if u32::try_from(source.num_rows()).is_err() {
            return Err(Error::invalid("a merge source holds fewer than 2^32 rows"));
        }
        let source_schema = source.schema();
        for key in on {
            same_type(schema, &source_schema, key)?;
        }
        let key_columns = columns(&source, on)?;
        let keys = converter(&key_columns)?;
        let index = key_index(&keys, &key_columns)?;
        let updates = match actions.matched {
            Matched::Patch => Updates::new(schema, &source, on, actions)?,
            Matched::Delete | Matched::Ignore => None,
        };
        if actions.not_matched == NotMatched::Insert {
            check_insertable(schema, &source_schema)?;
        }
        Ok(Self {
            source,
            on: on.iter().map(|key| (*key).to_owned()).collect(),
            keys,
            index,
            matched: actions.matched,
            not_matched: actions.not_matched,
            by_source: actions.not_matched_by_source,
            updates,
            table_schema: Arc::new(schema.clone()),
        })
    }

    /// Scan the table's keys, compared columns and row IDs, write what each batch changes, then
    /// append the unmatched source rows.
    async fn write(self, table: &Table) -> Result<MergedRows, Error> {
        let mut write = table.writer().await?;
        let mut matched_sources = vec![false; self.source.num_rows()];
        let (mut updated, mut deleted) = (0, 0);
        if self.reads_table() {
            let scan = table
                .scan(ScanOptions {
                    columns: Some(self.scanned_columns()),
                    filter: self.key_filter()?,
                    row_id_column: Some(ROW_ID_COLUMN.to_owned()),
                    ..Default::default()
                })
                .await?;
            let mut updates = self
                .updates
                .as_ref()
                .map(|updates| LazyPart::new(RowEdit::Patch, Arc::clone(&updates.schema)));
            let mut deletions = LazyPart::new(RowEdit::Delete, row_id_schema());
            let mut batches = scan.into_stream();
            while let Some(batch) = batches.next_batch().await {
                let edits = self.edit_batch(&batch?, &mut matched_sources)?;
                if let (Some(part), Some(rows)) = (updates.as_mut(), edits.updates) {
                    updated += rows.num_rows() as u64;
                    part.send(&mut write, rows).await?;
                }
                if let Some(rows) = edits.deletions {
                    deleted += rows.num_rows() as u64;
                    deletions.send(&mut write, rows).await?;
                }
            }
            if let Some(part) = updates {
                write.add_receipts(part.finish().await?);
            }
            write.add_receipts(deletions.finish().await?);
        }
        let inserts = self.inserted_rows(&matched_sources)?;
        let inserted = inserts.as_ref().map_or(0, |rows| rows.num_rows() as u64);
        if let Some(rows) = inserts {
            write.append(rows).await?;
        }
        if inserted + updated + deleted == 0 {
            return Ok(MergedRows {
                table: table.clone(),
                inserted,
                updated,
                deleted,
            });
        }
        Ok(MergedRows {
            table: write.finish().await?,
            inserted,
            updated,
            deleted,
        })
    }

    /// Whether any table row can be changed or can tell a source row it matched.
    fn reads_table(&self) -> bool {
        let changes_matches = self.matched == Matched::Delete || self.updates.is_some();
        self.by_source == BySource::Delete
            || (!self.index.is_empty()
                && (changes_matches || self.not_matched == NotMatched::Insert))
    }

    /// The keys, then the updated columns that are not keys.
    fn scanned_columns(&self) -> Vec<String> {
        let mut scanned = self.on.clone();
        if let Some(updates) = &self.updates {
            scanned.extend(
                updates
                    .updated_columns()
                    .filter(|column| !self.on.contains(column))
                    .cloned(),
            );
        }
        scanned
    }

    /// Bounds of the source's integer and string keys, which prune table files no source key
    /// can match. A merge that deletes unmatched table rows reads them all.
    fn key_filter(&self) -> Result<Option<Expression>, Error> {
        if self.by_source == BySource::Delete {
            return Ok(None);
        }
        let terms = self
            .on
            .iter()
            .filter_map(|key| {
                let column = self.source.column_by_name(key)?;
                let (min, max) = key_bounds(column.as_ref())?;
                Some(and(
                    gt_eq(col(key.as_str()), min),
                    lt_eq(col(key.as_str()), max),
                ))
            })
            .collect::<Vec<_>>();
        Ok(and_collect(terms).map(Expression::new).transpose()?)
    }

    /// The patches and deletions of one batch of table rows, marking the source
    /// rows it matched in `matched_sources`.
    fn edit_batch(
        &self,
        batch: &RecordBatch,
        matched_sources: &mut [bool],
    ) -> Result<BatchEdits, Error> {
        let row_ids = batch
            .column_by_name(ROW_ID_COLUMN)
            .ok_or_else(|| Error::Protocol(format!("the scan returned no `{ROW_ID_COLUMN}`")))?;
        let key_columns = columns(batch, &self.on)?;
        let keys = self
            .keys
            .convert_columns(&key_columns)
            .map_err(Error::arrow)?;
        let (mut changed, mut sources, mut deleted) = (Vec::new(), Vec::new(), Vec::new());
        for row in 0..batch.num_rows() {
            let source = key_columns
                .iter()
                .all(|key| key.is_valid(row))
                .then(|| self.index.get(keys.row(row).as_ref()))
                .flatten();
            let row = u32::try_from(row).map_err(Error::arrow)?;
            match source {
                Some(&source) => {
                    matched_sources[source] = true;
                    match self.matched {
                        Matched::Patch => {
                            changed.push(row);
                            sources.push(source as u32);
                        }
                        Matched::Delete => deleted.push(row),
                        Matched::Ignore => {}
                    }
                }
                None if self.by_source == BySource::Delete => deleted.push(row),
                None => {}
            }
        }
        let updates = match &self.updates {
            Some(updates) => {
                updates.changed_rows(batch, row_ids, changed, sources, &self.source)?
            }
            None => None,
        };
        Ok(BatchEdits {
            updates,
            deletions: row_id_rows(row_ids, deleted)?,
        })
    }

    /// The source rows to append: those no table row matched, in source order, with the
    /// table's schema.
    fn inserted_rows(&self, matched_sources: &[bool]) -> Result<Option<RecordBatch>, Error> {
        if self.not_matched == NotMatched::Ignore {
            return Ok(None);
        }
        let unmatched = (0..self.source.num_rows() as u32)
            .filter(|row| !matched_sources[*row as usize])
            .collect::<Vec<_>>();
        if unmatched.is_empty() {
            return Ok(None);
        }
        let indices = UInt32Array::from(unmatched);
        let columns =
            self.table_schema
                .fields()
                .iter()
                .map(|field| {
                    let column = self.source.column_by_name(field.name()).ok_or_else(|| {
                        Error::invalid(format!("source has no `{}`", field.name()))
                    })?;
                    take(column.as_ref(), &indices, None).map_err(Error::arrow)
                })
                .collect::<Result<Vec<_>, _>>()?;
        RecordBatch::try_new(Arc::clone(&self.table_schema), columns)
            .map(Some)
            .map_err(Error::arrow)
    }
}

impl Updates {
    /// The columns a match of `source` rows writes, or `None` when it writes none.
    fn new(
        schema: &Schema,
        source: &RecordBatch,
        on: &[&str],
        actions: &MergeActions,
    ) -> Result<Option<Self>, Error> {
        let source_schema = source.schema();
        let names = match &actions.columns {
            Some(names) => {
                if let Some(key) = names.iter().find(|name| on.contains(&name.as_str())) {
                    return Err(Error::invalid(format!(
                        "a merge does not update its key `{key}`"
                    )));
                }
                names.clone()
            }
            None => source_schema
                .fields()
                .iter()
                .map(|field| field.name().clone())
                .filter(|name| !on.contains(&name.as_str()))
                .collect(),
        };
        if names.is_empty() {
            return Ok(None);
        }
        let mut fields = vec![row_id_field()];
        for name in &names {
            fields.push(Arc::new(same_type(schema, &source_schema, name)?.clone()));
        }
        let source_columns = columns(source, &names)?;
        let converter = converter(&source_columns)?;
        let values = converter
            .convert_columns(&source_columns)
            .map_err(Error::arrow)?;
        Ok(Some(Self {
            schema: Arc::new(Schema::new(fields)),
            converter,
            values,
        }))
    }

    /// The updated columns, after the row ID column.
    fn updated_columns(&self) -> impl Iterator<Item = &String> {
        self.schema.fields()[1..].iter().map(|field| field.name())
    }

    /// The row IDs of the `table_rows` of `batch` whose values differ from their matching
    /// `source_rows`, with the source's values.
    fn changed_rows(
        &self,
        batch: &RecordBatch,
        row_ids: &ArrayRef,
        table_rows: Vec<u32>,
        source_rows: Vec<u32>,
        source: &RecordBatch,
    ) -> Result<Option<RecordBatch>, Error> {
        if table_rows.is_empty() {
            return Ok(None);
        }
        let names = self.updated_columns().cloned().collect::<Vec<_>>();
        let current = self
            .converter
            .convert_columns(&columns(batch, &names)?)
            .map_err(Error::arrow)?;
        let (table_rows, source_rows): (Vec<u32>, Vec<u32>) = table_rows
            .into_iter()
            .zip(source_rows)
            .filter(|(table, source)| {
                current.row(*table as usize) != self.values.row(*source as usize)
            })
            .unzip();
        if table_rows.is_empty() {
            return Ok(None);
        }
        let (table_rows, source_rows) = (
            UInt32Array::from(table_rows),
            UInt32Array::from(source_rows),
        );
        let mut values = vec![take(row_ids.as_ref(), &table_rows, None).map_err(Error::arrow)?];
        for column in columns(source, &names)? {
            values.push(take(column.as_ref(), &source_rows, None).map_err(Error::arrow)?);
        }
        RecordBatch::try_new(Arc::clone(&self.schema), values)
            .map(Some)
            .map_err(Error::arrow)
    }
}

/// One part of an edit whose rows a scan produces. Its write request starts with its first
/// rows, so a part without rows sends none, and runs beside the scan.
struct LazyPart {
    edit: RowEdit,
    schema: SchemaRef,
    request: Option<StartedPart>,
}

type StartedPart = (
    mpsc::Sender<Result<RecordBatch, ScandalError>>,
    JoinHandle<Result<Vec<FragmentReceipt>, Error>>,
);

/// Batches a part's request may hold before the scan waits for it.
const PART_BUFFER: usize = 2;

impl LazyPart {
    fn new(edit: RowEdit, schema: SchemaRef) -> Self {
        Self {
            edit,
            schema,
            request: None,
        }
    }

    /// Send `rows` to the part's request, starting it with the first rows.
    async fn send(&mut self, write: &mut TableWrite, rows: RecordBatch) -> Result<(), Error> {
        let (sender, request) = match self.request.take() {
            Some(started) => started,
            None => {
                let (sender, receiver) = mpsc::channel(PART_BUFFER);
                let request = write.start_edit(
                    self.edit,
                    Arc::clone(&self.schema),
                    receiver_stream(receiver),
                )?;
                (sender, spawn(request))
            }
        };
        if sender.send(Ok(rows)).await.is_err() {
            return Err(match joined(request).await {
                Err(error) => error,
                Ok(_) => Error::Protocol("a write request ended before reading its rows".into()),
            });
        }
        self.request = Some((sender, request));
        Ok(())
    }

    /// End the part's rows and return the receipts its request wrote.
    async fn finish(self) -> Result<Vec<FragmentReceipt>, Error> {
        let Some((sender, request)) = self.request else {
            return Ok(Vec::new());
        };
        drop(sender);
        joined(request).await
    }
}

async fn joined(
    request: JoinHandle<Result<Vec<FragmentReceipt>, Error>>,
) -> Result<Vec<FragmentReceipt>, Error> {
    request
        .await
        .map_err(|error| Error::Protocol(format!("a write request stopped: {error}")))?
}

fn receiver_stream(
    receiver: mpsc::Receiver<Result<RecordBatch, ScandalError>>,
) -> RecordBatchStream {
    stream::unfold(receiver, |mut receiver| async move {
        receiver.recv().await.map(|batch| (batch, receiver))
    })
    .boxed()
}

fn row_id_field() -> Arc<Field> {
    Arc::new(Field::new(ROW_ID_COLUMN, DataType::UInt64, false))
}

fn row_id_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![row_id_field()]))
}

/// The row IDs of `rows`, as rows to delete, or `None` when there are none.
fn row_id_rows(row_ids: &ArrayRef, rows: Vec<u32>) -> Result<Option<RecordBatch>, Error> {
    if rows.is_empty() {
        return Ok(None);
    }
    let ids = take(row_ids.as_ref(), &UInt32Array::from(rows), None).map_err(Error::arrow)?;
    RecordBatch::try_new(row_id_schema(), vec![ids])
        .map(Some)
        .map_err(Error::arrow)
}

/// The columns `names` of `batch`, in that order.
fn columns(batch: &RecordBatch, names: &[impl AsRef<str>]) -> Result<Vec<ArrayRef>, Error> {
    names
        .iter()
        .map(|name| {
            batch
                .column_by_name(name.as_ref())
                .cloned()
                .ok_or_else(|| Error::invalid(format!("no column `{}`", name.as_ref())))
        })
        .collect()
}

/// The table's field `column`, which the source must hold with the same type.
fn same_type<'a>(table: &'a Schema, source: &Schema, column: &str) -> Result<&'a Field, Error> {
    let field = table
        .field_with_name(column)
        .map_err(|_| Error::invalid(format!("the table has no column `{column}`")))?;
    let given = source
        .field_with_name(column)
        .map_err(|_| Error::invalid(format!("the merge source has no column `{column}`")))?;
    if field.data_type() != given.data_type() {
        return Err(Error::invalid(format!(
            "the merge source's `{column}` holds {}, not the stored {}",
            given.data_type(),
            field.data_type()
        )));
    }
    Ok(field)
}

/// Check that the source holds exactly the table's columns with their types, so its rows can
/// be appended.
fn check_insertable(table: &Schema, source: &Schema) -> Result<(), Error> {
    for field in table.fields() {
        same_type(table, source, field.name())?;
    }
    if let Some(extra) = source
        .fields()
        .iter()
        .find(|field| table.field_with_name(field.name()).is_err())
    {
        return Err(Error::invalid(format!(
            "inserting needs the table's columns only; the source also holds `{}`",
            extra.name()
        )));
    }
    Ok(())
}

fn converter(columns: &[ArrayRef]) -> Result<RowConverter, Error> {
    RowConverter::new(
        columns
            .iter()
            .map(|column| SortField::new(column.data_type().clone()))
            .collect(),
    )
    .map_err(Error::arrow)
}

/// The source row of each non-null key; a key that two rows hold is an error.
fn key_index(
    converter: &RowConverter,
    keys: &[ArrayRef],
) -> Result<HashMap<Box<[u8]>, usize>, Error> {
    let rows = converter.convert_columns(keys).map_err(Error::arrow)?;
    let mut index = HashMap::with_capacity(rows.num_rows());
    for row in 0..rows.num_rows() {
        if !keys.iter().all(|key| key.is_valid(row)) {
            continue;
        }
        if let Some(earlier) = index.insert(Box::from(rows.row(row).as_ref()), row) {
            return Err(Error::invalid(format!(
                "the merge source holds one key in rows {earlier} and {row}"
            )));
        }
    }
    Ok(index)
}

/// The least and greatest non-null values of an integer or string key, as literals.
fn key_bounds(column: &dyn Array) -> Option<(VortexExpression, VortexExpression)> {
    match column.data_type() {
        DataType::Int8 => primitive_bounds::<Int8Type>(column),
        DataType::Int16 => primitive_bounds::<Int16Type>(column),
        DataType::Int32 => primitive_bounds::<Int32Type>(column),
        DataType::Int64 => primitive_bounds::<Int64Type>(column),
        DataType::UInt8 => primitive_bounds::<UInt8Type>(column),
        DataType::UInt16 => primitive_bounds::<UInt16Type>(column),
        DataType::UInt32 => primitive_bounds::<UInt32Type>(column),
        DataType::UInt64 => primitive_bounds::<UInt64Type>(column),
        DataType::Utf8 => string_bounds(column.as_string::<i32>().iter()),
        DataType::LargeUtf8 => string_bounds(column.as_string::<i64>().iter()),
        DataType::Utf8View => string_bounds(column.as_string_view().iter()),
        _ => None,
    }
}

fn primitive_bounds<T: ArrowPrimitiveType>(
    column: &dyn Array,
) -> Option<(VortexExpression, VortexExpression)>
where
    T::Native: Into<Scalar> + PartialOrd,
{
    let (min, max) = bounds(column.as_primitive::<T>().iter().flatten())?;
    Some((lit(min), lit(max)))
}

fn string_bounds<'a>(
    values: impl Iterator<Item = Option<&'a str>>,
) -> Option<(VortexExpression, VortexExpression)> {
    let (min, max) = bounds(values.flatten())?;
    Some((lit(min), lit(max)))
}

fn bounds<T: PartialOrd + Copy>(values: impl Iterator<Item = T>) -> Option<(T, T)> {
    values.fold(None, |bounds, value| match bounds {
        None => Some((value, value)),
        Some((min, max)) => Some((
            if value < min { value } else { min },
            if value > max { value } else { max },
        )),
    })
}

#[cfg(test)]
mod tests {
    use arrow_array::{Int64Array, StringArray, UInt64Array};

    use super::*;

    fn table_schema() -> Schema {
        Schema::new(vec![
            Field::new("id", DataType::Int64, true),
            Field::new("label", DataType::Utf8, true),
        ])
    }

    fn source(ids: Vec<Option<i64>>, labels: Vec<&str>) -> RecordBatch {
        RecordBatch::try_new(
            Arc::new(table_schema()),
            vec![
                Arc::new(Int64Array::from(ids)),
                Arc::new(StringArray::from(labels)),
            ],
        )
        .unwrap()
    }

    /// Table rows as a scan with row IDs returns them.
    fn table_rows(ids: Vec<Option<i64>>, labels: Vec<&str>, row_ids: Vec<u64>) -> RecordBatch {
        let schema = Schema::new(vec![
            Field::new("id", DataType::Int64, true),
            Field::new("label", DataType::Utf8, true),
            Field::new(ROW_ID_COLUMN, DataType::UInt64, false),
        ]);
        RecordBatch::try_new(
            Arc::new(schema),
            vec![
                Arc::new(Int64Array::from(ids)),
                Arc::new(StringArray::from(labels)),
                Arc::new(UInt64Array::from(row_ids)),
            ],
        )
        .unwrap()
    }

    fn row_ids(rows: &RecordBatch) -> Vec<u64> {
        rows.column_by_name(ROW_ID_COLUMN)
            .unwrap()
            .as_primitive::<UInt64Type>()
            .values()
            .to_vec()
    }

    #[test]
    fn matches_patch_changed_rows_and_insert_unmatched_source_rows() {
        let merge = Merge::new(
            &table_schema(),
            source(
                vec![Some(1), Some(2), None, Some(4)],
                vec!["a", "B", "n", "d"],
            ),
            &["id"],
            &MergeActions::default(),
        )
        .unwrap();
        let mut matched = vec![false; 4];
        let edits = merge
            .edit_batch(
                &table_rows(
                    vec![Some(1), Some(2), Some(2), None, Some(3)],
                    vec!["a", "b", "x", "n", "c"],
                    vec![10, 11, 12, 13, 14],
                ),
                &mut matched,
            )
            .unwrap();
        // Row 10 already holds "a"; both rows keyed 2 take "B"; null keys match nothing.
        let updates = edits.updates.unwrap();
        assert_eq!(row_ids(&updates), [11, 12]);
        assert_eq!(
            updates
                .column(1)
                .as_string::<i32>()
                .iter()
                .collect::<Vec<_>>(),
            [Some("B"), Some("B")]
        );
        assert!(edits.deletions.is_none());
        assert_eq!(matched, [true, true, false, false]);
        let inserted = merge.inserted_rows(&matched).unwrap().unwrap();
        assert_eq!(
            inserted
                .column(1)
                .as_string::<i32>()
                .iter()
                .collect::<Vec<_>>(),
            [Some("n"), Some("d")]
        );
        assert_eq!(inserted.schema().as_ref(), &table_schema());
    }

    #[test]
    fn deletes_follow_the_matched_and_by_source_actions() {
        let actions = MergeActions {
            matched: Matched::Delete,
            columns: None,
            not_matched: NotMatched::Ignore,
            not_matched_by_source: BySource::Delete,
        };
        let merge = Merge::new(
            &table_schema(),
            source(vec![Some(1)], vec!["a"]),
            &["id"],
            &actions,
        )
        .unwrap();
        assert!(merge.key_filter().unwrap().is_none());
        let mut matched = vec![false];
        let edits = merge
            .edit_batch(
                &table_rows(
                    vec![Some(1), Some(2), None],
                    vec!["a", "b", "c"],
                    vec![7, 8, 9],
                ),
                &mut matched,
            )
            .unwrap();
        assert!(edits.updates.is_none());
        assert_eq!(row_ids(&edits.deletions.unwrap()), [7, 8, 9]);
        assert!(merge.inserted_rows(&matched).unwrap().is_none());
    }

    #[test]
    fn sources_with_duplicate_keys_or_other_types_are_refused() {
        let actions = MergeActions::default();
        let duplicate = source(vec![Some(1), Some(1)], vec!["a", "b"]);
        assert!(Merge::new(&table_schema(), duplicate, &["id"], &actions).is_err());
        let nulls = source(vec![None, None], vec!["a", "b"]);
        assert!(Merge::new(&table_schema(), nulls, &["id"], &actions).is_ok());
        let retyped = Schema::new(vec![
            Field::new("id", DataType::Int32, true),
            Field::new("label", DataType::Utf8, true),
        ]);
        assert!(
            Merge::new(
                &retyped,
                source(vec![Some(1)], vec!["a"]),
                &["id"],
                &actions
            )
            .is_err()
        );
        let keyed_update = MergeActions {
            columns: Some(vec!["id".to_owned()]),
            ..MergeActions::default()
        };
        assert!(
            Merge::new(
                &table_schema(),
                source(vec![Some(1)], vec!["a"]),
                &["id"],
                &keyed_update
            )
            .is_err()
        );
    }

    #[test]
    fn integer_and_string_keys_bound_the_scan() {
        let merge = Merge::new(
            &table_schema(),
            source(vec![Some(4), None, Some(-2)], vec!["a", "b", "c"]),
            &["id", "label"],
            &MergeActions {
                matched: Matched::Ignore,
                ..MergeActions::default()
            },
        )
        .unwrap();
        assert!(merge.key_filter().unwrap().is_some());
        assert_eq!(merge.scanned_columns(), ["id", "label"]);
        let ids = Int64Array::from(vec![Some(4), None, Some(-2)]);
        let (min, max) = key_bounds(&ids).unwrap();
        assert_eq!((min, max), (lit(-2i64), lit(4i64)));
        assert!(key_bounds(&Int64Array::from(vec![None::<i64>])).is_none());
    }
}