diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 1faa32b5f60a..71d0f81c83b5 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -41,6 +41,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 @@ -816,6 +817,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, wheel-build-cudf-streaming, 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: [build-details, wheel-build-cudf] permissions: diff --git a/ci/run_cudf_polars_pytests.sh b/ci/run_cudf_polars_pytests.sh index e1f629881684..be764be4078b 100755 --- a/ci/run_cudf_polars_pytests.sh +++ b/ci/run_cudf_polars_pytests.sh @@ -11,4 +11,6 @@ TIMEOUT_TOOL_PATH="$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/timeout_with_st cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_polars/ python "${TIMEOUT_TOOL_PATH}" --enable-python 5400 \ - python -m pytest --cache-clear -p no:benchmark "$@" tests + python -m pytest --cache-clear -p no:benchmark "$@" 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 new file mode 100755 index 000000000000..9b5a333d7a08 --- /dev/null +++ b/ci/run_cudf_polars_tpc.sh @@ -0,0 +1,50 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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}")" +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}" + +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 \ + --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \ + > "${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)" \ + "$(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 "Check GPU usage" +nvidia-smi + +rapids-logger "Running TPC-H and TPC-DS validation tests" + +cd python/cudf_polars + +python -m pytest tests/streaming/test_tpch.py tests/streaming/test_tpcds.py \ + --iterations 2 \ + -v diff --git a/dependencies.yaml b/dependencies.yaml index 7db198d1aaa3..d61c5bb56a5f 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -475,6 +475,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 matrix: @@ -724,6 +732,10 @@ files: polars_compat_version: ["1.35", "1.36", "1.37", "1.38", "1.39", "1.40", "1.41", "1.42"] includes: - test_cudf_polars_compat + test_cudf_polars_tpch: + output: none + includes: + - test_cudf_polars_tpch channels: - rapidsai-nightly - rapidsai @@ -1321,6 +1333,15 @@ dependencies: packages: # Used by timeout_with_stack.py utility - gdb + test_cudf_polars_tpch: + common: + - output_types: [conda, requirements, pyproject] + packages: + - duckdb + - pyarrow + # added lower bound pinning to get pytest-subtests for free + - pytest>=9 + - tpchgen-cli>=3.0.0 test_python_narwhals: common: - output_types: [conda, requirements, constraints, pyproject] diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py index 772e30175f5b..77d2874fda44 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py @@ -324,6 +324,10 @@ 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.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/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index 19ac159b76e0..a5e6980e741c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -31,18 +31,10 @@ if e.name is not None and not e.name.startswith("cudf_polars"): raise + 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.""" @@ -75,12 +67,18 @@ class PDSDSQueries(metaclass=PDSDSQueriesMeta): q_impl: str name: str = "pdsds" + num_queries: int = 99 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", + } # See comments for EXPECTED_CASTS and EXPECTED_CASTS_DECIMAL # in cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py # for more details. @@ -111,18 +109,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("itemrevenue").cast(pl.Decimal(18, 2)), - 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))], @@ -239,6 +228,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), @@ -331,7 +325,16 @@ class PDSDSDuckDBQueries(PDSDSQueries): if __name__ == "__main__": - parser = build_parser(num_queries=99) + # 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: os.environ["POLARS_MAX_THREADS"] = os.environ.get("POLARS_MAX_THREADS", "1") 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 5d5a71c9b361..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 @@ -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/NVIDIA/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/q20.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py index c8282098b01e..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 @@ -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,13 @@ 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") + # 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 - / 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/q24.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q24.py index 04590757c9fa..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 @@ -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/NVIDIA/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( @@ -159,11 +153,8 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .group_by(["c_last_name", "c_first_name", "s_store_name"]) .agg( # Polars sum() returns 0 for all-null groups; SQL returns NULL. - # See https://github.com/NVIDIA/cudf/issues/19560. - pl.when(pl.col("netpaid").count() > 0) - .then(pl.col("netpaid").sum()) - .otherwise(None) - .alias("paid") + # See https://github.com/rapidsai/cudf/issues/19560. + 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..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 @@ -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,43 +204,65 @@ def build_quarter_agg( .join(ws3, on="ca_county", suffix="_ws3") .with_columns( [ - # Calculate ratios with null handling + # 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") / 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"), ] ) .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")) + .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) ) - & - # 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")) + .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/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/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/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..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 @@ -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,21 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation + # 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( - 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 +279,21 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation + # 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( - 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 +340,21 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .otherwise(None) ).alias("return_ratio"), # Currency ratio calculation + # 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( - 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..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 @@ -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,15 @@ 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) + 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 +114,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 +131,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 be80304cd47b..2950f22a8cb2 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 @@ -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,9 @@ 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/NVIDIA/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") - ) + # Polars sum() returns 0 for all-null groups; SQL returns NULL. + # See https://github.com/NVIDIA/cudf/issues/19560. + .agg(sql_sum(pl.col(str(ch["ext_col"]))).alias("total_sales")) .select(["i_item_id", "total_sales"]) ) for ch in channels @@ -186,14 +182,9 @@ 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/NVIDIA/cudf/issues/19560. - pl.when(pl.col("total_sales").count() > 0) - .then(pl.col("total_sales").sum()) - .otherwise(None) - .alias("total_sales") - ) + # Polars sum() returns 0 for all-null groups; SQL returns NULL. + # See https://github.com/NVIDIA/cudf/issues/19560. + .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 9a60220e67bf..134c09265c55 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: @@ -134,11 +135,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 @@ -364,18 +365,9 @@ def build_cross_sales_for_year(target_year: int) -> pl.LazyFrame: pl.len().alias("cnt"), # Polars sum() returns 0 for all-null groups; SQL returns NULL. # See https://github.com/NVIDIA/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/q87.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py index 1de0babdc64b..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 @@ -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,15 @@ 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) + 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/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/pdsh.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py index f51d05d6c23f..9d08e8adb328 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 @@ -36,18 +36,10 @@ # We want to be able to import pdsh in a CPU-only environment. COUNT_DTYPE = None # type: ignore[assignment] + 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. @@ -120,9 +112,13 @@ 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 + # 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]] = {} @property def duckdb_queries(self) -> PDSHDuckDBQueries: @@ -1798,7 +1794,16 @@ def q22(run_config: RunConfig) -> str: if __name__ == "__main__": - parser = build_parser(num_queries=22) + # 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: os.environ["POLARS_MAX_THREADS"] = os.environ.get("POLARS_MAX_THREADS", "1") diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 9401cae81157..73f5415babf0 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -26,7 +26,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 @@ -34,13 +34,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 @@ -66,6 +73,7 @@ except ImportError: pynvml = None + try: import cudf_polars.dsl.tracing import cudf_polars.quent @@ -95,7 +103,7 @@ from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.explain import SerializablePlan -POLARS_VALIDATION_OPTIONS = { +POLARS_VALIDATION_OPTIONS: dict[str, Any] = { "check_row_order": True, "check_column_order": True, "check_dtypes": True, @@ -105,11 +113,52 @@ } -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, } @@ -527,6 +576,7 @@ class RunConfig: # Run parameters iterations: int + sleep_between_iterations: float = 0 io_mode: Literal["cold", "lukewarm", "hot"] = "lukewarm" collect_traces: bool = False # All streaming/rapidsmpf/engine knobs @@ -637,14 +687,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: @@ -670,10 +720,11 @@ 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, + sleep_between_iterations=args.sleep_between_iterations, io_mode=args.io_mode, collect_traces=args.collect_traces, streaming_options=streaming_options, @@ -798,7 +849,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, *, @@ -807,15 +858,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( @@ -856,7 +907,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.""" @@ -875,7 +926,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() @@ -1001,7 +1052,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, @@ -1009,7 +1060,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, 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) @@ -1035,16 +1088,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) @@ -1064,7 +1117,7 @@ def run_polars_query( query_result: QueryResult, benchmark: Any, run_config: RunConfig, - args: argparse.Namespace, + run_options: RunOptions, engine: pl.GPUEngine | None, numeric_type: str, date_type: str, @@ -1074,11 +1127,13 @@ def run_polars_query( """Run all iterations for a single query. Caller must wrap in try/except.""" 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 + ) 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 ): @@ -1116,11 +1171,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") @@ -1129,14 +1184,14 @@ def run_polars_query( validation_failed = False record: SuccessRecord | FailedRecord - for i in range(args.iterations): - if i > 0 and args.sleep_between_iterations > 0: + for i in range(run_config.iterations): + if i > 0 and run_config.sleep_between_iterations > 0: print( - f"==> Sleeping {args.sleep_between_iterations} seconds " + f"==> Sleeping {run_config.sleep_between_iterations} seconds " "between iterations", flush=True, ) - time.sleep(args.sleep_between_iterations) + time.sleep(run_config.sleep_between_iterations) if _HAS_STRUCTLOG and run_config.collect_traces: setup_logging(q_id, i) @@ -1159,7 +1214,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, @@ -1205,7 +1260,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, @@ -1224,6 +1279,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") @@ -1237,11 +1297,17 @@ def _run_query_loop( ) ) + known_failures: dict[int, str] = { + **getattr(benchmark, "EXPECTED_FAILURES_TPCDS", {}), + **getattr(benchmark, "EXPECTED_FAILURES_TPCH", {}), + } plan = None try: query_result: QueryResult = getattr(benchmark, f"q{q_id}")(run_config) - 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: # If this fails during serialization, we have issues. But we'd # rather see what the issues are with execution than query serialization, # so ignore exceptions here. @@ -1253,7 +1319,7 @@ def _run_query_loop( query_result=query_result, benchmark=benchmark, run_config=run_config, - args=args, + run_options=run_options, engine=engine, numeric_type=numeric_type, date_type=date_type, @@ -1261,9 +1327,12 @@ def _run_query_loop( plan=plan, ) 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, @@ -1279,12 +1348,24 @@ 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): + 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) @@ -1293,14 +1374,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]], serializable_engine_config: dict[str, Any], ) -> 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 @@ -1320,8 +1401,9 @@ def _finalize_benchmark_run( ) if not validation_failures and not query_failures: print("✅ All validated queries passed.") - args.output.write(json.dumps(serializable_engine_config)) - args.output.write("\n") + if run_options.output is not None: + run_options.output.write(json.dumps(serializable_engine_config)) + run_options.output.write("\n") sys.exit(1 if (query_failures or validation_failures) else 0) @@ -1333,9 +1415,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, @@ -1343,7 +1426,7 @@ def run_polars_cpu( ) run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) _finalize_benchmark_run( - args, + run_options, run_config, validation_failures, query_failures, @@ -1360,6 +1443,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, @@ -1371,7 +1455,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, @@ -1380,7 +1464,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_options, run_config, validation_failures, query_failures, @@ -1399,6 +1483,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) @@ -1430,7 +1515,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, @@ -1453,7 +1538,7 @@ def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: collect_traces=run_config.collect_traces, ) _finalize_benchmark_run( - args, + run_options, run_config, validation_failures, query_failures, @@ -1472,6 +1557,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) @@ -1495,7 +1581,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, @@ -1512,7 +1598,7 @@ def run_polars_ray( collect_traces=run_config.collect_traces, ) _finalize_benchmark_run( - args, + run_options, run_config, validation_failures, query_failures, @@ -1533,6 +1619,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) @@ -1562,7 +1649,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 @@ -1580,7 +1667,7 @@ def run_polars_dask( if dask_client is not None: dask_client.close() _finalize_benchmark_run( - args, + run_options, run_config, validation_failures, query_failures, @@ -1886,15 +1973,23 @@ def execute_duckdb_query( f"CREATE OR REPLACE VIEW {name} AS " f"SELECT * FROM parquet_scan('{pattern}');" ) - return conn.execute(query).pl() + result = conn.sql(query).pl() + assert isinstance(result, pl.DataFrame) + return result 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}") @@ -1917,7 +2012,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() @@ -1930,21 +2025,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( @@ -2013,24 +2109,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, @@ -2059,11 +2139,50 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: parser.add_argument( "--suffix", type=str, - default=".parquet", + default=None, help=textwrap.dedent("""\ 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(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 {num_queries})"""), + ) + _add_dataset_args(parser) parser.add_argument( "--frontend", required=True, @@ -2212,12 +2331,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, diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml index ea7546fdd31b..7c46277c1c4a 100644 --- a/python/cudf_polars/pyproject.toml +++ b/python/cudf_polars/pyproject.toml @@ -67,6 +67,12 @@ ray = [ dask = [ "rapids-dask-dependency==26.10.*,>=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", + "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`. [project.urls] Homepage = "https://github.com/NVIDIA/cudf" diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index 6965993241df..4f903c59488e 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -336,6 +336,23 @@ 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: + iterations = request.config.getoption("--iterations") + if iterations < 1: + raise pytest.UsageError("--iterations must be >= 1") + return iterations + + def pytest_configure(config: pytest.Config): config.addinivalue_line( "markers", diff --git a/python/cudf_polars/tests/streaming/conftest.py b/python/cudf_polars/tests/streaming/conftest.py new file mode 100644 index 000000000000..02eca624c686 --- /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 TYPE_CHECKING, Any + +import pytest + +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 ( + POLARS_VALIDATION_OPTIONS, + RunOptions, + ValidationMethod, + _add_dataset_args, +) + +if TYPE_CHECKING: + from collections.abc import Generator + + +class PytestGroupAdapter: + """Adapts a pytest option group to the argparse group interface.""" + + 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) -> PytestGroupAdapter: + return self + + +def pytest_addoption(parser: pytest.Parser) -> None: + group = parser.getgroup("TPC benchmark options") + adapter = PytestGroupAdapter(group) + StreamingOptions._add_cli_args(adapter) # type: ignore[arg-type] + _add_dataset_args(adapter) # type: ignore[arg-type] + + +def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: + if "q_id" not in metafunc.fixturenames: + return + 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", + 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 or action.dest == "help": + continue + ns[action.dest] = request.config.getoption(action.dest) + 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: + 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, +) -> Generator[SPMDEngine, None, None]: + 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_tpc_cli.py b/python/cudf_polars/tests/streaming/test_tpc_cli.py new file mode 100644 index 000000000000..686ae1b818dd --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_tpc_cli.py @@ -0,0 +1,76 @@ +# 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 + +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 + + +@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: Path) -> 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" + 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 new file mode 100644 index 000000000000..84a91fbda299 --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_tpcds.py @@ -0,0 +1,158 @@ +# 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 +import os +from pathlib import Path +from typing import TYPE_CHECKING + +import duckdb +import pytest + +from cudf_polars.streaming.benchmarks.pdsds import PDSDSPolarsQueries +from cudf_polars.streaming.benchmarks.utils import ( + FailedRecord, + RunConfig, + 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 + + 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." +) +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") 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") + 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) + + +@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, PDSDSPolarsQueries.num_queries + 1)), + query_set="pdsds", + 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"), + 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() + ) + + query_result = getattr(PDSDSPolarsQueries, f"q{q_id}")(tpcds_run_config) + + with ctx: + qr = run_polars_query( + q_id=q_id, + query_result=query_result, + 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) + 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}"): + 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..5ac0a313585d --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_tpch.py @@ -0,0 +1,154 @@ +# 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 pathlib import Path +from typing import TYPE_CHECKING + +import pytest + +from cudf_polars.streaming.benchmarks.pdsh import PDSHQueries +from cudf_polars.streaming.benchmarks.utils import ( + FailedRecord, + RunConfig, + 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 + + from cudf_polars.engine.spmd import SPMDEngine + from cudf_polars.streaming.benchmarks.utils import RunOptions, ValidationMethod + +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") + 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) + + +@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, PDSHQueries.num_queries + 1)), + query_set="pdsh", + dataset_path=Path(tpch_data_dir), + scale_factor=request.config.getoption("scale") or 1.0, + 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() + ) + + query_result = getattr(PDSHQueries, f"q{q_id}")(tpch_run_config) + + with ctx: + qr = run_polars_query( + q_id=q_id, + query_result=query_result, + 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) + 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}"): + 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")