spiral/
table_write.rs

1//! A table write that uploads large blobs before the rows that reference them, and edits its
2//! base version's rows by row ID.
3
4use std::sync::Arc;
5
6use arrow_schema::{DataType, Schema, SchemaRef};
7use futures::future::BoxFuture;
8use uuid::Uuid;
9
10use crate::client::Connection;
11use crate::registration::register;
12use crate::scandal::RecordBatchStream;
13use crate::scandal::types::FragmentReceipt;
14use crate::scandal_backend::{AttemptWrite, RowEdit, reader_stream};
15use crate::{Blob, BlobPack, BlobStream, Error, IntoArrowStream, Project, Table, TableIdentity};
16
17/// The column that names rows by row ID: in the rows that [`TableWrite::patch_rows`] and
18/// [`TableWrite::delete_rows`] take, and in the scans that the SDK's own edits and change reads
19/// open. It is a non-null `uint64` column.
20pub const ROW_ID_COLUMN: &str = "s:row_id";
21
22/// One write of a new table's first version or of a table's next version, under its own write
23/// attempt. Large payloads upload into the attempt first, as blobs; rows appended afterwards
24/// carry them, encoded with [`Blob::array`]. A write of a table's next version can also patch
25/// and delete the base version's rows by row ID. [`TableWrite::finish`] returns the
26/// written version, one version for every part, which is published only once a
27/// [`Commit`](crate::Commit) stages it and publishes.
28pub struct TableWrite {
29    /// The table this write continues, or `None` for a new table.
30    base: Option<Table>,
31    connection: Arc<Connection>,
32    write: AttemptWrite,
33}
34
35impl Project {
36    /// Begin the first version of a new table with rows of `schema`.
37    pub async fn table_writer(&self, schema: SchemaRef) -> Result<TableWrite, Error> {
38        let write = self
39            .client
40            .connection
41            .tables()?
42            .begin_create(self.id.as_str(), &schema)
43            .await?;
44        Ok(TableWrite {
45            base: None,
46            connection: Arc::clone(&self.client.connection),
47            write,
48        })
49    }
50}
51
52impl Table {
53    /// Begin the next version of this table: its rows, edited by the patches and deletions the
54    /// write takes, followed by every appended row.
55    pub async fn writer(&self) -> Result<TableWrite, Error> {
56        let write = self
57            .connection
58            .tables()?
59            .begin_write(self.identity(), false)
60            .await?;
61        Ok(TableWrite {
62            base: Some(self.clone()),
63            connection: Arc::clone(&self.connection),
64            write,
65        })
66    }
67}
68
69impl TableWrite {
70    /// The write attempt whose directory holds this write's files and uploads.
71    pub fn attempt(&self) -> Uuid {
72        self.write.attempt()
73    }
74
75    /// Upload `input` as one blob of this write, however large, in parts of 64 MiB sent several
76    /// at once. Each part is retried on its own, and a failed upload is aborted. The blob's
77    /// digest is the BLAKE3 digest of the whole payload.
78    pub async fn upload_blob(&self, media_type: &str, input: BlobStream) -> Result<Blob, Error> {
79        self.write.upload_blob(media_type, input).await
80    }
81
82    /// Begin a pack of payloads to store as blobs of this write, packed together as the inline
83    /// payloads of appended rows are.
84    pub fn blob_pack(&self) -> BlobPack<'_> {
85        BlobPack::new(&self.write)
86    }
87
88    /// Append rows, which may carry blobs this write uploaded.
89    pub async fn append(&mut self, input: impl IntoArrowStream) -> Result<(), Error> {
90        self.write.append(input.into_write_input()?).await
91    }
92
93    /// Record new values for existing rows beside the data: each row names its target in a
94    /// [`ROW_ID_COLUMN`] column, and its other columns are columns of the table, with their
95    /// stored types, whose values it replaces. A row ID appears at most once per write.
96    pub async fn patch_rows(&mut self, rows: impl IntoArrowStream) -> Result<(), Error> {
97        self.edit_input(RowEdit::Patch, rows).await
98    }
99
100    /// Delete existing rows, named by their row IDs in a [`ROW_ID_COLUMN`] column, which is the
101    /// input's only column. Each row must be visible in the base version.
102    pub async fn delete_rows(&mut self, ids: impl IntoArrowStream) -> Result<(), Error> {
103        self.edit_input(RowEdit::Delete, ids).await
104    }
105
106    /// Prepare the version from every part this write took and register it.
107    pub async fn finish(self) -> Result<Table, Error> {
108        let result = register(&self.connection.arbor, self.write.prepare().await?).await?;
109        match self.base {
110            Some(base) => base.successor(result),
111            None => Ok(Table::new(
112                self.connection,
113                TableIdentity {
114                    resource_id: result.resource_id,
115                    version_id: result.version_id,
116                },
117            )),
118        }
119    }
120
121    async fn edit_input(
122        &mut self,
123        edit: RowEdit,
124        input: impl IntoArrowStream,
125    ) -> Result<(), Error> {
126        let reader = input.into_arrow_stream()?;
127        self.edit_rows(edit, reader.schema(), reader_stream(reader))
128            .await
129    }
130
131    /// Write `input` as one part of `edit` and keep its receipts.
132    pub(crate) async fn edit_rows(
133        &mut self,
134        edit: RowEdit,
135        schema: SchemaRef,
136        input: RecordBatchStream,
137    ) -> Result<(), Error> {
138        let receipts = self.start_edit(edit, schema, input)?.await?;
139        self.write.add_receipts(receipts);
140        Ok(())
141    }
142
143    /// Start writing `input` as one part of `edit`, beside other parts. Keep the receipts the
144    /// request returns with [`Self::add_receipts`].
145    pub(crate) fn start_edit(
146        &mut self,
147        edit: RowEdit,
148        schema: SchemaRef,
149        input: RecordBatchStream,
150    ) -> Result<BoxFuture<'static, Result<Vec<FragmentReceipt>, Error>>, Error> {
151        if self.base.is_none() {
152            return Err(Error::invalid("a new table has no rows to patch or delete"));
153        }
154        check_edit_schema(edit, &schema)?;
155        self.write.edit_base()?;
156        Ok(self.write.edit_rows(edit, ROW_ID_COLUMN, schema, input))
157    }
158
159    /// Keep the receipts of a part started with [`Self::start_edit`].
160    pub(crate) fn add_receipts(&mut self, receipts: Vec<FragmentReceipt>) {
161        self.write.add_receipts(receipts);
162    }
163}
164
165/// Check that rows of `edit` name their targets in a `uint64` [`ROW_ID_COLUMN`] and carry the
166/// columns the edit needs: values for patches, nothing else for deletions.
167pub(crate) fn check_edit_schema(edit: RowEdit, schema: &Schema) -> Result<(), Error> {
168    let row_ids = schema.field_with_name(ROW_ID_COLUMN).map_err(|_| {
169        Error::invalid(format!(
170            "rows that address existing rows need a `{ROW_ID_COLUMN}` column"
171        ))
172    })?;
173    if row_ids.data_type() != &DataType::UInt64 {
174        return Err(Error::invalid(format!(
175            "`{ROW_ID_COLUMN}` holds uint64 row IDs, not {}",
176            row_ids.data_type()
177        )));
178    }
179    let values = schema.fields().len() - 1;
180    match edit {
181        RowEdit::Delete if values > 0 => Err(Error::invalid(format!(
182            "deletions hold only the `{ROW_ID_COLUMN}` column"
183        ))),
184        RowEdit::Patch if values == 0 => Err(Error::invalid(
185            "patched rows hold at least one column besides their row IDs",
186        )),
187        _ => Ok(()),
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use arrow_schema::Field;
194
195    use super::*;
196
197    fn schema(fields: &[(&str, DataType)]) -> Schema {
198        Schema::new(
199            fields
200                .iter()
201                .map(|(name, data_type)| Field::new(*name, data_type.clone(), false))
202                .collect::<Vec<_>>(),
203        )
204    }
205
206    #[test]
207    fn edit_rows_name_their_targets_and_carry_the_columns_their_edit_needs() {
208        let ids = schema(&[(ROW_ID_COLUMN, DataType::UInt64)]);
209        let values = schema(&[(ROW_ID_COLUMN, DataType::UInt64), ("label", DataType::Utf8)]);
210        assert!(check_edit_schema(RowEdit::Delete, &ids).is_ok());
211        assert!(check_edit_schema(RowEdit::Delete, &values).is_err());
212        assert!(check_edit_schema(RowEdit::Patch, &values).is_ok());
213        assert!(check_edit_schema(RowEdit::Patch, &ids).is_err());
214        let signed = schema(&[(ROW_ID_COLUMN, DataType::Int64)]);
215        assert!(check_edit_schema(RowEdit::Delete, &signed).is_err());
216        assert!(check_edit_schema(RowEdit::Patch, &schema(&[("label", DataType::Utf8)])).is_err());
217    }
218}
spiral/table_write.rs

spiral/
table_write.rs
