MaverickMCP

""" Data chunking utilities for memory-efficient processing of large datasets. Provides streaming, batching, and generator-based approaches for handling large DataFrames. """ import logging import math from collections.abc import Callable, Generator from typing import Any, Literal import numpy as np import pandas as pd from maverick_mcp.utils.memory_profiler import ( force_garbage_collection, get_dataframe_memory_usage, memory_context, optimize_dataframe, ) logger = logging.getLogger(__name__) # Default chunk size configurations DEFAULT_CHUNK_SIZE_MB = 50.0 MAX_CHUNK_SIZE_MB = 200.0 MIN_ROWS_PER_CHUNK = 100 class DataChunker: """Advanced data chunking utility with multiple strategies.""" def __init__( self, chunk_size_mb: float = DEFAULT_CHUNK_SIZE_MB, min_rows_per_chunk: int = MIN_ROWS_PER_CHUNK, optimize_chunks: bool = True, auto_gc: bool = True, ): """Initialize data chunker. Args: chunk_size_mb: Target chunk size in megabytes min_rows_per_chunk: Minimum rows per chunk optimize_chunks: Whether to optimize chunk memory usage auto_gc: Whether to automatically run garbage collection """ self.chunk_size_mb = min(chunk_size_mb, MAX_CHUNK_SIZE_MB) self.chunk_size_bytes = int(self.chunk_size_mb * 1024 * 1024) self.min_rows_per_chunk = min_rows_per_chunk self.optimize_chunks = optimize_chunks self.auto_gc = auto_gc logger.debug( f"DataChunker initialized: {self.chunk_size_mb}MB chunks, " f"min {self.min_rows_per_chunk} rows" ) def estimate_chunk_size(self, df: pd.DataFrame) -> tuple[int, int]: """Estimate optimal chunk size for a DataFrame. Args: df: DataFrame to analyze Returns: Tuple of (rows_per_chunk, estimated_chunks) """ total_memory = df.memory_usage(deep=True).sum() memory_per_row = total_memory / len(df) if len(df) > 0 else 0 if memory_per_row == 0: return len(df), 1 # Calculate rows per chunk based on memory target rows_per_chunk = max( self.min_rows_per_chunk, int(self.chunk_size_bytes / memory_per_row) ) # Ensure we don't exceed the DataFrame size rows_per_chunk = min(rows_per_chunk, len(df)) estimated_chunks = math.ceil(len(df) / rows_per_chunk) logger.debug( f"Estimated chunking: {rows_per_chunk} rows/chunk, " f"{estimated_chunks} chunks total" ) return rows_per_chunk, estimated_chunks def chunk_by_rows( self, df: pd.DataFrame, rows_per_chunk: int = None ) -> Generator[pd.DataFrame, None, None]: """Chunk DataFrame by number of rows. Args: df: DataFrame to chunk rows_per_chunk: Rows per chunk (auto-estimated if None) Yields: DataFrame chunks """ if rows_per_chunk is None: rows_per_chunk, _ = self.estimate_chunk_size(df) total_chunks = math.ceil(len(df) / rows_per_chunk) logger.debug( f"Chunking {len(df)} rows into {total_chunks} chunks " f"of ~{rows_per_chunk} rows each" ) for i, start_idx in enumerate(range(0, len(df), rows_per_chunk)): end_idx = min(start_idx + rows_per_chunk, len(df)) chunk = df.iloc[start_idx:end_idx].copy() if self.optimize_chunks: chunk = optimize_dataframe(chunk) logger.debug( f"Yielding chunk {i + 1}/{total_chunks}: rows {start_idx}-{end_idx - 1}" ) yield chunk # Cleanup after yielding if self.auto_gc: del chunk if i % 5 == 0: # GC every 5 chunks force_garbage_collection() def chunk_by_memory(self, df: pd.DataFrame) -> Generator[pd.DataFrame, None, None]: """Chunk DataFrame by memory size. Args: df: DataFrame to chunk Yields: DataFrame chunks """ total_memory = df.memory_usage(deep=True).sum() if total_memory <= self.chunk_size_bytes: if self.optimize_chunks: df = optimize_dataframe(df) yield df return # Use row-based chunking with memory-based estimation yield from self.chunk_by_rows(df) def chunk_by_date( self, df: pd.DataFrame, freq: Literal["D", "W", "M", "Q", "Y"] = "M", date_column: str = None, ) -> Generator[pd.DataFrame, None, None]: """Chunk DataFrame by date periods. Args: df: DataFrame to chunk (must have datetime index or date_column) freq: Frequency for chunking (D=daily, W=weekly, M=monthly, etc.) date_column: Name of date column (uses index if None) Yields: DataFrame chunks by date periods """ if date_column: if date_column not in df.columns: raise ValueError(f"Date column '{date_column}' not found") elif not isinstance(df.index, pd.DatetimeIndex): raise ValueError( "DataFrame must have datetime index or specify date_column" ) # Group by period period_groups = df.groupby( pd.Grouper(key=date_column, freq=freq) if date_column else pd.Grouper(freq=freq) ) total_periods = len(period_groups) logger.debug(f"Chunking by {freq} periods: {total_periods} chunks") for i, (period, group) in enumerate(period_groups): if len(group) == 0: continue if self.optimize_chunks: group = optimize_dataframe(group) logger.debug( f"Yielding period chunk {i + 1}/{total_periods}: " f"{period} ({len(group)} rows)" ) yield group if self.auto_gc and i % 3 == 0: # GC every 3 periods force_garbage_collection() def process_in_chunks( self, df: pd.DataFrame, processor: Callable[[pd.DataFrame], Any], combiner: Callable[[list], Any] = None, chunk_method: Literal["rows", "memory", "date"] = "memory", **chunk_kwargs, ) -> Any: """Process DataFrame in chunks and combine results. Args: df: DataFrame to process processor: Function to apply to each chunk combiner: Function to combine results (default: list) chunk_method: Chunking method to use **chunk_kwargs: Additional arguments for chunking method Returns: Combined results """ results = [] # Select chunking method if chunk_method == "rows": chunk_generator = self.chunk_by_rows(df, **chunk_kwargs) elif chunk_method == "memory": chunk_generator = self.chunk_by_memory(df) elif chunk_method == "date": chunk_generator = self.chunk_by_date(df, **chunk_kwargs) else: raise ValueError(f"Unknown chunk method: {chunk_method}") with memory_context("chunk_processing"): for i, chunk in enumerate(chunk_generator): try: with memory_context(f"chunk_{i}"): result = processor(chunk) results.append(result) except Exception as e: logger.error(f"Error processing chunk {i}: {e}") raise # Combine results if combiner: return combiner(results) elif results and isinstance(results[0], pd.DataFrame): # Auto-combine DataFrames return pd.concat(results, ignore_index=True) else: return results class StreamingDataProcessor: """Streaming data processor for very large datasets.""" def __init__(self, chunk_size_mb: float = DEFAULT_CHUNK_SIZE_MB): """Initialize streaming processor. Args: chunk_size_mb: Chunk size in MB """ self.chunk_size_mb = chunk_size_mb self.chunker = DataChunker(chunk_size_mb=chunk_size_mb) def stream_from_csv( self, filepath: str, processor: Callable[[pd.DataFrame], Any], chunksize: int = None, **read_kwargs, ) -> Generator[Any, None, None]: """Stream process CSV file in chunks. Args: filepath: Path to CSV file processor: Function to process each chunk chunksize: Rows per chunk (auto-estimated if None) **read_kwargs: Additional arguments for pd.read_csv Yields: Processed results for each chunk """ # Estimate chunk size if not provided if chunksize is None: # Read a sample to estimate memory usage sample = pd.read_csv(filepath, nrows=1000, **read_kwargs) memory_per_row = sample.memory_usage(deep=True).sum() / len(sample) chunksize = max(100, int(self.chunker.chunk_size_bytes / memory_per_row)) del sample force_garbage_collection() logger.info(f"Streaming CSV with {chunksize} rows per chunk") chunk_reader = pd.read_csv(filepath, chunksize=chunksize, **read_kwargs) for i, chunk in enumerate(chunk_reader): with memory_context(f"csv_chunk_{i}"): # Optimize chunk if needed if self.chunker.optimize_chunks: chunk = optimize_dataframe(chunk) result = processor(chunk) yield result # Clean up del chunk if i % 5 == 0: force_garbage_collection() def stream_from_database( self, query: str, connection, processor: Callable[[pd.DataFrame], Any], chunksize: int = None, ) -> Generator[Any, None, None]: """Stream process database query results in chunks. Args: query: SQL query connection: Database connection processor: Function to process each chunk chunksize: Rows per chunk Yields: Processed results for each chunk """ if chunksize is None: chunksize = 10000 # Default for database queries logger.info(f"Streaming database query with {chunksize} rows per chunk") chunk_reader = pd.read_sql(query, connection, chunksize=chunksize) for i, chunk in enumerate(chunk_reader): with memory_context(f"db_chunk_{i}"): if self.chunker.optimize_chunks: chunk = optimize_dataframe(chunk) result = processor(chunk) yield result del chunk if i % 3 == 0: force_garbage_collection() def optimize_dataframe_dtypes( df: pd.DataFrame, aggressive: bool = False, categorical_threshold: float = 0.5 ) -> pd.DataFrame: """Optimize DataFrame data types for memory efficiency. Args: df: DataFrame to optimize aggressive: Use aggressive optimizations (may lose precision) categorical_threshold: Threshold for categorical conversion Returns: Optimized DataFrame """ logger.debug(f"Optimizing DataFrame dtypes: {df.shape}") initial_memory = df.memory_usage(deep=True).sum() df_opt = df.copy() for col in df_opt.columns: col_type = df_opt[col].dtype try: if col_type == "object": # Convert string columns to categorical if beneficial unique_count = df_opt[col].nunique() total_count = len(df_opt[col]) if unique_count / total_count < categorical_threshold: df_opt[col] = df_opt[col].astype("category") logger.debug(f"Converted {col} to categorical") elif "int" in str(col_type): # Downcast integers c_min = df_opt[col].min() c_max = df_opt[col].max() if c_min >= np.iinfo(np.int8).min and c_max <= np.iinfo(np.int8).max: df_opt[col] = df_opt[col].astype(np.int8) elif ( c_min >= np.iinfo(np.int16).min and c_max <= np.iinfo(np.int16).max ): df_opt[col] = df_opt[col].astype(np.int16) elif ( c_min >= np.iinfo(np.int32).min and c_max <= np.iinfo(np.int32).max ): df_opt[col] = df_opt[col].astype(np.int32) elif "float" in str(col_type) and col_type == "float64": # Downcast float64 to float32 if no precision loss if aggressive: # Check if conversion preserves data temp = df_opt[col].astype(np.float32) if np.allclose( df_opt[col].fillna(0), temp.fillna(0), rtol=1e-6, equal_nan=True ): df_opt[col] = temp logger.debug(f"Converted {col} to float32") except Exception as e: logger.debug(f"Could not optimize column {col}: {e}") continue final_memory = df_opt.memory_usage(deep=True).sum() memory_saved = initial_memory - final_memory if memory_saved > 0: logger.info( f"DataFrame optimization saved {memory_saved / (1024**2):.2f}MB " f"({memory_saved / initial_memory * 100:.1f}% reduction)" ) return df_opt def create_memory_efficient_dataframe( data: dict | list, optimize: bool = True, categorical_columns: list[str] = None ) -> pd.DataFrame: """Create a memory-efficient DataFrame from data. Args: data: Data to create DataFrame from optimize: Whether to optimize dtypes categorical_columns: Columns to convert to categorical Returns: Memory-optimized DataFrame """ with memory_context("creating_dataframe"): df = pd.DataFrame(data) if categorical_columns: for col in categorical_columns: if col in df.columns: df[col] = df[col].astype("category") if optimize: df = optimize_dataframe_dtypes(df) return df def batch_process_large_dataframe( df: pd.DataFrame, operation: Callable, batch_size: int = None, combine_results: bool = True, ) -> Any: """Process large DataFrame in batches to manage memory. Args: df: Large DataFrame to process operation: Function to apply to each batch batch_size: Size of each batch (auto-estimated if None) combine_results: Whether to combine batch results Returns: Combined results or list of batch results """ chunker = DataChunker() if batch_size: chunk_generator = chunker.chunk_by_rows(df, batch_size) else: chunk_generator = chunker.chunk_by_memory(df) results = [] with memory_context("batch_processing"): for i, batch in enumerate(chunk_generator): logger.debug(f"Processing batch {i + 1}") with memory_context(f"batch_{i}"): result = operation(batch) results.append(result) if combine_results and results: if isinstance(results[0], pd.DataFrame): return pd.concat(results, ignore_index=True) elif isinstance(results[0], int | float): return sum(results) elif isinstance(results[0], list): return [item for sublist in results for item in sublist] return results class LazyDataFrame: """Lazy evaluation wrapper for large DataFrames.""" def __init__(self, data_source: str | pd.DataFrame, chunk_size_mb: float = 50.0): """Initialize lazy DataFrame. Args: data_source: File path or DataFrame chunk_size_mb: Chunk size for processing """ self.data_source = data_source self.chunker = DataChunker(chunk_size_mb=chunk_size_mb) self._cached_info = None def get_info(self) -> dict[str, Any]: """Get DataFrame information without loading full data.""" if self._cached_info: return self._cached_info if isinstance(self.data_source, str): # Read just the header and a sample sample = pd.read_csv(self.data_source, nrows=100) total_rows = sum(1 for _ in open(self.data_source)) - 1 # Subtract header self._cached_info = { "columns": sample.columns.tolist(), "dtypes": sample.dtypes.to_dict(), "estimated_rows": total_rows, "sample_memory_mb": sample.memory_usage(deep=True).sum() / (1024**2), } else: self._cached_info = get_dataframe_memory_usage(self.data_source) return self._cached_info def apply_chunked(self, operation: Callable) -> Any: """Apply operation in chunks.""" if isinstance(self.data_source, str): processor = StreamingDataProcessor(self.chunker.chunk_size_mb) results = list(processor.stream_from_csv(self.data_source, operation)) else: results = self.chunker.process_in_chunks(self.data_source, operation) return results def to_optimized_dataframe(self) -> pd.DataFrame: """Load and optimize the full DataFrame.""" if isinstance(self.data_source, str): df = pd.read_csv(self.data_source) else: df = self.data_source.copy() return optimize_dataframe_dtypes(df) # Utility functions for common operations def chunked_concat( dataframes: list[pd.DataFrame], chunk_size: int = 10 ) -> pd.DataFrame: """Concatenate DataFrames in chunks to manage memory. Args: dataframes: List of DataFrames to concatenate chunk_size: Number of DataFrames to concat at once Returns: Concatenated DataFrame """ if not dataframes: return pd.DataFrame() if len(dataframes) <= chunk_size: return pd.concat(dataframes, ignore_index=True) # Process in chunks results = [] for i in range(0, len(dataframes), chunk_size): chunk = dataframes[i : i + chunk_size] with memory_context(f"concat_chunk_{i // chunk_size}"): result = pd.concat(chunk, ignore_index=True) results.append(result) # Clean up chunk for df in chunk: del df force_garbage_collection() # Final concatenation with memory_context("final_concat"): final_result = pd.concat(results, ignore_index=True) return final_result def memory_efficient_groupby( df: pd.DataFrame, group_col: str, agg_func: Callable, chunk_size_mb: float = 50.0 ) -> pd.DataFrame: """Perform memory-efficient groupby operations. Args: df: DataFrame to group group_col: Column to group by agg_func: Aggregation function chunk_size_mb: Chunk size in MB Returns: Aggregated DataFrame """ if group_col not in df.columns: raise ValueError(f"Group column '{group_col}' not found") chunker = DataChunker(chunk_size_mb=chunk_size_mb) results = [] def process_chunk(chunk): return chunk.groupby(group_col).apply(agg_func).reset_index() results = chunker.process_in_chunks(df, process_chunk) # Combine and re-aggregate results combined = pd.concat(results, ignore_index=True) final_result = combined.groupby(group_col).apply(agg_func).reset_index() return final_result