Skip to main content

ursa_engine/database/
index_file.rs

1//! One stored graph index: `club.ursa.d/indexes/<hash>.idx`.
2//!
3//! A file holds one built index, the CSR topology and the id map, under the
4//! full key it was built for (the edge plan's fingerprint and the identity of
5//! every source file), followed by optional **chunks**: parts of the index
6//! that are built later, or that only some queries need, such as the
7//! in-adjacency. Version 2, all integers little-endian.
8//!
9//! The **main part** comes first:
10//!
11//! ```text
12//! offset  size  field
13//!      0     8  magic            "URSAIDX\0"
14//!      8     4  version          2
15//!     12     4  flags            bit 0: string ids
16//!     16     8  n_nodes
17//!     24     8  n_edges
18//!     32     8  key_len          bytes of the key that follows the header
19//!     40     4  n_sections
20//!     44     4  reserved (zero)
21//!     48     8  main_len         bytes of the main part, trailer included
22//!     56     8  n_dropped        edge rows left out (see below); zero before
23//!                                #311, and for a graph with no vertex table
24//!     64     …  key              JSON: {"fingerprint": …, "files": [{…}, …]}
25//!      …     …  section table    at the next multiple of 8; 24 bytes a section:
26//!                                kind u32, element width u32, offset u64, count u64
27//!      …     …  sections         each at a multiple of 64 from the file's start
28//!      …    16  trailer          at the next multiple of 8: main_len u64,
29//!                                CRC-32 of every byte before the trailer u32,
30//!                                "XDIU"
31//! ```
32//!
33//! Its sections, in this order: the out-adjacency's offsets (`u64`,
34//! `n_nodes + 1`), targets (`u32`, `n_edges`) and edge ids (`u32`, `n_edges`);
35//! then the ids, either `i64` values (`n_nodes`) or, for string ids,
36//! Arrow-style `i32` offsets (`n_nodes + 1`, from 0) and UTF-8 bytes.
37//!
38//! A graph with a vertex table leaves out the edge rows whose source or
39//! destination is not one of its keys (#311). When it left some out, the
40//! header counts them (`n_dropped`) and one more section (kind 11, bytes)
41//! says which rows of the edge plan were kept: one bit per row,
42//! `n_edges + n_dropped` of them, least significant bit first, as an Arrow
43//! boolean buffer. Its key includes the vertex table's, so no file written
44//! before #311 is ever read for such a graph; a file of a graph without one
45//! is read as before, with nothing left out.
46//!
47//! **Chunks** follow the main part, one after another, each appended in
48//! place once what it holds exists (see `IndexStore`):
49//!
50//! ```text
51//!      0    16  header           "CHNK", kind u32, end u64 (the chunk's end
52//!                                offset in the file)
53//!     16     …  sections         each at a multiple of 64 from the file's start
54//!      …     …  name             bytes naming what the chunk holds (may be empty)
55//!      …     …  section table    at the next multiple of 8, as in the main part
56//!      …    48  trailer          kind u32, n_sections u32, start u64,
57//!                                name offset u64, name length u64,
58//!                                data CRC u32 (the header and sections),
59//!                                meta CRC u32 (name, table and the 36 bytes
60//!                                of the trailer before it), "KNHC", zero u32
61//! ```
62//!
63//! Chunk kind 1 is the in-adjacency (no name; its offsets, targets and edge
64//! ids, as the out-adjacency's). Chunk kind 2 is a set of evaluated weights:
65//! its name is the weights' key (the weighted edge plan's fingerprint and the
66//! identity of the files it reads), and it holds one `f64` section (kind 10,
67//! `n_edges` values in edge-row order) per weight expression, in order. A reader of the main part never reads the
68//! chunks' arrays: it lists the chunks (their headers, names and tables, a
69//! few small reads), and reads a chunk's sections only when they are asked
70//! for, checking the main part's trailer and the chunk's data CRC first. The
71//! in-adjacency is read this way the first time an algorithm asks for it
72//! (`Topology::with_incoming_loader`).
73//!
74//! **Appending.** A chunk is written after the last chunk that checks out,
75//! so a chunk cut short by a crash is overwritten by the next append, and a
76//! reader stops listing at the first chunk that does not check out.
77//! Appending never changes the bytes of the main part or of an earlier
78//! chunk, so a reader in another process (`--readonly`) never sees them
79//! change.
80//!
81//! **Why raw arrays, not Arrow IPC.** The arrays have different lengths and
82//! no schema: an IPC stream would need one record batch per array, or padding
83//! them into one table, and its metadata says nothing a section table does
84//! not. Raw little-endian arrays at 64-byte-aligned offsets are what a later
85//! memory-mapped reader wants (each section is a ready `&[u32]`/`&[u64]` on a
86//! little-endian machine), need no dependency, and the string ids are already
87//! in Arrow's layout, so they become a `StringArray` without a copy once
88//! mapped. The main part is read once, front to back, decoding each section
89//! into its array as it goes ([`read_file`]).
90//!
91//! **Corruption.** The main part is read only if its length matches the
92//! trailer, the CRC matches, the version is 2, the key matches the key asked
93//! for, every section lies within it, and the arrays form a valid topology and
94//! id map (`ursa_core::Topology::from_parts`, `IdMap::from_user_ids`).
95//! Anything else is an error, which the store treats as a miss. A chunk that
96//! does not check out is ignored, and what it held is rebuilt when needed.
97//!
98//! **Version 1** files (the in-adjacency, if any, inside the main part) are
99//! not read: they are a miss like any other unreadable file, so the index is
100//! built again once and stored as version 2 in their place.
101
102use std::fs::{File, OpenOptions};
103use std::io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
104use std::path::Path;
105use std::sync::atomic::{AtomicU64, Ordering};
106use std::sync::Arc;
107
108use datafusion::arrow::array::{Array, Int64Array, StringArray};
109use datafusion::arrow::buffer::{BooleanBuffer, Buffer, OffsetBuffer, ScalarBuffer};
110use serde_json::{json, Value};
111use ursa_core::{Adjacency, IdMap, Topology};
112
113use crate::graph::FileIdentity;
114
115/// The first bytes of every index file.
116pub const MAGIC: [u8; 8] = *b"URSAIDX\0";
117/// The last bytes of every index file.
118const END_MAGIC: [u8; 4] = *b"XDIU";
119/// The index file version this build reads and writes.
120pub const VERSION: u32 = 2;
121/// The first bytes of every chunk.
122const CHUNK_MAGIC: [u8; 4] = *b"CHNK";
123/// The last bytes but four of every chunk's trailer.
124const CHUNK_END_MAGIC: [u8; 4] = *b"KNHC";
125const CHUNK_HEADER_LEN: u64 = 16;
126const CHUNK_TRAILER_LEN: u64 = 48;
127
128/// The chunk kind of an in-adjacency.
129pub const CHUNK_INCOMING: u32 = 1;
130/// The chunk kind of a set of evaluated weights.
131pub const CHUNK_WEIGHTS: u32 = 2;
132
133const HEADER_LEN: u64 = 64;
134const TRAILER_LEN: u64 = 16;
135const SECTION_ENTRY_LEN: u64 = 24;
136const SECTION_ALIGN: u64 = 64;
137
138const FLAG_STRING_IDS: u32 = 1;
139
140const OUT_OFFSETS: u32 = 1;
141const OUT_TARGETS: u32 = 2;
142const OUT_EDGE_IDS: u32 = 3;
143const IN_OFFSETS: u32 = 4;
144const IN_TARGETS: u32 = 5;
145const IN_EDGE_IDS: u32 = 6;
146const IDS_INT64: u32 = 7;
147const IDS_STR_OFFSETS: u32 = 8;
148const IDS_STR_BYTES: u32 = 9;
149const WEIGHTS_F64: u32 = 10;
150const KEPT_ROWS: u32 = 11;
151
152/// What an index is stored under: the edge plan's fingerprint and the
153/// identity of every file it reads.
154#[derive(Debug, Clone, PartialEq, Eq)]
155pub struct IndexKey {
156    pub fingerprint: String,
157    pub files: Vec<FileIdentity>,
158}
159
160impl IndexKey {
161    pub fn new(fingerprint: String, files: Vec<FileIdentity>) -> Self {
162        Self { fingerprint, files }
163    }
164
165    /// The key as JSON, as the file header and the manifest hold it.
166    pub fn to_json(&self) -> Value {
167        let files: Vec<Value> = self
168            .files
169            .iter()
170            .map(|f| {
171                json!({
172                    "store": f.store,
173                    "location": f.location,
174                    "size": f.size,
175                    "modified": f.modified,
176                    "e_tag": f.e_tag,
177                })
178            })
179            .collect();
180        json!({ "fingerprint": self.fingerprint, "files": files })
181    }
182
183    /// The key from its JSON, or `None` if it is not one.
184    pub fn from_json(value: &Value) -> Option<Self> {
185        let fingerprint = value.get("fingerprint")?.as_str()?.to_string();
186        let files = value
187            .get("files")?
188            .as_array()?
189            .iter()
190            .map(|f| {
191                Some(FileIdentity {
192                    store: f.get("store")?.as_str()?.to_string(),
193                    location: f.get("location")?.as_str()?.to_string(),
194                    size: f.get("size")?.as_u64()?,
195                    modified: f.get("modified")?.as_i64()?,
196                    e_tag: match f.get("e_tag")? {
197                        Value::Null => None,
198                        Value::String(s) => Some(s.clone()),
199                        _ => return None,
200                    },
201                })
202            })
203            .collect::<Option<Vec<_>>>()?;
204        Some(Self { fingerprint, files })
205    }
206
207    /// A stable 64-bit hash of the key (FNV-1a over its JSON), which names
208    /// its file. The file's header holds the full key, so a collision is
209    /// only a miss.
210    pub fn hash(&self) -> u64 {
211        let mut h: u64 = 0xcbf2_9ce4_8422_2325;
212        for b in self.to_json().to_string().bytes() {
213            h ^= u64::from(b);
214            h = h.wrapping_mul(0x0000_0100_0000_01b3);
215        }
216        h
217    }
218}
219
220/// One array of the file: borrowed from the topology being written, or a
221/// part of the id column (which the id map hands out as an `Arc`).
222enum Data<'a> {
223    U64(&'a [u64]),
224    F64(&'a [f64]),
225    U32(&'a [u32]),
226    /// The values of an `Int64Array`.
227    Ints(Int64Array),
228    /// The offsets of a `StringArray`, written less their first value so
229    /// they start at 0.
230    StrOffsets(StringArray),
231    /// The bytes of a `StringArray` its offsets span.
232    StrBytes(StringArray),
233    /// Bytes, as they are.
234    Bytes(Buffer),
235}
236
237impl Data<'_> {
238    fn width(&self) -> u64 {
239        match self {
240            Data::U64(_) | Data::F64(_) | Data::Ints(_) => 8,
241            Data::U32(_) | Data::StrOffsets(_) => 4,
242            Data::StrBytes(_) | Data::Bytes(_) => 1,
243        }
244    }
245
246    fn count(&self) -> u64 {
247        (match self {
248            Data::U64(v) => v.len(),
249            Data::F64(v) => v.len(),
250            Data::U32(v) => v.len(),
251            Data::Ints(a) => a.len(),
252            Data::StrOffsets(a) => a.len() + 1,
253            Data::StrBytes(a) => str_bytes(a).len(),
254            Data::Bytes(b) => b.len(),
255        }) as u64
256    }
257}
258
259/// The bytes of `strings` its offsets span.
260fn str_bytes(strings: &StringArray) -> &[u8] {
261    let offsets = strings.value_offsets();
262    let (first, last) = (offsets[0] as usize, offsets[offsets.len() - 1] as usize);
263    &strings.value_data()[first..last]
264}
265
266fn align(n: u64, to: u64) -> u64 {
267    n.div_ceil(to) * to
268}
269
270/// An index laid out for writing: its header fields, key, sections and where
271/// each goes.
272pub struct Encoded<'a> {
273    flags: u32,
274    n_nodes: u64,
275    n_edges: u64,
276    n_dropped: u64,
277    key: Vec<u8>,
278    sections: Vec<(u32, Data<'a>, u64)>,
279    table_offset: u64,
280    trailer_offset: u64,
281}
282
283impl<'a> Encoded<'a> {
284    /// Lay out the main part of `topology` and `ids` under `key` (not the
285    /// in-adjacency: see [`EncodedChunk::incoming`]). Nothing is written yet.
286    pub fn new(key: &IndexKey, topology: &'a Topology, ids: &'a IdMap) -> io::Result<Self> {
287        Self::with_dropped(key, topology, ids, 0, None)
288    }
289
290    /// [`Self::new`], for an index that left out `n_dropped` edge rows,
291    /// `kept` saying which it kept (see the module docs).
292    pub fn with_dropped(
293        key: &IndexKey,
294        topology: &'a Topology,
295        ids: &'a IdMap,
296        n_dropped: u64,
297        kept: Option<&BooleanBuffer>,
298    ) -> io::Result<Self> {
299        let id_array = ids.user_id_array();
300        let mut flags = 0;
301        let out = topology.out();
302        let mut sections: Vec<(u32, Data<'a>)> = vec![
303            (OUT_OFFSETS, Data::U64(&out.offsets)),
304            (OUT_TARGETS, Data::U32(&out.targets)),
305            (OUT_EDGE_IDS, Data::U32(&out.edge_ids)),
306        ];
307        // Cloning an Arrow array clones `Arc`s to its buffers, not the ids.
308        if let Some(ints) = id_array.as_any().downcast_ref::<Int64Array>() {
309            sections.push((IDS_INT64, Data::Ints(ints.clone())));
310        } else if let Some(strings) = id_array.as_any().downcast_ref::<StringArray>() {
311            flags |= FLAG_STRING_IDS;
312            sections.push((IDS_STR_OFFSETS, Data::StrOffsets(strings.clone())));
313            sections.push((IDS_STR_BYTES, Data::StrBytes(strings.clone())));
314        } else {
315            return Err(io::Error::other(format!(
316                "cannot store node ids of type {}",
317                id_array.data_type()
318            )));
319        }
320        let n_dropped = match kept {
321            Some(kept) if n_dropped > 0 => {
322                sections.push((KEPT_ROWS, Data::Bytes(kept.sliced())));
323                n_dropped
324            }
325            _ => 0,
326        };
327
328        let key = key.to_json().to_string().into_bytes();
329        let table_offset = align(HEADER_LEN + key.len() as u64, 8);
330        let mut at = table_offset + SECTION_ENTRY_LEN * sections.len() as u64;
331        let sections = sections
332            .into_iter()
333            .map(|(kind, data)| {
334                at = align(at, SECTION_ALIGN);
335                let offset = at;
336                at += data.width() * data.count();
337                (kind, data, offset)
338            })
339            .collect();
340        Ok(Self {
341            flags,
342            n_nodes: topology.n_nodes() as u64,
343            n_edges: topology.n_edges() as u64,
344            n_dropped,
345            key,
346            sections,
347            table_offset,
348            trailer_offset: align(at, 8),
349        })
350    }
351
352    /// The size of the main part, in bytes.
353    pub fn file_size(&self) -> u64 {
354        self.trailer_offset + TRAILER_LEN
355    }
356
357    /// Write the main part to `out`. Returns its identity.
358    pub fn write_to(&self, out: impl Write) -> io::Result<MainPart> {
359        let mut sink = Sink {
360            out,
361            crc: crc32fast::Hasher::new(),
362            at: 0,
363        };
364        let mut header = Vec::with_capacity(HEADER_LEN as usize);
365        header.extend_from_slice(&MAGIC);
366        header.extend_from_slice(&VERSION.to_le_bytes());
367        header.extend_from_slice(&self.flags.to_le_bytes());
368        header.extend_from_slice(&self.n_nodes.to_le_bytes());
369        header.extend_from_slice(&self.n_edges.to_le_bytes());
370        header.extend_from_slice(&(self.key.len() as u64).to_le_bytes());
371        header.extend_from_slice(&(self.sections.len() as u32).to_le_bytes());
372        header.extend_from_slice(&[0; 4]);
373        header.extend_from_slice(&self.file_size().to_le_bytes());
374        header.extend_from_slice(&self.n_dropped.to_le_bytes());
375        header.resize(HEADER_LEN as usize, 0);
376        sink.put(&header)?;
377        sink.put(&self.key)?;
378        sink.pad_to(self.table_offset)?;
379        for (kind, data, offset) in &self.sections {
380            sink.put(&section_entry(*kind, data, *offset))?;
381        }
382        for (_, data, offset) in &self.sections {
383            sink.pad_to(*offset)?;
384            sink.put_data(data)?;
385        }
386        sink.pad_to(self.trailer_offset)?;
387        let crc = sink.crc.clone().finalize();
388        let mut trailer = Vec::with_capacity(TRAILER_LEN as usize);
389        trailer.extend_from_slice(&self.file_size().to_le_bytes());
390        trailer.extend_from_slice(&crc.to_le_bytes());
391        trailer.extend_from_slice(&END_MAGIC);
392        sink.put(&trailer)?;
393        sink.out.flush()?;
394        Ok(MainPart {
395            len: self.file_size(),
396            crc,
397        })
398    }
399}
400
401/// The identity of a file's main part: its length and CRC. A chunk is read
402/// from, or appended to, a file only while its main part is still this one.
403#[derive(Debug, Clone, Copy, PartialEq, Eq)]
404pub struct MainPart {
405    pub len: u64,
406    pub crc: u32,
407}
408
409/// A chunk laid out for writing at a given offset of a file.
410pub struct EncodedChunk<'a> {
411    kind: u32,
412    name: Vec<u8>,
413    start: u64,
414    sections: Vec<(u32, Data<'a>, u64)>,
415    name_offset: u64,
416    table_offset: u64,
417}
418
419impl<'a> EncodedChunk<'a> {
420    /// The chunk of the in-adjacency `incoming`, to be written at `start`.
421    pub fn incoming(incoming: &'a Adjacency, start: u64) -> Self {
422        Self::new(
423            CHUNK_INCOMING,
424            Vec::new(),
425            vec![
426                (IN_OFFSETS, Data::U64(&incoming.offsets)),
427                (IN_TARGETS, Data::U32(&incoming.targets)),
428                (IN_EDGE_IDS, Data::U32(&incoming.edge_ids)),
429            ],
430            start,
431        )
432    }
433
434    /// The chunk of the weight columns `columns` (one value per edge row
435    /// each), stored under `key`, to be written at `start`.
436    pub fn weights(key: &str, columns: &[&'a [f64]], start: u64) -> Self {
437        Self::new(
438            CHUNK_WEIGHTS,
439            key.as_bytes().to_vec(),
440            columns
441                .iter()
442                .map(|column| (WEIGHTS_F64, Data::F64(column)))
443                .collect(),
444            start,
445        )
446    }
447
448    fn new(kind: u32, name: Vec<u8>, sections: Vec<(u32, Data<'a>)>, start: u64) -> Self {
449        let mut at = start + CHUNK_HEADER_LEN;
450        let sections = sections
451            .into_iter()
452            .map(|(kind, data)| {
453                at = align(at, SECTION_ALIGN);
454                let offset = at;
455                at += data.width() * data.count();
456                (kind, data, offset)
457            })
458            .collect();
459        let name_offset = at;
460        let table_offset = align(name_offset + name.len() as u64, 8);
461        Self {
462            kind,
463            name,
464            start,
465            sections,
466            name_offset,
467            table_offset,
468        }
469    }
470
471    fn trailer_offset(&self) -> u64 {
472        self.table_offset + SECTION_ENTRY_LEN * self.sections.len() as u64
473    }
474
475    /// The offset just past the chunk: the file's length once it is written.
476    pub fn end(&self) -> u64 {
477        self.trailer_offset() + CHUNK_TRAILER_LEN
478    }
479
480    /// The chunk's size, in bytes.
481    pub fn size(&self) -> u64 {
482        self.end() - self.start
483    }
484
485    /// Write the chunk to `out`, which is at its start.
486    pub fn write_to(&self, out: impl Write) -> io::Result<()> {
487        let mut sink = Sink {
488            out,
489            crc: crc32fast::Hasher::new(),
490            at: self.start,
491        };
492        let mut header = Vec::with_capacity(CHUNK_HEADER_LEN as usize);
493        header.extend_from_slice(&CHUNK_MAGIC);
494        header.extend_from_slice(&self.kind.to_le_bytes());
495        header.extend_from_slice(&self.end().to_le_bytes());
496        sink.put(&header)?;
497        for (_, data, offset) in &self.sections {
498            sink.pad_to(*offset)?;
499            sink.put_data(data)?;
500        }
501        sink.pad_to(self.name_offset)?;
502        let data_crc = std::mem::replace(&mut sink.crc, crc32fast::Hasher::new()).finalize();
503        sink.put(&self.name)?;
504        sink.pad_to(self.table_offset)?;
505        for (kind, data, offset) in &self.sections {
506            sink.put(&section_entry(*kind, data, *offset))?;
507        }
508        let mut trailer = Vec::with_capacity(CHUNK_TRAILER_LEN as usize);
509        trailer.extend_from_slice(&self.kind.to_le_bytes());
510        trailer.extend_from_slice(&(self.sections.len() as u32).to_le_bytes());
511        trailer.extend_from_slice(&self.start.to_le_bytes());
512        trailer.extend_from_slice(&self.name_offset.to_le_bytes());
513        trailer.extend_from_slice(&(self.name.len() as u64).to_le_bytes());
514        trailer.extend_from_slice(&data_crc.to_le_bytes());
515        sink.crc.update(&trailer);
516        let meta_crc = sink.crc.clone().finalize();
517        trailer.extend_from_slice(&meta_crc.to_le_bytes());
518        trailer.extend_from_slice(&CHUNK_END_MAGIC);
519        trailer.extend_from_slice(&[0; 4]);
520        sink.out.write_all(&trailer)?;
521        sink.out.flush()
522    }
523}
524
525/// A section table entry.
526fn section_entry(kind: u32, data: &Data<'_>, offset: u64) -> Vec<u8> {
527    let mut entry = Vec::with_capacity(SECTION_ENTRY_LEN as usize);
528    entry.extend_from_slice(&kind.to_le_bytes());
529    entry.extend_from_slice(&(data.width() as u32).to_le_bytes());
530    entry.extend_from_slice(&offset.to_le_bytes());
531    entry.extend_from_slice(&data.count().to_le_bytes());
532    entry
533}
534
535/// A writer that keeps a running CRC and position.
536struct Sink<W> {
537    out: W,
538    crc: crc32fast::Hasher,
539    at: u64,
540}
541
542impl<W: Write> Sink<W> {
543    fn put(&mut self, bytes: &[u8]) -> io::Result<()> {
544        self.crc.update(bytes);
545        self.at += bytes.len() as u64;
546        self.out.write_all(bytes)
547    }
548
549    fn pad_to(&mut self, offset: u64) -> io::Result<()> {
550        debug_assert!(offset >= self.at);
551        let zeros = [0u8; SECTION_ALIGN as usize];
552        while self.at < offset {
553            let n = ((offset - self.at) as usize).min(zeros.len());
554            self.put(&zeros[..n])?;
555        }
556        Ok(())
557    }
558
559    /// Write one section's array.
560    fn put_data(&mut self, data: &Data<'_>) -> io::Result<()> {
561        match data {
562            Data::U64(v) => self.put_all(v, |x| x.to_le_bytes()),
563            Data::F64(v) => self.put_all(v, |x| x.to_le_bytes()),
564            Data::U32(v) => self.put_all(v, |x| x.to_le_bytes()),
565            Data::Ints(a) => self.put_all(a.values(), |x| x.to_le_bytes()),
566            Data::StrOffsets(a) => {
567                let offsets = a.value_offsets();
568                let base = offsets[0];
569                self.put_all(offsets, |x| (x - base).to_le_bytes())
570            }
571            Data::StrBytes(a) => self.put(str_bytes(a)),
572            Data::Bytes(b) => self.put(b.as_slice()),
573        }
574    }
575
576    /// Write `values` little-endian, a block at a time.
577    fn put_all<T: Copy, const N: usize>(
578        &mut self,
579        values: &[T],
580        le: impl Fn(T) -> [u8; N],
581    ) -> io::Result<()> {
582        let mut block = Vec::with_capacity(64 * 1024);
583        for chunk in values.chunks(64 * 1024 / N) {
584            block.clear();
585            for &v in chunk {
586                block.extend_from_slice(&le(v));
587            }
588            self.put(&block)?;
589        }
590        Ok(())
591    }
592}
593
594/// Why a file could not be read as an index.
595pub type DecodeError = String;
596
597/// The edge rows an index left out, and which it kept: nothing left out and
598/// `None` for a graph with no vertex table (see the module docs).
599#[derive(Debug, Clone, Default)]
600pub struct Dropped {
601    pub n_dropped: u64,
602    pub kept: Option<BooleanBuffer>,
603}
604
605/// An index read from its file.
606pub struct Stored {
607    /// With a loader for the in-adjacency, if the file has its chunk.
608    pub topology: Topology,
609    pub ids: IdMap,
610    /// The edge rows left out, and which were kept (see the module docs).
611    pub dropped: Dropped,
612    pub main: MainPart,
613    /// The chunks that check out, in file order.
614    pub chunks: Vec<ChunkInfo>,
615}
616
617/// Whether `chunks` has one of `kind` named `name`.
618pub fn has_chunk(chunks: &[ChunkInfo], kind: u32, name: &[u8]) -> bool {
619    chunks.iter().any(|c| c.kind == kind && c.name == name)
620}
621
622/// How many in-adjacencies have been read from index files, in this process.
623static INCOMING_READS: AtomicU64 = AtomicU64::new(0);
624
625/// How many in-adjacencies this process has read from index files: for
626/// tests that an algorithm read one, or did not.
627pub fn incoming_reads() -> u64 {
628    INCOMING_READS.load(Ordering::Relaxed)
629}
630
631/// Read the index file at `path`, which must be stored under `key`.
632///
633/// The main part is read once, front to back, straight into the arrays the
634/// index is made of, so a load holds the index and a small buffer, not the
635/// file's bytes as well. The key is checked before the sections are read, so
636/// a file for another key costs only its header; the checksum covers every
637/// byte of the main part and is checked before any array is used. The chunks
638/// are only listed; the in-adjacency's is read when the topology is first
639/// asked for it.
640pub fn read_file(path: &Path, key: &IndexKey) -> Result<Stored, DecodeError> {
641    let mut file = File::open(path).map_err(|e| e.to_string())?;
642    let len = file.metadata().map_err(|e| e.to_string())?.len();
643    let (topology, ids, main, dropped) =
644        decode_from(BufReader::with_capacity(BLOCK, &file), len, key)?;
645    let (chunks, _) = list_chunks(&mut file, main.len, len);
646    let topology = match chunks.iter().find(|c| c.kind == CHUNK_INCOMING) {
647        Some(chunk) => {
648            let (path, chunk) = (path.to_path_buf(), chunk.clone());
649            let (n_nodes, n_edges) = (topology.n_nodes(), topology.n_edges());
650            topology.with_incoming_loader(Box::new(move || {
651                let incoming = read_incoming(&path, main, &chunk, n_nodes, n_edges).ok()?;
652                INCOMING_READS.fetch_add(1, Ordering::Relaxed);
653                Some(incoming)
654            }))
655        }
656        None => topology,
657    };
658    Ok(Stored {
659        topology,
660        ids,
661        dropped,
662        main,
663        chunks,
664    })
665}
666
667/// Read the main part of the index in `bytes`, which must be stored under
668/// `key`. Chunks are not read.
669pub fn decode(bytes: &[u8], key: &IndexKey) -> Result<(Topology, IdMap), DecodeError> {
670    decode_from(bytes, bytes.len() as u64, key).map(|(topology, ids, _, _)| (topology, ids))
671}
672
673/// One chunk of a file, as listed: where it is and what it holds.
674#[derive(Debug, Clone, PartialEq, Eq)]
675pub struct ChunkInfo {
676    pub kind: u32,
677    pub name: Vec<u8>,
678    pub start: u64,
679    pub end: u64,
680    /// The end of its header and sections, which the data CRC covers.
681    data_end: u64,
682    data_crc: u32,
683    /// (kind, width, offset, count), in file order.
684    sections: Vec<(u32, u64, u64, u64)>,
685}
686
687/// The chunks of a file of `len` bytes whose main part is `main_len` bytes,
688/// in file order, as far as they check out; and the offset where the next
689/// chunk goes (the end of the last chunk that checks out). Only the chunks'
690/// headers, names, tables and trailers are read.
691pub fn list_chunks<R: Read + Seek>(file: &mut R, main_len: u64, len: u64) -> (Vec<ChunkInfo>, u64) {
692    let mut chunks = Vec::new();
693    let mut at = main_len;
694    while let Some(chunk) = chunk_at(file, at, len) {
695        at = chunk.end;
696        chunks.push(chunk);
697    }
698    (chunks, at)
699}
700
701fn u32_at(bytes: &[u8], at: usize) -> u32 {
702    u32::from_le_bytes(bytes[at..at + 4].try_into().expect("4 bytes"))
703}
704
705/// The chunk at `start` of a file of `len` bytes, if one is there whole and
706/// its header, name, table and trailer check out.
707fn chunk_at<R: Read + Seek>(file: &mut R, start: u64, len: u64) -> Option<ChunkInfo> {
708    let read = |file: &mut R, at: u64, n: u64| -> Option<Vec<u8>> {
709        file.seek(SeekFrom::Start(at)).ok()?;
710        let mut buf = vec![0; usize::try_from(n).ok()?];
711        file.read_exact(&mut buf).ok()?;
712        Some(buf)
713    };
714    if len.saturating_sub(start) < CHUNK_HEADER_LEN + CHUNK_TRAILER_LEN {
715        return None;
716    }
717    let header = read(file, start, CHUNK_HEADER_LEN)?;
718    let end = u64_at(&header, 8);
719    if header[..4] != CHUNK_MAGIC || end > len || end < start + CHUNK_HEADER_LEN + CHUNK_TRAILER_LEN
720    {
721        return None;
722    }
723    let kind = u32_at(&header, 4);
724    let trailer_offset = end - CHUNK_TRAILER_LEN;
725    let trailer = read(file, trailer_offset, CHUNK_TRAILER_LEN)?;
726    let n_sections = u32_at(&trailer, 4) as u64;
727    let name_offset = u64_at(&trailer, 16);
728    let name_len = u64_at(&trailer, 24);
729    let table_offset = align(name_offset.checked_add(name_len)?, 8);
730    if trailer[40..44] != CHUNK_END_MAGIC
731        || u32_at(&trailer, 0) != kind
732        || u64_at(&trailer, 8) != start
733        || name_offset < start + CHUNK_HEADER_LEN
734        || n_sections
735            .checked_mul(SECTION_ENTRY_LEN)?
736            .checked_add(table_offset)?
737            != trailer_offset
738    {
739        return None;
740    }
741    let meta = read(file, name_offset, trailer_offset - name_offset)?;
742    let mut crc = crc32fast::Hasher::new();
743    crc.update(&meta);
744    crc.update(&trailer[..36]);
745    if crc.finalize() != u32_at(&trailer, 36) {
746        return None;
747    }
748    let name = meta[..name_len as usize].to_vec();
749    let table = &meta[(table_offset - name_offset) as usize..];
750    let mut sections = Vec::new();
751    for entry in table.chunks_exact(SECTION_ENTRY_LEN as usize) {
752        let (kind, width) = (u32_at(entry, 0), u32_at(entry, 4) as u64);
753        let (offset, count) = (u64_at(entry, 8), u64_at(entry, 16));
754        width
755            .checked_mul(count)
756            .and_then(|n| n.checked_add(offset))
757            .filter(|&e| offset >= start + CHUNK_HEADER_LEN && e <= name_offset)?;
758        sections.push((kind, width, offset, count));
759    }
760    sections.sort_by_key(|&(_, _, offset, _)| offset);
761    if sections
762        .windows(2)
763        .any(|w| w[0].2 + w[0].1 * w[0].3 > w[1].2)
764    {
765        return None;
766    }
767    Some(ChunkInfo {
768        kind,
769        name,
770        start,
771        end,
772        data_end: name_offset,
773        data_crc: u32_at(&trailer, 32),
774        sections,
775    })
776}
777
778/// Check that the file `file` still has the main part `main`.
779fn check_main(file: &mut File, main: MainPart) -> Result<(), DecodeError> {
780    let len = file.metadata().map_err(|e| e.to_string())?.len();
781    if len < main.len || main.len < HEADER_LEN + TRAILER_LEN {
782        return Err("the file's main part has changed".into());
783    }
784    let mut trailer = [0u8; TRAILER_LEN as usize];
785    file.seek(SeekFrom::Start(main.len - TRAILER_LEN))
786        .and_then(|_| file.read_exact(&mut trailer))
787        .map_err(|e| e.to_string())?;
788    if u64_at(&trailer, 0) != main.len || u32_at(&trailer, 8) != main.crc {
789        return Err("the file's main part has changed".into());
790    }
791    Ok(())
792}
793
794/// Read the sections of `chunk` from the file at `path`, whose main part
795/// must still be `main`. Every byte the data CRC covers is read and checked.
796pub fn read_chunk(
797    path: &Path,
798    main: MainPart,
799    chunk: &ChunkInfo,
800) -> Result<Vec<(u32, Section)>, DecodeError> {
801    let mut file = File::open(path).map_err(|e| e.to_string())?;
802    check_main(&mut file, main)?;
803    if file.metadata().map_err(|e| e.to_string())?.len() < chunk.end {
804        return Err("the file is shorter than the chunk".into());
805    }
806    file.seek(SeekFrom::Start(chunk.start))
807        .map_err(|e| e.to_string())?;
808    let mut src = Source {
809        inner: BufReader::with_capacity(BLOCK, file),
810        crc: crc32fast::Hasher::new(),
811        at: chunk.start,
812    };
813    let mut sections = Vec::with_capacity(chunk.sections.len());
814    for &(kind, width, offset, count) in &chunk.sections {
815        src.skip_to(offset)?;
816        sections.push((kind, src.section(kind, width, offset, count)?));
817    }
818    src.skip_to(chunk.data_end)?;
819    if src.crc.finalize() != chunk.data_crc {
820        return Err("the chunk's checksum does not match".into());
821    }
822    Ok(sections)
823}
824
825/// Read the in-adjacency in `chunk` of the file at `path`, for a topology
826/// of `n_nodes` nodes and `n_edges` edges.
827fn read_incoming(
828    path: &Path,
829    main: MainPart,
830    chunk: &ChunkInfo,
831    n_nodes: usize,
832    n_edges: usize,
833) -> Result<Adjacency, DecodeError> {
834    let mut sections = read_chunk(path, main, chunk)?;
835    let mut take = |kind: u32| {
836        let i = sections
837            .iter()
838            .position(|(k, _)| *k == kind)
839            .ok_or_else(|| format!("section {kind} is missing"))?;
840        Ok::<_, DecodeError>(sections.swap_remove(i).1)
841    };
842    match (take(IN_OFFSETS)?, take(IN_TARGETS)?, take(IN_EDGE_IDS)?) {
843        (Section::U64(offsets), Section::U32(targets), Section::U32(edge_ids))
844            if offsets.len() == n_nodes + 1
845                && targets.len() == n_edges
846                && edge_ids.len() == n_edges =>
847        {
848            Ok(Adjacency {
849                offsets,
850                targets,
851                edge_ids,
852            })
853        }
854        _ => Err("the in-adjacency is malformed".into()),
855    }
856}
857
858/// Read the weight columns in `chunk` of the file at `path`, for a topology
859/// of `n_edges` edges: one per section, in order.
860pub fn read_weights(
861    path: &Path,
862    main: MainPart,
863    chunk: &ChunkInfo,
864    n_edges: usize,
865) -> Result<Vec<Vec<f64>>, DecodeError> {
866    read_chunk(path, main, chunk)?
867        .into_iter()
868        .map(|(_, section)| match section {
869            Section::F64(values) if values.len() == n_edges => Ok(values),
870            _ => Err("the weights are malformed".to_string()),
871        })
872        .collect()
873}
874
875/// Append the chunk `encode` lays out (at the offset it is given) to the
876/// file at `path`, whose main part must still be `main`, after its last
877/// chunk that checks out; unless the file already has a chunk of `kind`
878/// named `name`. Returns the file's new length, or `None` if the chunk was
879/// there already.
880pub fn append_chunk<'a>(
881    path: &Path,
882    main: MainPart,
883    kind: u32,
884    name: &[u8],
885    encode: impl FnOnce(u64) -> EncodedChunk<'a>,
886) -> io::Result<Option<u64>> {
887    let mut file = OpenOptions::new().read(true).write(true).open(path)?;
888    check_main(&mut file, main).map_err(io::Error::other)?;
889    let len = file.metadata()?.len();
890    let (chunks, at) = list_chunks(&mut file, main.len, len);
891    if has_chunk(&chunks, kind, name) {
892        return Ok(None);
893    }
894    let chunk = encode(at);
895    debug_assert_eq!((chunk.kind, chunk.name.as_slice()), (kind, name));
896    // What follows the last chunk that checks out (one cut short by a
897    // crash) is dropped first, so the file never ends in bytes that do not.
898    if len > at {
899        file.set_len(at)?;
900    }
901    file.seek(SeekFrom::Start(at))?;
902    let mut out = BufWriter::with_capacity(1 << 20, &file);
903    chunk.write_to(&mut out)?;
904    out.into_inner().map_err(|e| e.into_error())?;
905    file.sync_all()?;
906    Ok(Some(chunk.end()))
907}
908
909/// The most a read asks for at once, in bytes.
910const BLOCK: usize = 1 << 20;
911
912/// A reader that keeps a running CRC and position.
913struct Source<R> {
914    inner: R,
915    crc: crc32fast::Hasher,
916    at: u64,
917}
918
919impl<R: Read> Source<R> {
920    fn fill(&mut self, buf: &mut [u8]) -> Result<(), DecodeError> {
921        self.inner.read_exact(buf).map_err(|e| match e.kind() {
922            io::ErrorKind::UnexpectedEof => "the file is truncated".to_string(),
923            _ => e.to_string(),
924        })?;
925        self.crc.update(buf);
926        self.at += buf.len() as u64;
927        Ok(())
928    }
929
930    fn bytes(&mut self, n: usize) -> Result<Vec<u8>, DecodeError> {
931        let mut buf = vec![0; n];
932        self.fill(&mut buf)?;
933        Ok(buf)
934    }
935
936    /// Read and discard up to `offset`: padding, or a section this build
937    /// does not read.
938    fn skip_to(&mut self, offset: u64) -> Result<(), DecodeError> {
939        let mut buf = vec![0; BLOCK.min(offset.saturating_sub(self.at) as usize)];
940        while self.at < offset {
941            let n = ((offset - self.at) as usize).min(buf.len());
942            self.fill(&mut buf[..n])?;
943        }
944        Ok(())
945    }
946
947    /// `count` little-endian values, a block at a time.
948    fn values<T, const N: usize>(
949        &mut self,
950        count: usize,
951        le: fn([u8; N]) -> T,
952    ) -> Result<Vec<T>, DecodeError> {
953        let mut out = Vec::with_capacity(count);
954        let mut buf = vec![0; BLOCK.min(count * N)];
955        while out.len() < count {
956            let n = ((count - out.len()) * N).min(buf.len());
957            self.fill(&mut buf[..n])?;
958            out.extend(read(&buf[..n], le));
959        }
960        Ok(out)
961    }
962}
963
964impl<R: Read> Source<R> {
965    /// The section of `kind` at `offset` (where the source is): `count`
966    /// values of `width` bytes. A section this build does not read, or a
967    /// known one of the wrong width, is skipped and comes back empty.
968    fn section(
969        &mut self,
970        kind: u32,
971        width: u64,
972        offset: u64,
973        count: u64,
974    ) -> Result<Section, DecodeError> {
975        let n = usize::try_from(count).map_err(|e| e.to_string())?;
976        Ok(match (kind, width) {
977            (OUT_OFFSETS | IN_OFFSETS, 8) => Section::U64(self.values(n, u64::from_le_bytes)?),
978            (OUT_TARGETS | OUT_EDGE_IDS | IN_TARGETS | IN_EDGE_IDS, 4) => {
979                Section::U32(self.values(n, u32::from_le_bytes)?)
980            }
981            (IDS_INT64, 8) => Section::I64(self.values(n, i64::from_le_bytes)?),
982            (WEIGHTS_F64, 8) => Section::F64(self.values(n, f64::from_le_bytes)?),
983            (IDS_STR_OFFSETS, 4) => Section::I32(self.values(n, i32::from_le_bytes)?),
984            (IDS_STR_BYTES | KEPT_ROWS, 1) => Section::Bytes(self.bytes(n)?),
985            _ => {
986                self.skip_to(offset + width * count)?;
987                Section::Bytes(Vec::new())
988            }
989        })
990    }
991}
992
993/// One section, read.
994pub enum Section {
995    U64(Vec<u64>),
996    F64(Vec<f64>),
997    U32(Vec<u32>),
998    I64(Vec<i64>),
999    I32(Vec<i32>),
1000    Bytes(Vec<u8>),
1001}
1002
1003/// Read the index from `reader`, a file of `len` bytes stored under `key`.
1004fn decode_from(
1005    reader: impl Read,
1006    len: u64,
1007    key: &IndexKey,
1008) -> Result<(Topology, IdMap, MainPart, Dropped), DecodeError> {
1009    if len < HEADER_LEN + TRAILER_LEN {
1010        return Err(format!("{len} bytes is too short for an index file"));
1011    }
1012    let mut src = Source {
1013        inner: reader,
1014        crc: crc32fast::Hasher::new(),
1015        at: 0,
1016    };
1017    let header = src.bytes(HEADER_LEN as usize)?;
1018    if header[..8] != MAGIC {
1019        return Err("not an Ursa index file".into());
1020    }
1021    let version = u32::from_le_bytes(header[8..12].try_into().expect("4 bytes"));
1022    if version != VERSION {
1023        return Err(format!(
1024            "index format version {version}; this build reads {VERSION}"
1025        ));
1026    }
1027    let flags = u32::from_le_bytes(header[12..16].try_into().expect("4 bytes"));
1028    let n_nodes = usize::try_from(u64_at(&header, 16)).map_err(|e| e.to_string())?;
1029    let n_edges = usize::try_from(u64_at(&header, 24)).map_err(|e| e.to_string())?;
1030    let key_len = u64_at(&header, 32);
1031    let n_sections = u32::from_le_bytes(header[40..44].try_into().expect("4 bytes")) as u64;
1032    let main_len = u64_at(&header, 48);
1033    let n_dropped = u64_at(&header, 56);
1034    if main_len > len || main_len < HEADER_LEN + TRAILER_LEN {
1035        return Err("the file is truncated".into());
1036    }
1037
1038    let body = main_len - TRAILER_LEN;
1039    let key_end = HEADER_LEN
1040        .checked_add(key_len)
1041        .filter(|&end| end <= body)
1042        .ok_or("the key runs past the end of the file")?;
1043    let stored = src.bytes((key_end - HEADER_LEN) as usize)?;
1044    let stored: Value =
1045        serde_json::from_slice(&stored).map_err(|e| format!("the key is not JSON: {e}"))?;
1046    if IndexKey::from_json(&stored).as_ref() != Some(key) {
1047        return Err("the file holds an index for another key".into());
1048    }
1049
1050    let table = align(key_end, 8);
1051    let table_end = n_sections
1052        .checked_mul(SECTION_ENTRY_LEN)
1053        .and_then(|n| n.checked_add(table))
1054        .filter(|&end| end <= body)
1055        .ok_or("the section table runs past the end of the file")?;
1056    src.skip_to(table)?;
1057    let entries = src.bytes((table_end - table) as usize)?;
1058    // (kind, width, offset, count), in file order. Sections are read in one
1059    // pass, so they may not overlap.
1060    let mut layout: Vec<(u32, u64, u64, u64)> = Vec::new();
1061    for entry in entries.chunks_exact(SECTION_ENTRY_LEN as usize) {
1062        let kind = u32::from_le_bytes(entry[0..4].try_into().expect("4 bytes"));
1063        let width = u32::from_le_bytes(entry[4..8].try_into().expect("4 bytes")) as u64;
1064        let offset = u64_at(entry, 8);
1065        let count = u64_at(entry, 16);
1066        width
1067            .checked_mul(count)
1068            .and_then(|n| n.checked_add(offset))
1069            .filter(|&end| offset >= table_end && end <= body)
1070            .ok_or("a section runs past the end of the file")?;
1071        layout.push((kind, width, offset, count));
1072    }
1073    layout.sort_by_key(|&(_, _, offset, _)| offset);
1074    if layout.windows(2).any(|w| w[0].2 + w[0].1 * w[0].3 > w[1].2) {
1075        return Err("the sections overlap".into());
1076    }
1077
1078    let mut sections: Vec<(u32, u64, u64, Section)> = Vec::new();
1079    for (kind, width, offset, count) in layout {
1080        src.skip_to(offset)?;
1081        // A section this build does not read, or a known one of the wrong
1082        // width (malformed, below), comes back empty.
1083        let data = src.section(kind, width, offset, count)?;
1084        sections.push((kind, width, count, data));
1085    }
1086    src.skip_to(body)?;
1087    let crc = src.crc.clone().finalize();
1088    let mut trailer = [0u8; TRAILER_LEN as usize];
1089    src.fill(&mut trailer)?;
1090    if trailer[12..] != END_MAGIC || u64_at(&trailer, 0) != main_len {
1091        return Err("the file is truncated".into());
1092    }
1093    if u32::from_le_bytes(trailer[8..12].try_into().expect("4 bytes")) != crc {
1094        return Err("the checksum does not match".into());
1095    }
1096
1097    let main = MainPart { len: main_len, crc };
1098    let mut take = |kind: u32, width: u64, count: Option<usize>| -> Result<Section, DecodeError> {
1099        let mut found = sections.iter().filter(|(k, _, _, _)| *k == kind);
1100        let (_, w, c, _) = found
1101            .next()
1102            .ok_or_else(|| format!("section {kind} is missing"))?;
1103        if found.next().is_some() || *w != width {
1104            return Err(format!("section {kind} is malformed"));
1105        }
1106        if count.is_some_and(|count| *c != count as u64) {
1107            return Err(format!("section {kind} has the wrong length"));
1108        }
1109        let i = sections
1110            .iter()
1111            .position(|(k, _, _, _)| *k == kind)
1112            .expect("found above");
1113        Ok(sections.swap_remove(i).3)
1114    };
1115    let malformed = |kind: u32| format!("section {kind} is malformed");
1116    let mut adjacency = |kinds: [u32; 3]| -> Result<Adjacency, DecodeError> {
1117        match (
1118            take(kinds[0], 8, Some(n_nodes + 1))?,
1119            take(kinds[1], 4, Some(n_edges))?,
1120            take(kinds[2], 4, Some(n_edges))?,
1121        ) {
1122            (Section::U64(offsets), Section::U32(targets), Section::U32(edge_ids)) => {
1123                Ok(Adjacency {
1124                    offsets,
1125                    targets,
1126                    edge_ids,
1127                })
1128            }
1129            _ => Err(malformed(kinds[0])),
1130        }
1131    };
1132    let out = adjacency([OUT_OFFSETS, OUT_TARGETS, OUT_EDGE_IDS])?;
1133    let topology = Topology::from_parts(n_nodes, out, None)?;
1134
1135    let ids: Arc<dyn Array> = if flags & FLAG_STRING_IDS == 0 {
1136        let Section::I64(values) = take(IDS_INT64, 8, Some(n_nodes))? else {
1137            return Err(malformed(IDS_INT64));
1138        };
1139        Arc::new(Int64Array::from(values))
1140    } else {
1141        let (Section::I32(offsets), Section::Bytes(data)) = (
1142            take(IDS_STR_OFFSETS, 4, Some(n_nodes + 1))?,
1143            take(IDS_STR_BYTES, 1, None)?,
1144        ) else {
1145            return Err(malformed(IDS_STR_OFFSETS));
1146        };
1147        if offsets[0] != 0
1148            || offsets.windows(2).any(|w| w[0] > w[1])
1149            || offsets[n_nodes] as usize != data.len()
1150        {
1151            return Err("the string ids' offsets are malformed".into());
1152        }
1153        let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets));
1154        Arc::new(
1155            StringArray::try_new(offsets, Buffer::from_vec(data), None)
1156                .map_err(|e| format!("the string ids are malformed: {e}"))?,
1157        )
1158    };
1159    let ids = IdMap::from_user_ids(ids.as_ref()).map_err(|e| e.to_string())?;
1160    if ids.len() != n_nodes {
1161        return Err("the id map does not match the topology".into());
1162    }
1163    let dropped = match n_dropped {
1164        0 => Dropped::default(),
1165        n_dropped => {
1166            let rows = usize::try_from(n_dropped)
1167                .ok()
1168                .and_then(|d| d.checked_add(n_edges))
1169                .ok_or("the kept rows are malformed")?;
1170            let Section::Bytes(bits) = take(KEPT_ROWS, 1, Some(rows.div_ceil(8)))? else {
1171                return Err(malformed(KEPT_ROWS));
1172            };
1173            let kept = BooleanBuffer::new(Buffer::from_vec(bits), 0, rows);
1174            if kept.count_set_bits() != n_edges {
1175                return Err("the kept rows do not match the edges".into());
1176            }
1177            Dropped {
1178                n_dropped,
1179                kept: Some(kept),
1180            }
1181        }
1182    };
1183    Ok((topology, ids, main, dropped))
1184}
1185
1186fn u64_at(bytes: &[u8], at: usize) -> u64 {
1187    u64::from_le_bytes(bytes[at..at + 8].try_into().expect("8 bytes"))
1188}
1189
1190/// The little-endian values in `bytes`.
1191fn read<T, const N: usize>(
1192    bytes: &[u8],
1193    le: fn([u8; N]) -> T,
1194) -> impl Iterator<Item = T> + use<'_, T, N> {
1195    bytes
1196        .chunks_exact(N)
1197        .map(move |c| le(c.try_into().expect("N bytes")))
1198}
1199
1200#[cfg(test)]
1201mod tests {
1202    use super::*;
1203    use datafusion::arrow::array::ArrayRef;
1204
1205    fn key() -> IndexKey {
1206        IndexKey::new(
1207            "Projection: src, dst\n  TableScan: ties".into(),
1208            vec![FileIdentity {
1209                store: "file:///".into(),
1210                location: "data/edges.parquet".into(),
1211                size: 1234,
1212                modified: 1_700_000_000_000_000_000,
1213                e_tag: Some("\"abc\"".into()),
1214            }],
1215        )
1216    }
1217
1218    fn built(src: ArrayRef, dst: ArrayRef) -> (Arc<Topology>, Arc<IdMap>) {
1219        ursa_io::build_topology_batches(&[(src.as_ref(), dst.as_ref())]).unwrap()
1220    }
1221
1222    fn encode(key: &IndexKey, topology: &Topology, ids: &IdMap) -> Vec<u8> {
1223        let encoded = Encoded::new(key, topology, ids).unwrap();
1224        let mut bytes = Vec::new();
1225        encoded.write_to(&mut bytes).unwrap();
1226        assert_eq!(bytes.len() as u64, encoded.file_size());
1227        bytes
1228    }
1229
1230    fn same(a: &Topology, b: &Topology) {
1231        assert_eq!(a.n_nodes(), b.n_nodes());
1232        assert_eq!(a.out().offsets, b.out().offsets);
1233        assert_eq!(a.out().targets, b.out().targets);
1234        assert_eq!(a.out().edge_ids, b.out().edge_ids);
1235    }
1236
1237    #[test]
1238    fn int_and_string_indexes_round_trip() {
1239        let ints = |v: &[i64]| Arc::new(Int64Array::from(v.to_vec())) as ArrayRef;
1240        let (topology, ids) = built(ints(&[10, 20, 10, 30]), ints(&[20, 30, 30, 10]));
1241        let bytes = encode(&key(), &topology, &ids);
1242        assert_eq!(&bytes[..8], b"URSAIDX\0");
1243        let (back, back_ids) = decode(&bytes, &key()).unwrap();
1244        same(&topology, &back);
1245        assert!(back.incoming_if_built().is_none());
1246        assert_eq!(
1247            back_ids.user_id_array().as_ref(),
1248            ids.user_id_array().as_ref()
1249        );
1250
1251        // The in-adjacency is never part of the main part.
1252        topology.incoming();
1253        assert_eq!(encode(&key(), &topology, &ids), bytes);
1254
1255        let strs = |v: &[&str]| Arc::new(StringArray::from(v.to_vec())) as ArrayRef;
1256        let (topology, ids) = built(strs(&["a", "bé", "a"]), strs(&["bé", "c", "c"]));
1257        let bytes = encode(&key(), &topology, &ids);
1258        let (back, back_ids) = decode(&bytes, &key()).unwrap();
1259        same(&topology, &back);
1260        assert_eq!(
1261            back_ids.user_id_array().as_ref(),
1262            ids.user_id_array().as_ref()
1263        );
1264
1265        // Sections are 64-byte aligned.
1266        let n_sections = u32::from_le_bytes(bytes[40..44].try_into().unwrap()) as usize;
1267        let table = align(HEADER_LEN + u64_at(&bytes, 32), 8) as usize;
1268        for i in 0..n_sections {
1269            assert_eq!(u64_at(&bytes, table + 24 * i + 8) % 64, 0);
1270        }
1271    }
1272
1273    #[test]
1274    fn the_rows_a_vertex_table_left_out_round_trip() {
1275        let ints = |v: &[i64]| Arc::new(Int64Array::from(v.to_vec())) as ArrayRef;
1276        let (topology, ids) = built(ints(&[1, 2]), ints(&[2, 3]));
1277        // Rows 0 and 3 of four were kept, rows 1 and 2 left out.
1278        let kept = BooleanBuffer::from(vec![true, false, false, true]);
1279        let encoded = Encoded::with_dropped(&key(), &topology, &ids, 2, Some(&kept)).unwrap();
1280        let mut bytes = Vec::new();
1281        encoded.write_to(&mut bytes).unwrap();
1282        let (back, _, _, dropped) = decode_from(&bytes[..], bytes.len() as u64, &key()).unwrap();
1283        same(&topology, &back);
1284        assert_eq!(dropped.n_dropped, 2);
1285        assert_eq!(dropped.kept, Some(kept));
1286        // A file with nothing left out has no such section, as before #311.
1287        let plain = encode(&key(), &topology, &ids);
1288        let (_, _, _, none) = decode_from(&plain[..], plain.len() as u64, &key()).unwrap();
1289        assert_eq!((none.n_dropped, none.kept), (0, None));
1290    }
1291
1292    #[test]
1293    fn an_empty_index_round_trips() {
1294        let (topology, ids) =
1295            ursa_io::topology::empty_topology(&datafusion::arrow::datatypes::DataType::Utf8);
1296        let (back, back_ids) = decode(&encode(&key(), &topology, &ids), &key()).unwrap();
1297        assert_eq!((back.n_nodes(), back.n_edges()), (0, 0));
1298        assert!(back_ids.is_empty());
1299    }
1300
1301    #[test]
1302    fn corruption_truncation_and_another_key_are_errors() {
1303        let ints = |v: &[i64]| Arc::new(Int64Array::from(v.to_vec())) as ArrayRef;
1304        let (topology, ids) = built(ints(&[1, 2, 3]), ints(&[2, 3, 1]));
1305        let bytes = encode(&key(), &topology, &ids);
1306        for cut in [0, 10, 64, bytes.len() / 2, bytes.len() - 1] {
1307            assert!(decode(&bytes[..cut], &key()).is_err(), "cut at {cut}");
1308        }
1309        for at in [9, 70, bytes.len() / 2, bytes.len() - 20] {
1310            let mut flipped = bytes.clone();
1311            flipped[at] ^= 0x40;
1312            assert!(decode(&flipped, &key()).is_err(), "flipped at {at}");
1313        }
1314        let mut other = key();
1315        other.files[0].size += 1;
1316        assert_eq!(
1317            decode(&bytes, &other).unwrap_err(),
1318            "the file holds an index for another key"
1319        );
1320        assert!(decode(
1321            b"not an index at all, just some text..........................................",
1322            &key()
1323        )
1324        .is_err());
1325    }
1326
1327    /// Write the main part of `topology` to `path`; its identity.
1328    fn write_main(path: &Path, topology: &Topology, ids: &IdMap) -> MainPart {
1329        let encoded = Encoded::new(&key(), topology, ids).unwrap();
1330        let mut file = File::create(path).unwrap();
1331        encoded.write_to(&mut file).unwrap()
1332    }
1333
1334    #[test]
1335    fn the_in_adjacency_is_appended_and_read_on_first_demand() {
1336        let dir = tempfile::tempdir().unwrap();
1337        let path = dir.path().join("x.idx");
1338        let ints = |v: &[i64]| Arc::new(Int64Array::from(v.to_vec())) as ArrayRef;
1339        let (topology, ids) = built(
1340            ints(&[10, 20, 10, 30, 20, 40]),
1341            ints(&[20, 30, 30, 10, 40, 10]),
1342        );
1343        let main = write_main(&path, &topology, &ids);
1344        assert_eq!(main.len, std::fs::metadata(&path).unwrap().len());
1345
1346        // No chunk: the transpose is built, nothing is read.
1347        let reads = incoming_reads();
1348        let stored = read_file(&path, &key()).unwrap();
1349        assert!(stored.chunks.is_empty());
1350        assert_eq!(
1351            stored.topology.incoming().targets,
1352            topology.incoming().targets
1353        );
1354        assert_eq!(incoming_reads(), reads);
1355
1356        // Appended once; a second append finds it there.
1357        let incoming = topology.incoming().clone();
1358        let end = append_chunk(&path, main, CHUNK_INCOMING, b"", |at| {
1359            EncodedChunk::incoming(&incoming, at)
1360        })
1361        .unwrap()
1362        .unwrap();
1363        assert_eq!(end, std::fs::metadata(&path).unwrap().len());
1364        let again = append_chunk(&path, main, CHUNK_INCOMING, b"", |at| {
1365            EncodedChunk::incoming(&incoming, at)
1366        });
1367        assert_eq!(again.unwrap(), None);
1368
1369        let stored = read_file(&path, &key()).unwrap();
1370        assert_eq!(stored.chunks.len(), 1);
1371        assert_eq!(stored.main, main);
1372        same(&topology, &stored.topology);
1373        // Listed, not read, until asked for.
1374        assert!(stored.topology.incoming_if_built().is_none());
1375        let back = stored.topology.incoming();
1376        assert!(incoming_reads() > reads);
1377        assert_eq!(
1378            (&back.offsets, &back.targets, &back.edge_ids),
1379            (&incoming.offsets, &incoming.targets, &incoming.edge_ids)
1380        );
1381        // The main part reads the same with or without the chunk.
1382        let bytes = std::fs::read(&path).unwrap();
1383        let (main_only, _) = decode(&bytes, &key()).unwrap();
1384        same(&topology, &main_only);
1385    }
1386
1387    #[test]
1388    fn a_damaged_chunk_is_ignored_and_replaced() {
1389        let dir = tempfile::tempdir().unwrap();
1390        let path = dir.path().join("x.idx");
1391        let ints = |v: &[i64]| Arc::new(Int64Array::from(v.to_vec())) as ArrayRef;
1392        let (topology, ids) = built(ints(&[1, 2, 3, 1]), ints(&[2, 3, 1, 3]));
1393        let main = write_main(&path, &topology, &ids);
1394        let incoming = topology.incoming().clone();
1395        let end = append_chunk(&path, main, CHUNK_INCOMING, b"", |at| {
1396            EncodedChunk::incoming(&incoming, at)
1397        })
1398        .unwrap()
1399        .unwrap();
1400        let good = std::fs::read(&path).unwrap();
1401
1402        // A flipped byte in the chunk's arrays: listed, but its data CRC
1403        // fails, so the transpose is built instead, the same.
1404        let mut bytes = good.clone();
1405        bytes[main.len as usize + 64] ^= 0x01;
1406        std::fs::write(&path, &bytes).unwrap();
1407        let stored = read_file(&path, &key()).unwrap();
1408        assert_eq!(stored.chunks.len(), 1);
1409        assert!(read_chunk(&path, main, &stored.chunks[0]).is_err());
1410        assert_eq!(stored.topology.incoming().targets, incoming.targets);
1411        assert_eq!(stored.topology.incoming().edge_ids, incoming.edge_ids);
1412
1413        // A chunk cut short (a crash mid-append): not listed, and the next
1414        // append writes over it.
1415        std::fs::write(&path, &good[..good.len() - 10]).unwrap();
1416        assert!(read_file(&path, &key()).unwrap().chunks.is_empty());
1417        let again = append_chunk(&path, main, CHUNK_INCOMING, b"", |at| {
1418            EncodedChunk::incoming(&incoming, at)
1419        })
1420        .unwrap();
1421        assert_eq!(again, Some(end));
1422        assert_eq!(std::fs::read(&path).unwrap(), good);
1423
1424        // Another main part: nothing is appended or read.
1425        let other = MainPart {
1426            crc: main.crc ^ 1,
1427            ..main
1428        };
1429        assert!(append_chunk(&path, other, CHUNK_INCOMING, b"", |at| {
1430            EncodedChunk::incoming(&incoming, at)
1431        })
1432        .is_err());
1433        let chunk = &read_file(&path, &key()).unwrap().chunks[0];
1434        assert!(read_chunk(&path, other, chunk).is_err());
1435        assert!(read_chunk(&path, main, chunk).is_ok());
1436    }
1437
1438    #[test]
1439    fn weight_chunks_round_trip_by_key() {
1440        let dir = tempfile::tempdir().unwrap();
1441        let path = dir.path().join("x.idx");
1442        let ints = |v: &[i64]| Arc::new(Int64Array::from(v.to_vec())) as ArrayRef;
1443        let (topology, ids) = built(ints(&[1, 2, 3]), ints(&[2, 3, 1]));
1444        let main = write_main(&path, &topology, &ids);
1445        let (a, b) = (vec![0.5, 1.0, f64::MAX], vec![2.0, 0.0, 1e-300]);
1446        for (key, columns) in [("w1", vec![&a[..]]), ("w2", vec![&a[..], &b[..]])] {
1447            let appended = append_chunk(&path, main, CHUNK_WEIGHTS, key.as_bytes(), |at| {
1448                EncodedChunk::weights(key, &columns, at)
1449            });
1450            assert!(appended.unwrap().is_some());
1451        }
1452        let stored = read_file(&path, &key()).unwrap();
1453        assert_eq!(stored.chunks.len(), 2);
1454        let chunk = |name: &[u8]| stored.chunks.iter().find(|c| c.name == name).unwrap();
1455        assert_eq!(
1456            read_weights(&path, main, chunk(b"w1"), 3).unwrap(),
1457            std::slice::from_ref(&a)
1458        );
1459        assert_eq!(read_weights(&path, main, chunk(b"w2"), 3).unwrap(), [a, b]);
1460        // The wrong edge count is an error, never a short column.
1461        assert!(read_weights(&path, main, chunk(b"w1"), 4).is_err());
1462    }
1463
1464    #[test]
1465    fn a_version_1_file_is_not_read() {
1466        let ints = |v: &[i64]| Arc::new(Int64Array::from(v.to_vec())) as ArrayRef;
1467        let (topology, ids) = built(ints(&[1, 2]), ints(&[2, 3]));
1468        let mut bytes = encode(&key(), &topology, &ids);
1469        bytes[8..12].copy_from_slice(&1u32.to_le_bytes());
1470        assert_eq!(
1471            decode(&bytes, &key()).unwrap_err(),
1472            "index format version 1; this build reads 2"
1473        );
1474    }
1475
1476    #[test]
1477    fn keys_round_trip_through_json() {
1478        let k = key();
1479        assert_eq!(IndexKey::from_json(&k.to_json()), Some(k.clone()));
1480        let mut no_tag = k.clone();
1481        no_tag.files[0].e_tag = None;
1482        assert_eq!(IndexKey::from_json(&no_tag.to_json()), Some(no_tag.clone()));
1483        assert_ne!(no_tag.hash(), k.hash());
1484    }
1485}