diff --git a/benchmark_data_tools/duckdb_utils.py b/benchmark_data_tools/duckdb_utils.py index 0e16fc2a..a53a49ab 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,32 @@ 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 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() + 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_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}" + 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..2e3dc31f 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 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")) _HIGH_CARD_NDV_THRESHOLD = 0.99 @@ -99,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" @@ -140,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. @@ -190,42 +191,72 @@ def write_metadata(args): def generate_data_files_with_duckdb(args): - init_benchmark_tables(args.benchmark_type, args.scale_factor) - - 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") - - tables = duckdb.sql("SHOW TABLES").fetchall() - for (table_name,) in tables: + # 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, + 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() + + 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, + args.num_threads, ) -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, + 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) + 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 = 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" + 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): @@ -309,7 +340,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": []} @@ -391,7 +422,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 b011a066..c73e9f11 100644 --- a/benchmark_data_tools/requirements.txt +++ b/benchmark_data_tools/requirements.txt @@ -1,7 +1,7 @@ certifi==2025.8.3 charset-normalizer==3.4.3 click==8.2.1 -duckdb==1.3.2 +duckdb==1.5.5 idna==3.10 presto-python-client==0.8.4 requests==2.32.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..69c48ba7 --- /dev/null +++ b/benchmark_data_tools/row_group_sizing.py @@ -0,0 +1,97 @@ +# 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 copy_to_parquet, get_select_query, init_benchmark_tables + +# DuckDB Parquet sizing: https://duckdb.org/docs/current/data/parquet/tips +_ROW_GROUP_GRANULARITY = 2048 +_STAGE1_ROWS = 122_880 + +_MAX_PROBE_SCALE_FACTOR = 10 +_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. + 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: 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 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): + """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) + return max(rows, 1) * _ROW_GROUP_GRANULARITY 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 new file mode 100644 index 00000000..f5f2bcf8 --- /dev/null +++ b/benchmark_data_tools/tests/multi_file_partitioning_test.py @@ -0,0 +1,65 @@ +# 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 duckdb_utils import init_benchmark_tables +from generate_data_files import generate_data_files + + +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 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 + # 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, + expected_row_counts, + ) + + +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 + 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) + ] + 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 + + assert split_table_count > 0, "no table was split, so the partitioning path was not tested"