From 475b7cdb2d2a5e63b0200be14dfceb3ca32604e8 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 8 Jun 2026 20:38:43 -0700 Subject: [PATCH 01/31] Run Polars TPC benchmarks in CI --- .github/workflows/pr.yaml | 18 +++ ci/run_cudf_polars_tpc.sh | 65 +++++++++++ dependencies.yaml | 10 ++ .../cudf_polars/streaming/benchmarks/utils.py | 2 +- python/cudf_polars/pyproject.toml | 4 + python/cudf_polars/tests/conftest.py | 14 +++ .../cudf_polars/tests/streaming/test_tpcds.py | 110 ++++++++++++++++++ .../cudf_polars/tests/streaming/test_tpch.py | 101 ++++++++++++++++ 8 files changed, 323 insertions(+), 1 deletion(-) create mode 100755 ci/run_cudf_polars_tpc.sh create mode 100644 python/cudf_polars/tests/streaming/test_tpcds.py create mode 100644 python/cudf_polars/tests/streaming/test_tpch.py diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 031e26e4a048..dedb9de9b16f 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -35,6 +35,7 @@ jobs: - wheel-build-cudf-polars - wheel-tests-cudf-polars - cudf-polars-polars-tests + - tpc-tests-cudf-polars - wheel-build-dask-cudf - wheel-tests-dask-cudf - devcontainer @@ -678,6 +679,23 @@ jobs: matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))])) build_type: pull-request script: "ci/test_cudf_polars_polars_tests.sh" + tpc-tests-cudf-polars: + needs: [wheel-build-cudf-polars, changed-files] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main + if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels && fromJSON(needs.changed-files.outputs.changed_file_groups).neither_cudf_nor_dask_cudf + with: + # This selects "ARCH=amd64 + the latest supported Python + CUDA". + matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))])) + build_type: pull-request + container-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g --ulimit=nofile=1000000:1000000" + script: "ci/run_cudf_polars_tpc.sh" wheel-build-dask-cudf: needs: wheel-build-cudf permissions: diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh new file mode 100755 index 000000000000..cc9dc734ad1a --- /dev/null +++ b/ci/run_cudf_polars_tpc.sh @@ -0,0 +1,65 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +source rapids-init-pip + +rapids-logger "Download wheels" + +RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")" +CUDF_POLARS_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="cudf_polars_${RAPIDS_PY_CUDA_SUFFIX}" RAPIDS_PY_WHEEL_PURE="1" rapids-download-wheels-from-github python) +LIBCUDF_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) +PYLIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" pylibcudf --stable --cuda "$RAPIDS_CUDA_VERSION")") + +rapids-generate-pip-constraints py_test_cudf_polars "${PIP_CONSTRAINT}" + +rapids-logger "Installing cudf_polars and TPC test dependencies" + +TPCH_REQUIREMENTS=$(mktemp --suffix=.txt) +rapids-dependency-file-generator \ + --config dependencies.yaml \ + --file-key test_cudf_polars_tpch \ + --output requirements \ + > "${TPCH_REQUIREMENTS}" + +rapids-pip-retry install \ + -v \ + --prefer-binary \ + --constraint "${PIP_CONSTRAINT}" \ + "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test]" \ + "$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ + "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ + -r "${TPCH_REQUIREMENTS}" + +rapids-logger "Generating TPC-H data at SF=0.01" + +export TPCH_DATA_DIR +TPCH_DATA_DIR=$(mktemp -d) +tpchgen-cli -s 0.01 --format=parquet --parts=4 --output-dir="${TPCH_DATA_DIR}" + +rapids-logger "Generating TPC-DS data at SF=0.01" + +export TPCDS_DATA_DIR +TPCDS_DATA_DIR=$(mktemp -d) +python3 - <=0.0.0a0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +tpch = [ + "duckdb", + "tpchgen-cli", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. [project.urls] Homepage = "https://github.com/rapidsai/cudf" diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index 14181e45c46c..89ee56233734 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -338,6 +338,20 @@ def timeout_seconds() -> int: return 30 +def pytest_addoption(parser: pytest.Parser) -> None: + parser.addoption( + "--iterations", + type=int, + default=1, + help="Number of times to collect each TPC query result and validate it.", + ) + + +@pytest.fixture(scope="session") +def tpc_iterations(request: pytest.FixtureRequest) -> int: + return request.config.getoption("--iterations") + + def pytest_configure(config: pytest.Config): config.addinivalue_line( "markers", diff --git a/python/cudf_polars/tests/streaming/test_tpcds.py b/python/cudf_polars/tests/streaming/test_tpcds.py new file mode 100644 index 000000000000..202b3ecaf412 --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_tpcds.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""TPC-DS validation tests for the cudf-polars rapidsmpf streaming engine. + +Data is expected at the path given by the TPCDS_DATA_DIR environment variable, +with one parquet file per table named .parquet. + +Each query is run with the streaming engine and validated against DuckDB on +the same data. Qualification parameters (TPC-DS specification Appendix B) are +used so query parameters are fixed and independent of scale factor. +""" + +from __future__ import annotations + +import os +import types +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Generator + +import pytest + +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.spmd import SPMDEngine +from cudf_polars.streaming.benchmarks.asserts import assert_tpch_result_equal +from cudf_polars.streaming.benchmarks.pdsds import ( + PDSDSDuckDBQueries, + PDSDSPolarsQueries, +) +from cudf_polars.streaming.benchmarks.utils import ( + POLARS_VALIDATION_OPTIONS, + execute_duckdb_query, +) + +QUERIES = PDSDSPolarsQueries() +DUCKDB_QUERIES = PDSDSDuckDBQueries() +TPCDS_TARGET_PARTITION_SIZE = 2_000_000 + + +@pytest.fixture(scope="session") +def tpcds_data_dir() -> Path: + data_dir = os.environ.get("TPCDS_DATA_DIR") + if data_dir is None: + pytest.skip("TPCDS_DATA_DIR environment variable not set") + return Path(data_dir) + + +@pytest.fixture(scope="session") +def tpcds_run_config(tpcds_data_dir: Path) -> types.SimpleNamespace: + return types.SimpleNamespace( + dataset_path=tpcds_data_dir, + suffix=".parquet", + query_set="pdsds", + scale_factor=1, + qualification=True, + ) + + +@pytest.fixture(scope="session") +def tpcds_engine() -> Generator[SPMDEngine, None, None]: + options = StreamingOptions( + target_partition_size=TPCDS_TARGET_PARTITION_SIZE, + raise_on_fail=True, + fallback_mode="raise", + allow_gpu_sharing=True, + ) + with SPMDEngine.from_options(options) as engine: + yield engine + + +@pytest.mark.parametrize( + "query_id", range(1, 100), ids=[f"q{i}" for i in range(1, 100)] +) +def test_tpcds( + query_id: int, + tpcds_run_config: types.SimpleNamespace, + tpcds_engine: SPMDEngine, + tpc_iterations: int, +) -> None: + query_result = getattr(QUERIES, f"q{query_id}")(tpcds_run_config) + + sql = getattr(DUCKDB_QUERIES, f"q{query_id}")(tpcds_run_config) + expected = execute_duckdb_query( + sql, + tpcds_run_config.dataset_path, + suffix=tpcds_run_config.suffix, + query_set="pdsds", + ) + + casts = [ + *QUERIES.EXPECTED_CASTS.get(query_id, []), + *QUERIES.EXPECTED_CASTS_DECIMAL.get(query_id, []), + ] + if casts: + expected = expected.with_columns(*casts) + + for _ in range(tpc_iterations): + gpu_result = query_result.frame.collect(engine=tpcds_engine) + assert_tpch_result_equal( + gpu_result, + expected, + sort_by=query_result.sort_by, + limit=query_result.limit, + nulls_last=query_result.nulls_last, + sort_keys=query_result.sort_keys, + **POLARS_VALIDATION_OPTIONS, + ) diff --git a/python/cudf_polars/tests/streaming/test_tpch.py b/python/cudf_polars/tests/streaming/test_tpch.py new file mode 100644 index 000000000000..70a988c225ce --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_tpch.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""TPC-H validation tests for the cudf-polars rapidsmpf streaming engine. + +Data is expected at the path given by the TPCH_DATA_DIR environment variable. +Generate it with tpchgen-cli before running: + + tpchgen-cli -s 0.01 --format=parquet --parts=4 --output-dir=/tmp/tpch + +Each query is run with the streaming engine and validated against DuckDB on +the same data. +""" + +from __future__ import annotations + +import os +import types +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Generator + +import pytest + +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.spmd import SPMDEngine +from cudf_polars.streaming.benchmarks.asserts import assert_tpch_result_equal +from cudf_polars.streaming.benchmarks.pdsh import PDSHDuckDBQueries, PDSHQueries +from cudf_polars.streaming.benchmarks.utils import ( + POLARS_VALIDATION_OPTIONS, + execute_duckdb_query, +) + +QUERIES = PDSHQueries() +DUCKDB_QUERIES = PDSHDuckDBQueries() +TPCH_TARGET_PARTITION_SIZE = 2_000_000 + + +@pytest.fixture(scope="session") +def tpch_data_dir() -> Path: + data_dir = os.environ.get("TPCH_DATA_DIR") + if data_dir is None: + pytest.skip("TPCH_DATA_DIR environment variable not set") + return Path(data_dir) + + +@pytest.fixture(scope="session") +def tpch_run_config(tpch_data_dir: Path) -> types.SimpleNamespace: + return types.SimpleNamespace( + dataset_path=tpch_data_dir, + suffix="/*.parquet", + query_set="pdsh", + ) + + +@pytest.fixture(scope="session") +def tpch_engine() -> Generator[SPMDEngine, None, None]: + options = StreamingOptions( + target_partition_size=TPCH_TARGET_PARTITION_SIZE, + raise_on_fail=True, + fallback_mode="raise", + allow_gpu_sharing=True, + ) + with SPMDEngine.from_options(options) as engine: + yield engine + + +@pytest.mark.parametrize("query_id", range(1, 23), ids=[f"q{i}" for i in range(1, 23)]) +def test_tpch( + query_id: int, + tpch_run_config: types.SimpleNamespace, + tpch_engine: SPMDEngine, + tpc_iterations: int, +) -> None: + query_result = getattr(QUERIES, f"q{query_id}")(tpch_run_config) + + sql = getattr(DUCKDB_QUERIES, f"q{query_id}")(tpch_run_config) + expected = execute_duckdb_query( + sql, tpch_run_config.dataset_path, suffix=tpch_run_config.suffix + ) + + casts = [ + *QUERIES.EXPECTED_CASTS.get(query_id, []), + *QUERIES.EXPECTED_CASTS_DECIMAL.get(query_id, []), + ] + if casts: + expected = expected.with_columns(*casts) + + for _ in range(tpc_iterations): + gpu_result = query_result.frame.collect(engine=tpch_engine) + assert_tpch_result_equal( + gpu_result, + expected, + sort_by=query_result.sort_by, + limit=query_result.limit, + nulls_last=query_result.nulls_last, + sort_keys=query_result.sort_keys, + **POLARS_VALIDATION_OPTIONS, + ) From c72714b2f70f01b4930051b8bc586a73fd7bb6ff Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 25 Jun 2026 19:37:21 +0000 Subject: [PATCH 02/31] move tpcds gen script to python, use existing pdsh and pdsds python scripts --- ci/run_cudf_polars_tpc.sh | 36 +++--- .../cudf_polars/tests/streaming/test_tpcds.py | 110 ------------------ .../cudf_polars/tests/streaming/test_tpch.py | 101 ---------------- 3 files changed, 19 insertions(+), 228 deletions(-) delete mode 100644 python/cudf_polars/tests/streaming/test_tpcds.py delete mode 100644 python/cudf_polars/tests/streaming/test_tpch.py diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index cc9dc734ad1a..8028f3df0c54 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -43,23 +43,25 @@ rapids-logger "Generating TPC-DS data at SF=0.01" export TPCDS_DATA_DIR TPCDS_DATA_DIR=$(mktemp -d) -python3 - <.parquet. - -Each query is run with the streaming engine and validated against DuckDB on -the same data. Qualification parameters (TPC-DS specification Appendix B) are -used so query parameters are fixed and independent of scale factor. -""" - -from __future__ import annotations - -import os -import types -from pathlib import Path -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from collections.abc import Generator - -import pytest - -from cudf_polars.engine.options import StreamingOptions -from cudf_polars.engine.spmd import SPMDEngine -from cudf_polars.streaming.benchmarks.asserts import assert_tpch_result_equal -from cudf_polars.streaming.benchmarks.pdsds import ( - PDSDSDuckDBQueries, - PDSDSPolarsQueries, -) -from cudf_polars.streaming.benchmarks.utils import ( - POLARS_VALIDATION_OPTIONS, - execute_duckdb_query, -) - -QUERIES = PDSDSPolarsQueries() -DUCKDB_QUERIES = PDSDSDuckDBQueries() -TPCDS_TARGET_PARTITION_SIZE = 2_000_000 - - -@pytest.fixture(scope="session") -def tpcds_data_dir() -> Path: - data_dir = os.environ.get("TPCDS_DATA_DIR") - if data_dir is None: - pytest.skip("TPCDS_DATA_DIR environment variable not set") - return Path(data_dir) - - -@pytest.fixture(scope="session") -def tpcds_run_config(tpcds_data_dir: Path) -> types.SimpleNamespace: - return types.SimpleNamespace( - dataset_path=tpcds_data_dir, - suffix=".parquet", - query_set="pdsds", - scale_factor=1, - qualification=True, - ) - - -@pytest.fixture(scope="session") -def tpcds_engine() -> Generator[SPMDEngine, None, None]: - options = StreamingOptions( - target_partition_size=TPCDS_TARGET_PARTITION_SIZE, - raise_on_fail=True, - fallback_mode="raise", - allow_gpu_sharing=True, - ) - with SPMDEngine.from_options(options) as engine: - yield engine - - -@pytest.mark.parametrize( - "query_id", range(1, 100), ids=[f"q{i}" for i in range(1, 100)] -) -def test_tpcds( - query_id: int, - tpcds_run_config: types.SimpleNamespace, - tpcds_engine: SPMDEngine, - tpc_iterations: int, -) -> None: - query_result = getattr(QUERIES, f"q{query_id}")(tpcds_run_config) - - sql = getattr(DUCKDB_QUERIES, f"q{query_id}")(tpcds_run_config) - expected = execute_duckdb_query( - sql, - tpcds_run_config.dataset_path, - suffix=tpcds_run_config.suffix, - query_set="pdsds", - ) - - casts = [ - *QUERIES.EXPECTED_CASTS.get(query_id, []), - *QUERIES.EXPECTED_CASTS_DECIMAL.get(query_id, []), - ] - if casts: - expected = expected.with_columns(*casts) - - for _ in range(tpc_iterations): - gpu_result = query_result.frame.collect(engine=tpcds_engine) - assert_tpch_result_equal( - gpu_result, - expected, - sort_by=query_result.sort_by, - limit=query_result.limit, - nulls_last=query_result.nulls_last, - sort_keys=query_result.sort_keys, - **POLARS_VALIDATION_OPTIONS, - ) diff --git a/python/cudf_polars/tests/streaming/test_tpch.py b/python/cudf_polars/tests/streaming/test_tpch.py deleted file mode 100644 index 70a988c225ce..000000000000 --- a/python/cudf_polars/tests/streaming/test_tpch.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -"""TPC-H validation tests for the cudf-polars rapidsmpf streaming engine. - -Data is expected at the path given by the TPCH_DATA_DIR environment variable. -Generate it with tpchgen-cli before running: - - tpchgen-cli -s 0.01 --format=parquet --parts=4 --output-dir=/tmp/tpch - -Each query is run with the streaming engine and validated against DuckDB on -the same data. -""" - -from __future__ import annotations - -import os -import types -from pathlib import Path -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from collections.abc import Generator - -import pytest - -from cudf_polars.engine.options import StreamingOptions -from cudf_polars.engine.spmd import SPMDEngine -from cudf_polars.streaming.benchmarks.asserts import assert_tpch_result_equal -from cudf_polars.streaming.benchmarks.pdsh import PDSHDuckDBQueries, PDSHQueries -from cudf_polars.streaming.benchmarks.utils import ( - POLARS_VALIDATION_OPTIONS, - execute_duckdb_query, -) - -QUERIES = PDSHQueries() -DUCKDB_QUERIES = PDSHDuckDBQueries() -TPCH_TARGET_PARTITION_SIZE = 2_000_000 - - -@pytest.fixture(scope="session") -def tpch_data_dir() -> Path: - data_dir = os.environ.get("TPCH_DATA_DIR") - if data_dir is None: - pytest.skip("TPCH_DATA_DIR environment variable not set") - return Path(data_dir) - - -@pytest.fixture(scope="session") -def tpch_run_config(tpch_data_dir: Path) -> types.SimpleNamespace: - return types.SimpleNamespace( - dataset_path=tpch_data_dir, - suffix="/*.parquet", - query_set="pdsh", - ) - - -@pytest.fixture(scope="session") -def tpch_engine() -> Generator[SPMDEngine, None, None]: - options = StreamingOptions( - target_partition_size=TPCH_TARGET_PARTITION_SIZE, - raise_on_fail=True, - fallback_mode="raise", - allow_gpu_sharing=True, - ) - with SPMDEngine.from_options(options) as engine: - yield engine - - -@pytest.mark.parametrize("query_id", range(1, 23), ids=[f"q{i}" for i in range(1, 23)]) -def test_tpch( - query_id: int, - tpch_run_config: types.SimpleNamespace, - tpch_engine: SPMDEngine, - tpc_iterations: int, -) -> None: - query_result = getattr(QUERIES, f"q{query_id}")(tpch_run_config) - - sql = getattr(DUCKDB_QUERIES, f"q{query_id}")(tpch_run_config) - expected = execute_duckdb_query( - sql, tpch_run_config.dataset_path, suffix=tpch_run_config.suffix - ) - - casts = [ - *QUERIES.EXPECTED_CASTS.get(query_id, []), - *QUERIES.EXPECTED_CASTS_DECIMAL.get(query_id, []), - ] - if casts: - expected = expected.with_columns(*casts) - - for _ in range(tpc_iterations): - gpu_result = query_result.frame.collect(engine=tpch_engine) - assert_tpch_result_equal( - gpu_result, - expected, - sort_by=query_result.sort_by, - limit=query_result.limit, - nulls_last=query_result.nulls_last, - sort_keys=query_result.sort_keys, - **POLARS_VALIDATION_OPTIONS, - ) From 4dc4c165a14f532992db8c5dd1ef0517b6677c82 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 25 Jun 2026 20:03:13 +0000 Subject: [PATCH 03/31] copyright --- python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py | 2 +- python/cudf_polars/tests/conftest.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 973c755e137e..a4048cb3846d 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Benchmark utilities for the RapidsMPF SPMD and Ray frontends.""" diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index 89ee56233734..d2209a377bc6 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations From 79dc8b89e472d3d7979146455889fbb3c7409869 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 2 Jul 2026 20:29:22 +0000 Subject: [PATCH 04/31] update artifact names --- .github/workflows/pr.yaml | 2 +- ci/run_cudf_polars_tpc.sh | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 1d2e385ba65c..353270e1cb27 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -698,7 +698,7 @@ jobs: build_type: pull-request script: "ci/test_cudf_polars_polars_tests.sh" tpc-tests-cudf-polars: - needs: [wheel-build-cudf-polars, changed-files] + needs: [wheel-build-cudf-polars, wheel-build-cudf-streaming, changed-files] permissions: actions: read contents: read diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index 8028f3df0c54..4826042f78f8 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -9,9 +9,11 @@ source rapids-init-pip rapids-logger "Download wheels" RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")" -CUDF_POLARS_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="cudf_polars_${RAPIDS_PY_CUDA_SUFFIX}" RAPIDS_PY_WHEEL_PURE="1" rapids-download-wheels-from-github python) -LIBCUDF_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) -PYLIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" pylibcudf --stable --cuda "$RAPIDS_CUDA_VERSION")") +LIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-artifact-name wheel_cpp libcudf cudf --cuda "$RAPIDS_CUDA_VERSION")") +PYLIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-artifact-name wheel_python pylibcudf cudf --stable --cuda "$RAPIDS_CUDA_VERSION")") +CUDF_POLARS_WHEELHOUSE=$(rapids-download-from-github "$(rapids-artifact-name wheel_python cudf-polars cudf --pure --arch any --cuda "$RAPIDS_CUDA_VERSION")") +LIBCUDF_STREAMING_WHEELHOUSE=$(rapids-download-from-github "$(rapids-artifact-name wheel_cpp libcudf-streaming cudf --cuda "$RAPIDS_CUDA_VERSION")") +CUDF_STREAMING_WHEELHOUSE=$(rapids-download-from-github "$(rapids-artifact-name wheel_python cudf-streaming cudf --stable --cuda "$RAPIDS_CUDA_VERSION")") rapids-generate-pip-constraints py_test_cudf_polars "${PIP_CONSTRAINT}" @@ -31,6 +33,8 @@ rapids-pip-retry install \ "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test]" \ "$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ + "$(echo "${LIBCUDF_STREAMING_WHEELHOUSE}"/libcudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ + "$(echo "${CUDF_STREAMING_WHEELHOUSE}"/cudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ -r "${TPCH_REQUIREMENTS}" rapids-logger "Generating TPC-H data at SF=0.01" From 63dda74580a45d969f6f2cd22f3adb8cedb0f849 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 6 Jul 2026 10:51:14 +0000 Subject: [PATCH 05/31] add --matrix --- ci/run_cudf_polars_tpc.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index 4826042f78f8..28076b6f7100 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -24,6 +24,7 @@ rapids-dependency-file-generator \ --config dependencies.yaml \ --file-key test_cudf_polars_tpch \ --output requirements \ + --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \ > "${TPCH_REQUIREMENTS}" rapids-pip-retry install \ From 8a333f874d48c68c22f75d5eae68e387ba6eef9f Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 6 Jul 2026 12:50:48 +0000 Subject: [PATCH 06/31] add tpcds datagen file --- ci/generate_tpcds_data.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 ci/generate_tpcds_data.py diff --git a/ci/generate_tpcds_data.py b/ci/generate_tpcds_data.py new file mode 100644 index 000000000000..2aeb5bc93301 --- /dev/null +++ b/ci/generate_tpcds_data.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate TPC-DS data at a given scale factor using DuckDB.""" + +from __future__ import annotations + +import argparse +import os + +import duckdb + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--scale", type=float, default=0.01, help="Scale factor." + ) + parser.add_argument( + "--output-dir", + default=os.environ.get("TPCDS_DATA_DIR"), + help="Output directory. Defaults to TPCDS_DATA_DIR environment variable.", + ) + args = parser.parse_args() + + if args.output_dir is None: + parser.error("--output-dir is required (or set TPCDS_DATA_DIR).") + + conn = duckdb.connect() + conn.execute(f"INSTALL tpcds; LOAD tpcds; CALL dsdgen(sf={args.scale});") + for table in conn.execute("SHOW TABLES").df()["name"]: + conn.execute( + f"COPY {table} TO '{args.output_dir}/{table}.parquet' (FORMAT PARQUET)" + ) + + +if __name__ == "__main__": + main() From 87c8460736a06225fbf66efe0f93a837eefb21ba Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 6 Jul 2026 12:53:38 +0000 Subject: [PATCH 07/31] tpcgen-cli cmd clap sub-cmd --- ci/run_cudf_polars_tpc.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index 28076b6f7100..1445f3c2f994 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -42,7 +42,7 @@ rapids-logger "Generating TPC-H data at SF=0.01" export TPCH_DATA_DIR TPCH_DATA_DIR=$(mktemp -d) -tpchgen-cli -s 0.01 --format=parquet --parts=4 --output-dir="${TPCH_DATA_DIR}" +tpchgen-cli parquet -s 0.01 --parts=4 --output-dir="${TPCH_DATA_DIR}" rapids-logger "Generating TPC-DS data at SF=0.01" From 22b4131cdec5564733c6500f51f7870e37a2957b Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 6 Jul 2026 18:58:28 +0000 Subject: [PATCH 08/31] remove implicit pyarrow dependency, fix SQL Q64 --- .../streaming/benchmarks/pdsds_queries/q64.py | 12 ++++++------ .../cudf_polars/streaming/benchmarks/utils.py | 7 ++++++- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py index b7046d2e65ac..659eb4a9ce52 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 64.""" @@ -134,11 +134,11 @@ def duckdb_impl(run_config: RunConfig) -> str: cs1.s1, cs1.s2, cs1.s3, - cs2.s1, - cs2.s2, - cs2.s3, - cs2.syear, - cs2.cnt + cs2.s1 AS s1_1, + cs2.s2 AS s2_1, + cs2.s3 AS s3_1, + cs2.syear AS syear_1, + cs2.cnt AS cnt_1 FROM cross_sales cs1, cross_sales cs2 WHERE cs1.item_sk = cs2.item_sk diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 72c927afb428..3c3382b4383f 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -1216,6 +1216,11 @@ def _finalize_benchmark_run( f"{len(validation_failures)} queries failed validation: " f"{sorted(set(validation_failures))}" ) + elif query_failures: + print( + f"⚠️ {len({q for q, _ in query_failures})} queries failed to run; " + "validation was skipped." + ) else: print("✅ All validated queries passed.") args.output.write(json.dumps(run_config.serialize(engine=engine))) @@ -1710,7 +1715,7 @@ def execute_duckdb_query( f"CREATE OR REPLACE VIEW {name} AS " f"SELECT * FROM parquet_scan('{pattern}');" ) - return conn.execute(query).pl() + return pl.from_arrow(conn.sql(query)) def run_duckdb(duckdb_queries_cls: Any, args: argparse.Namespace) -> None: From 6cdfe356d96bb296bd0a8eeca4adf265148154b8 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 6 Jul 2026 19:24:55 +0000 Subject: [PATCH 09/31] add TODO about tpcds gen --- ci/generate_tpcds_data.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ci/generate_tpcds_data.py b/ci/generate_tpcds_data.py index 2aeb5bc93301..db55034977b0 100644 --- a/ci/generate_tpcds_data.py +++ b/ci/generate_tpcds_data.py @@ -26,6 +26,7 @@ def main() -> None: if args.output_dir is None: parser.error("--output-dir is required (or set TPCDS_DATA_DIR).") + # TODO: switch to the Rust TPC-DS generator conn = duckdb.connect() conn.execute(f"INSTALL tpcds; LOAD tpcds; CALL dsdgen(sf={args.scale});") for table in conn.execute("SHOW TABLES").df()["name"]: From 4b444a51e9e683cf7dc04971bed334c1d6b42569 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 6 Jul 2026 19:28:00 +0000 Subject: [PATCH 10/31] mypy --- python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 3c3382b4383f..61b6df5249a5 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -1715,7 +1715,9 @@ def execute_duckdb_query( f"CREATE OR REPLACE VIEW {name} AS " f"SELECT * FROM parquet_scan('{pattern}');" ) - return pl.from_arrow(conn.sql(query)) + result = pl.from_arrow(conn.sql(query)) + assert isinstance(result, pl.DataFrame) + return result def run_duckdb(duckdb_queries_cls: Any, args: argparse.Namespace) -> None: From aedb32ac2e4d382acc6851928922c9b2e60abf28 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 6 Jul 2026 22:11:03 +0000 Subject: [PATCH 11/31] increase verbosity, switch to SF1 --- ci/run_cudf_polars_tpc.sh | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index 1445f3c2f994..4bf6133d9a94 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -38,11 +38,11 @@ rapids-pip-retry install \ "$(echo "${CUDF_STREAMING_WHEELHOUSE}"/cudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ -r "${TPCH_REQUIREMENTS}" -rapids-logger "Generating TPC-H data at SF=0.01" +rapids-logger "Generating TPC-H data at SF=1" export TPCH_DATA_DIR TPCH_DATA_DIR=$(mktemp -d) -tpchgen-cli parquet -s 0.01 --parts=4 --output-dir="${TPCH_DATA_DIR}" +tpchgen-cli parquet -s 1 --parts=4 --output-dir="${TPCH_DATA_DIR}" rapids-logger "Generating TPC-DS data at SF=0.01" @@ -59,7 +59,14 @@ python -m cudf_polars.streaming.benchmarks.pdsh all \ --suffix "/*.parquet" \ --frontend spmd \ --validate-against duckdb \ - --iterations 2 + --iterations 2 \ + --debug \ + --print-results \ + --explain \ + --explain-logical \ + --explain-partition-plan \ + --rapidsmpf-log DEBUG \ + --rapidsmpf-statistics rapids-logger "Running TPC-DS validation tests" From 75b3d51e06ddcbb50ef3ad4da4b3fc8889ba50a7 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 8 Jul 2026 15:58:01 +0000 Subject: [PATCH 12/31] fixes for decimals --- ci/run_cudf_polars_tpc.sh | 15 ++++-- dependencies.yaml | 1 + .../cudf_polars/streaming/benchmarks/pdsds.py | 6 ++- .../benchmarks/pdsds_queries/__init__.py | 25 ++++++++- .../streaming/benchmarks/pdsds_queries/q1.py | 12 ++--- .../streaming/benchmarks/pdsds_queries/q12.py | 17 ++---- .../streaming/benchmarks/pdsds_queries/q13.py | 9 ++-- .../streaming/benchmarks/pdsds_queries/q18.py | 20 +++---- .../streaming/benchmarks/pdsds_queries/q2.py | 13 ++--- .../streaming/benchmarks/pdsds_queries/q24.py | 17 ++---- .../streaming/benchmarks/pdsds_queries/q31.py | 8 +-- .../streaming/benchmarks/pdsds_queries/q32.py | 7 ++- .../streaming/benchmarks/pdsds_queries/q35.py | 19 ++++--- .../streaming/benchmarks/pdsds_queries/q48.py | 12 ++--- .../streaming/benchmarks/pdsds_queries/q49.py | 47 ++++++++++++++--- .../streaming/benchmarks/pdsds_queries/q5.py | 38 +++++++++----- .../streaming/benchmarks/pdsds_queries/q51.py | 33 ++++++------ .../streaming/benchmarks/pdsds_queries/q56.py | 21 ++------ .../streaming/benchmarks/pdsds_queries/q59.py | 52 ++++++------------- .../streaming/benchmarks/pdsds_queries/q64.py | 18 ++----- .../streaming/benchmarks/pdsds_queries/q67.py | 4 +- .../streaming/benchmarks/pdsds_queries/q72.py | 11 ++-- .../streaming/benchmarks/pdsds_queries/q76.py | 5 +- .../streaming/benchmarks/pdsds_queries/q78.py | 48 ++++------------- .../streaming/benchmarks/pdsds_queries/q92.py | 16 +++--- .../cudf_polars/streaming/benchmarks/utils.py | 2 +- 26 files changed, 229 insertions(+), 247 deletions(-) diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index 4bf6133d9a94..ca8b6d2487f4 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -44,11 +44,20 @@ export TPCH_DATA_DIR TPCH_DATA_DIR=$(mktemp -d) tpchgen-cli parquet -s 1 --parts=4 --output-dir="${TPCH_DATA_DIR}" -rapids-logger "Generating TPC-DS data at SF=0.01" +rapids-logger "Generating TPC-DS data at SF=1" export TPCDS_DATA_DIR TPCDS_DATA_DIR=$(mktemp -d) -python3 "$(dirname "$0")/generate_tpcds_data.py" --scale 0.01 --output-dir "${TPCDS_DATA_DIR}" +python3 "$(dirname "$0")/generate_tpcds_data.py" --scale 1 --output-dir "${TPCDS_DATA_DIR}" + +# Blackwell GPUs (compute capability >= 10.0) have decimal overflow issues in GPU +# aggregations; pre-convert decimals to float to work around rapidsai/cudf#23150. +COMPUTE_CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d '.') +if [[ "${COMPUTE_CAP}" -ge 100 ]]; then + rapids-logger "Blackwell GPU detected (sm_${COMPUTE_CAP}): converting decimals to float" + python3 "$(dirname "$0")/convert_tpc_decimals.py" --data-dir "${TPCH_DATA_DIR}" + python3 "$(dirname "$0")/convert_tpc_decimals.py" --data-dir "${TPCDS_DATA_DIR}" +fi rapids-logger "Running TPC-H validation tests" @@ -72,7 +81,7 @@ rapids-logger "Running TPC-DS validation tests" python -m cudf_polars.streaming.benchmarks.pdsds all \ --path "${TPCDS_DATA_DIR}" \ - --scale 0.01 \ + --scale 1 \ --qualification \ --frontend spmd \ --validate-against duckdb \ diff --git a/dependencies.yaml b/dependencies.yaml index b2373b7a2802..393e1d55fd7a 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -1188,6 +1188,7 @@ dependencies: - output_types: [conda, requirements, pyproject] packages: - duckdb + - pyarrow - tpchgen-cli test_python_narwhals: common: diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index 19ac159b76e0..dd45544b0133 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -112,7 +112,6 @@ class PDSDSPolarsQueries(PDSDSQueries): ], 19: [pl.col("ext_price").cast(pl.Decimal(18, 2))], 20: [ - pl.col("itemrevenue").cast(pl.Decimal(18, 2)), pl.col("revenueratio").cast(pl.Decimal(38, 2)), ], 24: [pl.col("paid").cast(pl.Decimal(18, 2))], @@ -239,6 +238,11 @@ class PDSDSPolarsQueries(PDSDSQueries): pl.col("inv_before").cast(pl.Int32), pl.col("inv_after").cast(pl.Int32), ], + 29: [ + pl.col("store_sales_quantity").cast(pl.Int64), + pl.col("store_returns_quantity").cast(pl.Int64), + pl.col("catalog_sales_quantity").cast(pl.Int64), + ], 34: [pl.col("cnt").cast(COUNT_DTYPE)], 35: [ pl.col("cnt1").cast(COUNT_DTYPE), diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/__init__.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/__init__.py index ae37fc2b5098..e5da3f49c8d6 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/__init__.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/__init__.py @@ -1,4 +1,27 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """DuckDB and Polars queries.""" + +from __future__ import annotations + +import polars as pl + + +def sql_sum(expr: str | pl.Expr) -> pl.Expr: + """ + Sum that returns NULL for all-null/empty groups, matching SQL SUM semantics. + + Polars sum() returns 0 for all-null or empty groups; SQL returns NULL. + See https://github.com/rapidsai/cudf/issues/19560. + + Parameters + ---------- + expr + Column name or expression to sum. If a string, wraps in ``pl.col``. + Pass a conditional expression (e.g. ``pl.when(...).then(...).otherwise(None)``) + to implement SQL ``SUM(CASE WHEN ... END)`` without ``.filter()`` inside + a groupby, which is not supported on GPU. + """ + e = pl.col(expr) if isinstance(expr, str) else expr + return pl.when(e.count() > 0).then(e.sum()).otherwise(None) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q1.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q1.py index a5aebfa87dda..87d1a2315e32 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q1.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q1.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 1.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -74,14 +75,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ) .filter(pl.col("d_year") == year) .group_by(["sr_customer_sk", "sr_store_sk"]) - .agg( - # Polars sum() returns 0 for all-null groups; SQL returns NULL. - # See https://github.com/rapidsai/cudf/issues/19560. - pl.when(pl.col("sr_return_amt").count() > 0) - .then(pl.col("sr_return_amt").sum()) - .otherwise(None) - .alias("ctr_total_return") - ) + .agg(sql_sum("sr_return_amt").alias("ctr_total_return")) .rename( { "sr_customer_sk": "ctr_customer_sk", diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q12.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q12.py index 8ca5b6cbeb50..7b9b016b2c3c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q12.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q12.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 12.""" @@ -11,6 +11,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -95,23 +96,13 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .group_by( ["i_item_id", "i_item_desc", "i_category", "i_class", "i_current_price"] ) - .agg( - [ - pl.when(pl.col("ws_ext_sales_price").count() > 0) - .then(pl.col("ws_ext_sales_price").sum()) - .otherwise(None) - .alias("itemrevenue") - ] - ) + .agg([sql_sum("ws_ext_sales_price").alias("itemrevenue")]) .with_columns( [ ( pl.col("itemrevenue") * 100 - / pl.when(pl.col("itemrevenue").count() > 0) - .then(pl.col("itemrevenue").sum()) - .otherwise(None) - .over("i_class") + / sql_sum("itemrevenue").over("i_class") ).alias("revenueratio") ] ) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q13.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q13.py index e7c0366aab67..a100fb7f8e19 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q13.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q13.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 13.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -158,9 +159,9 @@ def polars_impl(run_config: RunConfig) -> QueryResult: pl.col("ss_ext_wholesale_cost") .mean() .alias("avg(ss_ext_wholesale_cost)"), - pl.col("ss_ext_wholesale_cost") - .sum() - .alias("sum(ss_ext_wholesale_cost)"), + sql_sum("ss_ext_wholesale_cost").alias( + "sum(ss_ext_wholesale_cost)" + ), ] ) ), diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q18.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q18.py index 9c0784cbd052..4ba3cea66f58 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q18.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q18.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 18.""" @@ -72,7 +72,6 @@ def duckdb_impl(run_config: RunConfig) -> str: def level( # noqa: D103 base_query: pl.LazyFrame, agg_exprs: list[pl.Expr], - null_sentinel: str, group_cols: list[str], ) -> pl.LazyFrame: if group_cols: @@ -85,7 +84,7 @@ def level( # noqa: D103 if c not in group_cols ] if missing: - lf = lf.with_columns([pl.lit(null_sentinel).alias(c) for c in missing]) + lf = lf.with_columns([pl.lit(None, dtype=pl.String).alias(c) for c in missing]) return lf.select( [ "i_item_id", @@ -117,7 +116,6 @@ def polars_impl(run_config: RunConfig) -> QueryResult: es = params["es"] gen = params["gen"] - null_sentinel = "NULL" catalog_sales = get_data( run_config.dataset_path, "catalog_sales", run_config.suffix ) @@ -189,18 +187,17 @@ def polars_impl(run_config: RunConfig) -> QueryResult: pl.col("cd_dep_count").mean().alias("agg7"), ] + # ROLLUP(i_item_id, ca_country, ca_state, ca_county) drops columns from the right: + # level1: all four, level2: drop ca_county, level3: drop ca_state, level4: drop ca_country, level5: grand total level1 = level( base_query, agg_exprs, - null_sentinel, ["i_item_id", "ca_country", "ca_state", "ca_county"], ) - level2 = level( - base_query, agg_exprs, null_sentinel, ["ca_country", "ca_state", "ca_county"] - ) - level3 = level(base_query, agg_exprs, null_sentinel, ["ca_country", "ca_state"]) - level4 = level(base_query, agg_exprs, null_sentinel, ["ca_country"]) - level5 = level(base_query, agg_exprs, null_sentinel, []) + level2 = level(base_query, agg_exprs, ["i_item_id", "ca_country", "ca_state"]) + level3 = level(base_query, agg_exprs, ["i_item_id", "ca_country"]) + level4 = level(base_query, agg_exprs, ["i_item_id"]) + level5 = level(base_query, agg_exprs, []) sort_by = { "ca_country": False, @@ -213,7 +210,6 @@ def polars_impl(run_config: RunConfig) -> QueryResult: return QueryResult( frame=( pl.concat([level1, level2, level3, level4, level5]) - .filter(pl.col("i_item_id") != null_sentinel) .sort(sort_by.keys(), nulls_last=True) .limit(limit) ), diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py index 2753c64d5c6d..717368b72621 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 2.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -166,11 +167,11 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .group_by("d_week_seq") .agg( [ - pl.when(pl.col("d_day_name") == day) - .then(pl.col("sales_price")) - .otherwise(None) - .sum() - .alias(name) + sql_sum( + pl.when(pl.col("d_day_name") == day) + .then(pl.col("sales_price")) + .otherwise(None) + ).alias(name) for day, name in zip(days, day_cols, strict=True) ] ) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q24.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q24.py index fe506bf767cf..67f77024ae35 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q24.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q24.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 @@ -11,6 +11,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -137,14 +138,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: "i_size", ] ) - .agg( - # Polars sum() returns 0 for all-null groups; SQL returns NULL. - # See https://github.com/rapidsai/cudf/issues/19560. - pl.when(pl.col(amountone).count() > 0) - .then(pl.col(amountone).sum()) - .otherwise(None) - .alias("netpaid") - ) + .agg(sql_sum(pl.col(amountone)).alias("netpaid")) ) threshold_table = ssales.select( @@ -160,10 +154,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .agg( # Polars sum() returns 0 for all-null groups; SQL returns NULL. # See https://github.com/rapidsai/cudf/issues/19560. - pl.when(pl.col("netpaid").count() > 0) - .then(pl.col("netpaid").sum()) - .otherwise(None) - .alias("paid") + sql_sum("netpaid").alias("paid") ) .join(threshold_table, how="cross") .filter(pl.col("paid") > pl.col("threshold")) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py index 95351f103b66..20a3c47e0bf7 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 31.""" @@ -204,7 +204,6 @@ def build_quarter_agg( .join(ws3, on="ca_county", suffix="_ws3") .with_columns( [ - # Calculate ratios with null handling pl.when(pl.col("web_sales") > 0) .then(pl.col("web_sales_ws2") / pl.col("web_sales")) .otherwise(None) @@ -224,7 +223,6 @@ def build_quarter_agg( ] ) .filter( - # First condition: web_q1_q2 > store_q1_q2 ( pl.when(pl.col("web_sales") > 0) .then(pl.col("web_sales_ws2") / pl.col("web_sales")) @@ -233,9 +231,7 @@ def build_quarter_agg( .then(pl.col("store_sales_q2") / pl.col("store_sales")) .otherwise(None) ) - & - # Second condition: web_q2_q3 > store_q2_q3 - ( + & ( pl.when(pl.col("web_sales_ws2") > 0) .then(pl.col("web_sales_ws3") / pl.col("web_sales_ws2")) .otherwise(None) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q32.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q32.py index 9134f3366c2d..7b9b483184b4 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q32.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q32.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 32.""" @@ -11,6 +11,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -94,9 +95,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: & (pl.col("d_date").is_between(start_date, end_date)) & (pl.col("cs_ext_discount_amt") > pl.col("threshold_discount")) ) - .select( - [pl.col("cs_ext_discount_amt").sum().alias("excess discount amount")] - ) + .select([sql_sum("cs_ext_discount_amt").alias("excess discount amount")]) .limit(100) ), sort_by=[], diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q35.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q35.py index 9cb77a9e9f2e..24855680aee7 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q35.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q35.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 35.""" @@ -119,6 +119,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: aggone = params["aggone"] aggtwo = params["aggtwo"] aggthree = params["aggthree"] + aggthree_suffix = "_1" if aggthree in (aggone, aggtwo) else "" customer = get_data(run_config.dataset_path, "customer", run_config.suffix) customer_address = get_data( @@ -194,7 +195,11 @@ def polars_impl(run_config: RunConfig) -> QueryResult: pl.len().alias("cnt1"), _get_agg_expr("cd_dep_count", aggone, f"{aggone}(cd_dep_count)"), _get_agg_expr("cd_dep_count", aggtwo, f"{aggtwo}(cd_dep_count)"), - _get_agg_expr("cd_dep_count", aggthree, f"{aggthree}(cd_dep_count)_1"), + _get_agg_expr( + "cd_dep_count", + aggthree, + f"{aggthree}(cd_dep_count){aggthree_suffix}", + ), pl.len().alias("cnt2"), _get_agg_expr( "cd_dep_employed_count", aggone, f"{aggone}(cd_dep_employed_count)" @@ -205,7 +210,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: _get_agg_expr( "cd_dep_employed_count", aggthree, - f"{aggthree}(cd_dep_employed_count)_1", + f"{aggthree}(cd_dep_employed_count){aggthree_suffix}", ), pl.len().alias("cnt3"), _get_agg_expr( @@ -217,7 +222,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: _get_agg_expr( "cd_dep_college_count", aggthree, - f"{aggthree}(cd_dep_college_count)_1", + f"{aggthree}(cd_dep_college_count){aggthree_suffix}", ), ] ) @@ -230,17 +235,17 @@ def polars_impl(run_config: RunConfig) -> QueryResult: "cnt1", f"{aggone}(cd_dep_count)", f"{aggtwo}(cd_dep_count)", - f"{aggthree}(cd_dep_count)_1", + f"{aggthree}(cd_dep_count){aggthree_suffix}", "cd_dep_employed_count", "cnt2", f"{aggone}(cd_dep_employed_count)", f"{aggtwo}(cd_dep_employed_count)", - f"{aggthree}(cd_dep_employed_count)_1", + f"{aggthree}(cd_dep_employed_count){aggthree_suffix}", "cd_dep_college_count", "cnt3", f"{aggone}(cd_dep_college_count)", f"{aggtwo}(cd_dep_college_count)", - f"{aggthree}(cd_dep_college_count)_1", + f"{aggthree}(cd_dep_college_count){aggthree_suffix}", ] ) .sort(sort_by.keys(), nulls_last=True) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q48.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q48.py index 3a600f6584e3..59ca693baf96 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q48.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q48.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 48.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -121,14 +122,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .join(customer_address, left_on="ss_addr_sk", right_on="ca_address_sk") .join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") .filter((pl.col("d_year") == year) & demo_filter & geo_filter) - .select( - [ - pl.when(pl.col("ss_quantity").count() > 0) - .then(pl.col("ss_quantity").sum()) - .otherwise(None) - .alias("sum(ss_quantity)") - ] - ) + .select([sql_sum("ss_quantity").alias("sum(ss_quantity)")]) ), sort_by=[], limit=None, diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py index ee7ddfa3aae3..cef34bd05667 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 49.""" @@ -218,11 +218,22 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation + # TODO: Polars Decimal/Decimal division returns Decimal, but DuckDB SQL + # promotes Decimal/Decimal to Float64. We cast to Float64 here to match. + # Other queries with Decimal arithmetic may have similar semantic differences. ( pl.when(pl.col("ws_net_paid").drop_nulls().count() > 0) .then( - pl.col("wr_return_amt").fill_null(0).sum().round(4) - / pl.col("ws_net_paid").fill_null(0).sum().round(4) + pl.col("wr_return_amt") + .fill_null(0) + .sum() + .cast(pl.Decimal(15, 4)) + .cast(pl.Float64) + / pl.col("ws_net_paid") + .fill_null(0) + .sum() + .cast(pl.Decimal(15, 4)) + .cast(pl.Float64) ) .otherwise(None) ).alias("currency_ratio"), @@ -269,11 +280,22 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation + # TODO: Polars Decimal/Decimal division returns Decimal, but DuckDB SQL + # promotes Decimal/Decimal to Float64. We cast to Float64 here to match. + # Other queries with Decimal arithmetic may have similar semantic differences. ( pl.when(pl.col("cs_net_paid").drop_nulls().count() > 0) .then( - pl.col("cr_return_amount").fill_null(0).sum().round(4) - / pl.col("cs_net_paid").fill_null(0).sum().round(4) + pl.col("cr_return_amount") + .fill_null(0) + .sum() + .cast(pl.Decimal(15, 4)) + .cast(pl.Float64) + / pl.col("cs_net_paid") + .fill_null(0) + .sum() + .cast(pl.Decimal(15, 4)) + .cast(pl.Float64) ) .otherwise(None) ).alias("currency_ratio"), @@ -320,11 +342,22 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation + # TODO: Polars Decimal/Decimal division returns Decimal, but DuckDB SQL + # promotes Decimal/Decimal to Float64. We cast to Float64 here to match. + # Other queries with Decimal arithmetic may have similar semantic differences. ( pl.when(pl.col("ss_net_paid").drop_nulls().count() > 0) .then( - pl.col("sr_return_amt").fill_null(0).sum().round(4) - / pl.col("ss_net_paid").fill_null(0).sum().round(4) + pl.col("sr_return_amt") + .fill_null(0) + .sum() + .cast(pl.Decimal(15, 4)) + .cast(pl.Float64) + / pl.col("ss_net_paid") + .fill_null(0) + .sum() + .cast(pl.Decimal(15, 4)) + .cast(pl.Float64) ) .otherwise(None) ).alias("currency_ratio"), diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q5.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q5.py index e8c3216d1534..f8df00adb902 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q5.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q5.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 5.""" @@ -11,6 +11,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -202,10 +203,10 @@ def _channel_agg( .join(entity, left_on="entity_sk", right_on=entity_join_key) .group_by(entity_id_col) .agg( - pl.col("sales_price").sum().alias("sales"), - pl.col("profit").sum().alias("profit"), - pl.col("return_amt").sum().alias("returns1"), - pl.col("net_loss").sum().alias("profit_loss"), + sql_sum("sales_price").alias("sales"), + sql_sum("profit").alias("profit"), + sql_sum("return_amt").alias("returns1"), + sql_sum("net_loss").alias("profit_loss"), ) ) @@ -326,15 +327,28 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ) all_channels = pl.concat([store_channel, catalog_channel, web_channel]) + agg_exprs = [ + sql_sum("sales").alias("sales"), + sql_sum("returns1").alias("returns1"), + sql_sum("profit").alias("profit"), + ] + per_id = all_channels.group_by(["channel", "id"]).agg(agg_exprs) + per_channel = ( + all_channels.group_by("channel") + .agg(agg_exprs) + .with_columns(pl.lit(None, dtype=pl.String).alias("id")) + .select(["channel", "id", "sales", "returns1", "profit"]) + ) + grand_total = all_channels.select( + pl.lit(None, dtype=pl.String).alias("channel"), + pl.lit(None, dtype=pl.String).alias("id"), + *agg_exprs, + ) + return QueryResult( frame=( - all_channels.group_by(["channel", "id"]) - .agg( - pl.col("sales").sum().alias("sales"), - pl.col("returns1").sum().alias("returns1"), - pl.col("profit").sum().alias("profit"), - ) - .sort(["channel", "id"]) + pl.concat([per_id, per_channel, grand_total]) + .sort(["channel", "id"], nulls_last=True) .limit(100) ), sort_by=[("channel", False), ("id", False)], diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py index 35f5814f869c..44097ca00f6c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 51.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -96,6 +97,18 @@ def polars_impl(run_config: RunConfig) -> QueryResult: store_sales = get_data(run_config.dataset_path, "store_sales", run_config.suffix) date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) + # Polars cum_sum() emits NULL at null positions; SQL SUM() OVER (UNBOUNDED PRECEDING) + # skips NULLs and carries the running total forward. forward_fill() corrects that. + + def _cume_sales(partition_col: str) -> pl.Expr: + return ( + pl.col("daily_sum") + .cum_sum() + .forward_fill() + .over(partition_by=partition_col, order_by="d_date") + .alias("cume_sales") + ) + # web_v1: daily sums -> cumulative sum per (item, ordered by date) web_v1 = ( web_sales.join(date_dim, left_on="ws_sold_date_sk", right_on="d_date_sk") @@ -104,13 +117,8 @@ def polars_impl(run_config: RunConfig) -> QueryResult: & pl.col("ws_item_sk").is_not_null() ) .group_by(["ws_item_sk", "d_date"]) - .agg(pl.col("ws_sales_price").sum().alias("daily_sum")) - .with_columns( - pl.col("daily_sum") - .cum_sum() - .over(partition_by="ws_item_sk", order_by="d_date") - .alias("cume_sales") - ) + .agg(sql_sum("ws_sales_price").alias("daily_sum")) + .with_columns(_cume_sales("ws_item_sk")) .select( pl.col("ws_item_sk").alias("item_sk"), "d_date", @@ -126,13 +134,8 @@ def polars_impl(run_config: RunConfig) -> QueryResult: & pl.col("ss_item_sk").is_not_null() ) .group_by(["ss_item_sk", "d_date"]) - .agg(pl.col("ss_sales_price").sum().alias("daily_sum")) - .with_columns( - pl.col("daily_sum") - .cum_sum() - .over(partition_by="ss_item_sk", order_by="d_date") - .alias("cume_sales") - ) + .agg(sql_sum("ss_sales_price").alias("daily_sum")) + .with_columns(_cume_sales("ss_item_sk")) .select( pl.col("ss_item_sk").alias("item_sk"), "d_date", diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q56.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q56.py index 646ad01c8f7a..e36b3c2e8954 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q56.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q56.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 56.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -167,14 +168,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: & (pl.col("ca_gmt_offset") == gmt_offset) ) .group_by("i_item_id") - .agg( - # Polars sum() returns 0 for all-null groups; SQL returns NULL. - # See https://github.com/rapidsai/cudf/issues/19560. - pl.when(pl.col(str(ch["ext_col"])).count() > 0) - .then(pl.col(str(ch["ext_col"])).sum()) - .otherwise(None) - .alias("total_sales") - ) + .agg(sql_sum(pl.col(str(ch["ext_col"]))).alias("total_sales")) .select(["i_item_id", "total_sales"]) ) for ch in channels @@ -186,14 +180,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: frame=( pl.concat(per_channel) .group_by("i_item_id") - .agg( - # Polars sum() returns 0 for all-null groups; SQL returns NULL. - # See https://github.com/rapidsai/cudf/issues/19560. - pl.when(pl.col("total_sales").count() > 0) - .then(pl.col("total_sales").sum()) - .otherwise(None) - .alias("total_sales") - ) + .agg(sql_sum("total_sales").alias("total_sales")) .select(["i_item_id", "total_sales"]) .sort(sort_by.keys(), nulls_last=True) .limit(limit) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q59.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q59.py index 7e5d4bd62d1d..680d592ec12b 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q59.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q59.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 59.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -135,41 +136,20 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ).select(["d_week_seq", "ss_store_sk", "d_day_name", "ss_sales_price"]) wss = base.group_by(["d_week_seq", "ss_store_sk"]).agg( [ - pl.when(pl.col("d_day_name") == "Sunday") - .then(pl.col("ss_sales_price")) - .otherwise(None) - .sum() - .alias("sun_sales"), - pl.when(pl.col("d_day_name") == "Monday") - .then(pl.col("ss_sales_price")) - .otherwise(None) - .sum() - .alias("mon_sales"), - pl.when(pl.col("d_day_name") == "Tuesday") - .then(pl.col("ss_sales_price")) - .otherwise(None) - .sum() - .alias("tue_sales"), - pl.when(pl.col("d_day_name") == "Wednesday") - .then(pl.col("ss_sales_price")) - .otherwise(None) - .sum() - .alias("wed_sales"), - pl.when(pl.col("d_day_name") == "Thursday") - .then(pl.col("ss_sales_price")) - .otherwise(None) - .sum() - .alias("thu_sales"), - pl.when(pl.col("d_day_name") == "Friday") - .then(pl.col("ss_sales_price")) - .otherwise(None) - .sum() - .alias("fri_sales"), - pl.when(pl.col("d_day_name") == "Saturday") - .then(pl.col("ss_sales_price")) - .otherwise(None) - .sum() - .alias("sat_sales"), + sql_sum( + pl.when(pl.col("d_day_name") == day) + .then(pl.col("ss_sales_price")) + .otherwise(None) + ).alias(alias) + for day, alias in [ + ("Sunday", "sun_sales"), + ("Monday", "mon_sales"), + ("Tuesday", "tue_sales"), + ("Wednesday", "wed_sales"), + ("Thursday", "thu_sales"), + ("Friday", "fri_sales"), + ("Saturday", "sat_sales"), + ] ] ) wss_enriched = wss.join( diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py index 659eb4a9ce52..2a6821975e27 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -362,20 +363,9 @@ def build_cross_sales_for_year(target_year: int) -> pl.LazyFrame: .agg( [ pl.len().alias("cnt"), - # Polars sum() returns 0 for all-null groups; SQL returns NULL. - # See https://github.com/rapidsai/cudf/issues/19560. - pl.when(pl.col("ss_wholesale_cost").count() > 0) - .then(pl.col("ss_wholesale_cost").sum()) - .otherwise(None) - .alias("s1"), - pl.when(pl.col("ss_list_price").count() > 0) - .then(pl.col("ss_list_price").sum()) - .otherwise(None) - .alias("s2"), - pl.when(pl.col("ss_coupon_amt").count() > 0) - .then(pl.col("ss_coupon_amt").sum()) - .otherwise(None) - .alias("s3"), + sql_sum("ss_wholesale_cost").alias("s1"), + sql_sum("ss_list_price").alias("s2"), + sql_sum("ss_coupon_amt").alias("s3"), ] ) .select( diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q67.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q67.py index 8f0c18a6e59f..e9cd56aa50e2 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q67.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q67.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 67.""" @@ -265,7 +265,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ranked = rollup_data.with_columns( pl.col("sumsales") - .rank(method="dense", descending=True) + .rank(method="min", descending=True) .over("i_category") .alias("rk") ) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q72.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q72.py index f264053438c7..f4cdebd8bfea 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q72.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q72.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 72.""" @@ -126,13 +126,8 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .join(week_seqs_2002, left_on="d_week_seq", right_on="d1_week_seq") .rename({"d_date_sk": "d2_date_sk", "d_week_seq": "d2_week_seq"}) ) - d3_dates = ( - date_dim.filter( - pl.col("d_year").is_not_null() - & (pl.col("d_year").is_in([2001, 2002, 2003])) - ) - .select(["d_date_sk", "d_date"]) - .rename({"d_date_sk": "d3_date_sk", "d_date": "d3_date"}) + d3_dates = date_dim.select(["d_date_sk", "d_date"]).rename( + {"d_date_sk": "d3_date_sk", "d_date": "d3_date"} ) filtered_cd = customer_demographics.filter( pl.col("cd_marital_status").is_not_null() & (pl.col("cd_marital_status") == ms) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q76.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q76.py index 211987885921..a9b80108f29c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q76.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q76.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 76.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -174,7 +175,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .agg( [ pl.len().cast(pl.Int64).alias("sales_cnt"), - pl.col("ext_sales_price").sum().alias("sales_amt"), + sql_sum("ext_sales_price").alias("sales_amt"), ] ) .select( diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q78.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q78.py index 7e2a9019779a..22e994b83e44 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q78.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q78.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 78.""" @@ -10,6 +10,7 @@ import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -145,18 +146,9 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .group_by(["ws_item_sk", "ws_bill_customer_sk"]) .agg( [ - pl.when(pl.col("ws_quantity").count() > 0) - .then(pl.col("ws_quantity").sum()) - .otherwise(None) - .alias("ws_qty"), - pl.when(pl.col("ws_wholesale_cost").count() > 0) - .then(pl.col("ws_wholesale_cost").sum()) - .otherwise(None) - .alias("ws_wc"), - pl.when(pl.col("ws_sales_price").count() > 0) - .then(pl.col("ws_sales_price").sum()) - .otherwise(None) - .alias("ws_sp"), + sql_sum("ws_quantity").alias("ws_qty"), + sql_sum("ws_wholesale_cost").alias("ws_wc"), + sql_sum("ws_sales_price").alias("ws_sp"), ] ) .rename({"ws_bill_customer_sk": "ws_customer_sk"}) @@ -172,18 +164,9 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .group_by(["cs_item_sk", "cs_bill_customer_sk"]) .agg( [ - pl.when(pl.col("cs_quantity").count() > 0) - .then(pl.col("cs_quantity").sum()) - .otherwise(None) - .alias("cs_qty"), - pl.when(pl.col("cs_wholesale_cost").count() > 0) - .then(pl.col("cs_wholesale_cost").sum()) - .otherwise(None) - .alias("cs_wc"), - pl.when(pl.col("cs_sales_price").count() > 0) - .then(pl.col("cs_sales_price").sum()) - .otherwise(None) - .alias("cs_sp"), + sql_sum("cs_quantity").alias("cs_qty"), + sql_sum("cs_wholesale_cost").alias("cs_wc"), + sql_sum("cs_sales_price").alias("cs_sp"), ] ) .rename({"cs_bill_customer_sk": "cs_customer_sk"}) @@ -199,18 +182,9 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .group_by(["ss_item_sk", "ss_customer_sk"]) .agg( [ - pl.when(pl.col("ss_quantity").count() > 0) - .then(pl.col("ss_quantity").sum()) - .otherwise(None) - .alias("ss_qty"), - pl.when(pl.col("ss_wholesale_cost").count() > 0) - .then(pl.col("ss_wholesale_cost").sum()) - .otherwise(None) - .alias("ss_wc"), - pl.when(pl.col("ss_sales_price").count() > 0) - .then(pl.col("ss_sales_price").sum()) - .otherwise(None) - .alias("ss_sp"), + sql_sum("ss_quantity").alias("ss_qty"), + sql_sum("ss_wholesale_cost").alias("ss_wc"), + sql_sum("ss_sales_price").alias("ss_sp"), ] ) ) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q92.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q92.py index d187daeb4b86..5be7c940b2c0 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q92.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q92.py @@ -1,16 +1,17 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 92.""" from __future__ import annotations -from datetime import datetime +from datetime import datetime, timedelta from typing import TYPE_CHECKING import polars as pl from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_queries import sql_sum from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: @@ -67,9 +68,10 @@ def polars_impl(run_config: RunConfig) -> QueryResult: web_sales = get_data(run_config.dataset_path, "web_sales", run_config.suffix) item = get_data(run_config.dataset_path, "item", run_config.suffix) date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) - start_date_py = datetime.strptime(date, "%Y-%m-%d") - start_date = pl.lit(start_date_py, dtype=pl.Datetime("us")) - end_date = start_date + pl.duration(days=90) + start_date_py = datetime.strptime(date, "%Y-%m-%d").date() + end_date_py = start_date_py + timedelta(days=90) + start_date = pl.lit(start_date_py, dtype=pl.Date) + end_date = pl.lit(end_date_py, dtype=pl.Date) avg_discounts = ( web_sales.join( date_dim, left_on="ws_sold_date_sk", right_on="d_date_sk", how="inner" @@ -99,9 +101,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: & (pl.col("d_date") <= end_date) & (pl.col("ws_ext_discount_amt") > pl.col("threshold_discount")) ) - .select( - [pl.col("ws_ext_discount_amt").sum().alias("Excess Discount Amount")] - ) + .select([sql_sum("ws_ext_discount_amt").alias("Excess Discount Amount")]) .sort("Excess Discount Amount", nulls_last=True) .limit(100) ), diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 61b6df5249a5..c293898211d9 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -1715,7 +1715,7 @@ def execute_duckdb_query( f"CREATE OR REPLACE VIEW {name} AS " f"SELECT * FROM parquet_scan('{pattern}');" ) - result = pl.from_arrow(conn.sql(query)) + result = conn.sql(query).pl() assert isinstance(result, pl.DataFrame) return result From ae109924e7d548e2d27be8bf0bb7a735b1030adc Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 8 Jul 2026 22:43:08 +0000 Subject: [PATCH 13/31] validate more queries, xfail some, len and window function fix --- ci/convert_tpc_decimals.py | 44 +++++++++++++ .../cudf_polars/dsl/expressions/rolling.py | 63 +++++++++++++++++-- .../cudf_polars/cudf_polars/dsl/translate.py | 11 +++- .../cudf_polars/streaming/benchmarks/pdsds.py | 16 +++-- .../streaming/benchmarks/pdsds_queries/q20.py | 14 ++--- .../streaming/benchmarks/pdsds_queries/q31.py | 43 ++++++++++--- .../streaming/benchmarks/pdsds_queries/q38.py | 39 ++++++------ .../streaming/benchmarks/pdsds_queries/q51.py | 3 - .../streaming/benchmarks/pdsds_queries/q87.py | 57 ++++------------- .../cudf_polars/streaming/benchmarks/utils.py | 27 ++++++-- 10 files changed, 213 insertions(+), 104 deletions(-) create mode 100644 ci/convert_tpc_decimals.py diff --git a/ci/convert_tpc_decimals.py b/ci/convert_tpc_decimals.py new file mode 100644 index 000000000000..d99e55b50f8c --- /dev/null +++ b/ci/convert_tpc_decimals.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cast Decimal columns to Float64 in a TPC parquet dataset directory. + +Workaround for https://github.com/rapidsai/cudf/issues/23150. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path + +import polars as pl + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--data-dir", + default=os.environ.get("TPC_DATA_DIR"), + help="Directory containing parquet files. Defaults to TPC_DATA_DIR environment variable.", + ) + args = parser.parse_args() + + if args.data_dir is None: + parser.error("--data-dir is required (or set TPC_DATA_DIR).") + + for path in sorted(Path(args.data_dir).rglob("*.parquet")): + lf = pl.scan_parquet(path) + decimal_cols = [ + name + for name, dtype in lf.collect_schema().items() + if isinstance(dtype, pl.Decimal) + ] + if decimal_cols: + lf.with_columns( + pl.col(decimal_cols).cast(pl.Float64) + ).sink_parquet(path) + + +if __name__ == "__main__": + main() diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index 7f9c2868c04b..41d92d1d94fd 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py @@ -416,16 +416,42 @@ def _( # type: ignore[no-untyped-def] cum_named = op.named_exprs order_index = op.order_index + _fill_policy_map = { + "forward": plc.replace.ReplacePolicy.PRECEDING, + "backward": plc.replace.ReplacePolicy.FOLLOWING, + } + + # For fill_null_with_strategy(cum_sum(col)), ne.value is the outer wrapper; + # the actual data column lives one level deeper. + def _data_child(ne: expr.NamedExpr) -> expr.Expr: + v = ne.value + if ( + isinstance(v, expr.UnaryFunction) + and v.name == "fill_null_with_strategy" + ): + return v.children[0].children[0] + return v.children[0] + + def _fill_policy(ne: expr.NamedExpr) -> plc.replace.ReplacePolicy | None: + v = ne.value + if ( + isinstance(v, expr.UnaryFunction) + and v.name == "fill_null_with_strategy" + ): + return _fill_policy_map.get(v.options[0]) + return None + requests: list[plc.groupby.GroupByRequest] = [] out_names: list[str] = [] out_dtypes: list[DataType] = [] + fill_policies: list[plc.replace.ReplacePolicy | None] = [] # Instead of calling self._gather_columns, let's call plc.copying.gather directly # since we need plc.Column objects, not cudf_polars Column objects val_cols: Sequence[plc.Column] if order_index is not None: plc_cols = [ - ne.value.children[0].evaluate(df, context=ExecutionContext.FRAME).obj + _data_child(ne).evaluate(df, context=ExecutionContext.FRAME).obj for ne in cum_named ] val_cols = plc.copying.gather( @@ -436,7 +462,7 @@ def _( # type: ignore[no-untyped-def] ).columns() else: val_cols = [ - ne.value.children[0].evaluate(df, context=ExecutionContext.FRAME).obj + _data_child(ne).evaluate(df, context=ExecutionContext.FRAME).obj for ne in cum_named ] agg = plc.aggregation.sum() @@ -445,12 +471,31 @@ def _( # type: ignore[no-untyped-def] requests.append(plc.groupby.GroupByRequest(val_col, [agg])) out_names.append(ne.name) out_dtypes.append(ne.value.dtype) + fill_policies.append(_fill_policy(ne)) local_grouper = op.local_grouper assert isinstance(local_grouper, plc.groupby.GroupBy) _, tables = local_grouper.scan(requests) - return out_names, out_dtypes, tables + # Apply forward/backward fill per-partition for cum_sum wrapped in + # fill_null_with_strategy. Create a fresh grouper (same keys/sort order) + # so the scan's internal state doesn't interfere with replace_nulls. + any_fill = any(p is not None for p in fill_policies) + fill_grouper = ( + GroupedWindow._sorted_grouper(op.by_cols_for_scan) + if any_fill and op.by_cols_for_scan is not None + else None + ) + + result_tables = [] + for tbl, policy in zip(tables, fill_policies, strict=True): + if policy is not None and fill_grouper is not None: + _, filled = fill_grouper.replace_nulls(tbl, [policy]) + result_tables.append(filled) + else: + result_tables.append(tbl) + + return out_names, out_dtypes, result_tables def _reorder_to_input( self, @@ -514,7 +559,17 @@ def _split_named_expr( for ne in self.named_aggs: v = ne.value - if isinstance(v, expr.UnaryFunction) and v.name in unary_window_ops: + if ( + isinstance(v, expr.UnaryFunction) + and v.name == "fill_null_with_strategy" + and isinstance(v.children[0], expr.UnaryFunction) + and v.children[0].name in unary_window_ops + ): + # fill_null_with_strategy(window_fn(...)): + # route to the inner window function's bucket so it runs per-partition; + # the fill_null post-step is applied after the scan in that op's handler. + unary_window_ops[v.children[0].name].append(ne) + elif isinstance(v, expr.UnaryFunction) and v.name in unary_window_ops: unary_window_ops[v.name].append(ne) else: reductions.append(ne) diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index 7967b1e867ce..a57db3a24eea 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/python/cudf_polars/cudf_polars/dsl/translate.py @@ -1106,7 +1106,16 @@ def _( ] child_deps = [ - v.children[0] + # fill_null_with_strategy(inner_window(col)) is routed to inner_window's + # bucket; unwrap one extra level so child_deps tracks the actual data col. + v.children[0].children[0] + if ( + isinstance(v, expr.UnaryFunction) + and v.name == "fill_null_with_strategy" + and isinstance(v.children[0], expr.UnaryFunction) + and v.children[0].name in {"rank", "cum_sum"} + ) + else v.children[0] for ne in named_aggs for v in (ne.value,) if isinstance(v, expr.Agg) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index dd45544b0133..e013554daca0 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -81,6 +81,12 @@ class PDSDSPolarsQueries(PDSDSQueries): """Polars Queries.""" q_impl = "polars_impl" + # Queries expected to fail on GPU due to known bugs. Keys are query numbers; + # values are reasons for the failures. These queries will be skipped in GPU runs. + EXPECTED_FAILURES_TPCDS: ClassVar[dict[int, str]] = { + 5: "GPU execution failure (packed data cannot be empty): https://github.com/rapidsai/cudf/issues/22073", + 9: "GPU cross join with empty left table: https://github.com/rapidsai/cudf/issues/22824", + } # See comments for EXPECTED_CASTS and EXPECTED_CASTS_DECIMAL # in cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py # for more details. @@ -111,17 +117,9 @@ class PDSDSPolarsQueries(PDSDSQueries): pl.col("total net profit").cast(pl.Decimal(18, 2)), ], 19: [pl.col("ext_price").cast(pl.Decimal(18, 2))], - 20: [ - pl.col("revenueratio").cast(pl.Decimal(38, 2)), - ], + 20: [pl.col("revenueratio").cast(pl.Decimal(38, 2))], 24: [pl.col("paid").cast(pl.Decimal(18, 2))], 30: [pl.col("ctr_total_return").cast(pl.Decimal(18, 2))], - 31: [ - pl.col("web_q1_q2_increase").cast(pl.Decimal(38, 2)), - pl.col("store_q1_q2_increase").cast(pl.Decimal(38, 2)), - pl.col("web_q2_q3_increase").cast(pl.Decimal(38, 2)), - pl.col("store_q2_q3_increase").cast(pl.Decimal(38, 2)), - ], 32: [pl.col("excess discount amount").cast(pl.Decimal(18, 2))], 33: [pl.col("total_sales").cast(pl.Decimal(18, 2))], 42: [pl.col("sum(ss_ext_sales_price)").cast(pl.Decimal(18, 2))], diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py index c8282098b01e..f63999c2946c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 20.""" @@ -104,15 +104,11 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .agg([pl.col("cs_ext_sales_price").sum().alias("itemrevenue")]) .with_columns( [ - # Handle case where itemrevenue is 0 - should result in NULL like SQL - pl.when(pl.col("itemrevenue") == 0.0) - .then(None) - .otherwise( - pl.col("itemrevenue") + ( + pl.col("itemrevenue").cast(pl.Float64) * 100 - / pl.col("itemrevenue").sum().over("i_class") - ) - .alias("revenueratio") + / pl.col("itemrevenue").sum().over("i_class").cast(pl.Float64) + ).alias("revenueratio") ] ) .sort(sort_by.keys(), nulls_last=True) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py index 20a3c47e0bf7..bc4c4ac7587c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py @@ -204,20 +204,35 @@ def build_quarter_agg( .join(ws3, on="ca_county", suffix="_ws3") .with_columns( [ + # TODO: DuckDB SQL promotes Decimal/Decimal to Float64; Polars keeps it as + # Decimal. Cast to Float64 to match DuckDB return type and filter semantics. + # Check other queries with Decimal arithmetic for similar semantic differences. pl.when(pl.col("web_sales") > 0) - .then(pl.col("web_sales_ws2") / pl.col("web_sales")) + .then( + pl.col("web_sales_ws2").cast(pl.Float64) + / pl.col("web_sales").cast(pl.Float64) + ) .otherwise(None) .alias("web_q1_q2_increase"), pl.when(pl.col("store_sales") > 0) - .then(pl.col("store_sales_q2") / pl.col("store_sales")) + .then( + pl.col("store_sales_q2").cast(pl.Float64) + / pl.col("store_sales").cast(pl.Float64) + ) .otherwise(None) .alias("store_q1_q2_increase"), pl.when(pl.col("web_sales_ws2") > 0) - .then(pl.col("web_sales_ws3") / pl.col("web_sales_ws2")) + .then( + pl.col("web_sales_ws3").cast(pl.Float64) + / pl.col("web_sales_ws2").cast(pl.Float64) + ) .otherwise(None) .alias("web_q2_q3_increase"), pl.when(pl.col("store_sales_q2") > 0) - .then(pl.col("store_sales_q3") / pl.col("store_sales_q2")) + .then( + pl.col("store_sales_q3").cast(pl.Float64) + / pl.col("store_sales_q2").cast(pl.Float64) + ) .otherwise(None) .alias("store_q2_q3_increase"), ] @@ -225,18 +240,30 @@ def build_quarter_agg( .filter( ( pl.when(pl.col("web_sales") > 0) - .then(pl.col("web_sales_ws2") / pl.col("web_sales")) + .then( + pl.col("web_sales_ws2").cast(pl.Float64) + / pl.col("web_sales").cast(pl.Float64) + ) .otherwise(None) > pl.when(pl.col("store_sales") > 0) - .then(pl.col("store_sales_q2") / pl.col("store_sales")) + .then( + pl.col("store_sales_q2").cast(pl.Float64) + / pl.col("store_sales").cast(pl.Float64) + ) .otherwise(None) ) & ( pl.when(pl.col("web_sales_ws2") > 0) - .then(pl.col("web_sales_ws3") / pl.col("web_sales_ws2")) + .then( + pl.col("web_sales_ws3").cast(pl.Float64) + / pl.col("web_sales_ws2").cast(pl.Float64) + ) .otherwise(None) > pl.when(pl.col("store_sales_q2") > 0) - .then(pl.col("store_sales_q3") / pl.col("store_sales_q2")) + .then( + pl.col("store_sales_q3").cast(pl.Float64) + / pl.col("store_sales_q2").cast(pl.Float64) + ) .otherwise(None) ) ) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q38.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q38.py index cd1a091e39b1..15a02961a857 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q38.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q38.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 38.""" @@ -102,29 +102,28 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .select(["c_last_name", "c_first_name", "d_date"]) .unique() ) - # Find INTERSECT of all three using a different approach - # Combine all three and find tuples that appear exactly 3 times - all_customers = pl.concat( - [ - store_customers.with_columns(pl.lit("store").alias("source")), - catalog_customers.with_columns(pl.lit("catalog").alias("source")), - web_customers.with_columns(pl.lit("web").alias("source")), - ] - ) - # Find combinations that appear in all three sources - intersect_final = ( - all_customers.group_by(["c_last_name", "c_first_name", "d_date"]) - .agg(pl.col("source").n_unique().alias("source_count")) - .filter(pl.col("source_count") == 3) - .select(["c_last_name", "c_first_name", "d_date"]) + # Implement INTERSECT via semi-joins: keep only store rows that also appear + # in catalog and web (nulls_equal so NULL keys match, as SQL INTERSECT does). + intersect_final = store_customers.join( + catalog_customers, + on=["c_last_name", "c_first_name", "d_date"], + how="semi", + nulls_equal=True, + ).join( + web_customers, + on=["c_last_name", "c_first_name", "d_date"], + how="semi", + nulls_equal=True, ) limit = 100 - # Count the final result + # Count the final result. + # Use pl.col("d_date").len() instead of pl.len() to avoid the zero-column + # streaming chunk bug (https://github.com/rapidsai/cudf/issues/21428). return QueryResult( frame=( - intersect_final - # Cast -> Int64 to match DuckDB - .select([pl.len().cast(pl.Int64).alias("count_star()")]).limit(limit) + intersect_final.select( + [pl.col("d_date").len().cast(pl.Int64).alias("count_star()")] + ).limit(limit) ), sort_by=[], limit=limit, diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py index 44097ca00f6c..270ae19c1546 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py @@ -97,9 +97,6 @@ def polars_impl(run_config: RunConfig) -> QueryResult: store_sales = get_data(run_config.dataset_path, "store_sales", run_config.suffix) date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) - # Polars cum_sum() emits NULL at null positions; SQL SUM() OVER (UNBOUNDED PRECEDING) - # skips NULLs and carries the running total forward. forward_fill() corrects that. - def _cume_sales(partition_col: str) -> pl.Expr: return ( pl.col("daily_sum") diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py index 1de0babdc64b..0f6c5bc39009 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Query 87.""" @@ -130,52 +130,17 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .select(["c_last_name", "c_first_name", "d_date"]) .unique() ) - store_customers_sentinel = store_customers.with_columns( - [ - pl.col("c_last_name").fill_null("NULL_SENTINEL_LAST"), - pl.col("c_first_name").fill_null("NULL_SENTINEL_FIRST"), - ] - ) - catalog_customers_sentinel = catalog_customers.with_columns( - [ - pl.col("c_last_name").fill_null("NULL_SENTINEL_LAST"), - pl.col("c_first_name").fill_null("NULL_SENTINEL_FIRST"), - ] - ) - web_customers_sentinel = web_customers.with_columns( - [ - pl.col("c_last_name").fill_null("NULL_SENTINEL_LAST"), - pl.col("c_first_name").fill_null("NULL_SENTINEL_FIRST"), - ] - ) - result_after_first_except = store_customers_sentinel.join( - catalog_customers_sentinel, - on=["c_last_name", "c_first_name", "d_date"], - how="anti", - ).unique() - result_after_second_except = ( - result_after_first_except.join( - web_customers_sentinel, - on=["c_last_name", "c_first_name", "d_date"], - how="anti", - ) - .with_columns( - [ - pl.when(pl.col("c_last_name") == "NULL_SENTINEL_LAST") - .then(None) - .otherwise(pl.col("c_last_name")) - .alias("c_last_name"), - pl.when(pl.col("c_first_name") == "NULL_SENTINEL_FIRST") - .then(None) - .otherwise(pl.col("c_first_name")) - .alias("c_first_name"), - ] - ) - .unique() - ) + join_keys = ["c_last_name", "c_first_name", "d_date"] + result = store_customers.join( + catalog_customers, on=join_keys, how="anti", nulls_equal=True + ).join(web_customers, on=join_keys, how="anti", nulls_equal=True) return QueryResult( - frame=result_after_second_except.select( - [pl.len().cast(pl.Int64).alias("count_star()")] + # Use pl.col("d_date").len() instead of pl.len() to avoid the zero-column + # streaming chunk bug (https://github.com/rapidsai/cudf/issues/21428): + # Polars >=1.41 projects to 0 columns before len(), causing Len.do_evaluate + # to see df.num_rows == 0 and return 0 for every chunk. + frame=result.select( + [pl.col("d_date").len().cast(pl.Int64).alias("count_star()")] ), sort_by=[], limit=None, diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index ecba385e208e..be4e0fdbd765 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -1185,6 +1185,10 @@ def _run_query_loop( ) ) + known_failures: dict[int, str] = getattr( + benchmark, "EXPECTED_FAILURES_TPCDS", {} + ) + try: result = run_polars_query( q_id=q_id, @@ -1197,9 +1201,12 @@ def _run_query_loop( prepare_validation_result=prepare_validation_result, ) except Exception: - print(f"❌ query={q_id} failed (setup or execution)!") + if q_id in known_failures: + print(f"⚠️ query={q_id} failed (known issue): {known_failures[q_id]}") + else: + print(f"❌ query={q_id} failed (setup or execution)!") + query_failures.append((q_id, -1)) print(traceback.format_exc()) - query_failures.append((q_id, -1)) record = FailedRecord( query=q_id, iteration=-1, @@ -1215,9 +1222,21 @@ def _run_query_loop( records[q_id] = result.query_records if result.plan is not None: plans[q_id] = result.plan - query_failures.extend(result.iteration_failures) + for iteration_failure in result.iteration_failures: + if iteration_failure[0] in known_failures: + print( + f"⚠️ query={iteration_failure[0]} iteration {iteration_failure[1]} failed " + f"(known issue): {known_failures[iteration_failure[0]]}" + ) + else: + query_failures.append(iteration_failure) if result.validation_failed: - validation_failures.append(q_id) + if q_id in known_failures: + print( + f"⚠️ query={q_id} failed validation (known issue): {known_failures[q_id]}" + ) + else: + validation_failures.append(q_id) all_partition_plan_rows.extend(result.partition_plan_rows) if all_partition_plan_rows and getattr(args, "explain_partition_plan", False): From 1c88b6234930bd0ac72cb34eae05d9ed40fa0a57 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 9 Jul 2026 11:41:13 +0000 Subject: [PATCH 14/31] dont convert to floats to workaround groupby-sum agg bug --- ci/convert_tpc_decimals.py | 44 ------------------- ci/run_cudf_polars_tpc.sh | 9 ---- .../cudf_polars/streaming/benchmarks/pdsh.py | 16 ++++++- .../cudf_polars/streaming/benchmarks/utils.py | 26 +++++++++-- 4 files changed, 38 insertions(+), 57 deletions(-) delete mode 100644 ci/convert_tpc_decimals.py diff --git a/ci/convert_tpc_decimals.py b/ci/convert_tpc_decimals.py deleted file mode 100644 index d99e55b50f8c..000000000000 --- a/ci/convert_tpc_decimals.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Cast Decimal columns to Float64 in a TPC parquet dataset directory. - -Workaround for https://github.com/rapidsai/cudf/issues/23150. -""" - -from __future__ import annotations - -import argparse -import os -from pathlib import Path - -import polars as pl - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--data-dir", - default=os.environ.get("TPC_DATA_DIR"), - help="Directory containing parquet files. Defaults to TPC_DATA_DIR environment variable.", - ) - args = parser.parse_args() - - if args.data_dir is None: - parser.error("--data-dir is required (or set TPC_DATA_DIR).") - - for path in sorted(Path(args.data_dir).rglob("*.parquet")): - lf = pl.scan_parquet(path) - decimal_cols = [ - name - for name, dtype in lf.collect_schema().items() - if isinstance(dtype, pl.Decimal) - ] - if decimal_cols: - lf.with_columns( - pl.col(decimal_cols).cast(pl.Float64) - ).sink_parquet(path) - - -if __name__ == "__main__": - main() diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index ca8b6d2487f4..c9ba33040b4b 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -50,15 +50,6 @@ export TPCDS_DATA_DIR TPCDS_DATA_DIR=$(mktemp -d) python3 "$(dirname "$0")/generate_tpcds_data.py" --scale 1 --output-dir "${TPCDS_DATA_DIR}" -# Blackwell GPUs (compute capability >= 10.0) have decimal overflow issues in GPU -# aggregations; pre-convert decimals to float to work around rapidsai/cudf#23150. -COMPUTE_CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d '.') -if [[ "${COMPUTE_CAP}" -ge 100 ]]; then - rapids-logger "Blackwell GPU detected (sm_${COMPUTE_CAP}): converting decimals to float" - python3 "$(dirname "$0")/convert_tpc_decimals.py" --data-dir "${TPCH_DATA_DIR}" - python3 "$(dirname "$0")/convert_tpc_decimals.py" --data-dir "${TPCDS_DATA_DIR}" -fi - rapids-logger "Running TPC-H validation tests" cd python/cudf_polars diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py index f51d05d6c23f..7ed6e57731cf 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py @@ -15,7 +15,7 @@ import os from datetime import date -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import polars as pl @@ -25,6 +25,7 @@ _CPU_ENGINES, QueryResult, RunConfig, + _is_blackwell_gpu, build_parser, get_data, parse_args, @@ -36,6 +37,10 @@ # We want to be able to import pdsh in a CPU-only environment. COUNT_DTYPE = None # type: ignore[assignment] + def _is_blackwell_gpu() -> bool: + return False + + if TYPE_CHECKING: from cudf_polars.streaming.benchmarks.utils import RunConfig @@ -123,6 +128,15 @@ class PDSHQueries: EXPECTED_CASTS = EXPECTED_CASTS EXPECTED_CASTS_DECIMAL = EXPECTED_CASTS_DECIMAL EXPECTED_CASTS_TIMESTAMP = EXPECTED_CASTS_TIMESTAMP + # Queries expected to fail on GPU due to known bugs. Keys are query numbers; + # values are reasons for the failures. These queries will be skipped in GPU runs. + EXPECTED_FAILURES_TPCH: ClassVar[dict[int, str]] = ( + { + 1: "Q1 incorrect result on Blackwell: https://github.com/rapidsai/cudf/issues/23150", + } + if _is_blackwell_gpu() + else {} + ) @property def duckdb_queries(self) -> PDSHDuckDBQueries: diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index be4e0fdbd765..8371bfae8239 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -65,6 +65,25 @@ except ImportError: pynvml = None + +def _is_blackwell_gpu() -> bool: + """Return True if any visible GPU is Blackwell (SM 10.x+) architecture.""" + # TODO: switch to cuda.core.system for GPU arch detection once available; + # see https://github.com/rapidsai/cudf/pull/22305 + if pynvml is None: + return False + try: + pynvml.nvmlInit() + for i in range(pynvml.nvmlDeviceGetCount()): + handle = pynvml.nvmlDeviceGetHandleByIndex(i) + major, _ = pynvml.nvmlDeviceGetCudaComputeCapability(handle) + if major >= 10: + return True + except Exception: + pass + return False + + try: import cudf_polars.dsl.tracing import cudf_polars.quent @@ -1185,9 +1204,10 @@ def _run_query_loop( ) ) - known_failures: dict[int, str] = getattr( - benchmark, "EXPECTED_FAILURES_TPCDS", {} - ) + known_failures: dict[int, str] = { + **getattr(benchmark, "EXPECTED_FAILURES_TPCDS", {}), + **getattr(benchmark, "EXPECTED_FAILURES_TPCH", {}), + } try: result = run_polars_query( From 096f223b0956951b61d2be6e4739150a45119348 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 9 Jul 2026 13:46:17 +0000 Subject: [PATCH 15/31] xfail more queries due to groupby-sum agg bug --- .../cudf_polars/streaming/benchmarks/pdsds.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index e013554daca0..5425bfa0a538 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -23,6 +23,7 @@ from cudf_polars.streaming.benchmarks.utils import ( COUNT_DTYPE, _CPU_ENGINES, + _is_blackwell_gpu, build_parser, parse_args, run_polars, @@ -31,6 +32,10 @@ if e.name is not None and not e.name.startswith("cudf_polars"): raise + def _is_blackwell_gpu() -> bool: + return False + + if TYPE_CHECKING: from types import ModuleType @@ -86,6 +91,14 @@ class PDSDSPolarsQueries(PDSDSQueries): EXPECTED_FAILURES_TPCDS: ClassVar[dict[int, str]] = { 5: "GPU execution failure (packed data cannot be empty): https://github.com/rapidsai/cudf/issues/22073", 9: "GPU cross join with empty left table: https://github.com/rapidsai/cudf/issues/22824", + **( + { + 2: "decimal128 groupby-sum incorrect result on Blackwell: https://github.com/rapidsai/cudf/issues/23150", + 43: "decimal128 groupby-sum incorrect result on Blackwell: https://github.com/rapidsai/cudf/issues/23150", + } + if _is_blackwell_gpu() + else {} + ), } # See comments for EXPECTED_CASTS and EXPECTED_CASTS_DECIMAL # in cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py From 670a78cf8195a692b1c3a9a1bc0c83e34a7bb9fc Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 20 Jul 2026 23:07:23 +0000 Subject: [PATCH 16/31] update query comments --- .../streaming/benchmarks/pdsds_queries/q20.py | 2 ++ .../streaming/benchmarks/pdsds_queries/q31.py | 5 ++--- .../streaming/benchmarks/pdsds_queries/q49.py | 15 ++++++--------- .../streaming/benchmarks/pdsds_queries/q87.py | 4 +--- 4 files changed, 11 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py index f63999c2946c..73a18b260aa2 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py @@ -104,6 +104,8 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .agg([pl.col("cs_ext_sales_price").sum().alias("itemrevenue")]) .with_columns( [ + # DuckDB SQL promotes Decimal/Decimal to Float64; Polars keeps it as + # Decimal. Cast to Float64 to match DuckDB return type. ( pl.col("itemrevenue").cast(pl.Float64) * 100 diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py index bc4c4ac7587c..87cc8ea80ff1 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py @@ -204,9 +204,8 @@ def build_quarter_agg( .join(ws3, on="ca_county", suffix="_ws3") .with_columns( [ - # TODO: DuckDB SQL promotes Decimal/Decimal to Float64; Polars keeps it as - # Decimal. Cast to Float64 to match DuckDB return type and filter semantics. - # Check other queries with Decimal arithmetic for similar semantic differences. + # DuckDB SQL promotes Decimal/Decimal to Float64; Polars keeps it as + # Decimal. Cast to Float64 to match DuckDB return type. pl.when(pl.col("web_sales") > 0) .then( pl.col("web_sales_ws2").cast(pl.Float64) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py index cef34bd05667..a383966dd7a6 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py @@ -218,9 +218,8 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation - # TODO: Polars Decimal/Decimal division returns Decimal, but DuckDB SQL - # promotes Decimal/Decimal to Float64. We cast to Float64 here to match. - # Other queries with Decimal arithmetic may have similar semantic differences. + # DuckDB SQL promotes Decimal/Decimal to Float64; Polars keeps it as + # Decimal. Cast to Float64 to match DuckDB return type. ( pl.when(pl.col("ws_net_paid").drop_nulls().count() > 0) .then( @@ -280,9 +279,8 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation - # TODO: Polars Decimal/Decimal division returns Decimal, but DuckDB SQL - # promotes Decimal/Decimal to Float64. We cast to Float64 here to match. - # Other queries with Decimal arithmetic may have similar semantic differences. + # DuckDB SQL promotes Decimal/Decimal to Float64; Polars keeps it as + # Decimal. Cast to Float64 to match DuckDB return type. ( pl.when(pl.col("cs_net_paid").drop_nulls().count() > 0) .then( @@ -342,9 +340,8 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation - # TODO: Polars Decimal/Decimal division returns Decimal, but DuckDB SQL - # promotes Decimal/Decimal to Float64. We cast to Float64 here to match. - # Other queries with Decimal arithmetic may have similar semantic differences. + # DuckDB SQL promotes Decimal/Decimal to Float64; Polars keeps it as + # Decimal. Cast to Float64 to match DuckDB return type. ( pl.when(pl.col("ss_net_paid").drop_nulls().count() > 0) .then( diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py index 0f6c5bc39009..abd88eb669dd 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py @@ -136,9 +136,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ).join(web_customers, on=join_keys, how="anti", nulls_equal=True) return QueryResult( # Use pl.col("d_date").len() instead of pl.len() to avoid the zero-column - # streaming chunk bug (https://github.com/rapidsai/cudf/issues/21428): - # Polars >=1.41 projects to 0 columns before len(), causing Len.do_evaluate - # to see df.num_rows == 0 and return 0 for every chunk. + # streaming chunk bug (https://github.com/rapidsai/cudf/issues/21428) frame=result.select( [pl.col("d_date").len().cast(pl.Int64).alias("count_star()")] ), From b296f57bbfff530730e8307410815ed18e420654 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 21 Jul 2026 14:33:39 +0000 Subject: [PATCH 17/31] Fix stale EXPECTED_FAILURES entries and lukewarm cache drop semantics --- .../cudf_polars/streaming/benchmarks/asserts.py | 2 ++ .../cudf_polars/streaming/benchmarks/pdsds.py | 13 ------------- .../cudf_polars/streaming/benchmarks/pdsh.py | 12 +----------- 3 files changed, 3 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py index 772e30175f5b..189fae10dac9 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py @@ -324,6 +324,8 @@ def sort_for_comparison(df: pl.DataFrame) -> pl.DataFrame: pl.col(col).lt(val - 2 * abs_tol) | pl.col(col).gt(val + 2 * abs_tol) ) + elif val is None: + filter_exprs.append(pl.lit(False)) else: if desc: # then "before" means "greater than" diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index 5425bfa0a538..f42c22d38867 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -23,7 +23,6 @@ from cudf_polars.streaming.benchmarks.utils import ( COUNT_DTYPE, _CPU_ENGINES, - _is_blackwell_gpu, build_parser, parse_args, run_polars, @@ -32,9 +31,6 @@ if e.name is not None and not e.name.startswith("cudf_polars"): raise - def _is_blackwell_gpu() -> bool: - return False - if TYPE_CHECKING: from types import ModuleType @@ -90,15 +86,6 @@ class PDSDSPolarsQueries(PDSDSQueries): # values are reasons for the failures. These queries will be skipped in GPU runs. EXPECTED_FAILURES_TPCDS: ClassVar[dict[int, str]] = { 5: "GPU execution failure (packed data cannot be empty): https://github.com/rapidsai/cudf/issues/22073", - 9: "GPU cross join with empty left table: https://github.com/rapidsai/cudf/issues/22824", - **( - { - 2: "decimal128 groupby-sum incorrect result on Blackwell: https://github.com/rapidsai/cudf/issues/23150", - 43: "decimal128 groupby-sum incorrect result on Blackwell: https://github.com/rapidsai/cudf/issues/23150", - } - if _is_blackwell_gpu() - else {} - ), } # See comments for EXPECTED_CASTS and EXPECTED_CASTS_DECIMAL # in cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py index 7ed6e57731cf..5e0301dcee57 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py @@ -25,7 +25,6 @@ _CPU_ENGINES, QueryResult, RunConfig, - _is_blackwell_gpu, build_parser, get_data, parse_args, @@ -37,9 +36,6 @@ # We want to be able to import pdsh in a CPU-only environment. COUNT_DTYPE = None # type: ignore[assignment] - def _is_blackwell_gpu() -> bool: - return False - if TYPE_CHECKING: from cudf_polars.streaming.benchmarks.utils import RunConfig @@ -130,13 +126,7 @@ class PDSHQueries: EXPECTED_CASTS_TIMESTAMP = EXPECTED_CASTS_TIMESTAMP # Queries expected to fail on GPU due to known bugs. Keys are query numbers; # values are reasons for the failures. These queries will be skipped in GPU runs. - EXPECTED_FAILURES_TPCH: ClassVar[dict[int, str]] = ( - { - 1: "Q1 incorrect result on Blackwell: https://github.com/rapidsai/cudf/issues/23150", - } - if _is_blackwell_gpu() - else {} - ) + EXPECTED_FAILURES_TPCH: ClassVar[dict[int, str]] = {} @property def duckdb_queries(self) -> PDSHDuckDBQueries: From 73aeb3e4105ad762d935a8633773d597213db19f Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 21 Jul 2026 14:33:54 +0000 Subject: [PATCH 18/31] Refactor benchmark runner to use RunOptions dataclass instead of argparse.Namespace --- .../cudf_polars/streaming/benchmarks/utils.py | 240 ++++++++++++------ 1 file changed, 156 insertions(+), 84 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index d462fb24cea8..0729e2671ae7 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -25,7 +25,7 @@ from datetime import UTC, datetime from pathlib import Path from statistics import mean -from typing import TYPE_CHECKING, Any, Literal +from typing import IO, TYPE_CHECKING, Any, Literal import nvtx @@ -33,13 +33,20 @@ __all__: list[str] = [ "COUNT_DTYPE", + "FailedRecord", "QueryResult", + "QueryRunResult", "RunConfig", + "RunOptions", + "SuccessRecord", + "ValidationMethod", + "_add_dataset_args", "build_parser", "get_data", "parse_args", "run_duckdb", "run_polars", + "run_polars_query", ] # The dtype for count() aggregations depends on the presence @@ -119,11 +126,52 @@ def _is_blackwell_gpu() -> bool: } -def get_validation_options(args: Any) -> dict[str, Any]: - """Get validation options dict from parsed arguments.""" +@dataclasses.dataclass(kw_only=True) +class RunOptions: + """ + Options controlling a benchmark run, decoupled from argparse. + + Construct directly for programmatic / test callers, or use + :meth:`from_args` to build from a parsed :class:`argparse.Namespace`. + """ + + debug: bool = False + explain: bool = False + explain_logical: bool = False + explain_partition_plan: bool = False + print_plans: bool = False + print_results: bool = False + summarize: bool = False + output: IO[str] | None = None + output_expected_directory: Path | None = None + results_directory: Path | None = None + validation_abs_tol: float = POLARS_VALIDATION_OPTIONS["abs_tol"] + + @classmethod + def from_args(cls, args: argparse.Namespace) -> RunOptions: + """Create a RunOptions from a parsed argparse.Namespace.""" + return cls( + debug=getattr(args, "debug", False), + explain=getattr(args, "explain", False), + explain_logical=getattr(args, "explain_logical", False), + explain_partition_plan=getattr(args, "explain_partition_plan", False), + print_plans=getattr(args, "print_plans", False), + print_results=getattr(args, "print_results", False), + summarize=getattr(args, "summarize", False), + output=getattr(args, "output", None), + output_expected_directory=getattr(args, "output_expected_directory", None), + results_directory=getattr(args, "results_directory", None), + validation_abs_tol=getattr( + args, "validation_abs_tol", POLARS_VALIDATION_OPTIONS["abs_tol"] + ), + ) + + +def get_validation_options(run_options: RunOptions) -> dict[str, Any]: + """Get validation options dict from RunOptions.""" return { **POLARS_VALIDATION_OPTIONS, - "abs_tol": args.validation_abs_tol, + "abs_tol": run_options.validation_abs_tol, } @@ -774,7 +822,7 @@ def get_executor_options( def print_query_plan( q_id: int, q: pl.LazyFrame, - args: argparse.Namespace, + run_options: RunOptions, run_config: RunConfig, engine: None | pl.GPUEngine = None, *, @@ -783,15 +831,15 @@ def print_query_plan( """Print the query plan.""" logical_plan = plan = None if run_config.frontend == "polars-cpu": - if args.explain_logical: + if run_options.explain_logical: logical_plan = q.explain() - if args.explain: + if run_options.explain: plan = q.show_graph(engine="streaming", plan_stage="physical") elif CUDF_POLARS_AVAILABLE: assert isinstance(engine, pl.GPUEngine) - if args.explain_logical: + if run_options.explain_logical: logical_plan = explain_query(q, engine, physical=False) - if args.explain and run_config.frontend in _STREAMING_FRONTENDS: + if run_options.explain and run_config.frontend in _STREAMING_FRONTENDS: plan = explain_query(q, engine) else: raise RuntimeError( @@ -832,7 +880,7 @@ def execute_query( i: int, q: pl.LazyFrame, run_config: RunConfig, - args: argparse.Namespace, + run_options: RunOptions, engine: None | pl.GPUEngine = None, ) -> tuple[pl.DataFrame, float]: """Execute a query with NVTX annotation.""" @@ -851,7 +899,7 @@ def execute_query( elif CUDF_POLARS_AVAILABLE: assert isinstance(engine, pl.GPUEngine) - if args.debug: + if run_options.debug: translator = Translator(q._ldf.visit(), engine) ir = translator.translate_ir() context = IRExecutionContext() @@ -960,7 +1008,7 @@ def run_polars_query_iteration( iteration: int, q: pl.LazyFrame, run_config: RunConfig, - args: argparse.Namespace, + run_options: RunOptions, engine: pl.GPUEngine | None, expected: pl.DataFrame | None, query_result: Any, @@ -968,7 +1016,7 @@ def run_polars_query_iteration( result_casts: list[pl.Expr] | None = None, ) -> SuccessRecord: """Run a single query iteration. Caller must wrap in try/except.""" - result, duration = execute_query(q_id, iteration, q, run_config, args, engine) + result, duration = execute_query(q_id, iteration, q, run_config, run_options, engine) if expected is not None and prepare_validation_result is not None: result = prepare_validation_result(result) @@ -990,16 +1038,16 @@ def run_polars_query_iteration( limit=query_result.limit, nulls_last=query_result.nulls_last, sort_keys=query_result.sort_keys, - **get_validation_options(args), + **get_validation_options(run_options), ) else: validation_result = None - if args.print_results: + if run_options.print_results: print(result) - if args.results_directory is not None and iteration == 0: - results_dir = Path(args.results_directory) + if run_options.results_directory is not None and iteration == 0: + results_dir = Path(run_options.results_directory) results_dir.mkdir(parents=True, exist_ok=True) output_path = results_dir / f"q_{q_id:02d}.parquet" result.write_parquet(output_path) @@ -1017,7 +1065,7 @@ def run_polars_query( q_id: int, benchmark: Any, run_config: RunConfig, - args: argparse.Namespace, + run_options: RunOptions, engine: pl.GPUEngine | None, numeric_type: str, date_type: str, @@ -1027,16 +1075,16 @@ def run_polars_query( query_result: QueryResult = getattr(benchmark, f"q{q_id}")(run_config) q = query_result.frame - print_query_plan(q_id, q, args, run_config, engine, print_plans=args.print_plans) + print_query_plan(q_id, q, run_options, run_config, engine, print_plans=run_options.print_plans) plan = None - if (args.explain or args.explain_logical) and engine is not None: + if (run_options.explain or run_options.explain_logical) and engine is not None: from cudf_polars.streaming.explain import serialize_query plan = serialize_query(q, engine) part_plan_rows = [] if ( - getattr(args, "explain_partition_plan", False) + run_options.explain_partition_plan and engine is not None and run_config.frontend in _STREAMING_FRONTENDS ): @@ -1074,11 +1122,11 @@ def run_polars_query( case baseline: raise ValueError(f"Invalid baseline: {baseline}") - if args.output_expected_directory is not None: + if run_options.output_expected_directory is not None: assert expected is not None, ( "Expected result must be computed before writing to disk." ) - expected_dir = Path(args.output_expected_directory) + expected_dir = Path(run_options.output_expected_directory) expected_dir.mkdir(parents=True, exist_ok=True) expected.write_parquet(expected_dir / f"q_{q_id:02d}.parquet") @@ -1087,7 +1135,7 @@ def run_polars_query( validation_failed = False record: SuccessRecord | FailedRecord - for i in range(args.iterations): + for i in range(run_config.iterations): if _HAS_STRUCTLOG and run_config.collect_traces: setup_logging(q_id, i) if isinstance(engine, StreamingEngine): @@ -1109,7 +1157,7 @@ def run_polars_query( iteration=i, q=q, run_config=run_config, - args=args, + run_options=run_options, engine=engine, expected=expected, query_result=query_result, @@ -1155,7 +1203,7 @@ def run_polars_query( def _run_query_loop( benchmark: Any, - args: argparse.Namespace, + run_options: RunOptions, run_config: RunConfig, engine: pl.GPUEngine | None, numeric_type: str, @@ -1174,6 +1222,11 @@ def _run_query_loop( query_failures: list[tuple[int, int]] = [] all_partition_plan_rows: list = [] + # lukewarm: drop once before query 1 so the run starts from a known cold + # state, then let the cache warm naturally across queries and iterations. + if run_config.io_mode == "lukewarm": + drop_file_page_cache_recursively(run_config.dataset_path) + for q_id in run_config.queries: if engine is not None: quent_context = engine.config["executor_options"].get("quent_context") @@ -1197,7 +1250,7 @@ def _run_query_loop( q_id=q_id, benchmark=benchmark, run_config=run_config, - args=args, + run_options=run_options, engine=engine, numeric_type=numeric_type, date_type=date_type, @@ -1242,7 +1295,7 @@ def _run_query_loop( validation_failures.append(q_id) all_partition_plan_rows.extend(result.partition_plan_rows) - if all_partition_plan_rows and getattr(args, "explain_partition_plan", False): + if all_partition_plan_rows and run_options.explain_partition_plan: from cudf_polars.streaming.explain import format_partition_plan_table print(format_partition_plan_table(all_partition_plan_rows), flush=True) @@ -1251,14 +1304,14 @@ def _run_query_loop( def _finalize_benchmark_run( - args: argparse.Namespace, + run_options: RunOptions, run_config: RunConfig, validation_failures: list[int], query_failures: list[tuple[int, int]], engine: StreamingEngine | None, ) -> None: """Summarize, serialize, and exit after a benchmark run.""" - if args.summarize: + if run_options.summarize: run_config.summarize() if ( run_config.validation_method is not None @@ -1278,8 +1331,9 @@ def _finalize_benchmark_run( ) else: print("✅ All validated queries passed.") - args.output.write(json.dumps(run_config.serialize(engine=engine))) - args.output.write("\n") + if run_options.output is not None: + run_options.output.write(json.dumps(run_config.serialize(engine=engine))) + run_options.output.write("\n") sys.exit(1 if (query_failures or validation_failures) else 0) @@ -1291,9 +1345,10 @@ def run_polars_cpu( date_type: str, ) -> None: """Run benchmark queries using the Polars CPU streaming engine.""" + run_options = RunOptions.from_args(args) records, plans, validation_failures, query_failures = _run_query_loop( benchmark, - args, + run_options, run_config, engine=None, numeric_type=numeric_type, @@ -1301,7 +1356,7 @@ def run_polars_cpu( ) run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=None + run_options, run_config, validation_failures, query_failures, engine=None ) @@ -1314,6 +1369,7 @@ def run_polars_in_memory( date_type: str, ) -> None: """Run benchmark queries using a single-process GPU in-memory engine.""" + run_options = RunOptions.from_args(args) engine_options = { **run_config.streaming_options.to_engine_options(), "parquet_options": parquet_options, @@ -1325,7 +1381,7 @@ def run_polars_in_memory( ) records, plans, validation_failures, query_failures = _run_query_loop( benchmark, - args, + run_options, run_config, engine=engine, numeric_type=numeric_type, @@ -1334,7 +1390,7 @@ def run_polars_in_memory( run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) run_config = _consolidate_logs(run_config, engine=None) _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=None + run_options, run_config, validation_failures, query_failures, engine=None ) @@ -1349,6 +1405,7 @@ def run_polars_spmd( """Run benchmark queries using SPMD execution via the ``rrun`` launcher.""" from cudf_polars.engine.spmd import SPMDEngine + run_options = RunOptions.from_args(args) executor_options = get_executor_options(run_config, benchmark=benchmark) # "cluster" is reserved — SPMDEngine sets it executor_options.pop("cluster", None) @@ -1380,7 +1437,7 @@ def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: run_config = dataclasses.replace(run_config, n_workers=engine.nranks) records, plans, validation_failures, query_failures = _run_query_loop( benchmark, - args, + run_options, run_config, engine, numeric_type, @@ -1401,7 +1458,7 @@ def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: collect_traces=run_config.collect_traces, ) _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=engine + run_options, run_config, validation_failures, query_failures, engine=engine ) @@ -1416,6 +1473,7 @@ def run_polars_ray( """Run benchmark queries using Ray actor-based distributed execution.""" from cudf_polars.engine.ray import RayEngine + run_options = RunOptions.from_args(args) executor_options = get_executor_options(run_config, benchmark=benchmark) # "cluster" is reserved — RayEngine sets it executor_options.pop("cluster", None) @@ -1439,7 +1497,7 @@ def run_polars_ray( run_config = dataclasses.replace(run_config, n_workers=engine.nranks) records, plans, validation_failures, query_failures = _run_query_loop( benchmark, - args, + run_options, run_config, engine, numeric_type, @@ -1454,7 +1512,7 @@ def run_polars_ray( collect_traces=run_config.collect_traces, ) _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=engine + run_options, run_config, validation_failures, query_failures, engine=engine ) @@ -1471,6 +1529,7 @@ def run_polars_dask( from cudf_polars.engine.dask import DaskEngine + run_options = RunOptions.from_args(args) executor_options = get_executor_options(run_config, benchmark=benchmark) # "cluster" is reserved — DaskEngine sets it executor_options.pop("cluster", None) @@ -1500,7 +1559,7 @@ def run_polars_dask( ) as engine: run_config = dataclasses.replace(run_config, n_workers=engine.nranks) records, plans, validation_failures, query_failures = _run_query_loop( - benchmark, args, run_config, engine, numeric_type, date_type + benchmark, run_options, run_config, engine, numeric_type, date_type ) run_config = dataclasses.replace( run_config, records=dict(records), plans=plans @@ -1516,7 +1575,7 @@ def run_polars_dask( if dask_client is not None: dask_client.close() _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=engine + run_options, run_config, validation_failures, query_failures, engine=engine ) @@ -1826,9 +1885,15 @@ def execute_duckdb_query( def run_duckdb(duckdb_queries_cls: Any, args: argparse.Namespace) -> None: """Run the benchmark with DuckDB.""" vars(args).update({"query_set": duckdb_queries_cls.name}) + run_options = RunOptions.from_args(args) run_config = RunConfig.from_args(args) records: defaultdict[int, list[SuccessRecord | FailedRecord]] = defaultdict(list) + # lukewarm: drop once before query 1 so the run starts from a known cold + # state, then let the cache warm naturally across queries and iterations. + if run_config.io_mode == "lukewarm": + drop_file_page_cache_recursively(run_config.dataset_path) + for q_id in run_config.queries: try: get_q = getattr(duckdb_queries_cls, f"q{q_id}") @@ -1851,7 +1916,7 @@ def run_duckdb(duckdb_queries_cls: Any, args: argparse.Namespace) -> None: print(f"DuckDB Executing: {q_id}") records[q_id] = [] - for i in range(args.iterations): + for i in range(run_config.iterations): if run_config.io_mode == "cold": drop_file_page_cache_recursively(run_config.dataset_path) t0 = time.time() @@ -1864,21 +1929,22 @@ def run_duckdb(duckdb_queries_cls: Any, args: argparse.Namespace) -> None: ) t1 = time.time() record = SuccessRecord(query=q_id, iteration=i, duration=t1 - t0) - if args.print_results: + if run_options.print_results: print(result) print(f"Query {q_id} - Iteration {i} finished in {record.duration:0.4f}s") records[q_id].append(record) - if i == 0 and args.output_expected_directory is not None: - expected_dir = Path(args.output_expected_directory) + if i == 0 and run_options.output_expected_directory is not None: + expected_dir = Path(run_options.output_expected_directory) expected_dir.mkdir(parents=True, exist_ok=True) result.write_parquet(expected_dir / f"q_{q_id:02d}.parquet") run_config = dataclasses.replace(run_config, records=dict(records)) - if args.summarize: + if run_options.summarize: run_config.summarize() - args.output.write(json.dumps(run_config.serialize(engine=None))) - args.output.write("\n") + if run_options.output is not None: + run_options.output.write(json.dumps(run_config.serialize(engine=None))) + run_options.output.write("\n") def check_input_data_type( @@ -1947,24 +2013,8 @@ def parse(query: str | int) -> list[int]: return parse -def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: - """Build the argument parser for PDS-H/PDS-DS benchmarks.""" - from cudf_polars.engine.options import StreamingOptions - - parser = argparse.ArgumentParser( - prog="Cudf-Polars PDS-H/PDS-DS Benchmarks", - formatter_class=argparse.RawTextHelpFormatter, - ) - parser.add_argument( - "query", - type=_query_type(num_queries), - help=textwrap.dedent("""\ - Query to run. One of the following: - - A single number (e.g. 11) - - A comma-separated list of query numbers (e.g. 1,3,7) - - A range of query numbers (e.g. 1-11,23-34) - - The string 'all' to run all queries (1 through 22)"""), - ) +def _add_dataset_args(parser: argparse.ArgumentParser) -> None: + """Register dataset path and format arguments on *parser*.""" parser.add_argument( "--path", type=str, @@ -1998,6 +2048,45 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: File suffix for input table files. Default: .parquet"""), ) + parser.add_argument( + "--io-mode", + dest="io_mode", + default="lukewarm", + choices=["cold", "lukewarm", "hot"], + help=textwrap.dedent("""\ + Cache state control for each timed iteration: + - cold : Drop Linux page cache before each iteration (requires kvikio) + - lukewarm : Drop once before the first query, then let cache warm naturally (default) + - hot : One untimed warmup iteration to populate cache before measured runs"""), + ) + parser.add_argument( + "--validation-abs-tol", + dest="validation_abs_tol", + type=float, + default=POLARS_VALIDATION_OPTIONS["abs_tol"], + help=f"Absolute tolerance for validation comparisons (default: {POLARS_VALIDATION_OPTIONS['abs_tol']}).", + ) + + +def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: + """Build the argument parser for PDS-H/PDS-DS benchmarks.""" + from cudf_polars.engine.options import StreamingOptions + + parser = argparse.ArgumentParser( + prog="Cudf-Polars PDS-H/PDS-DS Benchmarks", + formatter_class=argparse.RawTextHelpFormatter, + ) + parser.add_argument( + "query", + type=_query_type(num_queries), + help=textwrap.dedent("""\ + Query to run. One of the following: + - A single number (e.g. 11) + - A comma-separated list of query numbers (e.g. 1,3,7) + - A range of query numbers (e.g. 1-11,23-34) + - The string 'all' to run all queries (1 through 22)"""), + ) + _add_dataset_args(parser) parser.add_argument( "--frontend", required=True, @@ -2037,17 +2126,6 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: type=int, help="Number of times to run the same query.", ) - parser.add_argument( - "--io-mode", - dest="io_mode", - default="lukewarm", - choices=["cold", "lukewarm", "hot"], - help=textwrap.dedent("""\ - Cache state control for each timed iteration: - - cold : Drop Linux page cache before each iteration (requires kvikio) - - lukewarm : No cache manipulation; OS cache state unchanged (default) - - hot : One untimed warmup iteration to populate cache before measured runs"""), - ) parser.add_argument( "--collect-traces", action=argparse.BooleanOptionalAction, @@ -2150,12 +2228,6 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: default=None, help="Optional directory to write expected results as parquet files.", ) - parser.add_argument( - "--validation-abs-tol", - type=float, - default=0.01, - help="Absolute tolerance for assert_frame_equal validation. Default: 0.01", - ) parser.add_argument( "--extra-info", type=json.loads, From dde3b2fca83537bea05b045c0803ef126db117a0 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 21 Jul 2026 14:34:06 +0000 Subject: [PATCH 19/31] Add TPC-H and TPC-DS pytest validation tests for the streaming engine --- .../cudf_polars/tests/streaming/conftest.py | 122 +++++++++++++++ .../cudf_polars/tests/streaming/test_tpcds.py | 143 ++++++++++++++++++ .../cudf_polars/tests/streaming/test_tpch.py | 136 +++++++++++++++++ 3 files changed, 401 insertions(+) create mode 100644 python/cudf_polars/tests/streaming/conftest.py create mode 100644 python/cudf_polars/tests/streaming/test_tpcds.py create mode 100644 python/cudf_polars/tests/streaming/test_tpch.py diff --git a/python/cudf_polars/tests/streaming/conftest.py b/python/cudf_polars/tests/streaming/conftest.py new file mode 100644 index 000000000000..1aa9ead5f2ad --- /dev/null +++ b/python/cudf_polars/tests/streaming/conftest.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +from typing import Any + +import pytest + +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.spmd import SPMDEngine +from cudf_polars.streaming.benchmarks.utils import ( + RunOptions, + ValidationMethod, + _add_dataset_args, +) + + +class _PytestGroupShim: + """Makes a pytest option group look like an argparse group. + + Allows :meth:`StreamingOptions._add_cli_args` to register its options + directly onto the pytest parser without duplicating the definitions. + """ + + def __init__(self, group: Any) -> None: + self._group = group + + def add_argument(self, *args: Any, **kwargs: Any) -> None: + if kwargs.get("action") is argparse.BooleanOptionalAction: + dest = kwargs.get("dest") or args[0].lstrip("-").replace("-", "_") + self._group.addoption( + *args, + dest=dest, + action="store_true", + default=None, + help=kwargs.get("help", ""), + ) + for opt in args: + self._group.addoption( + f"--no-{opt.lstrip('-')}", + dest=dest, + action="store_false", + ) + else: + kwargs.pop("metavar", None) + self._group.addoption(*args, **kwargs) + + def add_argument_group(self, *args: Any, **kwargs: Any) -> _PytestGroupShim: + return self + + +def pytest_addoption(parser: pytest.Parser) -> None: + group = parser.getgroup("TPC benchmark options") + shim = _PytestGroupShim(group) + StreamingOptions._add_cli_args(shim) # type: ignore[arg-type] + _add_dataset_args(shim) # type: ignore[arg-type] + + +_TPC_QUERY_COUNTS = {"tpch": 22, "tpcds": 99} + + +def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: + if "q_id" not in metafunc.fixturenames: + return + num_queries = next( + (v for k, v in _TPC_QUERY_COUNTS.items() if k in metafunc.function.__name__), + None, + ) + if num_queries is None: + return + metafunc.parametrize( + "q_id", + range(1, num_queries + 1), + ids=[f"q{i:02d}" for i in range(1, num_queries + 1)], + ) + + +@pytest.fixture(scope="session") +def tpc_streaming_options(request: pytest.FixtureRequest) -> StreamingOptions: + ref = argparse.ArgumentParser() + StreamingOptions._add_cli_args(ref) + _add_dataset_args(ref) + ns: dict[str, Any] = {"raise_on_fail": True} + for action in ref._actions: + if not action.option_strings: + continue + try: + ns[action.dest] = request.config.getoption(action.dest) + except ValueError: + pass + return StreamingOptions._from_argparse(argparse.Namespace(**ns)) + + +@pytest.fixture(scope="session") +def tpc_run_options(request: pytest.FixtureRequest) -> RunOptions: + return RunOptions( + validation_abs_tol=request.config.getoption("validation_abs_tol"), + ) + + +@pytest.fixture(scope="session") +def tpc_validation_method(tpc_run_options: RunOptions) -> ValidationMethod: + from cudf_polars.streaming.benchmarks.utils import POLARS_VALIDATION_OPTIONS + + return ValidationMethod( + expected_source="duckdb", + comparison_method="polars", + comparison_options={**POLARS_VALIDATION_OPTIONS, "abs_tol": tpc_run_options.validation_abs_tol}, + expected_location=None, + ) + + +@pytest.fixture(scope="session") +def tpc_spmd_engine(tpc_streaming_options: StreamingOptions) -> SPMDEngine: + with SPMDEngine( + rapidsmpf_options=tpc_streaming_options.to_rapidsmpf_options(), + executor_options=tpc_streaming_options.to_executor_options(), + engine_options=tpc_streaming_options.to_engine_options(), + ) as engine: + yield engine diff --git a/python/cudf_polars/tests/streaming/test_tpcds.py b/python/cudf_polars/tests/streaming/test_tpcds.py new file mode 100644 index 000000000000..dcdf887416d3 --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_tpcds.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TPC-DS validation tests for the streaming GPU engine.""" + +from __future__ import annotations + +import contextlib +from typing import TYPE_CHECKING + +import pytest + +import duckdb + +from cudf_polars.engine.spmd import SPMDEngine +from cudf_polars.streaming.benchmarks.pdsds import PDSDSPolarsQueries +from cudf_polars.streaming.benchmarks.utils import ( + FailedRecord, + RunConfig, + RunOptions, + ValidationMethod, + check_input_data_type, + run_polars_query, +) +from cudf_polars.testing.engine_utils import warns_on_spmd + +if TYPE_CHECKING: + from pytest_subtests import SubTests + +TPCDS_SUFFIX = ".parquet" + +CONDITIONAL_JOIN_NOT_SUPPORTED = "ConditionalJoin not supported for multiple partitions." +SORT_NOT_SUPPORTED = "sort currently only supports column names as `by` keys." + +EXPECTED_WARNINGS: dict[int, str] = { + 14: CONDITIONAL_JOIN_NOT_SUPPORTED, + 23: CONDITIONAL_JOIN_NOT_SUPPORTED, + 24: CONDITIONAL_JOIN_NOT_SUPPORTED, + 36: SORT_NOT_SUPPORTED, + 70: SORT_NOT_SUPPORTED, + 86: SORT_NOT_SUPPORTED, +} + + +@pytest.fixture(scope="session") +def tpcds_data_dir( + request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory +) -> str: + path = request.config.getoption("path") + if path is not None: + return path + scale = request.config.getoption("scale") or 1.0 + data_dir = tmp_path_factory.mktemp("tpcds") + conn = duckdb.connect() + conn.execute( + f"INSTALL tpcds; LOAD tpcds; CALL dsdgen(sf={scale});" + ) + for table in conn.execute("SHOW TABLES").df()["name"]: + conn.execute( + f"COPY {table} TO '{data_dir}/{table}.parquet' (FORMAT PARQUET)" + ) + return str(data_dir) + + +@pytest.fixture(scope="session") +def tpcds_run_config( + request: pytest.FixtureRequest, + tpcds_data_dir: str, + tpc_iterations: int, + tpc_validation_method: ValidationMethod, +) -> RunConfig: + return RunConfig( + engine_name="cudf-polars", + queries=list(range(1, 100)), + query_set="pdsds", + dataset_path=tpcds_data_dir, + scale_factor=request.config.getoption("scale") or 1.0, + suffix=request.config.getoption("suffix") or TPCDS_SUFFIX, + qualification=request.config.getoption("qualification"), + frontend="spmd", + iterations=tpc_iterations, + io_mode=request.config.getoption("io_mode"), + validation_method=tpc_validation_method, + command_line="", + capture_env_vars="", + ) + + +@pytest.fixture(scope="session") +def tpcds_numeric_type(tpcds_run_config: RunConfig) -> str: + numeric_type, _ = check_input_data_type(tpcds_run_config) + return numeric_type + + +@pytest.fixture(scope="session") +def tpcds_date_type(tpcds_run_config: RunConfig) -> str: + _, date_type = check_input_data_type(tpcds_run_config) + return date_type + + +def test_tpcds_query( + request: pytest.FixtureRequest, + subtests: SubTests, + q_id: int, + tpcds_run_config: RunConfig, + tpcds_numeric_type: str, + tpcds_date_type: str, + tpc_spmd_engine: SPMDEngine, + tpc_run_options: RunOptions, +) -> None: + reason = PDSDSPolarsQueries.EXPECTED_FAILURES_TPCDS.get(q_id) + if reason is not None: + request.applymarker(pytest.mark.xfail(reason=reason)) + + warning = EXPECTED_WARNINGS.get(q_id) + ctx = ( + warns_on_spmd(tpc_spmd_engine, UserWarning, match=warning) + if warning is not None + else contextlib.nullcontext() + ) + + with ctx: + qr = run_polars_query( + q_id=q_id, + benchmark=PDSDSPolarsQueries, + run_config=tpcds_run_config, + run_options=tpc_run_options, + engine=tpc_spmd_engine, + numeric_type=tpcds_numeric_type, + date_type=tpcds_date_type, + ) + + if reason is not None: + record = qr.query_records[0] + if isinstance(record, FailedRecord): + raise RuntimeError(record.traceback) + else: + for record in qr.query_records: + with subtests.test(msg=f"iter{record.iteration}"): + if isinstance(record, FailedRecord): + pytest.fail(record.traceback) + elif record.validation_result is not None and record.validation_result.status == "Failed": + pytest.fail(record.validation_result.message or "Validation failed") diff --git a/python/cudf_polars/tests/streaming/test_tpch.py b/python/cudf_polars/tests/streaming/test_tpch.py new file mode 100644 index 000000000000..cbe0ed7cdded --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_tpch.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TPC-H validation tests for the streaming GPU engine.""" + +from __future__ import annotations + +import contextlib +import subprocess +from typing import TYPE_CHECKING + +import pytest + +from cudf_polars.engine.spmd import SPMDEngine +from cudf_polars.streaming.benchmarks.pdsh import PDSHQueries +from cudf_polars.streaming.benchmarks.utils import ( + FailedRecord, + RunConfig, + RunOptions, + ValidationMethod, + check_input_data_type, + run_polars_query, +) +from cudf_polars.testing.engine_utils import warns_on_spmd + +if TYPE_CHECKING: + from pytest_subtests import SubTests + +TPCH_SUFFIX = "/*.parquet" + +EXPECTED_WARNINGS: dict[int, str] = { + 11: "ConditionalJoin not supported for multiple partitions.", + 22: "ConditionalJoin not supported for multiple partitions.", +} + + +@pytest.fixture(scope="session") +def tpch_data_dir( + request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory +) -> str: + path = request.config.getoption("path") + if path is not None: + return path + data_dir = tmp_path_factory.mktemp("tpch") + subprocess.run( + [ + "tpchgen-cli", + "parquet", + "-s", + "1", + "--parts=4", + f"--output-dir={data_dir}", + ], + check=True, + ) + return str(data_dir) + + +@pytest.fixture(scope="session") +def tpch_run_config( + request: pytest.FixtureRequest, + tpch_data_dir: str, + tpc_iterations: int, + tpc_validation_method: ValidationMethod, +) -> RunConfig: + return RunConfig( + engine_name="cudf-polars", + queries=list(range(1, 23)), + query_set="pdsh", + dataset_path=tpch_data_dir, + scale_factor=1, + suffix=request.config.getoption("suffix") or TPCH_SUFFIX, + frontend="spmd", + iterations=tpc_iterations, + io_mode=request.config.getoption("io_mode"), + validation_method=tpc_validation_method, + command_line="", + capture_env_vars="", + ) + + +@pytest.fixture(scope="session") +def tpch_numeric_type(tpch_run_config: RunConfig) -> str: + numeric_type, _ = check_input_data_type(tpch_run_config) + return numeric_type + + +@pytest.fixture(scope="session") +def tpch_date_type(tpch_run_config: RunConfig) -> str: + _, date_type = check_input_data_type(tpch_run_config) + return date_type + + +def test_tpch_query( + request: pytest.FixtureRequest, + subtests: SubTests, + q_id: int, + tpch_run_config: RunConfig, + tpch_numeric_type: str, + tpch_date_type: str, + tpc_spmd_engine: SPMDEngine, + tpc_run_options: RunOptions, +) -> None: + reason = PDSHQueries.EXPECTED_FAILURES_TPCH.get(q_id) + if reason is not None: + request.applymarker(pytest.mark.xfail(reason=reason)) + + warning = EXPECTED_WARNINGS.get(q_id) + ctx = ( + warns_on_spmd(tpc_spmd_engine, UserWarning, match=warning) + if warning is not None + else contextlib.nullcontext() + ) + + with ctx: + qr = run_polars_query( + q_id=q_id, + benchmark=PDSHQueries, + run_config=tpch_run_config, + run_options=tpc_run_options, + engine=tpc_spmd_engine, + numeric_type=tpch_numeric_type, + date_type=tpch_date_type, + ) + + if reason is not None: + record = qr.query_records[0] + if isinstance(record, FailedRecord): + raise RuntimeError(record.traceback) + else: + for record in qr.query_records: + with subtests.test(msg=f"iter{record.iteration}"): + if isinstance(record, FailedRecord): + pytest.fail(record.traceback) + elif record.validation_result is not None and record.validation_result.status == "Failed": + pytest.fail(record.validation_result.message or "Validation failed") From 95da769e72c8ec4eca8b408a08d68746c675eae4 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 21 Jul 2026 14:34:51 +0000 Subject: [PATCH 20/31] Run TPC validation tests via pytest in CI --- ci/run_cudf_polars_pytests.sh | 4 +++- ci/run_cudf_polars_tpc.sh | 26 +++----------------------- 2 files changed, 6 insertions(+), 24 deletions(-) diff --git a/ci/run_cudf_polars_pytests.sh b/ci/run_cudf_polars_pytests.sh index 82d1ccd4879f..696084ee3680 100755 --- a/ci/run_cudf_polars_pytests.sh +++ b/ci/run_cudf_polars_pytests.sh @@ -8,4 +8,6 @@ set -euo pipefail # Support invoking run_cudf_polars_pytests.sh outside the script directory cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_polars/ -python -m pytest --cache-clear "$@" tests +python -m pytest --cache-clear "$@" tests \ + --ignore=tests/streaming/test_tpch.py \ + --ignore=tests/streaming/test_tpcds.py diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index c9ba33040b4b..8a1d5aa3793c 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -50,30 +50,10 @@ export TPCDS_DATA_DIR TPCDS_DATA_DIR=$(mktemp -d) python3 "$(dirname "$0")/generate_tpcds_data.py" --scale 1 --output-dir "${TPCDS_DATA_DIR}" -rapids-logger "Running TPC-H validation tests" +rapids-logger "Running TPC-H and TPC-DS validation tests" cd python/cudf_polars -python -m cudf_polars.streaming.benchmarks.pdsh all \ - --path "${TPCH_DATA_DIR}" \ - --suffix "/*.parquet" \ - --frontend spmd \ - --validate-against duckdb \ +python -m pytest tests/streaming/test_tpch.py tests/streaming/test_tpcds.py \ --iterations 2 \ - --debug \ - --print-results \ - --explain \ - --explain-logical \ - --explain-partition-plan \ - --rapidsmpf-log DEBUG \ - --rapidsmpf-statistics - -rapids-logger "Running TPC-DS validation tests" - -python -m cudf_polars.streaming.benchmarks.pdsds all \ - --path "${TPCDS_DATA_DIR}" \ - --scale 1 \ - --qualification \ - --frontend spmd \ - --validate-against duckdb \ - --iterations 2 + -v From 6bfdcb7c8ae1e7a1a9c469251c26380258d69d02 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 21 Jul 2026 14:56:09 +0000 Subject: [PATCH 21/31] pin tpchgen-cli, pre-commit --- ci/run_cudf_polars_pytests.sh | 2 +- dependencies.yaml | 2 +- .../streaming/benchmarks/asserts.py | 2 +- .../cudf_polars/streaming/benchmarks/utils.py | 16 +++++--- python/cudf_polars/pyproject.toml | 2 +- .../cudf_polars/tests/streaming/conftest.py | 38 +++++++++++-------- .../cudf_polars/tests/streaming/test_tpcds.py | 29 +++++++------- .../cudf_polars/tests/streaming/test_tpch.py | 14 ++++--- 8 files changed, 60 insertions(+), 45 deletions(-) diff --git a/ci/run_cudf_polars_pytests.sh b/ci/run_cudf_polars_pytests.sh index 696084ee3680..4ea040c0d37d 100755 --- a/ci/run_cudf_polars_pytests.sh +++ b/ci/run_cudf_polars_pytests.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/dependencies.yaml b/dependencies.yaml index 29e75e88b4f1..679bb9808eab 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -1189,7 +1189,7 @@ dependencies: packages: - duckdb - pyarrow - - tpchgen-cli + - tpchgen-cli>=3.0.0 test_python_narwhals: common: - output_types: [conda, requirements, pyproject] diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py index 189fae10dac9..1c1c6592e9da 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py @@ -325,7 +325,7 @@ def sort_for_comparison(df: pl.DataFrame) -> pl.DataFrame: | pl.col(col).gt(val + 2 * abs_tol) ) elif val is None: - filter_exprs.append(pl.lit(False)) + filter_exprs.append(pl.lit(value=False)) else: if desc: # then "before" means "greater than" diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 0729e2671ae7..1a6ff511f4e9 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -657,14 +657,14 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: validation_method = ValidationMethod( expected_source="duckdb-disk", comparison_method="polars", - comparison_options=get_validation_options(args), + comparison_options=get_validation_options(RunOptions.from_args(args)), expected_location=args.validate_directory, ) elif args.validate_against is not None: validation_method = ValidationMethod( args.validate_against, comparison_method="polars", - comparison_options=get_validation_options(args), + comparison_options=get_validation_options(RunOptions.from_args(args)), expected_location=None, ) else: @@ -690,7 +690,7 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: query_set=name, dataset_path=path, scale_factor=scale_factor, - suffix=args.suffix, + suffix=args.suffix if args.suffix is not None else ".parquet", qualification=args.qualification, frontend=args.frontend, iterations=args.iterations, @@ -1016,7 +1016,9 @@ def run_polars_query_iteration( result_casts: list[pl.Expr] | None = None, ) -> SuccessRecord: """Run a single query iteration. Caller must wrap in try/except.""" - result, duration = execute_query(q_id, iteration, q, run_config, run_options, engine) + result, duration = execute_query( + q_id, iteration, q, run_config, run_options, engine + ) if expected is not None and prepare_validation_result is not None: result = prepare_validation_result(result) @@ -1075,7 +1077,9 @@ def run_polars_query( query_result: QueryResult = getattr(benchmark, f"q{q_id}")(run_config) q = query_result.frame - print_query_plan(q_id, q, run_options, run_config, engine, print_plans=run_options.print_plans) + print_query_plan( + q_id, q, run_options, run_config, engine, print_plans=run_options.print_plans + ) plan = None if (run_options.explain or run_options.explain_logical) and engine is not None: from cudf_polars.streaming.explain import serialize_query @@ -2043,7 +2047,7 @@ def _add_dataset_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--suffix", type=str, - default=".parquet", + default=None, help=textwrap.dedent("""\ File suffix for input table files. Default: .parquet"""), diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml index 1650d960a884..1ac936e74720 100644 --- a/python/cudf_polars/pyproject.toml +++ b/python/cudf_polars/pyproject.toml @@ -68,7 +68,7 @@ dask = [ ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. tpch = [ "duckdb", - "tpchgen-cli", + "tpchgen-cli>=3.0.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. [project.urls] diff --git a/python/cudf_polars/tests/streaming/conftest.py b/python/cudf_polars/tests/streaming/conftest.py index 1aa9ead5f2ad..00a0c6111774 100644 --- a/python/cudf_polars/tests/streaming/conftest.py +++ b/python/cudf_polars/tests/streaming/conftest.py @@ -4,7 +4,7 @@ from __future__ import annotations import argparse -from typing import Any +from typing import TYPE_CHECKING, Any import pytest @@ -16,12 +16,16 @@ _add_dataset_args, ) +if TYPE_CHECKING: + from collections.abc import Generator -class _PytestGroupShim: - """Makes a pytest option group look like an argparse group. - Allows :meth:`StreamingOptions._add_cli_args` to register its options - directly onto the pytest parser without duplicating the definitions. +class PytestGroupAdapter: + """Adapts a pytest option group to the argparse group interface. + + Allows :meth:`StreamingOptions._add_cli_args` and :func:`_add_dataset_args` + to register their options directly onto the pytest parser without duplicating + the definitions. """ def __init__(self, group: Any) -> None: @@ -47,15 +51,15 @@ def add_argument(self, *args: Any, **kwargs: Any) -> None: kwargs.pop("metavar", None) self._group.addoption(*args, **kwargs) - def add_argument_group(self, *args: Any, **kwargs: Any) -> _PytestGroupShim: + def add_argument_group(self, *args: Any, **kwargs: Any) -> PytestGroupAdapter: return self def pytest_addoption(parser: pytest.Parser) -> None: group = parser.getgroup("TPC benchmark options") - shim = _PytestGroupShim(group) - StreamingOptions._add_cli_args(shim) # type: ignore[arg-type] - _add_dataset_args(shim) # type: ignore[arg-type] + adapter = PytestGroupAdapter(group) + StreamingOptions._add_cli_args(adapter) # type: ignore[arg-type] + _add_dataset_args(adapter) # type: ignore[arg-type] _TPC_QUERY_COUNTS = {"tpch": 22, "tpcds": 99} @@ -84,12 +88,9 @@ def tpc_streaming_options(request: pytest.FixtureRequest) -> StreamingOptions: _add_dataset_args(ref) ns: dict[str, Any] = {"raise_on_fail": True} for action in ref._actions: - if not action.option_strings: + if not action.option_strings or action.dest == "help": continue - try: - ns[action.dest] = request.config.getoption(action.dest) - except ValueError: - pass + ns[action.dest] = request.config.getoption(action.dest) return StreamingOptions._from_argparse(argparse.Namespace(**ns)) @@ -107,13 +108,18 @@ def tpc_validation_method(tpc_run_options: RunOptions) -> ValidationMethod: return ValidationMethod( expected_source="duckdb", comparison_method="polars", - comparison_options={**POLARS_VALIDATION_OPTIONS, "abs_tol": tpc_run_options.validation_abs_tol}, + comparison_options={ + **POLARS_VALIDATION_OPTIONS, + "abs_tol": tpc_run_options.validation_abs_tol, + }, expected_location=None, ) @pytest.fixture(scope="session") -def tpc_spmd_engine(tpc_streaming_options: StreamingOptions) -> SPMDEngine: +def tpc_spmd_engine( + tpc_streaming_options: StreamingOptions, +) -> Generator[SPMDEngine, None, None]: with SPMDEngine( rapidsmpf_options=tpc_streaming_options.to_rapidsmpf_options(), executor_options=tpc_streaming_options.to_executor_options(), diff --git a/python/cudf_polars/tests/streaming/test_tpcds.py b/python/cudf_polars/tests/streaming/test_tpcds.py index dcdf887416d3..e28c7b208813 100644 --- a/python/cudf_polars/tests/streaming/test_tpcds.py +++ b/python/cudf_polars/tests/streaming/test_tpcds.py @@ -6,19 +6,16 @@ from __future__ import annotations import contextlib +from pathlib import Path from typing import TYPE_CHECKING -import pytest - import duckdb +import pytest -from cudf_polars.engine.spmd import SPMDEngine from cudf_polars.streaming.benchmarks.pdsds import PDSDSPolarsQueries from cudf_polars.streaming.benchmarks.utils import ( FailedRecord, RunConfig, - RunOptions, - ValidationMethod, check_input_data_type, run_polars_query, ) @@ -27,9 +24,14 @@ if TYPE_CHECKING: from pytest_subtests import SubTests + from cudf_polars.engine.spmd import SPMDEngine + from cudf_polars.streaming.benchmarks.utils import RunOptions, ValidationMethod + TPCDS_SUFFIX = ".parquet" -CONDITIONAL_JOIN_NOT_SUPPORTED = "ConditionalJoin not supported for multiple partitions." +CONDITIONAL_JOIN_NOT_SUPPORTED = ( + "ConditionalJoin not supported for multiple partitions." +) SORT_NOT_SUPPORTED = "sort currently only supports column names as `by` keys." EXPECTED_WARNINGS: dict[int, str] = { @@ -52,13 +54,9 @@ def tpcds_data_dir( scale = request.config.getoption("scale") or 1.0 data_dir = tmp_path_factory.mktemp("tpcds") conn = duckdb.connect() - conn.execute( - f"INSTALL tpcds; LOAD tpcds; CALL dsdgen(sf={scale});" - ) + conn.execute(f"INSTALL tpcds; LOAD tpcds; CALL dsdgen(sf={scale});") for table in conn.execute("SHOW TABLES").df()["name"]: - conn.execute( - f"COPY {table} TO '{data_dir}/{table}.parquet' (FORMAT PARQUET)" - ) + conn.execute(f"COPY {table} TO '{data_dir}/{table}.parquet' (FORMAT PARQUET)") return str(data_dir) @@ -73,7 +71,7 @@ def tpcds_run_config( engine_name="cudf-polars", queries=list(range(1, 100)), query_set="pdsds", - dataset_path=tpcds_data_dir, + dataset_path=Path(tpcds_data_dir), scale_factor=request.config.getoption("scale") or 1.0, suffix=request.config.getoption("suffix") or TPCDS_SUFFIX, qualification=request.config.getoption("qualification"), @@ -139,5 +137,8 @@ def test_tpcds_query( with subtests.test(msg=f"iter{record.iteration}"): if isinstance(record, FailedRecord): pytest.fail(record.traceback) - elif record.validation_result is not None and record.validation_result.status == "Failed": + elif ( + record.validation_result is not None + and record.validation_result.status == "Failed" + ): pytest.fail(record.validation_result.message or "Validation failed") diff --git a/python/cudf_polars/tests/streaming/test_tpch.py b/python/cudf_polars/tests/streaming/test_tpch.py index cbe0ed7cdded..207786dc737d 100644 --- a/python/cudf_polars/tests/streaming/test_tpch.py +++ b/python/cudf_polars/tests/streaming/test_tpch.py @@ -7,17 +7,15 @@ import contextlib import subprocess +from pathlib import Path from typing import TYPE_CHECKING import pytest -from cudf_polars.engine.spmd import SPMDEngine from cudf_polars.streaming.benchmarks.pdsh import PDSHQueries from cudf_polars.streaming.benchmarks.utils import ( FailedRecord, RunConfig, - RunOptions, - ValidationMethod, check_input_data_type, run_polars_query, ) @@ -26,6 +24,9 @@ if TYPE_CHECKING: from pytest_subtests import SubTests + from cudf_polars.engine.spmd import SPMDEngine + from cudf_polars.streaming.benchmarks.utils import RunOptions, ValidationMethod + TPCH_SUFFIX = "/*.parquet" EXPECTED_WARNINGS: dict[int, str] = { @@ -67,7 +68,7 @@ def tpch_run_config( engine_name="cudf-polars", queries=list(range(1, 23)), query_set="pdsh", - dataset_path=tpch_data_dir, + dataset_path=Path(tpch_data_dir), scale_factor=1, suffix=request.config.getoption("suffix") or TPCH_SUFFIX, frontend="spmd", @@ -132,5 +133,8 @@ def test_tpch_query( with subtests.test(msg=f"iter{record.iteration}"): if isinstance(record, FailedRecord): pytest.fail(record.traceback) - elif record.validation_result is not None and record.validation_result.status == "Failed": + elif ( + record.validation_result is not None + and record.validation_result.status == "Failed" + ): pytest.fail(record.validation_result.message or "Validation failed") From 03207a17f82d67dc3913051eb855d3cb23470f48 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 21 Jul 2026 15:00:55 +0000 Subject: [PATCH 22/31] Add pytest>=9 lower bound to TPC test deps to get subtests for free --- dependencies.yaml | 1 + python/cudf_polars/pyproject.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/dependencies.yaml b/dependencies.yaml index 679bb9808eab..8e9fb18a2ee7 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -1189,6 +1189,7 @@ dependencies: packages: - duckdb - pyarrow + - pytest>=9 - tpchgen-cli>=3.0.0 test_python_narwhals: common: diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml index 1ac936e74720..848ece7874af 100644 --- a/python/cudf_polars/pyproject.toml +++ b/python/cudf_polars/pyproject.toml @@ -68,6 +68,7 @@ dask = [ ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. tpch = [ "duckdb", + "pytest>=9", "tpchgen-cli>=3.0.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. From cd072052f2e47c9eb4a05f8fd278a1c9786477da Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 21 Jul 2026 15:04:21 +0000 Subject: [PATCH 23/31] pytest pinning --- dependencies.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/dependencies.yaml b/dependencies.yaml index 8e9fb18a2ee7..3af4f4085681 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -1189,6 +1189,7 @@ dependencies: packages: - duckdb - pyarrow + # added lower bound pinning to get pytest-subtests for free - pytest>=9 - tpchgen-cli>=3.0.0 test_python_narwhals: From 82825c5a2d210fa28a2a82fdd85ed63a98033973 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 23 Jul 2026 18:24:40 +0000 Subject: [PATCH 24/31] address reviews --- ci/run_cudf_polars_tpc.sh | 3 +++ dependencies.yaml | 8 ++++++++ .../streaming/benchmarks/asserts.py | 4 +++- .../cudf_polars/streaming/benchmarks/utils.py | 18 ------------------ python/cudf_polars/pyproject.toml | 1 + .../cudf_polars/tests/streaming/test_tpch.py | 4 ++-- 6 files changed, 17 insertions(+), 21 deletions(-) diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index 8a1d5aa3793c..c914fb5be6c6 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -38,6 +38,9 @@ rapids-pip-retry install \ "$(echo "${CUDF_STREAMING_WHEELHOUSE}"/cudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ -r "${TPCH_REQUIREMENTS}" +rapids-logger "Check GPU usage" +nvidia-smi + rapids-logger "Generating TPC-H data at SF=1" export TPCH_DATA_DIR diff --git a/dependencies.yaml b/dependencies.yaml index b05f45b16f12..bff82b9694d6 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -410,6 +410,14 @@ files: - test_python_common - test_python_cudf_polars - cudf_polars_trace + py_tpch_cudf_polars: + output: pyproject + pyproject_dir: python/cudf_polars + extras: + table: project.optional-dependencies + key: tpch + includes: + - test_cudf_polars_tpch py_trace_cudf_polars: output: pyproject pyproject_dir: python/cudf_polars diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py index 1c1c6592e9da..77d2874fda44 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py @@ -325,7 +325,9 @@ def sort_for_comparison(df: pl.DataFrame) -> pl.DataFrame: | pl.col(col).gt(val + 2 * abs_tol) ) elif val is None: - filter_exprs.append(pl.lit(value=False)) + filter_exprs.append( + pl.col(col).is_not_null() if nulls_last else pl.lit(value=False) + ) else: if desc: # then "before" means "greater than" diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 1a6ff511f4e9..db92d4ff2ca9 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -73,24 +73,6 @@ pynvml = None -def _is_blackwell_gpu() -> bool: - """Return True if any visible GPU is Blackwell (SM 10.x+) architecture.""" - # TODO: switch to cuda.core.system for GPU arch detection once available; - # see https://github.com/rapidsai/cudf/pull/22305 - if pynvml is None: - return False - try: - pynvml.nvmlInit() - for i in range(pynvml.nvmlDeviceGetCount()): - handle = pynvml.nvmlDeviceGetHandleByIndex(i) - major, _ = pynvml.nvmlDeviceGetCudaComputeCapability(handle) - if major >= 10: - return True - except Exception: - pass - return False - - try: import cudf_polars.dsl.tracing import cudf_polars.quent diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml index 21cf5331d18e..08cec18c014a 100644 --- a/python/cudf_polars/pyproject.toml +++ b/python/cudf_polars/pyproject.toml @@ -68,6 +68,7 @@ dask = [ ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. tpch = [ "duckdb", + "pyarrow", "pytest>=9", "tpchgen-cli>=3.0.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. diff --git a/python/cudf_polars/tests/streaming/test_tpch.py b/python/cudf_polars/tests/streaming/test_tpch.py index 207786dc737d..3d473f702090 100644 --- a/python/cudf_polars/tests/streaming/test_tpch.py +++ b/python/cudf_polars/tests/streaming/test_tpch.py @@ -48,7 +48,7 @@ def tpch_data_dir( "tpchgen-cli", "parquet", "-s", - "1", + str(request.config.getoption("scale") or 1.0), "--parts=4", f"--output-dir={data_dir}", ], @@ -69,7 +69,7 @@ def tpch_run_config( queries=list(range(1, 23)), query_set="pdsh", dataset_path=Path(tpch_data_dir), - scale_factor=1, + scale_factor=request.config.getoption("scale") or 1.0, suffix=request.config.getoption("suffix") or TPCH_SUFFIX, frontend="spmd", iterations=tpc_iterations, From 74cdd1cfb3776a4f0247ea8e1ac0f7150537ca4b Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 23 Jul 2026 18:44:53 +0000 Subject: [PATCH 25/31] address reviews --- ci/generate_tpcds_data.py | 9 +++++---- ci/run_cudf_polars_tpc.sh | 2 +- .../cudf_polars/streaming/benchmarks/pdsds.py | 3 ++- .../cudf_polars/streaming/benchmarks/pdsh.py | 3 ++- python/cudf_polars/tests/streaming/conftest.py | 16 ++++++++-------- python/cudf_polars/tests/streaming/test_tpcds.py | 8 ++++++-- 6 files changed, 24 insertions(+), 17 deletions(-) diff --git a/ci/generate_tpcds_data.py b/ci/generate_tpcds_data.py index db55034977b0..571c16d9259c 100644 --- a/ci/generate_tpcds_data.py +++ b/ci/generate_tpcds_data.py @@ -28,11 +28,12 @@ def main() -> None: # TODO: switch to the Rust TPC-DS generator conn = duckdb.connect() - conn.execute(f"INSTALL tpcds; LOAD tpcds; CALL dsdgen(sf={args.scale});") + conn.execute("INSTALL tpcds") + conn.execute("LOAD tpcds") + conn.execute("CALL dsdgen(sf=$1)", [args.scale]) for table in conn.execute("SHOW TABLES").df()["name"]: - conn.execute( - f"COPY {table} TO '{args.output_dir}/{table}.parquet' (FORMAT PARQUET)" - ) + path = f"{args.output_dir}/{table}.parquet" + conn.execute(f"COPY {table} TO $1 (FORMAT PARQUET)", [path]) if __name__ == "__main__": diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index c914fb5be6c6..e26a8f601d78 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -45,7 +45,7 @@ rapids-logger "Generating TPC-H data at SF=1" export TPCH_DATA_DIR TPCH_DATA_DIR=$(mktemp -d) -tpchgen-cli parquet -s 1 --parts=4 --output-dir="${TPCH_DATA_DIR}" +python3 "$(dirname "$0")/generate_tpch_data.py" --scale 1 --output-dir "${TPCH_DATA_DIR}" rapids-logger "Generating TPC-DS data at SF=1" diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index f42c22d38867..56801a09d562 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -76,6 +76,7 @@ class PDSDSQueries(metaclass=PDSDSQueriesMeta): q_impl: str name: str = "pdsds" + num_queries: int = 99 class PDSDSPolarsQueries(PDSDSQueries): @@ -333,7 +334,7 @@ class PDSDSDuckDBQueries(PDSDSQueries): if __name__ == "__main__": - parser = build_parser(num_queries=99) + parser = build_parser(num_queries=PDSDSQueries.num_queries) args = parse_args(parser=parser) if args.frontend not in _CPU_ENGINES: os.environ["POLARS_MAX_THREADS"] = os.environ.get("POLARS_MAX_THREADS", "1") diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py index 5e0301dcee57..74b89ef71341 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py @@ -121,6 +121,7 @@ class PDSHQueries: """PDS-H query definitions.""" name: str = "pdsh" + num_queries: int = 22 EXPECTED_CASTS = EXPECTED_CASTS EXPECTED_CASTS_DECIMAL = EXPECTED_CASTS_DECIMAL EXPECTED_CASTS_TIMESTAMP = EXPECTED_CASTS_TIMESTAMP @@ -1802,7 +1803,7 @@ def q22(run_config: RunConfig) -> str: if __name__ == "__main__": - parser = build_parser(num_queries=22) + parser = build_parser(num_queries=PDSHQueries.num_queries) args = parse_args(parser=parser) if args.frontend not in _CPU_ENGINES: os.environ["POLARS_MAX_THREADS"] = os.environ.get("POLARS_MAX_THREADS", "1") diff --git a/python/cudf_polars/tests/streaming/conftest.py b/python/cudf_polars/tests/streaming/conftest.py index 644b88621e39..141b9e6a5224 100644 --- a/python/cudf_polars/tests/streaming/conftest.py +++ b/python/cudf_polars/tests/streaming/conftest.py @@ -10,6 +10,8 @@ from cudf_polars.engine.options import StreamingOptions from cudf_polars.engine.spmd import SPMDEngine +from cudf_polars.streaming.benchmarks.pdsds import PDSDSPolarsQueries +from cudf_polars.streaming.benchmarks.pdsh import PDSHQueries from cudf_polars.streaming.benchmarks.utils import ( RunOptions, ValidationMethod, @@ -57,17 +59,15 @@ def pytest_addoption(parser: pytest.Parser) -> None: _add_dataset_args(adapter) # type: ignore[arg-type] -_TPC_QUERY_COUNTS = {"tpch": 22, "tpcds": 99} - - def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: if "q_id" not in metafunc.fixturenames: return - num_queries = next( - (v for k, v in _TPC_QUERY_COUNTS.items() if k in metafunc.function.__name__), - None, - ) - if num_queries is None: + name = metafunc.function.__name__ + if "tpch" in name: + num_queries = PDSHQueries.num_queries + elif "tpcds" in name: + num_queries = PDSDSPolarsQueries.num_queries + else: return metafunc.parametrize( "q_id", diff --git a/python/cudf_polars/tests/streaming/test_tpcds.py b/python/cudf_polars/tests/streaming/test_tpcds.py index e28c7b208813..065085507848 100644 --- a/python/cudf_polars/tests/streaming/test_tpcds.py +++ b/python/cudf_polars/tests/streaming/test_tpcds.py @@ -54,9 +54,13 @@ def tpcds_data_dir( scale = request.config.getoption("scale") or 1.0 data_dir = tmp_path_factory.mktemp("tpcds") conn = duckdb.connect() - conn.execute(f"INSTALL tpcds; LOAD tpcds; CALL dsdgen(sf={scale});") + conn.execute("INSTALL tpcds") + conn.execute("LOAD tpcds") + conn.execute("CALL dsdgen(sf=$1)", [scale]) for table in conn.execute("SHOW TABLES").df()["name"]: - conn.execute(f"COPY {table} TO '{data_dir}/{table}.parquet' (FORMAT PARQUET)") + conn.execute( + f"COPY {table} TO $1 (FORMAT PARQUET)", [f"{data_dir}/{table}.parquet"] + ) return str(data_dir) From d2ebd235232d67a3dad5a49166a8a107a165c5b5 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 23 Jul 2026 19:01:01 +0000 Subject: [PATCH 26/31] add arg parsing tests --- .../tests/streaming/test_tpc_cli.py | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) create mode 100644 python/cudf_polars/tests/streaming/test_tpc_cli.py diff --git a/python/cudf_polars/tests/streaming/test_tpc_cli.py b/python/cudf_polars/tests/streaming/test_tpc_cli.py new file mode 100644 index 000000000000..9d6819934742 --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_tpc_cli.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Smoke tests for the TPC benchmark CLI argument parsing.""" + +from __future__ import annotations + +import pytest + +from cudf_polars.streaming.benchmarks.pdsds import PDSDSPolarsQueries +from cudf_polars.streaming.benchmarks.pdsh import PDSHQueries +from cudf_polars.streaming.benchmarks.utils import RunConfig, parse_args + + +@pytest.mark.parametrize("query_id", [1, PDSHQueries.num_queries]) +def test_tpch_cli_parse(query_id: int) -> None: + args = parse_args( + [ + str(query_id), + "--path", + "/data/tpch", + "--scale", + "1", + "--suffix", + "/*.parquet", + "--frontend", + "spmd", + "--validate-against", + "duckdb", + "--iterations", + "2", + "--io-mode", + "lukewarm", + ], + num_queries=PDSHQueries.num_queries, + ) + assert args.query == [query_id] + assert args.iterations == 2 + assert args.frontend == "spmd" + assert args.io_mode == "lukewarm" + assert args.suffix == "/*.parquet" + + +@pytest.mark.parametrize("query_id", [1, PDSDSPolarsQueries.num_queries]) +def test_tpcds_cli_parse(query_id: int, tmp_path: pytest.TempPath) -> None: + args = parse_args( + [ + str(query_id), + "--path", + str(tmp_path), + "--scale", + "1", + "--qualification", + "--frontend", + "spmd", + "--validate-against", + "duckdb", + "--iterations", + "2", + "--io-mode", + "lukewarm", + ], + num_queries=PDSDSPolarsQueries.num_queries, + ) + vars(args).update({"query_set": PDSDSPolarsQueries.name}) + run_config = RunConfig.from_args(args) + assert run_config.queries == [query_id] + assert run_config.iterations == 2 + assert run_config.qualification + assert run_config.frontend == "spmd" From fa9a0a4258f5a69479c1c0d25c5510e07d4047c3 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 23 Jul 2026 21:06:39 +0000 Subject: [PATCH 27/31] pre-commit check --- python/cudf_polars/tests/streaming/test_tpc_cli.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/streaming/test_tpc_cli.py b/python/cudf_polars/tests/streaming/test_tpc_cli.py index 9d6819934742..d1509374c26d 100644 --- a/python/cudf_polars/tests/streaming/test_tpc_cli.py +++ b/python/cudf_polars/tests/streaming/test_tpc_cli.py @@ -5,8 +5,13 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import pytest +if TYPE_CHECKING: + from pathlib import Path + from cudf_polars.streaming.benchmarks.pdsds import PDSDSPolarsQueries from cudf_polars.streaming.benchmarks.pdsh import PDSHQueries from cudf_polars.streaming.benchmarks.utils import RunConfig, parse_args @@ -42,7 +47,7 @@ def test_tpch_cli_parse(query_id: int) -> None: @pytest.mark.parametrize("query_id", [1, PDSDSPolarsQueries.num_queries]) -def test_tpcds_cli_parse(query_id: int, tmp_path: pytest.TempPath) -> None: +def test_tpcds_cli_parse(query_id: int, tmp_path: Path) -> None: args = parse_args( [ str(query_id), From 7ebca1bb501ee4118eaa27fb959cbdbc62ac6450 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Sat, 22 Aug 2026 18:16:26 +0000 Subject: [PATCH 28/31] address reviews --- .../cudf_polars/streaming/benchmarks/utils.py | 4 +-- python/cudf_polars/tests/conftest.py | 5 ++- .../cudf_polars/tests/streaming/conftest.py | 3 +- .../tests/streaming/test_tpc_cli.py | 1 + .../cudf_polars/tests/streaming/test_tpcds.py | 27 ++++++++------ .../cudf_polars/tests/streaming/test_tpch.py | 35 ++++++++++++------- 6 files changed, 48 insertions(+), 27 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 9d55d85b190d..893ebb4f49cc 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -2099,12 +2099,12 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: parser.add_argument( "query", type=_query_type(num_queries), - help=textwrap.dedent("""\ + help=textwrap.dedent(f"""\ Query to run. One of the following: - A single number (e.g. 11) - A comma-separated list of query numbers (e.g. 1,3,7) - A range of query numbers (e.g. 1-11,23-34) - - The string 'all' to run all queries (1 through 22)"""), + - The string 'all' to run all queries (1 through {num_queries})"""), ) _add_dataset_args(parser) parser.add_argument( diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index e8fd89e3b1c8..4f903c59488e 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -347,7 +347,10 @@ def pytest_addoption(parser: pytest.Parser) -> None: @pytest.fixture(scope="session") def tpc_iterations(request: pytest.FixtureRequest) -> int: - return request.config.getoption("--iterations") + iterations = request.config.getoption("--iterations") + if iterations < 1: + raise pytest.UsageError("--iterations must be >= 1") + return iterations def pytest_configure(config: pytest.Config): diff --git a/python/cudf_polars/tests/streaming/conftest.py b/python/cudf_polars/tests/streaming/conftest.py index 141b9e6a5224..02eca624c686 100644 --- a/python/cudf_polars/tests/streaming/conftest.py +++ b/python/cudf_polars/tests/streaming/conftest.py @@ -13,6 +13,7 @@ from cudf_polars.streaming.benchmarks.pdsds import PDSDSPolarsQueries from cudf_polars.streaming.benchmarks.pdsh import PDSHQueries from cudf_polars.streaming.benchmarks.utils import ( + POLARS_VALIDATION_OPTIONS, RunOptions, ValidationMethod, _add_dataset_args, @@ -98,8 +99,6 @@ def tpc_run_options(request: pytest.FixtureRequest) -> RunOptions: @pytest.fixture(scope="session") def tpc_validation_method(tpc_run_options: RunOptions) -> ValidationMethod: - from cudf_polars.streaming.benchmarks.utils import POLARS_VALIDATION_OPTIONS - return ValidationMethod( expected_source="duckdb", comparison_method="polars", diff --git a/python/cudf_polars/tests/streaming/test_tpc_cli.py b/python/cudf_polars/tests/streaming/test_tpc_cli.py index d1509374c26d..686ae1b818dd 100644 --- a/python/cudf_polars/tests/streaming/test_tpc_cli.py +++ b/python/cudf_polars/tests/streaming/test_tpc_cli.py @@ -73,3 +73,4 @@ def test_tpcds_cli_parse(query_id: int, tmp_path: Path) -> None: assert run_config.iterations == 2 assert run_config.qualification assert run_config.frontend == "spmd" + assert run_config.io_mode == "lukewarm" diff --git a/python/cudf_polars/tests/streaming/test_tpcds.py b/python/cudf_polars/tests/streaming/test_tpcds.py index 973815f2ba41..84a91fbda299 100644 --- a/python/cudf_polars/tests/streaming/test_tpcds.py +++ b/python/cudf_polars/tests/streaming/test_tpcds.py @@ -6,6 +6,7 @@ from __future__ import annotations import contextlib +import os from pathlib import Path from typing import TYPE_CHECKING @@ -48,19 +49,20 @@ def tpcds_data_dir( request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory ) -> str: - path = request.config.getoption("path") + path = request.config.getoption("path") or os.environ.get("TPCDS_DATA_DIR") if path is not None: return path scale = request.config.getoption("scale") or 1.0 data_dir = tmp_path_factory.mktemp("tpcds") - conn = duckdb.connect() - conn.execute("INSTALL tpcds") - conn.execute("LOAD tpcds") - conn.execute("CALL dsdgen(sf=$1)", [scale]) - for table in conn.execute("SHOW TABLES").df()["name"]: - conn.execute( - f"COPY {table} TO $1 (FORMAT PARQUET)", [f"{data_dir}/{table}.parquet"] - ) + with duckdb.connect() as conn: + conn.execute("INSTALL tpcds") + conn.execute("LOAD tpcds") + conn.execute("CALL dsdgen(sf=$1)", [scale]) + for table in conn.execute("SHOW TABLES").df()["name"]: + conn.execute( + f"COPY {table} TO $1 (FORMAT PARQUET)", + [f"{data_dir}/{table}.parquet"], + ) return str(data_dir) @@ -73,7 +75,7 @@ def tpcds_run_config( ) -> RunConfig: return RunConfig( engine_name="cudf-polars", - queries=list(range(1, 100)), + queries=list(range(1, PDSDSPolarsQueries.num_queries + 1)), query_set="pdsds", dataset_path=Path(tpcds_data_dir), scale_factor=request.config.getoption("scale") or 1.0, @@ -139,6 +141,11 @@ def test_tpcds_query( record = qr.query_records[0] if isinstance(record, FailedRecord): raise RuntimeError(record.traceback) + if ( + record.validation_result is not None + and record.validation_result.status == "Failed" + ): + raise RuntimeError(record.validation_result.message or "Validation failed") else: for record in qr.query_records: with subtests.test(msg=f"iter{record.iteration}"): diff --git a/python/cudf_polars/tests/streaming/test_tpch.py b/python/cudf_polars/tests/streaming/test_tpch.py index 34140b52ca41..5ac0a313585d 100644 --- a/python/cudf_polars/tests/streaming/test_tpch.py +++ b/python/cudf_polars/tests/streaming/test_tpch.py @@ -43,17 +43,23 @@ def tpch_data_dir( if path is not None: return path data_dir = tmp_path_factory.mktemp("tpch") - subprocess.run( - [ - "tpchgen-cli", - "parquet", - "-s", - str(request.config.getoption("scale") or 1.0), - "--parts=4", - f"--output-dir={data_dir}", - ], - check=True, - ) + try: + subprocess.run( + [ + "tpchgen-cli", + "parquet", + "-s", + str(request.config.getoption("scale") or 1.0), + "--parts=4", + f"--output-dir={data_dir}", + ], + check=True, + timeout=1800, + ) + except FileNotFoundError as e: + raise RuntimeError( + "tpchgen-cli is not installed. Install it to generate TPC-H test data." + ) from e return str(data_dir) @@ -66,7 +72,7 @@ def tpch_run_config( ) -> RunConfig: return RunConfig( engine_name="cudf-polars", - queries=list(range(1, 23)), + queries=list(range(1, PDSHQueries.num_queries + 1)), query_set="pdsh", dataset_path=Path(tpch_data_dir), scale_factor=request.config.getoption("scale") or 1.0, @@ -131,6 +137,11 @@ def test_tpch_query( record = qr.query_records[0] if isinstance(record, FailedRecord): raise RuntimeError(record.traceback) + if ( + record.validation_result is not None + and record.validation_result.status == "Failed" + ): + raise RuntimeError(record.validation_result.message or "Validation failed") else: for record in qr.query_records: with subtests.test(msg=f"iter{record.iteration}"): From 53bf6adf85549e975156b2c8ed797e38f6a122bf Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 24 Aug 2026 18:31:15 +0000 Subject: [PATCH 29/31] add datagen script --- ci/generate_tpch_data.py | 46 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 ci/generate_tpch_data.py diff --git a/ci/generate_tpch_data.py b/ci/generate_tpch_data.py new file mode 100644 index 000000000000..c64f3bad6c10 --- /dev/null +++ b/ci/generate_tpch_data.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate TPC-H data at a given scale factor using tpchgen-cli.""" + +from __future__ import annotations + +import argparse +import os +import subprocess + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--scale", type=float, default=1.0, help="Scale factor." + ) + parser.add_argument( + "--parts", type=int, default=4, help="Number of parts per table." + ) + parser.add_argument( + "--output-dir", + default=os.environ.get("TPCH_DATA_DIR"), + help="Output directory. Defaults to TPCH_DATA_DIR environment variable.", + ) + args = parser.parse_args() + + if args.output_dir is None: + parser.error("--output-dir is required (or set TPCH_DATA_DIR).") + + subprocess.run( + [ + "tpchgen-cli", + "parquet", + "-s", + str(args.scale), + f"--parts={args.parts}", + f"--output-dir={args.output_dir}", + ], + check=True, + ) + + +if __name__ == "__main__": + main() From 84c3d23cc49a3ae0b97d49f375e9e7b0d0ae769f Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 24 Aug 2026 18:55:15 +0000 Subject: [PATCH 30/31] just use test fixtures for data generation --- ci/generate_tpcds_data.py | 40 ---------------------------------- ci/generate_tpch_data.py | 46 --------------------------------------- ci/run_cudf_polars_tpc.sh | 12 ---------- 3 files changed, 98 deletions(-) delete mode 100644 ci/generate_tpcds_data.py delete mode 100644 ci/generate_tpch_data.py diff --git a/ci/generate_tpcds_data.py b/ci/generate_tpcds_data.py deleted file mode 100644 index 571c16d9259c..000000000000 --- a/ci/generate_tpcds_data.py +++ /dev/null @@ -1,40 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Generate TPC-DS data at a given scale factor using DuckDB.""" - -from __future__ import annotations - -import argparse -import os - -import duckdb - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--scale", type=float, default=0.01, help="Scale factor." - ) - parser.add_argument( - "--output-dir", - default=os.environ.get("TPCDS_DATA_DIR"), - help="Output directory. Defaults to TPCDS_DATA_DIR environment variable.", - ) - args = parser.parse_args() - - if args.output_dir is None: - parser.error("--output-dir is required (or set TPCDS_DATA_DIR).") - - # TODO: switch to the Rust TPC-DS generator - conn = duckdb.connect() - conn.execute("INSTALL tpcds") - conn.execute("LOAD tpcds") - conn.execute("CALL dsdgen(sf=$1)", [args.scale]) - for table in conn.execute("SHOW TABLES").df()["name"]: - path = f"{args.output_dir}/{table}.parquet" - conn.execute(f"COPY {table} TO $1 (FORMAT PARQUET)", [path]) - - -if __name__ == "__main__": - main() diff --git a/ci/generate_tpch_data.py b/ci/generate_tpch_data.py deleted file mode 100644 index c64f3bad6c10..000000000000 --- a/ci/generate_tpch_data.py +++ /dev/null @@ -1,46 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Generate TPC-H data at a given scale factor using tpchgen-cli.""" - -from __future__ import annotations - -import argparse -import os -import subprocess - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--scale", type=float, default=1.0, help="Scale factor." - ) - parser.add_argument( - "--parts", type=int, default=4, help="Number of parts per table." - ) - parser.add_argument( - "--output-dir", - default=os.environ.get("TPCH_DATA_DIR"), - help="Output directory. Defaults to TPCH_DATA_DIR environment variable.", - ) - args = parser.parse_args() - - if args.output_dir is None: - parser.error("--output-dir is required (or set TPCH_DATA_DIR).") - - subprocess.run( - [ - "tpchgen-cli", - "parquet", - "-s", - str(args.scale), - f"--parts={args.parts}", - f"--output-dir={args.output_dir}", - ], - check=True, - ) - - -if __name__ == "__main__": - main() diff --git a/ci/run_cudf_polars_tpc.sh b/ci/run_cudf_polars_tpc.sh index e26a8f601d78..9b5a333d7a08 100755 --- a/ci/run_cudf_polars_tpc.sh +++ b/ci/run_cudf_polars_tpc.sh @@ -41,18 +41,6 @@ rapids-pip-retry install \ rapids-logger "Check GPU usage" nvidia-smi -rapids-logger "Generating TPC-H data at SF=1" - -export TPCH_DATA_DIR -TPCH_DATA_DIR=$(mktemp -d) -python3 "$(dirname "$0")/generate_tpch_data.py" --scale 1 --output-dir "${TPCH_DATA_DIR}" - -rapids-logger "Generating TPC-DS data at SF=1" - -export TPCDS_DATA_DIR -TPCDS_DATA_DIR=$(mktemp -d) -python3 "$(dirname "$0")/generate_tpcds_data.py" --scale 1 --output-dir "${TPCDS_DATA_DIR}" - rapids-logger "Running TPC-H and TPC-DS validation tests" cd python/cudf_polars From 6a1c5bb91a358e427a18666c4de02a430c7106f5 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 26 Aug 2026 15:43:58 +0000 Subject: [PATCH 31/31] move module scoped settings under __name__ == __main__ --- .../cudf_polars/streaming/benchmarks/pdsds.py | 18 +++++++++--------- .../cudf_polars/streaming/benchmarks/pdsh.py | 18 +++++++++--------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index 56801a09d562..a5e6980e741c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -35,15 +35,6 @@ if TYPE_CHECKING: from types import ModuleType -# Without this setting, the first IO task to run -# on each worker takes ~15 sec extra -os.environ["KVIKIO_COMPAT_MODE"] = os.environ.get("KVIKIO_COMPAT_MODE", "on") -os.environ["KVIKIO_NTHREADS"] = os.environ.get("KVIKIO_NTHREADS", "8") -# TODO: consider raising the rapidsmpf built-in default from 1 to 8. -os.environ["RAPIDSMPF_NUM_STREAMING_THREADS"] = os.environ.get( - "RAPIDSMPF_NUM_STREAMING_THREADS", "8" -) - def valid_query(name: str) -> bool: """Return True for valid query names eg. 'q9', 'q65', etc.""" @@ -334,6 +325,15 @@ class PDSDSDuckDBQueries(PDSDSQueries): if __name__ == "__main__": + # Without this setting, the first IO task to run + # on each worker takes ~15 sec extra + os.environ["KVIKIO_COMPAT_MODE"] = os.environ.get("KVIKIO_COMPAT_MODE", "on") + os.environ["KVIKIO_NTHREADS"] = os.environ.get("KVIKIO_NTHREADS", "8") + # TODO: consider raising the rapidsmpf built-in default from 1 to 8. + os.environ["RAPIDSMPF_NUM_STREAMING_THREADS"] = os.environ.get( + "RAPIDSMPF_NUM_STREAMING_THREADS", "8" + ) + parser = build_parser(num_queries=PDSDSQueries.num_queries) args = parse_args(parser=parser) if args.frontend not in _CPU_ENGINES: diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py index 74b89ef71341..9d08e8adb328 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py @@ -40,15 +40,6 @@ if TYPE_CHECKING: from cudf_polars.streaming.benchmarks.utils import RunConfig -# Without this setting, the first IO task to run -# on each worker takes ~15 sec extra -os.environ["KVIKIO_COMPAT_MODE"] = os.environ.get("KVIKIO_COMPAT_MODE", "on") -os.environ["KVIKIO_NTHREADS"] = os.environ.get("KVIKIO_NTHREADS", "8") -# TODO: consider raising the rapidsmpf built-in default from 1 to 8. -os.environ["RAPIDSMPF_NUM_STREAMING_THREADS"] = os.environ.get( - "RAPIDSMPF_NUM_STREAMING_THREADS", "8" -) - # The pre-computed expected results come from DuckDB, which has # different casting rules than Polars. For example, in polars # Series[Decimal].mean() returns a Float64, while DuckDB returns a Decimal. @@ -1803,6 +1794,15 @@ def q22(run_config: RunConfig) -> str: if __name__ == "__main__": + # Without this setting, the first IO task to run + # on each worker takes ~15 sec extra + os.environ["KVIKIO_COMPAT_MODE"] = os.environ.get("KVIKIO_COMPAT_MODE", "on") + os.environ["KVIKIO_NTHREADS"] = os.environ.get("KVIKIO_NTHREADS", "8") + # TODO: consider raising the rapidsmpf built-in default from 1 to 8. + os.environ["RAPIDSMPF_NUM_STREAMING_THREADS"] = os.environ.get( + "RAPIDSMPF_NUM_STREAMING_THREADS", "8" + ) + parser = build_parser(num_queries=PDSHQueries.num_queries) args = parse_args(parser=parser) if args.frontend not in _CPU_ENGINES: