Generic support for loading vectors, too

Author: vigna

cli/src/lib.rs

@@ -20,10 +20,11 @@
 
 use anyhow::{Context, Result, anyhow, bail, ensure};
 use clap::{Args, CommandFactory, Parser, Subcommand, ValueEnum};
-use common_traits::{ToBytes, UnsignedInt};
+use common_traits::{AsBytes, FromBytes, ToBytes, UnsignedInt};
 use dsi_bitstream::dispatch::Codes;
+use epserde::deser::Deserialize;
 use epserde::ser::Serialize;
-use std::io::{BufWriter, Write};
+use std::io::{BufRead, BufReader, BufWriter, Read, Write};
 use std::path::{Path, PathBuf};
 use std::time::Duration;
 use std::time::SystemTime;
@@ -280,6 +281,135 @@ impl FloatVectorFormat {
 
         Ok(())
     }
+
+    /// Loads float values from the specified `path` using the format defined
+    /// by `self`.
+    pub fn load<F>(&self, path: impl AsRef<Path>) -> Result<Vec<F>>
+    where
+        F: FromBytes + std::str::FromStr + Copy,
+        <F as AsBytes>::Bytes: for<'a> TryFrom<&'a [u8]>,
+        <F as std::str::FromStr>::Err: std::error::Error + Send + Sync + 'static,
+        Vec<F>: epserde::deser::Deserialize,
+    {
+        let path = path.as_ref();
+        let path_display = path.display();
+
+        match self {
+            FloatVectorFormat::Epserde => {
+                log::info!("Loading ε-serde format from {}", path_display);
+                Ok(unsafe {
+                    <Vec<F>>::load_full(path)
+                        .with_context(|| format!("Could not load vector from {}", path_display))?
+                })
+            }
+            FloatVectorFormat::Java => {
+                log::info!("Loading Java format from {}", path_display);
+                let file = std::fs::File::open(path)
+                    .with_context(|| format!("Could not open {}", path_display))?;
+                let file_len = file.metadata()?.len() as usize;
+                let byte_size = size_of::<F>();
+                ensure!(
+                    file_len % byte_size == 0,
+                    "File size ({}) is not a multiple of {} bytes",
+                    file_len,
+                    byte_size
+                );
+                let n = file_len / byte_size;
+                let mut reader = BufReader::new(file);
+                let mut result = Vec::with_capacity(n);
+                let mut buf = vec![0u8; byte_size];
+                for i in 0..n {
+                    reader.read_exact(&mut buf).with_context(|| {
+                        format!("Could not read value at index {i} from {}", path_display)
+                    })?;
+                    let bytes = buf.as_slice().into();
+                    result.push(F::from_be_bytes(bytes));
+                }
+                Ok(result)
+            }
+            FloatVectorFormat::Ascii => {
+                log::info!("Loading ASCII format from {}", path_display);
+                let file = std::fs::File::open(path)
+                    .with_context(|| format!("Could not open {}", path_display))?;
+                let reader = BufReader::new(file);
+                reader
+                    .lines()
+                    .enumerate()
+                    .filter(|(_, line)| line.as_ref().map_or(true, |l| !l.trim().is_empty()))
+                    .map(|(i, line)| {
+                        let line = line.with_context(|| {
+                            format!("Error reading line {} of {}", i + 1, path_display)
+                        })?;
+                        line.trim().parse::<F>().map_err(|e| {
+                            anyhow!("Error parsing line {} of {}: {}", i + 1, path_display, e)
+                        })
+                    })
+                    .collect()
+            }
+            FloatVectorFormat::Json => {
+                log::info!("Loading JSON format from {}", path_display);
+                let file = std::fs::File::open(path)
+                    .with_context(|| format!("Could not open {}", path_display))?;
+                let mut reader = BufReader::new(file);
+                let mut result = Vec::new();
+                let mut byte = [0u8; 1];
+
+                // Skip whitespace and opening bracket
+                loop {
+                    reader
+                        .read_exact(&mut byte)
+                        .with_context(|| format!("Unexpected end of file in {}", path_display))?;
+                    match byte[0] {
+                        b'[' => break,
+                        b if b.is_ascii_whitespace() => continue,
+                        _ => bail!("Expected '[' at start of JSON array in {}", path_display),
+                    }
+                }
+
+                // Parse comma-separated values until ']'
+                let mut token = String::new();
+                let mut index = 0usize;
+                loop {
+                    reader
+                        .read_exact(&mut byte)
+                        .with_context(|| format!("Unexpected end of file in {}", path_display))?;
+                    match byte[0] {
+                        b']' => {
+                            let trimmed = token.trim();
+                            if !trimmed.is_empty() {
+                                result.push(trimmed.parse::<F>().map_err(|e| {
+                                    anyhow!(
+                                        "Error parsing element {} of {}: {}",
+                                        index + 1,
+                                        path_display,
+                                        e
+                                    )
+                                })?);
+                            }
+                            break;
+                        }
+                        b',' => {
+                            let trimmed = token.trim();
+                            result.push(trimmed.parse::<F>().map_err(|e| {
+                                anyhow!(
+                                    "Error parsing element {} of {}: {}",
+                                    index + 1,
+                                    path_display,
+                                    e
+                                )
+                            })?);
+                            token.clear();
+                            index += 1;
+                        }
+                        c => {
+                            token.push(c as char);
+                        }
+                    }
+                }
+                Ok(result)
+            }
+        }
+    }
 }
 
 #[derive(Debug, Clone, Copy, ValueEnum)]

cli/src/rank/pagerank.rs

@@ -5,12 +5,11 @@
  */
 
 use crate::{FloatVectorFormat, GlobalArgs, GranularityArgs, NumThreadsArg, get_thread_pool};
-use anyhow::{Context, Result, ensure};
+use anyhow::{Result, ensure};
 use clap::Parser;
 use dsi_bitstream::prelude::*;
 use dsi_progress_logger::{ProgressLog, concurrent_progress_logger, progress_logger};
 use predicates::prelude::*;
-use std::io::BufRead;
 use std::path::PathBuf;
 use webgraph::graphs::bvgraph::get_endianness;
 use webgraph::prelude::BvGraph;
@@ -67,9 +66,13 @@ pub struct CliArgs {
     pub threshold: f64,
 
     #[arg(short, long)]
-    /// Path to a preference (personalization) vector (one f64 per line).
+    /// Path to a preference (personalization) vector.
     pub preference: Option<PathBuf>,
 
+    #[arg(long, value_enum, default_value_t = FloatVectorFormat::Ascii)]
+    /// The input format for the preference vector.
+    pub preference_fmt: FloatVectorFormat,
+
     #[arg(short, long, value_enum, default_value_t = CliMode::StronglyPreferential)]
     /// The PageRank mode.
     pub mode: CliMode,
@@ -127,7 +130,11 @@ pub fn pagerank<E: Endianness>(global_args: GlobalArgs, args: CliArgs) -> Result
     );
     let transpose = BvGraph::with_basename(&args.transpose).load()?;
 
-    let preference = args.preference.as_ref().map(load_f64_vector).transpose()?;
+    let preference: Option<Vec<f64>> = args
+        .preference
+        .as_ref()
+        .map(|path| args.preference_fmt.load(path))
+        .transpose()?;
 
     // Build stopping predicate
     let mut predicate = L1Norm::try_from(args.threshold)?.boxed();
@@ -156,27 +163,3 @@ pub fn pagerank<E: Endianness>(global_args: GlobalArgs, args: CliArgs) -> Result
 
     Ok(())
 }
-
-/// Reads a text file containing one `f64` per line.
-fn load_f64_vector(path: impl AsRef<std::path::Path>) -> Result<Vec<f64>> {
-    let path = path.as_ref();
-    let file = std::fs::File::open(path)
-        .with_context(|| format!("Could not open vector file {}", path.display()))?;
-    let reader = std::io::BufReader::new(file);
-    reader
-        .lines()
-        .enumerate()
-        .map(|(i, line)| {
-            let line = line
-                .with_context(|| format!("Error reading line {} of {}", i + 1, path.display()))?;
-            line.trim().parse::<f64>().with_context(|| {
-                format!(
-                    "Error parsing line {} of {}: {:?}",
-                    i + 1,
-                    path.display(),
-                    line
-                )
-            })
-        })
-        .collect()
-}