nodedb_physical/physical_plan/vector/
expand.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! The point writes a vector-primary page expands to.
4//!
5//! A transaction stages each row of a page on its own, so each row gets its
6//! own overlay image, undo entry and redo entry. A key the page names twice
7//! is in the overlay when its later row stages, so that row is refused,
8//! skipped or merged as the page judges it.
9
10use super::op::VectorOp;
11use super::write::{VectorDirectRow, VectorDirectWriteIntent};
12
13impl VectorOp {
14    /// One `DirectInsert` / `DirectInsertIfAbsent` / `DirectUpsert` per row
15    /// of a `DirectBatchWrite` page, in page order, by the page's intent.
16    /// `None` for every other op.
17    pub fn direct_batch_point_writes(&self) -> Option<Vec<VectorOp>> {
18        let VectorOp::DirectBatchWrite {
19            collection,
20            field,
21            intent,
22            rows,
23            quantization,
24            storage_dtype,
25            payload_indexes,
26            returning,
27            rls_filters,
28            on_conflict_updates,
29            rls_write_check,
30        } = self
31        else {
32            return None;
33        };
34        Some(
35            rows.iter()
36                .map(|row| {
37                    let VectorDirectRow {
38                        surrogate,
39                        pk_bytes,
40                        vector,
41                        payload,
42                    } = row.clone();
43                    match intent {
44                        VectorDirectWriteIntent::Insert => VectorOp::DirectInsert {
45                            collection: collection.clone(),
46                            field: field.clone(),
47                            surrogate,
48                            pk_bytes,
49                            vector,
50                            payload,
51                            quantization: *quantization,
52                            storage_dtype: *storage_dtype,
53                            payload_indexes: payload_indexes.clone(),
54                            returning: returning.clone(),
55                            rls_filters: rls_filters.clone(),
56                        },
57                        VectorDirectWriteIntent::InsertIfAbsent => VectorOp::DirectInsertIfAbsent {
58                            collection: collection.clone(),
59                            field: field.clone(),
60                            surrogate,
61                            pk_bytes,
62                            vector,
63                            payload,
64                            quantization: *quantization,
65                            storage_dtype: *storage_dtype,
66                            payload_indexes: payload_indexes.clone(),
67                            returning: returning.clone(),
68                            rls_filters: rls_filters.clone(),
69                        },
70                        VectorDirectWriteIntent::Upsert => VectorOp::DirectUpsert {
71                            collection: collection.clone(),
72                            field: field.clone(),
73                            surrogate,
74                            pk_bytes,
75                            vector,
76                            payload,
77                            quantization: *quantization,
78                            storage_dtype: *storage_dtype,
79                            payload_indexes: payload_indexes.clone(),
80                            returning: returning.clone(),
81                            rls_filters: rls_filters.clone(),
82                            on_conflict_updates: on_conflict_updates.clone(),
83                            rls_write_check: rls_write_check.clone(),
84                        },
85                    }
86                })
87                .collect(),
88        )
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use nodedb_types::{DatabaseId, QualifiedCollection, RlsWriteCheck, Surrogate};
95
96    use super::*;
97
98    fn page(intent: VectorDirectWriteIntent) -> VectorOp {
99        let row = |s: u32, pk: &str| VectorDirectRow {
100            surrogate: Surrogate::new(s),
101            pk_bytes: pk.as_bytes().to_vec(),
102            vector: vec![0.5, 0.25],
103            payload: Vec::new(),
104        };
105        VectorOp::DirectBatchWrite {
106            collection: QualifiedCollection::new(DatabaseId::DEFAULT, "vecs"),
107            field: "emb".into(),
108            intent,
109            rows: vec![row(1, "a"), row(2, "b")],
110            quantization: nodedb_types::VectorQuantization::None,
111            storage_dtype: nodedb_types::VectorStorageDtype::F32,
112            payload_indexes: Vec::new(),
113            returning: None,
114            rls_filters: Vec::new(),
115            on_conflict_updates: Vec::new(),
116            rls_write_check: RlsWriteCheck::NoPolicyApplies,
117        }
118    }
119
120    #[test]
121    fn a_page_expands_to_one_point_write_per_row_in_order() {
122        let writes = page(VectorDirectWriteIntent::InsertIfAbsent)
123            .direct_batch_point_writes()
124            .expect("a page expands");
125        let keys: Vec<Vec<u8>> = writes
126            .iter()
127            .map(|write| match write {
128                VectorOp::DirectInsertIfAbsent { pk_bytes, .. } => pk_bytes.clone(),
129                other => panic!("expected DirectInsertIfAbsent, got {other:?}"),
130            })
131            .collect();
132        assert_eq!(keys, vec![b"a".to_vec(), b"b".to_vec()]);
133        assert!(matches!(
134            page(VectorDirectWriteIntent::Upsert)
135                .direct_batch_point_writes()
136                .as_deref(),
137            Some([VectorOp::DirectUpsert { .. }, VectorOp::DirectUpsert { .. }])
138        ));
139    }
140
141    #[test]
142    fn a_point_write_does_not_expand() {
143        let point = page(VectorDirectWriteIntent::Insert)
144            .direct_batch_point_writes()
145            .and_then(|writes| writes.into_iter().next())
146            .expect("a row");
147        assert!(point.direct_batch_point_writes().is_none());
148    }
149}
nodedb_physical/physical_plan/vector/expand.rs

nodedb_physical/physical_plan/vector/
expand.rs
