ruwex 0.1.0

//! Document sinks: wikiextractor-compatible sharded output files
//! (`OUTPUT/AA/wiki_00`, rolling over at a size limit) or stdout.

use std::fs::{self, File};
use std::io::{self, BufWriter, Stdout, Write};
use std::path::{Path, PathBuf};

use bzip2::Compression;
use bzip2::write::BzEncoder;

/// Where rendered documents go. `write_doc` receives one complete document at
/// a time so implementations can split output at document boundaries.
pub trait DocSink {
    fn write_doc(&mut self, doc: &str) -> io::Result<()>;
    fn finish(&mut self) -> io::Result<()>;
}

pub struct StdoutSink(BufWriter<Stdout>);

impl StdoutSink {
    pub fn new() -> Self {
        Self(BufWriter::new(io::stdout()))
    }
}

impl Default for StdoutSink {
    fn default() -> Self {
        Self::new()
    }
}

impl DocSink for StdoutSink {
    fn write_doc(&mut self, doc: &str) -> io::Result<()> {
        self.0.write_all(doc.as_bytes())
    }

    fn finish(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

enum ShardStream {
    Plain(BufWriter<File>),
    Bz2(Box<BzEncoder<BufWriter<File>>>),
}

impl ShardStream {
    fn write_all(&mut self, bytes: &[u8]) -> io::Result<()> {
        match self {
            Self::Plain(w) => w.write_all(bytes),
            Self::Bz2(w) => w.write_all(bytes),
        }
    }

    fn finish(self) -> io::Result<()> {
        match self {
            Self::Plain(mut w) => w.flush(),
            Self::Bz2(w) => w.finish()?.flush(),
        }
    }
}

/// Writes documents into `root/AA/wiki_00`, `wiki_01`, … `ZZ/wiki_99`,
/// starting a new file whenever the current one would exceed `max_bytes`
/// (counted uncompressed, like wikiextractor). A document larger than the
/// limit still goes into a single file.
pub struct ShardedWriter {
    root: PathBuf,
    max_bytes: u64,
    compress: bool,
    dir_index: usize,
    file_index: usize,
    shard: Option<(ShardStream, u64)>,
}

impl ShardedWriter {
    const MAX_DIRS: usize = 26 * 26;
    const FILES_PER_DIR: usize = 100;

    /// `max_bytes == 0` means no size limit. The first output file is created
    /// immediately, matching wikiextractor.
    pub fn create(root: &Path, max_bytes: u64, compress: bool) -> io::Result<Self> {
        let mut writer = Self {
            root: root.to_path_buf(),
            max_bytes: if max_bytes == 0 { u64::MAX } else { max_bytes },
            compress,
            dir_index: 0,
            file_index: 0,
            shard: None,
        };
        writer.open_next()?;
        Ok(writer)
    }

    fn dir_name(index: usize) -> String {
        format!(
            "{}{}",
            (b'A' + (index / 26) as u8) as char,
            (b'A' + (index % 26) as u8) as char
        )
    }

    fn open_next(&mut self) -> io::Result<()> {
        if self.dir_index >= Self::MAX_DIRS {
            return Err(io::Error::other(
                "output exceeds ZZ/wiki_99; increase --bytes",
            ));
        }
        let dir = self.root.join(Self::dir_name(self.dir_index));
        fs::create_dir_all(&dir)?;
        let mut name = format!("wiki_{:02}", self.file_index);
        if self.compress {
            name.push_str(".bz2");
        }
        let file = BufWriter::new(File::create(dir.join(name))?);
        let stream = if self.compress {
            ShardStream::Bz2(Box::new(BzEncoder::new(file, Compression::best())))
        } else {
            ShardStream::Plain(file)
        };
        self.shard = Some((stream, 0));

        self.file_index += 1;
        if self.file_index == Self::FILES_PER_DIR {
            self.file_index = 0;
            self.dir_index += 1;
        }
        Ok(())
    }

    fn close_shard(&mut self) -> io::Result<()> {
        match self.shard.take() {
            Some((stream, _)) => stream.finish(),
            None => Ok(()),
        }
    }
}

impl DocSink for ShardedWriter {
    fn write_doc(&mut self, doc: &str) -> io::Result<()> {
        if self.shard.is_none() {
            self.open_next()?;
        }
        let written = self.shard.as_ref().expect("shard just opened").1;
        if written > 0 && written + doc.len() as u64 > self.max_bytes {
            self.close_shard()?;
            self.open_next()?;
        }
        let (stream, written) = self.shard.as_mut().expect("shard just opened");
        stream.write_all(doc.as_bytes())?;
        *written += doc.len() as u64;
        Ok(())
    }

    fn finish(&mut self) -> io::Result<()> {
        self.close_shard()
    }
}

#[cfg(test)]
mod tests {
    use std::io::Read;

    use bzip2::read::MultiBzDecoder;

    use super::*;

    #[test]
    fn dir_names_follow_wikiextractor() {
        assert_eq!(ShardedWriter::dir_name(0), "AA");
        assert_eq!(ShardedWriter::dir_name(25), "AZ");
        assert_eq!(ShardedWriter::dir_name(26), "BA");
        assert_eq!(ShardedWriter::dir_name(30), "BE");
        assert_eq!(ShardedWriter::dir_name(675), "ZZ");
    }

    #[test]
    fn splits_at_size_limit_and_rolls_over_directories() {
        let dir = tempfile::tempdir().unwrap();
        let mut writer = ShardedWriter::create(dir.path(), 10, false).unwrap();
        // 8 bytes per doc: two fit nowhere together under the 10-byte limit.
        for i in 0..3 {
            writer.write_doc(&format!("doc {i} !\n")).unwrap();
        }
        writer.finish().unwrap();

        let aa = dir.path().join("AA");
        assert_eq!(fs::read_to_string(aa.join("wiki_00")).unwrap(), "doc 0 !\n");
        assert_eq!(fs::read_to_string(aa.join("wiki_01")).unwrap(), "doc 1 !\n");
        assert_eq!(fs::read_to_string(aa.join("wiki_02")).unwrap(), "doc 2 !\n");
    }

    #[test]
    fn oversized_doc_goes_into_single_file() {
        let dir = tempfile::tempdir().unwrap();
        let mut writer = ShardedWriter::create(dir.path(), 4, false).unwrap();
        writer.write_doc("bigger than the limit").unwrap();
        writer.write_doc("next").unwrap();
        writer.finish().unwrap();

        let aa = dir.path().join("AA");
        assert_eq!(
            fs::read_to_string(aa.join("wiki_00")).unwrap(),
            "bigger than the limit"
        );
        assert_eq!(fs::read_to_string(aa.join("wiki_01")).unwrap(), "next");
    }

    #[test]
    fn compressed_output_roundtrips() {
        let dir = tempfile::tempdir().unwrap();
        let mut writer = ShardedWriter::create(dir.path(), 0, true).unwrap();
        writer.write_doc("hello ").unwrap();
        writer.write_doc("world\n").unwrap();
        writer.finish().unwrap();

        let path = dir.path().join("AA").join("wiki_00.bz2");
        let mut decoder = MultiBzDecoder::new(File::open(path).unwrap());
        let mut content = String::new();
        decoder.read_to_string(&mut content).unwrap();
        assert_eq!(content, "hello world\n");
    }

    #[test]
    fn first_file_exists_even_when_empty() {
        let dir = tempfile::tempdir().unwrap();
        let mut writer = ShardedWriter::create(dir.path(), 1024, false).unwrap();
        writer.finish().unwrap();
        assert!(dir.path().join("AA").join("wiki_00").is_file());
    }
}