1use std::collections::BTreeMap;
17use std::io::{self, Read, Write};
18
19use roaring::RoaringBitmap;
20
21use crate::ordinal::{Ordinal, to_usize};
22
23const MAGIC: &[u8; 8] = b"FTATTR\0\0";
24const VERSION: u32 = 1;
25const TAG_CONCEPT: u8 = 0;
26const TAG_NUMBER: u8 = 1;
27const TAG_STRING: u8 = 2;
28
29#[derive(Debug, thiserror::Error)]
31pub enum AttributesError {
32 #[error("attribute row ({from}, {kind}) is out of range for {nodes} nodes and {types} types")]
34 OutOfRange {
35 from: u32,
37 kind: u32,
39 nodes: u32,
41 types: u32,
43 },
44 #[error("too many attribute rows")]
46 TooMany(#[source] std::num::TryFromIntError),
47 #[error("attributes I/O failed")]
49 Io(#[from] io::Error),
50 #[error("not an attributes artifact")]
52 Magic,
53 #[error("attributes layout version {found}, expected {expected}")]
55 Version {
56 found: u32,
58 expected: u32,
60 },
61 #[error("the attribute arrays are inconsistent")]
63 Inconsistent,
64 #[error("an attribute value is not UTF-8")]
66 Text(#[from] std::string::FromUtf8Error),
67}
68
69#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
71pub enum Value {
72 Concept(Ordinal),
74 Number(String),
76 String(String),
78}
79
80#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
82pub struct Edge {
83 pub source: Ordinal,
85 pub group: u32,
87 pub kind: u32,
89 pub value: Value,
91}
92
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
95pub enum ValueRef<'a> {
96 Concept(Ordinal),
98 Number(&'a str),
100 String(&'a str),
102}
103
104#[derive(Debug, Clone, Copy, PartialEq, Eq)]
106pub struct Row<'a> {
107 pub group: u32,
109 pub kind: u32,
111 pub value: ValueRef<'a>,
113}
114
115#[derive(Debug, Clone, PartialEq, Eq, Default)]
117struct Inverted {
118 targets: Vec<u32>,
120 offsets: Vec<u32>,
122 sources: Vec<u32>,
124 all_sources: RoaringBitmap,
126}
127
128impl Inverted {
129 fn build(mut pairs: Vec<(u32, u32)>, all_sources: RoaringBitmap) -> Self {
130 pairs.sort_unstable();
131 pairs.dedup();
132 let mut targets = Vec::new();
133 let mut offsets = Vec::new();
134 let mut sources = Vec::with_capacity(pairs.len());
135 for (target, source) in pairs {
136 if targets.last() != Some(&target) {
137 targets.push(target);
138 offsets.push(u32::try_from(sources.len()).unwrap_or(u32::MAX));
139 }
140 sources.push(source);
141 }
142 offsets.push(u32::try_from(sources.len()).unwrap_or(u32::MAX));
143 Self {
144 targets,
145 offsets,
146 sources,
147 all_sources,
148 }
149 }
150
151 fn size_in_bytes(&self) -> usize {
152 crate::footprint::vector(&self.targets)
153 .saturating_add(crate::footprint::vector(&self.offsets))
154 .saturating_add(crate::footprint::vector(&self.sources))
155 .saturating_add(crate::footprint::bitmap(&self.all_sources))
156 }
157
158 fn sources(&self, target: u32) -> &[u32] {
159 let Ok(index) = self.targets.binary_search(&target) else {
160 return &[];
161 };
162 match (
163 self.offsets.get(index),
164 self.offsets.get(index.saturating_add(1)),
165 ) {
166 (Some(&start), Some(&end)) => self
167 .sources
168 .get(to_usize(start)..to_usize(end))
169 .unwrap_or_default(),
170 _ => &[],
171 }
172 }
173}
174
175#[derive(Debug, Clone, PartialEq, Eq, Default)]
177pub struct Attributes {
178 types: Vec<u64>,
180 offsets: Vec<u32>,
182 groups: Vec<u32>,
183 kinds: Vec<u32>,
184 tags: Vec<u8>,
185 payloads: Vec<u32>,
186 strings: Vec<String>,
187 inverted: Vec<Inverted>,
189}
190
191impl Attributes {
192 #[must_use]
194 pub fn rows_size_in_bytes(&self) -> usize {
195 crate::footprint::vector(&self.types)
196 .saturating_add(crate::footprint::vector(&self.offsets))
197 .saturating_add(crate::footprint::vector(&self.groups))
198 .saturating_add(crate::footprint::vector(&self.kinds))
199 .saturating_add(crate::footprint::vector(&self.tags))
200 .saturating_add(crate::footprint::vector(&self.payloads))
201 .saturating_add(crate::footprint::strings(&self.strings))
202 }
203
204 #[must_use]
209 pub fn inverted_size_in_bytes(&self) -> usize {
210 self.inverted.iter().fold(
211 crate::footprint::vector(&self.inverted),
212 |total, inverted| total.saturating_add(inverted.size_in_bytes()),
213 )
214 }
215
216 #[must_use]
218 pub fn size_in_bytes(&self) -> usize {
219 self.rows_size_in_bytes()
220 .saturating_add(self.inverted_size_in_bytes())
221 }
222
223 pub fn build(
231 nodes: u32,
232 types: Vec<u64>,
233 mut edges: Vec<Edge>,
234 ) -> Result<Self, AttributesError> {
235 let type_count = u32::try_from(types.len()).map_err(AttributesError::TooMany)?;
236 Self::check_range(nodes, type_count, &edges)?;
237 edges.sort_unstable();
238 edges.dedup();
239 u32::try_from(edges.len()).map_err(AttributesError::TooMany)?;
240 let mut attributes = Self::pack(nodes, types, &edges)?;
241 attributes.derive();
242 Ok(attributes)
243 }
244
245 fn check_range(nodes: u32, type_count: u32, edges: &[Edge]) -> Result<(), AttributesError> {
248 for edge in edges {
249 let target_ok = match edge.value {
250 Value::Concept(target) => target.index() < nodes,
251 Value::Number(_) | Value::String(_) => true,
252 };
253 if edge.source.index() >= nodes || edge.kind >= type_count || !target_ok {
254 return Err(AttributesError::OutOfRange {
255 from: edge.source.index(),
256 kind: edge.kind,
257 nodes,
258 types: type_count,
259 });
260 }
261 }
262 Ok(())
263 }
264
265 fn pack(nodes: u32, types: Vec<u64>, edges: &[Edge]) -> Result<Self, AttributesError> {
268 let mut interned: BTreeMap<String, u32> = BTreeMap::new();
269 let mut strings = Vec::new();
270 let mut intern = |text: &str| -> Result<u32, AttributesError> {
271 if let Some(&index) = interned.get(text) {
272 return Ok(index);
273 }
274 let index = u32::try_from(strings.len()).map_err(AttributesError::TooMany)?;
275 strings.push(text.to_owned());
276 interned.insert(text.to_owned(), index);
277 Ok(index)
278 };
279 let mut offsets = Vec::with_capacity(to_usize(nodes).saturating_add(1));
280 let (mut groups, mut kinds, mut tags, mut payloads) =
281 (Vec::new(), Vec::new(), Vec::new(), Vec::new());
282 let mut cursor = 0_usize;
283 for node in 0..nodes {
284 offsets.push(u32::try_from(groups.len()).unwrap_or(u32::MAX));
285 while let Some(edge) = edges.get(cursor) {
286 if edge.source.index() != node {
287 break;
288 }
289 groups.push(edge.group);
290 kinds.push(edge.kind);
291 let (tag, payload) = match &edge.value {
292 Value::Concept(target) => (TAG_CONCEPT, target.index()),
293 Value::Number(text) => (TAG_NUMBER, intern(text)?),
294 Value::String(text) => (TAG_STRING, intern(text)?),
295 };
296 tags.push(tag);
297 payloads.push(payload);
298 cursor = cursor.saturating_add(1);
299 }
300 }
301 offsets.push(u32::try_from(groups.len()).unwrap_or(u32::MAX));
302 Ok(Self {
303 types,
304 offsets,
305 groups,
306 kinds,
307 tags,
308 payloads,
309 strings,
310 inverted: Vec::new(),
311 })
312 }
313
314 fn derive(&mut self) {
316 let mut pairs: Vec<Vec<(u32, u32)>> = vec![Vec::new(); self.types.len()];
317 let mut all: Vec<RoaringBitmap> = vec![RoaringBitmap::new(); self.types.len()];
318 for node in 0..self.nodes() {
319 for row in self.rows(Ordinal::new(node)) {
320 let kind = to_usize(row.kind);
321 if let Some(set) = all.get_mut(kind) {
322 set.insert(node);
323 }
324 if let (ValueRef::Concept(target), Some(list)) = (row.value, pairs.get_mut(kind)) {
325 list.push((target.index(), node));
326 }
327 }
328 }
329 self.inverted = pairs
330 .into_iter()
331 .zip(all)
332 .map(|(pairs, all_sources)| Inverted::build(pairs, all_sources))
333 .collect();
334 }
335
336 #[must_use]
338 pub fn types(&self) -> &[u64] {
339 &self.types
340 }
341
342 #[must_use]
344 pub fn kind(&self, sctid: u64) -> Option<u32> {
345 self.types
346 .iter()
347 .position(|&t| t == sctid)
348 .and_then(|i| u32::try_from(i).ok())
349 }
350
351 #[must_use]
353 pub fn nodes(&self) -> u32 {
354 u32::try_from(self.offsets.len().saturating_sub(1)).unwrap_or(u32::MAX)
355 }
356
357 #[must_use]
359 pub fn edges(&self) -> usize {
360 self.groups.len()
361 }
362
363 pub fn rows(&self, source: Ordinal) -> impl Iterator<Item = Row<'_>> + '_ {
365 let index = to_usize(source.index());
366 let (start, end) = match (
367 self.offsets.get(index),
368 self.offsets.get(index.saturating_add(1)),
369 ) {
370 (Some(&s), Some(&e)) => (to_usize(s), to_usize(e)),
371 _ => (0, 0),
372 };
373 (start..end).filter_map(move |i| {
374 Some(Row {
375 group: *self.groups.get(i)?,
376 kind: *self.kinds.get(i)?,
377 value: match (*self.tags.get(i)?, *self.payloads.get(i)?) {
378 (TAG_CONCEPT, target) => ValueRef::Concept(Ordinal::new(target)),
379 (TAG_NUMBER, text) => ValueRef::Number(self.strings.get(to_usize(text))?),
380 (_, text) => ValueRef::String(self.strings.get(to_usize(text))?),
381 },
382 })
383 })
384 }
385
386 #[must_use]
388 pub fn sources(&self, kind: u32, target: Ordinal) -> &[u32] {
389 self.inverted
390 .get(to_usize(kind))
391 .map_or(&[], |inverted| inverted.sources(target.index()))
392 }
393
394 #[must_use]
396 pub fn sources_of_kind(&self, kind: u32) -> Option<&RoaringBitmap> {
397 self.inverted.get(to_usize(kind)).map(|i| &i.all_sources)
398 }
399
400 #[must_use]
402 pub fn targets_of_kind(&self, kind: u32) -> &[u32] {
403 self.inverted
404 .get(to_usize(kind))
405 .map_or(&[], |i| i.targets.as_slice())
406 }
407
408 pub fn write_to(&self, out: &mut impl Write) -> Result<(), AttributesError> {
414 let count = |len: usize| u32::try_from(len).map_err(AttributesError::TooMany);
415 out.write_all(MAGIC)?;
416 out.write_all(&VERSION.to_le_bytes())?;
417 out.write_all(&count(self.types.len())?.to_le_bytes())?;
418 for sctid in &self.types {
419 out.write_all(&sctid.to_le_bytes())?;
420 }
421 out.write_all(&self.nodes().to_le_bytes())?;
422 write_u32s(out, &self.offsets)?;
423 out.write_all(&count(self.groups.len())?.to_le_bytes())?;
424 write_u32s(out, &self.groups)?;
425 write_u32s(out, &self.kinds)?;
426 out.write_all(&self.tags)?;
427 write_u32s(out, &self.payloads)?;
428 out.write_all(&count(self.strings.len())?.to_le_bytes())?;
429 for text in &self.strings {
430 out.write_all(&count(text.len())?.to_le_bytes())?;
431 out.write_all(text.as_bytes())?;
432 }
433 Ok(())
434 }
435
436 pub fn read_from(input: &mut impl Read) -> Result<Self, AttributesError> {
443 let mut magic = [0_u8; 8];
444 input.read_exact(&mut magic)?;
445 if &magic != MAGIC {
446 return Err(AttributesError::Magic);
447 }
448 let version = read_u32(input)?;
449 if version != VERSION {
450 return Err(AttributesError::Version {
451 found: version,
452 expected: VERSION,
453 });
454 }
455 let type_count = read_u32(input)?;
456 let mut types = Vec::with_capacity(to_usize(type_count));
457 for _ in 0..type_count {
458 let mut long = [0_u8; 8];
459 input.read_exact(&mut long)?;
460 types.push(u64::from_le_bytes(long));
461 }
462 let nodes = read_u32(input)?;
463 let offsets = read_u32s(input, to_usize(nodes).saturating_add(1))?;
464 let rows = to_usize(read_u32(input)?);
465 let groups = read_u32s(input, rows)?;
466 let kinds = read_u32s(input, rows)?;
467 let mut tags = vec![0_u8; rows];
468 input.read_exact(&mut tags)?;
469 let payloads = read_u32s(input, rows)?;
470 let string_count = read_u32(input)?;
471 let mut strings = Vec::with_capacity(to_usize(string_count));
472 for _ in 0..string_count {
473 let len = to_usize(read_u32(input)?);
474 let mut bytes = vec![0_u8; len];
475 input.read_exact(&mut bytes)?;
476 strings.push(String::from_utf8(bytes)?);
477 }
478 let consistent = offsets.last().is_some_and(|&l| to_usize(l) == rows)
479 && offsets.windows(2).all(|w| w.first() <= w.get(1))
480 && kinds.iter().all(|&k| to_usize(k) < types.len())
481 && tags
482 .iter()
483 .zip(&payloads)
484 .all(|(&tag, &payload)| match tag {
485 TAG_CONCEPT => payload < nodes,
486 TAG_NUMBER | TAG_STRING => to_usize(payload) < strings.len(),
487 _ => false,
488 });
489 if !consistent {
490 return Err(AttributesError::Inconsistent);
491 }
492 let mut attributes = Self {
493 types,
494 offsets,
495 groups,
496 kinds,
497 tags,
498 payloads,
499 strings,
500 inverted: Vec::new(),
501 };
502 attributes.derive();
503 Ok(attributes)
504 }
505}
506
507fn write_u32s(out: &mut impl Write, values: &[u32]) -> Result<(), AttributesError> {
508 for value in values {
509 out.write_all(&value.to_le_bytes())?;
510 }
511 Ok(())
512}
513
514fn read_u32(input: &mut impl Read) -> Result<u32, AttributesError> {
515 let mut bytes = [0_u8; 4];
516 input.read_exact(&mut bytes)?;
517 Ok(u32::from_le_bytes(bytes))
518}
519
520fn read_u32s(input: &mut impl Read, count: usize) -> Result<Vec<u32>, AttributesError> {
521 let mut bytes = vec![0_u8; count.saturating_mul(4)];
522 input.read_exact(&mut bytes)?;
523 Ok(bytes
524 .as_chunks::<4>()
525 .0
526 .iter()
527 .map(|chunk| u32::from_le_bytes(*chunk))
528 .collect())
529}
530
531#[cfg(test)]
532mod tests {
533 use super::{Attributes, AttributesError, Edge, Value, ValueRef};
534 use crate::ordinal::Ordinal;
535
536 fn sample() -> Attributes {
537 let edge = |source: u32, group: u32, kind: u32, value: Value| Edge {
538 source: Ordinal::new(source),
539 group,
540 kind,
541 value,
542 };
543 Attributes::build(
544 5,
545 vec![100, 200],
546 vec![
547 edge(1, 1, 0, Value::Concept(Ordinal::new(3))),
548 edge(1, 1, 1, Value::Number(String::from("4"))),
549 edge(2, 0, 0, Value::Concept(Ordinal::new(3))),
550 edge(2, 2, 1, Value::Number(String::from("4"))),
551 edge(2, 0, 0, Value::Concept(Ordinal::new(3))),
552 edge(4, 0, 1, Value::String(String::from("blue"))),
553 ],
554 )
555 .expect("builds")
556 }
557
558 #[test]
559 fn rows_are_grouped_and_the_inverted_index_answers_by_type_and_target() {
560 let attributes = sample();
561 assert_eq!(attributes.edges(), 5, "the duplicate row is dropped");
562 let cat: Vec<_> = attributes.rows(Ordinal::new(1)).collect();
563 assert_eq!(cat.len(), 2);
564 assert_eq!(cat[0].group, 1);
565 assert_eq!(cat[0].value, ValueRef::Concept(Ordinal::new(3)));
566 assert_eq!(cat[1].value, ValueRef::Number("4"));
567 assert_eq!(attributes.sources(0, Ordinal::new(3)), [1, 2]);
568 assert_eq!(attributes.sources(1, Ordinal::new(3)), []);
569 assert_eq!(attributes.targets_of_kind(0), [3]);
570 assert_eq!(
571 attributes
572 .sources_of_kind(1)
573 .expect("kind")
574 .iter()
575 .collect::<Vec<_>>(),
576 [1, 2, 4]
577 );
578 assert_eq!(attributes.kind(200), Some(1));
579 assert_eq!(attributes.kind(300), None);
580 assert_eq!(
581 attributes.rows(Ordinal::new(4)).next().expect("row").value,
582 ValueRef::String("blue")
583 );
584 assert!(attributes.rows(Ordinal::new(0)).next().is_none());
585 }
586
587 #[test]
588 fn the_layout_round_trips_and_refuses_bad_input() {
589 let attributes = sample();
590 let mut bytes = Vec::new();
591 attributes.write_to(&mut bytes).expect("writes");
592 let again = Attributes::read_from(&mut bytes.as_slice()).expect("reads");
593 assert_eq!(again, attributes);
594 assert!(matches!(
595 Attributes::read_from(&mut b"nope".as_slice()),
596 Err(AttributesError::Io(_))
597 ));
598 assert!(matches!(
599 Attributes::build(
600 2,
601 vec![1],
602 vec![Edge {
603 source: Ordinal::new(1),
604 group: 0,
605 kind: 1,
606 value: Value::Concept(Ordinal::new(0)),
607 }]
608 ),
609 Err(AttributesError::OutOfRange { kind: 1, .. })
610 ));
611 let mut wrong = bytes.clone();
612 wrong[0] = b'X';
613 assert!(matches!(
614 Attributes::read_from(&mut wrong.as_slice()),
615 Err(AttributesError::Magic)
616 ));
617 }
618}