spiral/

table_edit.rs

//! Row updates and schema changes, each written as a table's next version.

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

use arrow_array::RecordBatch;
use arrow_schema::{DataType, FieldRef, SchemaRef};
use futures::{StreamExt as _, TryStreamExt as _, future, stream};

use crate::registration::register;
use crate::scandal::api::{
    AddColumn, ArrowDataType, ArrowField, ColumnPath as WireColumnPath, MoveColumn, RenameColumn,
    SchemaChange as WireSchemaChange, WidenType,
};
use crate::scandal_backend::{RowEdit, input_error};
use crate::table_write::ROW_ID_COLUMN;
use crate::{BatchStream, Error, Expression, NamedExpression, ScanOptions, Table};

/// Rows updated in a table, and the version with their new values.
#[derive(Clone, Debug)]
pub struct UpdatedRows {
    /// The version with the new values; the original table when no row matched.
    pub table: Table,
    /// How many visible rows were updated.
    pub rows: u64,
}

/// A top-level column, or a field nested inside one, named by its segments.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct ColumnPath(Vec<String>);

impl ColumnPath {
    /// The field reached by `segments`: a top-level column, then a field of it, and so on.
    pub fn nested(segments: impl IntoIterator<Item = impl Into<String>>) -> Self {
        Self(segments.into_iter().map(Into::into).collect())
    }

    /// The path's segments, from the top-level column.
    pub fn segments(&self) -> &[String] {
        &self.0
    }

    fn into_wire(self) -> Result<Vec<String>, Error> {
        if self.0.is_empty() || self.0.iter().any(String::is_empty) {
            return Err(Error::invalid("a column path has nonempty segments"));
        }
        Ok(self.0)
    }
}

impl From<&str> for ColumnPath {
    fn from(column: &str) -> Self {
        Self(vec![column.to_owned()])
    }
}

impl From<String> for ColumnPath {
    fn from(column: String) -> Self {
        Self(vec![column])
    }
}

/// One change to a table's schema. Changes write metadata only, and older versions keep their
/// schema. Arbor's rules apply: an added column or field is nullable, `NOT NULL` may relax but
/// not tighten, numbers widen only without loss, and partition and sort source fields cannot be
/// dropped.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SchemaChange {
    /// Add a nullable `field` after the column or field `after`, or last when absent.
    AddColumn {
        field: FieldRef,
        after: Option<ColumnPath>,
    },
    /// Drop the column or field at `path`.
    DropColumn { path: ColumnPath },
    /// Rename the column or field at `path` to `name`, keeping its values.
    RenameColumn { path: ColumnPath, name: String },
    /// Let the column or field at `path` hold nulls.
    SetNullable { path: ColumnPath },
    /// Change the type of the column or field at `path` to `data_type`, converting without loss,
    /// such as `int32` to `int64`.
    WidenType {
        path: ColumnPath,
        data_type: DataType,
    },
    /// Move the column or field at `path` after `after`, or first when absent.
    MoveColumn {
        path: ColumnPath,
        after: Option<ColumnPath>,
    },
}

impl SchemaChange {
    /// The contract form of this change, which names paths and carries Arrow JSON.
    pub(crate) fn into_wire(self) -> Result<WireSchemaChange, Error> {
        Ok(match self {
            Self::AddColumn { field, after } => {
                if !field.is_nullable() {
                    return Err(Error::invalid(format!(
                        "added column `{}` must be nullable",
                        field.name()
                    )));
                }
                WireSchemaChange::AddColumn(AddColumn {
                    field: ArrowField(serde_json::to_value(field.as_ref()).map_err(Error::arrow)?),
                    after: after.map(ColumnPath::into_wire).transpose()?,
                })
            }
            Self::DropColumn { path } => WireSchemaChange::DropColumn(WireColumnPath {
                path: path.into_wire()?,
            }),
            Self::RenameColumn { path, name } => WireSchemaChange::RenameColumn(RenameColumn {
                path: path.into_wire()?,
                name,
            }),
            Self::SetNullable { path } => WireSchemaChange::SetNullable(WireColumnPath {
                path: path.into_wire()?,
            }),
            Self::WidenType { path, data_type } => WireSchemaChange::WidenType(WidenType {
                path: path.into_wire()?,
                data_type: ArrowDataType(serde_json::to_value(&data_type).map_err(Error::arrow)?),
            }),
            Self::MoveColumn { path, after } => WireSchemaChange::MoveColumn(MoveColumn {
                path: path.into_wire()?,
                after: after.map(ColumnPath::into_wire).transpose()?,
            }),
        })
    }
}

impl Table {
    /// This version's Arrow schema. It plans a scan but reads no rows.
    pub async fn schema(&self) -> Result<SchemaRef, Error> {
        Ok(self.scan(ScanOptions::default()).await?.schema())
    }

    /// Write the next version of this table with `changes` applied to its schema, in order,
    /// without writing data files. Every change is checked against this version before
    /// anything is written; with no changes the result is this table.
    pub async fn alter_schema(
        &self,
        changes: impl IntoIterator<Item = SchemaChange>,
    ) -> Result<Table, Error> {
        let changes = changes
            .into_iter()
            .map(SchemaChange::into_wire)
            .collect::<Result<Vec<_>, _>>()?;
        if changes.is_empty() {
            return Ok(self.clone());
        }
        let write = self
            .connection
            .tables()?
            .begin_edit(self.identity(), changes)
            .await?;
        let result = register(&self.connection.arbor, write.prepare().await?).await?;
        self.successor(result)
    }

    /// Write the next version of this table with the values of `set` in the visible rows that
    /// match `filter`, or in every row without a filter, recorded as patches beside the data.
    ///
    /// Each assignment names an existing column and must evaluate to its stored type; nothing
    /// is cast. Rows keep their row IDs. When no row matches, nothing is written and the result
    /// holds this table.
    pub async fn patch_rows(
        &self,
        filter: Option<Expression>,
        set: Vec<NamedExpression>,
    ) -> Result<UpdatedRows, Error> {
        self.update_rows(filter, set).await
    }

    /// Scan the matching rows' row IDs and new values, and write them as patches.
    async fn update_rows(
        &self,
        filter: Option<Expression>,
        set: Vec<NamedExpression>,
    ) -> Result<UpdatedRows, Error> {
        if set.is_empty() {
            return Err(Error::invalid("an update assigns at least one column"));
        }
        let scan = self
            .scan(ScanOptions {
                projection: Some(set),
                filter,
                row_id_column: Some(ROW_ID_COLUMN.to_owned()),
                ..Default::default()
            })
            .await?;
        let schema = scan.schema();
        let mut batches = scan.into_stream();
        let Some(first) = first_rows(&mut batches).await? else {
            return Ok(UpdatedRows {
                table: self.clone(),
                rows: 0,
            });
        };
        let rows = Arc::new(AtomicU64::new(first.num_rows() as u64));
        let counted = Arc::clone(&rows);
        let input = stream::once(future::ready(Ok(first)))
            .chain(batches.inspect_ok(move |batch| {
                counted.fetch_add(batch.num_rows() as u64, Ordering::Relaxed);
            }))
            .map_err(input_error)
            .boxed();
        let mut write = self.writer().await?;
        write.edit_rows(RowEdit::Patch, schema, input).await?;
        Ok(UpdatedRows {
            table: write.finish().await?,
            rows: rows.load(Ordering::Relaxed),
        })
    }
}

/// The first batch of `batches` that holds rows, or `None` when none does.
pub(crate) async fn first_rows(batches: &mut BatchStream) -> Result<Option<RecordBatch>, Error> {
    while let Some(batch) = batches.next_batch().await {
        let batch = batch?;
        if batch.num_rows() > 0 {
            return Ok(Some(batch));
        }
    }
    Ok(None)
}

#[cfg(test)]
mod tests {
    use arrow_schema::Field;
    use serde_json::json;

    use super::*;

    #[test]
    fn schema_changes_name_paths_and_carry_arrow_json() {
        let added = SchemaChange::AddColumn {
            field: Arc::new(Field::new("score", DataType::Float32, true)),
            after: Some("label".into()),
        }
        .into_wire()
        .unwrap();
        let WireSchemaChange::AddColumn(added) = added else {
            panic!("an added column stays an added column");
        };
        assert_eq!(added.after, Some(vec!["label".to_owned()]));
        assert_eq!(added.field.0["name"], json!("score"));
        assert_eq!(added.field.0["nullable"], json!(true));

        let widened = SchemaChange::WidenType {
            path: ColumnPath::nested(["pose", "x"]),
            data_type: DataType::Float64,
        }
        .into_wire()
        .unwrap();
        assert_eq!(
            widened,
            WireSchemaChange::WidenType(WidenType {
                path: vec!["pose".to_owned(), "x".to_owned()],
                data_type: ArrowDataType(json!("Float64")),
            })
        );

        let required = SchemaChange::AddColumn {
            field: Arc::new(Field::new("score", DataType::Float32, false)),
            after: None,
        };
        assert!(required.into_wire().is_err());
        let unnamed = SchemaChange::DropColumn {
            path: ColumnPath::nested(Vec::<String>::new()),
        };
        assert!(unnamed.into_wire().is_err());
    }
}