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