sovai 0.2.45__tar.gz → 0.2.46__tar.gz

{sovai-0.2.45 → sovai-0.2.46}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sovai
-Version: 0.2.45
+Version: 0.2.46
 Summary: python for asset management
 License: Proprietary License
 Author: Derek Snow

{sovai-0.2.45 → sovai-0.2.46}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sovai"
-version = "0.2.45"
+version = "0.2.46"
 description = "python for asset management"
 authors = ["Derek Snow <d.snow@outlook.com>"]
 license = "Proprietary License"

{sovai-0.2.45 → sovai-0.2.46}/sovai/get_data.py

@@ -308,7 +308,9 @@ ENDPOINT_TO_TICKER = {
    "/news/pronouns": "",
    "/news/word_count": "",
    "/news/sentiment_score": "", # was None before?
+
    "/insider/trading": "",
+
    "/wikipedia/views": "",
    "/accounting/weekly": "",
    "/visas/h1b": "",
@@ -346,6 +348,8 @@ ENDPOINT_TO_TICKER = {
     
     "/clinical/trials": "",
 
+
+
     "/spending/awards": "",
     "/spending/compensation": "",
     "/spending/competition": "",
@@ -399,6 +403,8 @@ CLIENT_SIDE_ENDPOINTS_S3_PART_HIGH = {
 
    "accounting/weekly",
 
+   "insider/trading",
+
    "ratios/normal",
    "ratios/relative",
 
@@ -414,15 +420,42 @@ CLIENT_SIDE_ENDPOINTS_S3_PART_HIGH = {
 
    "breakout",
 
-   "corprisk/risks",
+    "corprisk/risks",
+    "corprisk/accounting",
+    "corprisk/events",
+    "corprisk/misstatements",
 
    "visas/h1b",
 
-   "wikipedia/views"
+   "wikipedia/views",
+
+    "short/volume",
+    "short/maker",
+    "short/over_shorted",
+
+    "institutional/trading"
 
 }
 
 
+
+ENDPOINT_ALIASES = {
+    "clinical/predict": {
+        "target_endpoint": "clinical/trials",
+        "columns": [
+            'ticker', 'date', 'source', 'subsidiary', 'sponsor', 
+            'trial_id', 'official_title', 'success_prediction', 
+            'economic_effect', 'duration_prediction', 'success_composite'
+        ]
+    }
+    # --- Add other aliases here in the future ---
+    # "some/alias": {
+    #     "target_endpoint": "real/endpoint",
+    #     "columns": ["col1", "col2", "date", "ticker"] # Ensure essential columns are included
+    # }
+}
+
+
 # Define handlers for different endpoint types
 ENDPOINT_HANDLERS = [
    (CLIENT_SIDE_ENDPOINTS_GCS, client_side_frame, "Grabbing GCS client side"),
@@ -1095,6 +1128,38 @@ def data(
     params = params or {}
     params = map_synonyms(params)
 
+    # --- MODIFIED SECTION: Alias Handling ---
+    original_endpoint_request = endpoint # Keep track for logging
+    forced_columns = None # Will store columns if an alias forces them
+    normalized_req_endpoint = normalize_endpoint(endpoint) 
+
+    if normalized_req_endpoint in ENDPOINT_ALIASES:
+        alias_config = ENDPOINT_ALIASES[normalized_req_endpoint]
+        target_endpoint = alias_config['target_endpoint']
+        forced_columns_list = alias_config['columns']
+        verbose_mode.log(f"Endpoint alias detected: '{original_endpoint_request}' maps to '{target_endpoint}' with predefined columns.")
+        
+        endpoint = target_endpoint # *Overwrite* endpoint for the rest of the function
+        columns = forced_columns_list # *Overwrite* columns parameter
+        
+        # Ensure the 'columns' value in the main 'params' dict is also updated
+        # This is crucial for cache key generation and potentially for client-side handlers
+        # Convert list to string if needed for _prepare_params consistency
+        params['columns'] = ",".join(forced_columns_list) 
+
+    else:
+        # If not an alias, handle the columns parameter normally
+        # Convert list to string if provided as list initially
+        if isinstance(columns, list):
+             columns_str = ",".join(columns)
+        else:
+             columns_str = columns # It's already a string or None
+        # Update params dict if columns were passed directly to the function
+        if columns is not None:
+            params['columns'] = columns_str
+    # --- END MODIFIED SECTION ---
+
+
     # Force full_history if date range is specified
     if start_date is not None or end_date is not None:
         full_history = True
@@ -1102,6 +1167,14 @@ def data(
 
     # Update params with prepared values
     try:
+
+        # Make sure to use the potentially updated 'columns' value from params dict
+        # if it exists, otherwise use the function argument (which might be None or overridden)
+        current_columns = params.get("columns", columns) 
+        # Convert back to string if it was forced as a list earlier
+        if isinstance(current_columns, list):
+            current_columns = ",".join(current_columns)
+
         params.update(
             _prepare_params(
                 tickers=params.get("tickers", tickers),
@@ -1150,6 +1223,10 @@ def data(
 
         # Try client-side handlers for specific endpoints
         handler_result = None
+
+        # Get the columns value to pass to handlers (use the string version from params)
+        handler_columns_arg = params.get('columns') 
+
         for endpoint_set, handler_func, message in ENDPOINT_HANDLERS:
             if (
                 normalized_endpoint in endpoint_set and
@@ -1160,7 +1237,7 @@ def data(
                 verbose_mode.log(f"Calling client-side handler with tickers={tickers}, start_date={start_date}, end_date={end_date}")
                 try:
                     handler_result = handler_func(
-                        normalized_endpoint, tickers, columns, start_date, end_date
+                        normalized_endpoint, tickers, handler_columns_arg, start_date, end_date
                     )
                     _query_cache[cache_key] = handler_result
                     result_data = handler_result
@@ -1241,11 +1318,15 @@ def data(
                         logger.error(error_msg)
                         raise Exception(error_msg)
                     
+                    filter_columns_arg = columns
+                    if isinstance(columns, str): # filter_data might prefer list
+                        filter_columns_arg = [c.strip() for c in columns.split(',')] if columns else None
+
                     # Apply filters
                     verbose_mode.log("Applying filters to data")
                     data_result = filter_data(
                         data_result, 
-                        columns=columns, 
+                        columns=filter_columns_arg, 
                         start_date=start_date, 
                         end_date=end_date, 
                         use_polars=use_polars
@@ -1331,6 +1412,22 @@ def data(
 
     result_data = process_dataframe(result_data)
 
+
+    # --- Final Check: Ensure alias columns are applied if data wasn't filtered before ---
+    # This is a safety net in case a data path (e.g., direct API return without 'all tickers')
+    # didn't explicitly filter columns based on the overridden parameter.
+    if forced_columns is not None and result_data is not None and isinstance(result_data, (pd.DataFrame, pl.DataFrame)):
+         # Check if columns need filtering (some might be missing if already filtered)
+         current_cols = list(result_data.columns)
+         cols_to_keep = [col for col in forced_columns_list if col in current_cols]
+         if set(cols_to_keep) != set(current_cols):
+             verbose_mode.log(f"Applying final column filter for alias: {cols_to_keep}")
+             if isinstance(result_data, pd.DataFrame):
+                 result_data = result_data[cols_to_keep]
+             elif isinstance(result_data, pl.DataFrame):
+                 result_data = result_data.select(cols_to_keep)
+                 
+
     return result_data
 
 

{sovai-0.2.45 → sovai-0.2.46}/sovai/utils/client_side_s3_part_high.py

@@ -32,16 +32,23 @@ from sovai.tools.authentication import authentication
 # Logging Configuration
 # =========================================================================
 
+# Near top of load_frame_s3_partitioned_high.py
 logger = logging.getLogger(__name__)
 if not logger.handlers:
-    logging.basicConfig(
-        level=logging.INFO,
-        format='%(asctime)s - %(levelname)s - %(message)s',
-        handlers=[
-            logging.FileHandler("data_operations.log"),
-            logging.StreamHandler()
-        ]
-    )
+    log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - [%(funcName)s] - %(message)s') # Added function name
+
+    # File Handler
+    file_handler = logging.FileHandler("data_operations.log")
+    file_handler.setFormatter(log_formatter)
+
+    # Console Handler (add this)
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(log_formatter)
+
+    logger.addHandler(file_handler)
+    logger.addHandler(console_handler) # Add console output
+    logger.setLevel(logging.INFO) # Set level (INFO or DEBUG)
+    # To see DEBUG level messages, set logger.setLevel(logging.DEBUG)
 
 # =========================================================================
 # Filesystem Access
@@ -295,10 +302,13 @@ class PartitionFinder:
 # =========================================================================
 # Data Loading
 # =========================================================================
+# =========================================================================
+# Data Loading
+# =========================================================================
 
 class DataLoader:
     """Core data loading functionality with filtering and parallelism."""
-    
+
     @staticmethod
     def load_partition(
         path: str,
@@ -308,77 +318,127 @@ class DataLoader:
         end_date: Optional[str] = None
     ) -> pd.DataFrame:
         """
-        Load and filter data from a single partition.
-        
+        Load and filter data from a single partition. Added logging for debugging column selection.
+
         Args:
             path: S3 path to the partition
             ticker_filter: Optional list of tickers to filter by
             columns: Optional list of columns to load
             start_date: Optional start date for filtering
             end_date: Optional end date for filtering
-            
+
         Returns:
             Filtered pandas DataFrame
         """
         fs = get_s3_filesystem()
-        
+        logger.info(f"--- Loading Partition Start: {path} ---")
+        logger.debug(f"Parameters received: ticker_filter={ticker_filter}, columns={columns}, start_date={start_date}, end_date={end_date}")
+
         try:
             # Create dataset
-            dataset = ds.dataset(path, filesystem=fs, format='parquet')
-            
+            logger.debug(f"[{path}] Creating dataset object...")
+            dataset = ds.dataset(path, filesystem=fs, format='parquet', partitioning="hive") # Assuming Hive partitioning
+            logger.info(f"[{path}] Initial dataset schema: {dataset.schema}")
+
             # Build filter expressions
             filters = DataLoader._build_date_filters(start_date, end_date)
-            
-            # Apply filters to dataset
+            logger.debug(f"[{path}] Built date filters: {filters}")
+
+            # Apply filters to dataset before reading columns
             if filters:
                 filter_expr = filters[0]
                 for f in filters[1:]:
                     filter_expr = filter_expr & f
-                dataset = dataset.filter(filter_expr)
-            
-            # Load data
-            table = dataset.to_table(columns=columns, use_threads=True)
+                try:
+                    logger.info(f"[{path}] Applying filter expression: {filter_expr}")
+                    dataset = dataset.filter(filter_expr)
+                    # Note: Checking row count here might be slow as it forces some reading
+                    # fragment_count = len(list(dataset.get_fragments()))
+                    # logger.info(f"[{path}] Filter applied. Found {fragment_count} fragments. Schema after filtering: {dataset.schema}")
+                    logger.info(f"[{path}] Filter applied. Schema after filtering: {dataset.schema}")
+                except Exception as filter_e:
+                    logger.error(f"[{path}] Error applying filters: {filter_e}", exc_info=True)
+                    logger.info(f"--- Loading Partition Failed (Filter Error): {path} ---")
+                    return pd.DataFrame()
+
+            # Load data using to_table
+            logger.info(f"[{path}] Calling dataset.to_table() with columns: {columns}")
+            try:
+                # Explicitly handle None columns case for logging clarity
+                cols_to_load = columns if columns else "ALL"
+                table = dataset.to_table(columns=columns, use_threads=True) # Pass None or list
+                logger.info(f"[{path}] dataset.to_table(columns={cols_to_load}) successful.")
+                logger.info(f"[{path}] Resulting Arrow Table shape: {table.shape}")
+                logger.debug(f"[{path}] Resulting Arrow Table columns: {table.column_names}")
+
+                # Check if table is empty AFTER the potentially selective read
+                if table.num_rows == 0:
+                    logger.warning(f"[{path}] Arrow Table is empty after to_table(columns={cols_to_load}). "
+                                   f"This might be expected if date filter removed all rows, "
+                                   f"or problematic if columns were requested but not found/read.")
+                    # Proceed to convert empty table to empty DataFrame
+
+            except Exception as to_table_e:
+                logger.error(f"[{path}] Error during dataset.to_table(columns={columns}): {to_table_e}", exc_info=True)
+                logger.error(f"[{path}] Dataset schema at time of error: {dataset.schema}")
+                logger.info(f"--- Loading Partition Failed (to_table Error): {path} ---")
+                return pd.DataFrame()
+
+            # Convert to pandas
+            logger.debug(f"[{path}] Converting Arrow Table to pandas DataFrame...")
             df = table.to_pandas()
-            
-            # Apply ticker filter if needed
+            logger.info(f"[{path}] Converted to pandas DataFrame. Shape: {df.shape}")
+            if df.empty and table.num_rows > 0:
+                 logger.warning(f"[{path}] Conversion to pandas resulted in empty DataFrame, but Arrow table had {table.num_rows} rows.")
+            logger.debug(f"[{path}] Pandas DataFrame columns: {list(df.columns)}")
+
+            # Apply ticker filter if needed (post-load)
             if ticker_filter and 'ticker' in df.columns:
+                logger.info(f"[{path}] Applying post-load ticker filter: {ticker_filter}")
+                initial_rows = len(df)
                 df = df[df['ticker'].isin(ticker_filter)]
-                
+                logger.info(f"[{path}] Ticker filter applied. Rows reduced from {initial_rows} to {len(df)}")
+            elif ticker_filter:
+                 logger.warning(f"[{path}] Ticker filter {ticker_filter} requested but 'ticker' column not found in loaded columns: {list(df.columns)}.")
+
+            logger.info(f"--- Loading Partition End: {path} --- Returning DataFrame shape: {df.shape}")
             return df
+
         except Exception as e:
-            logger.error(f"Error loading data from {path}: {e}")
+            # Catch-all for other errors (e.g., ds.dataset creation failed)
+            logger.error(f"Unhandled error processing partition {path}: {e}", exc_info=True)
+            logger.info(f"--- Loading Partition Failed (Unhandled Error): {path} ---")
             return pd.DataFrame()
-    
+
+
     @staticmethod
     def _build_date_filters(
-        start_date: Optional[str], 
+        start_date: Optional[str],
         end_date: Optional[str]
     ) -> List[ds.Expression]:
         """
-        Build PyArrow date filter expressions.
-        
-        Args:
-            start_date: Optional start date string
-            end_date: Optional end date string
-            
-        Returns:
-            List of PyArrow filter expressions
+        Build PyArrow date filter expressions. (Ensure 'date' column exists)
         """
         filters = []
-        
-        # Date filters
-        if start_date:
-            start_date_obj = pd.to_datetime(start_date).date()
-            start_date_pa = pa.scalar(start_date_obj, type=pa.date32())
-            filters.append(ds.field('date') >= start_date_pa)
-            
-        if end_date:
-            end_date_obj = pd.to_datetime(end_date).date()
-            end_date_pa = pa.scalar(end_date_obj, type=pa.date32())
-            filters.append(ds.field('date') <= end_date_pa)
-            
+        try:
+            # Date filters - Assuming 'date' column exists and is date32 or similar
+            if start_date:
+                start_date_obj = pd.to_datetime(start_date).date()
+                # Use pa.date32() explicitly assuming the Parquet schema uses it
+                start_date_pa = pa.scalar(start_date_obj, type=pa.date32())
+                filters.append(ds.field('date') >= start_date_pa)
+
+            if end_date:
+                end_date_obj = pd.to_datetime(end_date).date()
+                end_date_pa = pa.scalar(end_date_obj, type=pa.date32())
+                filters.append(ds.field('date') <= end_date_pa)
+        except Exception as e:
+             logger.error(f"Failed to build date filters (start={start_date}, end={end_date}): {e}", exc_info=True)
+             # Decide if you want to raise or return empty filters
+             # return [] # Return empty filters if parsing fails?
+
         return filters
-    
+
     @staticmethod
     def load_data_parallel(
         tasks: List[Tuple],
@@ -386,34 +446,42 @@ class DataLoader:
     ) -> List[pd.DataFrame]:
         """
         Process loading tasks in parallel with progress tracking.
-        
-        Args:
-            tasks: List of (path, ticker_filter, columns, start_date, end_date) tuples
-            max_workers: Maximum number of concurrent loading threads
-            
-        Returns:
-            List of loaded DataFrames
         """
         results = []
-        
+        if not tasks:
+            logger.warning("load_data_parallel received no tasks.")
+            return results
+
+        logger.info(f"Starting parallel load of {len(tasks)} tasks with max_workers={max_workers}")
+
         # Process tasks in parallel
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
-            future_to_task = {
-                executor.submit(DataLoader.load_partition, path, ticker_filter, columns, start_date, end_date): path
-                for path, ticker_filter, columns, start_date, end_date in tasks
+            # Prepare arguments properly for submit
+            future_to_path = {
+                executor.submit(DataLoader.load_partition, path, ticker_filter, task_columns, task_start_date, task_end_date): path
+                for path, ticker_filter, task_columns, task_start_date, task_end_date in tasks # Unpack task tuple
             }
-            
-            for future in tqdm(as_completed(future_to_task), total=len(tasks), desc="Loading data"):
+
+            # Use tqdm for progress bar
+            for future in tqdm(as_completed(future_to_path), total=len(tasks), desc="Loading partitions", unit="partition"):
+                path = future_to_path[future]
                 try:
-                    df = future.result()
-                    if not df.empty:
-                        results.append(df)
+                    df_partition = future.result() # Get result from load_partition
+                    if df_partition is not None and not df_partition.empty:
+                        logger.debug(f"Successfully loaded non-empty DataFrame from {path}, shape: {df_partition.shape}")
+                        results.append(df_partition)
+                    elif df_partition is not None:
+                         logger.info(f"Loaded empty DataFrame from {path}")
+                    else:
+                         # Should not happen if load_partition always returns a DataFrame
+                         logger.warning(f"Load task for {path} returned None.")
+
                 except Exception as e:
-                    path = future_to_task[future]
-                    logger.error(f"Task failed for {path}: {e}")
-                    
-        return results
+                    # This catches errors raised *after* submit, e.g., during future.result() processing if not caught inside load_partition
+                    logger.error(f"Task failed for partition {path} during result retrieval: {e}", exc_info=True)
 
+        logger.info(f"Parallel loading finished. Collected results from {len(results)} non-empty partitions.")
+        return results
 # =========================================================================
 # Main Loading Interface
 # =========================================================================
@@ -698,7 +766,55 @@ ENDPOINT_CONFIG = {
         "ticker_path": "sovai/sovai-wiki/views/partitioned/ticker",
         "date_path": "sovai/sovai-wiki/views/partitioned/date",
         "has_year": False
-    }
+    },
+
+    "insider/trading": {
+        "ticker_path": "sovai/sovai-insider/partitioned/ticker",
+        "date_path": "sovai/sovai-insider/partitioned/date",
+        "has_year": False
+    },
+
+    "corprisk/risks": {
+        "ticker_path": "sovai/sovai-flags/risks/partitioned/ticker",
+        "date_path": "sovai/sovai-flags/risks/partitioned/date",
+        "has_year": False
+    },
+    "corprisk/accounting": {
+        "ticker_path": "sovai/sovai-flags/accounting/partitioned/ticker",
+        "date_path": "sovai/sovai-flags/accounting/partitioned/date",
+        "has_year": False
+    },
+    "corprisk/events": {
+        "ticker_path": "sovai/sovai-flags/events/partitioned/ticker",
+        "date_path": "sovai/sovai-flags/events/partitioned/date",
+        "has_year": False
+    },
+    "corprisk/misstatements": {
+        "ticker_path": "sovai/sovai-flags/misstatements/partitioned/ticker",
+        "date_path": "sovai/sovai-flags/misstatements/partitioned/date",
+        "has_year": False
+    },
+    "short/volume": {
+        "ticker_path": "sovai/sovai-short/short_volume_weekly/partitioned/ticker",
+        "date_path": "sovai/sovai-short/short_volume_weekly/partitioned/date",
+        "has_year": False
+    },
+    ## probably more to do with short later.
+    "short/maker": {
+        "ticker_path": "sovai/sovai-short/over_shorted/partitioned/ticker",
+        "date_path": "sovai/sovai-short/over_shorted/partitioned/date",
+        "has_year": False
+    },
+    "short/over_shorted": {
+        "ticker_path": "sovai/sovai-short/over_shorted/partitioned/ticker",
+        "date_path": "sovai/sovai-short/over_shorted/partitioned/date",
+        "has_year": False
+    },
+    "institutional/trading": {
+        "ticker_path": "sovai/sovai-institutional/trading/partitioned/ticker",
+        "date_path": "sovai/sovai-institutional/trading/partitioned/date",
+        "has_year": False
+    },
 
 }
 
@@ -788,10 +904,15 @@ def load_frame_s3_partitioned_high(
         tickers=tickers,
         start_date=start_date,
         end_date=end_date,
-        columns=columns,
+        columns=[col.strip() for col in columns.split(',')] ,
         max_workers=8,
         post_process=post_process
     )
+
+    # print(columns)
+    # columns = [col.strip() for col in columns.split(',')] 
+
+    # df = df[columns] if columns else df
     
     # Convert to CustomDataFrame if available
     if HAS_CUSTOM_DATAFRAME: