sdmxlib 0.19.0__tar.gz → 0.19.2__tar.gz

{sdmxlib-0.19.0 → sdmxlib-0.19.2}/PKG-INFO

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

{sdmxlib-0.19.0 → sdmxlib-0.19.2}/pyproject.toml

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

{sdmxlib-0.19.0 → sdmxlib-0.19.2}/src/sdmxlib/storage/writers.py

@@ -16,6 +16,7 @@ from collections.abc import Sequence
 from typing import Any
 
 import duckdb
+import pyarrow as pa
 
 from sdmxlib.model.annotations import Annotations
 from sdmxlib.model.binding import DataflowBinding
@@ -69,27 +70,19 @@ def _istring_to_map(s: InternationalString | None) -> dict[str, str]:
     return {lang: s[lang] for lang in s}
 
 
-def _sql_literal(v: Any) -> str:
-    """Format a Python value as a DuckDB SQL literal."""
-    if v is None:
-        return "NULL"
-    if isinstance(v, bool):
-        return "TRUE" if v else "FALSE"
-    if isinstance(v, (int, float)):
-        return str(v)
-    if isinstance(v, dict):
-        if not v:
-            return "MAP([], [])"
-        keys = ", ".join(f"'{k}'" for k in v)
-        vals = ", ".join(f"'{str(val).replace(chr(39), chr(39) + chr(39))}'" for val in v.values())
-        return f"MAP([{keys}], [{vals}])"
-    if isinstance(v, list):
-        if not v:
-            return "[]"
-        items = ", ".join(f"'{str(x).replace(chr(39), chr(39) + chr(39))}'" for x in v)
-        return f"[{items}]"
-    s = str(v).replace("'", "''")
-    return f"'{s}'"
+def _to_arrow_array(values: list[Any]) -> pa.Array:
+    """Build a pyarrow Array from a column of Python values.
+
+    Dict values become MAP(VARCHAR, VARCHAR); other types use pyarrow's
+    default inference, which handles None, scalars, and list[str] fine.
+    """
+    sample = next((v for v in values if v is not None), None)
+    if isinstance(sample, dict):
+        return pa.array(
+            [None if d is None else list(d.items()) for d in values],
+            type=pa.map_(pa.string(), pa.string()),
+        )
+    return pa.array(values)
 
 
 def _execute_bulk_insert(
@@ -98,15 +91,22 @@ def _execute_bulk_insert(
     cols: Sequence[str],
     rows: Sequence[Sequence[Any]],
 ) -> None:
-    """Insert many rows in a single INSERT statement."""
+    """Insert many rows via Arrow columnar ingestion.
+
+    Avoids the multi-GB SQL parser allocation that VALUES-with-literals
+    causes for large bulk inserts (issue #66).
+    """
     if not rows:
         return
-    col_str = ", ".join(cols)
-    lit = _sql_literal
-    values = ", ".join("(" + ", ".join(lit(v) for v in row) + ")" for row in rows)
-    conn.execute(f"""
-        INSERT INTO {table} ({col_str}) VALUES {values}
-    """)
+    columns = list(zip(*rows, strict=True))
+    arrays = [_to_arrow_array(list(col)) for col in columns]
+    arrow_table = pa.Table.from_arrays(arrays, names=list(cols))
+    conn.register("__sdmx_bulk_insert", arrow_table)
+    try:
+        col_str = ", ".join(cols)
+        conn.execute(f"INSERT INTO {table} ({col_str}) SELECT * FROM __sdmx_bulk_insert")
+    finally:
+        conn.unregister("__sdmx_bulk_insert")
 
 
 # ── Spine + text + annotations helpers ──────────────────────────────────
@@ -188,44 +188,96 @@ def _put_localized_text(
         )
 
 
+class _AnnotationBatch:
+    """Accumulator for annotation + annotation_text rows.
+
+    Collects rows across an entire artefact write and flushes them with a
+    single bulk-allocated id batch. Replaces the per-annotation
+    ``INSERT … RETURNING id`` roundtrip pattern that made persisting a
+    78k-code codelist take ~130 s (issue #66).
+    """
+
+    __slots__ = ("ann_rows", "text_rows")
+
+    def __init__(self) -> None:
+        # Each row: [owner_urn, ordinal, annotation_id, title, annotation_type, url]
+        self.ann_rows: list[list[Any]] = []
+        # Each entry: (local_idx_into_ann_rows, lang, text). The local
+        # index is resolved to a real annotation.id at flush time.
+        self.text_rows: list[tuple[int, str, str]] = []
+
+    def add(self, owner_urn: str, annotations: Annotations | None) -> None:
+        if not annotations:
+            return
+        for ordinal, ann in enumerate(annotations):
+            local_idx = len(self.ann_rows)
+            self.ann_rows.append([owner_urn, ordinal, ann.id, ann.title, ann.type, ann.url])
+            if ann.text:
+                for lang in ann.text:
+                    self.text_rows.append((local_idx, lang, ann.text[lang]))
+
+    def flush(self, conn: duckdb.DuckDBPyConnection) -> None:
+        if not self.ann_rows:
+            return
+        n = len(self.ann_rows)
+        id_rows = conn.execute("SELECT nextval('annotation_id_seq') FROM range(?)", [n]).fetchall()
+        ids = [r[0] for r in id_rows]
+
+        full_ann_rows = [[ids[i], *row] for i, row in enumerate(self.ann_rows)]
+        _execute_bulk_insert(
+            conn,
+            "annotation",
+            ["id", "owner_urn", "ordinal", "annotation_id", "title", "annotation_type", "url"],
+            full_ann_rows,
+        )
+
+        if self.text_rows:
+            text_rows = [[ids[idx], lang, text] for idx, lang, text in self.text_rows]
+            _execute_bulk_insert(
+                conn,
+                "annotation_text",
+                ["annotation_id", "lang", "text"],
+                text_rows,
+            )
+
+        self.ann_rows.clear()
+        self.text_rows.clear()
+
+
 def _put_annotations(
     conn: duckdb.DuckDBPyConnection,
     owner_urn: str,
     annotations: Annotations | None,
 ) -> None:
-    """Insert annotation + annotation_text rows."""
-    if not annotations:
-        return
-    for ordinal, ann in enumerate(annotations):
-        row = conn.execute(
-            """
-            INSERT INTO annotation (owner_urn, ordinal, annotation_id, title, annotation_type, url)
-            VALUES (?, ?, ?, ?, ?, ?)
-            RETURNING id
-            """,
-            [owner_urn, ordinal, ann.id, ann.title, ann.type, ann.url],
-        ).fetchone()
-        ann_db_id = row[0]  # type: ignore[index]  # RETURNING always returns
-        if ann.text:
-            for lang in ann.text:
-                conn.execute(
-                    """
-                    INSERT INTO annotation_text (annotation_id, lang, text)
-                    VALUES (?, ?, ?)
-                    """,
-                    [ann_db_id, lang, ann.text[lang]],
-                )
+    """Insert annotation + annotation_text rows for one owner.
+
+    For artefacts with many annotated child items (e.g. codelists), prefer
+    building an :class:`_AnnotationBatch` directly and calling ``flush``
+    once at the end of the write — that avoids per-owner roundtrips.
+    """
+    batch = _AnnotationBatch()
+    batch.add(owner_urn, annotations)
+    batch.flush(conn)
 
 
 def _put_artefact_text_and_annotations(
     conn: duckdb.DuckDBPyConnection,
     urn: str,
     artefact: Any,
+    batch: _AnnotationBatch | None = None,
 ) -> None:
-    """Write localized_text and annotations for an artefact."""
+    """Write localized_text and annotations for an artefact.
+
+    If ``batch`` is provided, annotations are queued on it (caller flushes);
+    otherwise they are flushed immediately.
+    """
     _put_localized_text(conn, urn, "name", getattr(artefact, "name", None))
     _put_localized_text(conn, urn, "description", getattr(artefact, "description", None))
-    _put_annotations(conn, urn, getattr(artefact, "annotations", None))
+    anns = getattr(artefact, "annotations", None)
+    if batch is not None:
+        batch.add(urn, anns)
+    else:
+        _put_annotations(conn, urn, anns)
 
 
 # ── Reference graph ─────────────────────────────────────────────────────
@@ -330,6 +382,7 @@ def _delete_artefact_cascade(conn: duckdb.DuckDBPyConnection, urn: str) -> None:
     conn.execute("DELETE FROM concept WHERE scheme_urn = ?", p)
     conn.execute("DELETE FROM category WHERE scheme_urn = ?", p)
     conn.execute("DELETE FROM hierarchy_node WHERE hierarchy_urn = ?", p)
+    conn.execute("DELETE FROM hierarchy_level WHERE hierarchy_urn = ?", p)
     conn.execute("DELETE FROM agency WHERE scheme_urn = ?", p)
     conn.execute("DELETE FROM data_provider WHERE scheme_urn = ?", p)
     # Contact cleanup for agencies
@@ -393,7 +446,8 @@ def put_codelist(conn: duckdb.DuckDBPyConnection, cl: Codelist, *, source_regist
 
     _delete_artefact_cascade(conn, urn)
     _put_spine(conn, urn, "Codelist", cl, source_registry=source_registry)
-    _put_artefact_text_and_annotations(conn, urn, cl)
+    ann_batch = _AnnotationBatch()
+    _put_artefact_text_and_annotations(conn, urn, cl, ann_batch)
 
     refs: list[tuple[str, str, str, int | None]] = []
 
@@ -417,8 +471,7 @@ def put_codelist(conn: duckdb.DuckDBPyConnection, cl: Codelist, *, source_regist
                 ancestors,
             ]
             all_rows.append(row)
-            # Write annotations for each code
-            _put_annotations(conn, code_urn, code.annotations)
+            ann_batch.add(code_urn, code.annotations)
 
         _execute_bulk_insert(
             conn,
@@ -437,6 +490,7 @@ def put_codelist(conn: duckdb.DuckDBPyConnection, cl: Codelist, *, source_regist
             all_rows,
         )
 
+    ann_batch.flush(conn)
     _put_refs(conn, refs)
 
 
@@ -537,7 +591,8 @@ def put_data_structure(
 
     _delete_artefact_cascade(conn, urn)
     _put_spine(conn, urn, "DataStructure", dsd, source_registry=source_registry)
-    _put_artefact_text_and_annotations(conn, urn, dsd)
+    ann_batch = _AnnotationBatch()
+    _put_artefact_text_and_annotations(conn, urn, dsd, ann_batch)
 
     components: list[tuple[str, Any]] = [("Dimension", dim) for dim in dsd.dimensions]
     if dsd.time_dimension is not None:
@@ -554,10 +609,11 @@ def put_data_structure(
         # Component owner URN follows the same item-URN convention as Codes
         # and Concepts (`<dsd_urn-with-component-class>.<component_id>`).
         comp_urn = _make_item_urn(dsd, urn, comp_type, comp.id)
-        _put_annotations(conn, comp_urn, getattr(comp, "annotations", None))
+        ann_batch.add(comp_urn, getattr(comp, "annotations", None))
 
     if comp_rows:
         _execute_bulk_insert(conn, "dsd_component", _DSD_COMPONENT_COLS, comp_rows)
+    ann_batch.flush(conn)
 
     if dsd.groups:
         for group in dsd.groups:
@@ -602,7 +658,8 @@ def put_concept_scheme(
 
     _delete_artefact_cascade(conn, urn)
     _put_spine(conn, urn, "ConceptScheme", cs, source_registry=source_registry)
-    _put_artefact_text_and_annotations(conn, urn, cs)
+    ann_batch = _AnnotationBatch()
+    _put_artefact_text_and_annotations(conn, urn, cs, ann_batch)
 
     refs: list[tuple[str, str, str, int | None]] = []
 
@@ -643,7 +700,7 @@ def put_concept_scheme(
                     repr_max_length,
                 ]
             )
-            _put_annotations(conn, concept_urn_str, concept.annotations)
+            ann_batch.add(concept_urn_str, concept.annotations)
 
         _execute_bulk_insert(
             conn,
@@ -665,6 +722,7 @@ def put_concept_scheme(
             all_rows,
         )
 
+    ann_batch.flush(conn)
     _put_refs(conn, refs)
 
 
@@ -674,7 +732,8 @@ def put_hierarchy(conn: duckdb.DuckDBPyConnection, h: Hierarchy, *, source_regis
 
     _delete_artefact_cascade(conn, urn)
     _put_spine(conn, urn, "Hierarchy", h, source_registry=source_registry)
-    _put_artefact_text_and_annotations(conn, urn, h)
+    ann_batch = _AnnotationBatch()
+    _put_artefact_text_and_annotations(conn, urn, h, ann_batch)
 
     conn.execute(
         "INSERT INTO hierarchy_meta (urn, is_leveled) VALUES (?, ?)",
@@ -740,7 +799,7 @@ def put_hierarchy(conn: duckdb.DuckDBPyConnection, h: Hierarchy, *, source_regis
             if code_urn:
                 refs.append((urn, code_urn, "code", position))
 
-            _put_annotations(conn, node_urn, node.annotations)
+            ann_batch.add(node_urn, node.annotations)
 
         _execute_bulk_insert(
             conn,
@@ -763,6 +822,8 @@ def put_hierarchy(conn: duckdb.DuckDBPyConnection, h: Hierarchy, *, source_regis
         )
         _put_refs(conn, refs)
 
+    ann_batch.flush(conn)
+
 
 def _flatten_hcodes(
     nodes: Any,
@@ -825,7 +886,8 @@ def put_category_scheme(
 
     _delete_artefact_cascade(conn, urn)
     _put_spine(conn, urn, "CategoryScheme", cs, source_registry=source_registry)
-    _put_artefact_text_and_annotations(conn, urn, cs)
+    ann_batch = _AnnotationBatch()
+    _put_artefact_text_and_annotations(conn, urn, cs, ann_batch)
 
     # Flatten category tree
     flat: list[tuple[str, str | None, Any]] = []
@@ -851,7 +913,7 @@ def put_category_scheme(
                     ancestors,
                 ]
             )
-            _put_annotations(conn, cat_urn, cat.annotations)
+            ann_batch.add(cat_urn, cat.annotations)
 
         _execute_bulk_insert(
             conn,
@@ -870,6 +932,8 @@ def put_category_scheme(
             all_rows,
         )
 
+    ann_batch.flush(conn)
+
 
 def _flatten_categories(
     categories: Any,
@@ -941,7 +1005,8 @@ def put_agency_scheme(
 
     _delete_artefact_cascade(conn, urn)
     _put_spine(conn, urn, "AgencyScheme", ascheme, source_registry=source_registry)
-    _put_artefact_text_and_annotations(conn, urn, ascheme)
+    ann_batch = _AnnotationBatch()
+    _put_artefact_text_and_annotations(conn, urn, ascheme, ann_batch)
 
     if ascheme.agencies:
         all_rows: list[list[Any]] = []
@@ -979,7 +1044,7 @@ def put_agency_scheme(
                     ],
                 )
 
-            _put_annotations(conn, ag_urn, ag.annotations)
+            ann_batch.add(ag_urn, ag.annotations)
 
         _execute_bulk_insert(
             conn,
@@ -988,6 +1053,8 @@ def put_agency_scheme(
             all_rows,
         )
 
+    ann_batch.flush(conn)
+
 
 def put_data_provider_scheme(
     conn: duckdb.DuckDBPyConnection, dps: "DataProviderScheme", *, source_registry: str | None = None
@@ -997,7 +1064,8 @@ def put_data_provider_scheme(
 
     _delete_artefact_cascade(conn, urn)
     _put_spine(conn, urn, "DataProviderScheme", dps, source_registry=source_registry)
-    _put_artefact_text_and_annotations(conn, urn, dps)
+    ann_batch = _AnnotationBatch()
+    _put_artefact_text_and_annotations(conn, urn, dps, ann_batch)
 
     if dps.providers:
         all_rows: list[list[Any]] = []
@@ -1012,7 +1080,7 @@ def put_data_provider_scheme(
                     _istring_to_map(provider.description),
                 ]
             )
-            _put_annotations(conn, p_urn, provider.annotations)
+            ann_batch.add(p_urn, provider.annotations)
 
         _execute_bulk_insert(
             conn,
@@ -1021,6 +1089,8 @@ def put_data_provider_scheme(
             all_rows,
         )
 
+    ann_batch.flush(conn)
+
 
 def put_structure_set(
     conn: duckdb.DuckDBPyConnection, ss: "StructureSet", *, source_registry: str | None = None
@@ -1112,7 +1182,8 @@ def put_metadata_structure(
 
     _delete_artefact_cascade(conn, urn)
     _put_spine(conn, urn, "MetadataStructure", msd, source_registry=source_registry)
-    _put_artefact_text_and_annotations(conn, urn, msd)
+    ann_batch = _AnnotationBatch()
+    _put_artefact_text_and_annotations(conn, urn, msd, ann_batch)
 
     conn.execute("INSERT INTO metadata_structure (urn) VALUES (?)", [urn])
 
@@ -1177,10 +1248,11 @@ def put_metadata_structure(
                     ancestors,
                 ]
             )
-            _put_annotations(conn, attr_urn, attr.annotations)
+            ann_batch.add(attr_urn, attr.annotations)
 
         _execute_bulk_insert(conn, "metadata_attribute", _METADATA_ATTRIBUTE_COLS, rows)
 
+    ann_batch.flush(conn)
     _put_refs(conn, refs)