lance-nd 0.1.0

//! Write streams of nd record batches to a Lance table.

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

use arrow::array::RecordBatch;
use arrow::datatypes::Schema;
use futures::{Stream, StreamExt};
use lance::Dataset;
use lance_file::version::LanceFileVersion;
use nd_arrow_array::NdRecordBatch;

use crate::cast::{Casting, cast_to};
use crate::error::{Error, Result};
use crate::format::add_format;
use crate::options::{ColumnOptions, column_defaults, with_column_options};
use crate::schema::conform_batch;
use crate::stats::{add_stats, has_stats};
use crate::store::StorageOptions;
use crate::writer::{NdWriter, open_existing};

/// Lance file version for new tables. Version 2.2 can add nested fields,
/// which new attribute keys need.
pub(crate) const FILE_VERSION: LanceFileVersion = LanceFileVersion::V2_2;

/// How a write treats the table at its uri.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum WriteMode {
    /// Create a new table. Fails when a table exists.
    #[default]
    Create,
    /// Append to the table. [`write()`] creates it when it does not exist.
    Append,
    /// Replace the table, or create it when it does not exist.
    Overwrite,
}

/// Options for [`write()`] and [`append()`].
#[derive(Debug, Clone)]
pub struct WriteOptions {
    /// How [`write()`] treats an existing table. [`append()`] ignores it.
    pub mode: WriteMode,
    /// Most rows that a writer encodes as one group.
    pub max_rows_per_group: usize,
    /// Most nd values that a writer encodes as one group. A group ends after
    /// the row that reaches it.
    pub max_values_per_group: usize,
    /// Encoding defaults for the columns of a new table. The table stores
    /// them and gives them to each column that a later write adds.
    /// For an existing table, they must be `None` or equal to the stored defaults.
    pub column_defaults: Option<ColumnOptions>,
    /// Encoding options per column name, over the defaults. They apply when
    /// this write adds the column. Lance keeps them for later appends.
    pub column_options: HashMap<String, ColumnOptions>,
    /// Store the min and max of each variable per row, see `FORMAT.md` section 7.
    /// Only a new table or an overwrite reads it. An append writes stats
    /// when the table has them.
    pub stats: bool,
    /// How a write treats a variable or attribute type that differs from the
    /// table, see [`Casting`].
    pub casting: Casting,
    /// [`crate::NdWriter`] buffers rows until they have about this many
    /// bytes on disk, then writes them as one Lance fragment. The writer
    /// learns the ratio of disk to memory bytes from the fragments it wrote.
    /// Before the first fragment, it counts bytes in memory. The buffer
    /// holds at most 8 times this size in memory. A row is never split.
    /// [`write()`] ignores it.
    pub target_fragment_bytes: usize,
    /// The most fragments that [`crate::NdWriter`] writes at the same time.
    pub write_parallelism: usize,
    /// Lance storage options, for example the credentials of an object
    /// store, see [`crate::StorageOptions`]. [`append()`] uses the options of its
    /// table when this is empty.
    pub storage_options: StorageOptions,
}

impl Default for WriteOptions {
    fn default() -> Self {
        Self {
            mode: WriteMode::Create,
            max_rows_per_group: 1024,
            max_values_per_group: 1 << 26,
            column_defaults: None,
            column_options: HashMap::new(),
            stats: false,
            casting: Casting::No,
            target_fragment_bytes: 128 << 20,
            write_parallelism: 4,
            storage_options: StorageOptions::new(),
        }
    }
}

/// Write a stream of rows to the table at `uri` in one commit. Each nd
/// record batch is one row.
///
/// `WriteMode::Append` creates the table when it does not exist. A failed
/// row commits nothing. Fails when the stream has no rows and the table
/// does not exist.
#[tracing::instrument(skip_all, fields(uri = %uri, mode = ?options.mode), err)]
pub async fn write(
    uri: &str,
    rows: impl Stream<Item = Result<NdRecordBatch>>,
    options: &WriteOptions,
) -> Result<Dataset> {
    let writer = match options.mode {
        WriteMode::Append
            if open_existing(uri, &options.storage_options)
                .await?
                .is_some() =>
        {
            NdWriter::append(uri, options).await?
        }
        WriteMode::Append => {
            let options = WriteOptions {
                mode: WriteMode::Create,
                ..options.clone()
            };
            NdWriter::create(uri, &options).await?
        }
        WriteMode::Create | WriteMode::Overwrite => NdWriter::create(uri, options).await?,
    };
    write_all(writer, rows).await
}

/// Append a stream of rows to the table of `dataset` in one commit. Each
/// nd record batch is one row. A failed row commits nothing.
///
/// The append starts from the latest version of the table. `dataset` then
/// holds the new version.
#[tracing::instrument(skip_all, fields(uri = %dataset.uri()), err)]
pub async fn append(
    dataset: &mut Dataset,
    rows: impl Stream<Item = Result<NdRecordBatch>>,
    options: &WriteOptions,
) -> Result<()> {
    let mut options = options.clone();
    if options.storage_options.is_empty() {
        options.storage_options = dataset
            .initial_storage_options()
            .cloned()
            .unwrap_or_default();
    }
    let writer = NdWriter::append(dataset.uri(), &options).await?;
    *dataset = write_all(writer, rows).await?;
    Ok(())
}

async fn write_all(
    mut writer: NdWriter,
    rows: impl Stream<Item = Result<NdRecordBatch>>,
) -> Result<Dataset> {
    // Boxed: the writer futures are large, and would grow the future of
    // each caller beyond the compiler's layout depth.
    if let Err(error) = Box::pin(writer.write(rows)).await {
        writer.abort();
        return Err(error);
    }
    Box::pin(writer.finish()).await
}

/// The stored column defaults. Fails when `options` has other defaults.
pub(crate) fn table_defaults(dataset: &Dataset, options: &WriteOptions) -> Result<ColumnOptions> {
    let stored = column_defaults(dataset)?;
    match &options.column_defaults {
        Some(defaults) if *defaults != stored => Err(Error::InvalidArgument(
            "the table has other column defaults, use set_column_defaults to change them"
                .to_string(),
        )),
        _ => Ok(stored),
    }
}

/// `true` when the table stores stats.
pub(crate) fn table_has_stats(dataset: &Dataset) -> bool {
    has_stats(&Schema::from(dataset.schema()))
}

/// `batch` of nd columns cast to the types of `table` as
/// `options.casting` allows, with stats when `stats` is `true`, and with the
/// encoding metadata of its columns.
pub(crate) fn prepare_batch(
    mut batch: RecordBatch,
    defaults: &ColumnOptions,
    options: &WriteOptions,
    stats: bool,
    table: Option<&Schema>,
) -> Result<RecordBatch> {
    // Before the stats, so that they get the table types too.
    if let Some(table) = table {
        batch = cast_to(batch, table, options.casting)?;
    }
    if stats {
        batch = add_stats(&batch)?;
    }
    let mut metadata = batch.schema().metadata().clone();
    metadata.extend(defaults.to_table_metadata());
    add_format(&mut metadata);
    let schema = with_column_options(&batch.schema(), defaults, &options.column_options)
        .with_metadata(metadata);
    conform_batch(&batch, &Arc::new(schema))
}

/// The next rows for one commit, or `None` at the end of the stream.
pub(crate) async fn next_group(
    rows: &mut (impl Stream<Item = Result<NdRecordBatch>> + Unpin),
    options: &WriteOptions,
) -> Result<Option<Vec<NdRecordBatch>>> {
    let mut group = Vec::new();
    let mut values = 0;
    while group.len() < options.max_rows_per_group.max(1) && values < options.max_values_per_group {
        let Some(row) = rows.next().await else {
            break;
        };
        let row = row?;
        values += row.columns().iter().map(|c| c.len()).sum::<usize>();
        group.push(row);
    }
    Ok((!group.is_empty()).then_some(group))
}