ruwex 0.1.0

Fast Rust rewrite of wikiextractor: extract and clean text from Wikimedia XML dumps
Documentation
//! Opening a dump file and picking the best way to read it.

use std::fs::File;
use std::io::{self, BufRead, BufReader};
use std::path::{Path, PathBuf};

use bzip2::read::MultiBzDecoder;

use crate::dump::multistream;
use crate::error::Result;

const BUFFER_SIZE: usize = 64 * 1024;

/// A source of dump pages. Multistream dumps (with their index file present)
/// get a parallel seek-and-decompress fast path; everything else is read
/// sequentially.
pub enum PageSource {
    Sequential(Box<dyn BufRead + Send>),
    Multistream(MultistreamSource),
}

pub struct MultistreamSource {
    pub path: PathBuf,
    /// Start offsets of the page streams (header stream excluded).
    pub offsets: Vec<u64>,
}

impl PageSource {
    /// Opens a dump file: plain XML, `.bz2`, or multistream `.bz2` with its
    /// index file alongside.
    pub fn open(path: &Path) -> Result<Self> {
        if path.extension().is_some_and(|e| e == "bz2") {
            if let Some(index_path) = multistream::index_path_for(path) {
                let offsets = multistream::read_stream_offsets(&index_path)?;
                log::info!(
                    "using multistream index {} ({} streams)",
                    index_path.display(),
                    offsets.len()
                );
                return Ok(Self::Multistream(MultistreamSource {
                    path: path.to_path_buf(),
                    offsets,
                }));
            }
            let file = BufReader::with_capacity(BUFFER_SIZE, File::open(path)?);
            Ok(Self::Sequential(Box::new(BufReader::with_capacity(
                BUFFER_SIZE,
                MultiBzDecoder::new(file),
            ))))
        } else {
            Ok(Self::Sequential(Box::new(BufReader::with_capacity(
                BUFFER_SIZE,
                File::open(path)?,
            ))))
        }
    }

    /// Reads an uncompressed dump from standard input.
    pub fn stdin() -> Self {
        Self::Sequential(Box::new(BufReader::with_capacity(BUFFER_SIZE, io::stdin())))
    }
}