diff --git a/cpp/examples/hybrid_scan_io/CMakeLists.txt b/cpp/examples/hybrid_scan_io/CMakeLists.txt index b0d12692fa5e..90fae1005891 100644 --- a/cpp/examples/hybrid_scan_io/CMakeLists.txt +++ b/cpp/examples/hybrid_scan_io/CMakeLists.txt @@ -50,6 +50,7 @@ add_hybrid_scan_example(hybrid_scan_io hybrid_scan_io.cpp) add_hybrid_scan_example(hybrid_scan_pipeline hybrid_scan_pipeline.cpp) add_hybrid_scan_example(hybrid_scan_multifile_single_step hybrid_scan_multifile_single_step.cpp) add_hybrid_scan_example(hybrid_scan_multifile_two_step hybrid_scan_multifile_two_step.cpp) +add_hybrid_scan_example(mint1t_hybrid_scan mint1t_hybrid_scan.cpp) # Install the example.parquet file install(FILES ${CMAKE_CURRENT_LIST_DIR}/example.parquet diff --git a/cpp/examples/hybrid_scan_io/io_utils.cpp b/cpp/examples/hybrid_scan_io/io_utils.cpp index f21bb5a47645..fec6f43ba930 100644 --- a/cpp/examples/hybrid_scan_io/io_utils.cpp +++ b/cpp/examples/hybrid_scan_io/io_utils.cpp @@ -3,6 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "io_utils.hpp" + #include #include #include @@ -39,3 +41,73 @@ fetch_byte_ranges_async(cudf::io::datasource& datasource, // Using libcudf utility but may have custom implementation in the future return cudf::io::parquet::fetch_byte_ranges_to_device_async(datasource, byte_ranges, stream, mr); } + +multifile_inputs::multifile_inputs(cudf::io::source_info const& source_info) + : datasources{cudf::io::make_datasources(source_info)} +{ + datasource_refs.reserve(datasources.size()); + std::transform(datasources.begin(), + datasources.end(), + std::back_inserter(datasource_refs), + [](auto const& datasource) { return std::ref(*datasource); }); +} + +void multifile_inputs::fetch_footers() +{ + footer_buffers = cudf::io::parquet::fetch_footers_to_host(datasource_refs); + footer_byte_spans.clear(); + footer_byte_spans.reserve(footer_buffers.size()); + std::transform(footer_buffers.begin(), + footer_buffers.end(), + std::back_inserter(footer_byte_spans), + [](auto const& buffer) { return cudf::host_span{*buffer}; }); +} + +std::vector> group_byte_ranges_by_source( + std::pair, std::vector> const& + byte_ranges_and_source_map, + std::size_t num_sources) +{ + auto const& [byte_ranges, source_map] = byte_ranges_and_source_map; + CUDF_EXPECTS(byte_ranges.size() == source_map.size(), "Invalid source map size"); + + auto byte_ranges_per_source = + std::vector>(num_sources); + for (auto range_index = std::size_t{0}; range_index < byte_ranges.size(); ++range_index) { + auto const source_index = source_map[range_index]; + CUDF_EXPECTS( + source_index >= 0 and static_cast(source_index) < byte_ranges_per_source.size(), + "Invalid source index"); + byte_ranges_per_source[source_index].push_back(byte_ranges[range_index]); + } + return byte_ranges_per_source; +} + +multisource_device_data fetch_multisource_device_data( + multifile_inputs const& inputs, + std::pair, std::vector> const& + byte_ranges_and_source_map, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto const byte_ranges_per_source = + group_byte_ranges_by_source(byte_ranges_and_source_map, inputs.datasources.size()); + return fetch_multisource_device_data(inputs, byte_ranges_per_source, stream, mr); +} + +multisource_device_data fetch_multisource_device_data( + multifile_inputs const& inputs, + std::vector> const& byte_ranges_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto [buffers, per_source_spans, tasks] = cudf::io::parquet::fetch_byte_ranges_to_device_async( + inputs.datasource_refs, byte_ranges_per_source, stream, mr); + tasks.get(); + + auto flat_spans = std::vector>{}; + for (auto const& source_spans : per_source_spans) { + flat_spans.insert(flat_spans.end(), source_spans.begin(), source_spans.end()); + } + return {std::move(buffers), std::move(per_source_spans), std::move(flat_spans)}; +} diff --git a/cpp/examples/hybrid_scan_io/io_utils.hpp b/cpp/examples/hybrid_scan_io/io_utils.hpp index 3a977b16616b..5814a416ee87 100644 --- a/cpp/examples/hybrid_scan_io/io_utils.hpp +++ b/cpp/examples/hybrid_scan_io/io_utils.hpp @@ -57,3 +57,56 @@ fetch_byte_ranges_async(cudf::io::datasource& datasource, cudf::host_span byte_ranges, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); + +/** + * @brief Owns the datasources, footer buffers, and byte spans for a multifile read. + * + * Call `fetch_footers()` after construction. Keeping datasource and footer setup separate allows + * examples to time those operations independently. + */ +struct multifile_inputs { + explicit multifile_inputs(cudf::io::source_info const& source_info); + + void fetch_footers(); + + std::vector> datasources; + std::vector> datasource_refs; + std::vector> footer_buffers; + std::vector> footer_byte_spans; +}; + +/** + * @brief Owns multifile device buffers and the corresponding per-source and flattened spans. + */ +struct multisource_device_data { + std::vector buffers; + std::vector>> per_source_spans; + std::vector> flat_spans; +}; + +/** + * @brief Regroups flattened byte ranges using the source map returned by Hybrid Scan. + */ +[[nodiscard]] std::vector> group_byte_ranges_by_source( + std::pair, std::vector> const& + byte_ranges_and_source_map, + std::size_t num_sources); + +/** + * @brief Fetches source-mapped multifile byte ranges and flattens the resulting device spans. + */ +[[nodiscard]] multisource_device_data fetch_multisource_device_data( + multifile_inputs const& inputs, + std::pair, std::vector> const& + byte_ranges_and_source_map, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/** + * @brief Fetches source-grouped multifile byte ranges. + */ +[[nodiscard]] multisource_device_data fetch_multisource_device_data( + multifile_inputs const& inputs, + std::vector> const& byte_ranges_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); diff --git a/cpp/examples/hybrid_scan_io/mint1t_hybrid_scan.cpp b/cpp/examples/hybrid_scan_io/mint1t_hybrid_scan.cpp new file mode 100644 index 000000000000..5ea368ffb5da --- /dev/null +++ b/cpp/examples/hybrid_scan_io/mint1t_hybrid_scan.cpp @@ -0,0 +1,832 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "common_utils.hpp" +#include "io_utils.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using clock_type = std::chrono::steady_clock; +using cudf::io::parquet::experimental::hybrid_scan_multifile; +using cudf::io::parquet::experimental::use_data_page_mask; + +auto constexpr original_dataset_root = "/MINT1T_benchmarking_subset"; +auto constexpr rewritten_dataset_root = "/MINT1T_rewritten"; +auto constexpr default_refs = + "/MINT1T_benchmarking_subset/parquet/interleaved_image_url_refs/image_url_refs.parquet"; +auto constexpr default_image_dir = + "/MINT1T_rewritten/parquet/mint_1t_html_images_stable_row_ids/data"; +auto constexpr default_expected = "/MINT1T_rewritten/output/image_url_payloads.parquet"; +auto constexpr writable_root = "/MINT1T_rewritten"; + +std::vector const payload_columns{ + "image", "image_format", "mime_type", "image_size_bytes", "md5", "sha256", "width", "height"}; + +struct arguments { + std::filesystem::path dataset_root{rewritten_dataset_root}; + std::filesystem::path refs{default_refs}; + std::filesystem::path image_dir{default_image_dir}; + std::optional limit; + std::size_t pass_read_limit{}; + bool pass_read_limit_set{false}; + bool use_page_mask{true}; + bool use_sparse_page_io{true}; + bool drop_cache{false}; + bool validate{false}; + std::filesystem::path expected_output{default_expected}; + std::optional output; +}; + +struct reference { + std::string image_parquet; + uint32_t row_offset; + + bool operator<(reference const& other) const + { + return std::tie(image_parquet, row_offset) < std::tie(other.image_parquet, other.row_offset); + } + + bool operator==(reference const& other) const + { + return image_parquet == other.image_parquet and row_offset == other.row_offset; + } +}; + +struct timing_data { + std::vector> stages; + + void add_seconds(std::string name, double seconds) + { + std::cout << "TIMING " << std::left << std::setw(38) << name << std::right << std::fixed + << std::setprecision(6) << seconds << " s\n" + << std::flush; + stages.emplace_back(std::move(name), seconds); + } + + void add(std::string name, clock_type::time_point start) + { + auto const seconds = std::chrono::duration(clock_type::now() - start).count(); + auto const existing = std::find_if( + stages.begin(), stages.end(), [&](auto const& item) { return item.first == name; }); + if (existing == stages.end()) { + add_seconds(std::move(name), seconds); + } else { + existing->second += seconds; + } + } + + [[nodiscard]] double get(std::string_view name) const + { + auto const it = std::find_if( + stages.begin(), stages.end(), [&](auto const& item) { return item.first == name; }); + return it == stages.end() ? 0.0 : it->second; + } + + [[nodiscard]] double sum(std::span names) const + { + return std::accumulate( + names.begin(), names.end(), 0.0, [&](double total, auto name) { return total + get(name); }); + } + + void print() const + { + std::cout << "\nStage timings:\n"; + for (auto const& [name, seconds] : stages) { + std::cout << " " << std::left << std::setw(38) << name << std::right << std::fixed + << std::setprecision(6) << seconds << " s\n"; + } + } +}; + +struct cache_drop_result { + bool attempted{}; + std::size_t files{}; + std::size_t errors{}; +}; + +void print_usage() +{ + std::cout + << "Usage: mint1t_hybrid_scan [options]\n\n" + << " --dataset-root PATH Dataset root; sets the default image-shard directory\n" + << " --refs PATH Reference Parquet (default: original manifest)\n" + << " --image-parquet-dir PATH Directory containing image Parquet shards\n" + << " --limit N Use the first N manifest rows\n" + << " --pass-read-limit-mib N Per-pass limit; default is 105% of GPU memory, 0 unbounded\n" + << " --use-data-page-mask YES|NO Toggle payload data-page pruning (default YES)\n" + << " --use-sparse-page-io YES|NO Toggle sparse physical payload I/O (default YES)\n" + << " --best-effort-drop-cache Apply POSIX_FADV_DONTNEED before timed source setup\n" + << " --validate Compare with the projected Python output\n" + << " --expected-output PATH Validation Parquet used by --validate\n" + << " --output PATH Write the payload-only C++ result under /MINT1T_rewritten\n" + << " -h, --help Show this message\n"; +} + +arguments parse_args(int argc, char const** argv) +{ + arguments args; + auto image_dir_set = false; + auto dataset_root_set = false; + auto expected_output_set = false; + auto require_value = [&](int& index, std::string_view option) -> std::string { + if (++index >= argc) { + throw std::invalid_argument("Missing value for " + std::string{option}); + } + return argv[index]; + }; + + for (int index = 1; index < argc; ++index) { + auto const option = std::string_view{argv[index]}; + if (option == "-h" or option == "--help") { + print_usage(); + std::exit(0); + } else if (option == "--dataset-root") { + args.dataset_root = require_value(index, option); + dataset_root_set = true; + } else if (option == "--refs") { + args.refs = require_value(index, option); + } else if (option == "--image-parquet-dir") { + args.image_dir = require_value(index, option); + image_dir_set = true; + } else if (option == "--limit") { + auto const value = std::stoll(require_value(index, option)); + CUDF_EXPECTS(value >= 0 and value <= std::numeric_limits::max(), + "Invalid --limit"); + args.limit = static_cast(value); + } else if (option == "--pass-read-limit-mib") { + auto const value = std::stoull(require_value(index, option)); + CUDF_EXPECTS(value <= std::numeric_limits::max() / (1024 * 1024), + "Invalid --pass-read-limit-mib"); + args.pass_read_limit = value * 1024 * 1024; + args.pass_read_limit_set = true; + } else if (option == "--use-data-page-mask") { + args.use_page_mask = get_boolean(require_value(index, option)); + } else if (option == "--use-sparse-page-io") { + args.use_sparse_page_io = get_boolean(require_value(index, option)); + } else if (option == "--best-effort-drop-cache") { + args.drop_cache = true; + } else if (option == "--validate") { + args.validate = true; + } else if (option == "--expected-output") { + args.expected_output = require_value(index, option); + args.validate = true; + expected_output_set = true; + } else if (option == "--output") { + args.output = require_value(index, option); + } else { + throw std::invalid_argument("Unknown option: " + std::string{option}); + } + } + if (not image_dir_set) { + args.image_dir = + args.dataset_root / "parquet/mint_1t_html_images_stable_row_ids/data"; + } + if (dataset_root_set and not expected_output_set) { + auto const output_dir = + args.dataset_root == std::filesystem::path{original_dataset_root} + ? std::filesystem::path{writable_root} / "original" + : std::filesystem::path{writable_root} / "output"; + args.expected_output = output_dir / "image_url_payloads.parquet"; + } + return args; +} + +void require_writable_output_path(std::filesystem::path const& path) +{ + auto const output = std::filesystem::absolute(path).lexically_normal(); + auto const root = std::filesystem::path{writable_root}.lexically_normal(); + auto const rel = output.lexically_relative(root); + CUDF_EXPECTS(not rel.empty() and *rel.begin() != "..", + "Output paths must be under " + root.string()); +} + +std::vector refs_to_host(cudf::table_view refs, rmm::cuda_stream_view stream) +{ + CUDF_EXPECTS(refs.num_columns() == 2, "Unexpected reference table column count"); + CUDF_EXPECTS(refs.column(0).null_count() == 0 and refs.column(1).null_count() == 0, + "Null reference fields are not supported"); + + auto const strings = cudf::strings_column_view{refs.column(0)}; + auto host_offsets = std::vector(strings.size() + 1); + CUDF_CUDA_TRY(cudaMemcpyAsync(host_offsets.data(), + strings.offsets().data() + strings.offset(), + host_offsets.size() * sizeof(cudf::size_type), + cudaMemcpyDeviceToHost, + stream.value())); + stream.synchronize(); + + auto const first_char = host_offsets.front(); + auto const chars_size = host_offsets.back() - first_char; + auto host_chars = std::vector(chars_size); + CUDF_CUDA_TRY(cudaMemcpyAsync(host_chars.data(), + strings.chars_begin(stream) + first_char, + host_chars.size(), + cudaMemcpyDeviceToHost, + stream.value())); + + auto host_row_offsets = std::vector(refs.num_rows()); + CUDF_CUDA_TRY(cudaMemcpyAsync(host_row_offsets.data(), + refs.column(1).data() + refs.column(1).offset(), + host_row_offsets.size() * sizeof(uint32_t), + cudaMemcpyDeviceToHost, + stream.value())); + stream.synchronize(); + + auto result = std::vector{}; + result.reserve(refs.num_rows()); + for (cudf::size_type row = 0; row < refs.num_rows(); ++row) { + auto const begin = host_offsets[row] - first_char; + auto const end = host_offsets[row + 1] - first_char; + result.push_back({std::string{host_chars.data() + begin, static_cast(end - begin)}, + host_row_offsets[row]}); + } + return result; +} + +cache_drop_result drop_file_cache(std::vector const& paths) +{ + cache_drop_result result{.attempted = true}; + for (auto const& path : paths) { + auto const fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { + ++result.errors; + continue; + } + if (::posix_fadvise(fd, 0, 0, POSIX_FADV_DONTNEED) == 0) { + ++result.files; + } else { + ++result.errors; + } + ::close(fd); + } + return result; +} + +std::unique_ptr make_device_column(std::span host_data, + cudf::data_type type, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto result = + cudf::make_numeric_column(type, host_data.size(), cudf::mask_state::UNALLOCATED, stream, mr); + CUDF_CUDA_TRY(cudaMemcpyAsync(result->mutable_view().data(), + host_data.data(), + host_data.size(), + cudaMemcpyHostToDevice, + stream.value())); + stream.synchronize(); + return result; +} + +std::unique_ptr make_gather_map(std::span host_data, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto result = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, + host_data.size(), + cudf::mask_state::UNALLOCATED, + stream, + mr); + CUDF_CUDA_TRY(cudaMemcpyAsync(result->mutable_view().data(), + host_data.data(), + host_data.size_bytes(), + cudaMemcpyHostToDevice, + stream.value())); + stream.synchronize(); + return result; +} + +uint64_t payload_bytes(cudf::table_view table, rmm::cuda_stream_view stream) +{ + CUDF_EXPECTS(table.num_columns() == static_cast(payload_columns.size()), + "Unexpected payload table schema"); + auto const& sizes = table.column(3); + auto host_sizes = std::vector(table.num_rows()); + CUDF_CUDA_TRY(cudaMemcpyAsync(host_sizes.data(), + sizes.data() + sizes.offset(), + host_sizes.size() * sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream.value())); + stream.synchronize(); + return std::accumulate(host_sizes.begin(), host_sizes.end(), uint64_t{0}); +} + +void write_output(std::filesystem::path const& path, cudf::table_view table) +{ + require_writable_output_path(path); + std::filesystem::create_directories(path.parent_path()); + auto metadata = cudf::io::table_input_metadata{table}; + for (std::size_t index = 0; index < payload_columns.size(); ++index) { + metadata.column_metadata[index].set_name(payload_columns[index]); + } + metadata.column_metadata.front().set_output_as_binary(true); + auto options = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{path.string()}, table) + .metadata(std::move(metadata)) + .compression(cudf::io::compression_type::ZSTD) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + cudf::io::write_parquet(options); +} + +std::unique_ptr read_expected(std::filesystem::path const& path) +{ + auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info{path.string()}) + .column_names(payload_columns) + .build(); + return std::move(cudf::io::read_parquet(options).tbl); +} + +void print_json(timing_data const& timings, + cache_drop_result const& cache, + bool use_page_mask, + bool use_sparse_page_io, + std::size_t files, + std::size_t row_groups, + cudf::size_type rows, + uint64_t selected_payload_bytes, + uint64_t requested_bytes, + std::size_t output_allocated_bytes, + double benchmark_elapsed, + double end_to_end_elapsed) +{ + auto const payload_mib = selected_payload_bytes / double{1024 * 1024}; + auto const fetch_time = timings.get("payload byte-range fetch"); + auto const decode_time = timings.get("payload materialization/decode"); + + std::cout << "\n{" + << R"("method":"hybrid_scan_multifile_payload_mask",)" + << "\"use_data_page_mask\":" << std::boolalpha << use_page_mask << "," + << "\"use_sparse_page_io\":" << use_sparse_page_io << "," + << "\"elapsed_s\":" << benchmark_elapsed << "," + << "\"end_to_end_elapsed_s\":" << end_to_end_elapsed << "," + << "\"rows\":" << rows << "," + << "\"payload_bytes\":" << selected_payload_bytes << "," + << "\"payload_mib\":" << payload_mib << "," + << "\"rows_per_s\":" << (benchmark_elapsed > 0 ? rows / benchmark_elapsed : 0) << "," + << "\"payload_mib_per_s\":" + << (benchmark_elapsed > 0 ? payload_mib / benchmark_elapsed : 0) << "," + << "\"referenced_files\":" << files << "," + << "\"row_groups_read\":" << row_groups << "," + << "\"compressed_payload_bytes_requested\":" << requested_bytes << "," + << "\"output_allocated_bytes\":" << output_allocated_bytes << "," + << "\"payload_fetch_mib_per_s\":" + << (fetch_time > 0 ? requested_bytes / double{1024 * 1024} / fetch_time : 0) << "," + << "\"materialization_mib_per_s\":" + << (decode_time > 0 ? output_allocated_bytes / double{1024 * 1024} / decode_time : 0) + << "," + << R"("best_effort_cache_drop":{"attempted":)" << cache.attempted + << ",\"files\":" << cache.files << ",\"errors\":" << cache.errors << "}," + << "\"stage_seconds\":{"; + for (std::size_t index = 0; index < timings.stages.size(); ++index) { + if (index != 0) { std::cout << ","; } + std::cout << "\"" << timings.stages[index].first << "\":" << timings.stages[index].second; + } + std::cout << "}}\n"; +} + +} // namespace + +int main(int argc, char const** argv) +{ + auto const program_start = clock_type::now(); + auto timings = timing_data{}; + auto args = parse_args(argc, argv); + if (not args.pass_read_limit_set) { + auto free_bytes = std::size_t{}; + auto total_bytes = std::size_t{}; + CUDF_CUDA_TRY(cudaMemGetInfo(&free_bytes, &total_bytes)); + args.pass_read_limit = total_bytes + total_bytes / 20; + } + auto const stream = cudf::get_default_stream(); + auto resource = create_memory_resource(false); + auto stats_mr = rmm::mr::statistics_resource_adaptor{resource}; + rmm::mr::set_current_device_resource(stats_mr); + + auto start = clock_type::now(); + auto refs_options = + cudf::io::parquet_reader_options::builder(cudf::io::source_info{args.refs.string()}) + .column_names({"image_parquet", "_row_offset"}) + .build(); + if (args.limit.has_value()) { refs_options.set_num_rows(args.limit.value()); } + auto refs_table = std::move(cudf::io::read_parquet(refs_options, stream).tbl); + stream.synchronize(); + timings.add("reference parquet read", start); + + start = clock_type::now(); + auto refs = refs_to_host(refs_table->view(), stream); + std::stable_sort(refs.begin(), refs.end()); + CUDF_EXPECTS(not refs.empty(), "No references to read"); + + auto unique_refs = std::vector{}; + auto gather_map = std::vector{}; + unique_refs.reserve(refs.size()); + gather_map.reserve(refs.size()); + for (auto const& ref : refs) { + if (unique_refs.empty() or not(unique_refs.back() == ref)) { unique_refs.push_back(ref); } + gather_map.push_back(static_cast(unique_refs.size() - 1)); + } + + auto source_names = std::vector{}; + auto source_offsets = std::vector>{}; + for (auto const& ref : unique_refs) { + if (source_names.empty() or source_names.back() != ref.image_parquet) { + source_names.push_back(ref.image_parquet); + source_offsets.emplace_back(); + } + source_offsets.back().push_back(ref.row_offset); + } + auto source_paths = std::vector{}; + auto source_files = std::vector{}; + source_paths.reserve(source_names.size()); + source_files.reserve(source_names.size()); + for (auto const& name : source_names) { + auto const path = args.image_dir / name; + CUDF_EXPECTS(std::filesystem::is_regular_file(path), "Missing source file: " + path.string()); + source_paths.push_back(path); + source_files.push_back(path.string()); + } + timings.add("reference host extraction/sort/group", start); + + auto cache_result = cache_drop_result{}; + start = clock_type::now(); + if (args.drop_cache) { + auto cache_paths = source_paths; + cache_paths.insert(cache_paths.begin(), args.refs); + cache_result = drop_file_cache(cache_paths); + } + timings.add("best-effort cache drop", start); + + auto const benchmark_start = clock_type::now(); + + start = clock_type::now(); + auto inputs = multifile_inputs{cudf::io::source_info{source_files}}; + timings.add("datasource construction", start); + + start = clock_type::now(); + inputs.fetch_footers(); + timings.add("footer fetch", start); + + auto options = cudf::io::parquet_reader_options::builder().column_names(payload_columns).build(); + start = clock_type::now(); + auto reader = hybrid_scan_multifile{inputs.footer_byte_spans, options}; + timings.add("hybrid reader construction", start); + + start = clock_type::now(); + auto const metadatas = reader.parquet_metadatas(); + auto row_groups = std::vector>(source_names.size()); + auto host_row_mask = std::vector{}; + std::size_t selected_row_groups{}; + for (std::size_t source = 0; source < source_names.size(); ++source) { + auto const& metadata = metadatas[source]; + auto const& offsets = source_offsets[source]; + auto offset_index = std::size_t{0}; + auto file_row_start = uint64_t{0}; + for (std::size_t rg = 0; rg < metadata.row_groups.size() and offset_index < offsets.size(); + ++rg) { + auto const rows = metadata.row_groups[rg].num_rows; + CUDF_EXPECTS(rows >= 0, "Negative row count in Parquet metadata"); + auto const rows_unsigned = static_cast(rows); + auto const file_row_stop = file_row_start + rows_unsigned; + if (offsets[offset_index] < file_row_start) { + throw std::logic_error("References are not ordered"); + } + if (offsets[offset_index] < file_row_stop) { + row_groups[source].push_back(static_cast(rg)); + ++selected_row_groups; + auto const mask_start = host_row_mask.size(); + host_row_mask.resize(mask_start + static_cast(rows_unsigned), uint8_t{0}); + while (offset_index < offsets.size() and offsets[offset_index] < file_row_stop) { + auto const local_offset = + static_cast(offsets[offset_index] - file_row_start); + host_row_mask[mask_start + local_offset] = uint8_t{1}; + ++offset_index; + } + } + file_row_start = file_row_stop; + } + CUDF_EXPECTS(offset_index == offsets.size(), + "Out-of-range row offset in " + source_names[source]); + } + CUDF_EXPECTS( + host_row_mask.size() <= static_cast(std::numeric_limits::max()), + "Selected row groups exceed cudf column size limit"); + timings.add("row-group planning/host mask", start); + + start = clock_type::now(); + auto page_ranges = reader.page_index_byte_ranges(); + auto missing_page = std::find_if( + page_ranges.begin(), page_ranges.end(), [](auto const& range) { return range.is_empty(); }); + CUDF_EXPECTS(missing_page == page_ranges.end(), "A referenced source has no Parquet page index"); + timings.add("page-index range planning", start); + + start = clock_type::now(); + auto page_buffers = + cudf::io::parquet::fetch_page_indexes_to_host(inputs.datasource_refs, page_ranges); + auto page_spans = std::vector>{}; + page_spans.reserve(page_buffers.size()); + std::transform(page_buffers.begin(), + page_buffers.end(), + std::back_inserter(page_spans), + [](auto const& buffer) { return cudf::host_span{*buffer}; }); + timings.add("page-index fetch", start); + + start = clock_type::now(); + reader.setup_page_indexes(page_spans); + timings.add("page-index setup", start); + + start = clock_type::now(); + auto const indexed_metadatas = reader.parquet_metadatas(); + for (std::size_t source = 0; source < row_groups.size(); ++source) { + for (auto const rg : row_groups[source]) { + auto const& chunks = indexed_metadatas[source].row_groups[rg].columns; + for (auto const& name : payload_columns) { + auto const chunk = std::find_if(chunks.begin(), chunks.end(), [&](auto const& candidate) { + auto const& path = candidate.meta_data.path_in_schema; + return not path.empty() and path.back() == name; + }); + CUDF_EXPECTS(chunk != chunks.end(), + "Missing payload column " + name + " in " + source_names[source]); + CUDF_EXPECTS(chunk->offset_index.has_value(), + "Missing offset index for " + source_names[source] + " row group " + + std::to_string(rg) + " column " + name); + } + } + } + timings.add("selected index validation", start); + + start = clock_type::now(); + auto row_mask = + make_device_column(host_row_mask, cudf::data_type{cudf::type_id::BOOL8}, stream, stats_mr); + timings.add("row-mask host-to-device", start); + + start = clock_type::now(); + auto const passes = reader.construct_row_group_passes(row_groups, args.pass_read_limit); + timings.add("row-group pass construction", start); + + auto requested_bytes = uint64_t{0}; + auto selected_payload_bytes = uint64_t{0}; + auto output_allocated_bytes = std::size_t{0}; + auto global_mask_start = cudf::size_type{0}; + auto pass_tables = std::vector>{}; + auto const retain_output = + args.validate or args.output.has_value() or refs.size() != unique_refs.size(); + + auto range_planning_seconds = double{0}; + auto range_fetch_seconds = double{0}; + auto chunking_setup_seconds = double{0}; + auto materialize_seconds = double{0}; + auto stat_seconds = double{0}; + + std::cout << "PROGRESS constructed " << passes.size() << " multifile passes (limit " + << args.pass_read_limit / double{1024 * 1024} << " MiB)\n" + << std::flush; + for (std::size_t pass_index = 0; pass_index < passes.size(); ++pass_index) { + auto const& pass = passes[pass_index]; + auto const pass_start = clock_type::now(); + auto const pass_row_groups = + std::accumulate(pass.begin(), + pass.end(), + std::size_t{0}, + [](auto count, auto const& source_row_groups) { + return count + source_row_groups.size(); + }); + auto const pass_rows = reader.total_rows_in_row_groups(pass); + auto const pass_mask = cudf::column_view( + row_mask->type(), pass_rows, row_mask->view().data(), nullptr, 0, global_mask_start); + std::cout << "PROGRESS pass " << pass_index + 1 << "/" << passes.size() << " started (" + << pass_row_groups << " row groups, " << pass_rows << " input rows)\n" + << std::flush; + + auto payload_data = multisource_device_data{}; + if (args.use_page_mask and args.use_sparse_page_io) { + start = clock_type::now(); + auto const payload_ranges = reader.payload_column_chunks_byte_ranges( + pass, pass_mask, use_data_page_mask::YES, options, stream); + for (auto const& source_ranges : payload_ranges) { + requested_bytes = + std::accumulate(source_ranges.begin(), + source_ranges.end(), + requested_bytes, + [](uint64_t sum, auto const& range) { return sum + range.size(); }); + } + range_planning_seconds += std::chrono::duration(clock_type::now() - start).count(); + + start = clock_type::now(); + payload_data = fetch_multisource_device_data(inputs, payload_ranges, stream, stats_mr); + range_fetch_seconds += std::chrono::duration(clock_type::now() - start).count(); + + start = clock_type::now(); + reader.setup_chunking_for_payload_columns(0, + 0, + pass, + pass_mask, + use_data_page_mask::YES, + payload_data.per_source_spans, + options, + stream, + stats_mr); + } else { + start = clock_type::now(); + auto payload_ranges = reader.payload_column_chunks_byte_ranges(pass, options); + requested_bytes += + std::accumulate(payload_ranges.first.begin(), + payload_ranges.first.end(), + uint64_t{0}, + [](uint64_t sum, auto const& range) { return sum + range.size(); }); + range_planning_seconds += std::chrono::duration(clock_type::now() - start).count(); + + start = clock_type::now(); + payload_data = fetch_multisource_device_data(inputs, payload_ranges, stream, stats_mr); + range_fetch_seconds += std::chrono::duration(clock_type::now() - start).count(); + + start = clock_type::now(); + reader.setup_chunking_for_payload_columns(0, + 0, + pass, + pass_mask, + args.use_page_mask ? use_data_page_mask::YES + : use_data_page_mask::NO, + payload_data.flat_spans, + options, + stream, + stats_mr); + } + stream.synchronize(); + chunking_setup_seconds += std::chrono::duration(clock_type::now() - start).count(); + + auto current_tables = std::vector>{}; + start = clock_type::now(); + while (reader.has_next_table_chunk()) { + current_tables.push_back(reader.materialize_payload_columns_chunk(pass_mask).tbl); + } + stream.synchronize(); + materialize_seconds += std::chrono::duration(clock_type::now() - start).count(); + + if (retain_output) { + std::move(current_tables.begin(), current_tables.end(), std::back_inserter(pass_tables)); + } else { + start = clock_type::now(); + for (auto const& table : current_tables) { + selected_payload_bytes += payload_bytes(table->view(), stream); + output_allocated_bytes += table->alloc_size(); + } + stream.synchronize(); + stat_seconds += std::chrono::duration(clock_type::now() - start).count(); + } + global_mask_start += pass_rows; + auto const pass_seconds = + std::chrono::duration(clock_type::now() - pass_start).count(); + std::cout << "PROGRESS pass " << pass_index + 1 << "/" << passes.size() << " completed in " + << std::fixed << std::setprecision(3) << pass_seconds << " s\n" + << std::flush; + } + CUDF_EXPECTS(std::cmp_equal(global_mask_start, host_row_mask.size()), + "Row-group passes do not span the global row mask"); + + timings.add_seconds("payload byte-range planning", range_planning_seconds); + timings.add_seconds("payload byte-range fetch", range_fetch_seconds); + timings.add_seconds("payload chunking setup", chunking_setup_seconds); + timings.add_seconds("payload materialization/decode", materialize_seconds); + + start = clock_type::now(); + auto output = [&]() -> std::unique_ptr { + if (not retain_output) { return nullptr; } + CUDF_EXPECTS(not pass_tables.empty(), "Payload materialization produced no table chunks"); + if (pass_tables.size() == 1) { return std::move(pass_tables.front()); } + auto views = std::vector{}; + views.reserve(pass_tables.size()); + std::transform( + pass_tables.begin(), pass_tables.end(), std::back_inserter(views), [](auto const& table) { + return table->view(); + }); + return cudf::concatenate(views, stream, stats_mr); + }(); + stream.synchronize(); + timings.add("payload pass concatenation", start); + + start = clock_type::now(); + if (refs.size() != unique_refs.size()) { + CUDF_EXPECTS(output != nullptr, "Duplicate references require a retained output table"); + auto device_gather_map = make_gather_map(gather_map, stream, stats_mr); + output = cudf::gather(output->view(), + device_gather_map->view(), + cudf::out_of_bounds_policy::DONT_CHECK, + cudf::negative_index_policy::NOT_ALLOWED, + stream, + stats_mr); + stream.synchronize(); + } + timings.add("duplicate-order gather", start); + + if (retain_output) { + start = clock_type::now(); + selected_payload_bytes = payload_bytes(output->view(), stream); + output_allocated_bytes = output->alloc_size(); + stream.synchronize(); + stat_seconds = std::chrono::duration(clock_type::now() - start).count(); + } + timings.add_seconds("payload byte/stat extraction", stat_seconds); + + start = clock_type::now(); + std::unique_ptr expected; + if (args.validate) { + CUDF_EXPECTS(std::filesystem::is_regular_file(args.expected_output), + "Expected output does not exist: " + args.expected_output.string()); + expected = read_expected(args.expected_output); + stream.synchronize(); + } + timings.add("expected-output read", start); + + start = clock_type::now(); + if (args.validate) { + auto const equal = + cudf::tables_equal(output->view(), expected->view(), cudf::null_equality::EQUAL, stream); + stream.synchronize(); + CUDF_EXPECTS(equal, "Hybrid Scan output differs from Python output"); + } + timings.add("table comparison", start); + + auto const benchmark_elapsed = + std::chrono::duration(clock_type::now() - benchmark_start).count(); + + start = clock_type::now(); + if (args.output.has_value()) { write_output(args.output.value(), output->view()); } + stream.synchronize(); + timings.add("output parquet write", start); + + start = clock_type::now(); + expected.reset(); + output.reset(); + pass_tables.clear(); + row_mask.reset(); + stream.synchronize(); + timings.add("result cleanup", start); + + auto const end_to_end_elapsed = + std::chrono::duration(clock_type::now() - program_start).count(); + timings.print(); + std::cout << "\nSummary:\n" + << " setup + payload benchmark: " << benchmark_elapsed << " s\n" + << " end-to-end: " << end_to_end_elapsed << " s\n" + << " referenced files: " << source_names.size() << "\n" + << " selected row groups: " << selected_row_groups << "\n" + << " selected rows: " << refs.size() << "\n" + << " selected image bytes: " << selected_payload_bytes << "\n" + << " requested compressed bytes: " << requested_bytes << "\n" + << " use data page mask: " << std::boolalpha << args.use_page_mask << "\n" + << " use sparse page I/O: " + << (args.use_page_mask and args.use_sparse_page_io) << "\n"; + print_json(timings, + cache_result, + args.use_page_mask, + args.use_page_mask and args.use_sparse_page_io, + source_names.size(), + selected_row_groups, + static_cast(refs.size()), + selected_payload_bytes, + requested_bytes, + output_allocated_bytes, + benchmark_elapsed, + end_to_end_elapsed); + std::cout << "Peak device memory: " << stats_mr.get_bytes_counter().peak / double{1024 * 1024} + << " MiB\n"; + return 0; +} diff --git a/cpp/tests/io/experimental/hybrid_scan_composer.cpp b/cpp/tests/io/experimental/hybrid_scan_composer.cpp index 7f39bc7121c7..3f3d7e0c3522 100644 --- a/cpp/tests/io/experimental/hybrid_scan_composer.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_composer.cpp @@ -319,6 +319,109 @@ std::tuple, std::unique_ptr> chunked_h return std::tuple{std::move(filter_table), std::move(payload_table)}; } +std::tuple, std::unique_ptr> sparse_chunked_hybrid_scan( + cudf::io::datasource& datasource, + cudf::ast::operation const& filter_expression, + std::optional> const& payload_column_names, + bool case_sensitive_names, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr, + rmm::mr::aligned_resource_adaptor& aligned_mr) +{ + auto options = cudf::io::parquet_reader_options::builder() + .filter(filter_expression) + .case_sensitive_names(case_sensitive_names) + .build(); + if (payload_column_names.has_value()) { options.set_column_names(payload_column_names.value()); } + + auto const reader = setup_reader(datasource, options); + auto reader_ref = std::ref(*reader); + auto const filtered_row_group_indices = + apply_hybrid_scan_filters(datasource, reader_ref, options, stream, mr); + auto const current_row_group_indices = + cudf::host_span(filtered_row_group_indices); + auto row_mask = + options.get_filter().has_value() + ? reader->build_row_mask_with_page_index_stats(current_row_group_indices, options, stream, mr) + : reader->build_all_true_row_mask(current_row_group_indices, stream, mr); + + auto filter_tables = std::vector>{}; + auto payload_tables = std::vector>{}; + std::size_t rows_materialized = 0; + auto const materialize_pass = [&](cudf::host_span row_group_indices) { + auto const rows_in_pass = reader->total_rows_in_row_groups(row_group_indices); + auto* null_mask = row_mask->nullable() ? row_mask->mutable_view().null_mask() : nullptr; + auto const slice_null_count = + cudf::null_count(null_mask, rows_materialized, rows_materialized + rows_in_pass, stream); + auto row_mask_view = cudf::mutable_column_view(row_mask->type(), + rows_in_pass, + row_mask->mutable_view().data(), + null_mask, + slice_null_count, + rows_materialized); + + auto const filter_byte_ranges = + reader->filter_column_chunks_byte_ranges(row_group_indices, options); + auto [filter_buffers, filter_data, filter_tasks] = + cudf::io::parquet::fetch_byte_ranges_to_device_async( + datasource, filter_byte_ranges, stream, mr); + filter_tasks.get(); + reader->setup_chunking_for_filter_columns( + 1024, + 10240, + row_group_indices, + row_mask_view, + cudf::io::parquet::experimental::use_data_page_mask::YES, + filter_data, + options, + stream, + mr); + while (reader->has_next_table_chunk()) { + filter_tables.push_back(reader->materialize_filter_columns_chunk(row_mask_view).tbl); + } + + auto const payload_page_ranges = reader->payload_column_chunks_byte_ranges( + row_group_indices, + row_mask_view, + cudf::io::parquet::experimental::use_data_page_mask::YES, + options, + stream); + auto [payload_buffers, payload_data, payload_tasks] = + cudf::io::parquet::fetch_byte_ranges_to_device_async( + datasource, payload_page_ranges, stream, mr); + payload_tasks.get(); + auto const page_data_per_source = std::vector>>{ + {payload_data.begin(), payload_data.end()}}; + reader->setup_chunking_for_payload_columns( + 1024, + 10240, + row_group_indices, + row_mask_view, + cudf::io::parquet::experimental::use_data_page_mask::YES, + page_data_per_source, + options, + stream, + mr); + while (reader->has_next_table_chunk()) { + payload_tables.push_back(reader->materialize_payload_columns_chunk(row_mask_view).tbl); + } + + rows_materialized += rows_in_pass; + }; + + if (current_row_group_indices.size() > 1) { + auto const row_group_split = current_row_group_indices.size() / 2; + materialize_pass(current_row_group_indices.subspan(0, row_group_split)); + materialize_pass(current_row_group_indices.subspan( + row_group_split, current_row_group_indices.size() - row_group_split)); + } else { + materialize_pass(current_row_group_indices); + } + + return std::tuple{concatenate_tables(std::move(filter_tables), stream, mr), + concatenate_tables(std::move(payload_tables), stream, mr)}; +} + std::unique_ptr hybrid_scan_single_step( cudf::io::datasource& datasource, cudf::ast::operation const& filter_expression, diff --git a/cpp/tests/io/experimental/hybrid_scan_composer.hpp b/cpp/tests/io/experimental/hybrid_scan_composer.hpp index ca7439277e01..32b2f421bfd3 100644 --- a/cpp/tests/io/experimental/hybrid_scan_composer.hpp +++ b/cpp/tests/io/experimental/hybrid_scan_composer.hpp @@ -58,6 +58,31 @@ std::tuple, std::unique_ptr> chunked_h rmm::device_async_resource_ref mr, rmm::mr::aligned_resource_adaptor& aligned_mr); +/** + * @brief Read parquet file with chunked hybrid scan and sparse page-level payload I/O + * + * Filter columns use full-column-chunk I/O. Payload ranges are planned only after filter + * materialization updates the row mask, then only retained data pages are fetched. + * + * @param datasource Input datasource + * @param filter_expression Filter expression + * @param payload_column_names List of paths of select payload column names, if any + * @param case_sensitive_names Whether column names are case sensitive + * @param stream CUDA stream for hybrid scan reader + * @param mr Device memory resource + * @param aligned_mr Device memory resource to allocate aligned memory for bloom filters + * + * @return Tuple of filter and payload tables + */ +std::tuple, std::unique_ptr> sparse_chunked_hybrid_scan( + cudf::io::datasource& datasource, + cudf::ast::operation const& filter_expression, + std::optional> const& payload_column_names, + bool case_sensitive_names, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr, + rmm::mr::aligned_resource_adaptor& aligned_mr); + /** * @brief Read parquet file with the hybrid scan reader in a single step * diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp index 47c8164468b6..157546ee9518 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp @@ -168,6 +168,398 @@ TEST_F(HybridScanMultifileTest, MaterializeListsOfStrings) test_hybrid_scan_multifile({col0, *col1, *col2, *col3, *col4}, false); } +TEST_F(HybridScanMultifileTest, PageLevelDictionaryPayloadByteReduction) +{ + auto col0 = testdata::ascending(); + + auto payload_values = std::vector(num_ordered_rows); + for (auto i = std::size_t{0}; i < payload_values.size(); ++i) { + payload_values[i] = "dictionary value " + std::to_string(i % 8); + } + auto col1 = cudf::test::strings_column_wrapper(payload_values.begin(), payload_values.end()); + + // A page-aligned threshold retains two of four data pages. The writer's ALWAYS dictionary policy + // requires the page-I/O path to retain the dictionary while requesting fewer bytes than the + // legacy full-column-chunk path. + auto constexpr threshold = uint32_t{2 * page_size_for_ordered_tests / 100}; + test_hybrid_scan_multifile({col0, col1}, true, threshold, true); +} + +TEST_F(HybridScanMultifileTest, PageLevelStringsSeparatedByPrunedPages) +{ + auto filter_values = cudf::detail::make_counting_transform_iterator( + cudf::size_type{0}, [](auto i) { return (i / page_size_for_ordered_tests) % 2 == 0; }); + auto filter = + cudf::test::fixed_width_column_wrapper(filter_values, filter_values + num_ordered_rows); + + auto payload_values = std::vector(num_ordered_rows); + for (auto i = std::size_t{0}; i < payload_values.size(); ++i) { + payload_values[i] = "payload value " + std::to_string(i); + } + auto payload = cudf::test::strings_column_wrapper(payload_values.begin(), payload_values.end()); + auto table = cudf::table_view{{filter, payload}}; + + auto metadata = cudf::io::table_input_metadata(table); + metadata.column_metadata[0].set_name("filter"); + metadata.column_metadata[1].set_name("payload"); + + auto parquet_buffers = std::vector>(2); + for (auto& parquet_buffer : parquet_buffers) { + auto options = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&parquet_buffer}, table) + .metadata(metadata) + .row_group_size_rows(num_ordered_rows) + .max_page_size_rows(page_size_for_ordered_tests) + .max_page_size_bytes(64 * 1024 * 1024) + .compression(cudf::io::compression_type::NONE) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN); + cudf::io::write_parquet(options); + } + + auto const filter_ref = cudf::ast::column_name_reference("filter"); + auto filter_expression = cudf::ast::operation(cudf::ast::ast_operator::IDENTITY, filter_ref); + auto source_info = build_source_info(parquet_buffers); + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + auto expected_options = + cudf::io::parquet_reader_options::builder(source_info).filter(filter_expression).build(); + auto expected = cudf::io::read_parquet(expected_options, stream, mr); + + auto const [filter_result, payload_result] = + page_level_chunked_hybrid_scan_multifile(source_info, filter_expression, {}, true, stream, mr); + + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->select({0}), filter_result->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->select({1}), payload_result->view()); +} + +TEST_F(HybridScanMultifileTest, PageLevelAsymmetricSourceRowGroupOrdering) +{ + auto constexpr rows_per_group = 2 * page_size_for_ordered_tests; + auto constexpr rows_source_0 = num_ordered_rows; + auto constexpr rows_source_1 = num_ordered_rows; + auto constexpr rows_per_page = rows_per_group / 4; + + auto source_0_filter_values = cudf::detail::make_counting_transform_iterator( + 0, [](auto i) { return static_cast(i / 100); }); + auto source_1_filter_values = cudf::detail::make_counting_transform_iterator( + 0, [](auto i) { return static_cast((num_ordered_rows - i) / 100); }); + auto source_0_filter = cudf::test::fixed_width_column_wrapper( + source_0_filter_values, source_0_filter_values + rows_source_0); + auto source_1_filter = cudf::test::fixed_width_column_wrapper( + source_1_filter_values, source_1_filter_values + rows_source_1); + + auto source_0_payload_values = std::vector(rows_source_0); + auto source_1_payload_values = std::vector(rows_source_1); + for (auto i = std::size_t{0}; i < source_0_payload_values.size(); ++i) { + source_0_payload_values[i] = "source 0 dictionary value " + std::to_string(i % 8); + } + for (auto i = std::size_t{0}; i < source_1_payload_values.size(); ++i) { + source_1_payload_values[i] = "source 1 dictionary value " + std::to_string(i % 8); + } + auto source_0_payload = cudf::test::strings_column_wrapper(source_0_payload_values.begin(), + source_0_payload_values.end()); + auto source_1_payload = cudf::test::strings_column_wrapper(source_1_payload_values.begin(), + source_1_payload_values.end()); + auto const source_0_table = cudf::table_view{{source_0_filter, source_0_payload}}; + auto const source_1_table = cudf::table_view{{source_1_filter, source_1_payload}}; + + auto parquet_buffers = std::vector>(2); + auto const write_source = [&](auto const& table, auto& buffer) { + cudf::io::table_input_metadata metadata(table); + metadata.column_metadata[0].set_name("col0"); + auto options = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, table) + .metadata(metadata) + .row_group_size_rows(rows_per_group) + .max_page_size_rows(rows_per_page) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN); + cudf::io::write_parquet(options); + }; + write_source(source_0_table, parquet_buffers[0]); + write_source(source_1_table, parquet_buffers[1]); + + auto constexpr threshold = uint32_t{75}; + auto scalar = cudf::numeric_scalar(threshold); + auto literal = cudf::ast::literal(scalar); + auto col_ref = cudf::ast::column_name_reference("col0"); + auto filter_expression = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref, literal); + + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + auto const source_info = build_source_info(parquet_buffers); + auto const expected = cudf::io::read_parquet( + cudf::io::parquet_reader_options::builder(source_info).filter(filter_expression), stream, mr); + + auto const [filter_table, payload_table] = + page_level_chunked_hybrid_scan_multifile(source_info, filter_expression, {}, true, stream, mr); + + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->select({0}), filter_table->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->select({1}), payload_table->view()); + auto const [requested_payload_bytes, full_payload_bytes] = + payload_byte_range_sizes(source_info, filter_expression, true, stream, mr); + EXPECT_GT(requested_payload_bytes, 0); + EXPECT_LT(requested_payload_bytes, full_payload_bytes); +} + +TEST_F(HybridScanMultifileTest, PageLevelPlainEncodingExactCoalescedRanges) +{ + auto parquet_buffers = make_plain_payload_parquet_buffers(); + auto const source_info = build_source_info(parquet_buffers); + auto inputs = multifile_inputs(source_info); + + auto scalar = cudf::numeric_scalar(0); + auto literal = cudf::ast::literal(scalar); + auto col_ref_0 = cudf::ast::column_name_reference("col0"); + auto filter_expression = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref_0, literal); + auto options = cudf::io::parquet_reader_options::builder() + .column_names({"col1"}) + .filter(filter_expression) + .build(); + auto reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + setup_page_indexes(reader, inputs); + + auto const row_groups = reader.all_row_groups(options); + auto const metadatas = reader.parquet_metadatas(); + auto selected_rows = + std::vector(reader.total_rows_in_row_groups(row_groups), uint8_t{0}); + auto expected_ranges = + std::vector>(metadatas.size()); + + std::size_t source_row_offset = 0; + for (std::size_t source_idx = 0; source_idx < metadatas.size(); ++source_idx) { + auto const& metadata = metadatas[source_idx]; + ASSERT_EQ(metadata.row_groups.size(), 1); + auto const& payload_chunk = metadata.row_groups.front().columns[1]; + EXPECT_NE(std::find(payload_chunk.meta_data.encodings.begin(), + payload_chunk.meta_data.encodings.end(), + cudf::io::parquet::Encoding::PLAIN), + payload_chunk.meta_data.encodings.end()); + EXPECT_EQ(std::find(payload_chunk.meta_data.encodings.begin(), + payload_chunk.meta_data.encodings.end(), + cudf::io::parquet::Encoding::RLE_DICTIONARY), + payload_chunk.meta_data.encodings.end()); + EXPECT_EQ(payload_chunk.meta_data.dictionary_page_offset, 0); + ASSERT_TRUE(payload_chunk.offset_index.has_value()); + + auto const& pages = payload_chunk.offset_index->page_locations; + ASSERT_GE(pages.size(), 4); + ASSERT_EQ(pages[1].offset + pages[1].compressed_page_size, pages[2].offset); + auto const selected_begin = pages[1].first_row_index; + auto const selected_end = pages[3].first_row_index; + ASSERT_GE(selected_begin, 0); + ASSERT_LE(selected_end, metadata.row_groups.front().num_rows); + std::fill(selected_rows.begin() + source_row_offset + selected_begin, + selected_rows.begin() + source_row_offset + selected_end, + uint8_t{1}); + + expected_ranges[source_idx].emplace_back( + pages[1].offset, pages[2].offset + pages[2].compressed_page_size - pages[1].offset); + source_row_offset += metadata.row_groups.front().num_rows; + } + ASSERT_EQ(source_row_offset, selected_rows.size()); + auto row_mask = + cudf::test::fixed_width_column_wrapper(selected_rows.begin(), selected_rows.end()) + .release(); + + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + auto const page_ranges = reader.payload_column_chunks_byte_ranges( + row_groups, row_mask->view(), use_data_page_mask::YES, options, stream); + expect_byte_ranges_equal(expected_ranges, page_ranges); + + auto page_data = fetch_multisource_device_data(inputs, page_ranges, stream, mr); + reader.setup_chunking_for_payload_columns(256 * 1024, + 1024 * 1024, + row_groups, + row_mask->view(), + use_data_page_mask::YES, + page_data.per_source_spans, + options, + stream, + mr); + auto payload_chunks = std::vector>{}; + while (reader.has_next_table_chunk()) { + payload_chunks.push_back( + std::move(reader.materialize_payload_columns_chunk(row_mask->view()).tbl)); + } + auto actual = concatenate_tables(std::move(payload_chunks), stream, mr); + + auto const full = + cudf::io::read_parquet(cudf::io::parquet_reader_options::builder(source_info), stream, mr); + auto const expected = cudf::apply_boolean_mask(full.tbl->select({1}), row_mask->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), actual->view()); +} + +TEST_F(HybridScanMultifileTest, PageLevelAllFalseMaskHasNoRanges) +{ + auto parquet_buffers = make_plain_payload_parquet_buffers(); + auto const source_info = build_source_info(parquet_buffers); + auto inputs = multifile_inputs(source_info); + + auto scalar = cudf::numeric_scalar(0); + auto literal = cudf::ast::literal(scalar); + auto col_ref_0 = cudf::ast::column_name_reference("col0"); + auto filter_expression = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref_0, literal); + auto options = cudf::io::parquet_reader_options::builder() + .column_names({"col1"}) + .filter(filter_expression) + .build(); + auto reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + setup_page_indexes(reader, inputs); + + auto const row_groups = reader.all_row_groups(options); + auto false_values = cuda::make_constant_iterator(false); + auto row_mask = cudf::test::fixed_width_column_wrapper( + false_values, false_values + reader.total_rows_in_row_groups(row_groups)) + .release(); + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + + auto const page_ranges = reader.payload_column_chunks_byte_ranges( + row_groups, row_mask->view(), use_data_page_mask::YES, options, stream); + ASSERT_EQ(page_ranges.size(), parquet_buffers.size()); + EXPECT_TRUE(std::all_of( + page_ranges.begin(), page_ranges.end(), [](auto const& ranges) { return ranges.empty(); })); + + auto const empty_page_data = + std::vector>>(parquet_buffers.size()); + reader.setup_chunking_for_payload_columns(0, + 0, + row_groups, + row_mask->view(), + use_data_page_mask::YES, + empty_page_data, + options, + stream, + mr); + ASSERT_TRUE(reader.has_next_table_chunk()); + auto const result = reader.materialize_payload_columns_chunk(row_mask->view()); + EXPECT_EQ(result.tbl->num_rows(), 0); + EXPECT_EQ(result.tbl->num_columns(), 1); + EXPECT_EQ(result.metadata.num_input_row_groups, 2); + EXPECT_FALSE(reader.has_next_table_chunk()); +} + +TEST_F(HybridScanMultifileTest, PageLevelNoMaskFallbackAndPlanLifecycle) +{ + auto parquet_buffers = make_plain_payload_parquet_buffers(); + auto const source_info = build_source_info(parquet_buffers); + auto inputs = multifile_inputs(source_info); + + auto scalar = cudf::numeric_scalar(0); + auto literal = cudf::ast::literal(scalar); + auto col_ref_0 = cudf::ast::column_name_reference("col0"); + auto filter_expression = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref_0, literal); + auto options = cudf::io::parquet_reader_options::builder() + .column_names({"col1"}) + .filter(filter_expression) + .build(); + auto reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + setup_page_indexes(reader, inputs); + + auto const row_groups = reader.all_row_groups(options); + auto const full_ranges = group_byte_ranges_by_source( + reader.payload_column_chunks_byte_ranges(row_groups, options), parquet_buffers.size()); + auto true_values = cuda::make_constant_iterator(true); + auto row_mask = cudf::test::fixed_width_column_wrapper( + true_values, true_values + reader.total_rows_in_row_groups(row_groups)) + .release(); + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + + auto const planned_ranges = reader.payload_column_chunks_byte_ranges( + row_groups, row_mask->view(), use_data_page_mask::NO, options, stream); + expect_byte_ranges_equal(full_ranges, planned_ranges); + EXPECT_THROW(static_cast(reader.payload_column_chunks_byte_ranges( + row_groups, row_mask->view(), use_data_page_mask::NO, options, stream)), + cudf::logic_error); + + auto page_data = fetch_multisource_device_data(inputs, planned_ranges, stream, mr); + reader.setup_chunking_for_payload_columns(0, + 0, + row_groups, + row_mask->view(), + use_data_page_mask::NO, + page_data.per_source_spans, + options, + stream, + mr); + EXPECT_THROW(reader.setup_chunking_for_payload_columns(0, + 0, + row_groups, + row_mask->view(), + use_data_page_mask::NO, + page_data.per_source_spans, + options, + stream, + mr), + cudf::logic_error); +} + +TEST_F(HybridScanMultifileTest, PageLevelRejectsInvalidFetchedSpans) +{ + auto parquet_buffers = make_plain_payload_parquet_buffers(); + auto const source_info = build_source_info(parquet_buffers); + auto inputs = multifile_inputs(source_info); + + auto scalar = cudf::numeric_scalar(0); + auto literal = cudf::ast::literal(scalar); + auto col_ref_0 = cudf::ast::column_name_reference("col0"); + auto filter_expression = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref_0, literal); + auto options = cudf::io::parquet_reader_options::builder() + .column_names({"col1"}) + .filter(filter_expression) + .build(); + auto reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + setup_page_indexes(reader, inputs); + + auto const row_groups = reader.all_row_groups(options); + auto selected_values = cudf::detail::make_counting_transform_iterator( + cudf::size_type{0}, [](auto i) { return (i % num_ordered_rows) < num_ordered_rows / 2; }); + auto row_mask = cudf::test::fixed_width_column_wrapper( + selected_values, selected_values + reader.total_rows_in_row_groups(row_groups)) + .release(); + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + + auto page_ranges = reader.payload_column_chunks_byte_ranges( + row_groups, row_mask->view(), use_data_page_mask::YES, options, stream); + auto page_data = fetch_multisource_device_data(inputs, page_ranges, stream, mr); + auto bad_count = page_data.per_source_spans; + auto count_source = std::find_if( + bad_count.begin(), bad_count.end(), [](auto const& spans) { return not spans.empty(); }); + ASSERT_NE(count_source, bad_count.end()); + count_source->pop_back(); + EXPECT_THROW( + reader.setup_chunking_for_payload_columns( + 0, 0, row_groups, row_mask->view(), use_data_page_mask::YES, bad_count, options, stream, mr), + cudf::logic_error); + + page_ranges = reader.payload_column_chunks_byte_ranges( + row_groups, row_mask->view(), use_data_page_mask::YES, options, stream); + auto bad_size = page_data.per_source_spans; + auto size_source = std::find_if( + bad_size.begin(), bad_size.end(), [](auto const& spans) { return not spans.empty(); }); + ASSERT_NE(size_source, bad_size.end()); + ASSERT_GT(size_source->front().size(), 1); + size_source->front() = + cudf::device_span{size_source->front().data(), size_source->front().size() - 1}; + EXPECT_THROW( + reader.setup_chunking_for_payload_columns( + 0, 0, row_groups, row_mask->view(), use_data_page_mask::YES, bad_size, options, stream, mr), + cudf::logic_error); +} + TEST_F(HybridScanMultifileTest, PrependIndexColumns) { using T = int32_t; diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index 6607c941e51c..c3bf41f66463 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -509,7 +509,7 @@ TEST_F(HybridScanTest, ConsecutivePrunedPageOffsets) std::unique_ptr filter_table; std::unique_ptr payload_table; ASSERT_NO_THROW(std::tie(filter_table, payload_table) = - hybrid_scan(*datasource, filter, {}, true, stream, mr, aligned_mr)); + sparse_chunked_hybrid_scan(*datasource, filter, {}, true, stream, mr, aligned_mr)); CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view().select({0}), filter_table->view()); CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view().select({1, 2, 3, 4}), diff --git a/python/pylibcudf/pylibcudf/io/experimental/__init__.pxd b/python/pylibcudf/pylibcudf/io/experimental/__init__.pxd index 87cc217ebf94..b36338238833 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/__init__.pxd +++ b/python/pylibcudf/pylibcudf/io/experimental/__init__.pxd @@ -3,5 +3,6 @@ from pylibcudf.io.experimental.hybrid_scan cimport ( FileMetaData, + HybridScanMultiFile, HybridScanReader, ) diff --git a/python/pylibcudf/pylibcudf/io/experimental/__init__.py b/python/pylibcudf/pylibcudf/io/experimental/__init__.py index 6c64231eb1e9..5f31ee5efae2 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/__init__.py +++ b/python/pylibcudf/pylibcudf/io/experimental/__init__.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from pylibcudf.io.experimental.hybrid_scan import ( + HybridScanMultiFile, HybridScanReader, UseDataPageMask, ) @@ -9,6 +10,7 @@ __all__ = [ "FileMetaData", # backwards compatibility + "HybridScanMultiFile", "HybridScanReader", "UseDataPageMask", ] diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd index a19cd7db8bf8..91a26bb53d29 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd @@ -14,6 +14,7 @@ from pylibcudf.io.parquet_metadata cimport FileMetaData as c_FileMetaData from pylibcudf.io.types cimport TableWithMetadata from pylibcudf.libcudf.io.hybrid_scan cimport ( hybrid_scan_reader as cpp_hybrid_scan_reader, + hybrid_scan_multifile as cpp_hybrid_scan_multifile, use_data_page_mask, ) from pylibcudf.libcudf.io.hybrid_scan cimport const_uint8_t @@ -27,3 +28,9 @@ cdef class HybridScanReader: cdef unique_ptr[cpp_hybrid_scan_reader] c_obj cdef Stream _stream cdef DeviceMemoryResource mr + +cdef class HybridScanMultiFile: + cdef unique_ptr[cpp_hybrid_scan_multifile] c_obj + cdef Stream _stream + cdef DeviceMemoryResource mr + cdef object _page_data_keepalive diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi index f95dc8b054d3..b1229306862e 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi @@ -142,3 +142,54 @@ class HybridScanReader: pass_read_limit: int, ) -> list[list[int]]: ... def has_next_table_chunk(self) -> bool: ... + +class HybridScanMultiFile: + @staticmethod + def from_parquet_metadatas( + parquet_metadatas: list[FileMetaData], + options: ParquetReaderOptions, + ) -> HybridScanMultiFile: ... + def parquet_metadatas(self) -> list[FileMetaData]: ... + def page_index_byte_ranges(self) -> list[ByteRangeInfo]: ... + def setup_page_indexes(self, page_index_bytes: list[Buffer]) -> None: ... + def all_row_groups( + self, options: ParquetReaderOptions + ) -> list[list[int]]: ... + def total_rows_in_row_groups( + self, row_group_indices: list[list[int]] + ) -> int: ... + def build_all_true_row_mask( + self, + row_group_indices: list[list[int]], + stream: CudaStreamLike | None = None, + mr: DeviceMemoryResource | None = None, + ) -> Column: ... + def payload_column_chunks_byte_ranges( + self, + row_group_indices: list[list[int]], + row_mask: Column, + mask_data_pages: UseDataPageMask, + options: ParquetReaderOptions, + stream: CudaStreamLike | None = None, + ) -> list[list[ByteRangeInfo]]: ... + def setup_chunking_for_payload_columns( + self, + chunk_read_limit: int, + pass_read_limit: int, + row_group_indices: list[list[int]], + row_mask: Column, + mask_data_pages: UseDataPageMask, + page_data_per_source: list[list], + options: ParquetReaderOptions, + stream: CudaStreamLike | None = None, + mr: DeviceMemoryResource | None = None, + ) -> None: ... + def materialize_payload_columns_chunk( + self, row_mask: Column + ) -> TableWithMetadata: ... + def construct_row_group_passes( + self, + row_group_indices: list[list[int]], + pass_read_limit: int, + ) -> list[list[list[int]]]: ... + def has_next_table_chunk(self) -> bool: ... diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index ed878c2647bb..e2f74156afa7 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx @@ -22,8 +22,13 @@ from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view, mutable_column_view from pylibcudf.libcudf.io.hybrid_scan cimport ( const_device_span_const_uint8_t, + const_FileMetaData, + const_host_span_const_uint8_t, const_size_type, const_uint8_t, + const_vector_device_span_const_uint8_t, + const_vector_size_type, + hybrid_scan_multifile as cpp_hybrid_scan_multifile, hybrid_scan_reader as cpp_hybrid_scan_reader, use_data_page_mask as cpp_use_data_page_mask, ) @@ -46,7 +51,12 @@ import pylibcudf.libcudf.io.hybrid_scan UseDataPageMask = pylibcudf.libcudf.io.hybrid_scan.use_data_page_mask -__all__ = ["FileMetaData", "HybridScanReader", "UseDataPageMask"] +__all__ = [ + "FileMetaData", + "HybridScanMultiFile", + "HybridScanReader", + "UseDataPageMask", +] cdef device_span[const_uint8_t] _get_device_span(object obj) except *: @@ -60,6 +70,18 @@ cdef device_span[const_uint8_t] _get_device_span(object obj) except *: obj.size) +cdef vector[vector[size_type]] _get_row_groups(object row_groups) except *: + """Convert Python per-source row-group indices to C++ vectors.""" + cdef vector[vector[size_type]] result + cdef vector[size_type] source_row_groups + for source in row_groups: + for row_group in source: + source_row_groups.push_back(row_group) + result.push_back(source_row_groups) + source_row_groups.clear() + return result + + cdef class HybridScanReader: """Experimental Parquet reader optimized for highly selective filters. @@ -857,4 +879,230 @@ cdef class HybridScanReader: return self.c_obj.get()[0].has_next_table_chunk() +cdef class HybridScanMultiFile: + """Experimental Hybrid Scan reader for multiple Parquet sources. + + The per-source inputs and outputs use source order. A row mask spans all + selected rows in source order, and then in row-group order within a source. + This API is experimental. + """ + + @staticmethod + def from_parquet_metadatas(object parquet_metadatas, ParquetReaderOptions options): + """Create a reader from one ``FileMetaData`` per Parquet source.""" + cdef HybridScanMultiFile reader = HybridScanMultiFile.__new__( + HybridScanMultiFile + ) + cdef vector[cpp_FileMetaData] metadatas + cdef object metadata + for metadata in parquet_metadatas: + if not isinstance(metadata, FileMetaData): + raise TypeError( + "parquet_metadatas must contain only FileMetaData objects" + ) + metadatas.push_back((metadata).c_obj) + if metadatas.empty(): + raise ValueError("parquet_metadatas must not be empty") + reader.c_obj = make_unique[cpp_hybrid_scan_multifile]( + host_span[const_FileMetaData]( + metadatas.data(), metadatas.size() + ), + options.c_obj, + ) + return reader + + def parquet_metadatas(self): + """Return one ``FileMetaData`` object per source.""" + cdef vector[cpp_FileMetaData] metadatas = ( + self.c_obj.get()[0].parquet_metadatas() + ) + cdef cpp_FileMetaData metadata + cdef list result = [] + for metadata in metadatas: + result.append(c_FileMetaData.from_cpp(metadata)) + return result + + def page_index_byte_ranges(self): + """Return the page-index byte range for each source.""" + cdef vector[byte_range_info] ranges = ( + self.c_obj.get()[0].page_index_byte_ranges() + ) + return [ByteRangeInfo(r.offset(), r.size()) for r in ranges] + + def setup_page_indexes(self, list page_index_bytes): + """Install one host page-index buffer for each source.""" + cdef vector[host_span[const_uint8_t]] spans + cdef const uint8_t[::1] page_index + for page_index in page_index_bytes: + if len(page_index) == 0: + spans.push_back( + host_span[const_uint8_t](0, 0) + ) + else: + spans.push_back( + host_span[const_uint8_t](&page_index[0], len(page_index)) + ) + self.c_obj.get()[0].setup_page_indexes( + host_span[const_host_span_const_uint8_t]( + spans.data(), spans.size() + ) + ) + + def all_row_groups(self, ParquetReaderOptions options): + """Return row-group indices for every source.""" + cdef vector[vector[size_type]] row_groups = ( + self.c_obj.get()[0].all_row_groups(options.c_obj) + ) + return [list(row_groups[source]) for source in range(row_groups.size())] + + def total_rows_in_row_groups(self, object row_group_indices): + """Return total selected rows across all sources.""" + cdef vector[vector[size_type]] indices = _get_row_groups( + row_group_indices + ) + return self.c_obj.get()[0].total_rows_in_row_groups( + host_span[const_vector_size_type]( + indices.data(), indices.size() + ) + ) + + def build_all_true_row_mask( + self, + object row_group_indices, + object stream=None, + DeviceMemoryResource mr=None, + ): + """Build an all-true BOOL8 mask spanning selected rows.""" + cdef vector[vector[size_type]] indices = _get_row_groups( + row_group_indices + ) + cdef Stream _stream = _get_stream(stream) + mr = _get_memory_resource(mr) + cdef unique_ptr[column] c_result = ( + self.c_obj.get()[0].build_all_true_row_mask( + host_span[const_vector_size_type]( + indices.data(), indices.size() + ), + _stream.view().value(), + mr.get_mr(), + ) + ) + return Column.from_libcudf(move(c_result), _stream, mr) + + def payload_column_chunks_byte_ranges( + self, + object row_group_indices, + Column row_mask, + cpp_use_data_page_mask mask_data_pages, + ParquetReaderOptions options, + object stream=None, + ): + """Plan payload page ranges, grouped by source.""" + cdef vector[vector[size_type]] indices = _get_row_groups( + row_group_indices + ) + cdef Stream _stream = _get_stream(stream) + cdef column_view mask_view = row_mask.view() + cdef vector[vector[byte_range_info]] ranges = ( + self.c_obj.get()[0].payload_column_chunks_byte_ranges( + host_span[const_vector_size_type]( + indices.data(), indices.size() + ), + mask_view, + mask_data_pages, + options.c_obj, + _stream.view().value(), + ) + ) + return [ + [ + ByteRangeInfo(byte_range.offset(), byte_range.size()) + for byte_range in ranges[source] + ] + for source in range(ranges.size()) + ] + + def setup_chunking_for_payload_columns( + self, + size_t chunk_read_limit, + size_t pass_read_limit, + object row_group_indices, + Column row_mask, + cpp_use_data_page_mask mask_data_pages, + list page_data_per_source, + ParquetReaderOptions options, + object stream=None, + DeviceMemoryResource mr=None, + ): + """Configure payload chunking from source-grouped device page data.""" + cdef vector[vector[size_type]] indices = _get_row_groups( + row_group_indices + ) + cdef vector[vector[device_span[const_uint8_t]]] source_spans + cdef vector[device_span[const_uint8_t]] spans + cdef object source_data + cdef object data + for source_data in page_data_per_source: + for data in source_data: + spans.push_back(_get_device_span(data)) + source_spans.push_back(spans) + spans.clear() + + self._stream = _get_stream(stream) + self.mr = _get_memory_resource(mr) + self._page_data_keepalive = page_data_per_source + cdef column_view mask_view = row_mask.view() + self.c_obj.get()[0].setup_chunking_for_payload_columns( + chunk_read_limit, + pass_read_limit, + host_span[const_vector_size_type]( + indices.data(), indices.size() + ), + mask_view, + mask_data_pages, + host_span[const_vector_device_span_const_uint8_t]( + source_spans.data(), + source_spans.size(), + ), + options.c_obj, + self._stream.view().value(), + self.mr.get_mr(), + ) + + def materialize_payload_columns_chunk(self, Column row_mask): + """Materialize the next configured payload output chunk.""" + cdef column_view mask_view = row_mask.view() + cdef table_with_metadata c_result = ( + self.c_obj.get()[0].materialize_payload_columns_chunk(mask_view) + ) + return TableWithMetadata.from_libcudf(c_result, self._stream, self.mr) + + def construct_row_group_passes( + self, object row_group_indices, size_t pass_read_limit + ): + """Partition per-source row groups into bounded-memory passes.""" + cdef vector[vector[size_type]] indices = _get_row_groups( + row_group_indices + ) + cdef vector[vector[vector[size_type]]] passes = ( + self.c_obj.get()[0].construct_row_group_passes( + host_span[const_vector_size_type]( + indices.data(), indices.size() + ), + pass_read_limit, + ) + ) + return [ + [ + list(row_groups) + for row_groups in passes[pass_index] + ] + for pass_index in range(passes.size()) + ] + + def has_next_table_chunk(self): + """Return whether a configured chunked read has output remaining.""" + return self.c_obj.get()[0].has_next_table_chunk() + + UseDataPageMask.__str__ = UseDataPageMask.__repr__ diff --git a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd index 36201d545de6..2c62a19d007c 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd @@ -21,7 +21,12 @@ from rmm.librmm.memory_resource cimport device_async_resource_ref ctypedef const uint8_t const_uint8_t ctypedef const size_type const_size_type +ctypedef const FileMetaData const_FileMetaData ctypedef const device_span[const_uint8_t] const_device_span_const_uint8_t +ctypedef const vector[size_type] const_vector_size_type +ctypedef const vector[device_span[const_uint8_t]] const_vector_device_span_const_uint8_t +ctypedef host_span[const_uint8_t] host_span_const_uint8_t +ctypedef const host_span_const_uint8_t const_host_span_const_uint8_t cdef extern from "cudf/io/experimental/hybrid_scan.hpp" \ namespace "cudf::io::parquet::experimental" nogil: @@ -174,3 +179,66 @@ cdef extern from "cudf/io/experimental/hybrid_scan.hpp" \ ) except +libcudf_exception_handler bool has_next_table_chunk() except +libcudf_exception_handler + + +cdef extern from "cudf/io/experimental/hybrid_scan_multifile.hpp" \ + namespace "cudf::io::parquet::experimental" nogil: + + cdef cppclass hybrid_scan_multifile: + hybrid_scan_multifile( + host_span[const_FileMetaData] parquet_metadata, + const parquet_reader_options& options + ) except +libcudf_exception_handler + + vector[FileMetaData] parquet_metadatas() except +libcudf_exception_handler + + vector[byte_range_info] page_index_byte_ranges() except +libcudf_exception_handler + + void setup_page_indexes( + host_span[const_host_span_const_uint8_t] page_index_bytes + ) except +libcudf_exception_handler + + vector[vector[size_type]] all_row_groups( + const parquet_reader_options& options + ) except +libcudf_exception_handler + + size_type total_rows_in_row_groups( + host_span[const_vector_size_type] row_group_indices + ) except +libcudf_exception_handler + + unique_ptr[column] build_all_true_row_mask( + host_span[const_vector_size_type] row_group_indices, + cudaStream_t stream, + device_async_resource_ref mr + ) except +libcudf_exception_handler + + vector[vector[byte_range_info]] payload_column_chunks_byte_ranges( + host_span[const_vector_size_type] row_group_indices, + const column_view& row_mask, + use_data_page_mask mask_data_pages, + const parquet_reader_options& options, + cudaStream_t stream + ) except +libcudf_exception_handler + + void setup_chunking_for_payload_columns( + size_t chunk_read_limit, + size_t pass_read_limit, + host_span[const_vector_size_type] row_group_indices, + const column_view& row_mask, + use_data_page_mask mask_data_pages, + host_span[const_vector_device_span_const_uint8_t] page_data_per_source, + const parquet_reader_options& options, + cudaStream_t stream, + device_async_resource_ref mr + ) except +libcudf_exception_handler + + table_with_metadata materialize_payload_columns_chunk( + const column_view& row_mask + ) except +libcudf_exception_handler + + vector[vector[vector[size_type]]] construct_row_group_passes( + host_span[const_vector_size_type] row_group_indices, + size_t pass_read_limit + ) except +libcudf_exception_handler + + bool has_next_table_chunk() except +libcudf_exception_handler diff --git a/python/pylibcudf/tests/io/test_experimental_hybrid_scan.py b/python/pylibcudf/tests/io/test_experimental_hybrid_scan.py index 1c6dc4855b85..e373708006b8 100644 --- a/python/pylibcudf/tests/io/test_experimental_hybrid_scan.py +++ b/python/pylibcudf/tests/io/test_experimental_hybrid_scan.py @@ -781,3 +781,99 @@ def test_hybrid_scan_metadata_with_page_index( assert row_mask is not None assert row_mask.size() > 0 assert row_mask.type().id() == plc.types.TypeId.BOOL8 + + +@pytest.mark.parametrize("stream", [None, Stream()]) +def test_hybrid_scan_multifile_payload_page_scan( + tmp_path, + simple_parquet_table: pa.Table, + row_group_size: int, + stream: Stream | None, +) -> None: + """Read payload pages from multiple sources using the multifile reader.""" + paths = [tmp_path / f"source-{i}.parquet" for i in range(2)] + for index, path in enumerate(paths): + table = simple_parquet_table.slice( + index * 500, 500 + ) + pq.write_table( + table, + path, + row_group_size=row_group_size, + use_dictionary=True, + write_page_index=True, + ) + + source = plc.io.SourceInfo([str(path) for path in paths]) + options = plc.io.parquet.ParquetReaderOptions.builder(source).build() + options.set_column_names(["col1"]) + metadatas = plc.io.parquet_metadata.read_parquet_footers(source) + reader = plc.io.experimental.HybridScanMultiFile.from_parquet_metadatas( + metadatas, options + ) + + assert [metadata.num_rows for metadata in reader.parquet_metadatas()] == [ + 500, + 500, + ] + row_groups = reader.all_row_groups(options) + assert row_groups == [[0, 1], [0, 1]] + assert reader.total_rows_in_row_groups(row_groups) == 1000 + + page_ranges = reader.page_index_byte_ranges() + page_indexes = [] + for path, byte_range in zip(paths, page_ranges, strict=True): + with path.open("rb") as file: + file.seek(byte_range.offset) + page_indexes.append(file.read(byte_range.size)) + reader.setup_page_indexes(page_indexes) + + row_mask = reader.build_all_true_row_mask(row_groups, stream) + assert row_mask.size() == 1000 + + passes = reader.construct_row_group_passes(row_groups, 0) + assert passes == [row_groups] + page_ranges_per_source = reader.payload_column_chunks_byte_ranges( + passes[0], + row_mask, + UseDataPageMask.YES, + options, + stream, + ) + assert len(page_ranges_per_source) == len(paths) + assert all(ranges for ranges in page_ranges_per_source) + + page_data_per_source = [] + for path, byte_ranges in zip(paths, page_ranges_per_source, strict=True): + source_data = [] + with path.open("rb") as file: + for byte_range in byte_ranges: + file.seek(byte_range.offset) + source_data.append( + plc.gpumemoryview( + rmm.DeviceBuffer.to_device( + file.read(byte_range.size), + plc.utils._get_stream(stream), + ) + ) + ) + page_data_per_source.append(source_data) + + synchronize_stream(stream) + reader.setup_chunking_for_payload_columns( + 0, + 0, + passes[0], + row_mask, + UseDataPageMask.YES, + page_data_per_source, + options, + stream, + ) + output_rows = 0 + while reader.has_next_table_chunk(): + output = reader.materialize_payload_columns_chunk(row_mask) + assert output.tbl.num_columns() == 1 + output_rows += output.tbl.num_rows() + synchronize_stream(stream) + assert output_rows == 1000