Skip to main content

cloudiful_docling_convert/processor/
chunking.rs

1use crate::document::PdfConvertOptions;
2use crate::models::{ChunkMetadata, PdfInfo};
3
4pub fn build_pdf_chunk_plan(options: &PdfConvertOptions, pdf_info: &PdfInfo) -> Vec<ChunkMetadata> {
5    if options.split_by_bookmark {
6        pdf_info.calculate_hybrid_ranges(options.pages_per_file)
7    } else if options.split_input {
8        pdf_info.calculate_page_ranges(options.pages_per_file)
9    } else {
10        vec![ChunkMetadata {
11            start_page: 1,
12            end_page: pdf_info.total_pages,
13            bookmark_index: None,
14            bookmark_title: None,
15        }]
16    }
17}
18
19#[cfg(test)]
20mod tests {
21    use super::*;
22
23    #[test]
24    fn pdf_chunk_plan_stays_pdf_specific() {
25        let pdf_info = PdfInfo {
26            total_pages: 10,
27            outlines: Vec::new(),
28        };
29        let plan = build_pdf_chunk_plan(
30            &PdfConvertOptions {
31                pages_per_file: 4,
32                split_input: true,
33                split_by_bookmark: false,
34                chunking: false,
35                batch_size: 2,
36            },
37            &pdf_info,
38        );
39
40        let page_ranges: Vec<(u32, u32)> = plan
41            .into_iter()
42            .map(|chunk| (chunk.start_page, chunk.end_page))
43            .collect();
44        assert_eq!(page_ranges, vec![(1, 4), (5, 8), (9, 10)]);
45    }
46}