Experimental encoding of provenance multiblock files as parquet

python/convert_multiblock.py

@@ -0,0 +1,13 @@
+import click
+
+from lsst.pipe.base.quantum_graph._convert import convert_multiblock_to_parquet
+
+
+@click.command
+@click.argument("path")
+def run(path: str) -> None:
+    convert_multiblock_to_parquet(path, "output.zip")
+
+
+if __name__ == "__main__":
+    run()

python/lsst/pipe/base/quantum_graph/_convert.py

@@ -0,0 +1,84 @@
+import tempfile
+import zipfile
+from collections.abc import Iterator
+from contextlib import contextmanager
+from itertools import batched
+
+import pyarrow as pa
+import zstandard
+from duckdb import ColumnExpression, DuckDBPyConnection, connect
+
+from ._multiblock import MultiblockReader
+
+
+def convert_multiblock_to_parquet(input_zip_path: str, output_zip_path: str) -> None:
+    with zipfile.ZipFile(input_zip_path, "r") as zf:
+        if (cdict_path := zipfile.Path(zf, "compression_dict")).exists():
+            cdict = zstandard.ZstdCompressionDict(cdict_path.read_bytes())
+        decompressor = zstandard.ZstdDecompressor(cdict)
+        with initialize_duckdb_connection() as db:
+            for filename, id_field, variant_encode, row_group_size in (
+                ("quanta", "quantum_id", True, 122880),
+                ("datasets", "dataset_id", True, 122880),
+                # These don't really work as Parquet -- the individual values
+                # are too large, so we run out of memory easily when writing.
+                # ("metadata", None, False, 2048),
+                # ("logs", None, False, 2048),
+            ):
+                reader = _get_record_batch_reader(filename, zf, decompressor)
+                sql = db.from_arrow(reader)
+                if variant_encode:
+                    sql = sql.select("(data::JSON)::VARIANT as data")
+                if id_field is not None:
+                    sql = sql.select(
+                        ColumnExpression(f"data.{id_field}").cast("UUID").alias("id"),
+                        "data",
+                    ).order("id")
+                sql.to_parquet(f"{filename}.parquet", compression="zstd", row_group_size=row_group_size)
+
+
+@contextmanager
+def initialize_duckdb_connection(memory_limit: str = "16GB") -> Iterator[DuckDBPyConnection]:
+    """Set up an in-memory DuckDB database, applying a memory limit and using
+    the system tempdir for its working directory.
+    """
+    with (
+        tempfile.TemporaryDirectory(suffix=".duckdb") as tmpdir,
+        connect(
+            config={
+                # DuckDB will use up to 80% of system memory if you don't
+                # explicitly set it.
+                "memory_limit": memory_limit,
+                # DuckDB will use the current working directory if you don't
+                # set it, which is frequently a small quota-limited volume on
+                # USDF.
+                "temp_directory": tmpdir,
+                # For paranoia -- some of the scratch volumes at USDF are
+                # petabyte-sized.
+                "max_temp_directory_size": "128GB",
+            }
+        ) as conn,
+    ):
+        yield conn
+
+
+def _fetch_json_bytes(
+    filename: str, zip: zipfile.ZipFile, decompressor: zstandard.ZstdDecompressor
+) -> Iterator[bytes]:
+    for data in MultiblockReader.read_all_bytes_in_zip(zip, filename, int_size=8, page_size=8 * 1024 * 1024):
+        yield decompressor.decompress(data)
+
+
+def _to_batches(schema: pa.Schema, it: Iterator[bytes]) -> Iterator[pa.RecordBatch]:
+    for batch in batched(it, 1000):
+        array = pa.array(batch)
+        yield pa.RecordBatch.from_arrays([array], schema=schema)
+
+
+def _get_record_batch_reader(
+    filename: str, zip: zipfile.ZipFile, decompressor: zstandard.ZstdDecompressor
+) -> pa.RecordBatchReader:
+    schema = pa.schema([("data", pa.string())])
+    return pa.RecordBatchReader.from_batches(
+        schema, _to_batches(schema, _fetch_json_bytes(filename, zip, decompressor))
+    )