lakehouse/crates/vectord/src/chunker.rs

/// Text chunking strategies for embedding.
/// Chunks need to be small enough for the embedding model (typically <512 tokens)
/// but large enough to carry meaning.

/// A chunk of text with metadata pointing back to its source.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TextChunk {
    /// Source identifier (dataset name, filename, etc.)
    pub source: String,
    /// Row or document ID within the source
    pub doc_id: String,
    /// Chunk index within the document (0, 1, 2, ...)
    pub chunk_idx: u32,
    /// The actual text content
    pub text: String,
}

/// Split text into overlapping chunks.
/// - `chunk_size`: target characters per chunk (not tokens — chars are a good proxy)
/// - `overlap`: characters of overlap between consecutive chunks
pub fn chunk_text(
    text: &str,
    source: &str,
    doc_id: &str,
    chunk_size: usize,
    overlap: usize,
) -> Vec<TextChunk> {
    let text = text.trim();
    if text.is_empty() {
        return vec![];
    }

    // Short text — single chunk
    if text.len() <= chunk_size {
        return vec![TextChunk {
            source: source.to_string(),
            doc_id: doc_id.to_string(),
            chunk_idx: 0,
            text: text.to_string(),
        }];
    }

    // Work with char indices, not byte indices, to handle Unicode safely
    let chars: Vec<char> = text.chars().collect();
    let char_len = chars.len();
    let mut chunks = Vec::new();
    let mut start = 0usize; // char index
    let mut idx = 0u32;

    while start < char_len {
        let end = (start + chunk_size).min(char_len);
        let chunk_str: String = chars[start..end].iter().collect();

        let actual_end = if end < char_len {
            // Look for last sentence boundary
            if let Some(pos) = chunk_str.rfind(". ") {
                start + chunk_str[..pos + 2].chars().count()
            } else if let Some(pos) = chunk_str.rfind('\n') {
                start + chunk_str[..pos + 1].chars().count()
            } else if let Some(pos) = chunk_str.rfind(' ') {
                start + chunk_str[..pos + 1].chars().count()
            } else {
                end
            }
        } else {
            end
        };

        let chunk: String = chars[start..actual_end].iter().collect();
        let trimmed = chunk.trim();
        if !trimmed.is_empty() {
            chunks.push(TextChunk {
                source: source.to_string(),
                doc_id: doc_id.to_string(),
                chunk_idx: idx,
                text: trimmed.to_string(),
            });
            idx += 1;
        }

        if actual_end >= char_len { break; }
        start = if actual_end > overlap { actual_end - overlap } else { actual_end };
    }

    chunks
}

/// Chunk a dataset's text column. Returns all chunks from all rows.
pub fn chunk_column(
    source: &str,
    doc_ids: &[String],
    texts: &[String],
    chunk_size: usize,
    overlap: usize,
) -> Vec<TextChunk> {
    let mut all_chunks = Vec::new();
    for (doc_id, text) in doc_ids.iter().zip(texts.iter()) {
        let chunks = chunk_text(text, source, doc_id, chunk_size, overlap);
        all_chunks.extend(chunks);
    }
    all_chunks
}

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

    #[test]
    fn short_text_single_chunk() {
        let chunks = chunk_text("Hello world", "test", "1", 500, 50);
        assert_eq!(chunks.len(), 1);
        assert_eq!(chunks[0].text, "Hello world");
    }

    #[test]
    fn long_text_multiple_chunks() {
        let text = "First sentence. Second sentence. Third sentence. Fourth sentence. Fifth sentence. Sixth sentence. Seventh sentence. Eighth sentence.";
        let chunks = chunk_text(text, "test", "1", 50, 10);
        assert!(chunks.len() > 1);
        // All text should be covered
        for chunk in &chunks {
            assert!(!chunk.text.is_empty());
        }
    }

    #[test]
    fn chunk_preserves_source() {
        let chunks = chunk_text("Some text here", "candidates", "CAND-001", 500, 50);
        assert_eq!(chunks[0].source, "candidates");
        assert_eq!(chunks[0].doc_id, "CAND-001");
    }
}