aiecs 1.1.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

aiecs/tools/statistics/data_loader_tool.py

@@ -0,0 +1,518 @@
+"""
+Data Loader Tool - Universal data loading from multiple file formats
+
+This tool provides comprehensive data loading capabilities with:
+- Auto-detection of file formats
+- Multiple loading strategies (full, streaming, chunked, lazy)
+- Data quality validation on load
+- Schema inference and validation
+- Support for CSV, Excel, JSON, Parquet, and other formats
+"""
+
+import os
+import logging
+from typing import Dict, Any, List, Optional, Union, Iterator
+from enum import Enum
+from pathlib import Path
+
+import pandas as pd
+import numpy as np
+from pydantic import BaseModel, Field, ValidationError, ConfigDict
+
+from aiecs.tools.base_tool import BaseTool
+from aiecs.tools import register_tool
+
+
+class DataSourceType(str, Enum):
+    """Supported data source types"""
+    CSV = "csv"
+    EXCEL = "excel"
+    JSON = "json"
+    PARQUET = "parquet"
+    FEATHER = "feather"
+    HDF5 = "hdf5"
+    STATA = "stata"
+    SAS = "sas"
+    SPSS = "spss"
+    AUTO = "auto"
+
+
+class LoadStrategy(str, Enum):
+    """Data loading strategies"""
+    FULL_LOAD = "full_load"
+    STREAMING = "streaming"
+    CHUNKED = "chunked"
+    LAZY = "lazy"
+    INCREMENTAL = "incremental"
+
+
+
+
+class DataLoaderError(Exception):
+    """Base exception for DataLoader errors"""
+    pass
+
+
+class FileFormatError(DataLoaderError):
+    """Raised when file format is unsupported or invalid"""
+    pass
+
+
+class SchemaValidationError(DataLoaderError):
+    """Raised when schema validation fails"""
+    pass
+
+
+class DataQualityError(DataLoaderError):
+    """Raised when data quality issues are detected"""
+    pass
+
+
+@register_tool('data_loader')
+class DataLoaderTool(BaseTool):
+    """
+    Universal data loading tool that can:
+    1. Load data from multiple file formats
+    2. Auto-detect data formats and schemas
+    3. Handle large datasets with streaming
+    4. Validate data quality on load
+    
+    Integrates with pandas_tool for core data operations.
+    """
+    
+    # Configuration schema
+    class Config(BaseModel):
+        """Configuration for the data loader tool"""
+        model_config = ConfigDict(env_prefix="DATA_LOADER_")
+        
+        max_file_size_mb: int = Field(
+            default=500,
+            description="Maximum file size in megabytes"
+        )
+        default_chunk_size: int = Field(
+            default=10000,
+            description="Default chunk size for chunked loading"
+        )
+        max_memory_usage_mb: int = Field(
+            default=2000,
+            description="Maximum memory usage in megabytes"
+        )
+        enable_schema_inference: bool = Field(
+            default=True,
+            description="Whether to enable automatic schema inference"
+        )
+        enable_quality_validation: bool = Field(
+            default=True,
+            description="Whether to enable data quality validation"
+        )
+        default_encoding: str = Field(
+            default="utf-8",
+            description="Default text encoding for file operations"
+        )
+    
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize DataLoaderTool with settings.
+        
+        Args:
+            config: Optional configuration overrides
+        """
+        super().__init__(config)
+        
+        # Parse configuration
+        self.config = self.Config(**(config or {}))
+        
+        self.logger = logging.getLogger(__name__)
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
+            self.logger.addHandler(handler)
+        self.logger.setLevel(logging.INFO)
+        
+        # Initialize external tools
+        self._init_external_tools()
+    
+    def _init_external_tools(self):
+        """Initialize external task tools"""
+        self.external_tools = {}
+        
+        # Initialize PandasTool for data operations
+        try:
+            from aiecs.tools.task_tools.pandas_tool import PandasTool
+            self.external_tools['pandas'] = PandasTool()
+            self.logger.info("PandasTool initialized successfully")
+        except ImportError:
+            self.logger.warning("PandasTool not available")
+            self.external_tools['pandas'] = None
+    
+    # Schema definitions
+    class LoadDataSchema(BaseModel):
+        """Schema for load_data operation"""
+        source: str = Field(description="Path to data source file")
+        source_type: Optional[DataSourceType] = Field(default=DataSourceType.AUTO, description="Data source type")
+        strategy: LoadStrategy = Field(default=LoadStrategy.FULL_LOAD, description="Loading strategy")
+        schema: Optional[Dict[str, Any]] = Field(default=None, description="Expected schema for validation")
+        validation_rules: Optional[Dict[str, Any]] = Field(default=None, description="Data quality validation rules")
+        nrows: Optional[int] = Field(default=None, description="Number of rows to load")
+        chunk_size: Optional[int] = Field(default=None, description="Chunk size for chunked loading")
+        encoding: Optional[str] = Field(default=None, description="File encoding")
+    
+    class DetectFormatSchema(BaseModel):
+        """Schema for detect_format operation"""
+        source: str = Field(description="Path to data source file")
+    
+    class ValidateSchemaSchema(BaseModel):
+        """Schema for validate_schema operation"""
+        data: Union[Dict[str, Any], List[Dict[str, Any]]] = Field(description="Data to validate")
+        schema: Dict[str, Any] = Field(description="Expected schema")
+    
+    class StreamDataSchema(BaseModel):
+        """Schema for stream_data operation"""
+        source: str = Field(description="Path to data source file")
+        chunk_size: int = Field(default=10000, description="Chunk size for streaming")
+        source_type: Optional[DataSourceType] = Field(default=DataSourceType.AUTO, description="Data source type")
+    
+    def load_data(
+        self,
+        source: str,
+        source_type: DataSourceType = DataSourceType.AUTO,
+        strategy: LoadStrategy = LoadStrategy.FULL_LOAD,
+        schema: Optional[Dict[str, Any]] = None,
+        validation_rules: Optional[Dict[str, Any]] = None,
+        nrows: Optional[int] = None,
+        chunk_size: Optional[int] = None,
+        encoding: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Load data from source with automatic format detection.
+        
+        Args:
+            source: Path to data source file
+            source_type: Type of data source (auto-detected if not specified)
+            strategy: Loading strategy to use
+            schema: Expected schema for validation
+            validation_rules: Data quality validation rules
+            nrows: Number of rows to load (None for all)
+            chunk_size: Chunk size for chunked loading
+            encoding: File encoding
+            
+        Returns:
+            Dict containing:
+                - data: Loaded DataFrame or data structure
+                - metadata: Metadata about loaded data
+                - quality_report: Quality assessment results
+                
+        Raises:
+            DataLoaderError: If loading fails
+            FileFormatError: If format is unsupported
+        """
+        try:
+            # Validate source exists
+            if not os.path.exists(source):
+                raise DataLoaderError(f"Source file not found: {source}")
+            
+            # Detect format if auto
+            if source_type == DataSourceType.AUTO:
+                source_type = self._detect_format(source)
+            
+            # Check file size
+            file_size_mb = os.path.getsize(source) / (1024 * 1024)
+            if file_size_mb > self.config.max_file_size_mb:
+                self.logger.warning(f"File size {file_size_mb:.2f}MB exceeds recommended limit")
+            
+            # Load data based on strategy
+            if strategy == LoadStrategy.FULL_LOAD:
+                data = self._load_full(source, source_type, nrows, encoding)
+            elif strategy == LoadStrategy.CHUNKED:
+                data = self._load_chunked(source, source_type, chunk_size or self.config.default_chunk_size, encoding)
+            elif strategy == LoadStrategy.STREAMING:
+                data = self._load_streaming(source, source_type, chunk_size or self.config.default_chunk_size, encoding)
+            elif strategy == LoadStrategy.LAZY:
+                data = self._load_lazy(source, source_type, encoding)
+            else:
+                raise DataLoaderError(f"Unsupported loading strategy: {strategy}")
+            
+            # Generate metadata
+            metadata = self._generate_metadata(data, source, source_type)
+            
+            # Validate schema if provided
+            if schema and self.config.enable_schema_inference:
+                schema_valid = self._validate_schema_internal(data, schema)
+                metadata['schema_valid'] = schema_valid
+            
+            # Validate quality if enabled
+            quality_report = {}
+            if self.config.enable_quality_validation and isinstance(data, pd.DataFrame):
+                quality_report = self._validate_quality(data, validation_rules)
+            
+            self.logger.info(f"Successfully loaded data from {source}")
+            
+            return {
+                'data': data,
+                'metadata': metadata,
+                'quality_report': quality_report,
+                'source': source,
+                'source_type': source_type.value,
+                'strategy': strategy.value
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error loading data from {source}: {e}")
+            raise DataLoaderError(f"Failed to load data: {e}")
+    
+    def detect_format(self, source: str) -> Dict[str, Any]:
+        """
+        Detect file format from source.
+        
+        Args:
+            source: Path to data source file
+            
+        Returns:
+            Dict containing detected format information
+        """
+        try:
+            detected_type = self._detect_format(source)
+            
+            return {
+                'source': source,
+                'detected_type': detected_type.value,
+                'file_extension': Path(source).suffix.lower(),
+                'confidence': 'high'
+            }
+        except Exception as e:
+            self.logger.error(f"Error detecting format: {e}")
+            raise FileFormatError(f"Failed to detect format: {e}")
+    
+    def validate_schema(self, data: Union[Dict[str, Any], List[Dict[str, Any]]], schema: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Validate data against expected schema.
+        
+        Args:
+            data: Data to validate
+            schema: Expected schema definition
+            
+        Returns:
+            Dict containing validation results
+        """
+        try:
+            # Convert to DataFrame if needed
+            if isinstance(data, list):
+                df = pd.DataFrame(data)
+            elif isinstance(data, dict):
+                df = pd.DataFrame([data])
+            else:
+                df = data
+            
+            is_valid = self._validate_schema_internal(df, schema)
+            
+            issues = []
+            if not is_valid:
+                # Check column presence
+                expected_columns = set(schema.get('columns', {}).keys())
+                actual_columns = set(df.columns)
+                missing = expected_columns - actual_columns
+                extra = actual_columns - expected_columns
+                
+                if missing:
+                    issues.append(f"Missing columns: {missing}")
+                if extra:
+                    issues.append(f"Extra columns: {extra}")
+            
+            return {
+                'valid': is_valid,
+                'issues': issues,
+                'expected_columns': list(schema.get('columns', {}).keys()),
+                'actual_columns': list(df.columns)
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error validating schema: {e}")
+            raise SchemaValidationError(f"Schema validation failed: {e}")
+    
+    def stream_data(
+        self,
+        source: str,
+        chunk_size: int = 10000,
+        source_type: DataSourceType = DataSourceType.AUTO
+    ) -> Dict[str, Any]:
+        """
+        Stream data in chunks for large files.
+        
+        Args:
+            source: Path to data source file
+            chunk_size: Size of each chunk
+            source_type: Type of data source
+            
+        Returns:
+            Dict containing streaming iterator information
+        """
+        try:
+            if source_type == DataSourceType.AUTO:
+                source_type = self._detect_format(source)
+            
+            # Create iterator based on format
+            if source_type == DataSourceType.CSV:
+                iterator = pd.read_csv(source, chunksize=chunk_size)
+            elif source_type == DataSourceType.JSON:
+                iterator = pd.read_json(source, lines=True, chunksize=chunk_size)
+            else:
+                raise FileFormatError(f"Streaming not supported for format: {source_type}")
+            
+            return {
+                'iterator': iterator,
+                'chunk_size': chunk_size,
+                'source_type': source_type.value,
+                'message': 'Streaming iterator created successfully'
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error creating stream: {e}")
+            raise DataLoaderError(f"Failed to create stream: {e}")
+    
+    # Internal helper methods
+    
+    def _detect_format(self, source: str) -> DataSourceType:
+        """Detect file format from extension"""
+        ext = Path(source).suffix.lower()
+        
+        format_map = {
+            '.csv': DataSourceType.CSV,
+            '.xlsx': DataSourceType.EXCEL,
+            '.xls': DataSourceType.EXCEL,
+            '.json': DataSourceType.JSON,
+            '.parquet': DataSourceType.PARQUET,
+            '.feather': DataSourceType.FEATHER,
+            '.h5': DataSourceType.HDF5,
+            '.hdf': DataSourceType.HDF5,
+            '.dta': DataSourceType.STATA,
+            '.sas7bdat': DataSourceType.SAS,
+            '.sav': DataSourceType.SPSS
+        }
+        
+        detected = format_map.get(ext)
+        if not detected:
+            raise FileFormatError(f"Unsupported file format: {ext}")
+        
+        return detected
+    
+    def _load_full(self, source: str, source_type: DataSourceType, nrows: Optional[int], encoding: Optional[str]) -> pd.DataFrame:
+        """Load entire dataset into memory"""
+        encoding = encoding or self.config.default_encoding
+        
+        if source_type == DataSourceType.CSV:
+            return pd.read_csv(source, nrows=nrows, encoding=encoding)
+        elif source_type == DataSourceType.EXCEL:
+            return pd.read_excel(source, nrows=nrows)
+        elif source_type == DataSourceType.JSON:
+            return pd.read_json(source, nrows=nrows, encoding=encoding)
+        elif source_type == DataSourceType.PARQUET:
+            return pd.read_parquet(source)
+        elif source_type == DataSourceType.FEATHER:
+            return pd.read_feather(source)
+        elif source_type == DataSourceType.HDF5:
+            return pd.read_hdf(source)
+        elif source_type == DataSourceType.STATA:
+            df = pd.read_stata(source)
+            if nrows:
+                return df.head(nrows)
+            return df
+        elif source_type == DataSourceType.SAS:
+            return pd.read_sas(source)
+        elif source_type == DataSourceType.SPSS:
+            try:
+                import pyreadstat
+                df, meta = pyreadstat.read_sav(source)
+                return df
+            except ImportError:
+                raise DataLoaderError("pyreadstat required for SPSS files")
+        else:
+            raise FileFormatError(f"Unsupported format for full load: {source_type}")
+    
+    def _load_chunked(self, source: str, source_type: DataSourceType, chunk_size: int, encoding: Optional[str]) -> pd.DataFrame:
+        """Load data in chunks and combine"""
+        encoding = encoding or self.config.default_encoding
+        chunks = []
+        
+        if source_type == DataSourceType.CSV:
+            for chunk in pd.read_csv(source, chunksize=chunk_size, encoding=encoding):
+                chunks.append(chunk)
+        elif source_type == DataSourceType.JSON:
+            for chunk in pd.read_json(source, lines=True, chunksize=chunk_size, encoding=encoding):
+                chunks.append(chunk)
+        else:
+            raise FileFormatError(f"Chunked loading not supported for: {source_type}")
+        
+        return pd.concat(chunks, ignore_index=True)
+    
+    def _load_streaming(self, source: str, source_type: DataSourceType, chunk_size: int, encoding: Optional[str]) -> Iterator[pd.DataFrame]:
+        """Create streaming iterator"""
+        encoding = encoding or self.config.default_encoding
+        
+        if source_type == DataSourceType.CSV:
+            return pd.read_csv(source, chunksize=chunk_size, encoding=encoding)
+        elif source_type == DataSourceType.JSON:
+            return pd.read_json(source, lines=True, chunksize=chunk_size, encoding=encoding)
+        else:
+            raise FileFormatError(f"Streaming not supported for: {source_type}")
+    
+    def _load_lazy(self, source: str, source_type: DataSourceType, encoding: Optional[str]) -> Any:
+        """Create lazy loading wrapper"""
+        # For now, return full load with warning
+        self.logger.warning("Lazy loading not fully implemented, using full load")
+        return self._load_full(source, source_type, None, encoding)
+    
+    def _generate_metadata(self, data: Any, source: str, source_type: DataSourceType) -> Dict[str, Any]:
+        """Generate metadata about loaded data"""
+        if isinstance(data, pd.DataFrame):
+            return {
+                'rows': len(data),
+                'columns': len(data.columns),
+                'column_names': list(data.columns),
+                'dtypes': {col: str(dtype) for col, dtype in data.dtypes.items()},
+                'memory_usage_mb': data.memory_usage(deep=True).sum() / (1024 * 1024),
+                'file_size_mb': os.path.getsize(source) / (1024 * 1024)
+            }
+        else:
+            return {
+                'type': str(type(data)),
+                'file_size_mb': os.path.getsize(source) / (1024 * 1024)
+            }
+    
+    def _validate_schema_internal(self, data: pd.DataFrame, schema: Dict[str, Any]) -> bool:
+        """Internal schema validation"""
+        if 'columns' not in schema:
+            return True
+        
+        expected_columns = set(schema['columns'].keys())
+        actual_columns = set(data.columns)
+        
+        return expected_columns.issubset(actual_columns)
+    
+    def _validate_quality(self, data: pd.DataFrame, validation_rules: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+        """Validate data quality"""
+        quality_report = {
+            'total_rows': len(data),
+            'total_columns': len(data.columns),
+            'missing_values': data.isnull().sum().to_dict(),
+            'duplicate_rows': data.duplicated().sum(),
+            'quality_score': 1.0
+        }
+        
+        # Calculate quality score
+        missing_ratio = data.isnull().sum().sum() / (len(data) * len(data.columns)) if len(data) > 0 else 0
+        duplicate_ratio = quality_report['duplicate_rows'] / len(data) if len(data) > 0 else 0
+        
+        quality_score = 1.0 - (missing_ratio * 0.5 + duplicate_ratio * 0.5)
+        quality_report['quality_score'] = max(0.0, min(1.0, quality_score))
+        
+        # Add issues list
+        issues = []
+        if missing_ratio > 0.1:
+            issues.append(f"High missing value ratio: {missing_ratio:.2%}")
+        if duplicate_ratio > 0.05:
+            issues.append(f"High duplicate ratio: {duplicate_ratio:.2%}")
+        
+        quality_report['issues'] = issues
+        
+        return quality_report
+