From fe0e78e20d75272d5e2a8338bf252b32929c4e80 Mon Sep 17 00:00:00 2001 From: Abigale Kim Date: Wed, 22 Jul 2026 07:09:30 -0500 Subject: [PATCH 01/16] RMM Library Linking Error Bug Fix (#23341) Closes #23348 Fixes build/CI by linking RMM to cuDF testing libraries. Authors: - Abigale Kim (https://github.com/abigalekim) Approvers: - Kyle Edwards (https://github.com/KyleFromNVIDIA) - Muhammad Haseeb (https://github.com/mhaseeb123) - Bradley Dice (https://github.com/bdice) URL: https://github.com/rapidsai/cudf/pull/23341 --- cpp/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 4b5887d4ec19..8cb8ade06821 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1416,7 +1416,7 @@ if(CUDF_BUILD_TESTUTIL) ) target_link_libraries( - cudftestutil INTERFACE $ Threads::Threads cudf + cudftestutil INTERFACE $ Threads::Threads cudf rmm::rmm cudftest_default_stream $ ) From 2fb0ed734f4a8cb5868a8e62406ba484a77b1a88 Mon Sep 17 00:00:00 2001 From: Aryan Srivastava Date: Wed, 22 Jul 2026 18:55:41 +0530 Subject: [PATCH 02/16] Fix Parquet stats pinned host copy lifetime (#23219) Closes #23178. Fixes a lifetime hazard in Parquet statistics filtering where `stats_caster_base::host_column::to_device` could return while asynchronous copies were still reading buffers owned by the temporary host column. This keeps the host-owned value and null-mask buffers alive until the stream has consumed them, covering both fixed-width and string statistics. Authors: - Aryan Srivastava (https://github.com/aryansri05) - https://github.com/nursexams-lgtm - Muhammad Haseeb (https://github.com/mhaseeb123) Approvers: - Vukasin Milovanovic (https://github.com/vuule) - Muhammad Haseeb (https://github.com/mhaseeb123) - Basit Ayantunde (https://github.com/lamarrr) URL: https://github.com/rapidsai/cudf/pull/23219 --- cpp/src/io/parquet/stats_filter_helpers.hpp | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/cpp/src/io/parquet/stats_filter_helpers.hpp b/cpp/src/io/parquet/stats_filter_helpers.hpp index ff85b9febff0..60af924ea259 100644 --- a/cpp/src/io/parquet/stats_filter_helpers.hpp +++ b/cpp/src/io/parquet/stats_filter_helpers.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -280,21 +280,22 @@ class stats_caster_base { { if constexpr (std::is_same_v) { auto [d_chars, d_offsets, _] = make_strings_children(val, chars, stream, mr); + auto null_mask_buffer = rmm::device_buffer{ + null_mask.data(), cudf::bitmask_allocation_size_bytes(val.size()), stream, mr}; + stream.synchronize(); return cudf::make_strings_column( val.size(), std::make_unique(std::move(d_offsets), rmm::device_buffer{0, stream, mr}, 0), d_chars.release(), null_count, - rmm::device_buffer{ - null_mask.data(), cudf::bitmask_allocation_size_bytes(val.size()), stream, mr}); + std::move(null_mask_buffer)); } + auto data = cudf::detail::make_device_uvector_async(val, stream, mr); + auto null_mask_buffer = rmm::device_buffer{ + null_mask.data(), cudf::bitmask_allocation_size_bytes(val.size()), stream, mr}; + stream.synchronize(); return std::make_unique( - dtype, - val.size(), - cudf::detail::make_device_uvector_async(val, stream, mr).release(), - rmm::device_buffer{ - null_mask.data(), cudf::bitmask_allocation_size_bytes(val.size()), stream, mr}, - null_count); + dtype, val.size(), data.release(), std::move(null_mask_buffer), null_count); } }; }; From f367be29a4e64acaadb260d62b1df4b2cc2d801f Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 22 Jul 2026 11:58:46 -0500 Subject: [PATCH 03/16] Bump cudf-polars test suite timeouts (#23399) We've observed some cudf-polars runs taking >1 hour to complete, while runs on different, but similar, hardware take <15 minutes. We've resolved most of the outright hangs now. Under the assumption that it's preferable to wait for a job to complete rather than restarting, we'll bump the timeout here. Authors: - Tom Augspurger (https://github.com/TomAugspurger) Approvers: - Vyas Ramasubramani (https://github.com/vyasr) URL: https://github.com/rapidsai/cudf/pull/23399 --- ci/run_cudf_polars_polars_tests.sh | 4 ++-- ci/run_cudf_polars_pytests.sh | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ci/run_cudf_polars_polars_tests.sh b/ci/run_cudf_polars_polars_tests.sh index 3778e1a9035a..3073b61812d7 100755 --- a/ci/run_cudf_polars_polars_tests.sh +++ b/ci/run_cudf_polars_polars_tests.sh @@ -64,7 +64,7 @@ DESELECTED_TESTS_STR=$(printf -- " --deselect %s" "${DESELECTED_TESTS[@]}") # shellcheck disable=SC2086 # Fail fast (-x) because failed tests pollute the state echo "Run polars tests with injected in-memory GPU engine" -python "${TIMEOUT_TOOL_PATH}" --enable-python 3600 \ +python "${TIMEOUT_TOOL_PATH}" --enable-python 4800 \ python -m pytest \ --import-mode=importlib \ --cache-clear \ @@ -85,7 +85,7 @@ python "${TIMEOUT_TOOL_PATH}" --enable-python 3600 \ echo "Run polars tests with injected SPMD GPU engine, small blocksize" CUDF_POLARS__EXECUTOR__TARGET_PARTITION_SIZE=805306368 \ CUDF_POLARS__EXECUTOR__FALLBACK_MODE=silent \ -python "${TIMEOUT_TOOL_PATH}" --enable-python 3600 \ +python "${TIMEOUT_TOOL_PATH}" --enable-python 4800 \ python -m pytest \ --import-mode=importlib \ --cache-clear \ diff --git a/ci/run_cudf_polars_pytests.sh b/ci/run_cudf_polars_pytests.sh index 3fac6910c5a8..c7cdb580a4d7 100755 --- a/ci/run_cudf_polars_pytests.sh +++ b/ci/run_cudf_polars_pytests.sh @@ -10,5 +10,5 @@ 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 3600 \ +python "${TIMEOUT_TOOL_PATH}" --enable-python 4800 \ python -m pytest --cache-clear "$@" tests From 629dcc45e0108d8b8f40acdb9075c13add72c890 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 22 Jul 2026 12:06:28 -0500 Subject: [PATCH 04/16] Add cuFile wheel dependencies for Python tests (#23392) The `ci-imgs` change in https://github.com/rapidsai/ci-imgs/pull/408 switched `citestwheel` images from CUDA `devel` to `base`, which no longer bundles `libcufile.so` with the container. Now we must install cuFile wheels in Python test environments that use CUDA wheels. CUDA 12 uses `nvidia-cufile-cu12` with some caveats in the pinnings, because cuFile wheels weren't available until 12.6.3 so the `[cufile]` extra isn't reliable. CUDA 13 uses the `cuda-toolkit[cufile]` extra. Authors: - Bradley Dice (https://github.com/bdice) - Gil Forsyth (https://github.com/gforsyth) Approvers: - Gil Forsyth (https://github.com/gforsyth) URL: https://github.com/rapidsai/cudf/pull/23392 --- dependencies.yaml | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/dependencies.yaml b/dependencies.yaml index c2041b67228e..b2d75476716b 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -774,21 +774,27 @@ dependencies: use_cuda_wheels: "true" packages: - cuda-toolkit==12.2.* + # Pin the oldest available cuFile wheel, released with CUDA Toolkit 12.6.3. + - nvidia-cufile-cu12==1.11.1.6.* - matrix: cuda: "12.5" use_cuda_wheels: "true" packages: - cuda-toolkit==12.5.* + # Pin the oldest available cuFile wheel, released with CUDA Toolkit 12.6.3. + - nvidia-cufile-cu12==1.11.1.6.* - matrix: cuda: "12.8" use_cuda_wheels: "true" packages: - cuda-toolkit==12.8.* + - nvidia-cufile-cu12==1.13.1.3.* - matrix: cuda: "12.9" use_cuda_wheels: "true" packages: - cuda-toolkit==12.9.* + - nvidia-cufile-cu12>=1.13.1.3 - matrix: cuda: "13.0" use_cuda_wheels: "true" @@ -977,12 +983,14 @@ dependencies: use_cuda_wheels: "true" packages: - cuda-toolkit[nvcc,nvrtc]==12.* + # cuda-toolkit before 12.6.3 does not have the `cufile` extra. + - nvidia-cufile-cu12 - matrix: cuda: "13.*" cuda_suffixed: "true" use_cuda_wheels: "true" packages: - - cuda-toolkit[nvcc,nvrtc]==13.* + - cuda-toolkit[cufile,nvcc,nvrtc]==13.* - {matrix: null, packages: []} run_cudf_polars: common: From b6b58b37bc14838565f8064f6521db1751695ce7 Mon Sep 17 00:00:00 2001 From: paul-aiyedun <53453937+paul-aiyedun@users.noreply.github.com> Date: Wed, 22 Jul 2026 10:33:01 -0700 Subject: [PATCH 05/16] Include sources.jar and javadoc.jar in the Java Maven repo output (#23389) Align the `java-build` and `java-gather` artifact set with what Maven Central requires for a release (main classifier JAR, sources.jar, javadoc.jar, POM). Authors: - https://github.com/paul-aiyedun Approvers: - Tim Liu (https://github.com/NvTimLiu) URL: https://github.com/rapidsai/cudf/pull/23389 --- java/ci/README.md | 18 ++++-- java/ci/assemble_maven_repo.sh | 47 ++++++++++++-- java/ci/build_cudf_java_jar_in_container.sh | 68 +++++++++++++++++---- 3 files changed, 110 insertions(+), 23 deletions(-) diff --git a/java/ci/README.md b/java/ci/README.md index 219e08e5eaff..45a019f7d408 100644 --- a/java/ci/README.md +++ b/java/ci/README.md @@ -40,13 +40,16 @@ so plain `rm -rf` works. --cuda-version 12.9 ``` -This compiles the JNI layer against the static libcudf from Step 1 and emits a -single classifier JAR (e.g. `cudf-26.08.0-SNAPSHOT-cuda12.jar`) plus its POM -into a classifier-named subdirectory under `--output-dir`: +This compiles the JNI layer against the static libcudf from Step 1 and emits +the classifier JAR (e.g. `cudf-26.08.0-SNAPSHOT-cuda12.jar`), a +classifier-independent sources jar and javadoc jar, and the POM into a +classifier-named subdirectory under `--output-dir`: ``` /tmp/jars/cuda12/ cudf-26.08.0-SNAPSHOT-cuda12.jar + cudf-26.08.0-SNAPSHOT-sources.jar + cudf-26.08.0-SNAPSHOT-javadoc.jar cudf-26.08.0-SNAPSHOT.pom ``` @@ -68,14 +71,17 @@ inside the container. ``` This walks every subdirectory of `--jars-dir` (each subdir name IS the -classifier), gathers the per-classifier JAR and shared POM, derives the -artifact version from the JAR filenames (requiring a single unique version -across subdirs), and lays them out as: +classifier), gathers the per-classifier JAR, one shared sources jar, one +shared javadoc jar, and the shared POM, derives the artifact version from +the JAR filenames (requiring a single unique version across subdirs), and +lays them out as: ``` /tmp/maven-repo/ai/rapids/cudf/26.08.0-SNAPSHOT/ cudf-26.08.0-SNAPSHOT-cuda12.jar cudf-26.08.0-SNAPSHOT-cuda13.jar + cudf-26.08.0-SNAPSHOT-sources.jar + cudf-26.08.0-SNAPSHOT-javadoc.jar cudf-26.08.0-SNAPSHOT.pom ``` diff --git a/java/ci/assemble_maven_repo.sh b/java/ci/assemble_maven_repo.sh index b1d3e65a457c..96bf8c8f33f3 100755 --- a/java/ci/assemble_maven_repo.sh +++ b/java/ci/assemble_maven_repo.sh @@ -6,12 +6,17 @@ # Maven-repository-layout directory. # # Input: --jars-dir contains one subdirectory per classifier, each holding -# exactly one cudf--.jar and a cudf-.pom. Subdir -# names ARE the classifier names, and the artifact version is derived from -# the JAR filenames (all subdirs must agree). +# - exactly one cudf--.jar +# - one cudf--sources.jar +# - one cudf--javadoc.jar +# - one cudf-.pom +# Subdir names ARE the classifier names, and the artifact version is derived +# from the JAR filenames (all subdirs must agree). # # Output layout: # /ai/rapids/cudf//cudf--.jar +# /ai/rapids/cudf//cudf--sources.jar +# /ai/rapids/cudf//cudf--javadoc.jar # /ai/rapids/cudf//cudf-.pom set -e @@ -36,8 +41,12 @@ Gathers per-classifier cuDF Java JARs into a single Maven-repository-layout tree REQUIRED: -j, --jars-dir Parent directory containing one subdirectory per - classifier (each holding cudf--.jar - and cudf-.pom). Subdir name is the classifier. + classifier. Each subdir must hold: + cudf--.jar + cudf--sources.jar + cudf--javadoc.jar + cudf-.pom + Subdir name is the classifier. -o, --output-dir Directory to receive the combined Maven-repository layout. OPTIONS: @@ -48,6 +57,8 @@ EXAMPLE: # given /tmp/jars/{cuda12,cuda13}/ inputs, produces: # /tmp/maven-repo/ai/rapids/cudf//cudf--cuda12.jar # /tmp/maven-repo/ai/rapids/cudf//cudf--cuda13.jar + # /tmp/maven-repo/ai/rapids/cudf//cudf--sources.jar + # /tmp/maven-repo/ai/rapids/cudf//cudf--javadoc.jar # /tmp/maven-repo/ai/rapids/cudf//cudf-.pom EOF @@ -155,6 +166,32 @@ if [[ -z "${FIRST_VERSION}" ]]; then exit 1 fi +# Sources and javadoc jars are classifier-independent (pure Java, no arch or +# cuda variation). Every classifier subdir produces byte-equivalent copies; +# pick the lexicographically first subdir's copy as canonical. Fail fast if +# any subdir is missing either file - that indicates -Prelease or +# -Pjavadoc-jdk17 did not activate for that classifier's build. +DEST_DIR="${OUTPUT_DIR}/${GROUP_PATH}/${ARTIFACT_ID}/${FIRST_VERSION}" +FIRST_CLASSIFIER_SUBDIR="" +for subdir in "${JARS_DIR}"/*/; do + if [[ -z ${FIRST_CLASSIFIER_SUBDIR} ]]; then + FIRST_CLASSIFIER_SUBDIR=${subdir} + fi + for required in "cudf-${FIRST_VERSION}-sources.jar" \ + "cudf-${FIRST_VERSION}-javadoc.jar"; do + if [[ ! -f ${subdir}${required} ]]; then + echo "Error: ${required} missing from ${subdir} (expected in every classifier subdir)" >&2 + exit 1 + fi + done +done + +for shared in "cudf-${FIRST_VERSION}-sources.jar" \ + "cudf-${FIRST_VERSION}-javadoc.jar"; do + cp -f "${FIRST_CLASSIFIER_SUBDIR}${shared}" "${DEST_DIR}/" + echo " + ${shared}" +done + # POM is identical across subdirs; copy the first one found. POM_SRC="" for subdir in "${JARS_DIR}"/*/; do diff --git a/java/ci/build_cudf_java_jar_in_container.sh b/java/ci/build_cudf_java_jar_in_container.sh index 98fce46c4f4a..e02ea3171508 100755 --- a/java/ci/build_cudf_java_jar_in_container.sh +++ b/java/ci/build_cudf_java_jar_in_container.sh @@ -6,9 +6,10 @@ # # This script runs inside the rapidsai/ci-conda container launched by # java/ci/build_cudf_java_jar.sh. It generates the build_java conda toolchain -# environment, compiles the JNI layer against a prebuilt static libcudf -# (mounted at /libcudf), and packages the cuDF Java JAR. The resulting -# classifier JAR and its POM are copied to /output. /output and +# environment, installs a JDK 17 side-prefix for the javadoc-jdk17 profile, +# compiles the JNI layer against a prebuilt static libcudf (mounted at +# /libcudf), and packages the cuDF Java JAR. The resulting classifier JAR, +# sources JAR, javadoc JAR, and POM are copied to /output. /output and # /repo/java/target are chowned to HOST_UID:HOST_GID on exit so the host user # owns the outputs. # @@ -63,6 +64,15 @@ conda activate build_java rapids-print-env +# The `javadoc-jdk17` profile in java/pom.xml points +# at ${env.JDK17_HOME}/bin/javadoc. The build_java env +# above provides only JDK 8 (mvn's own JVM), so install JDK 17 into a +# dedicated prefix that JDK17_HOME can point to. The primary mvn JVM stays +# on JDK 8; only the javadoc binary is invoked from this prefix. +rapids-logger "Installing JDK 17 into /opt/jdk17 for javadoc-jdk17 profile" +rapids-mamba-retry create --yes --prefix /opt/jdk17 openjdk=17.* +export JDK17_HOME=/opt/jdk17 + if [[ -z ${CUDACXX} ]]; then export CUDACXX="${CONDA_PREFIX}/bin/nvcc" fi @@ -72,12 +82,22 @@ fi BUILD_ARG=( -B + # Prefix every log line with HH:mm:ss.SSS so the elapsed time of individual + # plugin executions is recorded. + "-Dorg.slf4j.simpleLogger.showDateTime=true" + "-Dorg.slf4j.simpleLogger.dateTimeFormat=HH:mm:ss.SSS" "-Dmaven.repo.local=/tmp/.m2" "-Dparallel.level=${PARALLEL_LEVEL}" "-DskipTests=true" "-DCUDF_USE_PER_THREAD_DEFAULT_STREAM=ON" "-DCUDF_JNI_LIBCUDF_STATIC=ON" "-DUSE_GDS=OFF" + # -Prelease produces the sources.jar file via maven-source-plugin; + # -Pjavadoc-jdk17 produces the javadoc.jar file via maven-javadoc-plugin + # running against ${env.JDK17_HOME}/bin/javadoc. Both are required by + # Maven Central for every published release. + "-Prelease" + "-Pjavadoc-jdk17" ) if [[ -n ${CMAKE_CUDA_ARCHITECTURES} ]]; then @@ -95,27 +115,51 @@ rapids-logger "Packaging cuDF Java JAR version ${CUDF_VERSION} (libcudf: ${CUDF_ # the scratch dir before each container launch to guarantee target/ starts empty. CUDF_INSTALL_DIR="${CUDF_INSTALL_DIR}" mvn package "${BUILD_ARG[@]}" +mkdir -p "${OUTPUT_DIR}" + +# Order matters: *-test-sources.jar must precede *-sources.jar because +# bash's case picks the first matching pattern, and *-sources.jar would +# also match *-test-sources.jar. MAIN_JAR="" for candidate in target/cudf-"${CUDF_VERSION}"-*.jar; do case "${candidate}" in - *-tests.jar|*-sources.jar|*-javadoc.jar) + *-tests.jar|*-test-sources.jar) continue ;; + *-sources.jar|*-javadoc.jar) + cp -f "${candidate}" "${OUTPUT_DIR}/" + ;; + *) + if [[ -f ${candidate} ]]; then + if [[ -n ${MAIN_JAR} ]]; then + echo "Error: multiple main classifier JARs matched under target/" >&2 + ls -l target/ >&2 || true + exit 1 + fi + MAIN_JAR=${candidate} + fi + ;; esac - if [[ -f ${candidate} ]]; then - MAIN_JAR=${candidate} - break - fi done if [[ -z ${MAIN_JAR} ]]; then - echo "Error: no cuDF classifier JAR produced under target/" - ls -l target/ || true + echo "Error: no cuDF classifier JAR produced under target/" >&2 + ls -l target/ >&2 || true exit 1 fi -mkdir -p "${OUTPUT_DIR}" +# Assert the release-profile artifacts landed. A missing file here means +# -Prelease or -Pjavadoc-jdk17 did not activate, or JDK17_HOME did not +# resolve to a usable javadoc binary. +for required in "${OUTPUT_DIR}/cudf-${CUDF_VERSION}-sources.jar" \ + "${OUTPUT_DIR}/cudf-${CUDF_VERSION}-javadoc.jar"; do + if [[ ! -f ${required} ]]; then + echo "Error: expected ${required} not found (mvn -Prelease -Pjavadoc-jdk17 did not produce it)" >&2 + exit 1 + fi +done + cp -f "${MAIN_JAR}" "${OUTPUT_DIR}/" cp -f pom.xml "${OUTPUT_DIR}/cudf-${CUDF_VERSION}.pom" -rapids-logger "Emitted $(basename "${MAIN_JAR}") + cudf-${CUDF_VERSION}.pom to ${OUTPUT_DIR}" +rapids-logger "Emitted $(basename "${MAIN_JAR}"), cudf-${CUDF_VERSION}-sources.jar, cudf-${CUDF_VERSION}-javadoc.jar, and cudf-${CUDF_VERSION}.pom to ${OUTPUT_DIR}" From 8b4c4a54b0221511cc36fe7e15808626c0a9e057 Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Wed, 22 Jul 2026 17:24:47 -0500 Subject: [PATCH 06/16] Preserve MultiIndex column fidelity through ColumnAccessor round trips (#23365) Split out of #23255 (2/6). Rebuilding a frame's pandas columns `MultiIndex` from tuples re-sorts the levels and re-infers their dtypes, losing the exact source layout: unsorted explicit level orders (which change the behavior of pandas operations that work on level codes, e.g. legacy `stack(sort=True)` after a fast-to-slow conversion under `cudf.pandas`), categorical/object/int64 level dtypes (int64 levels with missing entries upcast to float64), and NaN column labels (fresh `float('nan')` objects hash unequal, so lookups miss). - Prime the cached `to_pandas_index` with the exact source `pd.MultiIndex` at `DataFrame` construction and propagate it through accessor copies. - Restore recorded per-level dtypes in `to_pandas_index` when the cast round-trips losslessly. - Match NaN-containing column labels under pandas' all-NaNs-equal semantics. - Read level dtypes off `MultiIndex.levels` (`get_level_values` materializes missing entries as NaN and upcasts), also for `cudf.MultiIndex` columns. - Keep hierarchical columns through DataFrame binops when only level dtypes differ (restored `Int8` vs `int64` fails `Index.equals`). Fixes 13 pandas-tests (constructor dict-NaN-key, concat keys with specific levels, groupby ordered multi-func aggregate, MultiIndex loc, and several `test_stack_unstack.py` cases); their xfail entries are removed. Attribution verified by running each node id against an isolated build containing only this change (pass) and a clean build (fail). Independent of the other #23255 split PRs, but the stack (3/6) and unstack (4/6) PRs depend on this one. Authors: - GALI PREM SAGAR (https://github.com/galipremsagar) Approvers: - Vyas Ramasubramani (https://github.com/vyasr) URL: https://github.com/rapidsai/cudf/pull/23365 --- python/cudf/cudf/core/column_accessor.py | 109 +++++++++++++++++- python/cudf/cudf/core/dataframe.py | 27 ++++- .../pandas/scripts/pandas-testing-plugin.py | 13 --- .../private_objects/test_column_accessor.py | 75 +++++++++++- 4 files changed, 207 insertions(+), 17 deletions(-) diff --git a/python/cudf/cudf/core/column_accessor.py b/python/cudf/cudf/core/column_accessor.py index 5944e5486fe3..9c480812c046 100644 --- a/python/cudf/cudf/core/column_accessor.py +++ b/python/cudf/cudf/core/column_accessor.py @@ -35,6 +35,36 @@ def _is_bool(val: Any) -> bool: return isinstance(val, (bool, np.bool_)) +def _is_nan_scalar(val: Any) -> bool: + return isinstance(val, (float, np.floating)) and val != val + + +def _label_contains_nan(label: Any) -> bool: + if isinstance(label, tuple): + return any(_is_nan_scalar(lv) for lv in label) + return _is_nan_scalar(label) + + +def _canonicalize_nan_label(label: Any) -> Any: + """Map float NaN elements of a label to the np.nan singleton. + + No NaN object ever compares equal to any NaN (including itself), so + equality alone can never match a NaN-containing label. dict lookups and + tuple comparison, however, short-circuit on object *identity* before + trying ``==`` — and that shortcut is exactly what this canonicalization + targets: with every NaN mapped to the one ``np.nan`` object, two + canonicalized labels holding NaN in the same position match by identity. + This restores pandas' all-NaNs-are-equal label semantics for labels + round-tripped through a pandas Index, which materializes fresh NaN + objects on iteration. + """ + if isinstance(label, tuple): + return tuple(np.nan if _is_nan_scalar(lv) else lv for lv in label) + if _is_nan_scalar(label): + return np.nan + return label + + class _NestedGetItemDict(dict): """A dictionary whose __getitem__ method accesses nested dicts. @@ -102,6 +132,10 @@ class ColumnAccessor(MutableMapping): verify : bool, optional For non ColumnAccessor inputs, whether to verify column length and data.values() are all Columns + pandas_index : pd.Index, optional + The source pandas index the keys were taken from, if any. + A matching pd.MultiIndex primes the ``to_pandas_index`` cache + (see ``_prime_to_pandas_index``). """ _data: dict[Hashable, ColumnBase] @@ -117,6 +151,7 @@ def __init__( label_dtype: DtypeObj | None = None, verify: bool = True, level_dtypes: tuple[DtypeObj, ...] | None = None, + pandas_index: pd.Index | None = None, ) -> None: if isinstance(data, ColumnAccessor): self._data = data._data @@ -125,6 +160,11 @@ def __init__( self.rangeindex: bool = data.rangeindex self.label_dtype: DtypeObj | None = data.label_dtype self._level_dtypes = data._level_dtypes + if "to_pandas_index" in data.__dict__: + # carry over the primed/cached pandas index: it holds + # fidelity (e.g. explicit unsorted level order) that a + # rebuild from tuples would lose + self.to_pandas_index = data.__dict__["to_pandas_index"] elif isinstance(data, MutableMapping): # This code path is performance-critical for copies and should be # modified with care. @@ -157,12 +197,50 @@ def __init__( raise ValueError( f"data must be a ColumnAccessor or MutableMapping, not {type(data).__name__}" ) + if pandas_index is not None: + self._prime_to_pandas_index(pandas_index) + + def _prime_to_pandas_index(self, index: pd.Index) -> None: + """Prime the cached ``to_pandas_index`` with the exact source index. + + Rebuilding a pandas MultiIndex from the stored tuple labels re-sorts + its levels, losing an explicit unsorted level layout (the level + order affects pandas operations that work on level codes, e.g. + legacy ``stack(sort=True)``). Keeping the source MultiIndex itself + preserves that fidelity. Only a hierarchical columns axis whose + length matches the data is primed; anything else is ignored. + """ + if ( + self.multiindex + and isinstance(index, pd.MultiIndex) + and len(self._data) == len(index) + ): + self.to_pandas_index = index def __iter__(self) -> Iterator: return iter(self._data) def __getitem__(self, key: Hashable) -> ColumnBase: - return self._data[key] + try: + return self._data[key] + except KeyError: + if _label_contains_nan(key): + # NaN labels lose object identity when round-tripped through + # a pandas Index; retry with NaNs canonicalized so all NaNs + # match by identity, as pandas label semantics require. + canon = _canonicalize_nan_label(key) + for existing in self._data: + c = _canonicalize_nan_label(existing) + try: + match = c is canon or bool(c == canon) + except TypeError: + # e.g. a pd.NA label: its comparisons return pd.NA, + # whose truthiness raises. Ambiguity is not a match; + # keep scanning for a genuine NaN label. + match = False + if match: + return self._data[existing] + raise def __setitem__(self, key: Hashable, value: ColumnBase) -> None: self.set_by_label(key, value) @@ -322,6 +400,35 @@ def to_pandas_index(self) -> pd.Index: self.names, names=self.level_names, ) + if ( + self._level_dtypes is not None + and len(self._level_dtypes) == result.nlevels + ): + # ``from_tuples`` re-infers every level dtype from the + # materialized labels, degrading e.g. categorical levels + # to str, object levels to str once mixed-type labels are + # selected away, and int64 levels with missing entries to + # float64. Restore each preserved level dtype when the + # cast is lossless (round-trips to the inferred values). + new_levels = [] + changed = False + for lvl, level_dtype in zip( + result.levels, self._level_dtypes, strict=True + ): + if lvl.dtype != level_dtype: + try: + cast_lvl = lvl.astype(level_dtype) + except (TypeError, ValueError): + pass + else: + # missing-aware equality: ``==`` would treat + # NaN entries as unequal to themselves + if cast_lvl.astype(lvl.dtype).equals(lvl): + lvl = cast_lvl + changed = True + new_levels.append(lvl) + if changed: + result = result.set_levels(new_levels) else: # Determine if we can return a RangeIndex if self.rangeindex: diff --git a/python/cudf/cudf/core/dataframe.py b/python/cudf/cudf/core/dataframe.py index 02cdd9da3473..a9439c730348 100644 --- a/python/cudf/cudf/core/dataframe.py +++ b/python/cudf/cudf/core/dataframe.py @@ -615,7 +615,12 @@ def _pd_index_level_dtypes(idx) -> tuple | None: dtype cannot be inferred from zero entries). """ if isinstance(idx, pd.MultiIndex): - return tuple(idx.get_level_values(i).dtype for i in range(idx.nlevels)) + # use the levels, not get_level_values: materializing a level that + # has missing entries (-1 codes) upcasts e.g. int64 to float64 + return tuple(level.dtype for level in idx.levels) + if isinstance(idx, cudf.MultiIndex): + # the per-row columns share their dtype with the levels + return tuple(dtype for _, dtype in idx._dtypes) return None @@ -837,7 +842,7 @@ def _array_to_column_accessor( columns_labels = columns else: columns_labels = pd.RangeIndex(data.shape[1]) - return ColumnAccessor( + ca = ColumnAccessor( { column_label: as_column(data[:, i], nan_as_null=nan_as_null) for column_label, i in zip( @@ -850,7 +855,9 @@ def _array_to_column_accessor( label_dtype=columns_labels.dtype, level_names=tuple(columns_labels.names), level_dtypes=_pd_index_level_dtypes(columns_labels), + pandas_index=columns_labels, ) + return ca @_performance_tracking @@ -1280,6 +1287,7 @@ def __init__( level_names=tuple(columns.names), label_dtype=columns.dtype, level_dtypes=_pd_index_level_dtypes(columns), + pandas_index=columns, ) elif isinstance(data, Mapping): # Note: We excluded ColumnAccessor already above @@ -1352,6 +1360,12 @@ def __init__( if dtype: self._data = self.astype(dtype)._data + final_pd_columns = ( + second_columns if second_columns is not None else columns + ) + if isinstance(final_pd_columns, pd.Index): + self._data._prime_to_pandas_index(final_pd_columns) + @classmethod def _from_data( # type: ignore[override] cls, @@ -2554,6 +2568,11 @@ def _fill_same_ca_attributes( ) elif self._data._level_names == other._data._level_names: ca_attributes["level_names"] = self._data.level_names + if self._data.multiindex == other._data.multiindex: + # equal labels can still fail the ``equals`` check above + # on level-dtype differences (e.g. Int8 vs int64); the + # result keeps hierarchical columns like pandas + ca_attributes["multiindex"] = self._data.multiindex elif isinstance(other, (dict, Mapping)): # Need to fail early on host mapping types because we ultimately # convert everything to a dict. @@ -3237,6 +3256,7 @@ def columns(self, columns): rangeindex = False label_dtype = None level_names = None + level_dtypes = None if isinstance(columns, (pd.MultiIndex, cudf.MultiIndex)): multiindex = True if isinstance(columns, cudf.MultiIndex): @@ -3246,6 +3266,7 @@ def columns(self, columns): if pd_columns.nunique(dropna=False) != len(pd_columns): raise ValueError("Duplicate column names are not allowed") level_names = list(pd_columns.names) + level_dtypes = _pd_index_level_dtypes(pd_columns) elif isinstance(columns, (Index, ColumnBase, Series)): level_names = (getattr(columns, "name", None),) rangeindex = isinstance(columns, cudf.RangeIndex) @@ -3282,7 +3303,9 @@ def columns(self, columns): level_names=level_names, label_dtype=label_dtype, rangeindex=rangeindex, + level_dtypes=level_dtypes, verify=False, + pandas_index=pd_columns, ) def _set_columns_like(self, other: ColumnAccessor) -> None: diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index 3e219a23e8ec..98b9d91de5d4 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -1391,9 +1391,6 @@ def pytest_unconfigure(config): "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_cast": "TODO: Add a reason for failure", "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_multiindex": "TODO: Add a reason for failure", "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_nan_key[None]": "TODO: Add a reason for failure", - "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_nan_key[nan0]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='nan') are different", - "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_nan_key[nan1]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='nan') are different", - "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_nan_key_and_columns": "TODO: Add a reason for failure", "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_with_index": "TODO: Add a reason for failure", "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_with_index_and_columns": "TODO: Add a reason for failure", "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_with_none": "AssertionError: assert nan is None", @@ -1720,7 +1717,6 @@ def pytest_unconfigure(config): "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_in_multiindex_columns[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_in_multiindex_columns[True]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_level[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_order_with_unsorted_levels_multi_row_2[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_multiple[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_multiple[True]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_preserve_names[False]": "TODO: Add a reason for failure", @@ -1729,10 +1725,6 @@ def pytest_unconfigure(config): "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_wrong_level_name[True-unstack]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_unstack_preserve_types": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_unstack_with_missing_int_cast_to_float": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=first]": "AssertionError: Attributes of DataFrame.iloc[:, 0] (column name='('value', nan)') are different", - "tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=last]": "AssertionError: Attributes of DataFrame.iloc[:, 3] (column name='('value', nan)') are different", - "tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=second]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='('value', nan)') are different", - "tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=third]": "AssertionError: Attributes of DataFrame.iloc[:, 2] (column name='('value', nan)') are different", "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_asof": "TODO: Add a reason for failure", "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_equals_subclass": "TODO: Add a reason for failure", "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_frame_subclassing_and_slicing": "TODO: Add a reason for failure", @@ -1763,7 +1755,6 @@ def pytest_unconfigure(config): "tests/groupby/aggregate/test_aggregate.py::test_agg_str_with_kwarg_axis_1_raises[nunique]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_aggregate.py::test_groupby_aggregate_directory[size]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_aggregate.py::test_groupby_aggregate_empty_key_empty_return": "TODO: Add a reason for failure", - "tests/groupby/aggregate/test_aggregate.py::test_order_aggregate_multiple_funcs": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-prod-large_int-False]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-prod-large_int-True]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-sum-large_int-False]": "TODO: Add a reason for failure", @@ -2874,7 +2865,6 @@ def pytest_unconfigure(config): "tests/reshape/concat/test_categorical.py::TestCategoricalConcat::test_categorical_index_upcast": "TODO: Add a reason for failure", "tests/reshape/concat/test_categorical.py::TestCategoricalConcat::test_concat_categorical_datetime": "TODO: Add a reason for failure", "tests/reshape/concat/test_concat.py::TestConcatenate::test_concat_copy": "TODO: Add a reason for failure", - "tests/reshape/concat/test_concat.py::TestConcatenate::test_concat_keys_specific_levels": "TODO: Add a reason for failure", "tests/reshape/concat/test_concat.py::TestConcatenate::test_concat_order": "TODO: Add a reason for failure", "tests/reshape/concat/test_concat.py::test_concat_empty_and_non_empty_frame_regression": "TODO: Add a reason for failure", "tests/reshape/concat/test_concat.py::test_concat_ignore_empty_object_float[None-datetime64[ns]]": "AssertionError: Attributes of DataFrame.iloc[:, 0] (column name='foo') are different", @@ -3051,8 +3041,6 @@ def pytest_unconfigure(config): "tests/reshape/test_melt.py::TestWideToLong::test_raise_of_column_name_value": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivot::test_pivot_index_is_none": "AssertionError: DataFrame.index are different", "tests/reshape/test_pivot.py::TestPivotTable::test_categorical_pivot_index_ordering[False]": "TODO: Add a reason for failure", - "tests/reshape/test_pivot.py::TestPivotTable::test_daily": "TODO: Add a reason for failure", - "tests/reshape/test_pivot.py::TestPivotTable::test_monthly": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_complex_aggfunc": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_datetime_tz": "ValueError: Length of names must match number of levels in MultiIndex.", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_index_with_nan[False]": "AssertionError: DataFrame.index are different", @@ -3874,7 +3862,6 @@ def pytest_unconfigure(config): "tests/window/test_timeseries_window.py::TestRollingTS::test_rolling_on_decreasing_index[us]": "TODO: Add a reason for failure", "tests/window/test_win_type.py::test_cmov_window_corner[None]": "TODO: Add a reason for failure", "tests/window/test_win_type.py::test_win_type_not_implemented": "TODO: Add a reason for failure", - "tests/indexing/multiindex/test_loc.py::test_loc_getitem_duplicates_multiindex_empty_indexer[columns_indexer1]": "AssertionError: DataFrame.columns level [0] are different", } # Keep keys in alphabeical order diff --git a/python/cudf/cudf/tests/private_objects/test_column_accessor.py b/python/cudf/cudf/tests/private_objects/test_column_accessor.py index 8ab1a092579b..5570b0880b11 100644 --- a/python/cudf/cudf/tests/private_objects/test_column_accessor.py +++ b/python/cudf/cudf/tests/private_objects/test_column_accessor.py @@ -1,7 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import numpy as np import pandas as pd import pytest @@ -379,3 +380,75 @@ def test_not_rangeindex_and_multiindex(): def test_data_values_not_column_raises(): with pytest.raises(ValueError): ColumnAccessor({"a": [1]}) + + +def test_to_pandas_index_preserves_unsorted_level_order(): + # rebuilding the columns MultiIndex from tuples would re-sort the + # levels; the exact source level layout must round-trip (level order + # affects pandas operations that work on codes, e.g. legacy + # stack(sort=True)) + pmi = pd.MultiIndex( + levels=[["b", "a"], ["y", "x"]], + codes=[[0, 0, 1], [0, 1, 0]], + names=["outer", "inner"], + ) + gdf = cudf.DataFrame([[1, 2, 3]], columns=pmi) + pd.testing.assert_index_equal(gdf._data.to_pandas_index, pmi, exact=True) + + # the primed index survives accessor copies + copied = ColumnAccessor(gdf._data) + pd.testing.assert_index_equal(copied.to_pandas_index, pmi, exact=True) + + +def test_to_pandas_index_restores_int_level_with_missing_entries(): + # an int64 level whose codes contain -1 (missing) materializes as + # float64 via from_tuples; the recorded level dtype must be restored + pmi = pd.MultiIndex( + levels=[["a", "b"], pd.Index([10, 2], dtype="int64")], + codes=[[0, 1], [0, -1]], + ) + gdf = cudf.DataFrame([[1, 2]], columns=pmi) + # a fresh .copy() carries no primed cache, so to_pandas_index must + # rebuild from the tuple labels and restore the recorded level dtype + result = gdf._data.copy().to_pandas_index + assert result.levels[1].dtype == pmi.levels[1].dtype + pd.testing.assert_index_equal(result, pmi, exact=True) + + +def test_getitem_nan_label_any_nan_object(): + # NaN labels round-tripped through a pandas Index materialize fresh + # float('nan') objects, which hash/compare unequal; lookup must match + # them under pandas' all-NaNs-equal label semantics + ca = ColumnAccessor({np.nan: as_column([1, 2]), "x": as_column([3, 4])}) + fresh_nan = float("nan") + assert fresh_nan is not np.nan + assert_eq(ca[fresh_nan], ca[np.nan]) + with pytest.raises(KeyError): + ca["missing"] + + +@pytest.mark.parametrize( + "nan", [np.float16("nan"), np.float32("nan"), np.float64("nan")] +) +def test_getitem_nan_label_numpy_floating(nan): + # NumPy floating NaNs (only np.float64 subclasses python float) must + # canonicalize like python floats, on both sides of the lookup + ca = ColumnAccessor( + {("a", nan): as_column([1]), ("b", "x"): as_column([2])}, + multiindex=True, + ) + assert_eq(ca[("a", float("nan"))], ca[("a", nan)]) + assert_eq(ca[("a", np.nan)], ca[("a", nan)]) + + +def test_getitem_nan_label_skips_ambiguous_pd_na_label(): + # comparing against a pd.NA label yields pd.NA, whose truthiness + # raises; the NaN retry must treat that as a non-match and still find + # the genuine NaN label further on + ca = ColumnAccessor( + {("a", pd.NA): as_column([1]), ("a", np.nan): as_column([2])}, + multiindex=True, + ) + assert_eq(ca[("a", float("nan"))], ca[("a", np.nan)]) + with pytest.raises(KeyError): + ca[("b", float("nan"))] From c16794fc631b2ccc03e4e68b40a62df473ca0279 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:26:59 -0700 Subject: [PATCH 07/16] Experimental Parquet reader supports retention vectors (#23402) This PR enables the experimental parquet reader to now also support retention vectors in addition to deletion vectors Authors: - Muhammad Haseeb (https://github.com/mhaseeb123) Approvers: - Vukasin Milovanovic (https://github.com/vuule) - Jihoon Son (https://github.com/jihoonson) URL: https://github.com/rapidsai/cudf/pull/23402 --- .../cudf/io/experimental/deletion_vectors.hpp | 4 + .../parquet/experimental/deletion_vectors.cu | 25 +- .../io/parquet_deletion_vectors_test.cpp | 231 +++++++++++------- 3 files changed, 163 insertions(+), 97 deletions(-) diff --git a/cpp/include/cudf/io/experimental/deletion_vectors.hpp b/cpp/include/cudf/io/experimental/deletion_vectors.hpp index af3026a671d8..fbab01669b67 100644 --- a/cpp/include/cudf/io/experimental/deletion_vectors.hpp +++ b/cpp/include/cudf/io/experimental/deletion_vectors.hpp @@ -48,6 +48,9 @@ struct deletion_vector_info { std::vector row_group_offsets; /// Number of rows in each row group to be read from the Parquet source(s) std::vector row_group_num_rows; + + /// Whether the roaring bitmaps represent retention vectors + bool are_retention_vectors = false; }; /** @@ -147,6 +150,7 @@ class chunked_parquet_reader { std::queue _deletion_vector_row_counts; size_t _start_row; bool _is_unspecified_row_group_data; + bool _are_retentions; rmm::cuda_stream_view _stream; rmm::device_async_resource_ref _mr; rmm::device_async_resource_ref _table_mr; diff --git a/cpp/src/io/parquet/experimental/deletion_vectors.cu b/cpp/src/io/parquet/experimental/deletion_vectors.cu index 1d5ada49c69a..7af0cd3b3cf5 100644 --- a/cpp/src/io/parquet/experimental/deletion_vectors.cu +++ b/cpp/src/io/parquet/experimental/deletion_vectors.cu @@ -104,11 +104,15 @@ namespace detail { deletion_vector_row_counts, stream, cudf::get_current_device_resource_ref()); + + auto const mask_type = deletion_vector_info.are_retention_vectors + ? cudf::detail::mask_type::RETENTION + : cudf::detail::mask_type::DELETION; + // Filter the table using the deletion vector - return table_with_metadata{ + return cudf::io::table_with_metadata{ // Supply user-provided mr to apply deletion mask to allocate output table's memory - cudf::detail::apply_mask( - table_with_index->view(), row_mask->view(), cudf::detail::mask_type::DELETION, stream, mr), + cudf::detail::apply_mask(table_with_index->view(), row_mask->view(), mask_type, stream, mr), std::move(metadata)}; } @@ -161,7 +165,7 @@ namespace detail { dv_row_counts_queue.push(deletion_vector_row_counts[i]); } - size_t deleted_rows = 0; + size_t matched_rows = 0; size_t remaining_rows = num_rows; size_t start_row = 0; @@ -176,14 +180,15 @@ namespace detail { is_row_group_data_unspecified, stream, cudf::get_current_device_resource_ref()); - deleted_rows += compute_partial_deleted_row_count( + matched_rows += compute_partial_deleted_row_count( row_index_column->view(), dv_queue, dv_row_counts_queue, stream); start_row += chunk_rows; remaining_rows -= chunk_rows; } - return deleted_rows; + // Bitmap hits are deleted rows for deletion vectors, retained rows for retention vectors + return deletion_vector_info.are_retention_vectors ? num_rows - matched_rows : matched_rows; } } // namespace detail @@ -200,6 +205,7 @@ chunked_parquet_reader::chunked_parquet_reader(std::size_t chunk_read_limit, rmm::device_async_resource_ref mr) : _start_row{0}, _is_unspecified_row_group_data{deletion_vector_info.row_group_offsets.empty()}, + _are_retentions{deletion_vector_info.are_retention_vectors}, _stream{stream}, _mr{mr}, // Use default mr for the internal chunked reader and row index column if we will @@ -317,10 +323,11 @@ table_with_metadata chunked_parquet_reader::read_chunk() _deletion_vector_row_counts, _stream, cudf::get_current_device_resource_ref()); - return table_with_metadata{ + auto const mask_type = + _are_retentions ? cudf::detail::mask_type::RETENTION : cudf::detail::mask_type::DELETION; + return cudf::io::table_with_metadata{ // Supply user-provided mr to apply deletion mask to allocate output table's memory - cudf::detail::apply_mask( - table_with_index->view(), row_mask->view(), cudf::detail::mask_type::DELETION, _stream, _mr), + cudf::detail::apply_mask(table_with_index->view(), row_mask->view(), mask_type, _stream, _mr), std::move(metadata)}; } diff --git a/cpp/tests/io/parquet_deletion_vectors_test.cpp b/cpp/tests/io/parquet_deletion_vectors_test.cpp index a41652105e52..c261fe172be7 100644 --- a/cpp/tests/io/parquet_deletion_vectors_test.cpp +++ b/cpp/tests/io/parquet_deletion_vectors_test.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -136,14 +136,17 @@ auto build_expected_row_indices(cudf::host_span row_group_off * @param num_rows Number of rows in the table * @param deletion_probability The probability of a row being deleted * @param row_indices Host vector of row indices + * @param are_retention_vectors Whether to add retained, rather than deleted, row indices to the + * bitmap * - * @return A pair of a deletion vector and a host row mask vector + * @return A pair of a roaring bitmap and a host row mask vector */ -auto build_deletion_vector_and_expected_row_mask(cudf::size_type num_rows, - float deletion_probability, - cudf::host_span row_indices, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +auto build_roaring_bitmap_and_expected_row_mask(cudf::size_type num_rows, + float deletion_probability, + cudf::host_span row_indices, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr, + bool are_retention_vectors = false) { static constexpr auto seed = 0xbaLL; std::mt19937 engine{seed}; @@ -164,8 +167,9 @@ auto build_deletion_vector_and_expected_row_mask(cudf::size_type num_rows, std::for_each(cuda::counting_iterator(0), cuda::counting_iterator(num_rows), [&](auto row_idx) { - // Insert provided host row index if the row is deleted in the row mask - if (not expected_row_mask[row_idx]) { + // For retention vectors, retain rows selected in the row mask. For deletion + // vectors, delete the remaining rows. + if (expected_row_mask[row_idx] == are_retention_vectors) { roaring::api::roaring64_bitmap_add_bulk( deletion_vector, &roaring64_context, row_indices[row_idx]); } @@ -222,7 +226,7 @@ std::unique_ptr build_expected_table( * well as the `cudf::io::parquet::experimental::chunked_parquet_reader` */ template -void test_read_parquet_and_apply_deletion_vector( +void test_read_parquet_and_apply_mask( cudf::host_span parquet_buffer, cudf::io::parquet::experimental::deletion_vector_info const& deletion_vector_info, cudf::table_view const& input_table_view, @@ -245,6 +249,7 @@ void test_read_parquet_and_apply_deletion_vector( .deletion_vector_row_counts = std::vector(num_concat, input_table_view.num_rows()), .row_group_offsets = {}, .row_group_num_rows = {}, + .are_retention_vectors = deletion_vector_info.are_retention_vectors, }; // Vector to hold the Parquet buffer spans @@ -296,8 +301,13 @@ void test_read_parquet_and_apply_deletion_vector( local_expected_row_indices, cudf::type_id::UINT64, stream, mr); auto [local_deletion_vector, local_expected_row_mask_column] = - build_deletion_vector_and_expected_row_mask( - num_input_rows, deletion_probability, local_expected_row_indices, stream, mr); + build_roaring_bitmap_and_expected_row_mask( + num_input_rows, + deletion_probability, + local_expected_row_indices, + stream, + mr, + final_deletion_vector_info.are_retention_vectors); // Insert the expected table, the corresponding deletion vector and its data span tables.emplace_back(build_expected_table(input_table_view, @@ -383,37 +393,58 @@ TEST_F(ParquetDeletionVectorsTest, NoRowIndexColumn) auto expected_row_index_column = build_column_from_host_data( expected_row_indices, cudf::type_id::UINT64, stream, mr); - // Build deletion vector and the expected row mask column - auto [deletion_vector, expected_row_mask_column] = build_deletion_vector_and_expected_row_mask( - num_rows, deletion_probability, expected_row_indices, stream, mr); - // Use num_concat = 1 here since the row index column is simply a sequence and input table // concatenation won't properly reset it. auto constexpr num_concat = 1; - cudf::io::parquet::experimental::deletion_vector_info deletion_vector_info{ - .serialized_roaring_bitmaps = {deletion_vector}, - .deletion_vector_row_counts = {input_table->view().num_rows()}}; - test_read_parquet_and_apply_deletion_vector(parquet_buffer, - deletion_vector_info, - input_table->view(), - expected_row_mask_column->view(), - expected_row_index_column->view(), - stream, - mr); + + // Build deletion vector and the expected row mask column + { + auto [deletion_vector, expected_row_mask_column] = build_roaring_bitmap_and_expected_row_mask( + num_rows, deletion_probability, expected_row_indices, stream, mr); + auto deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{ + .serialized_roaring_bitmaps = {deletion_vector}, + .deletion_vector_row_counts = {input_table->view().num_rows()}}; + test_read_parquet_and_apply_mask(parquet_buffer, + deletion_vector_info, + input_table->view(), + expected_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); + } + + // Retention vectors retain the rows present in the bitmap. + { + auto constexpr are_retention_vectors = true; + auto [retention_vector, retention_row_mask_column] = build_roaring_bitmap_and_expected_row_mask( + num_rows, deletion_probability, expected_row_indices, stream, mr, are_retention_vectors); + auto retention_vector_info = cudf::io::parquet::experimental::deletion_vector_info{ + .serialized_roaring_bitmaps = {retention_vector}, + .deletion_vector_row_counts = {input_table->view().num_rows()}, + .are_retention_vectors = are_retention_vectors}; + test_read_parquet_and_apply_mask(parquet_buffer, + retention_vector_info, + input_table->view(), + retention_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); + } + // Test no row index column and no deletion vector { // Build expected row mask column containing all true values auto expected_row_mask = thrust::host_vector(num_rows, true); auto expected_row_mask_column = build_column_from_host_data(expected_row_mask, cudf::type_id::BOOL8, stream, mr); - deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{}; - test_read_parquet_and_apply_deletion_vector(parquet_buffer, - deletion_vector_info, - input_table->view(), - expected_row_mask_column->view(), - expected_row_index_column->view(), - stream, - mr); + auto deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{}; + test_read_parquet_and_apply_mask(parquet_buffer, + deletion_vector_info, + input_table->view(), + expected_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); } } @@ -477,7 +508,7 @@ TEST_F(ParquetDeletionVectorsTest, CustomRowIndexColumn) expected_row_indices, cudf::type_id::UINT64, stream, mr); // Build deletion vector and the expected row mask column - auto [deletion_vector, expected_row_mask_column] = build_deletion_vector_and_expected_row_mask( + auto [deletion_vector, expected_row_mask_column] = build_roaring_bitmap_and_expected_row_mask( num_rows, deletion_probability, expected_row_indices, stream, mr); // Don't concatenate the input table and test with single deletion vector @@ -486,38 +517,58 @@ TEST_F(ParquetDeletionVectorsTest, CustomRowIndexColumn) .deletion_vector_row_counts = {input_table->view().num_rows()}, .row_group_offsets = row_group_offsets, .row_group_num_rows = row_group_num_rows}; - test_read_parquet_and_apply_deletion_vector<1>(parquet_buffer, - deletion_vector_info, - input_table->view(), - expected_row_mask_column->view(), - expected_row_index_column->view(), - stream, - mr); + test_read_parquet_and_apply_mask<1>(parquet_buffer, + deletion_vector_info, + input_table->view(), + expected_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); // Concatenate input table and test with multiple deletion vectors - test_read_parquet_and_apply_deletion_vector<4>(parquet_buffer, - deletion_vector_info, - input_table->view(), - expected_row_mask_column->view(), - expected_row_index_column->view(), - stream, - mr); + test_read_parquet_and_apply_mask<4>(parquet_buffer, + deletion_vector_info, + input_table->view(), + expected_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); // Concatenate input table and test with many deletion vectors (>= stream fork threshold of 8) - test_read_parquet_and_apply_deletion_vector<8>(parquet_buffer, - deletion_vector_info, - input_table->view(), - expected_row_mask_column->view(), - expected_row_index_column->view(), - stream, - mr); - test_read_parquet_and_apply_deletion_vector<16>(parquet_buffer, - deletion_vector_info, - input_table->view(), - expected_row_mask_column->view(), - expected_row_index_column->view(), - stream, - mr); + test_read_parquet_and_apply_mask<8>(parquet_buffer, + deletion_vector_info, + input_table->view(), + expected_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); + test_read_parquet_and_apply_mask<16>(parquet_buffer, + deletion_vector_info, + input_table->view(), + expected_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); + + // Retention vectors with custom row indices. + { + auto constexpr are_retention_vectors = true; + auto [retention_vector, retention_row_mask_column] = build_roaring_bitmap_and_expected_row_mask( + num_rows, deletion_probability, expected_row_indices, stream, mr, are_retention_vectors); + deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{ + .serialized_roaring_bitmaps = {retention_vector}, + .deletion_vector_row_counts = {input_table->view().num_rows()}, + .row_group_offsets = row_group_offsets, + .row_group_num_rows = row_group_num_rows, + .are_retention_vectors = are_retention_vectors}; + test_read_parquet_and_apply_mask<1>(parquet_buffer, + deletion_vector_info, + input_table->view(), + retention_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); + } // Test custom row index column and no deletion vector { @@ -528,13 +579,13 @@ TEST_F(ParquetDeletionVectorsTest, CustomRowIndexColumn) auto deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{ .row_group_offsets = row_group_offsets, .row_group_num_rows = row_group_num_rows}; - test_read_parquet_and_apply_deletion_vector<1>(parquet_buffer, - deletion_vector_info, - input_table->view(), - expected_row_mask_column->view(), - expected_row_index_column->view(), - stream, - mr); + test_read_parquet_and_apply_mask<1>(parquet_buffer, + deletion_vector_info, + input_table->view(), + expected_row_mask_column->view(), + expected_row_index_column->view(), + stream, + mr); } } @@ -562,24 +613,28 @@ TEST_F(DeletionVectorsCountTests, NoRowIndex) auto row_indices = thrust::host_vector(num_rows); std::iota(row_indices.begin(), row_indices.end(), size_t{0}); - auto [deletion_vector, expected_row_mask_column] = build_deletion_vector_and_expected_row_mask( - num_rows, deletion_probability, row_indices, stream, mr); + for (auto const are_retention_vectors : {false, true}) { + auto [deletion_vector, expected_row_mask_column] = build_roaring_bitmap_and_expected_row_mask( + num_rows, deletion_probability, row_indices, stream, mr, are_retention_vectors); - auto const expected_row_mask = cudf::detail::make_host_vector( - cudf::device_span(expected_row_mask_column->view().data(), num_rows), stream); - auto const expected_deleted = - std::count(expected_row_mask.begin(), expected_row_mask.end(), false); + auto const expected_row_mask = cudf::detail::make_host_vector( + cudf::device_span(expected_row_mask_column->view().data(), num_rows), + stream); + auto const expected_deleted = + std::count(expected_row_mask.begin(), expected_row_mask.end(), false); - auto deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{ - .serialized_roaring_bitmaps = {deletion_vector}, - .deletion_vector_row_counts = {num_rows}, - .row_group_offsets = {}, - .row_group_num_rows = {}}; - - for (auto chunk_size : {num_rows, num_rows / 2}) { - auto const result = cudf::io::parquet::experimental::compute_num_deleted_rows( - deletion_vector_info, chunk_size, stream); - EXPECT_EQ(result, expected_deleted); + auto deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{ + .serialized_roaring_bitmaps = {deletion_vector}, + .deletion_vector_row_counts = {num_rows}, + .row_group_offsets = {}, + .row_group_num_rows = {}, + .are_retention_vectors = are_retention_vectors}; + + for (auto chunk_size : {num_rows, num_rows / 2}) { + auto const result = cudf::io::parquet::experimental::compute_num_deleted_rows( + deletion_vector_info, chunk_size, stream); + EXPECT_EQ(result, expected_deleted); + } } } @@ -628,7 +683,7 @@ TEST_F(DeletionVectorsCountTests, CustomRowIndex) auto expected_row_indices = build_expected_row_indices(row_group_offsets, row_group_num_rows, num_rows); - auto [deletion_vector, expected_row_mask_column] = build_deletion_vector_and_expected_row_mask( + auto [deletion_vector, expected_row_mask_column] = build_roaring_bitmap_and_expected_row_mask( num_rows, deletion_probability, expected_row_indices, stream, mr); auto const expected_row_mask = cudf::detail::make_host_vector( @@ -681,7 +736,7 @@ TEST_F(DeletionVectorsCountTests, MultipleDeletionVectors) auto local_indices = cudf::host_span(expected_row_indices.data() + span_start, num_rows_per_dv); - auto [dv, mask_col] = build_deletion_vector_and_expected_row_mask( + auto [dv, mask_col] = build_roaring_bitmap_and_expected_row_mask( num_rows_per_dv, deletion_probability, local_indices, stream, mr); auto const host_mask = cudf::detail::make_host_vector( From f6d1081e5e0fa4ce31b5100e5d6f74cb865f8b20 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 22 Jul 2026 17:27:35 -0700 Subject: [PATCH 08/16] Add array indexing support to Parquet variant field extraction (#22895) ## Summary Contributes to https://github.com/rapidsai/cudf/issues/22312 Extends the JSONPath-like path used by `get_variant_field`/`extract_variant_field` to descend into array values via zero-based `[N]` index steps, in addition to the existing object-key descent. This enables extraction of array elements from Parquet VARIANT columns (e.g. `$.a[0]`, `$[0].field`). Leading zeroes are allowed, whitespace characters are not. Authors: - Vukasin Milovanovic (https://github.com/vuule) Approvers: - Yunsong Wang (https://github.com/PointKernel) - Simon Eves (https://github.com/simoneves) - Abigale Kim (https://github.com/abigalekim) - Nghia Truong (https://github.com/ttnghia) URL: https://github.com/rapidsai/cudf/pull/22895 --- cpp/include/cudf/io/experimental/variant.hpp | 35 +++-- .../parquet/experimental/variant_extract.cu | 109 ++++++++++++- .../io/parquet/experimental/variant_path.cpp | 62 +++++--- .../io/parquet/experimental/variant_path.hpp | 21 +-- .../io/experimental/variant_extract_test.cpp | 143 +++++++++++++++++- .../java/ai/rapids/cudf/VariantUtilsTest.java | 4 +- 6 files changed, 318 insertions(+), 56 deletions(-) diff --git a/cpp/include/cudf/io/experimental/variant.hpp b/cpp/include/cudf/io/experimental/variant.hpp index 4ee7747dc055..b868a2493941 100644 --- a/cpp/include/cudf/io/experimental/variant.hpp +++ b/cpp/include/cudf/io/experimental/variant.hpp @@ -29,34 +29,33 @@ namespace io::parquet::experimental { */ /** - * @brief Extract the raw VARIANT-encoded bytes of a nested object field by JSONPath-like path. - * - * Walks `path` step by step, descending into object values (`basic_type == 2`) at each name step. - * Returns a `list` column containing the raw encoded bytes of the value at the end of - * the path for each row. - * - * Null is produced when the struct row is null, a name step's key is absent from the dictionary, - * or the current value is not an object (`basic_type != 2`). + * @brief Extract the raw VARIANT-encoded bytes of a nested field by JSONPath-like path. * * Path grammar: - * path := "$"? first_step ("." name)* - * first := name | "." name - * name := [^.\[]+ // any byte except '.' (step separator) and '[' (reserved) + * path := "$"? first_step step* + * first := name | "." name | "[" index "]" + * step := "." name | "[" index "]" + * name := any sequence of bytes other than '.' or '[' + * index := non-negative base-10 integer (leading zeros are allowed, e.g. "[01]" == "[1]") * * Examples: - * "x" -> top-level field "x" (leading $ optional) - * "$.foo" -> top-level field "foo" - * "$.foo.bar" -> object descent foo -> bar + * "x" -> top-level field "x" (leading $ optional) + * "$.foo" -> top-level field "foo" + * "$.foo.bar" -> object descent foo -> bar + * "$[0]" -> first element of a top-level array + * "$.a[0].b" -> object key "a" -> first array element -> object key "b" * * @param variant_column Struct column (VARIANT materialization) with `list` children * (`metadata`, `value`), plus optional shredded siblings - * @param path JSONPath-like path string identifying the target object field + * @param path JSONPath-like path string identifying the target field * @param stream CUDA stream * @param mr Device memory resource - * @return `list` column with the extracted field's encoded bytes + * @return `list` column with the extracted value's encoded bytes. A row is null when the + * input row is null, a name is absent, an index is out of bounds, or a step does not match + * the current value. * - * @throws std::invalid_argument on empty path or malformed syntax (including bracket steps, - * which require array-indexing support that is not yet implemented) + * @throws std::invalid_argument on empty path or malformed syntax (`[*]` wildcards, negative + * indices, out-of-range indices, and quoted names inside `[...]` are not supported) */ [[nodiscard]] std::unique_ptr get_variant_field( column_view const& variant_column, diff --git a/cpp/src/io/parquet/experimental/variant_extract.cu b/cpp/src/io/parquet/experimental/variant_extract.cu index aabc55449405..1580af555ee2 100644 --- a/cpp/src/io/parquet/experimental/variant_extract.cu +++ b/cpp/src/io/parquet/experimental/variant_extract.cu @@ -33,6 +33,7 @@ #include #include +#include #include #include #include @@ -266,7 +267,7 @@ __device__ cuda::std::optional find_key_in_metadata(device_span(num_entries.value()) + 1) * offset_size; - if (offsets_bytes > static_cast(meta_len - offsets_start)) { + if (cuda::std::cmp_greater(offsets_bytes, meta_len - offsets_start)) { return cuda::std::nullopt; } @@ -369,6 +370,64 @@ __device__ device_span locate_object_field(device_span locate_array_element(device_span value, + size_type index) +{ + if (index < 0) { return {}; } + + auto const value_size = static_cast(value.size()); + if (value_size < 1) { return {}; } + uint8_t const value_metadata = value[0]; + if (variant_basic_type(value_metadata) != basic_type::array) { return {}; } + + int const value_header = variant_value_header(value_metadata); + [[maybe_unused]] auto const [offset_size, _, num_elements_size] = + decode_object_array_header(value_header, false); + + size_type position = 1; + auto const num_elements_value = narrow_cast(read_uint64(value, position, num_elements_size)); + if (!num_elements_value.has_value()) { return {}; } + auto const num_elements = num_elements_value.value(); + if (index >= num_elements) { return {}; } + position += num_elements_size; + + size_type const offsets_start = position; + // Computed in 64-bit because (num_elements + 1) * offset_size can exceed the signed `size_type` + // range (which would be UB); the check below then rejects any array that overruns the value blob. + auto const offsets_bytes = (static_cast(num_elements) + 1) * offset_size; + if (cuda::std::cmp_greater(offsets_bytes, value_size - offsets_start)) { return {}; } + size_type const values_base = offsets_start + static_cast(offsets_bytes); + auto const values_extent = value_size - values_base; + + auto const start_offset_pos = offsets_start + static_cast(index) * offset_size; + auto const end_offset_pos = offsets_start + (static_cast(index) + 1) * offset_size; + if (cuda::std::cmp_greater(end_offset_pos + offset_size, value_size)) { return {}; } + + auto const start_offset = read_uint64(value, start_offset_pos, offset_size); + auto const end_offset = read_uint64(value, end_offset_pos, offset_size); + if (!start_offset.has_value() || !end_offset.has_value()) { return {}; } + auto const element_start = *start_offset; + auto const element_end = *end_offset; + if (element_end < element_start || cuda::std::cmp_greater(element_end, values_extent)) { + return {}; + } + return value.subspan(values_base + element_start, element_end - element_start); +} + // The fixed-width signed integers a VARIANT value can be cast to: INT{8,16,32,64}. Matches the // exact width types (not e.g. __int128) since those are the only variant primitive int headers. template @@ -401,18 +460,56 @@ __device__ inline cuda::std::optional decode_int(device_span e return cudf::io::unaligned_load(enc.data() + 1); } +// Parse an array-index step token of the form "[]" into its zero-based index. Returns nullopt +// for any malformed token or an index that does not fit in `size_type` (such an index is out of +// range for any array, so the caller treats it as a missing element). +__device__ cuda::std::optional parse_index_step(cudf::string_view step) +{ + auto const step_size = step.size_bytes(); + auto const* step_data = step.data(); + if (step_size < 3 || step_data[0] != '[' || step_data[step_size - 1] != ']') { + return cuda::std::nullopt; + } + + // Accumulate directly in `size_type`; the checked-arithmetic helpers reject the token if the + // running value overflows, which means the index is out of range for any array and the caller + // treats it as a missing element. + size_type index = 0; + for (size_type k = 1; k < step_size - 1; ++k) { + char const c = step_data[k]; + if (c < '0' || c > '9') { return cuda::std::nullopt; } + if (cuda::mul_overflow(index, index, size_type{10}) || + cuda::add_overflow(index, index, static_cast(c - '0'))) { + return cuda::std::nullopt; + } + } + return index; +} + +// Walk a path of object-key or array-index steps level by level starting at `val` and return +// the span of the final value (subspan of `val`). Returns an empty span on failure. +// +// Each path step is encoded in the `path` strings column as either: +// - "" -> descend into an object by dictionary key, or +// - "[]" -> descend into an array by zero-based integer index. +// The step kind is inferred from the first byte (`'['` means index). __device__ device_span resolve_path(device_span meta, device_span val, column_device_view path) { device_span sub_val = val; for (size_type i = 0; i < path.size(); ++i) { - auto const name = path.element(i); - - auto const field_id = find_key_in_metadata(meta, name); - if (!field_id.has_value()) { return {}; } + auto const step = path.element(i); - sub_val = locate_object_field(sub_val, field_id.value()); + if (step.size_bytes() >= 1 && step.data()[0] == '[') { + auto const index = parse_index_step(step); + if (!index.has_value()) { return {}; } + sub_val = locate_array_element(sub_val, index.value()); + } else { + auto const field_id = find_key_in_metadata(meta, step); + if (!field_id.has_value()) { return {}; } + sub_val = locate_object_field(sub_val, field_id.value()); + } if (sub_val.empty()) { return {}; } } return sub_val; diff --git a/cpp/src/io/parquet/experimental/variant_path.cpp b/cpp/src/io/parquet/experimental/variant_path.cpp index 891ba19bd9e8..58dc519e7965 100644 --- a/cpp/src/io/parquet/experimental/variant_path.cpp +++ b/cpp/src/io/parquet/experimental/variant_path.cpp @@ -1,17 +1,19 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "variant_path.hpp" +#include #include +#include #include -#include #include #include #include +#include #include namespace cudf::io::parquet::experimental::detail { @@ -21,12 +23,6 @@ namespace { // Dot-notation field names accept any byte except the structural characters '.' and '['. [[nodiscard]] constexpr bool is_name_char(char c) { return c != '.' && c != '['; } -[[noreturn]] void throw_parse_error(std::string_view path, std::size_t pos, std::string_view msg) -{ - CUDF_FAIL(std::format("invalid variant path \"{}\" at position {}: {}", path, pos, msg), - std::invalid_argument); -} - // Reads a maximal run of name characters from the front of `tail`. [[nodiscard]] std::string read_unquoted_name(std::string_view tail) { @@ -37,6 +33,34 @@ namespace { return std::string{tail.substr(0, n)}; } +// Reads a bracket step "[]" from the front of `tail`. +// The returned token keeps its brackets (e.g. "[42]"). +[[nodiscard]] std::string read_bracket_step(std::string_view tail) +{ + CUDF_EXPECTS(!tail.empty() && tail.front() == '[', + "expected '[' to open variant path index", + std::invalid_argument); + + // Consume the maximal run of decimal digits. + std::size_t n = 1; + while (n < tail.size() && tail[n] >= '0' && tail[n] <= '9') { + ++n; + } + CUDF_EXPECTS( + n != 1, "expected non-negative integer after '[' in variant path", std::invalid_argument); + + // Reject indices that cannot be a valid array position (don't fit in cudf::size_type) + cudf::size_type index = 0; + auto const result = std::from_chars(tail.data() + 1, tail.data() + n, index); + CUDF_EXPECTS( + result.ec == std::errc{}, "variant path index is out of range", std::invalid_argument); + + CUDF_EXPECTS(n < tail.size() && tail[n] == ']', + "expected ']' to close variant path index", + std::invalid_argument); + return std::string{tail.substr(0, n + 1)}; // include the closing ']' +} + } // namespace std::vector parse_variant_path(std::string_view path) @@ -51,17 +75,21 @@ std::vector parse_variant_path(std::string_view path) bool first = true; while (pos < len) { char const c = path[pos]; - if (c == '.') { - ++pos; - if (pos >= len || !is_name_char(path[pos])) { - throw_parse_error(path, pos - 1, "trailing '.' with no field name"); + if (c == '[') { + steps.emplace_back(read_bracket_step(path.substr(pos))); + } else { + if (c == '.') { + ++pos; + CUDF_EXPECTS(pos < len && is_name_char(path[pos]), + "trailing '.' with no field name", + std::invalid_argument); + } else { + // Neither a '.'/'[' step nor a valid leading name (e.g. a stray ']' or a name after a step) + CUDF_EXPECTS( + first && is_name_char(c), "unexpected character in variant path", std::invalid_argument); } - } else if (!(first && is_name_char(c))) { - // Neither a '.' step nor a valid leading name (e.g. a bracket step like "[0]" or "foo[1]") - throw_parse_error(path, pos, "unexpected character in variant path"); + steps.emplace_back(read_unquoted_name(path.substr(pos))); } - - steps.emplace_back(read_unquoted_name(path.substr(pos))); pos += steps.back().size(); first = false; } diff --git a/cpp/src/io/parquet/experimental/variant_path.hpp b/cpp/src/io/parquet/experimental/variant_path.hpp index 3cad1d1fc4af..760b36a7962c 100644 --- a/cpp/src/io/parquet/experimental/variant_path.hpp +++ b/cpp/src/io/parquet/experimental/variant_path.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -12,18 +12,21 @@ namespace cudf::io::parquet::experimental::detail { /** - * @brief Parse a JSONPath-like VARIANT path string into an ordered sequence of object-key steps. + * @brief Parse a JSONPath-like VARIANT path string into an ordered sequence of steps. * - * Grammar — object descent only: - * path := "$"? first_step ("." name)* - * first := name | "." name + * Grammar — object descent and array indexing: + * path := "$"? first_step (("." name) | index)* + * first := name | "." name | index * name := [^.\[]+ + * index := "[" [0-9]+ "]" * - * Names accept any byte except '.' (step separator) and '[' (start of a bracket step, - * reserved for future array indexing and quoted-name syntax). + * A step is either an object-key name or an array index. Names accept any byte except '.' (step + * separator) and '[' (start of an index step). Index steps hold a non-negative integer and are + * returned with their brackets kept (e.g. "[42]"), which is how downstream consumers tell an index + * step apart from an object key. * - * @throws std::invalid_argument on empty path or malformed syntax (including bracket steps, - * which require array-indexing support that is not yet implemented) + * @throws std::invalid_argument on an empty path or malformed syntax (e.g. a non-integer, negative, + * or out-of-range array index, an unterminated '[', or a trailing '.') */ [[nodiscard]] std::vector parse_variant_path(std::string_view path); diff --git a/cpp/tests/io/experimental/variant_extract_test.cpp b/cpp/tests/io/experimental/variant_extract_test.cpp index 051c3d94e240..0e8c099f554a 100644 --- a/cpp/tests/io/experimental/variant_extract_test.cpp +++ b/cpp/tests/io/experimental/variant_extract_test.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -490,9 +490,20 @@ TEST_F(ExtractVariantFieldTest, SyntaxErrors) { auto col = wrap_single_variant(build_metadata({}), enc_int32(1)); auto stream = cudf::test::get_default_stream(); - // Only object-key descent is supported — array indexing, bracket steps, and quoted keys should - // throw, alongside malformed paths. - for (auto const* bad : {"$..a", "$.a[0]", "$.a[", "$.a[]", "$.", "$['x']", "$.a[*]"}) { + // Object-key descent and array-index steps are supported; wildcards, quoted keys, negative + // indices, out-of-range indices, and other malformed bracket forms must throw. + for (auto const* bad : {"$..a", + "$.a[", + "$.a[]", + "$.", + "$['x']", + "$.a[*]", + "$.a[-1]", + "$.a[+1]", + "$.a[ 1]", + "$.a[01x]", + "$.a[1", + "$.a[99999999999999999999]"}) { EXPECT_THROW( static_cast(cudf::io::parquet::experimental::get_variant_field(col, bad, stream)), std::invalid_argument) @@ -500,6 +511,130 @@ TEST_F(ExtractVariantFieldTest, SyntaxErrors) } } +TEST_F(ExtractVariantFieldTest, ApacheArrayPrimitiveIndexing) +{ + // array_primitive encodes the int8 array [2, 1, 5, 9]; index into it via "[N]" steps. + auto col = make_apache_variant(avf::array_primitive); + auto stream = cudf::test::get_default_stream(); + auto const i8 = cudf::data_type{cudf::type_id::INT8}; + auto const get = [&](char const* path) { + return cudf::io::parquet::experimental::extract_variant_field(col, path, i8, stream); + }; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[0]"), + cudf::test::fixed_width_column_wrapper{int8_t{2}}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[2]"), + cudf::test::fixed_width_column_wrapper{int8_t{5}}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[3]"), + cudf::test::fixed_width_column_wrapper{int8_t{9}}); + + // Leading zeros are allowed + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[00]"), + cudf::test::fixed_width_column_wrapper{int8_t{2}}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[01]"), + cudf::test::fixed_width_column_wrapper{int8_t{1}}); + + // Out-of-bounds index resolves to null. + cudf::test::fixed_width_column_wrapper const null_expected({0}, {false}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[4]"), null_expected); + + // Exercise 2-, 3-, and 4-byte offsets with a four-byte element count. + for (uint8_t offset_size = 2; offset_size <= 4; ++offset_size) { + auto const value_header = static_cast(0x04 | (offset_size - 1)); + std::vector value{static_cast(0x03 | (value_header << 2)), 1, 0, 0, 0}; + value.insert(value.end(), offset_size, 0); // offsets[0] + value.push_back(2); // offsets[1] + value.insert(value.end(), offset_size - 1, 0); + value.insert(value.end(), {0x0c, 42}); // INT8(42) + + auto wide_col = wrap_single_variant(build_metadata({}), value); + auto got = cudf::io::parquet::experimental::extract_variant_field(wide_col, "$[0]", i8, stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, + cudf::test::fixed_width_column_wrapper{int8_t{42}}); + } +} + +TEST_F(ExtractVariantFieldTest, ArrayIndexingTypeMismatchAndBounds) +{ + // array_primitive is the int8 array [2, 1, 5, 9]. An object-key step against an array, an + // out-of-bounds index, and an index step against a non-array element all resolve to null. + auto col = make_apache_variant(avf::array_primitive); + auto stream = cudf::test::get_default_stream(); + auto const i8 = cudf::data_type{cudf::type_id::INT8}; + cudf::test::fixed_width_column_wrapper const null_expected({0}, {false}); + + // Object-key descent into an array value: no such key -> null. + auto key_on_array = + cudf::io::parquet::experimental::extract_variant_field(col, "$.foo", i8, stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*key_on_array, null_expected); + + // Index step against a primitive element (after first descending into it): non-array -> null. + auto index_on_primitive = + cudf::io::parquet::experimental::extract_variant_field(col, "$[0][0]", i8, stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*index_on_primitive, null_expected); +} + +TEST_F(ExtractVariantFieldTest, EmptyArrayIndexing) +{ + auto col = make_apache_variant(avf::array_empty); + auto stream = cudf::test::get_default_stream(); + auto const i8 = cudf::data_type{cudf::type_id::INT8}; + cudf::test::fixed_width_column_wrapper const null_expected({0}, {false}); + + for (auto const* path : {"$[0]", "$[1]"}) { + SCOPED_TRACE(std::string{"path: "} + path); + auto got = cudf::io::parquet::experimental::extract_variant_field(col, path, i8, stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, null_expected); + } + + // Truncated counts/tables, decreasing offsets, and offsets beyond the values region yield null. + for (auto const& value : std::vector>{{0x13}, + {0x03, 0x01, 0x00}, + {0x03, 0x01, 0x02, 0x01, 0x0c, 42}, + {0x03, 0x01, 0x00, 0x03, 0x0c, 42}}) { + auto malformed_col = wrap_single_variant(build_metadata({}), value); + auto got = + cudf::io::parquet::experimental::extract_variant_field(malformed_col, "$[0]", i8, stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, null_expected); + } +} + +TEST_F(ExtractVariantFieldTest, MixedObjectArrayTraversal) +{ + // array_nested encodes: + // [ {id:1, thing:{names:["Contrarian","Spider"]}}, + // null, + // {id:2, names:["Apple","Ray",null], type:"if"} ] + auto col = make_apache_variant(avf::array_nested); + auto stream = cudf::test::get_default_stream(); + + auto const check_str = [&](char const* path, char const* expected) { + SCOPED_TRACE(std::string{"path: "} + path); + auto got = cudf::io::parquet::experimental::extract_variant_field( + col, path, cudf::data_type{cudf::type_id::STRING}, stream); + cudf::test::strings_column_wrapper const expected_col({expected}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected_col); + }; + auto const check_null = [&](char const* path) { + SCOPED_TRACE(std::string{"path: "} + path); + auto got = cudf::io::parquet::experimental::extract_variant_field( + col, path, cudf::data_type{cudf::type_id::STRING}, stream); + cudf::test::strings_column_wrapper const null_col({""}, {false}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, null_col); + }; + + check_str("$[2].type", "if"); + check_str("$[0].thing.names[0]", "Contrarian"); + check_str("$[0].thing.names[1]", "Spider"); + check_str("$[2].names[0]", "Apple"); + check_str("$[2].names[1]", "Ray"); + + check_null("$[1].id"); // element 1 is a JSON null + check_null("$[0].name"); // element 0 has no "name" key (it has "thing") + check_null("$[2].names[2]"); // third name is null + check_null("$[2].names[3]"); // out-of-bounds array index +} + TEST_F(ExtractVariantFieldTest, LargeDictionaryAndObjectScan) { auto const keys = make_numeric_keys(50); diff --git a/java/src/test/java/ai/rapids/cudf/VariantUtilsTest.java b/java/src/test/java/ai/rapids/cudf/VariantUtilsTest.java index 67bded7528b4..d9c9633805ed 100644 --- a/java/src/test/java/ai/rapids/cudf/VariantUtilsTest.java +++ b/java/src/test/java/ai/rapids/cudf/VariantUtilsTest.java @@ -411,9 +411,9 @@ void emptyPathThrows() { void malformedPathThrows() { try (ColumnVector variant = makeXyzVariantColumn()) { assertThrows(CudfException.class, - () -> VariantUtils.getVariantFieldValue(variant, "$.x[0]")); + () -> VariantUtils.getVariantFieldValue(variant, "$.x[")); assertThrows(CudfException.class, - () -> VariantUtils.extractVariantField(variant, "$.x[0]", DType.INT32)); + () -> VariantUtils.extractVariantField(variant, "$.x[", DType.INT32)); } } From f03e9cd66d0f8e4ddaa63febccd39d199db7b44b Mon Sep 17 00:00:00 2001 From: Allen Xu Date: Thu, 23 Jul 2026 09:57:27 +0800 Subject: [PATCH 09/16] [BUG] Fix ORC chunked writer root row statistics (#23118) Fix ORC root column statistics when an `orc_chunked_writer` receives multiple non-empty tables. `persisted_statistics::persist` replaced `num_rows` on every write, so the root `numberOfValues` stored in the footer reflected only the final table. Accumulate the row count across writes instead, while retaining the existing early return for an empty current table. This caused Spark ORC aggregate pushdown to return the number of output files instead of the row count for `COUNT` on a top-level struct written by the RAPIDS Accelerator. Related issue: https://github.com/NVIDIA/cudf-spark/issues/15186. Validation: - Before the fix, `OrcChunkedWriterTest.RootStatisticsAccumulateRows` failed with root `number_of_values = 1`, expected `6`. - After the fix, the targeted test passed. - Full `ORC_TEST`: `197 tests from 36 test suites`, `197 passed`, `4 disabled`. Authors: - Allen Xu (https://github.com/wjxiz1992) Approvers: - Chong Gao (https://github.com/res-life) - Vukasin Milovanovic (https://github.com/vuule) - Muhammad Haseeb (https://github.com/mhaseeb123) URL: https://github.com/rapidsai/cudf/pull/23118 --- cpp/src/io/orc/writer_impl.cu | 8 ++++---- cpp/src/io/orc/writer_impl.hpp | 7 ++++--- cpp/tests/io/orc_test.cpp | 24 ++++++++++++++++++++++++ 3 files changed, 32 insertions(+), 7 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 25600d14f44a..62ff8f165c3e 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -388,15 +388,15 @@ intermediate_statistics::intermediate_statistics(orc_table_view const& table, }); } -void persisted_statistics::persist(int num_table_rows, +void persisted_statistics::persist(uint64_t num_table_rows, single_write_mode write_mode, intermediate_statistics&& intermediate_stats, rmm::cuda_stream_view stream) { + col_types = std::move(intermediate_stats.col_types); + num_rows += num_table_rows; + if (num_table_rows == 0) { return; } stats_dtypes = std::move(intermediate_stats.stats_dtypes); - col_types = std::move(intermediate_stats.col_types); - num_rows = num_table_rows; - if (num_rows == 0) { return; } if (write_mode == single_write_mode::NO) { // persist the strings in the chunks into a string pool and update pointers diff --git a/cpp/src/io/orc/writer_impl.hpp b/cpp/src/io/orc/writer_impl.hpp index 0454cbea61f0..6e320c02ec25 100644 --- a/cpp/src/io/orc/writer_impl.hpp +++ b/cpp/src/io/orc/writer_impl.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -169,7 +170,7 @@ struct persisted_statistics { num_rows = 0; } - void persist(int num_table_rows, + void persist(uint64_t num_table_rows, single_write_mode write_mode, intermediate_statistics&& intermediate_stats, rmm::cuda_stream_view stream); @@ -179,7 +180,7 @@ struct persisted_statistics { std::vector> string_pools; std::vector stats_dtypes; std::vector col_types; - int num_rows = 0; + uint64_t num_rows = 0; }; /** diff --git a/cpp/tests/io/orc_test.cpp b/cpp/tests/io/orc_test.cpp index 6159facde017..f000bde0ff5d 100644 --- a/cpp/tests/io/orc_test.cpp +++ b/cpp/tests/io/orc_test.cpp @@ -707,6 +707,23 @@ TEST_F(OrcChunkedWriterTest, SimpleTable) CUDF_TEST_EXPECT_TABLES_EQUAL(*result.tbl, *full_table); } +TEST_F(OrcChunkedWriterTest, RootStatisticsAccumulateRows) +{ + auto table1 = create_random_fixed_table(1, 5, false); + auto table2 = create_random_fixed_table(1, 1, false); + auto table3 = create_random_fixed_table(1, 0, false); + + auto filepath = temp_env->get_temp_filepath("ChunkedRootStatistics.orc"); + cudf::io::chunked_orc_writer_options opts = + cudf::io::chunked_orc_writer_options::builder(cudf::io::sink_info{filepath}); + cudf::io::orc_chunked_writer(opts).write(*table1).write(*table2).write(*table3); + + auto const stats = cudf::io::read_parsed_orc_statistics(cudf::io::source_info{filepath}); + ASSERT_FALSE(stats.file_stats.empty()); + ASSERT_TRUE(stats.file_stats.front().number_of_values.has_value()); + EXPECT_EQ(*stats.file_stats.front().number_of_values, 6); +} + TEST_F(OrcChunkedWriterTest, LargeTables) { srand(31337); @@ -2176,6 +2193,13 @@ TEST_F(OrcReaderTest, SizeTypeRowsOverflow) EXPECT_EQ(metadata.num_rows(), total_rows); EXPECT_EQ(metadata.num_stripes(), total_rows / 1'000'000); + auto const stats = + cudf::io::read_parsed_orc_statistics(cudf::io::source_info{cudf::host_span{ + reinterpret_cast(out_buffer.data()), out_buffer.size()}}); + ASSERT_FALSE(stats.file_stats.empty()); + ASSERT_TRUE(stats.file_stats.front().number_of_values.has_value()); + EXPECT_EQ(*stats.file_stats.front().number_of_values, static_cast(total_rows)); + constexpr auto num_rows_to_read = 1'000'000; auto const num_rows_to_skip = metadata.num_rows() - num_rows_to_read; From 29625641474ed385dd2ed75c554834e503a66149 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 23 Jul 2026 09:47:04 -0500 Subject: [PATCH 10/16] Bump cudf-polars CI timeouts (#23410) These two nightly jobs were scheduled on the same (slow) runner and made it to ~95% before being timed out: - https://github.com/rapidsai/cudf/actions/runs/29991014913/job/89153660965 - https://github.com/rapidsai/cudf/actions/runs/29991014913/job/89153660991 I've bumped the timeout slightly more. Authors: - Tom Augspurger (https://github.com/TomAugspurger) Approvers: - Lawrence Mitchell (https://github.com/wence-) - James Lamb (https://github.com/jameslamb) URL: https://github.com/rapidsai/cudf/pull/23410 --- ci/run_cudf_polars_polars_tests.sh | 4 ++-- ci/run_cudf_polars_pytests.sh | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ci/run_cudf_polars_polars_tests.sh b/ci/run_cudf_polars_polars_tests.sh index 3073b61812d7..7c3c138c3a60 100755 --- a/ci/run_cudf_polars_polars_tests.sh +++ b/ci/run_cudf_polars_polars_tests.sh @@ -64,7 +64,7 @@ DESELECTED_TESTS_STR=$(printf -- " --deselect %s" "${DESELECTED_TESTS[@]}") # shellcheck disable=SC2086 # Fail fast (-x) because failed tests pollute the state echo "Run polars tests with injected in-memory GPU engine" -python "${TIMEOUT_TOOL_PATH}" --enable-python 4800 \ +python "${TIMEOUT_TOOL_PATH}" --enable-python 5400 \ python -m pytest \ --import-mode=importlib \ --cache-clear \ @@ -85,7 +85,7 @@ python "${TIMEOUT_TOOL_PATH}" --enable-python 4800 \ echo "Run polars tests with injected SPMD GPU engine, small blocksize" CUDF_POLARS__EXECUTOR__TARGET_PARTITION_SIZE=805306368 \ CUDF_POLARS__EXECUTOR__FALLBACK_MODE=silent \ -python "${TIMEOUT_TOOL_PATH}" --enable-python 4800 \ +python "${TIMEOUT_TOOL_PATH}" --enable-python 5400 \ python -m pytest \ --import-mode=importlib \ --cache-clear \ diff --git a/ci/run_cudf_polars_pytests.sh b/ci/run_cudf_polars_pytests.sh index c7cdb580a4d7..afd5561b551b 100755 --- a/ci/run_cudf_polars_pytests.sh +++ b/ci/run_cudf_polars_pytests.sh @@ -10,5 +10,5 @@ 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 4800 \ +python "${TIMEOUT_TOOL_PATH}" --enable-python 5400 \ python -m pytest --cache-clear "$@" tests From 55317dce07961a249341b674296f91bad37bc947 Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Thu, 23 Jul 2026 11:33:04 -0500 Subject: [PATCH 11/16] Fix DataFrame.stack level resolution, ordering, and dtype preservation (#23370) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split out of #23255 (3/6), superseding it. **Depends on #23365 (MultiIndex column fidelity)** — 14 of the 28 un-xfailed pandas-tests need both fixes, so this PR's pandas-tests job goes green once #23365 merges. - Resolve `level` positionally: integer column-level *names* no longer collide with level *positions* (pandas' `Index.get_level_values` resolves integers by name first, so frames with integer level names returned data from the wrong level). - Validate out-of-bounds integer levels (`IndexError`) and duplicated level names (`ValueError`) with pandas' messages; negative out-of-bounds levels previously wrapped around silently. - Build the stacked level keys from the column MultiIndex's own levels/codes so per-level dtypes survive: int64 levels with missing entries no longer upcast to float64, and categorical levels stay categorical through the pylibcudf `tile` step (which only sees codes). - Emit stacked keys in appearance order, matching pandas. This replaces the argsort-based reordering, which misaligned column data for non-involution column permutations (e.g. a 3-cycle) and NaN keys; pandas legacy stack sorts multi-level keys by level *codes*, not values. - Attach pandas-faithful levels/codes to the result index eagerly (the original index contributes its own levels/codes; flat indexes and the tiled level get appearance-order factorization) so a later `unstack` restores the original row/column order; the legacy `dropna` path preserves them by masking codes instead of gathering the index. Fixes 28 pandas-tests; their xfail entries are removed. Three classic categorical unstack params are un-xfailed (fixed by this change together with #23365). Attribution verified per node id against isolated builds: 14 pass with only this change, 14 need this plus #23365. Authors: - GALI PREM SAGAR (https://github.com/galipremsagar) Approvers: - Matthew Murray (https://github.com/Matt711) URL: https://github.com/rapidsai/cudf/pull/23370 --- python/cudf/cudf/core/dataframe.py | 263 +++++++++++++----- python/cudf/cudf/core/indexed_frame.py | 13 +- python/cudf/cudf/core/multiindex.py | 12 +- .../pandas/scripts/pandas-testing-plugin.py | 28 -- python/cudf/cudf/tests/reshape/test_stack.py | 47 +++- .../cudf/cudf/tests/reshape/test_unstack.py | 9 +- 6 files changed, 265 insertions(+), 107 deletions(-) diff --git a/python/cudf/cudf/core/dataframe.py b/python/cudf/cudf/core/dataframe.py index a9439c730348..a5211ff186c5 100644 --- a/python/cudf/cudf/core/dataframe.py +++ b/python/cudf/cudf/core/dataframe.py @@ -47,10 +47,12 @@ is_decimal128_dtype, is_dict_like, is_dtype_equal, + is_integer, is_list_like, is_scalar, ) from cudf.core import indexing_utils, reshape +from cudf.core.algorithms import factorize from cudf.core.column import ( CategoricalColumn, ColumnBase, @@ -89,6 +91,7 @@ ) from cudf.core.indexed_frame import ( IndexedFrame, + _check_duplicate_level_names, _FrameIndexer, _indices_from_labels, doc_reset_index_template, @@ -8119,6 +8122,11 @@ def stack( "https://github.com/pandas-dev/pandas/issues/53515" ) + _check_duplicate_level_names( + [lv for lv in level if not is_integer(lv)], + self._data.level_names, + ) + # Compute the columns to stack based on specified levels level_indices: list[int] = [] @@ -8133,10 +8141,24 @@ def stack( "mixture of both." ) else: - # Must be a list of positions, normalize negative positions - level_indices = [ - lv + self._data.nlevels if lv < 0 else lv for lv in level - ] + # Must be a list of positions; normalize negative positions + # and validate bounds to match pandas MultiIndex._get_level_number + nlevels = self._data.nlevels + for lv in level: + if lv < 0: + if lv + nlevels < 0: + raise IndexError( + f"Too many levels: Index has only {nlevels} " + f"levels, {lv} is not a valid level number" + ) + level_indices.append(lv + nlevels) + else: + if lv >= nlevels: + raise IndexError( + f"Too many levels: Index has only {nlevels} " + f"levels, not {lv + 1}" + ) + level_indices.append(lv) unnamed_levels_indices = [ i for i in range(self._data.nlevels) if i not in level_indices @@ -8144,48 +8166,130 @@ def stack( has_unnamed_levels = len(unnamed_levels_indices) > 0 column_name_idx = self._data.to_pandas_index - # Construct new index from the levels specified by `level` - named_levels = pd.MultiIndex.from_arrays( - [column_name_idx.get_level_values(lv) for lv in level_indices] + # pandas' Index.get_level_values resolves an integer argument by + # name first: if a level is *named* that integer, that level is + # returned regardless of position. All lookups below use positional + # indices, so strip the names to force positional resolution and + # re-attach the real names afterwards. + nameless_column_name_idx = column_name_idx.set_names( + [None] * column_name_idx.nlevels ) + # Construct new index from the levels specified by `level` + if isinstance(column_name_idx, pd.MultiIndex): + # build from codes/levels to keep the level dtypes: materializing + # via get_level_values/from_arrays turns missing entries into NaN + # and upcasts e.g. int64 levels to float64 + named_levels = pd.MultiIndex( + levels=[column_name_idx.levels[i] for i in level_indices], + codes=[column_name_idx.codes[i] for i in level_indices], + names=[column_name_idx.names[i] for i in level_indices], + verify_integrity=False, + ) + else: + named_levels = pd.MultiIndex.from_arrays( + [ + nameless_column_name_idx.get_level_values(lv).rename( + column_name_idx.names[lv] + ) + for lv in level_indices + ] + ) # Since `level` may only specify a subset of all levels, `unique()` is - # required to remove duplicates. In pandas, the order of the keys in - # the specified levels are always sorted. + # required to remove duplicates. In pandas legacy stack, the keys of + # the specified levels are sorted by their level *codes* when the + # columns have multiple levels (flat column labels keep their + # original order): level order is preserved even for unsorted levels + # and missing labels (code -1) come first. unique_named_levels = named_levels.unique() - if not future_stack: - unique_named_levels = unique_named_levels.sort_values() + if not future_stack and self._data.nlevels > 1: + unique_named_levels = unique_named_levels.take( + np.lexsort(tuple(reversed(unique_named_levels.codes))) + ) # Each index from the original dataframe should repeat by the number # of unique values in the named_levels repeated_index = self.index.repeat(len(unique_named_levels)) # Each column name should tile itself by len(df) times - cols = [ - as_column(unique_named_levels.get_level_values(i)) - for i in range(unique_named_levels.nlevels) - ] + nameless_unique_named_levels = unique_named_levels.set_names( + [None] * unique_named_levels.nlevels + ) + cols = [] + for i in range(unique_named_levels.nlevels): + if future_stack: + # pandas future stack materializes the level values (a + # level with missing entries becomes e.g. float64 with NaN) + cols.append( + as_column(nameless_unique_named_levels.get_level_values(i)) + ) + else: + # pandas legacy stack keeps the original level dtype and + # represents missing entries as nulls (-1 codes) + level_col = as_column(unique_named_levels.levels[i]) + level_codes = np.asarray(unique_named_levels.codes[i]).astype( + "int64" + ) + level_codes[level_codes == -1] = np.iinfo(SIZE_TYPE_DTYPE).min + cols.append( + level_col.take(as_column(level_codes), nullify=True) + ) with access_columns(*cols, mode="read", scope="internal"): plc_table = plc.reshape.tile( plc.Table([col.plc_column for col in cols]), self.shape[0], ) tiled_index = [ - ColumnBase.create(plc, dtype=dtype_from_pylibcudf_column(plc)) - for plc in plc_table.columns() + ColumnBase.create(plc_col, dtype=src_col.dtype) + for src_col, plc_col in zip( + cols, plc_table.columns(), strict=True + ) ] - # Assemble the final index - new_index_columns = [*repeated_index._columns, *tiled_index] + # Assemble the final index — build levels/codes first so the + # MultiIndex can be constructed in one step via _simple_new. + # Codes/levels are attached eagerly, matching how pandas' stack builds + # the result MultiIndex, so a later unstack can restore the original + # row/column order (lazy materialization would sort the levels): + # the original index contributes its own levels/codes (repeated); + # a flat original index and the tiled stacked level(s) get + # appearance-order factorization. index_names = [*self.index.names, *unique_named_levels.names] - new_index = MultiIndex._from_data(dict(enumerate(new_index_columns))) - # Materialize the levels in order of first appearance (rather than the - # default sorted order) so that converting the result to pandas keeps - # the level order pandas' own ``stack`` produces. Otherwise a later - # ``unstack``/``to_pandas`` would lexicographically reorder the pivoted - # axis (e.g. ``"foo_10"`` before ``"foo_2"``). - new_index._maybe_materialize_codes_and_levels(sort=False) - new_index.names = index_names + new_levels: list[cudf.Index] = [] + new_codes: list[ColumnBase] = [] + n_tile = len(unique_named_levels) + if isinstance(self.index, MultiIndex): + src = self.index._maybe_materialize_codes_and_levels() + for src_level, src_code in zip( + src._levels, + src._codes, + strict=True, + ): + new_levels.append(src_level) + new_codes.append( + Index._from_column(src_code.astype(np.dtype(np.int64))) + .repeat(n_tile) + ._column + ) + else: + code, cats = factorize(self.index) + new_levels.append(cats) + new_codes.append( + Index._from_column(as_column(code).astype(np.dtype(np.int64))) + .repeat(n_tile) + ._column + ) + for tiled_col in tiled_index: + code, cats = factorize(Index._from_column(tiled_col)) + new_codes.append(as_column(code).astype(np.dtype(np.int64))) + new_levels.append(cats) + new_index_columns = [*repeated_index._columns, *tiled_index] + new_index = MultiIndex._simple_new( + ColumnAccessor(dict(enumerate(new_index_columns))), + new_levels, + new_codes, + pd.core.indexes.frozen.FrozenList(index_names), + ) # Compute the column indices that serves as the input for # `interleave_columns` @@ -8194,41 +8298,49 @@ def stack( ) if has_unnamed_levels: - unnamed_level_values = pd.MultiIndex.from_arrays( - list( - map( - column_name_idx.get_level_values, - unnamed_levels_indices, - ) - ) + # the columns axis has multiple levels here, so column_name_idx + # is always a pd.MultiIndex; build from codes/levels to keep the + # level dtypes and to resolve the levels positionally + unnamed_level_values = pd.MultiIndex( + levels=[ + column_name_idx.levels[i] for i in unnamed_levels_indices + ], + codes=[ + column_name_idx.codes[i] for i in unnamed_levels_indices + ], + names=[ + column_name_idx.names[i] for i in unnamed_levels_indices + ], + verify_integrity=False, ) def unnamed_group_generator(): if has_unnamed_levels: - for _, grpdf in column_idx_df.groupby(by=unnamed_level_values): + # sort=False iterates groups in first-appearance order, i.e. + # exactly ``unnamed_level_values.unique()`` order (also for + # NaN-containing tuple keys, which sorted groupby would + # reorder via codes), so the stacked columns can be zipped + # 1:1 with those keys when assembling the result. + for _, grpdf in column_idx_df.groupby( + by=unnamed_level_values, sort=False, dropna=False + ): # When stacking part of the levels, some combinations # of keys may not be present in this group but can be # present in others. Reindexing with the globally computed # `unique_named_levels` assigns -1 to these key # combinations, representing an all-null column that # is used in the subsequent libcudf call. - if future_stack: - yield grpdf.reindex( - unique_named_levels, axis=0, fill_value=-1 - ).values - else: - yield ( - grpdf.reindex( - unique_named_levels, axis=0, fill_value=-1 - ) - .sort_index() - .values - ) + # ``reindex`` returns rows in target order, so the + # legacy path needs no further sorting (the target was + # already sorted above). + yield grpdf.reindex( + unique_named_levels, axis=0, fill_value=-1 + ).values else: - if future_stack: + if future_stack or self._data.nlevels == 1: yield column_idx_df.values else: - yield column_idx_df.sort_index().values + yield column_idx_df.reindex(unique_named_levels).values # For each of the group constructed from the unnamed levels, # invoke `interleave_columns` to stack the values. @@ -8284,23 +8396,35 @@ def unnamed_group_generator(): unnamed_level_values = unnamed_level_values.get_level_values(0) unnamed_level_values = unnamed_level_values.unique() - data = ColumnAccessor( - dict( - zip( - unnamed_level_values, - [ - stacked[i] - for i in unnamed_level_values.argsort().argsort() - ] - if not future_stack - else [ - stacked[i] for i in unnamed_level_values.argsort() - ], - strict=True, + if isinstance(unnamed_level_values, pd.MultiIndex): + # build the labels from levels/codes to preserve scalar + # types: iterating a MultiIndex materializes e.g. an int64 + # level containing a missing entry as float + keys: list[tuple[Any, ...]] = [ + tuple( + unnamed_level_values.levels[j][c] + if c != -1 + else np.nan + for j, c in enumerate(row) ) - ), + for row in zip(*unnamed_level_values.codes, strict=True) + ] + else: + keys = unnamed_level_values + + # ``stacked`` is in group first-appearance order (groupby with + # sort=False above), which is exactly the order of + # ``unnamed_level_values.unique()``: zip 1:1. + data = ColumnAccessor( + dict(zip(keys, stacked, strict=True)), isinstance(unnamed_level_values, pd.MultiIndex), unnamed_level_values.names, + label_dtype=( + None + if isinstance(unnamed_level_values, pd.MultiIndex) + else unnamed_level_values.dtype + ), + level_dtypes=_pd_index_level_dtypes(unnamed_level_values), ) result = DataFrame._from_data( @@ -8308,7 +8432,18 @@ def unnamed_group_generator(): ) if not future_stack and dropna: - return result.dropna(how="all") + # Compute the row mask explicitly so the eagerly-attached + # codes can be subset alongside the data; pandas keeps the full + # pre-drop level set through dropna. + # _apply_boolean_mask propagates pre-set levels/codes on the + # index automatically. + if isinstance(result, Series): + keep = result.notna() + else: + keep = ~result.isna().all(axis=1) + return result._apply_boolean_mask( + BooleanMask(keep._column, len(result)) + ) else: return result diff --git a/python/cudf/cudf/core/indexed_frame.py b/python/cudf/cudf/core/indexed_frame.py index 7a5b3696649c..fc24fc877f75 100644 --- a/python/cudf/cudf/core/indexed_frame.py +++ b/python/cudf/cudf/core/indexed_frame.py @@ -4749,7 +4749,7 @@ def _apply_boolean_mask(self, boolean_mask: BooleanMask, keep_index=True): plc.Table([col.plc_column for col in cols]), mask_col.plc_column, ) - return self._from_columns_like_self( + result = self._from_columns_like_self( [ ColumnBase.create(col, dtype) for col, dtype in zip( @@ -4759,6 +4759,17 @@ def _apply_boolean_mask(self, boolean_mask: BooleanMask, keep_index=True): column_names=self._column_names, index_names=self.index.names if keep_index else None, ) + if ( + keep_index + and isinstance(self.index, MultiIndex) + and self.index._levels is not None + ): + result.index._levels = self.index._levels + result.index._codes = [ + code.apply_boolean_mask(boolean_mask.column) + for code in self.index._codes + ] + return result def _pandas_repr_compatible(self, nan_rep=None) -> Self: """Return Self but with columns prepared for a pandas-like repr.""" diff --git a/python/cudf/cudf/core/multiindex.py b/python/cudf/cudf/core/multiindex.py index c9c621ce7b98..b846f52ecf02 100644 --- a/python/cudf/cudf/core/multiindex.py +++ b/python/cudf/cudf/core/multiindex.py @@ -2109,14 +2109,14 @@ def _level_index_from_level(self, level) -> int: except ValueError: if not is_integer(level): raise KeyError(f"Level {level} not found") - if level < 0: - level += self.nlevels - if level >= self.nlevels: + norm = level + self.nlevels if level < 0 else level + if not 0 <= norm < self.nlevels: + # matches pandas MultiIndex._get_level_number raise IndexError( - f"Level {level} out of bounds. " - f"Index has {self.nlevels} levels." + f"Too many levels: Index has only {self.nlevels} " + f"levels, {level} is not a valid level number" ) from None - return level + return norm @_performance_tracking def get_indexer(self, target, method=None, limit=None, tolerance=None): diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index 98b9d91de5d4..630d2622ee8b 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -1672,26 +1672,6 @@ def pytest_unconfigure(config): "tests/frame/test_reductions.py::test_reduction_axis_none_returns_scalar[float64-True-mean]": "TODO: Add a reason for failure", "tests/frame/test_reductions.py::test_reduction_axis_none_returns_scalar[float64-True-median]": "TODO: Add a reason for failure", "tests/frame/test_reductions.py::test_reduction_axis_none_returns_scalar[float64-True-skew]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_int_level_names[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_int_level_names[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_ints[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_ints[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_mixed_level[False]": "AssertionError: DataFrame.columns are different", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_mixed_level[True]": "AssertionError: DataFrame.columns are different", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_mixed_levels[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_mixed_levels[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_multi_preserve_categorical_dtype[False-labels0-data0-False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_multi_preserve_categorical_dtype[False-labels0-data0-True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_multi_preserve_categorical_dtype[False-labels1-data1-False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_multi_preserve_categorical_dtype[False-labels1-data1-True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_multi_preserve_categorical_dtype[True-labels0-data0-False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_multi_preserve_categorical_dtype[True-labels0-data0-True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_multi_preserve_categorical_dtype[True-labels1-data1-False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_multi_preserve_categorical_dtype[True-labels1-data1-True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_preserve_categorical_dtype[False-False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_preserve_categorical_dtype[False-True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_preserve_categorical_dtype[True-False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_preserve_categorical_dtype[True-True]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_unstack[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_unstack[True]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_bool": "AssertionError: DataFrame.iloc[:, 0] (column name='('col', 'c')') are different", @@ -1704,19 +1684,12 @@ def pytest_unconfigure(config): "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_swaplevel_sortlevel[0]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_swaplevel_sortlevel[baz]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_unused_levels": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_multi_level_stack_categorical[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_multi_level_stack_categorical[True]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack[True]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_level_name[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_level_name[True]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_multiple_bug[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_multiple_bug[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_multiple_out_of_bounds[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_multiple_out_of_bounds[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_in_multiindex_columns[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_in_multiindex_columns[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_level[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_multiple[False]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_multiple[True]": "TODO: Add a reason for failure", "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_preserve_names[False]": "TODO: Add a reason for failure", @@ -1817,7 +1790,6 @@ def pytest_unconfigure(config): "tests/groupby/test_api.py::test_tab_completion": "TODO: Add a reason for failure", "tests/groupby/test_apply.py::test_apply_with_date_in_multiindex_does_not_convert_to_timestamp": "cudf stores datetime.date values as datetime64; the date type identity is lost on the GPU round trip", "tests/groupby/test_apply.py::test_positional_slice_groups_datetimelike": "the frame and its column Series are converted to pandas independently on fallback, losing the CoW block identity pandas' is_in_obj grouper check requires", - "tests/groupby/test_categorical.py::test_describe_categorical_columns": "cudf's multi-level groupby aggregation and stack() drop the categorical column-index dtype", "tests/groupby/test_cumulative.py::test_groupby_cumprod_nan_influences_other_columns": "TODO: Add a reason for failure", "tests/groupby/test_cumulative.py::test_numpy_compat[cumprod]": "TODO: Add a reason for failure", "tests/groupby/test_cumulative.py::test_numpy_compat[cumsum]": "TODO: Add a reason for failure", diff --git a/python/cudf/cudf/tests/reshape/test_stack.py b/python/cudf/cudf/tests/reshape/test_stack.py index ac08ece520e2..640da145bd02 100644 --- a/python/cudf/cudf/tests/reshape/test_stack.py +++ b/python/cudf/cudf/tests/reshape/test_stack.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import numpy as np @@ -165,3 +165,48 @@ def test_df_stack_multiindex_column_axis_pd_example(level): got = gdf.stack(level=level, future_stack=True) assert_eq(expect, got) + + +def test_df_stack_int_level_names_resolved_positionally(): + # integer level *names* must not hijack positional level lookup + # (pandas' get_level_values resolves integer arguments by name first) + columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y")], names=[1, 0]) + pdf = pd.DataFrame([[1, 2], [3, 4]], columns=columns) + gdf = cudf.from_pandas(pdf) + for level in (0, 1): + assert_eq( + pdf.stack(level=level, future_stack=True), + gdf.stack(level=level, future_stack=True), + ) + + +@pytest.mark.parametrize("level", [2, -3]) +def test_df_stack_out_of_bounds_level_raises(level): + columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y")]) + gdf = cudf.DataFrame([[1, 2]], columns=columns) + with pytest.raises(IndexError, match="Too many levels"): + gdf.stack(level=level) + + +def test_df_stack_duplicate_level_name_raises(): + columns = pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y")], names=["c", "c"] + ) + gdf = cudf.DataFrame([[1, 2]], columns=columns) + with pytest.raises(ValueError, match="occurs multiple times"): + gdf.stack(level="c") + + +def test_df_stack_unsorted_column_permutation_appearance_order(): + # a 3-cycle column permutation: the previous argsort-based reordering + # misaligned column data for non-involution permutations; stacked keys + # are emitted in appearance order like pandas + columns = pd.MultiIndex.from_tuples( + [("b", 1), ("c", 2), ("a", 3)], names=["l0", "l1"] + ) + pdf = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=columns) + gdf = cudf.from_pandas(pdf) + assert_eq( + pdf.stack("l0", future_stack=True), + gdf.stack("l0", future_stack=True), + ) diff --git a/python/cudf/cudf/tests/reshape/test_unstack.py b/python/cudf/cudf/tests/reshape/test_unstack.py index 0edf16202d31..0ca6930443f2 100644 --- a/python/cudf/cudf/tests/reshape/test_unstack.py +++ b/python/cudf/cudf/tests/reshape/test_unstack.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import re @@ -75,12 +75,7 @@ def test_unstack_multiindex(level): [ pd.Index(range(0, 5), name=None), pd.Index(range(0, 5), name="row_index"), - pytest.param( - pd.CategoricalIndex(["d", "e", "f", "g", "h"]), - marks=pytest.mark.xfail( - reason="Categorical column indexes not supported" - ), - ), + pd.CategoricalIndex(["d", "e", "f", "g", "h"]), ], ) @pytest.mark.parametrize( From f0a24f4df40fcbedaf0b4367341f0e91cee9faf5 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Thu, 23 Jul 2026 17:49:57 +0100 Subject: [PATCH 12/16] Rework approach to cudf-streaming bloom filter sizing (#23067) Rather than asking for a fitting number of blocks for an size, just require the user to provide the size of the filter directly. Since it must be a multiple of the storage type, also offer an aligned_size utility that takes an input and makes that happen. Additionally, share the new cuco parametric filter policy across cudf and cudf-streaming. Authors: - Lawrence Mitchell (https://github.com/wence-) - Vyas Ramasubramani (https://github.com/vyasr) - Tom Augspurger (https://github.com/TomAugspurger) Approvers: - Tom Augspurger (https://github.com/TomAugspurger) - Muhammad Haseeb (https://github.com/mhaseeb123) - Yunsong Wang (https://github.com/PointKernel) URL: https://github.com/rapidsai/cudf/pull/23067 --- cpp/include/cudf/reduction/bloom_filter.cuh | 27 ++++++ .../benchmarks/streaming/ndsh/q03.cpp | 6 +- .../benchmarks/streaming/ndsh/q04.cpp | 6 +- .../benchmarks/streaming/ndsh/q21.cpp | 6 +- .../include/cudf_streaming/bloom_filter.hpp | 37 ++++---- .../detail/device_bloom_filter.hpp | 40 +++++---- cpp/libcudf_streaming/src/bloom_filter.cpp | 38 ++++++--- .../src/detail/device_bloom_filter.cu | 58 +++++++------ cpp/libcudf_streaming/tests/CMakeLists.txt | 6 +- .../tests/streaming/test_bloom_filter.cu | 85 +++++++++++++++++++ .../streaming/test_bloom_filter_config.cpp | 41 +++++++++ cpp/src/io/parquet/bloom_filter_reader.cu | 13 +-- cpp/tests/io/parquet_bloom_filter_test.cu | 17 +--- .../cudf_polars/streaming/actor_graph/join.py | 6 +- .../cudf_streaming/bloom_filter.pxd | 11 +-- .../cudf_streaming/bloom_filter.pyi | 4 +- .../cudf_streaming/bloom_filter.pyx | 21 ++--- .../cudf_streaming/tests/test_bloom_filter.py | 17 +++- 18 files changed, 309 insertions(+), 130 deletions(-) create mode 100644 cpp/include/cudf/reduction/bloom_filter.cuh create mode 100644 cpp/libcudf_streaming/tests/streaming/test_bloom_filter.cu create mode 100644 cpp/libcudf_streaming/tests/streaming/test_bloom_filter_config.cpp diff --git a/cpp/include/cudf/reduction/bloom_filter.cuh b/cpp/include/cudf/reduction/bloom_filter.cuh new file mode 100644 index 000000000000..c8c7b3d56df8 --- /dev/null +++ b/cpp/include/cudf/reduction/bloom_filter.cuh @@ -0,0 +1,27 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +namespace cudf { + +/** + * @brief Policy describing the Apache Arrow Block-Split Bloom Filter layout. + * + * Uses cuco's `parametric_filter_policy` with the Apache Arrow layout: 256-bit blocks (8 x + * `uint32_t`), 8 fingerprint bits per key, fully horizontal add (Theta=8), and fully vertical + * contains (Phi=8). This layout is bit-compatible with Apache Arrow. + * + * @tparam Hash The hash function used to generate a hash for each key. + */ +template +using arrow_filter_policy = + cuco::parametric_filter_policy; + +} // namespace cudf diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp index af94c8992bf1..900cc69f3856 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp @@ -357,8 +357,8 @@ int main(int argc, char** argv) int device; RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); - auto const num_filter_blocks = - cudf_streaming::bloom_filter::fitting_num_blocks(static_cast(l2size)); + auto const filter_size = + cudf_streaming::bloom_filter::aligned_size(static_cast(l2size) * 2 / 3); for (int i = 0; i < arguments.num_iterations; i++) { int op_id{0}; @@ -406,7 +406,7 @@ int main(int argc, char** argv) actors.push_back(fanout_bounded( ctx, comm, customer_x_orders, bloom_filter_input, {0}, customer_x_orders_input)); auto bloom_filter = - cudf_streaming::bloom_filter(ctx, comm, cudf::DEFAULT_HASH_SEED, num_filter_blocks); + cudf_streaming::bloom_filter(ctx, comm, cudf::DEFAULT_HASH_SEED, filter_size); actors.push_back(bloom_filter.build( bloom_filter_input, bloom_filter_output, static_cast(10 * i + op_id++))); // Out: l_orderkey, l_extendedprice, l_discount diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp index 515fa9bff1dc..86bc8482507e 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp @@ -256,8 +256,8 @@ int main(int argc, char** argv) int device; RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); - auto const num_filter_blocks = - cudf_streaming::bloom_filter::fitting_num_blocks(static_cast(l2size)); + auto const filter_size = + cudf_streaming::bloom_filter::aligned_size(static_cast(l2size) * 2 / 3); for (int i = 0; i < arguments.num_iterations; i++) { rapidsmpf::OpID op_id{0}; @@ -302,7 +302,7 @@ int main(int argc, char** argv) // Build bloom filter from filtered orders' o_orderkey auto bloom_filter_output = ctx->create_channel(); auto bloom_filter = - cudf_streaming::bloom_filter(ctx, comm, cudf::DEFAULT_HASH_SEED, num_filter_blocks); + cudf_streaming::bloom_filter(ctx, comm, cudf::DEFAULT_HASH_SEED, filter_size); actors.push_back(bloom_filter.build( bloom_filter_input, bloom_filter_output, static_cast(10 * i + op_id++))); diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp index 635b4b0a7f5a..7eaa08229187 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp @@ -499,8 +499,8 @@ int main(int argc, char** argv) int device; RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); - auto const num_filter_blocks = - cudf_streaming::bloom_filter::fitting_num_blocks(static_cast(l2size)); + auto const filter_size = + cudf_streaming::bloom_filter::aligned_size(static_cast(l2size) * 2 / 3); for (int i = 0; i < arguments.num_iterations; i++) { int op_id{0}; std::vector actors; @@ -670,7 +670,7 @@ int main(int argc, char** argv) rapidsmpf::streaming::actor::FanoutPolicy::UNBOUNDED)); auto bloom_output = ctx->create_channel(); auto bloom_filter = - cudf_streaming::bloom_filter(ctx, comm, cudf::DEFAULT_HASH_SEED, num_filter_blocks); + cudf_streaming::bloom_filter(ctx, comm, cudf::DEFAULT_HASH_SEED, filter_size); // Select the relevant key column(s) and build filter. actors.push_back(populate_bloom_filter(ctx, comm, diff --git a/cpp/libcudf_streaming/include/cudf_streaming/bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/bloom_filter.hpp index 7ea1ec567b35..9e87cc832b2f 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/bloom_filter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/bloom_filter.hpp @@ -43,18 +43,25 @@ struct bloom_filter { * @param ctx Streaming context. * @param comm Communicator for the collective operation. * @param seed Hash seed used when hashing values into the filter. - * @param num_filter_blocks Number of blocks in the filter. + * @param filter_size Filter storage size in bytes. Must be positive and satisfy + * `aligned_size(filter_size) == filter_size`, and must not exceed the maximum size supported by + * the filter policy. + * + * @throws std::logic_error If `filter_size` is zero, incorrectly aligned, or exceeds the policy + * maximum. */ explicit bloom_filter(std::shared_ptr ctx, std::shared_ptr comm, std::uint64_t seed, - std::size_t num_filter_blocks) noexcept - : ctx_{std::move(ctx)}, - comm_{std::move(comm)}, - seed_{seed}, - num_filter_blocks_{num_filter_blocks} - { - } + std::size_t filter_size); + + /** + * @brief Find the largest valid filter size no greater than a byte count. + * + * @param size Byte count to align. + * @return Largest valid filter size less than or equal to `size`. + */ + [[nodiscard]] static std::size_t aligned_size(std::size_t size) noexcept; /** * @brief Gets the communicator associated with this bloom_filter. @@ -100,22 +107,10 @@ struct bloom_filter { std::shared_ptr ch_out, std::vector keys); - /** - * @brief Compute number of filter blocks that fit in the given L2 cache size. - * - * @param l2size L2 cache size in bytes. - * @return Number of filter blocks that fit. - */ - [[nodiscard]] static std::size_t fitting_num_blocks(std::size_t l2size) noexcept - { - using StorageType = std::uint32_t; - return (l2size * 2) / (3 * sizeof(StorageType)); - } - private: std::shared_ptr ctx_{}; std::shared_ptr comm_{}; std::uint64_t seed_{}; - std::size_t num_filter_blocks_{}; + std::size_t filter_size_{}; }; } // namespace cudf_streaming diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp index 3fcc8f5ddaf1..435087c60ffb 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp @@ -28,40 +28,58 @@ struct device_bloom_filter { /** * @brief Create a filter. * - * @param num_blocks Number of blocks in the filter. + * @param filter_size Filter storage size in bytes. Must be a positive multiple of the filter + * block size and no greater than the maximum supported by the filter policy. * @param seed Seed used for hashing each value. * @param storage Storage to view as a bloom filter, must be appropriately * initialized. */ - device_bloom_filter(std::size_t num_blocks, std::uint64_t seed, void* storage); + device_bloom_filter(std::size_t filter_size, std::uint64_t seed, void* storage); /** * @brief Create a read-only filter. * - * @param num_blocks Number of blocks in the filter. + * @param filter_size Filter storage size in bytes. Must be a positive multiple of the filter + * block size and no greater than the maximum supported by the filter policy. * @param seed Seed used for hashing each value. * @param storage View of storage, must be appropriately initialized. * * @return A const-qualified bloom filter viewing the underlying storage. */ - static device_bloom_filter const view(std::size_t num_blocks, + static device_bloom_filter const view(std::size_t filter_size, std::uint64_t seed, void const* storage); /** * @brief Create uninitialized storage for a filter. * - * @param num_blocks Number of blocks in the filter. + * @param filter_size Filter storage size in bytes. Must be a positive multiple of the filter + * block size and no greater than the maximum supported by the filter policy. * @param stream CUDA stream for device operations. * @param mr Memory resource for allocations. * * @return Unique pointer to a device buffer containing storage for the requested - * number of filter blocks. + * filter size. */ - static std::unique_ptr storage(std::size_t num_blocks, + static std::unique_ptr storage(std::size_t filter_size, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); + /** + * @brief Find the largest valid filter size no greater than a byte count. + * + * @param size Byte count to align. + * @return Largest valid filter size less than or equal to `size`. + */ + [[nodiscard]] static std::size_t aligned_size(std::size_t size) noexcept; + + /** + * Return the largest storage size supported by the filter policy. + * + * Maximum valid filter size in bytes. + */ + [[nodiscard]] static std::size_t max_size() noexcept; + /** * @brief Add values to the filter. * @@ -111,14 +129,6 @@ struct device_bloom_filter { */ [[nodiscard]] std::size_t size() const noexcept; - /** - * @brief @return Number of blocks to use if the filter should fit in a given L2 cache - * size. - * - * @param l2size Size of the L2 cache in bytes. - */ - [[nodiscard]] static std::size_t fitting_num_blocks(std::size_t l2size) noexcept; - private: std::size_t num_blocks_; ///< Number of blocks used in the filter. std::uint64_t seed_; ///< Seed used when hashing values. diff --git a/cpp/libcudf_streaming/src/bloom_filter.cpp b/cpp/libcudf_streaming/src/bloom_filter.cpp index 78acf8431f05..c3411c25a6cb 100644 --- a/cpp/libcudf_streaming/src/bloom_filter.cpp +++ b/cpp/libcudf_streaming/src/bloom_filter.cpp @@ -21,6 +21,24 @@ namespace cudf_streaming { +std::size_t bloom_filter::aligned_size(std::size_t size) noexcept +{ + return detail::device_bloom_filter::aligned_size(size); +} + +bloom_filter::bloom_filter(std::shared_ptr ctx, + std::shared_ptr comm, + std::uint64_t seed, + std::size_t filter_size) + : ctx_{std::move(ctx)}, comm_{std::move(comm)}, seed_{seed}, filter_size_{filter_size} +{ + RAPIDSMPF_EXPECTS(filter_size_ > 0, "Bloom filter storage size must be positive"); + RAPIDSMPF_EXPECTS(filter_size_ == aligned_size(filter_size_), + "Bloom filter storage size must be a multiple of the filter block size"); + RAPIDSMPF_EXPECTS(filter_size_ <= detail::device_bloom_filter::max_size(), + "Bloom filter storage exceeds the maximum size supported by its policy"); +} + rapidsmpf::streaming::Actor bloom_filter::build( std::shared_ptr ch_in, std::shared_ptr ch_out, @@ -35,10 +53,9 @@ rapidsmpf::streaming::Actor bloom_filter::build( auto filter_stream = br->stream_pool()->get_stream(); rapidsmpf::CudaEvent event; auto storage = - cudf_streaming::detail::device_bloom_filter::storage(num_filter_blocks_, filter_stream, mr); + cudf_streaming::detail::device_bloom_filter::storage(filter_size_, filter_stream, mr); RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(storage->data(), 0, storage->size(), filter_stream)); - auto filter = - cudf_streaming::detail::device_bloom_filter(num_filter_blocks_, seed_, storage->data()); + auto filter = cudf_streaming::detail::device_bloom_filter(filter_size_, seed_, storage->data()); rapidsmpf::CudaEvent build_event; build_event.record(filter_stream); while (!ch_out->is_shutdown()) { @@ -69,15 +86,15 @@ rapidsmpf::streaming::Actor bloom_filter::build( comm_, br->move(std::move(storage), filter_stream), br->move( - cudf_streaming::detail::device_bloom_filter::storage(num_filter_blocks_, filter_stream, mr), + cudf_streaming::detail::device_bloom_filter::storage(filter_size_, filter_stream, mr), filter_stream), tag, - [num_blocks = num_filter_blocks_, seed = seed_](rapidsmpf::Buffer const* left, - rapidsmpf::Buffer* right) { + [filter_size = filter_size_, seed = seed_](rapidsmpf::Buffer const* left, + rapidsmpf::Buffer* right) { right->write_access([&](std::byte* out_bytes, rmm::cuda_stream_view stream) { auto const in = - cudf_streaming::detail::device_bloom_filter::view(num_blocks, seed, left->data()); - cudf_streaming::detail::device_bloom_filter(num_blocks, seed, out_bytes) + cudf_streaming::detail::device_bloom_filter::view(filter_size, seed, left->data()); + cudf_streaming::detail::device_bloom_filter(filter_size, seed, out_bytes) .merge(in, stream); }); }); @@ -102,9 +119,8 @@ rapidsmpf::streaming::Actor bloom_filter::apply( "Bloom filter channel contained more than one message"); auto stream = storage.stream(); rapidsmpf::CudaEvent event; - auto filter = - cudf_streaming::detail::device_bloom_filter(num_filter_blocks_, seed_, storage.data()); - auto meta = co_await ch_in->receive_metadata(); + auto filter = cudf_streaming::detail::device_bloom_filter(filter_size_, seed_, storage.data()); + auto meta = co_await ch_in->receive_metadata(); if (!meta.empty()) { co_await ch_out->send_metadata(std::move(meta)); } while (!ch_out->is_shutdown()) { auto msg = co_await ch_in->receive(); diff --git a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu index 0525f139babf..ad3a2217ba13 100644 --- a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu +++ b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu @@ -18,7 +18,8 @@ #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wsign-conversion" #endif -#include +#include + #include #include #include @@ -33,6 +34,7 @@ #include +#include #include #include @@ -53,47 +55,50 @@ namespace cudf_streaming::detail { namespace { using KeyType = std::uint64_t; -using BloomFilterRefType = - cuco::bloom_filter_ref, - cuco::thread_scope_device, - cuco::parametric_filter_policy, - std::uint32_t, - 8, - 8, - 8, - 1, - 1, - 8, - false, - false>>; -using StorageType = BloomFilterRefType::filter_block_type; +using BloomFilterPolicy = cudf::arrow_filter_policy>; +using BloomFilterRefType = cuco::bloom_filter_ref, + cuco::thread_scope_device, + BloomFilterPolicy>; +using StorageType = BloomFilterRefType::filter_block_type; + +std::size_t num_blocks(std::size_t filter_size) +{ + RAPIDSMPF_EXPECTS(filter_size >= sizeof(StorageType), + "Bloom filter storage must contain at least one filter block"); + RAPIDSMPF_EXPECTS(filter_size == device_bloom_filter::aligned_size(filter_size), + "Bloom filter storage size must be a multiple of the filter block size"); + auto const blocks = filter_size / sizeof(StorageType); + RAPIDSMPF_EXPECTS(blocks <= BloomFilterPolicy::max_filter_blocks, + "Bloom filter storage exceeds the maximum size supported by its policy"); + return blocks; +} } // namespace -device_bloom_filter::device_bloom_filter(std::size_t num_blocks, std::uint64_t seed, void* storage) - : num_blocks_{num_blocks}, seed_{seed}, storage_{storage} +device_bloom_filter::device_bloom_filter(std::size_t filter_size, std::uint64_t seed, void* storage) + : num_blocks_{num_blocks(filter_size)}, seed_{seed}, storage_{storage} { RAPIDSMPF_EXPECTS( reinterpret_cast(storage_) % std::alignment_of_v == 0, "Allocation for bloom filter is not aligned."); } -device_bloom_filter const device_bloom_filter::view(std::size_t num_blocks, +device_bloom_filter const device_bloom_filter::view(std::size_t filter_size, std::uint64_t seed, void const* storage) { // const-cast is safe because the returned object is also const and therefore can't // call methods that throw away constness. - return device_bloom_filter(num_blocks, seed, const_cast(storage)); + return device_bloom_filter(filter_size, seed, const_cast(storage)); } -std::unique_ptr device_bloom_filter::storage(std::size_t num_blocks, +std::unique_ptr device_bloom_filter::storage(std::size_t filter_size, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { return std::make_unique( - num_blocks * sizeof(StorageType), std::alignment_of_v, stream, mr); + num_blocks(filter_size) * sizeof(StorageType), std::alignment_of_v, stream, mr); } void device_bloom_filter::add(cudf::table_view const& values_to_hash, @@ -135,9 +140,14 @@ rmm::device_uvector device_bloom_filter::contains(cudf::table_view const& return result; } -std::size_t device_bloom_filter::fitting_num_blocks(std::size_t l2size) noexcept +std::size_t device_bloom_filter::aligned_size(std::size_t size) noexcept +{ + return rmm::align_down(size, std::alignment_of_v); +} + +std::size_t device_bloom_filter::max_size() noexcept { - return (l2size * 2) / (3 * sizeof(StorageType)); + return BloomFilterPolicy::max_filter_blocks * sizeof(StorageType); } void* device_bloom_filter::data() noexcept { return storage_; } diff --git a/cpp/libcudf_streaming/tests/CMakeLists.txt b/cpp/libcudf_streaming/tests/CMakeLists.txt index 2fef944da8a3..1aae79f1f4bd 100644 --- a/cpp/libcudf_streaming/tests/CMakeLists.txt +++ b/cpp/libcudf_streaming/tests/CMakeLists.txt @@ -59,7 +59,7 @@ target_compile_options( ) target_link_libraries( libcudf_streaming_test_sources - PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf::cudftestutil cudf::cudftestutil_impl + PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf::cudftestutil cudf::cudftestutil_impl cuco::cuco PUBLIC GTest::gmock GTest::gtest ) # cudf::cudftestutil_impl injects cudf test-utility .cu sources (via INTERFACE_SOURCES) that are @@ -70,7 +70,9 @@ target_compile_options( ) target_sources( libcudf_streaming_test_sources - PRIVATE streaming/test_table_chunk.cpp + PRIVATE streaming/test_bloom_filter.cu + streaming/test_bloom_filter_config.cpp + streaming/test_table_chunk.cpp streaming/test_read_parquet.cpp streaming/test_channel_metadata.cpp streaming/test_partition.cpp diff --git a/cpp/libcudf_streaming/tests/streaming/test_bloom_filter.cu b/cpp/libcudf_streaming/tests/streaming/test_bloom_filter.cu new file mode 100644 index 000000000000..b8fc5b2e978a --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_bloom_filter.cu @@ -0,0 +1,85 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace { + +using policy_type = cudf::arrow_filter_policy>; + +__global__ void block_index_kernel(std::uint32_t upper_hash, + std::size_t num_blocks, + std::uint32_t* result) +{ + *result = policy_type{}.block_index(upper_hash, cuco::extent{num_blocks}); +} + +TEST(BloomFilterPolicyTest, UsesBlocksBeyondFormerArrowLimit) +{ + constexpr auto arrow_max_blocks = std::size_t{4'194'304}; + constexpr auto num_blocks = arrow_max_blocks + 1; + constexpr auto upper_hash = std::numeric_limits::max(); + auto const stream = cudf::get_default_stream(); + rmm::device_scalar index{0, stream}; + + block_index_kernel<<<1, 1, 0, stream.value()>>>(upper_hash, num_blocks, index.data()); + CUDF_CHECK_CUDA(stream.value()); + + EXPECT_EQ(index.value(stream), arrow_max_blocks); +} + +TEST(DeviceBloomFilterTest, RejectsStorageBeyondPolicyLimit) +{ + using filter_ref_type = cuco::bloom_filter_ref, + cuco::thread_scope_device, + policy_type>; + constexpr auto block_size = sizeof(filter_ref_type::filter_block_type); + constexpr auto too_large = (policy_type::max_filter_blocks + std::size_t{1}) * block_size; + auto const stream = cudf::get_default_stream(); + + EXPECT_THROW(cudf_streaming::detail::device_bloom_filter::storage( + too_large, stream, cudf::get_current_device_resource_ref()), + std::logic_error); +} + +TEST(DeviceBloomFilterTest, RequiresAlignedStorageSize) +{ + constexpr auto unaligned_size = std::size_t{65}; + constexpr auto aligned_size = std::size_t{64}; + auto const stream = cudf::get_default_stream(); + + EXPECT_THROW(cudf_streaming::detail::device_bloom_filter::storage( + unaligned_size, stream, cudf::get_current_device_resource_ref()), + std::logic_error); + + auto storage = cudf_streaming::detail::device_bloom_filter::storage( + aligned_size, stream, cudf::get_current_device_resource_ref()); + + EXPECT_EQ(storage->size(), aligned_size); + + auto const filter = cudf_streaming::detail::device_bloom_filter{aligned_size, 0, storage->data()}; + EXPECT_EQ(filter.size(), aligned_size); +} + +} // namespace diff --git a/cpp/libcudf_streaming/tests/streaming/test_bloom_filter_config.cpp b/cpp/libcudf_streaming/tests/streaming/test_bloom_filter_config.cpp new file mode 100644 index 000000000000..e50807fbc222 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_bloom_filter_config.cpp @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include +#include +#include + +namespace { + +TEST(BloomFilterTest, AlignsStorageSize) +{ + EXPECT_EQ(cudf_streaming::bloom_filter::aligned_size(31), 0); + EXPECT_EQ(cudf_streaming::bloom_filter::aligned_size(32), 32); + EXPECT_EQ(cudf_streaming::bloom_filter::aligned_size(65), 64); +} + +TEST(BloomFilterTest, RequiresAlignedStorageSize) +{ + auto make_filter = [](std::size_t filter_size) { + return cudf_streaming::bloom_filter{std::shared_ptr{}, + std::shared_ptr{}, + 0, + filter_size}; + }; + + EXPECT_THROW(make_filter(0), std::logic_error); + EXPECT_THROW(make_filter(65), std::logic_error); + EXPECT_THROW(make_filter(cudf_streaming::bloom_filter::aligned_size( + std::numeric_limits::max())), + std::logic_error); + EXPECT_NO_THROW(make_filter(64)); +} + +} // namespace diff --git a/cpp/src/io/parquet/bloom_filter_reader.cu b/cpp/src/io/parquet/bloom_filter_reader.cu index e0591f2d2115..e303ba0165f9 100644 --- a/cpp/src/io/parquet/bloom_filter_reader.cu +++ b/cpp/src/io/parquet/bloom_filter_reader.cu @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -24,7 +25,6 @@ #include #include -#include #include #include #include @@ -49,16 +49,7 @@ namespace { * @tparam Key The type of the values to generate a fingerprint for. */ template -using arrow_filter_policy = cuco::parametric_filter_policy, - std::uint32_t, - 8, - 8, - 8, - 1, - 1, - 8, - false, - false>; +using arrow_filter_policy = cudf::arrow_filter_policy>; /** * @brief Converts bloom filter membership results (for each column chunk) to a device column. diff --git a/cpp/tests/io/parquet_bloom_filter_test.cu b/cpp/tests/io/parquet_bloom_filter_test.cu index 4de68ff3c991..7fbb21714eff 100644 --- a/cpp/tests/io/parquet_bloom_filter_test.cu +++ b/cpp/tests/io/parquet_bloom_filter_test.cu @@ -11,12 +11,12 @@ #include #include #include +#include #include #include #include -#include #include @@ -26,19 +26,8 @@ class ParquetBloomFilterTest : public cudf::test::BaseFixture {}; TEST_F(ParquetBloomFilterTest, TestStrings) { - using key_type = StringType; - // Apache Arrow Block-Split Bloom Filter layout, hashing keys with cudf's `XXHash_64` (matching - // `cudf::io::parquet::detail::arrow_filter_policy`). - using policy_type = cuco::parametric_filter_policy, - std::uint32_t, - 8, - 8, - 8, - 1, - 1, - 8, - false, - false>; + using key_type = StringType; + using policy_type = cudf::arrow_filter_policy>; using word_type = policy_type::word_type; std::size_t constexpr num_filter_blocks = 4; diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 46cb66a04f92..a984e69957d1 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -806,9 +806,9 @@ def make_filter_tasks( ch_left = context.create_channel() bloom_apply_output = ch_left - # TODO: configure based on GPU L2 size - nblocks = BloomFilter.fitting_num_blocks(32 * 1024 * 1024) - filter = BloomFilter(context, comm, LIBCUDF_DEFAULT_HASH_SEED, nblocks) + # TODO: Make the filter size configurable. + filter_size = 32 * 1024 * 1024 + filter = BloomFilter(context, comm, LIBCUDF_DEFAULT_HASH_SEED, filter_size) filter_tasks: list[Coroutine[Any, Any, None]] = [] chs_to_shutdown = [ bloom_build_output, diff --git a/python/cudf_streaming/cudf_streaming/bloom_filter.pxd b/python/cudf_streaming/cudf_streaming/bloom_filter.pxd index 4e1c64db46bd..84345a7ac165 100644 --- a/python/cudf_streaming/cudf_streaming/bloom_filter.pxd +++ b/python/cudf_streaming/cudf_streaming/bloom_filter.pxd @@ -5,6 +5,7 @@ from libc.stddef cimport size_t from libc.stdint cimport uint64_t from libcpp.memory cimport shared_ptr, unique_ptr +from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.communicator.communicator cimport Communicator, cpp_Communicator from rapidsmpf.streaming.core.context cimport cpp_Context @@ -15,15 +16,15 @@ cdef extern from "" nogil: shared_ptr[cpp_Context] ctx, shared_ptr[cpp_Communicator] comm, uint64_t seed, - size_t num_filter_blocks, - ) noexcept + size_t filter_size, + ) except +ex_handler const shared_ptr[cpp_Communicator]& comm() noexcept cdef extern from "" nogil: - size_t cpp_fitting_num_blocks \ - "cudf_streaming::bloom_filter::fitting_num_blocks"( - size_t l2size + size_t cpp_aligned_size \ + "cudf_streaming::bloom_filter::aligned_size"( + size_t size ) noexcept diff --git a/python/cudf_streaming/cudf_streaming/bloom_filter.pyi b/python/cudf_streaming/cudf_streaming/bloom_filter.pyi index ea879d755771..0c8f6ab0ff4c 100644 --- a/python/cudf_streaming/cudf_streaming/bloom_filter.pyi +++ b/python/cudf_streaming/cudf_streaming/bloom_filter.pyi @@ -26,12 +26,12 @@ class BloomFilter: ctx: Context, comm: Communicator, seed: int, - num_filter_blocks: int, + filter_size: int, ) -> None: ... @property def comm(self) -> Communicator: ... @staticmethod - def fitting_num_blocks(l2size: int) -> int: ... + def aligned_size(size: int) -> int: ... async def build( self, ctx: Context, diff --git a/python/cudf_streaming/cudf_streaming/bloom_filter.pyx b/python/cudf_streaming/cudf_streaming/bloom_filter.pyx index 5116a8aa44e2..23b4c6509423 100644 --- a/python/cudf_streaming/cudf_streaming/bloom_filter.pyx +++ b/python/cudf_streaming/cudf_streaming/bloom_filter.pyx @@ -118,8 +118,9 @@ cdef class BloomFilter: The communicator the bloom filter construction is collective over. seed Seed used for hashing values into the bloom filter. - num_filter_blocks - Number of blocks used to size the filter. + filter_size + Filter storage size in bytes. Must be positive and satisfy + ``BloomFilter.aligned_size(filter_size) == filter_size``. """ def __init__( @@ -127,7 +128,7 @@ cdef class BloomFilter: Context ctx not None, Communicator comm not None, uint64_t seed, - size_t num_filter_blocks, + size_t filter_size, ): self._comm = comm with nogil: @@ -135,7 +136,7 @@ cdef class BloomFilter: ctx._handle, comm._handle, seed, - num_filter_blocks, + filter_size, ) def __dealloc__(self): @@ -154,22 +155,22 @@ cdef class BloomFilter: return self._comm @staticmethod - def fitting_num_blocks(size_t l2size): + def aligned_size(size_t size): """ - Return the number of blocks needed to fit within an L2 cache size. + Return the largest valid filter size no greater than a byte count. Parameters ---------- - l2size - Size of the L2 cache in bytes. + size + Byte count to align. Returns ------- - Number of blocks to use in the filter. + Largest valid filter size less than or equal to ``size``. """ cdef size_t ret with nogil: - ret = cpp_fitting_num_blocks(l2size) + ret = cpp_aligned_size(size) return ret async def build( diff --git a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py index d00e7926e50e..50089c2f4250 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py @@ -42,6 +42,17 @@ def make_table( ) +def test_aligned_size() -> None: + assert BloomFilter.aligned_size(31) == 0 + assert BloomFilter.aligned_size(32) == 32 + assert BloomFilter.aligned_size(65) == 64 + + +def test_requires_aligned_size(context: Context, comm: Communicator) -> None: + with pytest.raises(RuntimeError, match="must be a multiple"): + BloomFilter(context, comm, seed=0, filter_size=65) + + @define_actor() async def add_metadata( ctx: Context, ch_in: Channel[TableChunk], ch_out: Channel[TableChunk] @@ -94,13 +105,13 @@ def run_bloom_filter_pipeline( probe_table: TableChunk, *, seed: int = 42, - l2size: int = 1 << 20, + filter_size: int = 1 << 20, ) -> list[Message]: bloom = BloomFilter( context, comm, seed=seed, - num_filter_blocks=BloomFilter.fitting_num_blocks(l2size), + filter_size=filter_size, ) build_msg = Message(0, build_table) @@ -185,7 +196,7 @@ def test_bloom_filter_build_exception_no_shutdown( context, comm, seed=42, - num_filter_blocks=BloomFilter.fitting_num_blocks(1 << 20), + filter_size=(1 << 20), ) ch_in: Channel[TableChunk] = context.create_channel() ch_out: Channel[BloomFilterChunk] = context.create_channel() From 1ef4d74fb7fb558519516e28f96aa32747749253 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 23 Jul 2026 11:50:42 -0500 Subject: [PATCH 13/16] Disable join filter pushdown rewrites by default (#23407) This changes the cudf-polars default configuration to disable the join filter pushdown rewrite added in https://github.com/rapidsai/cudf/pull/22996. It can be re-enabled through an env var with CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN=1, or through python by passing `join_filter_pushdown=JoinFilterPushdownOptions()` I'm proposing to disable this for the 26.08 release based on some slowdowns in our pds-h benchmarks at smaller scales (SF-1K on a single H100 GPU). image Eventually, I think that the join filter pushdown optimization should be on by default. But I'd like to understand a bit more about the slowdowns observed in that benchmark before turning it on by default. Authors: - Tom Augspurger (https://github.com/TomAugspurger) Approvers: - Vyas Ramasubramani (https://github.com/vyasr) URL: https://github.com/rapidsai/cudf/pull/23407 --- python/cudf_polars/cudf_polars/utils/config.py | 5 ++++- python/cudf_polars/tests/test_config.py | 6 +++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 8a265d875b39..e24270e99d11 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -752,6 +752,9 @@ class StreamingExecutor: Options controlling the logical join-domain prefilter rewrite. See :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions` for more. ``None`` disables the rewrite. + + Enable through environment variables with + ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN=1``. max_io_threads Maximum number of IO threads. Default is 4. This controls the parallelism of IO operations when reading data. @@ -1102,7 +1105,7 @@ def from_polars_engine( ) if user_join_filter_pushdown is None: env_join_filter_pushdown = os.environ.get( - "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "1" + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "0" ) if not _bool_converter(env_join_filter_pushdown): user_executor_options["join_filter_pushdown"] = None diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 1d1395b48f0f..622a0ab3ff07 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -615,9 +615,7 @@ def test_dynamic_planning_defaults() -> None: assert config.executor.dynamic_planning.join_prefilter_threshold == 0.5 assert config.executor.dynamic_planning.join_prefilter_max_key_columns == 1 assert not config.executor.dynamic_planning.join_prefilter_trace - assert config.executor.join_filter_pushdown is not None - assert config.executor.join_filter_pushdown.threshold == 0.5 - assert not config.executor.join_filter_pushdown.trace + assert config.executor.join_filter_pushdown is None def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -642,6 +640,7 @@ def test_dynamic_planning_sample_chunk_count_from_env( def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "1") monkeypatch.setenv( "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_THRESHOLD", "0.25" ) @@ -665,6 +664,7 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non def test_join_filter_pushdown_options_from_env( monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "1") monkeypatch.setenv( "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__THRESHOLD", "0.125" ) From b1fbddacd432a8fd66833cd1f325635c407bc6bd Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Thu, 23 Jul 2026 13:27:24 -0500 Subject: [PATCH 14/16] Match pandas ordering, dtypes, and validation in unstack/_pivot (#23368) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split out of #23255 (4/6). **Depends on #23365 (fidelity), the stack PR (#23370), and #23366 (groupby agg)** — 8 of the 29 un-xfailed pandas-tests need those fixes too, so this PR's pandas-tests job goes green once they merge. - Order result rows/columns by the removed level's codes (level order preserved, missing keys first) instead of sorted values with nulls last, by encoding the integer code columns instead of the level values. - Propagate the source frame's column-axis level names into the result instead of hardcoding `None`; also fixes the `ValueError: Length of names must match number of levels` crash when unstacking MultiIndex-column frames. - Promote integer source columns to float64 when the reshape introduces missing cells (pandas' block semantics), gated on `mode.pandas_compatible`; `pivot_table`/`crosstab` opt out via a module-private `_unstack` parameter when `fill_value` fills the cells afterwards. - Preserve unused categories of the removed level in the result's column levels (pandas GH 17845); also fixes a libcudf `Column sizes don't match` crash for indexes with unused categorical categories. - Validate the level on flat-index frames (`KeyError`) and duplicated index names (`ValueError`) like pandas; `pivot` with `values=` drops the original columns-axis names. Fixes 29 pandas-tests; their xfail entries are removed, three remaining `test_stack_unstack.py` entries get real failure reasons, and two classic categorical unstack params are un-xfailed. Attribution verified per node id against isolated builds: 21 pass with only this change, 4 need the stack PR, 2 need stack+fidelity, 2 need the groupby-agg PR. Authors: - GALI PREM SAGAR (https://github.com/galipremsagar) Approvers: - Matthew Roeschke (https://github.com/mroeschke) URL: https://github.com/rapidsai/cudf/pull/23368 --- python/cudf/cudf/core/multiindex.py | 25 +-- python/cudf/cudf/core/reshape.py | 208 ++++++++++++++++-- .../pandas/scripts/pandas-testing-plugin.py | 35 +-- .../tests/dataframe/methods/test_swaplevel.py | 21 +- .../cudf/tests/reshape/test_pivot_table.py | 31 ++- .../cudf/cudf/tests/reshape/test_unstack.py | 14 +- 6 files changed, 257 insertions(+), 77 deletions(-) diff --git a/python/cudf/cudf/core/multiindex.py b/python/cudf/cudf/core/multiindex.py index b846f52ecf02..1004863131e8 100644 --- a/python/cudf/cudf/core/multiindex.py +++ b/python/cudf/cudf/core/multiindex.py @@ -1738,21 +1738,18 @@ def swaplevel(self, i=-2, j=-1) -> Self: ('aa', 'b')], ) """ - name_i = self._column_names[i] if isinstance(i, int) else i - name_j = self._column_names[j] if isinstance(j, int) else j - to_swap = {name_i, name_j} - new_data = {} + lvl_i = self._level_index_from_level(i) + lvl_j = self._level_index_from_level(j) + order = list(range(self.nlevels)) + order[lvl_i], order[lvl_j] = order[lvl_j], order[lvl_i] + keys = self._column_names + columns = self._columns # TODO: Preserve self._codes and self._levels if set - for k, v in self._column_labels_and_values: - if k not in to_swap: - new_data[k] = v - elif k == name_i: - new_data[name_j] = self._data[name_j] - elif k == name_j: - new_data[name_i] = self._data[name_i] - midx = type(self)._from_data(new_data) - if all(n is None for n in self.names): - midx = midx.set_names(self.names) + midx = type(self)._from_data({keys[k]: columns[k] for k in order}) + # the accessor keys may be positional stand-ins (e.g. for + # duplicated or unset level names), so carry the level names + # over from ``self.names`` rather than deriving them from keys + midx.names = [self.names[k] for k in order] return midx @_performance_tracking diff --git a/python/cudf/cudf/core/reshape.py b/python/cudf/cudf/core/reshape.py index e8736f789377..b521a34a0e87 100644 --- a/python/cudf/cudf/core/reshape.py +++ b/python/cudf/cudf/core/reshape.py @@ -11,7 +11,7 @@ import cudf from cudf.api.extensions import no_default -from cudf.api.types import is_list_like, is_scalar +from cudf.api.types import is_integer, is_list_like, is_scalar from cudf.core.column import ( ColumnBase, as_column, @@ -938,8 +938,11 @@ def get_dummies( def _pivot( col_accessor: ColumnAccessor, - index: Index | MultiIndex, - columns: Index | MultiIndex, + index_labels: Index | MultiIndex, + index_idx: ColumnBase, + columns_labels: Index | MultiIndex, + columns_idx: ColumnBase, + promote_ints_on_missing: bool = False, ) -> DataFrame: """ Reorganize the values of the DataFrame according to the given @@ -947,14 +950,22 @@ def _pivot( Parameters ---------- - col_accessor : DataFrame - index : Index - Index labels of the result - columns : Index - Column labels of the result + col_accessor : ColumnAccessor + Values to pivot into the result's columns. + index_labels : Index + Distinct index keys; row labels of the result. + index_idx : ColumnBase + Position of each source row's key within ``index_labels``. + columns_labels : Index + Distinct column keys; labels of the result's new column level(s). + columns_idx : ColumnBase + Position of each source row's key within ``columns_labels``. + promote_ints_on_missing : bool + Promote integer source columns to float64 when the reshape + introduces missing cells, as pandas' unstack does. The unstack + and pivot paths want this; pivot_table/crosstab fill missing + cells afterwards and keep the integer dtype. """ - columns_labels, columns_idx = columns._encode() - index_labels, index_idx = index._encode() column_labels = columns_labels.to_pandas().to_flat_index() result = {} @@ -964,12 +975,26 @@ def as_tuple(x): return x if isinstance(x, tuple) else (x,) nrows = len(index_labels) + promote_ints = promote_ints_on_missing and cudf.get_option( + "mode.pandas_compatible" + ) for col_label, col in col_accessor.items(): names = [ as_tuple(col_label) + as_tuple(name) for name in column_labels ] new_size = nrows * len(names) scatter_map = (columns_idx * np.int32(nrows)) + index_idx + if ( + promote_ints + and new_size > len(col) + and isinstance(col.dtype, np.dtype) + and col.dtype.kind in "iu" + ): + # pandas builds one 2-D values block per source column and + # promotes the whole block to float64 when the reshape + # introduces missing entries, so even gap-free result + # columns become float64 + col = col.astype(np.dtype(np.float64)) target_col = column_empty(row_count=new_size, dtype=col.dtype) target_col[scatter_map] = col result.update( @@ -984,16 +1009,82 @@ def as_tuple(x): ) ) - # the result of pivot always has a MultiIndex + # the result of pivot always has a MultiIndex; the leading level(s) + # come from the source frame's column labels, so preserve their names ca = ColumnAccessor( result, multiindex=True, - level_names=(None, *columns._column_names), + level_names=( + *col_accessor.level_names, + *columns_labels.names, + ), verify=False, ) return cudf.DataFrame._from_data(ca, index=index_labels) +def _unstack_encode_by_codes( + mi: MultiIndex, level +) -> tuple[Index | MultiIndex, ColumnBase, Index | MultiIndex, ColumnBase]: + """Encode unstack keys ordered by the MultiIndex level codes. + + libcudf's ``encode`` orders distinct keys by sorted value with nulls + last, but pandas' unstack orders keys by the index's level codes: the + level order is preserved and missing entries (code -1) come first. + Encoding the integer code columns instead of the level values yields + exactly that order. + """ + lvl_idx = mi._level_index_from_level(level) + mi._maybe_materialize_codes_and_levels() + names = mi.names + + def encode_side(sel: list[int]) -> tuple[Index | MultiIndex, ColumnBase]: + code_cols = [] + for i in sel: + code = mi._codes[i].astype(np.dtype(np.int64)).copy() # type: ignore[index] + # Normalize the NA sentinel (``MultiIndex.__init__`` stores + # ``iinfo(SIZE_TYPE_DTYPE).min``, lazy factorization stores -1) + # so the missing-key group encodes as one key that sorts first. + code[code < 0] = -1 + code_cols.append(code) + code_frame = cudf.DataFrame._from_data( + ColumnAccessor(dict(enumerate(code_cols)), verify=False) + ) + key_codes, idx = code_frame._encode() + labels_data = {} + out_levels = [] + out_codes = [] + for j, i in enumerate(sel): + kc = key_codes._columns[j].astype(np.dtype(np.int64)) + out_levels.append(mi._levels[i]) # type: ignore[index] + out_codes.append(kc) + gather_codes = kc.copy() + gather_codes[gather_codes == -1] = np.iinfo(SIZE_TYPE_DTYPE).min + # key by position: level names may be duplicated or None + labels_data[j] = mi._levels[i]._column.take( # type: ignore[index] + gather_codes, nullify=True + ) + if len(labels_data) == 1: + labels: Index | MultiIndex = cudf.Index._from_column( + next(iter(labels_data.values())), name=names[sel[0]] + ) + else: + mi_labels = cudf.MultiIndex._from_data(labels_data) + mi_labels.names = [names[i] for i in sel] + # carry the original level objects and the keys' codes so that + # a subsequent unstack/stack keeps ordering by the original + # levels, exactly like pandas (which reuses the level objects) + mi_labels._levels = out_levels + mi_labels._codes = out_codes + labels = mi_labels + return labels, idx + + remaining = [i for i in range(mi.nlevels) if i != lvl_idx] + index_labels, index_idx = encode_side(remaining) + columns_labels, columns_idx = encode_side([lvl_idx]) + return index_labels, index_idx, columns_labels, columns_idx + + def pivot( data: DataFrame, columns=None, index=no_default, values=no_default ) -> DataFrame: @@ -1125,8 +1216,20 @@ def pivot( if len(columns_index) != len(columns_index.drop_duplicates()): raise ValueError("Duplicate index-column pairs found. Cannot reshape.") + selection = data._data.select_by_label(cols_to_select) + if values is not no_default: + # pandas rebuilds the columns axis from ``values`` and drops the + # original columns-axis name(s) + selection._level_names = (None,) * selection.nlevels + columns_labels, columns_idx = column_data._encode() + index_labels, index_idx = index_data._encode() result = _pivot( - data._data.select_by_label(cols_to_select), index_data, column_data + selection, + index_labels, + index_idx, + columns_labels, + columns_idx, + promote_ints_on_missing=True, ) result._attrs = data.attrs @@ -1229,6 +1332,23 @@ def unstack(df, level, fill_value=None, sort: bool = True): 2 7 dtype: int64 """ + return _unstack(df, level, fill_value=fill_value, sort=sort) + + +def _unstack( + df, + level, + fill_value=None, + sort: bool = True, + promote_ints_on_missing: bool = True, +): + """``unstack`` implementation. + + ``promote_ints_on_missing`` promotes integer source columns to float64 + when the reshape introduces missing cells, like pandas' unstack. + ``pivot_table`` (and thereby ``crosstab``) disables it because those fill + the missing cells afterwards and keep the integer dtype. + """ if not isinstance(df, cudf.DataFrame): raise ValueError("`df` should be a cudf Dataframe object.") @@ -1256,7 +1376,14 @@ def unstack(df, level, fill_value=None, sort: bool = True): if not is_scalar(level): if not level: return df + if len(level) == 1: + # pandas normalizes a length-1 list-like level to a scalar + level = level[0] if not isinstance(df.index, cudf.MultiIndex): + if not is_integer(level): + # pandas validates non-integer levels against the flat index + # name and raises KeyError on a mismatch + df.index._validate_index_level(level) dtype = df._columns[0].dtype if any(col_dtype != dtype for _, col_dtype in df._dtypes): raise ValueError( @@ -1271,10 +1398,22 @@ def unstack(df, level, fill_value=None, sort: bool = True): res._attrs = df.attrs return res else: - index = df.index.droplevel(level) + from cudf.core.indexed_frame import _check_duplicate_level_names + + specified = [level] if is_scalar(level) else list(level) + _check_duplicate_level_names( + [lv for lv in specified if not is_integer(lv)], + df.index.names, + ) if is_scalar(level): - columns = df.index.get_level_values(level) + # order rows/columns by the removed level's codes (pandas + # semantics: level order preserved, missing entries first), + # not by sorted level values + index_labels, index_idx, columns_labels, columns_idx = ( + _unstack_encode_by_codes(df.index, level) + ) else: + index = df.index.droplevel(level) new_names = [] ca_data = {} for lev in level: @@ -1285,8 +1424,37 @@ def unstack(df, level, fill_value=None, sort: bool = True): ColumnAccessor(ca_data, verify=False) ) columns.names = new_names - result = _pivot(df, index, columns) + columns_labels, columns_idx = columns._encode() + index_labels, index_idx = index._encode() + result = _pivot( + df._data, + index_labels, + index_idx, + columns_labels, + columns_idx, + promote_ints_on_missing=promote_ints_on_missing, + ) result._attrs = df.attrs + if is_scalar(level) and result._data.multiindex: + # pandas keeps unused categories of the removed level + # ("removed_level_full", pandas GH 17845) in + # result.columns.levels even though no columns are created + # for them. + _, level_idx = df.index._level_to_ca_label(level) + full_level = df.index.levels[level_idx].to_pandas() + pdi = result._data.to_pandas_index + level_values = pdi.get_level_values(-1) + new_codes = full_level.get_indexer(level_values) + # -1 codes for NA labels are pandas' canonical missing + # representation; only bail out when a non-NA label failed + # to map into the full level + if ((new_codes >= 0) | pd.isna(level_values)).all(): + result._data.to_pandas_index = pd.MultiIndex( + levels=[*pdi.levels[:-1], full_level], + codes=[*pdi.codes[:-1], new_codes], + names=pdi.names, + verify_integrity=False, + ) if result.index.nlevels == 1: result.index = result.index.get_level_values(result.index.names[0]) return result @@ -1600,7 +1768,13 @@ def pivot_table( to_unstack.append(i) else: to_unstack.append(name) - table = agged.unstack(to_unstack) + table = _unstack( + agged, + to_unstack, + # pandas keeps the integer dtype when the missing cells are + # filled afterwards, and promotes to float64 when they are not + promote_ints_on_missing=fill_value is None, + ) if fill_value is not None: table = table.fillna(fill_value) diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index 630d2622ee8b..27c8a46ecbe9 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -1302,7 +1302,6 @@ def pytest_unconfigure(config): "tests/frame/methods/test_shift.py::TestDataFrameShift::test_shift_dt64values_axis1_invalid_fill[datetime64[us]-False]": "AssertionError: assert 1 == 2", "tests/frame/methods/test_shift.py::TestDataFrameShift::test_shift_dt64values_axis1_invalid_fill[timedelta64[us]-False]": "AssertionError: assert 1 == 2", "tests/frame/methods/test_shift.py::TestDataFrameShift::test_shift_dt64values_int_fill_deprecated": "TODO: Add a reason for failure", - "tests/frame/methods/test_sort_index.py::TestDataFrameSortIndex::test_sort_index_intervalindex": "TODO: Add a reason for failure", "tests/frame/methods/test_sort_index.py::TestDataFrameSortIndex::test_sort_index_nan": "TODO: Add a reason for failure", "tests/frame/methods/test_sort_values.py::TestDataFrameSortValues::test_sort_by_column_named_none": "AssertionError: DataFrame.index are different", "tests/frame/methods/test_sort_values.py::TestDataFrameSortValues::test_sort_values_by_empty_list": "TODO: Add a reason for failure", @@ -1672,32 +1671,9 @@ def pytest_unconfigure(config): "tests/frame/test_reductions.py::test_reduction_axis_none_returns_scalar[float64-True-mean]": "TODO: Add a reason for failure", "tests/frame/test_reductions.py::test_reduction_axis_none_returns_scalar[float64-True-median]": "TODO: Add a reason for failure", "tests/frame/test_reductions.py::test_reduction_axis_none_returns_scalar[float64-True-skew]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_unstack[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_stack_unstack[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_bool": "AssertionError: DataFrame.iloc[:, 0] (column name='('col', 'c')') are different", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_multi_level_rows_and_cols": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_nan_index2": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_nan_index3": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_non_unique_index_names[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_non_unique_index_names[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_not_consolidated": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_swaplevel_sortlevel[0]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_swaplevel_sortlevel[baz]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_unused_levels": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_level_name[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_level_name[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_multiple_bug[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_multiple_bug[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_multiple[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_multiple[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_preserve_names[False]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_preserve_names[True]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_wrong_level_name[False-unstack]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_wrong_level_name[True-unstack]": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_unstack_preserve_types": "TODO: Add a reason for failure", - "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_unstack_with_missing_int_cast_to_float": "TODO: Add a reason for failure", + "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_bool": "cudf converts null bools to None where pandas' unstack upcasts to object with np.nan", + "tests/frame/test_stack_unstack.py::TestDataFrameReshape::test_unstack_not_consolidated": "Asserts DataFrame._mgr block layout (pandas internals)", + "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_unstack_with_missing_int_cast_to_float": "Asserts DataFrame._mgr block layout (pandas internals)", "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_asof": "TODO: Add a reason for failure", "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_equals_subclass": "TODO: Add a reason for failure", "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_frame_subclassing_and_slicing": "TODO: Add a reason for failure", @@ -2964,10 +2940,8 @@ def pytest_unconfigure(config): "tests/reshape/merge/test_multi.py::TestMergeMulti::test_left_join_multi_index[True-False]": "AssertionError: DataFrame.iloc[:, 4] (column name='5th') are different", "tests/reshape/merge/test_multi.py::TestMergeMulti::test_left_join_multi_index[True-True]": "TODO: Add a reason for failure", "tests/reshape/test_crosstab.py::TestCrosstab::test_crosstab_duplicate_names": "TODO: Add a reason for failure", - "tests/reshape/test_crosstab.py::TestCrosstab::test_crosstab_multiple": "TODO: Add a reason for failure", "tests/reshape/test_crosstab.py::TestCrosstab::test_crosstab_no_overlap": "TODO: Add a reason for failure", "tests/reshape/test_crosstab.py::TestCrosstab::test_crosstab_with_categorial_columns": "TODO: Add a reason for failure", - "tests/reshape/test_crosstab.py::TestCrosstab::test_crosstab_with_empties": "TODO: Add a reason for failure", "tests/reshape/test_cut.py::test_bins[array]": "TODO: Add a reason for failure", "tests/reshape/test_cut.py::test_bins[list]": "TODO: Add a reason for failure", "tests/reshape/test_cut.py::test_bins_from_interval_index": "TODO: Add a reason for failure", @@ -3013,12 +2987,10 @@ def pytest_unconfigure(config): "tests/reshape/test_melt.py::TestWideToLong::test_raise_of_column_name_value": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivot::test_pivot_index_is_none": "AssertionError: DataFrame.index are different", "tests/reshape/test_pivot.py::TestPivotTable::test_categorical_pivot_index_ordering[False]": "TODO: Add a reason for failure", - "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_complex_aggfunc": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_datetime_tz": "ValueError: Length of names must match number of levels in MultiIndex.", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_index_with_nan[False]": "AssertionError: DataFrame.index are different", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_index_with_nan[True]": "AssertionError: DataFrame.index are different", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_multi_functions": "TODO: Add a reason for failure", - "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_no_level_overlap": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_string_as_func": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_string_func_vs_func[f3-f_numpy3]": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_string_func_vs_func[f4-f_numpy4]": "TODO: Add a reason for failure", @@ -3029,7 +3001,6 @@ def pytest_unconfigure(config): "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_table_nocols": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_table_not_series": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_table_with_iterator_values": "TODO: Add a reason for failure", - "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_table_with_mixed_nested_tuples": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_with_categorical[False-False]": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_with_categorical[False-None]": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_with_categorical[False-True]": "TODO: Add a reason for failure", diff --git a/python/cudf/cudf/tests/dataframe/methods/test_swaplevel.py b/python/cudf/cudf/tests/dataframe/methods/test_swaplevel.py index c48e0603a1ed..b5e048f94952 100644 --- a/python/cudf/cudf/tests/dataframe/methods/test_swaplevel.py +++ b/python/cudf/cudf/tests/dataframe/methods/test_swaplevel.py @@ -1,7 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import pandas as pd import pytest import cudf @@ -37,6 +38,24 @@ def test_dataframe_swaplevel_axis_0(): assert_eq(cdf.swaplevel("a", "b"), cdf.swaplevel("b", "a")) +def test_dataframe_swaplevel_stacked_names(): + # the index produced by an unstack/stack round trip keys its levels + # positionally while carrying real level names; swaplevel must + # preserve the names, not re-derive them from the internal keys + pidx = pd.MultiIndex.from_product( + [[2000, 2001], [1, 2], [1, 15]], names=["year", "month", "day"] + ) + pdf = pd.DataFrame( + {"A": range(8), "B": range(8, 16)}, index=pidx, dtype="float64" + ) + cdf = cudf.DataFrame(pdf) + + expected = pdf.unstack(1).stack(future_stack=True).swaplevel(1, 2) + result = cdf.unstack(1).stack(future_stack=True).swaplevel(1, 2) + assert list(result.index.names) == ["year", "month", "day"] + assert_eq(expected.sort_index(), result.sort_index()) + + def test_dataframe_swaplevel_TypeError(): cdf = cudf.DataFrame( {"a": [1, 2, 3], "c": [10, 20, 30]}, index=["x", "y", "z"] diff --git a/python/cudf/cudf/tests/reshape/test_pivot_table.py b/python/cudf/cudf/tests/reshape/test_pivot_table.py index 368a8b3ffbed..347a73f84090 100644 --- a/python/cudf/cudf/tests/reshape/test_pivot_table.py +++ b/python/cudf/cudf/tests/reshape/test_pivot_table.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import numpy as np @@ -94,3 +94,32 @@ def test_pivot_table_scalar_index_columns(index, columns): values="D", index=index, columns=columns, aggfunc="sum" ) assert_eq(result, expected) + + +@pytest.mark.parametrize("aggfunc", ["sum", "min", "max"]) +@pytest.mark.parametrize("fill_value", [None, 0]) +def test_pivot_table_sparse_int_fill_value(aggfunc, fill_value): + # pandas promotes integer values to float64 when the reshape leaves + # missing cells unfilled (fill_value=None), and keeps the integer + # dtype when they are filled + data = { + "i": ["r1", "r1", "r2"], + "c": ["a", "b", "a"], + "v": [1, 2, 3], + } + with cudf.option_context("mode.pandas_compatible", True): + result = cudf.DataFrame(data).pivot_table( + index="i", + columns="c", + values="v", + aggfunc=aggfunc, + fill_value=fill_value, + ) + expected = pd.DataFrame(data).pivot_table( + index="i", + columns="c", + values="v", + aggfunc=aggfunc, + fill_value=fill_value, + ) + assert_eq(expected, result) diff --git a/python/cudf/cudf/tests/reshape/test_unstack.py b/python/cudf/cudf/tests/reshape/test_unstack.py index 0ca6930443f2..3a2462c4b49a 100644 --- a/python/cudf/cudf/tests/reshape/test_unstack.py +++ b/python/cudf/cudf/tests/reshape/test_unstack.py @@ -14,20 +14,10 @@ "level", [ 0, - pytest.param( - 1, - marks=pytest.mark.xfail( - reason="Categorical column indexes not supported" - ), - ), + 1, 2, "foo", - pytest.param( - "bar", - marks=pytest.mark.xfail( - reason="Categorical column indexes not supported" - ), - ), + "bar", "baz", [], pytest.param( From 9b00d888c77f2cfb6c2f67b810d1047e4d4a7f1c Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 23 Jul 2026 21:54:16 +0200 Subject: [PATCH 15/16] Add profitability guards for join filter pushdown (#22997) Teach the join filter pushdown planner to estimate the source-scan cost needed to build candidate domains. Candidate selection now prefers cheaper domain producers before smaller estimated output rows, and rejects pushdown when the domain-build cost is too large relative to the target being reduced. Also prevent a single-source domain from being stacked onto a target source that is already below the filtered side of a semi join. This heuristic does a better job of keeping profitable domain choices while avoiding cases where pushing down a filter would add an additional high-cardinality source onto a join target that is already significantly filtered. Additionally, extends trace metadata with estimated target, domain, and constraint costs. Material results of this change running NDSH SF30K on 8xNVL4 nodes are (previous results come from the change in https://github.com/rapidsai/cudf/pull/22996): * **Q9 improved runtime: 11.77s lukewarm, 7.45s hot** (previously 47.12s lukewarm, 32.68s hot) * Q5 unchaned: 9.55s lukewarm, 5.31s hot (previously 9.35s lukewarm, 5.14s hot) * **NDSH SF30K for all 22 queries: 118.87s lukewarm, 111.63s hot** Authors: - Peter Andreas Entschev (https://github.com/pentschev) - Lawrence Mitchell (https://github.com/wence-) Approvers: - Gil Forsyth (https://github.com/gforsyth) - Tom Augspurger (https://github.com/TomAugspurger) URL: https://github.com/rapidsai/cudf/pull/22997 --- ci/run_cudf_polars_polars_tests.sh | 3 +- ci/test_python_other.sh | 5 +- ci/test_wheel_cudf_polars.sh | 1 - .../cudf_polars/cudf_polars/dsl/traversal.py | 27 ++- .../streaming/join_filter_pushdown.py | 178 +++++++++++++---- .../streaming/test_join_filter_pushdown.py | 184 +++++++++++++++++- 6 files changed, 344 insertions(+), 54 deletions(-) diff --git a/ci/run_cudf_polars_polars_tests.sh b/ci/run_cudf_polars_polars_tests.sh index 7c3c138c3a60..1c6a542afbf4 100755 --- a/ci/run_cudf_polars_polars_tests.sh +++ b/ci/run_cudf_polars_polars_tests.sh @@ -62,7 +62,7 @@ DESELECTED_TESTS_STR=$(printf -- " --deselect %s" "${DESELECTED_TESTS[@]}") # Don't quote the `DESELECTED_...` variable because `pytest` can't handle # multiple quoted arguments inline # shellcheck disable=SC2086 -# Fail fast (-x) because failed tests pollute the state +# Fail fast (-x) rather than trying to continue because failed tests pollute the state echo "Run polars tests with injected in-memory GPU engine" python "${TIMEOUT_TOOL_PATH}" --enable-python 5400 \ python -m pytest \ @@ -90,7 +90,6 @@ python "${TIMEOUT_TOOL_PATH}" --enable-python 5400 \ --import-mode=importlib \ --cache-clear \ -x \ - -v \ -m "" \ -p cudf_polars.testing.inject_gpu_engine \ -W ignore::ResourceWarning \ diff --git a/ci/test_python_other.sh b/ci/test_python_other.sh index 0efebe1afafe..ecb83d66df22 100755 --- a/ci/test_python_other.sh +++ b/ci/test_python_other.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -42,8 +42,9 @@ timeout 30m ./ci/run_custreamz_pytests.sh \ --cov-report=term rapids-logger "pytest cudf-polars" +# Fail fast (-x) rather than trying to continue because failed tests pollute the state ./ci/run_cudf_polars_pytests.sh \ - -vv \ + -x \ --junitxml="${RAPIDS_TESTS_DIR}/junit-cudf-polars.xml" \ --numprocesses=4 \ --dist=worksteal \ diff --git a/ci/test_wheel_cudf_polars.sh b/ci/test_wheel_cudf_polars.sh index d9df9c05d4eb..1869ec8800a9 100755 --- a/ci/test_wheel_cudf_polars.sh +++ b/ci/test_wheel_cudf_polars.sh @@ -91,7 +91,6 @@ for version in "${VERSIONS[@]}"; do # Fail fast (-x) rather than trying to continue because failed tests pollute the state ./ci/run_cudf_polars_pytests.sh \ - -vv \ "${COVERAGE_ARGS[@]}" \ --numprocesses=4 \ --dist=worksteal \ diff --git a/python/cudf_polars/cudf_polars/dsl/traversal.py b/python/cudf_polars/cudf_polars/dsl/traversal.py index 095a9719597b..654b6fba17f4 100644 --- a/python/cudf_polars/cudf_polars/dsl/traversal.py +++ b/python/cudf_polars/cudf_polars/dsl/traversal.py @@ -1,11 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Traversal and visitor utilities for nodes.""" from __future__ import annotations -from collections import deque +from collections import Counter, deque from typing import TYPE_CHECKING, Generic from cudf_polars.typing import ( @@ -15,19 +15,40 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Generator, MutableMapping, Sequence + from collections.abc import Callable, Generator, Mapping, MutableMapping, Sequence from cudf_polars.typing import GenericTransformer, NodeT __all__: list[str] = [ "CachingVisitor", + "collect_refcount", "make_recursive", + "post_traversal", "reuse_if_unchanged", "traversal", ] +def collect_refcount(nodes: Sequence[NodeT]) -> Mapping[NodeT, int]: + """ + Determine reference counts of all nodes in a DAG. + + Parameters + ---------- + nodes + Sequence of root nodes + + Returns + ------- + Mapping from nodes to frequency of occurrence in the DAG. + """ + refcount = Counter(nodes) + for node in traversal(nodes): + refcount.update(node.children) + return refcount + + def traversal(nodes: Sequence[NodeT]) -> Generator[NodeT, None, None]: """ Pre-order traversal of nodes in an expression. diff --git a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py index 5ba74515cf1b..b099e1584ea6 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -38,6 +38,10 @@ A node on a column lineage, together with the column name at that node and its edge path from the join input. So termed because it "produces" the key values participating in the join. +``source cost`` + An estimate of the cost required to materialize a producer. This guards + against treating a small intermediate result as a cheap domain when + producing it requires scanning large inputs. ``constraint domain`` Selective values of another join key from the target input, used to reduce the domain before deriving the values that will filter the target. @@ -49,10 +53,10 @@ constraint domain, then projects the reduced domain's key used to filter the target. -Plan rewrite has three stages. ``analyze_plan`` gathers row estimates, -selective nodes, and column value-domain lineages. Candidate selection -consumes those facts and returns a decision. ``apply_candidate`` then -constructs the selected semi-join rewrite. +Plan rewrite has three stages. ``analyze_plan`` gathers row estimates, source +scan facts, selective nodes, and column value-domain lineages. +Candidate selection consumes those facts and returns a decision. +``apply_candidate`` then constructs the selected semi-join rewrite. Row estimates, selectivity propagation, thresholds, and candidate scores are only heuristics for deciding whether a safe rewrite is likely to improve @@ -86,6 +90,7 @@ from cudf_polars.dsl.tracing import Scope, log from cudf_polars.dsl.traversal import ( CachingVisitor, + collect_refcount, post_traversal, reuse_if_unchanged, traversal, @@ -104,6 +109,17 @@ from cudf_polars.utils.config import ConfigOptions, StreamingExecutor +DomainScore: TypeAlias = tuple[int, int, int] + + +@dataclass(frozen=True) +class SourceFacts: + """Source-derived facts for an IR node.""" + + cost: int | None + is_single_source: bool + + @dataclass(frozen=True) class _Producer: """A subtree and its bound column names at an insertion point.""" @@ -111,6 +127,8 @@ class _Producer: node: IR columns: tuple[str, ...] rows: int + cost: int + is_single_source: bool path: tuple[int, ...] = () """Child-edge path from the candidate root to ``node``.""" @@ -119,6 +137,11 @@ def column(self) -> str: """First bound column in the producer.""" return self.columns[0] + @property + def domain_score(self) -> DomainScore: + """Scoring function for a domain.""" + return (self.cost, self.rows, len(self.node.schema)) + @dataclass(frozen=True) class SimpleCandidate: @@ -132,9 +155,9 @@ class SimpleCandidate: domain_key: expr.Col @property - def score(self) -> tuple[int, int, int]: - """Rank after composite candidates, then by domain size.""" - return (1, self.domain.rows, self.domain.rows) + def score(self) -> tuple[int, DomainScore]: + """Rank after composite candidates, then by domain cost.""" + return (1, self.domain.domain_score) @dataclass(frozen=True) @@ -152,16 +175,16 @@ class CompositeCandidate: target_constraint_key: expr.Col @property - def score(self) -> tuple[int, int, int]: - """Prefer smaller constraint and domain inputs.""" - return (0, self.constraint_domain.rows, self.domain.rows) + def score(self) -> tuple[int, DomainScore, DomainScore]: + """Prefer cheaper constraint and domain inputs.""" + return (0, self.constraint_domain.domain_score, self.domain.domain_score) Candidate: TypeAlias = SimpleCandidate | CompositeCandidate DecisionReason: TypeAlias = Literal[ "applied", "maintain_order", - "no_selective_domain", + "no_profitable_domain", "non_column_join_key", "not_inner_join", "sliced_join", @@ -181,8 +204,10 @@ class PlanFacts: """Facts derived in one bottom-up traversal of an IR DAG.""" row_estimates: Mapping[IR, int | None] + source_facts: Mapping[IR, SourceFacts] selective_nodes: frozenset[IR] column_lineages: Mapping[ColumnRef, ColumnLineage] + refcounts: Mapping[IR, int] class _RewriteState(TypedDict): @@ -210,8 +235,11 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: Gather facts about the plan. """ row_estimates: dict[IR, int | None] = {} + source_facts: dict[IR, SourceFacts] = {} + source_nodes: dict[IR, frozenset[IR]] = {} selective_nodes: set[IR] = set() column_lineages: dict[ColumnRef, ColumnLineage] = {} + refcounts = collect_refcount([ir]) for node in post_traversal([ir]): if isinstance(node, (Scan, DataFrameScan)): @@ -236,6 +264,23 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: rows = max(child_estimates, default=None) row_estimates[node] = rows + if isinstance(node, (Scan, DataFrameScan)): + sources: frozenset[IR] = frozenset((node,)) + else: + sources = frozenset( + source for child in node.children for source in source_nodes[child] + ) + source_nodes[node] = sources + source_rows = [ + source_rows + for source in sources + if (source_rows := row_estimates[source]) is not None and source_rows > 0 + ] + source_facts[node] = SourceFacts( + cost=sum(source_rows) if source_rows else rows, + is_single_source=len(sources) == 1, + ) + if ( (isinstance(node, Scan) and node.predicate is not None) or isinstance(node, Filter) @@ -263,26 +308,33 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: return PlanFacts( row_estimates=row_estimates, + source_facts=source_facts, selective_nodes=frozenset(selective_nodes), column_lineages=column_lineages, + refcounts=refcounts, ) -def blocks_pushdown(node: IR) -> bool: +def blocks_pushdown(node: IR, facts: PlanFacts) -> bool: """ Return whether a node blocks filter pushdown. Parameters ---------- node - Node to check + Node to check. + facts + Facts about the plan. Returns ------- bool True if a semijoin cannot be pushed past this node, otherwise False. """ - return ( + # TODO: Need better cost model to handle nodes that are shared. Pushing + # a filter into a shared node will typically mean that it is no longer + # shared, since the same filter will not come from every consumer. + return facts.refcounts[node] > 1 or ( # TODO: Distinct and Rolling only block pushdown in some # circumstances, but we'd need to make the logic more complicated: # - We can push through distinct if the filter applies to the columns @@ -327,7 +379,7 @@ def semijoin_pushdown_candidates( yield lineage.column, path source = lineage.source source_child_index = lineage.source_child_index - if blocks_pushdown(lineage.column.node) or source is None: + if blocks_pushdown(lineage.column.node, facts) or source is None: return assert source_child_index is not None path = (*path, source_child_index) @@ -526,7 +578,7 @@ def _select_candidate( ) if not candidates: - return Decision(reason="no_selective_domain") + return Decision(reason="no_profitable_domain") return Decision(reason="applied", candidate=min(candidates, key=lambda c: c.score)) @@ -559,6 +611,12 @@ def _simple_candidates( continue if contains_node(target.node, domain.node): continue + if domain.is_single_source and has_filtering_semi_ancestor( + target_child, target.path + ): + continue + if not domain_cost_is_small(domain, target, threshold): + continue yield SimpleCandidate( target_side=target_side, target=target, @@ -617,6 +675,10 @@ def _composite_candidates( target.node, constraint_domain.node ): continue + if not domain_cost_is_small(domain, target, threshold): + continue + if not domain_cost_is_small(constraint_domain, domain, threshold): + continue yield CompositeCandidate( target_side=target_side, target=target, @@ -692,6 +754,27 @@ def _make_semi_join( ) +def make_producer( + node: IR, + columns: tuple[str, ...], + path: tuple[int, ...], + facts: PlanFacts, +) -> _Producer | None: + """Construct a producer from gathered plan facts, if fully estimated.""" + rows = facts.row_estimates.get(node) + source = facts.source_facts[node] + if rows is None or rows <= 0 or source.cost is None: + return None + return _Producer( + node=node, + columns=columns, + rows=rows, + cost=source.cost, + is_single_source=source.is_single_source, + path=path, + ) + + def _smallest_key_producer( root: IR, column: str, @@ -700,28 +783,25 @@ def _smallest_key_producer( require_selective: bool, exclude: IR | None = None, ) -> _Producer | None: - candidates = [] + producers = [] for reference, path in semijoin_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name if node is exclude: continue - rows = facts.row_estimates.get(node) - if rows is None or rows <= 0: - continue if require_selective and node not in facts.selective_nodes: continue - candidates.append( - (rows, len(node.schema), _Producer(node, (bound_column,), rows, path)) - ) - if not candidates: + producer = make_producer(node, (bound_column,), path, facts) + if producer is not None: + producers.append(producer) + if not producers: return None - return min(candidates, key=lambda item: (item[0], item[1]))[2] + return min(producers, key=lambda p: p.domain_score) def _smallest_node_containing_all( root: IR, columns: Sequence[str], facts: PlanFacts ) -> _Producer | None: - candidates = [] + producers = [] lineages: list[ColumnLineage] = [] for column in columns: lineage = facts.column_lineages.get(ColumnRef(root, column)) @@ -736,16 +816,10 @@ def _smallest_node_containing_all( if any(lineage.column.node != node for lineage in lineages[1:]): break bound_columns = tuple(lineage.column.name for lineage in lineages) - rows = facts.row_estimates.get(node) - if rows is not None and rows > 0: - candidates.append( - ( - rows, - len(node.schema), - _Producer(node, bound_columns, rows, path), - ) - ) - if blocks_pushdown(node): + producer = make_producer(node, bound_columns, path, facts) + if producer is not None: + producers.append(producer) + if blocks_pushdown(node, facts): break source_child_index = lineages[0].source_child_index if source_child_index is None or any( @@ -758,9 +832,9 @@ def _smallest_node_containing_all( break path = (*path, source_child_index) lineages = sources - if not candidates: + if not producers: return None - return min(candidates, key=lambda item: (item[0], item[1]))[2] + return min(producers, key=lambda p: p.domain_score) def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | None: @@ -768,13 +842,13 @@ def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | fallback_candidates = [] for reference, path in semijoin_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name - rows = facts.row_estimates.get(node) - if rows is None or rows <= 0: + producer = make_producer(node, (bound_column,), path, facts) + if producer is None: continue item = ( - rows, + producer.rows, len(node.schema), - _Producer(node, (bound_column,), rows, path), + producer, ) if isinstance(node, (Scan, DataFrameScan)): source_candidates.append(item) @@ -809,6 +883,23 @@ def contains_node(root: IR, needle: IR) -> bool: return needle in traversal([root]) +def domain_cost_is_small( + domain: _Producer, target: _Producer, threshold: float +) -> bool: + """Return whether building a domain is cheap enough for its target.""" + return domain.cost / target.rows <= threshold + + +def has_filtering_semi_ancestor(root: IR, path: Sequence[int]) -> bool: + """Return whether a selected child edge is below a filtering semi join.""" + node = root + for child_index in path: + if isinstance(node, Join) and node.options[0] == "Semi" and child_index == 0: + return True + node = node.children[child_index] + return False + + def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: join_filter_pushdown: dict[str, Any] = { "considered": True, @@ -830,6 +921,8 @@ def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: "domain_key": candidate.domain_key.name, "estimated_target_rows": candidate.target.rows, "estimated_domain_rows": candidate.domain.rows, + "estimated_target_cost": candidate.target.cost, + "estimated_domain_cost": candidate.domain.cost, "target_node_type": type(candidate.target.node).__name__, "domain_node_type": type(candidate.domain.node).__name__, } @@ -839,6 +932,7 @@ def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: { "constraint_key": candidate.target_constraint_key.name, "estimated_constraint_rows": candidate.constraint_domain.rows, + "estimated_constraint_cost": candidate.constraint_domain.cost, } ) log("Join Filter Pushdown", **record) diff --git a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py index 3dfe42df7fd3..7601788400db 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -10,6 +10,7 @@ import polars as pl from cudf_polars import Translator +from cudf_polars.dsl.expr import Col from cudf_polars.dsl.ir import Cache, DataFrameScan, Distinct, Join, Select, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.column_domain import ColumnRef @@ -94,6 +95,13 @@ def dataframe_scan(ir: IR, column: str) -> DataFrameScan: return match +def join_key_names(join: Join) -> tuple[str, ...]: + """Return the column names used on the left of a simple-column join.""" + names = tuple(key.value.name for key in join.left_on if isinstance(key.value, Col)) + assert len(names) == len(join.left_on) + return names + + @pytest.fixture def simple_query() -> pl.LazyFrame: """Return a query with a small selective join domain.""" @@ -257,7 +265,7 @@ def test_no_simple_filter_pushdown_when_domain_is_not_selective( ConfigOptions.from_polars_engine(engine), ) - assert decision == Decision(reason="no_selective_domain") + assert decision == Decision(reason="no_profitable_domain") assert optimized is root assert not find_joins(optimized, "Semi") assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -330,6 +338,120 @@ def test_composite_filter_pushdown_constrains_domain_first( assert_gpu_result_equal(query, engine=engine, check_row_order=False) +def test_prefilter_uses_cheaper_source_domain_and_skips_expensive_domain( + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": range(60), + "part_active": [True] * 30 + [False] * 30, + } + ) + .filter("part_active") + .select("p_partkey") + ) + partsupp = pl.LazyFrame( + { + "ps_partkey": [i % 60 for i in range(120)], + "ps_suppkey": [i % 30 for i in range(120)], + } + ) + supplier = pl.LazyFrame({"s_suppkey": range(30)}) + lineitem = pl.LazyFrame( + { + "l_partkey": [i % 60 for i in range(1_800)], + "l_suppkey": [i % 30 for i in range(1_800)], + "l_orderkey": [i % 900 for i in range(1_800)], + } + ) + orders = pl.LazyFrame({"o_orderkey": range(900)}) + query = ( + part.join(partsupp, left_on="p_partkey", right_on="ps_partkey") + .join(supplier, left_on="ps_suppkey", right_on="s_suppkey") + .join( + lineitem, + left_on=("p_partkey", "ps_suppkey"), + right_on=("l_partkey", "l_suppkey"), + ) + .join(orders, left_on="l_orderkey", right_on="o_orderkey") + ) + root = translate_query(query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + part_ir = dataframe_scan(root, "p_partkey") + supplier_ir = dataframe_scan(root, "s_suppkey") + lineitem_ir = dataframe_scan(root, "l_orderkey") + orders_ir = dataframe_scan(root, "o_orderkey") + semis = find_joins(optimized, "Semi") + partkey_semis = [ + semi + for semi in semis + if semi.children[0] is lineitem_ir and join_key_names(semi) == ("l_partkey",) + ] + assert partkey_semis + assert not any(semi.children[0] is orders_ir for semi in semis) + assert contains_node(partkey_semis[0].children[1], part_ir) + assert not contains_node(partkey_semis[0].children[1], supplier_ir) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_source_only_domain_does_not_stack_on_prefiltered_source( + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": range(60), + "part_active": [True] * 30 + [False] * 30, + } + ) + .filter("part_active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame( + { + "l_partkey": [i % 60 for i in range(1_800)], + "l_orderkey": [i % 150 for i in range(1_800)], + } + ) + orders = ( + pl.LazyFrame( + { + "o_orderkey": range(150), + "order_active": [True] * 75 + [False] * 75, + } + ) + .filter("order_active") + .select("o_orderkey") + ) + query = part.join(lineitem, left_on="p_partkey", right_on="l_partkey").join( + orders, left_on="l_orderkey", right_on="o_orderkey" + ) + root = translate_query(query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + lineitem_ir = dataframe_scan(root, "l_partkey") + lineitem_semis = [ + semi + for semi in find_joins(optimized, "Semi") + if semi.children[0] is lineitem_ir + ] + assert any(join_key_names(semi) == ("l_partkey",) for semi in lineitem_semis) + assert not any(join_key_names(semi) == ("l_orderkey",) for semi in lineitem_semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + def test_derived_selectivity_propagates_through_rewritten_children( engine: SPMDEngine, ) -> None: @@ -432,7 +554,7 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking( orders_ir = dataframe_scan(root, "o_orderkey") semis = find_joins(optimized, "Semi") assert sum(semi.children[0] is lineitem_ir for semi in semis) == 1 - assert any(semi.children[0] is orders_ir for semi in semis) + assert not any(semi.children[0] is orders_ir for semi in semis) assert not any( isinstance(semi.children[0], Join) and semi.children[0].options[0] == "Semi" for semi in semis @@ -568,8 +690,10 @@ def test_composite_domain_columns_follow_renames(engine: SPMDEngine) -> None: analyzed = analyze_plan(renamed, StatsCollector()) facts = PlanFacts( row_estimates={renamed: 20, source: 10}, + source_facts=analyzed.source_facts, selective_nodes=analyzed.selective_nodes, column_lineages=analyzed.column_lineages, + refcounts=analyzed.refcounts, ) producer = _smallest_node_containing_all( renamed, ("domain_key", "domain_constraint"), facts @@ -739,7 +863,7 @@ def test_target_replacement_does_not_rewrite_shared_domain_side( assert domain_ir.children[0] is shared_ir semis = find_joins(filtered, "Semi") assert len(semis) == 1 - assert semis[0].children[0] is dataframe_scan(root, "target_key") + assert semis[0].children[0] is shared_ir assert not find_joins(unfiltered_domain, "Semi") assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -792,7 +916,59 @@ def test_target_prefilter_rewrites_only_selected_self_join_edge( filtered_semis = find_joins(filtered, "Semi") assert len(filtered_semis) == 1 assert not find_joins(unfiltered, "Semi") - assert any(filtered_semis[0].children[0] is node for node in traversal([source_ir])) + # The shared node is a valid insertion point, but its children are not: + # Only this consumer should be wrapped by the semi-join. + assert filtered_semis[0].children[0] is source_ir + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_internal_prefilter_rewrites_shared_subplan_once( + engine: SPMDEngine, +) -> None: + domain = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key") + ) + target = pl.LazyFrame( + { + "target_key": [i % 10 for i in range(20)], + "value": range(20), + } + ) + shared = domain.join( + target, + left_on="domain_key", + right_on="target_key", + ) + query = shared.join(shared, on="domain_key", suffix="_right") + translated = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(translated, Join) + shared_cache = translated.children[0] + assert isinstance(shared_cache, Cache) + assert translated.children[1] is shared_cache + original_shared = shared_cache.children[0] + assert len(find_joins(original_shared, "Inner")) == 1 + target_ir = dataframe_scan(original_shared, "target_key") + + optimized = optimize_with_stats( + translated, + ConfigOptions.from_polars_engine(engine), + StatsCollector(), + ) + + assert isinstance(optimized, Join) + rewritten_left, rewritten_right = optimized.children + assert rewritten_left is rewritten_right + assert rewritten_left is not original_shared + (internal_semi,) = find_joins(rewritten_left, "Semi") + assert internal_semi.children[0] is target_ir assert_gpu_result_equal(query, engine=engine, check_row_order=False) From 5beaa5954688fcb12236ffb434e192ea2c77db30 Mon Sep 17 00:00:00 2001 From: Jihoon Son Date: Mon, 27 Jul 2026 06:31:33 -0700 Subject: [PATCH 16/16] JNI bridge for the retention filter support for deletion vectors (#23423) First of all, I apologize for making a PR targeting the 26.08 branch. I understand the code freeze has taken effect since yesterday, but wish this change can still be shipped in the 26.08 release. ## Description This PR adds the JNI bridge for the recent change in https://github.com/rapidsai/cudf/pull/23402 to support the retention filter for the deletion vector. Previously, it was assumed that deletion vectors represent only the rows deleted. With https://github.com/rapidsai/cudf/pull/23402 and this change, now Java applications can pass the correct filter type for deletion vectors. ## Checklist - [x] I am familiar with the [Contributing Guidelines](https://github.com/rapidsai/cudf/blob/HEAD/CONTRIBUTING.md). - [x] New or existing tests cover these changes. - [x] The documentation is up to date with these changes. --- .../java/ai/rapids/cudf/DeletionVector.java | 30 ++++++++-- .../src/main/native/src/DeletionVectorJni.cpp | 28 ++++++--- .../rapids/cudf/DeletionVectorTableTest.java | 58 ++++++++++++++----- 3 files changed, 89 insertions(+), 27 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/DeletionVector.java b/java/src/main/java/ai/rapids/cudf/DeletionVector.java index 16eeee386e5e..92f82b0b561d 100644 --- a/java/src/main/java/ai/rapids/cudf/DeletionVector.java +++ b/java/src/main/java/ai/rapids/cudf/DeletionVector.java @@ -1,6 +1,6 @@ /* * - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * */ @@ -44,6 +44,8 @@ public static class DeletionVectorInfo { */ public final HostMemoryBuffer serializedBitmap; + public final boolean isRetention; + /** * Row index offsets for each row group to read. Can be null to read all row groups. */ @@ -59,8 +61,9 @@ public static class DeletionVectorInfo { */ public final int totalNumRows; - public DeletionVectorInfo(HostMemoryBuffer serializedBitmap, long[] rowGroupOffsets, int[] rowGroupNumRows) { + public DeletionVectorInfo(HostMemoryBuffer serializedBitmap, boolean isRetention, long[] rowGroupOffsets, int[] rowGroupNumRows) { this.serializedBitmap = serializedBitmap; + this.isRetention = isRetention; this.rowGroupOffsets = rowGroupOffsets; this.rowGroupNumRows = rowGroupNumRows; this.totalNumRows = computeTotalNumRows(); @@ -81,6 +84,16 @@ private int computeTotalNumRows() { } } + private static boolean getDeletionVectorTypes(DeletionVectorInfo[] deletionVectorInfos) { + boolean isRetention = deletionVectorInfos[0].isRetention; + for (DeletionVectorInfo info : deletionVectorInfos) { + if (info.isRetention != isRetention) { + throw new IllegalArgumentException("All DeletionVectorInfo objects must have the same isRetention value."); + } + } + return isRetention; + } + /** * Reads a Parquet file with deletion vector support. * @@ -174,6 +187,7 @@ private static Table readParquet(ParquetOptions opts, List deletionVectorRowCountsList = new ArrayList<>(deletionVectorInfos.length); List rowGroupOffsetsList = new ArrayList<>(deletionVectorInfos.length); List rowGroupNumRowsList = new ArrayList<>(deletionVectorInfos.length); + boolean areRetentionVectors = getDeletionVectorTypes(deletionVectorInfos); if (deletionVectorInfos != null) { for (DeletionVectorInfo info : deletionVectorInfos) { serializedBitmapList.add(info.serializedBitmap); @@ -202,7 +216,8 @@ private static Table readParquet(ParquetOptions opts, bitmapAddrsSizes, deletionVectorRowCounts, rowGroupOffsets, - rowGroupNumRows); + rowGroupNumRows, + areRetentionVectors); return new Table(columnHandles); } @@ -325,6 +340,7 @@ private ParquetChunkedReader(long chunkSizeByteLimit, long passReadLimit, List deletionVectorRowCountsList = new ArrayList<>(deletionVectorInfos.length); List rowGroupOffsetsList = new ArrayList<>(deletionVectorInfos.length); List rowGroupNumRowsList = new ArrayList<>(deletionVectorInfos.length); + boolean areRetentionVectors = getDeletionVectorTypes(deletionVectorInfos); if (deletionVectorInfos != null) { for (DeletionVectorInfo info : deletionVectorInfos) { serializedBitmapList.add(info.serializedBitmap); @@ -347,7 +363,7 @@ private ParquetChunkedReader(long chunkSizeByteLimit, long passReadLimit, long[] handles = createParquetChunkedReader(chunkSizeByteLimit, passReadLimit, opts.getIncludeColumnNames(), opts.getReadBinaryAsString(), inputFilePaths, dataBufferAddrsSizes, rowGroups, opts.timeUnit().typeId.getNativeId(), - bitmapAddrsSizes, deletionVectorRowCounts, rowGroupOffsets, rowGroupNumRows); + bitmapAddrsSizes, deletionVectorRowCounts, rowGroupOffsets, rowGroupNumRows, areRetentionVectors); readerHandle = handles[0]; if (readerHandle == 0) { throw new IllegalStateException("Cannot create native chunked Parquet reader object."); @@ -431,7 +447,8 @@ private static native long[] readParquet(String[] filterColumnNames, long[] serializedRoaring64, int[] deletionVectorRowCounts, long[] rowGroupOffsets, - int[] rowGroupNumRows) + int[] rowGroupNumRows, + boolean areRetentionVectors) throws CudfException; private static native long[] createParquetChunkedReader(long chunkReadLimit, @@ -445,7 +462,8 @@ private static native long[] createParquetChunkedReader(long chunkReadLimit, long[] serializedRoaringBitmaps, int[] deletionVectorRowCounts, long[] rowGroupOffsets, - int[] rowGroupNumRows) + int[] rowGroupNumRows, + boolean areRetentionVectors) throws CudfException; private static native boolean parquetChunkedReaderHasNext(long readerHandle) throws CudfException; diff --git a/java/src/main/native/src/DeletionVectorJni.cpp b/java/src/main/native/src/DeletionVectorJni.cpp index 4da3f0a9c130..9b7a41910da9 100644 --- a/java/src/main/native/src/DeletionVectorJni.cpp +++ b/java/src/main/native/src/DeletionVectorJni.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -82,7 +82,8 @@ std::unique_ptr make_dele jlongArray const& serialized_roaring64, jintArray const& deletion_vector_row_counts, jlongArray const& row_group_offsets, - jintArray const& row_group_num_rows) + jintArray const& row_group_num_rows, + jboolean are_retention_vectors) { cudf::jni::native_jlongArray n_serialized_roaring64(env, serialized_roaring64); std::vector> serialized_bitmaps = @@ -97,6 +98,7 @@ std::unique_ptr make_dele dv_info->deletion_vector_row_counts = n_deletion_vector_row_counts.to_vector(); dv_info->row_group_num_rows = n_row_group_num_rows.to_vector(); dv_info->row_group_offsets.reserve(n_row_group_offsets.size()); + dv_info->are_retention_vectors = are_retention_vectors; std::transform(n_row_group_offsets.begin(), n_row_group_offsets.end(), std::back_inserter(dv_info->row_group_offsets), @@ -141,7 +143,8 @@ Java_ai_rapids_cudf_DeletionVector_readParquet(JNIEnv* env, jlongArray serialized_roaring64, jintArray deletion_vector_row_counts, jlongArray row_group_offsets, - jintArray row_group_num_rows) + jintArray row_group_num_rows, + jboolean are_retention_vectors) { bool read_buffer = true; if (addrs_and_sizes == nullptr) { @@ -175,8 +178,12 @@ Java_ai_rapids_cudf_DeletionVector_readParquet(JNIEnv* env, cudf::io::parquet_reader_options opts = make_parquet_reader_options( env, filter_col_names, col_binary_read, row_groups, std::move(source), unit); - auto dv_info = make_deletion_vector_info( - env, serialized_roaring64, deletion_vector_row_counts, row_group_offsets, row_group_num_rows); + auto dv_info = make_deletion_vector_info(env, + serialized_roaring64, + deletion_vector_row_counts, + row_group_offsets, + row_group_num_rows, + are_retention_vectors); auto tbl = cudf::io::parquet::experimental::read_parquet(opts, *dv_info).tbl; return cudf::jni::convert_table_for_return(env, tbl); @@ -231,7 +238,8 @@ Java_ai_rapids_cudf_DeletionVector_createParquetChunkedReader(JNIEnv* env, jlongArray serialized_roaring64, jintArray deletion_vector_row_counts, jlongArray row_group_offsets, - jintArray row_group_num_rows) + jintArray row_group_num_rows, + jboolean are_retention_vectors) { bool read_buffer = true; if (addrs_sizes == nullptr) { @@ -265,8 +273,12 @@ Java_ai_rapids_cudf_DeletionVector_createParquetChunkedReader(JNIEnv* env, cudf::io::parquet_reader_options opts = make_parquet_reader_options( env, filter_col_names, col_binary_read, row_groups, std::move(source), unit); - auto dv_info = make_deletion_vector_info( - env, serialized_roaring64, deletion_vector_row_counts, row_group_offsets, row_group_num_rows); + auto dv_info = make_deletion_vector_info(env, + serialized_roaring64, + deletion_vector_row_counts, + row_group_offsets, + row_group_num_rows, + are_retention_vectors); // Create the chunked reader with pass read limit and multiple deletion vectors auto reader = new cudf::io::parquet::experimental::chunked_parquet_reader( diff --git a/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java b/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java index 4016d1029eb0..29400f5f3346 100644 --- a/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java +++ b/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java @@ -1,6 +1,6 @@ /* * - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * */ @@ -11,6 +11,8 @@ import ai.rapids.cudf.DeletionVector.DeletionVectorInfo; import ai.rapids.cudf.DeletionVector.ParquetChunkedReader; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import java.io.File; import java.io.IOException; @@ -18,6 +20,7 @@ import static ai.rapids.cudf.AssertUtils.assertTableTypes; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; class DeletionVectorTableTest extends CudfTestBase { @@ -29,8 +32,9 @@ class DeletionVectorTableTest extends CudfTestBase { private static final int DELETED_ROWS_COUNT2 = 3959; private static final int DELETED_ROWS_COUNT2_RGS_1_AND_3 = 1974; - @Test - void testReadParquetReadAllRowGroups() throws IOException { + @ParameterizedTest(name = "isRetention={0}") + @CsvSource({"false", "true"}) + void testReadParquetReadAllRowGroups(boolean isRetention) throws IOException { ParquetOptions opts = ParquetOptions.builder() .includeColumn("loan_id") .includeColumn("zip") @@ -40,10 +44,12 @@ void testReadParquetReadAllRowGroups() throws IOException { byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE1); try (HostMemoryBufferArray array = TableTestUtils.buffersFrom(data); HostMemoryBufferArray bitmapArray = TableTestUtils.buffersFrom(new byte[][] { bitmapData })) { - DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], null, null); + DeletionVectorInfo dvInfo = new DeletionVectorInfo( + bitmapArray.buffers[0], isRetention, null, null); try (Table table = DeletionVector.readParquet(opts, array.buffers, new DeletionVectorInfo[] { dvInfo })) { - long rows = table.getRowCount(); - assertEquals(1000 - DELETED_ROWS_COUNT1, rows); + long expectedRows = + isRetention ? DELETED_ROWS_COUNT1 : 1000 - DELETED_ROWS_COUNT1; + assertEquals(expectedRows, table.getRowCount()); assertTableTypes(new DType[]{DType.UINT64, DType.INT64, DType.INT32, DType.INT32}, table); } } @@ -59,7 +65,7 @@ void testReadParquetReadSomeRowGroups() throws IOException { HostMemoryBufferArray bitmapArray = TableTestUtils.buffersFrom(new byte[][] { bitmapData })) { long[] rowGroupOffsets = Arrays.stream(rowGroups[0]).mapToLong(i -> i * 10000L).toArray(); int[] rowGroupNumRows = Arrays.stream(rowGroups[0]).map(i -> 10000).toArray(); - DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], rowGroupOffsets, rowGroupNumRows); + DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], false, rowGroupOffsets, rowGroupNumRows); try (Table table = DeletionVector.readParquet(opts, array.buffers, rowGroups, new DeletionVectorInfo[] { dvInfo })) { long rows = table.getRowCount(); assertEquals(20000 - DELETED_ROWS_COUNT2_RGS_1_AND_3, rows); @@ -76,7 +82,7 @@ void testChunkedReadParquetAllRowGroups() throws Exception { try (HostMemoryBufferArray array = TableTestUtils.buffersFrom(data); HostMemoryBufferArray bitmapArray = TableTestUtils.buffersFrom(new byte[][] { bitmapData })) { ParquetOptions opts = ParquetOptions.DEFAULT; - DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], null, null); + DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], false, null, null); try (ParquetChunkedReader reader = DeletionVector.newParquetChunkedReader(240000, 0, opts, array.buffers, new DeletionVectorInfo[] { dvInfo })) { int numChunks = 0; @@ -94,8 +100,9 @@ void testChunkedReadParquetAllRowGroups() throws Exception { } } - @Test - void testChunkedReadParquetSomeRowGroups() throws Exception { + @ParameterizedTest(name = "isRetention={0}") + @CsvSource({"false", "true"}) + void testChunkedReadParquetSomeRowGroups(boolean isRetention) throws Exception { byte[][] data = TableTestUtils.sliceBytes(TableTestUtils.arrayFrom(TEST_FILE2), 2); byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE2); int[][] rowGroups = new int[][] { {1, 3} }; @@ -104,7 +111,8 @@ void testChunkedReadParquetSomeRowGroups() throws Exception { ParquetOptions opts = ParquetOptions.DEFAULT; long[] rowGroupOffsets = Arrays.stream(rowGroups[0]).mapToLong(i -> i * 10000L).toArray(); int[] rowGroupNumRows = Arrays.stream(rowGroups[0]).map(i -> 10000).toArray(); - DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], rowGroupOffsets, rowGroupNumRows); + DeletionVectorInfo dvInfo = new DeletionVectorInfo( + bitmapArray.buffers[0], isRetention, rowGroupOffsets, rowGroupNumRows); try (ParquetChunkedReader reader = DeletionVector.newParquetChunkedReader(120000, 0, opts, array.buffers, rowGroups, new DeletionVectorInfo[] { dvInfo })) { int numChunks = 0; @@ -117,7 +125,10 @@ void testChunkedReadParquetSomeRowGroups() throws Exception { } } assertEquals(2, numChunks); - assertEquals(20000 - DELETED_ROWS_COUNT2_RGS_1_AND_3, totalRows); + long expectedRows = isRetention + ? DELETED_ROWS_COUNT2_RGS_1_AND_3 + : 20000 - DELETED_ROWS_COUNT2_RGS_1_AND_3; + assertEquals(expectedRows, totalRows); } } } @@ -132,7 +143,7 @@ void testChunkedReadParquetMultiFiles() throws Exception { ParquetOptions opts = ParquetOptions.DEFAULT; long[] rowGroupOffsets = new long[] { 30000L, 10000L }; int[] rowGroupNumRows = new int[] { 10000, 10000 }; - DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], rowGroupOffsets, rowGroupNumRows); + DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], false, rowGroupOffsets, rowGroupNumRows); try (ParquetChunkedReader reader = DeletionVector.newParquetChunkedReader(120000, 0, opts, new String[] { TEST_FILE2.getAbsolutePath(), TEST_FILE2.getAbsolutePath() @@ -152,4 +163,25 @@ void testChunkedReadParquetMultiFiles() throws Exception { } } } + + @Test + void testMixedDeletionAndRetentionVectorsRejected() throws IOException { + byte[][] data = TableTestUtils.sliceBytes(TableTestUtils.arrayFrom(TEST_FILE1), 10); + byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE1); + try (HostMemoryBufferArray array = TableTestUtils.buffersFrom(data); + HostMemoryBufferArray bitmapArray = + TableTestUtils.buffersFrom(new byte[][] { bitmapData })) { + DeletionVectorInfo[] mixedVectorInfos = new DeletionVectorInfo[] { + new DeletionVectorInfo(bitmapArray.buffers[0], false, null, null), + new DeletionVectorInfo(bitmapArray.buffers[0], true, null, null) + }; + + assertThrows(IllegalArgumentException.class, + () -> DeletionVector.readParquet( + ParquetOptions.DEFAULT, array.buffers, mixedVectorInfos)); + assertThrows(IllegalArgumentException.class, + () -> DeletionVector.newParquetChunkedReader( + 0, 0, ParquetOptions.DEFAULT, array.buffers, mixedVectorInfos)); + } + } }