Smart Batching Search

A high-performance semantic search system that reduces API queries by 67-99% (varies by topic specificity) through intelligent company grouping and parallel execution.

This module provides a two-step pipeline for efficient semantic search:

1. Planning -- organize the search via smart batching and return a chunk upper bound estimate.
2. Execution -- run the plan with proportional sampling to preserve the result distribution.

Key Benefits

• 67-99% query reduction -- search 4,732 companies with only 17-3,699 queries (varies by topic).
• Parallel execution -- rate-limited concurrent requests with semaphore-controlled concurrency.
• Proportional sampling -- retrieve a percentage of results while preserving the distribution across baskets.
• Category filtering -- restrict searches to specific document categories (news, transcripts, filings, ...).
• Deterministic plans -- every request stays within service limits; plans can be saved and replayed.
• Scalable -- handles universes with 10,000+ companies efficiently.

Installation

Install the package from PyPI (Python 3.11+):

pip install bigdata-smart-batching

With uv:

uv add bigdata-smart-batching

Development

To work on this repository locally, from the project root:

uv sync

Environment Setup

Set up environment variables:

export BIGDATA_API_KEY="your_api_key_here"
export BIGDATA_API_BASE_URL="https://api.bigdata.com"  # Optional, defaults to this

Or create a .env file:

BIGDATA_API_KEY=your_api_key_here
BIGDATA_API_BASE_URL=https://api.bigdata.com

Universe (CSV path or entity list)

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.

When using a CSV, two layouts are supported:

1. Header row with an id column (optional extra columns such as name are ignored):

id,name
B8EF97,Example Corp A
BB07E4,Example Corp B
3461CF,Example Corp C

You can also pass a plain list instead of a CSV path:

plan = plan_search(
    universe=["B8EF97", "BB07E4", "3461CF"],
    start_date="2023-01-01",
    end_date="2023-12-31",
    text="earnings revenue profit",
)

Quick Start

from bigdata_smart_batching import (
    plan_search,
    execute_search,
    deduplicate_documents,
    convert_to_dataframe,
)

plan = plan_search(
    universe="id_name_mapping_us_top_3000.csv",
    start_date="2023-01-01",
    end_date="2023-12-31",
    text="earnings revenue profit",
    api_key="your_api_key",  # or set BIGDATA_API_KEY env var
)

print(f"Chunk upper bound estimate: {plan['chunk_upper_bound_estimate']:,}")

results_raw = execute_search(
    search_plan=plan,
    chunk_percentage=0.1,
    requests_per_minute=100,
)

results = deduplicate_documents(results_raw)
print(f"Retrieved {len(results)} documents (deduplicated)")

df = convert_to_dataframe(results)  # one row per chunk

Category Filtering

Restrict searches to a specific set of document categories using the category argument. It accepts either a CategoryFilter dataclass or a plain dict.

from bigdata_smart_batching import CategoryFilter, CategoryMode, plan_search

plan = plan_search(
    universe="id_name_mapping_us_top_3000.csv",
    start_date="2023-01-01",
    end_date="2023-12-31",
    text="merger acquisition",
    category=CategoryFilter(
        mode=CategoryMode.INCLUDE,
        values=("news_premium", "transcripts"),
    ),
)

# Equivalent dict form:
plan = plan_search(
    universe="id_name_mapping_us_top_3000.csv",
    start_date="2023-01-01",
    end_date="2023-12-31",
    text="merger acquisition",
    category={"mode": "INCLUDE", "values": ["news_premium", "transcripts"]},
)

Valid category values are exposed via VALID_CATEGORY_VALUES and include: expert_interviews, filings, my_files, news, news_premium, news_public, podcasts, research, research_academic_journals, research_investment_research, transcripts.

Save and Load Plans

from bigdata_smart_batching import plan_search, execute_search, save_plan, load_plan

plan = plan_search(
    universe="id_name_mapping_us_top_3000.csv",
    start_date="2023-01-01",
    end_date="2023-12-31",
    text="merger acquisition",
)
save_plan(plan, "my_search_plan.json")

plan = load_plan("my_search_plan.json")
raw_10 = execute_search(plan, chunk_percentage=0.1)
raw_50 = execute_search(plan, chunk_percentage=0.5)

How It Works

Architecture Overview

Step 1: PLANNING
  Universe  -->  Co-mention volumes  -->  Entity groups (by volume)  -->  Sub-period split  -->  Search plan

Step 2: EXECUTION
  Proportional sampling  -->  Parallel search (rate limited)  -->  Collect & aggregate

Planning (plan_search())

1. Loads the universe of companies from a CSV path or inline list.
2. Queries the co-mention endpoint to get chunk volumes per company (split into sub-ranges of up to 365 days, summed per company).
3. Splits date ranges by volume when a company exceeds the chunk limit, honoring min_period_days.
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).
5. Classifies each group by volume_bucket (high, medium, low, very_low).
6. Returns a plan with the chunk upper bound estimate and per-basket query configurations.

Execution (execute_search())

1. Calculates proportional chunks per basket, with a minimum of 1 chunk per basket when expected_chunks > 0.
2. Executes searches in parallel with sliding-window rate limiting and a concurrency semaphore.
3. Optionally performs a second pass, removing already-retrieved entities from each basket and re-querying for the remainder.
4. Enriches each chunk with entity_ids and primary_entity_id, and returns the raw documents.

API Reference

plan_search()

execute_search()

Helper Functions

• deduplicate_documents(documents) -- merge duplicate documents by id, combining their chunks.
• convert_to_dataframe(raw_results) -- convert documents to a DataFrame with one row per chunk.
• resolve_universe(universe) -- load entity IDs from a CSV path or validate a list[str].
• load_universe_from_csv(csv_path) -- load entity IDs from a CSV only.
• save_plan(plan, path) / load_plan(path) -- persist plans as JSON.
• validate_category_filter(category) -- validate and normalize a category filter to the API dict shape.
• 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.

Public Types

• CategoryFilter -- frozen dataclass with mode: CategoryMode and values: tuple[str, ...].
• CategoryMode -- StrEnum with INCLUDE and EXCLUDE values.
• CategoryFilterInput -- type alias: CategoryFilter | dict[str, str | list[str]].
• UniverseInput -- type alias: str | list[str].
• VALID_CATEGORY_VALUES -- frozenset[str] of allowed category strings.

Testing

# Run all tests
uv run pytest

# With coverage
uv run pytest --cov=bigdata_smart_batching --cov-report=term-missing

# Specific test file
uv run pytest tests/test_validation.py -v

Project Structure

bigdata-smart-batching/
├── pyproject.toml
├── README.md
├── CHANGELOG.md
├── LICENSE
├── Makefile
├── .python-version
├── src/
│   └── bigdata_smart_batching/
│       ├── __init__.py
│       ├── smart_batching.py
│       ├── smart_batching_config.py
│       ├── search_function.py
│       └── output_converter.py
├── examples/
│   └── example.py
└── tests/
    ├── __init__.py
    ├── test_config.py
    ├── test_output_converter.py
    ├── test_rate_limiter.py
    └── test_validation.py

Configuration

Environment Variables

• BIGDATA_API_KEY -- required; your Bigdata API key.
• BIGDATA_API_BASE_URL -- optional; API base URL (default: https://api.bigdata.com).

Default Settings

• requests_per_minute: 100
• max_workers: 40
• MAX_CHUNKS_PER_BASKET: 1000
• MAX_ENTITIES_IN_ANY_OF: 500
• volume_query_mode: "three_pass"

Volume Buckets

Companies (and basket groups) are classified by total chunk volume. Ranges use [lower, upper) semantics:

Ties are resolved using VOLUME_BUCKET_PRIORITY: high > medium > low > very_low.

Changelog

See CHANGELOG.md for release notes. The package is currently at version 1.0.1.

License

This project is part of Bigdata.com.

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.