spiral-sdk 0.1.11

//! Values for new fields over the rows of one exact version of a source table.

use std::collections::HashSet;
use std::sync::Arc;

use arrow_array::RecordBatch;
use arrow_schema::{Field, Schema, SchemaRef};
use futures::future::BoxFuture;
use futures::{StreamExt as _, TryStreamExt as _, stream};
use uuid::Uuid;

use crate::registration::register;
use crate::scandal::RecordBatchStream;
use crate::scandal::api::AttachedColumns;
use crate::scandal_backend::{AttemptWrite, RowEdit, input_error, reader_stream};
use crate::table_write::{ROW_ID_COLUMN, check_edit_schema};
use crate::{
    Error, IntoArrowStream, NamedExpression, Scan, ScanOptions, ScanSplit, SchemaChange, Table,
    WorkItem,
};

/// Fragments whose values [`Table::add_columns`] writes at once.
const FRAGMENTS_IN_FLIGHT: usize = 4;

/// Values for new nullable fields over the rows of one exact version of a table, from
/// [`Table::column_writer`], which join its schema in its next version.
///
/// Values cover every row of a fragment with [`Self::attach_columns`] or chosen rows with
/// [`Self::patch_rows`]; rows without values read `NULL`. Either way the values go into new files
/// beside the fragment's existing ones, which are never rewritten. [`ColumnWrite::finish`]
/// returns the table's next version, which is published only once a [`Commit`](crate::Commit)
/// stages it and publishes.
pub struct ColumnWrite {
    source: Table,
    fields: SchemaRef,
    write: AttemptWrite,
}

impl ColumnWrite {
    /// The exact source version whose rows the values cover.
    pub fn source(&self) -> &Table {
        &self.source
    }

    /// The fields this write gives values for.
    pub fn fields(&self) -> SchemaRef {
        Arc::clone(&self.fields)
    }

    /// Attach values of every field to fragment `fragment_id` of the source: one stream of
    /// batches with a value for every row the fragment keeps, in stored order, which is the
    /// order a scan of the fragment's splits returns. The values are written as they arrive,
    /// into new files attached to the fragment in the next version; the fragment's existing
    /// files are left as they are. Rows the source deletes read `NULL`. A write attaches each
    /// fragment once, and not to rows it also patches: [`ColumnWrite::finish`] refuses both.
    pub async fn attach_columns(
        &mut self,
        fragment_id: Uuid,
        values: impl IntoArrowStream,
    ) -> Result<(), Error> {
        let reader = values.into_arrow_stream()?;
        let schema = reader.schema();
        let fragment = self
            .fragment_values(fragment_id, schema, reader_stream(reader))
            .await?
            .await?;
        self.write.add_receipts([fragment.receipt]);
        Ok(())
    }

    /// Write values for some rows, recorded beside their fragments: each row names its target in
    /// a [`ROW_ID_COLUMN`] column, and its other columns are some of [`Self::fields`].
    pub async fn patch_rows(&mut self, rows: impl IntoArrowStream) -> Result<(), Error> {
        let reader = rows.into_arrow_stream()?;
        let schema = reader.schema();
        check_edit_schema(RowEdit::Patch, &schema)?;
        check_patched_fields(&self.fields, &schema)?;
        let receipts = self
            .write
            .edit_rows(RowEdit::Patch, ROW_ID_COLUMN, schema, reader_stream(reader))
            .await?;
        self.write.add_receipts(receipts);
        Ok(())
    }

    /// Attach values to every fragment, as [`Self::attach_columns`] does, with the values `scan`,
    /// a scan of the source whose outputs are this write's fields, computes over it. Several
    /// fragments are written at once.
    async fn attach_scanned_columns(&mut self, scan: &Scan) -> Result<(), Error> {
        let fields = Arc::clone(&self.fields);
        let work = fragment_work(scan.splits())?;
        let written = stream::iter(work)
            .map(|(fragment_id, work)| {
                let (this, fields) = (&*self, &fields);
                async move {
                    let batches = scan.read_work(work).await?;
                    let schema = Arc::clone(fields);
                    let input = batches
                        .map(move |read| {
                            let columns = read?.batch.columns().to_vec();
                            RecordBatch::try_new(Arc::clone(&schema), columns).map_err(Error::arrow)
                        })
                        .map_err(input_error)
                        .boxed();
                    this.fragment_values(fragment_id, Arc::clone(fields), input)
                        .await?
                        .await
                }
            })
            .buffer_unordered(FRAGMENTS_IN_FLIGHT)
            .try_collect::<Vec<_>>()
            .await?;
        self.write
            .add_receipts(written.into_iter().map(|fragment| fragment.receipt));
        Ok(())
    }

    /// Start writing one fragment's values beside other fragments. Keep the receipt the request
    /// returns.
    async fn fragment_values(
        &self,
        fragment_id: Uuid,
        schema: SchemaRef,
        input: RecordBatchStream,
    ) -> Result<BoxFuture<'static, Result<AttachedColumns, Error>>, Error> {
        check_fragment_fields(&self.fields, &schema)?;
        Ok(self.write.attach_columns(fragment_id, schema, input))
    }

    /// Prepare the table's next version, with the new columns and every value written, and
    /// register it.
    pub async fn finish(self) -> Result<Table, Error> {
        let result = register(&self.source.connection.arbor, self.write.prepare().await?).await?;
        self.source.successor(result)
    }
}

impl Table {
    /// Begin writing values for new nullable `fields` of this table, which join its schema in
    /// its next version. Rows the write gives no values read `NULL`, so values may follow over
    /// several versions.
    pub async fn column_writer(&self, fields: SchemaRef) -> Result<ColumnWrite, Error> {
        if fields.fields().is_empty() {
            return Err(Error::invalid("a column write adds at least one field"));
        }
        let changes = fields
            .fields()
            .iter()
            .map(|field| {
                SchemaChange::AddColumn {
                    field: Arc::clone(field),
                    after: None,
                }
                .into_wire()
            })
            .collect::<Result<Vec<_>, _>>()?;
        let write = self
            .connection
            .tables()?
            .begin_edit(self.identity(), changes)
            .await?;
        Ok(ColumnWrite {
            source: self.clone(),
            fields,
            write,
        })
    }

    /// Write the next version of this table with a new nullable column for each of `columns`,
    /// filled for every row with the expression's value over that row.
    ///
    /// Each fragment is read in stored order with the expressions as its projection and written
    /// whole, several fragments at once.
    pub async fn add_columns(&self, columns: Vec<NamedExpression>) -> Result<Table, Error> {
        if columns.is_empty() {
            return Err(Error::invalid(
                "adding columns needs at least one expression",
            ));
        }
        let scan = self
            .scan(ScanOptions {
                projection: Some(columns),
                ..Default::default()
            })
            .await?;
        let mut write = self.column_writer(nullable(&scan.schema())).await?;
        write.attach_scanned_columns(&scan).await?;
        write.finish().await
    }
}

/// `schema` with every field nullable, as new columns are.
pub(crate) fn nullable(schema: &Schema) -> SchemaRef {
    Arc::new(Schema::new(
        schema
            .fields()
            .iter()
            .map(|field| Field::clone(field).with_nullable(true))
            .collect::<Vec<_>>(),
    ))
}

/// The work that reads each fragment of a scan whole: the fragment's ID and its splits, in
/// stored order.
fn fragment_work(splits: &[ScanSplit]) -> Result<Vec<(Uuid, Vec<WorkItem>)>, Error> {
    let mut fragments: Vec<(Uuid, Vec<WorkItem>)> = Vec::new();
    for (split, planned) in 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")
        })?;
        let item = WorkItem {
            split,
            rows: planned.rows.clone(),
        };
        match fragments.last_mut() {
            Some((last, items)) if *last == fragment_id => items.push(item),
            _ => fragments.push((fragment_id, vec![item])),
        }
    }
    let mut seen = HashSet::new();
    if let Some((repeated, _)) = fragments.iter().find(|(id, _)| !seen.insert(*id)) {
        return Err(Error::Protocol(format!(
            "the splits of fragment {repeated} are not adjacent"
        )));
    }
    Ok(fragments)
}

/// Check that a fragment's values hold every one of `fields`, in order, with its type.
fn check_fragment_fields(fields: &Schema, values: &Schema) -> Result<(), Error> {
    let matches = fields.fields().len() == values.fields().len()
        && fields
            .fields()
            .iter()
            .zip(values.fields())
            .all(|(field, value)| {
                field.name() == value.name() && field.data_type() == value.data_type()
            });
    if !matches {
        return Err(Error::invalid(format!(
            "a fragment's values hold exactly the written fields {:?}",
            field_names(fields)
        )));
    }
    Ok(())
}

/// Check that patched rows hold a row ID column and some of `fields`, each with its type.
fn check_patched_fields(fields: &Schema, rows: &Schema) -> Result<(), Error> {
    for value in rows.fields() {
        if value.name() == ROW_ID_COLUMN {
            continue;
        }
        let field = fields.field_with_name(value.name()).map_err(|_| {
            Error::invalid(format!(
                "patched column `{}` is not one of the written fields {:?}",
                value.name(),
                field_names(fields)
            ))
        })?;
        if field.data_type() != value.data_type() {
            return Err(Error::invalid(format!(
                "patched column `{}` holds {}, not {}",
                value.name(),
                value.data_type(),
                field.data_type()
            )));
        }
    }
    Ok(())
}

fn field_names(schema: &Schema) -> Vec<&str> {
    schema
        .fields()
        .iter()
        .map(|field| field.name().as_str())
        .collect()
}

#[cfg(test)]
mod tests {
    use std::ops::Range;

    use arrow_schema::DataType;

    use super::*;

    fn split(fragment_id: Option<Uuid>, rows: Range<u64>) -> ScanSplit {
        ScanSplit {
            fragment: 0,
            rows,
            estimated_rows: None,
            estimated_bytes: None,
            fragment_id,
        }
    }

    #[test]
    fn each_fragment_is_read_whole_from_its_adjacent_splits() {
        let (first, second) = (Uuid::now_v7(), Uuid::now_v7());
        let work = fragment_work(&[
            split(Some(first), 0..10),
            split(Some(first), 10..20),
            split(Some(second), 0..5),
        ])
        .unwrap();
        assert_eq!(work.len(), 2);
        assert_eq!(work[0].0, first);
        assert_eq!(
            work[0].1.iter().map(|item| item.split).collect::<Vec<_>>(),
            [0, 1]
        );
        assert_eq!(work[1].1[0].rows, 0..5);
        assert!(fragment_work(&[split(None, 0..1)]).is_err());
        assert!(
            fragment_work(&[
                split(Some(first), 0..1),
                split(Some(second), 0..1),
                split(Some(first), 1..2),
            ])
            .is_err()
        );
    }

    #[test]
    fn values_hold_the_written_fields_with_their_types() {
        let fields = Schema::new(vec![
            Field::new("score", DataType::Float32, true),
            Field::new("label", DataType::Utf8, true),
        ]);
        let required = Schema::new(vec![
            Field::new("score", DataType::Float32, false),
            Field::new("label", DataType::Utf8, false),
        ]);
        assert!(check_fragment_fields(&fields, &required).is_ok());
        let partial = Schema::new(vec![Field::new("score", DataType::Float32, true)]);
        assert!(check_fragment_fields(&fields, &partial).is_err());
        let patched = Schema::new(vec![
            Field::new(ROW_ID_COLUMN, DataType::UInt64, false),
            Field::new("label", DataType::Utf8, true),
        ]);
        assert!(check_patched_fields(&fields, &patched).is_ok());
        let retyped = Schema::new(vec![
            Field::new(ROW_ID_COLUMN, DataType::UInt64, false),
            Field::new("label", DataType::LargeUtf8, true),
        ]);
        assert!(check_patched_fields(&fields, &retyped).is_err());
    }
}