bagelquant-data 0.1.0__py3-none-any.whl

bagelquant_data/__init__.py

@@ -0,0 +1,32 @@
+"""Source-agnostic data lake framework for BagelQuant research."""
+
+from bagelquant_data.core import (
+    BagelQuantDataError,
+    DataSource,
+    DatasetNotFoundError,
+    DatasetSpec,
+    DatasetSpecError,
+    DuplicateResolutionError,
+    SourceNotFoundError,
+    ValidationError,
+    stable_bucket,
+)
+from bagelquant_data.finance import FinancialFieldKind, FinancialFieldSpec
+from bagelquant_data.management import DataLake
+from bagelquant_data.sources.tushare import TushareSource
+
+__all__ = [
+    "BagelQuantDataError",
+    "DataLake",
+    "DataSource",
+    "DatasetNotFoundError",
+    "DatasetSpec",
+    "DatasetSpecError",
+    "DuplicateResolutionError",
+    "FinancialFieldKind",
+    "FinancialFieldSpec",
+    "SourceNotFoundError",
+    "TushareSource",
+    "ValidationError",
+    "stable_bucket",
+]

bagelquant_data/cli/main.py

@@ -0,0 +1,30 @@
+"""Thin CLI for the Python-first data lake API."""
+
+from __future__ import annotations
+
+import argparse
+
+from bagelquant_data import DataLake
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="bagelquant-data")
+    parser.add_argument("--root", default="data")
+    sub = parser.add_subparsers(dest="command", required=True)
+    sub.add_parser("status")
+    datasets = sub.add_parser("dataset-list")
+    datasets.add_argument("--source")
+    sub.add_parser("source-list")
+    args = parser.parse_args(argv)
+    lake = DataLake.open(args.root)
+    if args.command == "status":
+        print(lake.status.summary())
+    elif args.command == "dataset-list":
+        print(lake.datasets.list(args.source))
+    elif args.command == "source-list":
+        print(lake.sources.list())
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

bagelquant_data/core/__init__.py

@@ -0,0 +1,37 @@
+"""Core source-agnostic framework primitives."""
+
+from bagelquant_data.core.dataset import DatasetSpec
+from bagelquant_data.core.exceptions import (
+    BagelQuantDataError,
+    ConfigurationError,
+    DatasetNotFoundError,
+    DatasetSpecError,
+    DestructiveOperationError,
+    DuplicateResolutionError,
+    SourceNotFoundError,
+    ValidationError,
+)
+from bagelquant_data.core.hashing import frame_content_hash, stable_bucket, stable_record_hash
+from bagelquant_data.core.registry import FrameworkRegistries, Registry, default_registries
+from bagelquant_data.core.request import RequestContext
+from bagelquant_data.core.source import DataSource
+
+__all__ = [
+    "BagelQuantDataError",
+    "ConfigurationError",
+    "DataSource",
+    "DatasetNotFoundError",
+    "DatasetSpec",
+    "DatasetSpecError",
+    "DestructiveOperationError",
+    "DuplicateResolutionError",
+    "FrameworkRegistries",
+    "Registry",
+    "RequestContext",
+    "SourceNotFoundError",
+    "ValidationError",
+    "default_registries",
+    "frame_content_hash",
+    "stable_bucket",
+    "stable_record_hash",
+]

bagelquant_data/core/dataset.py

@@ -0,0 +1,170 @@
+"""Dataset specification model."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from bagelquant_data.core.exceptions import DatasetSpecError
+
+
+@dataclass(frozen=True, slots=True)
+class DatasetSpec:
+    """Declarative canonical dataset behavior."""
+
+    name: str
+    source: str
+    source_dataset: str
+    category: str
+    field_mapping: dict[str, str]
+    required_columns: tuple[str, ...]
+    primary_key: tuple[str, ...] | None = None
+    business_key: tuple[str, ...] | None = None
+    asset_column: str | None = None
+    time_column: str | None = None
+    period_column: str | None = None
+    request_planner: str = "snapshot"
+    request_options: dict[str, Any] = field(default_factory=dict)
+    normalizer: str = "standard"
+    deduplication: str = "exact_record_hash"
+    partition_strategy: str = "single_file"
+    partition_options: dict[str, Any] = field(default_factory=dict)
+    update_mode: str = "upsert"
+    sort_columns: tuple[str, ...] = ()
+    point_in_time: bool = False
+    reference: bool = False
+    enabled: bool = True
+
+    @classmethod
+    def from_mapping(cls, value: dict[str, Any]) -> "DatasetSpec":
+        """Build a specification from parsed config data."""
+
+        required = ("name", "source", "source_dataset", "category")
+        missing = [key for key in required if key not in value]
+        if missing:
+            raise DatasetSpecError(f"Dataset spec missing required keys: {missing}")
+        return cls(
+            name=str(value["name"]),
+            source=str(value["source"]),
+            source_dataset=str(value["source_dataset"]),
+            category=str(value["category"]),
+            field_mapping=dict(value.get("field_mapping") or {}),
+            required_columns=_tuple(value.get("required_columns")),
+            primary_key=_optional_tuple(value.get("primary_key")),
+            business_key=_optional_tuple(value.get("business_key")),
+            asset_column=_optional_str(value.get("asset_column")),
+            time_column=_optional_str(value.get("time_column")),
+            period_column=_optional_str(value.get("period_column")),
+            request_planner=str(value.get("request_planner") or "snapshot"),
+            request_options=dict(value.get("request_options") or {}),
+            normalizer=str(value.get("normalizer") or "standard"),
+            deduplication=str(value.get("deduplication") or "exact_record_hash"),
+            partition_strategy=str(value.get("partition_strategy") or "single_file"),
+            partition_options=dict(value.get("partition_options") or {}),
+            update_mode=str(value.get("update_mode") or "upsert"),
+            sort_columns=_tuple(value.get("sort_columns")),
+            point_in_time=bool(value.get("point_in_time", False)),
+            reference=bool(value.get("reference", False)),
+            enabled=bool(value.get("enabled", True)),
+        )
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> "DatasetSpec":
+        """Load a dataset spec from a small YAML file.
+
+        The project intentionally avoids a YAML runtime dependency. This parser
+        supports the simple mappings and lists used by bundled dataset specs.
+        """
+
+        return cls.from_mapping(_parse_simple_yaml(Path(path).read_text()))
+
+    @property
+    def key(self) -> tuple[str, str]:
+        """Return the `(source, name)` lookup key."""
+
+        return (self.source, self.name)
+
+
+def _tuple(value: Any) -> tuple[str, ...]:
+    if value is None:
+        return ()
+    if isinstance(value, str):
+        return (value,)
+    return tuple(str(item) for item in value)
+
+
+def _optional_tuple(value: Any) -> tuple[str, ...] | None:
+    result = _tuple(value)
+    return result or None
+
+
+def _optional_str(value: Any) -> str | None:
+    return None if value is None else str(value)
+
+
+def _parse_simple_yaml(text: str) -> dict[str, Any]:
+    rows = [line.rstrip() for line in text.splitlines()]
+    root: dict[str, Any] = {}
+    stack: list[tuple[int, dict[str, Any] | list[Any]]] = [(-1, root)]
+    for raw in rows:
+        if not raw.strip() or raw.lstrip().startswith("#"):
+            continue
+        indent = len(raw) - len(raw.lstrip(" "))
+        line = raw.strip()
+        while stack and indent <= stack[-1][0]:
+            stack.pop()
+        parent = stack[-1][1]
+        if line.startswith("- "):
+            if not isinstance(parent, list):
+                raise DatasetSpecError(f"Unsupported YAML list location: {line}")
+            parent.append(_yaml_scalar(line[2:]))
+            continue
+        key, sep, value = line.partition(":")
+        if not sep:
+            raise DatasetSpecError(f"Unsupported YAML line: {line}")
+        key = key.strip()
+        value = value.strip()
+        if value == "":
+            container: dict[str, Any] | list[Any]
+            next_list = _next_content_is_list(rows, raw)
+            container = [] if next_list else {}
+            if isinstance(parent, dict):
+                parent[key] = container
+            else:
+                raise DatasetSpecError(f"Unsupported nested YAML key: {key}")
+            stack.append((indent, container))
+        elif isinstance(parent, dict):
+            parent[key] = _yaml_scalar(value)
+    return root
+
+
+def _next_content_is_list(rows: list[str], current: str) -> bool:
+    index = rows.index(current)
+    current_indent = len(current) - len(current.lstrip(" "))
+    for row in rows[index + 1 :]:
+        if not row.strip() or row.lstrip().startswith("#"):
+            continue
+        indent = len(row) - len(row.lstrip(" "))
+        return indent > current_indent and row.strip().startswith("- ")
+    return False
+
+
+def _yaml_scalar(value: str) -> Any:
+    if value in {"true", "True"}:
+        return True
+    if value in {"false", "False"}:
+        return False
+    if value in {"null", "None", "~"}:
+        return None
+    if value.startswith("[") and value.endswith("]"):
+        inner = value[1:-1].strip()
+        return [] if not inner else [_yaml_scalar(part.strip()) for part in inner.split(",")]
+    if (value.startswith('"') and value.endswith('"')) or (
+        value.startswith("'") and value.endswith("'")
+    ):
+        return value[1:-1]
+    try:
+        return int(value)
+    except ValueError:
+        return value

bagelquant_data/core/deduplication.py

@@ -0,0 +1,40 @@
+"""Deduplication strategies."""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+import polars as pl
+
+from bagelquant_data.core.dataset import DatasetSpec
+
+
+class DeduplicationStrategy(Protocol):
+    """Deduplicate records for a dataset."""
+
+    def apply(self, frame: pl.LazyFrame, spec: DatasetSpec) -> pl.LazyFrame:
+        """Return deduplicated records."""
+        ...
+
+
+class NoDeduplication:
+    """Leave records unchanged."""
+
+    def apply(self, frame: pl.LazyFrame, spec: DatasetSpec) -> pl.LazyFrame:
+        return frame
+
+
+class ExactRecordHashDeduplication:
+    """Drop exact duplicate rows."""
+
+    def apply(self, frame: pl.LazyFrame, spec: DatasetSpec) -> pl.LazyFrame:
+        return frame.unique(maintain_order=True)
+
+
+class PrimaryKeyLastDeduplication:
+    """Keep the last row for each primary key."""
+
+    def apply(self, frame: pl.LazyFrame, spec: DatasetSpec) -> pl.LazyFrame:
+        if not spec.primary_key:
+            return frame.unique(maintain_order=True)
+        return frame.unique(subset=list(spec.primary_key), keep="last", maintain_order=True)

bagelquant_data/core/exceptions.py

@@ -0,0 +1,39 @@
+"""Framework exceptions."""
+
+from __future__ import annotations
+
+
+class BagelQuantDataError(Exception):
+    """Base package error."""
+
+
+class ConfigurationError(BagelQuantDataError):
+    """Configuration is invalid or incomplete."""
+
+
+class DatasetSpecError(ConfigurationError):
+    """Dataset specification is invalid."""
+
+
+class DatasetNotFoundError(BagelQuantDataError):
+    """Requested dataset is not registered or has no canonical data."""
+
+
+class SourceNotFoundError(BagelQuantDataError):
+    """Requested source is not registered."""
+
+
+class DataSourceError(BagelQuantDataError):
+    """Source adapter failed."""
+
+
+class ValidationError(BagelQuantDataError):
+    """Data failed validation."""
+
+
+class DuplicateResolutionError(BagelQuantDataError):
+    """A single-value panel cannot be produced without resolving duplicates."""
+
+
+class DestructiveOperationError(BagelQuantDataError):
+    """A destructive operation was requested without explicit confirmation."""

bagelquant_data/core/hashing.py

@@ -0,0 +1,32 @@
+"""Stable hashing helpers."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Iterable
+
+import polars as pl
+
+
+def stable_bucket(asset_id: str, bucket_count: int) -> int:
+    """Return a deterministic asset bucket."""
+
+    digest = hashlib.blake2b(asset_id.encode("utf-8"), digest_size=8).digest()
+    return int.from_bytes(digest, byteorder="big") % bucket_count
+
+
+def stable_record_hash(values: dict[str, object]) -> str:
+    """Hash a record using stable JSON encoding."""
+
+    payload = json.dumps(values, sort_keys=True, default=str, separators=(",", ":"))
+    return hashlib.blake2b(payload.encode("utf-8"), digest_size=16).hexdigest()
+
+
+def frame_content_hash(frame: pl.DataFrame, columns: Iterable[str] | None = None) -> str:
+    """Hash a dataframe deterministically after sorting selected columns."""
+
+    selected = list(columns or frame.columns)
+    rows = frame.select(selected).sort(selected).to_dicts()
+    payload = json.dumps(rows, sort_keys=True, default=str, separators=(",", ":"))
+    return hashlib.blake2b(payload.encode("utf-8"), digest_size=16).hexdigest()

bagelquant_data/core/normalization.py

@@ -0,0 +1,67 @@
+"""Canonical normalization."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+import polars as pl
+
+from bagelquant_data.core.dataset import DatasetSpec
+
+
+@dataclass(frozen=True, slots=True)
+class NormalizeContext:
+    """Normalization context."""
+
+    source: str
+    dataset: str
+    run_id: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class NormalizeResult:
+    """Accepted and rejected normalized records."""
+
+    accepted: pl.LazyFrame
+    rejected: pl.LazyFrame
+
+
+class Normalizer(Protocol):
+    """Dataset normalizer protocol."""
+
+    def normalize(
+        self, frame: pl.LazyFrame, spec: DatasetSpec, context: NormalizeContext
+    ) -> NormalizeResult:
+        """Normalize source rows."""
+        ...
+
+
+class StandardNormalizer:
+    """Map configured source fields into canonical columns."""
+
+    def normalize(
+        self, frame: pl.LazyFrame, spec: DatasetSpec, context: NormalizeContext
+    ) -> NormalizeResult:
+        lf = frame.rename(spec.field_mapping)
+        expressions: list[pl.Expr] = [
+            pl.lit(context.source).alias("source"),
+            pl.lit(spec.source_dataset).alias("source_dataset"),
+        ]
+        if spec.asset_column and spec.asset_column in lf.collect_schema().names():
+            expressions.append(pl.col(spec.asset_column).cast(pl.String).alias("asset_id"))
+        if spec.time_column and spec.time_column in lf.collect_schema().names():
+            expressions.append(_date_expr(spec.time_column).alias("time"))
+        if spec.period_column and spec.period_column in lf.collect_schema().names():
+            expressions.append(_date_expr(spec.period_column).alias("period"))
+        accepted = lf.with_columns(expressions)
+        rejected = accepted.filter(pl.lit(False))
+        return NormalizeResult(accepted=accepted, rejected=rejected)
+
+
+def _date_expr(column: str) -> pl.Expr:
+    return (
+        pl.when(pl.col(column).cast(pl.String).str.len_chars() == 8)
+        .then(pl.col(column).cast(pl.String).str.strptime(pl.Date, "%Y%m%d", strict=False))
+        .otherwise(pl.col(column).cast(pl.Date, strict=False))
+    )

bagelquant_data/core/partitioning.py

@@ -0,0 +1,76 @@
+"""Partition strategies."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Protocol
+
+import polars as pl
+
+from bagelquant_data.core.dataset import DatasetSpec
+from bagelquant_data.core.hashing import stable_bucket
+
+
+class PartitionStrategy(Protocol):
+    """Derive partition values and paths."""
+
+    def derive_columns(self, frame: pl.LazyFrame, spec: DatasetSpec) -> pl.LazyFrame:
+        """Add partition columns."""
+        ...
+
+    def paths_for_query(self, spec: DatasetSpec, query: object) -> list[Path]:
+        """Return candidate partition paths."""
+        ...
+
+    def path_for_values(self, spec: DatasetSpec, values: dict[str, object]) -> Path:
+        """Return partition path for values."""
+        ...
+
+
+class SingleFilePartition:
+    """One canonical file per dataset."""
+
+    def derive_columns(self, frame: pl.LazyFrame, spec: DatasetSpec) -> pl.LazyFrame:
+        return frame
+
+    def paths_for_query(self, spec: DatasetSpec, query: object) -> list[Path]:
+        return [Path("data.parquet")]
+
+    def path_for_values(self, spec: DatasetSpec, values: dict[str, object]) -> Path:
+        return Path("data.parquet")
+
+
+class YearMonthPartition:
+    """Partition by year and month of canonical time."""
+
+    def derive_columns(self, frame: pl.LazyFrame, spec: DatasetSpec) -> pl.LazyFrame:
+        return frame.with_columns(
+            pl.col("time").dt.year().cast(pl.Int16).alias("year"),
+            pl.col("time").dt.month().cast(pl.Int8).alias("month"),
+        )
+
+    def paths_for_query(self, spec: DatasetSpec, query: object) -> list[Path]:
+        return []
+
+    def path_for_values(self, spec: DatasetSpec, values: dict[str, object]) -> Path:
+        return Path(f"year={values['year']}") / f"month={int(str(values['month'])):02d}" / "data.parquet"
+
+
+class YearBucketPartition:
+    """Partition by year(time) and stable asset bucket."""
+
+    def derive_columns(self, frame: pl.LazyFrame, spec: DatasetSpec) -> pl.LazyFrame:
+        bucket_count = int(spec.partition_options.get("bucket_count", 32))
+        return frame.with_columns(
+            pl.col("time").dt.year().cast(pl.Int16).alias("year"),
+            pl.col("asset_id")
+            .cast(pl.String)
+            .map_elements(lambda value: stable_bucket(value, bucket_count), return_dtype=pl.Int16)
+            .alias("bucket"),
+        )
+
+    def paths_for_query(self, spec: DatasetSpec, query: object) -> list[Path]:
+        return []
+
+    def path_for_values(self, spec: DatasetSpec, values: dict[str, object]) -> Path:
+        return Path(f"year={values['year']}") / f"bucket={int(str(values['bucket'])):02d}" / "data.parquet"

bagelquant_data/core/registry.py

@@ -0,0 +1,67 @@
+"""Plugin registries."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Generic, TypeVar
+
+from bagelquant_data.core.deduplication import (
+    ExactRecordHashDeduplication,
+    NoDeduplication,
+    PrimaryKeyLastDeduplication,
+)
+from bagelquant_data.core.normalization import StandardNormalizer
+from bagelquant_data.core.partitioning import (
+    SingleFilePartition,
+    YearBucketPartition,
+    YearMonthPartition,
+)
+from bagelquant_data.core.validation import FrameworkValidator
+
+T = TypeVar("T")
+
+
+@dataclass
+class Registry(Generic[T]):
+    """Named object registry."""
+
+    _items: dict[str, T] = field(default_factory=dict)
+
+    def register(self, name: str, value: T) -> None:
+        self._items[name] = value
+
+    def get(self, name: str) -> T:
+        try:
+            return self._items[name]
+        except KeyError as exc:
+            raise KeyError(f"Unknown registry item: {name}") from exc
+
+    def list(self) -> tuple[str, ...]:
+        return tuple(sorted(self._items))
+
+
+@dataclass
+class FrameworkRegistries:
+    """All extension registries used by the framework."""
+
+    sources: Registry[object] = field(default_factory=lambda: Registry[object]())
+    normalizers: Registry[object] = field(default_factory=lambda: Registry[object]())
+    validators: Registry[object] = field(default_factory=lambda: Registry[object]())
+    partition_strategies: Registry[object] = field(default_factory=lambda: Registry[object]())
+    deduplication_strategies: Registry[object] = field(default_factory=lambda: Registry[object]())
+    financial_fields: Registry[object] = field(default_factory=lambda: Registry[object]())
+
+
+def default_registries() -> FrameworkRegistries:
+    """Return registries with built-in plugins installed."""
+
+    registries = FrameworkRegistries()
+    registries.normalizers.register("standard", StandardNormalizer())
+    registries.validators.register("framework", FrameworkValidator())
+    registries.partition_strategies.register("single_file", SingleFilePartition())
+    registries.partition_strategies.register("year_month", YearMonthPartition())
+    registries.partition_strategies.register("year_bucket", YearBucketPartition())
+    registries.deduplication_strategies.register("none", NoDeduplication())
+    registries.deduplication_strategies.register("exact_record_hash", ExactRecordHashDeduplication())
+    registries.deduplication_strategies.register("primary_key_last", PrimaryKeyLastDeduplication())
+    return registries

bagelquant_data/core/request.py

@@ -0,0 +1,21 @@
+"""Request planning models."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+from bagelquant_data.core.types import DateLike
+
+
+@dataclass(frozen=True, slots=True)
+class RequestContext:
+    """Context passed to request planners and source adapters."""
+
+    source: str
+    dataset: str
+    start: DateLike | None = None
+    end: DateLike | None = None
+    assets: Sequence[str] | None = None
+    options: Mapping[str, Any] = field(default_factory=dict)

bagelquant_data/core/source.py

@@ -0,0 +1,38 @@
+"""Source adapter protocol."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from typing import Any, Protocol
+
+import polars as pl
+
+from bagelquant_data.core.dataset import DatasetSpec
+from bagelquant_data.core.request import RequestContext
+
+
+class DataSource(Protocol):
+    """Generic external source adapter."""
+
+    @property
+    def name(self) -> str:
+        """Source name."""
+        ...
+
+    def configure(self, **options: Any) -> None:
+        """Configure credentials and runtime options."""
+        ...
+
+    def test_connection(self) -> None:
+        """Raise when the source cannot be reached."""
+        ...
+
+    def fetch(self, source_dataset: str, request: Mapping[str, Any]) -> pl.DataFrame:
+        """Fetch one source response."""
+        ...
+
+    def plan_requests(
+        self, dataset: DatasetSpec, context: RequestContext
+    ) -> Iterable[Mapping[str, Any]]:
+        """Plan source requests for a dataset."""
+        ...

bagelquant_data/core/types.py

@@ -0,0 +1,10 @@
+"""Shared type aliases."""
+
+from __future__ import annotations
+
+from datetime import date, datetime
+from os import PathLike
+from typing import TypeAlias
+
+DateLike: TypeAlias = str | date | datetime
+PathLikeStr: TypeAlias = str | PathLike[str]