pdfvec/chunker/

mod.rs

//! Text chunking strategies for vectorization pipelines.
//!
//! This module provides [`Chunker`] for splitting extracted text into
//! appropriately-sized segments for embedding and RAG pipelines.
//!
//! # Strategies
//!
//! - [`ChunkStrategy::Fixed`]: Split at character boundaries with optional overlap
//! - [`ChunkStrategy::Paragraph`]: Split on paragraph boundaries (double newlines)
//! - [`ChunkStrategy::Sentence`]: Split on sentence boundaries using Unicode segmentation
//!
//! # Example
//!
//! ```
//! use pdfvec::{Chunker, ChunkStrategy};
//!
//! let text = "First sentence. Second sentence.\n\nNew paragraph here.";
//!
//! // Fixed-size chunks with overlap
//! let chunker = Chunker::new(ChunkStrategy::Fixed)
//!     .chunk_size(20)
//!     .overlap(5);
//!
//! for chunk in chunker.chunks(text) {
//!     println!("{}: {}", chunk.index(), chunk.text());
//! }
//! ```

mod fixed;
mod paragraph;
mod sentence;

use crate::Chunk;

/// Strategy for splitting text into chunks.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ChunkStrategy {
    /// Fixed-size chunks with optional overlap.
    #[default]
    Fixed,
    /// Split on paragraph boundaries (double newlines).
    Paragraph,
    /// Split on sentence boundaries using Unicode segmentation.
    Sentence,
}

/// Text chunker with configurable strategy and parameters.
///
/// # Example
///
/// ```
/// use pdfvec::{Chunker, ChunkStrategy};
///
/// let chunker = Chunker::new(ChunkStrategy::Fixed)
///     .chunk_size(512)
///     .overlap(50);
///
/// let text = "Your document text here...";
/// let chunks: Vec<_> = chunker.chunks(text).collect();
/// ```
#[derive(Debug, Clone)]
pub struct Chunker {
    strategy: ChunkStrategy,
    chunk_size: usize,
    overlap: usize,
    min_chunk_size: usize,
}

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

impl Chunker {
    /// Creates a new chunker with the specified strategy.
    ///
    /// # Example
    ///
    /// ```
    /// use pdfvec::{Chunker, ChunkStrategy};
    ///
    /// let chunker = Chunker::new(ChunkStrategy::Paragraph);
    /// ```
    #[must_use]
    pub fn new(strategy: ChunkStrategy) -> Self {
        Self {
            strategy,
            chunk_size: 512,
            overlap: 0,
            min_chunk_size: 100,
        }
    }

    /// Sets the target chunk size in characters.
    ///
    /// For [`ChunkStrategy::Fixed`], this is the exact size (except for the last chunk).
    /// For other strategies, this is the target size for merging small segments.
    ///
    /// Default: 512
    ///
    /// # Example
    ///
    /// ```
    /// use pdfvec::{Chunker, ChunkStrategy};
    ///
    /// let chunker = Chunker::new(ChunkStrategy::Fixed).chunk_size(1024);
    /// ```
    #[must_use]
    pub fn chunk_size(mut self, size: usize) -> Self {
        self.chunk_size = size.max(1);
        self
    }

    /// Sets the overlap between consecutive chunks in characters.
    ///
    /// Overlap helps maintain context across chunk boundaries. Only applies
    /// to [`ChunkStrategy::Fixed`].
    ///
    /// Default: 0
    ///
    /// # Example
    ///
    /// ```
    /// use pdfvec::{Chunker, ChunkStrategy};
    ///
    /// let chunker = Chunker::new(ChunkStrategy::Fixed)
    ///     .chunk_size(100)
    ///     .overlap(20);
    /// ```
    #[must_use]
    pub fn overlap(mut self, chars: usize) -> Self {
        self.overlap = chars;
        self
    }

    /// Sets the minimum chunk size for paragraph and sentence strategies.
    ///
    /// Chunks smaller than this will be merged with adjacent chunks.
    ///
    /// Default: 100
    ///
    /// # Example
    ///
    /// ```
    /// use pdfvec::{Chunker, ChunkStrategy};
    ///
    /// let chunker = Chunker::new(ChunkStrategy::Paragraph)
    ///     .min_chunk_size(50);
    /// ```
    #[must_use]
    pub fn min_chunk_size(mut self, size: usize) -> Self {
        self.min_chunk_size = size;
        self
    }

    /// Returns an iterator over chunks of the input text.
    ///
    /// # Example
    ///
    /// ```
    /// use pdfvec::{Chunker, ChunkStrategy};
    ///
    /// let text = "Hello world. How are you?";
    /// let chunker = Chunker::new(ChunkStrategy::Sentence);
    ///
    /// for chunk in chunker.chunks(text) {
    ///     println!("Chunk {}: {}", chunk.index(), chunk.text());
    /// }
    /// ```
    #[must_use]
    pub fn chunks<'a>(&self, text: &'a str) -> Box<dyn Iterator<Item = Chunk<'a>> + 'a> {
        match self.strategy {
            ChunkStrategy::Fixed => Box::new(fixed::FixedChunker::new(
                text,
                self.chunk_size,
                self.overlap,
            )),
            ChunkStrategy::Paragraph => Box::new(paragraph::ParagraphChunker::new(
                text,
                self.chunk_size,
                self.min_chunk_size,
            )),
            ChunkStrategy::Sentence => Box::new(sentence::SentenceChunker::new(
                text,
                self.chunk_size,
                self.min_chunk_size,
            )),
        }
    }

    /// Returns the configured chunk size.
    #[must_use]
    pub fn get_chunk_size(&self) -> usize {
        self.chunk_size
    }

    /// Returns the configured overlap.
    #[must_use]
    pub fn get_overlap(&self) -> usize {
        self.overlap
    }

    /// Returns the configured strategy.
    #[must_use]
    pub fn strategy(&self) -> ChunkStrategy {
        self.strategy
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn chunker_default() {
        let chunker = Chunker::default();
        assert_eq!(chunker.strategy(), ChunkStrategy::Fixed);
        assert_eq!(chunker.get_chunk_size(), 512);
        assert_eq!(chunker.get_overlap(), 0);
    }

    #[test]
    fn chunker_builder() {
        let chunker = Chunker::new(ChunkStrategy::Paragraph)
            .chunk_size(1024)
            .overlap(100)
            .min_chunk_size(50);

        assert_eq!(chunker.strategy(), ChunkStrategy::Paragraph);
        assert_eq!(chunker.get_chunk_size(), 1024);
        assert_eq!(chunker.get_overlap(), 100);
    }

    #[test]
    fn chunk_size_minimum_is_one() {
        let chunker = Chunker::new(ChunkStrategy::Fixed).chunk_size(0);
        assert_eq!(chunker.get_chunk_size(), 1);
    }
}