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