gigatoken 0.1.0

data/src/input/mod.rs

@@ -0,0 +1,333 @@
+//! Layered input abstraction: Resource → Document → (Pretoken via pretokenize module).
+//!
+//! - **Resource**: a handle to a contiguous byte buffer (file, bytes, string).
+//! - **DocumentIter**: splits a byte buffer on a configurable separator, yielding documents.
+//! - Parallel chunking: `Resource::par_document_chunks` returns N document iterators
+//!   with chunk boundaries aligned to separator positions.
+
+use memchr::memmem;
+use memmap2::Mmap;
+use std::path::Path;
+
+pub(crate) mod decompress;
+pub mod file_source;
+pub mod jsonl;
+pub mod parquet;
+
+// ---------------------------------------------------------------------------
+// Document — owned/borrowed byte buffer used by jsonl iterator
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Clone)]
+pub struct Document<'a>(std::borrow::Cow<'a, [u8]>);
+
+impl<'a> From<&'a [u8]> for Document<'a> {
+    fn from(value: &'a [u8]) -> Self {
+        Document(std::borrow::Cow::Borrowed(value))
+    }
+}
+
+impl From<Vec<u8>> for Document<'_> {
+    fn from(value: Vec<u8>) -> Self {
+        Document(value.into())
+    }
+}
+
+impl AsRef<[u8]> for Document<'_> {
+    fn as_ref(&self) -> &[u8] {
+        &self.0
+    }
+}
+
+// ---------------------------------------------------------------------------
+// DocRef — lightweight reference used internally by the pretokenizer
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct DocRef<'a>(pub &'a [u8]);
+
+impl<'a> From<&'a [u8]> for DocRef<'a> {
+    fn from(value: &'a [u8]) -> Self {
+        DocRef(value)
+    }
+}
+
+impl<'a> std::ops::Deref for DocRef<'a> {
+    type Target = &'a [u8];
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Resource trait
+// ---------------------------------------------------------------------------
+
+/// A contiguous byte buffer that can be split into documents and parallel chunks.
+pub trait Resource: Sync {
+    fn as_bytes(&self) -> &[u8];
+
+    /// Iterate documents by splitting on `separator`.
+    /// If separator is empty, the entire buffer is one document.
+    fn documents<'a>(&'a self, separator: &'a [u8]) -> DocumentIter<'a> {
+        DocumentIter::new(self.as_bytes(), separator)
+    }
+
+    /// Split into `n` chunk iterators, each yielding documents.
+    /// Chunk boundaries are aligned to separator positions so no document
+    /// is split across chunks.
+    fn par_document_chunks<'a>(
+        &'a self,
+        separator: &'a [u8],
+        n: usize,
+    ) -> Vec<DocumentIter<'a>> {
+        par_document_chunks(self.as_bytes(), separator, n)
+    }
+}
+
+// Blanket implementations
+
+impl Resource for [u8] {
+    fn as_bytes(&self) -> &[u8] {
+        self
+    }
+}
+
+impl Resource for Vec<u8> {
+    fn as_bytes(&self) -> &[u8] {
+        self
+    }
+}
+
+impl Resource for str {
+    fn as_bytes(&self) -> &[u8] {
+        str::as_bytes(self)
+    }
+}
+
+impl Resource for String {
+    fn as_bytes(&self) -> &[u8] {
+        str::as_bytes(self)
+    }
+}
+
+impl Resource for Mmap {
+    fn as_bytes(&self) -> &[u8] {
+        self
+    }
+}
+
+// ---------------------------------------------------------------------------
+// MmappedFile
+// ---------------------------------------------------------------------------
+
+/// Owns a memory-mapped file and implements `Resource`.
+pub struct MmappedFile {
+    mmap: Mmap,
+}
+
+impl MmappedFile {
+    pub fn open(path: impl AsRef<Path>) -> Result<Self, std::io::Error> {
+        let file = std::fs::File::open(path)?;
+        let mmap = unsafe { Mmap::map(&file)? };
+        Ok(Self { mmap })
+    }
+}
+
+impl Resource for MmappedFile {
+    fn as_bytes(&self) -> &[u8] {
+        &self.mmap
+    }
+}
+
+// ---------------------------------------------------------------------------
+// DocumentIter
+// ---------------------------------------------------------------------------
+
+/// Zero-copy iterator that splits a byte slice on a separator, yielding documents.
+/// Empty documents (consecutive separators) are skipped.
+pub struct DocumentIter<'a> {
+    bytes: &'a [u8],
+    separator: &'a [u8],
+    /// Prebuilt searcher for `separator`, constructed once instead of per
+    /// yielded document.
+    finder: memmem::Finder<'a>,
+    position: usize,
+    end: usize,
+    finished: bool,
+}
+
+impl<'a> DocumentIter<'a> {
+    pub fn new(bytes: &'a [u8], separator: &'a [u8]) -> Self {
+        Self::new_range(bytes, separator, 0, bytes.len())
+    }
+
+    fn new_range(bytes: &'a [u8], separator: &'a [u8], start: usize, end: usize) -> Self {
+        Self {
+            bytes,
+            separator,
+            finder: memmem::Finder::new(separator),
+            position: start,
+            end,
+            finished: false,
+        }
+    }
+}
+
+impl<'a> Iterator for DocumentIter<'a> {
+    type Item = &'a [u8];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if self.finished || self.position >= self.end {
+                return None;
+            }
+
+            let search_range = &self.bytes[self.position..self.end];
+
+            // Empty separator means "no splitting" — yield entire remainder as one document
+            if self.separator.is_empty() {
+                self.finished = true;
+                return if search_range.is_empty() {
+                    None
+                } else {
+                    Some(search_range)
+                };
+            }
+
+            match self.finder.find(search_range) {
+                Some(offset) => {
+                    let doc = &self.bytes[self.position..self.position + offset];
+                    self.position += offset + self.separator.len();
+                    // Skip empty documents
+                    if !doc.is_empty() {
+                        return Some(doc);
+                    }
+                }
+                None => {
+                    // No more separators; yield the remainder
+                    self.finished = true;
+                    return if search_range.is_empty() {
+                        None
+                    } else {
+                        Some(search_range)
+                    };
+                }
+            }
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Parallel document chunking
+// ---------------------------------------------------------------------------
+
+/// Split `bytes` into `n` document iterators with chunk boundaries aligned
+/// to separator positions. Each iterator covers a disjoint range.
+fn par_document_chunks<'a>(
+    bytes: &'a [u8],
+    separator: &'a [u8],
+    n: usize,
+) -> Vec<DocumentIter<'a>> {
+    if n <= 1 || bytes.is_empty() || separator.is_empty() {
+        return vec![DocumentIter::new(bytes, separator)];
+    }
+
+    let chunk_size = bytes.len() / n;
+    let finder = memmem::Finder::new(separator);
+
+    let mut boundaries = Vec::with_capacity(n + 1);
+    boundaries.push(0usize);
+
+    for i in 1..n {
+        let target = i * chunk_size;
+        // Scan forward from target to find the next separator
+        match finder.find(&bytes[target..]) {
+            Some(offset) => {
+                let boundary = target + offset + separator.len();
+                boundaries.push(boundary);
+            }
+            None => {
+                // No separator found; remaining data goes in the last chunk
+                break;
+            }
+        }
+    }
+    boundaries.push(bytes.len());
+    boundaries.dedup();
+
+    boundaries
+        .windows(2)
+        .map(|w| DocumentIter::new_range(bytes, separator, w[0], w[1]))
+        .collect()
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_document_iter_basic() {
+        let data = b"hello<|endoftext|>world<|endoftext|>foo";
+        let sep = b"<|endoftext|>";
+        let docs: Vec<&[u8]> = data.as_slice().documents(sep).collect();
+        assert_eq!(docs, vec![b"hello".as_slice(), b"world", b"foo"]);
+    }
+
+    #[test]
+    fn test_document_iter_no_separator() {
+        let data = b"hello world";
+        let docs: Vec<&[u8]> = data.as_slice().documents(b"<|endoftext|>").collect();
+        assert_eq!(docs, vec![b"hello world".as_slice()]);
+    }
+
+    #[test]
+    fn test_document_iter_empty_separator() {
+        let data = b"hello world";
+        let docs: Vec<&[u8]> = data.as_slice().documents(b"").collect();
+        assert_eq!(docs, vec![b"hello world".as_slice()]);
+    }
+
+    #[test]
+    fn test_document_iter_consecutive_separators() {
+        let data = b"a<SEP><SEP>b";
+        let docs: Vec<&[u8]> = data.as_slice().documents(b"<SEP>").collect();
+        assert_eq!(docs, vec![b"a".as_slice(), b"b"]);
+    }
+
+    #[test]
+    fn test_document_iter_separator_at_edges() {
+        let data = b"<SEP>hello<SEP>";
+        let docs: Vec<&[u8]> = data.as_slice().documents(b"<SEP>").collect();
+        assert_eq!(docs, vec![b"hello".as_slice()]);
+    }
+
+    #[test]
+    fn test_par_document_chunks_basic() {
+        let sep = b"<|endoftext|>";
+        let parts: Vec<&str> = (0..100)
+            .map(|i| if i % 10 == 9 { "doc" } else { "word " })
+            .collect();
+        let data = parts.join(std::str::from_utf8(sep).unwrap());
+        let bytes = data.as_bytes();
+
+        let chunks = bytes.par_document_chunks(sep, 4);
+        // All documents should be found across all chunks
+        let all_docs: Vec<&[u8]> = chunks.into_iter().flatten().collect();
+        let single_docs: Vec<&[u8]> = bytes.documents(sep).collect();
+        assert_eq!(all_docs, single_docs);
+    }
+
+    #[test]
+    fn test_par_chunks_single_thread() {
+        let data = b"a<SEP>b<SEP>c";
+        let chunks = data.as_slice().par_document_chunks(b"<SEP>", 1);
+        assert_eq!(chunks.len(), 1);
+        let docs: Vec<&[u8]> = chunks.into_iter().flatten().collect();
+        assert_eq!(docs, vec![b"a".as_slice(), b"b", b"c"]);
+    }
+}

data/src/input/parquet.rs

@@ -0,0 +1,303 @@
+//! Parquet input: one document per row, text taken from a string or binary
+//! column. Parquet pages are decoded through the arrow reader, so documents
+//! are materialized as owned buffers rather than borrowed from an mmap like
+//! the byte-stream formats. Null rows become empty documents so results
+//! stay row-aligned with the source table. Row groups are the parallel work
+//! units; documents always come back in row order.
+
+use std::collections::HashMap;
+use std::fs::File;
+use std::io;
+use std::path::Path;
+
+use arrow_array::cast::AsArray;
+use arrow_array::types::{ByteArrayType, ByteViewType};
+use arrow_array::{Array, GenericByteArray, GenericByteViewArray, RecordBatch};
+use arrow_schema::DataType;
+use parquet::arrow::ProjectionMask;
+use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
+use rayon::prelude::*;
+use rustc_hash::FxBuildHasher;
+
+use crate::pretokenize::pretokenize_as_iter;
+
+fn parquet_err(path: &Path, err: impl std::fmt::Display) -> io::Error {
+    io::Error::new(
+        io::ErrorKind::InvalidData,
+        format!("{}: {err}", path.display()),
+    )
+}
+
+fn open_builder(path: &Path) -> io::Result<ParquetRecordBatchReaderBuilder<File>> {
+    let file = File::open(path)
+        .map_err(|e| io::Error::new(e.kind(), format!("{}: {e}", path.display())))?;
+    ParquetRecordBatchReaderBuilder::try_new(file).map_err(|e| parquet_err(path, e))
+}
+
+pub fn n_row_groups(path: &Path) -> io::Result<usize> {
+    Ok(open_builder(path)?.metadata().num_row_groups())
+}
+
+/// Reader over `row_groups` (all when None) that decodes only `column`.
+/// `column` must be a top-level field of the file's schema.
+fn open_reader(
+    path: &Path,
+    column: &str,
+    row_groups: Option<Vec<usize>>,
+) -> io::Result<ParquetRecordBatchReader> {
+    let builder = open_builder(path)?;
+    if builder.schema().index_of(column).is_err() {
+        let available: Vec<&str> = builder
+            .schema()
+            .fields()
+            .iter()
+            .map(|f| f.name().as_str())
+            .collect();
+        return Err(parquet_err(
+            path,
+            format!(
+                "no column {column:?} (available: {})",
+                available.join(", ")
+            ),
+        ));
+    }
+    let mask = ProjectionMask::columns(builder.parquet_schema(), [column]);
+    let builder = builder.with_projection(mask);
+    let builder = match row_groups {
+        Some(rgs) => builder.with_row_groups(rgs),
+        None => builder,
+    };
+    builder.build().map_err(|e| parquet_err(path, e))
+}
+
+/// Visit each row of the batch's single projected column as document bytes.
+/// Null rows visit as empty slices, keeping documents row-aligned.
+fn for_each_row(
+    batch: &RecordBatch,
+    path: &Path,
+    column: &str,
+    f: &mut impl FnMut(&[u8]),
+) -> io::Result<()> {
+    fn visit_bytes<T: ByteArrayType>(arr: &GenericByteArray<T>, f: &mut impl FnMut(&[u8]))
+    where
+        T::Native: AsRef<[u8]>,
+    {
+        for i in 0..arr.len() {
+            if arr.is_null(i) {
+                f(b"");
+            } else {
+                f(arr.value(i).as_ref());
+            }
+        }
+    }
+    fn visit_views<T: ByteViewType>(arr: &GenericByteViewArray<T>, f: &mut impl FnMut(&[u8]))
+    where
+        T::Native: AsRef<[u8]>,
+    {
+        for i in 0..arr.len() {
+            if arr.is_null(i) {
+                f(b"");
+            } else {
+                f(arr.value(i).as_ref());
+            }
+        }
+    }
+
+    let col = batch.column(0);
+    match col.data_type() {
+        DataType::Utf8 => visit_bytes(col.as_string::<i32>(), f),
+        DataType::LargeUtf8 => visit_bytes(col.as_string::<i64>(), f),
+        DataType::Utf8View => visit_views(col.as_string_view(), f),
+        DataType::Binary => visit_bytes(col.as_binary::<i32>(), f),
+        DataType::LargeBinary => visit_bytes(col.as_binary::<i64>(), f),
+        DataType::BinaryView => visit_views(col.as_binary_view(), f),
+        other => {
+            return Err(parquet_err(
+                path,
+                format!(
+                    "column {column:?} has unsupported type {other} \
+                     (expected a string or binary column)"
+                ),
+            ));
+        }
+    }
+    Ok(())
+}
+
+/// Visit every document (row) of `column` in the given row groups (all when
+/// None), in row order, on the calling thread.
+pub fn for_each_doc(
+    path: &Path,
+    column: &str,
+    row_groups: Option<Vec<usize>>,
+    mut f: impl FnMut(&[u8]),
+) -> io::Result<()> {
+    for batch in open_reader(path, column, row_groups)? {
+        let batch = batch.map_err(|e| parquet_err(path, e))?;
+        for_each_row(&batch, path, column, &mut f)?;
+    }
+    Ok(())
+}
+
+/// All documents of `column` in `path`, one owned buffer per row, in row
+/// order. Parallel over row groups when `parallel`; strictly on the calling
+/// thread otherwise (the serial encode paths must never touch rayon).
+pub fn read_docs(path: &Path, column: &str, parallel: bool) -> io::Result<Vec<Vec<u8>>> {
+    if !parallel {
+        let mut docs = Vec::new();
+        for_each_doc(path, column, None, |d| docs.push(d.to_vec()))?;
+        return Ok(docs);
+    }
+    let per_group: Vec<Vec<Vec<u8>>> = (0..n_row_groups(path)?)
+        .into_par_iter()
+        .map(|rg| {
+            let mut docs = Vec::new();
+            for_each_doc(path, column, Some(vec![rg]), |d| docs.push(d.to_vec()))?;
+            Ok(docs)
+        })
+        .collect::<io::Result<_>>()?;
+    Ok(per_group.into_iter().flatten().collect())
+}
+
+/// Parallel pretokenization of `column`, one rayon task per row group.
+/// Documents are visited batch-by-batch without materializing the file.
+pub fn pretokenize_par(
+    path: &Path,
+    column: &str,
+) -> io::Result<HashMap<Vec<u8>, usize, FxBuildHasher>> {
+    (0..n_row_groups(path)?)
+        .into_par_iter()
+        .map(|rg| {
+            let mut counts: HashMap<Vec<u8>, usize, FxBuildHasher> = HashMap::default();
+            for_each_doc(path, column, Some(vec![rg]), |doc| {
+                for pretoken in pretokenize_as_iter(doc) {
+                    *counts.entry(pretoken.as_ref().to_vec()).or_default() += 1;
+                }
+            })?;
+            Ok(counts)
+        })
+        .try_reduce(HashMap::default, |mut acc, counts| {
+            if acc.is_empty() {
+                return Ok(counts);
+            }
+            for (k, v) in counts {
+                *acc.entry(k).or_default() += v;
+            }
+            Ok(acc)
+        })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow_array::{BinaryArray, Int64Array, StringArray};
+    use parquet::arrow::ArrowWriter;
+    use parquet::file::properties::WriterProperties;
+    use std::sync::Arc;
+
+    /// Write a single-row-group-capped parquet file with a nullable string
+    /// "text" column, a binary "raw" column, and an int "id" column.
+    fn write_fixture(path: &Path, texts: &[Option<&str>], max_row_group_size: usize) {
+        let schema = Arc::new(arrow_schema::Schema::new(vec![
+            arrow_schema::Field::new("id", DataType::Int64, false),
+            arrow_schema::Field::new("text", DataType::Utf8, true),
+            arrow_schema::Field::new("raw", DataType::Binary, true),
+        ]));
+        let ids = Int64Array::from((0..texts.len() as i64).collect::<Vec<_>>());
+        let text_col = StringArray::from(texts.to_vec());
+        let raw_col = BinaryArray::from(
+            texts
+                .iter()
+                .map(|t| t.map(|s| s.as_bytes()))
+                .collect::<Vec<_>>(),
+        );
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(ids), Arc::new(text_col), Arc::new(raw_col)],
+        )
+        .unwrap();
+        let props = WriterProperties::builder()
+            .set_max_row_group_row_count(Some(max_row_group_size))
+            .build();
+        let file = File::create(path).unwrap();
+        let mut writer = ArrowWriter::try_new(file, schema, Some(props)).unwrap();
+        writer.write(&batch).unwrap();
+        writer.close().unwrap();
+    }
+
+    const TEXTS: [Option<&str>; 7] = [
+        Some("The quick brown fox"),
+        Some("jumps over the lazy dog."),
+        None,
+        Some("She sells seashells"),
+        Some(""),
+        Some("by the seashore."),
+        Some("Peter Piper picked a peck"),
+    ];
+
+    fn expected_docs() -> Vec<Vec<u8>> {
+        TEXTS
+            .iter()
+            .map(|t| t.unwrap_or("").as_bytes().to_vec())
+            .collect()
+    }
+
+    #[test]
+    fn test_read_docs_row_order_and_nulls() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("docs.parquet");
+        write_fixture(&path, &TEXTS, 1024);
+        assert_eq!(read_docs(&path, "text", false).unwrap(), expected_docs());
+    }
+
+    #[test]
+    fn test_read_docs_parallel_multiple_row_groups() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("docs.parquet");
+        write_fixture(&path, &TEXTS, 2); // 4 row groups for 7 rows
+        assert!(n_row_groups(&path).unwrap() > 1);
+        assert_eq!(read_docs(&path, "text", true).unwrap(), expected_docs());
+        assert_eq!(read_docs(&path, "text", false).unwrap(), expected_docs());
+    }
+
+    #[test]
+    fn test_read_binary_column() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("docs.parquet");
+        write_fixture(&path, &TEXTS, 3);
+        assert_eq!(read_docs(&path, "raw", true).unwrap(), expected_docs());
+    }
+
+    #[test]
+    fn test_missing_column_lists_available() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("docs.parquet");
+        write_fixture(&path, &TEXTS, 1024);
+        let err = read_docs(&path, "content", false).unwrap_err().to_string();
+        assert!(err.contains("no column \"content\""), "{err}");
+        assert!(err.contains("text"), "{err}");
+    }
+
+    #[test]
+    fn test_non_string_column_rejected() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("docs.parquet");
+        write_fixture(&path, &TEXTS, 1024);
+        let err = read_docs(&path, "id", false).unwrap_err().to_string();
+        assert!(err.contains("unsupported type"), "{err}");
+    }
+
+    #[test]
+    fn test_pretokenize_matches_per_doc_counts() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("docs.parquet");
+        write_fixture(&path, &TEXTS, 2);
+        let mut expected: HashMap<Vec<u8>, usize, FxBuildHasher> = HashMap::default();
+        for doc in expected_docs() {
+            for pretoken in pretokenize_as_iter(&doc) {
+                *expected.entry(pretoken.as_ref().to_vec()).or_default() += 1;
+            }
+        }
+        assert_eq!(pretokenize_par(&path, "text").unwrap(), expected);
+    }
+}