bigdata-smart-batching 1.0.1__tar.gz → 1.0.2__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.
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/CHANGELOG.md +6 -0
- bigdata_smart_batching-1.0.2/PKG-INFO +340 -0
- bigdata_smart_batching-1.0.2/README.md +324 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/pyproject.toml +1 -1
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/uv.lock +1 -1
- bigdata_smart_batching-1.0.1/PKG-INFO +0 -257
- bigdata_smart_batching-1.0.1/README.md +0 -241
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/.gitignore +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/LICENSE +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/Makefile +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/docs/smart-batching-architecture.md +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/examples/example.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/examples/example_overlapping.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/src/bigdata_smart_batching/__init__.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/src/bigdata_smart_batching/output_converter.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/src/bigdata_smart_batching/search_function.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/src/bigdata_smart_batching/smart_batching.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/src/bigdata_smart_batching/smart_batching_config.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/tests/__init__.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/tests/test_config.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/tests/test_output_converter.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/tests/test_rate_limiter.py +0 -0
- {bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/tests/test_validation.py +0 -0
|
@@ -5,6 +5,12 @@ All notable changes to this project are documented in this file.
|
|
|
5
5
|
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
|
|
6
6
|
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
|
7
7
|
|
|
8
|
+
## [1.0.2] - 2026-04-17
|
|
9
|
+
|
|
10
|
+
### Fixed
|
|
11
|
+
|
|
12
|
+
- Corrected README file
|
|
13
|
+
|
|
8
14
|
## [1.0.1] - 2026-04-17
|
|
9
15
|
|
|
10
16
|
### Fixed
|
|
@@ -0,0 +1,340 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: bigdata-smart-batching
|
|
3
|
+
Version: 1.0.2
|
|
4
|
+
Summary: High-performance semantic search with intelligent company grouping and parallel execution
|
|
5
|
+
Project-URL: homepage, https://bigdata.com/api
|
|
6
|
+
Author-email: "Bigdata.com" <support@ravenpack.com>
|
|
7
|
+
License: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Requires-Python: <4.0,>=3.11
|
|
10
|
+
Requires-Dist: matplotlib>=3.7
|
|
11
|
+
Requires-Dist: numpy>=1.24
|
|
12
|
+
Requires-Dist: pandas>=2.0
|
|
13
|
+
Requires-Dist: python-dotenv>=1.2.2
|
|
14
|
+
Requires-Dist: requests>=2.31.0
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
|
|
17
|
+
# Smart Batching Search
|
|
18
|
+
|
|
19
|
+
A high-performance semantic search system that reduces API queries by **67-99%** (varies by topic specificity) through intelligent company grouping and parallel execution.
|
|
20
|
+
|
|
21
|
+
This module provides a two-step pipeline for efficient semantic search:
|
|
22
|
+
|
|
23
|
+
1. **Planning** -- organize the search via smart batching and return a chunk upper bound estimate.
|
|
24
|
+
2. **Execution** -- run the plan with proportional sampling to preserve the result distribution.
|
|
25
|
+
|
|
26
|
+
## Key Benefits
|
|
27
|
+
|
|
28
|
+
- **67-99% query reduction** -- search 4,732 companies with only 17-3,699 queries (varies by topic).
|
|
29
|
+
- **Parallel execution** -- rate-limited concurrent requests with semaphore-controlled concurrency.
|
|
30
|
+
- **Proportional sampling** -- retrieve a percentage of results while preserving the distribution across baskets.
|
|
31
|
+
- **Category filtering** -- restrict searches to specific document categories (news, transcripts, filings, ...).
|
|
32
|
+
- **Deterministic plans** -- every request stays within service limits; plans can be saved and replayed.
|
|
33
|
+
- **Scalable** -- handles universes with 10,000+ companies efficiently.
|
|
34
|
+
|
|
35
|
+
## Installation
|
|
36
|
+
|
|
37
|
+
Install the package from PyPI (Python 3.11+):
|
|
38
|
+
|
|
39
|
+
```bash
|
|
40
|
+
pip install bigdata-smart-batching
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
With [uv](https://docs.astral.sh/uv/):
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
uv add bigdata-smart-batching
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
### Development
|
|
50
|
+
|
|
51
|
+
To work on this repository locally, from the project root:
|
|
52
|
+
|
|
53
|
+
```bash
|
|
54
|
+
uv sync
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
### Environment Setup
|
|
58
|
+
|
|
59
|
+
Set up environment variables:
|
|
60
|
+
|
|
61
|
+
```bash
|
|
62
|
+
export BIGDATA_API_KEY="your_api_key_here"
|
|
63
|
+
export BIGDATA_API_BASE_URL="https://api.bigdata.com" # Optional, defaults to this
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
Or create a `.env` file:
|
|
67
|
+
|
|
68
|
+
```
|
|
69
|
+
BIGDATA_API_KEY=your_api_key_here
|
|
70
|
+
BIGDATA_API_BASE_URL=https://api.bigdata.com
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
## Universe (CSV path or entity list)
|
|
74
|
+
|
|
75
|
+
`plan_search()` takes a **`universe`** argument: either a path to a UTF-8 CSV file, or a `list[str]` of entity IDs. IDs must match the identifiers used by the Bigdata API for your dataset.
|
|
76
|
+
|
|
77
|
+
When using a CSV, two layouts are supported:
|
|
78
|
+
|
|
79
|
+
**1. Header row with an `id` column** (optional extra columns such as `name` are ignored):
|
|
80
|
+
|
|
81
|
+
```csv
|
|
82
|
+
id,name
|
|
83
|
+
B8EF97,Example Corp A
|
|
84
|
+
BB07E4,Example Corp B
|
|
85
|
+
3461CF,Example Corp C
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
You can also pass a plain list instead of a CSV path:
|
|
89
|
+
|
|
90
|
+
```python
|
|
91
|
+
plan = plan_search(
|
|
92
|
+
universe=["B8EF97", "BB07E4", "3461CF"],
|
|
93
|
+
start_date="2023-01-01",
|
|
94
|
+
end_date="2023-12-31",
|
|
95
|
+
text="earnings revenue profit",
|
|
96
|
+
)
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
## Quick Start
|
|
100
|
+
|
|
101
|
+
```python
|
|
102
|
+
from bigdata_smart_batching import (
|
|
103
|
+
plan_search,
|
|
104
|
+
execute_search,
|
|
105
|
+
deduplicate_documents,
|
|
106
|
+
convert_to_dataframe,
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
plan = plan_search(
|
|
110
|
+
universe="id_name_mapping_us_top_3000.csv",
|
|
111
|
+
start_date="2023-01-01",
|
|
112
|
+
end_date="2023-12-31",
|
|
113
|
+
text="earnings revenue profit",
|
|
114
|
+
api_key="your_api_key", # or set BIGDATA_API_KEY env var
|
|
115
|
+
)
|
|
116
|
+
|
|
117
|
+
print(f"Chunk upper bound estimate: {plan['chunk_upper_bound_estimate']:,}")
|
|
118
|
+
|
|
119
|
+
results_raw = execute_search(
|
|
120
|
+
search_plan=plan,
|
|
121
|
+
chunk_percentage=0.1,
|
|
122
|
+
requests_per_minute=100,
|
|
123
|
+
)
|
|
124
|
+
|
|
125
|
+
results = deduplicate_documents(results_raw)
|
|
126
|
+
print(f"Retrieved {len(results)} documents (deduplicated)")
|
|
127
|
+
|
|
128
|
+
df = convert_to_dataframe(results) # one row per chunk
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
### Category Filtering
|
|
132
|
+
|
|
133
|
+
Restrict searches to a specific set of document categories using the `category`
|
|
134
|
+
argument. It accepts either a `CategoryFilter` dataclass or a plain dict.
|
|
135
|
+
|
|
136
|
+
```python
|
|
137
|
+
from bigdata_smart_batching import CategoryFilter, CategoryMode, plan_search
|
|
138
|
+
|
|
139
|
+
plan = plan_search(
|
|
140
|
+
universe="id_name_mapping_us_top_3000.csv",
|
|
141
|
+
start_date="2023-01-01",
|
|
142
|
+
end_date="2023-12-31",
|
|
143
|
+
text="merger acquisition",
|
|
144
|
+
category=CategoryFilter(
|
|
145
|
+
mode=CategoryMode.INCLUDE,
|
|
146
|
+
values=("news_premium", "transcripts"),
|
|
147
|
+
),
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
# Equivalent dict form:
|
|
151
|
+
plan = plan_search(
|
|
152
|
+
universe="id_name_mapping_us_top_3000.csv",
|
|
153
|
+
start_date="2023-01-01",
|
|
154
|
+
end_date="2023-12-31",
|
|
155
|
+
text="merger acquisition",
|
|
156
|
+
category={"mode": "INCLUDE", "values": ["news_premium", "transcripts"]},
|
|
157
|
+
)
|
|
158
|
+
```
|
|
159
|
+
|
|
160
|
+
Valid category values are exposed via `VALID_CATEGORY_VALUES` and include:
|
|
161
|
+
`expert_interviews`, `filings`, `my_files`, `news`, `news_premium`,
|
|
162
|
+
`news_public`, `podcasts`, `research`, `research_academic_journals`,
|
|
163
|
+
`research_investment_research`, `transcripts`.
|
|
164
|
+
|
|
165
|
+
### Save and Load Plans
|
|
166
|
+
|
|
167
|
+
```python
|
|
168
|
+
from bigdata_smart_batching import plan_search, execute_search, save_plan, load_plan
|
|
169
|
+
|
|
170
|
+
plan = plan_search(
|
|
171
|
+
universe="id_name_mapping_us_top_3000.csv",
|
|
172
|
+
start_date="2023-01-01",
|
|
173
|
+
end_date="2023-12-31",
|
|
174
|
+
text="merger acquisition",
|
|
175
|
+
)
|
|
176
|
+
save_plan(plan, "my_search_plan.json")
|
|
177
|
+
|
|
178
|
+
plan = load_plan("my_search_plan.json")
|
|
179
|
+
raw_10 = execute_search(plan, chunk_percentage=0.1)
|
|
180
|
+
raw_50 = execute_search(plan, chunk_percentage=0.5)
|
|
181
|
+
```
|
|
182
|
+
|
|
183
|
+
## How It Works
|
|
184
|
+
|
|
185
|
+
### Architecture Overview
|
|
186
|
+
|
|
187
|
+
```
|
|
188
|
+
Step 1: PLANNING
|
|
189
|
+
Universe --> Co-mention volumes --> Entity groups (by volume) --> Sub-period split --> Search plan
|
|
190
|
+
|
|
191
|
+
Step 2: EXECUTION
|
|
192
|
+
Proportional sampling --> Parallel search (rate limited) --> Collect & aggregate
|
|
193
|
+
```
|
|
194
|
+
|
|
195
|
+
### Planning (`plan_search()`)
|
|
196
|
+
|
|
197
|
+
1. Loads the universe of companies from a CSV path or inline list.
|
|
198
|
+
2. Queries the co-mention endpoint to get chunk volumes per company (split into sub-ranges of up to 365 days, summed per company).
|
|
199
|
+
3. Splits date ranges by volume when a company exceeds the chunk limit, honoring `min_period_days`.
|
|
200
|
+
4. Sorts companies by chunk volume and greedy-packs them into entity groups under `MAX_CHUNKS_PER_BASKET` (1000) and `MAX_ENTITIES_IN_ANY_OF` (500).
|
|
201
|
+
5. Classifies each group by `volume_bucket` (`high`, `medium`, `low`, `very_low`).
|
|
202
|
+
6. Returns a plan with the chunk upper bound estimate and per-basket query configurations.
|
|
203
|
+
|
|
204
|
+
### Execution (`execute_search()`)
|
|
205
|
+
|
|
206
|
+
1. Calculates proportional chunks per basket, with a minimum of 1 chunk per basket when `expected_chunks > 0`.
|
|
207
|
+
2. Executes searches in parallel with sliding-window rate limiting and a concurrency semaphore.
|
|
208
|
+
3. Optionally performs a **second pass**, removing already-retrieved entities from each basket and re-querying for the remainder.
|
|
209
|
+
4. Enriches each chunk with `entity_ids` and `primary_entity_id`, and returns the raw documents.
|
|
210
|
+
|
|
211
|
+
## API Reference
|
|
212
|
+
|
|
213
|
+
### `plan_search()`
|
|
214
|
+
|
|
215
|
+
| Parameter | Type | Default | Description |
|
|
216
|
+
|-----------|------|---------|-------------|
|
|
217
|
+
| `universe` | `str \| list[str]` | required | CSV path or list of entity IDs |
|
|
218
|
+
| `start_date` | `str` | required | Start date (YYYY-MM-DD) |
|
|
219
|
+
| `end_date` | `str` | required | End date (YYYY-MM-DD) |
|
|
220
|
+
| `text` | `str \| None` | `None` | Search query text; omitted from all payloads when `None` |
|
|
221
|
+
| `api_key` | `str \| None` | env var | API key |
|
|
222
|
+
| `api_base_url` | `str \| None` | env var | API base URL |
|
|
223
|
+
| `volume_query_mode` | `str` | `"three_pass"` | `"three_pass"` or `"iterative"` |
|
|
224
|
+
| `max_iterations_per_batch` | `int` | `10` | Max iterations per batch in `"iterative"` mode |
|
|
225
|
+
| `apply_volume_splits` | `bool` | `True` | Use volume time series for period splitting |
|
|
226
|
+
| `min_period_days` | `int` | `30` | Minimum days per sub-period |
|
|
227
|
+
| `min_entities_per_basket` | `int` | `1` | Lower bound on entities per basket |
|
|
228
|
+
| `category` | `CategoryFilter \| dict \| None` | `None` | Optional category filter (`mode` + `values`) |
|
|
229
|
+
| `content_diversification` | `bool` | `True` | Enable/disable content diversification in ranking |
|
|
230
|
+
|
|
231
|
+
### `execute_search()`
|
|
232
|
+
|
|
233
|
+
| Parameter | Type | Default | Description |
|
|
234
|
+
|-----------|------|---------|-------------|
|
|
235
|
+
| `search_plan` | `dict` | required | Plan from `plan_search()` |
|
|
236
|
+
| `chunk_percentage` | `float` | required | 0.0 to 1.0 sampling ratio |
|
|
237
|
+
| `requests_per_minute` | `int` | `100` | Rate limit |
|
|
238
|
+
| `api_key` | `str \| None` | env var | API key |
|
|
239
|
+
| `api_base_url` | `str \| None` | env var | API base URL |
|
|
240
|
+
| `max_workers` | `int` | `40` | Parallel workers |
|
|
241
|
+
| `overwrite_chunks_per_basket` | `int \| None` | `None` | Override the proportional chunk allocation |
|
|
242
|
+
| `second_pass` | `bool` | `False` | Re-query each basket for entities not yet retrieved |
|
|
243
|
+
|
|
244
|
+
### Helper Functions
|
|
245
|
+
|
|
246
|
+
- `deduplicate_documents(documents)` -- merge duplicate documents by `id`, combining their chunks.
|
|
247
|
+
- `convert_to_dataframe(raw_results)` -- convert documents to a DataFrame with one row per chunk.
|
|
248
|
+
- `resolve_universe(universe)` -- load entity IDs from a CSV path or validate a `list[str]`.
|
|
249
|
+
- `load_universe_from_csv(csv_path)` -- load entity IDs from a CSV only.
|
|
250
|
+
- `save_plan(plan, path)` / `load_plan(path)` -- persist plans as JSON.
|
|
251
|
+
- `validate_category_filter(category)` -- validate and normalize a category filter to the API dict shape.
|
|
252
|
+
- `build_comention_payload(...)`, `build_volume_payload(...)`, `build_search_query(...)` -- low-level payload builders used by the planner; useful if you need to construct custom requests.
|
|
253
|
+
|
|
254
|
+
### Public Types
|
|
255
|
+
|
|
256
|
+
- `CategoryFilter` -- frozen dataclass with `mode: CategoryMode` and `values: tuple[str, ...]`.
|
|
257
|
+
- `CategoryMode` -- `StrEnum` with `INCLUDE` and `EXCLUDE` values.
|
|
258
|
+
- `CategoryFilterInput` -- type alias: `CategoryFilter | dict[str, str | list[str]]`.
|
|
259
|
+
- `UniverseInput` -- type alias: `str | list[str]`.
|
|
260
|
+
- `VALID_CATEGORY_VALUES` -- `frozenset[str]` of allowed category strings.
|
|
261
|
+
|
|
262
|
+
## Testing
|
|
263
|
+
|
|
264
|
+
```bash
|
|
265
|
+
# Run all tests
|
|
266
|
+
uv run pytest
|
|
267
|
+
|
|
268
|
+
# With coverage
|
|
269
|
+
uv run pytest --cov=bigdata_smart_batching --cov-report=term-missing
|
|
270
|
+
|
|
271
|
+
# Specific test file
|
|
272
|
+
uv run pytest tests/test_validation.py -v
|
|
273
|
+
```
|
|
274
|
+
|
|
275
|
+
## Project Structure
|
|
276
|
+
|
|
277
|
+
```
|
|
278
|
+
bigdata-smart-batching/
|
|
279
|
+
├── pyproject.toml
|
|
280
|
+
├── README.md
|
|
281
|
+
├── CHANGELOG.md
|
|
282
|
+
├── LICENSE
|
|
283
|
+
├── Makefile
|
|
284
|
+
├── .python-version
|
|
285
|
+
├── src/
|
|
286
|
+
│ └── bigdata_smart_batching/
|
|
287
|
+
│ ├── __init__.py
|
|
288
|
+
│ ├── smart_batching.py
|
|
289
|
+
│ ├── smart_batching_config.py
|
|
290
|
+
│ ├── search_function.py
|
|
291
|
+
│ └── output_converter.py
|
|
292
|
+
├── examples/
|
|
293
|
+
│ └── example.py
|
|
294
|
+
└── tests/
|
|
295
|
+
├── __init__.py
|
|
296
|
+
├── test_config.py
|
|
297
|
+
├── test_output_converter.py
|
|
298
|
+
├── test_rate_limiter.py
|
|
299
|
+
└── test_validation.py
|
|
300
|
+
```
|
|
301
|
+
|
|
302
|
+
## Configuration
|
|
303
|
+
|
|
304
|
+
### Environment Variables
|
|
305
|
+
|
|
306
|
+
- `BIGDATA_API_KEY` -- required; your Bigdata API key.
|
|
307
|
+
- `BIGDATA_API_BASE_URL` -- optional; API base URL (default: `https://api.bigdata.com`).
|
|
308
|
+
|
|
309
|
+
### Default Settings
|
|
310
|
+
|
|
311
|
+
- `requests_per_minute`: 100
|
|
312
|
+
- `max_workers`: 40
|
|
313
|
+
- `MAX_CHUNKS_PER_BASKET`: 1000
|
|
314
|
+
- `MAX_ENTITIES_IN_ANY_OF`: 500
|
|
315
|
+
- `volume_query_mode`: `"three_pass"`
|
|
316
|
+
|
|
317
|
+
### Volume Buckets
|
|
318
|
+
|
|
319
|
+
Companies (and basket groups) are classified by total chunk volume. Ranges use
|
|
320
|
+
`[lower, upper)` semantics:
|
|
321
|
+
|
|
322
|
+
| Bucket | Range |
|
|
323
|
+
|--------|-------|
|
|
324
|
+
| `high` | 500+ chunks |
|
|
325
|
+
| `medium` | 100-499 chunks |
|
|
326
|
+
| `low` | 1-99 chunks |
|
|
327
|
+
| `very_low` | 0 chunks (headlines only) |
|
|
328
|
+
|
|
329
|
+
Ties are resolved using `VOLUME_BUCKET_PRIORITY`: `high > medium > low > very_low`.
|
|
330
|
+
|
|
331
|
+
## Changelog
|
|
332
|
+
|
|
333
|
+
See [CHANGELOG.md](CHANGELOG.md) for release notes. The package is currently at
|
|
334
|
+
version **1.0.1**.
|
|
335
|
+
|
|
336
|
+
## License
|
|
337
|
+
|
|
338
|
+
This project is part of Bigdata.com.
|
|
339
|
+
|
|
340
|
+
**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,324 @@
|
|
|
1
|
+
# Smart Batching Search
|
|
2
|
+
|
|
3
|
+
A high-performance semantic search system that reduces API queries by **67-99%** (varies by topic specificity) through intelligent company grouping and parallel execution.
|
|
4
|
+
|
|
5
|
+
This module provides a two-step pipeline for efficient semantic search:
|
|
6
|
+
|
|
7
|
+
1. **Planning** -- organize the search via smart batching and return a chunk upper bound estimate.
|
|
8
|
+
2. **Execution** -- run the plan with proportional sampling to preserve the result distribution.
|
|
9
|
+
|
|
10
|
+
## Key Benefits
|
|
11
|
+
|
|
12
|
+
- **67-99% query reduction** -- search 4,732 companies with only 17-3,699 queries (varies by topic).
|
|
13
|
+
- **Parallel execution** -- rate-limited concurrent requests with semaphore-controlled concurrency.
|
|
14
|
+
- **Proportional sampling** -- retrieve a percentage of results while preserving the distribution across baskets.
|
|
15
|
+
- **Category filtering** -- restrict searches to specific document categories (news, transcripts, filings, ...).
|
|
16
|
+
- **Deterministic plans** -- every request stays within service limits; plans can be saved and replayed.
|
|
17
|
+
- **Scalable** -- handles universes with 10,000+ companies efficiently.
|
|
18
|
+
|
|
19
|
+
## Installation
|
|
20
|
+
|
|
21
|
+
Install the package from PyPI (Python 3.11+):
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
pip install bigdata-smart-batching
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
With [uv](https://docs.astral.sh/uv/):
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
uv add bigdata-smart-batching
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
### Development
|
|
34
|
+
|
|
35
|
+
To work on this repository locally, from the project root:
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
uv sync
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
### Environment Setup
|
|
42
|
+
|
|
43
|
+
Set up environment variables:
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
export BIGDATA_API_KEY="your_api_key_here"
|
|
47
|
+
export BIGDATA_API_BASE_URL="https://api.bigdata.com" # Optional, defaults to this
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
Or create a `.env` file:
|
|
51
|
+
|
|
52
|
+
```
|
|
53
|
+
BIGDATA_API_KEY=your_api_key_here
|
|
54
|
+
BIGDATA_API_BASE_URL=https://api.bigdata.com
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
## Universe (CSV path or entity list)
|
|
58
|
+
|
|
59
|
+
`plan_search()` takes a **`universe`** argument: either a path to a UTF-8 CSV file, or a `list[str]` of entity IDs. IDs must match the identifiers used by the Bigdata API for your dataset.
|
|
60
|
+
|
|
61
|
+
When using a CSV, two layouts are supported:
|
|
62
|
+
|
|
63
|
+
**1. Header row with an `id` column** (optional extra columns such as `name` are ignored):
|
|
64
|
+
|
|
65
|
+
```csv
|
|
66
|
+
id,name
|
|
67
|
+
B8EF97,Example Corp A
|
|
68
|
+
BB07E4,Example Corp B
|
|
69
|
+
3461CF,Example Corp C
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
You can also pass a plain list instead of a CSV path:
|
|
73
|
+
|
|
74
|
+
```python
|
|
75
|
+
plan = plan_search(
|
|
76
|
+
universe=["B8EF97", "BB07E4", "3461CF"],
|
|
77
|
+
start_date="2023-01-01",
|
|
78
|
+
end_date="2023-12-31",
|
|
79
|
+
text="earnings revenue profit",
|
|
80
|
+
)
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
## Quick Start
|
|
84
|
+
|
|
85
|
+
```python
|
|
86
|
+
from bigdata_smart_batching import (
|
|
87
|
+
plan_search,
|
|
88
|
+
execute_search,
|
|
89
|
+
deduplicate_documents,
|
|
90
|
+
convert_to_dataframe,
|
|
91
|
+
)
|
|
92
|
+
|
|
93
|
+
plan = plan_search(
|
|
94
|
+
universe="id_name_mapping_us_top_3000.csv",
|
|
95
|
+
start_date="2023-01-01",
|
|
96
|
+
end_date="2023-12-31",
|
|
97
|
+
text="earnings revenue profit",
|
|
98
|
+
api_key="your_api_key", # or set BIGDATA_API_KEY env var
|
|
99
|
+
)
|
|
100
|
+
|
|
101
|
+
print(f"Chunk upper bound estimate: {plan['chunk_upper_bound_estimate']:,}")
|
|
102
|
+
|
|
103
|
+
results_raw = execute_search(
|
|
104
|
+
search_plan=plan,
|
|
105
|
+
chunk_percentage=0.1,
|
|
106
|
+
requests_per_minute=100,
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
results = deduplicate_documents(results_raw)
|
|
110
|
+
print(f"Retrieved {len(results)} documents (deduplicated)")
|
|
111
|
+
|
|
112
|
+
df = convert_to_dataframe(results) # one row per chunk
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
### Category Filtering
|
|
116
|
+
|
|
117
|
+
Restrict searches to a specific set of document categories using the `category`
|
|
118
|
+
argument. It accepts either a `CategoryFilter` dataclass or a plain dict.
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
from bigdata_smart_batching import CategoryFilter, CategoryMode, plan_search
|
|
122
|
+
|
|
123
|
+
plan = plan_search(
|
|
124
|
+
universe="id_name_mapping_us_top_3000.csv",
|
|
125
|
+
start_date="2023-01-01",
|
|
126
|
+
end_date="2023-12-31",
|
|
127
|
+
text="merger acquisition",
|
|
128
|
+
category=CategoryFilter(
|
|
129
|
+
mode=CategoryMode.INCLUDE,
|
|
130
|
+
values=("news_premium", "transcripts"),
|
|
131
|
+
),
|
|
132
|
+
)
|
|
133
|
+
|
|
134
|
+
# Equivalent dict form:
|
|
135
|
+
plan = plan_search(
|
|
136
|
+
universe="id_name_mapping_us_top_3000.csv",
|
|
137
|
+
start_date="2023-01-01",
|
|
138
|
+
end_date="2023-12-31",
|
|
139
|
+
text="merger acquisition",
|
|
140
|
+
category={"mode": "INCLUDE", "values": ["news_premium", "transcripts"]},
|
|
141
|
+
)
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
Valid category values are exposed via `VALID_CATEGORY_VALUES` and include:
|
|
145
|
+
`expert_interviews`, `filings`, `my_files`, `news`, `news_premium`,
|
|
146
|
+
`news_public`, `podcasts`, `research`, `research_academic_journals`,
|
|
147
|
+
`research_investment_research`, `transcripts`.
|
|
148
|
+
|
|
149
|
+
### Save and Load Plans
|
|
150
|
+
|
|
151
|
+
```python
|
|
152
|
+
from bigdata_smart_batching import plan_search, execute_search, save_plan, load_plan
|
|
153
|
+
|
|
154
|
+
plan = plan_search(
|
|
155
|
+
universe="id_name_mapping_us_top_3000.csv",
|
|
156
|
+
start_date="2023-01-01",
|
|
157
|
+
end_date="2023-12-31",
|
|
158
|
+
text="merger acquisition",
|
|
159
|
+
)
|
|
160
|
+
save_plan(plan, "my_search_plan.json")
|
|
161
|
+
|
|
162
|
+
plan = load_plan("my_search_plan.json")
|
|
163
|
+
raw_10 = execute_search(plan, chunk_percentage=0.1)
|
|
164
|
+
raw_50 = execute_search(plan, chunk_percentage=0.5)
|
|
165
|
+
```
|
|
166
|
+
|
|
167
|
+
## How It Works
|
|
168
|
+
|
|
169
|
+
### Architecture Overview
|
|
170
|
+
|
|
171
|
+
```
|
|
172
|
+
Step 1: PLANNING
|
|
173
|
+
Universe --> Co-mention volumes --> Entity groups (by volume) --> Sub-period split --> Search plan
|
|
174
|
+
|
|
175
|
+
Step 2: EXECUTION
|
|
176
|
+
Proportional sampling --> Parallel search (rate limited) --> Collect & aggregate
|
|
177
|
+
```
|
|
178
|
+
|
|
179
|
+
### Planning (`plan_search()`)
|
|
180
|
+
|
|
181
|
+
1. Loads the universe of companies from a CSV path or inline list.
|
|
182
|
+
2. Queries the co-mention endpoint to get chunk volumes per company (split into sub-ranges of up to 365 days, summed per company).
|
|
183
|
+
3. Splits date ranges by volume when a company exceeds the chunk limit, honoring `min_period_days`.
|
|
184
|
+
4. Sorts companies by chunk volume and greedy-packs them into entity groups under `MAX_CHUNKS_PER_BASKET` (1000) and `MAX_ENTITIES_IN_ANY_OF` (500).
|
|
185
|
+
5. Classifies each group by `volume_bucket` (`high`, `medium`, `low`, `very_low`).
|
|
186
|
+
6. Returns a plan with the chunk upper bound estimate and per-basket query configurations.
|
|
187
|
+
|
|
188
|
+
### Execution (`execute_search()`)
|
|
189
|
+
|
|
190
|
+
1. Calculates proportional chunks per basket, with a minimum of 1 chunk per basket when `expected_chunks > 0`.
|
|
191
|
+
2. Executes searches in parallel with sliding-window rate limiting and a concurrency semaphore.
|
|
192
|
+
3. Optionally performs a **second pass**, removing already-retrieved entities from each basket and re-querying for the remainder.
|
|
193
|
+
4. Enriches each chunk with `entity_ids` and `primary_entity_id`, and returns the raw documents.
|
|
194
|
+
|
|
195
|
+
## API Reference
|
|
196
|
+
|
|
197
|
+
### `plan_search()`
|
|
198
|
+
|
|
199
|
+
| Parameter | Type | Default | Description |
|
|
200
|
+
|-----------|------|---------|-------------|
|
|
201
|
+
| `universe` | `str \| list[str]` | required | CSV path or list of entity IDs |
|
|
202
|
+
| `start_date` | `str` | required | Start date (YYYY-MM-DD) |
|
|
203
|
+
| `end_date` | `str` | required | End date (YYYY-MM-DD) |
|
|
204
|
+
| `text` | `str \| None` | `None` | Search query text; omitted from all payloads when `None` |
|
|
205
|
+
| `api_key` | `str \| None` | env var | API key |
|
|
206
|
+
| `api_base_url` | `str \| None` | env var | API base URL |
|
|
207
|
+
| `volume_query_mode` | `str` | `"three_pass"` | `"three_pass"` or `"iterative"` |
|
|
208
|
+
| `max_iterations_per_batch` | `int` | `10` | Max iterations per batch in `"iterative"` mode |
|
|
209
|
+
| `apply_volume_splits` | `bool` | `True` | Use volume time series for period splitting |
|
|
210
|
+
| `min_period_days` | `int` | `30` | Minimum days per sub-period |
|
|
211
|
+
| `min_entities_per_basket` | `int` | `1` | Lower bound on entities per basket |
|
|
212
|
+
| `category` | `CategoryFilter \| dict \| None` | `None` | Optional category filter (`mode` + `values`) |
|
|
213
|
+
| `content_diversification` | `bool` | `True` | Enable/disable content diversification in ranking |
|
|
214
|
+
|
|
215
|
+
### `execute_search()`
|
|
216
|
+
|
|
217
|
+
| Parameter | Type | Default | Description |
|
|
218
|
+
|-----------|------|---------|-------------|
|
|
219
|
+
| `search_plan` | `dict` | required | Plan from `plan_search()` |
|
|
220
|
+
| `chunk_percentage` | `float` | required | 0.0 to 1.0 sampling ratio |
|
|
221
|
+
| `requests_per_minute` | `int` | `100` | Rate limit |
|
|
222
|
+
| `api_key` | `str \| None` | env var | API key |
|
|
223
|
+
| `api_base_url` | `str \| None` | env var | API base URL |
|
|
224
|
+
| `max_workers` | `int` | `40` | Parallel workers |
|
|
225
|
+
| `overwrite_chunks_per_basket` | `int \| None` | `None` | Override the proportional chunk allocation |
|
|
226
|
+
| `second_pass` | `bool` | `False` | Re-query each basket for entities not yet retrieved |
|
|
227
|
+
|
|
228
|
+
### Helper Functions
|
|
229
|
+
|
|
230
|
+
- `deduplicate_documents(documents)` -- merge duplicate documents by `id`, combining their chunks.
|
|
231
|
+
- `convert_to_dataframe(raw_results)` -- convert documents to a DataFrame with one row per chunk.
|
|
232
|
+
- `resolve_universe(universe)` -- load entity IDs from a CSV path or validate a `list[str]`.
|
|
233
|
+
- `load_universe_from_csv(csv_path)` -- load entity IDs from a CSV only.
|
|
234
|
+
- `save_plan(plan, path)` / `load_plan(path)` -- persist plans as JSON.
|
|
235
|
+
- `validate_category_filter(category)` -- validate and normalize a category filter to the API dict shape.
|
|
236
|
+
- `build_comention_payload(...)`, `build_volume_payload(...)`, `build_search_query(...)` -- low-level payload builders used by the planner; useful if you need to construct custom requests.
|
|
237
|
+
|
|
238
|
+
### Public Types
|
|
239
|
+
|
|
240
|
+
- `CategoryFilter` -- frozen dataclass with `mode: CategoryMode` and `values: tuple[str, ...]`.
|
|
241
|
+
- `CategoryMode` -- `StrEnum` with `INCLUDE` and `EXCLUDE` values.
|
|
242
|
+
- `CategoryFilterInput` -- type alias: `CategoryFilter | dict[str, str | list[str]]`.
|
|
243
|
+
- `UniverseInput` -- type alias: `str | list[str]`.
|
|
244
|
+
- `VALID_CATEGORY_VALUES` -- `frozenset[str]` of allowed category strings.
|
|
245
|
+
|
|
246
|
+
## Testing
|
|
247
|
+
|
|
248
|
+
```bash
|
|
249
|
+
# Run all tests
|
|
250
|
+
uv run pytest
|
|
251
|
+
|
|
252
|
+
# With coverage
|
|
253
|
+
uv run pytest --cov=bigdata_smart_batching --cov-report=term-missing
|
|
254
|
+
|
|
255
|
+
# Specific test file
|
|
256
|
+
uv run pytest tests/test_validation.py -v
|
|
257
|
+
```
|
|
258
|
+
|
|
259
|
+
## Project Structure
|
|
260
|
+
|
|
261
|
+
```
|
|
262
|
+
bigdata-smart-batching/
|
|
263
|
+
├── pyproject.toml
|
|
264
|
+
├── README.md
|
|
265
|
+
├── CHANGELOG.md
|
|
266
|
+
├── LICENSE
|
|
267
|
+
├── Makefile
|
|
268
|
+
├── .python-version
|
|
269
|
+
├── src/
|
|
270
|
+
│ └── bigdata_smart_batching/
|
|
271
|
+
│ ├── __init__.py
|
|
272
|
+
│ ├── smart_batching.py
|
|
273
|
+
│ ├── smart_batching_config.py
|
|
274
|
+
│ ├── search_function.py
|
|
275
|
+
│ └── output_converter.py
|
|
276
|
+
├── examples/
|
|
277
|
+
│ └── example.py
|
|
278
|
+
└── tests/
|
|
279
|
+
├── __init__.py
|
|
280
|
+
├── test_config.py
|
|
281
|
+
├── test_output_converter.py
|
|
282
|
+
├── test_rate_limiter.py
|
|
283
|
+
└── test_validation.py
|
|
284
|
+
```
|
|
285
|
+
|
|
286
|
+
## Configuration
|
|
287
|
+
|
|
288
|
+
### Environment Variables
|
|
289
|
+
|
|
290
|
+
- `BIGDATA_API_KEY` -- required; your Bigdata API key.
|
|
291
|
+
- `BIGDATA_API_BASE_URL` -- optional; API base URL (default: `https://api.bigdata.com`).
|
|
292
|
+
|
|
293
|
+
### Default Settings
|
|
294
|
+
|
|
295
|
+
- `requests_per_minute`: 100
|
|
296
|
+
- `max_workers`: 40
|
|
297
|
+
- `MAX_CHUNKS_PER_BASKET`: 1000
|
|
298
|
+
- `MAX_ENTITIES_IN_ANY_OF`: 500
|
|
299
|
+
- `volume_query_mode`: `"three_pass"`
|
|
300
|
+
|
|
301
|
+
### Volume Buckets
|
|
302
|
+
|
|
303
|
+
Companies (and basket groups) are classified by total chunk volume. Ranges use
|
|
304
|
+
`[lower, upper)` semantics:
|
|
305
|
+
|
|
306
|
+
| Bucket | Range |
|
|
307
|
+
|--------|-------|
|
|
308
|
+
| `high` | 500+ chunks |
|
|
309
|
+
| `medium` | 100-499 chunks |
|
|
310
|
+
| `low` | 1-99 chunks |
|
|
311
|
+
| `very_low` | 0 chunks (headlines only) |
|
|
312
|
+
|
|
313
|
+
Ties are resolved using `VOLUME_BUCKET_PRIORITY`: `high > medium > low > very_low`.
|
|
314
|
+
|
|
315
|
+
## Changelog
|
|
316
|
+
|
|
317
|
+
See [CHANGELOG.md](CHANGELOG.md) for release notes. The package is currently at
|
|
318
|
+
version **1.0.1**.
|
|
319
|
+
|
|
320
|
+
## License
|
|
321
|
+
|
|
322
|
+
This project is part of Bigdata.com.
|
|
323
|
+
|
|
324
|
+
**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.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "bigdata-smart-batching"
|
|
3
|
-
version = "1.0.
|
|
3
|
+
version = "1.0.2"
|
|
4
4
|
description = "High-performance semantic search with intelligent company grouping and parallel execution"
|
|
5
5
|
readme = "README.md"
|
|
6
6
|
authors = [{ name = "Bigdata.com", email = "support@ravenpack.com" }]
|
|
@@ -1,257 +0,0 @@
|
|
|
1
|
-
Metadata-Version: 2.4
|
|
2
|
-
Name: bigdata-smart-batching
|
|
3
|
-
Version: 1.0.1
|
|
4
|
-
Summary: High-performance semantic search with intelligent company grouping and parallel execution
|
|
5
|
-
Project-URL: homepage, https://bigdata.com/api
|
|
6
|
-
Author-email: "Bigdata.com" <support@ravenpack.com>
|
|
7
|
-
License: MIT
|
|
8
|
-
License-File: LICENSE
|
|
9
|
-
Requires-Python: <4.0,>=3.11
|
|
10
|
-
Requires-Dist: matplotlib>=3.7
|
|
11
|
-
Requires-Dist: numpy>=1.24
|
|
12
|
-
Requires-Dist: pandas>=2.0
|
|
13
|
-
Requires-Dist: python-dotenv>=1.2.2
|
|
14
|
-
Requires-Dist: requests>=2.31.0
|
|
15
|
-
Description-Content-Type: text/markdown
|
|
16
|
-
|
|
17
|
-
# Smart Batching Search
|
|
18
|
-
|
|
19
|
-
A high-performance semantic search system that reduces API queries by **67-99%** (varies by topic specificity) through intelligent company grouping and parallel execution.
|
|
20
|
-
|
|
21
|
-
This module provides a two-step system for efficient semantic search:
|
|
22
|
-
1. **Planning**: Organize search using smart batching and return Chunk upper bound estimate
|
|
23
|
-
2. **Execution**: Perform search with proportional sampling to preserve distribution
|
|
24
|
-
|
|
25
|
-
## Key Benefits
|
|
26
|
-
|
|
27
|
-
- **67-99% Query Reduction**: Search 4,732 companies with only 17-3,699 queries (varies by topic)
|
|
28
|
-
- **Parallel Execution**: Rate-limited concurrent requests with semaphore control
|
|
29
|
-
- **Proportional Sampling**: Retrieve percentage of results while preserving distribution
|
|
30
|
-
- **Production Ready**: Comprehensive error handling, retries, and logging
|
|
31
|
-
- **Scalable**: Efficiently handles universes with 10,000+ companies
|
|
32
|
-
|
|
33
|
-
## Installation
|
|
34
|
-
|
|
35
|
-
Install the package from PyPI (Python 3.11+):
|
|
36
|
-
|
|
37
|
-
```bash
|
|
38
|
-
pip install bigdata-smart-batching
|
|
39
|
-
```
|
|
40
|
-
|
|
41
|
-
With [uv](https://docs.astral.sh/uv/):
|
|
42
|
-
|
|
43
|
-
```bash
|
|
44
|
-
uv add bigdata-smart-batching
|
|
45
|
-
```
|
|
46
|
-
|
|
47
|
-
### Development
|
|
48
|
-
|
|
49
|
-
To work on this repository locally, from the project root:
|
|
50
|
-
|
|
51
|
-
```bash
|
|
52
|
-
uv sync
|
|
53
|
-
```
|
|
54
|
-
|
|
55
|
-
### Environment Setup
|
|
56
|
-
|
|
57
|
-
Set up environment variables:
|
|
58
|
-
|
|
59
|
-
```bash
|
|
60
|
-
export BIGDATA_API_KEY="your_api_key_here"
|
|
61
|
-
export BIGDATA_API_BASE_URL="https://api.bigdata.com" # Optional, defaults to this
|
|
62
|
-
```
|
|
63
|
-
|
|
64
|
-
Or create a `.env` file:
|
|
65
|
-
|
|
66
|
-
```
|
|
67
|
-
BIGDATA_API_KEY=your_api_key_here
|
|
68
|
-
BIGDATA_API_BASE_URL=https://api.bigdata.com
|
|
69
|
-
```
|
|
70
|
-
|
|
71
|
-
## Universe (CSV path or entity list)
|
|
72
|
-
|
|
73
|
-
`plan_search()` takes a **`universe`** argument: either a path to a UTF-8 CSV file, or a `list[str]` of entity IDs. IDs must match the identifiers used by the Bigdata API for your dataset.
|
|
74
|
-
|
|
75
|
-
When using a CSV, two layouts are supported:
|
|
76
|
-
|
|
77
|
-
**1. Header row with an `id` column** (optional extra columns such as `name` are ignored):
|
|
78
|
-
|
|
79
|
-
```csv
|
|
80
|
-
id,name
|
|
81
|
-
B8EF97,Example Corp A
|
|
82
|
-
BB07E4,Example Corp B
|
|
83
|
-
3461CF,Example Corp C
|
|
84
|
-
```
|
|
85
|
-
|
|
86
|
-
## Quick Start
|
|
87
|
-
|
|
88
|
-
```python
|
|
89
|
-
from bigdata_smart_batching import (
|
|
90
|
-
plan_search,
|
|
91
|
-
execute_search,
|
|
92
|
-
deduplicate_documents,
|
|
93
|
-
convert_to_dataframe,
|
|
94
|
-
)
|
|
95
|
-
|
|
96
|
-
# Step 1: Plan the search
|
|
97
|
-
plan = plan_search(
|
|
98
|
-
text="earnings revenue profit",
|
|
99
|
-
universe="id_name_mapping_us_top_3000.csv",
|
|
100
|
-
start_date="2023-01-01",
|
|
101
|
-
end_date="2023-12-31",
|
|
102
|
-
api_key="your_api_key", # or set BIGDATA_API_KEY env var
|
|
103
|
-
)
|
|
104
|
-
|
|
105
|
-
print(f"Chunk upper bound estimate: {plan['chunk_upper_bound_estimate']:,}")
|
|
106
|
-
|
|
107
|
-
# Step 2: Execute search with 10% of total chunks (preserves distribution)
|
|
108
|
-
results_raw = execute_search(
|
|
109
|
-
search_plan=plan,
|
|
110
|
-
chunk_percentage=0.1,
|
|
111
|
-
requests_per_minute=100,
|
|
112
|
-
)
|
|
113
|
-
|
|
114
|
-
# Step 3: Deduplicate and convert to DataFrame
|
|
115
|
-
results = deduplicate_documents(results_raw)
|
|
116
|
-
print(f"Retrieved {len(results)} documents (deduplicated)")
|
|
117
|
-
|
|
118
|
-
df = convert_to_dataframe(results) # one row per chunk
|
|
119
|
-
```
|
|
120
|
-
|
|
121
|
-
### Save and Load Plans
|
|
122
|
-
|
|
123
|
-
```python
|
|
124
|
-
from bigdata_smart_batching import plan_search, execute_search, save_plan, load_plan
|
|
125
|
-
|
|
126
|
-
# Create and save a plan
|
|
127
|
-
plan = plan_search(
|
|
128
|
-
text="merger acquisition",
|
|
129
|
-
universe="id_name_mapping_us_top_3000.csv",
|
|
130
|
-
start_date="2023-01-01",
|
|
131
|
-
end_date="2023-12-31",
|
|
132
|
-
)
|
|
133
|
-
save_plan(plan, "my_search_plan.json")
|
|
134
|
-
|
|
135
|
-
# Later: reload and run with different sampling
|
|
136
|
-
plan = load_plan("my_search_plan.json")
|
|
137
|
-
raw_10 = execute_search(plan, chunk_percentage=0.1)
|
|
138
|
-
raw_50 = execute_search(plan, chunk_percentage=0.5)
|
|
139
|
-
```
|
|
140
|
-
|
|
141
|
-
## How It Works
|
|
142
|
-
|
|
143
|
-
### Architecture Overview
|
|
144
|
-
|
|
145
|
-
```
|
|
146
|
-
Step 1: PLANNING
|
|
147
|
-
Universe CSV --> Co-mention API Query --> Basket Creation --> Search Plan
|
|
148
|
-
|
|
149
|
-
Step 2: EXECUTION
|
|
150
|
-
Proportional Sampling --> Parallel Search (Rate Limited) --> Collect & Aggregate
|
|
151
|
-
```
|
|
152
|
-
|
|
153
|
-
### Planning (`plan_search()`)
|
|
154
|
-
|
|
155
|
-
1. Loads the universe of companies from a CSV path or inline list
|
|
156
|
-
2. Queries the comention endpoint to get chunk volumes per company
|
|
157
|
-
3. Splits date ranges by volume when a company exceeds the chunk limit
|
|
158
|
-
4. Creates optimized baskets grouped by volume
|
|
159
|
-
5. Returns a plan with Chunk upper bound estimate and basket configurations
|
|
160
|
-
|
|
161
|
-
### Execution (`execute_search()`)
|
|
162
|
-
|
|
163
|
-
1. Calculates proportional chunks per basket
|
|
164
|
-
2. Ensures minimum of 1 chunk per basket (if expected > 0)
|
|
165
|
-
3. Executes searches in parallel with rate limiting and semaphore
|
|
166
|
-
4. Collects and returns document results
|
|
167
|
-
|
|
168
|
-
## API Reference
|
|
169
|
-
|
|
170
|
-
### `plan_search()`
|
|
171
|
-
|
|
172
|
-
| Parameter | Type | Default | Description |
|
|
173
|
-
|-----------|------|---------|-------------|
|
|
174
|
-
| `text` | `str` | required | Search query text |
|
|
175
|
-
| `universe` | `str \| list[str]` | required | CSV path or list of entity IDs |
|
|
176
|
-
| `start_date` | `str` | required | Start date (YYYY-MM-DD) |
|
|
177
|
-
| `end_date` | `str` | required | End date (YYYY-MM-DD) |
|
|
178
|
-
| `api_key` | `str` | env var | API key |
|
|
179
|
-
| `api_base_url` | `str` | env var | API base URL |
|
|
180
|
-
| `volume_query_mode` | `str` | `"three_pass"` | `"three_pass"` or `"iterative"` |
|
|
181
|
-
| `apply_volume_splits` | `bool` | `True` | Use volume time series for period splitting |
|
|
182
|
-
| `min_period_days` | `int` | `30` | Minimum days per sub-period |
|
|
183
|
-
|
|
184
|
-
### `execute_search()`
|
|
185
|
-
|
|
186
|
-
| Parameter | Type | Default | Description |
|
|
187
|
-
|-----------|------|---------|-------------|
|
|
188
|
-
| `search_plan` | `Dict` | required | Plan from `plan_search()` |
|
|
189
|
-
| `chunk_percentage` | `float` | required | 0.0 to 1.0 sampling ratio |
|
|
190
|
-
| `requests_per_minute` | `int` | `100` | Rate limit |
|
|
191
|
-
| `api_key` | `str` | env var | API key |
|
|
192
|
-
| `max_workers` | `int` | `40` | Parallel workers |
|
|
193
|
-
|
|
194
|
-
### Helper Functions
|
|
195
|
-
|
|
196
|
-
- **`deduplicate_documents(documents)`** -- Merges duplicate documents by `id`
|
|
197
|
-
- **`resolve_universe(universe)`** -- Loads entity IDs from a CSV path or validates a `list[str]`
|
|
198
|
-
- **`load_universe_from_csv(csv_path)`** -- Loads entity IDs from CSV only
|
|
199
|
-
- **`convert_to_dataframe(raw_results)`** -- Converts documents to DataFrame (one row per chunk)
|
|
200
|
-
- **`save_plan(plan, path)`** / **`load_plan(path)`** -- Persist plans as JSON
|
|
201
|
-
- **`portfolio_backtesting_pipeline(...)`** -- Long-short portfolio backtesting
|
|
202
|
-
|
|
203
|
-
## Testing
|
|
204
|
-
|
|
205
|
-
```bash
|
|
206
|
-
# Run all tests
|
|
207
|
-
uv run pytest
|
|
208
|
-
|
|
209
|
-
# With coverage
|
|
210
|
-
uv run pytest --cov=bigdata_smart_batching --cov-report=term-missing
|
|
211
|
-
|
|
212
|
-
# Specific test file
|
|
213
|
-
uv run pytest tests/test_validation.py -v
|
|
214
|
-
```
|
|
215
|
-
|
|
216
|
-
## Project Structure
|
|
217
|
-
|
|
218
|
-
```
|
|
219
|
-
bigdata-smart-batching/
|
|
220
|
-
├── pyproject.toml
|
|
221
|
-
├── README.md
|
|
222
|
-
├── .python-version
|
|
223
|
-
├── src/
|
|
224
|
-
│ └── bigdata_smart_batching/
|
|
225
|
-
│ ├── __init__.py
|
|
226
|
-
│ ├── smart_batching.py
|
|
227
|
-
│ ├── smart_batching_config.py
|
|
228
|
-
│ ├── search_function.py
|
|
229
|
-
│ ├── output_converter.py
|
|
230
|
-
│ └── portfolio_backtesting.py
|
|
231
|
-
└── tests/
|
|
232
|
-
├── __init__.py
|
|
233
|
-
├── test_config.py
|
|
234
|
-
├── test_output_converter.py
|
|
235
|
-
├── test_validation.py
|
|
236
|
-
└── test_rate_limiter.py
|
|
237
|
-
```
|
|
238
|
-
|
|
239
|
-
## Configuration
|
|
240
|
-
|
|
241
|
-
### Environment Variables
|
|
242
|
-
|
|
243
|
-
- `BIGDATA_API_KEY`: Required -- Your Bigdata API key
|
|
244
|
-
- `BIGDATA_API_BASE_URL`: Optional -- API base URL (default: `https://api.bigdata.com`)
|
|
245
|
-
|
|
246
|
-
### Default Settings
|
|
247
|
-
|
|
248
|
-
- `requests_per_minute`: 100
|
|
249
|
-
- `max_workers`: 40
|
|
250
|
-
- `max_chunks_per_basket`: 1000
|
|
251
|
-
- `volume_query_mode`: `"three_pass"`
|
|
252
|
-
|
|
253
|
-
## License
|
|
254
|
-
|
|
255
|
-
This project is part of Bigdata.com.
|
|
256
|
-
|
|
257
|
-
**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.
|
|
@@ -1,241 +0,0 @@
|
|
|
1
|
-
# Smart Batching Search
|
|
2
|
-
|
|
3
|
-
A high-performance semantic search system that reduces API queries by **67-99%** (varies by topic specificity) through intelligent company grouping and parallel execution.
|
|
4
|
-
|
|
5
|
-
This module provides a two-step system for efficient semantic search:
|
|
6
|
-
1. **Planning**: Organize search using smart batching and return Chunk upper bound estimate
|
|
7
|
-
2. **Execution**: Perform search with proportional sampling to preserve distribution
|
|
8
|
-
|
|
9
|
-
## Key Benefits
|
|
10
|
-
|
|
11
|
-
- **67-99% Query Reduction**: Search 4,732 companies with only 17-3,699 queries (varies by topic)
|
|
12
|
-
- **Parallel Execution**: Rate-limited concurrent requests with semaphore control
|
|
13
|
-
- **Proportional Sampling**: Retrieve percentage of results while preserving distribution
|
|
14
|
-
- **Production Ready**: Comprehensive error handling, retries, and logging
|
|
15
|
-
- **Scalable**: Efficiently handles universes with 10,000+ companies
|
|
16
|
-
|
|
17
|
-
## Installation
|
|
18
|
-
|
|
19
|
-
Install the package from PyPI (Python 3.11+):
|
|
20
|
-
|
|
21
|
-
```bash
|
|
22
|
-
pip install bigdata-smart-batching
|
|
23
|
-
```
|
|
24
|
-
|
|
25
|
-
With [uv](https://docs.astral.sh/uv/):
|
|
26
|
-
|
|
27
|
-
```bash
|
|
28
|
-
uv add bigdata-smart-batching
|
|
29
|
-
```
|
|
30
|
-
|
|
31
|
-
### Development
|
|
32
|
-
|
|
33
|
-
To work on this repository locally, from the project root:
|
|
34
|
-
|
|
35
|
-
```bash
|
|
36
|
-
uv sync
|
|
37
|
-
```
|
|
38
|
-
|
|
39
|
-
### Environment Setup
|
|
40
|
-
|
|
41
|
-
Set up environment variables:
|
|
42
|
-
|
|
43
|
-
```bash
|
|
44
|
-
export BIGDATA_API_KEY="your_api_key_here"
|
|
45
|
-
export BIGDATA_API_BASE_URL="https://api.bigdata.com" # Optional, defaults to this
|
|
46
|
-
```
|
|
47
|
-
|
|
48
|
-
Or create a `.env` file:
|
|
49
|
-
|
|
50
|
-
```
|
|
51
|
-
BIGDATA_API_KEY=your_api_key_here
|
|
52
|
-
BIGDATA_API_BASE_URL=https://api.bigdata.com
|
|
53
|
-
```
|
|
54
|
-
|
|
55
|
-
## Universe (CSV path or entity list)
|
|
56
|
-
|
|
57
|
-
`plan_search()` takes a **`universe`** argument: either a path to a UTF-8 CSV file, or a `list[str]` of entity IDs. IDs must match the identifiers used by the Bigdata API for your dataset.
|
|
58
|
-
|
|
59
|
-
When using a CSV, two layouts are supported:
|
|
60
|
-
|
|
61
|
-
**1. Header row with an `id` column** (optional extra columns such as `name` are ignored):
|
|
62
|
-
|
|
63
|
-
```csv
|
|
64
|
-
id,name
|
|
65
|
-
B8EF97,Example Corp A
|
|
66
|
-
BB07E4,Example Corp B
|
|
67
|
-
3461CF,Example Corp C
|
|
68
|
-
```
|
|
69
|
-
|
|
70
|
-
## Quick Start
|
|
71
|
-
|
|
72
|
-
```python
|
|
73
|
-
from bigdata_smart_batching import (
|
|
74
|
-
plan_search,
|
|
75
|
-
execute_search,
|
|
76
|
-
deduplicate_documents,
|
|
77
|
-
convert_to_dataframe,
|
|
78
|
-
)
|
|
79
|
-
|
|
80
|
-
# Step 1: Plan the search
|
|
81
|
-
plan = plan_search(
|
|
82
|
-
text="earnings revenue profit",
|
|
83
|
-
universe="id_name_mapping_us_top_3000.csv",
|
|
84
|
-
start_date="2023-01-01",
|
|
85
|
-
end_date="2023-12-31",
|
|
86
|
-
api_key="your_api_key", # or set BIGDATA_API_KEY env var
|
|
87
|
-
)
|
|
88
|
-
|
|
89
|
-
print(f"Chunk upper bound estimate: {plan['chunk_upper_bound_estimate']:,}")
|
|
90
|
-
|
|
91
|
-
# Step 2: Execute search with 10% of total chunks (preserves distribution)
|
|
92
|
-
results_raw = execute_search(
|
|
93
|
-
search_plan=plan,
|
|
94
|
-
chunk_percentage=0.1,
|
|
95
|
-
requests_per_minute=100,
|
|
96
|
-
)
|
|
97
|
-
|
|
98
|
-
# Step 3: Deduplicate and convert to DataFrame
|
|
99
|
-
results = deduplicate_documents(results_raw)
|
|
100
|
-
print(f"Retrieved {len(results)} documents (deduplicated)")
|
|
101
|
-
|
|
102
|
-
df = convert_to_dataframe(results) # one row per chunk
|
|
103
|
-
```
|
|
104
|
-
|
|
105
|
-
### Save and Load Plans
|
|
106
|
-
|
|
107
|
-
```python
|
|
108
|
-
from bigdata_smart_batching import plan_search, execute_search, save_plan, load_plan
|
|
109
|
-
|
|
110
|
-
# Create and save a plan
|
|
111
|
-
plan = plan_search(
|
|
112
|
-
text="merger acquisition",
|
|
113
|
-
universe="id_name_mapping_us_top_3000.csv",
|
|
114
|
-
start_date="2023-01-01",
|
|
115
|
-
end_date="2023-12-31",
|
|
116
|
-
)
|
|
117
|
-
save_plan(plan, "my_search_plan.json")
|
|
118
|
-
|
|
119
|
-
# Later: reload and run with different sampling
|
|
120
|
-
plan = load_plan("my_search_plan.json")
|
|
121
|
-
raw_10 = execute_search(plan, chunk_percentage=0.1)
|
|
122
|
-
raw_50 = execute_search(plan, chunk_percentage=0.5)
|
|
123
|
-
```
|
|
124
|
-
|
|
125
|
-
## How It Works
|
|
126
|
-
|
|
127
|
-
### Architecture Overview
|
|
128
|
-
|
|
129
|
-
```
|
|
130
|
-
Step 1: PLANNING
|
|
131
|
-
Universe CSV --> Co-mention API Query --> Basket Creation --> Search Plan
|
|
132
|
-
|
|
133
|
-
Step 2: EXECUTION
|
|
134
|
-
Proportional Sampling --> Parallel Search (Rate Limited) --> Collect & Aggregate
|
|
135
|
-
```
|
|
136
|
-
|
|
137
|
-
### Planning (`plan_search()`)
|
|
138
|
-
|
|
139
|
-
1. Loads the universe of companies from a CSV path or inline list
|
|
140
|
-
2. Queries the comention endpoint to get chunk volumes per company
|
|
141
|
-
3. Splits date ranges by volume when a company exceeds the chunk limit
|
|
142
|
-
4. Creates optimized baskets grouped by volume
|
|
143
|
-
5. Returns a plan with Chunk upper bound estimate and basket configurations
|
|
144
|
-
|
|
145
|
-
### Execution (`execute_search()`)
|
|
146
|
-
|
|
147
|
-
1. Calculates proportional chunks per basket
|
|
148
|
-
2. Ensures minimum of 1 chunk per basket (if expected > 0)
|
|
149
|
-
3. Executes searches in parallel with rate limiting and semaphore
|
|
150
|
-
4. Collects and returns document results
|
|
151
|
-
|
|
152
|
-
## API Reference
|
|
153
|
-
|
|
154
|
-
### `plan_search()`
|
|
155
|
-
|
|
156
|
-
| Parameter | Type | Default | Description |
|
|
157
|
-
|-----------|------|---------|-------------|
|
|
158
|
-
| `text` | `str` | required | Search query text |
|
|
159
|
-
| `universe` | `str \| list[str]` | required | CSV path or list of entity IDs |
|
|
160
|
-
| `start_date` | `str` | required | Start date (YYYY-MM-DD) |
|
|
161
|
-
| `end_date` | `str` | required | End date (YYYY-MM-DD) |
|
|
162
|
-
| `api_key` | `str` | env var | API key |
|
|
163
|
-
| `api_base_url` | `str` | env var | API base URL |
|
|
164
|
-
| `volume_query_mode` | `str` | `"three_pass"` | `"three_pass"` or `"iterative"` |
|
|
165
|
-
| `apply_volume_splits` | `bool` | `True` | Use volume time series for period splitting |
|
|
166
|
-
| `min_period_days` | `int` | `30` | Minimum days per sub-period |
|
|
167
|
-
|
|
168
|
-
### `execute_search()`
|
|
169
|
-
|
|
170
|
-
| Parameter | Type | Default | Description |
|
|
171
|
-
|-----------|------|---------|-------------|
|
|
172
|
-
| `search_plan` | `Dict` | required | Plan from `plan_search()` |
|
|
173
|
-
| `chunk_percentage` | `float` | required | 0.0 to 1.0 sampling ratio |
|
|
174
|
-
| `requests_per_minute` | `int` | `100` | Rate limit |
|
|
175
|
-
| `api_key` | `str` | env var | API key |
|
|
176
|
-
| `max_workers` | `int` | `40` | Parallel workers |
|
|
177
|
-
|
|
178
|
-
### Helper Functions
|
|
179
|
-
|
|
180
|
-
- **`deduplicate_documents(documents)`** -- Merges duplicate documents by `id`
|
|
181
|
-
- **`resolve_universe(universe)`** -- Loads entity IDs from a CSV path or validates a `list[str]`
|
|
182
|
-
- **`load_universe_from_csv(csv_path)`** -- Loads entity IDs from CSV only
|
|
183
|
-
- **`convert_to_dataframe(raw_results)`** -- Converts documents to DataFrame (one row per chunk)
|
|
184
|
-
- **`save_plan(plan, path)`** / **`load_plan(path)`** -- Persist plans as JSON
|
|
185
|
-
- **`portfolio_backtesting_pipeline(...)`** -- Long-short portfolio backtesting
|
|
186
|
-
|
|
187
|
-
## Testing
|
|
188
|
-
|
|
189
|
-
```bash
|
|
190
|
-
# Run all tests
|
|
191
|
-
uv run pytest
|
|
192
|
-
|
|
193
|
-
# With coverage
|
|
194
|
-
uv run pytest --cov=bigdata_smart_batching --cov-report=term-missing
|
|
195
|
-
|
|
196
|
-
# Specific test file
|
|
197
|
-
uv run pytest tests/test_validation.py -v
|
|
198
|
-
```
|
|
199
|
-
|
|
200
|
-
## Project Structure
|
|
201
|
-
|
|
202
|
-
```
|
|
203
|
-
bigdata-smart-batching/
|
|
204
|
-
├── pyproject.toml
|
|
205
|
-
├── README.md
|
|
206
|
-
├── .python-version
|
|
207
|
-
├── src/
|
|
208
|
-
│ └── bigdata_smart_batching/
|
|
209
|
-
│ ├── __init__.py
|
|
210
|
-
│ ├── smart_batching.py
|
|
211
|
-
│ ├── smart_batching_config.py
|
|
212
|
-
│ ├── search_function.py
|
|
213
|
-
│ ├── output_converter.py
|
|
214
|
-
│ └── portfolio_backtesting.py
|
|
215
|
-
└── tests/
|
|
216
|
-
├── __init__.py
|
|
217
|
-
├── test_config.py
|
|
218
|
-
├── test_output_converter.py
|
|
219
|
-
├── test_validation.py
|
|
220
|
-
└── test_rate_limiter.py
|
|
221
|
-
```
|
|
222
|
-
|
|
223
|
-
## Configuration
|
|
224
|
-
|
|
225
|
-
### Environment Variables
|
|
226
|
-
|
|
227
|
-
- `BIGDATA_API_KEY`: Required -- Your Bigdata API key
|
|
228
|
-
- `BIGDATA_API_BASE_URL`: Optional -- API base URL (default: `https://api.bigdata.com`)
|
|
229
|
-
|
|
230
|
-
### Default Settings
|
|
231
|
-
|
|
232
|
-
- `requests_per_minute`: 100
|
|
233
|
-
- `max_workers`: 40
|
|
234
|
-
- `max_chunks_per_basket`: 1000
|
|
235
|
-
- `volume_query_mode`: `"three_pass"`
|
|
236
|
-
|
|
237
|
-
## License
|
|
238
|
-
|
|
239
|
-
This project is part of Bigdata.com.
|
|
240
|
-
|
|
241
|
-
**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.
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/docs/smart-batching-architecture.md
RENAMED
|
File without changes
|
|
File without changes
|
{bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/examples/example_overlapping.py
RENAMED
|
File without changes
|
{bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/src/bigdata_smart_batching/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{bigdata_smart_batching-1.0.1 → bigdata_smart_batching-1.0.2}/tests/test_output_converter.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|