xlr8 0.1.2__py3-none-any.whl

xlr8/schema/encoder.py

@@ -0,0 +1,235 @@
+"""
+Value encoder for XLR8 (Python reference implementation).
+
+NOTE: This module provides a pure Python implementation for reference
+and testing. The Rust backend (encode_any_values_to_arrow) provides
+an optimized version that is much faster.
+
+================================================================================
+DATA FLOW - POLYMORPHIC VALUE ENCODING (Types.Any)
+================================================================================
+
+This module handles encoding/decoding of polymorphic values (Types.Any) to/from
+union structs. This is the "bitmap struct" pattern from cetolib.
+
+THE PROBLEM:
+────────────────────────────────────────────────────────────────────────────────
+
+MongoDB "value" fields often contain mixed types:
+  doc1: {"value": 42.5}      # float
+  doc2: {"value": 100}       # int
+  doc3: {"value": "active"}  # string
+  doc4: {"value": true}      # bool
+
+Parquet columns must be homogeneous (one type per column).
+How do we store mixed types?
+
+THE SOLUTION - UNION STRUCT:
+────────────────────────────────────────────────────────────────────────────────
+
+Store as a struct with ONE field populated, others null:
+
+┌─────────────────────────────────────────────────────────────────────────────┐
+│ ENCODE: Python value -> Struct                                               │
+│                                                                             │
+│ encode_any(42.5) returns:                                                   │
+│ {                                                                           │
+│     "float_value": 42.5,      ← VALUE IS HERE                               │
+│     "int_value": null,                                                      │
+│     "string_value": null,                                                   │
+│     "bool_value": null,                                                     │
+│     "datetime_value": null,                                                 │
+│     "objectid_value": null,                                                 │
+│     "null_value": null,                                                     │
+│ }                                                                           │
+│                                                                             │
+│ encode_any("active") returns:                                               │
+│ {                                                                           │
+│     "float_value": null,                                                    │
+│     "int_value": null,                                                      │
+│     "string_value": "active",  ← VALUE IS HERE                              │
+│     "bool_value": null,                                                     │
+│     "datetime_value": null,                                                 │
+│     "objectid_value": null,                                                 │
+│     "null_value": null,                                                     │
+│ }                                                                           │
+└─────────────────────────────────────────────────────────────────────────────┘
+
+┌─────────────────────────────────────────────────────────────────────────────┐
+│ DECODE: Struct -> Python value                                               │
+│                                                                             │
+│ decode_any({"float_value": 42.5, ...others null}) returns: 42.5             │
+│ decode_any({"string_value": "active", ...others null}) returns: "active"    │
+│                                                                             │
+│ Algorithm: Check each field in order, return first non-null                 │
+└─────────────────────────────────────────────────────────────────────────────┘
+
+TYPE MAPPING:
+────────────────────────────────────────────────────────────────────────────────
+
+Python type      -> Struct field
+──────────────────────────────────
+None             -> null_value: True
+bool             -> bool_value (CHECK BEFORE int!)
+int              -> int_value
+float            -> float_value
+str              -> string_value
+datetime         -> datetime_value
+ObjectId         -> objectid_value (as string)
+other            -> string_value (JSON serialized)
+
+NOTE: bool must be checked BEFORE int because isinstance(True, int) is True!
+
+================================================================================
+"""
+
+import json
+from datetime import datetime
+from typing import Any as AnyPython
+from typing import Dict
+
+from bson import ObjectId
+
+__all__ = [
+    "ValueEncoder",
+]
+
+
+class ValueEncoder:
+    """
+    Encodes and decodes values according to schema types.
+
+    For Types.Any fields, encodes Python values into union structs
+    where only one field is populated based on the value's type.
+
+    Example:
+        encoder = ValueEncoder()
+
+        # Encode different types
+        encoder.encode_any(42.5)      # {"float_value": 42.5, ...others null}
+        encoder.encode_any("hello")   # {"string_value": "hello", ...others null}
+        encoder.encode_any(True)      # {"bool_value": True, ...others null}
+
+        # Decode back
+        struct = {"float_value": 42.5, "int_value": None, ...}
+        encoder.decode_any(struct)    # Returns: 42.5
+    """
+
+    @staticmethod
+    def encode_any(value: AnyPython) -> Dict[str, AnyPython]:
+        """
+        Encode a polymorphic value to union struct.
+
+        Maps Python types to appropriate struct fields:
+        - None -> null_value: True
+        - bool -> bool_value
+        - int -> int_value
+        - float -> float_value
+        - str -> string_value
+        - datetime -> datetime_value
+        - ObjectId -> objectid_value (as string)
+        - other -> string_value (JSON serialized)
+
+        Args:
+            value: Python value to encode
+
+        Returns:
+            Dict with one field populated, others None
+        """
+        result: Dict[str, AnyPython] = {
+            "float_value": None,
+            "int_value": None,
+            "string_value": None,
+            "bool_value": None,
+            "datetime_value": None,
+            "objectid_value": None,
+            "null_value": None,
+        }
+
+        if value is None:
+            result["null_value"] = True
+        elif isinstance(value, bool):
+            # Check bool BEFORE int (bool is subclass of int in Python)
+            result["bool_value"] = value
+        elif isinstance(value, int):
+            result["int_value"] = value
+        elif isinstance(value, float):
+            result["float_value"] = value
+        elif isinstance(value, str):
+            result["string_value"] = value
+        elif isinstance(value, datetime):
+            result["datetime_value"] = value
+        elif isinstance(value, ObjectId):
+            result["objectid_value"] = str(value)
+        else:
+            # Fallback: JSON serialize complex types
+            try:
+                result["string_value"] = json.dumps(value, default=str)
+            except (TypeError, ValueError):
+                result["string_value"] = str(value)
+
+        return result
+
+    @staticmethod
+    def decode_any(struct_value: Dict[str, AnyPython]) -> AnyPython:
+        """
+        Decode union struct back to Python value.
+
+        Checks fields in priority order and returns the first non-null value.
+
+        Args:
+            struct_value: Dict with union struct fields
+
+        Returns:
+            Decoded Python value
+        """
+        if struct_value.get("null_value"):
+            return None
+
+        # Check in order of specificity
+        if struct_value.get("float_value") is not None:
+            return struct_value["float_value"]
+
+        if struct_value.get("int_value") is not None:
+            return struct_value["int_value"]
+
+        if struct_value.get("bool_value") is not None:
+            return struct_value["bool_value"]
+
+        if struct_value.get("datetime_value") is not None:
+            return struct_value["datetime_value"]
+
+        if struct_value.get("objectid_value") is not None:
+            return ObjectId(struct_value["objectid_value"])
+
+        if struct_value.get("string_value") is not None:
+            return struct_value["string_value"]
+
+        # All fields None (shouldn't happen with valid data)
+        return None
+
+    @staticmethod
+    def encode_batch(values: list) -> list:
+        """
+        Encode a batch of values.
+
+        Args:
+            values: List of Python values
+
+        Returns:
+            List of encoded structs
+        """
+        return [ValueEncoder.encode_any(v) for v in values]
+
+    @staticmethod
+    def decode_batch(struct_values: list) -> list:
+        """
+        Decode a batch of struct values.
+
+        Args:
+            struct_values: List of union structs
+
+        Returns:
+            List of decoded Python values
+        """
+        return [ValueEncoder.decode_any(s) for s in struct_values]

xlr8/schema/schema.py

@@ -0,0 +1,265 @@
+"""
+Schema definition for XLR8.
+
+Schema describes the structure of MongoDB documents and how they map to Arrow/Parquet.
+"""
+
+from typing import Dict, List
+
+import pyarrow as pa
+
+from .types import Any as AnyType
+from .types import BaseType, DateTime, Timestamp
+
+
+class Schema:
+    """
+    Defines the structure of MongoDB documents for XLR8 acceleration.
+
+    Schema is required to:
+    - Convert MongoDB documents to Arrow tables
+    - Store data efficiently in Parquet
+    - Reconstruct DataFrames with correct types
+
+    Example:
+        ```python
+        schema = Schema(
+            time_field="timestamp",
+            fields={
+                "timestamp": Types.Timestamp("ns", tz="UTC"),
+                "sensor_id": Types.String(),
+                "value": Types.Float(),
+                "metadata": Types.Any,  # Polymorphic
+            }
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        time_field: str,
+        fields: Dict[str, BaseType],
+        avg_doc_size_bytes: int = 500,
+        flatten_nested: bool = True,
+    ):
+        """
+        Create a schema definition.
+
+        Args:
+            time_field: Name of the timestamp field (required for chunking)
+            fields: Dict mapping field name to XLR8 type
+            avg_doc_size_bytes: Average BSON document size in bytes (default: 500)
+                Used for memory-aware batch sizing and execution planning
+            flatten_nested: If True, flatten nested paths like "metadata.user_id"
+
+        Raises:
+            ValueError: If time_field not in fields or not Timestamp type
+        """
+        self.time_field = time_field
+        self.fields = fields
+        self.avg_doc_size_bytes = avg_doc_size_bytes
+        self.flatten_nested = flatten_nested
+        self._validate()
+
+    SPEC_VERSION = 1
+
+    def _validate(self):
+        """Validate schema configuration."""
+        if self.time_field not in self.fields:
+            raise ValueError(
+                f"time_field '{self.time_field}' must be present in fields. "
+                f"Available fields: {list(self.fields.keys())}"
+            )
+
+        time_field_type = self.fields[self.time_field]
+        if not isinstance(time_field_type, (Timestamp, DateTime)):
+            raise ValueError(
+                f"time_field '{self.time_field}' must be Timestamp or DateTime type, "
+                f"got {type(time_field_type).__name__}"
+            )
+
+    def to_arrow_schema(self) -> pa.Schema:
+        """
+        Convert to PyArrow schema.
+
+        Returns:
+            PyArrow schema object
+        """
+        return pa.schema(
+            [(name, field_type.to_arrow()) for name, field_type in self.fields.items()]
+        )
+
+    def get_any_fields(self) -> List[str]:
+        """
+        Get list of fields with Types.Any (polymorphic types).
+
+        Returns:
+            List of field names that are Any type
+        """
+        return [
+            name
+            for name, field_type in self.fields.items()
+            if isinstance(field_type, AnyType)
+        ]
+
+    def get_field_names(self) -> List[str]:
+        """
+        Get all field names in schema.
+
+        Returns:
+            List of field names
+        """
+        return list(self.fields.keys())
+
+    def has_field(self, name: str) -> bool:
+        """
+        Check if field exists in schema.
+
+        Args:
+            name: Field name to check
+
+        Returns:
+            True if field exists
+        """
+        return name in self.fields
+
+    def get_field_type(self, name: str) -> BaseType:
+        """
+        Get type for a field.
+
+        Args:
+            name: Field name
+
+        Returns:
+            XLR8 type object
+
+        Raises:
+            KeyError: If field not in schema
+        """
+        return self.fields[name]
+
+    def to_spec(self) -> Dict[str, object]:
+        """Export schema to a JSON-serializable specification.
+
+        Converts Python schema objects to a plain dict that can be:
+        - Saved to disk (JSON/YAML)
+        - Transmitted over network
+        - Consumed by native backends (e.g., Rust)
+        - Reconstructed later using from_spec()
+
+        The spec format is intentionally generic and uses introspection to
+        automatically handle any user-defined Types.* classes without
+        hardcoding each type. This means you can add new type classes and
+        they'll automatically work with serialization/deserialization.
+
+        Returns:
+            Dict containing schema version, time field, and field specifications
+
+        Example:
+            >>> schema = Schema(
+            ...     time_field="ts",
+            ...     fields={"ts": Timestamp("ms"), "value": Float()}
+            ... )
+            >>> spec = schema.to_spec()
+            >>> # Later: schema2 = Schema.from_spec(spec)
+        """
+        from . import types as Types  # local to avoid cycles
+
+        fields_spec: List[Dict[str, object]] = []
+        for name, f in self.fields.items():
+            entry: Dict[str, object] = {"name": name}
+
+            if isinstance(f, (Types.Timestamp, Types.DateTime)):
+                # Both Timestamp and DateTime serialize the same way
+                # DateTime is just a convenience wrapper that defaults to "ms"
+                entry.update(
+                    {
+                        "kind": "timestamp",
+                        "unit": getattr(f, "unit", "ms"),
+                        "tz": getattr(f, "tz", "UTC") or "UTC",
+                    }
+                )
+            elif isinstance(f, Types.ObjectId):
+                entry.update({"kind": "objectid"})
+            elif isinstance(f, Types.Any):
+                # Preserve Any() union/bitmap layout; the concrete
+                # encoder decides how to materialize this.
+                any_layout: Dict[str, object] = {
+                    "variants": [
+                        {"name": "int64", "id": 0},
+                        {"name": "float64", "id": 1},
+                        {"name": "bool", "id": 2},
+                        {"name": "string", "id": 3},
+                        {"name": "timestamp_ms_utc", "id": 4},
+                        {"name": "json_blob", "id": 5},
+                    ],
+                }
+
+                # If the Any type exposes explicit bitmap/payload
+                # field naming, surface that; otherwise let the
+                # backend choose sensible defaults.
+                bitmap_field = getattr(f, "bitmap_field_name", None)
+                payload_field = getattr(f, "payload_field_name", None)
+                if bitmap_field is not None:
+                    any_layout["bitmap_field"] = bitmap_field
+                if payload_field is not None:
+                    any_layout["payload_field"] = payload_field
+
+                entry.update({"kind": "any", "any_layout": any_layout})
+            elif isinstance(f, Types.Int):
+                entry.update({"kind": "int64"})
+            elif isinstance(f, Types.Float):
+                entry.update({"kind": "float64"})
+            elif isinstance(f, Types.String):
+                entry.update({"kind": "string"})
+            elif isinstance(f, Types.Bool):
+                entry.update({"kind": "bool"})
+            elif isinstance(f, Types.List):
+                # List type - serialize with element type info
+                # Map element type to kind string
+                elem_type = f.element_type
+                if isinstance(elem_type, Types.Float):
+                    elem_kind = "float"
+                elif isinstance(elem_type, Types.Int):
+                    elem_kind = "int"
+                elif isinstance(elem_type, Types.String):
+                    elem_kind = "string"
+                elif isinstance(elem_type, Types.Bool):
+                    elem_kind = "bool"
+                elif isinstance(elem_type, Types.DateTime):
+                    elem_kind = "datetime"
+                elif isinstance(elem_type, Types.ObjectId):
+                    elem_kind = "objectid"
+                else:
+                    raise ValueError(
+                        f"Unsupported List element type: {type(elem_type).__name__}. "
+                        f"Supported types: Float, Int, String, Bool, DateTime, ObjectId"
+                    )
+                entry.update({"kind": f"list:{elem_kind}"})
+            else:
+                # Conservative fallback: treat as json_blob-backed Any.
+                entry.update(
+                    {
+                        "kind": "any",
+                        "any_layout": {
+                            "variants": [{"name": "json_blob", "id": 0}],
+                        },
+                    }
+                )
+
+            fields_spec.append(entry)
+
+        return {
+            "version": self.SPEC_VERSION,
+            "time_field": self.time_field,
+            "avg_doc_size_bytes": self.avg_doc_size_bytes,
+            "fields": fields_spec,
+        }
+
+    def __repr__(self) -> str:
+        field_lines = [
+            f"  {name}: {field_type}" for name, field_type in self.fields.items()
+        ]
+        return (
+            f"Schema(time_field='{self.time_field}',\n" + "\n".join(field_lines) + "\n)"
+        )