lance_nd/
write.rs

1//! Write streams of nd record batches to a Lance table.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use arrow::array::RecordBatch;
7use arrow::datatypes::Schema;
8use futures::{Stream, StreamExt};
9use lance::Dataset;
10use lance_file::version::LanceFileVersion;
11use nd_arrow_array::NdRecordBatch;
12
13use crate::cast::{Casting, cast_to};
14use crate::error::{Error, Result};
15use crate::format::add_format;
16use crate::options::{ColumnOptions, column_defaults, with_column_options};
17use crate::schema::conform_batch;
18use crate::stats::{add_stats, has_stats};
19use crate::store::StorageOptions;
20use crate::writer::{NdWriter, open_existing};
21
22/// Lance file version for new tables. Version 2.2 can add nested fields,
23/// which new attribute keys need.
24pub(crate) const FILE_VERSION: LanceFileVersion = LanceFileVersion::V2_2;
25
26/// How a write treats the table at its uri.
27#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
28pub enum WriteMode {
29    /// Create a new table. Fails when a table exists.
30    #[default]
31    Create,
32    /// Append to the table. [`write()`] creates it when it does not exist.
33    Append,
34    /// Replace the table, or create it when it does not exist.
35    Overwrite,
36}
37
38/// Options for [`write()`] and [`append()`].
39#[derive(Debug, Clone)]
40pub struct WriteOptions {
41    /// How [`write()`] treats an existing table. [`append()`] ignores it.
42    pub mode: WriteMode,
43    /// Most rows that a writer encodes as one group.
44    pub max_rows_per_group: usize,
45    /// Most nd values that a writer encodes as one group. A group ends after
46    /// the row that reaches it.
47    pub max_values_per_group: usize,
48    /// Encoding defaults for the columns of a new table. The table stores
49    /// them and gives them to each column that a later write adds.
50    /// For an existing table, they must be `None` or equal to the stored defaults.
51    pub column_defaults: Option<ColumnOptions>,
52    /// Encoding options per column name, over the defaults. They apply when
53    /// this write adds the column. Lance keeps them for later appends.
54    pub column_options: HashMap<String, ColumnOptions>,
55    /// Store the min and max of each variable per row, see `FORMAT.md` section 7.
56    /// Only a new table or an overwrite reads it. An append writes stats
57    /// when the table has them.
58    pub stats: bool,
59    /// How a write treats a variable or attribute type that differs from the
60    /// table, see [`Casting`].
61    pub casting: Casting,
62    /// [`crate::NdWriter`] buffers rows until they have about this many
63    /// bytes on disk, then writes them as one Lance fragment. The writer
64    /// learns the ratio of disk to memory bytes from the fragments it wrote.
65    /// Before the first fragment, it counts bytes in memory. The buffer
66    /// holds at most 8 times this size in memory. A row is never split.
67    /// [`write()`] ignores it.
68    pub target_fragment_bytes: usize,
69    /// The most fragments that [`crate::NdWriter`] writes at the same time.
70    pub write_parallelism: usize,
71    /// Lance storage options, for example the credentials of an object
72    /// store, see [`crate::StorageOptions`]. [`append()`] uses the options of its
73    /// table when this is empty.
74    pub storage_options: StorageOptions,
75}
76
77impl Default for WriteOptions {
78    fn default() -> Self {
79        Self {
80            mode: WriteMode::Create,
81            max_rows_per_group: 1024,
82            max_values_per_group: 1 << 26,
83            column_defaults: None,
84            column_options: HashMap::new(),
85            stats: false,
86            casting: Casting::No,
87            target_fragment_bytes: 128 << 20,
88            write_parallelism: 4,
89            storage_options: StorageOptions::new(),
90        }
91    }
92}
93
94/// Write a stream of rows to the table at `uri` in one commit. Each nd
95/// record batch is one row.
96///
97/// `WriteMode::Append` creates the table when it does not exist. A failed
98/// row commits nothing. Fails when the stream has no rows and the table
99/// does not exist.
100#[tracing::instrument(skip_all, fields(uri = %uri, mode = ?options.mode), err)]
101pub async fn write(
102    uri: &str,
103    rows: impl Stream<Item = Result<NdRecordBatch>>,
104    options: &WriteOptions,
105) -> Result<Dataset> {
106    let writer = match options.mode {
107        WriteMode::Append
108            if open_existing(uri, &options.storage_options)
109                .await?
110                .is_some() =>
111        {
112            NdWriter::append(uri, options).await?
113        }
114        WriteMode::Append => {
115            let options = WriteOptions {
116                mode: WriteMode::Create,
117                ..options.clone()
118            };
119            NdWriter::create(uri, &options).await?
120        }
121        WriteMode::Create | WriteMode::Overwrite => NdWriter::create(uri, options).await?,
122    };
123    write_all(writer, rows).await
124}
125
126/// Append a stream of rows to the table of `dataset` in one commit. Each
127/// nd record batch is one row. A failed row commits nothing.
128///
129/// The append starts from the latest version of the table. `dataset` then
130/// holds the new version.
131#[tracing::instrument(skip_all, fields(uri = %dataset.uri()), err)]
132pub async fn append(
133    dataset: &mut Dataset,
134    rows: impl Stream<Item = Result<NdRecordBatch>>,
135    options: &WriteOptions,
136) -> Result<()> {
137    let mut options = options.clone();
138    if options.storage_options.is_empty() {
139        options.storage_options = dataset
140            .initial_storage_options()
141            .cloned()
142            .unwrap_or_default();
143    }
144    let writer = NdWriter::append(dataset.uri(), &options).await?;
145    *dataset = write_all(writer, rows).await?;
146    Ok(())
147}
148
149async fn write_all(
150    mut writer: NdWriter,
151    rows: impl Stream<Item = Result<NdRecordBatch>>,
152) -> Result<Dataset> {
153    // Boxed: the writer futures are large, and would grow the future of
154    // each caller beyond the compiler's layout depth.
155    if let Err(error) = Box::pin(writer.write(rows)).await {
156        writer.abort();
157        return Err(error);
158    }
159    Box::pin(writer.finish()).await
160}
161
162/// The stored column defaults. Fails when `options` has other defaults.
163pub(crate) fn table_defaults(dataset: &Dataset, options: &WriteOptions) -> Result<ColumnOptions> {
164    let stored = column_defaults(dataset)?;
165    match &options.column_defaults {
166        Some(defaults) if *defaults != stored => Err(Error::InvalidArgument(
167            "the table has other column defaults, use set_column_defaults to change them"
168                .to_string(),
169        )),
170        _ => Ok(stored),
171    }
172}
173
174/// `true` when the table stores stats.
175pub(crate) fn table_has_stats(dataset: &Dataset) -> bool {
176    has_stats(&Schema::from(dataset.schema()))
177}
178
179/// `batch` of nd columns cast to the types of `table` as
180/// `options.casting` allows, with stats when `stats` is `true`, and with the
181/// encoding metadata of its columns.
182pub(crate) fn prepare_batch(
183    mut batch: RecordBatch,
184    defaults: &ColumnOptions,
185    options: &WriteOptions,
186    stats: bool,
187    table: Option<&Schema>,
188) -> Result<RecordBatch> {
189    // Before the stats, so that they get the table types too.
190    if let Some(table) = table {
191        batch = cast_to(batch, table, options.casting)?;
192    }
193    if stats {
194        batch = add_stats(&batch)?;
195    }
196    let mut metadata = batch.schema().metadata().clone();
197    metadata.extend(defaults.to_table_metadata());
198    add_format(&mut metadata);
199    let schema = with_column_options(&batch.schema(), defaults, &options.column_options)
200        .with_metadata(metadata);
201    conform_batch(&batch, &Arc::new(schema))
202}
203
204/// The next rows for one commit, or `None` at the end of the stream.
205pub(crate) async fn next_group(
206    rows: &mut (impl Stream<Item = Result<NdRecordBatch>> + Unpin),
207    options: &WriteOptions,
208) -> Result<Option<Vec<NdRecordBatch>>> {
209    let mut group = Vec::new();
210    let mut values = 0;
211    while group.len() < options.max_rows_per_group.max(1) && values < options.max_values_per_group {
212        let Some(row) = rows.next().await else {
213            break;
214        };
215        let row = row?;
216        values += row.columns().iter().map(|c| c.len()).sum::<usize>();
217        group.push(row);
218    }
219    Ok((!group.is_empty()).then_some(group))
220}
lance_nd/write.rs

lance_nd/
write.rs
