xlr8 0.1.7b2__cp313-cp313-macosx_10_12_x86_64.whl

xlr8/collection/cursor.pyi

@@ -0,0 +1,104 @@
+"""Type stubs for XLR8 Cursor.
+
+Inherits from PyMongoCursor for full IDE autocomplete of PyMongo methods.
+Adds 4 XLR8-specific methods for DataFrame/Polars conversion.
+"""
+
+from __future__ import annotations
+
+from datetime import date, datetime, timedelta
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import pandas as pd
+import polars as pl
+import pyarrow as pa
+from pymongo.cursor import Cursor as PyMongoCursor
+
+class XLR8Cursor(PyMongoCursor):
+    """PyMongo-compatible cursor with optional acceleration.
+
+    Inherits all PyMongo cursor methods. Adds 4 XLR8-specific methods:
+    - to_dataframe(): Convert to Pandas DataFrame with acceleration
+    - to_polars(): Convert to Polars DataFrame with acceleration
+    - to_dataframe_batches(): Memory-efficient batch streaming
+    - stream_to_callback(): Partitioned PyArrow table callbacks
+    """
+
+    def __init__(
+        self,
+        collection: Any,
+        query_filter: Dict[str, Any],
+        projection: Optional[Dict[str, Any]] = None,
+        skip: int = 0,
+        limit: int = 0,
+        sort: Optional[List[Tuple[str, int]]] = None,
+        batch_size: int = 1000,
+    ) -> None: ...
+
+    # XLR8-specific accelerated methods
+    def to_dataframe(
+        self,
+        accelerate: bool = True,
+        cache_read: bool = True,
+        cache_write: bool = True,
+        start_date: Optional[Union[datetime, date, str]] = None,
+        end_date: Optional[Union[datetime, date, str]] = None,
+        coerce: Literal["raise", "error"] = "raise",
+        max_workers: int = 4,
+        chunking_granularity: Optional[timedelta] = None,
+        row_group_size: Optional[int] = None,
+        flush_ram_limit_mb: int = 512,
+    ) -> pd.DataFrame: ...
+    def to_polars(
+        self,
+        accelerate: bool = True,
+        cache_read: bool = True,
+        cache_write: bool = True,
+        start_date: Optional[Union[datetime, date, str]] = None,
+        end_date: Optional[Union[datetime, date, str]] = None,
+        coerce: Literal["raise", "error"] = "raise",
+        max_workers: int = 4,
+        chunking_granularity: Optional[timedelta] = None,
+        row_group_size: Optional[int] = None,
+        any_type_strategy: Literal["float", "string", "keep_struct"] = "float",
+        flush_ram_limit_mb: int = 512,
+    ) -> pl.DataFrame: ...
+    def to_dataframe_batches(
+        self,
+        batch_size: int = 10000,
+        cache_read: bool = True,
+        cache_write: bool = True,
+        start_date: Optional[Union[datetime, date, str]] = None,
+        end_date: Optional[Union[datetime, date, str]] = None,
+        coerce: Literal["raise", "error"] = "raise",
+        max_workers: int = 4,
+        chunking_granularity: Optional[timedelta] = None,
+        row_group_size: Optional[int] = None,
+        flush_ram_limit_mb: int = 512,
+    ) -> Generator[pd.DataFrame, None, None]: ...
+    def stream_to_callback(
+        self,
+        callback: Callable[[pa.Table, Dict[str, Any]], None],
+        *,
+        partition_time_delta: timedelta,
+        partition_by: Optional[Union[str, List[str]]] = None,
+        any_type_strategy: Literal["float", "string", "keep_struct"] = "float",
+        max_workers: int = 4,
+        chunking_granularity: Optional[timedelta] = None,
+        row_group_size: Optional[int] = None,
+        flush_ram_limit_mb: int = 512,
+        cache_read: bool = True,
+        cache_write: bool = True,
+    ) -> Dict[str, Any]: ...
+    def raw_cursor(self) -> PyMongoCursor: ...
+    def explain_acceleration(self) -> Dict[str, Any]: ...

xlr8/collection/wrapper.py

@@ -0,0 +1,399 @@
+"""
+XLR8 collection wrapper with PyMongo compatibility.
+
+================================================================================
+DATA FLOW - COLLECTION WRAPPER
+================================================================================
+
+This module wraps pymongo.collection.Collection to provide the `accelerate()`
+function - the main entry point for users.
+
+TYPICAL USAGE FLOW:
+────────────────────────────────────────────────────────────────────────────────
+
+1. USER WRAPS A COLLECTION:
+┌─────────────────────────────────────────────────────────────────────────────┐
+│  from xlr8 import accelerate, Schema, Types                                 │
+│                                                                             │
+│  schema = Schema(                                                           │
+│      time_field="timestamp",                                                │
+│      fields={                                                               │
+│          "timestamp": Types.Timestamp("ms", tz="UTC"),                      │
+│          "metadata.device_id": Types.ObjectId(),                            │
+│          "metadata.sensor_id": Types.ObjectId(),                         │
+│          "value": Types.Any(),  # Polymorphic - can be int, float, str etc..│
+│      }                                                                      │
+│  )                                                                          │
+│                                                                             │
+│  xlr8_col = accelerate(pymongo_collection, schema, mongo_uri)               │
+└─────────────────────────────────────────────────────────────────────────────┘
+
+2. USER CALLS find() - RETURNS XLR8Cursor (NOT PYMONGO CURSOR):
+┌─────────────────────────────────────────────────────────────────────────────┐
+│  cursor = xlr8_col.find({                                                   │
+│      "timestamp": {"$gte": start, "$lt": end},                              │
+│      "metadata.device_id": ObjectId("64a..."),                              │
+│  })                                                                         │
+│  # cursor is XLR8Cursor, wrapping the query params                          │
+└─────────────────────────────────────────────────────────────────────────────┘
+
+3. USER CALLS to_dataframe() - TRIGGERS ACCELERATION:
+┌─────────────────────────────────────────────────────────────────────────────┐
+│  df = cursor.to_dataframe()                                                 │
+│  # This triggers:                                                           │
+│  # 1. Query analysis (can we chunk by time?)                                │
+│  # 2. Check cache (have we fetched this before?)                            │
+│  # 3. Parallel fetch via Rust async backend                               │
+│  # 4. Stream to Parquet cache                                               │
+│  # 5. Read back and return DataFrame                                        │
+└─────────────────────────────────────────────────────────────────────────────┘
+
+KEY CONFIG OPTIONS:
+────────────────────────────────────────────────────────────────────────────────
+- schema: Required for type-aware encoding (especially Types.Any)
+- mongo_uri: Required for accelerated execution (workers create connections)
+- cache_dir: Where to store Parquet cache (default: .xlr8_cache)
+
+PER-QUERY OPTIONS (via to_dataframe):
+────────────────────────────────────────────────────────────────────────────────
+- max_workers: Number of parallel workers (default: 4)
+- flush_ram_limit_mb: RAM budget for batch sizing (default: 512)
+- chunking_granularity: Time chunk size (e.g., timedelta(days=7))
+
+================================================================================
+"""
+
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from pymongo.collection import Collection as PyMongoCollection
+
+from xlr8.collection.cursor import XLR8Cursor
+from xlr8.schema import Schema
+
+
+class XLR8Collection:
+    """
+    PyMongo-compatible collection wrapper with acceleration.
+
+    Drop-in replacement for pymongo.collection.Collection that transparently
+    accelerates analytical queries through parallel execution and caching.
+
+    All write operations (insert, update, delete) pass through to PyMongo.
+    Read operations (find, aggregate) can be accelerated if:
+    - Schema is provided
+    - Query has time-range predicates
+    - Query doesn't use complex operators ($nor, $where, etc.)
+
+    Example:
+        >>> import pymongo
+        >>> from xlr8 import XLR8Collection, Schema, Types
+        >>>
+        >>> # Create schema
+        >>> schema = Schema(
+        ...     time_field="timestamp",
+        ...     fields={
+        ...         "timestamp": Types.Timestamp(),
+        ...         "value": Types.Float(),
+        ...         "sensor_id": Types.String(),
+        ...     }
+        ... )
+        >>>
+        >>> # Wrap collection with mongo_uri for accelerated execution
+        >>> client = pymongo.MongoClient("mongodb://localhost:27017")
+        >>> pymongo_col = client.mydb.mycollection
+        >>> col = XLR8Collection(pymongo_col, schema=schema, mongo_uri="mongodb://localhost:27017")
+        >>>
+        >>> # Use like regular PyMongo
+        >>> cursor = col.find({"timestamp": {"$gte": start, "$lt": end}})
+        >>> df = cursor.to_dataframe(flush_ram_limit_mb=2000)
+    """
+
+    def __init__(
+        self,
+        pymongo_collection,
+        schema: Optional[Schema] = None,
+        mongo_uri: Union[str, Callable[[], str], None] = None,
+        approx_document_size_bytes: int = 500,
+    ):
+        """
+        Initialize XLR8 collection wrapper.
+
+        Args:
+            pymongo_collection: PyMongo Collection instance
+            schema: Optional schema definition for acceleration
+            mongo_uri: MongoDB connection string (str) or callable that returns one.
+                       Required for accelerated execution. Can be:
+                       - A string: "mongodb://localhost:27017"
+                       - A callable: lambda: os.environ["MONGODB_URI"]
+            approx_document_size_bytes: Approximate size of each document in bytes
+                (default: 500). Used for memory budget calculations.
+
+        Note:
+            Cache directory is auto-managed based on query hash.
+            flush_ram_limit_mb and max_workers are parameters of to_dataframe(),
+            to_polars(), etc. for per-query control.
+        """
+        self._pymongo_collection = pymongo_collection
+        self._schema = schema
+        self._mongo_uri = mongo_uri
+        self._approx_document_size_bytes = approx_document_size_bytes
+
+    def raw_collection(self) -> PyMongoCollection:
+        """
+        Get direct access to underlying PyMongo collection.
+
+        This is an escape hatch for power users who need direct access to PyMongo
+        collection methods that may not be available through delegation.
+
+        Returns:
+            pymongo.collection.Collection: The underlying PyMongo collection
+
+        Example:
+            >>> xlr8_col = accelerate(collection, schema=schema)
+            >>> xlr8_col.raw_collection().watch()  # Use MongoDB change streams
+            >>> xlr8_col.raw_collection().list_indexes()  # Direct PyMongo access
+        """
+        return self._pymongo_collection
+
+    # PyMongo pass-through properties
+    @property
+    def name(self) -> str:
+        """Collection name."""
+        return self._pymongo_collection.name
+
+    @property
+    def full_name(self) -> str:
+        """Full collection name (database.collection)."""
+        return self._pymongo_collection.full_name
+
+    @property
+    def database(self):
+        """Parent database."""
+        return self._pymongo_collection.database
+
+    # Public accessor properties for cursor usage
+
+    @property
+    def schema(self):
+        """Schema definition for acceleration."""
+        return self._schema
+
+    @property
+    def pymongo_collection(self):
+        """Underlying PyMongo collection instance."""
+        return self._pymongo_collection
+
+    @property
+    def mongo_uri(self):
+        """MongoDB connection URI for accelerated execution."""
+        return self._mongo_uri
+
+    @property
+    def approx_document_size_bytes(self) -> int:
+        """Approximate size of each document in bytes."""
+        return self._approx_document_size_bytes
+
+    def __getattr__(self, name: str):
+        """
+        Delegate unknown methods to PyMongo collection.
+
+        Why:
+            Provides full PyMongo compatibility without manually implementing
+            every collection method (insert, update, delete, indexes, etc.).
+
+        Example:
+            >>> xlr8_col.insert_one({...})  # Works via delegation
+            >>> xlr8_col.create_index("timestamp")  # Works via delegation
+            >>> count = xlr8_col.count_documents({})  # Works via delegation
+        """
+        return getattr(self._pymongo_collection, name)
+
+    # Read operations (can be accelerated)
+    def find(
+        self,
+        filter: Optional[Dict[str, Any]] = None,
+        projection: Optional[Dict[str, Any]] = None,
+        skip: int = 0,
+        limit: int = 0,
+        sort: Optional[List[tuple]] = None,
+        batch_size: int = 1000,
+        **kwargs,
+    ) -> XLR8Cursor:
+        """
+        Query collection with optional acceleration.
+
+        Returns XLR8Cursor which is PyMongo-compatible but can accelerate
+        to_dataframe() / to_polars() conversions.
+
+
+        DATA FLOW EXAMPLE:
+
+        INPUT (filter parameter):
+        {
+            "$or": [
+                {"metadata.sensor_id": ObjectId("64a...")},
+                {"metadata.sensor_id": ObjectId("64b...")},
+            ],
+            "timestamp": {"$gte": datetime(2024,1,1), "$lt": datetime(...)}
+        }
+
+        OUTPUT: XLR8Cursor object containing:
+        - _filter: The query dict (unchanged)
+        - _collection: Reference back to this XLR8Collection
+        - _projection, _skip, _limit, _sort: Query modifiers
+
+        NEXT STEP: User calls cursor.to_dataframe() which triggers:
+        1. Query analysis in analysis/brackets.py
+        2. Execution planning in execution/planner.py
+        3. Parallel fetch in execution/worker.py
+
+        Args:
+            filter: Query filter dict
+            projection: Field projection dict
+            skip: Number of documents to skip
+            limit: Maximum documents to return
+            sort: Sort specification
+            batch_size: Batch size for iteration
+            **kwargs: Additional PyMongo cursor options
+
+        Returns:
+            XLR8Cursor instance
+
+        Example:
+            >>> # Simple query
+            >>> cursor = col.find({"status": "active"})
+            >>>
+            >>> # Query with time range (accelerated)
+            >>> cursor = col.find({
+            ...     "timestamp": {"$gte": start, "$lt": end},
+            ...     "sensor_id": "sensor_1"
+            ... })
+            >>> df = cursor.to_dataframe()
+        """
+        if filter is None:
+            filter = {}
+
+        return XLR8Cursor(
+            collection=self,
+            query_filter=filter,
+            projection=projection,
+            skip=skip,
+            limit=limit,
+            sort=sort,
+            batch_size=batch_size,
+        )
+
+    # XLR8-specific methods
+
+    def set_schema(self, schema: Schema) -> None:
+        """
+        Set or update schema for acceleration.
+
+        Args:
+            schema: Schema definition
+        """
+        self._schema = schema
+
+    def get_schema(self) -> Optional[Schema]:
+        """
+        Get current schema.
+
+        Returns:
+            Schema or None
+        """
+        return self._schema
+
+
+def accelerate(
+    pymongo_collection: PyMongoCollection,
+    schema: Schema,
+    mongo_uri: Union[str, Callable[[], str]],
+    approx_document_size_bytes: int = 500,
+) -> XLR8Collection:
+    """
+    Convenience function to wrap a PyMongo collection with acceleration.
+
+
+    DATA FLOW EXAMPLE - MAIN ENTRY POINT:
+
+    INPUT:
+    - pymongo_collection: client["main"]["sensorData"]
+    - schema: Schema(time_field="timestamp", fields={...})
+    - mongo_uri: Connection string used by accelerated workers
+
+    Example:
+    accelerate(
+        collection,
+        schema,
+        mongo_uri="mongodb://localhost:27017",  # Or callable
+    )
+
+    OUTPUT: XLR8Collection wrapper that:
+    - Wraps pymongo collection for transparent pass-through
+    - Stores schema for type-aware Parquet encoding
+    - Stores mongo_uri for workers to create their own connections
+
+    WHAT HAPPENS NEXT:
+    1. User calls: xlr8_col.find({...})
+    2. Returns XLR8Cursor (wraps query params)
+    3. User calls: cursor.to_dataframe()
+    4. Workers use mongo_uri to create their own connections
+
+
+    Args:
+        pymongo_collection: PyMongo Collection instance
+        schema: Schema definition
+        mongo_uri: MongoDB connection string (str) or callable that returns one.
+                   Required for accelerated execution. Can be:
+                   - A string: "mongodb://localhost:27017"
+                   - A callable: lambda: os.environ["MONGODB_URI"]
+        approx_document_size_bytes: Approximate size of each document in bytes
+            (default: 500). Used for memory budget calculations.
+
+    Returns:
+        XLR8Collection wrapper
+
+    Note:
+        Cache directory is auto-managed based on query hash.
+        flush_ram_limit_mb and max_workers are parameters of to_dataframe(),
+        to_polars(), etc. for per-query control.
+
+    Example:
+        >>> import pymongo
+        >>> from xlr8 import accelerate, Schema, Types
+        >>>
+        >>> # Connection string or callable
+        >>> MONGO_URI = "mongodb://localhost:27017"
+        >>> # OR: get_uri = lambda: os.environ["MONGODB_URI"]
+        >>>
+        >>> client = pymongo.MongoClient(MONGO_URI)
+        >>> col = client.mydb.sensor_logs
+        >>>
+        >>> schema = Schema(
+        ...     time_field="timestamp",
+        ...     fields={
+        ...         "timestamp": Types.Timestamp(),
+        ...         "sensor_id": Types.String(),
+        ...         "value": Types.Float(),
+        ...     },
+        ... )
+        >>>
+        >>> # Pass mongo_uri for accelerated workers
+        >>> accelerated_col = accelerate(col, schema, mongo_uri=MONGO_URI)
+        >>>
+        >>> # max_workers and flush_ram_limit_mb are per-query
+        >>> from datetime import timedelta
+        >>> df = accelerated_col.find({
+        ...     "timestamp": {"$gte": start, "$lt": end}
+        ... }).to_dataframe(
+        ...     max_workers=8,
+        ...     chunking_granularity=timedelta(days=1),
+        ...     flush_ram_limit_mb=2000,
+        ... )
+    """
+    return XLR8Collection(
+        pymongo_collection=pymongo_collection,
+        schema=schema,
+        mongo_uri=mongo_uri,
+        approx_document_size_bytes=approx_document_size_bytes,
+    )

xlr8/collection/wrapper.pyi

@@ -0,0 +1,61 @@
+"""Type stubs for XLR8 Collection.
+
+Inherits from PyMongoCollection for full IDE autocomplete.
+Only overrides find() to return XLR8Cursor instead of PyMongo cursor.
+"""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+from pymongo.collection import Collection as PyMongoCollection
+
+from .cursor import XLR8Cursor
+
+class XLR8Collection(PyMongoCollection):
+    """PyMongo-compatible collection with optional acceleration.
+
+    Inherits all PyMongo collection methods. Only find() is overridden
+    to return XLR8Cursor for accelerated DataFrame/Polars conversion.
+    """
+
+    def __init__(
+        self,
+        pymongo_collection: PyMongoCollection,
+        schema: Optional[Any] = None,
+        mongo_uri: Optional[str] = None,
+        approx_document_size_bytes: int = 500,
+    ) -> None: ...
+
+    # Override find() to return XLR8Cursor
+    def find(
+        self,
+        filter: Optional[Dict[str, Any]] = None,
+        projection: Optional[Dict[str, Any]] = None,
+        skip: int = 0,
+        limit: int = 0,
+        sort: Optional[List[Tuple[str, int]]] = None,
+        batch_size: int = 1000,
+        **kwargs: Any,
+    ) -> XLR8Cursor: ...
+
+    # XLR8-specific methods
+    def raw_collection(self) -> PyMongoCollection: ...
+    def set_schema(self, schema: Any) -> None: ...
+    def get_schema(self) -> Optional[Any]: ...
+    def clear_cache(self) -> None: ...
+
+    # Properties
+    @property
+    def schema(self) -> Optional[Any]: ...
+    @property
+    def pymongo_collection(self) -> PyMongoCollection: ...
+    @property
+    def mongo_uri(self) -> Optional[str]: ...
+    @property
+    def approx_document_size_bytes(self) -> int: ...
+
+def accelerate(
+    pymongo_collection: PyMongoCollection,
+    schema: Any,
+    mongo_uri: Any,
+    approx_document_size_bytes: int = 500,
+) -> XLR8Collection: ...

xlr8/constants.py

@@ -0,0 +1,24 @@
+"""
+XLR8 constants and configuration values.
+
+Centralized constants to avoid magic numbers scattered throughout codebase.
+All tuneable performance parameters should be defined here.
+"""
+
+# =============================================================================
+# PARQUET FILE SETTINGS
+# =============================================================================
+
+# Default row group size for compression can be altered via argument passed
+# to the special cursor methods e.g to_dataframe
+PARQUET_ROW_GROUP_SIZE = 100_000
+
+# Default compression codec for Parquet files
+DEFAULT_COMPRESSION = "zstd"
+
+# =============================================================================
+# BATCH PROCESSING
+# =============================================================================
+
+# Default batch size for DataFrame operations
+DEFAULT_BATCH_SIZE = 10_000

xlr8/execution/__init__.py

@@ -0,0 +1,43 @@
+"""
+Execution engine for parallel query execution via Rust backend.
+
+All parallel execution now goes through the Rust backend for GIL-free performance.
+
+Components:
+- executor: High-level parallel execution (execute_parallel_stream_to_cache)
+- callback: Partitioned streaming for data lake population
+- planner: Memory-aware execution planning and worker configuration
+
+Python handles:
+- Query planning and bracketing
+- Memory budget calculations
+- Result reading and DataFrame construction
+
+Rust backend handles:
+- Parallel MongoDB fetches (GIL-free)
+- BSON decoding and Arrow encoding
+- Memory-aware buffering
+- Parquet writing
+"""
+
+from .callback import PartitionWorkItem, execute_partitioned_callback
+from .executor import execute_parallel_stream_to_cache
+from .planner import (
+    Backend,
+    BackendConfig,
+    ExecutionPlan,
+    build_execution_plan,
+)
+
+__all__ = [
+    # Executor
+    "execute_parallel_stream_to_cache",
+    # Callback
+    "PartitionWorkItem",
+    "execute_partitioned_callback",
+    # Planner
+    "Backend",
+    "BackendConfig",
+    "ExecutionPlan",
+    "build_execution_plan",
+]