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