diff --git a/python/convert_multiblock.py b/python/convert_multiblock.py new file mode 100644 index 000000000..0dcf1b658 --- /dev/null +++ b/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() diff --git a/python/lsst/pipe/base/quantum_graph/_convert.py b/python/lsst/pipe/base/quantum_graph/_convert.py new file mode 100644 index 000000000..0bc49aa6e --- /dev/null +++ b/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)) + )