spiral-sdk 0.1.10

//! The Arrow framing of the scandal API, shared by every transport. The scandal contract
//! (`spiraldb/services/scandal.yaml`) defines the routes and JSON requests; rows travel as
//! Arrow IPC streams, and a `ReadScan` batch carries its progress in two trailing columns.

use std::sync::Arc;

use arrow_array::types::Int32Type;
use arrow_array::{
    Array as _, DictionaryArray, Int32Array, RecordBatch, RecordBatchOptions, StringArray,
    UInt64Array,
};
use arrow_ipc::writer::StreamWriter;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use bytes::Bytes;
use futures::stream::{self, BoxStream};
use futures::{StreamExt as _, TryStreamExt as _};

use super::error::ScandalError;
use super::types::ScanMetricsSnapshot;
use super::{RecordBatchStream, SplitBatch};

/// Liveness. Not `/healthz`: Cloud Run reserves paths ending in `z` and never forwards them.
pub const HEALTH: &str = "/health";

/// The audience of a scandal capability: the RFC 8707 resource for which a client exchanges
/// its access token, and whose result every request except [`HEALTH`] carries as its bearer
/// token.
pub const CAPABILITY_AUDIENCE: &str = "https://scandal.spiraldb.dev";

/// The Arrow IPC streaming format's media type.
pub const ARROW_STREAM: &str = "application/vnd.apache.arrow.stream";
/// The column after each `ReadScan` batch's rows holding the next stored row it reached.
pub const NEXT_ROW_COLUMN: &str = "scandal.next_row";
/// The column after [`NEXT_ROW_COLUMN`] holding a split read's counters, on its last batch.
pub const METRICS_COLUMN: &str = "scandal.metrics";

/// Encode batches as an Arrow IPC stream. Only a stream that ends without error receives
/// the end marker, so a reader detects truncation.
pub fn encode_ipc(
    schema: SchemaRef,
    batches: RecordBatchStream,
) -> BoxStream<'static, Result<Bytes, ScandalError>> {
    let writer = match StreamWriter::try_new(Vec::new(), &schema) {
        Ok(writer) => writer,
        Err(error) => return stream::once(async move { Err(internal(error)) }).boxed(),
    };
    stream::try_unfold(Some((batches, writer, true)), |state| async move {
        let Some((mut batches, mut writer, first)) = state else {
            return Ok(None);
        };
        if first {
            let header = Bytes::from(std::mem::take(writer.get_mut()));
            return Ok(Some((header, Some((batches, writer, false)))));
        }
        match batches.try_next().await? {
            Some(batch) => {
                writer.write(&batch).map_err(internal)?;
                let bytes = Bytes::from(std::mem::take(writer.get_mut()));
                Ok(Some((bytes, Some((batches, writer, false)))))
            }
            None => {
                writer.finish().map_err(internal)?;
                Ok(Some((Bytes::from(std::mem::take(writer.get_mut())), None)))
            }
        }
    })
    .boxed()
}

/// The schema of a `ReadScan` response: the split's rows, then [`NEXT_ROW_COLUMN`] and
/// [`METRICS_COLUMN`].
///
/// Each trailing column is a dictionary of at most one value that every row's key names, so
/// the value travels with its batch even when the batch holds no rows.
pub fn read_scan_schema(rows: &Schema) -> SchemaRef {
    let dictionary = |values| DataType::Dictionary(Box::new(DataType::Int32), Box::new(values));
    let mut fields = rows.fields().iter().cloned().collect::<Vec<_>>();
    fields.push(Arc::new(Field::new(
        NEXT_ROW_COLUMN,
        dictionary(DataType::UInt64),
        false,
    )));
    fields.push(Arc::new(Field::new(
        METRICS_COLUMN,
        dictionary(DataType::Utf8),
        true,
    )));
    Arc::new(Schema::new(fields))
}

/// Append a split batch's next stored row and counters as single-value dictionaries.
pub fn attach_progress(batch: SplitBatch, schema: &SchemaRef) -> Result<RecordBatch, ScandalError> {
    let rows = batch.batch.num_rows();
    let next_row = DictionaryArray::<Int32Type>::try_new(
        Int32Array::from(vec![0; rows]),
        Arc::new(UInt64Array::from(vec![batch.next_row])),
    )
    .map_err(internal)?;
    let metrics = match batch.metrics {
        Some(metrics) => DictionaryArray::<Int32Type>::try_new(
            Int32Array::from(vec![0; rows]),
            Arc::new(StringArray::from(vec![
                serde_json::to_string(&metrics).map_err(internal)?,
            ])),
        ),
        None => DictionaryArray::<Int32Type>::try_new(
            Int32Array::from(vec![None; rows]),
            Arc::new(StringArray::from(Vec::<String>::new())),
        ),
    }
    .map_err(internal)?;
    let mut columns = batch.batch.columns().to_vec();
    columns.push(Arc::new(next_row));
    columns.push(Arc::new(metrics));
    RecordBatch::try_new_with_options(
        Arc::clone(schema),
        columns,
        &RecordBatchOptions::new().with_row_count(Some(rows)),
    )
    .map_err(internal)
}

/// Separate a `ReadScan` batch's rows from its next stored row and counters.
pub fn detach_progress(batch: RecordBatch) -> Result<SplitBatch, ScandalError> {
    let rows = trailing_columns(batch.schema_ref())?;
    let value = |index: usize| {
        batch
            .column(index)
            .as_any()
            .downcast_ref::<DictionaryArray<Int32Type>>()
            .map(|dictionary| Arc::clone(dictionary.values()))
    };
    let next_row = value(rows)
        .and_then(|values| values.as_any().downcast_ref::<UInt64Array>().cloned())
        .filter(|values| values.len() == 1 && values.null_count() == 0)
        .map(|values| values.value(0))
        .ok_or_else(|| internal("scandal batch's next stored row is malformed"))?;
    let metrics = value(rows + 1)
        .and_then(|values| values.as_any().downcast_ref::<StringArray>().cloned())
        .filter(|values| values.len() <= 1 && values.null_count() == 0)
        .ok_or_else(|| internal("scandal batch's counters are malformed"))?;
    let metrics = metrics
        .iter()
        .flatten()
        .next()
        .map(serde_json::from_str::<ScanMetricsSnapshot>)
        .transpose()
        .map_err(internal)?;
    let data = batch
        .project(&(0..rows).collect::<Vec<_>>())
        .map_err(internal)?;
    Ok(SplitBatch {
        batch: data,
        next_row,
        metrics,
    })
}

/// The row schema of a `ReadScan` response, without its trailing progress columns.
pub fn read_scan_rows(schema: &Schema) -> Result<SchemaRef, ScandalError> {
    let rows = trailing_columns(schema)?;
    Ok(Arc::new(Schema::new(schema.fields()[..rows].to_vec())))
}

/// The number of row columns before [`NEXT_ROW_COLUMN`] and [`METRICS_COLUMN`].
fn trailing_columns(schema: &Schema) -> Result<usize, ScandalError> {
    schema
        .fields()
        .len()
        .checked_sub(2)
        .filter(|rows| {
            schema.field(*rows).name() == NEXT_ROW_COLUMN
                && schema.field(rows + 1).name() == METRICS_COLUMN
        })
        .ok_or_else(|| internal("scandal stream lacks its progress columns"))
}

fn internal(message: impl ToString) -> ScandalError {
    ScandalError::Internal {
        message: message.to_string(),
    }
}

#[cfg(test)]
mod tests {
    use arrow_array::Int64Array;
    use spiral_contract::safelong::SafeUnsignedLong;

    use super::*;

    #[test]
    fn progress_round_trips_beside_split_rows() {
        let rows = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            false,
        )]));
        let schema = read_scan_schema(&rows);
        let batch =
            RecordBatch::try_new(rows, vec![Arc::new(Int64Array::from(vec![4, 5]))]).unwrap();
        let metrics = ScanMetricsSnapshot {
            output_rows: SafeUnsignedLong(2),
            ..Default::default()
        };
        for (batch, metrics) in [(batch.clone(), None), (batch.slice(0, 0), Some(metrics))] {
            let attached = attach_progress(
                SplitBatch {
                    batch: batch.clone(),
                    next_row: u64::MAX,
                    metrics: metrics.clone(),
                },
                &schema,
            )
            .unwrap();
            let detached = detach_progress(attached).unwrap();
            assert_eq!(
                (detached.batch, detached.next_row, detached.metrics),
                (batch, u64::MAX, metrics)
            );
        }
        assert_eq!(
            read_scan_rows(&schema).unwrap().fields().len(),
            1,
            "progress columns are not row columns"
        );
    }
}