sovai 0.2.25__tar.gz → 0.2.27__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {sovai-0.2.25 → sovai-0.2.27}/PKG-INFO +1 -1
- {sovai-0.2.25 → sovai-0.2.27}/pyproject.toml +15 -3
- {sovai-0.2.25 → sovai-0.2.27}/sovai/__init__.py +3 -1
- {sovai-0.2.25 → sovai-0.2.27}/sovai/get_data.py +19 -3
- sovai-0.2.27/sovai/tools/authentication/authentication.py +98 -0
- sovai-0.2.27/sovai/utils/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/client_side.py +1 -0
- sovai-0.2.27/sovai/utils/client_side_s3.py +212 -0
- {sovai-0.2.25 → sovai-0.2.27}/LICENSE +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/README.md +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/api_config.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/assets/dropdown.css +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/assets/features_mapping.parq +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/basic_auth.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/computations/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/computations/functions.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/errors/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/errors/sovai_errors.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/anomalies.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/ask_df_llm.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/change_point_generator.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/clustering.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/core_kshape.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/cum_sum_torch.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/cusum.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/dimensionality_reduction.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/feature_extraction.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/feature_importance.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/feature_neutralizer.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/filter_df.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/fractional_differencing.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/nowcasting.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/pairwise.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/pandas_extensions.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/parallel_functions.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/pfa_feature_selector.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/regime_change.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/regime_change_pca.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/shapley_global_importance.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/shapley_importance.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/signal_evaluation.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/technical_indicators.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/time_decomposition.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/extensions/weight_optimization.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/get_compute.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/get_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/get_reports.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/get_tools.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/accounting/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/accounting/accounting_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/allocation/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/allocation/allocation_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/bankruptcy/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/bankruptcy/bankruptcy_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/breakout/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/breakout/breakout_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/corp_risk/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/corp_risk/corp_risk_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/earnings_surprise/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/earnings_surprise/earnings_surprise_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/insider/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/insider/insider_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/institutional/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/institutional/institutional_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/news/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/news/news_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/ratios/__init__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/plots/ratios/ratios_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/reports/accounting/__inti__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/reports/accounting/accounting_balance_sheet.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/reports/bankruptcy/__inti__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/reports/bankruptcy/bankruptcy_monthly_top.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/reports/general/__inti__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/reports/general/general_plots.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/reports/news/__inti__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/reports/news/news_econometric_analysis.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/studies/nowcasting.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/token_auth.py +0 -0
- {sovai-0.2.25/sovai/tools/sec → sovai-0.2.27/sovai/tools/authentication}/__inti__.py +0 -0
- /sovai-0.2.25/sovai/utils/__init__.py → /sovai-0.2.27/sovai/tools/sec/__inti__.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/tools/sec/llm_code_generator.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/tools/sec/sec_10_k_8_k_filings.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/tools/sec/sec_edgar_search.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/converter.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/datetime_formats.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/file_management.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/get_tickers.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/helpers.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/pandas_extensions copy.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/plot.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/port_manager.py +0 -0
- {sovai-0.2.25 → sovai-0.2.27}/sovai/utils/stream.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[tool.poetry]
|
|
2
2
|
name = "sovai"
|
|
3
|
-
version = "0.2.
|
|
3
|
+
version = "0.2.27"
|
|
4
4
|
description = "python for asset management"
|
|
5
5
|
authors = ["Derek Snow <d.snow@outlook.com>"]
|
|
6
6
|
license = "Proprietary License"
|
|
@@ -40,6 +40,8 @@ python-dotenv = ">=0.10"
|
|
|
40
40
|
polars = ">=0.20.30"
|
|
41
41
|
boto3 = ">=1.20"
|
|
42
42
|
joblib = ">=1.0"
|
|
43
|
+
rapidfuzz = "^3.9.7"
|
|
44
|
+
edgar-tool = "^1.3.3"
|
|
43
45
|
scipy = {version = ">=1.0", optional = true}
|
|
44
46
|
matplotlib = {version = ">=3.0", optional = true}
|
|
45
47
|
plotly = {version = ">=5.0", optional = true}
|
|
@@ -59,8 +61,7 @@ lightgbm = {version = ">=4.5.0", optional = true}
|
|
|
59
61
|
ipywidgets = {version = ">=8.1.3", optional = true}
|
|
60
62
|
polars-talib = {version = "0.1.3", optional = true}
|
|
61
63
|
dash-bootstrap-components = {version = ">=1.6.0", optional = true}
|
|
62
|
-
|
|
63
|
-
edgar-tool = "^1.3.3"
|
|
64
|
+
|
|
64
65
|
|
|
65
66
|
[tool.poetry.extras]
|
|
66
67
|
full = [
|
|
@@ -85,6 +86,17 @@ full = [
|
|
|
85
86
|
"dash-bootstrap-components"
|
|
86
87
|
]
|
|
87
88
|
|
|
89
|
+
[tool.poetry.group.dev]
|
|
90
|
+
optional = true
|
|
91
|
+
|
|
92
|
+
[tool.poetry.group.dev.dependencies]
|
|
93
|
+
jinja2 = "^3.1.4"
|
|
94
|
+
traitlets = "^5.14.3"
|
|
95
|
+
jupyterlab = "^4.2.5"
|
|
96
|
+
pytest = "^7.0"
|
|
97
|
+
black = "^22.0"
|
|
98
|
+
|
|
99
|
+
|
|
88
100
|
[build-system]
|
|
89
101
|
requires = ["poetry-core"]
|
|
90
102
|
build-backend = "poetry.core.masonry.api"
|
|
@@ -6,8 +6,10 @@ from .get_data import data
|
|
|
6
6
|
from .basic_auth import basic_auth
|
|
7
7
|
from .token_auth import token_auth
|
|
8
8
|
from .utils.file_management import save_or_update_tickers, save_or_update_codes
|
|
9
|
+
|
|
10
|
+
|
|
9
11
|
import warnings
|
|
10
|
-
warnings.filterwarnings("ignore"
|
|
12
|
+
warnings.filterwarnings("ignore")
|
|
11
13
|
|
|
12
14
|
|
|
13
15
|
__version__ = "0.1.25"
|
|
@@ -21,6 +21,7 @@ from sovai.utils.converter import convert_data2df
|
|
|
21
21
|
from sovai.utils.stream import stream_data, stream_data_pyarrow
|
|
22
22
|
from sovai.utils.datetime_formats import datetime_format
|
|
23
23
|
from sovai.utils.client_side import client_side_frame
|
|
24
|
+
from sovai.utils.client_side_s3 import load_frame_s3
|
|
24
25
|
|
|
25
26
|
## Note run these for TEST - IT COULD BE HERE some silent problems
|
|
26
27
|
try:
|
|
@@ -262,6 +263,8 @@ endpoint_to_ticker = {
|
|
|
262
263
|
"/news/polarity_score": "",
|
|
263
264
|
"/macro/features": "",
|
|
264
265
|
"/congress": "",
|
|
266
|
+
"/market/closeadj": "",
|
|
267
|
+
|
|
265
268
|
|
|
266
269
|
|
|
267
270
|
|
|
@@ -279,7 +282,7 @@ def normalize_endpoint(endpoint):
|
|
|
279
282
|
return endpoint.strip("/").strip()
|
|
280
283
|
|
|
281
284
|
|
|
282
|
-
|
|
285
|
+
client_side_endpoints_gcs = {
|
|
283
286
|
"ratios/relative",
|
|
284
287
|
"market/prices",
|
|
285
288
|
"market/closeadj",
|
|
@@ -289,6 +292,10 @@ client_side_endpoints = {
|
|
|
289
292
|
"lobbying/public",
|
|
290
293
|
}
|
|
291
294
|
|
|
295
|
+
client_side_endpoints_s3 = {
|
|
296
|
+
"sec/10k"
|
|
297
|
+
}
|
|
298
|
+
|
|
292
299
|
# In your main function or wherever this logic is implemented:
|
|
293
300
|
|
|
294
301
|
|
|
@@ -575,16 +582,25 @@ def data(
|
|
|
575
582
|
normalized_endpoint = normalize_endpoint(endpoint)
|
|
576
583
|
|
|
577
584
|
if (
|
|
578
|
-
normalized_endpoint in
|
|
585
|
+
normalized_endpoint in client_side_endpoints_gcs
|
|
579
586
|
and tickers is not None
|
|
580
587
|
and frequency is None
|
|
581
588
|
):
|
|
582
|
-
verbose_mode.vprint("Grabbing client side")
|
|
589
|
+
verbose_mode.vprint("Grabbing gcs client side")
|
|
583
590
|
_query_cache[cache_key] = client_side_frame(
|
|
584
591
|
normalized_endpoint, tickers, columns, start_date, end_date
|
|
585
592
|
)
|
|
586
593
|
return _query_cache[cache_key]
|
|
587
594
|
|
|
595
|
+
elif (
|
|
596
|
+
normalized_endpoint in client_side_endpoints_s3
|
|
597
|
+
and tickers is not None
|
|
598
|
+
and frequency is None
|
|
599
|
+
):
|
|
600
|
+
verbose_mode.vprint("Grabbing s3 client side")
|
|
601
|
+
_query_cache[cache_key] = load_frame_s3(
|
|
602
|
+
normalized_endpoint, tickers, columns, start_date, end_date)
|
|
603
|
+
return _query_cache[cache_key]
|
|
588
604
|
try:
|
|
589
605
|
if tickers is not None and not is_all(tickers):
|
|
590
606
|
params = ticker_mapper(params,verbose)
|
|
@@ -0,0 +1,98 @@
|
|
|
1
|
+
import sovai as sov
|
|
2
|
+
from sovai.api_config import ApiConfig
|
|
3
|
+
import requests
|
|
4
|
+
import pickle
|
|
5
|
+
import base64
|
|
6
|
+
from pyarrow.fs import S3FileSystem
|
|
7
|
+
|
|
8
|
+
def get_s3_filesystem_pickle(storage_provider: str, verbose: bool = False):
|
|
9
|
+
"""
|
|
10
|
+
Retrieve the S3FileSystem for the specified storage provider using pickle.
|
|
11
|
+
|
|
12
|
+
:param storage_provider: Either 'wasabi' or 'digitalocean'
|
|
13
|
+
:param verbose: If True, print additional information about the request
|
|
14
|
+
:return: S3FileSystem object, or None if the request fails
|
|
15
|
+
"""
|
|
16
|
+
if storage_provider.lower() not in ['wasabi', 'digitalocean']:
|
|
17
|
+
raise ValueError("storage_provider must be either 'wasabi' or 'digitalocean'")
|
|
18
|
+
|
|
19
|
+
url = f"{ApiConfig.base_url}/credentials/get_s3_filesystem_pickle"
|
|
20
|
+
headers = {
|
|
21
|
+
"Authorization": f"Bearer {ApiConfig.token}",
|
|
22
|
+
"Content-Type": "application/json"
|
|
23
|
+
}
|
|
24
|
+
params = {"storage_provider": storage_provider.lower()}
|
|
25
|
+
|
|
26
|
+
try:
|
|
27
|
+
response = requests.get(url, headers=headers, params=params)
|
|
28
|
+
response.raise_for_status()
|
|
29
|
+
|
|
30
|
+
encoded = response.content
|
|
31
|
+
pickled = base64.b64decode(encoded)
|
|
32
|
+
s3fs = pickle.loads(pickled)
|
|
33
|
+
|
|
34
|
+
return s3fs
|
|
35
|
+
|
|
36
|
+
except requests.RequestException as e:
|
|
37
|
+
if verbose:
|
|
38
|
+
print(f"Failed to retrieve S3FileSystem for {storage_provider}: {e}")
|
|
39
|
+
return None
|
|
40
|
+
|
|
41
|
+
def get_s3_filesystem_json(storage_provider: str, verbose: bool = False):
|
|
42
|
+
"""
|
|
43
|
+
Retrieve the S3FileSystem for the specified storage provider using JSON.
|
|
44
|
+
|
|
45
|
+
:param storage_provider: Either 'wasabi' or 'digitalocean'
|
|
46
|
+
:param verbose: If True, print additional information about the request
|
|
47
|
+
:return: S3FileSystem object, or None if the request fails
|
|
48
|
+
"""
|
|
49
|
+
if storage_provider.lower() not in ['wasabi', 'digitalocean']:
|
|
50
|
+
raise ValueError("storage_provider must be either 'wasabi' or 'digitalocean'")
|
|
51
|
+
|
|
52
|
+
url = f"{ApiConfig.base_url}/credentials/get_s3_filesystem_json"
|
|
53
|
+
headers = {
|
|
54
|
+
"Authorization": f"Bearer {ApiConfig.token}",
|
|
55
|
+
"Content-Type": "application/json"
|
|
56
|
+
}
|
|
57
|
+
params = {"storage_provider": storage_provider.lower()}
|
|
58
|
+
|
|
59
|
+
try:
|
|
60
|
+
response = requests.get(url, headers=headers, params=params)
|
|
61
|
+
response.raise_for_status()
|
|
62
|
+
|
|
63
|
+
cred = response.json()
|
|
64
|
+
|
|
65
|
+
if verbose:
|
|
66
|
+
print(f"Successfully retrieved S3 credentials for {storage_provider}")
|
|
67
|
+
|
|
68
|
+
s3fs = S3FileSystem(
|
|
69
|
+
access_key=cred['access_key'],
|
|
70
|
+
secret_key=cred['secret_key'],
|
|
71
|
+
endpoint_override=cred['endpoint']
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
try:
|
|
75
|
+
s3fs.ls('/')
|
|
76
|
+
if verbose:
|
|
77
|
+
print(f"Successfully verified S3FileSystem for {storage_provider}")
|
|
78
|
+
return s3fs
|
|
79
|
+
except Exception as e:
|
|
80
|
+
if verbose:
|
|
81
|
+
print(f"Failed to verify S3FileSystem: {str(e)}")
|
|
82
|
+
return None
|
|
83
|
+
|
|
84
|
+
except requests.RequestException as e:
|
|
85
|
+
if verbose:
|
|
86
|
+
print(f"Failed to retrieve S3 credentials for {storage_provider}: {e}")
|
|
87
|
+
return None
|
|
88
|
+
|
|
89
|
+
# Usage example:
|
|
90
|
+
# s3fs = get_s3_filesystem_pickle('wasabi', verbose=True)
|
|
91
|
+
# or
|
|
92
|
+
# s3fs = get_s3_filesystem_json('wasabi', verbose=True)
|
|
93
|
+
# if s3fs:
|
|
94
|
+
# # Use s3fs in your operations
|
|
95
|
+
# files = s3fs.ls('/')
|
|
96
|
+
# print(files)
|
|
97
|
+
# else:
|
|
98
|
+
# print("Failed to get S3FileSystem")
|
|
File without changes
|
|
@@ -16,6 +16,7 @@ except ImportError:
|
|
|
16
16
|
CustomDataFrame = pd.DataFrame # Fallback to regular DataFrame
|
|
17
17
|
|
|
18
18
|
|
|
19
|
+
## might do the same thing as token_auth, you might have recreated it.
|
|
19
20
|
def verify_token(verbose=False):
|
|
20
21
|
url = f"{ApiConfig.base_url}/verify/token/{ApiConfig.token}" # Updated endpoint
|
|
21
22
|
try:
|
|
@@ -0,0 +1,212 @@
|
|
|
1
|
+
import pyarrow.dataset as ds
|
|
2
|
+
from pyarrow.fs import S3FileSystem
|
|
3
|
+
import pyarrow as pa
|
|
4
|
+
import pandas as pd
|
|
5
|
+
from functools import lru_cache
|
|
6
|
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
7
|
+
from sovai.tools.authentication import authentication
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
# Try to import CustomDataFrame, use regular DataFrame if not available
|
|
11
|
+
try:
|
|
12
|
+
from sovai.extensions.pandas_extensions import CustomDataFrame
|
|
13
|
+
HAS_CUSTOM_DATAFRAME = True
|
|
14
|
+
except ImportError:
|
|
15
|
+
HAS_CUSTOM_DATAFRAME = False
|
|
16
|
+
CustomDataFrame = pd.DataFrame # Fallback to regular DataFrame
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@lru_cache(maxsize=2)
|
|
20
|
+
def get_cached_s3_filesystem(storage_provider):
|
|
21
|
+
return authentication.get_s3_filesystem_pickle(storage_provider, verbose=True)
|
|
22
|
+
|
|
23
|
+
def process_partition(storage_provider, base_path, identifier_column, identifier=None, columns=None, filters=None):
|
|
24
|
+
s3 = get_cached_s3_filesystem(storage_provider)
|
|
25
|
+
|
|
26
|
+
if identifier:
|
|
27
|
+
base_path += f"/{identifier_column}={identifier}"
|
|
28
|
+
|
|
29
|
+
dataset = ds.dataset(base_path, filesystem=s3, format='parquet')
|
|
30
|
+
|
|
31
|
+
if filters:
|
|
32
|
+
operator_map = {
|
|
33
|
+
'==': lambda a, b: a == b,
|
|
34
|
+
'!=': lambda a, b: a != b,
|
|
35
|
+
'<': lambda a, b: a < b,
|
|
36
|
+
'<=': lambda a, b: a <= b,
|
|
37
|
+
'>': lambda a, b: a > b,
|
|
38
|
+
'>=': lambda a, b: a >= b,
|
|
39
|
+
'in': lambda a, b: a.isin(b),
|
|
40
|
+
'not in': lambda a, b: ~a.isin(b),
|
|
41
|
+
'like': lambda a, b: a.match_substring(b),
|
|
42
|
+
}
|
|
43
|
+
schema = dataset.schema
|
|
44
|
+
filter_expr = None
|
|
45
|
+
for col, op, val in filters:
|
|
46
|
+
if op not in operator_map:
|
|
47
|
+
raise ValueError(f"Unsupported operator '{op}' in filters.")
|
|
48
|
+
field = ds.field(col)
|
|
49
|
+
field_type = schema.field(col).type
|
|
50
|
+
if pa.types.is_timestamp(field_type) or pa.types.is_date(field_type):
|
|
51
|
+
val = pd.to_datetime(val)
|
|
52
|
+
elif pa.types.is_integer(field_type):
|
|
53
|
+
val = int(val)
|
|
54
|
+
elif pa.types.is_floating(field_type):
|
|
55
|
+
val = float(val)
|
|
56
|
+
condition = operator_map[op](field, val)
|
|
57
|
+
filter_expr = condition if filter_expr is None else filter_expr & condition
|
|
58
|
+
else:
|
|
59
|
+
filter_expr = None
|
|
60
|
+
|
|
61
|
+
table = dataset.to_table(columns=columns, filter=filter_expr, use_threads=True)
|
|
62
|
+
df = table.to_pandas(use_threads=True)
|
|
63
|
+
|
|
64
|
+
if identifier:
|
|
65
|
+
df[identifier_column] = identifier
|
|
66
|
+
|
|
67
|
+
return df
|
|
68
|
+
|
|
69
|
+
import numpy as np
|
|
70
|
+
|
|
71
|
+
def client_side_s3_frame(config, identifiers, columns, start_date, end_date):
|
|
72
|
+
storage_provider = "digitalocean" # Default to DigitalOcean
|
|
73
|
+
base_path = config["storage_provider"][storage_provider]
|
|
74
|
+
identifier_column = config.get("identifier_column")
|
|
75
|
+
|
|
76
|
+
filters = []
|
|
77
|
+
if start_date:
|
|
78
|
+
filters.append(('date', '>=', start_date))
|
|
79
|
+
if end_date:
|
|
80
|
+
filters.append(('date', '<=', end_date))
|
|
81
|
+
|
|
82
|
+
if not identifiers:
|
|
83
|
+
# Load the entire database
|
|
84
|
+
return process_partition(storage_provider, base_path, identifier_column, columns=columns, filters=filters)
|
|
85
|
+
else:
|
|
86
|
+
identifiers = map_identifier(identifiers, input_identifier='ticker', output_identifier=identifier_column)
|
|
87
|
+
|
|
88
|
+
# Use numpy.isscalar to check if identifiers is a single value
|
|
89
|
+
if np.isscalar(identifiers):
|
|
90
|
+
# Single identifier
|
|
91
|
+
identifier = identifiers
|
|
92
|
+
return process_partition(storage_provider, base_path, identifier_column, identifier=identifier, columns=columns, filters=filters)
|
|
93
|
+
|
|
94
|
+
# Multiple identifiers
|
|
95
|
+
with ThreadPoolExecutor() as executor:
|
|
96
|
+
futures = []
|
|
97
|
+
for identifier in identifiers:
|
|
98
|
+
futures.append(executor.submit(process_partition, storage_provider, base_path, identifier_column, identifier, columns, filters))
|
|
99
|
+
|
|
100
|
+
results = []
|
|
101
|
+
for future in as_completed(futures):
|
|
102
|
+
results.append(future.result())
|
|
103
|
+
|
|
104
|
+
df = pd.concat(results, ignore_index=True)
|
|
105
|
+
return df
|
|
106
|
+
|
|
107
|
+
# if 'filing_date' in df.columns:
|
|
108
|
+
# df['date'] = pd.to_datetime(df['filing_date'])
|
|
109
|
+
# elif 'date' not in df.columns:
|
|
110
|
+
# raise ValueError("Neither 'filing_date' nor 'date' column found in the data")
|
|
111
|
+
|
|
112
|
+
# df.set_index([identifier_column, 'date'], inplace=True)
|
|
113
|
+
# df.sort_index(inplace=True)
|
|
114
|
+
|
|
115
|
+
return df
|
|
116
|
+
|
|
117
|
+
import pandas as pd
|
|
118
|
+
import numpy as np
|
|
119
|
+
|
|
120
|
+
def map_identifier(input_values, input_identifier='ticker', output_identifier='cik', trailing_zero_removal=True):
|
|
121
|
+
"""
|
|
122
|
+
Map from one identifier to another using the df_codes DataFrame.
|
|
123
|
+
Handles string, single-item list, and multiple-item lists as input.
|
|
124
|
+
Uses vectorization for improved performance.
|
|
125
|
+
|
|
126
|
+
:param input_values: The input identifier value(s) to map from (string, list, or array-like)
|
|
127
|
+
:param input_identifier: The type of input identifier (default: 'ticker')
|
|
128
|
+
:param output_identifier: The type of output identifier (default: 'cik')
|
|
129
|
+
:param trailing_zero_removal: If True, convert CIK to int (default: True)
|
|
130
|
+
:return: Mapped identifier value(s), same type as input
|
|
131
|
+
"""
|
|
132
|
+
df_codes = pd.read_parquet("data/codes.parq")
|
|
133
|
+
|
|
134
|
+
# Check if the identifier columns exist in df_codes
|
|
135
|
+
if input_identifier not in df_codes.columns or output_identifier not in df_codes.columns:
|
|
136
|
+
raise ValueError(f"Identifier '{input_identifier}' or '{output_identifier}' not found in the DataFrame")
|
|
137
|
+
|
|
138
|
+
# Convert input to numpy array for consistent handling
|
|
139
|
+
if isinstance(input_values, str):
|
|
140
|
+
input_array = np.array([input_values])
|
|
141
|
+
input_is_string = True
|
|
142
|
+
else:
|
|
143
|
+
input_array = np.array(input_values)
|
|
144
|
+
input_is_string = False
|
|
145
|
+
|
|
146
|
+
# Ensure all input values are strings
|
|
147
|
+
input_array = np.char.strip(input_array.astype(str))
|
|
148
|
+
|
|
149
|
+
# Create a dictionary for fast lookup
|
|
150
|
+
mapping_dict = dict(zip(df_codes[input_identifier], df_codes[output_identifier]))
|
|
151
|
+
|
|
152
|
+
# Vectorized mapping
|
|
153
|
+
result = np.vectorize(mapping_dict.get)(input_array)
|
|
154
|
+
|
|
155
|
+
# Convert to int if trailing_zero_removal is True and output is CIK
|
|
156
|
+
if trailing_zero_removal and output_identifier.lower() == 'cik':
|
|
157
|
+
result = np.where(result == None, 0, result) # Replace None with 0
|
|
158
|
+
result = result.astype(np.int32) # Convert to int64
|
|
159
|
+
|
|
160
|
+
# Return result in the same format as input
|
|
161
|
+
if input_is_string:
|
|
162
|
+
return result[0]
|
|
163
|
+
elif len(input_array) == 1:
|
|
164
|
+
return result[0]
|
|
165
|
+
else:
|
|
166
|
+
return result.tolist()
|
|
167
|
+
|
|
168
|
+
def load_frame_s3(endpoint, tickers=None, columns=None, start_date=None, end_date=None):
|
|
169
|
+
|
|
170
|
+
endpoint_config = {
|
|
171
|
+
"sec/10k": {
|
|
172
|
+
"storage_provider": {
|
|
173
|
+
"digitalocean": "sovai/sovai-filings/tenk",
|
|
174
|
+
"wasabi": "sovai-filings/tenk"
|
|
175
|
+
},
|
|
176
|
+
"identifier_column": "cik",
|
|
177
|
+
"partitioned": True,
|
|
178
|
+
|
|
179
|
+
},
|
|
180
|
+
"future/enpdoint": {
|
|
181
|
+
"storage_provider": {
|
|
182
|
+
"digitalocean": "sovai/sovai-filings/tenk",
|
|
183
|
+
"wasabi": "sovai-filings/tenk"
|
|
184
|
+
},
|
|
185
|
+
"identifier_column": "cik",
|
|
186
|
+
"partitioned": True,
|
|
187
|
+
}
|
|
188
|
+
}
|
|
189
|
+
|
|
190
|
+
if endpoint not in endpoint_config:
|
|
191
|
+
raise ValueError(f"Invalid endpoint: {endpoint}")
|
|
192
|
+
|
|
193
|
+
config = endpoint_config[endpoint]
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
if endpoint == "sec/10k":
|
|
197
|
+
df_frame = client_side_s3_frame(
|
|
198
|
+
config,
|
|
199
|
+
tickers,
|
|
200
|
+
columns,
|
|
201
|
+
start_date,
|
|
202
|
+
end_date
|
|
203
|
+
)
|
|
204
|
+
else:
|
|
205
|
+
# Implement logic for other endpoints if needed
|
|
206
|
+
raise NotImplementedError(f"Endpoint {endpoint} is not yet implemented")
|
|
207
|
+
|
|
208
|
+
if HAS_CUSTOM_DATAFRAME:
|
|
209
|
+
return CustomDataFrame(df_frame)
|
|
210
|
+
else:
|
|
211
|
+
return df_frame # Returns a regular pandas DataFrame if CustomDataFrame is not available
|
|
212
|
+
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|