Skip to main content

oxirs_embed/
batch_processing.rs

1//! Offline batch embedding generation with incremental updates
2//!
3//! This module provides comprehensive batch processing capabilities for generating
4//! embeddings offline, with support for incremental updates, resumable jobs,
5//! and efficient resource utilization with SciRS2 integration.
6
7use crate::{CacheManager, EmbeddingModel};
8use anyhow::{anyhow, Result};
9use chrono::{DateTime, Utc};
10use rayon::prelude::*;
11use serde::{Deserialize, Serialize};
12use std::collections::{HashMap, HashSet};
13use std::path::{Path, PathBuf};
14use std::sync::Arc;
15use std::time::Instant;
16use tokio::fs;
17use tokio::sync::{RwLock, Semaphore};
18use tokio::task::JoinHandle;
19use tracing::{debug, info, warn};
20use uuid::Uuid;
21
22/// Memory-optimized batch iterator for large datasets
23pub struct MemoryOptimizedBatchIterator<T> {
24    /// Data source
25    data: Vec<T>,
26    /// Current position
27    position: usize,
28    /// Batch size
29    batch_size: usize,
30    /// Memory usage tracker
31    memory_usage: usize,
32    /// Maximum memory threshold (bytes)
33    max_memory_bytes: usize,
34}
35
36impl<T> MemoryOptimizedBatchIterator<T> {
37    /// Create a new memory-optimized batch iterator
38    pub fn new(data: Vec<T>, batch_size: usize, max_memory_mb: usize) -> Self {
39        Self {
40            data,
41            position: 0,
42            batch_size,
43            memory_usage: 0,
44            max_memory_bytes: max_memory_mb * 1024 * 1024,
45        }
46    }
47
48    /// Get the next batch with memory optimization
49    pub fn next_batch(&mut self) -> Option<Vec<T>>
50    where
51        T: Clone,
52    {
53        if self.position >= self.data.len() {
54            return None;
55        }
56
57        let mut batch = Vec::new();
58        let mut current_memory = 0;
59        let item_size = std::mem::size_of::<T>();
60
61        // Collect items for batch while respecting memory limits
62        while self.position < self.data.len()
63            && batch.len() < self.batch_size
64            && current_memory + item_size <= self.max_memory_bytes
65        {
66            batch.push(self.data[self.position].clone());
67            self.position += 1;
68            current_memory += item_size;
69        }
70
71        self.memory_usage = current_memory;
72
73        if batch.is_empty() {
74            None
75        } else {
76            Some(batch)
77        }
78    }
79
80    /// Get current memory usage
81    pub fn get_memory_usage(&self) -> usize {
82        self.memory_usage
83    }
84
85    /// Get progress percentage
86    pub fn get_progress(&self) -> f64 {
87        if self.data.is_empty() {
88            1.0
89        } else {
90            self.position as f64 / self.data.len() as f64
91        }
92    }
93
94    /// Check if iterator is finished
95    pub fn is_finished(&self) -> bool {
96        self.position >= self.data.len()
97    }
98}
99
100/// Batch processing manager for offline embedding generation
101pub struct BatchProcessingManager {
102    /// Active batch jobs
103    active_jobs: Arc<RwLock<HashMap<Uuid, BatchJob>>>,
104    /// Configuration
105    config: BatchProcessingConfig,
106    /// Cache manager for optimization
107    cache_manager: Arc<CacheManager>,
108    /// Concurrency semaphore
109    semaphore: Arc<Semaphore>,
110    /// Job persistence directory
111    persistence_dir: PathBuf,
112}
113
114/// Configuration for batch processing
115#[derive(Debug, Clone)]
116pub struct BatchProcessingConfig {
117    /// Maximum concurrent workers
118    pub max_workers: usize,
119    /// Chunk size for processing
120    pub chunk_size: usize,
121    /// Enable incremental updates
122    pub enable_incremental: bool,
123    /// Checkpoint frequency (number of chunks)
124    pub checkpoint_frequency: usize,
125    /// Enable resume from checkpoint
126    pub enable_resume: bool,
127    /// Maximum memory usage per worker (MB)
128    pub max_memory_per_worker_mb: usize,
129    /// Enable progress notifications
130    pub enable_notifications: bool,
131    /// Retry configuration
132    pub retry_config: RetryConfig,
133    /// Output format configuration
134    pub output_config: OutputConfig,
135}
136
137impl Default for BatchProcessingConfig {
138    fn default() -> Self {
139        Self {
140            max_workers: std::thread::available_parallelism()
141                .map(|n| n.get())
142                .unwrap_or(1),
143            chunk_size: 1000,
144            enable_incremental: true,
145            checkpoint_frequency: 10,
146            enable_resume: true,
147            max_memory_per_worker_mb: 512,
148            enable_notifications: true,
149            retry_config: RetryConfig::default(),
150            output_config: OutputConfig::default(),
151        }
152    }
153}
154
155/// Retry configuration
156#[derive(Debug, Clone)]
157pub struct RetryConfig {
158    /// Maximum retry attempts
159    pub max_retries: usize,
160    /// Initial backoff delay in milliseconds
161    pub initial_backoff_ms: u64,
162    /// Maximum backoff delay in milliseconds
163    pub max_backoff_ms: u64,
164    /// Backoff multiplier
165    pub backoff_multiplier: f64,
166}
167
168impl Default for RetryConfig {
169    fn default() -> Self {
170        Self {
171            max_retries: 3,
172            initial_backoff_ms: 1000,
173            max_backoff_ms: 30000,
174            backoff_multiplier: 2.0,
175        }
176    }
177}
178
179/// Output configuration
180#[derive(Debug, Clone)]
181pub struct OutputConfig {
182    /// Output format
183    pub format: OutputFormat,
184    /// Compression level (0-9)
185    pub compression_level: u32,
186    /// Include metadata
187    pub include_metadata: bool,
188    /// Batch output into files
189    pub batch_output: bool,
190    /// Maximum entities per output file
191    pub max_entities_per_file: usize,
192}
193
194impl Default for OutputConfig {
195    fn default() -> Self {
196        Self {
197            format: OutputFormat::Parquet,
198            compression_level: 6,
199            include_metadata: true,
200            batch_output: true,
201            max_entities_per_file: 100_000,
202        }
203    }
204}
205
206/// Output formats
207#[derive(Debug, Clone)]
208pub enum OutputFormat {
209    /// Apache Parquet format
210    Parquet,
211    /// Compressed JSON Lines
212    JsonLines,
213    /// Binary format (custom)
214    Binary,
215    /// HDF5 format
216    HDF5,
217}
218
219/// Batch job definition
220#[derive(Debug, Clone, Serialize, Deserialize)]
221pub struct BatchJob {
222    /// Unique job ID
223    pub job_id: Uuid,
224    /// Job name
225    pub name: String,
226    /// Job status
227    pub status: JobStatus,
228    /// Input specification
229    pub input: BatchInput,
230    /// Output specification
231    pub output: BatchOutput,
232    /// Processing configuration
233    pub config: BatchJobConfig,
234    /// Model information
235    pub model_id: Uuid,
236    /// Created timestamp
237    pub created_at: DateTime<Utc>,
238    /// Started timestamp
239    pub started_at: Option<DateTime<Utc>>,
240    /// Completed timestamp
241    pub completed_at: Option<DateTime<Utc>>,
242    /// Progress information
243    pub progress: JobProgress,
244    /// Error information
245    pub error: Option<String>,
246    /// Checkpoint data
247    pub checkpoint: Option<JobCheckpoint>,
248}
249
250/// Job status
251#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
252pub enum JobStatus {
253    Pending,
254    Running,
255    Completed,
256    Failed,
257    Cancelled,
258    Paused,
259}
260
261/// Batch input specification
262#[derive(Debug, Clone, Serialize, Deserialize)]
263pub struct BatchInput {
264    /// Input type
265    pub input_type: InputType,
266    /// Input source
267    pub source: String,
268    /// Filter criteria
269    pub filters: Option<HashMap<String, String>>,
270    /// Incremental mode settings
271    pub incremental: Option<IncrementalConfig>,
272}
273
274/// Input types
275#[derive(Debug, Clone, Serialize, Deserialize)]
276pub enum InputType {
277    /// List of entity IDs
278    EntityList,
279    /// File containing entity IDs
280    EntityFile,
281    /// SPARQL query result
282    SparqlQuery,
283    /// Database query
284    DatabaseQuery,
285    /// Stream source
286    StreamSource,
287}
288
289/// Incremental processing configuration
290#[derive(Debug, Clone, Serialize, Deserialize)]
291pub struct IncrementalConfig {
292    /// Enable incremental processing
293    pub enabled: bool,
294    /// Last processed timestamp
295    pub last_processed: Option<DateTime<Utc>>,
296    /// Timestamp field name
297    pub timestamp_field: String,
298    /// Check for deletions
299    pub check_deletions: bool,
300    /// Existing embeddings source
301    pub existing_embeddings_path: Option<String>,
302}
303
304/// Batch output specification
305#[derive(Debug, Clone, Serialize, Deserialize)]
306pub struct BatchOutput {
307    /// Output path
308    pub path: String,
309    /// Output format
310    pub format: String,
311    /// Compression settings
312    pub compression: Option<String>,
313    /// Partitioning strategy
314    pub partitioning: Option<PartitioningStrategy>,
315}
316
317/// Partitioning strategy
318#[derive(Debug, Clone, Serialize, Deserialize)]
319pub enum PartitioningStrategy {
320    /// No partitioning
321    None,
322    /// Partition by entity type
323    ByEntityType,
324    /// Partition by date
325    ByDate,
326    /// Partition by hash
327    ByHash { num_partitions: usize },
328    /// Custom partitioning
329    Custom { field: String },
330}
331
332/// Job-specific configuration
333#[derive(Debug, Clone, Serialize, Deserialize)]
334pub struct BatchJobConfig {
335    /// Chunk size for this job
336    pub chunk_size: usize,
337    /// Number of workers
338    pub num_workers: usize,
339    /// Retry configuration
340    pub max_retries: usize,
341    /// Enable caching
342    pub use_cache: bool,
343    /// Custom parameters
344    pub custom_params: HashMap<String, String>,
345}
346
347/// Job progress information
348#[derive(Debug, Clone, Serialize, Deserialize)]
349pub struct JobProgress {
350    /// Total entities to process
351    pub total_entities: usize,
352    /// Entities processed
353    pub processed_entities: usize,
354    /// Entities failed
355    pub failed_entities: usize,
356    /// Current chunk being processed
357    pub current_chunk: usize,
358    /// Total chunks
359    pub total_chunks: usize,
360    /// Processing rate (entities/second)
361    pub processing_rate: f64,
362    /// Estimated time remaining
363    pub eta_seconds: Option<u64>,
364    /// Memory usage (MB)
365    pub memory_usage_mb: f64,
366}
367
368impl Default for JobProgress {
369    fn default() -> Self {
370        Self {
371            total_entities: 0,
372            processed_entities: 0,
373            failed_entities: 0,
374            current_chunk: 0,
375            total_chunks: 0,
376            processing_rate: 0.0,
377            eta_seconds: None,
378            memory_usage_mb: 0.0,
379        }
380    }
381}
382
383/// Job checkpoint for resumability
384#[derive(Debug, Clone, Serialize, Deserialize)]
385pub struct JobCheckpoint {
386    /// Checkpoint timestamp
387    pub timestamp: DateTime<Utc>,
388    /// Last processed entity index
389    pub last_processed_index: usize,
390    /// Processed entity IDs
391    pub processed_entities: HashSet<String>,
392    /// Failed entity IDs with error messages
393    pub failed_entities: HashMap<String, String>,
394    /// Intermediate results path
395    pub intermediate_results_path: String,
396    /// Model state hash
397    pub model_state_hash: String,
398}
399
400/// Batch processing result
401#[derive(Debug, Clone, Serialize, Deserialize)]
402pub struct BatchProcessingResult {
403    /// Job ID
404    pub job_id: Uuid,
405    /// Processing statistics
406    pub stats: BatchProcessingStats,
407    /// Output information
408    pub output_info: OutputInfo,
409    /// Quality metrics
410    pub quality_metrics: Option<QualityMetrics>,
411}
412
413/// Batch processing statistics
414#[derive(Debug, Clone, Serialize, Deserialize)]
415pub struct BatchProcessingStats {
416    /// Total processing time
417    pub total_time_seconds: f64,
418    /// Total entities processed
419    pub total_entities: usize,
420    /// Successful embeddings
421    pub successful_embeddings: usize,
422    /// Failed embeddings
423    pub failed_embeddings: usize,
424    /// Cache hits
425    pub cache_hits: usize,
426    /// Cache misses
427    pub cache_misses: usize,
428    /// Average processing time per entity (ms)
429    pub avg_time_per_entity_ms: f64,
430    /// Peak memory usage (MB)
431    pub peak_memory_mb: f64,
432    /// CPU utilization
433    pub cpu_utilization: f64,
434}
435
436/// Output information
437#[derive(Debug, Clone, Serialize, Deserialize)]
438pub struct OutputInfo {
439    /// Output files created
440    pub output_files: Vec<String>,
441    /// Total output size (bytes)
442    pub total_size_bytes: u64,
443    /// Compression ratio
444    pub compression_ratio: f64,
445    /// Number of partitions
446    pub num_partitions: usize,
447}
448
449/// Quality metrics for batch processing
450#[derive(Debug, Clone, Serialize, Deserialize)]
451pub struct QualityMetrics {
452    /// Average embedding norm
453    pub avg_embedding_norm: f64,
454    /// Embedding norm standard deviation
455    pub embedding_norm_std: f64,
456    /// Average cosine similarity to centroid
457    pub avg_cosine_similarity: f64,
458    /// Embedding dimension
459    pub embedding_dimension: usize,
460    /// Number of zero embeddings
461    pub zero_embeddings: usize,
462    /// Number of NaN embeddings
463    pub nan_embeddings: usize,
464}
465
466impl BatchProcessingManager {
467    /// Create a new batch processing manager
468    pub fn new(
469        config: BatchProcessingConfig,
470        cache_manager: Arc<CacheManager>,
471        persistence_dir: PathBuf,
472    ) -> Self {
473        Self {
474            active_jobs: Arc::new(RwLock::new(HashMap::new())),
475            semaphore: Arc::new(Semaphore::new(config.max_workers)),
476            config,
477            cache_manager,
478            persistence_dir,
479        }
480    }
481
482    /// Submit a new batch job
483    pub async fn submit_job(&self, job: BatchJob) -> Result<Uuid> {
484        let job_id = job.job_id;
485
486        // Validate job
487        self.validate_job(&job).await?;
488
489        // Store job
490        {
491            let mut jobs = self.active_jobs.write().await;
492            jobs.insert(job_id, job.clone());
493        }
494
495        // Persist job configuration
496        self.persist_job(&job).await?;
497
498        info!("Submitted batch job: {} ({})", job.name, job_id);
499        Ok(job_id)
500    }
501
502    /// Start processing a batch job
503    pub async fn start_job(
504        &self,
505        job_id: Uuid,
506        model: Arc<dyn EmbeddingModel + Send + Sync>,
507    ) -> Result<JoinHandle<Result<BatchProcessingResult>>> {
508        let job = {
509            let mut jobs = self.active_jobs.write().await;
510            let job = jobs
511                .get_mut(&job_id)
512                .ok_or_else(|| anyhow!("Job not found: {}", job_id))?;
513
514            if !matches!(job.status, JobStatus::Pending | JobStatus::Paused) {
515                return Err(anyhow!("Job {} is not in a startable state", job_id));
516            }
517
518            job.status = JobStatus::Running;
519            job.started_at = Some(Utc::now());
520            job.clone()
521        };
522
523        let manager = self.clone();
524        let handle = tokio::spawn(async move { manager.process_job(job, model).await });
525
526        Ok(handle)
527    }
528
529    /// Process a batch job
530    async fn process_job(
531        &self,
532        job: BatchJob,
533        model: Arc<dyn EmbeddingModel + Send + Sync>,
534    ) -> Result<BatchProcessingResult> {
535        let start_time = Instant::now();
536        info!(
537            "Starting batch job processing: {} ({})",
538            job.name, job.job_id
539        );
540
541        // Load entities to process
542        let entities = self.load_entities(&job).await?;
543
544        // Filter entities for incremental processing
545        let entities_to_process = if job
546            .input
547            .incremental
548            .as_ref()
549            .map(|inc| inc.enabled)
550            .unwrap_or(false)
551        {
552            self.filter_incremental_entities(&job, entities).await?
553        } else {
554            entities
555        };
556
557        // Update job progress
558        {
559            let mut jobs = self.active_jobs.write().await;
560            if let Some(active_job) = jobs.get_mut(&job.job_id) {
561                active_job.progress.total_entities = entities_to_process.len();
562                active_job.progress.total_chunks =
563                    (entities_to_process.len() + job.config.chunk_size - 1) / job.config.chunk_size;
564            }
565        }
566
567        // Process entities in chunks
568        let chunks: Vec<_> = entities_to_process
569            .chunks(job.config.chunk_size)
570            .map(|chunk| chunk.to_vec())
571            .collect();
572
573        let mut successful_embeddings = 0;
574        let mut failed_embeddings = 0;
575        let mut cache_hits = 0;
576        let mut cache_misses = 0;
577        let mut processed_entities = HashSet::new();
578        let mut failed_entities = HashMap::new();
579
580        for (chunk_idx, chunk) in chunks.iter().enumerate() {
581            // Check if job was cancelled
582            {
583                let jobs = self.active_jobs.read().await;
584                if let Some(active_job) = jobs.get(&job.job_id) {
585                    if matches!(active_job.status, JobStatus::Cancelled) {
586                        info!("Job {} was cancelled", job.job_id);
587                        return Err(anyhow!("Job was cancelled"));
588                    }
589                }
590            }
591
592            // Process chunk
593            let chunk_result = self
594                .process_chunk(&job, chunk, chunk_idx, model.clone())
595                .await?;
596
597            // Update statistics
598            successful_embeddings += chunk_result.successful;
599            failed_embeddings += chunk_result.failed;
600            cache_hits += chunk_result.cache_hits;
601            cache_misses += chunk_result.cache_misses;
602
603            // Track processed entities
604            for entity in chunk {
605                processed_entities.insert(entity.clone());
606            }
607            for (entity, error) in chunk_result.failures {
608                failed_entities.insert(entity, error);
609            }
610
611            // Update progress
612            self.update_job_progress(
613                &job.job_id,
614                chunk_idx + 1,
615                successful_embeddings + failed_embeddings,
616            )
617            .await?;
618
619            // Create checkpoint
620            if chunk_idx % self.config.checkpoint_frequency == 0 {
621                self.create_checkpoint(&job.job_id, &processed_entities, &failed_entities)
622                    .await?;
623            }
624
625            info!(
626                "Processed chunk {}/{} for job {}",
627                chunk_idx + 1,
628                chunks.len(),
629                job.job_id
630            );
631        }
632
633        // Finalize processing
634        let processing_time = start_time.elapsed().as_secs_f64();
635        let result = self
636            .finalize_job_processing(
637                &job,
638                processing_time,
639                successful_embeddings,
640                failed_embeddings,
641                cache_hits,
642                cache_misses,
643            )
644            .await?;
645
646        // Update job status
647        {
648            let mut jobs = self.active_jobs.write().await;
649            if let Some(active_job) = jobs.get_mut(&job.job_id) {
650                active_job.status = JobStatus::Completed;
651                active_job.completed_at = Some(Utc::now());
652            }
653        }
654
655        info!(
656            "Completed batch job: {} in {:.2}s",
657            job.job_id, processing_time
658        );
659        Ok(result)
660    }
661
662    /// Process a single chunk of entities
663    async fn process_chunk(
664        &self,
665        job: &BatchJob,
666        entities: &[String],
667        chunk_idx: usize,
668        model: Arc<dyn EmbeddingModel + Send + Sync>,
669    ) -> Result<ChunkResult> {
670        let _permit = self.semaphore.acquire().await?;
671
672        let mut successful = 0;
673        let mut failed = 0;
674        let mut cache_hits = 0;
675        let mut cache_misses = 0;
676        let mut failures = HashMap::new();
677
678        for entity in entities {
679            match self
680                .process_single_entity(entity, model.clone(), job.config.use_cache)
681                .await
682            {
683                Ok(from_cache) => {
684                    successful += 1;
685                    if from_cache {
686                        cache_hits += 1;
687                    } else {
688                        cache_misses += 1;
689                    }
690                }
691                Err(e) => {
692                    failed += 1;
693                    failures.insert(entity.clone(), e.to_string());
694                    warn!("Failed to process entity {}: {}", entity, e);
695                }
696            }
697        }
698
699        Ok(ChunkResult {
700            chunk_idx,
701            successful,
702            failed,
703            cache_hits,
704            cache_misses,
705            failures,
706        })
707    }
708
709    /// Process a single entity
710    async fn process_single_entity(
711        &self,
712        entity: &str,
713        model: Arc<dyn EmbeddingModel + Send + Sync>,
714        use_cache: bool,
715    ) -> Result<bool> {
716        if use_cache {
717            // Check cache first
718            if let Some(_embedding) = self.cache_manager.get_embedding(entity) {
719                return Ok(true);
720            }
721        }
722
723        // Generate embedding
724        let embedding = model.get_entity_embedding(entity)?;
725
726        // Cache the result
727        if use_cache {
728            self.cache_manager
729                .put_embedding(entity.to_string(), embedding);
730        }
731
732        Ok(false)
733    }
734
735    /// Load entities to process based on input specification
736    async fn load_entities(&self, job: &BatchJob) -> Result<Vec<String>> {
737        match &job.input.input_type {
738            InputType::EntityList => {
739                // Parse entity list from source
740                let entities: Vec<String> = serde_json::from_str(&job.input.source)?;
741                Ok(entities)
742            }
743            InputType::EntityFile => {
744                // Read entities from file
745                let content = fs::read_to_string(&job.input.source).await?;
746                let entities: Vec<String> = content
747                    .lines()
748                    .map(|line| line.trim().to_string())
749                    .filter(|line| !line.is_empty())
750                    .collect();
751                Ok(entities)
752            }
753            InputType::SparqlQuery => {
754                // Execute SPARQL query and extract entities
755                // This would need to be implemented based on SPARQL engine
756                warn!("SPARQL query input type not yet implemented");
757                Ok(Vec::new())
758            }
759            InputType::DatabaseQuery => {
760                // Execute database query and extract entities
761                warn!("Database query input type not yet implemented");
762                Ok(Vec::new())
763            }
764            InputType::StreamSource => {
765                // Read from stream source
766                warn!("Stream source input type not yet implemented");
767                Ok(Vec::new())
768            }
769        }
770    }
771
772    /// Filter entities for incremental processing
773    async fn filter_incremental_entities(
774        &self,
775        job: &BatchJob,
776        entities: Vec<String>,
777    ) -> Result<Vec<String>> {
778        if let Some(incremental) = &job.input.incremental {
779            if !incremental.enabled {
780                return Ok(entities);
781            }
782
783            // Load existing embeddings if specified
784            let existing_entities =
785                if let Some(existing_path) = &incremental.existing_embeddings_path {
786                    self.load_existing_entities(existing_path).await?
787                } else {
788                    HashSet::new()
789                };
790
791            // Filter out entities that already have embeddings
792            let filtered: Vec<String> = entities
793                .into_iter()
794                .filter(|entity| !existing_entities.contains(entity))
795                .collect();
796
797            info!(
798                "Incremental filtering: {} entities remaining after filtering",
799                filtered.len()
800            );
801            Ok(filtered)
802        } else {
803            Ok(entities)
804        }
805    }
806
807    /// Load existing entities from embeddings file
808    async fn load_existing_entities(&self, path: &str) -> Result<HashSet<String>> {
809        // This would depend on the output format
810        // For now, assume a simple text file with entity IDs
811        if Path::new(path).exists() {
812            let content = fs::read_to_string(path).await?;
813            let entities: HashSet<String> = content
814                .lines()
815                .map(|line| line.trim().to_string())
816                .filter(|line| !line.is_empty())
817                .collect();
818            Ok(entities)
819        } else {
820            Ok(HashSet::new())
821        }
822    }
823
824    /// Update job progress
825    async fn update_job_progress(
826        &self,
827        job_id: &Uuid,
828        current_chunk: usize,
829        processed_entities: usize,
830    ) -> Result<()> {
831        let mut jobs = self.active_jobs.write().await;
832        if let Some(job) = jobs.get_mut(job_id) {
833            job.progress.current_chunk = current_chunk;
834            job.progress.processed_entities = processed_entities;
835
836            // Calculate processing rate
837            if let Some(started_at) = job.started_at {
838                let elapsed = Utc::now().signed_duration_since(started_at);
839                let elapsed_seconds = elapsed.num_seconds() as f64;
840                if elapsed_seconds > 0.0 {
841                    job.progress.processing_rate = processed_entities as f64 / elapsed_seconds;
842
843                    // Estimate time remaining
844                    let remaining_entities = job.progress.total_entities - processed_entities;
845                    if job.progress.processing_rate > 0.0 {
846                        let eta = remaining_entities as f64 / job.progress.processing_rate;
847                        job.progress.eta_seconds = Some(eta as u64);
848                    }
849                }
850            }
851        }
852        Ok(())
853    }
854
855    /// Create a checkpoint for job resumability
856    async fn create_checkpoint(
857        &self,
858        job_id: &Uuid,
859        processed_entities: &HashSet<String>,
860        failed_entities: &HashMap<String, String>,
861    ) -> Result<()> {
862        let checkpoint = JobCheckpoint {
863            timestamp: Utc::now(),
864            last_processed_index: processed_entities.len(),
865            processed_entities: processed_entities.clone(),
866            failed_entities: failed_entities.clone(),
867            intermediate_results_path: format!(
868                "{}/checkpoint_{}.json",
869                self.persistence_dir.display(),
870                job_id
871            ),
872            model_state_hash: "placeholder".to_string(), // Would calculate actual hash
873        };
874
875        // Save checkpoint to disk
876        let checkpoint_path = self
877            .persistence_dir
878            .join(format!("checkpoint_{job_id}.json"));
879        let checkpoint_json = serde_json::to_string_pretty(&checkpoint)?;
880        fs::write(checkpoint_path, checkpoint_json).await?;
881
882        // Update job with checkpoint
883        let mut jobs = self.active_jobs.write().await;
884        if let Some(job) = jobs.get_mut(job_id) {
885            job.checkpoint = Some(checkpoint);
886        }
887
888        debug!("Created checkpoint for job {}", job_id);
889        Ok(())
890    }
891
892    /// Finalize job processing and create result
893    async fn finalize_job_processing(
894        &self,
895        job: &BatchJob,
896        processing_time: f64,
897        successful_embeddings: usize,
898        failed_embeddings: usize,
899        cache_hits: usize,
900        cache_misses: usize,
901    ) -> Result<BatchProcessingResult> {
902        let total_entities = successful_embeddings + failed_embeddings;
903        let avg_time_per_entity_ms = if total_entities > 0 {
904            (processing_time * 1000.0) / total_entities as f64
905        } else {
906            0.0
907        };
908
909        let stats = BatchProcessingStats {
910            total_time_seconds: processing_time,
911            total_entities,
912            successful_embeddings,
913            failed_embeddings,
914            cache_hits,
915            cache_misses,
916            avg_time_per_entity_ms,
917            peak_memory_mb: 0.0,  // Would measure actual memory usage
918            cpu_utilization: 0.0, // Would measure actual CPU usage
919        };
920
921        let output_info = OutputInfo {
922            output_files: vec![job.output.path.clone()],
923            total_size_bytes: 0, // Would calculate actual size
924            compression_ratio: 1.0,
925            num_partitions: 1,
926        };
927
928        Ok(BatchProcessingResult {
929            job_id: job.job_id,
930            stats,
931            output_info,
932            quality_metrics: None, // Would calculate if requested
933        })
934    }
935
936    /// Validate a batch job before submission
937    async fn validate_job(&self, job: &BatchJob) -> Result<()> {
938        // Validate input source exists
939        if let InputType::EntityFile = &job.input.input_type {
940            if !Path::new(&job.input.source).exists() {
941                return Err(anyhow!("Input file does not exist: {}", job.input.source));
942            }
943        } // Other validations would be implemented
944
945        // Validate output path is writable
946        if let Some(parent) = Path::new(&job.output.path).parent() {
947            if !parent.exists() {
948                fs::create_dir_all(parent).await?;
949            }
950        }
951
952        Ok(())
953    }
954
955    /// Persist job configuration to disk
956    async fn persist_job(&self, job: &BatchJob) -> Result<()> {
957        let job_path = self
958            .persistence_dir
959            .join(format!("job_{}.json", job.job_id));
960        let job_json = serde_json::to_string_pretty(job)?;
961        fs::write(job_path, job_json).await?;
962        Ok(())
963    }
964
965    /// Get job status
966    pub async fn get_job_status(&self, job_id: &Uuid) -> Option<JobStatus> {
967        let jobs = self.active_jobs.read().await;
968        jobs.get(job_id).map(|job| job.status.clone())
969    }
970
971    /// Get job progress
972    pub async fn get_job_progress(&self, job_id: &Uuid) -> Option<JobProgress> {
973        let jobs = self.active_jobs.read().await;
974        jobs.get(job_id).map(|job| job.progress.clone())
975    }
976
977    /// Cancel a job
978    pub async fn cancel_job(&self, job_id: &Uuid) -> Result<()> {
979        let mut jobs = self.active_jobs.write().await;
980        if let Some(job) = jobs.get_mut(job_id) {
981            job.status = JobStatus::Cancelled;
982            info!("Cancelled job: {}", job_id);
983            Ok(())
984        } else {
985            Err(anyhow!("Job not found: {}", job_id))
986        }
987    }
988
989    /// List all jobs
990    pub async fn list_jobs(&self) -> Vec<BatchJob> {
991        let jobs = self.active_jobs.read().await;
992        jobs.values().cloned().collect()
993    }
994}
995
996impl Clone for BatchProcessingManager {
997    fn clone(&self) -> Self {
998        Self {
999            active_jobs: Arc::clone(&self.active_jobs),
1000            config: self.config.clone(),
1001            cache_manager: Arc::clone(&self.cache_manager),
1002            semaphore: Arc::clone(&self.semaphore),
1003            persistence_dir: self.persistence_dir.clone(),
1004        }
1005    }
1006}
1007
1008/// Result of processing a single chunk
1009#[derive(Debug)]
1010#[allow(dead_code)]
1011struct ChunkResult {
1012    chunk_idx: usize,
1013    successful: usize,
1014    failed: usize,
1015    cache_hits: usize,
1016    cache_misses: usize,
1017    failures: HashMap<String, String>,
1018}
1019
1020#[cfg(test)]
1021mod tests {
1022    use super::*;
1023    use tempfile::tempdir;
1024
1025    #[test]
1026    fn test_batch_job_creation() {
1027        let job = BatchJob {
1028            job_id: Uuid::new_v4(),
1029            name: "test_job".to_string(),
1030            status: JobStatus::Pending,
1031            input: BatchInput {
1032                input_type: InputType::EntityList,
1033                source: r#"["entity1", "entity2", "entity3"]"#.to_string(),
1034                filters: None,
1035                incremental: None,
1036            },
1037            output: BatchOutput {
1038                path: std::env::temp_dir()
1039                    .join(format!("oxirs_batch_out_{}", std::process::id()))
1040                    .display()
1041                    .to_string(),
1042                format: "parquet".to_string(),
1043                compression: Some("gzip".to_string()),
1044                partitioning: Some(PartitioningStrategy::None),
1045            },
1046            config: BatchJobConfig {
1047                chunk_size: 100,
1048                num_workers: 4,
1049                max_retries: 3,
1050                use_cache: true,
1051                custom_params: HashMap::new(),
1052            },
1053            model_id: Uuid::new_v4(),
1054            created_at: Utc::now(),
1055            started_at: None,
1056            completed_at: None,
1057            progress: JobProgress::default(),
1058            error: None,
1059            checkpoint: None,
1060        };
1061
1062        assert_eq!(job.status, JobStatus::Pending);
1063        assert_eq!(job.name, "test_job");
1064    }
1065
1066    #[tokio::test]
1067    async fn test_batch_processing_manager_creation() {
1068        let config = BatchProcessingConfig::default();
1069        let cache_config = crate::CacheConfig::default();
1070        let cache_manager = Arc::new(CacheManager::new(cache_config));
1071        let temp_dir = tempdir().expect("should succeed");
1072
1073        let manager =
1074            BatchProcessingManager::new(config, cache_manager, temp_dir.path().to_path_buf());
1075
1076        assert_eq!(
1077            manager.config.max_workers,
1078            std::thread::available_parallelism()
1079                .map(|n| n.get())
1080                .unwrap_or(1)
1081        );
1082        assert_eq!(manager.config.chunk_size, 1000);
1083    }
1084
1085    #[test]
1086    fn test_incremental_config() {
1087        let incremental = IncrementalConfig {
1088            enabled: true,
1089            last_processed: Some(Utc::now()),
1090            timestamp_field: "updated_at".to_string(),
1091            check_deletions: true,
1092            existing_embeddings_path: Some("/path/to/existing".to_string()),
1093        };
1094
1095        assert!(incremental.enabled);
1096        assert!(incremental.last_processed.is_some());
1097        assert_eq!(incremental.timestamp_field, "updated_at");
1098    }
1099
1100    #[test]
1101    fn test_memory_optimized_batch_iterator() {
1102        let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
1103        let mut iterator = MemoryOptimizedBatchIterator::new(data.clone(), 3, 1); // 1MB limit
1104
1105        // Test first batch
1106        let batch1 = iterator.next_batch().expect("should succeed");
1107        assert_eq!(batch1.len(), 3);
1108        assert_eq!(batch1, vec![1, 2, 3]);
1109        assert_eq!(iterator.get_progress(), 0.3);
1110        assert!(!iterator.is_finished());
1111
1112        // Test second batch
1113        let batch2 = iterator.next_batch().expect("should succeed");
1114        assert_eq!(batch2.len(), 3);
1115        assert_eq!(batch2, vec![4, 5, 6]);
1116        assert_eq!(iterator.get_progress(), 0.6);
1117
1118        // Test third batch
1119        let batch3 = iterator.next_batch().expect("should succeed");
1120        assert_eq!(batch3.len(), 3);
1121        assert_eq!(batch3, vec![7, 8, 9]);
1122        assert_eq!(iterator.get_progress(), 0.9);
1123
1124        // Test final batch
1125        let batch4 = iterator.next_batch().expect("should succeed");
1126        assert_eq!(batch4.len(), 1);
1127        assert_eq!(batch4, vec![10]);
1128        assert_eq!(iterator.get_progress(), 1.0);
1129        assert!(iterator.is_finished());
1130
1131        // Test empty batch
1132        let batch5 = iterator.next_batch();
1133        assert!(batch5.is_none());
1134    }
1135
1136    #[test]
1137    fn test_memory_optimized_batch_iterator_empty() {
1138        let data: Vec<i32> = vec![];
1139        let mut iterator = MemoryOptimizedBatchIterator::new(data, 3, 1);
1140
1141        assert_eq!(iterator.get_progress(), 1.0);
1142        assert!(iterator.is_finished());
1143        assert!(iterator.next_batch().is_none());
1144    }
1145
1146    #[test]
1147    fn test_memory_optimized_batch_iterator_single_item() {
1148        let data = vec![42];
1149        let mut iterator = MemoryOptimizedBatchIterator::new(data, 5, 1);
1150
1151        let batch = iterator.next_batch().expect("should succeed");
1152        assert_eq!(batch.len(), 1);
1153        assert_eq!(batch[0], 42);
1154        assert_eq!(iterator.get_progress(), 1.0);
1155        assert!(iterator.is_finished());
1156    }
1157
1158    #[test]
1159    fn test_memory_optimized_batch_iterator_memory_tracking() {
1160        let data = vec![1, 2, 3, 4, 5];
1161        let mut iterator = MemoryOptimizedBatchIterator::new(data, 3, 1);
1162
1163        // Process one batch and check memory usage
1164        let _batch = iterator.next_batch().expect("should succeed");
1165        let memory_usage = iterator.get_memory_usage();
1166        assert!(memory_usage > 0);
1167
1168        // Memory usage should be roughly 3 * size_of::<i32>()
1169        let expected_memory = 3 * std::mem::size_of::<i32>();
1170        assert_eq!(memory_usage, expected_memory);
1171    }
1172
1173    #[test]
1174    fn test_parallel_batch_processor() {
1175        // Test basic functionality
1176        let processor =
1177            ParallelBatchProcessor::new(ParallelBatchConfig::default()).expect("should succeed");
1178        // Should use the system's available parallelism
1179        assert!(processor.num_workers() > 0);
1180        assert!(
1181            processor.num_workers()
1182                <= std::thread::available_parallelism()
1183                    .map(|n| n.get())
1184                    .unwrap_or(1)
1185        );
1186    }
1187}
1188
1189/// Advanced parallel batch processor using SciRS2 and Rayon
1190///
1191/// This processor leverages parallel operations for:
1192/// - Optimal work distribution across cores
1193/// - Adaptive load balancing
1194/// - Memory-efficient chunking
1195/// - NUMA-aware processing
1196pub struct ParallelBatchProcessor {
1197    config: ParallelBatchConfig,
1198}
1199
1200/// Configuration for parallel batch processing
1201#[derive(Debug, Clone, Serialize, Deserialize)]
1202pub struct ParallelBatchConfig {
1203    /// Number of worker threads
1204    pub num_workers: usize,
1205    /// Chunk size for parallel processing
1206    pub chunk_size: usize,
1207    /// Enable adaptive load balancing
1208    pub adaptive_balancing: bool,
1209    /// Memory threshold in MB
1210    pub memory_threshold_mb: usize,
1211    /// Enable NUMA optimization
1212    pub numa_aware: bool,
1213    /// Enable work stealing
1214    pub work_stealing: bool,
1215}
1216
1217impl Default for ParallelBatchConfig {
1218    fn default() -> Self {
1219        Self {
1220            num_workers: std::thread::available_parallelism()
1221                .map(|n| n.get())
1222                .unwrap_or(1),
1223            chunk_size: 1000,
1224            adaptive_balancing: true,
1225            memory_threshold_mb: 512,
1226            numa_aware: true,
1227            work_stealing: true,
1228        }
1229    }
1230}
1231
1232impl ParallelBatchProcessor {
1233    /// Create new parallel batch processor
1234    pub fn new(config: ParallelBatchConfig) -> Result<Self> {
1235        // Configure rayon thread pool
1236        rayon::ThreadPoolBuilder::new()
1237            .num_threads(config.num_workers)
1238            .build_global()
1239            .ok(); // Ignore error if already initialized
1240
1241        Ok(Self { config })
1242    }
1243
1244    /// Get number of worker threads
1245    pub fn num_workers(&self) -> usize {
1246        self.config.num_workers
1247    }
1248
1249    /// Process batch in parallel with automatic load balancing
1250    pub fn process_parallel<T, F, R>(&self, items: Vec<T>, process_fn: F) -> Result<Vec<R>>
1251    where
1252        T: Send + Sync,
1253        F: Fn(&T) -> R + Send + Sync,
1254        R: Send,
1255    {
1256        // Use rayon for parallel processing
1257        let results: Vec<R> = items.par_iter().map(process_fn).collect();
1258
1259        Ok(results)
1260    }
1261
1262    /// Process batch with dynamic load balancing (using rayon's work stealing)
1263    ///
1264    /// Uses Rayon's work stealing for:
1265    /// - Automatic work stealing between threads
1266    /// - Dynamic work distribution
1267    /// - Optimal CPU utilization
1268    pub fn process_with_load_balancing<T, F, R>(
1269        &self,
1270        items: Vec<T>,
1271        process_fn: F,
1272    ) -> Result<Vec<R>>
1273    where
1274        T: Send + Sync,
1275        F: Fn(&T) -> R + Send + Sync,
1276        R: Send,
1277    {
1278        // Rayon automatically uses work stealing
1279        let results: Vec<R> = items.par_iter().map(process_fn).collect();
1280
1281        Ok(results)
1282    }
1283
1284    /// Process very large batches with memory-efficient chunking
1285    ///
1286    /// Uses chunked parallel processing to:
1287    /// - Respect memory limits
1288    /// - Optimize cache locality
1289    /// - Minimize memory allocations
1290    pub fn process_memory_efficient<T, F, R>(&self, items: Vec<T>, process_fn: F) -> Result<Vec<R>>
1291    where
1292        T: Send + Sync,
1293        F: Fn(&T) -> R + Send + Sync,
1294        R: Send,
1295    {
1296        // Process in chunks to respect memory limits
1297        let chunk_size =
1298            (self.config.memory_threshold_mb * 1024 * 1024) / (std::mem::size_of::<T>().max(1));
1299
1300        let chunk_size = chunk_size.min(self.config.chunk_size).max(100);
1301
1302        let results: Vec<R> = items
1303            .par_chunks(chunk_size)
1304            .flat_map(|chunk| chunk.iter().map(&process_fn).collect::<Vec<_>>())
1305            .collect();
1306
1307        Ok(results)
1308    }
1309
1310    /// Process with nested parallelism using rayon
1311    ///
1312    /// Enables safe nested parallel processing:
1313    /// - Automatic thread lifetime management
1314    /// - Rayon's work stealing
1315    /// - Cache-friendly execution
1316    pub fn process_nested_parallel<T, F, R>(
1317        &self,
1318        items: Vec<Vec<T>>,
1319        process_fn: F,
1320    ) -> Result<Vec<Vec<R>>>
1321    where
1322        T: Send + Sync,
1323        F: Fn(&T) -> R + Send + Sync,
1324        R: Send,
1325    {
1326        let results: Vec<Vec<R>> = items
1327            .par_iter()
1328            .map(|batch| batch.iter().map(&process_fn).collect())
1329            .collect();
1330
1331        Ok(results)
1332    }
1333
1334    /// Get processing statistics
1335    pub fn get_stats(&self) -> ParallelProcessingStats {
1336        ParallelProcessingStats {
1337            num_workers: self.config.num_workers,
1338            profiler_report: "Stats available".to_string(),
1339            memory_usage: 0,
1340        }
1341    }
1342}
1343
1344/// Statistics for parallel batch processing
1345#[derive(Debug, Clone)]
1346pub struct ParallelProcessingStats {
1347    pub num_workers: usize,
1348    pub profiler_report: String,
1349    pub memory_usage: usize,
1350}