diff --git a/cpp/examples/hybrid_scan_io/hybrid_scan_composer.cpp b/cpp/examples/hybrid_scan_io/hybrid_scan_composer.cpp index dc7f724e827f..0fedbac5e6f5 100644 --- a/cpp/examples/hybrid_scan_io/hybrid_scan_composer.cpp +++ b/cpp/examples/hybrid_scan_io/hybrid_scan_composer.cpp @@ -426,22 +426,6 @@ std::unique_ptr hybrid_scan( } } -// Specialization for two-step read without page index -template - requires(not single_step_read and not use_page_index) -std::unique_ptr inline hybrid_scan( - io_source const& io_source, - std::optional filter_expression, - std::unordered_set const& filters, - bool verbose, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - static_assert(single_step_read or use_page_index, - "Hybrid scan requires parquet page index for two-step parquet read"); - return nullptr; -} - // Instantiations for hybrid_scan template template std::unique_ptr hybrid_scan( @@ -460,6 +444,14 @@ template std::unique_ptr hybrid_scan( cuda::stream_ref, rmm::device_async_resource_ref); +template std::unique_ptr hybrid_scan( + io_source const&, + std::optional, + std::unordered_set const&, + bool, + cuda::stream_ref, + rmm::device_async_resource_ref); + template std::unique_ptr hybrid_scan( io_source const&, std::optional, diff --git a/cpp/include/cudf/io/experimental/hybrid_scan.hpp b/cpp/include/cudf/io/experimental/hybrid_scan.hpp index 772b65e62fc9..333fbfcf2c05 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan.hpp @@ -564,7 +564,7 @@ class hybrid_scan_reader { * * @param row_group_indices Input row groups indices * @param column_chunk_data Device spans of column chunk data of filter columns - * @param[in,out] row_mask Mutable boolean column indicating surviving rows from page pruning + * @param[in,out] row_mask Mutable boolean column indicating surviving rows * @param mask_data_pages Whether to build and use a data page mask using the row mask * @param options Parquet reader options * @param stream CUDA stream used for device memory operations and kernel launches @@ -645,7 +645,7 @@ class hybrid_scan_reader { * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if * there is no limit * @param row_group_indices Input row groups indices - * @param row_mask Boolean column indicating which rows need to be read + * @param[in,out] row_mask Mutable boolean column indicating surviving rows * @param mask_data_pages Whether to build and use a data page mask using the row mask * @param column_chunk_data Device spans of column chunk data of filter columns * @param options Parquet reader options @@ -656,7 +656,7 @@ class hybrid_scan_reader { std::size_t chunk_read_limit, std::size_t pass_read_limit, std::span row_group_indices, - cudf::column_view const& row_mask, + cudf::mutable_column_view const& row_mask, use_data_page_mask mask_data_pages, std::span const> column_chunk_data, parquet_reader_options const& options, diff --git a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp index c75fa3d186d3..0cc290270ff5 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp @@ -274,7 +274,7 @@ class hybrid_scan_multifile { * @param column_chunk_data Flattened device spans of filter column chunk data returned in the * same order as `filter_column_chunks_byte_ranges` * @param[in,out] row_mask Mutable boolean column spanning all selected rows across all sources - * and indicating surviving rows from page pruning + * indicating surviving rows * @param mask_data_pages Whether to build and use a data page mask using the row mask * @param options Parquet reader options * @param stream CUDA stream used for device memory operations and kernel launches @@ -389,8 +389,8 @@ class hybrid_scan_multifile { * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if * there is no limit * @param row_group_indices Span of vectors of input row group indices, one per source - * @param row_mask Boolean column spanning all selected rows across all sources and indicating - * which rows need to be read + * @param[in,out] row_mask Mutable boolean column spanning all selected rows across all sources + * indicating surviving rows * @param mask_data_pages Whether to build and use a data page mask using the row mask * @param column_chunk_data Flattened device spans of filter column chunk data returned in the * same order as `filter_column_chunks_byte_ranges` @@ -402,7 +402,7 @@ class hybrid_scan_multifile { std::size_t chunk_read_limit, std::size_t pass_read_limit, cudf::host_span const> row_group_indices, - cudf::column_view const& row_mask, + cudf::mutable_column_view const& row_mask, use_data_page_mask mask_data_pages, cudf::host_span const> column_chunk_data, parquet_reader_options const& options, diff --git a/cpp/src/io/parquet/experimental/hybrid_scan.cpp b/cpp/src/io/parquet/experimental/hybrid_scan.cpp index 09faac8c261a..4f98d14ee92c 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan.cpp @@ -293,7 +293,7 @@ void hybrid_scan_reader::setup_chunking_for_filter_columns( std::size_t chunk_read_limit, std::size_t pass_read_limit, std::span row_group_indices, - cudf::column_view const& row_mask, + cudf::mutable_column_view const& row_mask, use_data_page_mask mask_data_pages, std::span const> column_chunk_data, parquet_reader_options const& options, diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu b/cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu index 970f790da3e3..4918b8b1c532 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu +++ b/cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu @@ -126,7 +126,18 @@ void hybrid_scan_reader_impl::setup_next_pass( set_sparse_pass_page_mask(column_chunk_data); } else { setup_compressed_data(column_chunk_data); - set_pass_page_mask(data_page_mask); + // When offset index is absent, compute and use the data page mask using the decoded page + // headers from `setup_compressed_data`. + auto const data_page_mask_pghdr = [&]() { + if (not _has_offset_index and not _row_mask.is_empty()) { + return compute_data_page_mask_with_page_headers(); + } + return thrust::host_vector{}; + }(); + set_pass_page_mask( + data_page_mask_pghdr.empty() + ? data_page_mask + : std::span{data_page_mask_pghdr.data(), data_page_mask_pghdr.size()}); } // detect malformed columns. diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp index 334f47f6670e..5134c3ee63d0 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp @@ -338,24 +338,18 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base { * Compute a vector of boolean vectors indicating which data pages need to be decoded to * construct each input column based on the row mask, one vector per column * - * @tparam ColumnView Type of the row mask column view - cudf::mutable_column_view for filter - * columns and cudf::column_view for payload columns - * - * @param row_mask Boolean column indicating which rows need to be read after page-pruning + * @param row_mask Non-nullable boolean column view indicating surviving rows * @param row_group_indices Input row groups indices * @param input_columns Input column information - * @param row_mask_offset Offset into the row mask column for the current pass * @param stream CUDA stream used for device memory operations and kernel launches * * @return Boolean vector indicating which data pages need to be decoded to produce * the output table based on the input row mask across all input columns */ - template [[nodiscard]] thrust::host_vector compute_data_page_mask( - ColumnView const& row_mask, + cudf::column_view const& row_mask, std::span const> row_group_indices, std::span input_columns, - cudf::size_type row_mask_offset, cuda::stream_ref stream) const; }; diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index a3090c1393fe..0dacccfc0e56 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -9,6 +9,7 @@ #include "hybrid_scan_helpers.hpp" #include "io/parquet/reader_impl_chunking_utils.cuh" #include "io/parquet/synthetic_column_helpers.hpp" +#include "page_index_filter_utils.hpp" #include #include @@ -655,8 +656,8 @@ hybrid_scan_reader_impl::payload_pages_byte_ranges( // Compute the data page mask auto const mask_size = mask_offsets.back(); - auto data_page_mask = _extended_metadata->compute_data_page_mask( - row_mask, row_group_indices, _input_columns, 0, stream); + auto data_page_mask = + _extended_metadata->compute_data_page_mask(row_mask, row_group_indices, _input_columns, stream); CUDF_EXPECTS(data_page_mask.empty() or data_page_mask.size() == mask_size, "Computed data page mask does not match offset indexes"); @@ -757,8 +758,9 @@ table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns( auto data_page_mask = thrust::host_vector{}; if (mask_data_pages == use_data_page_mask::YES) { + _row_mask = set_nulls_to_true(row_mask, stream); data_page_mask = _extended_metadata->compute_data_page_mask( - row_mask, row_group_indices, _input_columns, _row_mask_offset, stream); + _row_mask, row_group_indices, _input_columns, stream); } prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, data_page_mask); @@ -795,8 +797,9 @@ table_with_metadata hybrid_scan_reader_impl::materialize_payload_columns( auto data_page_mask = thrust::host_vector{}; if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) { + _row_mask = row_mask; data_page_mask = _extended_metadata->compute_data_page_mask( - row_mask, row_group_indices, _input_columns, _row_mask_offset, stream); + _row_mask, row_group_indices, _input_columns, stream); } prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, data_page_mask); @@ -834,7 +837,7 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns( std::size_t chunk_read_limit, std::size_t pass_read_limit, std::span const> row_group_indices, - cudf::column_view const& row_mask, + cudf::mutable_column_view const& row_mask, use_data_page_mask mask_data_pages, std::span const> column_chunk_data, parquet_reader_options const& options, @@ -866,8 +869,9 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns( auto data_page_mask = thrust::host_vector{}; if (mask_data_pages == use_data_page_mask::YES) { + _row_mask = set_nulls_to_true(row_mask, stream); data_page_mask = _extended_metadata->compute_data_page_mask( - row_mask, row_group_indices, _input_columns, _row_mask_offset, stream); + _row_mask, row_group_indices, _input_columns, stream); } prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, data_page_mask); @@ -926,8 +930,9 @@ void hybrid_scan_reader_impl::setup_chunking_for_payload_columns( auto data_page_mask = thrust::host_vector{}; if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) { + _row_mask = row_mask; data_page_mask = _extended_metadata->compute_data_page_mask( - row_mask, row_group_indices, _input_columns, _row_mask_offset, stream); + _row_mask, row_group_indices, _input_columns, stream); } prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, data_page_mask); @@ -1135,7 +1140,6 @@ bool hybrid_scan_reader_impl::has_next_table_chunk() void hybrid_scan_reader_impl::reset_internal_state() { - _row_mask_offset = 0; _file_itm_data = file_intermediate_data{}; _file_preprocessed = false; _has_offset_index = false; @@ -1159,6 +1163,10 @@ void hybrid_scan_reader_impl::reset_internal_state() _output_chunk_read_limit = 0; _strings_to_categorical = false; _reader_column_schema.reset(); + + _row_mask = column_view{}; + _row_mask_offset = 0; + _expr_conv = parquet_filter_normalizer{}; _mr = cudf::get_current_device_resource_ref(); } @@ -1380,9 +1388,9 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( // Prepend the source and row index columns to filter columns only if (read_columns_mode == read_columns_mode::FILTER_COLUMNS) { if (_options.prepend_row_index_column) { - out_columns.emplace( - out_columns.begin(), - synthesize_row_index_column(_file_itm_data.row_groups, read_info, _stream, _mr)); + out_columns.emplace(out_columns.begin(), + parquet::detail::synthesize_row_index_column( + _file_itm_data.row_groups, read_info, _stream, _mr)); out_metadata.schema_info.emplace(out_metadata.schema_info.begin(), column_name_info{.name = "row_index", .is_nullable = false}); } @@ -1506,6 +1514,73 @@ void hybrid_scan_reader_impl::set_pass_page_mask(std::span data_page mark_buffers_nullable_for_pruned_pages(); } +thrust::host_vector hybrid_scan_reader_impl::compute_data_page_mask_with_page_headers() +{ + auto const& pass = *_pass_itm_data; + + // Return an empty vector if all rows are required + if (are_all_rows_retained(_row_mask, _stream)) { return thrust::host_vector(0); } + + std::vector page_row_offsets; + page_row_offsets.reserve(pass.pages.size() * 2); + + // Maps each data page to its flat-page range; -1 keeps nested pages enabled. + std::vector row_range_map; + row_range_map.reserve(pass.pages.size()); + + cudf::size_type previous_chunk_idx = -1; + auto max_page_size = cudf::size_type{0}; + + for (auto const& page : pass.pages) { + // Ignore dictionary pages altogether + if (page.flags & parquet::detail::PAGEINFO_FLAGS_DICTIONARY) { continue; } + + auto const& chunk = pass.chunks[page.chunk_idx]; + + // Don't prune list column pages as rows may span page boundaries when offset index isn't + // present. + if (chunk.max_level[parquet::detail::level_type::REPETITION] > 0) { + row_range_map.push_back(-1); + continue; + } + + auto const page_start = chunk.start_row + page.chunk_row; + auto const page_end = page_start + page.num_rows; + max_page_size = std::max(max_page_size, page_end - page_start); + + // Starting a new column chunk. Push page start row + if (page.chunk_idx != previous_chunk_idx) { + page_row_offsets.push_back(page_start); + previous_chunk_idx = page.chunk_idx; + } + + // Push row range index and page end row + row_range_map.push_back(page_row_offsets.size() - 1); + page_row_offsets.push_back(page_end); + } + + auto data_page_mask = thrust::host_vector{}; + + // Compute the row range mask + CUDF_EXPECTS(std::cmp_equal(_row_mask.size(), pass.num_rows), + "Row mask must span across all rows in the pass"); + auto const row_range_mask = + compute_row_range_selection_mask(_row_mask, page_row_offsets, max_page_size, _stream); + + if (row_range_mask.empty()) { return data_page_mask; } + + CUDF_EXPECTS(row_range_mask.size() == page_row_offsets.size() - 1, + "Encountered invalid row range mask size"); + + data_page_mask.reserve(row_range_map.size()); + + // Scatter row range results while retaining list column pages. + for (auto const range_idx : row_range_map) { + data_page_mask.push_back(range_idx < 0 ? true : row_range_mask[range_idx]); + } + return data_page_mask; +} + void hybrid_scan_reader_impl::set_sparse_pass_page_mask( std::span const> page_data) { diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp index 19027cd8fcdd..5e30e2b578f0 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp @@ -253,7 +253,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { std::size_t chunk_read_limit, std::size_t pass_read_limit, std::span const> row_group_indices, - cudf::column_view const& row_mask, + cudf::mutable_column_view const& row_mask, use_data_page_mask mask_data_pages, std::span const> column_chunk_data, parquet_reader_options const& options, @@ -394,6 +394,11 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { */ void set_sparse_pass_page_mask(std::span const> page_data); + /** + * @brief Compute a data page mask from the decoded page headers. + */ + [[nodiscard]] thrust::host_vector compute_data_page_mask_with_page_headers(); + /** * @brief Mark output buffers nullable when page pruning synthesizes null rows */ @@ -626,6 +631,8 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { std::optional> _filter_columns_names; + cudf::column_view _row_mask{}; + std::vector _original_output_buffers_template; cudf::size_type _row_mask_offset{0}; diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp index 259435401aab..16360d0bc7d9 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp @@ -197,7 +197,7 @@ void hybrid_scan_multifile::setup_chunking_for_filter_columns( std::size_t chunk_read_limit, std::size_t pass_read_limit, cudf::host_span const> row_group_indices, - cudf::column_view const& row_mask, + cudf::mutable_column_view const& row_mask, use_data_page_mask mask_data_pages, cudf::host_span const> column_chunk_data, parquet_reader_options const& options, diff --git a/cpp/src/io/parquet/experimental/page_index_filter.cu b/cpp/src/io/parquet/experimental/page_index_filter.cu index 7ec4aa859f0c..6ec51293aa50 100644 --- a/cpp/src/io/parquet/experimental/page_index_filter.cu +++ b/cpp/src/io/parquet/experimental/page_index_filter.cu @@ -10,19 +10,15 @@ #include #include -#include #include #include #include #include -#include #include #include -#include #include #include #include -#include #include #include #include @@ -41,7 +37,6 @@ #include #include -#include #include namespace cudf::io::parquet::experimental::detail { @@ -577,14 +572,13 @@ struct page_stats_to_row_mask_converter : public page_stats_caster { auto const page_mask_nullmask = page_mask->null_count() - ? cudf::detail::make_host_vector_async( + ? cudf::detail::make_host_vector( cudf::device_span{ page_mask->view().null_mask(), static_cast(num_bitmask_words(page_mask->size()))}, stream) : cudf::detail::make_empty_host_vector(0, stream); - stream.sync(); auto [row_mask_data, row_mask_bitmask] = build_data_and_nullmask(page_mask->mutable_view(), page_mask_nullmask.data(), @@ -606,234 +600,6 @@ struct page_stats_to_row_mask_converter : public page_stats_caster { } }; -/* - * @brief Functor to build a Fenwick tree level from the previous level data - * - * @param tree_level_ptrs Pointers to the start of Fenwick tree level data - * @param prev_level Previous tree level - * @param prev_level_size Size of the previous tree level - * @param current_level_size Size of the current tree level - */ -struct build_fenwick_tree_level_functor { - bool** tree_level_ptrs; - cudf::size_type prev_level; - cudf::size_type prev_level_size; - cudf::size_type current_level_size; - - /** - * @brief Builds the next Fenwick tree level from the current level data - * by ORing two elements at the current level. - * - * elem_current_level[idx] = elem_prev_level[idx * 2] OR elem_prev_level[idx * 2 + 1]; - * - * @param current_level_idx Current tree level element index - */ - __device__ void operator()(cudf::size_type current_level_idx) const noexcept - { - auto const prev_level_ptr = tree_level_ptrs[prev_level]; - auto current_level_ptr = tree_level_ptrs[prev_level + 1]; - - // Handle the odd-sized remaining element if prev_level_size is odd - if (prev_level_size % 2 and current_level_idx == current_level_size - 1) { - current_level_ptr[current_level_idx] = prev_level_ptr[prev_level_size - 1]; - } else { - current_level_ptr[current_level_idx] = - prev_level_ptr[(current_level_idx * 2)] or prev_level_ptr[(current_level_idx * 2) + 1]; - } - } -}; - -/** - * @brief Functor to binary search a `true` value in the Fenwick tree in range [start, end) - * - * @param tree_level_ptrs Pointers to the start of Fenwick tree level data - * @param page_offsets Pointer to page offsets describing each search range i as [page_offsets[i], - * page_offsets[i+1)) - * @param num_ranges Number of search ranges - */ -struct search_fenwick_tree_functor { - bool** tree_level_ptrs; - cudf::size_type const* page_offsets; - cudf::size_type num_ranges; - - /** - * @brief Enum class to represent which range boundary we are currently processing - */ - enum class boundary : uint8_t { - START = 0, - END = 1, - }; - - /** - * @brief Checks if a value is a power of two - * - * @param value Value to check - * @return Boolean indicating if the value is a power of two - */ - __device__ bool inline constexpr is_power_of_two(cudf::size_type value) const noexcept - { - return (value & (value - 1)) == 0; - } - - /** - * @brief Finds the smallest power of two in the range [start, end). If no power of two is - * found, returns a zero. - * - * @param start Range start - * @param end Range end - * @return Largest power of two in the range [start, end) or a zero if no power of two is found - */ - __device__ cudf::size_type inline constexpr smallest_power_of_two_in_range( - cudf::size_type start, cudf::size_type end) const noexcept - { - start--; - start |= start >> 1; - start |= start >> 2; - start |= start >> 4; - start |= start >> 8; - start |= start >> 16; - auto const result = start + 1; - return result < end ? result : 0; - } - - /** - * @brief Finds the largest power of two in the range (start, end]. If no power of two is found, - * returns a zero. - * - * @param start Range start - * @param end Range end - * @return Largest power of two in the range (start, end] or a zero if no power of two is found - */ - __device__ size_type inline constexpr largest_power_of_two_in_range(size_type start, - size_type end) const noexcept - { - auto constexpr nbits = cudf::detail::size_in_bits() - 1; - auto const result = size_type{1} << (nbits - cuda::std::countl_zero(end)); - return result > start ? result : 0; - } - - /** - * @brief Aligns a range boundary to the next power-of-two block - * - * @tparam Boundary Current boundary type (START or END) - * @param start Range start - * @param end Range end - * @return A pair of the tree level and block size - */ - template - __device__ auto inline constexpr align_range_boundary(cudf::size_type start, - cudf::size_type end) const noexcept - { - if constexpr (Boundary == boundary::START) { - if (start == 0 or is_power_of_two(start)) { - auto const block_size = - cuda::std::max(start & -start, largest_power_of_two_in_range(start, end)); - auto const tree_level = cuda::std::countr_zero(block_size); - return cuda::std::pair{tree_level, block_size}; - } else { - auto const tree_level = cuda::std::countr_zero(start); - return cuda::std::pair{tree_level, size_type{1} << tree_level}; - } - } else { - auto block_size = end & -end; - if (start > 0 and is_power_of_two(end)) { - auto const next_alignment = cuda::std::max(smallest_power_of_two_in_range(start, end), - largest_power_of_two_in_range(0, end - start)); - block_size = end - next_alignment; - } - return cuda::std::pair{cuda::std::countr_zero(block_size), block_size}; - } - } - - /** - * @brief Queries the Fenwick tree for the given boundary position, tree level and block size - * - * @tparam Boundary Current boundary type (START or END) - * @param boundary_pos Current boundary position - * @param tree_level Corresponding tree level to query - * @param block_size Alignment block size of the current boundary - * @return Boolean indicating if a `true` value is found in the fenwick tree - */ - template - __device__ bool inline constexpr query_fenwick_tree(cudf::size_type boundary_pos, - cudf::size_type tree_level, - cudf::size_type block_size) const noexcept - { - if constexpr (Boundary == boundary::START) { - auto const mask_index = boundary_pos >> tree_level; - return tree_level_ptrs[tree_level][mask_index]; - } else { - auto const mask_index = (boundary_pos - block_size) >> tree_level; - return tree_level_ptrs[tree_level][mask_index]; - } - } - - /** - * @brief Searches the Fenwick tree to find a `true` value in range [start, end) - * - * Algorithm: While `start` < `end`, align `start` UP and `end` DOWN to the next power-of-two - * searchable tree block. For the two aligned blocks, query the fenwick tree at corresponding - * levels for a `true` value (larger block first). If found, return. Else, move the boundaries - * to their alignments. - * - * @param range_idx Index of the range to search - * @return Boolean indicating if a `true` value is found in the range - */ - __device__ bool operator()(cudf::size_type range_idx) const noexcept - { - // Retrieve start and end for the current range [start, end) - size_type start = page_offsets[range_idx]; - size_type end = page_offsets[range_idx + 1]; - - // Return early if the range is empty or invalid - if (start >= end or range_idx >= num_ranges) { return false; } - - // Binary search decomposition loop - while (start < end) { - // Find the largest power-of-two block that aligns `start` up - auto const [start_tree_level, start_block_size] = - align_range_boundary(start, end); - - // Find the largest power-of-two block that aligns `end` down - auto const [end_tree_level, end_block_size] = align_range_boundary(start, end); - - // Check the larger block first to minimize the number of queries - if (start_block_size >= end_block_size) { - // Check the `start` side alignment block first - if (start + start_block_size <= end) { - if (query_fenwick_tree(start, start_tree_level, start_block_size)) { - return true; - } - start += start_block_size; - } - // Check the `end` side alignment block if it's still in range - if (end - end_block_size >= start) { - if (query_fenwick_tree(end, end_tree_level, end_block_size)) { - return true; - } - end -= end_block_size; - } - } else { - // Check the `end` side alignment block first - if (end - end_block_size >= start) { - if (query_fenwick_tree(end, end_tree_level, end_block_size)) { - return true; - } - end -= end_block_size; - } - // Check the `start` side alignment block if it's still in range - if (start + start_block_size <= end) { - if (query_fenwick_tree(start, start_tree_level, start_block_size)) { - return true; - } - start += start_block_size; - } - } - } - return false; - } -}; - } // namespace std::unique_ptr aggregate_reader_metadata::build_row_mask_with_page_index_stats( @@ -978,12 +744,10 @@ std::unique_ptr aggregate_reader_metadata::build_row_mask_with_pag page_stats_table, stats_expr.get_stats_expr().get(), stream, mr); } -template thrust::host_vector aggregate_reader_metadata::compute_data_page_mask( - ColumnView const& row_mask, + cudf::column_view const& row_mask, std::span const> row_group_indices, std::span input_columns, - cudf::size_type row_mask_offset, cuda::stream_ref stream) const { CUDF_FUNC_RANGE(); @@ -998,19 +762,14 @@ thrust::host_vector aggregate_reader_metadata::compute_data_page_mask( std::invalid_argument); CUDF_EXPECTS( - std::cmp_less_equal(row_mask_offset + total_rows, row_mask.size()), + std::cmp_equal(total_rows, row_mask.size()), "Encountered a mismatch in number of rows in the row group pass and the row mask size", std::overflow_error); + CUDF_EXPECTS( + row_mask.null_count() == 0, "Row mask must not contain nulls", std::invalid_argument); - // Return an empty vector if all rows are invalid or all rows are required - if (std::cmp_equal(row_mask.null_count(row_mask_offset, row_mask_offset + total_rows, stream), - total_rows) or - cudf::detail::all_of(row_mask.template begin() + row_mask_offset, - row_mask.template begin() + row_mask_offset + total_rows, - cuda::std::identity{}, - stream)) { - return thrust::host_vector(0); - } + // Return an empty vector if all rows are required + if (are_all_rows_retained(row_mask, stream)) { return thrust::host_vector{}; } // Collect column schema indices from the input columns. auto column_schema_indices = std::vector(input_columns.size()); @@ -1024,8 +783,8 @@ thrust::host_vector aggregate_reader_metadata::compute_data_page_mask( page_index_presence(row_group_indices, column_schema_indices).second; if (not has_offset_index) { CUDF_LOG_WARN( - "Encountered missing Parquet offset index for one or more output columns. Skipping page " - "pruning."); + "Encountered missing Parquet offset index for one or more output columns. Skipping " + "page-index based pruning."); return thrust::host_vector(0); } @@ -1099,112 +858,26 @@ thrust::host_vector aggregate_reader_metadata::compute_data_page_mask( }); } - // Make sure all row_mask elements contain valid values even if they are nulls - if constexpr (cuda::std::is_same_v) { - if (row_mask.nullable() and row_mask.null_count() > 0) { - thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator(row_mask_offset), - cuda::counting_iterator(row_mask_offset + total_rows), - [row_mask = row_mask.template begin(), - null_mask = row_mask.null_mask()] __device__(auto const row_idx) { - if (not bit_is_set(null_mask, row_idx)) { row_mask[row_idx] = true; } - }); - } - } else { - CUDF_EXPECTS(not row_mask.nullable() or row_mask.null_count() == 0, - "Row mask must not contain nulls for payload columns"); - } - - auto const mr = cudf::get_current_device_resource_ref(); - - // Compute fenwick tree level offsets and total size (level 1 and higher) - auto const tree_level_offsets = compute_fenwick_tree_level_offsets(total_rows, max_page_size); - auto const num_levels = static_cast(tree_level_offsets.size()); - // Buffer to store Fenwick tree levels (level 1 and higher) data - auto tree_levels_data = rmm::device_uvector(tree_level_offsets.back(), stream, mr); - - // Pointers to each Fenwick tree level data - auto host_tree_level_ptrs = cudf::detail::make_pinned_vector_async(num_levels, stream); - // Zeroth level is just the row mask itself - host_tree_level_ptrs[0] = const_cast(row_mask.template begin()) + row_mask_offset; - std::for_each(cuda::counting_iterator{1}, - cuda::counting_iterator{num_levels}, - [&](auto const level_idx) { - host_tree_level_ptrs[level_idx] = - tree_levels_data.data() + tree_level_offsets[level_idx - 1]; - }); + auto data_page_mask = thrust::host_vector{}; - auto fenwick_tree_level_ptrs = - cudf::detail::make_device_uvector_async(host_tree_level_ptrs, stream, mr); - - // Build Fenwick tree levels (zeroth level is just the row mask itself) - auto prev_level_size = static_cast(total_rows); - std::for_each( - cuda::counting_iterator{0}, - cuda::counting_iterator{num_levels - 1}, - [&](auto const prev_level) { - auto const current_level_size = cudf::util::div_rounding_up_safe(prev_level_size, 2); - thrust::for_each( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - cuda::counting_iterator{current_level_size}, - build_fenwick_tree_level_functor{ - fenwick_tree_level_ptrs.data(), prev_level, prev_level_size, current_level_size}); - prev_level_size = current_level_size; - }); - - // Search the Fenwick tree to see if there's a surviving row in each page's row range - auto const num_ranges = static_cast(page_row_offsets.size() - 1); - rmm::device_uvector device_data_page_mask(num_ranges, stream, mr); - // Use a pinned bounce buffer to avoid pageable h2d copy - auto pinned_page_offsets = cudf::detail::make_pinned_vector( - cudf::host_span{page_row_offsets}, stream); - auto page_offsets = cudf::detail::make_device_uvector_async(pinned_page_offsets, stream, mr); - thrust::transform( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - cuda::counting_iterator{num_ranges}, - device_data_page_mask.begin(), - search_fenwick_tree_functor{fenwick_tree_level_ptrs.data(), page_offsets.data(), num_ranges}); - - // Copy over search results to host - auto host_results = cudf::detail::make_pinned_vector_async(device_data_page_mask, stream); - auto const total_pages = pinned_page_offsets.size() - num_columns; - auto data_page_mask = thrust::host_vector{}; - data_page_mask.reserve(total_pages); - auto host_results_iter = host_results.begin(); - stream.sync(); + auto const row_range_mask = + compute_row_range_selection_mask(row_mask, page_row_offsets, max_page_size, stream); + if (row_range_mask.empty()) { return data_page_mask; } + data_page_mask.reserve(page_row_offsets.size() - num_columns); // Discard results for invalid ranges. i.e. ranges starting at the last page of a column and // ending at the first page of the next column - auto num_pages_inserted = 0; std::for_each(cuda::counting_iterator{0}, cuda::counting_iterator{num_columns}, [&](auto col_idx) { auto const col_num_pages = col_page_offsets[col_idx + 1] - col_page_offsets[col_idx] - 1; - data_page_mask.insert(data_page_mask.begin() + num_pages_inserted, - host_results_iter, - host_results_iter + col_num_pages); - host_results_iter += col_num_pages + 1; - num_pages_inserted += col_num_pages; + auto const first_page_range = col_page_offsets[col_idx]; + data_page_mask.insert(data_page_mask.end(), + row_range_mask.begin() + first_page_range, + row_range_mask.begin() + first_page_range + col_num_pages); }); return data_page_mask; } -// Instantiate the templates with ColumnView as cudf::column_view and cudf::mutable_column_view -template thrust::host_vector aggregate_reader_metadata::compute_data_page_mask< - cudf::column_view>(cudf::column_view const& row_mask, - std::span const> row_group_indices, - std::span input_columns, - cudf::size_type row_mask_offset, - cuda::stream_ref stream) const; - -template thrust::host_vector aggregate_reader_metadata::compute_data_page_mask< - cudf::mutable_column_view>(cudf::mutable_column_view const& row_mask, - std::span const> row_group_indices, - std::span input_columns, - cudf::size_type row_mask_offset, - cuda::stream_ref stream) const; - } // namespace cudf::io::parquet::experimental::detail diff --git a/cpp/src/io/parquet/experimental/page_index_filter_utils.cu b/cpp/src/io/parquet/experimental/page_index_filter_utils.cu index 07c81c787afe..9df82e836c10 100644 --- a/cpp/src/io/parquet/experimental/page_index_filter_utils.cu +++ b/cpp/src/io/parquet/experimental/page_index_filter_utils.cu @@ -5,16 +5,24 @@ #include "page_index_filter_utils.hpp" +#include +#include +#include #include #include #include #include #include +#include #include +#include +#include #include -#include +#include +#include +#include #include #include @@ -22,6 +30,266 @@ namespace cudf::io::parquet::experimental::detail { +namespace { + +/* + * @brief Functor to build a Fenwick tree level from the previous level data + * + * @param tree_level_ptrs Pointers to the start of Fenwick tree level data + * @param prev_level Previous tree level + * @param prev_level_size Size of the previous tree level + * @param current_level_size Size of the current tree level + */ +struct build_fenwick_tree_level_functor { + bool** tree_level_ptrs; + cudf::size_type prev_level; + cudf::size_type prev_level_size; + cudf::size_type current_level_size; + + /** + * @brief Builds the next Fenwick tree level from the current level data + * by ORing two elements at the current level. + * + * elem_current_level[idx] = elem_prev_level[idx * 2] OR elem_prev_level[idx * 2 + 1]; + * + * @param current_level_idx Current tree level element index + */ + __device__ void operator()(cudf::size_type current_level_idx) const noexcept + { + auto const prev_level_ptr = tree_level_ptrs[prev_level]; + auto current_level_ptr = tree_level_ptrs[prev_level + 1]; + + // Handle the odd-sized remaining element if prev_level_size is odd + if (prev_level_size % 2 and current_level_idx == current_level_size - 1) { + current_level_ptr[current_level_idx] = prev_level_ptr[prev_level_size - 1]; + } else { + current_level_ptr[current_level_idx] = + prev_level_ptr[(current_level_idx * 2)] or prev_level_ptr[(current_level_idx * 2) + 1]; + } + } +}; + +/** + * @brief Functor to binary search a `true` value in the Fenwick tree in range [start, end) + * + * @param tree_level_ptrs Pointers to the start of Fenwick tree level data + * @param page_offsets Pointer to page offsets describing each search range i as [page_offsets[i], + * page_offsets[i+1)) + * @param num_ranges Number of search ranges + */ +struct search_fenwick_tree_functor { + bool** tree_level_ptrs; + cudf::size_type const* page_offsets; + cudf::size_type num_ranges; + + /** + * @brief Enum class to represent which range boundary we are currently processing + */ + enum class boundary : uint8_t { + START = 0, + END = 1, + }; + + /** + * @brief Checks if a value is a power of two + * + * @param value Value to check + * @return Boolean indicating if the value is a power of two + */ + __device__ bool inline constexpr is_power_of_two(cudf::size_type value) const noexcept + { + return (value & (value - 1)) == 0; + } + + /** + * @brief Finds the smallest power of two in the range [start, end). If no power of two is + * found, returns a zero. + * + * @param start Range start + * @param end Range end + * @return Largest power of two in the range [start, end) or a zero if no power of two is found + */ + __device__ cudf::size_type inline constexpr smallest_power_of_two_in_range( + cudf::size_type start, cudf::size_type end) const noexcept + { + start--; + start |= start >> 1; + start |= start >> 2; + start |= start >> 4; + start |= start >> 8; + start |= start >> 16; + auto const result = start + 1; + return result < end ? result : 0; + } + + /** + * @brief Finds the largest power of two in the range (start, end]. If no power of two is found, + * returns a zero. + * + * @param start Range start + * @param end Range end + * @return Largest power of two in the range (start, end] or a zero if no power of two is found + */ + __device__ size_type inline constexpr largest_power_of_two_in_range(size_type start, + size_type end) const noexcept + { + auto constexpr nbits = cudf::detail::size_in_bits() - 1; + auto const result = size_type{1} << (nbits - cuda::std::countl_zero(end)); + return result > start ? result : 0; + } + + /** + * @brief Aligns a range boundary to the next power-of-two block + * + * @tparam Boundary Current boundary type (START or END) + * @param start Range start + * @param end Range end + * @return A pair of the tree level and block size + */ + template + __device__ auto inline constexpr align_range_boundary(cudf::size_type start, + cudf::size_type end) const noexcept + { + if constexpr (Boundary == boundary::START) { + if (start == 0 or is_power_of_two(start)) { + auto const block_size = + cuda::std::max(start & -start, largest_power_of_two_in_range(start, end)); + auto const tree_level = cuda::std::countr_zero(block_size); + return cuda::std::pair{tree_level, block_size}; + } else { + auto const tree_level = cuda::std::countr_zero(start); + return cuda::std::pair{tree_level, size_type{1} << tree_level}; + } + } else { + auto block_size = end & -end; + if (start > 0 and is_power_of_two(end)) { + auto const next_alignment = cuda::std::max(smallest_power_of_two_in_range(start, end), + largest_power_of_two_in_range(0, end - start)); + block_size = end - next_alignment; + } + return cuda::std::pair{cuda::std::countr_zero(block_size), block_size}; + } + } + + /** + * @brief Queries the Fenwick tree for the given boundary position, tree level and block size + * + * @tparam Boundary Current boundary type (START or END) + * @param boundary_pos Current boundary position + * @param tree_level Corresponding tree level to query + * @param block_size Alignment block size of the current boundary + * @return Boolean indicating if a `true` value is found in the fenwick tree + */ + template + __device__ bool inline constexpr query_fenwick_tree(cudf::size_type boundary_pos, + cudf::size_type tree_level, + cudf::size_type block_size) const noexcept + { + if constexpr (Boundary == boundary::START) { + auto const mask_index = boundary_pos >> tree_level; + return tree_level_ptrs[tree_level][mask_index]; + } else { + auto const mask_index = (boundary_pos - block_size) >> tree_level; + return tree_level_ptrs[tree_level][mask_index]; + } + } + + /** + * @brief Searches the Fenwick tree to find a `true` value in range [start, end) + * + * Algorithm: While `start` < `end`, align `start` UP and `end` DOWN to the next power-of-two + * searchable tree block. For the two aligned blocks, query the fenwick tree at corresponding + * levels for a `true` value (larger block first). If found, return. Else, move the boundaries + * to their alignments. + * + * @param range_idx Index of the range to search + * @return Boolean indicating if a `true` value is found in the range + */ + __device__ bool operator()(cudf::size_type range_idx) const noexcept + { + // Retrieve start and end for the current range [start, end) + size_type start = page_offsets[range_idx]; + size_type end = page_offsets[range_idx + 1]; + + // Return early if the range is empty or invalid + if (start >= end or range_idx >= num_ranges) { return false; } + + // Binary search decomposition loop + while (start < end) { + // Find the largest power-of-two block that aligns `start` up + auto const [start_tree_level, start_block_size] = + align_range_boundary(start, end); + + // Find the largest power-of-two block that aligns `end` down + auto const [end_tree_level, end_block_size] = align_range_boundary(start, end); + + // Check the larger block first to minimize the number of queries + if (start_block_size >= end_block_size) { + // Check the `start` side alignment block first + if (start + start_block_size <= end) { + if (query_fenwick_tree(start, start_tree_level, start_block_size)) { + return true; + } + start += start_block_size; + } + // Check the `end` side alignment block if it's still in range + if (end - end_block_size >= start) { + if (query_fenwick_tree(end, end_tree_level, end_block_size)) { + return true; + } + end -= end_block_size; + } + } else { + // Check the `end` side alignment block first + if (end - end_block_size >= start) { + if (query_fenwick_tree(end, end_tree_level, end_block_size)) { + return true; + } + end -= end_block_size; + } + // Check the `start` side alignment block if it's still in range + if (start + start_block_size <= end) { + if (query_fenwick_tree(start, start_tree_level, start_block_size)) { + return true; + } + start += start_block_size; + } + } + } + return false; + } +}; + +/** + * @brief Computes the offsets of the Fenwick tree levels (level 1 and higher) until the tree level + * block size becomes larger than the maximum page (search range) size + * + * @param level0_size Size of the zeroth tree level (the row mask) + * @param max_page_size Maximum page (search range) size + * @return Fenwick tree level offsets + */ +std::vector compute_fenwick_tree_level_offsets(cudf::size_type level0_size, + cudf::size_type max_page_size) +{ + std::vector tree_level_offsets; + tree_level_offsets.push_back(0); + + cudf::size_type current_level_size = cudf::util::div_rounding_up_safe(level0_size, 2); + cudf::size_type current_level = 1; + + while (current_level_size > 0) { + auto const block_size = 1 << current_level; + if (std::cmp_greater(block_size, max_page_size)) { break; } + tree_level_offsets.push_back(tree_level_offsets.back() + current_level_size); + current_level_size = + current_level_size == 1 ? 0 : cudf::util::div_rounding_up_safe(current_level_size, 2); + current_level++; + } + return tree_level_offsets; +} + +} // namespace + std::pair, cudf::detail::host_vector> compute_page_row_offsets_and_colchunk_page_offsets( std::span per_file_metadata, @@ -158,24 +426,81 @@ rmm::device_uvector compute_page_indices_async( return page_indices; } -std::vector compute_fenwick_tree_level_offsets(cudf::size_type level0_size, - cudf::size_type max_page_size) +cudf::column_view set_nulls_to_true(cudf::mutable_column_view const& row_mask, + cuda::stream_ref stream) { - std::vector tree_level_offsets; - tree_level_offsets.push_back(0); + if (row_mask.has_nulls()) { + auto const d_row_mask = cudf::column_device_view::create(row_mask, stream); + auto const iter = cudf::detail::make_null_replacement_iterator(*d_row_mask, true, true); + thrust::copy(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + iter, + iter + row_mask.size(), + row_mask.begin()); + } - cudf::size_type current_level_size = cudf::util::div_rounding_up_safe(level0_size, 2); - cudf::size_type current_level = 1; + return cudf::column_view{ + row_mask.type(), row_mask.size(), row_mask.head(), nullptr, 0, row_mask.offset()}; +} - while (current_level_size > 0) { - auto const block_size = 1 << current_level; - if (std::cmp_greater(block_size, max_page_size)) { break; } - tree_level_offsets.push_back(tree_level_offsets.back() + current_level_size); - current_level_size = - current_level_size == 1 ? 0 : cudf::util::div_rounding_up_safe(current_level_size, 2); - current_level++; - } - return tree_level_offsets; +bool are_all_rows_retained(cudf::column_view const& row_mask, cuda::stream_ref stream) +{ + return cudf::detail::all_of( + row_mask.begin(), row_mask.end(), cuda::std::identity{}, stream); +} + +thrust::host_vector compute_row_range_selection_mask( + cudf::column_view const& row_mask, + std::span page_row_offsets, + cudf::size_type max_page_size, + cuda::stream_ref stream) +{ + // Need at least two offsets (or one range) to search the Fenwick tree + if (page_row_offsets.size() < 2) return thrust::host_vector{}; + + auto const total_rows = row_mask.size(); + auto const mr = cudf::get_current_device_resource_ref(); + auto const tree_level_offsets = compute_fenwick_tree_level_offsets(total_rows, max_page_size); + auto const num_levels = static_cast(tree_level_offsets.size()); + auto tree_levels_data = rmm::device_uvector(tree_level_offsets.back(), stream, mr); + auto host_tree_level_ptrs = cudf::detail::make_pinned_vector_async(num_levels, stream); + host_tree_level_ptrs[0] = const_cast(row_mask.begin()); + std::for_each(cuda::counting_iterator{1}, + cuda::counting_iterator{num_levels}, + [&](auto const level_idx) { + host_tree_level_ptrs[level_idx] = + tree_levels_data.data() + tree_level_offsets[level_idx - 1]; + }); + auto tree_level_ptrs = cudf::detail::make_device_uvector_async(host_tree_level_ptrs, stream, mr); + + auto prev_level_size = total_rows; + std::for_each( + cuda::counting_iterator{0}, + cuda::counting_iterator{num_levels - 1}, + [&](auto const prev_level) { + auto const current_level_size = cudf::util::div_rounding_up_safe(prev_level_size, 2); + thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + cuda::counting_iterator{0}, + cuda::counting_iterator{current_level_size}, + build_fenwick_tree_level_functor{ + tree_level_ptrs.data(), prev_level, prev_level_size, current_level_size}); + prev_level_size = current_level_size; + }); + + auto const num_ranges = static_cast(page_row_offsets.size() - 1); + auto device_results = rmm::device_uvector(num_ranges, stream, mr); + auto pinned_page_offsets = cudf::detail::make_pinned_vector(page_row_offsets, stream); + auto page_offsets = cudf::detail::make_device_uvector_async(pinned_page_offsets, stream, mr); + thrust::transform( + rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + cuda::counting_iterator{0}, + cuda::counting_iterator{num_ranges}, + device_results.begin(), + search_fenwick_tree_functor{tree_level_ptrs.data(), page_offsets.data(), num_ranges}); + + auto results = cudf::detail::make_pinned_vector_async(device_results, stream); + stream.sync(); + + return thrust::host_vector(results.begin(), results.end()); } } // namespace cudf::io::parquet::experimental::detail diff --git a/cpp/src/io/parquet/experimental/page_index_filter_utils.hpp b/cpp/src/io/parquet/experimental/page_index_filter_utils.hpp index 264510ff48fd..bc7b5eecfa5e 100644 --- a/cpp/src/io/parquet/experimental/page_index_filter_utils.hpp +++ b/cpp/src/io/parquet/experimental/page_index_filter_utils.hpp @@ -7,6 +7,7 @@ #include "io/parquet/reader_impl_helpers.hpp" +#include #include #include #include @@ -47,8 +48,7 @@ compute_page_row_offsets_and_colchunk_page_offsets( * @param per_file_metadata Span of parquet footer metadata * @param row_group_indices Span of input row group indices * @param schema_idx Column's schema index - * @return A pair of page row offsets and the size of the largest page in this - * column + * @return A pair of page row offsets and the size of the largest page in this column */ [[nodiscard]] std::pair, size_type> compute_page_row_offsets( cudf::host_span per_file_metadata, @@ -71,14 +71,39 @@ compute_page_row_offsets_and_colchunk_page_offsets( rmm::device_async_resource_ref mr); /** - * @brief Computes the offsets of the Fenwick tree levels (level 1 and higher) until the tree level - * block size becomes larger than the maximum page (search range) size + * @brief Sets nulls in the row mask to true and returns a non-nullable row mask view * - * @param level0_size Size of the zeroth tree level (the row mask) - * @param max_page_size Maximum page (search range) size - * @return Fenwick tree level offsets + * @param row_mask Mutable row mask column view + * @param stream CUDA stream used for device memory + * operations and kernel launches + * @return Non-nullable column view of the resolved row mask */ -[[nodiscard]] std::vector compute_fenwick_tree_level_offsets( - cudf::size_type level0_size, cudf::size_type max_page_size); +[[nodiscard]] cudf::column_view set_nulls_to_true(cudf::mutable_column_view const& row_mask, + cuda::stream_ref stream); + +/** + * @brief Checks whether every row is reatained by the boolean row mask + * + * @param retention_mask Boolean column indicating retained rows + * @param stream CUDA stream used for device memory operations and kernel launches + * @return Boolean indicating whether every row is retained + */ +[[nodiscard]] bool are_all_rows_retained(cudf::column_view const& retention_mask, + cuda::stream_ref stream); + +/** + * @brief Computes a mask indicating which row ranges contain at least one selected row + * + * @param row_mask Boolean column indicating selected rows + * @param page_row_offsets Page row offsets defining the row ranges + * @param max_page_size Size of the largest page row range + * @param stream CUDA stream used for device memory operations and kernel launches + * @return Boolean vector with one entry for each consecutive row range + */ +[[nodiscard]] thrust::host_vector compute_row_range_selection_mask( + cudf::column_view const& row_mask, + std::span page_row_offsets, + cudf::size_type max_page_size, + cuda::stream_ref stream); } // namespace cudf::io::parquet::experimental::detail diff --git a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp index a47f66cc77cd..666f84ced9c6 100644 --- a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp @@ -509,7 +509,8 @@ TEST_F(HybridScanFiltersTest, FilterRowGroupsWithComplexExpressions) auto input_row_group_indices = reader->all_row_groups(options); auto stats_filtered = reader->filter_row_groups_with_stats( input_row_group_indices, options, cudf::get_default_stream()); - EXPECT_EQ(stats_filtered.size(), 2); + auto const expected = std::vector{1, 2}; + EXPECT_EQ(stats_filtered, expected); } // Filter: NOT(NOT(col0 < 100) OR col0 > 150) @@ -908,10 +909,11 @@ TEST_F(HybridScanFiltersTest, OffsetIndexOnlyDataPageMask) cudf::apply_retention_mask(written_table->view(), row_mask_view, stream, mr); CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), result.tbl->view()); - // Without offset index, data-page pruning falls back to decoding all pages. + // Without an offset index, data-page pruning derives page ranges from decoded page headers. for (auto& row_group : metadata.row_groups) { for (auto& column : row_group.columns) { column.offset_index.reset(); + ASSERT_FALSE(column.offset_index.has_value()); } } auto no_index_reader = cudf::io::parquet::experimental::hybrid_scan_reader(metadata, options); @@ -932,6 +934,55 @@ TEST_F(HybridScanFiltersTest, OffsetIndexOnlyDataPageMask) CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), no_index_result.tbl->view()); } +TEST_F(HybridScanFiltersTest, NoOffsetIndexListColumns) +{ + // List pages cannot safely be mapped to rows from page headers alone, because a leaf page may + // begin in the middle of a list row. Verify the header-derived fallback retains correct output. + using T = uint32_t; + std::mt19937 gen(0xc0c0a); + auto list_col = make_list_str_column(gen, false, false); + auto scalar_col = testdata::ascending(); + auto list_table = cudf::table_view{{scalar_col, *list_col}}; + auto list_buffer = std::vector{}; + auto list_writer = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&list_buffer}, list_table) + .row_group_size_rows(num_ordered_rows) + .max_page_size_rows(page_size_for_ordered_tests) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + cudf::io::write_parquet(list_writer); + + auto const list_datasource = cudf::io::datasource::create(cudf::host_span( + reinterpret_cast(list_buffer.data()), list_buffer.size())); + auto const list_footer = cudf::io::parquet::fetch_footer_to_host(*list_datasource); + auto options = cudf::io::parquet_reader_options::builder().build(); + auto list_reader = cudf::io::parquet::experimental::hybrid_scan_reader(*list_footer, options); + auto const list_row_groups = list_reader.all_row_groups(options); + auto const list_row_mask_values = cudf::detail::make_counting_transform_iterator( + 0, [](auto const row) { return std::cmp_less(row, num_ordered_rows / 2); }); + auto list_row_mask = cudf::test::fixed_width_column_wrapper( + list_row_mask_values, list_row_mask_values + num_ordered_rows); + auto const list_row_mask_view = static_cast(list_row_mask); + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + auto const list_byte_ranges = + list_reader.payload_column_chunks_byte_ranges(list_row_groups, options); + auto [list_buffers, list_data, list_tasks] = cudf::io::parquet::fetch_byte_ranges_to_device_async( + *list_datasource, list_byte_ranges, stream, mr); + list_tasks.get(); + + auto const list_result = list_reader.materialize_payload_columns( + list_row_groups, + list_data, + list_row_mask_view, + cudf::io::parquet::experimental::use_data_page_mask::YES, + options, + stream, + mr); + auto const list_expected = cudf::apply_boolean_mask(list_table, list_row_mask_view, stream, mr); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(list_expected->view(), list_result.tbl->view()); +} + template struct TimestampPageFiltering : public HybridScanFiltersTest {}; diff --git a/java/src/main/java/ai/rapids/cudf/HybridScanReader.java b/java/src/main/java/ai/rapids/cudf/HybridScanReader.java index b783aa86cff4..ef8a31b6633c 100644 --- a/java/src/main/java/ai/rapids/cudf/HybridScanReader.java +++ b/java/src/main/java/ai/rapids/cudf/HybridScanReader.java @@ -39,10 +39,15 @@ * chunked reader pipeline. * *

The filter and payload materialization paths accept a boolean that toggles - * page-level pruning: skips decode of pages the filter (or row mask) proves empty, in - * exchange for a per-page stats scan and a carried row-mask column. Enable when the - * workload prunes many pages; requires prior {@link #setupPageIndex(HostMemoryBuffer)} to - * prune filter column pages using page-level statistics. + * page-level pruning: skips decode of pages the filter (or row mask) proves empty. Enable when the + * workload prunes many pages. Pruning requirements for the two column materializations differ: + *

    + *
  • With page-level pruning enabled, filter columns require + * {@link #setupPageIndex(HostMemoryBuffer)} to seed the row mask from page-index + * statistics.
    • + *
  • Once a row mask exists, filter and payload columns use page row boundaries from the + * {@code OffsetIndex} when available, otherwise from decoded page headers.
    • + *
* *

The reader is created with no filter expression installed. Filter-related APIs * behave as though nothing has been filtered out unless a filter is first supplied via @@ -197,8 +202,8 @@ public ByteRange pageIndexByteRange() { /** * Materialize the {@code ColumnIndex} / {@code OffsetIndex} structs (collectively, the page - * index) from the supplied bytes. Required before any filter or payload column materialization - * call with {@code usePageLevelPruning == true}. + * index) from the supplied bytes. Required before any filter column materialization with + * {@code usePageLevelPruning == true}. * * @param pageIndexBuffer host-resident page index bytes */ @@ -382,8 +387,7 @@ public FilterMaterializationResult materializeFilterColumns(int[] rowGroupIndice * returned by {@link #payloadColumnChunksByteRanges(int[])} * @param rowMask row mask (read-only) * @param usePageLevelPruning enable the data page mask to skip decode of pages the row - * mask proves empty; requires prior - * {@link #setupPageIndex(HostMemoryBuffer)} to avoid fall back path + * mask proves empty * @return the materialized payload column table */ public Table materializePayloadColumns(int[] rowGroupIndices, @@ -529,8 +533,7 @@ public ColumnVector takeFilterRowMask() { * @param rowGroupIndices row groups to read * @param rowMask row mask (read-only) * @param usePageLevelPruning enable the data page mask to skip decode of pages the row - * mask proves empty; requires prior - * {@link #setupPageIndex(HostMemoryBuffer)} to avoid fall back path + * mask proves empty * @param columnChunkData device buffers holding the payload column chunks, in the order * returned by {@link #payloadColumnChunksByteRanges(int[])} */ diff --git a/java/src/main/native/include/hybrid_scan_jni_internal.hpp b/java/src/main/native/include/hybrid_scan_jni_internal.hpp index b9c4ed2850ef..b4e986c027f0 100644 --- a/java/src/main/native/include/hybrid_scan_jni_internal.hpp +++ b/java/src/main/native/include/hybrid_scan_jni_internal.hpp @@ -74,7 +74,6 @@ struct row_group_span_holder { */ cudf::io::parquet_reader_options build_options(JNIEnv* env, jobjectArray j_column_names, - jbooleanArray j_read_binary_as_string, jint time_unit_type_id); row_group_span_holder make_row_group_span(JNIEnv* env, jintArray j_row_groups); diff --git a/java/src/main/native/src/HybridScanReaderJni.cpp b/java/src/main/native/src/HybridScanReaderJni.cpp index d73ba3493143..06830509cc50 100644 --- a/java/src/main/native/src/HybridScanReaderJni.cpp +++ b/java/src/main/native/src/HybridScanReaderJni.cpp @@ -41,7 +41,8 @@ Java_ai_rapids_cudf_HybridScanReader_createFromFooter(JNIEnv* env, { cudf::jni::auto_set_device(env); auto const len = checked_size_t(env, footer_length, "footerLength"); - auto opts = build_options(env, j_column_names, j_binary_as_str, time_unit_type_id); + (void)j_binary_as_str; + auto opts = build_options(env, j_column_names, time_unit_type_id); auto const* footer_ptr = reinterpret_cast(footer_address); cudf::host_span footer_bytes{footer_ptr, len}; auto wrapper = std::make_unique(footer_bytes, std::move(opts)); diff --git a/java/src/main/native/src/HybridScanReaderJniInternal.cpp b/java/src/main/native/src/HybridScanReaderJniInternal.cpp index e7974d2a650c..09fa5e9e4ac0 100644 --- a/java/src/main/native/src/HybridScanReaderJniInternal.cpp +++ b/java/src/main/native/src/HybridScanReaderJniInternal.cpp @@ -16,13 +16,8 @@ namespace hybrid_scan { cudf::io::parquet_reader_options build_options(JNIEnv* env, jobjectArray j_column_names, - jbooleanArray j_read_binary_as_string, jint time_unit_type_id) { - // The hybrid_scan_reader's options builder is constructed without a source_info because - // the reader works on already-parsed footer bytes (and on byte ranges fetched separately). - // Filter is not installed here; use HybridScanReader.setFilter (JNI setFilter) after - // construction. cudf::io::parquet_reader_options_builder builder; cudf::jni::native_jstringArray names(env, j_column_names); @@ -30,24 +25,6 @@ cudf::io::parquet_reader_options build_options(JNIEnv* env, builder = builder.column_names(names.as_cpp_vector()); } - // Translate Java's per-column "read binary as string" flags into the C++ schema override - // hooks. The reader_column_schema mechanism lets callers force binary→string conversion - // for the i-th projected column. - cudf::jni::native_jbooleanArray binary_as_str(env, j_read_binary_as_string); - if (!binary_as_str.is_null() && binary_as_str.size() > 0) { - std::vector schemas; - schemas.reserve(binary_as_str.size()); - for (int i = 0; i < binary_as_str.size(); ++i) { - cudf::io::reader_column_schema s; - s.set_convert_binary_to_strings(static_cast(binary_as_str[i])); - schemas.emplace_back(std::move(s)); - } - builder = builder.set_column_schema(std::move(schemas)); - binary_as_str.cancel(); - } - - // convert_strings_to_categories and ignore_missing_columns are fixed to match the - // standard cudf-java Parquet reader (see readParquet in TableJni.cpp). return builder.convert_strings_to_categories(false) .timestamp_type(cudf::data_type(static_cast(time_unit_type_id))) .ignore_missing_columns(true) diff --git a/java/src/main/native/src/HybridScanReaderJniMaterialize.cpp b/java/src/main/native/src/HybridScanReaderJniMaterialize.cpp index f6f671ae4e7d..172eb5acb947 100644 --- a/java/src/main/native/src/HybridScanReaderJniMaterialize.cpp +++ b/java/src/main/native/src/HybridScanReaderJniMaterialize.cpp @@ -183,7 +183,7 @@ Java_ai_rapids_cudf_HybridScanReader_setupChunkingForFilterColumns(JNIEnv* env, wrapper->reader->setup_chunking_for_filter_columns(chunk_limit, pass_limit, holder.span(), - row_mask_col->view(), + row_mask_col->mutable_view(), mode, spans, wrapper->options, diff --git a/java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java b/java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java index 5eed87e1e0e0..a36282e6b2ee 100644 --- a/java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java +++ b/java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java @@ -28,6 +28,8 @@ import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; import java.util.function.Consumer; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -38,6 +40,8 @@ public class HybridScanReaderTest extends CudfTestBase { private static final String[] DEFAULT_COLS = {"id", "zip_code", "num_units"}; + private static final HostColumnVector.ListType LIST_OF_INTS = + new HostColumnVector.ListType(true, new HostColumnVector.BasicType(true, DType.INT32)); private static final int[] ALL_ROW_GROUPS = {0, 1, 2}; // -------------------------------------------------------------------- @@ -593,6 +597,42 @@ void testMaterializePayloadColumnsExactRowCount(@TempDir Path tmp) throws IOExce } } + /** + * Verifies materializePayloadColumns() prunes pages from page-header row counts when the + * file has no page index: zip_code > 100,000 keeps only the last of row group 1's three + * pages, so the payload must be exactly the 19,999 rows with ids 100,001–119,999. + */ + @Test + void testMaterializePayloadColumnsPagePruningWithoutPageIndex(@TempDir Path tmp) + throws IOException { + try (OpenReader open = + OpenReader.multiPage(tmp).withFilter("zip_code", BinaryOperator.GREATER, 100000)) { + HybridScanReader reader = open.reader; + assertEquals(0L, reader.pageIndexByteRange().size(), + "Fixture must have no page index so the header-derived fallback is exercised"); + int[] survived = reader.filterRowGroupsWithStats(reader.allRowGroups()); + assertArrayEquals(new int[]{1}, survived, + "Group 0 (zip_code 0-59,999) cannot satisfy zip_code > 100,000"); + DeviceMemoryBuffer[] filterCols = copyRangesToDevice( + open.file, reader.filterColumnChunksByteRanges(survived)); + DeviceMemoryBuffer[] payloadCols = copyRangesToDevice( + open.file, reader.payloadColumnChunksByteRanges(survived)); + try (HybridScanReader.FilterMaterializationResult fr = + reader.materializeFilterColumns(survived, filterCols, false); + Table payload = reader.materializePayloadColumns(survived, payloadCols, + fr.rowMask(), true); + ColumnVector expectedIds = ColumnVector.fromInts( + IntStream.rangeClosed(100001, 119999).toArray())) { + assertEquals(2, payload.getNumberOfColumns(), "payload table contains id + num_units"); + assertEquals(19999L, payload.getRowCount()); + AssertUtils.assertColumnsAreEqual(expectedIds, payload.getColumn(0), "id"); + } finally { + closeAll(filterCols); + closeAll(payloadCols); + } + } + } + // -------------------------------------------------------------------- // Tests: materializeAllColumns() // -------------------------------------------------------------------- @@ -825,6 +865,92 @@ void testMaterializePayloadColumnsChunkExactTotal(@TempDir Path tmp) throws IOEx } } + /** + * Verifies the chunked payload pipeline drains the same 19,999 rows when page pruning + * falls back to page-header row counts on a file with no page index. + */ + @Test + void testMaterializePayloadColumnsChunkPagePruningWithoutPageIndex(@TempDir Path tmp) + throws IOException { + try (OpenReader open = + OpenReader.multiPage(tmp).withFilter("zip_code", BinaryOperator.GREATER, 100000)) { + HybridScanReader reader = open.reader; + assertEquals(0L, reader.pageIndexByteRange().size(), + "Fixture must have no page index so the header-derived fallback is exercised"); + int[] survived = reader.filterRowGroupsWithStats(reader.allRowGroups()); + DeviceMemoryBuffer[] filterCols = copyRangesToDevice( + open.file, reader.filterColumnChunksByteRanges(survived)); + DeviceMemoryBuffer[] payloadCols = copyRangesToDevice( + open.file, reader.payloadColumnChunksByteRanges(survived)); + try { + reader.setupChunkingForFilterColumns(0L, 0L, survived, false, filterCols); + while (reader.hasNextTableChunk()) { + reader.materializeFilterColumnsChunk().close(); + } + try (ColumnVector rowMask = reader.takeFilterRowMask()) { + assertEquals(60000L, rowMask.getRowCount(), "Mask spans row group 1"); + assertEquals(19999L, countTrue(rowMask), "zip_code 100,001-119,999 survive"); + reader.setupChunkingForPayloadColumns(0L, 0L, survived, rowMask, true, payloadCols); + long total = 0; + while (reader.hasNextTableChunk()) { + try (Table chunk = reader.materializePayloadColumnsChunk(rowMask)) { + assertEquals(2, chunk.getNumberOfColumns()); + total += chunk.getRowCount(); + } + } + assertEquals(19999L, total, + "Header-derived page pruning must not drop or duplicate selected rows"); + } + } finally { + closeAll(filterCols); + closeAll(payloadCols); + } + } + } + + /** + * Verifies that list payload columns remain readable when header-derived page pruning is used. + * List leaf pages may start mid-row, so they must not be pruned without an offset index. + */ + @Test + void testMaterializePayloadColumnsChunkListPagePruningWithoutPageIndex(@TempDir Path tmp) + throws IOException { + try (OpenReader open = OpenReader.multiPageWithList(tmp) + .withFilter("zip_code", BinaryOperator.GREATER, 30000)) { + HybridScanReader reader = open.reader; + assertEquals(0L, reader.pageIndexByteRange().size(), + "Fixture must have no page index so the header-derived fallback is exercised"); + int[] survived = reader.filterRowGroupsWithStats(reader.allRowGroups()); + DeviceMemoryBuffer[] filterCols = copyRangesToDevice( + open.file, reader.filterColumnChunksByteRanges(survived)); + DeviceMemoryBuffer[] payloadCols = copyRangesToDevice( + open.file, reader.payloadColumnChunksByteRanges(survived)); + try { + reader.setupChunkingForFilterColumns(0L, 0L, survived, false, filterCols); + while (reader.hasNextTableChunk()) { + reader.materializeFilterColumnsChunk().close(); + } + try (ColumnVector rowMask = reader.takeFilterRowMask()) { + assertEquals(60000L, rowMask.getRowCount(), "Mask spans the row group"); + assertEquals(29999L, countTrue(rowMask), "zip_code 30,001-59,999 survive"); + reader.setupChunkingForPayloadColumns(0L, 0L, survived, rowMask, true, payloadCols); + long total = 0; + while (reader.hasNextTableChunk()) { + try (Table chunk = reader.materializePayloadColumnsChunk(rowMask)) { + assertEquals(2, chunk.getNumberOfColumns()); + total += chunk.getRowCount(); + } + } + assertEquals(29999L, total, + "Header-derived pruning must retain all selected rows with list payloads"); + } + } finally { + closeAll(filterCols); + closeAll(payloadCols); + } + } + } + // -------------------------------------------------------------------- // Tests: setupChunkingForAllColumns() / materializeAllColumnsChunk() // @@ -1212,13 +1338,30 @@ static OpenReader pageIndex(Path tmp) throws IOException { return openFromFile(pq, DEFAULT_COLS); } + /** A single 100-row group, small enough that every column chunk holds one data page. */ static OpenReader rowGroupStats(Path tmp) throws IOException { File pq = tmp.resolve("fixture.parquet").toFile(); - writeRowGroupStatsParquet(pq); + writeNoPageIndexParquet(pq, 100, 1, + ParquetWriterOptions.StatisticsFrequency.ROWGROUP); + return openFromFile(pq, DEFAULT_COLS); + } + + /** Two 60,000-row groups, so each column chunk spans 3 data pages (20,000 rows each). */ + static OpenReader multiPage(Path tmp) throws IOException { + File pq = tmp.resolve("fixture.parquet").toFile(); + writeNoPageIndexParquet(pq, 60_000, 2, + ParquetWriterOptions.StatisticsFrequency.PAGE); return openFromFile(pq, DEFAULT_COLS); } - private static OpenReader openFromFile(File pq, String[] cols) throws IOException { + /** A 60,000-row group with list payloads spanning multiple leaf pages. */ + static OpenReader multiPageWithList(Path tmp) throws IOException { + File pq = tmp.resolve("fixture.parquet").toFile(); + writeNoPageIndexListParquet(pq); + return openFromFile(pq, "id", "zip_code", "list_values"); + } + + private static OpenReader openFromFile(File pq, String... cols) throws IOException { HostMemoryBuffer file = readFileToHostBuffer(pq); HostMemoryBuffer footer = null; HybridScanReader reader = null; @@ -1323,28 +1466,62 @@ private static int writeFixtureParquet(File path) { } /** - * Writes a small Parquet file with {@code ROWGROUP}-level statistics: row-group min/max - * are recorded but no page index (no {@code ColumnIndex}/{@code OffsetIndex}) is emitted. - * Includes a low-cardinality {@code num_units} column ({1, 2, 3} cycle) so the writer's - * ADAPTIVE dictionary policy emits a dictionary; this lets tests exercise the - * "no page index, dict exists" path (see - * {@link #testSecondaryFiltersByteRangesEmptyForRowGroupStats}). + * Writes a Parquet file with no page index; only {@code COLUMN} statistics emit one. Both + * {@code id} and {@code zip_code} hold the globally sequential row index, and + * {@code num_units} cycles over {1, 2, 3}, low-cardinality enough that the ADAPTIVE + * dictionary policy emits a dictionary (see + * {@link #testSecondaryFiltersByteRangesEmptyForRowGroupStats}). The writer caps a page at + * 20,000 rows, so exceed that per group for chunks spanning several pages. */ - private static void writeRowGroupStatsParquet(File path) { - int rows = 100; + private static void writeNoPageIndexParquet(File path, int rowsPerGroup, int numGroups, + ParquetWriterOptions.StatisticsFrequency stats) { ParquetWriterOptions opts = ParquetWriterOptions.builder() .withNonNullableColumns("id", "zip_code", "num_units") + .withRowGroupSizeRows(rowsPerGroup) + .withStatisticsFrequency(stats) + .build(); + try (TableWriter writer = Table.writeParquetChunked(opts, path)) { + for (int g = 0; g < numGroups; g++) { + int start = g * rowsPerGroup; + try (ColumnVector id = ColumnVector.fromInts( + IntStream.range(start, start + rowsPerGroup).toArray()); + ColumnVector zipCode = ColumnVector.fromInts( + IntStream.range(start, start + rowsPerGroup).toArray()); + ColumnVector numUnits = ColumnVector.fromInts( + IntStream.range(start, start + rowsPerGroup) + .map(i -> 1 + (i % 3)).toArray()); + Table t = new Table(id, zipCode, numUnits)) { + writer.write(t); + } + } + } + } + + /** + * Writes one 60,000-row group with a list payload column. Each row holds three values, so + * leaf-page boundaries can fall within a logical row. + */ + @SuppressWarnings("unchecked") + private static void writeNoPageIndexListParquet(File path) { + int rows = 60_000; + ParquetWriterOptions opts = ParquetWriterOptions.builder() + .withNonNullableColumns("id", "zip_code") + .withListColumn(ColumnWriterOptions.listBuilder("list_values", false) + .withNonNullableColumns("element") + .build()) .withRowGroupSizeRows(rows) - .withStatisticsFrequency(ParquetWriterOptions.StatisticsFrequency.ROWGROUP) + .withStatisticsFrequency(ParquetWriterOptions.StatisticsFrequency.PAGE) .build(); - try (TableWriter writer = Table.writeParquetChunked(opts, path); - ColumnVector id = ColumnVector.fromInts(IntStream.range(0, rows).toArray()); - ColumnVector zipCode = ColumnVector.fromInts( - IntStream.range(0, rows).map(i -> 10000 + i).toArray()); - ColumnVector numUnits = ColumnVector.fromInts( - IntStream.range(0, rows).map(i -> 1 + (i % 3)).toArray()); - Table t = new Table(id, zipCode, numUnits)) { - writer.write(t); + List[] listRows = (List[]) new List[rows]; + for (int i = 0; i < rows; i++) { + listRows[i] = Arrays.asList(i * 3, i * 3 + 1, i * 3 + 2); + } + try (ColumnVector id = ColumnVector.fromInts(IntStream.range(0, rows).toArray()); + ColumnVector zipCode = ColumnVector.fromInts(IntStream.range(0, rows).toArray()); + ColumnVector listValues = ColumnVector.fromLists(LIST_OF_INTS, listRows); + Table table = new Table(id, zipCode, listValues); + TableWriter writer = Table.writeParquetChunked(opts, path)) { + writer.write(table); } } diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index 0d092541d84b..336528b06ffa 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx @@ -831,7 +831,7 @@ cdef class HybridScanReader: row_group_indices : list[int] Input row group indices row_mask : Column - Boolean column indicating surviving rows + Mutable boolean column indicating surviving rows mask_data_pages : UseDataPageMask Whether to use a data page mask column_chunk_data : Sequence @@ -854,7 +854,7 @@ cdef class HybridScanReader: # keep reference to avoid use-after-free of device spans self._filter_chunk_data = column_chunk_data - cdef column_view mask_view = row_mask.view() + cdef mutable_column_view mask_view = row_mask.mutable_view() with nogil: self.c_obj.get()[0].setup_chunking_for_filter_columns( chunk_read_limit, diff --git a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd index 7a5aec269a56..04d84fb68a33 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd @@ -161,7 +161,7 @@ cdef extern from "cudf/io/experimental/hybrid_scan.hpp" \ size_t chunk_read_limit, size_t pass_read_limit, std_span[const_size_type] row_group_indices, - const column_view& row_mask, + const mutable_column_view& row_mask, use_data_page_mask mask_data_pages, std_span[const_device_span_const_uint8_t] column_chunk_data, const parquet_reader_options& options, diff --git a/python/pylibcudf/tests/io/test_experimental_hybrid_scan.py b/python/pylibcudf/tests/io/test_experimental_hybrid_scan.py index dd553d111971..73352fc80c5d 100644 --- a/python/pylibcudf/tests/io/test_experimental_hybrid_scan.py +++ b/python/pylibcudf/tests/io/test_experimental_hybrid_scan.py @@ -430,6 +430,73 @@ def test_hybrid_scan_materialize_columns( assert expected_arrow.equals(hybrid_arrow) +def test_hybrid_scan_payload_page_mask_without_page_index( + simple_parquet_bytes: bytes, + simple_hybrid_scan_reader: HybridScanReader, + simple_parquet_options: plc.io.parquet.ParquetReaderOptions, + simple_parquet_table: pa.Table, + num_rows: int, +) -> None: + """Test payload page pruning without a page index set up on the reader.""" + reader = simple_hybrid_scan_reader + row_groups = reader.all_row_groups(simple_parquet_options) + + # Keep the first half of the rows so the trailing data pages get pruned. + num_selected = num_rows // 2 + row_mask = plc.Column.from_arrow( + pa.array([i < num_selected for i in range(num_rows)], type=pa.bool_()) + ) + + payload_data = [ + plc.gpumemoryview( + rmm.DeviceBuffer.to_device( + simple_parquet_bytes[r.offset : r.offset + r.size], + plc.utils._get_stream(), + ) + ) + for r in reader.payload_column_chunks_byte_ranges( + row_groups, simple_parquet_options + ) + ] + synchronize_stream() + + # Chunks can disagree on field nullability, so compare row values only. + def to_rows(tbl: plc.Table) -> list: + return ( + tbl.to_arrow() + .rename_columns(simple_parquet_table.column_names) + .to_pylist() + ) + + expected_rows = simple_parquet_table.slice(0, num_selected).to_pylist() + + payload_result = reader.materialize_payload_columns( + row_groups, + payload_data, + row_mask, + UseDataPageMask.YES, + simple_parquet_options, + ) + synchronize_stream() + assert to_rows(payload_result.tbl) == expected_rows + + reader.setup_chunking_for_payload_columns( + 256, + 0, + row_groups, + row_mask, + UseDataPageMask.YES, + payload_data, + simple_parquet_options, + ) + chunked_rows = [] + while reader.has_next_table_chunk(): + chunk = reader.materialize_payload_columns_chunk(row_mask) + chunked_rows.extend(to_rows(chunk.tbl)) + synchronize_stream() + assert chunked_rows == expected_rows + + @pytest.mark.parametrize("stream", [None, Stream()]) def test_hybrid_scan_single_step_materialize( simple_parquet_bytes: bytes,