From dfac79eaef8ce66d70f49939117821e024dde596 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 17 Aug 2026 23:19:25 +0000 Subject: [PATCH] fix(backend): shrink output files via row groups, zstd and selective dict encoding Three changes reduce the size of the parquet files produced by the conversion backends. Row group sizing (DataFusion backend). pyarrow's ParquetWriter flushes one row group per write_table/write_batch call, and the default row group size is min(rows, 1 MiB) *per call*, not accumulated across calls. DataFusion's execute_stream() yields 8192-row batches, so streamed output previously degenerated into thousands of tiny row groups. Each row group compresses independently -- fragmenting zstd's compression window -- and carries its own column-chunk metadata (offsets, per-chunk statistics), inflating the file well beyond what one 122880-row group would occupy. Batches are now buffered and flushed as a single row group of _ROWS_PER_ROW_GROUP = 122880 rows, which also matches the row-group size of the LDBC input files (DuckDB's writer) for consistent chunk-level pruning and parallelism. Compression (DuckDB backend). cli.py's COPY now passes COMPRESSION ZSTD instead of DuckDB's snappy default, matching the codec used by the DataFusion backend and the LDBC input files. Dictionary encoding (DataFusion backend). pyarrow dictionary-encodes string columns by default, writing a dictionary page per row group. That can never pay off for columns whose values are effectively unique or random per row group -- CSR `target` neighbour ids and node primary keys -- since a dictionary of N distinct values costs roughly 2x the plain data. The new `plain_columns` argument writes those columns with PLAIN encoding while low-cardinality property columns keep dictionary encoding where it does pay off. --- icebug_format/_convert_datafusion.py | 54 ++++++++++++++++++++++++++-- icebug_format/cli.py | 8 +++-- 2 files changed, 57 insertions(+), 5 deletions(-) diff --git a/icebug_format/_convert_datafusion.py b/icebug_format/_convert_datafusion.py index 568f2a9..8ee8b79 100644 --- a/icebug_format/_convert_datafusion.py +++ b/icebug_format/_convert_datafusion.py @@ -40,6 +40,15 @@ _COMPRESSION = "zstd" +# Row groups are flushed on every ``ParquetWriter.write_table``/``write_batch`` +# call (the default row group size is ``min(rows, 1 MiB)`` *per call*, not +# accumulated across calls), so DataFusion's 8192-row stream batches would each +# become their own tiny row group, fragmenting zstd's compression window and +# row-group metadata. Buffer batches to this many rows and write each chunk as +# a single row group. 122880 matches the row-group size of the LDBC input +# parquet files (DuckDB's writer). +_ROWS_PER_ROW_GROUP = 122880 + # Arrow tables are registered with the engine in chunks of at most this many # rows. A single-chunk table would otherwise arrive as one huge record batch # and force operators (e.g. SortExec feeding the ROW_NUMBER window) to reserve @@ -95,7 +104,12 @@ def _build_context(datafusion, memory_limit: str | None): return datafusion.SessionContext(config=config, runtime=runtime) -def _stream_to_parquet(df, path: Path, target_to_uint64: bool = False) -> None: +def _stream_to_parquet( + df, + path: Path, + target_to_uint64: bool = False, + plain_columns: list[str] | None = None, +) -> None: """ Stream a DataFusion DataFrame into a Parquet file with icebug metadata. @@ -103,8 +117,31 @@ def _stream_to_parquet(df, path: Path, target_to_uint64: bool = False) -> None: rather than materialising the whole result, which is the point of the DataFusion backend. Falls back to ``to_arrow_table()`` only for empty results (where there is nothing to stream). + + ``plain_columns`` names columns whose values are effectively unique or + random per row group (CSR ``target`` neighbour ids, node primary keys). + pyarrow's default dictionary encoding writes a dictionary page per row + group that can never pay off for such columns (a dictionary of N distinct + values costs roughly 2x the plain data), inflating the file; those columns + are written with PLAIN encoding while any other (low-cardinality property) + columns keep dictionary encoding. """ writer = None + plain = set(plain_columns or ()) + buf: list[pa.RecordBatch] = [] + buf_rows = 0 + + def _flush() -> None: + """Write the buffered batches as one row group.""" + nonlocal buf, buf_rows + if not buf: + return + # row_group_size=None -> min(rows, 1 MiB) per call, i.e. exactly one + # row group of _ROWS_PER_ROW_GROUP rows (the tail chunk may be smaller). + writer.write_table(pa.Table.from_batches(buf)) + buf = [] + buf_rows = 0 + for batch in df.execute_stream(): rb = batch.to_pyarrow() if writer is None: @@ -113,13 +150,21 @@ def _stream_to_parquet(df, path: Path, target_to_uint64: bool = False) -> None: i = schema.get_field_index("target") schema = schema.set(i, pa.field(schema.field(i).name, pa.uint64())) schema = schema.with_metadata({"icebug_disk_version": ICEBUG_DISK_VERSION}) - writer = pq.ParquetWriter(path, schema, compression=_COMPRESSION) + # Enable dictionary encoding only where it can pay off; a listed + # use_dictionary of column names disables it for all others. + dict_cols = [c for c in schema.names if c not in plain] + writer = pq.ParquetWriter( + path, schema, compression=_COMPRESSION, use_dictionary=dict_cols + ) if target_to_uint64 and "target" in rb.schema.names: i = rb.schema.get_field_index("target") rb = rb.set_column( i, rb.schema.field(i).name, rb.column(i).cast(pa.uint64()) ) - writer.write_batch(rb) + buf.append(rb) + buf_rows += rb.num_rows + if buf_rows >= _ROWS_PER_ROW_GROUP: + _flush() if writer is None: table = df.to_arrow_table() if target_to_uint64 and "target" in table.schema.names: @@ -131,6 +176,7 @@ def _stream_to_parquet(df, path: Path, target_to_uint64: bool = False) -> None: ) _write_icebug_parquet(table, path, compression=_COMPRESSION) return + _flush() writer.close() @@ -346,6 +392,7 @@ def convert_graph( ctx.sql(idx_sql), out_dir / f"indices_{csr_rel_name(name)}.parquet", target_to_uint64=True, + plain_columns=["target"], ) # indptr: degrees from a streaming GROUP BY, then histogram + prefix sum. @@ -358,6 +405,7 @@ def convert_graph( _stream_to_parquet( ctx.sql(f"SELECT * FROM vertices ORDER BY {_q(pk)}"), out_dir / f"nodes_{name}.parquet", + plain_columns=[pk], ) _write_icebug_parquet( indptr, diff --git a/icebug_format/cli.py b/icebug_format/cli.py index ed77f7b..fea5a04 100644 --- a/icebug_format/cli.py +++ b/icebug_format/cli.py @@ -41,11 +41,15 @@ def _write_parquet_with_icebug_metadata( con, table_name: str, output_path: Path ) -> None: - """Export a DuckDB table to parquet with icebug_disk_version metadata.""" + """Export a DuckDB table to parquet with icebug_disk_version metadata. + + Uses zstd compression (DuckDB's COPY default is snappy) to match the + codec of the LDBC input files and keep the output compact. + """ con.execute( f""" COPY {table_name} TO ? - (FORMAT PARQUET, KV_METADATA {{ icebug_disk_version: '{ICEBUG_DISK_VERSION}' }}) + (FORMAT PARQUET, COMPRESSION ZSTD, KV_METADATA {{ icebug_disk_version: '{ICEBUG_DISK_VERSION}' }}) """, [str(output_path)], )