bigdata-smart-batching 0.1.0__py3-none-any.whl
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.
- bigdata_smart_batching/__init__.py +28 -0
- bigdata_smart_batching/output_converter.py +96 -0
- bigdata_smart_batching/search_function.py +1071 -0
- bigdata_smart_batching/smart_batching.py +1593 -0
- bigdata_smart_batching/smart_batching_config.py +39 -0
- bigdata_smart_batching-0.1.0.dist-info/METADATA +231 -0
- bigdata_smart_batching-0.1.0.dist-info/RECORD +8 -0
- bigdata_smart_batching-0.1.0.dist-info/WHEEL +4 -0
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Configuration constants for smart batching.
|
|
3
|
+
|
|
4
|
+
This module contains all configuration parameters for the smart batching system,
|
|
5
|
+
including API endpoints, limits, time periods, and file paths.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import os
|
|
9
|
+
|
|
10
|
+
# API Configuration
|
|
11
|
+
# Use environment variable or default to production API
|
|
12
|
+
API_BASE_URL = os.getenv("BIGDATA_API_BASE_URL", "https://api.bigdata.com")
|
|
13
|
+
COMENTION_ENDPOINT = "/v1/search/co-mentions/entities"
|
|
14
|
+
VOLUME_ENDPOINT = "/v1/search/volume"
|
|
15
|
+
MAX_ENTITIES_PER_QUERY = 1000 # Max entities the API can return in a single response
|
|
16
|
+
MAX_ENTITIES_IN_ANY_OF = 500 # Max entities we can send in any_of filter (API complexity limit)
|
|
17
|
+
MAX_CHUNKS_PER_BASKET = 1000 # Maximum chunks per basket to stay within query limits
|
|
18
|
+
|
|
19
|
+
# Authentication
|
|
20
|
+
# Note: API_KEY is not required at import time - it's passed to SmartBatchingPlanner
|
|
21
|
+
# This allows the module to be imported without errors
|
|
22
|
+
API_KEY = os.environ.get("BIGDATA_API_KEY") # Optional at import time
|
|
23
|
+
|
|
24
|
+
# Time Period Configuration
|
|
25
|
+
# Default date range: 2 years from January 2021 to December 2022
|
|
26
|
+
START_DATE = "2021-01-01"
|
|
27
|
+
END_DATE = "2022-12-31"
|
|
28
|
+
|
|
29
|
+
# Volume Buckets for basket creation
|
|
30
|
+
# Used to group companies by their chunk volume for efficient batching
|
|
31
|
+
VOLUME_BUCKETS = {
|
|
32
|
+
"high": (500, float("inf")), # 500-1000+ chunks
|
|
33
|
+
"medium": (100, 500), # 100-500 chunks
|
|
34
|
+
"low": (1, 100), # 1-100 chunks
|
|
35
|
+
"very_low": (0, 1), # 0 chunks (headlines only, no semantic content)
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
# File Paths
|
|
39
|
+
UNIVERSE_CSV_PATH = "us_top3000.csv" # Default universe file (example: US top 3000 companies)
|
|
@@ -0,0 +1,231 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: bigdata-smart-batching
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: High-performance semantic search with intelligent company grouping and parallel execution
|
|
5
|
+
License: MIT
|
|
6
|
+
Requires-Python: >=3.11
|
|
7
|
+
Requires-Dist: matplotlib>=3.7
|
|
8
|
+
Requires-Dist: numpy>=1.24
|
|
9
|
+
Requires-Dist: pandas>=2.0
|
|
10
|
+
Requires-Dist: python-dotenv>=1.2.2
|
|
11
|
+
Requires-Dist: requests>=2.31.0
|
|
12
|
+
Description-Content-Type: text/markdown
|
|
13
|
+
|
|
14
|
+
# Smart Batching Search
|
|
15
|
+
|
|
16
|
+
A high-performance semantic search system that reduces API queries by **67-99%** (varies by topic specificity) through intelligent company grouping and parallel execution.
|
|
17
|
+
|
|
18
|
+
This module provides a two-step system for efficient semantic search:
|
|
19
|
+
1. **Planning**: Organize search using smart batching and return total expected chunks
|
|
20
|
+
2. **Execution**: Perform search with proportional sampling to preserve distribution
|
|
21
|
+
|
|
22
|
+
## Key Benefits
|
|
23
|
+
|
|
24
|
+
- **67-99% Query Reduction**: Search 4,732 companies with only 17-3,699 queries (varies by topic)
|
|
25
|
+
- **Parallel Execution**: Rate-limited concurrent requests with semaphore control
|
|
26
|
+
- **Proportional Sampling**: Retrieve percentage of results while preserving distribution
|
|
27
|
+
- **Production Ready**: Comprehensive error handling, retries, and logging
|
|
28
|
+
- **Scalable**: Efficiently handles universes with 10,000+ companies
|
|
29
|
+
|
|
30
|
+
## Installation
|
|
31
|
+
|
|
32
|
+
### Using uv (Recommended)
|
|
33
|
+
|
|
34
|
+
```bash
|
|
35
|
+
# Clone or navigate to the project
|
|
36
|
+
cd bigdata-smart-batching
|
|
37
|
+
|
|
38
|
+
# Install in development mode
|
|
39
|
+
uv sync
|
|
40
|
+
|
|
41
|
+
# Or with dev dependencies
|
|
42
|
+
uv sync --all-extras
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
### Environment Setup
|
|
46
|
+
|
|
47
|
+
Set up environment variables:
|
|
48
|
+
|
|
49
|
+
```bash
|
|
50
|
+
export BIGDATA_API_KEY="your_api_key_here"
|
|
51
|
+
export BIGDATA_API_BASE_URL="https://api.bigdata.com" # Optional, defaults to this
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
Or create a `.env` file:
|
|
55
|
+
|
|
56
|
+
```
|
|
57
|
+
BIGDATA_API_KEY=your_api_key_here
|
|
58
|
+
BIGDATA_API_BASE_URL=https://api.bigdata.com
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
## Quick Start
|
|
62
|
+
|
|
63
|
+
```python
|
|
64
|
+
from bigdata_smart_batching import (
|
|
65
|
+
plan_search,
|
|
66
|
+
execute_search,
|
|
67
|
+
deduplicate_documents,
|
|
68
|
+
convert_to_dataframe,
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
# Step 1: Plan the search
|
|
72
|
+
plan = plan_search(
|
|
73
|
+
text="earnings revenue profit",
|
|
74
|
+
universe_csv_path="id_name_mapping_us_top_3000.csv",
|
|
75
|
+
start_date="2023-01-01",
|
|
76
|
+
end_date="2023-12-31",
|
|
77
|
+
api_key="your_api_key", # or set BIGDATA_API_KEY env var
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
print(f"Total expected chunks: {plan['total_expected_chunks']:,}")
|
|
81
|
+
|
|
82
|
+
# Step 2: Execute search with 10% of total chunks (preserves distribution)
|
|
83
|
+
results_raw = execute_search(
|
|
84
|
+
search_plan=plan,
|
|
85
|
+
chunk_percentage=0.1,
|
|
86
|
+
requests_per_minute=100,
|
|
87
|
+
)
|
|
88
|
+
|
|
89
|
+
# Step 3: Deduplicate and convert to DataFrame
|
|
90
|
+
results = deduplicate_documents(results_raw)
|
|
91
|
+
print(f"Retrieved {len(results)} documents (deduplicated)")
|
|
92
|
+
|
|
93
|
+
df = convert_to_dataframe(results) # one row per chunk
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
### Save and Load Plans
|
|
97
|
+
|
|
98
|
+
```python
|
|
99
|
+
from bigdata_smart_batching import plan_search, execute_search, save_plan, load_plan
|
|
100
|
+
|
|
101
|
+
# Create and save a plan
|
|
102
|
+
plan = plan_search(
|
|
103
|
+
text="merger acquisition",
|
|
104
|
+
universe_csv_path="id_name_mapping_us_top_3000.csv",
|
|
105
|
+
start_date="2023-01-01",
|
|
106
|
+
end_date="2023-12-31",
|
|
107
|
+
)
|
|
108
|
+
save_plan(plan, "my_search_plan.json")
|
|
109
|
+
|
|
110
|
+
# Later: reload and run with different sampling
|
|
111
|
+
plan = load_plan("my_search_plan.json")
|
|
112
|
+
raw_10 = execute_search(plan, chunk_percentage=0.1)
|
|
113
|
+
raw_50 = execute_search(plan, chunk_percentage=0.5)
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
## How It Works
|
|
117
|
+
|
|
118
|
+
### Architecture Overview
|
|
119
|
+
|
|
120
|
+
```
|
|
121
|
+
Step 1: PLANNING
|
|
122
|
+
Universe CSV --> Co-mention API Query --> Basket Creation --> Search Plan
|
|
123
|
+
|
|
124
|
+
Step 2: EXECUTION
|
|
125
|
+
Proportional Sampling --> Parallel Search (Rate Limited) --> Collect & Aggregate
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
### Planning (`plan_search()`)
|
|
129
|
+
|
|
130
|
+
1. Loads the universe of companies from CSV
|
|
131
|
+
2. Queries the comention endpoint to get chunk volumes per company
|
|
132
|
+
3. Splits date ranges by volume when a company exceeds the chunk limit
|
|
133
|
+
4. Creates optimized baskets grouped by volume
|
|
134
|
+
5. Returns a plan with total expected chunks and basket configurations
|
|
135
|
+
|
|
136
|
+
### Execution (`execute_search()`)
|
|
137
|
+
|
|
138
|
+
1. Calculates proportional chunks per basket
|
|
139
|
+
2. Ensures minimum of 1 chunk per basket (if expected > 0)
|
|
140
|
+
3. Executes searches in parallel with rate limiting and semaphore
|
|
141
|
+
4. Collects and returns document results
|
|
142
|
+
|
|
143
|
+
## API Reference
|
|
144
|
+
|
|
145
|
+
### `plan_search()`
|
|
146
|
+
|
|
147
|
+
| Parameter | Type | Default | Description |
|
|
148
|
+
|-----------|------|---------|-------------|
|
|
149
|
+
| `text` | `str` | required | Search query text |
|
|
150
|
+
| `universe_csv_path` | `str` | required | Path to CSV with entity IDs |
|
|
151
|
+
| `start_date` | `str` | required | Start date (YYYY-MM-DD) |
|
|
152
|
+
| `end_date` | `str` | required | End date (YYYY-MM-DD) |
|
|
153
|
+
| `api_key` | `str` | env var | API key |
|
|
154
|
+
| `api_base_url` | `str` | env var | API base URL |
|
|
155
|
+
| `volume_query_mode` | `str` | `"three_pass"` | `"three_pass"` or `"iterative"` |
|
|
156
|
+
| `apply_volume_splits` | `bool` | `True` | Use volume time series for period splitting |
|
|
157
|
+
| `min_period_days` | `int` | `30` | Minimum days per sub-period |
|
|
158
|
+
|
|
159
|
+
### `execute_search()`
|
|
160
|
+
|
|
161
|
+
| Parameter | Type | Default | Description |
|
|
162
|
+
|-----------|------|---------|-------------|
|
|
163
|
+
| `search_plan` | `Dict` | required | Plan from `plan_search()` |
|
|
164
|
+
| `chunk_percentage` | `float` | required | 0.0 to 1.0 sampling ratio |
|
|
165
|
+
| `requests_per_minute` | `int` | `100` | Rate limit |
|
|
166
|
+
| `api_key` | `str` | env var | API key |
|
|
167
|
+
| `max_workers` | `int` | `40` | Parallel workers |
|
|
168
|
+
|
|
169
|
+
### Helper Functions
|
|
170
|
+
|
|
171
|
+
- **`deduplicate_documents(documents)`** -- Merges duplicate documents by `id`
|
|
172
|
+
- **`load_universe_from_csv(csv_path)`** -- Loads entity IDs from CSV
|
|
173
|
+
- **`convert_to_dataframe(raw_results)`** -- Converts documents to DataFrame (one row per chunk)
|
|
174
|
+
- **`save_plan(plan, path)`** / **`load_plan(path)`** -- Persist plans as JSON
|
|
175
|
+
- **`portfolio_backtesting_pipeline(...)`** -- Long-short portfolio backtesting
|
|
176
|
+
|
|
177
|
+
## Testing
|
|
178
|
+
|
|
179
|
+
```bash
|
|
180
|
+
# Run all tests
|
|
181
|
+
uv run pytest
|
|
182
|
+
|
|
183
|
+
# With coverage
|
|
184
|
+
uv run pytest --cov=bigdata_smart_batching --cov-report=term-missing
|
|
185
|
+
|
|
186
|
+
# Specific test file
|
|
187
|
+
uv run pytest tests/test_validation.py -v
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
## Project Structure
|
|
191
|
+
|
|
192
|
+
```
|
|
193
|
+
bigdata-smart-batching/
|
|
194
|
+
├── pyproject.toml
|
|
195
|
+
├── README.md
|
|
196
|
+
├── .python-version
|
|
197
|
+
├── src/
|
|
198
|
+
│ └── bigdata_smart_batching/
|
|
199
|
+
│ ├── __init__.py
|
|
200
|
+
│ ├── smart_batching.py
|
|
201
|
+
│ ├── smart_batching_config.py
|
|
202
|
+
│ ├── search_function.py
|
|
203
|
+
│ ├── output_converter.py
|
|
204
|
+
│ └── portfolio_backtesting.py
|
|
205
|
+
└── tests/
|
|
206
|
+
├── __init__.py
|
|
207
|
+
├── test_config.py
|
|
208
|
+
├── test_output_converter.py
|
|
209
|
+
├── test_validation.py
|
|
210
|
+
└── test_rate_limiter.py
|
|
211
|
+
```
|
|
212
|
+
|
|
213
|
+
## Configuration
|
|
214
|
+
|
|
215
|
+
### Environment Variables
|
|
216
|
+
|
|
217
|
+
- `BIGDATA_API_KEY`: Required -- Your Bigdata API key
|
|
218
|
+
- `BIGDATA_API_BASE_URL`: Optional -- API base URL (default: `https://api.bigdata.com`)
|
|
219
|
+
|
|
220
|
+
### Default Settings
|
|
221
|
+
|
|
222
|
+
- `requests_per_minute`: 100
|
|
223
|
+
- `max_workers`: 40
|
|
224
|
+
- `max_chunks_per_basket`: 1000
|
|
225
|
+
- `volume_query_mode`: `"three_pass"`
|
|
226
|
+
|
|
227
|
+
## License
|
|
228
|
+
|
|
229
|
+
This project is part of the Bigdata.com and WorldQuant Challenge.
|
|
230
|
+
|
|
231
|
+
**Disclaimer**: This software is provided "as is" without warranty of any kind, express or implied. The authors and contributors assume no responsibility for the accuracy, completeness, or usefulness of any information, results, or processes provided. This software is for educational and research purposes only and is not intended to be used as financial advice.
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
bigdata_smart_batching/__init__.py,sha256=t19V51mbiS-gO5IJFC1p2Jk3FdplKfUFvWmGOrtAAAs,620
|
|
2
|
+
bigdata_smart_batching/output_converter.py,sha256=_YRlgtkw79nD2OzIKnISD4GEowOZrmizxDgSBGuUhBw,3427
|
|
3
|
+
bigdata_smart_batching/search_function.py,sha256=gdmIq1SpQIKFRCwIoRRXY6TnEHc-mBlRrcmd5VVHwug,41743
|
|
4
|
+
bigdata_smart_batching/smart_batching.py,sha256=5eECKFXyMeT94c4o37wFWAV1-2wCzbB6gsLmMFZh67w,69119
|
|
5
|
+
bigdata_smart_batching/smart_batching_config.py,sha256=wIRmKfUed3TgSEaEO3JIx66NIHwZwd9q_kw3eY_2oSM,1543
|
|
6
|
+
bigdata_smart_batching-0.1.0.dist-info/METADATA,sha256=J4nwxKtwSFOE4_zTyDQDizWZFBOn5bUfh2GjWAxWTFQ,7169
|
|
7
|
+
bigdata_smart_batching-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
|
|
8
|
+
bigdata_smart_batching-0.1.0.dist-info/RECORD,,
|