sdmxlib 0.16.1__tar.gz → 0.16.2__tar.gz

{sdmxlib-0.16.1 → sdmxlib-0.16.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sdmxlib
-Version: 0.16.1
+Version: 0.16.2
 Summary: SDMX structural metadata library for Python
 Keywords: sdmx,statistics,metadata,datastructure
 Author: gabrielgellner

{sdmxlib-0.16.1 → sdmxlib-0.16.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sdmxlib"
-version = "0.16.1"
+version = "0.16.2"
 description = "SDMX structural metadata library for Python"
 readme = "README.md"
 license = { text = "Apache-2.0" }

{sdmxlib-0.16.1 → sdmxlib-0.16.2}/src/sdmxlib/model/code_query.py

@@ -355,6 +355,7 @@ class CodeQuery:
     predicates: tuple[Predicate, ...] = ()
     sort_keys: tuple[_SortKey, ...] = ()
     limit_n: int | None = None
+    annotation_types: tuple[str, ...] = ()  # () = no annotations loaded
 
     # ── Composition ─────────────────────────────────────────────────────
 
@@ -377,6 +378,21 @@ class CodeQuery:
         """Return a new query capped at ``n`` rows."""
         return attrs.evolve(self, limit_n=n)
 
+    def with_annotations(self, *types: str) -> CodeQuery:
+        """Return a new query that hydrates the listed annotation types on collect.
+
+        By default ``.collect()`` returns ``Code`` objects with empty
+        annotations — loading them eagerly is expensive on large
+        codelists (one round-trip + ``Annotation`` per row). Pass the
+        annotation types you actually need:
+
+            cl.codes.with_annotations("geo_level").collect()
+
+        Pass no types to clear; pass ``"*"`` (a single asterisk) to load
+        every annotation type for matched codes.
+        """
+        return attrs.evolve(self, annotation_types=types)
+
     # ── Materialisation ─────────────────────────────────────────────────
 
     def collect(self) -> list[Code]:
@@ -384,14 +400,23 @@ class CodeQuery:
 
         ``Code.name`` and ``Code.description`` carry the full
         ``InternationalString``; no language is selected here.
+
+        Annotations are **not** loaded by default. Call
+        ``.with_annotations(*types)`` first to hydrate the types you
+        need — loading every annotation for every code adds a second
+        SQL round-trip and per-row Python object construction that
+        dominates on large codelists.
         """
         sql, params = _compile(self)
         rows = self._conn.execute(sql, params).fetchall()
         if not rows:
             return []
-        urns = [r[4] for r in rows]
-        anns_by_urn = _fetch_annotations_for(self._conn, urns)
-        return [_row_to_code(row, anns_by_urn.get(row[4], Annotations())) for row in rows]
+        if self.annotation_types:
+            urns = [r[4] for r in rows]
+            anns_by_urn = _fetch_annotations_for(self._conn, urns, types=self.annotation_types)
+            return [_row_to_code(row, anns_by_urn.get(row[4], Annotations())) for row in rows]
+        empty = Annotations()
+        return [_row_to_code(row, empty) for row in rows]
 
     def ids(self) -> list[str]:
         """Compile to SQL, execute, and return just the code ids."""
@@ -427,8 +452,7 @@ class CodeQuery:
             ).fetchone()
             if row is None:
                 raise KeyError(key)
-            anns = _fetch_annotations_for(self._conn, [row[4]]).get(row[4], Annotations())
-            return _row_to_code(row, anns)
+            return _row_to_code(row, self._anns_for_one(row[4]))
 
         offset = key if key >= 0 else len(self) + key
         row = self._conn.execute(
@@ -443,8 +467,21 @@ class CodeQuery:
         ).fetchone()
         if row is None:
             raise IndexError(key)
-        anns = _fetch_annotations_for(self._conn, [row[4]]).get(row[4], Annotations())
-        return _row_to_code(row, anns)
+        return _row_to_code(row, self._anns_for_one(row[4]))
+
+    def _anns_for_one(self, code_urn: str) -> Annotations:
+        """Return annotations for a single code, honouring ``annotation_types``.
+
+        Single-code reads (``cl.codes["A"]``) are cheap enough that we
+        load whichever types the query is configured for. Default
+        ``annotation_types=()`` skips the round-trip entirely.
+        """
+        if not self.annotation_types:
+            return Annotations()
+        return _fetch_annotations_for(self._conn, [code_urn], types=self.annotation_types).get(
+            code_urn,
+            Annotations(),
+        )
 
     def __contains__(self, key: object) -> bool:
         if isinstance(key, str):
@@ -807,20 +844,30 @@ def _row_to_code(row: Sequence[Any], annotations: Annotations) -> Code:
 def _fetch_annotations_for(
     conn: duckdb.DuckDBPyConnection,
     code_urns: list[str],
+    *,
+    types: tuple[str, ...] = ("*",),
 ) -> dict[str, Annotations]:
-    """Bulk-load annotations + their text for many code URNs in two queries."""
+    """Bulk-load annotations + their text for many code URNs in two queries.
+
+    ``types`` filters which annotation types to load. ``("*",)`` (default)
+    loads every type. Pass an explicit tuple of types to project just
+    those, e.g. ``types=("geo_level",)``.
+    """
     if not code_urns:
         return {}
     placeholders = ", ".join("?" for _ in code_urns)
-    ann_rows = conn.execute(
-        f"""
-        SELECT id, owner_urn, annotation_id, title, annotation_type, url, ordinal
-        FROM sdmx.annotation
-        WHERE owner_urn IN ({placeholders})
-        ORDER BY owner_urn, ordinal
-        """,
-        list(code_urns),
-    ).fetchall()
+    sql_parts = [
+        "SELECT id, owner_urn, annotation_id, title, annotation_type, url, ordinal",
+        "FROM sdmx.annotation",
+        f"WHERE owner_urn IN ({placeholders})",
+    ]
+    sql_params: list[object] = list(code_urns)
+    if types and "*" not in types:
+        type_placeholders = ", ".join("?" for _ in types)
+        sql_parts.append(f"AND annotation_type IN ({type_placeholders})")
+        sql_params.extend(types)
+    sql_parts.append("ORDER BY owner_urn, ordinal")
+    ann_rows = conn.execute(" ".join(sql_parts), sql_params).fetchall()
     if not ann_rows:
         return {}
 

{sdmxlib-0.16.1 → sdmxlib-0.16.2}/src/sdmxlib/polars.py

@@ -19,7 +19,8 @@ directly for standalone use::
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import weakref
+from typing import TYPE_CHECKING, Any
 
 import polars as pl
 
@@ -168,6 +169,79 @@ def validation_report(
     )
 
 
+# ── Label-map cache ─────────────────────────────────────────────────────
+#
+# slpl.label() and _label_view() are hot paths. Building the id→label
+# dict by iterating ``codelist.codes`` materialises a Code object per
+# row, which dominates on large codelists (CL_GEOGRAPHY: 78k codes →
+# ~16s per call). We cache the dict per (codelist instance, lang) pair.
+#
+# Identity is the codelist *object* (weakref so the cache doesn't pin
+# instances). When the codelist is garbage-collected the cache entry
+# goes with it. A fresh ingest produces a new instance and a fresh
+# cache slot, so re-ingestion is automatically correct without a
+# manual invalidation call.
+
+_LABEL_MAP_CACHE: weakref.WeakKeyDictionary[object, dict[str, dict[str, str]]] = weakref.WeakKeyDictionary()
+
+
+def clear_label_cache() -> None:
+    """Drop the cached id→label maps used by ``label`` / ``_label_view``.
+
+    The cache is normally invalidated automatically when the underlying
+    codelist instance is garbage-collected. Call this explicitly in
+    tests or other contexts that want to force a rebuild.
+    """
+    _LABEL_MAP_CACHE.clear()
+
+
+def _component_label_map(component: object, lang: str) -> dict[str, str] | None:
+    """Return ``{code_id: label}`` for a resolved enumerated component.
+
+    Fast path: when ``component.representation()`` is a
+    :class:`~sdmxlib.local.LazyCodelist` (or anything with a
+    ``code_map(*, lang)`` method), use direct-SQL projection — single
+    DuckDB query, no Code/InternationalString construction. Falls back
+    to in-memory iteration for plain :class:`Codelist` instances.
+
+    Returns ``None`` when the component is not enumerated or its
+    representation Ref is unresolved.
+    """
+    from typing import cast  # noqa: PLC0415
+
+    rep = getattr(component, "representation", None)
+    if not (isinstance(rep, Ref) and rep.resolved):
+        return None
+    cl: object = rep()
+    try:
+        per_lang = _LABEL_MAP_CACHE.get(cl)
+    except TypeError:
+        # Some codelist instances (e.g. attrs frozen) aren't weakly
+        # referenceable. Skip caching for those.
+        per_lang = None
+    if per_lang is not None and lang in per_lang:
+        return per_lang[lang]
+    code_map = getattr(cl, "code_map", None)
+    mapping: dict[str, str]
+    if callable(code_map):
+        # Fast path: direct SQL projection (LazyCodelist).
+        result = cast("dict[str, str]", code_map(lang=lang))
+        mapping = dict(result)
+    else:
+        # Fallback: in-memory codelist iteration.
+        codes = cast("Any", cl).codes
+        mapping = {c.id: c.name.get(lang) or c.id for c in codes}
+    try:
+        if per_lang is None:
+            _LABEL_MAP_CACHE[cl] = {lang: mapping}
+        else:
+            per_lang[lang] = mapping
+    except TypeError:
+        # Codelist isn't weakly referenceable — return uncached.
+        pass
+    return mapping
+
+
 def _label_view(
     lf: pl.LazyFrame,
     dsd: DataStructure,
@@ -182,19 +256,20 @@ def _label_view(
     ``pl.Enum`` (DSD-defined order preserved). Non-enumerated columns are
     unchanged.
     """
-    enum_map: dict[str, tuple[list[str], list[str]]] = {}
+    enum_map: dict[str, dict[str, str]] = {}
     md = [dsd.measure_dimension] if dsd.measure_dimension is not None else []
     for comp in [*dsd.dimensions, *md, *dsd.attributes, *dsd.measures]:
-        if not (isinstance(comp.representation, Ref) and comp.representation.resolved):
+        mapping = _component_label_map(comp, lang)
+        if mapping is None:
             continue
-        codes = [c.id for c in comp.representation().codes]  # type: ignore[union-attr]
-        labels = [c.name.get(lang) or c.id for c in comp.representation().codes]  # type: ignore[union-attr]
-        enum_map[comp.id] = (codes, labels)
+        enum_map[comp.id] = mapping
 
     select_exprs = []
     for col_name in lf.collect_schema().names():
         if col_name in enum_map:
-            codes, labels = enum_map[col_name]
+            mapping = enum_map[col_name]
+            codes = list(mapping.keys())
+            labels = list(mapping.values())
             unique_labels = list(dict.fromkeys(labels))
             select_exprs.append(
                 pl.col(col_name)
@@ -222,6 +297,10 @@ def label(
     appends a ``{col}{suffix}`` column with the code label in ``lang``.
     Codes with no label fall back to the code id.
 
+    Label maps are cached per ``(codelist_urn, lang)`` after the first
+    call — call :func:`clear_label_cache` after re-ingesting a codelist
+    in long-running processes.
+
     Example:
         labelled = slpl.label(pl.scan_csv("data.csv"), dsd, lang="en").collect()
     """
@@ -230,9 +309,9 @@ def label(
     for comp in [*dsd.dimensions, *mds, *dsd.attributes, *dsd.measures]:
         if comp.id not in schema_cols:
             continue
-        if not (isinstance(comp.representation, Ref) and comp.representation.resolved):
+        mapping = _component_label_map(comp, lang)
+        if mapping is None:
             continue
-        mapping = {c.id: c.name.get(lang) or c.id for c in comp.representation().codes}  # type: ignore[union-attr]
         old = list(mapping.keys())
         new = list(mapping.values())
         label_col = f"{comp.id}{suffix}"

{sdmxlib-0.16.1 → sdmxlib-0.16.2}/src/sdmxlib/storage/lazy.py

@@ -35,6 +35,7 @@ class LazyCodelist:
     """
 
     __slots__ = (
+        "__weakref__",  # required so slpl.label's per-codelist cache can hold a weakref
         "_agency",
         "_conn",
         "_description",