Phase 1: Align the TPC-DS generation path with the TPC-H path

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

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"

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

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

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.

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