1use 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
105pub const MAGIC: [u8; 8] = *b"URSAIDX\0";
107const END_MAGIC: [u8; 4] = *b"XDIU";
109pub const VERSION: u32 = 2;
111const CHUNK_MAGIC: [u8; 4] = *b"CHNK";
113const CHUNK_END_MAGIC: [u8; 4] = *b"KNHC";
115const CHUNK_HEADER_LEN: u64 = 16;
116const CHUNK_TRAILER_LEN: u64 = 48;
117
118pub const CHUNK_INCOMING: u32 = 1;
120pub 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#[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 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 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 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
209enum Data<'a> {
212 U64(&'a [u64]),
213 F64(&'a [f64]),
214 U32(&'a [u32]),
215 Ints(Int64Array),
217 StrOffsets(StringArray),
220 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
245fn 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
256pub 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 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 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 pub fn file_size(&self) -> u64 {
319 self.trailer_offset + TRAILER_LEN
320 }
321
322 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(§ion_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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
368pub struct MainPart {
369 pub len: u64,
370 pub crc: u32,
371}
372
373pub 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 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 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 pub fn end(&self) -> u64 {
441 self.trailer_offset() + CHUNK_TRAILER_LEN
442 }
443
444 pub fn size(&self) -> u64 {
446 self.end() - self.start
447 }
448
449 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(§ion_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
489fn 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
499struct 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 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 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
557pub type DecodeError = String;
559
560pub struct Stored {
562 pub topology: Topology,
564 pub ids: IdMap,
565 pub main: MainPart,
566 pub chunks: Vec<ChunkInfo>,
568}
569
570pub fn has_chunk(chunks: &[ChunkInfo], kind: u32, name: &[u8]) -> bool {
572 chunks.iter().any(|c| c.kind == kind && c.name == name)
573}
574
575static INCOMING_READS: AtomicU64 = AtomicU64::new(0);
577
578pub fn incoming_reads() -> u64 {
581 INCOMING_READS.load(Ordering::Relaxed)
582}
583
584pub 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
618pub 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#[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 data_end: u64,
633 data_crc: u32,
634 sections: Vec<(u32, u64, u64, u64)>,
636}
637
638pub 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
656fn 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
729fn 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
745pub 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
776fn 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
809pub 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
826pub 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 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
860const BLOCK: usize = 1 << 20;
862
863struct 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 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 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 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
944pub 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
954fn 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 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 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
1120fn 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 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 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 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 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 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 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 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 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 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 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 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}