additory 0.1.0a4__py3-none-any.whl → 0.1.1a1__py3-none-any.whl

additory/core/backends/arrow_bridge.py

@@ -1,483 +0,0 @@
-# enhanced_arrow_bridge.py
-# Universal Arrow bridge for cross-backend dataframe compatibility
-
-import os
-import gc
-import psutil
-from typing import Any, Dict, Optional, Tuple, Union
-from dataclasses import dataclass
-from datetime import datetime
-
-try:
-    import pyarrow as pa
-    import polars as pl
-    import pandas as pd
-    ARROW_AVAILABLE = True
-except ImportError as e:
-    ARROW_AVAILABLE = False
-    IMPORT_ERROR = str(e)
-    # Create dummy classes for type annotations
-    class pa:
-        Table = Any
-    class pl:
-        DataFrame = Any
-    class pd:
-        DataFrame = Any
-
-from ..logging import log_info, log_warning
-from .cudf_bridge import get_cudf_bridge
-
-
-@dataclass
-class ConversionStats:
-    """Statistics for Arrow bridge operations"""
-    conversions: int = 0
-    memory_used_mb: float = 0.0
-    total_rows_processed: int = 0
-    total_columns_processed: int = 0
-    conversion_time_ms: float = 0.0
-    cleanup_count: int = 0
-
-
-class ArrowBridgeError(Exception):
-    """Raised when Arrow bridge operations fail"""
-    pass
-
-
-class EnhancedArrowBridge:
-    """Universal data bridge using Apache Arrow for cross-backend compatibility"""
-    
-    def __init__(self):
-        if not ARROW_AVAILABLE:
-            raise ArrowBridgeError(f"Arrow dependencies not available: {IMPORT_ERROR}")
-        
-        # Memory management
-        self.memory_pool = pa.default_memory_pool()
-        self.memory_threshold_mb = 100  # Cleanup threshold
-        
-        # Statistics tracking
-        self.stats = ConversionStats()
-        
-        # Supported backends (order matters - check cuDF before pandas)
-        self.supported_backends = {
-            "cudf": self._detect_cudf,
-            "polars": self._detect_polars,
-            "pandas": self._detect_pandas
-        }
-        
-        log_info("[arrow_bridge] Enhanced Arrow Bridge initialized")
-    
-    def detect_backend(self, df: Any) -> str:
-        """
-        Automatically detect dataframe backend
-        
-        Args:
-            df: Input dataframe of unknown type
-            
-        Returns:
-            Backend name ("pandas", "polars", "cudf")
-            
-        Raises:
-            ArrowBridgeError: If backend cannot be detected
-        """
-        for backend_name, detector in self.supported_backends.items():
-            if detector(df):
-                log_info(f"[arrow_bridge] Detected backend: {backend_name}")
-                return backend_name
-        
-        # Try to get type information for better error message
-        df_type = type(df).__name__
-        df_module = getattr(type(df), '__module__', 'unknown')
-        
-        raise ArrowBridgeError(
-            f"Unsupported dataframe type: {df_type} from module {df_module}. "
-            f"Supported backends: {list(self.supported_backends.keys())}"
-        )
-    
-    def _detect_pandas(self, df: Any) -> bool:
-        """Detect if dataframe is pandas (but not cuDF)"""
-        # First check if it's cuDF (which also has pandas-like attributes)
-        if self._detect_cudf(df):
-            return False
-        # Then check for pandas attributes
-        return hasattr(df, 'iloc') and hasattr(df, 'loc') and hasattr(df, 'dtypes')
-    
-    def _detect_polars(self, df: Any) -> bool:
-        """Detect if dataframe is Polars"""
-        return (hasattr(df, 'lazy') and hasattr(df, 'collect') and hasattr(df, 'schema')) or \
-               str(type(df)).find('polars') != -1
-    
-    def _detect_cudf(self, df: Any) -> bool:
-        """Detect if dataframe is cuDF using enhanced detection"""
-        cudf_bridge = get_cudf_bridge()
-        return cudf_bridge.detect_cudf(df)
-    
-    def to_arrow(self, df: Any, backend_type: Optional[str] = None) -> pa.Table:
-        """
-        Convert any dataframe to Arrow table
-        
-        Args:
-            df: Input dataframe
-            backend_type: Backend type (auto-detected if None)
-            
-        Returns:
-            PyArrow Table
-            
-        Raises:
-            ArrowBridgeError: If conversion fails
-        """
-        start_time = datetime.now()
-        
-        try:
-            # Auto-detect backend if not specified
-            if backend_type is None:
-                backend_type = self.detect_backend(df)
-            
-            # Get memory usage before conversion
-            memory_before = self._get_memory_usage_mb()
-            
-            # Convert based on backend
-            if backend_type == "pandas":
-                arrow_table = self._pandas_to_arrow(df)
-            elif backend_type == "polars":
-                arrow_table = self._polars_to_arrow(df)
-            elif backend_type == "cudf":
-                arrow_table = self._cudf_to_arrow(df)
-            else:
-                raise ArrowBridgeError(f"Unsupported backend for conversion: {backend_type}")
-            
-            # Update statistics
-            memory_after = self._get_memory_usage_mb()
-            conversion_time = (datetime.now() - start_time).total_seconds() * 1000
-            
-            self.stats.conversions += 1
-            self.stats.memory_used_mb += max(0, memory_after - memory_before)
-            self.stats.total_rows_processed += arrow_table.num_rows
-            self.stats.total_columns_processed += arrow_table.num_columns
-            self.stats.conversion_time_ms += conversion_time
-            
-            log_info(f"[arrow_bridge] Converted {backend_type} to Arrow: "
-                    f"{arrow_table.num_rows} rows, {arrow_table.num_columns} columns "
-                    f"({conversion_time:.1f}ms)")
-            
-            return arrow_table
-            
-        except Exception as e:
-            raise ArrowBridgeError(f"Failed to convert {backend_type} to Arrow: {e}")
-    
-    def from_arrow(self, arrow_table: pa.Table, target_backend: str) -> Any:
-        """
-        Convert Arrow table back to target dataframe format
-        
-        Args:
-            arrow_table: PyArrow Table
-            target_backend: Target backend ("pandas", "polars", "cudf")
-            
-        Returns:
-            Dataframe in target format
-            
-        Raises:
-            ArrowBridgeError: If conversion fails
-        """
-        start_time = datetime.now()
-        
-        try:
-            # Convert based on target backend
-            if target_backend == "pandas":
-                result_df = self._arrow_to_pandas(arrow_table)
-            elif target_backend == "polars":
-                result_df = self._arrow_to_polars(arrow_table)
-            elif target_backend == "cudf":
-                result_df = self._arrow_to_cudf(arrow_table)
-            else:
-                raise ArrowBridgeError(f"Unsupported target backend: {target_backend}")
-            
-            conversion_time = (datetime.now() - start_time).total_seconds() * 1000
-            
-            log_info(f"[arrow_bridge] Converted Arrow to {target_backend}: "
-                    f"{arrow_table.num_rows} rows, {arrow_table.num_columns} columns "
-                    f"({conversion_time:.1f}ms)")
-            
-            return result_df
-            
-        except Exception as e:
-            raise ArrowBridgeError(f"Failed to convert Arrow to {target_backend}: {e}")
-    
-    def _pandas_to_arrow(self, df: pd.DataFrame) -> pa.Table:
-        """Convert pandas DataFrame to Arrow table"""
-        try:
-            # Use zero-copy conversion when possible
-            return pa.Table.from_pandas(df, preserve_index=False)
-        except Exception as e:
-            # Fallback: try with schema inference
-            try:
-                schema = pa.Schema.from_pandas(df)
-                return pa.Table.from_pandas(df, schema=schema, preserve_index=False)
-            except Exception as e2:
-                raise ArrowBridgeError(f"Pandas to Arrow conversion failed: {e2}")
-    
-    def _polars_to_arrow(self, df: pl.DataFrame) -> pa.Table:
-        """Convert Polars DataFrame to Arrow table"""
-        try:
-            return df.to_arrow()
-        except Exception as e:
-            raise ArrowBridgeError(f"Polars to Arrow conversion failed: {e}")
-    
-    def _cudf_to_arrow(self, df: Any) -> pa.Table:
-        """Convert cuDF DataFrame to Arrow table using enhanced bridge"""
-        cudf_bridge = get_cudf_bridge()
-        return cudf_bridge.cudf_to_arrow(df)
-    
-    def _arrow_to_pandas(self, arrow_table: pa.Table) -> pd.DataFrame:
-        """Convert Arrow table to pandas DataFrame"""
-        try:
-            return arrow_table.to_pandas()
-        except Exception as e:
-            raise ArrowBridgeError(f"Arrow to pandas conversion failed: {e}")
-    
-    def _arrow_to_polars(self, arrow_table: pa.Table) -> pl.DataFrame:
-        """Convert Arrow table to Polars DataFrame"""
-        try:
-            return pl.from_arrow(arrow_table)
-        except Exception as e:
-            raise ArrowBridgeError(f"Arrow to Polars conversion failed: {e}")
-    
-    def _arrow_to_cudf(self, arrow_table: pa.Table) -> Any:
-        """Convert Arrow table to cuDF DataFrame using enhanced bridge"""
-        cudf_bridge = get_cudf_bridge()
-        return cudf_bridge.arrow_to_cudf(arrow_table)
-    
-    def cleanup_arrow_memory(self) -> Dict[str, float]:
-        """
-        Clean up Arrow heap memory
-        
-        Returns:
-            Dictionary with cleanup statistics
-        """
-        # Get memory usage before cleanup
-        memory_before = self._get_memory_usage_mb()
-        arrow_allocated_before = self.memory_pool.bytes_allocated()
-        
-        # Force garbage collection
-        gc.collect()
-        
-        # Additional Arrow-specific cleanup
-        try:
-            # Clear any cached Arrow data
-            if hasattr(pa, 'jemalloc_memory_pool'):
-                # Use jemalloc pool if available for better memory management
-                pass
-        except Exception:
-            pass
-        
-        # Get memory usage after cleanup
-        memory_after = self._get_memory_usage_mb()
-        arrow_allocated_after = self.memory_pool.bytes_allocated()
-        
-        # Calculate cleanup statistics
-        memory_freed_mb = max(0, memory_before - memory_after)
-        arrow_freed_bytes = max(0, arrow_allocated_before - arrow_allocated_after)
-        
-        self.stats.cleanup_count += 1
-        
-        cleanup_stats = {
-            "memory_freed_mb": memory_freed_mb,
-            "arrow_freed_bytes": arrow_freed_bytes,
-            "memory_before_mb": memory_before,
-            "memory_after_mb": memory_after,
-            "cleanup_count": self.stats.cleanup_count
-        }
-        
-        if memory_freed_mb > 0 or arrow_freed_bytes > 0:
-            log_info(f"[arrow_bridge] Memory cleanup: {memory_freed_mb:.1f}MB system, "
-                    f"{arrow_freed_bytes} bytes Arrow heap freed")
-        
-        return cleanup_stats
-    
-    def _get_memory_usage_mb(self) -> float:
-        """Get current memory usage in MB"""
-        try:
-            process = psutil.Process(os.getpid())
-            return process.memory_info().rss / 1024 / 1024
-        except Exception:
-            return 0.0
-    
-    def should_cleanup(self) -> bool:
-        """Check if memory cleanup is needed"""
-        current_memory = self._get_memory_usage_mb()
-        return current_memory > self.memory_threshold_mb
-    
-    def get_memory_stats(self) -> Dict[str, Any]:
-        """
-        Get memory usage statistics
-        
-        Returns:
-            Dictionary with memory statistics
-        """
-        return {
-            "current_memory_mb": self._get_memory_usage_mb(),
-            "arrow_allocated_bytes": self.memory_pool.bytes_allocated(),
-            "memory_threshold_mb": self.memory_threshold_mb,
-            "total_memory_used_mb": self.stats.memory_used_mb,
-            "cleanup_needed": self.should_cleanup()
-        }
-    
-    def get_conversion_stats(self) -> ConversionStats:
-        """Get conversion statistics"""
-        return self.stats
-    
-    def reset_stats(self):
-        """Reset conversion statistics"""
-        self.stats = ConversionStats()
-        log_info("[arrow_bridge] Statistics reset")
-    
-    def set_memory_threshold(self, threshold_mb: float):
-        """Set memory cleanup threshold"""
-        self.memory_threshold_mb = threshold_mb
-        log_info(f"[arrow_bridge] Memory threshold set to {threshold_mb}MB")
-    
-    def convert_with_cleanup(self, df: Any, target_backend: str, 
-                           backend_type: Optional[str] = None) -> Any:
-        """
-        Convert dataframe with automatic memory cleanup
-        
-        Args:
-            df: Input dataframe
-            target_backend: Target backend for output
-            backend_type: Source backend (auto-detected if None)
-            
-        Returns:
-            Converted dataframe
-        """
-        try:
-            # Convert to Arrow
-            arrow_table = self.to_arrow(df, backend_type)
-            
-            # Convert to target format
-            result_df = self.from_arrow(arrow_table, target_backend)
-            
-            # Cleanup if needed
-            if self.should_cleanup():
-                self.cleanup_arrow_memory()
-            
-            return result_df
-            
-        except Exception as e:
-            # Always try to cleanup on error
-            self.cleanup_arrow_memory()
-            raise e
-    
-    def validate_arrow_table(self, arrow_table: pa.Table) -> bool:
-        """
-        Validate Arrow table structure
-        
-        Args:
-            arrow_table: PyArrow table to validate
-            
-        Returns:
-            True if table is valid
-        """
-        try:
-            # Basic validation
-            if arrow_table is None:
-                return False
-            
-            if not isinstance(arrow_table, pa.Table):
-                return False
-            
-            # Check if table has data
-            if arrow_table.num_rows == 0 and arrow_table.num_columns == 0:
-                log_warning("[arrow_bridge] Arrow table is empty")
-                return True  # Empty table is valid
-            
-            # Validate schema
-            schema = arrow_table.schema
-            if schema is None:
-                return False
-            
-            # Check for null schema fields
-            for field in schema:
-                if field is None:
-                    return False
-            
-            log_info(f"[arrow_bridge] Arrow table validation passed: "
-                    f"{arrow_table.num_rows} rows, {arrow_table.num_columns} columns")
-            
-            return True
-            
-        except Exception as e:
-            log_warning(f"[arrow_bridge] Arrow table validation failed: {e}")
-            return False
-    
-    def get_supported_backends(self) -> Dict[str, bool]:
-        """
-        Get list of supported backends and their availability
-        
-        Returns:
-            Dictionary mapping backend names to availability status
-        """
-        availability = {}
-        
-        # Check pandas
-        try:
-            import pandas
-            availability["pandas"] = True
-        except ImportError:
-            availability["pandas"] = False
-        
-        # Check Polars
-        try:
-            import polars
-            availability["polars"] = True
-        except ImportError:
-            availability["polars"] = False
-        
-        # Check cuDF
-        try:
-            import cudf
-            availability["cudf"] = True
-        except ImportError:
-            availability["cudf"] = False
-        
-        return availability
-    
-    def benchmark_conversion(self, df: Any, target_backend: str, 
-                           iterations: int = 3) -> Dict[str, float]:
-        """
-        Benchmark conversion performance
-        
-        Args:
-            df: Input dataframe
-            target_backend: Target backend
-            iterations: Number of iterations for benchmarking
-            
-        Returns:
-            Dictionary with benchmark results
-        """
-        import time
-        
-        times = []
-        source_backend = self.detect_backend(df)
-        
-        for i in range(iterations):
-            start_time = time.time()
-            
-            # Perform conversion
-            arrow_table = self.to_arrow(df)
-            result_df = self.from_arrow(arrow_table, target_backend)
-            
-            end_time = time.time()
-            times.append((end_time - start_time) * 1000)  # Convert to ms
-            
-            # Cleanup between iterations
-            self.cleanup_arrow_memory()
-        
-        return {
-            "source_backend": source_backend,
-            "target_backend": target_backend,
-            "iterations": iterations,
-            "min_time_ms": min(times),
-            "max_time_ms": max(times),
-            "avg_time_ms": sum(times) / len(times),
-            "total_time_ms": sum(times)
-        }