sdmxlib 0.8.0__py3-none-any.whl

sdmxlib/model/dataset.py

@@ -0,0 +1,337 @@
+"""Dataset — observation data paired with its DataStructure."""
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import polars as pl
+from attrs import define
+
+from sdmxlib.model.ref import Ref
+
+if TYPE_CHECKING:
+    from sdmxlib.model.dataflow import Dataflow
+    from sdmxlib.model.datastructure import DataStructure
+
+
+@define
+class Dataset:
+    """Tabular SDMX observation data paired with its structural metadata.
+
+    Always lazy — ``data`` is a ``pl.LazyFrame``. Use :meth:`collect` to
+    materialise it. All transformation methods return a new ``Dataset`` and
+    stay lazy until a terminal is called.
+
+    The ``structure`` must be a fully resolved :class:`~sdmxlib.model.datastructure.DataStructure`.
+    Enumerated dimensions need their codelists resolved for ``pl.Enum`` typing
+    and label/validation operations.
+
+    Example:
+        import sdmxlib.polars as slpl
+
+        ds = slpl.scan_csv("data.csv", dsd)
+        df = ds.label().collect()
+
+        ds2 = ds.filter(pl.col("FREQ") == "M")
+        df2 = ds2.collect()
+
+        ds.sink_parquet("output.parquet")  # streaming, no collect needed
+    """
+
+    structure: "DataStructure"
+    data: pl.LazyFrame
+    dataflow: "Ref[Dataflow] | None" = None
+    measure_dim_hint: str | None = None
+    """User-supplied hint for which regular dimension acts as the measure axis.
+
+    Set via ``measure_dim=`` on :meth:`~sdmxlib.api.registry.RestRegistry.data`
+    when the DSD lacks a formal ``<MeasureDimension>`` element but one of its
+    regular dimensions carries indicator codes (e.g. ``MEASURE``, ``INDICATOR``,
+    ``SUBJECT``). Controls :attr:`wide_dimension` and :meth:`pivot`.
+    """
+
+    # ── Schema / structure helpers ────────────────────────────────────────────
+
+    @property
+    def schema(self) -> dict[str, pl.DataType]:
+        """DSD-derived Polars schema for all components."""
+        from sdmxlib import polars as slpl  # noqa: PLC0415
+
+        return slpl.schema(self.structure)
+
+    @property
+    def dimensions(self) -> list[str]:
+        """Ordered dimension column ids.
+
+        Includes the measure dimension (if any) and TIME_PERIOD last.
+        """
+        ids = [d.id for d in self.structure.dimensions]
+        if self.structure.measure_dimension is not None:
+            ids.append(self.structure.measure_dimension.id)
+        if self.structure.time_dimension is not None:
+            ids.append(self.structure.time_dimension.id)
+        return ids
+
+    @property
+    def wide_dimension(self) -> str | None:
+        """The measure dimension id for this long-format flow, else ``None``.
+
+        Sourced from the DSD's formal ``<MeasureDimension>`` if present,
+        otherwise from the ``measure_dim=`` hint passed to
+        :meth:`~sdmxlib.api.registry.RestRegistry.data`.
+
+        When set, the data arrives in long format — one ``OBS_VALUE`` per
+        indicator per row. Call :meth:`pivot` to reshape to wide format.
+        """
+        md = self.structure.measure_dimension
+        if md is not None:
+            return md.id
+        return self.measure_dim_hint
+
+    @property
+    def measures(self) -> list[str]:
+        """Measure column ids (typically ``["OBS_VALUE"]``)."""
+        return [m.id for m in self.structure.measures]
+
+    # ── Wide / pivot ──────────────────────────────────────────────────────────
+
+    def pivot(self) -> pl.LazyFrame:
+        """Pivot long-format data to wide on the measure dimension.
+
+        The pivot column is resolved from (in order):
+
+        1. The DSD's formal ``<MeasureDimension>`` element.
+        2. The ``measure_dim=`` hint passed to
+           :meth:`~sdmxlib.api.registry.RestRegistry.data`.
+
+        If neither is available, raises ``ValueError`` with a list of
+        enumerated dimensions that could act as the measure axis. Re-fetch
+        with ``measure_dim=`` to select one.
+
+        The pivot column codes (from the resolved codelist, or discovered from
+        the data) become column names; all remaining dimensions form the index.
+        Returns a ``pl.LazyFrame`` — call ``.collect()`` to materialise or
+        ``.sink_parquet()`` to stream to a file. DSD coupling ends here.
+
+        Example:
+            # DSD has formal MeasureDimension — automatic
+            if ds.wide_dimension:
+                wide = ds.pivot().collect()
+
+            # Regular dimension acting as measure axis — hint required
+            ds = reg.data(flow, measure_dim="MEASURE").fetch()
+            wide = ds.pivot().collect()
+        """
+        md = self.structure.measure_dimension
+
+        if md is not None:
+            # ── Formal MeasureDimension path ──────────────────────────────────
+            if not (isinstance(md.representation, Ref) and md.representation.resolved):
+                msg = "pivot() requires the MeasureDimension codelist to be resolved."
+                raise ValueError(msg)
+            pivot_id = md.id
+            on_columns = [c.id for c in md.representation().codes]  # type: ignore[union-attr]
+            index = [d.id for d in self.structure.dimensions]
+
+        elif self.measure_dim_hint is not None:
+            # ── User-supplied hint path ────────────────────────────────────────
+            pivot_id = self.measure_dim_hint
+            hint_dim = next(
+                (d for d in self.structure.dimensions if d.id == pivot_id),
+                None,
+            )
+            if hint_dim is None:
+                msg = f"measure_dim {pivot_id!r} not found in this DSD's dimensions"
+                raise ValueError(msg)
+            rep = getattr(hint_dim, "representation", None)
+            if isinstance(rep, Ref) and rep.resolved:
+                on_columns = [c.id for c in rep().codes]  # type: ignore[union-attr]
+            else:
+                # Codelist not resolved — discover column values from the data
+                on_columns = self.data.select(pivot_id).unique().collect()[pivot_id].cast(pl.String).to_list()
+            index = [d.id for d in self.structure.dimensions if d.id != pivot_id]
+
+        else:
+            # ── No measure axis — raise with helpful candidates ────────────────
+            def _cl_id(d: object) -> str:
+                rep = getattr(d, "representation", None)
+                if isinstance(rep, Ref):
+                    return rep.urn.id
+                return "?"
+
+            enumerated = [
+                f"  {d.id:<20} (codelist {_cl_id(d)})"
+                for d in self.structure.dimensions
+                if isinstance(getattr(d, "representation", None), Ref) and d.representation.resolved  # type: ignore[union-attr]
+            ]
+            candidates = "\n".join(enumerated) if enumerated else "  (none — no enumerated dimensions found)"
+            msg = (
+                "pivot() requires a measure dimension — none found in this DSD.\n\n"
+                "Enumerated dimensions that could act as a measure axis:\n"
+                f"{candidates}\n\n"
+                "Re-fetch with measure_dim= to select one, e.g.:\n"
+                '    reg.data(..., measure_dim="MEASURE").fetch()'
+            )
+            raise ValueError(msg)
+
+        if self.structure.time_dimension is not None:
+            index.append(self.structure.time_dimension.id)
+        measures = self.measures
+        values = measures[0] if len(measures) == 1 else measures
+        return self.data.pivot(
+            on=pivot_id,
+            on_columns=pl.Series(on_columns),
+            index=index,
+            values=values,
+        )
+
+    # ── SDMX-aware expressions ────────────────────────────────────────────────
+
+    def _find_component(self, column: str) -> "object":
+        """Return the DSD component for ``column``. Raises ValueError if absent."""
+        for comp in self.structure.components:
+            if comp.id == column:
+                return comp
+        td = self.structure.time_dimension
+        if td is not None and td.id == column:
+            return td
+        msg = f"{column!r} is not a component in this Dataset's DSD"
+        raise ValueError(msg)
+
+    def code(self, column: str) -> pl.Expr:
+        """Return a Polars expression for a DSD component by code value.
+
+        Equivalent to ``pl.col(column)`` but validates that ``column`` is a
+        known DSD component, giving an early error on typos.
+
+        Example:
+            ds.filter(ds.code("FREQ") == "M")
+            ds.filter(ds.code("FREQ").is_in(["M", "Q"]))
+            ds.lazy().sort(ds.code("TIME_PERIOD"))
+        """
+        self._find_component(column)
+        return pl.col(column)
+
+    def label(self, column: str, lang: str = "en") -> pl.Expr:
+        """Return a Polars expression mapping codes to labels for a DSD component.
+
+        The expression replaces code values with human-readable labels from
+        the resolved codelist. Works in any Polars context that accepts an
+        expression — ``filter``, ``sort``, ``with_columns``, etc.
+
+        Args:
+            column: DSD component id. Must be enumerated (resolved codelist).
+            lang: Language for labels. Defaults to ``"en"``.
+
+        Raises:
+            ValueError: If ``column`` is not a DSD component or is not enumerated.
+
+        Example:
+            ds.filter(ds.label("FREQ") == "Monthly")
+            ds.filter(
+                (ds.label("FREQ") == "Monthly") & (ds.code("REF_AREA") == "AT")
+            )
+            ds.lazy().sort(ds.label("FREQ")).collect()
+        """
+        comp = self._find_component(column)
+        rep = getattr(comp, "representation", None)
+        if not (isinstance(rep, Ref) and rep.resolved):
+            msg = f"{column!r} is not an enumerated component — no codelist to map labels from"
+            raise ValueError(msg)
+        codes = [c.id for c in rep().codes]  # type: ignore[union-attr]
+        labels = [c.name.get(lang) or c.id for c in rep().codes]  # type: ignore[union-attr]
+        return pl.col(column).cast(pl.String).replace_strict(old=codes, new=labels, default=None)
+
+    # ── Lazy transformations ──────────────────────────────────────────────────
+
+    def filter(self, predicate: pl.Expr) -> "Dataset":
+        """Return a new Dataset with a Polars filter applied. Stays lazy.
+
+        Example:
+            ds.filter(pl.col("FREQ") == "M").collect()
+        """
+        return Dataset(
+            structure=self.structure,
+            data=self.data.filter(predicate),
+            dataflow=self.dataflow,
+            measure_dim_hint=self.measure_dim_hint,
+        )
+
+    # ── Validation ────────────────────────────────────────────────────────────
+
+    def validate(self) -> pl.DataFrame:
+        """Validate data against DSD codelists.
+
+        Returns a DataFrame of violations (empty if all values are valid)::
+
+            ┌──────────┬───────────────┬───────┐
+            │ column   │ invalid_value │ count │
+            └──────────┴───────────────┴───────┘
+
+        Example:
+            report = ds.validate()
+            if report.is_empty():
+                print("All values valid")
+        """
+        from sdmxlib import polars as slpl  # noqa: PLC0415
+
+        return slpl.validation_report(self.structure, self.data)
+
+    # ── Terminals / escapes ───────────────────────────────────────────────────
+
+    def lazy(self, *, with_labels: bool = False, lang: str = "en") -> pl.LazyFrame:
+        """Return the underlying ``pl.LazyFrame``, leaving the Dataset boundary.
+
+        Use this to apply arbitrary Polars operations (select, rename, join,
+        group_by, …) that are not expressible as DSD-preserving transforms.
+        The DSD is no longer coupled to the result — you are in Polars land.
+
+        Args:
+            with_labels: If ``True``, replace enumerated code columns with
+                         label columns (``{col}_label``). Values become
+                         human-readable strings as ``pl.Enum`` in DSD-defined
+                         order. Non-enumerated columns are unchanged.
+            lang: Language for labels. Defaults to ``"en"``.
+
+        Example:
+            ds.lazy().select("FREQ", "OBS_VALUE").collect()
+            ds.lazy().group_by("REF_AREA").agg(pl.col("OBS_VALUE").mean()).collect()
+            ds.lazy(with_labels=True).sort("FREQ_label").collect()
+        """
+        if not with_labels:
+            return self.data
+        from sdmxlib import polars as slpl  # noqa: PLC0415
+
+        return slpl._label_view(self.data, self.structure, lang=lang)  # noqa: SLF001
+
+    def collect(self, *, with_labels: bool = False, lang: str = "en") -> pl.DataFrame:
+        """Materialise the lazy frame. Terminal operation.
+
+        Args:
+            with_labels: If ``True``, replace enumerated code columns with
+                         label columns (``{col}_label``). Values become
+                         human-readable strings as ``pl.Enum`` in DSD-defined
+                         order. Non-enumerated columns are unchanged.
+            lang: Language for labels. Defaults to ``"en"``.
+
+        Example:
+            ds.collect()                       # FREQ: pl.Enum(["A", "M", "Q"])
+            ds.collect(with_labels=True)       # FREQ_label: pl.Enum(["Annual", ...])
+            ds.collect(with_labels=True, lang="fr")
+        """
+        return self.lazy(with_labels=with_labels, lang=lang).collect()
+
+    def sink_csv(self, path: str | Path) -> None:
+        """Sink observations to a CSV file. Streaming — no full collect needed.
+
+        Example:
+            ds.sink_csv("output.csv")
+        """
+        self.data.sink_csv(path)
+
+    def sink_parquet(self, path: str | Path) -> None:
+        """Sink observations to a Parquet file. Streaming — no full collect needed.
+
+        Example:
+            ds.sink_parquet("output.parquet")
+        """
+        self.data.sink_parquet(path)