kreuzberg 4.0.0.pre.rc.29 → 4.0.0.rc1

data/vendor/kreuzberg/src/extractors/pdf.rs

@@ -3,28 +3,19 @@
 use crate::Result;
 use crate::core::config::ExtractionConfig;
 use crate::plugins::{DocumentExtractor, Plugin};
-use crate::types::{ExtractionResult, Metadata, PageContent};
+use crate::types::{ExtractionResult, Metadata};
 use async_trait::async_trait;
-#[cfg(feature = "tokio-runtime")]
 use std::path::Path;
 
 #[cfg(feature = "pdf")]
 use crate::pdf::error::PdfError;
 #[cfg(feature = "ocr")]
 use crate::pdf::rendering::{PageRenderOptions, PdfRenderer};
-#[cfg(feature = "pdf")]
+#[cfg(all(feature = "pdf", feature = "ocr"))]
 use crate::types::Table;
 #[cfg(feature = "pdf")]
 use pdfium_render::prelude::*;
 
-#[cfg(feature = "pdf")]
-type PdfExtractionPhaseResult = (
-    crate::pdf::metadata::PdfExtractionMetadata,
-    String,
-    Vec<Table>,
-    Option<Vec<PageContent>>,
-);
-
 #[cfg(feature = "ocr")]
 const MIN_TOTAL_NON_WHITESPACE: usize = 64;
 #[cfg(feature = "ocr")]
@@ -146,37 +137,41 @@ fn evaluate_native_text_for_ocr(native_text: &str, page_count: Option<usize>) ->
 ///
 /// This function converts PDF character positions to HocrWord format,
 /// then uses the existing table reconstruction logic to detect tables.
-///
-/// Uses the shared PdfDocument reference (wrapped in Arc<RwLock<>> for thread-safety).
 #[cfg(all(feature = "pdf", feature = "ocr"))]
 fn extract_tables_from_document(
     document: &PdfDocument,
-    _metadata: &crate::pdf::metadata::PdfExtractionMetadata,
+    _metadata: &crate::pdf::metadata::PdfMetadata,
 ) -> Result<Vec<Table>> {
     use crate::ocr::table::{reconstruct_table, table_to_markdown};
     use crate::pdf::table::extract_words_from_page;
 
     let mut all_tables = Vec::new();
 
+    // Process each page
     for (page_index, page) in document.pages().iter().enumerate() {
-        let words = extract_words_from_page(&page, 0.0)?;
+        // Extract words with positions from the page
+        let words = extract_words_from_page(&page, 0.0)?; // Use 0.0 confidence for PDF (always high quality)
 
         if words.is_empty() {
             continue;
         }
 
+        // Use existing table reconstruction logic
+        // These thresholds match the defaults from TesseractConfig
         let column_threshold = 50;
         let row_threshold_ratio = 0.5;
 
-        let table_cells = reconstruct_table(&words, column_threshold, row_threshold_ratio);
+        // Reconstruct table from positioned words
+        let table_cells = reconstruct_table(&words, column_threshold, row_threshold_ratio, true);
 
         if !table_cells.is_empty() {
+            // Generate markdown representation
             let markdown = table_to_markdown(&table_cells);
 
             all_tables.push(Table {
                 cells: table_cells,
                 markdown,
-                page_number: page_index + 1,
+                page_number: page_index + 1, // 1-indexed
             });
         }
     }
@@ -188,47 +183,11 @@ fn extract_tables_from_document(
 #[cfg(all(feature = "pdf", not(feature = "ocr")))]
 fn extract_tables_from_document(
     _document: &PdfDocument,
-    _metadata: &crate::pdf::metadata::PdfExtractionMetadata,
+    _metadata: &crate::pdf::metadata::PdfMetadata,
 ) -> Result<Vec<crate::types::Table>> {
     Ok(vec![])
 }
 
-/// Helper function to assign tables and images to pages.
-///
-/// If page_contents is None, returns None (no per-page tracking enabled).
-/// Otherwise, iterates through tables and images, assigning them to pages based on page_number.
-///
-/// # Performance
-///
-/// Uses Arc::new to wrap tables and images, avoiding expensive copies.
-/// This reduces memory overhead by enabling zero-copy sharing of table/image data
-/// across multiple references (e.g., when the same table appears on multiple pages).
-fn assign_tables_and_images_to_pages(
-    mut page_contents: Option<Vec<PageContent>>,
-    tables: &[crate::types::Table],
-    images: &[crate::types::ExtractedImage],
-) -> Option<Vec<PageContent>> {
-    let pages = page_contents.take()?;
-
-    let mut updated_pages = pages;
-
-    for table in tables {
-        if let Some(page) = updated_pages.iter_mut().find(|p| p.page_number == table.page_number) {
-            page.tables.push(std::sync::Arc::new(table.clone()));
-        }
-    }
-
-    for image in images {
-        if let Some(page_num) = image.page_number
-            && let Some(page) = updated_pages.iter_mut().find(|p| p.page_number == page_num)
-        {
-            page.images.push(std::sync::Arc::new(image.clone()));
-        }
-    }
-
-    Some(updated_pages)
-}
-
 /// PDF document extractor using pypdfium2 and playa-pdf.
 pub struct PdfExtractor;
 
@@ -243,41 +202,6 @@ impl PdfExtractor {
         Self
     }
 
-    /// Extract text, metadata, and tables from a PDF document using a single shared instance.
-    ///
-    /// This method consolidates all PDF extraction phases (text, metadata, tables) into a single
-    /// operation using a single PdfDocument instance. This avoids redundant document parsing
-    /// and pdfium initialization overhead.
-    ///
-    /// # Performance
-    ///
-    /// By reusing a single document instance across all extraction phases, we eliminate:
-    /// - Duplicate document parsing overhead (25-40ms saved)
-    /// - Redundant pdfium bindings initialization
-    /// - Multiple page tree traversals
-    ///
-    /// Expected improvement: 20-30% faster PDF processing.
-    ///
-    /// # Returns
-    ///
-    /// A tuple containing:
-    /// - PDF metadata (title, authors, dates, page structure, etc.)
-    /// - Native extracted text (or empty if using OCR)
-    /// - Extracted tables (if OCR feature enabled)
-    /// - Per-page content (if page extraction configured)
-    #[cfg(feature = "pdf")]
-    fn extract_all_from_document(
-        document: &PdfDocument,
-        config: &ExtractionConfig,
-    ) -> Result<PdfExtractionPhaseResult> {
-        let (native_text, _boundaries, page_contents, pdf_metadata) =
-            crate::pdf::text::extract_text_and_metadata_from_pdf_document(document, Some(config))?;
-
-        let tables = extract_tables_from_document(document, &pdf_metadata)?;
-
-        Ok((pdf_metadata, native_text, tables, page_contents))
-    }
-
     /// Extract text from PDF using OCR.
     ///
     /// Renders all pages to images and processes them with OCR.
@@ -363,13 +287,6 @@ impl Plugin for PdfExtractor {
 
 #[async_trait]
 impl DocumentExtractor for PdfExtractor {
-    #[cfg_attr(feature = "otel", tracing::instrument(
-        skip(self, content, config),
-        fields(
-            extractor.name = self.name(),
-            content.size_bytes = content.len(),
-        )
-    ))]
     async fn extract_bytes(
         &self,
         content: &[u8],
@@ -377,26 +294,18 @@ impl DocumentExtractor for PdfExtractor {
         config: &ExtractionConfig,
     ) -> Result<ExtractionResult> {
         #[cfg(feature = "pdf")]
-        let (pdf_metadata, native_text, tables, page_contents) = {
-            #[cfg(target_arch = "wasm32")]
-            {
-                let pdfium = crate::pdf::bindings::bind_pdfium(PdfError::MetadataExtractionFailed, "initialize Pdfium")
-                    .map_err(|pdf_err| {
-                        if pdf_err.to_string().contains("WASM") || pdf_err.to_string().contains("Module") {
-                            crate::error::KreuzbergError::Parsing {
-                                message: "PDF extraction requires proper WASM module initialization. \
-                                     Ensure your WASM environment is set up with PDFium support. \
-                                     See: https://docs.kreuzberg.dev/wasm/pdf"
-                                    .to_string(),
-                                source: None,
-                            }
-                        } else {
-                            pdf_err.into()
-                        }
-                    })?;
-
-                let document = pdfium.load_pdf_from_byte_slice(content, None).map_err(|e| {
-                    let err_msg = crate::pdf::error::format_pdfium_error(e);
+        let (pdf_metadata, native_text, tables) = if crate::core::batch_mode::is_batch_mode() {
+            // Batch mode: Move PDF extraction to blocking thread pool to enable parallelism
+            let content_owned = content.to_vec();
+            tokio::task::spawn_blocking(move || {
+                let bindings = Pdfium::bind_to_library(Pdfium::pdfium_platform_library_name_at_path("./"))
+                    .or_else(|_| Pdfium::bind_to_system_library())
+                    .map_err(|e| PdfError::MetadataExtractionFailed(format!("Failed to initialize Pdfium: {}", e)))?;
+
+                let pdfium = Pdfium::new(bindings);
+
+                let document = pdfium.load_pdf_from_byte_slice(&content_owned, None).map_err(|e| {
+                    let err_msg = e.to_string();
                     if err_msg.contains("password") || err_msg.contains("Password") {
                         PdfError::PasswordRequired
                     } else {
@@ -404,79 +313,40 @@ impl DocumentExtractor for PdfExtractor {
                     }
                 })?;
 
-                Self::extract_all_from_document(&document, config)?
-            }
-            #[cfg(all(not(target_arch = "wasm32"), feature = "tokio-runtime"))]
-            {
-                if crate::core::batch_mode::is_batch_mode() {
-                    let content_owned = content.to_vec();
-                    let span = tracing::Span::current();
-                    let config_owned = config.clone();
-                    tokio::task::spawn_blocking(move || {
-                        let _guard = span.entered();
-
-                        let pdfium =
-                            crate::pdf::bindings::bind_pdfium(PdfError::MetadataExtractionFailed, "initialize Pdfium")?;
-
-                        let document = pdfium.load_pdf_from_byte_slice(&content_owned, None).map_err(|e| {
-                            let err_msg = crate::pdf::error::format_pdfium_error(e);
-                            if err_msg.contains("password") || err_msg.contains("Password") {
-                                PdfError::PasswordRequired
-                            } else {
-                                PdfError::InvalidPdf(err_msg)
-                            }
-                        })?;
-
-                        let (pdf_metadata, native_text, tables, page_contents) =
-                            Self::extract_all_from_document(&document, &config_owned)?;
-
-                        if let Some(page_cfg) = config_owned.pages.as_ref()
-                            && page_cfg.extract_pages
-                            && page_contents.is_none()
-                        {
-                            return Err(PdfError::ExtractionFailed(
-                                "Page extraction was configured but no page data was extracted in batch mode"
-                                    .to_string(),
-                            )
-                            .into());
-                        }
-
-                        Ok::<_, crate::error::KreuzbergError>((pdf_metadata, native_text, tables, page_contents))
-                    })
-                    .await
-                    .map_err(|e| crate::error::KreuzbergError::Other(format!("PDF extraction task failed: {}", e)))??
+                let metadata = crate::pdf::metadata::extract_metadata_from_document(&document)?;
+                let native_text = crate::pdf::text::extract_text_from_pdf_document(&document)?;
+
+                // Extract tables from native PDF text (when not using OCR)
+                let tables = extract_tables_from_document(&document, &metadata)?;
+
+                Ok::<_, crate::error::KreuzbergError>((metadata, native_text, tables))
+            })
+            .await
+            .map_err(|e| crate::error::KreuzbergError::Other(format!("PDF extraction task failed: {}", e)))??
+        } else {
+            // Single-file mode: Direct extraction (no spawn overhead)
+            let bindings = Pdfium::bind_to_library(Pdfium::pdfium_platform_library_name_at_path("./"))
+                .or_else(|_| Pdfium::bind_to_system_library())
+                .map_err(|e| PdfError::MetadataExtractionFailed(format!("Failed to initialize Pdfium: {}", e)))?;
+
+            let pdfium = Pdfium::new(bindings);
+
+            let document = pdfium.load_pdf_from_byte_slice(content, None).map_err(|e| {
+                let err_msg = e.to_string();
+                if err_msg.contains("password") || err_msg.contains("Password") {
+                    PdfError::PasswordRequired
                 } else {
-                    let pdfium =
-                        crate::pdf::bindings::bind_pdfium(PdfError::MetadataExtractionFailed, "initialize Pdfium")?;
-
-                    let document = pdfium.load_pdf_from_byte_slice(content, None).map_err(|e| {
-                        let err_msg = crate::pdf::error::format_pdfium_error(e);
-                        if err_msg.contains("password") || err_msg.contains("Password") {
-                            PdfError::PasswordRequired
-                        } else {
-                            PdfError::InvalidPdf(err_msg)
-                        }
-                    })?;
-
-                    Self::extract_all_from_document(&document, config)?
+                    PdfError::InvalidPdf(err_msg)
                 }
-            }
-            #[cfg(all(not(target_arch = "wasm32"), not(feature = "tokio-runtime")))]
-            {
-                let pdfium =
-                    crate::pdf::bindings::bind_pdfium(PdfError::MetadataExtractionFailed, "initialize Pdfium")?;
+            })?;
 
-                let document = pdfium.load_pdf_from_byte_slice(content, None).map_err(|e| {
-                    let err_msg = crate::pdf::error::format_pdfium_error(e);
-                    if err_msg.contains("password") || err_msg.contains("Password") {
-                        PdfError::PasswordRequired
-                    } else {
-                        PdfError::InvalidPdf(err_msg)
-                    }
-                })?;
+            let metadata = crate::pdf::metadata::extract_metadata_from_document(&document)?;
+            let native_text = crate::pdf::text::extract_text_from_pdf_document(&document)?;
 
-                Self::extract_all_from_document(&document, config)?
-            }
+            // Extract tables from native PDF text (when not using OCR)
+            let tables = extract_tables_from_document(&document, &metadata)?;
+
+            (metadata, native_text, tables)
         };
 
         #[cfg(feature = "ocr")]
@@ -487,19 +357,20 @@ impl DocumentExtractor for PdfExtractor {
                 native_text
             }
         } else if config.ocr.is_some() {
-            let decision = evaluate_native_text_for_ocr(&native_text, None);
+            let decision = evaluate_native_text_for_ocr(&native_text, pdf_metadata.page_count);
 
             if std::env::var("KREUZBERG_DEBUG_OCR").is_ok() {
                 eprintln!(
                     "[kreuzberg::pdf::ocr] fallback={} non_whitespace={} alnum={} meaningful_words={} \
-                     avg_non_whitespace={:.2} avg_alnum={:.2} alnum_ratio={:.3}",
+                     avg_non_whitespace={:.2} avg_alnum={:.2} alnum_ratio={:.3} pages={}",
                     decision.fallback,
                     decision.stats.non_whitespace,
                     decision.stats.alnum,
                     decision.stats.meaningful_words,
                     decision.avg_non_whitespace,
                     decision.avg_alnum,
-                    decision.stats.alnum_ratio
+                    decision.stats.alnum_ratio,
+                    pdf_metadata.page_count.unwrap_or(0)
                 );
             }
 
@@ -515,22 +386,7 @@ impl DocumentExtractor for PdfExtractor {
         #[cfg(not(feature = "ocr"))]
         let text = native_text;
 
-        #[cfg(feature = "pdf")]
-        if let Some(ref page_cfg) = config.pages
-            && page_cfg.insert_page_markers
-        {
-            let marker_placeholder = page_cfg.marker_format.replace("{page_num}", "");
-            if !marker_placeholder.is_empty() && !text.contains(&marker_placeholder) {
-                #[cfg(feature = "otel")]
-                tracing::warn!(
-                    "Page markers were configured but none found in extracted content. \
-                     This may indicate very short documents or incomplete extraction."
-                );
-            }
-        }
-
-        let images = if config.images.as_ref().map(|c| c.extract_images).unwrap_or(false) {
-            // Image extraction is enabled, extract images if present
+        let images = if config.images.is_some() {
             match crate::pdf::images::extract_images_from_pdf(content) {
                 Ok(pdf_images) => Some(
                     pdf_images
@@ -554,41 +410,23 @@ impl DocumentExtractor for PdfExtractor {
                         })
                         .collect(),
                 ),
-                // If extraction fails, return empty vector instead of None
-                Err(_) => Some(vec![]),
+                Err(_) => None,
             }
         } else {
-            // Image extraction is not enabled
             None
         };
 
-        let final_pages = assign_tables_and_images_to_pages(page_contents, &tables, images.as_deref().unwrap_or(&[]));
+        // Tables were extracted during metadata/text extraction phase
+        // (see extract_tables_from_document function below)
 
         Ok(ExtractionResult {
             content: text,
             mime_type: mime_type.to_string(),
             metadata: Metadata {
                 #[cfg(feature = "pdf")]
-                title: pdf_metadata.title.clone(),
-                #[cfg(feature = "pdf")]
-                subject: pdf_metadata.subject.clone(),
-                #[cfg(feature = "pdf")]
-                authors: pdf_metadata.authors.clone(),
-                #[cfg(feature = "pdf")]
-                keywords: pdf_metadata.keywords.clone(),
-                #[cfg(feature = "pdf")]
-                created_at: pdf_metadata.created_at.clone(),
-                #[cfg(feature = "pdf")]
-                modified_at: pdf_metadata.modified_at.clone(),
-                #[cfg(feature = "pdf")]
-                created_by: pdf_metadata.created_by.clone(),
-                #[cfg(feature = "pdf")]
-                pages: pdf_metadata.page_structure.clone(),
-                #[cfg(feature = "pdf")]
-                format: Some(crate::types::FormatMetadata::Pdf(pdf_metadata.pdf_specific)),
+                format: Some(crate::types::FormatMetadata::Pdf(pdf_metadata)),
                 ..Default::default()
             },
-            pages: final_pages,
             tables,
             detected_languages: None,
             chunks: None,
@@ -596,7 +434,6 @@ impl DocumentExtractor for PdfExtractor {
         })
     }
 
-    #[cfg(feature = "tokio-runtime")]
     async fn extract_file(&self, path: &Path, mime_type: &str, config: &ExtractionConfig) -> Result<ExtractionResult> {
         let bytes = tokio::fs::read(path).await?;
         self.extract_bytes(&bytes, mime_type, config).await
@@ -656,106 +493,4 @@ mod tests {
         let sample = " . , ; : -- -- ";
         assert!(evaluate_native_text_for_ocr(sample, Some(2)).fallback);
     }
-
-    #[tokio::test]
-    #[cfg(feature = "pdf")]
-    async fn test_pdf_batch_mode_validates_page_config_enabled() {
-        use crate::core::config::PageConfig;
-
-        let extractor = PdfExtractor::new();
-
-        let config = ExtractionConfig {
-            pages: Some(PageConfig {
-                extract_pages: true,
-                insert_page_markers: false,
-                marker_format: "<!-- PAGE {page_num} -->".to_string(),
-            }),
-            ..Default::default()
-        };
-
-        let pdf_path =
-            std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../test_documents/pdfs/google_doc_document.pdf");
-        if let Ok(content) = std::fs::read(pdf_path) {
-            let result = extractor.extract_bytes(&content, "application/pdf", &config).await;
-            assert!(
-                result.is_ok(),
-                "Failed to extract PDF with page config: {:?}",
-                result.err()
-            );
-
-            let extraction_result = result.unwrap();
-            assert!(
-                extraction_result.pages.is_some(),
-                "Pages should be extracted when extract_pages is true"
-            );
-        }
-    }
-
-    #[tokio::test]
-    #[cfg(feature = "pdf")]
-    async fn test_pdf_batch_mode_validates_page_config_disabled() {
-        let extractor = PdfExtractor::new();
-        let config = ExtractionConfig::default();
-
-        let pdf_path =
-            std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../test_documents/pdfs/google_doc_document.pdf");
-        if let Ok(content) = std::fs::read(pdf_path) {
-            let result = extractor.extract_bytes(&content, "application/pdf", &config).await;
-            assert!(
-                result.is_ok(),
-                "Failed to extract PDF without page config: {:?}",
-                result.err()
-            );
-
-            let extraction_result = result.unwrap();
-            assert!(
-                extraction_result.pages.is_none(),
-                "Pages should not be extracted when pages config is None"
-            );
-        }
-    }
-
-    #[tokio::test]
-    #[cfg(feature = "pdf")]
-    async fn test_pdf_page_marker_validation() {
-        use crate::core::config::PageConfig;
-
-        let extractor = PdfExtractor::new();
-
-        let config = ExtractionConfig {
-            pages: Some(PageConfig {
-                extract_pages: true,
-                insert_page_markers: true,
-                marker_format: "\n\n<!-- PAGE {page_num} -->\n\n".to_string(),
-            }),
-            ..Default::default()
-        };
-
-        let pdf_path =
-            std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../test_documents/pdfs/multi_page.pdf");
-        if let Ok(content) = std::fs::read(pdf_path) {
-            let result = extractor.extract_bytes(&content, "application/pdf", &config).await;
-            assert!(
-                result.is_ok(),
-                "Failed to extract PDF with page markers: {:?}",
-                result.err()
-            );
-
-            let extraction_result = result.unwrap();
-            let marker_placeholder = "<!-- PAGE ";
-            if extraction_result.content.len() > 100 {
-                assert!(
-                    extraction_result.content.contains(marker_placeholder),
-                    "Page markers should be inserted when configured and document has multiple pages"
-                );
-            }
-        }
-    }
-
-    #[test]
-    #[cfg(feature = "pdf")]
-    fn test_pdf_extractor_without_feature_pdf() {
-        let extractor = PdfExtractor::new();
-        assert_eq!(extractor.name(), "pdf-extractor");
-    }
 }