14. Batch Ingestion

Chapter 14 of 18 · 20 min

KEY INSIGHT

Process documents in batches of 100-500 to balance memory usage and throughput during large-scale indexing. Loading thousands of documents one at a time is slow. Loading all at once exhausts memory. Batch processing finds the balance. ```python from typing import List, Dict, Generator import time def batch_generator(items: List, batch_size: int) -> Generator[List, None, None]: """Yield batches of items.""" for i in range(0, len(items), batch_size): yield items[i:i + batch_size] def ingest_documents( engine: SemanticSearchEngine, documents: List[Dict], batch_size: int = 100 ) -> Dict: """Ingest documents in batches with progress reporting.""" total = len(documents) start_time = time.time() texts = [doc["text"] for doc in documents] ids = [doc["id"] for doc in documents] metadatas = [doc.get("metadata", {}) for doc in documents] processed = 0 for batch in batch_generator(range(total), batch_size): batch_texts = [texts[i] for i in batch] batch_ids = [ids[i] for i in batch] batch_metas = [metadatas[i] for i in batch] engine.index_documents( batch_texts, ids=batch_ids, metadatas=batch_metas ) processed += len(batch) elapsed = time.time() - start_time rate = processed / elapsed if elapsed > 0 else 0 print(f"Progress: {processed}/{total} ({100*processed/total:.1f}%) " f"- {rate:.1f} docs/sec") return { "total": total, "elapsed": time.time() - start_time } ``` ### Parallel Embedding Generation For CPU-bound embedding generation, use multiprocessing: ```python from sentence_transformers import SentenceTransformer from multiprocessing import Pool, cpu_count from functools import partial def encode_batch(texts: List[str], model_name: str) -> List: """Encode a batch of texts (worker function).""" model = SentenceTransformer(model_name) return model.encode(texts).tolist() def parallel_encode(texts: List[str], model_name: str, num_workers: int = None) -> List: """Encode texts in parallel using multiple processes.""" num_workers = num_workers or max(1, cpu_count() - 1) batch_size = 100 batches = list(batch_generator(texts, batch_size)) with Pool(num_workers) as pool: results = pool.map( partial(encode_batch, model_name=model_name), batches ) # Flatten results return [embedding for batch in results for embedding in batch] ``` ### Progress Tracking For long-running ingestion, track progress in a file: ```python import json from pathlib import Path def ingest_with_checkpoint( engine: SemanticSearchEngine, documents: List[Dict], checkpoint_file: str = ".ingestion_checkpoint.json", batch_size: int = 100 ): """Ingest documents with checkpointing for recovery.""" checkpoint_path = Path(checkpoint_file) # Load checkpoint if exists if checkpoint_path.exists(): with open(checkpoint_path) as f: checkpoint = json.load(f) start_index = checkpoint["indexed_count"] print(f"Resuming from checkpoint: {start_index} already indexed") else: start_index = 0 # Process remaining documents for i in range(start_index, len(documents), batch_size): batch = documents[i:i + batch_size] engine.index_documents( [d["text"] for d in batch], ids=[d["id"] for d in batch], metadatas=[d.get("metadata", {}) for d in batch] ) # Save checkpoint with open(checkpoint_path, 'w') as f: json.dump({"indexed_count": i + len(batch)}, f) ```

EXERCISE

Create 5000 synthetic documents. Ingest them with batch sizes of 50, 100, and 500. Time each run and compare throughput (documents/second).