Skip to main content

webgraph_cli/build/
ef.rs

1/*
2 * SPDX-FileCopyrightText: 2023 Inria
3 * SPDX-FileCopyrightText: 2023 Tommaso Fontana
4 *
5 * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
6 */
7
8use crate::GlobalArgs;
9use anyhow::{Context, Result};
10use clap::Parser;
11use dsi_bitstream::dispatch::factory::CodesReaderFactoryHelper;
12use dsi_bitstream::prelude::*;
13use dsi_progress_logger::prelude::*;
14use epserde::prelude::*;
15use log::info;
16use mmap_rs::MmapFlags;
17use std::fs::File;
18use std::io::{BufReader, BufWriter, Seek};
19use std::path::{Path, PathBuf};
20use sux::prelude::*;
21use webgraph::prelude::*;
22
23#[derive(Parser, Debug, Clone)]
24#[command(name = "ef", about = "Builds the Elias-Fano representation of the offsets of a graph.", long_about = None)]
25pub struct CliArgs {
26    /// The basename of the graph (or labels).
27    pub basename: PathBuf,
28    /// The number of nodes of the graph. When passed, we don't need to load the
29    /// ".properties" file. This allows to build Elias-Fano from the offsets of
30    /// something that might not be a graph but that has offsets, like labels.
31    /// For this reason, if passed, we will also try to read the ".labeloffsets"
32    /// file and then fallback to the usual ".offsets" file.
33    pub number_of_nodes: Option<usize>,
34}
35
36/// Returns the length in bits of the given file.
37fn file_len_bits(path: &Path) -> Result<usize> {
38    let mut file =
39        File::open(path).with_context(|| format!("Could not open {}", path.display()))?;
40    let len = 8 * file
41        .seek(std::io::SeekFrom::End(0))
42        .with_context(|| format!("Could not seek to end of {}", path.display()))?;
43    Ok(len as usize)
44}
45
46pub fn main(global_args: GlobalArgs, args: CliArgs) -> Result<()> {
47    let ef_path = args.basename.with_extension(EF_EXTENSION);
48    // check that ef_path is writable, this is the only portable way I found
49    // to check that the file is writable.
50    if ef_path.exists() && ef_path.metadata()?.permissions().readonly() {
51        return Err(anyhow::anyhow!(
52            "The file is not writable: {}",
53            ef_path.display()
54        ));
55    }
56
57    match get_endianness(&args.basename)?.as_str() {
58        #[cfg(feature = "be_bins")]
59        BE::NAME => build_elias_fano::<BE>(global_args, args),
60        #[cfg(feature = "le_bins")]
61        LE::NAME => build_elias_fano::<LE>(global_args, args),
62        e => panic!("Unknown endianness: {}", e),
63    }
64}
65
66pub fn build_elias_fano<E: Endianness + 'static>(
67    global_args: GlobalArgs,
68    args: CliArgs,
69) -> Result<()>
70where
71    for<'a> BufBitReader<E, MemWordReader<u32, &'a [u32]>>: CodesRead<E> + BitSeek,
72{
73    let mut pl = ProgressLogger::default();
74    pl.display_memory(true).item_name("offset");
75    if let Some(duration) = &global_args.log_interval {
76        pl.log_interval(*duration);
77    }
78
79    let basename = args.basename.clone();
80
81    // When number_of_nodes is provided and label offsets exist, use them
82    // instead of graph offsets.
83    if let Some(num_nodes) = args.number_of_nodes {
84        let label_offsets_path = basename.with_extension(LABELOFFSETS_EXTENSION);
85        if label_offsets_path.exists() {
86            let file_len = file_len_bits(&basename.with_extension(LABELS_EXTENSION))?;
87            pl.expected_updates(Some(num_nodes));
88            let mut efb = EliasFanoBuilder::new(num_nodes + 1, file_len);
89            info!("The label offsets file exists, reading it to build Elias-Fano");
90            let of = <MmapHelper<u32>>::mmap(label_offsets_path, MmapFlags::SEQUENTIAL)?;
91            build_elias_fano_from_offsets(
92                &global_args,
93                &args,
94                num_nodes,
95                of.new_reader(),
96                &mut pl,
97                &mut efb,
98            )?;
99            return serialize_elias_fano(&global_args, &args, efb, &mut pl);
100        }
101    }
102
103    // Standard graph case
104    let of_file_path = basename.with_extension(OFFSETS_EXTENSION);
105
106    let graph_path = basename.with_extension(GRAPH_EXTENSION);
107    info!("Getting size of graph at '{}'", graph_path.display());
108    let file_len = file_len_bits(&graph_path)?;
109    info!("Graph file size: {} bits", file_len);
110
111    // if the num_of_nodes is not present, read it from the properties file
112    // otherwise use the provided value, this is so we can build the Elias-Fano
113    // for offsets of any custom format that might not use the standard
114    // properties file
115    let num_nodes = args
116        .number_of_nodes
117        .map(Ok::<_, anyhow::Error>)
118        .unwrap_or_else(|| {
119            let properties_path = basename.with_extension(PROPERTIES_EXTENSION);
120            info!(
121                "Reading num_of_nodes from properties file at '{}'",
122                properties_path.display()
123            );
124            let f = File::open(&properties_path).with_context(|| {
125                format!(
126                    "Could not open properties file: {}",
127                    properties_path.display()
128                )
129            })?;
130            let map = java_properties::read(BufReader::new(f))?;
131            Ok(map.get("nodes").unwrap().parse::<usize>()?)
132        })?;
133    pl.expected_updates(Some(num_nodes));
134
135    let mut efb = EliasFanoBuilder::new(num_nodes + 1, file_len);
136
137    info!("Checking if offsets exists at '{}'", of_file_path.display());
138    // if the offset files exists, read it to build elias-fano
139    if of_file_path.exists() {
140        info!("The offsets file exists, reading it to build Elias-Fano");
141        let of = <MmapHelper<u32>>::mmap(of_file_path, MmapFlags::SEQUENTIAL)?;
142        build_elias_fano_from_offsets(
143            &global_args,
144            &args,
145            num_nodes,
146            of.new_reader(),
147            &mut pl,
148            &mut efb,
149        )?;
150    } else {
151        build_elias_fano_from_graph(&args, &mut pl, &mut efb)?;
152    }
153
154    serialize_elias_fano(&global_args, &args, efb, &mut pl)
155}
156
157pub fn build_elias_fano_from_graph(
158    args: &CliArgs,
159    pl: &mut impl ProgressLog,
160    efb: &mut EliasFanoBuilder,
161) -> Result<()> {
162    info!("The offsets file does not exists, reading the graph to build Elias-Fano");
163    match get_endianness(&args.basename)?.as_str() {
164        #[cfg(feature = "be_bins")]
165        BE::NAME => build_elias_fano_from_graph_with_endianness::<BE>(args, pl, efb),
166        #[cfg(feature = "le_bins")]
167        LE::NAME => build_elias_fano_from_graph_with_endianness::<LE>(args, pl, efb),
168        e => panic!("Unknown endianness: {}", e),
169    }
170}
171
172pub fn build_elias_fano_from_offsets<E: Endianness>(
173    _global_args: &GlobalArgs,
174    _args: &CliArgs,
175    num_nodes: usize,
176    mut reader: impl GammaRead<E>,
177    pl: &mut impl ProgressLog,
178    efb: &mut EliasFanoBuilder,
179) -> Result<()> {
180    info!("Building Elias-Fano from offsets...");
181
182    // progress bar
183    pl.start("Translating offsets to EliasFano...");
184    // read the graph a write the offsets
185    let mut offset = 0;
186    for _node_id in 0..num_nodes + 1 {
187        // write where
188        offset += reader.read_gamma().context("Could not read gamma")?;
189        efb.push(offset as _);
190        // decode the next nodes so we know where the next node_id starts
191        pl.light_update();
192    }
193    pl.done();
194    Ok(())
195}
196
197pub fn build_elias_fano_from_graph_with_endianness<E: Endianness>(
198    args: &CliArgs,
199    pl: &mut impl ProgressLog,
200    efb: &mut EliasFanoBuilder,
201) -> Result<()>
202where
203    MmapHelper<u32>: CodesReaderFactoryHelper<E>,
204    for<'a> LoadModeCodesReader<'a, E, Mmap>: BitSeek,
205{
206    let seq_graph =
207        webgraph::graphs::bvgraph::sequential::BvGraphSeq::with_basename(&args.basename)
208            .endianness::<E>()
209            .load()
210            .with_context(|| format!("Could not load graph at {}", args.basename.display()))?;
211    // otherwise directly read the graph
212    // progress bar
213    pl.start("Building EliasFano...");
214    // read the graph a write the offsets
215    let mut iter = seq_graph.offset_deg_iter();
216    for (new_offset, _degree) in iter.by_ref() {
217        // write where
218        efb.push(new_offset as _);
219        // decode the next nodes so we know where the next node_id starts
220        pl.light_update();
221    }
222    efb.push(iter.get_pos() as _);
223    Ok(())
224}
225
226pub fn serialize_elias_fano(
227    global_args: &GlobalArgs,
228    args: &CliArgs,
229    efb: EliasFanoBuilder,
230    pl: &mut impl ProgressLog,
231) -> Result<()> {
232    let ef = efb.build();
233    pl.done();
234
235    let mut pl = ProgressLogger::default();
236    pl.display_memory(true);
237    if let Some(duration) = &global_args.log_interval {
238        pl.log_interval(*duration);
239    }
240    pl.start("Building the Index over the ones in the high-bits...");
241    let ef: EF = unsafe { ef.map_high_bits(SelectAdaptConst::<_, _, 12, 4>::new) };
242    pl.done();
243
244    let mut pl = ProgressLogger::default();
245    pl.display_memory(true);
246    if let Some(duration) = &global_args.log_interval {
247        pl.log_interval(*duration);
248    }
249    pl.start("Writing to disk...");
250
251    let ef_path = args.basename.with_extension(EF_EXTENSION);
252    info!("Creating Elias-Fano at '{}'", ef_path.display());
253    let mut ef_file = BufWriter::new(
254        File::create(&ef_path)
255            .with_context(|| format!("Could not create {}", ef_path.display()))?,
256    );
257
258    // serialize and dump the schema to disk
259    unsafe {
260        ef.serialize(&mut ef_file)
261            .with_context(|| format!("Could not serialize EliasFano to {}", ef_path.display()))
262    }?;
263
264    pl.done();
265    Ok(())
266}