added creation of offsets during parallel compression + fmt + clippy

Author: zommiommy

src/graphs/bvgraph/codecs/factories.rs

@@ -188,7 +188,7 @@ impl<E: Endianness> MemoryFactory<E, Box<[u32]>> {
         bytes[file_len..].fill(0);
         Ok(Self {
             // Safety: the length is a multiple of 16.
-            data: unsafe { std::mem::transmute(bytes.into_boxed_slice()) },
+            data: unsafe { std::mem::transmute::<Box<[u8]>, Box<[u32]>>(bytes.into_boxed_slice()) },
             _marker: core::marker::PhantomData,
         })
     }

src/graphs/bvgraph/comp/impls.rs

@@ -12,7 +12,7 @@ use dsi_progress_logger::prelude::*;
 use lender::prelude::*;
 use std::fs::File;
 use std::io::{BufReader, BufWriter};
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 /// A queue that pulls jobs with ids in a contiguous initial segment of the
 /// natural numbers from an iterator out of order and implement an iterator in
@@ -42,15 +42,17 @@ impl<I: Iterator> TaskQueue<I> {
 
 impl<I: Iterator> Iterator for TaskQueue<I>
 where
-    I::Item: JobId + Copy,
+    I::Item: JobId,
 {
     type Item = I::Item;
 
     fn next(&mut self) -> Option<Self::Item> {
         loop {
-            if let Some(Some(item)) = self.jobs.get(self.next_id) {
-                self.next_id += 1;
-                return Some(*item);
+            if let Some(item) = self.jobs.get_mut(self.next_id) {
+                if item.is_some() {
+                    self.next_id += 1;
+                    return item.take();
+                }
             }
             if let Some(item) = self.iter.next() {
                 let id = item.id();
@@ -66,12 +68,15 @@ where
 }
 
 /// A compression job.
-#[derive(Debug, PartialEq, PartialOrd, Eq, Ord, Clone, Copy)]
+#[derive(Debug, PartialEq, PartialOrd, Eq, Ord, Clone)]
 struct Job {
     job_id: usize,
     first_node: usize,
     last_node: usize,
+    chunk_graph_path: PathBuf,
     written_bits: u64,
+    chunk_offsets_path: PathBuf,
+    offsets_written_bits: u64,
     num_arcs: u64,
 }
 
@@ -284,6 +289,7 @@ impl BVComp<()> {
         let basename = basename.as_ref();
 
         let graph_path = basename.with_extension(GRAPH_EXTENSION);
+        let offsets_path = basename.with_extension(OFFSETS_EXTENSION);
 
         let (tx, rx) = std::sync::mpsc::channel();
 
@@ -293,51 +299,70 @@ impl BVComp<()> {
             let cp_flags = &compression_flags;
 
             for (thread_id, mut thread_lender) in iter.enumerate() {
-                let file_path = thread_path(thread_id);
+                let tmp_path = thread_path(thread_id);
+                let chunk_graph_path = tmp_path.with_extension(GRAPH_EXTENSION);
+                let chunk_offsets_path = tmp_path.with_extension(OFFSETS_EXTENSION);
                 let tx = tx.clone();
                 // Spawn the thread
                 s.spawn(move |_| {
                     log::info!("Thread {} started", thread_id);
                     let first_node;
-
-                    let (mut bvcomp, mut written_bits) =
-                        if let Some((node_id, successors)) = thread_lender.next() {
+                    let mut bvcomp;
+                    let mut offsets_writer;
+                    let mut written_bits;
+                    let mut offsets_written_bits;
+
+                    match thread_lender.next() {
+                        None => return,
+                        Some((node_id, successors)) => {
                             first_node = node_id;
+
+                            offsets_writer = <BufBitWriter<BigEndian, _>>::new(<WordAdapter<usize, _>>::new(
+                                BufWriter::new(File::create(&chunk_offsets_path).unwrap()),
+                            ));
+
                             let writer = <BufBitWriter<E, _>>::new(<WordAdapter<usize, _>>::new(
-                                BufWriter::new(File::create(&file_path).unwrap()),
+                                BufWriter::new(File::create(&chunk_graph_path).unwrap()),
                             ));
                             let codes_encoder = <DynCodesEncoder<E, _>>::new(writer, cp_flags);
 
-                            let mut bvcomp = BVComp::new(
+                            bvcomp = BVComp::new(
                                 codes_encoder,
                                 cp_flags.compression_window,
                                 cp_flags.max_ref_count,
                                 cp_flags.min_interval_length,
                                 node_id,
                             );
-                            let written_bits = bvcomp.push(successors).unwrap();
-                            (bvcomp, written_bits)
-                        } else {
-                            return;
-                        };
+                            written_bits = bvcomp.push(successors).unwrap();
+                            offsets_written_bits = offsets_writer.write_gamma(written_bits).unwrap() as u64;
+                        }
+                    };
 
                     let mut last_node = first_node;
-                    written_bits += bvcomp
-                        .extend(thread_lender.inspect(|(x, _)| last_node = *x))
-                        .unwrap();
+                    let iter_nodes = thread_lender.inspect(|(x, _)| last_node = *x);
+                    for_! ( (_, succ) in iter_nodes {
+                        let node_bits = bvcomp.push(succ.into_iter()).unwrap();
+                        written_bits += node_bits;
+                        offsets_written_bits += offsets_writer.write_gamma(node_bits).unwrap() as u64;
+                    });
+
                     let num_arcs = bvcomp.arcs;
                     bvcomp.flush().unwrap();
-                    // TODO written_bits += bvcomp.flush().unwrap();
+
                     log::info!(
-                        "Finished Compression thread {} and wrote {} bits",
+                        "Finished Compression thread {} and wrote {} bits for the graph and {} bits for the offsets",
                         thread_id,
-                        written_bits
+                        written_bits,
+                        offsets_written_bits,
                     );
                     tx.send(Job {
                         job_id: thread_id,
                         first_node,
                         last_node,
+                        chunk_graph_path,
                         written_bits,
+                        chunk_offsets_path,
+                        offsets_written_bits,
                         num_arcs,
                     })
                     .unwrap()
@@ -346,15 +371,19 @@ impl BVComp<()> {
 
             drop(tx);
 
-            // setup the final bitstream from the end, because the first thread
-            // already wrote the first chunk
             let file = File::create(&graph_path)
                 .with_context(|| format!("Could not create graph {}", graph_path.display()))?;
-
-            let mut result_writer =
+            let mut graph_writer =
                 <BufBitWriter<E, _>>::new(<WordAdapter<usize, _>>::new(BufWriter::new(file)));
 
+            let file = File::create(&offsets_path)
+                .with_context(|| format!("Could not create offsets {}", offsets_path.display()))?;
+            let mut offsets_writer =
+                <BufBitWriter<BigEndian, _>>::new(<WordAdapter<usize, _>>::new(BufWriter::new(file)));
+            offsets_writer.write_gamma(0)?;
+
             let mut total_written_bits: u64 = 0;
+            let mut total_offsets_written_bits: u64 = 0;
             let mut total_arcs: u64 = 0;
 
             let mut next_node = 0;
@@ -364,7 +393,10 @@ impl BVComp<()> {
                 job_id,
                 first_node,
                 last_node,
+                chunk_graph_path,
                 written_bits,
+                chunk_offsets_path,
+                offsets_written_bits,
                 num_arcs,
             } in TaskQueue::new(rx.iter())
             {
@@ -378,36 +410,60 @@ impl BVComp<()> {
 
                 next_node = last_node + 1;
                 total_arcs += num_arcs;
-                // compute the path of the bitstream created by this thread
-                let file_path = thread_path(job_id);
                 log::info!(
                     "Copying {} [{}..{}) bits from {} to {}",
                     written_bits,
                     total_written_bits,
                     total_written_bits + written_bits,
-                    file_path.display(),
-                    basename.display()
+                    chunk_graph_path.display(),
+                    graph_path.display()
                 );
                 total_written_bits += written_bits;
 
                 let mut reader =
                     <BufBitReader<E, _>>::new(<WordAdapter<u32, _>>::new(BufReader::new(
-                        File::open(&file_path)
-                            .with_context(|| format!("Could not open {}", file_path.display()))?,
+                        File::open(&chunk_graph_path)
+                            .with_context(|| format!("Could not open {}", chunk_graph_path.display()))?,
                     )));
-                result_writer
+                graph_writer
                     .copy_from(&mut reader, written_bits)
                     .with_context(|| {
                         format!(
                             "Could not copy from {} to {}",
-                            file_path.display(),
+                            chunk_graph_path.display(),
                             graph_path.display()
                         )
                     })?;
+
+                log::info!(
+                    "Copying offsets {} [{}..{}) bits from {} to {}",
+                    written_bits,
+                    total_offsets_written_bits,
+                    total_offsets_written_bits + offsets_written_bits,
+                    chunk_offsets_path.display(),
+                    offsets_path.display()
+                );
+                total_offsets_written_bits += offsets_written_bits;
+
+                let mut reader =
+                    <BufBitReader<BigEndian, _>>::new(<WordAdapter<u32, _>>::new(BufReader::new(
+                        File::open(&chunk_offsets_path)
+                            .with_context(|| format!("Could not open {}", chunk_offsets_path.display()))?,
+                    )));
+                offsets_writer
+                    .copy_from(&mut reader, offsets_written_bits)
+                    .with_context(|| {
+                        format!(
+                            "Could not copy from {} to {}",
+                            chunk_offsets_path.display(),
+                            offsets_path.display()
+                        )
+                    })?;
             }
 
-            log::info!("Flushing the merged Compression bitstream");
-            result_writer.flush()?;
+            log::info!("Flushing the merged bitstreams");
+            graph_writer.flush()?;
+            offsets_writer.flush()?;
 
             log::info!("Writing the .properties file");
             let properties = compression_flags
@@ -427,6 +483,11 @@ impl BVComp<()> {
                 total_written_bits,
                 total_written_bits as f64 / total_arcs as f64
             );
+            log::info!(
+                "Created offsets file with {} bits for {:.4} bits/node",
+                total_offsets_written_bits,
+                total_offsets_written_bits as f64 / num_nodes as f64
+            );
 
             // cleanup the temp files
             std::fs::remove_dir_all(tmp_dir).with_context(|| {

src/utils/sort_pairs.rs

@@ -254,7 +254,12 @@ impl BatchIterator<()> {
         file_path: P,
         batch: &mut [(usize, usize)],
     ) -> anyhow::Result<Self> {
-        Self::new_from_vec_labeled(file_path, unsafe { core::mem::transmute(batch) }, &(), ())
+        Self::new_from_vec_labeled(
+            file_path,
+            unsafe { core::mem::transmute::<&mut [(usize, usize)], &mut [Triple<()>]>(batch) },
+            &(),
+            (),
+        )
     }
     /// Dumps the given triples in `file_path` and returns an iterator over
     /// them, assuming they are already sorted.
@@ -264,7 +269,7 @@ impl BatchIterator<()> {
     ) -> anyhow::Result<Self> {
         Self::new_from_vec_sorted_labeled(
             file_path,
-            unsafe { core::mem::transmute(batch) },
+            unsafe { core::mem::transmute::<&[(usize, usize)], &[Triple<()>]>(batch) },
             &(),
             (),
         )

tests/test_par_bvcomp.rs

@@ -5,6 +5,7 @@
  * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
  */
 
+use std::io::BufReader;
 use std::path::PathBuf;
 
 use anyhow::Result;
@@ -88,8 +89,32 @@ fn _test_par_bvcomp(basename: &str) -> Result<()> {
         }
 
         pr.done();
+
+        let offsets_path = tmp_basename.with_extension(OFFSETS_EXTENSION);
+        let mut offsets_reader =
+            <BufBitReader<BE, _>>::new(<WordAdapter<u32, _>>::new(BufReader::new(
+                std::fs::File::open(&offsets_path)
+                    .expect(&format!("Could not open {}", offsets_path.display())),
+            )));
+
+        let mut pr = ProgressLogger::default();
+        pr.display_memory(true)
+            .item_name("node")
+            .expected_updates(Some(graph.num_nodes()));
+        pr.start("Checking that the generated offsets are correct...");
+
+        let mut offset = 0;
+        for (real_offset, _degree) in comp_graph.offset_deg_iter().by_ref() {
+            let gap_offset = offsets_reader.read_gamma().unwrap();
+            offset += gap_offset;
+            assert_eq!(offset, real_offset);
+            pr.light_update();
+        }
+        pr.done();
+
         // cancel the file at the end
         std::fs::remove_file(tmp_basename.with_extension(GRAPH_EXTENSION))?;
+        std::fs::remove_file(tmp_basename.with_extension(OFFSETS_EXTENSION))?;
         std::fs::remove_file(tmp_basename.with_extension(PROPERTIES_EXTENSION))?;
         log::info!("\n");
     }