1use 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
22pub struct MemoryOptimizedBatchIterator<T> {
24 data: Vec<T>,
26 position: usize,
28 batch_size: usize,
30 memory_usage: usize,
32 max_memory_bytes: usize,
34}
35
36impl<T> MemoryOptimizedBatchIterator<T> {
37 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 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 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 pub fn get_memory_usage(&self) -> usize {
82 self.memory_usage
83 }
84
85 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 pub fn is_finished(&self) -> bool {
96 self.position >= self.data.len()
97 }
98}
99
100pub struct BatchProcessingManager {
102 active_jobs: Arc<RwLock<HashMap<Uuid, BatchJob>>>,
104 config: BatchProcessingConfig,
106 cache_manager: Arc<CacheManager>,
108 semaphore: Arc<Semaphore>,
110 persistence_dir: PathBuf,
112}
113
114#[derive(Debug, Clone)]
116pub struct BatchProcessingConfig {
117 pub max_workers: usize,
119 pub chunk_size: usize,
121 pub enable_incremental: bool,
123 pub checkpoint_frequency: usize,
125 pub enable_resume: bool,
127 pub max_memory_per_worker_mb: usize,
129 pub enable_notifications: bool,
131 pub retry_config: RetryConfig,
133 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#[derive(Debug, Clone)]
157pub struct RetryConfig {
158 pub max_retries: usize,
160 pub initial_backoff_ms: u64,
162 pub max_backoff_ms: u64,
164 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#[derive(Debug, Clone)]
181pub struct OutputConfig {
182 pub format: OutputFormat,
184 pub compression_level: u32,
186 pub include_metadata: bool,
188 pub batch_output: bool,
190 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#[derive(Debug, Clone)]
208pub enum OutputFormat {
209 Parquet,
211 JsonLines,
213 Binary,
215 HDF5,
217}
218
219#[derive(Debug, Clone, Serialize, Deserialize)]
221pub struct BatchJob {
222 pub job_id: Uuid,
224 pub name: String,
226 pub status: JobStatus,
228 pub input: BatchInput,
230 pub output: BatchOutput,
232 pub config: BatchJobConfig,
234 pub model_id: Uuid,
236 pub created_at: DateTime<Utc>,
238 pub started_at: Option<DateTime<Utc>>,
240 pub completed_at: Option<DateTime<Utc>>,
242 pub progress: JobProgress,
244 pub error: Option<String>,
246 pub checkpoint: Option<JobCheckpoint>,
248}
249
250#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
252pub enum JobStatus {
253 Pending,
254 Running,
255 Completed,
256 Failed,
257 Cancelled,
258 Paused,
259}
260
261#[derive(Debug, Clone, Serialize, Deserialize)]
263pub struct BatchInput {
264 pub input_type: InputType,
266 pub source: String,
268 pub filters: Option<HashMap<String, String>>,
270 pub incremental: Option<IncrementalConfig>,
272}
273
274#[derive(Debug, Clone, Serialize, Deserialize)]
276pub enum InputType {
277 EntityList,
279 EntityFile,
281 SparqlQuery,
283 DatabaseQuery,
285 StreamSource,
287}
288
289#[derive(Debug, Clone, Serialize, Deserialize)]
291pub struct IncrementalConfig {
292 pub enabled: bool,
294 pub last_processed: Option<DateTime<Utc>>,
296 pub timestamp_field: String,
298 pub check_deletions: bool,
300 pub existing_embeddings_path: Option<String>,
302}
303
304#[derive(Debug, Clone, Serialize, Deserialize)]
306pub struct BatchOutput {
307 pub path: String,
309 pub format: String,
311 pub compression: Option<String>,
313 pub partitioning: Option<PartitioningStrategy>,
315}
316
317#[derive(Debug, Clone, Serialize, Deserialize)]
319pub enum PartitioningStrategy {
320 None,
322 ByEntityType,
324 ByDate,
326 ByHash { num_partitions: usize },
328 Custom { field: String },
330}
331
332#[derive(Debug, Clone, Serialize, Deserialize)]
334pub struct BatchJobConfig {
335 pub chunk_size: usize,
337 pub num_workers: usize,
339 pub max_retries: usize,
341 pub use_cache: bool,
343 pub custom_params: HashMap<String, String>,
345}
346
347#[derive(Debug, Clone, Serialize, Deserialize)]
349pub struct JobProgress {
350 pub total_entities: usize,
352 pub processed_entities: usize,
354 pub failed_entities: usize,
356 pub current_chunk: usize,
358 pub total_chunks: usize,
360 pub processing_rate: f64,
362 pub eta_seconds: Option<u64>,
364 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#[derive(Debug, Clone, Serialize, Deserialize)]
385pub struct JobCheckpoint {
386 pub timestamp: DateTime<Utc>,
388 pub last_processed_index: usize,
390 pub processed_entities: HashSet<String>,
392 pub failed_entities: HashMap<String, String>,
394 pub intermediate_results_path: String,
396 pub model_state_hash: String,
398}
399
400#[derive(Debug, Clone, Serialize, Deserialize)]
402pub struct BatchProcessingResult {
403 pub job_id: Uuid,
405 pub stats: BatchProcessingStats,
407 pub output_info: OutputInfo,
409 pub quality_metrics: Option<QualityMetrics>,
411}
412
413#[derive(Debug, Clone, Serialize, Deserialize)]
415pub struct BatchProcessingStats {
416 pub total_time_seconds: f64,
418 pub total_entities: usize,
420 pub successful_embeddings: usize,
422 pub failed_embeddings: usize,
424 pub cache_hits: usize,
426 pub cache_misses: usize,
428 pub avg_time_per_entity_ms: f64,
430 pub peak_memory_mb: f64,
432 pub cpu_utilization: f64,
434}
435
436#[derive(Debug, Clone, Serialize, Deserialize)]
438pub struct OutputInfo {
439 pub output_files: Vec<String>,
441 pub total_size_bytes: u64,
443 pub compression_ratio: f64,
445 pub num_partitions: usize,
447}
448
449#[derive(Debug, Clone, Serialize, Deserialize)]
451pub struct QualityMetrics {
452 pub avg_embedding_norm: f64,
454 pub embedding_norm_std: f64,
456 pub avg_cosine_similarity: f64,
458 pub embedding_dimension: usize,
460 pub zero_embeddings: usize,
462 pub nan_embeddings: usize,
464}
465
466impl BatchProcessingManager {
467 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 pub async fn submit_job(&self, job: BatchJob) -> Result<Uuid> {
484 let job_id = job.job_id;
485
486 self.validate_job(&job).await?;
488
489 {
491 let mut jobs = self.active_jobs.write().await;
492 jobs.insert(job_id, job.clone());
493 }
494
495 self.persist_job(&job).await?;
497
498 info!("Submitted batch job: {} ({})", job.name, job_id);
499 Ok(job_id)
500 }
501
502 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 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 let entities = self.load_entities(&job).await?;
543
544 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 {
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 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 {
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 let chunk_result = self
594 .process_chunk(&job, chunk, chunk_idx, model.clone())
595 .await?;
596
597 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 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 self.update_job_progress(
613 &job.job_id,
614 chunk_idx + 1,
615 successful_embeddings + failed_embeddings,
616 )
617 .await?;
618
619 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 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 {
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 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 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 if let Some(_embedding) = self.cache_manager.get_embedding(entity) {
719 return Ok(true);
720 }
721 }
722
723 let embedding = model.get_entity_embedding(entity)?;
725
726 if use_cache {
728 self.cache_manager
729 .put_embedding(entity.to_string(), embedding);
730 }
731
732 Ok(false)
733 }
734
735 async fn load_entities(&self, job: &BatchJob) -> Result<Vec<String>> {
737 match &job.input.input_type {
738 InputType::EntityList => {
739 let entities: Vec<String> = serde_json::from_str(&job.input.source)?;
741 Ok(entities)
742 }
743 InputType::EntityFile => {
744 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 warn!("SPARQL query input type not yet implemented");
757 Ok(Vec::new())
758 }
759 InputType::DatabaseQuery => {
760 warn!("Database query input type not yet implemented");
762 Ok(Vec::new())
763 }
764 InputType::StreamSource => {
765 warn!("Stream source input type not yet implemented");
767 Ok(Vec::new())
768 }
769 }
770 }
771
772 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 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 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 async fn load_existing_entities(&self, path: &str) -> Result<HashSet<String>> {
809 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 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 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 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 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(), };
874
875 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 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 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, cpu_utilization: 0.0, };
920
921 let output_info = OutputInfo {
922 output_files: vec![job.output.path.clone()],
923 total_size_bytes: 0, 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, })
934 }
935
936 async fn validate_job(&self, job: &BatchJob) -> Result<()> {
938 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 } 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 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 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 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 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 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#[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); 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 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 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 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 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 let _batch = iterator.next_batch().expect("should succeed");
1165 let memory_usage = iterator.get_memory_usage();
1166 assert!(memory_usage > 0);
1167
1168 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 let processor =
1177 ParallelBatchProcessor::new(ParallelBatchConfig::default()).expect("should succeed");
1178 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
1189pub struct ParallelBatchProcessor {
1197 config: ParallelBatchConfig,
1198}
1199
1200#[derive(Debug, Clone, Serialize, Deserialize)]
1202pub struct ParallelBatchConfig {
1203 pub num_workers: usize,
1205 pub chunk_size: usize,
1207 pub adaptive_balancing: bool,
1209 pub memory_threshold_mb: usize,
1211 pub numa_aware: bool,
1213 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 pub fn new(config: ParallelBatchConfig) -> Result<Self> {
1235 rayon::ThreadPoolBuilder::new()
1237 .num_threads(config.num_workers)
1238 .build_global()
1239 .ok(); Ok(Self { config })
1242 }
1243
1244 pub fn num_workers(&self) -> usize {
1246 self.config.num_workers
1247 }
1248
1249 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 let results: Vec<R> = items.par_iter().map(process_fn).collect();
1258
1259 Ok(results)
1260 }
1261
1262 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 let results: Vec<R> = items.par_iter().map(process_fn).collect();
1280
1281 Ok(results)
1282 }
1283
1284 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 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 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 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#[derive(Debug, Clone)]
1346pub struct ParallelProcessingStats {
1347 pub num_workers: usize,
1348 pub profiler_report: String,
1349 pub memory_usage: usize,
1350}