Skip to main content

terminus_store/layer/
builder.rs

1use std::io;
2
3use bytes::{Bytes, BytesMut};
4use futures::TryStreamExt;
5use rayon::prelude::*;
6
7use super::layer::*;
8use crate::{chrono_log, storage::*};
9use tdb_succinct::util::{heap_sorted_iter, stream_iter_ok};
10use tdb_succinct::*;
11
12pub struct DictionarySetFileBuilder<F: 'static + FileLoad + FileStore> {
13    node_files: DictionaryFiles<F>,
14    predicate_files: DictionaryFiles<F>,
15    value_files: TypedDictionaryFiles<F>,
16    node_dictionary_builder: StringDictBufBuilder<BytesMut, BytesMut>,
17    predicate_dictionary_builder: StringDictBufBuilder<BytesMut, BytesMut>,
18    value_dictionary_builder: TypedDictBufBuilder<BytesMut, BytesMut, BytesMut, BytesMut>,
19}
20
21impl<F: 'static + FileLoad + FileStore> DictionarySetFileBuilder<F> {
22    pub async fn from_files(
23        node_files: DictionaryFiles<F>,
24        predicate_files: DictionaryFiles<F>,
25        value_files: TypedDictionaryFiles<F>,
26    ) -> io::Result<Self> {
27        let node_dictionary_builder = StringDictBufBuilder::new(BytesMut::new(), BytesMut::new());
28        let predicate_dictionary_builder =
29            StringDictBufBuilder::new(BytesMut::new(), BytesMut::new());
30        let value_dictionary_builder = TypedDictBufBuilder::new(
31            BytesMut::new(),
32            BytesMut::new(),
33            BytesMut::new(),
34            BytesMut::new(),
35        );
36
37        Ok(Self {
38            node_files,
39            predicate_files,
40            value_files,
41            node_dictionary_builder,
42            predicate_dictionary_builder,
43            value_dictionary_builder,
44        })
45    }
46
47    /// Add a node string.
48    ///
49    /// Panics if the given node string is not a lexical successor of the previous node string.
50    pub fn add_node(&mut self, node: &str) -> u64 {
51        let id = self
52            .node_dictionary_builder
53            .add(Bytes::copy_from_slice(node.as_bytes()));
54
55        id
56    }
57
58    pub fn add_node_bytes(&mut self, node: Bytes) -> u64 {
59        let id = self.node_dictionary_builder.add(node);
60
61        id
62    }
63
64    /// Add a predicate string.
65    ///
66    /// Panics if the given predicate string is not a lexical successor of the previous node string.
67    pub fn add_predicate(&mut self, predicate: &str) -> u64 {
68        let id = self
69            .predicate_dictionary_builder
70            .add(Bytes::copy_from_slice(predicate.as_bytes()));
71
72        id
73    }
74
75    pub fn add_predicate_bytes(&mut self, predicate: Bytes) -> u64 {
76        let id = self.predicate_dictionary_builder.add(predicate);
77
78        id
79    }
80
81    /// Add a value string.
82    ///
83    /// Panics if the given value string is not a lexical successor of the previous value string.
84    pub fn add_value(&mut self, value: TypedDictEntry) -> u64 {
85        let id = self.value_dictionary_builder.add(value);
86
87        id
88    }
89
90    /// Add nodes from an iterable.
91    ///
92    /// Panics if the nodes are not in lexical order, or if previous added nodes are a lexical succesor of any of these nodes.
93    pub fn add_nodes<I: 'static + IntoIterator<Item = String> + Unpin + Send + Sync>(
94        &mut self,
95        nodes: I,
96    ) -> Vec<u64>
97    where
98        <I as std::iter::IntoIterator>::IntoIter: Unpin + Send + Sync,
99    {
100        let mut ids = Vec::new();
101        for node in nodes {
102            let id = self.add_node(&node);
103            ids.push(id);
104        }
105
106        ids
107    }
108
109    pub fn add_nodes_bytes<I: 'static + IntoIterator<Item = Bytes> + Unpin + Send + Sync>(
110        &mut self,
111        nodes: I,
112    ) -> Vec<u64>
113    where
114        <I as std::iter::IntoIterator>::IntoIter: Unpin + Send + Sync,
115    {
116        let mut ids = Vec::new();
117        for node in nodes {
118            let id = self.add_node_bytes(node);
119            ids.push(id);
120        }
121
122        ids
123    }
124
125    /// Add predicates from an iterable.
126    ///
127    /// Panics if the predicates are not in lexical order, or if previous added predicates are a lexical succesor of any of these predicates.
128    pub fn add_predicates<I: 'static + IntoIterator<Item = String> + Unpin + Send + Sync>(
129        &mut self,
130        predicates: I,
131    ) -> Vec<u64>
132    where
133        <I as std::iter::IntoIterator>::IntoIter: Unpin + Send + Sync,
134    {
135        let mut ids = Vec::new();
136        for predicate in predicates {
137            let id = self.add_predicate(&predicate);
138            ids.push(id);
139        }
140
141        ids
142    }
143
144    pub fn add_predicates_bytes<I: 'static + IntoIterator<Item = Bytes> + Unpin + Send + Sync>(
145        &mut self,
146        predicates: I,
147    ) -> Vec<u64>
148    where
149        <I as std::iter::IntoIterator>::IntoIter: Unpin + Send + Sync,
150    {
151        let mut ids = Vec::new();
152        for predicate in predicates {
153            let id = self.add_predicate_bytes(predicate);
154            ids.push(id);
155        }
156
157        ids
158    }
159
160    /// Add values from an iterable.
161    ///
162    /// Panics if the values are not in lexical order, or if previous added values are a lexical succesor of any of these values.
163    pub fn add_values<I: 'static + IntoIterator<Item = TypedDictEntry> + Unpin + Send + Sync>(
164        &mut self,
165        values: I,
166    ) -> Vec<u64>
167    where
168        <I as std::iter::IntoIterator>::IntoIter: Unpin + Send + Sync,
169    {
170        let mut ids = Vec::new();
171        for value in values {
172            let id = self.add_value(value);
173            ids.push(id);
174        }
175
176        ids
177    }
178
179    pub async fn finalize(self) -> io::Result<()> {
180        let (mut node_offsets_buf, mut node_data_buf) = self.node_dictionary_builder.finalize();
181        let (mut predicate_offsets_buf, mut predicate_data_buf) =
182            self.predicate_dictionary_builder.finalize();
183        let (
184            mut value_types_present_buf,
185            mut value_type_offsets_buf,
186            mut value_offsets_buf,
187            mut value_data_buf,
188        ) = self.value_dictionary_builder.finalize();
189
190        self.node_files
191            .write_all_from_bufs(&mut node_data_buf, &mut node_offsets_buf)
192            .await?;
193        self.predicate_files
194            .write_all_from_bufs(&mut predicate_data_buf, &mut predicate_offsets_buf)
195            .await?;
196
197        self.value_files
198            .write_all_from_bufs(
199                &mut value_types_present_buf,
200                &mut value_type_offsets_buf,
201                &mut value_offsets_buf,
202                &mut value_data_buf,
203            )
204            .await?;
205
206        Ok(())
207    }
208}
209
210pub struct TripleFileBuilder<F: 'static + FileLoad + FileStore> {
211    subjects_file: Option<F>,
212    subjects: Option<Vec<u64>>,
213
214    s_p_adjacency_list_builder: AdjacencyListBuilder<F, F::Write, F::Write, F::Write>,
215    sp_o_adjacency_list_builder: AdjacencyListBuilder<F, F::Write, F::Write, F::Write>,
216    last_subject: u64,
217    last_predicate: u64,
218}
219
220impl<F: 'static + FileLoad + FileStore> TripleFileBuilder<F> {
221    pub async fn new(
222        s_p_adjacency_list_files: AdjacencyListFiles<F>,
223        sp_o_adjacency_list_files: AdjacencyListFiles<F>,
224        num_nodes: usize,
225        num_predicates: usize,
226        num_values: usize,
227        subjects_file: Option<F>,
228    ) -> io::Result<Self> {
229        let s_p_width = util::calculate_width(num_predicates as u64);
230        let sp_o_width = util::calculate_width((num_nodes + num_values) as u64);
231
232        let s_p_adjacency_list_builder = AdjacencyListBuilder::new(
233            s_p_adjacency_list_files.bitindex_files.bits_file,
234            s_p_adjacency_list_files
235                .bitindex_files
236                .blocks_file
237                .open_write()
238                .await?,
239            s_p_adjacency_list_files
240                .bitindex_files
241                .sblocks_file
242                .open_write()
243                .await?,
244            s_p_adjacency_list_files.nums_file.open_write().await?,
245            s_p_width,
246        )
247        .await?;
248
249        let sp_o_adjacency_list_builder = AdjacencyListBuilder::new(
250            sp_o_adjacency_list_files.bitindex_files.bits_file,
251            sp_o_adjacency_list_files
252                .bitindex_files
253                .blocks_file
254                .open_write()
255                .await?,
256            sp_o_adjacency_list_files
257                .bitindex_files
258                .sblocks_file
259                .open_write()
260                .await?,
261            sp_o_adjacency_list_files.nums_file.open_write().await?,
262            sp_o_width,
263        )
264        .await?;
265
266        let subjects = match subjects_file.is_some() {
267            true => Some(Vec::new()),
268            false => None,
269        };
270
271        Ok(Self {
272            subjects,
273            subjects_file,
274            s_p_adjacency_list_builder,
275            sp_o_adjacency_list_builder,
276            last_subject: 0,
277            last_predicate: 0,
278        })
279    }
280
281    /// Add the given subject, predicate and object.
282    ///
283    /// This will panic if a greater triple has already been added.
284    pub async fn add_triple(
285        &mut self,
286        subject: u64,
287        predicate: u64,
288        object: u64,
289    ) -> io::Result<()> {
290        if subject == 0 || predicate == 0 || object == 0 {
291            return Ok(());
292        }
293
294        if subject < self.last_subject {
295            panic!("layer builder got addition in wrong order (subject is {} while previously {} was pushed)", subject, self.last_subject)
296        } else if self.last_subject == subject && self.last_predicate == predicate {
297            // only the second adjacency list has to be pushed to
298            let count = self.s_p_adjacency_list_builder.count() + 1;
299
300            self.sp_o_adjacency_list_builder.push(count, object).await?;
301        } else {
302            // both list have to be pushed to
303            if self.subjects.is_some() && subject != self.last_subject {
304                self.subjects.as_mut().unwrap().push(subject);
305            }
306            let mapped_subject = self
307                .subjects
308                .as_ref()
309                .map(|s| s.len() as u64)
310                .unwrap_or(subject);
311            self.s_p_adjacency_list_builder
312                .push(mapped_subject, predicate)
313                .await?;
314            let count = self.s_p_adjacency_list_builder.count() + 1;
315
316            self.sp_o_adjacency_list_builder.push(count, object).await?;
317        }
318
319        self.last_subject = subject;
320        self.last_predicate = predicate;
321
322        Ok(())
323    }
324
325    /// Add the given triples.
326    ///
327    /// This will panic if a greater triple has already been added.
328    pub async fn add_id_triples<I: 'static + IntoIterator<Item = IdTriple>>(
329        &mut self,
330        triples: I,
331    ) -> io::Result<()> {
332        for triple in triples {
333            self.add_triple(triple.subject, triple.predicate, triple.object)
334                .await?;
335        }
336
337        Ok(())
338    }
339
340    pub async fn finalize(self) -> io::Result<()> {
341        self.s_p_adjacency_list_builder.finalize().await?;
342        self.sp_o_adjacency_list_builder.finalize().await?;
343
344        if let Some(subjects) = self.subjects {
345            // isn't this just last_subject?
346            let max_subject = if subjects.is_empty() {
347                0
348            } else {
349                subjects[subjects.len() - 1]
350            };
351
352            let subjects_width = util::calculate_width(max_subject);
353            let mut subjects_logarray_builder = LogArrayFileBuilder::new(
354                self.subjects_file.unwrap().open_write().await?,
355                subjects_width,
356            );
357
358            subjects_logarray_builder.push_vec(subjects).await?;
359            subjects_logarray_builder.finalize().await?;
360        };
361
362        Ok(())
363    }
364}
365
366const SINGLE_SORT_LIMIT: u64 = 0x8000_0000;
367pub async fn build_object_index_from_direct_files<
368    FLoad: 'static + FileLoad,
369    F: 'static + FileLoad + FileStore,
370>(
371    sp_o_nums_file: FLoad,
372    sp_o_bits_file: FLoad,
373    o_ps_files: AdjacencyListFiles<F>,
374    objects_file: Option<F>,
375) -> io::Result<()> {
376    chrono_log!("starting object index build");
377    let build_sparse_index = objects_file.is_some();
378    let (count, spo_width) = logarray_file_get_length_and_width(sp_o_nums_file.clone()).await?;
379    let mut aj_stream = adjacency_list_stream_pairs(sp_o_bits_file, sp_o_nums_file).await?;
380    let mut pairs = Vec::with_capacity(std::cmp::min(count, SINGLE_SORT_LIMIT) as usize);
381    let mut greatest_sp = 0;
382    chrono_log!("opened sp_o stream");
383    let mut tally: u64 = 0;
384    let mut temp_arrays: Vec<(LogArray, LogArray)> = Vec::new();
385    // gather up pars
386    while let Some((sp, object)) = aj_stream.try_next().await? {
387        greatest_sp = sp;
388        pairs.push((object, sp));
389        tally += 1;
390        if tally % 10000000 == 0 {
391            chrono_log!(
392                "collected {tally} pairs for o_ps index ({}%)",
393                (tally * 100 / count)
394            );
395        }
396
397        if tally % SINGLE_SORT_LIMIT == 0 {
398            chrono_log!("collect currently gathered elements into a logarray");
399            pairs.par_sort_unstable();
400            let mut sp_file = BytesMut::with_capacity(0);
401            let mut o_file = BytesMut::with_capacity(0);
402            let sp_width = util::calculate_width(greatest_sp);
403            let mut sp_logarray = LogArrayBufBuilder::new(&mut sp_file, sp_width);
404            sp_logarray.reserve(pairs.len());
405            let mut o_logarray = LogArrayBufBuilder::new(&mut o_file, spo_width);
406            o_logarray.reserve(pairs.len());
407            for (o, sp) in pairs.iter_mut() {
408                sp_logarray.push(*sp);
409                o_logarray.push(*o);
410            }
411            sp_logarray.finalize();
412            o_logarray.finalize();
413            temp_arrays.push((
414                LogArray::parse(sp_file.freeze()).unwrap(),
415                LogArray::parse(o_file.freeze()).unwrap(),
416            ));
417
418            pairs.clear();
419        }
420    }
421    chrono_log!("collected object pairs");
422
423    // par_sort_unstable unfortunately can run out of stack for very
424    // large sorts. If so, we have to do something else.
425    if pairs.len() as u64 > SINGLE_SORT_LIMIT {
426        chrono_log!("perform multi sort");
427        let mut tally: u64 = 0;
428        while tally < pairs.len() as u64 {
429            let end = std::cmp::min(count as usize, (tally + SINGLE_SORT_LIMIT) as usize);
430            let slice = &mut pairs[tally as usize..end];
431            slice.par_sort_unstable();
432            tally += SINGLE_SORT_LIMIT;
433        }
434        chrono_log!("perform final sort");
435        // we use the normal sort as it is fast for cases where you
436        // have a bunch of appended sorted slices.
437        pairs.sort();
438    } else {
439        chrono_log!("perform single sort");
440        pairs.par_sort_unstable();
441    }
442    chrono_log!("sorted object pairs");
443
444    let aj_width = util::calculate_width(greatest_sp);
445    let mut o_ps_adjacency_list_builder = AdjacencyListBuilder::new(
446        o_ps_files.bitindex_files.bits_file,
447        o_ps_files.bitindex_files.blocks_file.open_write().await?,
448        o_ps_files.bitindex_files.sblocks_file.open_write().await?,
449        o_ps_files.nums_file.open_write().await?,
450        aj_width,
451    )
452    .await?;
453
454    // now construct a sorted stream out of the part still in memory and the parts written out to files
455    let mut iters = Vec::with_capacity(temp_arrays.len() + 1);
456    for (sp_file, o_file) in temp_arrays {
457        let sp_iter = sp_file.iter();
458        let o_iter = o_file.iter();
459
460        let iter = o_iter.zip(sp_iter);
461        iters.push(itertools::Either::Left(iter));
462    }
463    iters.push(itertools::Either::Right(pairs.into_iter()));
464    let mut merged_iters = heap_sorted_iter(iters);
465
466    if build_sparse_index {
467        // a sparse index compresses the adjacency list so that all objects in use are remapped to form a continuous range.
468        // We need to iterate over the pairs, and write them out without gaps.
469
470        let mut objects = Vec::new();
471        let mut last_object = 0;
472        let mut object_ix = 0;
473        while let Some((object, sp)) = merged_iters.next() {
474            if object > last_object {
475                object_ix += 1;
476                last_object = object;
477
478                // keep track of all objects in use in a separate list
479                objects.push(object);
480            }
481
482            o_ps_adjacency_list_builder.push(object_ix, sp).await?;
483        }
484        let objects_width = util::calculate_width(last_object);
485
486        // write out the object list
487        let mut objects_builder =
488            LogArrayFileBuilder::new(objects_file.unwrap().open_write().await?, objects_width);
489        objects_builder.push_vec(objects).await?;
490        objects_builder.finalize().await?;
491    } else {
492        o_ps_adjacency_list_builder
493            .push_all(stream_iter_ok::<_, io::Error, _>(merged_iters))
494            .await?;
495    }
496    chrono_log!("added object pairs to adjacency list builder");
497
498    o_ps_adjacency_list_builder.finalize().await?;
499    chrono_log!("finalized object index");
500
501    Ok(())
502}
503
504pub async fn build_object_index<FLoad: 'static + FileLoad, F: 'static + FileLoad + FileStore>(
505    sp_o_files: AdjacencyListFiles<FLoad>,
506    o_ps_files: AdjacencyListFiles<F>,
507    objects_file: Option<F>,
508) -> io::Result<()> {
509    build_object_index_from_direct_files(
510        sp_o_files.nums_file,
511        sp_o_files.bitindex_files.bits_file,
512        o_ps_files,
513        objects_file,
514    )
515    .await
516}
517
518pub async fn build_predicate_index<FLoad: 'static + FileLoad, F: 'static + FileLoad + FileStore>(
519    source: FLoad,
520    destination_bits: F,
521    destination_blocks: F,
522    destination_sblocks: F,
523) -> io::Result<()> {
524    build_wavelet_tree_from_logarray(
525        source,
526        destination_bits,
527        destination_blocks,
528        destination_sblocks,
529    )
530    .await
531}
532
533pub async fn build_indexes<FLoad: 'static + FileLoad, F: 'static + FileLoad + FileStore>(
534    s_p_files: AdjacencyListFiles<FLoad>,
535    sp_o_files: AdjacencyListFiles<FLoad>,
536    o_ps_files: AdjacencyListFiles<F>,
537    objects_file: Option<F>,
538    wavelet_files: BitIndexFiles<F>,
539) -> io::Result<()> {
540    let object_index_task = tokio::spawn(build_object_index(sp_o_files, o_ps_files, objects_file));
541    let predicate_index_task = tokio::spawn(build_predicate_index(
542        s_p_files.nums_file,
543        wavelet_files.bits_file,
544        wavelet_files.blocks_file,
545        wavelet_files.sblocks_file,
546    ));
547
548    object_index_task.await??;
549    chrono_log!("built object index");
550    predicate_index_task.await??;
551    chrono_log!("built predicate index");
552
553    Ok(())
554}