spiral/
column_write.rs

1//! Values for new fields over the rows of one exact version of a source table.
2
3use std::collections::HashSet;
4use std::sync::Arc;
5
6use arrow_array::RecordBatch;
7use arrow_schema::{Field, Schema, SchemaRef};
8use futures::future::BoxFuture;
9use futures::{StreamExt as _, TryStreamExt as _, stream};
10use uuid::Uuid;
11
12use crate::registration::register;
13use crate::scandal::RecordBatchStream;
14use crate::scandal::api::AttachedColumns;
15use crate::scandal_backend::{AttemptWrite, RowEdit, input_error, reader_stream};
16use crate::table_write::{ROW_ID_COLUMN, check_edit_schema};
17use crate::{
18    Error, IntoArrowStream, NamedExpression, Scan, ScanOptions, ScanSplit, SchemaChange, Table,
19    WorkItem,
20};
21
22/// Fragments whose values [`Table::add_columns`] writes at once.
23const FRAGMENTS_IN_FLIGHT: usize = 4;
24
25/// Values for new nullable fields over the rows of one exact version of a table, from
26/// [`Table::column_writer`], which join its schema in its next version.
27///
28/// Values cover every row of a fragment with [`Self::attach_columns`] or chosen rows with
29/// [`Self::patch_rows`]; rows without values read `NULL`. Either way the values go into new files
30/// beside the fragment's existing ones, which are never rewritten. [`ColumnWrite::finish`]
31/// returns the table's next version, which is published only once a [`Commit`](crate::Commit)
32/// stages it and publishes.
33pub struct ColumnWrite {
34    source: Table,
35    fields: SchemaRef,
36    write: AttemptWrite,
37}
38
39impl ColumnWrite {
40    /// The exact source version whose rows the values cover.
41    pub fn source(&self) -> &Table {
42        &self.source
43    }
44
45    /// The fields this write gives values for.
46    pub fn fields(&self) -> SchemaRef {
47        Arc::clone(&self.fields)
48    }
49
50    /// Attach values of every field to fragment `fragment_id` of the source: one stream of
51    /// batches with a value for every row the fragment keeps, in stored order, which is the
52    /// order a scan of the fragment's splits returns. The values are written as they arrive,
53    /// into new files attached to the fragment in the next version; the fragment's existing
54    /// files are left as they are. Rows the source deletes read `NULL`. A write attaches each
55    /// fragment once, and not to rows it also patches: [`ColumnWrite::finish`] refuses both.
56    pub async fn attach_columns(
57        &mut self,
58        fragment_id: Uuid,
59        values: impl IntoArrowStream,
60    ) -> Result<(), Error> {
61        let reader = values.into_arrow_stream()?;
62        let schema = reader.schema();
63        let fragment = self
64            .fragment_values(fragment_id, schema, reader_stream(reader))
65            .await?
66            .await?;
67        self.write.add_receipts([fragment.receipt]);
68        Ok(())
69    }
70
71    /// Write values for some rows, recorded beside their fragments: each row names its target in
72    /// a [`ROW_ID_COLUMN`] column, and its other columns are some of [`Self::fields`].
73    pub async fn patch_rows(&mut self, rows: impl IntoArrowStream) -> Result<(), Error> {
74        let reader = rows.into_arrow_stream()?;
75        let schema = reader.schema();
76        check_edit_schema(RowEdit::Patch, &schema)?;
77        check_patched_fields(&self.fields, &schema)?;
78        let receipts = self
79            .write
80            .edit_rows(RowEdit::Patch, ROW_ID_COLUMN, schema, reader_stream(reader))
81            .await?;
82        self.write.add_receipts(receipts);
83        Ok(())
84    }
85
86    /// Attach values to every fragment, as [`Self::attach_columns`] does, with the values `scan`,
87    /// a scan of the source whose outputs are this write's fields, computes over it. Several
88    /// fragments are written at once.
89    async fn attach_scanned_columns(&mut self, scan: &Scan) -> Result<(), Error> {
90        let fields = Arc::clone(&self.fields);
91        let work = fragment_work(scan.splits())?;
92        let written = stream::iter(work)
93            .map(|(fragment_id, work)| {
94                let (this, fields) = (&*self, &fields);
95                async move {
96                    let batches = scan.read_work(work).await?;
97                    let schema = Arc::clone(fields);
98                    let input = batches
99                        .map(move |read| {
100                            let columns = read?.batch.columns().to_vec();
101                            RecordBatch::try_new(Arc::clone(&schema), columns).map_err(Error::arrow)
102                        })
103                        .map_err(input_error)
104                        .boxed();
105                    this.fragment_values(fragment_id, Arc::clone(fields), input)
106                        .await?
107                        .await
108                }
109            })
110            .buffer_unordered(FRAGMENTS_IN_FLIGHT)
111            .try_collect::<Vec<_>>()
112            .await?;
113        self.write
114            .add_receipts(written.into_iter().map(|fragment| fragment.receipt));
115        Ok(())
116    }
117
118    /// Start writing one fragment's values beside other fragments. Keep the receipt the request
119    /// returns.
120    async fn fragment_values(
121        &self,
122        fragment_id: Uuid,
123        schema: SchemaRef,
124        input: RecordBatchStream,
125    ) -> Result<BoxFuture<'static, Result<AttachedColumns, Error>>, Error> {
126        check_fragment_fields(&self.fields, &schema)?;
127        Ok(self.write.attach_columns(fragment_id, schema, input))
128    }
129
130    /// Prepare the table's next version, with the new columns and every value written, and
131    /// register it.
132    pub async fn finish(self) -> Result<Table, Error> {
133        let result = register(&self.source.connection.arbor, self.write.prepare().await?).await?;
134        self.source.successor(result)
135    }
136}
137
138impl Table {
139    /// Begin writing values for new nullable `fields` of this table, which join its schema in
140    /// its next version. Rows the write gives no values read `NULL`, so values may follow over
141    /// several versions.
142    pub async fn column_writer(&self, fields: SchemaRef) -> Result<ColumnWrite, Error> {
143        if fields.fields().is_empty() {
144            return Err(Error::invalid("a column write adds at least one field"));
145        }
146        let changes = fields
147            .fields()
148            .iter()
149            .map(|field| {
150                SchemaChange::AddColumn {
151                    field: Arc::clone(field),
152                    after: None,
153                }
154                .into_wire()
155            })
156            .collect::<Result<Vec<_>, _>>()?;
157        let write = self
158            .connection
159            .tables()?
160            .begin_edit(self.identity(), changes)
161            .await?;
162        Ok(ColumnWrite {
163            source: self.clone(),
164            fields,
165            write,
166        })
167    }
168
169    /// Write the next version of this table with a new nullable column for each of `columns`,
170    /// filled for every row with the expression's value over that row.
171    ///
172    /// Each fragment is read in stored order with the expressions as its projection and written
173    /// whole, several fragments at once.
174    pub async fn add_columns(&self, columns: Vec<NamedExpression>) -> Result<Table, Error> {
175        if columns.is_empty() {
176            return Err(Error::invalid(
177                "adding columns needs at least one expression",
178            ));
179        }
180        let scan = self
181            .scan(ScanOptions {
182                projection: Some(columns),
183                ..Default::default()
184            })
185            .await?;
186        let mut write = self.column_writer(nullable(&scan.schema())).await?;
187        write.attach_scanned_columns(&scan).await?;
188        write.finish().await
189    }
190}
191
192/// `schema` with every field nullable, as new columns are.
193pub(crate) fn nullable(schema: &Schema) -> SchemaRef {
194    Arc::new(Schema::new(
195        schema
196            .fields()
197            .iter()
198            .map(|field| Field::clone(field).with_nullable(true))
199            .collect::<Vec<_>>(),
200    ))
201}
202
203/// The work that reads each fragment of a scan whole: the fragment's ID and its splits, in
204/// stored order.
205fn fragment_work(splits: &[ScanSplit]) -> Result<Vec<(Uuid, Vec<WorkItem>)>, Error> {
206    let mut fragments: Vec<(Uuid, Vec<WorkItem>)> = Vec::new();
207    for (split, planned) in splits.iter().enumerate() {
208        let fragment_id = planned.fragment_id.ok_or_else(|| {
209            Error::invalid("the scandal does not name the fragments a scan's splits read")
210        })?;
211        let item = WorkItem {
212            split,
213            rows: planned.rows.clone(),
214        };
215        match fragments.last_mut() {
216            Some((last, items)) if *last == fragment_id => items.push(item),
217            _ => fragments.push((fragment_id, vec![item])),
218        }
219    }
220    let mut seen = HashSet::new();
221    if let Some((repeated, _)) = fragments.iter().find(|(id, _)| !seen.insert(*id)) {
222        return Err(Error::Protocol(format!(
223            "the splits of fragment {repeated} are not adjacent"
224        )));
225    }
226    Ok(fragments)
227}
228
229/// Check that a fragment's values hold every one of `fields`, in order, with its type.
230fn check_fragment_fields(fields: &Schema, values: &Schema) -> Result<(), Error> {
231    let matches = fields.fields().len() == values.fields().len()
232        && fields
233            .fields()
234            .iter()
235            .zip(values.fields())
236            .all(|(field, value)| {
237                field.name() == value.name() && field.data_type() == value.data_type()
238            });
239    if !matches {
240        return Err(Error::invalid(format!(
241            "a fragment's values hold exactly the written fields {:?}",
242            field_names(fields)
243        )));
244    }
245    Ok(())
246}
247
248/// Check that patched rows hold a row ID column and some of `fields`, each with its type.
249fn check_patched_fields(fields: &Schema, rows: &Schema) -> Result<(), Error> {
250    for value in rows.fields() {
251        if value.name() == ROW_ID_COLUMN {
252            continue;
253        }
254        let field = fields.field_with_name(value.name()).map_err(|_| {
255            Error::invalid(format!(
256                "patched column `{}` is not one of the written fields {:?}",
257                value.name(),
258                field_names(fields)
259            ))
260        })?;
261        if field.data_type() != value.data_type() {
262            return Err(Error::invalid(format!(
263                "patched column `{}` holds {}, not {}",
264                value.name(),
265                value.data_type(),
266                field.data_type()
267            )));
268        }
269    }
270    Ok(())
271}
272
273fn field_names(schema: &Schema) -> Vec<&str> {
274    schema
275        .fields()
276        .iter()
277        .map(|field| field.name().as_str())
278        .collect()
279}
280
281#[cfg(test)]
282mod tests {
283    use std::ops::Range;
284
285    use arrow_schema::DataType;
286
287    use super::*;
288
289    fn split(fragment_id: Option<Uuid>, rows: Range<u64>) -> ScanSplit {
290        ScanSplit {
291            fragment: 0,
292            rows,
293            estimated_rows: None,
294            estimated_bytes: None,
295            fragment_id,
296        }
297    }
298
299    #[test]
300    fn each_fragment_is_read_whole_from_its_adjacent_splits() {
301        let (first, second) = (Uuid::now_v7(), Uuid::now_v7());
302        let work = fragment_work(&[
303            split(Some(first), 0..10),
304            split(Some(first), 10..20),
305            split(Some(second), 0..5),
306        ])
307        .unwrap();
308        assert_eq!(work.len(), 2);
309        assert_eq!(work[0].0, first);
310        assert_eq!(
311            work[0].1.iter().map(|item| item.split).collect::<Vec<_>>(),
312            [0, 1]
313        );
314        assert_eq!(work[1].1[0].rows, 0..5);
315        assert!(fragment_work(&[split(None, 0..1)]).is_err());
316        assert!(
317            fragment_work(&[
318                split(Some(first), 0..1),
319                split(Some(second), 0..1),
320                split(Some(first), 1..2),
321            ])
322            .is_err()
323        );
324    }
325
326    #[test]
327    fn values_hold_the_written_fields_with_their_types() {
328        let fields = Schema::new(vec![
329            Field::new("score", DataType::Float32, true),
330            Field::new("label", DataType::Utf8, true),
331        ]);
332        let required = Schema::new(vec![
333            Field::new("score", DataType::Float32, false),
334            Field::new("label", DataType::Utf8, false),
335        ]);
336        assert!(check_fragment_fields(&fields, &required).is_ok());
337        let partial = Schema::new(vec![Field::new("score", DataType::Float32, true)]);
338        assert!(check_fragment_fields(&fields, &partial).is_err());
339        let patched = Schema::new(vec![
340            Field::new(ROW_ID_COLUMN, DataType::UInt64, false),
341            Field::new("label", DataType::Utf8, true),
342        ]);
343        assert!(check_patched_fields(&fields, &patched).is_ok());
344        let retyped = Schema::new(vec![
345            Field::new(ROW_ID_COLUMN, DataType::UInt64, false),
346            Field::new("label", DataType::LargeUtf8, true),
347        ]);
348        assert!(check_patched_fields(&fields, &retyped).is_err());
349    }
350}
spiral/column_write.rs

spiral/
column_write.rs
