kreuzberg 4.0.0.rc1

data/vendor/kreuzberg/src/extraction/table.rs

@@ -0,0 +1,328 @@
+//! Table conversion utilities.
+//!
+//! This module provides functions for converting tabular data between different formats.
+//! Currently supports converting Apache Arrow IPC format to Markdown tables using Polars.
+//!
+//! # Features
+//!
+//! - **Arrow IPC parsing**: Read tables from Arrow IPC binary format
+//! - **Markdown generation**: Convert DataFrames to clean Markdown tables
+//! - **Type-safe**: Handles all Polars data types safely
+//! - **Empty table handling**: Gracefully handles empty DataFrames
+//!
+//! # Supported Conversions
+//!
+//! - Arrow IPC → Markdown table
+//! - Polars DataFrame → Markdown table
+//!
+//! # Example
+//!
+//! ```rust,no_run
+//! use kreuzberg::extraction::table::table_from_arrow_to_markdown;
+//!
+//! # fn example() -> kreuzberg::Result<()> {
+//! // Convert Arrow IPC bytes to Markdown
+//! let arrow_bytes = vec![/* Arrow IPC data */];
+//! let markdown = table_from_arrow_to_markdown(&arrow_bytes)?;
+//!
+//! println!("Markdown table:\n{}", markdown);
+//! # Ok(())
+//! # }
+//! ```
+//!
+//! # Output Format
+//!
+//! The generated Markdown follows GitHub Flavored Markdown table syntax:
+//! ```markdown
+//! | Column1 | Column2 | Column3 |
+//! |------|------|------|
+//! | value1 | value2 | value3 |
+//! | value4 | value5 | value6 |
+//! ```
+
+use crate::error::{KreuzbergError, Result};
+use polars::prelude::*;
+use std::io::Cursor;
+
+/// Convert Arrow IPC bytes to markdown table format
+pub fn table_from_arrow_to_markdown(arrow_bytes: &[u8]) -> Result<String> {
+    let cursor = Cursor::new(arrow_bytes);
+    let df = IpcReader::new(cursor)
+        .finish()
+        .map_err(|e| KreuzbergError::parsing(format!("Failed to read Arrow IPC data: {}", e)))?;
+
+    dataframe_to_markdown(&df)
+}
+
+/// Convert a Polars DataFrame to markdown table format
+fn dataframe_to_markdown(df: &DataFrame) -> Result<String> {
+    if df.is_empty() {
+        return Ok(String::new());
+    }
+
+    let mut markdown = String::new();
+
+    markdown.push_str("| ");
+    for col_name in df.get_column_names() {
+        markdown.push_str(col_name);
+        markdown.push_str(" | ");
+    }
+    markdown.push('\n');
+
+    markdown.push('|');
+    for _ in 0..df.width() {
+        markdown.push_str("------|");
+    }
+    markdown.push('\n');
+
+    for row_idx in 0..df.height() {
+        markdown.push_str("| ");
+        for col in df.get_columns() {
+            let series = col.as_materialized_series();
+            let value = format_cell_value(series, row_idx)?;
+            markdown.push_str(&value);
+            markdown.push_str(" | ");
+        }
+        markdown.push('\n');
+    }
+
+    Ok(markdown)
+}
+
+fn format_cell_value(series: &Series, idx: usize) -> Result<String> {
+    let is_null_array = series.is_null();
+    if is_null_array.get(idx).unwrap_or(false) {
+        return Ok(String::new());
+    }
+
+    let value_str = match series.dtype() {
+        DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
+            let casted = series
+                .cast(&DataType::Int64)
+                .map_err(|e| KreuzbergError::parsing(format!("Failed to cast to i64: {}", e)))?;
+            casted
+                .i64()
+                .map_err(|e| KreuzbergError::parsing(format!("Failed to get i64 value: {}", e)))?
+                .get(idx)
+                .map(|v| v.to_string())
+                .unwrap_or_default()
+        }
+        DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => {
+            let casted = series
+                .cast(&DataType::UInt64)
+                .map_err(|e| KreuzbergError::parsing(format!("Failed to cast to u64: {}", e)))?;
+            casted
+                .u64()
+                .map_err(|e| KreuzbergError::parsing(format!("Failed to get u64 value: {}", e)))?
+                .get(idx)
+                .map(|v| v.to_string())
+                .unwrap_or_default()
+        }
+        DataType::Float32 | DataType::Float64 => {
+            let casted = series
+                .cast(&DataType::Float64)
+                .map_err(|e| KreuzbergError::parsing(format!("Failed to cast to f64: {}", e)))?;
+            casted
+                .f64()
+                .map_err(|e| KreuzbergError::parsing(format!("Failed to get f64 value: {}", e)))?
+                .get(idx)
+                .map(|v| format!("{:.2}", v))
+                .unwrap_or_default()
+        }
+        DataType::Boolean => series
+            .bool()
+            .map_err(|e| KreuzbergError::parsing(format!("Failed to get bool value: {}", e)))?
+            .get(idx)
+            .map(|v| v.to_string())
+            .unwrap_or_default(),
+        DataType::String => series
+            .str()
+            .map_err(|e| KreuzbergError::parsing(format!("Failed to get string value: {}", e)))?
+            .get(idx)
+            .map(|v| v.to_string())
+            .unwrap_or_default(),
+        _ => {
+            format!("{:?}", series.get(idx))
+        }
+    };
+
+    Ok(value_str)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::Cursor;
+
+    fn create_test_dataframe() -> DataFrame {
+        df!(
+            "name" => &["Alice", "Bob", "Charlie"],
+            "age" => &[30, 25, 35],
+            "score" => &[95.5, 87.3, 92.1]
+        )
+        .unwrap()
+    }
+
+    fn dataframe_to_arrow_bytes(df: &DataFrame) -> Vec<u8> {
+        let mut buffer = Vec::new();
+        let mut cursor = Cursor::new(&mut buffer);
+        let mut df_mut = df.clone();
+        IpcWriter::new(&mut cursor).finish(&mut df_mut).unwrap();
+        buffer
+    }
+
+    #[test]
+    fn test_dataframe_to_markdown_basic() {
+        let df = create_test_dataframe();
+        let markdown = dataframe_to_markdown(&df).unwrap();
+
+        assert!(markdown.contains("| name | age | score |"));
+        assert!(markdown.contains("|------|------|------|"));
+        assert!(markdown.contains("| Alice | 30 | 95.50 |"));
+        assert!(markdown.contains("| Bob | 25 | 87.30 |"));
+        assert!(markdown.contains("| Charlie | 35 | 92.10 |"));
+    }
+
+    #[test]
+    fn test_table_from_arrow_to_markdown() {
+        let df = create_test_dataframe();
+        let arrow_bytes = dataframe_to_arrow_bytes(&df);
+
+        let markdown = table_from_arrow_to_markdown(&arrow_bytes).unwrap();
+
+        assert!(markdown.contains("| name | age | score |"));
+        assert!(markdown.contains("| Alice | 30 | 95.50 |"));
+        assert!(markdown.contains("| Bob | 25 | 87.30 |"));
+        assert!(markdown.contains("| Charlie | 35 | 92.10 |"));
+    }
+
+    #[test]
+    fn test_empty_dataframe() {
+        let df = df!("col1" => Vec::<i32>::new()).unwrap();
+        let markdown = dataframe_to_markdown(&df).unwrap();
+        assert_eq!(markdown, "");
+    }
+
+    #[test]
+    fn test_dataframe_with_nulls() {
+        let s1 = Series::new("name".into(), &["Alice", "Bob", "Charlie"]);
+        let s2 = Series::new("value".into(), &[Some(1), None, Some(3)]);
+        let df = DataFrame::new(vec![s1.into(), s2.into()]).unwrap();
+
+        let markdown = dataframe_to_markdown(&df).unwrap();
+
+        assert!(markdown.contains("| name | value |"));
+        assert!(markdown.contains("| Alice | 1 |"));
+        assert!(markdown.contains("| Bob |  |"));
+        assert!(markdown.contains("| Charlie | 3 |"));
+    }
+
+    #[test]
+    fn test_dataframe_with_booleans() {
+        let df = df!(
+            "name" => &["Alice", "Bob"],
+            "active" => &[true, false]
+        )
+        .unwrap();
+
+        let markdown = dataframe_to_markdown(&df).unwrap();
+
+        assert!(markdown.contains("| name | active |"));
+        assert!(markdown.contains("| Alice | true |"));
+        assert!(markdown.contains("| Bob | false |"));
+    }
+
+    #[test]
+    fn test_dataframe_with_integers() {
+        let df = df!(
+            "id" => &[1i64, 2i64, 3i64],
+            "count" => &[100u64, 200u64, 300u64]
+        )
+        .unwrap();
+
+        let markdown = dataframe_to_markdown(&df).unwrap();
+
+        assert!(markdown.contains("| id | count |"));
+        assert!(markdown.contains("| 1 | 100 |"));
+        assert!(markdown.contains("| 2 | 200 |"));
+        assert!(markdown.contains("| 3 | 300 |"));
+    }
+
+    #[test]
+    fn test_single_column_dataframe() {
+        let df = df!("name" => &["Alice", "Bob", "Charlie"]).unwrap();
+        let markdown = dataframe_to_markdown(&df).unwrap();
+
+        assert!(markdown.contains("| name |"));
+        assert!(markdown.contains("|------|"));
+        assert!(markdown.contains("| Alice |"));
+        assert!(markdown.contains("| Bob |"));
+        assert!(markdown.contains("| Charlie |"));
+    }
+
+    #[test]
+    fn test_single_row_dataframe() {
+        let df = df!(
+            "name" => &["Alice"],
+            "age" => &[30]
+        )
+        .unwrap();
+
+        let markdown = dataframe_to_markdown(&df).unwrap();
+
+        assert!(markdown.contains("| name | age |"));
+        assert!(markdown.contains("| Alice | 30 |"));
+    }
+
+    #[test]
+    fn test_arrow_bytes_roundtrip() {
+        let original_df = df!(
+            "col1" => &[1, 2, 3],
+            "col2" => &["a", "b", "c"]
+        )
+        .unwrap();
+
+        let arrow_bytes = dataframe_to_arrow_bytes(&original_df);
+        let markdown = table_from_arrow_to_markdown(&arrow_bytes).unwrap();
+
+        assert!(markdown.contains("| col1 | col2 |"));
+        assert!(markdown.contains("| 1 | a |"));
+        assert!(markdown.contains("| 2 | b |"));
+        assert!(markdown.contains("| 3 | c |"));
+    }
+
+    #[test]
+    fn test_invalid_arrow_bytes() {
+        let invalid_bytes = vec![0u8; 10];
+        let result = table_from_arrow_to_markdown(&invalid_bytes);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_float_formatting() {
+        let df = df!(
+            "value" => &[1.234, 5.678, 9.012]
+        )
+        .unwrap();
+
+        let markdown = dataframe_to_markdown(&df).unwrap();
+
+        assert!(markdown.contains("| 1.23 |"));
+        assert!(markdown.contains("| 5.68 |"));
+        assert!(markdown.contains("| 9.01 |"));
+    }
+
+    #[test]
+    fn test_special_characters_in_strings() {
+        let df = df!(
+            "text" => &["Hello | World", "A & B", "C > D"]
+        )
+        .unwrap();
+
+        let markdown = dataframe_to_markdown(&df).unwrap();
+
+        assert!(markdown.contains("| Hello | World |"));
+        assert!(markdown.contains("| A & B |"));
+        assert!(markdown.contains("| C > D |"));
+    }
+}

data/vendor/kreuzberg/src/extraction/text.rs

@@ -0,0 +1,269 @@
+//! Plain text and Markdown extraction functions.
+//!
+//! This module provides memory-efficient streaming parsers for plain text and Markdown files.
+//! Key features:
+//!
+//! - **Streaming parsing**: Processes files line-by-line to handle multi-GB files
+//! - **Markdown support**: Extracts headers, links, and code blocks from Markdown
+//! - **Word/line counting**: Accurate statistics without loading entire file
+//! - **CRLF support**: Handles both Unix and Windows line endings
+//!
+//! # Example
+//!
+//! ```rust
+//! use kreuzberg::extraction::text::parse_text;
+//!
+//! # fn example() -> kreuzberg::Result<()> {
+//! let text = b"# Hello\n\nThis is [a link](https://example.com).";
+//! let result = parse_text(text, true)?; // true = is Markdown
+//!
+//! assert_eq!(result.line_count, 3);
+//! assert!(result.headers.unwrap().contains(&"Hello".to_string()));
+//! # Ok(())
+//! # }
+//! ```
+use once_cell::sync::Lazy;
+use regex::Regex;
+
+use crate::error::Result;
+use crate::types::TextExtractionResult;
+
+static MARKDOWN_HEADER: Lazy<Regex> =
+    Lazy::new(|| Regex::new(r"^#{1,6}\s*(.+)$").expect("Markdown header regex pattern is valid and should compile"));
+static MARKDOWN_LINK: Lazy<Regex> = Lazy::new(|| {
+    Regex::new(r"\[([^\]]+)\]\(([^)]+)\)").expect("Markdown link regex pattern is valid and should compile")
+});
+static CODE_BLOCK_DELIMITER: Lazy<Regex> = Lazy::new(|| {
+    Regex::new(r"^```(\w*)[\r]?$").expect("Code block delimiter regex pattern is valid and should compile")
+});
+
+pub fn parse_text(text_bytes: &[u8], is_markdown: bool) -> Result<TextExtractionResult> {
+    let text = String::from_utf8_lossy(text_bytes).into_owned();
+
+    let mut line_count = 0;
+    let mut word_count = 0;
+    let character_count = text.len();
+
+    let mut headers = Vec::new();
+    let mut links = Vec::new();
+    let mut code_blocks = Vec::new();
+    let mut in_code_block = false;
+    let mut current_code_lang = String::new();
+    let mut current_code = String::new();
+
+    for line in text.lines() {
+        line_count += 1;
+        word_count += line.split_whitespace().count();
+
+        if !is_markdown {
+            continue;
+        }
+
+        if CODE_BLOCK_DELIMITER.is_match(line) {
+            if in_code_block {
+                code_blocks.push((
+                    if current_code_lang.is_empty() {
+                        "plain".to_string()
+                    } else {
+                        current_code_lang.clone()
+                    },
+                    current_code.trim_end().to_string(),
+                ));
+                current_code.clear();
+                current_code_lang.clear();
+                in_code_block = false;
+            } else {
+                if let Some(caps) = CODE_BLOCK_DELIMITER.captures(line) {
+                    current_code_lang = caps.get(1).map(|m| m.as_str()).unwrap_or("").to_string();
+                }
+                in_code_block = true;
+            }
+            continue;
+        }
+
+        if in_code_block {
+            current_code.push_str(line);
+            current_code.push('\n');
+            continue;
+        }
+
+        if let Some(caps) = MARKDOWN_HEADER.captures(line)
+            && let Some(header) = caps.get(1)
+        {
+            headers.push(header.as_str().to_string());
+        }
+
+        for caps in MARKDOWN_LINK.captures_iter(line) {
+            if let (Some(text), Some(url)) = (caps.get(1), caps.get(2)) {
+                links.push((text.as_str().to_string(), url.as_str().to_string()));
+            }
+        }
+    }
+
+    Ok(TextExtractionResult {
+        content: text,
+        line_count,
+        word_count,
+        character_count,
+        headers: if headers.is_empty() { None } else { Some(headers) },
+        links: if links.is_empty() { None } else { Some(links) },
+        code_blocks: if code_blocks.is_empty() {
+            None
+        } else {
+            Some(code_blocks)
+        },
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_plain_text_basic() {
+        let text = b"Hello, World!\nThis is a test.\nThird line here.";
+        let result = parse_text(text, false).unwrap();
+        assert_eq!(result.content, "Hello, World!\nThis is a test.\nThird line here.");
+        assert_eq!(result.line_count, 3);
+        assert_eq!(result.word_count, 9);
+        assert_eq!(result.character_count, result.content.len());
+        assert!(result.headers.is_none());
+        assert!(result.links.is_none());
+        assert!(result.code_blocks.is_none());
+    }
+
+    #[test]
+    fn test_plain_text_empty() {
+        let text = b"";
+        let result = parse_text(text, false).unwrap();
+        assert_eq!(result.content, "");
+        assert_eq!(result.line_count, 0);
+        assert_eq!(result.word_count, 0);
+        assert_eq!(result.character_count, 0);
+    }
+
+    #[test]
+    fn test_markdown_headers() {
+        let text = b"# Header 1\n## Header 2\n### Header 3\n#NoSpace\n## Multiple  spaces";
+        let result = parse_text(text, true).unwrap();
+        assert_eq!(result.line_count, 5);
+        let headers = result.headers.unwrap();
+        assert_eq!(headers.len(), 5);
+        assert!(headers.contains(&"Header 1".to_string()));
+        assert!(headers.contains(&"Header 2".to_string()));
+        assert!(headers.contains(&"Header 3".to_string()));
+        assert!(headers.contains(&"NoSpace".to_string()));
+        assert!(headers.contains(&"Multiple  spaces".to_string()));
+    }
+
+    #[test]
+    fn test_markdown_links() {
+        let text =
+            b"Check [Google](https://google.com) and [GitHub](https://github.com).\n[Another](https://example.com)";
+        let result = parse_text(text, true).unwrap();
+        let links = result.links.unwrap();
+        assert_eq!(links.len(), 3);
+        assert!(links.contains(&("Google".to_string(), "https://google.com".to_string())));
+        assert!(links.contains(&("GitHub".to_string(), "https://github.com".to_string())));
+        assert!(links.contains(&("Another".to_string(), "https://example.com".to_string())));
+    }
+
+    #[test]
+    fn test_markdown_code_blocks() {
+        let text = b"```python\ndef hello():\n    print(\"Hello\")\n```\n\n```javascript\nconsole.log(\"Hi\");\n```\n\n```\nplain code\n```";
+        let result = parse_text(text, true).unwrap();
+        let code_blocks = result.code_blocks.unwrap();
+        assert_eq!(code_blocks.len(), 3);
+
+        let python_block = code_blocks.iter().find(|(lang, _)| lang == "python").unwrap();
+        assert!(python_block.1.contains("def hello()"));
+
+        let js_block = code_blocks.iter().find(|(lang, _)| lang == "javascript").unwrap();
+        assert!(js_block.1.contains("console.log"));
+
+        let plain_block = code_blocks.iter().find(|(lang, _)| lang == "plain").unwrap();
+        assert!(plain_block.1.contains("plain code"));
+    }
+
+    #[test]
+    fn test_markdown_code_blocks_crlf() {
+        let text = b"```python\r\ndef hello():\r\n    print(\"Hello\")\r\n```\r\n";
+        let result = parse_text(text, true).unwrap();
+        let code_blocks = result.code_blocks.unwrap();
+        assert_eq!(code_blocks.len(), 1);
+        assert_eq!(code_blocks[0].0, "python");
+        assert!(code_blocks[0].1.contains("def hello()"));
+    }
+
+    #[test]
+    fn test_markdown_complex() {
+        let text = b"# Documentation\n\n## Overview\nThis is a [test](https://example.com).\n\n```python\nx = 42\n```\n\n## Another\nMore [links](https://test.com).";
+        let result = parse_text(text, true).unwrap();
+        assert!(result.line_count > 0);
+        assert!(result.word_count > 0);
+
+        let headers = result.headers.unwrap();
+        assert_eq!(headers.len(), 3);
+
+        let links = result.links.unwrap();
+        assert_eq!(links.len(), 2);
+
+        let code_blocks = result.code_blocks.unwrap();
+        assert_eq!(code_blocks.len(), 1);
+    }
+
+    #[test]
+    fn test_unicode_content() {
+        let text = "Hello 世界 🌍\nUnicode test".as_bytes();
+        let result = parse_text(text, false).unwrap();
+        assert!(result.content.contains("世界"));
+        assert!(result.content.contains("🌍"));
+        assert_eq!(result.line_count, 2);
+    }
+
+    #[test]
+    fn test_word_count_accuracy() {
+        let text = b"One two three four five.\nSix seven eight.\nNine.";
+        let result = parse_text(text, false).unwrap();
+        assert_eq!(result.line_count, 3);
+        assert_eq!(result.word_count, 9);
+    }
+
+    #[test]
+    fn test_headers_not_in_code_blocks() {
+        let text = b"# Real Header\n```\n# Not a header\n```\n## Another Real";
+        let result = parse_text(text, true).unwrap();
+        let headers = result.headers.unwrap();
+        assert_eq!(headers.len(), 2);
+        assert!(headers.contains(&"Real Header".to_string()));
+        assert!(headers.contains(&"Another Real".to_string()));
+        assert!(!headers.iter().any(|h| h.contains("Not a header")));
+    }
+
+    #[test]
+    fn test_links_not_in_code_blocks() {
+        let text = b"[Real Link](https://real.com)\n```\n[Not Link](https://fake.com)\n```";
+        let result = parse_text(text, true).unwrap();
+        let links = result.links.unwrap();
+        assert_eq!(links.len(), 1);
+        assert_eq!(links[0].0, "Real Link");
+        assert_eq!(links[0].1, "https://real.com");
+    }
+
+    #[test]
+    fn test_empty_code_block_language() {
+        let text = b"```\ncode without language\n```";
+        let result = parse_text(text, true).unwrap();
+        let code_blocks = result.code_blocks.unwrap();
+        assert_eq!(code_blocks.len(), 1);
+        assert_eq!(code_blocks[0].0, "plain");
+    }
+
+    #[test]
+    fn test_large_text_streaming() {
+        let large_text = "Line\n".repeat(10000);
+        let result = parse_text(large_text.as_bytes(), false).unwrap();
+        assert_eq!(result.line_count, 10000);
+        assert_eq!(result.word_count, 10000);
+    }
+}