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