github-featured/run-llama--liteparse
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Merge pull request #429 from run-llama/jw/bounded-memory-parsing

Browse Source 
feat: optimize memory during rasterization/OCR
This commit is contained in:
Logan
2026-08-19 21:53:55 -06:00
committed by GitHub
parent 9a2bc04e98 5552e6c953
commit 59b63ede9b
4 changed files with 363 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/liteparse/src/extract.rs
+11
View File
@@ -75,6 +75,12 @@ pub(crate) struct ExtractedPages {
/// mutates the open PDFium document, so a caller that still needs the
/// original widget annotations must reopen the input.
pub flattened_form_widgets: bool,
/// The page numbers extraction actually flattened. Flattening is a
/// per-page decision, so any consumer reproducing it on a reopened
/// document (e.g. OCR raster rendering) must apply it to exactly these
/// pages — flattening a page extraction never touched hides that page's
/// non-widget annotations from the raster.
pub flattened_page_numbers: Vec<u32>,
}
/// Same as `extract_pages_from_document` but optionally also renders every
@@ -96,6 +102,7 @@ pub(crate) fn extract_pages_and_images(
let mut image_cache = ImageCache::default();
let mut image_error_count = 0u32;
let mut flattened_form_widgets = false;
let mut flattened_page_numbers: Vec<u32> = Vec::new();
// One FFI call keeps the per-page annotation walk off the hot path for
// every document without an AcroForm catalog, which is nearly all of them.
let document_has_form = document.form_type() != 0;
@@ -136,6 +143,9 @@ pub(crate) fn extract_pages_and_images(
&mut page_errors,
)? {
Some(extraction) => {
if extraction.flattened_form_widgets {
flattened_page_numbers.push(page_number);
}
pages.push(extraction.page);
images.extend(extraction.images);
image_error_count += extraction.image_error_count;
@@ -154,6 +164,7 @@ pub(crate) fn extract_pages_and_images(
images,
image_error_count,
flattened_form_widgets,
flattened_page_numbers,
})
}
crates/liteparse/src/ocr_merge.rs
+35 -5
View File
@@ -434,14 +434,22 @@ fn count_columns(region: &Region, total_items: usize) -> usize {
/// to avoid excessive memory usage.
pub(crate) const MAX_OCR_RENDER_LONG_EDGE_PX: f32 = 4096.0;
/// Render the pages in `pages[start..]` that need OCR, stopping once
/// `max_rasters` of them have been rendered (`0` means no limit). Returns the
/// rasters plus the index to resume scanning from, so a caller can process a
/// long document in bounded rounds.
#[allow(clippy::too_many_arguments)]
pub(crate) fn render_pages_for_ocr(
document: &Document,
pages: &[Page],
start: usize,
max_rasters: usize,
dpi: f32,
grayscale: bool,
render_form_fields: bool,
continue_on_page_error: bool,
) -> Result<Vec<RenderedPage>, LiteParseError> {
flatten_page_numbers: &std::collections::HashSet<u32>,
) -> Result<(Vec<RenderedPage>, usize), LiteParseError> {
let mut rendered = Vec::new();
// With `render_form_fields`, draw form-field appearances into the OCR
// raster so filled-in form values are visible to the OCR engine (matches
@@ -453,9 +461,25 @@ pub(crate) fn render_pages_for_ocr(
if let Some(form) = form.as_ref() {
form.run_document_actions();
}
for (idx, page) in pages.iter().enumerate() {
let mut next_start = pages.len();
for (idx, page) in pages.iter().enumerate().skip(start) {
let page_render = (|| -> Result<Option<RenderedPage>, LiteParseError> {
let page_obj = document.page((page.page_number - 1) as i32)?;
let page_index = (page.page_number - 1) as i32;
// Text extraction may have flattened THIS page's widget
// annotations into page content. When the caller hands us a
// freshly reopened document, re-apply the flatten on exactly the
// pages extraction flattened, so the raster shows the same
// content extraction saw. Never flatten any other page:
// flattening hides a page's non-widget annotations (stamps,
// highlights, free text) from the raster, losing their text.
let page_obj = if flatten_page_numbers.contains(&(page.page_number as u32)) {
match document.flatten_form_widgets(page_index)? {
Some(flattened_page) => flattened_page,
None => document.page(page_index)?,
}
} else {
document.page(page_index)?
};
let page_complexity = calculate_page_complexity(page, &page_obj)?;
if !page_complexity.needs_ocr {
@@ -491,7 +515,13 @@ pub(crate) fn render_pages_for_ocr(
}))
})();
match page_render {
Ok(Some(render)) => rendered.push(render),
Ok(Some(render)) => {
rendered.push(render);
if max_rasters > 0 && rendered.len() >= max_rasters {
next_start = idx + 1;
break;
}
}
Ok(None) => {}
// The page already extracted successfully, so a tolerant parse
// keeps its native text and only forgoes the OCR enrichment.
@@ -502,7 +532,7 @@ pub(crate) fn render_pages_for_ocr(
Err(error) => return Err(error),
}
}
Ok(rendered)
Ok((rendered, next_start))
}
/// Run OCR on pre-rendered page bitmaps and merge results into `pages`.
crates/liteparse/src/parser.rs
+66 -39
View File
@@ -581,7 +581,6 @@ impl LiteParse {
pages,
page_errors,
total_pages,
ocr_rendered,
outline,
mut images,
screenshots,
@@ -592,6 +591,9 @@ impl LiteParse {
producer,
doc_meta,
xfa_packets,
flattened_form_widgets,
flattened_page_numbers,
repaired_input,
) = {
let lib = Library::init();
#[cfg(not(target_arch = "wasm32"))]
@@ -673,6 +675,7 @@ impl LiteParse {
images,
image_error_count,
flattened_form_widgets,
flattened_page_numbers,
} = extracted;
// Reopening the input costs a full parse, so it is confined to the
// one consumer that genuinely needs live widget annotations: the
@@ -680,12 +683,14 @@ impl LiteParse {
// run document actions and paint computed field appearances that
// have no appearance stream to flatten.
//
// Plain rendering (OCR rasters, screenshots) does not need it —
// flattening promotes the widget appearances into page content,
// so the raster is the same either way. Complexity likewise runs
// on the flattened document by design (see `is_complex`).
// Plain rendering (screenshots) does not need it — flattening
// promotes the widget appearances into page content, so the
// raster is the same either way. Complexity likewise runs on the
// flattened document by design (see `is_complex`). OCR rasters
// render in bounded rounds after this section, each against a
// freshly reopened (hence pristine) document.
let needs_pristine_document = flattened_form_widgets
&& (self.config.ocr_enabled || self.config.extract_screenshots)
&& self.config.extract_screenshots
&& self.config.render_form_fields;
let pristine_document = needs_pristine_document
.then(|| extract::load_document_from_input(&lib, document_input, password))
@@ -697,28 +702,6 @@ impl LiteParse {
t_extract.duration_since(t0).as_secs_f64() * 1000.0,
pages.len()
));
let rendered = if self.config.ocr_enabled {
let r = ocr_merge::render_pages_for_ocr(
analysis_document,
&pages,
self.config.dpi,
ocr_grayscale,
self.config.render_form_fields,
self.config.continue_on_page_error,
)?;
log(&format!(
"[liteparse] ocr render: {:.1}ms ({} pages)",
web_time::Instant::now()
.duration_since(t_extract)
.as_secs_f64()
* 1000.0,
r.len()
));
r
} else {
Vec::new()
};
let complexity = if self.config.include_complexity {
let mut complexity = Vec::with_capacity(pages.len());
for page in &pages {
@@ -768,12 +751,13 @@ impl LiteParse {
} else {
Vec::new()
};
#[cfg(target_arch = "wasm32")]
let repaired_input: Option<crate::types::PdfInput> = None;
// `lib` is dropped here, releasing the PDFium lock.
(
pages,
page_errors,
total_pages,
rendered,
outline,
images,
screenshots,
@@ -784,22 +768,65 @@ impl LiteParse {
producer,
doc_meta,
xfa_packets,
flattened_form_widgets,
flattened_page_numbers,
repaired_input,
)
};
let mut pages = pages;
let t1 = web_time::Instant::now();
// OCR pass (engine resolved before the render block above).
if let Some(engine) = ocr_engine {
ocr_merge::ocr_and_merge_rendered(
&mut pages,
ocr_rendered,
engine,
&self.config.ocr_language,
self.config.num_workers,
self.config.ocr_failure_fatal,
)
.await?;
let round_rasters = self.config.num_workers.max(1);
// Extraction may have flattened SOME pages' form widgets into
// page content in its (now dropped) document instance; re-apply
// per round on exactly those pages so the rasters match what
// extraction saw. With `render_form_fields` the form environment
// paints the widgets instead, so no re-flatten is needed.
let reflatten_pages: std::collections::HashSet<u32> =
if flattened_form_widgets && !self.config.render_form_fields {
flattened_page_numbers.iter().copied().collect()
} else {
std::collections::HashSet::new()
};
let ocr_input = repaired_input.as_ref().unwrap_or(validated_input);
let mut round_start = 0usize;
while round_start < pages.len() {
let (rendered, next_start) = {
let lib = Library::init();
let document = extract::load_document_from_input(&lib, ocr_input, password)?;
ocr_merge::render_pages_for_ocr(
&document,
&pages,
round_start,
round_rasters,
self.config.dpi,
ocr_grayscale,
self.config.render_form_fields,
self.config.continue_on_page_error,
&reflatten_pages,
)?
// `lib` drops here, releasing the PDFium lock before the
// engine's async recognition below.
};
round_start = next_start;
if rendered.is_empty() {
// The scan reached the end without finding another page
// that needs OCR.
continue;
}
// `RenderedPage::idx` is absolute, so the whole slice is
// passed regardless of where this round started.
ocr_merge::ocr_and_merge_rendered(
&mut pages,
rendered,
engine.clone(),
&self.config.ocr_language,
self.config.num_workers,
self.config.ocr_failure_fatal,
)
.await?;
}
}
let t_ocr = web_time::Instant::now();
log(&format!(
crates/liteparse/tests/integration_test.rs
+251
View File
@@ -641,3 +641,254 @@ async fn test_filled_acroform_values_are_extracted_as_text() {
"widget text is extractable after flattening, so it is not annotation-only text"
);
}
/// A blank multi-page PDF: no text, so every page is text-poor and routes to
/// OCR. Pages take distinct sizes so each page's raster is uniquely
/// identifiable by the dimensions the OCR engine receives.
fn blank_pdf(page_sizes: &[(u32, u32)]) -> Vec<u8> {
let kids: Vec<String> = (0..page_sizes.len())
.map(|i| format!("{} 0 R", i + 3))
.collect();
let mut objects: Vec<Vec<u8>> = vec![
b"<< /Type /Catalog /Pages 2 0 R >>".to_vec(),
format!(
"<< /Type /Pages /Kids [{}] /Count {} >>",
kids.join(" "),
page_sizes.len()
)
.into_bytes(),
];
for (width, height) in page_sizes {
objects.push(
format!("<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {width} {height}] >>")
.into_bytes(),
);
}
let mut pdf = b"%PDF-1.7\n".to_vec();
let mut offsets = Vec::with_capacity(objects.len());
for (index, object) in objects.iter().enumerate() {
offsets.push(pdf.len());
pdf.extend_from_slice(format!("{} 0 obj\n", index + 1).as_bytes());
pdf.extend_from_slice(object);
pdf.extend_from_slice(b"\nendobj\n");
}
let xref = pdf.len();
pdf.extend_from_slice(format!("xref\n0 {}\n", objects.len() + 1).as_bytes());
pdf.extend_from_slice(b"0000000000 65535 f \n");
for offset in offsets {
pdf.extend_from_slice(format!("{offset:010} 00000 n \n").as_bytes());
}
pdf.extend_from_slice(
format!(
"trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{}\n%%EOF",
objects.len() + 1,
xref
)
.as_bytes(),
);
pdf
}
/// An OCR engine that reports each raster's dimensions as its recognized
/// text (so misrouted rasters are detectable per page), counts calls, and
/// tracks the concurrent-recognition high-water mark (so the round structure
/// itself is observable: serialized rounds of one page cap it at 1).
struct ProbeEngine {
calls: std::sync::Arc<std::sync::atomic::AtomicUsize>,
in_flight: std::sync::Arc<std::sync::atomic::AtomicUsize>,
peak_in_flight: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}
impl liteparse::ocr::OcrEngine for ProbeEngine {
fn name(&self) -> &str {
"probe"
}
fn recognize<'a, 'b: 'a, 'c: 'a>(
&'a self,
_image_data: &'c [u8],
width: u32,
height: u32,
_options: &'b liteparse::ocr::OcrOptions,
) -> std::pin::Pin<
Box<
dyn Future<
Output = Result<
Vec<liteparse::ocr::OcrResult>,
Box<dyn std::error::Error + Send + Sync>,
>,
> + Send
+ '_,
>,
> {
use std::sync::atomic::Ordering;
self.calls.fetch_add(1, Ordering::SeqCst);
let now_in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
self.peak_in_flight
.fetch_max(now_in_flight, Ordering::SeqCst);
let in_flight = self.in_flight.clone();
Box::pin(async move {
// Hold the slot briefly so overlapping recognitions overlap
// observably; a serialized round of one page keeps the peak at 1.
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
in_flight.fetch_sub(1, Ordering::SeqCst);
Ok(vec![liteparse::ocr::OcrResult {
text: format!("DIM{width}x{height}"),
bbox: [10.0, 10.0, 200.0, 40.0],
confidence: 0.99,
polygon: None,
}])
})
}
}
/// OCR runs in render→recognize rounds of `num_workers` pages. Whatever the
/// round size, every page must be recognized exactly once and each page's OCR
/// text must land on the page whose raster produced it — distinct page sizes
/// make a misroute visible, which is the failure mode the per-round document
/// reopen and form-widget re-flatten could introduce. `num_workers: 1` forces
/// one page per round (four rounds over four pages) and must also serialize
/// recognition; `num_workers: 4` covers the whole document in a single round
/// with overlapping recognition. Both must agree page-for-page.
#[tokio::test]
#[serial]
async fn test_ocr_rounds_cover_every_page_once() {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
// Distinct sizes so a raster delivered to the wrong page is detectable.
let page_sizes: [(u32, u32); 4] = [(612, 792), (400, 500), (300, 300), (500, 900)];
let expected_texts: Vec<String> = page_sizes
.iter()
.map(|(width, height)| {
let scale = 150.0 / 72.0;
format!(
"DIM{}x{}",
(*width as f32 * scale).round() as u32,
(*height as f32 * scale).round() as u32
)
})
.collect();
let run = |num_workers: usize| {
let calls = Arc::new(AtomicUsize::new(0));
let peak = Arc::new(AtomicUsize::new(0));
let engine = ProbeEngine {
calls: calls.clone(),
in_flight: Arc::new(AtomicUsize::new(0)),
peak_in_flight: peak.clone(),
};
let parser = LiteParse::new(LiteParseConfig {
ocr_enabled: true,
num_workers,
dpi: 150.0,
quiet: true,
..Default::default()
})
.with_ocr_engine(std::sync::Arc::new(engine));
(parser, calls, peak)
};
let assert_pages_carry_own_rasters = |result: &liteparse::ParseResult, label: &str| {
assert_eq!(result.pages.len(), page_sizes.len(), "{label}: page count");
for (page, expected) in result.pages.iter().zip(&expected_texts) {
let text: String = page
.text_items
.iter()
.map(|item| item.text.as_str())
.collect();
assert!(
text.contains(expected.as_str()),
"{label}: page {} carries {text:?}, expected {expected:?} — OCR text landed on the wrong page",
page.page_number
);
}
};
// One page per round: four rounds, recognition fully serialized.
let (parser, calls, peak) = run(1);
let serialized = parser
.parse_input(PdfInput::Bytes(blank_pdf(&page_sizes)))
.await
.expect("single-page-round OCR parse should succeed");
assert_pages_carry_own_rasters(&serialized, "rounds of 1");
assert_eq!(
calls.load(Ordering::SeqCst),
page_sizes.len(),
"one recognition per page"
);
assert_eq!(
peak.load(Ordering::SeqCst),
1,
"num_workers=1 must serialize recognition"
);
// Round wide enough for the whole document: one round, overlapping
// recognition, identical routing.
let (parser, calls, peak) = run(4);
let overlapped = parser
.parse_input(PdfInput::Bytes(blank_pdf(&page_sizes)))
.await
.expect("single-round OCR parse should succeed");
assert_pages_carry_own_rasters(&overlapped, "rounds of 4");
assert_eq!(calls.load(Ordering::SeqCst), page_sizes.len());
assert!(
peak.load(Ordering::SeqCst) > 1,
"num_workers=4 over 4 pages must overlap recognition within the round"
);
}
/// A round is bounded by rasters rendered, not by page span, so OCR-needing
/// pages that are sparsely scattered through a mostly-native-text document
/// still fill a round and recognize concurrently.
///
/// This guards a real regression: bounding the round by page span instead
/// made each round contain only the OCR-needing pages that happened to fall
/// inside its span — often one or two — which starved the worker pool and
/// serialized recognition. On this document that was a 2.4x wall-clock loss
/// at realistic OCR latency, with no test failing.
#[tokio::test]
#[serial]
async fn test_ocr_rounds_fill_across_sparse_pages() {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
let calls = Arc::new(AtomicUsize::new(0));
let peak = Arc::new(AtomicUsize::new(0));
let parser = LiteParse::new(LiteParseConfig {
ocr_enabled: true,
num_workers: 8,
dpi: 72.0, // Small rasters: this test is about scheduling, not pixels.
quiet: true,
..Default::default()
})
.with_ocr_engine(std::sync::Arc::new(ProbeEngine {
calls: calls.clone(),
in_flight: Arc::new(AtomicUsize::new(0)),
peak_in_flight: peak.clone(),
}));
let result = parser
.parse("../../demo/docs/apple-10k-2024.pdf")
.await
.expect("should parse the 10-K with OCR enabled");
let recognized = calls.load(Ordering::SeqCst);
assert!(
recognized > 8,
"expected the fixture to need OCR on more than one round's worth of pages, got {recognized}"
);
assert!(
result.pages.len() > recognized,
"fixture should be mostly native text, so OCR pages ({recognized}) must be sparse \
among its {} pages",
result.pages.len()
);
// The scan-ahead must gather a full round even though the OCR-needing
// pages are interleaved with native-text pages it skips.
assert_eq!(
peak.load(Ordering::SeqCst),
8,
"rounds must fill to num_workers across skipped pages; a lower peak means \
rounds are being cut short by page span"
);
}