Skip to main content

lora_store/memory/
snapshot.rs

1//! Snapshot payload helpers for the in-memory graph.
2//!
3//! `lora-store` no longer ships its own on-disk codec. The byte-level
4//! columnar format lives in `lora-snapshot`; this module just bridges
5//! between [`InMemoryGraph`] and the portable [`SnapshotPayload`]
6//! vocabulary.
7
8use std::collections::BTreeSet;
9use std::sync::Arc;
10
11use crate::dict::{Dicts, ValueDictionary};
12use crate::encoded::{self, StoredNode, StoredRel};
13use crate::{
14    Name, NodeId, NodeRecord, NodeRef, PropertyValue, RelRef, RelationshipId, RelationshipRecord,
15    SnapshotError, SnapshotMeta, SnapshotPayload, ValueRef,
16};
17
18use super::distinct_stats::DistinctBuilder;
19use super::graph::DeferredData;
20use super::index_catalog::StoredIndexEntity;
21use super::record_slab::BuiltLeaves;
22use super::{ConstraintDefinition, InMemoryGraph, IndexDefinition, VectorIndexSnapshot};
23
24/// Format-version stamp surfaced through [`SnapshotMeta::format_version`]
25/// for payloads produced via the inherent helpers below. Kept stable
26/// across `lora-snapshot` codec versions because the payload shape
27/// itself has not changed; only the on-disk encoding has.
28pub(super) const PAYLOAD_FORMAT_VERSION: u32 = 1;
29
30impl InMemoryGraph {
31    /// Return the portable graph-state payload. Callers downstream of
32    /// `lora-store` (typically `lora-database`) feed this into
33    /// `lora-snapshot` for byte-level encoding.
34    pub fn snapshot_payload(&self) -> SnapshotPayload {
35        let mut vector_indexes = self
36            .vector_indexes_read(StoredIndexEntity::Node)
37            .to_snapshots(StoredIndexEntity::Node);
38        vector_indexes.extend(
39            self.vector_indexes_read(StoredIndexEntity::Relationship)
40                .to_snapshots(StoredIndexEntity::Relationship),
41        );
42        SnapshotPayload {
43            next_node_id: self.next_node_id,
44            next_rel_id: self.next_rel_id,
45            nodes: self.iter_node_refs().map(|r| r.to_record()).collect(),
46            relationships: self.iter_rel_refs().map(|r| r.to_record()).collect(),
47            indexes: self.index_catalog_read().list(),
48            constraints: self.constraint_catalog_read().list(),
49            vector_indexes,
50        }
51    }
52
53    /// Replace the graph from a portable graph-state payload, preserving the
54    /// currently installed mutation recorder across the swap.
55    pub fn load_snapshot_payload(
56        &mut self,
57        payload: SnapshotPayload,
58    ) -> Result<SnapshotMeta, SnapshotError> {
59        let meta = SnapshotMeta {
60            format_version: PAYLOAD_FORMAT_VERSION,
61            node_count: payload.nodes.len(),
62            relationship_count: payload.relationships.len(),
63            wal_lsn: None,
64        };
65
66        validate_payload_ids(&payload)?;
67
68        let mut builder = SnapshotGraphBuilder::new(payload.next_node_id, payload.next_rel_id)?;
69        for node in payload.nodes {
70            builder.push_node(node)?;
71        }
72        for rel in payload.relationships {
73            builder.push_relationship(rel)?;
74        }
75        let mut rebuilt =
76            builder.finish(payload.indexes, payload.constraints, payload.vector_indexes)?;
77
78        // Preserve the existing recorder across the swap — observers of the
79        // store's identity should not be silently detached by a restore,
80        // same policy as `clear()`.
81        rebuilt.recorder = self.recorder.take();
82        *self = rebuilt;
83
84        Ok(meta)
85    }
86
87    /// Replace the graph with `restored`, one a snapshot decoder built
88    /// (`lora_snapshot::read_graph_snapshot`), preserving the currently
89    /// installed mutation recorder as [`Self::load_snapshot_payload`] does.
90    pub fn replace_with_restored(&mut self, mut restored: InMemoryGraph) {
91        restored.recorder = self.recorder.take();
92        *self = restored;
93    }
94
95    /// The node records, in id order, as borrowed views. With
96    /// [`Self::snapshot_relationships`] and [`Self::snapshot_catalogs`] this
97    /// is what a snapshot encoder reads instead of a [`SnapshotPayload`],
98    /// which would hold a decoded copy of every record at once. Each call
99    /// is a fresh pass, so an encoder that writes column by column can
100    /// walk the graph once per column.
101    pub fn snapshot_nodes(&self) -> impl Iterator<Item = NodeRef<'_>> + '_ {
102        self.iter_node_refs()
103    }
104
105    /// The relationship records, in id order, as borrowed views.
106    pub fn snapshot_relationships(&self) -> impl Iterator<Item = RelRef<'_>> + '_ {
107        self.iter_rel_refs()
108    }
109
110    /// Everything a snapshot holds besides the records.
111    pub fn snapshot_catalogs(&self) -> SnapshotCatalogs {
112        let mut vector_indexes = self
113            .vector_indexes_read(StoredIndexEntity::Node)
114            .to_snapshots(StoredIndexEntity::Node);
115        vector_indexes.extend(
116            self.vector_indexes_read(StoredIndexEntity::Relationship)
117                .to_snapshots(StoredIndexEntity::Relationship),
118        );
119        SnapshotCatalogs {
120            next_node_id: self.next_node_id,
121            next_rel_id: self.next_rel_id,
122            indexes: self.index_catalog_read().list(),
123            constraints: self.constraint_catalog_read().list(),
124            vector_indexes,
125        }
126    }
127}
128
129/// The parts of a snapshot that are not records: the id counters and the
130/// declared catalogs. Small next to the graph.
131#[derive(Debug, Clone)]
132pub struct SnapshotCatalogs {
133    pub next_node_id: NodeId,
134    pub next_rel_id: RelationshipId,
135    pub indexes: Vec<IndexDefinition>,
136    pub constraints: Vec<ConstraintDefinition>,
137    pub vector_indexes: Vec<VectorIndexSnapshot>,
138}
139
140/// Builds an [`InMemoryGraph`] from snapshot records one at a time.
141///
142/// A decoder that streams a snapshot pushes each record as it decodes it,
143/// so the records are never all held at once beside the graph they fill.
144/// Push every node before the relationships that reference it, then call
145/// [`Self::finish`] with the catalogs, which registers the indexes and
146/// constraints over the loaded data.
147pub struct SnapshotGraphBuilder {
148    graph: InMemoryGraph,
149    node_distinct: DistinctBuilder,
150    rel_distinct: DistinctBuilder,
151}
152
153/// A node record names one label twice: no write makes one, and the label
154/// index would hold the node twice.
155fn repeated_label(id: NodeId) -> SnapshotError {
156    SnapshotError::Decode(format!("node {id} has a repeated label in snapshot"))
157}
158
159impl SnapshotGraphBuilder {
160    pub fn new(next_node_id: NodeId, next_rel_id: RelationshipId) -> Result<Self, SnapshotError> {
161        validate_next_id("node", next_node_id)?;
162        validate_next_id("relationship", next_rel_id)?;
163        let mut graph = InMemoryGraph::new();
164        graph.next_node_id = next_node_id;
165        graph.next_rel_id = next_rel_id;
166        Ok(Self {
167            graph,
168            node_distinct: DistinctBuilder::default(),
169            rel_distinct: DistinctBuilder::default(),
170        })
171    }
172
173    pub fn push_node(&mut self, node: NodeRecord) -> Result<(), SnapshotError> {
174        let id = node.id;
175        validate_entity_id("node", id, self.graph.next_node_id)?;
176        if node
177            .labels
178            .iter()
179            .enumerate()
180            .any(|(i, label)| node.labels[..i].contains(label))
181        {
182            return Err(repeated_label(id));
183        }
184        // Distinct-value sketches are not persisted: count the record as a
185        // live create would, through the bulk builder.
186        for label in &node.labels {
187            self.node_distinct.add_all(label, &node.properties);
188        }
189        if self.graph.has_node_at(id) {
190            return Err(SnapshotError::Decode(format!(
191                "duplicate node id {id} in snapshot payload"
192            )));
193        }
194        self.graph
195            .put_node_checked(id, &node)
196            .map_err(SnapshotError::Decode)?;
197        for label in &node.labels {
198            self.graph.insert_node_label_index(id, label);
199        }
200        Ok(())
201    }
202
203    pub fn push_relationship(&mut self, rel: RelationshipRecord) -> Result<(), SnapshotError> {
204        validate_entity_id("relationship", rel.id, self.graph.next_rel_id)?;
205        validate_slot_id("relationship source node", rel.src)?;
206        validate_slot_id("relationship target node", rel.dst)?;
207        if self.graph.has_rel_at(rel.id) {
208            return Err(SnapshotError::Decode(format!(
209                "duplicate relationship id {} in snapshot payload",
210                rel.id
211            )));
212        }
213        if !self.graph.has_node_at(rel.src) {
214            return Err(SnapshotError::Decode(format!(
215                "relationship {} references missing source node {}",
216                rel.id, rel.src
217            )));
218        }
219        if !self.graph.has_node_at(rel.dst) {
220            return Err(SnapshotError::Decode(format!(
221                "relationship {} references missing target node {}",
222                rel.id, rel.dst
223            )));
224        }
225        self.graph
226            .put_rel_checked(rel.id, &rel)
227            .map_err(SnapshotError::Decode)?;
228        self.graph.attach_relationship(&rel);
229        self.rel_distinct.add_all(&rel.rel_type, &rel.properties);
230        Ok(())
231    }
232
233    pub fn finish(
234        self,
235        indexes: Vec<IndexDefinition>,
236        constraints: Vec<ConstraintDefinition>,
237        vector_indexes: Vec<VectorIndexSnapshot>,
238    ) -> Result<InMemoryGraph, SnapshotError> {
239        let Self {
240            graph: mut rebuilt,
241            node_distinct,
242            rel_distinct,
243        } = self;
244        rebuilt.distinct_stats.nodes = node_distinct.finish();
245        rebuilt.distinct_stats.relationships = rel_distinct.finish();
246        // No hash property index is built here. Restart reproduces the
247        // state of a fresh process: the declared RANGE indexes and the
248        // uniqueness/key constraints below activate (and backfill) the
249        // hash indexes they need through `register_index` /
250        // `register_constraint`. Implicit, lookup-activated indexes are
251        // rebuilt lazily by the first lookup that needs them.
252
253        let constraint_owned_indexes: BTreeSet<String> = constraints
254            .iter()
255            .filter_map(|def| {
256                def.owned_index
257                    .clone()
258                    .or_else(|| def.kind.requires_backing_index().then(|| def.name.clone()))
259            })
260            .collect();
261
262        // Re-register every user-visible index in the catalog. Going through
263        // `register_index` re-populates RANGE buckets and keeps the
264        // `populate_index_data` invariant aligned with the catalog —
265        // skipping it would leave RANGE indexes registered but never populated.
266        // Constraint-owned backing indexes are restored by re-registering the
267        // owning constraint below, which keeps catalog ownership explicit.
268        //
269        // Registration only fills the catalogs, in order; the indexes'
270        // data and the constraint scans are left in `deferred` and done
271        // together at the end, on every core.
272        let mut deferred = DeferredData::default();
273        for def in indexes {
274            if constraint_owned_indexes.contains(&def.name) {
275                continue;
276            }
277            // Errors here would mean the snapshot itself is corrupt or
278            // ambiguous; map them into Decode rather than panicking.
279            rebuilt
280                .register_index_deferring(
281                    crate::memory::IndexRequest {
282                        explicit_name: Some(def.name.clone()),
283                        kind: def.kind,
284                        entity: def.entity,
285                        label: def.label.clone(),
286                        additional_labels: def.additional_labels.clone(),
287                        properties: def.properties.clone(),
288                        options: def.options.clone(),
289                    },
290                    /*if_not_exists*/ true,
291                    /*record_event*/ true,
292                    Some(&mut deferred),
293                )
294                .map_err(|e| SnapshotError::Decode(format!("index `{}`: {e}", def.name)))?;
295        }
296
297        // Overlay persisted HNSW snapshots over the freshly-registered
298        // (and freshly-backfilled) vector indexes. This is the
299        // post-step that gives Phase 5 its raison d'être: instead of
300        // paying O(n log n) to re-insert every vector through the
301        // HNSW algorithm, we install the graph topology byte-for-byte.
302        // Snapshots from versions before this trailer round-trip with
303        // `vector_indexes = []` so the fallback path (the backfill
304        // that already ran inside `register_index`) handles them
305        // correctly.
306        for snap in vector_indexes {
307            let entity = snap.entity;
308            let mut registry = rebuilt.vector_indexes_write(entity);
309            if !registry.restore_snapshot(snap) {
310                // Catalog/snapshot mismatch — registry already
311                // contains the populate-built backend, which is the
312                // safe fallback. No further action.
313            }
314        }
315
316        // Re-register constraints. Uniqueness / key constraints recreate
317        // their own backing indexes as part of registration.
318        for def in constraints {
319            rebuilt
320                .register_constraint_deferring(
321                    crate::memory::ConstraintRequest {
322                        name: def.name.clone(),
323                        kind: def.kind.clone(),
324                        entity: def.entity,
325                        label: def.label.clone(),
326                        properties: def.properties.clone(),
327                        backing: def.backing,
328                    },
329                    /*if_not_exists*/ true,
330                    Some(&mut deferred),
331                )
332                .map_err(|e| SnapshotError::Decode(format!("constraint `{}`: {e}", def.name)))?;
333        }
334
335        rebuilt
336            .run_deferred(deferred)
337            .map_err(SnapshotError::Decode)?;
338        Ok(rebuilt)
339    }
340}
341
342fn validate_payload_ids(payload: &SnapshotPayload) -> Result<(), SnapshotError> {
343    validate_next_id("node", payload.next_node_id)?;
344    validate_next_id("relationship", payload.next_rel_id)?;
345
346    for node in &payload.nodes {
347        validate_entity_id("node", node.id, payload.next_node_id)?;
348    }
349    for rel in &payload.relationships {
350        validate_entity_id("relationship", rel.id, payload.next_rel_id)?;
351        validate_slot_id("relationship source node", rel.src)?;
352        validate_slot_id("relationship target node", rel.dst)?;
353    }
354
355    Ok(())
356}
357
358fn validate_next_id(kind: &str, next_id: u64) -> Result<(), SnapshotError> {
359    validate_slot_id(&format!("next {kind} id"), next_id)?;
360    if next_id == u64::MAX {
361        return Err(SnapshotError::Decode(format!(
362            "next {kind} id {next_id} leaves no allocatable id"
363        )));
364    }
365    Ok(())
366}
367
368fn validate_entity_id(kind: &str, id: u64, next_id: u64) -> Result<(), SnapshotError> {
369    validate_slot_id(kind, id)?;
370    if id >= next_id {
371        return Err(SnapshotError::Decode(format!(
372            "{kind} id {id} is not below next {kind} id {next_id}"
373        )));
374    }
375    Ok(())
376}
377
378fn validate_slot_id(label: &str, id: u64) -> Result<(), SnapshotError> {
379    let idx = usize::try_from(id).map_err(|_| {
380        SnapshotError::Decode(format!(
381            "{label} {id} does not fit in usize on this platform"
382        ))
383    })?;
384    idx.checked_add(1)
385        .ok_or_else(|| SnapshotError::Decode(format!("{label} {id} leaves no valid slab slot")))?;
386    Ok(())
387}
388
389// ---------------------------------------------------------------------------
390// Encoding on several threads
391// ---------------------------------------------------------------------------
392
393/// The dictionaries a graph numbers its records with, in number order.
394/// A snapshot writer names records by these numbers, and a
395/// [`SnapshotGraphBuilder`] primed with them ([`SnapshotGraphBuilder::prime`])
396/// encodes records on any thread.
397#[derive(Debug, Clone, Default, PartialEq, Eq)]
398pub struct SnapshotNames {
399    pub labels: Vec<Name>,
400    pub types: Vec<Name>,
401    pub keys: Vec<Name>,
402    /// Which short string values the graph stores by number.
403    pub values: ValueDictionary,
404}
405
406/// A stored node as a snapshot writer reads it: names by number in
407/// [`SnapshotNames`], values borrowed.
408#[derive(Clone, Copy)]
409pub struct SnapshotNode<'a> {
410    stored: StoredNode<'a>,
411}
412
413impl<'a> SnapshotNode<'a> {
414    pub fn id(self) -> NodeId {
415        self.stored.id
416    }
417
418    /// Numbers in [`SnapshotNames::labels`], in the node's order.
419    pub fn label_ids(self) -> impl ExactSizeIterator<Item = u32> + 'a {
420        let labels = self.stored.labels();
421        let count = labels.len();
422        ExactLen {
423            inner: labels.ids(),
424            remaining: count,
425        }
426    }
427
428    /// `(number in` [`SnapshotNames::keys`]`, value)` in key-name order.
429    pub fn properties(self) -> impl Iterator<Item = (u32, ValueRef<'a>)> + 'a {
430        self.stored.properties().iter_ids()
431    }
432
433    pub fn property_count(self) -> usize {
434        self.stored.properties().len()
435    }
436}
437
438/// A stored relationship as a snapshot writer reads it.
439#[derive(Clone, Copy)]
440pub struct SnapshotRelationship<'a> {
441    stored: StoredRel<'a>,
442}
443
444impl<'a> SnapshotRelationship<'a> {
445    pub fn id(self) -> RelationshipId {
446        self.stored.id
447    }
448
449    pub fn src(self) -> NodeId {
450        self.stored.src
451    }
452
453    pub fn dst(self) -> NodeId {
454        self.stored.dst
455    }
456
457    /// Number in [`SnapshotNames::types`].
458    pub fn type_id(self) -> u32 {
459        self.stored.type_id()
460    }
461
462    /// `(number in` [`SnapshotNames::keys`]`, value)` in key-name order.
463    pub fn properties(self) -> impl Iterator<Item = (u32, ValueRef<'a>)> + 'a {
464        self.stored.properties().iter_ids()
465    }
466
467    pub fn property_count(self) -> usize {
468        self.stored.properties().len()
469    }
470}
471
472/// An iterator of known length.
473struct ExactLen<I> {
474    inner: I,
475    remaining: usize,
476}
477
478impl<I: Iterator> Iterator for ExactLen<I> {
479    type Item = I::Item;
480
481    fn next(&mut self) -> Option<I::Item> {
482        let item = self.inner.next()?;
483        self.remaining -= 1;
484        Some(item)
485    }
486
487    fn size_hint(&self) -> (usize, Option<usize>) {
488        (self.remaining, Some(self.remaining))
489    }
490}
491
492impl<I: Iterator> ExactSizeIterator for ExactLen<I> {}
493
494impl InMemoryGraph {
495    /// The graph's dictionaries. Every name a record holds is in them,
496    /// and so may names no record holds any more.
497    pub fn snapshot_names(&self) -> SnapshotNames {
498        let names = |dict: &crate::dict::NameDict| {
499            (0..dict.len() as u32)
500                .map(|id| dict.name(id).clone())
501                .collect()
502        };
503        SnapshotNames {
504            labels: names(&self.dicts.labels),
505            types: names(&self.dicts.types),
506            keys: names(&self.dicts.keys),
507            values: self
508                .dicts
509                .values
510                .export(|key| self.dicts.keys.name(key).clone()),
511        }
512    }
513
514    /// The nodes with an id in `ids`, in id order, numbered as
515    /// [`Self::snapshot_names`] numbers them. Writers on several threads
516    /// each take a range.
517    pub fn snapshot_nodes_in(
518        &self,
519        ids: std::ops::Range<NodeId>,
520    ) -> impl Iterator<Item = SnapshotNode<'_>> + '_ {
521        let dicts = &self.dicts;
522        let start = usize::try_from(ids.start).unwrap_or(usize::MAX);
523        let end = ids.end;
524        self.nodes
525            .iter_from(start)
526            .map(|(i, bytes)| (i as NodeId, bytes))
527            .take_while(move |(id, _)| *id < end)
528            .map(move |(id, bytes)| SnapshotNode {
529                stored: StoredNode::new(id, bytes, dicts),
530            })
531    }
532
533    /// The relationships with an id in `ids`, in id order.
534    pub fn snapshot_relationships_in(
535        &self,
536        ids: std::ops::Range<RelationshipId>,
537    ) -> impl Iterator<Item = SnapshotRelationship<'_>> + '_ {
538        let dicts = &self.dicts;
539        let start = usize::try_from(ids.start).unwrap_or(usize::MAX);
540        let end = ids.end;
541        self.relationships
542            .iter_from(start)
543            .map(|(i, bytes)| (i as RelationshipId, bytes))
544            .take_while(move |(id, _)| *id < end)
545            .map(move |(id, bytes)| SnapshotRelationship {
546                stored: StoredRel::new(id, bytes, dicts),
547            })
548    }
549
550    /// One past the highest node slot (an upper bound on node ids).
551    pub fn snapshot_node_slots(&self) -> u64 {
552        self.nodes.len() as u64
553    }
554
555    /// One past the highest relationship slot.
556    pub fn snapshot_relationship_slots(&self) -> u64 {
557        self.relationships.len() as u64
558    }
559}
560
561/// Encodes snapshot records into a graph's stored form on any thread,
562/// for the [`SnapshotGraphBuilder`] that [primed](SnapshotGraphBuilder::prime)
563/// it. Cheap to clone.
564#[derive(Clone)]
565pub struct SnapshotEncoder {
566    inner: Arc<EncoderTables>,
567}
568
569struct EncoderTables {
570    /// The builder's dictionaries when it was primed. Encoding reads them
571    /// only; a value they do not number yet is stored in full.
572    dicts: Dicts,
573    /// Snapshot label / type / key number -> graph number.
574    labels: Vec<u32>,
575    types: Vec<u32>,
576    keys: Vec<u32>,
577    /// The snapshot's key names, and each one's place in name order:
578    /// records keep their properties in that order.
579    key_names: Vec<Arc<str>>,
580    key_ranks: Vec<u32>,
581    label_names: Vec<Name>,
582    type_names: Vec<Name>,
583    next_node_id: NodeId,
584    next_rel_id: RelationshipId,
585}
586
587/// Nodes encoded off the builder's thread, for
588/// [`SnapshotGraphBuilder::push_nodes`]. Fill with [`SnapshotEncoder::push_node`].
589pub struct NodeBatch {
590    /// The encoded records, in whole leaves of the node slab.
591    records: BuiltLeaves,
592    by_label: Vec<(u32, Vec<NodeId>)>,
593    distinct: DistinctBuilder,
594    /// Reused per record.
595    scratch: Vec<u8>,
596    labels: Vec<u32>,
597}
598
599/// Relationships encoded off the builder's thread, for
600/// [`SnapshotGraphBuilder::push_relationships`].
601pub struct RelationshipBatch {
602    /// The encoded records, in whole leaves of the relationship slab.
603    records: BuiltLeaves,
604    /// `(id, source, target, graph type number)` per relationship.
605    links: Vec<(RelationshipId, NodeId, NodeId, u32)>,
606    by_type: Vec<(u32, Vec<RelationshipId>)>,
607    distinct: DistinctBuilder,
608    scratch: Vec<u8>,
609}
610
611impl NodeBatch {
612    pub fn len(&self) -> usize {
613        self.records.len()
614    }
615
616    pub fn is_empty(&self) -> bool {
617        self.records.len() == 0
618    }
619}
620
621impl RelationshipBatch {
622    pub fn len(&self) -> usize {
623        self.records.len()
624    }
625
626    pub fn is_empty(&self) -> bool {
627        self.records.len() == 0
628    }
629}
630
631/// Add record `id` to a batch's leaves. Ids ascend within a batch, so
632/// a repeated or falling one is refused here.
633fn push_record(records: &mut BuiltLeaves, id: u64, bytes: &[u8]) -> Result<(), SnapshotError> {
634    let slot = usize::try_from(id)
635        .map_err(|_| SnapshotError::Decode(format!("id {id} does not fit in usize")))?;
636    if records.last_slot().is_some_and(|last| last >= slot) {
637        return Err(SnapshotError::Decode(format!(
638            "record ids in a snapshot batch must ascend (id {id})"
639        )));
640    }
641    records.push(slot, bytes);
642    Ok(())
643}
644
645/// `ids` for `number` in a batch's by-name lists, which hold few names.
646fn ids_for(lists: &mut Vec<(u32, Vec<u64>)>, number: u32) -> &mut Vec<u64> {
647    let at = match lists.iter().rposition(|(n, _)| *n == number) {
648        Some(at) => at,
649        None => {
650            lists.push((number, Vec::new()));
651            lists.len() - 1
652        }
653    };
654    &mut lists[at].1
655}
656
657impl SnapshotEncoder {
658    /// A batch for about `capacity` records.
659    pub fn node_batch(&self, capacity: usize) -> NodeBatch {
660        let _ = capacity;
661        NodeBatch {
662            records: BuiltLeaves::default(),
663            by_label: Vec::new(),
664            distinct: DistinctBuilder::default(),
665            scratch: Vec::new(),
666            labels: Vec::new(),
667        }
668    }
669
670    /// A batch for about `capacity` records. `capacity` comes from the
671    /// snapshot being read, so at most a frame's worth is reserved up front.
672    pub fn relationship_batch(&self, capacity: usize) -> RelationshipBatch {
673        RelationshipBatch {
674            records: BuiltLeaves::default(),
675            links: Vec::with_capacity(capacity.min(1 << 16)),
676            by_type: Vec::new(),
677            distinct: DistinctBuilder::default(),
678            scratch: Vec::new(),
679        }
680    }
681
682    /// Put `properties` (snapshot key numbers) in the order a record keeps
683    /// them, the last of a repeated key winning as in [`Properties`], and
684    /// renumber the keys for the graph.
685    fn order_properties(
686        &self,
687        properties: &mut Vec<(u32, PropertyValue)>,
688    ) -> Result<(), SnapshotError> {
689        let tables = &*self.inner;
690        let mut sorted = true;
691        let mut last = None;
692        for (key, _) in properties.iter() {
693            let rank = *tables
694                .key_ranks
695                .get(*key as usize)
696                .ok_or_else(|| SnapshotError::Decode("invalid property key number".into()))?;
697            if last.is_some_and(|last| last >= rank) {
698                sorted = false;
699            }
700            last = Some(rank);
701        }
702        if !sorted {
703            // Stable, so of two equal keys the later stays later; keep it.
704            properties.sort_by_key(|(key, _)| tables.key_ranks[*key as usize]);
705            let mut i = properties.len();
706            while i > 1 {
707                i -= 1;
708                let same = tables.key_ranks[properties[i].0 as usize]
709                    == tables.key_ranks[properties[i - 1].0 as usize];
710                if same {
711                    properties.remove(i - 1);
712                }
713            }
714        }
715        Ok(())
716    }
717
718    /// Encode a node: `labels` and the keys of `properties` are numbers in
719    /// the snapshot's [`SnapshotNames`]. Takes the values out of
720    /// `properties`, leaving it empty for the next record.
721    pub fn push_node(
722        &self,
723        batch: &mut NodeBatch,
724        id: NodeId,
725        labels: &[u32],
726        properties: &mut Vec<(u32, PropertyValue)>,
727    ) -> Result<(), SnapshotError> {
728        let tables = &*self.inner;
729        validate_entity_id("node", id, tables.next_node_id)?;
730        self.order_properties(properties)?;
731        batch.labels.clear();
732        for &label in labels {
733            let number = *tables
734                .labels
735                .get(label as usize)
736                .ok_or_else(|| SnapshotError::Decode("invalid label dictionary id".into()))?;
737            if batch.labels.contains(&number) {
738                return Err(repeated_label(id));
739            }
740            batch.labels.push(number);
741            ids_for(&mut batch.by_label, number).push(id);
742        }
743        for &label in labels {
744            let name = &tables.label_names[label as usize];
745            batch.distinct.add_all(
746                name,
747                properties
748                    .iter()
749                    .map(|(key, value)| (&tables.key_names[*key as usize], value)),
750            );
751        }
752        for (key, _) in properties.iter_mut() {
753            *key = tables.keys[*key as usize];
754        }
755        encoded::encode_node_ids(&batch.labels, properties, &tables.dicts, &mut batch.scratch);
756        properties.clear();
757        push_record(&mut batch.records, id, &batch.scratch)
758    }
759
760    /// Encode a relationship: `rel_type` and the keys of `properties` are
761    /// numbers in the snapshot's [`SnapshotNames`].
762    pub fn push_relationship(
763        &self,
764        batch: &mut RelationshipBatch,
765        id: RelationshipId,
766        src: NodeId,
767        dst: NodeId,
768        rel_type: u32,
769        properties: &mut Vec<(u32, PropertyValue)>,
770    ) -> Result<(), SnapshotError> {
771        let tables = &*self.inner;
772        validate_entity_id("relationship", id, tables.next_rel_id)?;
773        validate_slot_id("relationship source node", src)?;
774        validate_slot_id("relationship target node", dst)?;
775        let type_id = *tables
776            .types
777            .get(rel_type as usize)
778            .ok_or_else(|| SnapshotError::Decode("invalid relationship type id".into()))?;
779        self.order_properties(properties)?;
780        batch.distinct.add_all(
781            &tables.type_names[rel_type as usize],
782            properties
783                .iter()
784                .map(|(key, value)| (&tables.key_names[*key as usize], value)),
785        );
786        for (key, _) in properties.iter_mut() {
787            *key = tables.keys[*key as usize];
788        }
789        encoded::encode_rel_ids(
790            src,
791            dst,
792            type_id,
793            properties,
794            &tables.dicts,
795            &mut batch.scratch,
796        );
797        properties.clear();
798        push_record(&mut batch.records, id, &batch.scratch)?;
799        batch.links.push((id, src, dst, type_id));
800        ids_for(&mut batch.by_type, type_id).push(id);
801        Ok(())
802    }
803}
804
805impl SnapshotGraphBuilder {
806    /// Number the snapshot's names in the graph being built, and return
807    /// an encoder that turns records named by those numbers into the
808    /// graph's stored form on any thread. Call before pushing records.
809    pub fn prime(&mut self, names: &SnapshotNames) -> Result<SnapshotEncoder, SnapshotError> {
810        let dicts = &mut self.graph.dicts;
811        let labels = names
812            .labels
813            .iter()
814            .map(|name| dicts.labels.id_or_insert(name))
815            .collect();
816        let types = names
817            .types
818            .iter()
819            .map(|name| dicts.types.id_or_insert(name))
820            .collect();
821        let keys: Vec<u32> = names
822            .keys
823            .iter()
824            .map(|name| dicts.keys.id_or_insert(name))
825            .collect();
826        let mut by_name: Vec<u32> = (0..names.keys.len() as u32).collect();
827        by_name.sort_by(|a, b| {
828            names.keys[*a as usize]
829                .as_str()
830                .cmp(names.keys[*b as usize].as_str())
831        });
832        let mut key_ranks = vec![0u32; names.keys.len()];
833        let mut rank = 0u32;
834        for (i, &key) in by_name.iter().enumerate() {
835            if i > 0
836                && names.keys[by_name[i - 1] as usize].as_str() != names.keys[key as usize].as_str()
837            {
838                rank += 1;
839            }
840            key_ranks[key as usize] = rank;
841        }
842        let keys_dict = &mut dicts.keys;
843        dicts
844            .values
845            .import(&names.values, |name| keys_dict.id_or_insert(name))
846            .map_err(SnapshotError::Decode)?;
847        Ok(SnapshotEncoder {
848            inner: Arc::new(EncoderTables {
849                dicts: dicts.clone(),
850                labels,
851                types,
852                keys,
853                key_names: names
854                    .keys
855                    .iter()
856                    .map(|name| name.as_arc().clone())
857                    .collect(),
858                key_ranks,
859                label_names: names.labels.clone(),
860                type_names: names.types.clone(),
861                next_node_id: self.graph.next_node_id,
862                next_rel_id: self.graph.next_rel_id,
863            }),
864        })
865    }
866
867    /// Store nodes an encoder from [`Self::prime`] encoded.
868    pub fn push_nodes(&mut self, batch: NodeBatch) -> Result<(), SnapshotError> {
869        self.graph
870            .put_encoded_node_leaves(batch.records)
871            .map_err(SnapshotError::Decode)?;
872        for (label, ids) in &batch.by_label {
873            self.graph.extend_node_label_index(*label, ids);
874        }
875        self.node_distinct.merge(batch.distinct);
876        Ok(())
877    }
878
879    /// Store relationships an encoder from [`Self::prime`] encoded. Their
880    /// endpoints must already be pushed.
881    pub fn push_relationships(&mut self, batch: RelationshipBatch) -> Result<(), SnapshotError> {
882        self.push_relationship_batches(vec![batch])
883    }
884
885    /// [`Self::push_relationships`] for several batches, in order: their
886    /// adjacency is linked in one pass, on every core when the build
887    /// allows.
888    pub fn push_relationship_batches(
889        &mut self,
890        batches: Vec<RelationshipBatch>,
891    ) -> Result<(), SnapshotError> {
892        let mut links = Vec::new();
893        for batch in batches {
894            self.graph
895                .put_encoded_relationship_leaves(batch.records, &batch.links)
896                .map_err(SnapshotError::Decode)?;
897            for (rel_type, ids) in &batch.by_type {
898                self.graph.extend_relationship_type_index(*rel_type, ids);
899            }
900            self.rel_distinct.merge(batch.distinct);
901            if links.is_empty() {
902                links = batch.links;
903            } else {
904                links.extend(batch.links);
905            }
906        }
907        self.graph.link_relationships(&links);
908        Ok(())
909    }
910}