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