28829100046

Committed 06 Jul 2026 11:00PM UTC coverage: 72.759% (+3.5%) from 69.212%

Build # 28829100046

Build Type

Pull #171

github

Committed by

zommiommy

Commit Message

test: guard mmap-dependent new tests from Miri

The default and explicit LoadMmap load paths mprotect anonymous maps,
which Miri does not support; the zero-node roundtrip now runs only its
LoadMem loads under Miri. The new CLI integration tests are excluded
wholesale like the preexisting ones, as the CLI loads graphs and
Elias-Fano structures through mmap.

Pull Request Pull Request #171: Fixes from a full-workspace code review: correctness, error handling, robustness

Coverage Stats

356 of 482 new or added lines in 60 files covered. (73.86%)

9 existing lines in 6 files now uncovered.

8157 of 11211 relevant lines covered (72.76%)

49198900.83 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

81.82

/webgraph/src/utils/batch_codec/grouped_gaps.rs

/*
 * SPDX-FileCopyrightText: 2025 Tommaso Fontana
 * SPDX-FileCopyrightText: 2025 Sebastiano Vigna
 *
 * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
 */

use super::{BitReader, BitWriter};
use crate::traits::SortedIterator;
use crate::utils::{ArcMmapHelper, MmapHelper, Triple};
use crate::{
    traits::{BitDeserializer, BitSerializer},
    utils::{BatchCodec, humanize},
};

use std::sync::Arc;

use anyhow::{Context, Result};
use dsi_bitstream::prelude::*;
use mmap_rs::MmapFlags;
use rdst::*;

/// A codec for encoding and decoding batches of triples using grouped gap compression.
///
/// This codec encodes triples of the form `(src, dst, label)` by grouping edges
/// with the same source node, and encoding the gaps between consecutive sources
/// and destinations using a specified code (default: gamma). The outdegree
/// (number of edges for each source) is also encoded using the specified code.
///
/// # Type Parameters
///
/// * `SD` - A type implementing both [`BitSerializer`] and [`BitDeserializer`]
///   for the label type. Use [`BitSerDeser`] to combine separate serializer and
///   deserializer implementations.
/// * `OUTDEGREE_CODE` - Code used for encoding outdegrees (default: [ɣ]).
/// * `SRC_CODE` - Code used for encoding source gaps (default: [ɣ]).
/// * `DST_CODE` - Code used for encoding destination gaps (default: [ɣ]).
///
/// # Encoding Format
///
/// 1. The batch length is written using delta coding.
/// 2. For each group of triples with the same source:
///     - The gap from the previous source is encoded.
///     - The outdegree (number of edges for this source) is encoded.
///     - For each destination:
///         - The gap from the previous destination is encoded.
///         - The label is serialized.
///
/// `SD` must be [`Clone`] because we need one copy for each
/// [`GroupedGapsIter`], and there are possible scenarios in which the
/// deserializer might be stateful.
///
/// [ɣ]: dsi_bitstream::codes::gamma
/// [`BitSerDeser`]: crate::traits::BitSerDeser
#[derive(Clone, Debug)]
pub struct GroupedGapsCodec<
    E: Endianness = NE,
    SD: BitSerializer<E, BitWriter<E>>
        + BitDeserializer<E, BitReader<E>, DeserType = <SD as BitSerializer<E, BitWriter<E>>>::SerType>
        + Clone = (),
    const OUTDEGREE_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
    const SRC_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
    const DST_CODE: usize = { dsi_bitstream::dispatch::code_consts::DELTA },
    const DEDUP: bool = false,
> where
    BitReader<E>: BitRead<E>,
    BitWriter<E>: BitWrite<E>,
{
    /// Serializer and deserializer for the labels.
    pub sd: SD,

    _marker: core::marker::PhantomData<E>,
}

impl<
    E,
    SD,
    const OUTDEGREE_CODE: usize,
    const SRC_CODE: usize,
    const DST_CODE: usize,
    const DEDUP: bool,
> GroupedGapsCodec<E, SD, OUTDEGREE_CODE, SRC_CODE, DST_CODE, DEDUP>
where
    E: Endianness,
    SD: BitSerializer<E, BitWriter<E>>
        + BitDeserializer<E, BitReader<E>, DeserType = SD::SerType>
        + Send
        + Sync
        + Clone,
    BitReader<E>: BitRead<E>,
    BitWriter<E>: BitWrite<E>,
{
    /// Creates a new `GroupedGapsCodec` with the given serializer/deserializer.
    pub const fn new(sd: SD) -> Self {
        Self {
            sd,
            _marker: core::marker::PhantomData,
        }
    }
}

impl<
    E: Endianness,
    SD: BitSerializer<E, BitWriter<E>>
        + BitDeserializer<E, BitReader<E>, DeserType = SD::SerType>
        + Clone
        + Default,
    const OUTDEGREE_CODE: usize,
    const SRC_CODE: usize,
    const DST_CODE: usize,
    const DEDUP: bool,
> Default for GroupedGapsCodec<E, SD, OUTDEGREE_CODE, SRC_CODE, DST_CODE, DEDUP>
where
    BitReader<E>: BitRead<E>,
    BitWriter<E>: BitWrite<E>,
{
    fn default() -> Self {
        Self {
            sd: SD::default(),
            _marker: core::marker::PhantomData,
        }
    }
}

/// Statistics about the encoding performed by [`GroupedGapsCodec`].
#[derive(Debug, Clone, Copy)]
pub struct GroupedGapsStats {
    /// Total number of triples encoded
    pub total_triples: usize,
    /// Number of bits used for outdegrees
    pub outdegree_bits: usize,
    /// Number of bits used for source gaps
    pub src_bits: usize,
    /// Number of bits used for destination gaps
    pub dst_bits: usize,
    /// Number of bits used for labels
    pub labels_bits: usize,
}

impl core::fmt::Display for GroupedGapsStats {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(
            f,
            "outdegree: {}B ({:.3} bits / triple), src: {}B ({:.3} bits / triple), dst: {}B ({:.3} bits / triple), labels: {}B ({:.3} bits / triple)",
            humanize(self.outdegree_bits as f64 / 8.0),
            self.outdegree_bits as f64 / self.total_triples as f64,
            humanize(self.src_bits as f64 / 8.0),
            self.src_bits as f64 / self.total_triples as f64,
            humanize(self.dst_bits as f64 / 8.0),
            self.dst_bits as f64 / self.total_triples as f64,
            humanize(self.labels_bits as f64 / 8.0),
            self.labels_bits as f64 / self.total_triples as f64,
        )
    }
}

impl crate::utils::BatchStats for GroupedGapsStats {
    #[inline]
    fn total_triples(&self) -> usize {
        self.total_triples
    }
}

impl<
    E,
    SD,
    const OUTDEGREE_CODE: usize,
    const SRC_CODE: usize,
    const DST_CODE: usize,
    const DEDUP: bool,
> BatchCodec for GroupedGapsCodec<E, SD, OUTDEGREE_CODE, SRC_CODE, DST_CODE, DEDUP>
where
    E: Endianness,
    SD: BitSerializer<E, BitWriter<E>>
        + BitDeserializer<E, BitReader<E>, DeserType = SD::SerType>
        + Send
        + Sync
        + Clone,
    SD::SerType: Send + Sync + Copy + 'static, // needed by radix sort
    BitReader<E>: BitRead<E> + CodesRead<E>,
    BitWriter<E>: BitWrite<E> + CodesWrite<E>,
{
    type Label = SD::SerType;
    type DecodedBatch = GroupedGapsIter<E, SD, OUTDEGREE_CODE, SRC_CODE, DST_CODE>;
    type EncodedBatchStats = GroupedGapsStats;

    fn encode_batch(
        &self,
        path: impl AsRef<std::path::Path>,
        batch: &mut [((usize, usize), Self::Label)],
    ) -> Result<(usize, Self::EncodedBatchStats)> {
        let start = std::time::Instant::now();
        Triple::cast_batch_mut(batch).radix_sort_unstable();
        log::debug!("Sorted {} arcs in {:?}", batch.len(), start.elapsed());
        self.encode_sorted_batch(path, batch)
    }

    fn encode_sorted_batch(
        &self,
        path: impl AsRef<std::path::Path>,
        batch: &[((usize, usize), Self::Label)],
    ) -> Result<(usize, Self::EncodedBatchStats)> {
        debug_assert!(Triple::cast_batch(batch).is_sorted(), "Batch is not sorted");
        // create a bitstream to write to the file
        let file_path = path.as_ref();
        let mut stream = buf_bit_writer::from_path::<E, usize>(file_path).with_context(|| {
            format!(
                "Could not create BatchIterator temporary file {}",
                file_path.display()
            )
        })?;

        // Pre-count unique pairs for the length prefix when deduplicating
        let batch_len = if DEDUP {
            if batch.is_empty() {
                0
            } else {
                1 + batch.windows(2).filter(|w| w[0].0 != w[1].0).count()
            }
        } else {
            batch.len()
        };

        // prefix the stream with the length of the batch
        // we use a delta code since it'll be a big number most of the time
        stream
            .write_delta(batch_len as u64)
            .context("Could not write length")?;

        let mut stats = GroupedGapsStats {
            total_triples: batch_len,
            outdegree_bits: 0,
            src_bits: 0,
            dst_bits: 0,
            labels_bits: 0,
        };
        // dump the triples to the bitstream
        let mut prev_src = 0;
        let mut i = 0;
        while i < batch.len() {
            let ((src, _), _) = batch[i];
            // write the source gap as gamma
            stats.src_bits += ConstCode::<SRC_CODE>
                .write(&mut stream, (src - prev_src) as _)
                .with_context(|| format!("Could not write {src} after {prev_src}"))?;
            // figure out how many edges have this source
            let group_size = batch[i..].iter().take_while(|t| t.0.0 == src).count();

            // compute outdegree (unique destinations when deduplicating)
            let outdegree = if DEDUP {
                let group = &batch[i..i + group_size];
                if group.is_empty() {
                    0
                } else {
                    1 + group.windows(2).filter(|w| w[0].0.1 != w[1].0.1).count()
                }
            } else {
                group_size
            };

            // write the outdegree
            stats.outdegree_bits += ConstCode::<OUTDEGREE_CODE>
                .write(&mut stream, outdegree as _)
                .with_context(|| format!("Could not write outdegree {outdegree} for {src}"))?;

            // encode the destinations
            let mut prev_dst = 0;
            let mut last_written_dst: Option<usize> = None;
            let end = i + group_size;
            while i < end {
                let ((_, dst), label) = &batch[i];
                if DEDUP && last_written_dst == Some(*dst) {
                    i += 1;
                    continue;
                }
                if DEDUP {
                    last_written_dst = Some(*dst);
                }
                // write the destination gap as gamma
                stats.dst_bits += ConstCode::<DST_CODE>
                    .write(&mut stream, (dst - prev_dst) as _)
                    .with_context(|| format!("Could not write {dst} after {prev_dst}"))?;
                // write the label
                stats.labels_bits += self
                    .sd
                    .serialize(label, &mut stream)
                    .context("Could not serialize label")?;
                prev_dst = *dst;
                i += 1;
            }
            prev_src = src;
        }
        // flush the stream and reset the buffer
        stream.flush().context("Could not flush stream")?;

        let total_bits = stats.outdegree_bits + stats.src_bits + stats.dst_bits + stats.labels_bits;
        Ok((total_bits, stats))
    }

    fn decode_batch(&self, path: impl AsRef<std::path::Path>) -> Result<Self::DecodedBatch> {
        // open the file
        let mut stream = <BufBitReader<E, _>>::new(MemWordReader::new(ArcMmapHelper(Arc::new(
            MmapHelper::mmap(
                path.as_ref(),
                MmapFlags::TRANSPARENT_HUGE_PAGES | MmapFlags::SEQUENTIAL,
            )
            .with_context(|| format!("Could not mmap {}", path.as_ref().display()))?,
        ))));

        // read the length of the batch (first value in the stream)
        let len = stream.read_delta().context("Could not read length")? as usize;

        // create the iterator
        Ok(GroupedGapsIter {
            deserializer: self.sd.clone(),
            stream,
            len,
            current: 0,
            src: 0,
            dst_left: 0,
            prev_dst: 0,
        })
    }
}

/// An iterator over triples encoded with gaps, this is returned by [`GroupedGapsCodec`].
#[derive(Clone, Debug)]
pub struct GroupedGapsIter<
    E: Endianness = NE,
    D: BitDeserializer<E, BitReader<E>> = (),
    const OUTDEGREE_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
    const SRC_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
    const DST_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
> where
    BitReader<E>: BitRead<E>,
    BitWriter<E>: BitWrite<E>,
{
    /// Deserializer for the labels
    deserializer: D,
    /// Bitstream to read from
    stream: BitReader<E>,
    /// Length of the iterator (number of triples)
    len: usize,
    /// Current position in the iterator
    current: usize,
    /// Current source node
    src: usize,
    /// Number of destinations left for the current source
    dst_left: usize,
    /// Previous destination node
    prev_dst: usize,
}

// SAFETY: gaps are decoded in non-decreasing (src, dst) order.
unsafe impl<
    E: Endianness,
    D: BitDeserializer<E, BitReader<E>>,
    const OUTDEGREE_CODE: usize,
    const SRC_CODE: usize,
    const DST_CODE: usize,
> SortedIterator for GroupedGapsIter<E, D, OUTDEGREE_CODE, SRC_CODE, DST_CODE>
where
    BitReader<E>: BitRead<E> + CodesRead<E>,
    BitWriter<E>: BitWrite<E> + CodesWrite<E>,
{
}

impl<
    E: Endianness,
    D: BitDeserializer<E, BitReader<E>>,
    const OUTDEGREE_CODE: usize,
    const SRC_CODE: usize,
    const DST_CODE: usize,
> Iterator for GroupedGapsIter<E, D, OUTDEGREE_CODE, SRC_CODE, DST_CODE>
where
    BitReader<E>: BitRead<E> + CodesRead<E>,
    BitWriter<E>: BitWrite<E> + CodesWrite<E>,
{
    type Item = ((usize, usize), D::DeserType);
    fn next(&mut self) -> Option<Self::Item> {
        if self.current >= self.len {
            return None;
        }
        if self.dst_left == 0 {
            // read a new source
            let src_gap = ConstCode::<SRC_CODE>.read(&mut self.stream).ok()?;
            self.src += src_gap as usize;
            // read the outdegree
            self.dst_left = ConstCode::<OUTDEGREE_CODE>.read(&mut self.stream).ok()? as usize;
            self.prev_dst = 0;
        }

        let dst_gap = ConstCode::<DST_CODE>.read(&mut self.stream).ok()?;
        let label = self.deserializer.deserialize(&mut self.stream).ok()?;
        self.prev_dst += dst_gap as usize;
        self.current += 1;
        self.dst_left -= 1;
        Some(((self.src, self.prev_dst), label))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.len(), Some(self.len()))
    }

    fn count(self) -> usize {
        self.len()
    }
}

impl<
    E: Endianness,
    D: BitDeserializer<E, BitReader<E>>,
    const OUTDEGREE_CODE: usize,
    const SRC_CODE: usize,
    const DST_CODE: usize,
> ExactSizeIterator for GroupedGapsIter<E, D, OUTDEGREE_CODE, SRC_CODE, DST_CODE>
where
    BitReader<E>: BitRead<E> + CodesRead<E>,
    BitWriter<E>: BitWrite<E> + CodesWrite<E>,
{
    fn len(&self) -> usize {
        self.len - self.current
    }
}

impl<
    E: Endianness,
    D: BitDeserializer<E, BitReader<E>>,
    const OUTDEGREE_CODE: usize,
    const SRC_CODE: usize,
    const DST_CODE: usize,
> core::iter::FusedIterator for GroupedGapsIter<E, D, OUTDEGREE_CODE, SRC_CODE, DST_CODE>
where
    BitReader<E>: BitRead<E> + CodesRead<E>,
    BitWriter<E>: BitWrite<E> + CodesWrite<E>,
{
}

1	/*
2	* SPDX-FileCopyrightText: 2025 Tommaso Fontana
3	* SPDX-FileCopyrightText: 2025 Sebastiano Vigna
4	*
5	* SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
6	*/
7
8	use super::{BitReader, BitWriter};
9	use crate::traits::SortedIterator;
10	use crate::utils::{ArcMmapHelper, MmapHelper, Triple};
11	use crate::{
12	traits::{BitDeserializer, BitSerializer},
13	utils::{BatchCodec, humanize},
14	};
15
16	use std::sync::Arc;
17
18	use anyhow::{Context, Result};
19	use dsi_bitstream::prelude::*;
20	use mmap_rs::MmapFlags;
21	use rdst::*;
22
23	/// A codec for encoding and decoding batches of triples using grouped gap compression.
24	///
25	/// This codec encodes triples of the form `(src, dst, label)` by grouping edges
26	/// with the same source node, and encoding the gaps between consecutive sources
27	/// and destinations using a specified code (default: gamma). The outdegree
28	/// (number of edges for each source) is also encoded using the specified code.
29	///
30	/// # Type Parameters
31	///
32	/// * `SD` - A type implementing both [`BitSerializer`] and [`BitDeserializer`]
33	/// for the label type. Use [`BitSerDeser`] to combine separate serializer and
34	/// deserializer implementations.
35	/// * `OUTDEGREE_CODE` - Code used for encoding outdegrees (default: [ɣ]).
36	/// * `SRC_CODE` - Code used for encoding source gaps (default: [ɣ]).
37	/// * `DST_CODE` - Code used for encoding destination gaps (default: [ɣ]).
38	///
39	/// # Encoding Format
40	///
41	/// 1. The batch length is written using delta coding.
42	/// 2. For each group of triples with the same source:
43	/// - The gap from the previous source is encoded.
44	/// - The outdegree (number of edges for this source) is encoded.
45	/// - For each destination:
46	/// - The gap from the previous destination is encoded.
47	/// - The label is serialized.
48	///
49	/// `SD` must be [`Clone`] because we need one copy for each
50	/// [`GroupedGapsIter`], and there are possible scenarios in which the
51	/// deserializer might be stateful.
52	///
53	/// [ɣ]: dsi_bitstream::codes::gamma
54	/// [`BitSerDeser`]: crate::traits::BitSerDeser
55	#[derive(Clone, Debug)]
56	pub struct GroupedGapsCodec<
57	E: Endianness = NE,
58	SD: BitSerializer<E, BitWriter<E>>
59	+ BitDeserializer<E, BitReader<E>, DeserType = <SD as BitSerializer<E, BitWriter<E>>>::SerType>
60	+ Clone = (),
61	const OUTDEGREE_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
62	const SRC_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
63	const DST_CODE: usize = { dsi_bitstream::dispatch::code_consts::DELTA },
64	const DEDUP: bool = false,
65	> where
66	BitReader<E>: BitRead<E>,
67	BitWriter<E>: BitWrite<E>,
68	{
69	/// Serializer and deserializer for the labels.
70	pub sd: SD,
71
72	_marker: core::marker::PhantomData<E>,
73	}
74
75	impl<
76	E,
77	SD,
78	const OUTDEGREE_CODE: usize,
79	const SRC_CODE: usize,
80	const DST_CODE: usize,
81	const DEDUP: bool,
82	> GroupedGapsCodec<E, SD, OUTDEGREE_CODE, SRC_CODE, DST_CODE, DEDUP>
83	where
84	E: Endianness,
85	SD: BitSerializer<E, BitWriter<E>>
86	+ BitDeserializer<E, BitReader<E>, DeserType = SD::SerType>
87	+ Send
88	+ Sync
89	+ Clone,
90	BitReader<E>: BitRead<E>,
91	BitWriter<E>: BitWrite<E>,
92	{
93	/// Creates a new `GroupedGapsCodec` with the given serializer/deserializer.
94	pub const fn new(sd: SD) -> Self {	146✔
95	Self {
96	sd,
97	_marker: core::marker::PhantomData,
98	}
99	}
100	}
101
102	impl<
103	E: Endianness,
104	SD: BitSerializer<E, BitWriter<E>>
105	+ BitDeserializer<E, BitReader<E>, DeserType = SD::SerType>
106	+ Clone
107	+ Default,
108	const OUTDEGREE_CODE: usize,
109	const SRC_CODE: usize,
110	const DST_CODE: usize,
111	const DEDUP: bool,
112	> Default for GroupedGapsCodec<E, SD, OUTDEGREE_CODE, SRC_CODE, DST_CODE, DEDUP>
113	where
114	BitReader<E>: BitRead<E>,
115	BitWriter<E>: BitWrite<E>,
116	{
117	fn default() -> Self {	20✔
118	Self {
119	sd: SD::default(),	20✔
120	_marker: core::marker::PhantomData,
121	}
122	}
123	}
124
125	/// Statistics about the encoding performed by [`GroupedGapsCodec`].
126	#[derive(Debug, Clone, Copy)]
127	pub struct GroupedGapsStats {
128	/// Total number of triples encoded
129	pub total_triples: usize,
130	/// Number of bits used for outdegrees
131	pub outdegree_bits: usize,
132	/// Number of bits used for source gaps
133	pub src_bits: usize,
134	/// Number of bits used for destination gaps
135	pub dst_bits: usize,
136	/// Number of bits used for labels
137	pub labels_bits: usize,
138	}
139
140	impl core::fmt::Display for GroupedGapsStats {
141	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {	×
142	write!(	×
143	f,	×
144	"outdegree: {}B ({:.3} bits / triple), src: {}B ({:.3} bits / triple), dst: {}B ({:.3} bits / triple), labels: {}B ({:.3} bits / triple)",
145	humanize(self.outdegree_bits as f64 / 8.0),	×
146	self.outdegree_bits as f64 / self.total_triples as f64,	×
147	humanize(self.src_bits as f64 / 8.0),	×
148	self.src_bits as f64 / self.total_triples as f64,	×
149	humanize(self.dst_bits as f64 / 8.0),	×
150	self.dst_bits as f64 / self.total_triples as f64,	×
151	humanize(self.labels_bits as f64 / 8.0),	×
152	self.labels_bits as f64 / self.total_triples as f64,	×
153	)
154	}
155	}
156
157	impl crate::utils::BatchStats for GroupedGapsStats {
158	#[inline]
159	fn total_triples(&self) -> usize {	×
160	self.total_triples	×
161	}
162	}
163
164	impl<
165	E,
166	SD,
167	const OUTDEGREE_CODE: usize,
168	const SRC_CODE: usize,
169	const DST_CODE: usize,
170	const DEDUP: bool,
171	> BatchCodec for GroupedGapsCodec<E, SD, OUTDEGREE_CODE, SRC_CODE, DST_CODE, DEDUP>
172	where
173	E: Endianness,
174	SD: BitSerializer<E, BitWriter<E>>
175	+ BitDeserializer<E, BitReader<E>, DeserType = SD::SerType>
176	+ Send
177	+ Sync
178	+ Clone,
179	SD::SerType: Send + Sync + Copy + 'static, // needed by radix sort
180	BitReader<E>: BitRead<E> + CodesRead<E>,
181	BitWriter<E>: BitWrite<E> + CodesWrite<E>,
182	{
183	type Label = SD::SerType;
184	type DecodedBatch = GroupedGapsIter<E, SD, OUTDEGREE_CODE, SRC_CODE, DST_CODE>;
185	type EncodedBatchStats = GroupedGapsStats;
186
187	fn encode_batch(	778✔
188	&self,
189	path: impl AsRef<std::path::Path>,
190	batch: &mut [((usize, usize), Self::Label)],
191	) -> Result<(usize, Self::EncodedBatchStats)> {
192	let start = std::time::Instant::now();	1,556✔
193	Triple::cast_batch_mut(batch).radix_sort_unstable();	1,556✔
194	log::debug!("Sorted {} arcs in {:?}", batch.len(), start.elapsed());	778✔
195	self.encode_sorted_batch(path, batch)	3,112✔
196	}
197
198	fn encode_sorted_batch(	778✔
199	&self,
200	path: impl AsRef<std::path::Path>,
201	batch: &[((usize, usize), Self::Label)],
202	) -> Result<(usize, Self::EncodedBatchStats)> {
203	debug_assert!(Triple::cast_batch(batch).is_sorted(), "Batch is not sorted");	2,334✔
204	// create a bitstream to write to the file
205	let file_path = path.as_ref();	2,334✔
206	let mut stream = buf_bit_writer::from_path::<E, usize>(file_path).with_context(\|\| {	3,112✔
207	format!(	×
208	"Could not create BatchIterator temporary file {}",	×
209	file_path.display()	×
210	)
211	})?;
212
213	// Pre-count unique pairs for the length prefix when deduplicating
214	let batch_len = if DEDUP {	1,556✔
215	if batch.is_empty() {	14✔
UNCOV 216	0	×
217	} else {
218	1 + batch.windows(2).filter(\|w\| w[0].0 != w[1].0).count()	38✔
219	}
220	} else {
221	batch.len()	1,542✔
222	};
223
224	// prefix the stream with the length of the batch
225	// we use a delta code since it'll be a big number most of the time
226	stream	778✔
227	.write_delta(batch_len as u64)	1,556✔
228	.context("Could not write length")?;
229
230	let mut stats = GroupedGapsStats {
231	total_triples: batch_len,
232	outdegree_bits: 0,
233	src_bits: 0,
234	dst_bits: 0,
235	labels_bits: 0,
236	};
237	// dump the triples to the bitstream
238	let mut prev_src = 0;	1,556✔
239	let mut i = 0;	1,556✔
240	while i < batch.len() {	6,540,908✔
241	let ((src, _), _) = batch[i];	6,539,352✔
242	// write the source gap as gamma
243	stats.src_bits += ConstCode::<SRC_CODE>	3,269,676✔
244	.write(&mut stream, (src - prev_src) as _)	9,809,028✔
245	.with_context(\|\| format!("Could not write {src} after {prev_src}"))?;	3,269,676✔
246	// figure out how many edges have this source
247	let group_size = batch[i..].iter().take_while(\|t\| t.0.0 == src).count();	124,400,128✔
248
249	// compute outdegree (unique destinations when deduplicating)
250	let outdegree = if DEDUP {	6,539,352✔
251	let group = &batch[i..i + group_size];	36✔
252	if group.is_empty() {	18✔
253	0	×
254	} else {
255	1 + group.windows(2).filter(\|w\| w[0].0.1 != w[1].0.1).count()	42✔
256	}
257	} else {
258	group_size	3,269,667✔
259	};
260
261	// write the outdegree
262	stats.outdegree_bits += ConstCode::<OUTDEGREE_CODE>	3,269,676✔
263	.write(&mut stream, outdegree as _)	9,809,028✔
264	.with_context(\|\| format!("Could not write outdegree {outdegree} for {src}"))?;	3,269,676✔
265
266	// encode the destinations
267	let mut prev_dst = 0;	6,539,352✔
268	let mut last_written_dst: Option<usize> = None;	9,809,028✔
269	let end = i + group_size;	6,539,352✔
270	while i < end {	54,026,652✔
271	let ((_, dst), label) = &batch[i];	152,270,928✔
272	if DEDUP && last_written_dst == Some(*dst) {	50,756,988✔
273	i += 1;	3✔
274	continue;	3✔
275	}
276	if DEDUP {	50,756,982✔
277	last_written_dst = Some(*dst);	9✔
278	}
279	// write the destination gap as gamma
280	stats.dst_bits += ConstCode::<DST_CODE>	50,756,973✔
281	.write(&mut stream, (dst - prev_dst) as _)	152,270,919✔
282	.with_context(\|\| format!("Could not write {dst} after {prev_dst}"))?;	50,756,973✔
283	// write the label
284	stats.labels_bits += self	50,756,973✔
285	.sd	50,756,973✔
286	.serialize(label, &mut stream)	152,270,919✔
287	.context("Could not serialize label")?;	50,756,973✔
288	prev_dst = *dst;	50,756,973✔
289	i += 1;	50,756,973✔
290	}
291	prev_src = src;	3,269,676✔
292	}
293	// flush the stream and reset the buffer
294	stream.flush().context("Could not flush stream")?;	2,334✔
295
296	let total_bits = stats.outdegree_bits + stats.src_bits + stats.dst_bits + stats.labels_bits;	1,556✔
297	Ok((total_bits, stats))	778✔
298	}
299
300	fn decode_batch(&self, path: impl AsRef<std::path::Path>) -> Result<Self::DecodedBatch> {	778✔
301	// open the file
302	let mut stream = <BufBitReader<E, _>>::new(MemWordReader::new(ArcMmapHelper(Arc::new(	3,112✔
303	MmapHelper::mmap(	778✔
304	path.as_ref(),	1,556✔
305	MmapFlags::TRANSPARENT_HUGE_PAGES \| MmapFlags::SEQUENTIAL,	778✔
306	)
307	.with_context(\|\| format!("Could not mmap {}", path.as_ref().display()))?,	778✔
308	))));
309
310	// read the length of the batch (first value in the stream)
311	let len = stream.read_delta().context("Could not read length")? as usize;	3,112✔
312
313	// create the iterator
314	Ok(GroupedGapsIter {	778✔
315	deserializer: self.sd.clone(),	2,334✔
316	stream,	778✔
317	len,	778✔
318	current: 0,	778✔
319	src: 0,	778✔
320	dst_left: 0,	778✔
321	prev_dst: 0,	778✔
322	})
323	}
324	}
325
326	/// An iterator over triples encoded with gaps, this is returned by [`GroupedGapsCodec`].
327	#[derive(Clone, Debug)]
328	pub struct GroupedGapsIter<
329	E: Endianness = NE,
330	D: BitDeserializer<E, BitReader<E>> = (),
331	const OUTDEGREE_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
332	const SRC_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
333	const DST_CODE: usize = { dsi_bitstream::dispatch::code_consts::GAMMA },
334	> where
335	BitReader<E>: BitRead<E>,
336	BitWriter<E>: BitWrite<E>,
337	{
338	/// Deserializer for the labels
339	deserializer: D,
340	/// Bitstream to read from
341	stream: BitReader<E>,
342	/// Length of the iterator (number of triples)
343	len: usize,
344	/// Current position in the iterator
345	current: usize,
346	/// Current source node
347	src: usize,
348	/// Number of destinations left for the current source
349	dst_left: usize,
350	/// Previous destination node
351	prev_dst: usize,
352	}
353
354	// SAFETY: gaps are decoded in non-decreasing (src, dst) order.
355	unsafe impl<
356	E: Endianness,
357	D: BitDeserializer<E, BitReader<E>>,
358	const OUTDEGREE_CODE: usize,
359	const SRC_CODE: usize,
360	const DST_CODE: usize,
361	> SortedIterator for GroupedGapsIter<E, D, OUTDEGREE_CODE, SRC_CODE, DST_CODE>
362	where
363	BitReader<E>: BitRead<E> + CodesRead<E>,
364	BitWriter<E>: BitWrite<E> + CodesWrite<E>,
365	{
366	}
367
368	impl<
369	E: Endianness,
370	D: BitDeserializer<E, BitReader<E>>,
371	const OUTDEGREE_CODE: usize,
372	const SRC_CODE: usize,
373	const DST_CODE: usize,
374	> Iterator for GroupedGapsIter<E, D, OUTDEGREE_CODE, SRC_CODE, DST_CODE>
375	where
376	BitReader<E>: BitRead<E> + CodesRead<E>,
377	BitWriter<E>: BitWrite<E> + CodesWrite<E>,
378	{
379	type Item = ((usize, usize), D::DeserType);
380	fn next(&mut self) -> Option<Self::Item> {	50,757,750✔
381	if self.current >= self.len {	50,757,750✔
382	return None;	777✔
383	}
384	if self.dst_left == 0 {	50,756,973✔
385	// read a new source
386	let src_gap = ConstCode::<SRC_CODE>.read(&mut self.stream).ok()?;	16,348,375✔
387	self.src += src_gap as usize;	3,269,675✔
388	// read the outdegree
389	self.dst_left = ConstCode::<OUTDEGREE_CODE>.read(&mut self.stream).ok()? as usize;	13,078,700✔
390	self.prev_dst = 0;	3,269,675✔
391	}
392
393	let dst_gap = ConstCode::<DST_CODE>.read(&mut self.stream).ok()?;	253,784,865✔
394	let label = self.deserializer.deserialize(&mut self.stream).ok()?;	253,784,865✔
395	self.prev_dst += dst_gap as usize;	50,756,973✔
396	self.current += 1;	50,756,973✔
397	self.dst_left -= 1;	50,756,973✔
398	Some(((self.src, self.prev_dst), label))	50,756,973✔
399	}
400
401	fn size_hint(&self) -> (usize, Option<usize>) {	2✔
402	(self.len(), Some(self.len()))	6✔
403	}
404
405	fn count(self) -> usize {	×
406	self.len()	×
407	}
408	}
409
410	impl<
411	E: Endianness,
412	D: BitDeserializer<E, BitReader<E>>,
413	const OUTDEGREE_CODE: usize,
414	const SRC_CODE: usize,
415	const DST_CODE: usize,
416	> ExactSizeIterator for GroupedGapsIter<E, D, OUTDEGREE_CODE, SRC_CODE, DST_CODE>
417	where
418	BitReader<E>: BitRead<E> + CodesRead<E>,
419	BitWriter<E>: BitWrite<E> + CodesWrite<E>,
420	{
421	fn len(&self) -> usize {	5✔
422	self.len - self.current	5✔
423	}
424	}
425
426	impl<
427	E: Endianness,
428	D: BitDeserializer<E, BitReader<E>>,
429	const OUTDEGREE_CODE: usize,
430	const SRC_CODE: usize,
431	const DST_CODE: usize,
432	> core::iter::FusedIterator for GroupedGapsIter<E, D, OUTDEGREE_CODE, SRC_CODE, DST_CODE>
433	where
434	BitReader<E>: BitRead<E> + CodesRead<E>,
435	BitWriter<E>: BitWrite<E> + CodesWrite<E>,
436	{
437	}

vigna / webgraph-rs / 28829100046

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous