additory 0.1.0a1__py3-none-any.whl

additory/core/backends/arrow_bridge.py

@@ -0,0 +1,476 @@
+# 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)
+
+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)
+        }