diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 9dd90eb..1fa5600 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -24,9 +24,9 @@ src/lakebench/ │ └── clickbench/ # ClickBench: 43 queries on clickstream data │ ├── datagen/ -│ ├── _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) +│ ├── _tpcgen.py # Thin tpcgen-cli wrapper: builds the command, picks part counts, runs it +│ ├── tpch.py # TPCHDataGenerator (tpcgen-cli) +│ ├── tpcds.py # TPCDSDataGenerator (tpcgen-cli; DuckDB fallback) │ └── clickbench.py # Downloads dataset from ClickHouse host │ ├── engines/ @@ -143,9 +143,9 @@ Install only what you need: | `polars` | `polars`, `deltalake`, `pyarrow` | | `fabric_data_warehouse` | `sqlalchemy`, `pyodbc`, `pandas`, `requests` (needs ODBC Driver 18 on the host) | | `daft` | `daft`, `deltalake`, `pyarrow` | -| `tpcds_datagen` | Compatibility extra; generator is bundled in supported platform wheels | +| `tpcds_datagen` | `tpcgen-cli==4.0.0` | | `tpcds_duckdb_datagen` | Legacy DuckDB TPC-DS generator | -| `tpch_datagen` | Compatibility extra; generator is bundled in supported platform wheels | +| `tpch_datagen` | `tpcgen-cli==4.0.0` | | `sparkmeasure` | `sparkmeasure` | | `sail` | `pysail`, `pyspark[connect]`, `deltalake`, `pyarrow` | @@ -195,10 +195,10 @@ Set `sql_text` *before* executing, so the statement is still recorded when the e ## Native TPC Generator Format -TPC-H and TPC-DS support the generators' pipe-delimited text output alongside Parquet: `TPCxDataGenerator(output_format="native")` and `TPCx(input_format="native")`. TPC-DS native format support is currently in preview. TPC-DS uses subcommand/extension `dat`, TPC-H uses `tbl`; the currently bundled `tpcgen-cli tpcds dat` rejects `--num-threads` and runs single-threaded regardless of available cores. Part count controls output splitting, not parallel execution. Neither native subcommand accepts `--compression` or `--row-group-bytes`. +TPC-H and TPC-DS support the generators' pipe-delimited text output alongside Parquet: `TPCxDataGenerator(output_format="native")` and `TPCx(input_format="native")`. TPC-DS native format support is currently in preview. TPC-DS uses subcommand/extension `dat`, TPC-H uses `tbl`; both accept `--num-threads`. Neither native subcommand accepts `--compression` or `--row-group-bytes`. -The format has no header and no types, so reader schemas come from the benchmark's resolved DDL via `utils/schema_utils.py`. Native files are named the way the official tools name them: `.tbl`/`
.dat` for a single part, and `
.tbl.` (dbgen) or `
__.dat` (dsdgen) for multiple parts. Because dbgen's parallel names do not end in the extension, engines glob a benchmark-supplied `NATIVE_FILE_GLOB` pattern (`*.tbl*`, `*.dat`) rather than an extension. +The format has no header and no types, so reader schemas come from the benchmark's resolved DDL via `utils/schema_utils.py`. Generated native files keep tpcgen-cli's names, `
/
..tbl`/`.dat`. Engines glob a benchmark-supplied `NATIVE_FILE_GLOB` (`*.tbl*`, `*.dat`) so official dbgen parallel names (`
.tbl.`) also load. Every generated line ends with a trailing delimiter, so readers declare one extra trailing column (`TRAILING_DELIMITER_COLUMN`) and drop it; quoting must be disabled and empty fields must read as NULL. Engines implement `load_delimited_to_delta`; `Sail` subclasses `BaseEngine` rather than `Spark`, so it needs its own copy. Daft's CSV reader mislabels decimal precision, so decimals are read as text and cast. -Native part counts use `NATIVE_SIZE_FACTOR_DICT` (native bytes ÷ uncompressed Parquet bytes, measured per table at SF1). Do not derive these from `SF1000_SIZE_GB_DICT` at a different scale factor: that dict assumes linear scaling, which is false for TPC-DS's fixed-size dimensions. +Generators pass options straight to tpcgen-cli. LakeBench only chooses `--parts` per table: ceil(`SF1000_SIZE_GB` × SF / 1000 ÷ 256 MiB). tpcgen-cli writes no file for empty parts (TPC-DS tables under 1M source rows, TPC-H nation/region), so a table may have fewer files than parts. diff --git a/.github/workflows/native-wheels.yml b/.github/workflows/native-wheels.yml deleted file mode 100644 index 1c91f7e..0000000 --- a/.github/workflows/native-wheels.yml +++ /dev/null @@ -1,95 +0,0 @@ -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 diff --git a/README.md b/README.md index 163a39b..bdc459b 100644 --- a/README.md +++ b/README.md @@ -172,7 +172,7 @@ Select the optional dependency group and any non-Python prerequisites from the [ 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** 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. +- **TPC-H** and **TPC-DS** data generation is blazing fast via the unified Rust [`tpcgen-cli`](https://pypi.org/project/tpcgen-cli/) from the [tpcgen-rs](https://github.com/datafusion-contrib/tpcgen-rs) project, installed by the `tpch_datagen` / `tpcds_datagen` extras. _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 | TPC-H Duration (hh:mm:ss)| TPC-DS Duration (hh:mm:ss)| @@ -207,11 +207,10 @@ datagen.run() ``` _Notes:_ -- 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 `/
/*.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 -- `output_format="native"` emits the TPC generators' pipe-delimited text instead of Parquet, matching what the official `dsdgen`/`dbgen` tools produce. TPC-DS native format support is currently in preview. See [Native TPC Generator Format](#native-tpc-generator-format-dat--tbl) below.- Editable/source installations use the matching vendored binary from `native/tpcgen`; installed wheels always use their packaged binary. +- Output is written by `tpcgen-cli` as `/
/
..parquet`, for example `lineitem/lineitem.1.parquet`. Each table is split into roughly 256 MiB ZSTD(1) files based on its measured size; other codecs and the native format produce proportionally larger files. Small TPC-DS tables, and TPC-H `nation` and `region`, are always a single file. +- `compression` (default `ZSTD(1)`) and `target_row_group_size_mb` (default 128, approximate uncompressed size) are passed to `tpcgen-cli` as `--compression` and `--row-group-bytes`. +- To use the legacy implementation, install `lakebench[tpcds_duckdb_datagen]` on Python 3.10+ and pass `backend="duckdb"`. +- `output_format="native"` emits the TPC generators' pipe-delimited text instead of Parquet, matching what the official `dsdgen`/`dbgen` tools produce. TPC-DS native format support is currently in preview. See [Native TPC Generator Format](#native-tpc-generator-format-dat--tbl) below. - 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. @@ -240,14 +239,13 @@ benchmark.run(mode='load_and_query') ``` _Notes:_ -- Files are named exactly as the official tools name them, one folder per table. A single-part table uses the serial name `
.tbl` / `
.dat`. A multi-part table uses the parallel names: `dbgen`'s `
.tbl.` for TPC-H and `dsdgen`'s `
__.dat` for TPC-DS. -- Neither `dsdgen` nor `dbgen` splits output on its own; parts exist only because the operator runs the tool once per chunk, which is the normal way to generate large scale factors. LakeBench picks the part count for you from the estimated table size, so bigger scale factors naturally produce more parts, matching what a parallel `dsdgen`/`dbgen` run would leave on disk. +- Files are written one folder per table as `
/
..tbl` / `.dat`. Loading also accepts the official `dbgen`/`dsdgen` parallel names (`
.tbl.`, `
__.dat`). - The format carries no header and no types, so LakeBench derives each reader's schema from the benchmark's resolved DDL. As a result, native loads are always typed exactly as the DDL declares. The generator's Parquet output does not always agree with the DDL on integer width (for example TPC-H `n_nationkey` is `int64` in Parquet but `integer` in the DDL), so on engines that do not pre-create tables the two formats can differ in integer width. Values are identical. - Every generated line ends with a trailing delimiter. LakeBench reads one extra trailing column and drops it, so no reader needs a lenient mode. - Empty fields are read as `NULL`, and quoting is disabled so `"` is treated as ordinary data. -- Parquet-only options (`target_row_group_size_mb`, `compression`, and `compression_factor`) are rejected with `output_format="native"`. Automatic part counts use per-table native-to-uncompressed-Parquet size ratios measured at SF1. +- Parquet-only options (`target_row_group_size_mb` and `compression`) are rejected with `output_format="native"`. - Native-input support varies by engine; see the selected [engine page](docs/engines/README.md). -- `output_format="native"` requires `backend="rust"`. With the currently bundled generator, `tpcgen-cli tpcds dat` has no thread option, so TPC-DS native generation runs single-threaded regardless of the available core count and rejects `num_threads`. Part count controls output splitting but does not make the current TPC-DS native run parallel. `tpcgen-cli tpch tbl` does accept `--num-threads`, so TPC-H native honours it exactly like Parquet (defaulting to all available CPU cores). +- `output_format="native"` requires `backend="rust"`. `num_threads` applies to native output exactly as it does to Parquet (defaulting to all available CPU cores). #### Naming the Input Location @@ -371,7 +369,7 @@ LakeBench is released under the MIT License (see [`LICENSE`](LICENSE)). It also Most notably: - **ClickBench** queries and schema come from [ClickHouse/ClickBench](https://github.com/ClickHouse/ClickBench) (Alexey Milovidov and the ClickHouse team, 2022), which is published under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/). The queries are vendored verbatim from a pinned upstream commit; provenance and hashes are recorded in [`PROVENANCE.md`](src/lakebench/benchmarks/clickbench/resources/queries/canonical/PROVENANCE.md) and `source_manifest.json` alongside them. If the NonCommercial or ShareAlike terms matter for your use, review them before redistributing LakeBench or building on it. -- **tpcgen-rs** is bundled as a prebuilt binary in platform wheels under Apache 2.0. +- **tpcgen-rs** is installed from PyPI as `tpcgen-cli` (Apache 2.0) by the datagen extras; it is not redistributed. - **TPC-H / TPC-DS** tools kits, templates, and generator executables are *not* redistributed; only generated query text and a reproducibility manifest are checked in. TPC-H and TPC-DS are trademarks of the [Transaction Processing Performance Council](https://www.tpc.org), and LakeBench results are not audited, endorsed, or comparable to published TPC results. diff --git a/THIRD-PARTY-NOTICES.md b/THIRD-PARTY-NOTICES.md index 982bf62..5ecbc2f 100644 --- a/THIRD-PARTY-NOTICES.md +++ b/THIRD-PARTY-NOTICES.md @@ -67,12 +67,12 @@ their own legal counsel or contact the upstream project. - **Project:** tpcgen-rs — TPC data generation CLI - **Upstream:** https://github.com/datafusion-contrib/tpcgen-rs -- **Build source:** https://github.com/mwc360/tpcgen-rs -- **License:** Apache License 2.0 — see `native/tpcgen/LICENSE.tpcgen-rs` +- **Package:** [`tpcgen-cli`](https://pypi.org/project/tpcgen-cli/) on PyPI +- **License:** Apache License 2.0 -Platform wheels bundle a prebuilt `tpcgen-cli` binary under `native/tpcgen/`. -Per-platform build provenance, including upstream commit and binary checksum, is -recorded in `native/tpcgen//provenance.json`. +LakeBench does not redistribute `tpcgen-cli`. The optional `tpch_datagen` and +`tpcds_datagen` extras declare a dependency on the published `tpcgen-cli` +package, which is installed from PyPI under its own license. --- diff --git a/hatch_build.py b/hatch_build.py deleted file mode 100644 index 76d3d95..0000000 --- a/hatch_build.py +++ /dev/null @@ -1,57 +0,0 @@ -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" diff --git a/native/tpcgen/LICENSE.tpcgen-rs b/native/tpcgen/LICENSE.tpcgen-rs deleted file mode 100644 index f49a4e1..0000000 --- a/native/tpcgen/LICENSE.tpcgen-rs +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. \ No newline at end of file diff --git a/native/tpcgen/linux-x86_64/provenance.json b/native/tpcgen/linux-x86_64/provenance.json deleted file mode 100644 index 94f19f3..0000000 --- a/native/tpcgen/linux-x86_64/provenance.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "repository": "https://github.com/mwc360/tpcgen-rs", - "commit": "eed8a40b6d69b8ed6b5cce88eb3bb6fdcce76c60", - "cargo_lock_commit": "40cedfdc6ec7f47080660808897752436d662d0c", - "source_branch": "users/mcole/align-tpcds-schema-and-parts", - "pull_request": "https://github.com/datafusion-contrib/tpcgen-rs/pull/406", - "cli_version": "tpcgen-cli 0.1.0", - "platform": "linux", - "binary_name": "tpcgen-cli", - "binary_sha256": "9f4f80b691e59f1b4f0533dcd86b9cacd5f69b8ec6c1022c07415168f9e8348e", - "upstream_wheel": "tpcgen_cli-0.1.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl" -} diff --git a/native/tpcgen/linux-x86_64/tpcgen-cli b/native/tpcgen/linux-x86_64/tpcgen-cli deleted file mode 100644 index 4d70330..0000000 Binary files a/native/tpcgen/linux-x86_64/tpcgen-cli and /dev/null differ diff --git a/native/tpcgen/windows-x86_64/provenance.json b/native/tpcgen/windows-x86_64/provenance.json deleted file mode 100644 index 8551b26..0000000 --- a/native/tpcgen/windows-x86_64/provenance.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "repository": "https://github.com/mwc360/tpcgen-rs", - "commit": "eed8a40b6d69b8ed6b5cce88eb3bb6fdcce76c60", - "cargo_lock_commit": "40cedfdc6ec7f47080660808897752436d662d0c", - "source_branch": "users/mcole/align-tpcds-schema-and-parts", - "pull_request": "https://github.com/datafusion-contrib/tpcgen-rs/pull/406", - "cli_version": "tpcgen-cli 0.1.0", - "platform": "windows", - "binary_name": "tpcgen-cli.exe", - "binary_sha256": "716a8bdc4c391c492e62c67d1c15808ea61bf64172331810771bd7209d6fe157", - "upstream_wheel": "tpcgen_cli-0.1.0-py3-none-win_amd64.whl" -} diff --git a/native/tpcgen/windows-x86_64/tpcgen-cli.exe b/native/tpcgen/windows-x86_64/tpcgen-cli.exe deleted file mode 100644 index 5e7e519..0000000 Binary files a/native/tpcgen/windows-x86_64/tpcgen-cli.exe and /dev/null differ diff --git a/pyproject.toml b/pyproject.toml index 3dba671..1493138 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,10 +43,11 @@ polars = ["polars==1.44.1; python_version >= '3.10'", "deltalake==1.6.2; python_ # Daft is held on deltalake 1.5.x: 1.6.x emits DataFusion `string_view` parquet # from MERGE, and Daft's reader rejects Arrow Utf8View. daft = ["daft[deltalake]==0.7.21; python_version >= '3.10'", "deltalake>=1.5.0,<1.6.0; python_version >= '3.10'", "pyarrow>=16.0.0; python_version >= '3.10'"] -# The vendored tpcgen-cli serves both TPC-H and TPC-DS in platform wheels. -tpcds_datagen = [] -tpch_datagen = [] -# Temporary fallback until the unified upstream tpcgen-cli package is released. +# The unified tpcgen-cli generates both TPC-H and TPC-DS. It is pinned exactly +# because the part-count and row-group sizing tables are calibrated to its output. +tpcds_datagen = ["tpcgen-cli==4.0.0"] +tpch_datagen = ["tpcgen-cli==4.0.0"] +# Legacy DuckDB TPC-DS generator, selected with backend="duckdb". tpcds_duckdb_datagen = ["duckdb==1.5.5; python_version >= '3.10'"] sparkmeasure = ["sparkmeasure==0.28.0; python_version >= '3.9'"] spark = ["pyspark>=3.5.0,<4.0.0; python_version >= '3.9'", "delta-spark>=3.2.0,<4.0.0; python_version >= '3.9'"] @@ -79,8 +80,6 @@ packages = ["src/lakebench"] [tool.hatch.build.targets.wheel.force-include] "THIRD-PARTY-NOTICES.md" = "lakebench/THIRD-PARTY-NOTICES.md" -[tool.hatch.build.hooks.custom] - [dependency-groups] dev = [ "pytest>=7.0.0", diff --git a/scripts/build_linux_wheel.py b/scripts/build_linux_wheel.py deleted file mode 100644 index 550a719..0000000 --- a/scripts/build_linux_wheel.py +++ /dev/null @@ -1,136 +0,0 @@ -import argparse -import hashlib -import os -import shutil -import subprocess -import sys -from pathlib import Path - -import tomllib - -BUILD_IDENTIFIER = "cp311-manylinux_x86_64" -CIBUILDWHEEL_REQUIREMENT = "cibuildwheel==4.2.0" - - -def _project_configuration(project_root: Path) -> dict: - with (project_root / "pyproject.toml").open("rb") as file: - return tomllib.load(file) - - -def _require_docker() -> None: - if shutil.which("docker") is None: - raise RuntimeError( - "Docker is required to build a Linux wheel. Install and start Docker Desktop, then rerun this command." - ) - - result = subprocess.run( - ["docker", "info"], - capture_output=True, - text=True, - check=False, - ) - if result.returncode != 0: - details = result.stderr.strip() or result.stdout.strip() - raise RuntimeError( - f"Docker is installed but unavailable. Start Docker Desktop and rerun this command.\n{details}" - ) - - -def _require_uv() -> str: - uv = shutil.which("uv") - if uv is None: - raise RuntimeError("uv is required to prepare the offline Linux build environment.") - return uv - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def main() -> int: - parser = argparse.ArgumentParser(description="Build the LakeBench Linux x86_64 wheel using Docker.") - parser.add_argument( - "--output-dir", - type=Path, - default=Path("dist"), - help="Wheel output directory relative to the repository root (default: dist).", - ) - args = parser.parse_args() - - project_root = Path(__file__).resolve().parents[1] - output_dir = args.output_dir if args.output_dir.is_absolute() else project_root / args.output_dir - output_dir.mkdir(parents=True, exist_ok=True) - - _require_docker() - uv = _require_uv() - - configuration = _project_configuration(project_root) - version = configuration["project"]["version"] - wheel_pattern = f"lakebench-{version}-py3-none-manylinux_2_17_x86_64.whl" - current_linux_wheels = f"lakebench-{version}-py3-none-*manylinux*x86_64.whl" - for existing_wheel in output_dir.glob(current_linux_wheels): - existing_wheel.unlink() - - wheelhouse = project_root / ".cibuildwheel-build-requirements" - shutil.rmtree(wheelhouse, ignore_errors=True) - wheelhouse.mkdir() - try: - subprocess.run( - [ - uv, - "pip", - "install", - "--target", - str(wheelhouse), - "--no-compile", - *configuration["build-system"]["requires"], - ], - cwd=project_root, - check=True, - ) - - environment = os.environ.copy() - environment["CIBW_ENVIRONMENT"] = "PYTHONPATH=/project/.cibuildwheel-build-requirements" - environment["CIBW_BUILD_FRONTEND"] = "build; args: --no-isolation" - environment["CIBW_REPAIR_WHEEL_COMMAND_LINUX"] = "" - subprocess.run( - [ - uv, - "tool", - "run", - "--from", - CIBUILDWHEEL_REQUIREMENT, - "cibuildwheel", - "--only", - BUILD_IDENTIFIER, - "--output-dir", - str(output_dir), - str(project_root), - ], - cwd=project_root, - env=environment, - check=True, - ) - finally: - shutil.rmtree(wheelhouse, ignore_errors=True) - - wheels = sorted(output_dir.glob(wheel_pattern)) - if len(wheels) != 1: - raise RuntimeError(f"Expected one {wheel_pattern} artifact, found {len(wheels)}.") - - wheel = wheels[0] - print(f"Built {wheel}") - print(f"SHA-256: {_sha256(wheel)}") - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main()) - except (OSError, RuntimeError, subprocess.CalledProcessError) as error: - print(f"error: {error}", file=sys.stderr) - raise SystemExit(1) from error diff --git a/scripts/smoke_test_tpcgen_bundle.py b/scripts/smoke_test_tpcgen_bundle.py deleted file mode 100644 index 6eef77a..0000000 --- a/scripts/smoke_test_tpcgen_bundle.py +++ /dev/null @@ -1,77 +0,0 @@ -import tempfile -from pathlib import Path - -import pyarrow as pa -import pyarrow.parquet as pq - -from lakebench.datagen import TPCDSDataGenerator, TPCHDataGenerator - - -def main() -> None: - with tempfile.TemporaryDirectory() as temp_dir: - root = Path(temp_dir) - inventory_output = root / "inventory" - TPCDSDataGenerator( - scale_factor=1, - target_folder_uri=str(inventory_output), - target_row_group_size_mb=16, - table_list=["inventory"], - num_threads=1, - ).run() - - inventory_file = inventory_output / "inventory" / "inventory-00001.zstd.parquet" - if not inventory_file.is_file(): - raise RuntimeError(f"Expected inventory output is missing: {inventory_file}") - total_rows = pq.ParquetFile(inventory_file).metadata.num_rows - if total_rows == 0: - raise RuntimeError("Smoke test generated no inventory rows.") - - schema_output = root / "schema" - TPCDSDataGenerator( - scale_factor=0.01, - target_folder_uri=str(schema_output), - target_row_group_size_mb=16, - table_list=["customer", "reason", "store_sales"], - num_threads=1, - ).run() - - customer_schema = pq.read_schema(schema_output / "customer" / "customer-00001.zstd.parquet") - if customer_schema.field("c_customer_sk").type != pa.int32(): - raise RuntimeError("customer.c_customer_sk does not have the expected integer type.") - if customer_schema.field("c_last_review_date_sk").type != pa.int32(): - raise RuntimeError("customer.c_last_review_date_sk does not have the expected integer type.") - - store_sales_schema = pq.read_schema(schema_output / "store_sales" / "store_sales-00001.zstd.parquet") - if store_sales_schema.field("ss_sold_date_sk").type != pa.int32(): - raise RuntimeError("store_sales.ss_sold_date_sk does not have the expected integer type.") - if store_sales_schema.field("ss_ticket_number").type != pa.int64(): - raise RuntimeError("store_sales.ss_ticket_number does not have the expected long integer type.") - sales_price_type = store_sales_schema.field("ss_sales_price").type - if sales_price_type != pa.decimal128(7, 2): - raise RuntimeError("store_sales.ss_sales_price does not have the expected decimal type.") - - reason_rows = pq.ParquetFile(schema_output / "reason" / "reason-00001.zstd.parquet").metadata.num_rows - if reason_rows != 75: - raise RuntimeError(f"C-reference compatibility produced {reason_rows} reason rows instead of 75.") - - tpch_output = root / "tpch" - TPCHDataGenerator( - scale_factor=0.01, - target_folder_uri=str(tpch_output), - target_row_group_size_mb=16, - table_list=["region"], - num_threads=1, - ).run() - region_file = tpch_output / "region" / "region-00001.zstd.parquet" - region_rows = pq.ParquetFile(region_file).metadata.num_rows - if region_rows != 5: - raise RuntimeError(f"TPC-H generation produced {region_rows} region rows instead of 5.") - - print( - f"Validated unified TPC-H/TPC-DS generation, {total_rows} inventory rows, " - "automatic part sizing, critical schemas, and C-reference compatibility." - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/stage_tpcgen_binary.py b/scripts/stage_tpcgen_binary.py deleted file mode 100644 index e271598..0000000 --- a/scripts/stage_tpcgen_binary.py +++ /dev/null @@ -1,96 +0,0 @@ -import argparse -import hashlib -import json -import stat -import subprocess -import zipfile -from pathlib import Path - -UPSTREAM_REPOSITORY = "https://github.com/mwc360/tpcgen-rs" -UPSTREAM_COMMIT = "eed8a40b6d69b8ed6b5cce88eb3bb6fdcce76c60" -UPSTREAM_LOCK_COMMIT = "40cedfdc6ec7f47080660808897752436d662d0c" - - -def find_upstream_wheel(wheel_directory: Path) -> Path: - wheels = list(wheel_directory.glob("tpcgen_cli-*.whl")) - if len(wheels) != 1: - raise RuntimeError(f"Expected one tpcgen-cli wheel in {wheel_directory}, found {len(wheels)}.") - return wheels[0] - - -def extract_binary(wheel_path: Path, output_directory: Path, platform_name: str) -> Path: - binary_name = "tpcgen-cli.exe" if platform_name == "windows" else "tpcgen-cli" - with zipfile.ZipFile(wheel_path) as wheel: - candidates = [name for name in wheel.namelist() if name.endswith(f"/scripts/{binary_name}")] - if len(candidates) != 1: - raise RuntimeError(f"Expected one {binary_name} script in {wheel_path}, found {len(candidates)}.") - binary_path = output_directory / binary_name - binary_path.write_bytes(wheel.read(candidates[0])) - - binary_path.chmod(binary_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) - return binary_path - - -def run_smoke_test(binary_path: Path) -> str: - version = subprocess.run( - [str(binary_path), "--version"], - capture_output=True, - text=True, - check=True, - ).stdout.strip() - subprocess.run( - [str(binary_path), "tpcds", "--help"], - capture_output=True, - text=True, - check=True, - ) - return version - - -def sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as file_handle: - for chunk in iter(lambda: file_handle.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--wheel-directory", type=Path, required=True) - parser.add_argument("--upstream-root", type=Path, required=True) - parser.add_argument("--output-directory", type=Path, required=True) - parser.add_argument("--platform", choices=("linux", "windows"), required=True) - args = parser.parse_args() - - args.output_directory.mkdir(parents=True, exist_ok=True) - wheel_path = find_upstream_wheel(args.wheel_directory) - binary_path = extract_binary(wheel_path, args.output_directory, args.platform) - version = run_smoke_test(binary_path) - - license_source = args.upstream_root / "LICENSE" - license_target = args.output_directory / "LICENSE.tpcgen-rs" - license_target.write_bytes(license_source.read_bytes()) - - provenance = { - "repository": UPSTREAM_REPOSITORY, - "commit": UPSTREAM_COMMIT, - "cargo_lock_commit": UPSTREAM_LOCK_COMMIT, - "source_branch": "users/mcole/align-tpcds-schema-and-parts", - "pull_request": "https://github.com/datafusion-contrib/tpcgen-rs/pull/406", - "cli_version": version, - "platform": args.platform, - "binary_name": binary_path.name, - "binary_sha256": sha256(binary_path), - "upstream_wheel": wheel_path.name, - } - provenance_path = args.output_directory / "provenance.json" - provenance_path.write_text(json.dumps(provenance, indent=2) + "\n", encoding="utf-8") - - print(f"binary={binary_path.resolve()}") - print(f"provenance={provenance_path.resolve()}") - print(f"license={license_target.resolve()}") - - -if __name__ == "__main__": - main() diff --git a/src/lakebench/datagen/_tpcds_rs.py b/src/lakebench/datagen/_tpcds_rs.py deleted file mode 100644 index c26b5ce..0000000 --- a/src/lakebench/datagen/_tpcds_rs.py +++ /dev/null @@ -1,230 +0,0 @@ -from ._tpcgen_rs import _TpcgenRsDataGenerator - - -class _TPCDSRsDataGenerator(_TpcgenRsDataGenerator): - BENCHMARK_NAME = "tpcds" - BENCHMARK_LABEL = "TPC-DS" - COMMAND_ARGUMENTS = ("--compat", "c") - NATIVE_SUBCOMMAND = "dat" - NATIVE_FILE_EXTENSION = "dat" - # `tpcgen-cli tpcds dat` rejects --num-threads even though `tpcds parquet` accepts it. - NATIVE_SUPPORTS_NUM_THREADS = False - GEN_TABLE_REGISTRY = [ - "call_center", - "catalog_page", - "catalog_returns", - "catalog_sales", - "customer", - "customer_address", - "customer_demographics", - "date_dim", - "household_demographics", - "income_band", - "inventory", - "item", - "promotion", - "reason", - "ship_mode", - "store", - "store_returns", - "store_sales", - "time_dim", - "warehouse", - "web_page", - "web_returns", - "web_sales", - "web_site", - ] - SIZE_REFERENCE_SCALE_FACTOR = 10 - DEFINED_SCALE_FACTORS = (0, 1, 10, 100, 300, 1000, 3000, 10000, 30000, 100000) - # Final output row counts from the TPC-DS scaling model. Fact and return - # tables scale linearly and therefore do not need entries here. - ROW_COUNT_SCALING = { - "call_center": ("logarithmic", (0, 6, 24, 30, 36, 42, 48, 54, 60, 60)), - "catalog_page": ("static", (0, 11718, 12000, 20400, 26000, 30000, 36000, 40000, 46000, 50000)), - "customer": ( - "logarithmic", - (0, 100000, 500000, 2000000, 5000000, 12000000, 30000000, 65000000, 80000000, 100000000), - ), - "customer_address": ( - "logarithmic", - (0, 50000, 250000, 1000000, 2500000, 6000000, 15000000, 32500000, 40000000, 50000000), - ), - "customer_demographics": ( - "static", - (0, 1920800, 1920800, 1920800, 1920800, 1920800, 1920800, 1920800, 1920800, 1920800), - ), - "date_dim": ("static", (0, 73049, 73049, 73049, 73049, 73049, 73049, 73049, 73049, 73049)), - "household_demographics": ("static", (0, 7200, 7200, 7200, 7200, 7200, 7200, 7200, 7200, 7200)), - "income_band": ("static", (0, 20, 20, 20, 20, 20, 20, 20, 20, 20)), - "item": ( - "logarithmic", - (0, 18000, 102000, 204000, 264000, 300000, 360000, 402000, 462000, 502000), - ), - "promotion": ("logarithmic", (0, 300, 500, 1000, 1300, 1500, 1800, 2000, 2300, 2500)), - "reason": ("static", (0, 75, 75, 75, 75, 75, 75, 75, 75, 75)), - "ship_mode": ("static", (0, 20, 20, 20, 20, 20, 20, 20, 20, 20)), - "store": ("logarithmic", (0, 12, 102, 402, 804, 1002, 1350, 1500, 1704, 1902)), - "time_dim": ("static", (0, 86400, 86400, 86400, 86400, 86400, 86400, 86400, 86400, 86400)), - "warehouse": ("logarithmic", (0, 5, 10, 15, 17, 20, 22, 25, 27, 30)), - "web_page": ("logarithmic", (0, 60, 200, 2040, 2604, 3000, 3600, 4002, 4602, 5004)), - "web_site": ("logarithmic", (0, 30, 42, 24, 42, 54, 66, 78, 84, 96)), - } - # Measured from C-compatible SF10 output generated by the pinned tpcgen-cli - # build with ZSTD(1), then normalized linearly to SF1000 for storage. - SF1000_SIZE_GB_DICT = { - "call_center": 0.001222, - "catalog_page": 0.046315, - "catalog_returns": 7.501222, - "catalog_sales": 77.612198, - "customer": 1.954509, - "customer_address": 0.372212, - "customer_demographics": 0.533869, - "date_dim": 0.160342, - "household_demographics": 0.002795, - "income_band": 0.000129, - "inventory": 40.881562, - "item": 0.630520, - "promotion": 0.002376, - "reason": 0.000191, - "ship_mode": 0.000256, - "store": 0.001636, - "store_returns": 11.303521, - "store_sales": 94.989904, - "time_dim": 0.086069, - "warehouse": 0.000494, - "web_page": 0.000703, - "web_returns": 3.976669, - "web_sales": 35.409775, - "web_site": 0.001225, - } - #: Ratio of native pipe-delimited bytes to uncompressed Parquet bytes, - #: measured per table at SF1 with the bundled generator. - NATIVE_SIZE_FACTOR_DICT = { - "call_center": 0.177, - "catalog_page": 1.281, - "catalog_returns": 1.826, - "catalog_sales": 2.848, - "customer": 1.603, - "customer_address": 2.592, - "customer_demographics": 8.285, - "date_dim": 2.68, - "household_demographics": 3.165, - "income_band": 0.237, - "inventory": 6.492, - "item": 1.702, - "promotion": 1.159, - "reason": 0.586, - "ship_mode": 0.355, - "store": 0.299, - "store_returns": 1.986, - "store_sales": 2.859, - "time_dim": 1.735, - "warehouse": 0.123, - "web_page": 0.922, - "web_returns": 1.564, - "web_sales": 3.16, - "web_site": 0.676, - } - ZSTD1_COMPRESSION_FACTOR_DICT = { - "call_center": 1.166, - "catalog_page": 2.632, - "catalog_returns": 1.271, - "catalog_sales": 1.279, - "customer": 1.871, - "customer_address": 2.394, - "customer_demographics": 1.701, - "date_dim": 2.243, - "household_demographics": 1.629, - "income_band": 1.0, - "inventory": 1.026, - "item": 2.717, - "promotion": 2.144, - "reason": 3.299, - "ship_mode": 1.361, - "store": 1.562, - "store_returns": 1.267, - "store_sales": 1.344, - "time_dim": 3.194, - "warehouse": 1.017, - "web_page": 1.701, - "web_returns": 1.223, - "web_sales": 1.233, - "web_site": 1.374, - } - SNAPPY_COMPRESSION_FACTOR_DICT = { - "call_center": 1.073, - "catalog_page": 1.685, - "catalog_returns": 1.092, - "catalog_sales": 1.067, - "customer": 1.404, - "customer_address": 1.703, - "customer_demographics": 1.163, - "date_dim": 1.511, - "household_demographics": 1.122, - "income_band": 1.0, - "inventory": 1.007, - "item": 1.808, - "promotion": 1.470, - "reason": 1.930, - "ship_mode": 1.231, - "store": 1.262, - "store_returns": 1.109, - "store_sales": 1.120, - "time_dim": 1.863, - "warehouse": 1.107, - "web_page": 1.306, - "web_returns": 1.072, - "web_sales": 1.031, - "web_site": 1.152, - } - - def _estimated_table_size_gib(self, table_name: str) -> float: - reference_size_gib = self.SF1000_SIZE_GB_DICT[table_name] * (self.SIZE_REFERENCE_SCALE_FACTOR / 1000) - if table_name in self.ROW_COUNT_SCALING or table_name == "inventory": - scale_ratio = self._expected_row_count(table_name, self.scale_factor) / self._expected_row_count( - table_name, self.SIZE_REFERENCE_SCALE_FACTOR - ) - else: - scale_ratio = self.scale_factor / self.SIZE_REFERENCE_SCALE_FACTOR - zstd_size_gib = reference_size_gib * scale_ratio - return ( - zstd_size_gib - * self.ZSTD1_COMPRESSION_FACTOR_DICT[table_name] - / self.compression_factors_by_table[table_name] - ) - - @classmethod - def _expected_row_count(cls, table_name: str, scale_factor: float) -> int: - if table_name == "inventory": - item_rows = cls._expected_row_count("item", scale_factor) - item_ids = (item_rows // 6) * 3 - item_ids += (0, 1, 2, 2, 3, 3)[item_rows % 6] - return item_ids * cls._expected_row_count("warehouse", scale_factor) * 261 - - model, row_counts = cls.ROW_COUNT_SCALING[table_name] - if scale_factor in cls.DEFINED_SCALE_FACTORS: - return row_counts[cls.DEFINED_SCALE_FACTORS.index(scale_factor)] - if model == "static": - return row_counts[1] - - scale_slot = next( - index for index, defined_scale in enumerate(cls.DEFINED_SCALE_FACTORS) if scale_factor <= defined_scale - ) - lower_scale = cls.DEFINED_SCALE_FACTORS[scale_slot - 1] - upper_scale = cls.DEFINED_SCALE_FACTORS[scale_slot] - fraction = (scale_factor - lower_scale) / (upper_scale - lower_scale) - base_row_count = row_counts[0] if scale_factor < 1 else row_counts[1] - count = int(fraction * (row_counts[scale_slot] - row_counts[scale_slot - 1])) + base_row_count - return count or 1 - - def _native_output_file_name(self, table_name: str, part_number: int, part_count: int) -> str: - """Name native files the way official ``dsdgen`` does. - - Serial ``dsdgen`` writes ``
.dat``; a parallel run started with - ``-parallel -child `` writes - ``
__.dat``. - """ - if part_count == 1: - return f"{table_name}.{self.NATIVE_FILE_EXTENSION}" - return f"{table_name}_{part_number}_{part_count}.{self.NATIVE_FILE_EXTENSION}" diff --git a/src/lakebench/datagen/_tpcgen.py b/src/lakebench/datagen/_tpcgen.py new file mode 100644 index 0000000..a058134 --- /dev/null +++ b/src/lakebench/datagen/_tpcgen.py @@ -0,0 +1,135 @@ +import logging +import math +import shutil +import subprocess +import sysconfig +from collections import defaultdict +from pathlib import Path +from typing import Dict, List, Optional +from urllib.parse import urlparse +from urllib.request import url2pathname + +logger = logging.getLogger(__name__) + + +class TpcgenDataGenerator: + """Generate TPC data with the ``tpcgen-cli`` executable from the tpcgen-cli PyPI package. + + Output is written in tpcgen-cli's layout, one folder per table: + ``/
/
..``. + """ + + #: tpcgen-cli benchmark subcommand. + BENCHMARK = "" + #: tpcgen-cli subcommand and file extension of the native pipe-delimited format. + NATIVE_FORMAT = "" + #: Approximate ZSTD(1) Parquet size of each table at SF1000, in GB. Used only to + #: choose how many parts (files) to split each table into. + SF1000_SIZE_GB: Dict[str, float] = {} + TARGET_FILE_SIZE_MB = 256 + + def __init__( + self, + scale_factor: float, + target_folder_uri: str, + target_row_group_size_mb: Optional[int] = None, + compression: Optional[str] = None, + table_list: Optional[List[str]] = None, + num_threads: Optional[int] = None, + output_format: str = "parquet", + ) -> None: + if output_format not in ("parquet", "native"): + raise ValueError("output_format must be either 'parquet' or 'native'.") + if output_format == "native" and (target_row_group_size_mb is not None or compression is not None): + raise ValueError("target_row_group_size_mb and compression apply only to output_format='parquet'.") + table_list = list(table_list or self.SF1000_SIZE_GB) + unknown_tables = sorted(set(table_list) - set(self.SF1000_SIZE_GB)) + if unknown_tables: + raise ValueError(f"Unsupported {self.BENCHMARK} tables: {', '.join(unknown_tables)}") + + self.scale_factor = scale_factor + self.target_folder = _local_path(target_folder_uri) + self.target_row_group_size_mb = target_row_group_size_mb or 128 + self.compression = compression or "ZSTD(1)" + self.table_list = table_list + self.num_threads = num_threads + self.output_format = output_format + self.executable = _find_executable() + self.parts_by_table = { + table: max(1, math.ceil(self.SF1000_SIZE_GB[table] * scale_factor / 1000 * 1024 / self.TARGET_FILE_SIZE_MB)) + for table in table_list + } + + def run(self) -> None: + self.target_folder.mkdir(parents=True, exist_ok=True) + tables_by_parts = defaultdict(list) + for table in self.table_list: + # tpcgen-cli skips files that already exist; clear any previous output. + shutil.rmtree(self.target_folder / table, ignore_errors=True) + tables_by_parts[self.parts_by_table[table]].append(table) + for parts, tables in tables_by_parts.items(): + logger.info("Generating %s tables %s with %d part(s).", self.BENCHMARK, ", ".join(tables), parts) + _run(self._command(tables, parts)) + + def _command(self, tables: List[str], parts: int) -> List[str]: + command = [ + self.executable, + self.BENCHMARK, + "parquet" if self.output_format == "parquet" else self.NATIVE_FORMAT, + "--scale-factor", + str(self.scale_factor), + "--output-dir", + str(self.target_folder), + "--tables", + ",".join(tables), + "--parts", + str(parts), + "--no-progress", + ] + if self.output_format == "parquet": + command += ["--compression", self.compression, "--row-group-bytes", f"{self.target_row_group_size_mb}MiB"] + if self.num_threads is not None: + command += ["--num-threads", str(self.num_threads)] + return command + + +def _local_path(uri: str) -> Path: + parsed = urlparse(uri) + if parsed.scheme == "file": + return Path(url2pathname(f"//{parsed.netloc}{parsed.path}")).resolve() + if "://" in uri: + raise ValueError(f"tpcgen-cli writes to local or mounted paths only, not {parsed.scheme}:// URIs.") + return Path(uri).expanduser().resolve() + + +def _find_executable() -> str: + # Prefer the running interpreter's scripts directory: notebook kernels often + # do not have the environment's bin directory on PATH. + executable = shutil.which("tpcgen-cli", path=sysconfig.get_path("scripts")) or shutil.which("tpcgen-cli") + if executable is None: + raise ImportError( + "tpcgen-cli is not installed. Install it with `pip install lakebench[tpch_datagen]` or " + "`pip install lakebench[tpcds_datagen]`. TPC-DS can alternatively use backend='duckdb' " + "with lakebench[tpcds_duckdb_datagen] on Python 3.10 or newer." + ) + return executable + + +def _run(command: List[str]) -> None: + try: + result = subprocess.run(command, capture_output=True, text=True, check=True) + except subprocess.CalledProcessError as exc: + details = [f"tpcgen-cli failed with exit code {exc.returncode}: {' '.join(command)}"] + if exc.stdout: + details.append(f"stdout:\n{exc.stdout.rstrip()}") + if exc.stderr: + details.append(f"stderr:\n{exc.stderr.rstrip()}") + if "Failed to rename" in exc.stderr: + details.append( + "Mounted filesystems can fail under concurrent file renames. Retry with a lower " + "num_threads value, such as num_threads=8 or num_threads=16." + ) + raise RuntimeError("\n".join(details)) from exc + for output in (result.stdout, result.stderr): + if output: + logger.info(output.rstrip()) diff --git a/src/lakebench/datagen/_tpcgen_cli.py b/src/lakebench/datagen/_tpcgen_cli.py deleted file mode 100644 index fe22b0f..0000000 --- a/src/lakebench/datagen/_tpcgen_cli.py +++ /dev/null @@ -1,104 +0,0 @@ -import hashlib -import json -import os -import platform -import stat -import subprocess -from pathlib import Path -from typing import List, Optional, Tuple - - -class TpcgenCli: - """Resolve and invoke the unified tpcgen-cli executable.""" - - def __init__(self) -> None: - self.executable, self.provenance = self._resolve_executable() - self.binary_sha256 = self.provenance["binary_sha256"] - - @staticmethod - def _resolve_executable() -> Tuple[str, dict]: - binary_name = "tpcgen-cli.exe" if os.name == "nt" else "tpcgen-cli" - packaged_path = Path(__file__).parent / "_bin" / binary_name - if packaged_path.is_file(): - manifest = TpcgenCli._validate_packaged_binary(packaged_path) - TpcgenCli._ensure_executable(packaged_path) - return str(packaged_path), manifest - - source_tree_path = TpcgenCli._source_tree_binary(binary_name) - if source_tree_path is not None: - manifest = TpcgenCli._validate_packaged_binary(source_tree_path) - TpcgenCli._ensure_executable(source_tree_path) - return str(source_tree_path), manifest - - machine = platform.machine().lower() - raise ImportError( - "The bundled tpcgen-cli executable is unavailable for " - f"{platform.system()} {machine}. Install a supported LakeBench wheel " - "(Windows x86_64 or Linux x86_64). TPC-DS can alternatively use " - "backend='duckdb' with lakebench[tpcds_duckdb_datagen] on Python 3.10 or newer." - ) - - @staticmethod - def _source_tree_binary(binary_name: str) -> Optional[Path]: - platform_key = (platform.system().lower(), platform.machine().lower()) - platform_directories = { - ("windows", "amd64"): "windows-x86_64", - ("windows", "x86_64"): "windows-x86_64", - ("linux", "amd64"): "linux-x86_64", - ("linux", "x86_64"): "linux-x86_64", - } - platform_directory = platform_directories.get(platform_key) - if platform_directory is None: - return None - - repository_root = Path(__file__).parents[3] - binary_path = repository_root / "native" / "tpcgen" / platform_directory / binary_name - return binary_path if binary_path.is_file() else None - - @staticmethod - def _validate_packaged_binary(binary_path: Path) -> dict: - manifest_path = binary_path.parent / "provenance.json" - if not manifest_path.is_file(): - raise RuntimeError(f"Bundled tpcgen-cli provenance is missing: {manifest_path}") - - with manifest_path.open(encoding="utf-8") as manifest_file: - manifest = json.load(manifest_file) - - expected_hash = manifest.get("binary_sha256") - if not expected_hash: - raise RuntimeError("Bundled tpcgen-cli provenance does not contain a SHA-256 checksum.") - - actual_hash = TpcgenCli._sha256(binary_path) - if actual_hash != expected_hash: - raise RuntimeError(f"Bundled tpcgen-cli checksum mismatch: expected {expected_hash}, got {actual_hash}.") - return manifest - - @staticmethod - def _sha256(binary_path: Path) -> str: - digest = hashlib.sha256() - with binary_path.open("rb") as binary_file: - for chunk in iter(lambda: binary_file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - @staticmethod - def _ensure_executable(binary_path: Path) -> None: - binary_path.chmod(binary_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) - - def run(self, args: List[str]) -> subprocess.CompletedProcess: - command = [self.executable] + args - try: - return subprocess.run(command, capture_output=True, text=True, check=True) - except subprocess.CalledProcessError as exc: - details = [f"tpcgen-cli failed with exit code {exc.returncode}: {' '.join(command)}"] - if exc.stdout: - details.append(f"stdout:\n{exc.stdout.rstrip()}") - if exc.stderr: - details.append(f"stderr:\n{exc.stderr.rstrip()}") - if "Failed to rename" in exc.stderr and "os error 5" in exc.stderr: - details.append( - "The target mounted filesystem returned an I/O error while finalizing a generated file. " - "Remove the incomplete table directory and retry with a lower num_threads value, such as " - "num_threads=8 or num_threads=16, to reduce concurrent file creation and rename pressure." - ) - raise RuntimeError("\n".join(details)) from exc diff --git a/src/lakebench/datagen/_tpcgen_rs.py b/src/lakebench/datagen/_tpcgen_rs.py deleted file mode 100644 index 6f5fabe..0000000 --- a/src/lakebench/datagen/_tpcgen_rs.py +++ /dev/null @@ -1,339 +0,0 @@ -import logging -import math -import os -import re -import shutil -from collections import defaultdict -from pathlib import Path -from typing import Dict, List, Optional, Tuple -from urllib.parse import urlparse -from urllib.request import url2pathname - -from ._tpcgen_cli import TpcgenCli - -logger = logging.getLogger(__name__) - - -class _TpcgenRsDataGenerator: - BENCHMARK_NAME = "" - BENCHMARK_LABEL = "" - COMMAND_ARGUMENTS: Tuple[str, ...] = () - #: tpcgen-cli subcommand and file extension for the generator's native - #: pipe-delimited output. TPC-DS emits ``dat`` and TPC-H emits ``tbl``, - #: matching the official ``dsdgen`` and ``dbgen`` conventions. - NATIVE_SUBCOMMAND = "" - NATIVE_FILE_EXTENSION = "" - #: ``tpcgen-cli tpcds dat`` does not accept ``--num-threads`` even though the - #: parquet subcommand does. - NATIVE_SUPPORTS_NUM_THREADS = True - OUTPUT_FORMATS = ("parquet", "native") - PARQUET_ONLY_OPTIONS = ("target_row_group_size_mb", "compression", "compression_factor") - #: Ratio of native pipe-delimited bytes to uncompressed Parquet bytes, measured - #: per table at SF1. Both sides are measured at the same scale factor, so the - #: ratio is a property of the encoding and holds at any scale factor. - NATIVE_SIZE_FACTOR_DICT: Dict[str, float] = {} - GEN_TABLE_REGISTRY: List[str] = [] - SF1000_SIZE_GB_DICT: Dict[str, float] = {} - ZSTD1_COMPRESSION_FACTOR_DICT: Dict[str, float] = {} - SNAPPY_COMPRESSION_FACTOR_DICT: Dict[str, float] = {} - TARGET_FILE_SIZE_MAP = [ - (10, 128), - (1024, 256), - (5120, 512), - (10240, 1024), - ] - MINIMUM_FILE_SIZE_RATIO = 0.5 - MINIMUM_FILE_SIZE_ESTIMATION_HEADROOM = 1.10 - ROW_GROUP_ESTIMATION_HEADROOM = 1.05 - CLOUD_SCHEMES = {"s3", "gs", "gcs", "abfs", "abfss", "adl", "wasb", "wasbs"} - COMPRESSION_PATTERN = re.compile( - r"^(UNCOMPRESSED|SNAPPY|LZO|LZ4|LZ4_RAW|GZIP(?:\(\d+\))?|BROTLI(?:\(\d+\))?|ZSTD\(\d+\))$" - ) - - def __init__( - self, - scale_factor: float, - target_folder_uri: str, - target_row_group_size_mb: Optional[int] = None, - compression: Optional[str] = None, - table_list: Optional[List[str]] = None, - num_threads: Optional[int] = None, - compression_factor: Optional[float] = None, - output_format: str = "parquet", - ) -> None: - if not self.BENCHMARK_NAME or not self.BENCHMARK_LABEL: - raise TypeError("_TpcgenRsDataGenerator must be subclassed with benchmark metadata.") - if output_format not in self.OUTPUT_FORMATS: - raise ValueError(f"output_format must be one of: {', '.join(self.OUTPUT_FORMATS)}") - - self.output_format = output_format - if output_format == "native": - supplied_parquet_options = [ - name - for name, value in zip( - self.PARQUET_ONLY_OPTIONS, - (target_row_group_size_mb, compression, compression_factor), - ) - if value is not None - ] - if supplied_parquet_options: - raise ValueError( - "The following options apply only to output_format='parquet': " - + ", ".join(supplied_parquet_options) - ) - if num_threads is not None and not self.NATIVE_SUPPORTS_NUM_THREADS: - raise ValueError( - f"'num_threads' is not supported by tpcgen-cli {self.BENCHMARK_NAME} " - f"{self.NATIVE_SUBCOMMAND}, which has no thread option. Use 'parts' " - "parallelism instead, or generate Parquet." - ) - target_row_group_size_mb = 128 if target_row_group_size_mb is None else target_row_group_size_mb - # Native pipe-delimited text has no column compression, so the uncompressed - # estimate is the closest available proxy for its on-disk size. - if compression is None: - compression = "UNCOMPRESSED" if output_format == "native" else "ZSTD(1)" - - parsed_uri = urlparse(target_folder_uri) - uri_scheme = parsed_uri.scheme.lower() - if uri_scheme in self.CLOUD_SCHEMES: - raise ValueError( - f"{uri_scheme} protocol is not supported for {self.BENCHMARK_LABEL} Rust data generation. " - "Use a local filesystem path or a mounted storage path." - ) - if scale_factor <= 0: - raise ValueError("scale_factor must be greater than zero.") - if target_row_group_size_mb <= 0: - raise ValueError("target_row_group_size_mb must be greater than zero.") - if num_threads is not None and ( - isinstance(num_threads, bool) or not isinstance(num_threads, int) or num_threads <= 0 - ): - raise ValueError("num_threads must be a positive integer.") - - compression = compression.upper() - if not self.COMPRESSION_PATTERN.fullmatch(compression): - raise ValueError(f"Unsupported compression codec: {compression}") - if compression_factor is not None and compression_factor < 1: - raise ValueError("compression_factor must be greater than or equal to one.") - - requested_tables = table_list if table_list is not None else self.GEN_TABLE_REGISTRY - unknown_tables = sorted(set(requested_tables) - set(self.GEN_TABLE_REGISTRY)) - if unknown_tables: - raise ValueError(f"Unsupported {self.BENCHMARK_LABEL} tables: {', '.join(unknown_tables)}") - if not requested_tables: - raise ValueError("table_list must contain at least one table.") - - target_path = ( - Path(url2pathname(f"//{parsed_uri.netloc}{parsed_uri.path}")) - if uri_scheme == "file" - else Path(target_folder_uri) - ) - self.scale_factor = scale_factor - self.target_folder = target_path.expanduser().resolve() - self.target_row_group_size_mb = target_row_group_size_mb - self.compression = compression - self.table_list = list(dict.fromkeys(requested_tables)) - self.compression_factor = compression_factor - self.compression_factors_by_table = { - table_name: self._resolve_compression_factor(table_name, compression_factor) - for table_name in self.table_list - } - self.num_threads = num_threads if num_threads is not None else (os.cpu_count() or 1) - self.cli = TpcgenCli() - self.parts_by_table = {table_name: self._calculate_optimal_parts(table_name) for table_name in self.table_list} - - def run(self) -> None: - self._prepare_target() - for (part_count, compression_factor), table_names in self._group_tables_by_generation_settings().items(): - if self.output_format == "native": - logger.info( - "Generating %s tables %s as %s with %d part(s).", - self.BENCHMARK_LABEL, - ", ".join(table_names), - self.NATIVE_FILE_EXTENSION, - part_count, - ) - else: - logger.info( - "Generating %s tables %s with %d part(s), %.2f MB on-disk row groups, and a %.3f compression factor.", - self.BENCHMARK_LABEL, - ", ".join(table_names), - part_count, - self.target_row_group_size_mb, - compression_factor, - ) - result = self.cli.run(self._build_command(self.target_folder, table_names, part_count, compression_factor)) - if result.stdout: - logger.info(result.stdout.rstrip()) - if result.stderr: - logger.info(result.stderr.rstrip()) - - self._normalize_outputs() - missing_outputs = [str(path) for path in self._expected_outputs() if not path.is_file()] - if missing_outputs: - raise RuntimeError("tpcgen-cli did not create expected files: " + ", ".join(missing_outputs)) - - def _prepare_target(self) -> None: - self.target_folder.mkdir(parents=True, exist_ok=True) - for table_name in self.table_list: - root_file = self.target_folder / f"{table_name}.{self._source_file_extension()}" - if root_file.exists(): - root_file.unlink() - table_folder = self.target_folder / table_name - if table_folder.exists(): - shutil.rmtree(table_folder) - - def _source_file_extension(self) -> str: - return "parquet" if self.output_format == "parquet" else self.NATIVE_FILE_EXTENSION - - def _normalize_outputs(self) -> None: - extension = self._source_file_extension() - for table_name in self.table_list: - part_count = self.parts_by_table[table_name] - table_folder = self.target_folder / table_name - table_folder.mkdir(parents=True, exist_ok=True) - for part_number in range(1, part_count + 1): - source_candidates = [table_folder / f"{table_name}.{part_number}.{extension}"] - if part_count == 1: - source_candidates.insert(0, self.target_folder / f"{table_name}.{extension}") - source_file = next((path for path in source_candidates if path.is_file()), source_candidates[0]) - if source_file.is_file(): - shutil.move( - str(source_file), - str(table_folder / self._output_file_name(table_name, part_number)), - ) - - def _expected_outputs(self) -> List[Path]: - outputs = [] - for table_name in self.table_list: - part_count = self.parts_by_table[table_name] - table_folder = self.target_folder / table_name - outputs.extend( - table_folder / self._output_file_name(table_name, part_number) - for part_number in range(1, part_count + 1) - ) - return outputs - - def _output_file_name(self, table_name: str, part_number: int) -> str: - if self.output_format == "native": - return self._native_output_file_name(table_name, part_number, self.parts_by_table[table_name]) - compression_name = self.compression.partition("(")[0].lower() - return f"{table_name}-{part_number:05d}.{compression_name}.parquet" - - def _native_output_file_name(self, table_name: str, part_number: int, part_count: int) -> str: - """Return the file name the benchmark's official generator would use. - - Overridden per benchmark so native output is named exactly like - ``dsdgen``/``dbgen`` output. - """ - raise NotImplementedError(f"{type(self).__name__} does not define native output file naming.") - - def _group_tables_by_generation_settings(self) -> Dict[Tuple[int, float], List[str]]: - grouped_tables = defaultdict(list) - for table_name in self.table_list: - settings = ( - self.parts_by_table[table_name], - self.compression_factors_by_table[table_name], - ) - grouped_tables[settings].append(table_name) - return dict(grouped_tables) - - def _build_command( - self, - output_folder: Path, - table_names: List[str], - part_count: int, - compression_factor: Optional[float] = None, - ) -> List[str]: - if compression_factor is None: - factors = {self.compression_factors_by_table[table_name] for table_name in table_names} - if len(factors) != 1: - raise ValueError("Tables with different compression factors must be generated separately.") - compression_factor = factors.pop() - - if self.output_format == "native": - command = [ - self.BENCHMARK_NAME, - self.NATIVE_SUBCOMMAND, - "--scale-factor", - str(self.scale_factor), - "--output-dir", - str(output_folder), - "--tables", - ",".join(table_names), - *self.COMMAND_ARGUMENTS, - ] - if self.NATIVE_SUPPORTS_NUM_THREADS: - command += ["--num-threads", str(self.num_threads)] - return command + ["--parts", str(part_count), "--no-progress"] - - return [ - self.BENCHMARK_NAME, - "parquet", - "--scale-factor", - str(self.scale_factor), - "--output-dir", - str(output_folder), - "--tables", - ",".join(table_names), - *self.COMMAND_ARGUMENTS, - "--compression", - self.compression, - "--row-group-bytes", - str( - round( - self.target_row_group_size_mb - * compression_factor - * self.ROW_GROUP_ESTIMATION_HEADROOM - * 1024 - * 1024 - ) - ), - "--num-threads", - str(self.num_threads), - "--parts", - str(part_count), - "--no-progress", - ] - - def _calculate_optimal_parts(self, table_name: str) -> int: - scaled_size_gib = self._estimated_table_size_gib(table_name) - if scaled_size_gib < self.TARGET_FILE_SIZE_MAP[0][0]: - minimum_file_size_mb = ( - self.target_row_group_size_mb - * self.MINIMUM_FILE_SIZE_RATIO - * self.MINIMUM_FILE_SIZE_ESTIMATION_HEADROOM - ) - return max(math.floor(scaled_size_gib * 1024 / minimum_file_size_mb), 1) - target_file_size_mb = self._target_file_size_mb(scaled_size_gib) - return max(math.ceil(scaled_size_gib * 1024 / target_file_size_mb), 1) - - def _estimated_table_size_gib(self, table_name: str) -> float: - zstd_size_gib = self.SF1000_SIZE_GB_DICT[table_name] * (self.scale_factor / 1000) - estimated_gib = ( - zstd_size_gib - * self.ZSTD1_COMPRESSION_FACTOR_DICT[table_name] - / self.compression_factors_by_table[table_name] - ) - if self.output_format == "native": - estimated_gib *= self.NATIVE_SIZE_FACTOR_DICT[table_name] - return estimated_gib - - def _target_file_size_mb(self, scaled_size_gib: float) -> int: - for threshold_gib, target_size_mb in self.TARGET_FILE_SIZE_MAP: - if scaled_size_gib < threshold_gib: - return target_size_mb - return 1024 - - def _resolve_compression_factor(self, table_name: str, compression_factor: Optional[float]) -> float: - if compression_factor is not None: - return compression_factor - if self.compression == "UNCOMPRESSED": - return 1.0 - if self.compression.startswith("ZSTD("): - return self.ZSTD1_COMPRESSION_FACTOR_DICT[table_name] - if self.compression == "SNAPPY": - return self.SNAPPY_COMPRESSION_FACTOR_DICT[table_name] - raise ValueError( - f"compression_factor is required for {self.compression}; " - "built-in measured defaults are available only for ZSTD(N) and SNAPPY." - ) diff --git a/src/lakebench/datagen/_tpch_rs.py b/src/lakebench/datagen/_tpch_rs.py deleted file mode 100644 index 3ae815a..0000000 --- a/src/lakebench/datagen/_tpch_rs.py +++ /dev/null @@ -1,71 +0,0 @@ -from ._tpcgen_rs import _TpcgenRsDataGenerator - - -class _TPCHRsDataGenerator(_TpcgenRsDataGenerator): - BENCHMARK_NAME = "tpch" - BENCHMARK_LABEL = "TPC-H" - NATIVE_SUBCOMMAND = "tbl" - NATIVE_FILE_EXTENSION = "tbl" - GEN_TABLE_REGISTRY = [ - "customer", - "lineitem", - "nation", - "orders", - "part", - "partsupp", - "region", - "supplier", - ] - # Measured from SF10 output generated by the pinned tpcgen-cli build with - # ZSTD(1), then normalized linearly to SF1000. - SF1000_SIZE_GB_DICT = { - "customer": 7.755863, - "lineitem": 145.739593, - "nation": 0.000250, - "orders": 36.859225, - "part": 4.295145, - "partsupp": 27.735921, - "region": 0.000141, - "supplier": 0.550992, - } - #: Ratio of native pipe-delimited bytes to uncompressed Parquet bytes, - #: measured per table at SF1 with the bundled generator. - NATIVE_SIZE_FACTOR_DICT = { - "customer": 0.961, - "lineitem": 2.031, - "nation": 0.549, - "orders": 1.529, - "part": 1.733, - "partsupp": 1.064, - "region": 0.239, - "supplier": 0.86, - } - ZSTD1_COMPRESSION_FACTOR_DICT = { - "customer": 2.956, - "lineitem": 2.465, - "nation": 2.027, - "orders": 3.054, - "part": 2.947, - "partsupp": 3.800, - "region": 1.268, - "supplier": 2.702, - } - SNAPPY_COMPRESSION_FACTOR_DICT = { - "customer": 1.974, - "lineitem": 1.866, - "nation": 1.570, - "orders": 2.098, - "part": 2.154, - "partsupp": 2.789, - "region": 1.149, - "supplier": 1.858, - } - - def _native_output_file_name(self, table_name: str, part_number: int, part_count: int) -> str: - """Name native files the way official ``dbgen`` does. - - Serial ``dbgen`` writes ``
.tbl``; a parallel run started with - ``-C -S `` writes ``
.tbl.``. - """ - base_name = f"{table_name}.{self.NATIVE_FILE_EXTENSION}" - return base_name if part_count == 1 else f"{base_name}.{part_number}" diff --git a/src/lakebench/datagen/tpcds.py b/src/lakebench/datagen/tpcds.py index 72f9e33..6c05ce7 100644 --- a/src/lakebench/datagen/tpcds.py +++ b/src/lakebench/datagen/tpcds.py @@ -1,7 +1,7 @@ from typing import List, Optional from ._tpc import _TPCDataGenerator -from ._tpcds_rs import _TPCDSRsDataGenerator +from ._tpcgen import TpcgenDataGenerator class _TPCDSDuckDBDataGenerator(_TPCDataGenerator): @@ -9,51 +9,63 @@ class _TPCDSDuckDBDataGenerator(_TPCDataGenerator): GEN_TYPE = "tpcds" +class _TPCDSTpcgenDataGenerator(TpcgenDataGenerator): + BENCHMARK = "tpcds" + NATIVE_FORMAT = "dat" + SF1000_SIZE_GB = { + "call_center": 0.001222, + "catalog_page": 0.046315, + "catalog_returns": 7.501222, + "catalog_sales": 77.612198, + "customer": 1.954509, + "customer_address": 0.372212, + "customer_demographics": 0.533869, + "date_dim": 0.160342, + "household_demographics": 0.002795, + "income_band": 0.000129, + "inventory": 40.881562, + "item": 0.630520, + "promotion": 0.002376, + "reason": 0.000191, + "ship_mode": 0.000256, + "store": 0.001636, + "store_returns": 11.303521, + "store_sales": 94.989904, + "time_dim": 0.086069, + "warehouse": 0.000494, + "web_page": 0.000703, + "web_returns": 3.976669, + "web_sales": 35.409775, + "web_site": 0.001225, + } + + class TPCDSDataGenerator: """ - Generate TPC-DS data with the bundled Rust tpcgen-cli executable or the - legacy DuckDB generator. + Generate TPC-DS data with ``tpcgen-cli`` (installed by ``lakebench[tpcds_datagen]``) + or the legacy DuckDB generator. Parameters ---------- scale_factor : float - The scale factor for the data generation, which determines the size of the generated dataset. + Scale factor that determines the generated dataset size. target_folder_uri : str - The folder path where the generated Parquet data will be stored. A folder for each table will be created. + Local or mounted folder. Each table is written to its own subfolder. target_row_group_size_mb : int, default=128 - Target on-disk size of Parquet row groups in megabytes. The Rust - generator scales the upstream uncompressed-byte target by the - compression factor. + Approximate uncompressed Parquet row-group size, passed to ``--row-group-bytes``. compression : str, default="ZSTD(1)" - Parquet compression used by the Rust backend. + Parquet compression codec, for example ``"SNAPPY"`` or ``"ZSTD(3)"``. table_list : list of str, optional - TPC-DS tables to generate. The Rust backend generates all 24 tables by default. + TPC-DS tables to generate. All 24 tables are generated by default. num_threads : int, optional - Rust worker-thread count. Defaults to all available CPU cores. Use a - lower value such as 8 or 16 for large generations targeting mounted - filesystems. Rejected when ``output_format="native"``, because - the currently bundled ``tpcgen-cli tpcds dat`` has no thread option - and runs single-threaded regardless of available cores. Native part - count controls output splitting, not parallel execution. + Generator thread count. Defaults to all available CPU cores. Use a lower + value such as 8 or 16 for large generations targeting mounted filesystems. backend : {"rust", "duckdb"}, default="rust" - Data generator backend. DuckDB is retained as an explicit legacy fallback. - compression_factor : float, optional - Ratio of uncompressed to on-disk Parquet bytes used to translate the - requested row-group target and estimate physical table size. Measured - ZSTD levels use per-table factors measured with ``ZSTD(1)`` and - ``SNAPPY`` uses its own measured factors; other compressed codecs - require an explicit value. + ``"rust"`` uses tpcgen-cli. ``"duckdb"`` is the legacy DuckDB generator. output_format : {"parquet", "native"}, default="parquet" - ``"native"`` emits the generator's pipe-delimited ``.dat`` files, the - same layout the official ``dsdgen`` produces. TPC-DS native format - support is currently in preview. Parquet-only options - (``target_row_group_size_mb``, ``compression``, - ``compression_factor``) are rejected in that mode. Requires - ``backend="rust"``. - Methods - ------- - run() - Generates TPC-DS data in Parquet format based on the input scale factor and writes it to the target folder. + ``"native"`` writes the pipe-delimited ``.dat`` format of ``dsdgen`` + (preview). ``target_row_group_size_mb`` and ``compression`` apply only + to Parquet. Requires ``backend="rust"``. """ def __init__( @@ -65,7 +77,6 @@ def __init__( table_list: Optional[List[str]] = None, num_threads: Optional[int] = None, backend: str = "rust", - compression_factor: Optional[float] = None, output_format: str = "parquet", ) -> None: self.scale_factor = scale_factor @@ -75,25 +86,19 @@ def __init__( self.output_format = output_format if backend == "rust": - self._generator = _TPCDSRsDataGenerator( + self._generator = _TPCDSTpcgenDataGenerator( scale_factor=scale_factor, target_folder_uri=target_folder_uri, target_row_group_size_mb=target_row_group_size_mb, compression=compression, table_list=table_list, num_threads=num_threads, - compression_factor=compression_factor, output_format=output_format, ) elif backend == "duckdb": if output_format != "parquet": raise ValueError("output_format='native' is supported only by backend='rust'.") - rust_options = { - "compression": compression, - "table_list": table_list, - "num_threads": num_threads, - "compression_factor": compression_factor, - } + rust_options = {"compression": compression, "table_list": table_list, "num_threads": num_threads} specified_options = [name for name, value in rust_options.items() if value is not None] if specified_options: raise ValueError( diff --git a/src/lakebench/datagen/tpch.py b/src/lakebench/datagen/tpch.py index f0bd027..ee3a1a2 100644 --- a/src/lakebench/datagen/tpch.py +++ b/src/lakebench/datagen/tpch.py @@ -1,36 +1,39 @@ -from ._tpch_rs import _TPCHRsDataGenerator +from ._tpcgen import TpcgenDataGenerator -class TPCHDataGenerator(_TPCHRsDataGenerator): +class TPCHDataGenerator(TpcgenDataGenerator): """ - Generate TPC-H Parquet data with the bundled unified ``tpcgen-cli``. + Generate TPC-H data with ``tpcgen-cli`` (installed by ``lakebench[tpch_datagen]``). Parameters ---------- scale_factor : float Scale factor that determines the generated dataset size. target_folder_uri : str - Local or mounted folder where table directories are written. + Local or mounted folder. Each table is written to its own subfolder. target_row_group_size_mb : int, default=128 - Target on-disk Parquet row-group size in megabytes. + Approximate uncompressed Parquet row-group size, passed to ``--row-group-bytes``. compression : str, default="ZSTD(1)" - Parquet compression codec. + Parquet compression codec, for example ``"SNAPPY"`` or ``"ZSTD(3)"``. table_list : list of str, optional TPC-H tables to generate. All eight tables are generated by default. num_threads : int, optional - Rust worker-thread count. Defaults to all available CPU cores. Use a - lower value such as 8 or 16 for large generations targeting mounted - filesystems. Honoured for both ``output_format="parquet"`` and - ``output_format="native"``, since ``tpcgen-cli tpch tbl`` accepts - ``--num-threads``. - compression_factor : float, optional - Ratio of uncompressed to on-disk Parquet bytes. Measured per-table - ZSTD levels use per-table factors measured with ``ZSTD(1)`` and - ``SNAPPY`` uses its own measured factors; other compressed codecs - require an explicit value. + Generator thread count. Defaults to all available CPU cores. Use a lower + value such as 8 or 16 for large generations targeting mounted filesystems. output_format : {"parquet", "native"}, default="parquet" - ``"native"`` emits the generator's pipe-delimited ``.tbl`` files, the - same layout the official ``dbgen`` produces. Parquet-only options - (``target_row_group_size_mb``, ``compression``, ``compression_factor``) - are rejected in that mode. + ``"native"`` writes the pipe-delimited ``.tbl`` format of ``dbgen``. + ``target_row_group_size_mb`` and ``compression`` apply only to Parquet. """ + + BENCHMARK = "tpch" + NATIVE_FORMAT = "tbl" + SF1000_SIZE_GB = { + "customer": 7.755863, + "lineitem": 145.739593, + "nation": 0.000250, + "orders": 36.859225, + "part": 4.295145, + "partsupp": 27.735921, + "region": 0.000141, + "supplier": 0.550992, + } diff --git a/tests/test_tpc_native_format.py b/tests/test_tpc_native_format.py index 7aa77d1..640bf21 100644 --- a/tests/test_tpc_native_format.py +++ b/tests/test_tpc_native_format.py @@ -6,9 +6,8 @@ from lakebench.benchmarks.tpcds import TPCDS from lakebench.benchmarks.tpch import TPCH -from lakebench.datagen._tpcds_rs import _TPCDSRsDataGenerator -from lakebench.datagen._tpch_rs import _TPCHRsDataGenerator -from lakebench.datagen.tpcds import TPCDSDataGenerator +from lakebench.datagen.tpcds import _TPCDSTpcgenDataGenerator +from lakebench.datagen.tpch import TPCHDataGenerator from lakebench.engines.base import BaseEngine from lakebench.engines.duckdb import DuckDB from lakebench.engines.polars import Polars @@ -21,179 +20,11 @@ ) from .conftest import _uninitialized_engine -from .test_tpcds_datagen import fake_executable # noqa: F401 TPCDS_DDL = Path("src/lakebench/benchmarks/tpcds/resources/ddl/canonical/ddl_v4.0.0.simple.sql") TPCH_DDL = Path("src/lakebench/benchmarks/tpch/resources/ddl/canonical/ddl_v3.0.1.simple.sql") -def _mock_generation(generator, extension): - def create_outputs(args): - output_path = Path(args[args.index("--output-dir") + 1]) - parts = int(args[args.index("--parts") + 1]) - for table_name in args[args.index("--tables") + 1].split(","): - table_dir = output_path / table_name - table_dir.mkdir(parents=True, exist_ok=True) - for part_number in range(1, parts + 1): - (table_dir / f"{table_name}.{part_number}.{extension}").write_bytes(b"row|") - return subprocess.CompletedProcess(args, 0, "", "") - - generator.cli.run = Mock(side_effect=create_outputs) - return generator - - -def test_tpcds_native_uses_dat_subcommand_and_omits_parquet_arguments(tmp_path, fake_executable): # noqa: F811 - output_dir = tmp_path / "tpcds" - generator = _mock_generation( - _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(output_dir), - table_list=["store_sales"], - output_format="native", - ), - "dat", - ) - generator.run() - - args = generator.cli.run.call_args.args[0] - assert args[:2] == ["tpcds", "dat"] - assert args[args.index("--compat") + 1] == "c" - # `tpcgen-cli tpcds dat` rejects these Parquet-only arguments. - assert "--compression" not in args - assert "--row-group-bytes" not in args - assert "--num-threads" not in args - - -def test_tpch_native_uses_tbl_subcommand_and_keeps_num_threads(tmp_path, fake_executable): # noqa: F811 - output_dir = tmp_path / "tpch" - generator = _mock_generation( - _TPCHRsDataGenerator( - scale_factor=1, - target_folder_uri=str(output_dir), - table_list=["lineitem"], - num_threads=4, - output_format="native", - ), - "tbl", - ) - generator.run() - - args = generator.cli.run.call_args.args[0] - assert args[:2] == ["tpch", "tbl"] - assert args[args.index("--num-threads") + 1] == "4" - assert "--compression" not in args - assert "--row-group-bytes" not in args - - -def test_native_outputs_use_official_generator_file_names(tmp_path, fake_executable): # noqa: F811 - output_dir = tmp_path / "tpch" - generator = _mock_generation( - _TPCHRsDataGenerator( - scale_factor=1, - target_folder_uri=str(output_dir), - table_list=["nation"], - output_format="native", - ), - "tbl", - ) - generator.run() - - assert [path.name for path in (output_dir / "nation").glob("*")] == ["nation.tbl"] - - -def test_native_part_names_match_official_parallel_generators(): - tpch = _TPCHRsDataGenerator(scale_factor=1, target_folder_uri="/tmp/x", table_list=["orders"]) - tpcds = _TPCDSRsDataGenerator(scale_factor=1, target_folder_uri="/tmp/x", table_list=["store_sales"]) - - # Serial dbgen/dsdgen write a single unsuffixed file per table. - assert tpch._native_output_file_name("orders", 1, 1) == "orders.tbl" - assert tpcds._native_output_file_name("store_sales", 1, 1) == "store_sales.dat" - - # Parallel dbgen writes
.tbl.; dsdgen writes
__.dat. - assert tpch._native_output_file_name("orders", 3, 4) == "orders.tbl.3" - assert tpcds._native_output_file_name("store_sales", 3, 4) == "store_sales_3_4.dat" - - -def test_native_globs_match_every_official_part_name(): - from fnmatch import fnmatch - - for benchmark, generator, table in ( - (TPCH, _TPCHRsDataGenerator, "orders"), - (TPCDS, _TPCDSRsDataGenerator, "store_sales"), - ): - instance = generator(scale_factor=1, target_folder_uri="/tmp/x", table_list=[table]) - for part_count in (1, 4): - for part_number in range(1, part_count + 1): - name = instance._native_output_file_name(table, part_number, part_count) - assert fnmatch(name, benchmark.NATIVE_FILE_GLOB), (benchmark.__name__, name) - - -def test_native_sizing_applies_measured_expansion_factor(tmp_path, fake_executable): # noqa: F811 - shared = dict(scale_factor=1000, target_folder_uri=str(tmp_path / "t"), table_list=["lineitem"]) - parquet_generator = _TPCHRsDataGenerator(**shared) - native_generator = _TPCHRsDataGenerator(**shared, output_format="native") - - uncompressed_gib = ( - _TPCHRsDataGenerator.SF1000_SIZE_GB_DICT["lineitem"] - * _TPCHRsDataGenerator.ZSTD1_COMPRESSION_FACTOR_DICT["lineitem"] - ) - expected = uncompressed_gib * _TPCHRsDataGenerator.NATIVE_SIZE_FACTOR_DICT["lineitem"] - - assert native_generator._estimated_table_size_gib("lineitem") == pytest.approx(expected) - # Native text is larger than compressed Parquet, so it must be split more. - assert native_generator.parts_by_table["lineitem"] > parquet_generator.parts_by_table["lineitem"] - - -@pytest.mark.parametrize("option", ["target_row_group_size_mb", "compression", "compression_factor"]) -def test_native_rejects_parquet_only_options(tmp_path, fake_executable, option): # noqa: F811 - values = {"target_row_group_size_mb": 256, "compression": "SNAPPY", "compression_factor": 2.0} - with pytest.raises(ValueError, match="apply only to output_format='parquet'"): - _TPCHRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path / "t"), - output_format="native", - **{option: values[option]}, - ) - - -def test_tpcds_native_rejects_num_threads(tmp_path, fake_executable): # noqa: F811 - # `tpcgen-cli tpcds dat` writes serially and exposes no thread option, so an - # explicit num_threads must be rejected rather than silently ignored. - with pytest.raises(ValueError, match="'num_threads' is not supported"): - _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path / "t"), - output_format="native", - num_threads=4, - ) - - -def test_tpcds_parquet_still_accepts_num_threads(tmp_path, fake_executable): # noqa: F811 - generator = _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path / "t"), - output_format="parquet", - num_threads=4, - ) - - assert generator.num_threads == 4 - - -def test_native_output_format_is_rejected_for_duckdb_backend(tmp_path): - with pytest.raises(ValueError, match="supported only by backend='rust'"): - TPCDSDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path / "t"), - backend="duckdb", - output_format="native", - ) - - -def test_unknown_output_format_is_rejected(tmp_path, fake_executable): # noqa: F811 - with pytest.raises(ValueError, match="output_format must be one of"): - _TPCHRsDataGenerator(scale_factor=1, target_folder_uri=str(tmp_path / "t"), output_format="csv") - - def test_ddl_schema_derivation_covers_every_table(): tpcds_schemas = table_schemas_from_ddl(TPCDS_DDL.read_text()) tpch_schemas = table_schemas_from_ddl(TPCH_DDL.read_text()) @@ -222,14 +53,9 @@ def test_every_canonical_ddl_variant_derives_a_usable_schema(): schema_to_sql_types(columns, "duckdb") -def test_native_size_factors_cover_every_generated_table(): - for generator in (_TPCHRsDataGenerator, _TPCDSRsDataGenerator): - assert set(generator.NATIVE_SIZE_FACTOR_DICT) == set(generator.GEN_TABLE_REGISTRY) - - def test_benchmarks_declare_native_extensions_matching_generators(): - assert TPCDS.NATIVE_FILE_EXTENSION == _TPCDSRsDataGenerator.NATIVE_FILE_EXTENSION == "dat" - assert TPCH.NATIVE_FILE_EXTENSION == _TPCHRsDataGenerator.NATIVE_FILE_EXTENSION == "tbl" + assert TPCDS.NATIVE_FILE_EXTENSION == _TPCDSTpcgenDataGenerator.NATIVE_FORMAT == "dat" + assert TPCH.NATIVE_FILE_EXTENSION == TPCHDataGenerator.NATIVE_FORMAT == "tbl" assert TPCDS.NATIVE_FILE_GLOB == "*.dat" assert TPCH.NATIVE_FILE_GLOB == "*.tbl*" diff --git a/tests/test_tpcds_datagen.py b/tests/test_tpcds_datagen.py deleted file mode 100644 index a684727..0000000 --- a/tests/test_tpcds_datagen.py +++ /dev/null @@ -1,572 +0,0 @@ -import hashlib -import inspect -import json -import stat -import subprocess -from pathlib import Path -from unittest.mock import Mock - -import pytest - -from lakebench.datagen._tpcds_rs import _TPCDSRsDataGenerator -from lakebench.datagen._tpcgen_cli import TpcgenCli -from lakebench.datagen._tpcgen_rs import _TpcgenRsDataGenerator -from lakebench.datagen._tpch_rs import _TPCHRsDataGenerator -from lakebench.datagen.tpcds import TPCDSDataGenerator -from lakebench.datagen.tpch import TPCHDataGenerator - - -@pytest.fixture -def fake_executable(tmp_path, monkeypatch): - executable = tmp_path / "tpcgen-cli" - executable.write_bytes(b"fake") - provenance = { - "source": "test", - "binary_sha256": hashlib.sha256(b"fake").hexdigest(), - } - monkeypatch.setattr( - TpcgenCli, - "_resolve_executable", - staticmethod(lambda: (str(executable), provenance)), - ) - return str(executable) - - -def test_rust_generator_builds_automatic_multipart_command_and_outputs(tmp_path, fake_executable): - output_dir = tmp_path / "tpcds" - generator = _TPCDSRsDataGenerator( - scale_factor=10, - target_folder_uri=str(output_dir), - table_list=["store_sales"], - num_threads=1, - ) - - def create_outputs(args): - output_path = Path(args[args.index("--output-dir") + 1]) - table_dir = output_path / "store_sales" - table_dir.mkdir() - parts = int(args[args.index("--parts") + 1]) - for part_number in range(1, parts + 1): - (table_dir / f"store_sales.{part_number}.parquet").write_bytes(b"parquet") - return subprocess.CompletedProcess(args, 0, "", "") - - generator.cli.run = Mock(side_effect=create_outputs) - generator.run() - - args = generator.cli.run.call_args.args[0] - assert args[:2] == ["tpcds", "parquet"] - assert args[args.index("--tables") + 1] == "store_sales" - assert args[args.index("--parts") + 1] == "13" - assert args[args.index("--compat") + 1] == "c" - assert args[args.index("--num-threads") + 1] == "1" - assert Path(args[args.index("--output-dir") + 1]) == output_dir.resolve() - assert "--no-progress" in args - assert sorted(path.name for path in (output_dir / "store_sales").glob("*.parquet")) == [ - f"store_sales-{part_number:05d}.zstd.parquet" for part_number in range(1, 14) - ] - assert not list(output_dir.rglob("*.json")) - - -def test_public_generator_mirrors_tpch_generation_parameters(): - tpcds_parameters = list(inspect.signature(TPCDSDataGenerator).parameters) - tpch_parameters = list(inspect.signature(TPCHDataGenerator).parameters) - - assert tpcds_parameters[:6] == tpch_parameters[:6] - assert "executable" not in tpcds_parameters - assert "executable" not in tpch_parameters - - -def test_benchmark_variants_share_unified_generator(): - assert issubclass(_TPCDSRsDataGenerator, _TpcgenRsDataGenerator) - assert issubclass(_TPCHRsDataGenerator, _TpcgenRsDataGenerator) - - -def test_unified_generator_supports_file_uri(tmp_path, fake_executable): - generator = TPCHDataGenerator( - scale_factor=1, - target_folder_uri=tmp_path.as_uri(), - table_list=["region"], - ) - - assert generator.target_folder == tmp_path.resolve() - - -def test_tpch_uses_unified_cli_and_normalizes_outputs(tmp_path, fake_executable): - output_dir = tmp_path / "tpch" - generator = TPCHDataGenerator( - scale_factor=10, - target_folder_uri=str(output_dir), - table_list=["lineitem"], - num_threads=1, - ) - - def create_outputs(args): - output_path = Path(args[args.index("--output-dir") + 1]) - table_dir = output_path / "lineitem" - table_dir.mkdir() - parts = int(args[args.index("--parts") + 1]) - for part_number in range(1, parts + 1): - (table_dir / f"lineitem.{part_number}.parquet").write_bytes(b"parquet") - return subprocess.CompletedProcess(args, 0, "", "") - - generator.cli.run = Mock(side_effect=create_outputs) - generator.run() - - args = generator.cli.run.call_args.args[0] - assert args[:2] == ["tpch", "parquet"] - assert args[args.index("--tables") + 1] == "lineitem" - assert args[args.index("--parts") + 1] == "21" - assert "--compat" not in args - assert args[args.index("--num-threads") + 1] == "1" - assert sorted(path.name for path in (output_dir / "lineitem").glob("*.parquet")) == [ - f"lineitem-{part_number:05d}.zstd.parquet" for part_number in range(1, 22) - ] - - -def test_tpch_snappy_uses_measured_compression(tmp_path, fake_executable): - generator = TPCHDataGenerator( - scale_factor=10, - target_folder_uri=str(tmp_path), - target_row_group_size_mb=64, - compression="SNAPPY", - table_list=["lineitem"], - ) - - command = generator._build_command(tmp_path, ["lineitem"], generator.parts_by_table["lineitem"]) - - assert generator.compression_factors_by_table["lineitem"] == 1.866 - assert generator._estimated_table_size_gib("lineitem") == pytest.approx(1.924, rel=0.002) - assert generator.parts_by_table["lineitem"] == 56 - assert int(command[command.index("--row-group-bytes") + 1]) == round(64 * 1.866 * 1.05 * 1024 * 1024) - - -def test_tpch_single_part_directory_output_is_normalized(tmp_path, fake_executable): - output_dir = tmp_path / "tpch" - generator = TPCHDataGenerator( - scale_factor=0.01, - target_folder_uri=str(output_dir), - table_list=["region"], - ) - - def create_output(args): - table_dir = Path(args[args.index("--output-dir") + 1]) / "region" - table_dir.mkdir() - (table_dir / "region.1.parquet").write_bytes(b"parquet") - return subprocess.CompletedProcess(args, 0, "", "") - - generator.cli.run = Mock(side_effect=create_output) - generator.run() - - assert (output_dir / "region" / "region-00001.zstd.parquet").read_bytes() == b"parquet" - - -def test_automatic_parts_group_tables_by_target_size(monkeypatch, tmp_path, fake_executable): - output_dir = tmp_path / "tpcds" - generator = _TPCDSRsDataGenerator( - scale_factor=10, - target_folder_uri=str(output_dir), - table_list=["reason", "store_sales"], - compression_factor=2.0, - ) - - def create_outputs(args): - staging_dir = Path(args[args.index("--output-dir") + 1]) - tables = args[args.index("--tables") + 1].split(",") - parts = int(args[args.index("--parts") + 1]) - for table in tables: - if parts == 1: - (staging_dir / f"{table}.parquet").write_bytes(b"parquet") - else: - table_dir = staging_dir / table - table_dir.mkdir() - for part_number in range(1, parts + 1): - (table_dir / f"{table}.{part_number}.parquet").write_bytes(b"parquet") - return subprocess.CompletedProcess(args, 0, "", "") - - generator.cli.run = Mock(side_effect=create_outputs) - generator.run() - - assert generator.parts_by_table == {"reason": 1, "store_sales": 9} - assert generator.cli.run.call_count == 2 - commands = [call.args[0] for call in generator.cli.run.call_args_list] - assert { - (command[command.index("--tables") + 1], command[command.index("--parts") + 1]) for command in commands - } == {("reason", "1"), ("store_sales", "9")} - - -def test_default_zstd_row_group_target_uses_per_table_compression(tmp_path, fake_executable): - generator = _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path), - target_row_group_size_mb=64, - table_list=["inventory", "store_sales"], - ) - - grouped = generator._group_tables_by_generation_settings() - - assert set(grouped) == {(1, 1.026), (2, 1.344)} - inventory_command = generator._build_command(tmp_path, ["inventory"], 1) - store_sales_command = generator._build_command(tmp_path, ["store_sales"], 1) - assert int(inventory_command[inventory_command.index("--row-group-bytes") + 1]) == round( - 64 * 1.026 * 1.05 * 1024 * 1024 - ) - assert int(store_sales_command[store_sales_command.index("--row-group-bytes") + 1]) == round( - 64 * 1.344 * 1.05 * 1024 * 1024 - ) - - -def test_snappy_uses_measured_compression_for_row_groups_and_parts(tmp_path, fake_executable): - generator = _TPCDSRsDataGenerator( - scale_factor=10, - target_folder_uri=str(tmp_path), - target_row_group_size_mb=64, - compression="SNAPPY", - table_list=["store_sales"], - ) - - command = generator._build_command(tmp_path, ["store_sales"], generator.parts_by_table["store_sales"]) - - assert generator.compression_factors_by_table["store_sales"] == 1.12 - assert generator._estimated_table_size_gib("store_sales") == pytest.approx(1.140, rel=0.002) - assert generator.parts_by_table["store_sales"] == 33 - assert generator._output_file_name("store_sales", 1) == "store_sales-00001.snappy.parquet" - assert int(command[command.index("--row-group-bytes") + 1]) == round(64 * 1.12 * 1.05 * 1024 * 1024) - - -def test_all_zstd_levels_use_zstd1_measured_compression(tmp_path, fake_executable): - generator = _TPCDSRsDataGenerator( - scale_factor=10, - target_folder_uri=str(tmp_path), - target_row_group_size_mb=64, - compression="ZSTD(9)", - table_list=["store_sales"], - ) - - command = generator._build_command(tmp_path, ["store_sales"], generator.parts_by_table["store_sales"]) - - assert generator.compression_factors_by_table["store_sales"] == 1.344 - assert generator.parts_by_table["store_sales"] == 27 - assert command[command.index("--compression") + 1] == "ZSTD(9)" - assert int(command[command.index("--row-group-bytes") + 1]) == round(64 * 1.344 * 1.05 * 1024 * 1024) - - -@pytest.mark.parametrize( - ("scaled_size_gib", "expected_target_mib"), - [ - (9.999, 128), - (10, 256), - (1023.999, 256), - (1024, 512), - (5119.999, 512), - (5120, 1024), - ], -) -def test_target_file_size_thresholds(tmp_path, fake_executable, scaled_size_gib, expected_target_mib): - generator = _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path), - table_list=["reason"], - ) - - assert generator._target_file_size_mb(scaled_size_gib) == expected_target_mib - - -@pytest.mark.parametrize( - ("table_name", "scale_factor", "expected_parts"), - [ - ("catalog_sales", 3, 3), - ("catalog_sales", 10, 11), - ("store_sales", 3, 4), - ("store_sales", 10, 13), - ("web_sales", 10, 5), - ("web_sales", 20, 10), - ("item", 1000, 1), - ("customer", 1000, 6), - ("customer_address", 1000, 1), - ("customer_demographics", 1000, 1), - ("inventory", 1000, 34), - ], -) -def test_automatic_parts_use_tpcds_row_count_scaling( - tmp_path, - fake_executable, - table_name, - scale_factor, - expected_parts, -): - generator = _TPCDSRsDataGenerator( - scale_factor=scale_factor, - target_folder_uri=str(tmp_path), - table_list=[table_name], - ) - - assert generator.parts_by_table[table_name] == expected_parts - - -def test_small_tables_target_guarded_minimum_file_size(tmp_path, fake_executable): - generator = _TPCDSRsDataGenerator( - scale_factor=1000, - target_folder_uri=str(tmp_path), - table_list=["inventory"], - ) - - estimated_file_size_mb = ( - generator._estimated_table_size_gib("inventory") * 1024 / generator.parts_by_table["inventory"] - ) - - assert generator.parts_by_table["inventory"] == 34 - assert estimated_file_size_mb >= ( - 128 * generator.MINIMUM_FILE_SIZE_RATIO * generator.MINIMUM_FILE_SIZE_ESTIMATION_HEADROOM - ) - - -@pytest.mark.parametrize( - ("table_name", "scale_factor", "expected_size_gib"), - [ - ("inventory", 10, 0.408816), - ("inventory", 1000, 2.404798), - ("customer", 1000, 0.469082), - ("customer_demographics", 1000, 0.00533869), - ], -) -def test_table_size_uses_tpcds_row_count_scaling( - tmp_path, - fake_executable, - table_name, - scale_factor, - expected_size_gib, -): - generator = _TPCDSRsDataGenerator( - scale_factor=scale_factor, - target_folder_uri=str(tmp_path), - table_list=[table_name], - ) - - assert generator._estimated_table_size_gib(table_name) == pytest.approx(expected_size_gib, rel=1e-5) - - -@pytest.mark.parametrize( - ("table_name", "scale_factor", "expected_rows"), - [ - ("customer", 10, 500000), - ("customer", 1000, 12000000), - ("item", 1000, 300000), - ("warehouse", 1000, 20), - ("inventory", 1, 11745000), - ("inventory", 10, 133110000), - ("inventory", 1000, 783000000), - ], -) -def test_tpcds_expected_row_counts_match_scaling_model(table_name, scale_factor, expected_rows): - assert _TPCDSRsDataGenerator._expected_row_count(table_name, scale_factor) == expected_rows - - -def test_row_group_target_is_adjusted_for_compression(tmp_path, fake_executable): - generator = _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path), - target_row_group_size_mb=64, - table_list=["reason"], - compression_factor=3.25, - ) - - command = generator._build_command(tmp_path, ["reason"], 1) - - assert int(command[command.index("--row-group-bytes") + 1]) == round(64 * 3.25 * 1.05 * 1024 * 1024) - - -def test_explicit_num_threads_is_passed_to_cli(tmp_path, fake_executable): - generator = _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path), - table_list=["reason"], - num_threads=8, - ) - - command = generator._build_command(tmp_path, ["reason"], 1) - - assert command[command.index("--num-threads") + 1] == "8" - - -def test_num_threads_defaults_to_available_cpu_count(monkeypatch, tmp_path, fake_executable): - monkeypatch.setattr("lakebench.datagen._tpcgen_rs.os.cpu_count", lambda: 12) - generator = _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path), - table_list=["reason"], - ) - - command = generator._build_command(tmp_path, ["reason"], 1) - - assert command[command.index("--num-threads") + 1] == "12" - - -@pytest.mark.parametrize("num_threads", [0, -1, 1.5, True]) -def test_num_threads_must_be_a_positive_integer(tmp_path, fake_executable, num_threads): - with pytest.raises(ValueError, match="num_threads must be a positive integer"): - _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path), - num_threads=num_threads, - ) - - -def test_non_default_compression_requires_factor(tmp_path, fake_executable): - with pytest.raises(ValueError, match="compression_factor is required"): - _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path), - compression="GZIP(1)", - ) - - -def test_unpartitioned_output_is_normalized_to_table_directory(tmp_path, fake_executable): - output_dir = tmp_path / "tpcds" - generator = _TPCDSRsDataGenerator( - scale_factor=0.1, - target_folder_uri=str(output_dir), - table_list=["reason"], - ) - - def create_output(args): - staging_dir = Path(args[args.index("--output-dir") + 1]) - (staging_dir / "reason.parquet").write_bytes(b"parquet") - return subprocess.CompletedProcess(args, 0, "", "") - - generator.cli.run = Mock(side_effect=create_output) - generator.run() - - assert (output_dir / "reason" / "reason-00001.zstd.parquet").read_bytes() == b"parquet" - - -@pytest.mark.parametrize( - ("kwargs", "message"), - [ - ({"table_list": ["unknown"]}, "Unsupported TPC-DS tables"), - ], -) -def test_rust_generator_validates_options(tmp_path, fake_executable, kwargs, message): - with pytest.raises(ValueError, match=message): - _TPCDSRsDataGenerator( - scale_factor=1, - target_folder_uri=str(tmp_path / "output"), - **kwargs, - ) - - -def test_public_generator_dispatches_to_duckdb(monkeypatch, tmp_path): - duckdb_generator = Mock() - constructor = Mock(return_value=duckdb_generator) - monkeypatch.setattr("lakebench.datagen.tpcds._TPCDSDuckDBDataGenerator", constructor) - - generator = TPCDSDataGenerator(1, str(tmp_path), backend="duckdb") - generator.run() - - constructor.assert_called_once_with( - scale_factor=1, - target_folder_uri=str(tmp_path), - target_row_group_size_mb=128, - ) - duckdb_generator.run.assert_called_once_with() - assert generator.scale_factor == 1 - assert generator.target_folder_uri == str(tmp_path) - assert generator.backend == "duckdb" - - -def test_public_generator_rejects_unknown_backend(tmp_path): - with pytest.raises(ValueError, match="backend"): - TPCDSDataGenerator(1, str(tmp_path), backend="unknown") - - -def test_duckdb_backend_rejects_rust_only_options(tmp_path): - with pytest.raises(ValueError, match="compression, num_threads"): - TPCDSDataGenerator( - 1, - str(tmp_path), - backend="duckdb", - compression="SNAPPY", - num_threads=1, - ) - - -def test_packaged_binary_checksum_validation(tmp_path): - binary = tmp_path / "tpcgen-cli" - binary.write_bytes(b"binary") - checksum = hashlib.sha256(b"binary").hexdigest() - (tmp_path / "provenance.json").write_text( - json.dumps({"binary_sha256": checksum}), - encoding="utf-8", - ) - - TpcgenCli._validate_packaged_binary(binary) - - (tmp_path / "provenance.json").write_text( - json.dumps({"binary_sha256": "0" * 64}), - encoding="utf-8", - ) - with pytest.raises(RuntimeError, match="checksum mismatch"): - TpcgenCli._validate_packaged_binary(binary) - - -def test_packaged_binary_is_made_executable(): - binary = Mock() - binary.stat.return_value.st_mode = stat.S_IRUSR | stat.S_IWUSR - - TpcgenCli._ensure_executable(binary) - - binary.chmod.assert_called_once_with(stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) - - -def test_vendored_binary_checksums_match_provenance(): - bundle_root = Path(__file__).parents[1] / "native" / "tpcgen" - for platform_directory, binary_name in ( - ("windows-x86_64", "tpcgen-cli.exe"), - ("linux-x86_64", "tpcgen-cli"), - ): - binary_path = bundle_root / platform_directory / binary_name - manifest = TpcgenCli._validate_packaged_binary(binary_path) - - assert manifest["commit"] == "eed8a40b6d69b8ed6b5cce88eb3bb6fdcce76c60" - assert manifest["binary_name"] == binary_name - - -def test_cli_surfaces_subprocess_output(monkeypatch, fake_executable): - error = subprocess.CalledProcessError( - 2, - [fake_executable], - output="partial output", - stderr="bad option", - ) - monkeypatch.setattr(subprocess, "run", Mock(side_effect=error)) - - with pytest.raises(RuntimeError, match="partial output") as exc_info: - TpcgenCli().run(["tpcds", "parquet"]) - assert "bad option" in str(exc_info.value) - - -def test_cli_explains_mounted_filesystem_rename_errors(monkeypatch, fake_executable): - error = subprocess.CalledProcessError( - 1, - [fake_executable], - stderr='Failed to rename "part.parquet.inprogress" to "part.parquet": Input/output error (os error 5)', - ) - monkeypatch.setattr(subprocess, "run", Mock(side_effect=error)) - - with pytest.raises(RuntimeError, match="num_threads=8") as exc_info: - TpcgenCli().run(["tpcds", "parquet"]) - - assert "mounted filesystem" in str(exc_info.value) - - -def test_shared_cli_runner_accepts_future_tpch_command(monkeypatch, fake_executable): - run = Mock(return_value=subprocess.CompletedProcess([], 0, "", "")) - monkeypatch.setattr(subprocess, "run", run) - - TpcgenCli().run(["tpch", "parquet", "--scale-factor", "1"]) - - run.assert_called_once_with( - [fake_executable, "tpch", "parquet", "--scale-factor", "1"], - capture_output=True, - text=True, - check=True, - ) diff --git a/tests/test_tpcgen_datagen.py b/tests/test_tpcgen_datagen.py new file mode 100644 index 0000000..ee629e9 --- /dev/null +++ b/tests/test_tpcgen_datagen.py @@ -0,0 +1,190 @@ +import inspect +import os +import subprocess +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from lakebench.datagen import TPCDSDataGenerator, TPCHDataGenerator, _tpcgen +from lakebench.datagen.tpcds import _TPCDSTpcgenDataGenerator + + +@pytest.fixture +def fake_executable(monkeypatch): + monkeypatch.setattr(_tpcgen, "_find_executable", lambda: "tpcgen-cli") + return "tpcgen-cli" + + +@pytest.fixture +def commands(monkeypatch, fake_executable): + run = Mock() + monkeypatch.setattr(_tpcgen, "_run", run) + return lambda: [call.args[0] for call in run.call_args_list] + + +def _arg(command, flag): + return command[command.index(flag) + 1] + + +def test_tpch_parquet_command(tmp_path, commands): + TPCHDataGenerator(scale_factor=10, target_folder_uri=str(tmp_path), table_list=["lineitem"]).run() + + (command,) = commands() + assert command[:3] == ["tpcgen-cli", "tpch", "parquet"] + assert _arg(command, "--scale-factor") == "10" + assert _arg(command, "--output-dir") == str(tmp_path.resolve()) + assert _arg(command, "--tables") == "lineitem" + assert _arg(command, "--parts") == "6" + assert _arg(command, "--compression") == "ZSTD(1)" + assert _arg(command, "--row-group-bytes") == "128MiB" + assert "--num-threads" not in command + + +def test_tpcds_native_command(tmp_path, commands): + TPCDSDataGenerator( + scale_factor=1, + target_folder_uri=str(tmp_path), + table_list=["store_sales"], + num_threads=4, + output_format="native", + ).run() + + (command,) = commands() + assert command[1:3] == ["tpcds", "dat"] + assert _arg(command, "--num-threads") == "4" + assert "--compression" not in command + assert "--row-group-bytes" not in command + + +def test_parquet_options_are_passed_through(tmp_path, commands): + TPCHDataGenerator( + scale_factor=1, + target_folder_uri=str(tmp_path), + table_list=["region"], + compression="SNAPPY", + target_row_group_size_mb=64, + ).run() + + (command,) = commands() + assert _arg(command, "--compression") == "SNAPPY" + assert _arg(command, "--row-group-bytes") == "64MiB" + + +def test_parts_scale_with_table_size(fake_executable, tmp_path): + generator = TPCHDataGenerator(scale_factor=1000, target_folder_uri=str(tmp_path)) + + # 145.7 GB of lineitem at SF1000 in 256 MiB files. + assert generator.parts_by_table["lineitem"] == 583 + assert generator.parts_by_table["nation"] == 1 + + +def test_tables_with_equal_part_counts_share_one_command(tmp_path, commands): + TPCHDataGenerator(scale_factor=1, target_folder_uri=str(tmp_path)).run() + + tables = sorted(table for command in commands() for table in _arg(command, "--tables").split(",")) + assert tables == sorted(TPCHDataGenerator.SF1000_SIZE_GB) + assert len(commands()) == len(set(TPCHDataGenerator(1, str(tmp_path)).parts_by_table.values())) + + +def test_run_clears_previous_table_output(tmp_path, commands): + stale = tmp_path / "region" / "region.2.parquet" + stale.parent.mkdir() + stale.write_bytes(b"old") + + TPCHDataGenerator(scale_factor=1, target_folder_uri=str(tmp_path), table_list=["region"]).run() + + assert not stale.exists() + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"output_format": "csv"}, "output_format"), + ({"output_format": "native", "compression": "SNAPPY"}, "apply only to output_format='parquet'"), + ({"output_format": "native", "target_row_group_size_mb": 64}, "apply only to output_format='parquet'"), + ({"table_list": ["unknown"]}, "Unsupported tpch tables: unknown"), + ({"target_folder_uri": "abfss://ws@onelake/files"}, "local or mounted paths"), + ], +) +def test_invalid_options_are_rejected(tmp_path, fake_executable, kwargs, message): + kwargs = {"scale_factor": 1, "target_folder_uri": str(tmp_path), **kwargs} + with pytest.raises(ValueError, match=message): + TPCHDataGenerator(**kwargs) + + +def test_file_uri_target(tmp_path, fake_executable): + assert TPCHDataGenerator(1, tmp_path.as_uri()).target_folder == tmp_path.resolve() + + +def test_executable_is_found_in_interpreter_scripts_directory(monkeypatch, tmp_path): + executable = tmp_path / ("tpcgen-cli.exe" if os.name == "nt" else "tpcgen-cli") + executable.write_bytes(b"binary") + executable.chmod(0o755) + monkeypatch.setattr(_tpcgen.sysconfig, "get_path", lambda name: str(tmp_path)) + monkeypatch.setenv("PATH", "") + + assert Path(_tpcgen._find_executable()) == executable + + +def test_missing_executable_explains_how_to_install(monkeypatch, tmp_path): + monkeypatch.setattr(_tpcgen.sysconfig, "get_path", lambda name: str(tmp_path)) + monkeypatch.setenv("PATH", "") + + with pytest.raises(ImportError, match=r"lakebench\[tpch_datagen\]"): + _tpcgen._find_executable() + + +def test_cli_failure_surfaces_output(monkeypatch): + error = subprocess.CalledProcessError(2, ["tpcgen-cli"], output="partial output", stderr="bad option") + monkeypatch.setattr(subprocess, "run", Mock(side_effect=error)) + + with pytest.raises(RuntimeError, match="partial output") as exc_info: + _tpcgen._run(["tpcgen-cli"]) + assert "bad option" in str(exc_info.value) + + +def test_cli_rename_failure_suggests_fewer_threads(monkeypatch): + error = subprocess.CalledProcessError(1, ["tpcgen-cli"], stderr='Failed to rename "a.inprogress" to "a"') + monkeypatch.setattr(subprocess, "run", Mock(side_effect=error)) + + with pytest.raises(RuntimeError, match="num_threads=8"): + _tpcgen._run(["tpcgen-cli"]) + + +def test_public_generators_share_parameters(): + tpcds_parameters = list(inspect.signature(TPCDSDataGenerator).parameters) + tpch_parameters = list(inspect.signature(TPCHDataGenerator).parameters) + + assert tpcds_parameters[:6] == tpch_parameters[:6] + assert tpcds_parameters[-1] == tpch_parameters[-1] == "output_format" + + +def test_tpcds_dispatches_to_tpcgen(tmp_path, fake_executable): + generator = TPCDSDataGenerator(1, str(tmp_path)) + + assert isinstance(generator._generator, _TPCDSTpcgenDataGenerator) + + +def test_tpcds_dispatches_to_duckdb(monkeypatch, tmp_path): + duckdb_generator = Mock() + constructor = Mock(return_value=duckdb_generator) + monkeypatch.setattr("lakebench.datagen.tpcds._TPCDSDuckDBDataGenerator", constructor) + + TPCDSDataGenerator(1, str(tmp_path), backend="duckdb").run() + + constructor.assert_called_once_with(scale_factor=1, target_folder_uri=str(tmp_path), target_row_group_size_mb=128) + duckdb_generator.run.assert_called_once_with() + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"backend": "unknown"}, "backend"), + ({"backend": "duckdb", "compression": "SNAPPY", "num_threads": 1}, "compression, num_threads"), + ({"backend": "duckdb", "output_format": "native"}, "supported only by backend='rust'"), + ], +) +def test_tpcds_backend_options_are_validated(tmp_path, kwargs, message): + with pytest.raises(ValueError, match=message): + TPCDSDataGenerator(1, str(tmp_path), **kwargs) diff --git a/uv.lock b/uv.lock index ac0f882..448b48a 100644 --- a/uv.lock +++ b/uv.lock @@ -1470,9 +1470,15 @@ spark = [ sparkmeasure = [ { name = "sparkmeasure", marker = "python_full_version >= '3.9' or (extra == 'extra-9-lakebench-daft' and extra == 'extra-9-lakebench-duckdb') or (extra == 'extra-9-lakebench-daft' and extra == 'extra-9-lakebench-polars') or (extra == 'extra-9-lakebench-daft' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-databricks-spark' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-duckdb' and extra == 'extra-9-lakebench-fabric-data-warehouse') or (extra == 'extra-9-lakebench-duckdb' and extra == 'extra-9-lakebench-fabric-warehouse') or (extra == 'extra-9-lakebench-duckdb' and extra == 'extra-9-lakebench-fabricwarehouse') or (extra == 'extra-9-lakebench-fabric-data-warehouse' and extra == 'extra-9-lakebench-polars') or (extra == 'extra-9-lakebench-fabric-data-warehouse' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-fabric-warehouse' and extra == 'extra-9-lakebench-polars') or (extra == 'extra-9-lakebench-fabric-warehouse' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-fabricwarehouse' and extra == 'extra-9-lakebench-polars') or (extra == 'extra-9-lakebench-fabricwarehouse' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-sail' and extra == 'extra-9-lakebench-spark')" }, ] +tpcds-datagen = [ + { name = "tpcgen-cli" }, +] tpcds-duckdb-datagen = [ { name = "duckdb", marker = "python_full_version >= '3.10' or (extra == 'extra-9-lakebench-daft' and extra == 'extra-9-lakebench-duckdb') or (extra == 'extra-9-lakebench-daft' and extra == 'extra-9-lakebench-polars') or (extra == 'extra-9-lakebench-daft' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-databricks-spark' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-duckdb' and extra == 'extra-9-lakebench-fabric-data-warehouse') or (extra == 'extra-9-lakebench-duckdb' and extra == 'extra-9-lakebench-fabric-warehouse') or (extra == 'extra-9-lakebench-duckdb' and extra == 'extra-9-lakebench-fabricwarehouse') or (extra == 'extra-9-lakebench-fabric-data-warehouse' and extra == 'extra-9-lakebench-polars') or (extra == 'extra-9-lakebench-fabric-data-warehouse' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-fabric-warehouse' and extra == 'extra-9-lakebench-polars') or (extra == 'extra-9-lakebench-fabric-warehouse' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-fabricwarehouse' and extra == 'extra-9-lakebench-polars') or (extra == 'extra-9-lakebench-fabricwarehouse' and extra == 'extra-9-lakebench-sail') or (extra == 'extra-9-lakebench-sail' and extra == 'extra-9-lakebench-spark')" }, ] +tpch-datagen = [ + { name = "tpcgen-cli" }, +] [package.dev-dependencies] dev = [ @@ -1523,6 +1529,8 @@ requires-dist = [ { name = "sqlglot", marker = "python_full_version >= '3.9'", specifier = "==30.18.0" }, { name = "tenacity", marker = "python_full_version < '3.9'", specifier = ">=8.2.3,<9" }, { name = "tenacity", marker = "python_full_version >= '3.9'", specifier = "==9.1.2" }, + { name = "tpcgen-cli", marker = "extra == 'tpcds-datagen'", specifier = "==4.0.0" }, + { name = "tpcgen-cli", marker = "extra == 'tpch-datagen'", specifier = "==4.0.0" }, ] provides-extras = ["duckdb", "polars", "daft", "tpcds-datagen", "tpch-datagen", "tpcds-duckdb-datagen", "sparkmeasure", "spark", "databricks-sql-warehouse", "sail", "fabric-data-warehouse", "fabric-warehouse", "fabricwarehouse"] @@ -3100,6 +3108,27 @@ wheels = [ { url = "https://ms-feed-25.pkgs.visualstudio.com/6f084628-a36d-42cb-934d-057357e379dc/_packaging/49d7402f-07bb-4b18-a9ce-086e6e98a554/pypi/download/tomli/2.4/tomli-2.4.0-py3-none-any.whl", hash = "sha256:1f776e7d669ebceb01dee46484485f43a4048746235e683bcdffacdf1fb4785a", size = 14477, upload-time = "2026-08-12T00:29:03Z" }, ] +[[package]] +name = "tpcgen-cli" +version = "4.0.0" +source = { registry = "https://packagefeedproxy.microsoft.io/pypi/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/93/12/08345fe77dd9c7107a0adf8f2305fc8ca4c09ebb091529150cf5d021d885/tpcgen_cli-4.0.0.tar.gz", hash = "sha256:02eab9277b0c4826155ad20e9b42e25c00e1c9746763743fb4ebd8127d2232aa", size = 602248, upload-time = "2026-10-19T23:32:19Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/22/0e/c59dc9ff91abef24bca40b31757c989690ce4b1af0f536be6f07d0d66291/tpcgen_cli-4.0.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:e011f8a1ecdaf560901fe2d8a795ca00f9e7d3e7d0187247604e59c1e63e240a", size = 5623055, upload-time = "2026-10-19T23:31:52Z" }, + { url = "https://files.pythonhosted.org/packages/00/68/499518dc6c532bdbb8847a793252c84a352b623bd3f16253664cc24e8413/tpcgen_cli-4.0.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:2913936c31b6a28dd171af2e394f73b2c563f75f41154615813ed9b2d3988bf2", size = 5244347, upload-time = "2026-10-19T23:31:54Z" }, + { url = "https://files.pythonhosted.org/packages/80/59/b3518ce4eae1d5bd015508e7488d01c24b031cb4f9371427aa9151193896/tpcgen_cli-4.0.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2fb311272eb71f99156870d2177164823a4a39730982be9bc2448af8da84c96e", size = 5604637, upload-time = "2026-10-19T23:31:57Z" }, + { url = "https://files.pythonhosted.org/packages/6d/fb/d99b72e4537da7213ff2125176b7dba6c625b124ab45115e0b3e20708549/tpcgen_cli-4.0.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:787bd07b5c4433382b57ffa3610d2da3a58849f8bcc3614ac0f43e265dd73d0c", size = 6155287, upload-time = "2026-10-19T23:31:59Z" }, + { url = "https://files.pythonhosted.org/packages/25/09/4d881a9368cb4908c99f748d1707a34f4235e9c665daaee59353d7746976/tpcgen_cli-4.0.0-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4f2f7eaf2899a86cfcc0250a02ad5044b8fc05757ec14af95be991821221a0c0", size = 5909721, upload-time = "2026-10-19T23:32:01Z" }, + { url = "https://files.pythonhosted.org/packages/c4/ed/ea699bc0a57790ebe8ce49ed1c917c536b734d50e3352569758eaba5b04c/tpcgen_cli-4.0.0-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:276b70011b7a4bf223f337598e2b4435bd8eb7a6f5bc4ca798eba2b288574117", size = 6003460, upload-time = "2026-10-19T23:32:03Z" }, + { url = "https://files.pythonhosted.org/packages/32/9a/6ad50479c83e177ae31fb9b78e841aea56ff59bc7e97ba6fc4f0a4e76c1a/tpcgen_cli-4.0.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:490d693dad29ddc4ad6a41f96932a7ddd7b13b04546a54f8dd9c0468648aa17e", size = 5730466, upload-time = "2026-10-19T23:32:05Z" }, + { url = "https://files.pythonhosted.org/packages/57/e3/18d5e3e0fa29f150bc9459786977a5093ef67b99dc330b03252912926c2f/tpcgen_cli-4.0.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:61ab9e058603e2dd16b449b2f6de6515b03d1634539b913f8f27f4165422b017", size = 5609760, upload-time = "2026-10-19T23:32:07Z" }, + { url = "https://files.pythonhosted.org/packages/1b/8c/b2812b34b5f09d286623041649b71cca4f95b65828a715c4653ed1fe5ab9/tpcgen_cli-4.0.0-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:447142b034ee18269b03f84fafd8d501ec3680533048607d833647f9196bc65e", size = 6179679, upload-time = "2026-10-19T23:32:09Z" }, + { url = "https://files.pythonhosted.org/packages/eb/73/566e1ebbff0935a371b0d39b49c82e09dcc97177949705754b1802d6ddb9/tpcgen_cli-4.0.0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:764f4b68b51b99a451f162ad5d2827936bb13afb12875ef3a5f1a2a0697be89f", size = 5833075, upload-time = "2026-10-19T23:32:11Z" }, + { url = "https://files.pythonhosted.org/packages/92/87/4f57e02e702227e48c060577a5bca86f0606294db206ddc855c0eef62f62/tpcgen_cli-4.0.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:f19b43b4435ef084128669589a15d2f86098298c529d352ccf26e701816f7073", size = 5826634, upload-time = "2026-10-19T23:32:13Z" }, + { url = "https://files.pythonhosted.org/packages/d2/25/246c419bc1ebd4027548d8260d6c324112f3f43a2be21858242ae02cf42a/tpcgen_cli-4.0.0-py3-none-win32.whl", hash = "sha256:52a734be9e4fc4cb3d5fbe7653427dd7968add2e6ee92f8a22295367880061d7", size = 6219969, upload-time = "2026-10-19T23:32:15Z" }, + { url = "https://files.pythonhosted.org/packages/bf/3b/fc7e24c40db25a36d93b8b4748631207fc32a792c650d37a73762d31c193/tpcgen_cli-4.0.0-py3-none-win_amd64.whl", hash = "sha256:469d0a5a6797fa6e98d51a12ef7d16afa7b6f9622d3eb4805acc585725fcd25b", size = 6819294, upload-time = "2026-10-19T23:32:17Z" }, +] + [[package]] name = "tqdm" version = "4.67.3"