spiral-sdk 0.1.11

//! Rows that changed between two versions of a table, read from a change plan.
//!
//! A change plan comes from version metadata alone. Reading the changed rows composes scans of
//! both versions, each restricted to the plan's row IDs, into one [`Scan`].

use std::collections::{HashMap, HashSet};
use std::iter::repeat_n;
use std::mem::take;
use std::ops::Range;
use std::sync::Arc;

use arrow_array::cast::AsArray as _;
use arrow_array::types::UInt64Type;
use arrow_array::{ArrayRef, BooleanArray, RecordBatch, StringArray, UInt64Array, new_null_array};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use arrow_select::concat::concat_batches;
use arrow_select::filter::filter_record_batch;
use arrow_select::interleave::interleave;
use async_trait::async_trait;
use futures::stream::BoxStream;
use futures::{StreamExt as _, TryStreamExt as _, stream};
use uuid::Uuid;

use crate::catalog_types::{
    Commit as CatalogCommit, CommitNode, CommitSelector, CommitTarget, GraphTraversal,
    ProjectCommitId, ReadCommitLogRequest, ResourceChange,
};
use crate::scan::{ScanInfo, ScanSource, SplitFilter};
use crate::scandal::api::{ChangePlan as WireChangePlan, FragmentDelta as WireDelta, RowIdRange};
use crate::scandal::types::FilterClause;
use crate::table_write::ROW_ID_COLUMN;
use crate::{BatchStream, Error, Scan, ScanOptions, Snapshot, Table, WorkItem};

/// The column of a change read that says how each row changed: one of the
/// [`ChangeKind`] names.
pub const CHANGE_KIND_COLUMN: &str = "s:change_kind";

/// How a row of a change read changed.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum ChangeKind {
    /// The row was added: `insert`.
    Insert,
    /// The row was deleted: `delete`.
    Delete,
    /// The row's values before an update: `update_before`.
    UpdateBefore,
    /// The row's values after an update: `update_after`.
    UpdateAfter,
}

impl ChangeKind {
    /// The name [`CHANGE_KIND_COLUMN`] holds for this kind.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Insert => "insert",
            Self::Delete => "delete",
            Self::UpdateBefore => "update_before",
            Self::UpdateAfter => "update_after",
        }
    }
}

/// How one version of a table changed since an earlier one, read from version metadata.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ChangePlan {
    /// One entry for each fragment of the later version, in stored order.
    pub fragments: Vec<FragmentChange>,
    /// Fragments of the earlier version that the later one no longer has. Each of their rows
    /// was deleted or carried into an added fragment.
    pub removed_fragment_ids: Vec<Uuid>,
    /// Requested columns whose type differs from their type in the earlier version, including
    /// columns the earlier version lacks.
    pub retyped_columns: Vec<String>,
}

/// One fragment of the later version of a [`ChangePlan`].
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FragmentChange {
    /// The fragment's stable ID.
    pub fragment_id: Uuid,
    /// The fragment's stored row count, deleted rows included.
    pub row_count: u64,
    /// How the fragment relates to the earlier version.
    pub change: FragmentDelta,
}

/// How one fragment relates to the earlier version of a [`ChangePlan`]. Each set of row IDs is
/// sorted, disjoint, half-open ranges.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum FragmentDelta {
    /// A fragment of the earlier version, at the same positions. Its other rows are unchanged.
    Retained {
        /// Rows kept whose requested columns a patch or column write changed.
        patched: Vec<Range<u64>>,
        /// Rows deleted since.
        deleted: Vec<Range<u64>>,
    },
    /// A fragment written since the earlier version.
    Added {
        /// Rows added since, with row IDs past every row ID the earlier version allocated.
        appended: Vec<Range<u64>>,
        /// Rows of the earlier version that a rewrite, such as compaction, moved here.
        carried: Vec<Range<u64>>,
        /// Carried rows whose requested columns a patch changed.
        patched: Vec<Range<u64>>,
    },
}

impl TryFrom<WireChangePlan> for ChangePlan {
    type Error = Error;

    /// The change plan a scandal returned, as sorted row ID ranges.
    fn try_from(plan: WireChangePlan) -> Result<Self, Error> {
        let fragments = plan
            .fragments
            .into_iter()
            .map(|fragment| {
                let change = match fragment.change {
                    WireDelta::Retained(retained) => FragmentDelta::Retained {
                        patched: ranges(&retained.patched),
                        deleted: ranges(&retained.deleted),
                    },
                    WireDelta::Added(added) => FragmentDelta::Added {
                        appended: ranges(&added.appended),
                        carried: ranges(&added.carried),
                        patched: ranges(&added.patched),
                    },
                    WireDelta::Unknown { tag, .. } => {
                        return Err(Error::Protocol(format!("unknown fragment change `{tag}`")));
                    }
                };
                Ok(FragmentChange {
                    fragment_id: fragment.fragment_id,
                    row_count: fragment.row_count.0,
                    change,
                })
            })
            .collect::<Result<_, Error>>()?;
        Ok(Self {
            fragments,
            removed_fragment_ids: plan.removed_fragment_ids,
            retyped_columns: plan.retyped_columns,
        })
    }
}

fn ranges(wire: &[RowIdRange]) -> Vec<Range<u64>> {
    wire.iter()
        .map(|range| range.start.0..range.end.0)
        .collect()
}

/// Which rows a change read returns, and how.
#[derive(Clone, Debug, Default)]
pub struct ChangeOptions {
    /// The columns, projection and filter of each image. A change read sets the row IDs itself,
    /// and takes no limit, row range, split token or emit.
    pub scan: ScanOptions,
    /// Which images of each changed row to return.
    pub images: ChangeImages,
}

/// Which images of each changed row a change read returns. Every row carries its row ID in
/// [`ROW_ID_COLUMN`] and its kind in [`CHANGE_KIND_COLUMN`].
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum ChangeImages {
    /// Only the kind and row ID, read from the change plan without reading column values.
    /// Updates are reported as `update_after`. Takes no filter.
    Keys,
    /// Values of inserted rows and the new values of updated rows, as `update_after`; deleted
    /// rows carry only their row ID.
    #[default]
    After,
    /// Values of inserted and deleted rows, and adjacent `update_before` and `update_after`
    /// images of each update.
    BeforeAndAfter,
}

/// The changes one Catalog commit made to a table: the versions it advanced between, and the
/// changed rows.
pub struct TableCommitChanges {
    /// The commit that advanced the table.
    pub commit_id: ProjectCommitId,
    /// The table's version before the commit.
    pub since: Table,
    /// The table's version the commit published.
    pub table: Table,
    /// The changed rows, as [`Table::changes_since`] reads them.
    pub changes: Scan,
}

impl Table {
    /// How this version changed since `since`, from version metadata only: for each fragment,
    /// whether `since` has it at the same positions, with the row IDs patched and deleted since,
    /// or whether it was added since, with the row IDs appended since and those a rewrite
    /// carried over. `columns` names the columns whose changes count; all when absent. `since`
    /// must be this version or one of its ancestors.
    pub async fn plan_changes(
        &self,
        since: &Table,
        columns: Option<Vec<String>>,
    ) -> Result<ChangePlan, Error> {
        if since.identity().resource_id != self.identity().resource_id {
            return Err(Error::invalid("changes compare two versions of one table"));
        }
        let plan = self
            .connection
            .tables()?
            .plan_changes(self.identity(), since.identity().version_id, columns)
            .await?;
        ChangePlan::try_from(plan)
    }

    /// The rows that changed since `since`, with their [`ChangeKind`] and row ID, as one scan.
    ///
    /// A row a rewrite carried without changing it is not reported; across an overwrite of the
    /// whole table every old row is deleted and every new row inserted. The filter applies to
    /// each image, and an update is reported when either image matches. Before images are read
    /// from `since` with the same columns and must have this version's names and types there.
    /// The scan has no splits to read by work item.
    pub async fn changes_since(
        &self,
        since: &Table,
        options: ChangeOptions,
    ) -> Result<Scan, Error> {
        check_change_options(&options)?;
        let counted = match options.scan.projection {
            Some(_) => None,
            None => options.scan.columns.clone(),
        };
        let plan = self.plan_changes(since, counted).await?;
        let reader = ChangeReader::open(
            self.clone(),
            since.clone(),
            options.clone(),
            ChangedRows::from_plan(&plan),
        )
        .await?;
        let info = ScanInfo {
            source_fields: None,
            schema: reader.schema.as_ref().clone(),
            filter_clauses: reader.filter_clauses.clone(),
            write_layout: None,
            statistics: None,
        };
        Scan::new(
            self.identity(),
            options.scan,
            info,
            Arc::new(ChangeSource(Arc::new(reader))),
        )
    }
}

impl Snapshot {
    /// The changes to the table at `path` in each commit between `since` and this snapshot,
    /// walking first-parent Catalog history from the oldest commit. Each commit that advanced
    /// the table contributes one scan, read as [`Table::changes_since`] reads it. `since` must
    /// be a first-parent ancestor of this snapshot, and `path` must name the same table in
    /// both.
    pub async fn table_changes(
        &self,
        path: &str,
        since: &Snapshot,
        options: ChangeOptions,
    ) -> Result<Vec<TableCommitChanges>, Error> {
        if since.project.id.as_str() != self.project.id.as_str() {
            return Err(Error::invalid("changes compare two commits of one project"));
        }
        let resource_id = self.table(path).await?.identity().resource_id;
        if since.table(path).await?.identity().resource_id != resource_id {
            return Err(Error::invalid(format!(
                "`{path}` names another table at {}",
                since.commit_id
            )));
        }
        let mut changes = Vec::new();
        for node in self.first_parent_commits(since).await? {
            if node.resource_change_count == 0 {
                continue;
            }
            let commit = self
                .project
                .client
                .connection
                .catalog
                .read_journal_commit(&self.project.id, node.commit_id)
                .await?;
            let Some((before, after)) = table_advance(&commit, resource_id)? else {
                continue;
            };
            let before = self.project.client.table_by_id(resource_id, before);
            let after = self.project.client.table_by_id(resource_id, after);
            let scan = after.changes_since(&before, options.clone()).await?;
            changes.push(TableCommitChanges {
                commit_id: node.commit_id,
                since: before,
                table: after,
                changes: scan,
            });
        }
        Ok(changes)
    }

    /// The commits after `since` on this snapshot's first-parent chain, oldest first.
    async fn first_parent_commits(&self, since: &Snapshot) -> Result<Vec<CommitNode>, Error> {
        let catalog = &self.project.client.connection.catalog;
        let excluded = vec![CommitSelector::Commit(CommitTarget {
            commit_id: since.commit_id,
        })];
        let mut nodes = HashMap::new();
        let mut starts = vec![self.commit_id];
        while !starts.is_empty() {
            let log = catalog
                .read_commit_log(
                    &self.project.id,
                    &ReadCommitLogRequest {
                        starts: starts
                            .iter()
                            .map(|commit_id| {
                                CommitSelector::Commit(CommitTarget {
                                    commit_id: *commit_id,
                                })
                            })
                            .collect(),
                        exclude_ancestors_of: Some(excluded.clone()),
                        max_commits: None,
                        traversal: Some(GraphTraversal::FirstParent),
                    },
                )
                .await?;
            if log.commits.is_empty() && !log.continuation_commit_ids.is_empty() {
                return Err(Error::Protocol(
                    "the commit log continued without returning commits".into(),
                ));
            }
            for node in log.commits {
                nodes.insert(node.commit_id.0, node);
            }
            starts = log
                .continuation_commit_ids
                .into_iter()
                .filter(|commit_id| !nodes.contains_key(&commit_id.0))
                .collect();
        }
        let not_ancestor = || {
            Error::invalid(format!(
                "commit {} is not a first-parent ancestor of {}",
                since.commit_id, self.commit_id
            ))
        };
        let mut chain = Vec::new();
        let mut current = self.commit_id;
        while current != since.commit_id {
            let node = nodes.remove(&current.0).ok_or_else(not_ancestor)?;
            current = *node.parent_commit_ids.first().ok_or_else(not_ancestor)?;
            chain.push(node);
        }
        chain.reverse();
        Ok(chain)
    }
}

/// The versions a commit advanced the Table `resource_id` between, or `None` when it left the
/// table as it was.
fn table_advance(commit: &CatalogCommit, resource_id: Uuid) -> Result<Option<(Uuid, Uuid)>, Error> {
    for change in &commit.resources {
        match change {
            ResourceChange::Advance(advance) if advance.resource_id.0 == resource_id => {
                return Ok(Some((advance.expected.0, advance.replacement.0)));
            }
            ResourceChange::Create(create) if create.descriptor.resource_id.0 == resource_id => {
                return Err(Error::invalid(format!(
                    "commit {} created the table, so it has no earlier version",
                    commit.commit_id
                )));
            }
            ResourceChange::Retire(retire) if retire.resource_id.0 == resource_id => {
                return Err(Error::invalid(format!(
                    "commit {} retired the table",
                    commit.commit_id
                )));
            }
            _ => {}
        }
    }
    Ok(None)
}

fn check_change_options(options: &ChangeOptions) -> Result<(), Error> {
    let scan = &options.scan;
    if scan.limit.is_some()
        || scan.row_range.is_some()
        || scan.split_token.is_some()
        || scan.emit.is_some()
        || scan.row_ids.is_some()
    {
        return Err(Error::invalid(
            "a change read takes no limit, row range, split token, emit or row IDs",
        ));
    }
    if scan
        .row_id_column
        .as_ref()
        .is_some_and(|column| column != ROW_ID_COLUMN)
    {
        return Err(Error::invalid(format!(
            "a change read names row IDs `{ROW_ID_COLUMN}`"
        )));
    }
    if options.images == ChangeImages::Keys && scan.filter.is_some() {
        return Err(Error::invalid(
            "a change read of keys reads no values to filter",
        ));
    }
    Ok(())
}

/// The row IDs a change plan reports, by what happened to them.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
struct ChangedRows {
    inserted: Vec<Range<u64>>,
    updated: Vec<Range<u64>>,
    /// Rows deleted from fragments both versions have.
    deleted: Vec<Range<u64>>,
    /// Fragments of the earlier version that the later one dropped. Their rows were deleted
    /// unless carried.
    removed_fragments: HashSet<Uuid>,
    carried: Vec<Range<u64>>,
}

impl ChangedRows {
    fn from_plan(plan: &ChangePlan) -> Self {
        let mut rows = Self {
            removed_fragments: plan.removed_fragment_ids.iter().copied().collect(),
            ..Self::default()
        };
        for fragment in &plan.fragments {
            match &fragment.change {
                FragmentDelta::Retained { patched, deleted } => {
                    rows.updated.extend(patched.iter().cloned());
                    rows.deleted.extend(deleted.iter().cloned());
                }
                FragmentDelta::Added {
                    appended,
                    carried,
                    patched,
                } => {
                    rows.inserted.extend(appended.iter().cloned());
                    rows.carried.extend(carried.iter().cloned());
                    rows.updated.extend(patched.iter().cloned());
                }
            }
        }
        Self {
            inserted: union_ranges(rows.inserted),
            updated: union_ranges(rows.updated),
            deleted: union_ranges(rows.deleted),
            carried: union_ranges(rows.carried),
            removed_fragments: rows.removed_fragments,
        }
    }
}

/// Row IDs one image read selects, at most this many ranges at once.
const RANGES_PER_READ: usize = 4096;
/// Updates paired at once, by row ID, when reading before and after images.
const PAIRED_ROWS: u64 = 65_536;
/// Rows of each batch a change read builds from row IDs alone.
const ID_BATCH_ROWS: u64 = 65_536;

/// The version an image is read from.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Image {
    Before,
    After,
}

/// What an image read returns for each row.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Returns {
    /// The image's columns and the row ID.
    Values,
    /// The row ID alone.
    Ids,
}

type Batches = BoxStream<'static, Result<RecordBatch, Error>>;

/// The reads of one change read: what changed, and how to read each image.
struct ChangeReader {
    after: Table,
    before: Table,
    images: ChangeImages,
    /// The caller's columns, projection and filter, which every image read applies.
    scan: ScanOptions,
    rows: ChangedRows,
    /// The image columns, nullable, then the row ID and the change kind.
    schema: SchemaRef,
    filter_clauses: Vec<FilterClause>,
}

impl ChangeReader {
    /// Plan an image read of each version to learn the output schema and check the filter.
    async fn open(
        after: Table,
        before: Table,
        options: ChangeOptions,
        rows: ChangedRows,
    ) -> Result<Self, Error> {
        let mut reader = Self {
            after,
            before,
            images: options.images,
            scan: options.scan,
            rows,
            schema: output_schema(&[]),
            filter_clauses: Vec::new(),
        };
        if reader.images == ChangeImages::Keys {
            return Ok(reader);
        }
        let probe = reader.read_options(Some(Vec::new()), Returns::Values);
        let after_scan = reader.after.scan(probe.clone()).await?;
        let image = after_scan.schema();
        let image_fields = &image.fields()[..image.fields().len() - 1];
        if reader.images == ChangeImages::BeforeAndAfter {
            let before = reader.before.scan(probe).await?.schema();
            if !same_columns(&image, &before) {
                return Err(Error::invalid(
                    "before images need the earlier version to hold the read columns with this \
                     version's names and types",
                ));
            }
        }
        reader.schema = output_schema(
            &image_fields
                .iter()
                .map(|field| Field::clone(field).with_nullable(true))
                .collect::<Vec<_>>(),
        );
        reader.filter_clauses = after_scan.filter_clauses().to_vec();
        Ok(reader)
    }

    /// Every changed row: deletions, then insertions, then updates.
    fn read(self: Arc<Self>) -> Batches {
        let parts = match self.images {
            ChangeImages::Keys => vec![
                Arc::clone(&self).deleted_rows(Returns::Ids),
                Arc::clone(&self).id_rows(self.rows.inserted.clone(), ChangeKind::Insert),
                Arc::clone(&self).id_rows(self.rows.updated.clone(), ChangeKind::UpdateAfter),
            ],
            ChangeImages::After => vec![
                Arc::clone(&self).deleted_rows(Returns::Ids),
                Arc::clone(&self).inserted_rows(),
                Arc::clone(&self).updated_after(),
            ],
            ChangeImages::BeforeAndAfter => vec![
                Arc::clone(&self).deleted_rows(Returns::Values),
                Arc::clone(&self).inserted_rows(),
                Arc::clone(&self).updated_pairs(),
            ],
        };
        stream::iter(parts).flatten().boxed()
    }

    /// Rows deleted since the earlier version: those deleted from fragments both versions
    /// have, then those of dropped fragments that no rewrite carried.
    fn deleted_rows(self: Arc<Self>, returns: Returns) -> Batches {
        let deleted = self.rows.deleted.clone();
        let retained = if returns == Returns::Ids && self.scan.filter.is_none() {
            Arc::clone(&self).id_rows(deleted, ChangeKind::Delete)
        } else {
            Arc::clone(&self).image_rows(Image::Before, deleted, returns, ChangeKind::Delete)
        };
        retained.chain(self.dropped_rows(returns)).boxed()
    }

    fn inserted_rows(self: Arc<Self>) -> Batches {
        let inserted = self.rows.inserted.clone();
        self.image_rows(Image::After, inserted, Returns::Values, ChangeKind::Insert)
    }

    /// The new values of updated rows. With a filter, an update is reported when either image
    /// matches, so each group of updates first reads both images' matching row IDs.
    fn updated_after(self: Arc<Self>) -> Batches {
        let updated = self.rows.updated.clone();
        if self.scan.filter.is_none() {
            return self.image_rows(
                Image::After,
                updated,
                Returns::Values,
                ChangeKind::UpdateAfter,
            );
        }
        stream::iter(chunk_ranges(&updated, PAIRED_ROWS))
            .then(move |chunk| {
                let reader = Arc::clone(&self);
                async move {
                    let matching = reader.matching_ids(chunk).await?;
                    let after = reader.read_images(Image::After, matching).await?;
                    after
                        .into_iter()
                        .map(|batch| {
                            reader.output(&batch, kinds(ChangeKind::UpdateAfter, batch.num_rows()))
                        })
                        .collect::<Result<Vec<_>, _>>()
                }
            })
            .map_ok(|batches| stream::iter(batches.into_iter().map(Ok)))
            .try_flatten()
            .boxed()
    }

    /// Adjacent before and after images of each update, a group of updates at a time.
    fn updated_pairs(self: Arc<Self>) -> Batches {
        let updated = self.rows.updated.clone();
        stream::iter(chunk_ranges(&updated, PAIRED_ROWS))
            .then(move |chunk| {
                let reader = Arc::clone(&self);
                async move {
                    let matching = match reader.scan.filter {
                        Some(_) => reader.matching_ids(chunk).await?,
                        None => chunk,
                    };
                    let before = reader.read_images(Image::Before, matching.clone()).await?;
                    let after = reader.read_images(Image::After, matching).await?;
                    let (before, after) = (
                        concat_batches(&image_schema(&reader.schema, &before)?, &before)
                            .map_err(Error::arrow)?,
                        concat_batches(&image_schema(&reader.schema, &after)?, &after)
                            .map_err(Error::arrow)?,
                    );
                    let (paired, kinds) = pair_images(&before, &after)?;
                    reader.output(&paired, kinds)
                }
            })
            .boxed()
    }

    /// The rows of fragments the earlier version had and this one dropped, less those a
    /// rewrite carried, read from the earlier version.
    fn dropped_rows(self: Arc<Self>, returns: Returns) -> Batches {
        if self.rows.removed_fragments.is_empty() {
            return stream::empty().boxed();
        }
        stream::once(async move {
            let scan = self.before.scan(self.read_options(None, returns)).await?;
            let mut work = Vec::new();
            for (split, planned) in scan.splits().iter().enumerate() {
                let fragment_id = planned.fragment_id.ok_or_else(|| {
                    Error::invalid("the scandal does not name the fragments a scan's splits read")
                })?;
                if self.rows.removed_fragments.contains(&fragment_id) {
                    work.push(WorkItem {
                        split,
                        rows: planned.rows.clone(),
                    });
                }
            }
            if work.is_empty() {
                return Ok::<Batches, Error>(stream::empty().boxed());
            }
            let batches = scan.read_work(work).await?;
            Ok(batches
                .map(move |read| {
                    let read = read?;
                    let deleted = without_rows(&read.batch, &self.rows.carried)?;
                    self.output(&deleted, kinds(ChangeKind::Delete, deleted.num_rows()))
                })
                .boxed())
        })
        .try_flatten()
        .boxed()
    }

    /// Rows of `kind` built from their row IDs alone, without reading either version.
    fn id_rows(self: Arc<Self>, ids: Vec<Range<u64>>, kind: ChangeKind) -> Batches {
        stream::iter(chunk_ranges(&ids, ID_BATCH_ROWS))
            .map(move |chunk| {
                let ids = UInt64Array::from_iter_values(chunk.into_iter().flatten());
                let batch = RecordBatch::try_new(row_id_schema(), vec![Arc::new(ids)])
                    .map_err(Error::arrow)?;
                self.output(&batch, kinds(kind, batch.num_rows()))
            })
            .boxed()
    }

    /// Rows `ids` of one version, filtered, as rows of `kind`.
    fn image_rows(
        self: Arc<Self>,
        image: Image,
        ids: Vec<Range<u64>>,
        returns: Returns,
        kind: ChangeKind,
    ) -> Batches {
        stream::iter(
            ids.chunks(RANGES_PER_READ)
                .map(<[Range<u64>]>::to_vec)
                .collect::<Vec<_>>(),
        )
        .then(move |chunk| {
            let reader = Arc::clone(&self);
            async move {
                let options = reader.read_options(Some(chunk), returns);
                let batches = reader.table(image).scan(options).await?.into_stream();
                Ok::<_, Error>(batches.map(move |batch| {
                    let batch = batch?;
                    reader.output(&batch, kinds(kind, batch.num_rows()))
                }))
            }
        })
        .try_flatten()
        .boxed()
    }

    /// The row IDs among `ids` whose image in either version matches the filter.
    async fn matching_ids(&self, ids: Vec<Range<u64>>) -> Result<Vec<Range<u64>>, Error> {
        let mut matching = Vec::new();
        for image in [Image::Before, Image::After] {
            let options = self.read_options(Some(ids.clone()), Returns::Ids);
            let mut batches = self.table(image).scan(options).await?.into_stream();
            while let Some(batch) = batches.next_batch().await {
                matching.extend_from_slice(row_ids(&batch?)?.values());
            }
        }
        Ok(ranges_of(matching))
    }

    /// Every row `ids` of one version, unfiltered, with the image's columns.
    async fn read_images(
        &self,
        image: Image,
        ids: Vec<Range<u64>>,
    ) -> Result<Vec<RecordBatch>, Error> {
        if ids.is_empty() {
            return Ok(Vec::new());
        }
        let options = ScanOptions {
            filter: None,
            ..self.read_options(Some(ids), Returns::Values)
        };
        self.table(image)
            .scan(options)
            .await?
            .into_stream()
            .try_collect()
            .await
    }

    fn table(&self, image: Image) -> &Table {
        match image {
            Image::Before => &self.before,
            Image::After => &self.after,
        }
    }

    /// Scan options that read `row_ids` with the caller's filter, returning the image's
    /// columns or only row IDs.
    fn read_options(&self, row_ids: Option<Vec<Range<u64>>>, returns: Returns) -> ScanOptions {
        let (columns, projection) = match returns {
            Returns::Values => (self.scan.columns.clone(), self.scan.projection.clone()),
            Returns::Ids => (Some(Vec::new()), None),
        };
        ScanOptions {
            columns,
            projection,
            filter: self.scan.filter.clone(),
            batch_rows: self.scan.batch_rows,
            collect_metrics: self.scan.collect_metrics,
            row_ids,
            row_id_column: Some(ROW_ID_COLUMN.to_owned()),
            ..Default::default()
        }
    }

    /// `batch`, image columns then row IDs or row IDs alone, as output rows of `kinds`.
    fn output(&self, batch: &RecordBatch, kinds: ArrayRef) -> Result<RecordBatch, Error> {
        let image_fields = &self.schema.fields()[..self.schema.fields().len() - 2];
        let rows = batch.num_rows();
        let mut columns = match batch.num_columns() {
            count if count == image_fields.len() + 1 => batch.columns().to_vec(),
            1 => {
                let mut columns = image_fields
                    .iter()
                    .map(|field| new_null_array(field.data_type(), rows))
                    .collect::<Vec<_>>();
                columns.push(Arc::clone(batch.column(0)));
                columns
            }
            count => {
                return Err(Error::Protocol(format!(
                    "a change read's image has {count} columns, not {}",
                    image_fields.len() + 1
                )));
            }
        };
        columns.push(kinds);
        RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Error::arrow)
    }
}

/// A change read as the source of a [`Scan`].
struct ChangeSource(Arc<ChangeReader>);

#[async_trait]
impl ScanSource for ChangeSource {
    async fn read(&self, _narrowing: Option<Arc<dyn SplitFilter>>) -> Result<BatchStream, Error> {
        let schema = Arc::clone(&self.0.schema);
        Ok(BatchStream::from_stream(schema, Arc::clone(&self.0).read()))
    }
}

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

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

/// `image_fields`, then the row ID and the change kind.
fn output_schema(image_fields: &[Field]) -> SchemaRef {
    let mut fields = image_fields.to_vec();
    fields.push(row_id_field());
    fields.push(Field::new(CHANGE_KIND_COLUMN, DataType::Utf8, false));
    Arc::new(Schema::new(fields))
}

/// The schema of image batches, from the first of them or else the output's image columns.
fn image_schema(output: &Schema, batches: &[RecordBatch]) -> Result<SchemaRef, Error> {
    if let Some(batch) = batches.first() {
        return Ok(batch.schema());
    }
    let fields = &output.fields()[..output.fields().len() - 1];
    Ok(Arc::new(Schema::new(fields.to_vec())))
}

/// Whether two image schemas hold the same columns, by name and type.
fn same_columns(left: &Schema, right: &Schema) -> bool {
    left.fields().len() == right.fields().len()
        && left
            .fields()
            .iter()
            .zip(right.fields())
            .all(|(left, right)| {
                left.name() == right.name() && left.data_type() == right.data_type()
            })
}

fn kinds(kind: ChangeKind, rows: usize) -> ArrayRef {
    Arc::new(StringArray::from_iter_values(repeat_n(kind.as_str(), rows)))
}

/// The row IDs of an image batch, its last column.
fn row_ids(batch: &RecordBatch) -> Result<&UInt64Array, Error> {
    batch
        .columns()
        .last()
        .and_then(|column| column.as_primitive_opt::<UInt64Type>())
        .ok_or_else(|| Error::Protocol(format!("an image read returned no `{ROW_ID_COLUMN}`")))
}

/// Each update's before image followed by its after image, matched by row ID, and the kind of
/// each output row.
fn pair_images(
    before: &RecordBatch,
    after: &RecordBatch,
) -> Result<(RecordBatch, ArrayRef), Error> {
    let after_rows = row_ids(after)?
        .values()
        .iter()
        .enumerate()
        .map(|(row, id)| (*id, row))
        .collect::<HashMap<_, _>>();
    let mut indices = Vec::with_capacity(before.num_rows() * 2);
    for (row, id) in row_ids(before)?.values().iter().enumerate() {
        let paired = after_rows.get(id).ok_or_else(|| {
            Error::Protocol(format!(
                "updated row {id} has no image in the later version"
            ))
        })?;
        indices.push((0, row));
        indices.push((1, *paired));
    }
    if indices.len() != after.num_rows() * 2 {
        return Err(Error::Protocol(
            "an updated row has no image in the earlier version".into(),
        ));
    }
    let columns = (0..before.num_columns())
        .map(|column| {
            interleave(
                &[
                    before.column(column).as_ref(),
                    after.column(column).as_ref(),
                ],
                &indices,
            )
            .map_err(Error::arrow)
        })
        .collect::<Result<Vec<_>, _>>()?;
    let paired = RecordBatch::try_new(after.schema(), columns).map_err(Error::arrow)?;
    let kinds = StringArray::from_iter_values(indices.iter().map(|(image, _)| match image {
        0 => ChangeKind::UpdateBefore.as_str(),
        _ => ChangeKind::UpdateAfter.as_str(),
    }));
    Ok((paired, Arc::new(kinds)))
}

/// `batch` without the rows whose row IDs fall in the sorted, disjoint `ranges`.
fn without_rows(batch: &RecordBatch, ranges: &[Range<u64>]) -> Result<RecordBatch, Error> {
    if ranges.is_empty() {
        return Ok(batch.clone());
    }
    let keep = row_ids(batch)?
        .values()
        .iter()
        .map(|id| Some(!in_ranges(*id, ranges)))
        .collect::<BooleanArray>();
    filter_record_batch(batch, &keep).map_err(Error::arrow)
}

/// Whether `id` falls in the sorted, disjoint `ranges`.
fn in_ranges(id: u64, ranges: &[Range<u64>]) -> bool {
    let after = ranges.partition_point(|range| range.start <= id);
    after > 0 && ranges[after - 1].contains(&id)
}

/// `ranges` sorted, with overlapping and adjacent ranges joined and empty ones dropped.
fn union_ranges(mut ranges: Vec<Range<u64>>) -> Vec<Range<u64>> {
    ranges.retain(|range| range.start < range.end);
    ranges.sort_by_key(|range| range.start);
    let mut joined: Vec<Range<u64>> = Vec::with_capacity(ranges.len());
    for range in ranges {
        match joined.last_mut() {
            Some(last) if range.start <= last.end => last.end = last.end.max(range.end),
            _ => joined.push(range),
        }
    }
    joined
}

/// The row IDs `ids` as sorted, disjoint ranges.
fn ranges_of(mut ids: Vec<u64>) -> Vec<Range<u64>> {
    ids.sort_unstable();
    ids.dedup();
    union_ranges(ids.into_iter().map(|id| id..id + 1).collect())
}

/// The sorted, disjoint `ranges` in groups of at most `ids` row IDs, splitting ranges at group
/// boundaries.
fn chunk_ranges(ranges: &[Range<u64>], ids: u64) -> Vec<Vec<Range<u64>>> {
    let mut chunks = Vec::new();
    let (mut chunk, mut held) = (Vec::new(), 0);
    for range in ranges {
        let mut start = range.start;
        while start < range.end {
            let end = range.end.min(start + (ids - held));
            chunk.push(start..end);
            held += end - start;
            start = end;
            if held == ids {
                chunks.push(take(&mut chunk));
                held = 0;
            }
        }
    }
    if !chunk.is_empty() {
        chunks.push(chunk);
    }
    chunks
}

#[cfg(test)]
#[expect(
    clippy::single_range_in_vec_init,
    reason = "Row-ID sets are lists of ranges; one range is a one-element list."
)]
mod tests {
    use arrow_array::Int64Array;
    use arrow_array::types::Int64Type;
    use serde_json::json;

    use super::*;

    #[test]
    fn row_id_ranges_join_split_and_test_membership() {
        assert_eq!(
            union_ranges(vec![5..7, 0..2, 2..3, 6..9, 4..4]),
            [0..3, 5..9]
        );
        assert_eq!(ranges_of(vec![4, 1, 2, 2, 7]), [1..3, 4..5, 7..8]);
        assert_eq!(
            chunk_ranges(&[0..3, 10..15], 4),
            [vec![0..3, 10..11], vec![11..15]]
        );
        assert!(chunk_ranges(&[], 4).is_empty());
        let ranges = [2..4, 8..9];
        assert_eq!(
            (0..10)
                .filter(|id| in_ranges(*id, &ranges))
                .collect::<Vec<_>>(),
            [2, 3, 8]
        );
    }

    #[test]
    fn a_plan_sorts_its_row_ids_by_what_happened_to_them() {
        let range = |start: u64, end: u64| json!({"start": start, "end": end});
        let removed = Uuid::now_v7();
        let wire: WireChangePlan = serde_json::from_value(json!({
            "fragments": [
                {"fragmentId": Uuid::now_v7(), "rowCount": 10, "change": {
                    "type": "retained",
                    "retained": {"patched": [range(3, 4)], "deleted": [range(0, 2)]}
                }},
                {"fragmentId": Uuid::now_v7(), "rowCount": 6, "change": {
                    "type": "added",
                    "added": {
                        "appended": [range(100, 103)],
                        "carried": [range(20, 23)],
                        "patched": [range(21, 22)]
                    }
                }}
            ],
            "removedFragmentIds": [removed],
            "retypedColumns": []
        }))
        .unwrap();
        let plan = ChangePlan::try_from(wire).unwrap();
        let rows = ChangedRows::from_plan(&plan);
        assert_eq!(rows.inserted, [100..103]);
        assert_eq!(rows.updated, [3..4, 21..22]);
        assert_eq!(rows.deleted, [0..2]);
        assert_eq!(rows.carried, [20..23]);
        assert!(rows.removed_fragments.contains(&removed));
    }

    fn image(values: Vec<i64>, ids: Vec<u64>) -> RecordBatch {
        let schema = Schema::new(vec![
            Field::new("value", DataType::Int64, true),
            row_id_field(),
        ]);
        RecordBatch::try_new(
            Arc::new(schema),
            vec![
                Arc::new(Int64Array::from(values)),
                Arc::new(UInt64Array::from(ids)),
            ],
        )
        .unwrap()
    }

    #[test]
    fn updates_pair_their_images_by_row_id() {
        let (paired, kinds) = pair_images(
            &image(vec![1, 2], vec![7, 9]),
            &image(vec![20, 10], vec![9, 7]),
        )
        .unwrap();
        assert_eq!(
            paired.column(0).as_primitive::<Int64Type>().values(),
            &[1, 10, 2, 20]
        );
        assert_eq!(row_ids(&paired).unwrap().values(), &[7, 7, 9, 9]);
        assert_eq!(
            kinds
                .as_string::<i32>()
                .iter()
                .flatten()
                .collect::<Vec<_>>(),
            [
                "update_before",
                "update_after",
                "update_before",
                "update_after"
            ]
        );
        assert!(pair_images(&image(vec![1], vec![7]), &image(vec![2], vec![8])).is_err());
        assert!(pair_images(&image(vec![1], vec![7]), &image(vec![2, 3], vec![7, 8])).is_err());
    }

    #[test]
    fn dropped_fragments_lose_their_carried_rows() {
        let kept = without_rows(&image(vec![1, 2, 3], vec![20, 21, 30]), &[20..22]).unwrap();
        assert_eq!(row_ids(&kept).unwrap().values(), &[30]);
    }

    #[test]
    fn change_reads_set_their_own_row_ids() {
        let keys = ChangeOptions {
            images: ChangeImages::Keys,
            ..Default::default()
        };
        assert!(check_change_options(&keys).is_ok());
        let limited = ChangeOptions {
            scan: ScanOptions {
                limit: Some(1),
                ..Default::default()
            },
            ..Default::default()
        };
        assert!(check_change_options(&limited).is_err());
        let renamed = ChangeOptions {
            scan: ScanOptions {
                row_id_column: Some("id".into()),
                ..Default::default()
            },
            ..Default::default()
        };
        assert!(check_change_options(&renamed).is_err());
    }
}