sovai 0.2.35__tar.gz → 0.2.36__tar.gz

{sovai-0.2.35 → sovai-0.2.36}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sovai
-Version: 0.2.35
+Version: 0.2.36
 Summary: python for asset management
 License: Proprietary License
 Author: Derek Snow

{sovai-0.2.35 → sovai-0.2.36}/pyproject.toml

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

{sovai-0.2.35 → sovai-0.2.36}/sovai/get_data.py

@@ -20,8 +20,11 @@ from sovai.errors.sovai_errors import InvalidInputData
 from sovai.utils.converter import convert_data2df
 from sovai.utils.stream import stream_data, stream_data_pyarrow
 from sovai.utils.datetime_formats import datetime_format
+
 from sovai.utils.client_side import client_side_frame
 from sovai.utils.client_side_s3 import load_frame_s3
+from sovai.utils.client_side_s3_part_high import load_frame_s3_paritioned_high
+
 
 ## Note run these for TEST - IT COULD BE HERE some silent problems
 try:
@@ -296,6 +299,9 @@ def normalize_endpoint(endpoint):
 
 ## If you also want access by full file, please also add to endpoint_to_ticker
 
+
+
+# Define your endpoint sets
 client_side_endpoints_gcs = {
     "ratios/relative",
     "market/prices",
@@ -306,7 +312,6 @@ client_side_endpoints_gcs = {
     "lobbying/public",
 }
 
-## If you also want access by full file, please also add to endpoint_to_ticker
 client_side_endpoints_s3 = {
     "sec/10k",
     "trials/predict",
@@ -314,9 +319,18 @@ client_side_endpoints_s3 = {
     "trials/all/predict",
     "trials/all/decribe",
     "trials/all",
+}
 
+client_side_endpoints_s3_part_high = {
+    "patents/applications"
 }
 
+# Create a list of handlers
+handlers = [
+    (client_side_endpoints_gcs, client_side_frame, "Grabbing GCS client side"),
+    (client_side_endpoints_s3, load_frame_s3, "Grabbing S3 client side"),
+    (client_side_endpoints_s3_part_high, load_frame_s3_paritioned_high, "Grabbing S3 Partitioned High client side"),
+]
 
 
 # In your main function or wherever this logic is implemented:
@@ -604,26 +618,22 @@ def data(
 
     normalized_endpoint = normalize_endpoint(endpoint)
 
-    if (
-        normalized_endpoint in client_side_endpoints_gcs
-        and tickers is not None
-        and frequency is None
-    ):
-        verbose_mode.vprint("Grabbing gcs client side")
-        _query_cache[cache_key] = client_side_frame(
-            normalized_endpoint, tickers, columns, start_date, end_date
-        )
-        return _query_cache[cache_key]
+    # ACTAULLY I DON'T CARE, TICKERS CAN BE NONE FOR PATENTS
+
+    # Iterate through the handlers to find a matching handler
+    for endpoint_set, handler_func, message in handlers:
+        if (
+            normalized_endpoint in endpoint_set and
+            (tickers is not None or start_date is not None) and
+            frequency is None
+        ):
+            verbose_mode.vprint(message)
+            _query_cache[cache_key] = handler_func(
+                normalized_endpoint, tickers, columns, start_date, end_date
+            )
+            return _query_cache[cache_key]
+
 
-    elif (
-        normalized_endpoint in client_side_endpoints_s3
-        and tickers is not None
-        and frequency is None
-    ):
-        verbose_mode.vprint("Grabbing s3 client side")
-        _query_cache[cache_key] = load_frame_s3(
-            normalized_endpoint, tickers, columns, start_date, end_date)
-        return _query_cache[cache_key]
     try:
         if tickers is not None and not is_all(tickers):
             params = ticker_mapper(params,verbose)

{sovai-0.2.35 → sovai-0.2.36}/sovai/tools/authentication/authentication.py

@@ -7,7 +7,7 @@ from pyarrow.fs import S3FileSystem as S3File
 import s3fs
 
 
-def get_s3fs_filesystem_pickle(storage_provider: str, verbose: bool = False):
+def get_s3_filesystem_pickle(storage_provider: str, verbose: bool = False):
     """
     Retrieve the S3FileSystem for the specified storage provider using pickle.
     
@@ -70,8 +70,8 @@ def get_s3fs_filesystem_json(storage_provider: str, verbose: bool = False):
             print(f"Successfully retrieved S3 credentials for {storage_provider}")
 
         s3fsa = s3fs.S3FileSystem(
-            access_key=cred['access_key'],
-            secret_key=cred['secret_key'],
+            key=cred['access_key'],
+            secret=cred['secret_key'],
             client_kwargs={'endpoint_url': cred['endpoint']}
         )
 

sovai-0.2.36/sovai/utils/client_side_s3_part_high.py

@@ -0,0 +1,321 @@
+import pyarrow.dataset as ds
+from pyarrow.fs import S3FileSystem
+import pyarrow as pa
+import pandas as pd
+from functools import lru_cache
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from sovai.tools.authentication import authentication
+import os
+import pandas as pd
+import pyarrow as pa
+import pyarrow.dataset as ds
+import s3fs
+import logging
+import datetime
+from functools import lru_cache
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
+from pyarrow.fs import S3FileSystem
+
+
+# Try to import CustomDataFrame, use regular DataFrame if not available
+try:
+    from sovai.extensions.pandas_extensions import CustomDataFrame
+    HAS_CUSTOM_DATAFRAME = True
+except ImportError:
+    HAS_CUSTOM_DATAFRAME = False
+    CustomDataFrame = pd.DataFrame  # Fallback to regular DataFrame
+
+
+@lru_cache(maxsize=2)
+def get_cached_s3_filesystem(storage_provider):
+    return authentication.get_s3_filesystem_pickle(storage_provider, verbose=True)
+
+@lru_cache(maxsize=2)
+def get_cached_s3fs_filesystem(storage_provider):
+    return authentication.get_s3fs_filesystem_json(storage_provider, verbose=True)
+
+
+# Configure logging to output to both console and a log file
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("data_operations.log"),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger(__name__)
+
+# -------------------------------
+# Helper Functions
+# -------------------------------
+
+def construct_s3_path(type_name, provider, partition_type, ticker=None, publish_date=None, year=None, all_dates_after=False):
+    """
+    Construct the S3 path based on the partitioning scheme.
+
+    Parameters:
+    - type_name (str): Type of data (e.g., 'applications')
+    - provider (str): 'wasabi' or 'digitalocean'
+    - partition_type (str): 'date' or 'ticker'
+    - ticker (str, optional): Ticker symbol
+    - publish_date (str, optional): Publish date in 'YYYY-MM-DD' format
+    - year (int, optional): Year extracted from publish_date
+    - all_dates_after (bool, optional): If True, scan all date partitions after a certain date
+
+    Returns:
+    - str: Constructed S3 path
+    """
+    if provider=="digitalocean":
+        bucket ="sovai/sovai-patents-export"
+    else:
+        bucket = "sovai-patents-export"
+        
+
+    if partition_type == "date":
+        if all_dates_after:
+            # Return the parent directory to scan all date partitions
+            path = f"{type_name}/date/"
+        else:
+            if not publish_date:
+                raise ValueError("publish_date must be provided for date partitioning.")
+            path = f"{type_name}/date/date_partitioned={publish_date}/"
+    elif partition_type == "ticker":
+        if not ticker or not year:
+            raise ValueError("Both ticker and year must be provided for ticker partitioning.")
+        path = f"{type_name}/ticker/ticker_partitioned={ticker}/year={year}/"
+    else:
+        raise ValueError("Invalid partition_type. Choose 'date' or 'ticker'.")
+
+    return f"{bucket}/{path}"
+
+
+def list_date_partitions(provider, type_name):
+    s3_path = construct_s3_path(type_name, provider, partition_type="date", all_dates_after=True)
+    fs = get_cached_s3fs_filesystem(provider)
+    try:
+        partitions = fs.ls(s3_path)
+        dates = []
+        for partition in partitions:
+            basename = os.path.basename(partition.rstrip('/'))
+            if basename.startswith('date_partitioned='):
+                date_str = basename.split('=')[1]
+                try:
+                    date_obj = datetime.datetime.strptime(date_str, '%Y-%m-%d').date()
+                    dates.append((date_obj, partition))
+                except ValueError:
+                    logger.warning(f"Unable to parse date from partition {partition}")
+        return dates
+    except Exception as e:
+        logger.error(f"Error listing date partitions in {s3_path} ({provider}): {e}")
+        return []
+
+def load_parquet_from_s3(s3_path, provider, columns=None, start_date=None):
+    """
+    Load a Parquet file from S3 into a pandas DataFrame with optional date filtering.
+
+    Parameters:
+    - s3_path (str): S3 path to the Parquet file or directory
+    - provider (str): 'wasabi' or 'digitalocean'
+    - columns (list, optional): List of columns to select
+    - start_date (str, optional): Minimum date in 'YYYY-MM-DD' format to filter the data
+
+    Returns:
+    - pandas.DataFrame: Loaded DataFrame
+    """
+    
+    fs = get_cached_s3_filesystem(provider)
+    
+    try:
+        dataset = ds.dataset(s3_path, filesystem=fs, format='parquet')
+
+        if start_date:
+            # Convert start_date string to pyarrow date32 scalar
+            min_date_pa = pa.scalar(pd.to_datetime(start_date).date(), type=pa.date32())
+            # Use ds.field to construct the filter expression
+            filter_expr = ds.field('date') >= min_date_pa
+            dataset = dataset.filter(filter_expr)
+
+        table = dataset.to_table(columns=columns)
+
+        # Log the schema for debugging
+        logger.info(f"Schema for {s3_path} ({provider}):")
+        logger.info(table.schema)
+
+        df = table.to_pandas()
+
+        partition_cols = ['ticker', 'date']
+        existing_cols = [col for col in partition_cols if col in df.columns]
+        other_cols = [col for col in df.columns if col not in partition_cols]
+        df = df[existing_cols + other_cols]
+
+        return df
+    except Exception as e:
+        logger.error(f"Error loading data from {s3_path} ({provider}): {e}")
+        return pd.DataFrame()
+
+# -------------------------------
+# Loading Functions
+# -------------------------------
+
+def load_data_by_ticker(
+    type_name,
+    providers=['wasabi', 'digitalocean'],
+    tickers=None,
+    start_date=None,
+    columns=None,
+    max_workers=4
+):
+    """
+    Load data for specified tickers and/or dates from S3 storage providers.
+
+    Parameters:
+    - type_name (str): Type of data (e.g., 'applications', 'prediction_all', etc.)
+    - providers (list, optional): List of storage providers to load from ('wasabi', 'digitalocean')
+    - tickers (str or list, optional): Ticker symbol(s) to load data for
+    - start_date (str, optional): Minimum publishDate to filter data by (format 'YYYY-MM-DD') (only for date partitioning)
+    - columns (list, optional): List of columns to select
+    - max_workers (int, optional): Number of parallel threads
+
+    Returns:
+    - pandas.DataFrame: Concatenated DataFrame containing the loaded data
+    """
+    date_tasks = []
+    ticker_tasks = []
+    results = []
+
+    # Normalize input parameters
+    if tickers and isinstance(tickers, str):
+        tickers = [tickers]
+    if not tickers:
+        tickers = []
+    
+    # If both tickers and min_date are provided
+    if tickers and start_date:
+        try:
+            min_date_obj = datetime.datetime.strptime(start_date, '%Y-%m-%d').date()
+            min_year = min_date_obj.year
+            logger.info(f"Minimum date provided: {min_date_obj} (Year: {min_year})")
+        except ValueError:
+            logger.error(f"Invalid min_date format: {start_date}. Expected 'YYYY-MM-DD'.")
+            return pd.DataFrame()
+
+        for provider in providers:
+            for ticker in tickers:
+                # Determine the range of years to load based on start_date
+                # Assuming data starts from 2011
+                start_year = min_year
+                current_year = datetime.datetime.now().year
+                for year in range(start_year, current_year + 1):
+                    s3_path = construct_s3_path(
+                        type_name, provider, partition_type="ticker",
+                        ticker=ticker, year=year
+                    )
+                    ticker_tasks.append((s3_path, provider, columns, None))  # No additional filtering here
+
+    # If only min_date is provided
+    elif start_date and not tickers:
+        try:
+            min_date_obj = datetime.datetime.strptime(start_date, '%Y-%m-%d').date()
+            logger.info(f"Minimum date provided: {min_date_obj}")
+        except ValueError:
+            logger.error(f"Invalid start_date format: {start_date}. Expected 'YYYY-MM-DD'.")
+            return pd.DataFrame()
+
+        for provider in providers:
+            # List all date-based partitions
+            date_partitions = list_date_partitions(provider, type_name)
+            for date_obj, partition_path in date_partitions:
+                if date_obj >= min_date_obj:
+                    date_tasks.append((partition_path, provider, columns, None))
+
+    # If only tickers are provided
+    elif tickers and not start_date:
+        for provider in providers:
+            for ticker in tickers:
+                # Assuming data starts from 2011
+                start_year = 2010 ## Change this to 2008 eventually
+                current_year = datetime.datetime.now().year
+                for year in range(start_year, current_year + 1):
+                    s3_path = construct_s3_path(
+                        type_name, provider, partition_type="ticker",
+                        ticker=ticker, year=year
+                    )
+                    ticker_tasks.append((s3_path, provider, columns, None))
+
+    else:
+        logger.warning("No tickers or start_date provided. Please provide at least one.")
+        return pd.DataFrame()
+
+    # Combine all tasks
+    all_tasks = ticker_tasks + date_tasks
+
+    if not all_tasks:
+        logger.warning("No tasks to process. Exiting.")
+        return pd.DataFrame()
+
+    def load_task(task):
+        s3_path, provider, columns, ticker_filter = task
+        df = load_parquet_from_s3(s3_path, provider, columns, start_date=start_date if ticker_filter is None else None)
+        if not df.empty and ticker_filter:
+            if 'ticker' in df.columns:
+                df = df[df['ticker'].isin(ticker_filter)]
+            else:
+                logger.warning(f"'ticker' column not found in data from {s3_path} ({provider}). Skipping ticker filter.")
+        return df
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        future_to_task = {executor.submit(load_task, task): task for task in all_tasks}
+        for future in tqdm(as_completed(future_to_task), total=len(all_tasks), desc="Loading data"):
+            df = future.result()
+            if not df.empty:
+                results.append(df)
+
+    if results:
+        final_df = pd.concat(results, ignore_index=True)
+        logger.info(f"Successfully loaded data with {len(final_df)} records.")
+    else:
+        final_df = pd.DataFrame()
+        logger.warning("No data found.")
+
+    return final_df
+
+# -------------------------------
+# Example Usage
+# -------------------------------
+
+
+def load_frame_s3_paritioned_high(endpoint, tickers=None, columns=None, start_date=None, end_date=None):
+
+    endpoint_config = {
+        "patents/applications": {
+            "bucket": {
+                "digitalocean": "sovai/sovai-patents-export",
+                "wasabi": "sovai-patents-export"
+            },
+            "type": "applications"
+        },
+    }
+    
+    if endpoint not in endpoint_config:
+        raise ValueError(f"Invalid endpoint: {endpoint}")
+    
+    config = endpoint_config[endpoint]
+
+
+    loaded_df_both = load_data_by_ticker(
+        type_name=config["type"],
+        providers=['digitalocean'],
+        tickers=tickers,
+        start_date=start_date,
+        columns=columns,
+        max_workers=8  # Adjust based on your system's capabilities
+    )
+
+    
+    if HAS_CUSTOM_DATAFRAME:
+        return CustomDataFrame(loaded_df_both)
+    else:
+        return loaded_df_both  # Returns a regular pandas DataFrame if CustomDataFrame is not available
+

sovai-0.2.35/sovai/utils/client_side_s3 copy.py

@@ -1,206 +0,0 @@
-import pyarrow.dataset as ds
-from pyarrow.fs import S3FileSystem
-import pyarrow as pa
-import pandas as pd
-from functools import lru_cache
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from sovai.tools.authentication import authentication
-
-
-# Try to import CustomDataFrame, use regular DataFrame if not available
-try:
-    from sovai.extensions.pandas_extensions import CustomDataFrame
-    HAS_CUSTOM_DATAFRAME = True
-except ImportError:
-    HAS_CUSTOM_DATAFRAME = False
-    CustomDataFrame = pd.DataFrame  # Fallback to regular DataFrame
-
-
-@lru_cache(maxsize=2)
-def get_cached_s3_filesystem(storage_provider):
-    return authentication.get_s3_filesystem_pickle(storage_provider, verbose=True)
-
-def process_partition(storage_provider, base_path, identifier_column, identifier=None, columns=None, filters=None):
-    s3 = get_cached_s3_filesystem(storage_provider)
-    
-    if identifier:
-        base_path += f"/{identifier_column}={identifier}"
-    
-    dataset = ds.dataset(base_path, filesystem=s3, format='parquet')
-    
-    if filters:
-        operator_map = {
-            '==': lambda a, b: a == b,
-            '!=': lambda a, b: a != b,
-            '<': lambda a, b: a < b,
-            '<=': lambda a, b: a <= b,
-            '>': lambda a, b: a > b,
-            '>=': lambda a, b: a >= b,
-            'in': lambda a, b: a.isin(b),
-            'not in': lambda a, b: ~a.isin(b),
-            'like': lambda a, b: a.match_substring(b),
-        }
-        schema = dataset.schema
-        filter_expr = None
-        for col, op, val in filters:
-            if op not in operator_map:
-                raise ValueError(f"Unsupported operator '{op}' in filters.")
-            field = ds.field(col)
-            field_type = schema.field(col).type
-            if pa.types.is_timestamp(field_type) or pa.types.is_date(field_type):
-                val = pd.to_datetime(val)
-            elif pa.types.is_integer(field_type):
-                val = int(val)
-            elif pa.types.is_floating(field_type):
-                val = float(val)
-            condition = operator_map[op](field, val)
-            filter_expr = condition if filter_expr is None else filter_expr & condition
-    else:
-        filter_expr = None
-    
-    table = dataset.to_table(columns=columns, filter=filter_expr, use_threads=True)
-    df = table.to_pandas(use_threads=True)
-    
-    if identifier:
-        df[identifier_column] = identifier
-
-        df = df[[identifier_column] + [col for col in df.columns if col != identifier_column]]
-
-    
-    return df
-
-import numpy as np
-
-def client_side_s3_frame(config, identifiers, columns, start_date, end_date):
-    storage_provider = "digitalocean"  # Default to DigitalOcean
-    base_path = config["storage_provider"][storage_provider]
-    identifier_column = config.get("identifier_column")
-
-    filters = []
-    if start_date:
-        filters.append(('date', '>=', start_date))
-    if end_date:
-        filters.append(('date', '<=', end_date))
-
-    if not identifiers:
-        # Load the entire database
-        return process_partition(storage_provider, base_path, identifier_column, columns=columns, filters=filters)
-    else:
-        identifiers = map_identifier(identifiers, input_identifier='ticker', output_identifier=identifier_column)
-
-    # Use numpy.isscalar to check if identifiers is a single value
-    if np.isscalar(identifiers):
-        # Single identifier
-        identifier = identifiers
-        return process_partition(storage_provider, base_path, identifier_column, identifier=identifier, columns=columns, filters=filters)
-    
-    # Multiple identifiers
-    with ThreadPoolExecutor() as executor:
-        futures = []
-        for identifier in identifiers:
-            futures.append(executor.submit(process_partition, storage_provider, base_path, identifier_column, identifier, columns, filters))
-        
-        results = []
-        for future in as_completed(futures):
-            results.append(future.result())
-    
-    df = pd.concat(results, ignore_index=True)
-    return df
-
-    # if 'filing_date' in df.columns:
-    #     df['date'] = pd.to_datetime(df['filing_date'])
-    # elif 'date' not in df.columns:
-    #     raise ValueError("Neither 'filing_date' nor 'date' column found in the data")
-    
-    # df.set_index([identifier_column, 'date'], inplace=True)
-    # df.sort_index(inplace=True)
-    
-    return df
-
-import pandas as pd
-import numpy as np
-
-def map_identifier(input_values, input_identifier='ticker', output_identifier='cik', trailing_zero_removal=True):
-    """
-    Map from one identifier to another using the df_codes DataFrame.
-    Handles string, single-item list, and multiple-item lists as input.
-    Uses vectorization for improved performance.
-    
-    :param input_values: The input identifier value(s) to map from (string, list, or array-like)
-    :param input_identifier: The type of input identifier (default: 'ticker')
-    :param output_identifier: The type of output identifier (default: 'cik')
-    :param trailing_zero_removal: If True, convert CIK to int (default: True)
-    :return: Mapped identifier value(s), same type as input
-    """
-    df_codes = pd.read_parquet("data/codes.parq")
-
-    # Check if the identifier columns exist in df_codes
-    if input_identifier not in df_codes.columns or output_identifier not in df_codes.columns:
-        raise ValueError(f"Identifier '{input_identifier}' or '{output_identifier}' not found in the DataFrame")
-    
-    # Convert input to numpy array for consistent handling
-    if isinstance(input_values, str):
-        input_array = np.array([input_values])
-        input_is_string = True
-    else:
-        input_array = np.array(input_values)
-        input_is_string = False
-    
-    # Ensure all input values are strings
-    input_array = np.char.strip(input_array.astype(str))
-    
-    # Create a dictionary for fast lookup
-    mapping_dict = dict(zip(df_codes[input_identifier], df_codes[output_identifier]))
-    
-    # Vectorized mapping
-    result = np.vectorize(mapping_dict.get)(input_array)
-    
-    # Convert to int if trailing_zero_removal is True and output is CIK
-    if trailing_zero_removal and output_identifier.lower() == 'cik':
-        result = np.where(result == None, 0, result)  # Replace None with 0
-        result = result.astype(np.int32)  # Convert to int64
-    
-    # Return result in the same format as input
-    if input_is_string:
-        return result[0]
-    elif len(input_array) == 1:
-        return result[0]
-    else:
-        return result.tolist()
-
-
-def load_frame_s3(endpoint, tickers=None, columns=None, start_date=None, end_date=None):
-
-    endpoint_config = {
-
-        "patents/applications": {
-            "storage_provider": {
-                "digitalocean": "sovai/sovai-patents-export/applications",
-                "wasabi": "sovai-patents-export/applications"
-            },
-            "identifier_columns": ["date_partitioned", "ticker_partitioned"],
-            "partitioned": True,
-        },
-
-    }
-    
-    if endpoint not in endpoint_config:
-        raise ValueError(f"Invalid endpoint: {endpoint}")
-    
-    config = endpoint_config[endpoint]
-
-    
-    df_frame = client_side_s3_frame(
-        config,
-        tickers,
-        columns,
-        start_date,
-        end_date
-    )
-
-    
-    if HAS_CUSTOM_DATAFRAME:
-        return CustomDataFrame(df_frame)
-    else:
-        return df_frame  # Returns a regular pandas DataFrame if CustomDataFrame is not available
-