bagelquant-data 0.1.0__py3-none-any.whl

bagelquant_data/core/validation.py

@@ -0,0 +1,32 @@
+"""Framework validation."""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+import polars as pl
+
+from bagelquant_data.core.dataset import DatasetSpec
+from bagelquant_data.core.exceptions import ValidationError
+
+
+class Validator(Protocol):
+    """Validate canonical records."""
+
+    def validate(self, frame: pl.LazyFrame, spec: DatasetSpec) -> None:
+        """Raise on invalid data."""
+
+
+class FrameworkValidator:
+    """Basic schema validation shared by all datasets."""
+
+    def validate(self, frame: pl.LazyFrame, spec: DatasetSpec) -> None:
+        names = set(frame.collect_schema().names())
+        required = set(spec.required_columns)
+        if not spec.reference:
+            required.update({"asset_id", "time"})
+        missing = sorted(required - names)
+        if missing:
+            raise ValidationError(f"{spec.source}/{spec.name} missing columns: {missing}")
+        if spec.point_in_time and "period" not in names:
+            raise ValidationError(f"{spec.source}/{spec.name} PIT dataset requires period")

bagelquant_data/finance/__init__.py

@@ -0,0 +1,80 @@
+"""Generic financial transformation API."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import polars as pl
+
+from bagelquant_data.core.types import DateLike
+from bagelquant_data.finance.fields import FinancialFieldKind, FinancialFieldSpec
+from bagelquant_data.finance.flows import ytd_to_period
+from bagelquant_data.finance.point_in_time import asof
+from bagelquant_data.finance.ratios import ratio
+from bagelquant_data.finance.rolling import trailing
+from bagelquant_data.finance.shares import weighted_average
+from bagelquant_data.finance.stocks import average_stock
+from bagelquant_data.query.raw import RawQueryService
+
+
+class FinanceFacade:
+    """User-facing generic finance API."""
+
+    def __init__(self, raw_service: RawQueryService) -> None:
+        self._raw = raw_service
+
+    def field(
+        self,
+        dataset: str,
+        field: str,
+        *,
+        source: str,
+        start: DateLike | None = None,
+        end: DateLike | None = None,
+        assets: Sequence[str] | None = None,
+        value_name: str = "value",
+    ) -> pl.LazyFrame:
+        return self._raw.raw(
+            dataset,
+            source=source,
+            start=start,
+            end=end,
+            assets=assets,
+            columns=("asset_id", "time", "period", field),
+        ).select("asset_id", "time", "period", pl.col(field).alias(value_name))
+
+    def asof(self, data: pl.LazyFrame, observations: pl.LazyFrame, **kwargs: Any) -> pl.LazyFrame | pl.DataFrame:
+        return asof(data, observations, **kwargs)
+
+    def latest(
+        self,
+        dataset: str,
+        field: str,
+        *,
+        source: str,
+        observations: pl.LazyFrame,
+        value_name: str | None = None,
+        collect: bool = False,
+    ) -> pl.LazyFrame | pl.DataFrame:
+        events = self.field(dataset, field, source=source, value_name=value_name or field)
+        return asof(events, observations, value_column=value_name or field, output_name=value_name or field, collect=collect)
+
+    ytd_to_period = staticmethod(ytd_to_period)
+    trailing = staticmethod(trailing)
+    average_stock = staticmethod(average_stock)
+    weighted_average = staticmethod(weighted_average)
+    ratio = staticmethod(ratio)
+
+
+__all__ = [
+    "FinanceFacade",
+    "FinancialFieldKind",
+    "FinancialFieldSpec",
+    "asof",
+    "average_stock",
+    "ratio",
+    "trailing",
+    "weighted_average",
+    "ytd_to_period",
+]

bagelquant_data/finance/align.py

@@ -0,0 +1,5 @@
+"""Alignment helpers."""
+
+from bagelquant_data.finance.point_in_time import asof
+
+__all__ = ["asof"]

bagelquant_data/finance/fields.py

@@ -0,0 +1,30 @@
+"""Financial field metadata."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+
+class FinancialFieldKind(str, Enum):
+    """Generic financial field semantics."""
+
+    FLOW_YTD = "flow_ytd"
+    FLOW_PERIOD = "flow_period"
+    STOCK = "stock"
+    PER_SHARE = "per_share"
+    RATIO = "ratio"
+    COUNT = "count"
+    UNKNOWN = "unknown"
+
+
+@dataclass(frozen=True, slots=True)
+class FinancialFieldSpec:
+    """Financial field metadata separate from storage."""
+
+    source: str
+    dataset: str
+    field: str
+    kind: FinancialFieldKind
+    unit: str | None = None
+    currency: str | None = None

bagelquant_data/finance/flows.py

@@ -0,0 +1,24 @@
+"""Generic flow transformations."""
+
+from __future__ import annotations
+
+import polars as pl
+
+
+def ytd_to_period(
+    data: pl.LazyFrame,
+    *,
+    value_column: str = "value",
+    frequency: str = "quarter",
+    output_name: str | None = None,
+) -> pl.LazyFrame:
+    """Convert cumulative YTD flow values into period values."""
+
+    output = output_name or value_column
+    sorted_data = data.sort("asset_id", "period", "time")
+    year = pl.col("period").dt.year()
+    previous = pl.col(value_column).shift(1).over("asset_id", year)
+    period_value = pl.when(previous.is_null()).then(pl.col(value_column)).otherwise(pl.col(value_column) - previous)
+    return sorted_data.with_columns(period_value.alias(output)).select(
+        "asset_id", "time", "period", output
+    )

bagelquant_data/finance/periods.py

@@ -0,0 +1,11 @@
+"""Period helpers."""
+
+from __future__ import annotations
+
+import polars as pl
+
+
+def with_period_year(data: pl.LazyFrame) -> pl.LazyFrame:
+    """Add period year for downstream grouping."""
+
+    return data.with_columns(pl.col("period").dt.year().alias("period_year"))

bagelquant_data/finance/point_in_time.py

@@ -0,0 +1,27 @@
+"""Point-in-time alignment."""
+
+from __future__ import annotations
+
+import polars as pl
+
+
+def asof(
+    data: pl.LazyFrame,
+    observations: pl.LazyFrame,
+    *,
+    value_column: str = "value",
+    output_name: str | None = None,
+    collect: bool = False,
+) -> pl.LazyFrame | pl.DataFrame:
+    """Align latest event whose availability time is not after observation time."""
+
+    output = output_name or value_column
+    events = data.select("asset_id", "time", pl.col(value_column)).sort("asset_id", "time")
+    obs = observations.select("time", "asset_id").sort("asset_id", "time")
+    result = obs.join_asof(
+        events,
+        on="time",
+        by="asset_id",
+        strategy="backward",
+    ).select("time", "asset_id", pl.col(value_column).alias(output))
+    return result.collect() if collect else result

bagelquant_data/finance/ratios.py

@@ -0,0 +1,31 @@
+"""Generic ratio operations."""
+
+from __future__ import annotations
+
+import polars as pl
+
+
+def ratio(
+    numerator: pl.LazyFrame,
+    denominator: pl.LazyFrame,
+    *,
+    numerator_column: str,
+    denominator_column: str,
+    output_name: str = "value",
+    zero_policy: str = "null",
+) -> pl.LazyFrame:
+    """Join numerator and denominator and compute a generic ratio."""
+
+    keys = ["asset_id", "time"]
+    if "period" in numerator.collect_schema().names() and "period" in denominator.collect_schema().names():
+        keys.append("period")
+    joined = numerator.join(denominator, on=keys, how="inner", suffix="__den")
+    denominator_expr = pl.col(denominator_column)
+    if zero_policy == "raise":
+        zero_count = joined.filter(denominator_expr == 0).select(pl.len()).collect().item()
+        if zero_count:
+            raise ZeroDivisionError("Ratio denominator contains zero")
+    value = pl.when(denominator_expr == 0).then(None if zero_policy == "null" else float("nan")).otherwise(
+        pl.col(numerator_column) / denominator_expr
+    )
+    return joined.with_columns(value.alias(output_name)).select(*keys, output_name)

bagelquant_data/finance/rolling.py

@@ -0,0 +1,34 @@
+"""Generic rolling financial operations."""
+
+from __future__ import annotations
+
+import polars as pl
+
+
+def trailing(
+    data: pl.LazyFrame,
+    *,
+    value_column: str = "value",
+    periods: int,
+    operation: str,
+    output_name: str | None = None,
+    require_complete: bool = True,
+) -> pl.LazyFrame:
+    """Compute a trailing window over event-period rows."""
+
+    output = output_name or value_column
+    expr = {
+        "sum": pl.col(value_column).rolling_sum(periods),
+        "mean": pl.col(value_column).rolling_mean(periods),
+        "min": pl.col(value_column).rolling_min(periods),
+        "max": pl.col(value_column).rolling_max(periods),
+        "first": pl.col(value_column).rolling_map(lambda values: values[0], window_size=periods),
+        "last": pl.col(value_column).rolling_map(lambda values: values[-1], window_size=periods),
+    }[operation]
+    result = data.sort("asset_id", "period", "time").with_columns(
+        expr.over("asset_id").alias(output),
+        pl.len().rolling_sum(window_size=periods).over("asset_id").alias("__window_count"),
+    )
+    if require_complete:
+        result = result.filter(pl.col("__window_count") >= periods)
+    return result.drop("__window_count").select("asset_id", "time", "period", output)

bagelquant_data/finance/shares.py

@@ -0,0 +1,28 @@
+"""Generic weighted-average support."""
+
+from __future__ import annotations
+
+import polars as pl
+
+
+def weighted_average(
+    data: pl.LazyFrame,
+    *,
+    value_column: str,
+    effective_time_column: str,
+    period_start_column: str,
+    period_end_column: str,
+    output_name: str | None = None,
+) -> pl.LazyFrame:
+    """Compute a generic time-weighted average inside each asset/period."""
+
+    output = output_name or value_column
+    weighted = data.with_columns(
+        (
+            pl.col(period_end_column).cast(pl.Date).sub(pl.col(period_start_column).cast(pl.Date)).dt.total_days()
+        ).alias("__days")
+    ).with_columns((pl.col(value_column) * pl.col("__days")).alias("__weighted"))
+    return weighted.group_by("asset_id", "period").agg(
+        pl.max("time").alias("time"),
+        (pl.sum("__weighted") / pl.sum("__days")).alias(output),
+    ).select("asset_id", "time", "period", output)

bagelquant_data/finance/stocks.py

@@ -0,0 +1,27 @@
+"""Generic stock variable transformations."""
+
+from __future__ import annotations
+
+import polars as pl
+
+
+def average_stock(
+    data: pl.LazyFrame,
+    *,
+    value_column: str = "value",
+    periods: int = 4,
+    method: str = "endpoint",
+    output_name: str | None = None,
+) -> pl.LazyFrame:
+    """Average stock variables such as assets, equity, inventory, or shares."""
+
+    output = output_name or value_column
+    sorted_data = data.sort("asset_id", "period", "time")
+    if method == "endpoint":
+        lagged = pl.col(value_column).shift(periods).over("asset_id")
+        expr = ((pl.col(value_column) + lagged) / 2).alias(output)
+    elif method == "period_mean":
+        expr = pl.col(value_column).rolling_mean(periods).over("asset_id").alias(output)
+    else:
+        raise ValueError(f"Unsupported average_stock method: {method}")
+    return sorted_data.with_columns(expr).select("asset_id", "time", "period", output)

bagelquant_data/management/__init__.py

@@ -0,0 +1,5 @@
+"""Management API."""
+
+from bagelquant_data.management.lake import DataLake
+
+__all__ = ["DataLake"]

bagelquant_data/management/datasets.py

@@ -0,0 +1,63 @@
+"""Dataset management API."""
+
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+from typing import Any
+
+from bagelquant_data.core.dataset import DatasetSpec
+from bagelquant_data.core.exceptions import DatasetNotFoundError, DestructiveOperationError
+from bagelquant_data.storage.metadata import MetadataStore
+from bagelquant_data.storage.paths import LakePaths
+
+
+class DatasetManager:
+    """Register and inspect dataset specifications."""
+
+    def __init__(self, metadata: MetadataStore, paths: LakePaths) -> None:
+        self.metadata = metadata
+        self.paths = paths
+        self._specs: dict[tuple[str, str], DatasetSpec] = {}
+
+    def add(self, spec: DatasetSpec) -> None:
+        self.validate_spec(spec)
+        self._specs[spec.key] = spec
+        self.metadata.upsert_dataset(spec)
+
+    def add_from_yaml(self, path: str | Path) -> DatasetSpec:
+        spec = DatasetSpec.from_yaml(path)
+        self.add(spec)
+        return spec
+
+    def get(self, dataset: str, *, source: str) -> DatasetSpec:
+        key = (source, dataset)
+        if key in self._specs:
+            return self._specs[key]
+        row = self.metadata.get_dataset(source, dataset)
+        if row is None:
+            raise DatasetNotFoundError(f"Dataset is not registered: {source}/{dataset}")
+        spec = DatasetSpec.from_mapping(__import__("json").loads(row["spec_json"]))
+        self._specs[key] = spec
+        return spec
+
+    def list(self, source: str | None = None) -> list[dict[str, Any]]:
+        return self.metadata.list_datasets(source)
+
+    def enable(self, dataset: str, *, source: str) -> None:
+        self.metadata.set_dataset_enabled(source, dataset, True)
+
+    def disable(self, dataset: str, *, source: str) -> None:
+        self.metadata.set_dataset_enabled(source, dataset, False)
+
+    def validate_spec(self, spec: DatasetSpec) -> None:
+        if not spec.reference and ("asset_id" not in spec.required_columns or "time" not in spec.required_columns):
+            pass
+
+    def remove(self, dataset: str, *, source: str, delete_data: bool = False, confirm: bool = False) -> None:
+        if delete_data and not confirm:
+            raise DestructiveOperationError("Pass confirm=True to delete canonical data")
+        self.metadata.remove_dataset(source, dataset)
+        self._specs.pop((source, dataset), None)
+        if delete_data:
+            shutil.rmtree(self.paths.dataset_root(source, dataset), ignore_errors=True)

bagelquant_data/management/lake.py

@@ -0,0 +1,190 @@
+"""Public DataLake facade."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import polars as pl
+
+from bagelquant_data.core.dataset import DatasetSpec
+from bagelquant_data.core.exceptions import ConfigurationError, DatasetNotFoundError
+from bagelquant_data.core.registry import FrameworkRegistries, default_registries
+from bagelquant_data.core.request import RequestContext
+from bagelquant_data.finance import FinanceFacade
+from bagelquant_data.management.datasets import DatasetManager
+from bagelquant_data.management.sources import SourceManager
+from bagelquant_data.management.status import StatusManager
+from bagelquant_data.pipeline.ingest import IngestionPipeline, IngestionReport
+from bagelquant_data.pipeline.update import UpdateReport, combine_reports, update_dataset
+from bagelquant_data.query import QueryFacade
+from bagelquant_data.query.raw import RawQueryService
+from bagelquant_data.storage.metadata import MetadataStore
+from bagelquant_data.storage.parquet import ParquetStore
+from bagelquant_data.storage.paths import LakePaths
+from bagelquant_data.storage.rejected import RejectedStore
+from bagelquant_data.storage.staging import StagingStore
+
+
+class DataLake:
+    """Source-agnostic local data lake facade."""
+
+    def __init__(self, root: str | Path, registries: FrameworkRegistries | None = None) -> None:
+        self.paths = LakePaths.open(root)
+        self.paths.ensure()
+        self.registries = registries or default_registries()
+        self.metadata = MetadataStore(self.paths.database)
+        self.parquet = ParquetStore(self.paths, self.metadata)
+        self.sources = SourceManager(self.registries, self.metadata)
+        self.datasets = DatasetManager(self.metadata, self.paths)
+        self.status = StatusManager(self.metadata)
+        raw = RawQueryService(self.parquet, self.metadata)
+        self.query = QueryFacade(raw)
+        self.finance = FinanceFacade(raw)
+        self.update = UpdateManager(self)
+        self._pipeline = IngestionPipeline(
+            registries=self.registries,
+            parquet=self.parquet,
+            metadata=self.metadata,
+            staging=StagingStore(self.paths),
+            rejected=RejectedStore(self.paths),
+        )
+
+    @classmethod
+    def open(cls, root: str | Path = "data") -> "DataLake":
+        """Open or create a local data lake."""
+
+        return cls(root)
+
+    def ingest_frame(self, spec: DatasetSpec, frame: pl.DataFrame) -> IngestionReport:
+        """Convenience method for tests and local file adapters."""
+
+        self.datasets.add(spec)
+        return self._pipeline.ingest_frame(spec, frame, mode=spec.update_mode)
+
+
+@dataclass
+class UpdateManager:
+    """Public update API."""
+
+    lake: DataLake
+
+    def dataset(self, dataset: str, *, source: str, **kwargs: Any) -> IngestionReport:
+        spec = self.lake.datasets.get(dataset, source=source)
+        adapter = self.lake.sources.get(source)
+        if spec.request_planner == "by_asset" and not kwargs.get("assets"):
+            kwargs["assets"] = self._default_assets(source)
+        if spec.source == "tushare" and spec.category == "market":
+            kwargs["trade_dates"] = self._trade_dates(source, start=kwargs.get("start"), end=kwargs.get("end"))
+        context = _request_context(source=source, dataset=dataset, kwargs=kwargs)
+        return update_dataset(spec=spec, source_adapter=adapter, pipeline=self.lake._pipeline, context=context)
+
+    def datasets(self, datasets: list[str], *, source: str, **kwargs: Any) -> UpdateReport:
+        reports = [self.dataset(dataset, source=source, **kwargs) for dataset in datasets]
+        return combine_reports(source, reports)
+
+    def source(self, source: str, **kwargs: Any) -> UpdateReport:
+        names = [row["name"] for row in self.lake.datasets.list(source) if row["enabled"]]
+        return self.datasets(names, source=source, **kwargs)
+
+    def _default_assets(self, source: str) -> list[str]:
+        if source != "tushare":
+            raise ConfigurationError("assets=... is required for by_asset dataset updates")
+        try:
+            frame = self.lake.query.reference("stock_basic", source=source, collect=True)
+        except DatasetNotFoundError as exc:
+            raise ConfigurationError(
+                "Tushare by_asset updates require an asset universe. "
+                "Update/register stock_basic first or pass assets=[...] explicitly."
+            ) from exc
+        if isinstance(frame, pl.LazyFrame):
+            frame = frame.collect()
+        columns = frame.columns
+        column = "asset_id" if "asset_id" in columns else "ts_code" if "ts_code" in columns else None
+        if column is None:
+            raise ConfigurationError("tushare/stock_basic does not contain asset_id or ts_code")
+        return [str(value) for value in frame.get_column(column).drop_nulls().unique().sort().to_list()]
+
+    def _trade_dates(self, source: str, *, start: Any, end: Any) -> list[str]:
+        try:
+            frame = self.lake.query.reference("trade_cal", source=source, collect=True)
+        except DatasetNotFoundError as exc:
+            raise ConfigurationError(
+                "Tushare market updates require trade_cal. Update/register trade_cal first."
+            ) from exc
+        if isinstance(frame, pl.LazyFrame):
+            frame = frame.collect()
+        if frame.is_empty():
+            raise ConfigurationError("tushare/trade_cal is empty; update trade_cal before market datasets")
+        column = "time" if "time" in frame.columns else "cal_date" if "cal_date" in frame.columns else None
+        if column is None:
+            raise ConfigurationError("tushare/trade_cal does not contain time or cal_date")
+        dates = frame.with_columns(_calendar_date_expr(column).alias("_calendar_date"))
+        if start is not None:
+            dates = dates.filter(pl.col("_calendar_date") >= _date_literal(start))
+        if end is not None:
+            dates = dates.filter(pl.col("_calendar_date") <= _date_literal(end))
+        if "is_open" in dates.columns:
+            dates = dates.filter(pl.col("is_open").cast(pl.Int8, strict=False) == 1)
+        result = [
+            value.strftime("%Y-%m-%d")
+            for value in dates.select("_calendar_date").drop_nulls().unique().sort("_calendar_date").to_series().to_list()
+        ]
+        if not result:
+            raise ConfigurationError(
+                f"tushare/trade_cal has no open trading days between {start} and {end}"
+            )
+        return result
+
+
+def _request_context(source: str, dataset: str, kwargs: dict[str, Any]) -> RequestContext:
+    known = {
+        "start": kwargs.pop("start", None),
+        "end": kwargs.pop("end", None),
+        "assets": kwargs.pop("assets", None),
+    }
+    workers = kwargs.pop("workers", None)
+    batch_size = kwargs.pop("batch_size", None)
+    source_options = kwargs.pop("source_options", None)
+    progress = kwargs.pop("progress", None)
+    max_retries = kwargs.pop("max_retries", None)
+    retry_backoff_seconds = kwargs.pop("retry_backoff_seconds", None)
+    trade_dates = kwargs.pop("trade_dates", None)
+    if kwargs:
+        keys = ", ".join(sorted(kwargs))
+        raise ConfigurationError(f"Unsupported update option(s): {keys}")
+    options: dict[str, Any] = {}
+    if workers is not None:
+        options["workers"] = workers
+    if batch_size is not None:
+        options["batch_size"] = batch_size
+    if source_options is not None:
+        options["source_options"] = source_options
+    if progress is not None:
+        options["progress"] = progress
+    if max_retries is not None:
+        options["max_retries"] = max_retries
+    if retry_backoff_seconds is not None:
+        options["retry_backoff_seconds"] = retry_backoff_seconds
+    if trade_dates is not None:
+        options["trade_dates"] = trade_dates
+    return RequestContext(source=source, dataset=dataset, options=options, **known)
+
+
+def _calendar_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))
+    )
+
+
+def _date_literal(value: Any) -> pl.Expr:
+    if hasattr(value, "strftime"):
+        return pl.lit(value).cast(pl.Date, strict=False)
+    text = str(value)
+    if "T" in text:
+        text = text.split("T", maxsplit=1)[0]
+    return pl.lit(datetime.strptime(text[:10], "%Y-%m-%d").date())

bagelquant_data/management/sources.py

@@ -0,0 +1,57 @@
+"""Source management API."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from bagelquant_data.core.exceptions import SourceNotFoundError
+from bagelquant_data.core.registry import FrameworkRegistries
+from bagelquant_data.storage.metadata import MetadataStore
+
+
+class SourceManager:
+    """Register and configure source adapters."""
+
+    def __init__(self, registries: FrameworkRegistries, metadata: MetadataStore) -> None:
+        self.registries = registries
+        self.metadata = metadata
+
+    def register(self, source: object) -> None:
+        name = getattr(source, "name")
+        if callable(name):
+            name = name()
+        saved_options = self.metadata.source_options(str(name))
+        if saved_options and hasattr(source, "configure"):
+            source.configure(**saved_options)  # type: ignore[attr-defined]
+        self.registries.sources.register(str(name), source)
+        self.metadata.upsert_source(
+            str(name),
+            type(source).__name__,
+            configured=bool(saved_options),
+        )
+
+    def remove(self, name: str) -> None:
+        self.registries.sources._items.pop(name, None)
+        self.metadata.remove_source(name)
+
+    def list(self) -> list[dict[str, Any]]:
+        return self.metadata.list_sources()
+
+    def get(self, name: str) -> object:
+        try:
+            return self.registries.sources.get(name)
+        except KeyError as exc:
+            raise SourceNotFoundError(f"Source is not registered: {name}") from exc
+
+    def configure(self, name: str, **options: Any) -> None:
+        source = self.get(name)
+        source.configure(**options)  # type: ignore[attr-defined]
+        saved = self.metadata.source_options(name)
+        saved.update(options)
+        self.metadata.upsert_source(name, type(source).__name__, configured=True, options=saved)
+
+    def configure_tushare(self, token: str) -> None:
+        self.configure("tushare", token=token)
+
+    def test(self, name: str) -> None:
+        self.get(name).test_connection()  # type: ignore[attr-defined]