diff --git a/ci/run_cudf_polars_polars_tests.sh b/ci/run_cudf_polars_polars_tests.sh index 3778e1a9035a..1c6a542afbf4 100755 --- a/ci/run_cudf_polars_polars_tests.sh +++ b/ci/run_cudf_polars_polars_tests.sh @@ -62,9 +62,9 @@ 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 3600 \ +python "${TIMEOUT_TOOL_PATH}" --enable-python 5400 \ python -m pytest \ --import-mode=importlib \ --cache-clear \ @@ -85,12 +85,11 @@ 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 5400 \ python -m pytest \ --import-mode=importlib \ --cache-clear \ -x \ - -v \ -m "" \ -p cudf_polars.testing.inject_gpu_engine \ -W ignore::ResourceWarning \ diff --git a/ci/run_cudf_polars_pytests.sh b/ci/run_cudf_polars_pytests.sh index 3fac6910c5a8..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 3600 \ +python "${TIMEOUT_TOOL_PATH}" --enable-python 5400 \ python -m pytest --cache-clear "$@" tests 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/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 15ceb7fff93f..8d88feb7e18d 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1417,7 +1417,7 @@ if(CUDF_BUILD_TESTUTIL) ) target_link_libraries( - cudftestutil INTERFACE $ Threads::Threads cudf + cudftestutil INTERFACE $ Threads::Threads cudf rmm::rmm cudftest_default_stream $ ) 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/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/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/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/src/io/parquet/bloom_filter_reader.cu b/cpp/src/io/parquet/bloom_filter_reader.cu index eb7dfdc90b6b..3418c8ca9ada 100644 --- a/cpp/src/io/parquet/bloom_filter_reader.cu +++ b/cpp/src/io/parquet/bloom_filter_reader.cu @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -25,7 +26,6 @@ #include #include -#include #include #include #include @@ -53,16 +53,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/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/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/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); } }; }; 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/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; 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/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( diff --git a/dependencies.yaml b/dependencies.yaml index 5eb907f5649a..6caa4932c2e2 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: 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}" 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)); + } + } } 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)); } } 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 939e125fa8d8..92a43a6adedc 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, @@ -615,7 +618,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 +845,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 +858,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 +1290,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 +1363,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 +2571,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 +3259,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 +3269,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 +3306,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: @@ -8096,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] = [] @@ -8110,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 @@ -8121,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` @@ -8171,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. @@ -8261,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( @@ -8285,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..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 @@ -2109,14 +2106,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/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 3e219a23e8ec..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", @@ -1391,9 +1390,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", @@ -1675,64 +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_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", - "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_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_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", - "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::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_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", @@ -1763,7 +1704,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", @@ -1826,7 +1766,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", @@ -2874,7 +2813,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", @@ -3002,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", @@ -3051,14 +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_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", "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", @@ -3069,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", @@ -3874,7 +3805,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/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/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"))] 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_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..3a2462c4b49a 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 @@ -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( @@ -75,12 +65,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( 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/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_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/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/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) 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" ) 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()