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 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 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 pub fn add_value(&mut self, value: TypedDictEntry) -> u64 {
85 let id = self.value_dictionary_builder.add(value);
86
87 id
88 }
89
90 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}