feat: tpcgen-cli early access

.github/copilot-instructions.md

@@ -24,8 +24,9 @@ src/lakebench/
 │   └── clickbench/              # ClickBench: 43 queries on clickstream data
 │
 ├── datagen/
-│   ├── tpch.py                  # TPCHDataGenerator (uses tpchgen-rs, ~10x faster than alternatives)
-│   ├── tpcds.py                 # TPCDSDataGenerator (wraps DuckDB TPC-DS extension)
+│   ├── _tpcgen_rs.py            # Shared bundled tpcgen-cli generation, sizing, and output normalization
+│   ├── tpch.py                  # TPCHDataGenerator (bundled tpcgen-cli)
+│   ├── tpcds.py                 # TPCDSDataGenerator (bundled tpcgen-cli; DuckDB fallback)
 │   └── clickbench.py            # Downloads dataset from ClickHouse host
 │
 ├── engines/
@@ -135,8 +136,9 @@ Install only what you need:
 | `duckdb` | `duckdb`, `deltalake`, `pyarrow` |
 | `polars` | `polars`, `deltalake`, `pyarrow` |
 | `daft` | `daft`, `deltalake`, `pyarrow` |
-| `tpcds_datagen` | `duckdb`, `pyarrow` |
-| `tpch_datagen` | `tpchgen-cli` |
+| `tpcds_datagen` | Compatibility extra; generator is bundled in supported platform wheels |
+| `tpcds_duckdb_datagen` | Legacy DuckDB TPC-DS generator |
+| `tpch_datagen` | Compatibility extra; generator is bundled in supported platform wheels |
 | `sparkmeasure` | `sparkmeasure` |
 | `sail` | `pysail`, `pyspark[connect]`, `deltalake`, `pyarrow` |
 

.github/workflows/native-wheels.yml

@@ -0,0 +1,95 @@
+name: Native LakeBench wheels
+
+on:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+env:
+  TPCGEN_COMMIT: eed8a40b6d69b8ed6b5cce88eb3bb6fdcce76c60
+
+jobs:
+  build:
+    name: ${{ matrix.platform }}
+    runs-on: ${{ matrix.runner }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - platform: linux
+            runner: ubuntu-22.04
+            target: x86_64
+            wheel_tag: py3-none-manylinux_2_17_x86_64
+          - platform: windows
+            runner: windows-latest
+            target: x64
+            wheel_tag: py3-none-win_amd64
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Check out pinned tpcgen-rs
+        uses: actions/checkout@v4
+        with:
+          repository: mwc360/tpcgen-rs
+          ref: ${{ env.TPCGEN_COMMIT }}
+          path: .build/tpcgen-rs
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          python-version: "3.11"
+          enable-cache: true
+
+      - name: Build pinned tpcgen-cli wheel
+        uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b
+        with:
+          target: ${{ matrix.target }}
+          args: --locked --release --out dist
+          working-directory: .build/tpcgen-rs/tpcgen-cli
+          manylinux: ${{ matrix.platform == 'linux' && '2_17' || '' }}
+
+      - name: Stage and verify tpcgen-cli
+        id: stage
+        shell: pwsh
+        run: |
+          $lines = python scripts/stage_tpcgen_binary.py `
+            --wheel-directory .build/tpcgen-rs/tpcgen-cli/dist `
+            --upstream-root .build/tpcgen-rs `
+            --output-directory .build/tpcgen-bundle `
+            --platform ${{ matrix.platform }}
+          $lines | ForEach-Object {
+            Write-Host $_
+            Add-Content -Path $env:GITHUB_OUTPUT -Value $_
+          }
+
+      - name: Build self-contained LakeBench wheel
+        shell: pwsh
+        env:
+          LAKEBENCH_TPCGEN_BINARY: ${{ steps.stage.outputs.binary }}
+          LAKEBENCH_TPCGEN_PROVENANCE: ${{ steps.stage.outputs.provenance }}
+          LAKEBENCH_TPCGEN_LICENSE: ${{ steps.stage.outputs.license }}
+          LAKEBENCH_WHEEL_TAG: ${{ matrix.wheel_tag }}
+        run: uv build --wheel
+
+      - name: Install and smoke test LakeBench wheel
+        shell: pwsh
+        run: |
+          $wheel = (Get-ChildItem dist\lakebench-*.whl | Select-Object -First 1).FullName
+          uv venv .build/wheel-test
+          if ("${{ matrix.platform }}" -eq "windows") {
+            $python = ".build\wheel-test\Scripts\python.exe"
+          } else {
+            $python = ".build/wheel-test/bin/python"
+          }
+          uv pip install --python $python $wheel
+          & $python -c "from lakebench.datagen._tpcgen_cli import TpcgenCli; print(TpcgenCli().run(['tpcds', '--help']).returncode)"
+          & $python scripts/smoke_test_tpcgen_bundle.py
+
+      - name: Upload wheel
+        uses: actions/upload-artifact@v4
+        with:
+          name: lakebench-${{ matrix.platform }}-x86_64
+          path: dist/lakebench-*.whl
+          if-no-files-found: error

.github/workflows/tests.yml

@@ -53,24 +53,30 @@ jobs:
   integration-tests:
     name: integration (${{ matrix.engine }})
     runs-on: ubuntu-latest
+    continue-on-error: ${{ matrix.allow_failure }}
     strategy:
       fail-fast: false
       matrix:
         include:
           - engine: duckdb
+            allow_failure: false
             extras_flags: "--extra duckdb --extra tpcds_datagen --extra tpch_datagen"
             test_file: "tests/integration/test_duckdb.py"
           - engine: daft
+            allow_failure: true
             extras_flags: "--extra daft --extra tpcds_datagen --extra tpch_datagen"
             test_file: "tests/integration/test_daft.py"
           - engine: polars
+            allow_failure: false
             extras_flags: "--extra polars --extra tpcds_datagen --extra tpch_datagen"
             test_file: "tests/integration/test_polars.py"
           - engine: spark
+            allow_failure: false
             extras_flags: "--extra spark --extra tpcds_datagen --extra tpch_datagen"
             test_file: "tests/integration/test_spark.py"
             java: "17"
           - engine: sail
+            allow_failure: false
             extras_flags: "--extra sail --extra tpcds_datagen --extra tpch_datagen"
             test_file: "tests/integration/test_sail.py"
 

.gitignore

@@ -86,3 +86,4 @@ scratch/
 # uv lockfile - not committed; resolution is intentionally left floating so
 # each environment picks the newest packages available in its configured index.
 uv.lock
+.cibuildwheel-build-requirements/

README.md

@@ -167,22 +167,26 @@ pip install lakebench[duckdb,polars,tpcds_datagen,tpch_datagen,sparkmeasure]
 ```
 
 > _Note: the `daft` extra pins `deltalake` to 1.5.x (Daft cannot read the Arrow `Utf8View` parquet that `deltalake` 1.6.x emits from `MERGE`), so it must be installed in its own environment rather than alongside `duckdb`, `polars`, or `sail`._
+>
+> `tpch_datagen` and `tpcds_datagen` use the same self-contained Rust
+> `tpcgen-cli` binary bundled in the Windows x86_64 and Linux x86_64 LakeBench
+> wheels. The legacy DuckDB TPC-DS generator remains available separately
+> through `tpcds_duckdb_datagen`.
 
 ## Example Usage
 To run any LakeBench benchmark, first do a one time generation of the data required for the benchmark and scale of interest. LakeBench provides datagen classes to quickly generate parquet datasets required by the benchmarks.
 
 ### Data Generation
-- **TPC-H** data generation is provided via the (tpchgen-rs)[https://github.com/clflushopt/tpchgen-rs] project. The project is currently about 10x+ faster than the next closest method of generating TPC-H datasets. _The TPC-DS version of project is currently under development._
+- **TPC-H** and **TPC-DS** data generation is blazing fast via a pinned build of the unified Rust `tpcgen-cli` from the [tpcgen-rs](https://github.com/datafusion-contrib/tpcgen-rs) project. The temporary Windows x86_64 and manylinux 2.17 x86_64 executables are committed under `native/tpcgen`. LakeBench will migrate to the official `tpcgen-cli` Python package after it is released on PyPI.
 
     _The below are generation runtimes on a 64 v-core VM writing to OneLake. Scale factors below 1000 can easily be generated on a 2 v-core machine._
-    | Scale Factor | Duration (hh:mm:ss)|
-    |:------------:|:------------------:|
-    | 1            | 00:00:04           |
-    | 10           | 00:00:09           |
-    | 100          | 00:01:09           |
-    | 1000         | 00:10:15           |
-
-- **TPC-DS** data generation is provided via the DuckDB [TPC-DS](https://duckdb.org/docs/stable/core_extensions/tpcds) extension. The LakeBench wrapper around DuckDB adds support for writing out parquet files with a provided row-group target file size as normally the files generated by DuckDB are atypically small (i.e. 10MB) and are most suitable for ultra-small scale scenarios. LakeBench defaults to target 128MB row groups but can be configured via the `target_row_group_size_mb` parameter of both TPC-H and TPC-DS DataGenerator classes.
+    | Scale Factor | TPC-H Duration (hh:mm:ss)| TPC-DS Duration (hh:mm:ss)|
+    |:------------:|:------------------:|:------------------:|
+    | 1            | 00:00:20           | 00:00:24           |
+    | 10           | 00:00:34           | 00:01:01           |
+    | 100          | 00:01:26           | 00:02:17           |
+    | 1000         | 00:07:49           | 00:10:30           |
+
 - **ClickBench** data is downloaded directly from the Clickhouse host site.
 
 #### TPC-H Data Generation
@@ -208,26 +212,47 @@ datagen.run()
 ```
 
 _Notes:_
-- TPC-DS data up to SF1000 can be generated on a 32-vCore machine. 
-- TPC-H datasets are generated extremely fast (i.e. SF1000 in 10 minutes on an 64-vCore machine).
+- By default, each table is split automatically using its estimated total
+  compressed size: 128 MiB files below 10 GiB, 256 MiB below 1 TiB, 512 MiB
+  below 5 TiB, and 1 GiB for larger tables. Estimated physical size is
+  calculated directly from each table's SF1000 baseline for any supported
+  scale factor; part counts are always selected automatically.
+- `target_row_group_size_mb` is an on-disk compressed-size target. LakeBench
+  converts it to the uncompressed-byte value expected by `tpcgen-cli` using
+  benchmark- and table-specific ZSTD(1) or Snappy compression ratios measured
+  from SF10 output. All `ZSTD(N)` levels use the ZSTD(1) measurements for
+  planning while the requested compression level is passed through unchanged.
+  Automatic part counts are adjusted for the selected codec. The row-group
+  conversion includes a 5% planning margin for upstream's estimated
+  bytes-per-source-row model. Other compressed codecs require an explicit
+  `compression_factor`, which is used for both row groups and file estimates.
+  TPC-DS generation uses the upstream C-reference compatibility mode.
+- Output remains organized as `<root>/<table>/*.parquet`. Filenames include
+  the one-based part number and codec for quick inspection, for example
+  `lineitem/lineitem-00001.zstd.parquet` or
+  `store_sales/store_sales-00001.zstd.parquet`.
+- To use the legacy implementation, install
+  `lakebench[tpcds_duckdb_datagen]` on Python 3.10+ and pass
+  `backend="duckdb"`.
+- Editable/source installations use the matching vendored binary from
+  `native/tpcgen`; installed wheels always use their packaged binary.
+- Large generations targeting mounted filesystems can set `num_threads=8` or
+  `num_threads=16` to limit concurrent file creation and atomic renames. The
+  default remains all available CPU cores.
 - The ClickBench dataset (only 1 size) should download with partitioned files in ~ 1 minute and ~ 6 minutes as a single file. 
 
 #### Is BYO Data Supported?
-If you want to use you own TPC-DS, TPC-H, or ClickBench parquet datasets, that is fine and encouraged as long as they are to specification. The Databricks [spark-sql-perf](https://github.com/databricks/spark-sql-perf) repo which is commonly used to produce TPC-DS and TPC-H datasets for benchmarking Spark has two critical schema bugs (typos?) in their implementation. Rather than supporting the perpetuation of these typos, LakeBench sticks to the schema defined in the specs. An [issue](https://github.com/databricks/spark-sql-perf/issues/219) was raised for tracking if this gets fixed. These datasets need to be fixed before running LakeBench with any data generated from spark-sql-perf:
-1. The `c_last_review_date_sk` column in the TPC-DS `customer` table was named `c_last_review_date` (the **_sk** is missing) and it is generated as a string whereas the TPC-DS spec says this column is a Identity type which would map to a integer. The data value is still a surrogate key but the schema doesn't exactly match the specification.
-    _Fix via:_
-    ```python
-    df = spark.read.parquet(f".../customer/")
-    df = df.withColumn('c_last_review_date_sk', sf.col('c_last_review_date').cast('int')).drop('c_last_review_date')
-    df.write.mode('overwrite').parquet(f".../customer/")
-    ```
-1. The `s_tax_percentage` column in the TPC-DS `store` table was named with a typo: `s_tax_precentage` (is "**pre**centage" the precursor of a "**per**centage"??).
-    _Fix via:_
-    ```python
-    df = spark.read.parquet(f"..../store/")
-    df = df.withColumnRenamed('s_tax_precentage', 's_tax_percentage')
-    df.write.mode('overwrite').parquet(f"..../store/")
-    ```
+If you want to use your own TPC-DS, TPC-H, or ClickBench Parquet datasets, that is fine and encouraged as long as they are to specification. LakeBench keeps the canonical TPC-DS schema as its table and query contract, but automatically corrects these recognized legacy input names while loading Parquet:
+
+| Benchmark | Table | Legacy input name | Canonical LakeBench name |
+|---|---|---|---|
+| TPC-DS | `catalog_returns` | `cr_return_amount_inc_tax` | `cr_return_amt_inc_tax` |
+| TPC-DS | `income_band` | `ib_income_band_id` | `ib_income_band_sk` |
+| TPC-DS | `reason` | `r_reason_description` | `r_reason_desc` |
+| TPC-DS | `store` | `s_tax_precentage` | `s_tax_percentage` |
+| TPC-DS | `web_returns` | `wr_store_credit` | `wr_account_credit` |
+
+Canonical names are accepted unchanged. Input containing both names, or neither required name, is rejected as ambiguous or invalid. Loaded Delta tables always use the canonical name.
 
 ### Load-Time Statistics
 

examples/data_generation.ipynb

@@ -52,7 +52,6 @@
     "    scale_factor=0.1,\n",
     "    target_folder_uri=r\"C:\\lakebench_local_tests\\data\\source\\tpch\\sf1_parquet\",\n",
     "    compression=\"ZSTD(1)\",\n",
-    "    multithreading=True,\n",
     ")\n",
     "datagen.run()"
    ]

hatch_build.py

@@ -0,0 +1,57 @@
+import os
+import platform
+import stat
+from pathlib import Path
+
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
+
+
+class CustomBuildHook(BuildHookInterface):
+    VENDORED_BINARIES = {
+        ("windows", "amd64"): ("windows-x86_64/tpcgen-cli.exe", "py3-none-win_amd64"),
+        ("windows", "x86_64"): ("windows-x86_64/tpcgen-cli.exe", "py3-none-win_amd64"),
+        ("linux", "amd64"): ("linux-x86_64/tpcgen-cli", "py3-none-manylinux_2_17_x86_64"),
+        ("linux", "x86_64"): ("linux-x86_64/tpcgen-cli", "py3-none-manylinux_2_17_x86_64"),
+    }
+
+    def initialize(self, version, build_data):
+        if self.target_name != "wheel":
+            return
+
+        binary = os.environ.get("LAKEBENCH_TPCGEN_BINARY")
+        if not binary:
+            platform_key = (platform.system().lower(), platform.machine().lower())
+            vendored = self.VENDORED_BINARIES.get(platform_key)
+            if vendored is None:
+                return
+
+            relative_binary, default_wheel_tag = vendored
+            bundle_root = Path(self.root) / "native" / "tpcgen"
+            binary = bundle_root / relative_binary
+            provenance = binary.parent / "provenance.json"
+            license_file = bundle_root / "LICENSE.tpcgen-rs"
+            wheel_tag = default_wheel_tag
+        else:
+            wheel_tag = os.environ.get("LAKEBENCH_WHEEL_TAG")
+            provenance = os.environ.get("LAKEBENCH_TPCGEN_PROVENANCE")
+            license_file = os.environ.get("LAKEBENCH_TPCGEN_LICENSE")
+            if not wheel_tag or not provenance or not license_file:
+                raise RuntimeError(
+                    "Native wheel builds require LAKEBENCH_WHEEL_TAG, "
+                    "LAKEBENCH_TPCGEN_PROVENANCE, and LAKEBENCH_TPCGEN_LICENSE."
+                )
+
+        binary_path = Path(binary).resolve()
+        provenance_path = Path(provenance).resolve()
+        license_path = Path(license_file).resolve()
+        for required_path in (binary_path, provenance_path, license_path):
+            if not required_path.is_file():
+                raise RuntimeError(f"Native wheel input is missing: {required_path}")
+
+        binary_path.chmod(binary_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+
+        build_data["pure_python"] = False
+        build_data["tag"] = wheel_tag
+        build_data["force_include"][str(binary_path)] = f"lakebench/datagen/_bin/{binary_path.name}"
+        build_data["force_include"][str(provenance_path)] = "lakebench/datagen/_bin/provenance.json"
+        build_data["force_include"][str(license_path)] = "lakebench/datagen/_bin/LICENSE.tpcgen-rs"