From 356d9c74996b4d62dfa7d5192df09489ab6cf059 Mon Sep 17 00:00:00 2001 From: Qi Chen Date: Mon, 10 Aug 2026 09:06:52 +0200 Subject: [PATCH 1/5] Upgrade DuckDB --- benchmark_data_tools/requirements.txt | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/benchmark_data_tools/requirements.txt b/benchmark_data_tools/requirements.txt index b011a066..cc9381ea 100644 --- a/benchmark_data_tools/requirements.txt +++ b/benchmark_data_tools/requirements.txt @@ -1,7 +1,11 @@ certifi==2025.8.3 charset-normalizer==3.4.3 click==8.2.1 -duckdb==1.3.2 +# 1.5.5 is the first stable release that respects PARQUET_VERSION 'V2'; 1.3.2 ignored it +# silently and wrote v1.0. Its dsdgen is still single-threaded, which dominates TPC-DS +# generation (98s vs 5.6s at SF10). Parallel dsdgen is in nightly only -- swap this for +# `duckdb --pre` to try it, and re-pin once a stable release ships it. +duckdb==1.5.5 idna==3.10 presto-python-client==0.8.4 requests==2.32.4 From b04df60be373313a8a85a30861523e413e27a055 Mon Sep 17 00:00:00 2001 From: Qi Chen Date: Mon, 10 Aug 2026 11:46:51 +0200 Subject: [PATCH 2/5] Refactor DuckDB utility functions to accept connection parameter - Updated `init_benchmark_tables` to accept a DuckDB connection parameter, allowing for more flexible database interactions. - Introduced `get_select_query` and `get_column_projection` functions to handle SQL queries and column type conversions, improving data handling in the generation process. - Modified `generate_data_files_with_duckdb` to utilize the new connection parameter and updated logic for writing metadata and table partitions. - Enhanced tests to validate the new functionality and ensure compatibility with both TPCH and TPCDS benchmarks. --- benchmark_data_tools/duckdb_utils.py | 25 +++- benchmark_data_tools/generate_data_files.py | 79 ++++++++----- benchmark_data_tools/requirements.txt | 5 +- benchmark_data_tools/row_group_sizing.py | 109 ++++++++++++++++++ .../tests/multi_file_partitioning_test.py | 57 +++++++++ .../tests/parquet_file_metadata_test.py | 12 +- .../tests/row_group_size_test.py | 9 +- 7 files changed, 255 insertions(+), 41 deletions(-) create mode 100644 benchmark_data_tools/row_group_sizing.py create mode 100644 benchmark_data_tools/tests/multi_file_partitioning_test.py diff --git a/benchmark_data_tools/duckdb_utils.py b/benchmark_data_tools/duckdb_utils.py index 0e16fc2a..6d299ec1 100644 --- a/benchmark_data_tools/duckdb_utils.py +++ b/benchmark_data_tools/duckdb_utils.py @@ -10,8 +10,8 @@ def quote_ident(name: str) -> str: return '"' + name.replace('"', '""') + '"' -def init_benchmark_tables(benchmark_type, scale_factor): - tables = duckdb.sql("SHOW TABLES").fetchall() +def init_benchmark_tables(benchmark_type, scale_factor, conn=duckdb): + tables = conn.sql("SHOW TABLES").fetchall() assert len(tables) == 0 if benchmark_type == "tpch": @@ -20,7 +20,7 @@ def init_benchmark_tables(benchmark_type, scale_factor): assert benchmark_type == "tpcds" function_name = "dsdgen" - duckdb.sql(f"INSTALL {benchmark_type}; LOAD {benchmark_type}; CALL {function_name}(sf = {scale_factor});") + conn.sql(f"INSTALL {benchmark_type}; LOAD {benchmark_type}; CALL {function_name}(sf = {scale_factor});") def drop_benchmark_tables(): @@ -51,3 +51,22 @@ def create_table_from_sample(table_name, data_path): def is_decimal_column(column_type): return bool(re.match(r"^DECIMAL\(\d+,\d+\)$", column_type)) + + +def get_select_query(table_name, convert_decimals_to_floats, conn=duckdb): + if convert_decimals_to_floats: + column_metadata_rows = conn.query(f"DESCRIBE {table_name}").fetchall() + column_projections = [get_column_projection(column_metadata) for column_metadata in column_metadata_rows] + query = f"SELECT {','.join(column_projections)} FROM {table_name}" + else: + query = f"SELECT * FROM {table_name}" + return query + + +def get_column_projection(column_metadata): + col_name, col_type, *_ = column_metadata + if is_decimal_column(col_type): + projection = f"CAST({col_name} AS DOUBLE) AS {col_name}" + else: + projection = col_name + return projection diff --git a/benchmark_data_tools/generate_data_files.py b/benchmark_data_tools/generate_data_files.py index 828133f7..c1bf192e 100644 --- a/benchmark_data_tools/generate_data_files.py +++ b/benchmark_data_tools/generate_data_files.py @@ -12,7 +12,8 @@ from pathlib import Path import duckdb -from duckdb_utils import init_benchmark_tables, is_decimal_column +from duckdb_utils import get_select_query, init_benchmark_tables +from row_group_sizing import row_group_row_count_probe _INTEGER_TYPES = frozenset(("INTEGER", "BIGINT", "SMALLINT", "TINYINT", "HUGEINT", "INT")) _HIGH_CARD_NDV_THRESHOLD = 0.99 @@ -190,42 +191,58 @@ def write_metadata(args): def generate_data_files_with_duckdb(args): - init_benchmark_tables(args.benchmark_type, args.scale_factor) + # Run the probe while DuckDB materializes the target dataset. + with row_group_row_count_probe( + args.benchmark_type, + args.approx_row_group_bytes, + args.scale_factor, + args.convert_decimals_to_floats, + ) as probed_row_counts: + init_benchmark_tables(args.benchmark_type, args.scale_factor) + row_group_rows = probed_row_counts() - with open(f"{args.data_dir_path}/metadata.json", "w") as file: - json.dump({"scale_factor": args.scale_factor}, file, indent=2) - file.write("\n") + write_metadata(args) - tables = duckdb.sql("SHOW TABLES").fetchall() - for (table_name,) in tables: + tables = [t[0] for t in duckdb.sql("SHOW TABLES").fetchall()] + for table_name in tables: table_data_dir = f"{args.data_dir_path}/{table_name}" Path(table_data_dir).mkdir(exist_ok=False) - duckdb.sql( - f"COPY ({get_select_query(table_name, args.convert_decimals_to_floats)}) " - f"TO '{table_data_dir}/{table_name}.parquet' (FORMAT parquet)" + _write_table_partitions( + table_name, + table_data_dir, + row_group_rows.get(table_name), + args.convert_decimals_to_floats, + args.max_rows_per_file, ) -def get_select_query(table_name, convert_decimals_to_floats): - if convert_decimals_to_floats: - column_metadata_rows = duckdb.query(f"DESCRIBE {table_name}").fetchall() - column_projections = [ - get_column_projection(column_metadata, convert_decimals_to_floats) - for column_metadata in column_metadata_rows - ] - query = f"SELECT {','.join(column_projections)} FROM {table_name}" - else: - query = f"SELECT * FROM {table_name}" - return query - - -def get_column_projection(column_metadata, convert_decimals_to_floats): - col_name, col_type, *_ = column_metadata - if convert_decimals_to_floats and is_decimal_column(col_type): - projection = f"CAST({col_name} AS DOUBLE) AS {col_name}" - else: - projection = col_name - return projection +def _write_table_partitions( + table_name, + table_data_dir, + estimated_rows_per_row_group, + convert_decimals_to_floats, + max_rows_per_file, +): + """Write 1-indexed Parquet parts, respecting max_rows_per_file when set.""" + select_query = get_select_query(table_name, convert_decimals_to_floats) + row_count = duckdb.sql(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0] + if estimated_rows_per_row_group is None: + estimated_rows_per_row_group = row_count + # The estimate can exceed the target table's row count. + rows_per_row_group = min(estimated_rows_per_row_group, row_count) + num_partitions = math.ceil(row_count / max_rows_per_file) if max_rows_per_file else 1 + for part in range(num_partitions): + # Avoid a redundant LIMIT/OFFSET for a single part. + partition_query = ( + f"{select_query} LIMIT {max_rows_per_file} OFFSET {part * max_rows_per_file}" + if num_partitions > 1 + else select_query + ) + file_path = f"{table_data_dir}/{table_name}-{part + 1}.parquet" + duckdb.sql( + f"COPY ({partition_query}) TO '{file_path}' " + f"(FORMAT parquet, PARQUET_VERSION 'V2', ROW_GROUP_SIZE {rows_per_row_group})" + ) def get_tpchgen_codec_args(codec_defs, table_name): @@ -309,7 +326,7 @@ def build_default_codec_defs(): Achieved ~11% improvement over baseline with ~15% smaller dataset. """ with duckdb.connect() as conn: - conn.execute(f"INSTALL tpch; LOAD tpch; CALL dbgen(sf = {_SAMPLE_SF});") + init_benchmark_tables("tpch", _SAMPLE_SF, conn) tables = [row[0] for row in conn.execute("SHOW TABLES").fetchall()] config = {"tables": []} diff --git a/benchmark_data_tools/requirements.txt b/benchmark_data_tools/requirements.txt index cc9381ea..da2bb1cf 100644 --- a/benchmark_data_tools/requirements.txt +++ b/benchmark_data_tools/requirements.txt @@ -1,10 +1,7 @@ certifi==2025.8.3 charset-normalizer==3.4.3 click==8.2.1 -# 1.5.5 is the first stable release that respects PARQUET_VERSION 'V2'; 1.3.2 ignored it -# silently and wrote v1.0. Its dsdgen is still single-threaded, which dominates TPC-DS -# generation (98s vs 5.6s at SF10). Parallel dsdgen is in nightly only -- swap this for -# `duckdb --pre` to try it, and re-pin once a stable release ships it. +# DuckDB 1.5.5 supports Parquet V2; parallel dsdgen is not yet in a stable release. duckdb==1.5.5 idna==3.10 presto-python-client==0.8.4 diff --git a/benchmark_data_tools/row_group_sizing.py b/benchmark_data_tools/row_group_sizing.py new file mode 100644 index 00000000..cbc09bde --- /dev/null +++ b/benchmark_data_tools/row_group_sizing.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +import multiprocessing +import tempfile +from concurrent.futures import ProcessPoolExecutor +from contextlib import contextmanager +from pathlib import Path + +import duckdb +import pyarrow.parquet as pq +from duckdb_utils import get_select_query, init_benchmark_tables + +_ROW_GROUP_GRANULARITY = 2048 # DuckDB rounds ROW_GROUP_SIZE to its vector size +_MAX_PROBE_SCALE_FACTOR = 10 +_STAGE1_ROWS = 200_000 +_STAGE2_SLICE = 1.2 +_PROBE_MEMORY_LIMIT = "8GB" + + +@contextmanager +def row_group_row_count_probe( + benchmark_type, target_bytes, target_scale_factor, convert_decimals_to_floats +): + """Yield a callable that waits for row-count estimates from a spawned process.""" + # dbgen/dsdgen use process-global state and cannot run concurrently in one process. + with ProcessPoolExecutor( + max_workers=1, mp_context=multiprocessing.get_context("spawn") + ) as executor: + future = executor.submit( + get_row_group_row_counts, + benchmark_type, + target_bytes, + target_scale_factor, + convert_decimals_to_floats, + ) + yield future.result + + +def get_row_group_row_counts( + benchmark_type, target_bytes, target_scale_factor, convert_decimals_to_floats +): + """Measure rows per row group for every table in a throwaway dataset.""" + scale_factor = ( + target_scale_factor + if target_scale_factor <= 1 + else min(target_scale_factor / 10, _MAX_PROBE_SCALE_FACTOR) + ) + + row_counts = {} + with duckdb.connect() as conn: + conn.execute(f"SET memory_limit='{_PROBE_MEMORY_LIMIT}'") + init_benchmark_tables(benchmark_type, scale_factor, conn) + for (table_name,) in conn.execute("SHOW TABLES").fetchall(): + select_query = get_select_query(table_name, convert_decimals_to_floats, conn) + table_rows = conn.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0] + rows = _measure_row_group_rows(conn, select_query, table_rows, target_bytes) + if rows is not None: + row_counts[table_name] = rows + return row_counts + + +def _measure_row_group_rows(conn, select_query, table_rows, target_bytes): + """Rows per row group so that a row group weighs about target_bytes.""" + if table_rows == 0: + return None + + with tempfile.TemporaryDirectory() as tmp_dir: + probe_path = Path(tmp_dir) / "probe.parquet" + + # First pass: estimate bytes/row using DuckDB's default row-group size. + _write_probe(conn, f"{select_query} LIMIT {_STAGE1_ROWS}", probe_path) + rows = _rows_for_target(_bytes_per_row(probe_path), target_bytes) + if rows is None or rows >= table_rows: + # A second write cannot fill the estimated row group or improve the result. + return rows + + # Second pass: measure one full row group near the requested size. + slice_rows = int(_STAGE2_SLICE * rows) + _write_probe(conn, f"{select_query} LIMIT {slice_rows}", probe_path, rows) + refined = _rows_for_target(_bytes_per_row(probe_path), target_bytes) + + return rows if refined is None else refined + + +def _write_probe(conn, query, path, row_group_rows=None): + options = "FORMAT parquet, PARQUET_VERSION 'V2'" + if row_group_rows is not None: + options += f", ROW_GROUP_SIZE {row_group_rows}" + conn.sql(f"COPY ({query}) TO '{path}' ({options})") + + +def _bytes_per_row(path): + """Read bytes/row from full row groups, excluding a trailing partial group.""" + metadata = pq.ParquetFile(path).metadata + num_row_groups = metadata.num_row_groups + full = range(num_row_groups - 1) if num_row_groups > 1 else range(num_row_groups) + total_rows = sum(metadata.row_group(i).num_rows for i in full) + if total_rows == 0: + return None + return sum(metadata.row_group(i).total_byte_size for i in full) / total_rows + + +def _rows_for_target(bytes_per_row, target_bytes): + """Convert bytes to rows and round to DuckDB's 2,048-row granularity.""" + if not bytes_per_row: + return None + rows = int(round(target_bytes / bytes_per_row / _ROW_GROUP_GRANULARITY)) + return max(rows, 1) * _ROW_GROUP_GRANULARITY diff --git a/benchmark_data_tools/tests/multi_file_partitioning_test.py b/benchmark_data_tools/tests/multi_file_partitioning_test.py new file mode 100644 index 00000000..694f6456 --- /dev/null +++ b/benchmark_data_tools/tests/multi_file_partitioning_test.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import duckdb +import pyarrow.parquet as pq +from generate_data_files import generate_data_files + + +def test_max_rows_per_file_splits_tables_tpcds(setup_and_teardown): + """Validate that max_rows_per_file controls TPC-DS file partitioning. + + Verifies that: + - At least one table is split into multiple files + - Part files use contiguous 1-based names, including single-part tables + - Every part except the last contains exactly max_rows_per_file rows + - The last part contains between one and max_rows_per_file rows + - The parts contain the same total number of rows as the source table + """ + data_dir_path, args = setup_and_teardown + args.benchmark_type = "tpcds" + args.use_duckdb = True + # SF1 tables fit below the 100M default, so use a lower limit to exercise splitting. + args.max_rows_per_file = 500_000 + generate_data_files(args) + + assert_partitions_respect_max_rows_per_file(data_dir_path, args.max_rows_per_file) + + +def assert_partitions_respect_max_rows_per_file(data_dir_path, max_rows_per_file): + split_table_count = 0 + for table_dir in sorted(path for path in Path(data_dir_path).iterdir() if path.is_dir()): + table_name = table_dir.name + file_paths = list(table_dir.glob("*.parquet")) + num_parts = len(file_paths) + + # All tables use the same 1-indexed layout, including single-part tables. + assert {path.name for path in file_paths} == { + f"{table_name}-{part}.parquet" for part in range(1, num_parts + 1) + } + + rows_per_part = [ + pq.ParquetFile(table_dir / f"{table_name}-{part}.parquet").metadata.num_rows + for part in range(1, num_parts + 1) + ] + # Only the final part may contain fewer rows than the limit. + assert all(rows == max_rows_per_file for rows in rows_per_part[:-1]), rows_per_part + assert 0 < rows_per_part[-1] <= max_rows_per_file, rows_per_part + # Verify that the parts preserve the table's row count. + expected_rows = duckdb.sql(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0] + assert sum(rows_per_part) == expected_rows + + if num_parts > 1: + split_table_count += 1 + + assert split_table_count > 0, "no table was split, so the partitioning path was not tested" diff --git a/benchmark_data_tools/tests/parquet_file_metadata_test.py b/benchmark_data_tools/tests/parquet_file_metadata_test.py index 60879367..a18a3e41 100644 --- a/benchmark_data_tools/tests/parquet_file_metadata_test.py +++ b/benchmark_data_tools/tests/parquet_file_metadata_test.py @@ -2,14 +2,22 @@ # SPDX-License-Identifier: Apache-2.0 import pyarrow.parquet as pq +import pytest from generate_data_files import generate_data_files from .common_fixtures import get_all_parquet_relative_file_paths -def test_generated_files_use_v2_page_format(setup_and_teardown): - """Verify every generated Parquet file is written with the v2 format.""" +@pytest.mark.parametrize( + "benchmark_type,use_duckdb", + [("tpch", False), ("tpcds", True)], + ids=["tpch-tpchgen", "tpcds-duckdb"], +) +def test_generated_files_use_v2_page_format(setup_and_teardown, benchmark_type, use_duckdb): + """Verify that both generation paths write Parquet V2 files.""" data_dir_path, args = setup_and_teardown + args.benchmark_type = benchmark_type + args.use_duckdb = use_duckdb generate_data_files(args) for file_path in get_all_parquet_relative_file_paths(data_dir_path): diff --git a/benchmark_data_tools/tests/row_group_size_test.py b/benchmark_data_tools/tests/row_group_size_test.py index f91a5e42..e373cc6c 100644 --- a/benchmark_data_tools/tests/row_group_size_test.py +++ b/benchmark_data_tools/tests/row_group_size_test.py @@ -15,7 +15,12 @@ _MIN_ROW_GROUPS_FOR_SIZE_CHECK = 4 -def test_approx_row_group_bytes_parameter(setup_and_teardown): +@pytest.mark.parametrize( + "benchmark_type,use_duckdb", + [("tpch", False), ("tpcds", True)], + ids=["tpch-tpchgen", "tpcds-duckdb"], +) +def test_approx_row_group_bytes_parameter(setup_and_teardown, benchmark_type, use_duckdb): """Validate that the approx_row_group_bytes parameter controls row group sizing. Verifies that: @@ -25,6 +30,8 @@ def test_approx_row_group_bytes_parameter(setup_and_teardown): - Small tables are excluded from size checks (see _MIN_ROW_GROUPS_FOR_SIZE_CHECK) """ data_dir_path, args = setup_and_teardown + args.benchmark_type = benchmark_type + args.use_duckdb = use_duckdb args.approx_row_group_bytes = 1024 * 1024 generate_data_files(args) From 6763e52dac80ab3bb10206d9b2973c0ee5fc5517 Mon Sep 17 00:00:00 2001 From: Qi Chen Date: Tue, 11 Aug 2026 16:29:32 +0200 Subject: [PATCH 3/5] Clean code structure --- benchmark_data_tools/duckdb_utils.py | 7 +++ benchmark_data_tools/generate_data_files.py | 47 ++++++++++++------- benchmark_data_tools/requirements.txt | 1 - benchmark_data_tools/row_group_sizing.py | 19 +++----- .../tests/codec_definitions_test.py | 2 + benchmark_data_tools/tests/common_fixtures.py | 8 ++-- .../tests/multi_file_partitioning_test.py | 36 ++++++++------ .../tests/parquet_file_metadata_test.py | 12 +---- .../tests/row_group_size_test.py | 9 +--- 9 files changed, 73 insertions(+), 68 deletions(-) diff --git a/benchmark_data_tools/duckdb_utils.py b/benchmark_data_tools/duckdb_utils.py index 6d299ec1..2cfa3cc8 100644 --- a/benchmark_data_tools/duckdb_utils.py +++ b/benchmark_data_tools/duckdb_utils.py @@ -53,6 +53,13 @@ def is_decimal_column(column_type): return bool(re.match(r"^DECIMAL\(\d+,\d+\)$", column_type)) +def copy_to_parquet(select_query, file_path, row_group_rows=None, conn=duckdb): + options = "FORMAT parquet, PARQUET_VERSION 'V2'" + if row_group_rows is not None: + options += f", ROW_GROUP_SIZE {row_group_rows}" + conn.sql(f"COPY ({select_query}) TO '{file_path}' ({options})") + + def get_select_query(table_name, convert_decimals_to_floats, conn=duckdb): if convert_decimals_to_floats: column_metadata_rows = conn.query(f"DESCRIBE {table_name}").fetchall() diff --git a/benchmark_data_tools/generate_data_files.py b/benchmark_data_tools/generate_data_files.py index c1bf192e..d4b609f8 100644 --- a/benchmark_data_tools/generate_data_files.py +++ b/benchmark_data_tools/generate_data_files.py @@ -12,7 +12,7 @@ from pathlib import Path import duckdb -from duckdb_utils import get_select_query, init_benchmark_tables +from duckdb_utils import copy_to_parquet, get_select_query, init_benchmark_tables from row_group_sizing import row_group_row_count_probe _INTEGER_TYPES = frozenset(("INTEGER", "BIGINT", "SMALLINT", "TINYINT", "HUGEINT", "INT")) @@ -100,6 +100,8 @@ def generate_data_files(args): print("generating with duckdb") generate_data_files_with_duckdb(args) + write_metadata(args) + def generate_data_files_with_tpchgen(args, codec_defs): local_installs_bin = Path(__file__).resolve().parent / ".local_installs" / "bin" @@ -141,8 +143,6 @@ def generate_data_files_with_tpchgen(args, codec_defs): if args.verbose: print(f"Raw data created at: {raw_data_path}") - write_metadata(args) - # This dictionary maps each table to the number of partitions it should have based on it's # expected file size relative to the SF. @@ -201,8 +201,6 @@ def generate_data_files_with_duckdb(args): init_benchmark_tables(args.benchmark_type, args.scale_factor) row_group_rows = probed_row_counts() - write_metadata(args) - tables = [t[0] for t in duckdb.sql("SHOW TABLES").fetchall()] for table_name in tables: table_data_dir = f"{args.data_dir_path}/{table_name}" @@ -213,6 +211,7 @@ def generate_data_files_with_duckdb(args): row_group_rows.get(table_name), args.convert_decimals_to_floats, args.max_rows_per_file, + args.num_threads, ) @@ -222,6 +221,7 @@ def _write_table_partitions( estimated_rows_per_row_group, convert_decimals_to_floats, max_rows_per_file, + num_threads, ): """Write 1-indexed Parquet parts, respecting max_rows_per_file when set.""" select_query = get_select_query(table_name, convert_decimals_to_floats) @@ -230,19 +230,30 @@ def _write_table_partitions( estimated_rows_per_row_group = row_count # The estimate can exceed the target table's row count. rows_per_row_group = min(estimated_rows_per_row_group, row_count) - num_partitions = math.ceil(row_count / max_rows_per_file) if max_rows_per_file else 1 - for part in range(num_partitions): - # Avoid a redundant LIMIT/OFFSET for a single part. - partition_query = ( - f"{select_query} LIMIT {max_rows_per_file} OFFSET {part * max_rows_per_file}" - if num_partitions > 1 - else select_query - ) + num_partitions = max(math.ceil(row_count / max_rows_per_file), 1) if max_rows_per_file else 1 + + def write_partition(part): + partition_query = select_query + if num_partitions > 1: + start_row = part * max_rows_per_file + partition_query += f" WHERE rowid >= {start_row} AND rowid < {start_row + max_rows_per_file}" file_path = f"{table_data_dir}/{table_name}-{part + 1}.parquet" - duckdb.sql( - f"COPY ({partition_query}) TO '{file_path}' " - f"(FORMAT parquet, PARQUET_VERSION 'V2', ROW_GROUP_SIZE {rows_per_row_group})" - ) + conn = duckdb.cursor() + try: + copy_to_parquet(partition_query, file_path, rows_per_row_group, conn) + finally: + conn.close() + + max_workers = min(num_threads, num_partitions) + if max_workers == 1: + for part in range(num_partitions): + write_partition(part) + return + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [executor.submit(write_partition, part) for part in range(num_partitions)] + for future in futures: + future.result() def get_tpchgen_codec_args(codec_defs, table_name): @@ -408,7 +419,7 @@ def build_default_codec_defs(): type=int, required=False, default=4, - help="Number of threads to generate data with tpchgen", + help="Number of concurrent data generation tasks", ) parser.add_argument( "-v", "--verbose", action="store_true", required=False, default=False, help="Extra verbose logging" diff --git a/benchmark_data_tools/requirements.txt b/benchmark_data_tools/requirements.txt index da2bb1cf..c73e9f11 100644 --- a/benchmark_data_tools/requirements.txt +++ b/benchmark_data_tools/requirements.txt @@ -1,7 +1,6 @@ certifi==2025.8.3 charset-normalizer==3.4.3 click==8.2.1 -# DuckDB 1.5.5 supports Parquet V2; parallel dsdgen is not yet in a stable release. duckdb==1.5.5 idna==3.10 presto-python-client==0.8.4 diff --git a/benchmark_data_tools/row_group_sizing.py b/benchmark_data_tools/row_group_sizing.py index cbc09bde..8a89bb4c 100644 --- a/benchmark_data_tools/row_group_sizing.py +++ b/benchmark_data_tools/row_group_sizing.py @@ -9,12 +9,12 @@ import duckdb import pyarrow.parquet as pq -from duckdb_utils import get_select_query, init_benchmark_tables +from duckdb_utils import copy_to_parquet, get_select_query, init_benchmark_tables _ROW_GROUP_GRANULARITY = 2048 # DuckDB rounds ROW_GROUP_SIZE to its vector size _MAX_PROBE_SCALE_FACTOR = 10 -_STAGE1_ROWS = 200_000 -_STAGE2_SLICE = 1.2 +_STAGE1_ROWS = 122_880 # one DuckDB default row group +_STAGE2_SLICE = 1.2 # one full estimated group plus a bounded partial group _PROBE_MEMORY_LIMIT = "8GB" @@ -69,27 +69,20 @@ def _measure_row_group_rows(conn, select_query, table_rows, target_bytes): probe_path = Path(tmp_dir) / "probe.parquet" # First pass: estimate bytes/row using DuckDB's default row-group size. - _write_probe(conn, f"{select_query} LIMIT {_STAGE1_ROWS}", probe_path) + copy_to_parquet(f"{select_query} LIMIT {_STAGE1_ROWS}", probe_path, conn=conn) rows = _rows_for_target(_bytes_per_row(probe_path), target_bytes) if rows is None or rows >= table_rows: # A second write cannot fill the estimated row group or improve the result. return rows - # Second pass: measure one full row group near the requested size. + # Second pass: bytes/row changes with row-group size, so remeasure near the requested size. slice_rows = int(_STAGE2_SLICE * rows) - _write_probe(conn, f"{select_query} LIMIT {slice_rows}", probe_path, rows) + copy_to_parquet(f"{select_query} LIMIT {slice_rows}", probe_path, rows, conn) refined = _rows_for_target(_bytes_per_row(probe_path), target_bytes) return rows if refined is None else refined -def _write_probe(conn, query, path, row_group_rows=None): - options = "FORMAT parquet, PARQUET_VERSION 'V2'" - if row_group_rows is not None: - options += f", ROW_GROUP_SIZE {row_group_rows}" - conn.sql(f"COPY ({query}) TO '{path}' ({options})") - - def _bytes_per_row(path): """Read bytes/row from full row groups, excluding a trailing partial group.""" metadata = pq.ParquetFile(path).metadata diff --git a/benchmark_data_tools/tests/codec_definitions_test.py b/benchmark_data_tools/tests/codec_definitions_test.py index 49c807f9..4d56edae 100644 --- a/benchmark_data_tools/tests/codec_definitions_test.py +++ b/benchmark_data_tools/tests/codec_definitions_test.py @@ -15,6 +15,8 @@ TEST_NON_DEFAULT_COMPRESSION_PATH = TESTS_DIR / "test_codec_definitions_non_default_compression.json" TEST_INVALID_COMPRESSION_PATH = TESTS_DIR / "test_codec_definitions_invalid_compression.json" +pytestmark = pytest.mark.parametrize("setup_and_teardown", ["tpch"], indirect=True) + def test_default_codec_defs_applied(setup_and_teardown): """Generate data with default codec defs and verify encodings in parquet metadata. diff --git a/benchmark_data_tools/tests/common_fixtures.py b/benchmark_data_tools/tests/common_fixtures.py index 3c6478da..fb1727c6 100644 --- a/benchmark_data_tools/tests/common_fixtures.py +++ b/benchmark_data_tools/tests/common_fixtures.py @@ -33,12 +33,12 @@ class DataGenArgs: codec_definitions: str = None -@pytest.fixture -def setup_and_teardown(): - test_data_dir_path = os.path.abspath("./tpch_test") +@pytest.fixture(params=["tpch", "tpcds"]) +def setup_and_teardown(request): + test_data_dir_path = os.path.abspath(f"./{request.param}_test") try: args = DataGenArgs( - benchmark_type="tpch", + benchmark_type=request.param, data_dir_path=test_data_dir_path, scale_factor=1.0, # Setting convert_decimals_to_floats to True ensures that the diff --git a/benchmark_data_tools/tests/multi_file_partitioning_test.py b/benchmark_data_tools/tests/multi_file_partitioning_test.py index 694f6456..f5f2bcf8 100644 --- a/benchmark_data_tools/tests/multi_file_partitioning_test.py +++ b/benchmark_data_tools/tests/multi_file_partitioning_test.py @@ -5,30 +5,42 @@ import duckdb import pyarrow.parquet as pq +from duckdb_utils import init_benchmark_tables from generate_data_files import generate_data_files -def test_max_rows_per_file_splits_tables_tpcds(setup_and_teardown): - """Validate that max_rows_per_file controls TPC-DS file partitioning. +def test_max_rows_per_file_splits_tables(setup_and_teardown): + """Validate that max_rows_per_file controls file partitioning. Verifies that: - At least one table is split into multiple files - Part files use contiguous 1-based names, including single-part tables - - Every part except the last contains exactly max_rows_per_file rows - - The last part contains between one and max_rows_per_file rows + - Every part contains between one and max_rows_per_file rows - The parts contain the same total number of rows as the source table """ data_dir_path, args = setup_and_teardown - args.benchmark_type = "tpcds" - args.use_duckdb = True # SF1 tables fit below the 100M default, so use a lower limit to exercise splitting. args.max_rows_per_file = 500_000 + expected_row_counts = get_expected_row_counts(args.benchmark_type, args.scale_factor) generate_data_files(args) - assert_partitions_respect_max_rows_per_file(data_dir_path, args.max_rows_per_file) + assert_partitions_respect_max_rows_per_file( + data_dir_path, + args.max_rows_per_file, + expected_row_counts, + ) -def assert_partitions_respect_max_rows_per_file(data_dir_path, max_rows_per_file): +def get_expected_row_counts(benchmark_type, scale_factor): + with duckdb.connect() as conn: + init_benchmark_tables(benchmark_type, scale_factor, conn) + return { + table_name: conn.sql(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0] + for (table_name,) in conn.sql("SHOW TABLES").fetchall() + } + + +def assert_partitions_respect_max_rows_per_file(data_dir_path, max_rows_per_file, expected_row_counts): split_table_count = 0 for table_dir in sorted(path for path in Path(data_dir_path).iterdir() if path.is_dir()): table_name = table_dir.name @@ -44,12 +56,8 @@ def assert_partitions_respect_max_rows_per_file(data_dir_path, max_rows_per_file pq.ParquetFile(table_dir / f"{table_name}-{part}.parquet").metadata.num_rows for part in range(1, num_parts + 1) ] - # Only the final part may contain fewer rows than the limit. - assert all(rows == max_rows_per_file for rows in rows_per_part[:-1]), rows_per_part - assert 0 < rows_per_part[-1] <= max_rows_per_file, rows_per_part - # Verify that the parts preserve the table's row count. - expected_rows = duckdb.sql(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0] - assert sum(rows_per_part) == expected_rows + assert all(0 < rows <= max_rows_per_file for rows in rows_per_part), rows_per_part + assert sum(rows_per_part) == expected_row_counts[table_name] if num_parts > 1: split_table_count += 1 diff --git a/benchmark_data_tools/tests/parquet_file_metadata_test.py b/benchmark_data_tools/tests/parquet_file_metadata_test.py index a18a3e41..60879367 100644 --- a/benchmark_data_tools/tests/parquet_file_metadata_test.py +++ b/benchmark_data_tools/tests/parquet_file_metadata_test.py @@ -2,22 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 import pyarrow.parquet as pq -import pytest from generate_data_files import generate_data_files from .common_fixtures import get_all_parquet_relative_file_paths -@pytest.mark.parametrize( - "benchmark_type,use_duckdb", - [("tpch", False), ("tpcds", True)], - ids=["tpch-tpchgen", "tpcds-duckdb"], -) -def test_generated_files_use_v2_page_format(setup_and_teardown, benchmark_type, use_duckdb): - """Verify that both generation paths write Parquet V2 files.""" +def test_generated_files_use_v2_page_format(setup_and_teardown): + """Verify every generated Parquet file is written with the v2 format.""" data_dir_path, args = setup_and_teardown - args.benchmark_type = benchmark_type - args.use_duckdb = use_duckdb generate_data_files(args) for file_path in get_all_parquet_relative_file_paths(data_dir_path): diff --git a/benchmark_data_tools/tests/row_group_size_test.py b/benchmark_data_tools/tests/row_group_size_test.py index e373cc6c..f91a5e42 100644 --- a/benchmark_data_tools/tests/row_group_size_test.py +++ b/benchmark_data_tools/tests/row_group_size_test.py @@ -15,12 +15,7 @@ _MIN_ROW_GROUPS_FOR_SIZE_CHECK = 4 -@pytest.mark.parametrize( - "benchmark_type,use_duckdb", - [("tpch", False), ("tpcds", True)], - ids=["tpch-tpchgen", "tpcds-duckdb"], -) -def test_approx_row_group_bytes_parameter(setup_and_teardown, benchmark_type, use_duckdb): +def test_approx_row_group_bytes_parameter(setup_and_teardown): """Validate that the approx_row_group_bytes parameter controls row group sizing. Verifies that: @@ -30,8 +25,6 @@ def test_approx_row_group_bytes_parameter(setup_and_teardown, benchmark_type, us - Small tables are excluded from size checks (see _MIN_ROW_GROUPS_FOR_SIZE_CHECK) """ data_dir_path, args = setup_and_teardown - args.benchmark_type = benchmark_type - args.use_duckdb = use_duckdb args.approx_row_group_bytes = 1024 * 1024 generate_data_files(args) From 62ccf191ac829bb7586cd6ac10748fb09ce9bbcf Mon Sep 17 00:00:00 2001 From: Qi Chen Date: Wed, 12 Aug 2026 10:03:44 +0200 Subject: [PATCH 4/5] Refactor DuckDB utility functions and enhance data generation - Updated `get_select_query` to use `get_column_projection_with_decimals_as_double` for improved handling of decimal columns. - Added installation command for DuckDB in `generate_data_files_with_duckdb` to prevent concurrent installations. - Simplified row group sizing logic in `row_group_sizing.py` for better performance and clarity. --- benchmark_data_tools/duckdb_utils.py | 7 ++++-- benchmark_data_tools/generate_data_files.py | 3 +++ benchmark_data_tools/row_group_sizing.py | 25 +++++++++------------ 3 files changed, 18 insertions(+), 17 deletions(-) diff --git a/benchmark_data_tools/duckdb_utils.py b/benchmark_data_tools/duckdb_utils.py index 2cfa3cc8..a53a49ab 100644 --- a/benchmark_data_tools/duckdb_utils.py +++ b/benchmark_data_tools/duckdb_utils.py @@ -63,14 +63,17 @@ def copy_to_parquet(select_query, file_path, row_group_rows=None, conn=duckdb): def get_select_query(table_name, convert_decimals_to_floats, conn=duckdb): if convert_decimals_to_floats: column_metadata_rows = conn.query(f"DESCRIBE {table_name}").fetchall() - column_projections = [get_column_projection(column_metadata) for column_metadata in column_metadata_rows] + column_projections = [ + get_column_projection_with_decimals_as_double(column_metadata) + for column_metadata in column_metadata_rows + ] query = f"SELECT {','.join(column_projections)} FROM {table_name}" else: query = f"SELECT * FROM {table_name}" return query -def get_column_projection(column_metadata): +def get_column_projection_with_decimals_as_double(column_metadata): col_name, col_type, *_ = column_metadata if is_decimal_column(col_type): projection = f"CAST({col_name} AS DOUBLE) AS {col_name}" diff --git a/benchmark_data_tools/generate_data_files.py b/benchmark_data_tools/generate_data_files.py index d4b609f8..2e3dc31f 100644 --- a/benchmark_data_tools/generate_data_files.py +++ b/benchmark_data_tools/generate_data_files.py @@ -191,6 +191,9 @@ def write_metadata(args): def generate_data_files_with_duckdb(args): + # Avoid concurrent first-time extension installation across processes. + duckdb.sql(f"INSTALL {args.benchmark_type}") + # Run the probe while DuckDB materializes the target dataset. with row_group_row_count_probe( args.benchmark_type, diff --git a/benchmark_data_tools/row_group_sizing.py b/benchmark_data_tools/row_group_sizing.py index 8a89bb4c..4693b626 100644 --- a/benchmark_data_tools/row_group_sizing.py +++ b/benchmark_data_tools/row_group_sizing.py @@ -11,10 +11,11 @@ import pyarrow.parquet as pq from duckdb_utils import copy_to_parquet, get_select_query, init_benchmark_tables -_ROW_GROUP_GRANULARITY = 2048 # DuckDB rounds ROW_GROUP_SIZE to its vector size +# DuckDB Parquet sizing: https://duckdb.org/docs/current/data/parquet/tips +_ROW_GROUP_GRANULARITY = 2048 +_STAGE1_ROWS = 122_880 + _MAX_PROBE_SCALE_FACTOR = 10 -_STAGE1_ROWS = 122_880 # one DuckDB default row group -_STAGE2_SLICE = 1.2 # one full estimated group plus a bounded partial group _PROBE_MEMORY_LIMIT = "8GB" @@ -75,28 +76,22 @@ def _measure_row_group_rows(conn, select_query, table_rows, target_bytes): # A second write cannot fill the estimated row group or improve the result. return rows - # Second pass: bytes/row changes with row-group size, so remeasure near the requested size. - slice_rows = int(_STAGE2_SLICE * rows) - copy_to_parquet(f"{select_query} LIMIT {slice_rows}", probe_path, rows, conn) + # Second pass: bytes/row changes with row-group size, so remeasure near the requested size. + copy_to_parquet(f"{select_query} LIMIT {rows}", probe_path, rows, conn) refined = _rows_for_target(_bytes_per_row(probe_path), target_bytes) return rows if refined is None else refined def _bytes_per_row(path): - """Read bytes/row from full row groups, excluding a trailing partial group.""" - metadata = pq.ParquetFile(path).metadata - num_row_groups = metadata.num_row_groups - full = range(num_row_groups - 1) if num_row_groups > 1 else range(num_row_groups) - total_rows = sum(metadata.row_group(i).num_rows for i in full) - if total_rows == 0: - return None - return sum(metadata.row_group(i).total_byte_size for i in full) / total_rows + """Read bytes/row from the probe's only row group.""" + row_group = pq.ParquetFile(path).metadata.row_group(0) + return row_group.total_byte_size / row_group.num_rows def _rows_for_target(bytes_per_row, target_bytes): """Convert bytes to rows and round to DuckDB's 2,048-row granularity.""" if not bytes_per_row: return None - rows = int(round(target_bytes / bytes_per_row / _ROW_GROUP_GRANULARITY)) + rows = round(target_bytes / bytes_per_row / _ROW_GROUP_GRANULARITY) return max(rows, 1) * _ROW_GROUP_GRANULARITY From 56a7656d607c415c524d9d56c1e973bec4983fc1 Mon Sep 17 00:00:00 2001 From: Qi Chen Date: Wed, 12 Aug 2026 17:01:22 +0200 Subject: [PATCH 5/5] Update row group sizing function documentation for clarity - Revised the docstring of `_rows_for_target` to specify that it rounds to the nearest 2,048-row multiple instead of allowing DuckDB to round up, enhancing clarity for future developers. --- benchmark_data_tools/row_group_sizing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmark_data_tools/row_group_sizing.py b/benchmark_data_tools/row_group_sizing.py index 4693b626..69c48ba7 100644 --- a/benchmark_data_tools/row_group_sizing.py +++ b/benchmark_data_tools/row_group_sizing.py @@ -90,7 +90,7 @@ def _bytes_per_row(path): def _rows_for_target(bytes_per_row, target_bytes): - """Convert bytes to rows and round to DuckDB's 2,048-row granularity.""" + """Round to the nearest 2,048-row multiple instead of letting DuckDB round up.""" if not bytes_per_row: return None rows = round(target_bytes / bytes_per_row / _ROW_GROUP_GRANULARITY)