Packages
kreuzberg
4.5.3
4.10.2
4.10.1
4.10.0
4.9.9
4.9.7
4.9.5
4.9.4
4.9.3
4.9.2
4.9.1
4.8.6
4.8.5
4.8.4
4.8.3
4.8.2
4.8.1
4.8.0
4.7.4
4.7.3
4.7.2
4.7.1
4.7.0
4.6.3
4.6.2
4.6.1
4.6.0
4.5.4
4.5.3
4.5.2
4.5.1
4.4.6
4.4.5
4.4.4
4.4.3
4.4.2
4.4.1
4.4.0
4.3.8
4.3.7
4.3.6
4.3.5
4.3.4
4.3.3
4.3.2
4.3.0
4.2.15
4.2.14
4.2.13
4.2.12
4.2.11
4.2.10
4.2.9
4.2.8
4.2.7
4.2.6
4.2.5
4.2.4
4.2.3
4.2.2
4.2.1
4.2.0
4.1.2
4.1.1
4.1.0
4.0.8
4.0.7
4.0.6
4.0.4
4.0.3
4.0.2
4.0.1
4.0.0
4.0.0-rc.27
4.0.0-rc.26
High-performance document intelligence library with OCR support
Current section
Files
Jump to
Current section
Files
native/kreuzberg_rustler/src/batch.rs
//! Batch extraction NIFs
//!
//! This module provides Native Implemented Functions (NIFs) for batch document extraction,
//! processing multiple files or binary inputs efficiently.
use crate::atoms;
use crate::config::{parse_extraction_config, parse_file_extraction_config};
use crate::conversion::convert_extraction_result_to_term;
use rustler::{Binary, Encoder, Env, NifResult, Term};
use std::path::PathBuf;
// Constants for validation
const MAX_BINARY_SIZE: usize = 500 * 1024 * 1024; // 500MB
/// Batch extract text and data from multiple files with default configuration
///
/// # Arguments
/// * `paths` - Vec of file paths as strings
/// * `mime_type` - Optional string representing the MIME type for all files; if None, MIME type is detected per file
///
/// # Returns
/// * `{:ok, [result_map]}` - List of extraction result maps
/// * `{:error, reason}` - Error tuple with reason string
#[rustler::nif(schedule = "DirtyCpu")]
pub fn batch_extract_files<'a>(env: Env<'a>, paths: Vec<String>, mime_type: Option<String>) -> NifResult<Term<'a>> {
if paths.is_empty() {
return Ok((atoms::error(), "File paths list cannot be empty").encode(env));
}
let config = kreuzberg::core::config::ExtractionConfig::default();
let mime_ref = mime_type.as_deref();
let mut results = Vec::new();
// Process each file
for path in paths {
match kreuzberg::extract_file_sync(&path, mime_ref, &config) {
Ok(result) => match convert_extraction_result_to_term(env, &result) {
Ok(term) => results.push(term),
Err(e) => {
return Ok((atoms::error(), format!("Failed to encode result for '{}': {}", path, e)).encode(env));
}
},
Err(e) => return Ok((atoms::error(), format!("Extraction failed for '{}': {}", path, e)).encode(env)),
}
}
Ok((atoms::ok(), results).encode(env))
}
/// Batch extract text and data from multiple files with custom configuration
///
/// Supports optional per-file config overrides via `file_configs` parameter.
///
/// # Arguments
/// * `paths` - Vec of file paths as strings
/// * `mime_type` - Optional string representing the MIME type for all files; if None, MIME type is detected per file
/// * `options_term` - Term containing extraction options (as map or keyword list)
/// * `file_configs` - Optional vec of per-file config terms (nil or map), must match paths length
///
/// # Returns
/// * `{:ok, [result_map]}` - List of extraction result maps
/// * `{:error, reason}` - Error tuple with reason string
#[rustler::nif(schedule = "DirtyCpu")]
pub fn batch_extract_files_with_options<'a>(
env: Env<'a>,
paths: Vec<String>,
mime_type: Option<String>,
options_term: Term<'a>,
file_configs: Option<Vec<Term<'a>>>,
) -> NifResult<Term<'a>> {
if paths.is_empty() {
return Ok((atoms::error(), "File paths list cannot be empty").encode(env));
}
// Parse options from Elixir term to ExtractionConfig
let config = match parse_extraction_config(env, options_term) {
Ok(cfg) => cfg,
Err(e) => return Ok((atoms::error(), format!("Invalid options: {}", e)).encode(env)),
};
match file_configs {
Some(fc_list) => {
// Per-file config path: use batch_extract_file_sync with items
if paths.len() != fc_list.len() {
return Ok((
atoms::error(),
format!("Mismatch: {} paths but {} file configs", paths.len(), fc_list.len()),
)
.encode(env));
}
let mut items: Vec<(PathBuf, Option<kreuzberg::FileExtractionConfig>)> = Vec::with_capacity(paths.len());
for (idx, (path, fc_term)) in paths.into_iter().zip(fc_list.into_iter()).enumerate() {
let fc = match parse_file_extraction_config(env, fc_term) {
Ok(fc) => fc,
Err(e) => {
return Ok((atoms::error(), format!("Invalid file config at index {}: {}", idx, e)).encode(env));
}
};
items.push((PathBuf::from(path), fc));
}
let results = match kreuzberg::batch_extract_file_sync(items, &config) {
Ok(r) => r,
Err(e) => return Ok((atoms::error(), format!("Batch extraction failed: {}", e)).encode(env)),
};
let mut result_terms = Vec::with_capacity(results.len());
for (idx, result) in results.into_iter().enumerate() {
match convert_extraction_result_to_term(env, &result) {
Ok(term) => result_terms.push(term),
Err(e) => {
return Ok((
atoms::error(),
format!("Failed to encode result at index {}: {}", idx, e),
)
.encode(env));
}
}
}
Ok((atoms::ok(), result_terms).encode(env))
}
None => {
// No per-file configs: process each file individually (backwards compatible)
let mime_ref = mime_type.as_deref();
let mut results = Vec::new();
for path in paths {
match kreuzberg::extract_file_sync(&path, mime_ref, &config) {
Ok(result) => match convert_extraction_result_to_term(env, &result) {
Ok(term) => results.push(term),
Err(e) => {
return Ok(
(atoms::error(), format!("Failed to encode result for '{}': {}", path, e)).encode(env),
);
}
},
Err(e) => {
return Ok((atoms::error(), format!("Extraction failed for '{}': {}", path, e)).encode(env));
}
}
}
Ok((atoms::ok(), results).encode(env))
}
}
}
/// Batch extract text and data from multiple binary inputs with default configuration
///
/// # Arguments
/// * `data_list` - Vec of binary data inputs
/// * `mime_types` - Vec of MIME type strings (one per input)
///
/// # Returns
/// * `{:ok, [result_map]}` - List of extraction result maps
/// * `{:error, reason}` - Error tuple with reason string
#[rustler::nif(schedule = "DirtyCpu")]
pub fn batch_extract_bytes<'a>(
env: Env<'a>,
data_list: Vec<Binary<'a>>,
mime_types: Vec<String>,
) -> NifResult<Term<'a>> {
if data_list.is_empty() {
return Ok((atoms::error(), "Data list cannot be empty").encode(env));
}
if data_list.len() != mime_types.len() {
return Ok((
atoms::error(),
format!(
"Mismatch: {} data inputs but {} MIME types",
data_list.len(),
mime_types.len()
),
)
.encode(env));
}
let config = kreuzberg::core::config::ExtractionConfig::default();
let mut results = Vec::new();
// Process each binary input with its corresponding MIME type
for (idx, (data, mime_type)) in data_list.iter().zip(mime_types.iter()).enumerate() {
if data.is_empty() {
return Ok((atoms::error(), format!("Binary input at index {} cannot be empty", idx)).encode(env));
}
if data.len() > MAX_BINARY_SIZE {
return Ok((
atoms::error(),
format!("Binary input at index {} exceeds maximum size of 500MB", idx),
)
.encode(env));
}
match kreuzberg::extract_bytes_sync(data.as_slice(), mime_type, &config) {
Ok(result) => match convert_extraction_result_to_term(env, &result) {
Ok(term) => results.push(term),
Err(e) => {
return Ok((
atoms::error(),
format!("Failed to encode result at index {}: {}", idx, e),
)
.encode(env));
}
},
Err(e) => return Ok((atoms::error(), format!("Extraction failed at index {}: {}", idx, e)).encode(env)),
}
}
Ok((atoms::ok(), results).encode(env))
}
/// Batch extract text and data from multiple binary inputs with custom configuration
///
/// Supports optional per-file config overrides via `file_configs` parameter.
///
/// # Arguments
/// * `data_list` - Vec of binary data inputs
/// * `mime_types` - Vec of MIME type strings (one per input)
/// * `options_term` - Term containing extraction options (as map or keyword list)
/// * `file_configs` - Optional vec of per-file config terms (nil or map), must match data_list length
///
/// # Returns
/// * `{:ok, [result_map]}` - List of extraction result maps
/// * `{:error, reason}` - Error tuple with reason string
#[rustler::nif(schedule = "DirtyCpu")]
pub fn batch_extract_bytes_with_options<'a>(
env: Env<'a>,
data_list: Vec<Binary<'a>>,
mime_types: Vec<String>,
options_term: Term<'a>,
file_configs: Option<Vec<Term<'a>>>,
) -> NifResult<Term<'a>> {
if data_list.is_empty() {
return Ok((atoms::error(), "Data list cannot be empty").encode(env));
}
if data_list.len() != mime_types.len() {
return Ok((
atoms::error(),
format!(
"Mismatch: {} data inputs but {} MIME types",
data_list.len(),
mime_types.len()
),
)
.encode(env));
}
// Parse options from Elixir term to ExtractionConfig
let config = match parse_extraction_config(env, options_term) {
Ok(cfg) => cfg,
Err(e) => return Ok((atoms::error(), format!("Invalid options: {}", e)).encode(env)),
};
match file_configs {
Some(fc_list) => {
// Per-file config path: use batch_extract_bytes_sync with items
if data_list.len() != fc_list.len() {
return Ok((
atoms::error(),
format!(
"Mismatch: {} data inputs but {} file configs",
data_list.len(),
fc_list.len()
),
)
.encode(env));
}
let mut items: Vec<(Vec<u8>, String, Option<kreuzberg::FileExtractionConfig>)> =
Vec::with_capacity(data_list.len());
for (idx, ((data, mime_type), fc_term)) in data_list
.iter()
.zip(mime_types.into_iter())
.zip(fc_list.into_iter())
.enumerate()
{
if data.is_empty() {
return Ok((atoms::error(), format!("Binary input at index {} cannot be empty", idx)).encode(env));
}
if data.len() > MAX_BINARY_SIZE {
return Ok((
atoms::error(),
format!("Binary input at index {} exceeds maximum size of 500MB", idx),
)
.encode(env));
}
let fc = match parse_file_extraction_config(env, fc_term) {
Ok(fc) => fc,
Err(e) => {
return Ok((atoms::error(), format!("Invalid file config at index {}: {}", idx, e)).encode(env));
}
};
items.push((data.as_slice().to_vec(), mime_type, fc));
}
let results = match kreuzberg::batch_extract_bytes_sync(items, &config) {
Ok(r) => r,
Err(e) => return Ok((atoms::error(), format!("Batch extraction failed: {}", e)).encode(env)),
};
let mut result_terms = Vec::with_capacity(results.len());
for (idx, result) in results.into_iter().enumerate() {
match convert_extraction_result_to_term(env, &result) {
Ok(term) => result_terms.push(term),
Err(e) => {
return Ok((
atoms::error(),
format!("Failed to encode result at index {}: {}", idx, e),
)
.encode(env));
}
}
}
Ok((atoms::ok(), result_terms).encode(env))
}
None => {
// No per-file configs: process each individually (backwards compatible)
let mut results = Vec::new();
for (idx, (data, mime_type)) in data_list.iter().zip(mime_types.iter()).enumerate() {
if data.is_empty() {
return Ok((atoms::error(), format!("Binary input at index {} cannot be empty", idx)).encode(env));
}
if data.len() > MAX_BINARY_SIZE {
return Ok((
atoms::error(),
format!("Binary input at index {} exceeds maximum size of 500MB", idx),
)
.encode(env));
}
match kreuzberg::extract_bytes_sync(data.as_slice(), mime_type, &config) {
Ok(result) => match convert_extraction_result_to_term(env, &result) {
Ok(term) => results.push(term),
Err(e) => {
return Ok((
atoms::error(),
format!("Failed to encode result at index {}: {}", idx, e),
)
.encode(env));
}
},
Err(e) => {
return Ok((atoms::error(), format!("Extraction failed at index {}: {}", idx, e)).encode(env));
}
}
}
Ok((atoms::ok(), results).encode(env))
}
}
}