From 6f2d7737cd8484f1b302ebdde7a2cdaec2d8411c Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 20 Aug 2026 06:04:02 +0000 Subject: [PATCH 1/2] Support multi-field extraction from variant columns Pulling many fields out of one VARIANT column meant calling get_variant_field once per path, which re-resolved every shared prefix and re-launched the locate and copy passes per field. Add get_variant_fields and extract_variant_fields, which merge the requested paths into a prefix trie so a shared prefix is resolved once per row, then locate every path in one kernel and copy every (path, row) pair in one batched pass. --- .../parquet/experimental/variant/extract.cpp | 116 +++++ cpp/include/cudf/io/experimental/variant.hpp | 51 +++ .../parquet/experimental/variant_extract.cu | 396 +++++++++++++++++- .../io/parquet/experimental/variant_path.cpp | 101 +++++ .../io/parquet/experimental/variant_path.hpp | 46 ++ .../io/experimental/variant_extract_test.cpp | 198 +++++++++ 6 files changed, 898 insertions(+), 10 deletions(-) diff --git a/cpp/benchmarks/io/parquet/experimental/variant/extract.cpp b/cpp/benchmarks/io/parquet/experimental/variant/extract.cpp index c2ea942d8ec3..abe59646c29a 100644 --- a/cpp/benchmarks/io/parquet/experimental/variant/extract.cpp +++ b/cpp/benchmarks/io/parquet/experimental/variant/extract.cpp @@ -320,6 +320,41 @@ std::vector build_flat_object(int num_fields, return out; } +// Dictionary key of the i-th leaf field: "f00", "f01", ... +std::string field_key(int i) { return "f" + std::string(i < 10 ? "0" : "") + std::to_string(i); } + +// Build a VARIANT object from (field id, value) pairs. Ids must be ascending, which for a +// name-sorted dictionary is also name order, as the spec requires. Field offsets widen to 2 bytes +// once the values region outgrows a single byte. +std::vector build_object( + std::vector>> const& fields) +{ + constexpr std::size_t max_single_byte_offset = 255; + auto const values_bytes = + std::accumulate(fields.begin(), fields.end(), std::size_t{0}, [](auto acc, auto const& field) { + return acc + field.second.size(); + }); + int const offset_size = values_bytes > max_single_byte_offset ? 2 : 1; + + // object value_header: | is_large (1) | field_id_size-1 (2) | field_offset_size-1 (2) | + std::vector out{ + make_variant_header(variant_basic_type::OBJECT, static_cast(offset_size - 1)), + static_cast(fields.size())}; + for (auto const& [id, value] : fields) { + out.push_back(id); + } + std::size_t running = 0; + for (auto const& [id, value] : fields) { + append_le(out, running, offset_size); + running += value.size(); + } + append_le(out, running, offset_size); // sentinel offset after the last field + for (auto const& [id, value] : fields) { + out.insert(out.end(), value.begin(), value.end()); + } + return out; +} + // Build the JSONPath-like extraction path. // For nesting=2, type=array: "a.b[1]" // For nesting=3, type=string: "a.b.c" @@ -508,3 +543,84 @@ NVBENCH_BENCH(bench_variant_extract_fields) .add_int64_axis("num_fields", {1, 10, 100}) .add_string_axis("field_position", {"first", "last"}) .add_int64_axis("hit_rate", {20, 80}); + +// Compares extracting many fields in one batched call against looping the single-field API, with +// and without a prefix shared by all the requested paths. Type is fixed to int32_t; both APIs +// decode, so the two sides of the comparison are end-to-end equivalent. +static void bench_variant_extract_multi_field(nvbench::state& state) +{ + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + + auto const num_rows = static_cast(state.get_int64("num_rows")); + auto const num_fields = static_cast(state.get_int64("num_fields")); + auto const hit_rate = static_cast(state.get_int64("hit_rate")); + bool const shared_prefix = state.get_string("prefix") == "shared"; + bool const batched = state.get_string("api") == "batched"; + + // Dictionary: the shared parent key "a", the leaf keys f00..f{N-1}, and "z" for miss rows, in + // sorted order. Field ids are dictionary indices, so "a" is 0 and leaf `i` is `i + 1`. + std::vector keys{"a"}; + for (int i = 0; i < num_fields; ++i) { + keys.push_back(field_key(i)); + } + keys.emplace_back("z"); + auto const meta_blob = build_metadata(keys); + + auto const leaf = build_leaf_value(bench_variant_type::INT32); + std::vector>> leaf_fields; + for (int i = 0; i < num_fields; ++i) { + leaf_fields.emplace_back(static_cast(i + 1), leaf); + } + // Shared-prefix layout nests every leaf under "a"; the disjoint layout puts them at the top + // level. + auto const leaf_object = build_object(leaf_fields); + auto hit_val = shared_prefix ? build_object({{uint8_t{0}, leaf_object}}) : leaf_object; + // Miss rows hold an object keyed on "z" alone, so every path fails at its first step. + auto miss_val = build_object({{static_cast(num_fields + 1), leaf}}); + pad_to_equal_size(hit_val, miss_val); + + std::vector> meta_spans(num_rows, std::span{meta_blob}); + auto val_spans = fill_val_rows(num_rows, hit_val, miss_val, hit_rate); + auto col = build_variant_column(meta_spans, val_spans, stream, mr); + CUDF_CUDA_TRY(cudaStreamSynchronize(stream.value())); + + std::vector path_strings; + path_strings.reserve(num_fields); + for (int i = 0; i < num_fields; ++i) { + path_strings.push_back((shared_prefix ? "$.a." : "$.") + field_key(i)); + } + std::vector const paths(path_strings.begin(), path_strings.end()); + auto const target_type = cudf::data_type{cudf::type_id::INT32}; + std::vector const target_types(num_fields, target_type); + + auto const data_size = static_cast(num_rows) * (meta_blob.size() + hit_val.size()); + + auto mem_stats_logger = cudf::memory_stats_logger(); + mr = cudf::get_current_device_resource_ref(); + state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value())); + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) { + if (batched) { + std::ignore = cudf::io::parquet::experimental::extract_variant_fields( + col->view(), paths, target_types, stream, mr); + } else { + for (auto const& path : path_strings) { + std::ignore = cudf::io::parquet::experimental::extract_variant_field( + col->view(), path, target_type, stream, mr); + } + } + }); + + auto const time = state.get_summary("nv/cold/time/gpu/mean").get_float64("value"); + state.add_element_count(static_cast(data_size) / time, "bytes_per_second"); + state.add_buffer_size( + mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); +} + +NVBENCH_BENCH(bench_variant_extract_multi_field) + .set_name("bench_variant_extract_multi_field") + .add_int64_axis("num_rows", {262144, 2097152}) + .add_int64_axis("num_fields", {1, 4, 16, 64}) + .add_string_axis("prefix", {"shared", "disjoint"}) + .add_string_axis("api", {"batched", "looped"}) + .add_int64_axis("hit_rate", {80}); diff --git a/cpp/include/cudf/io/experimental/variant.hpp b/cpp/include/cudf/io/experimental/variant.hpp index 22e94d116328..c13589876987 100644 --- a/cpp/include/cudf/io/experimental/variant.hpp +++ b/cpp/include/cudf/io/experimental/variant.hpp @@ -8,9 +8,11 @@ #include #include #include +#include #include #include #include +#include #include @@ -110,6 +112,55 @@ namespace io::parquet::experimental { cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +/** + * @brief Extract the raw VARIANT-encoded bytes of several nested fields in one pass. + * + * Equivalent to calling `get_variant_field` once per path, but paths that share a prefix resolve + * that prefix once per row instead of once per path, and all paths are located by a single kernel + * and copied by a single pass. + * + * @param variant_column Struct column (VARIANT materialization) with `list` children + * (`metadata`, `value`), plus optional shredded siblings + * @param paths JSONPath-like path strings (see `get_variant_field` for syntax). Duplicate and + * overlapping paths are allowed + * @param stream CUDA stream + * @param mr Device memory resource + * @return Table of one `list` column per path, in the order the paths were given. Row + * nullability matches `get_variant_field` + * + * @throws std::invalid_argument if any path is empty or malformed + */ +[[nodiscard]] std::unique_ptr get_variant_fields( + column_view const& variant_column, + host_span paths, + cuda::stream_ref stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + +/** + * @brief Convenience wrapper: extract several nested fields by path and decode each into a typed + * column. + * + * Semantically equivalent to `extract_variant_field` per path, sharing the work of resolving common + * path prefixes. + * + * @param variant_column Struct column (VARIANT materialization) + * @param paths JSONPath-like path strings (see `get_variant_field` for syntax) + * @param desired_types Target type of each path's output column; parallels `paths`. Supported types + * are those of `cast_variant` + * @param stream CUDA stream + * @param mr Device memory resource + * @return Table of one column per path, in the order the paths were given + * + * @throws std::invalid_argument if any path is empty or malformed, if `paths` and `desired_types` + * differ in size, or if a desired type is unsupported + */ +[[nodiscard]] std::unique_ptr
extract_variant_fields( + column_view const& variant_column, + host_span paths, + host_span desired_types, + cuda::stream_ref stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + /** * @brief Return the logical type of each VARIANT value blob in a `list` column. * diff --git a/cpp/src/io/parquet/experimental/variant_extract.cu b/cpp/src/io/parquet/experimental/variant_extract.cu index 11507e6839d8..de4efd3e3de0 100644 --- a/cpp/src/io/parquet/experimental/variant_extract.cu +++ b/cpp/src/io/parquet/experimental/variant_extract.cu @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -35,6 +36,7 @@ #include #include #include +#include #include #include #include @@ -42,7 +44,9 @@ #include #include +#include #include +#include #include #include #include @@ -508,19 +512,25 @@ __device__ cuda::std::optional parse_index_step(cudf::string_view ste return index; } -// Walk a path of object-key or array-index steps level by level starting at `val` and return -// the span of the final value (subspan of `val`). Returns an empty span on failure. +// Walk the steps `path[step_begin : step_end]` level by level starting at `val` and return the span +// of the final value (subspan of `val`). Returns an empty span on failure. // // Each path step is encoded in the `path` strings column as either: // - "" -> descend into an object by dictionary key, or // - "[]" -> descend into an array by zero-based integer index. // The step kind is inferred from the first byte (`'['` means index). -__device__ device_span resolve_path(device_span meta, - device_span val, - column_device_view path) +__device__ device_span resolve_steps(device_span meta, + device_span val, + column_device_view path, + size_type step_begin, + size_type step_end) { + // An empty starting value cannot resolve anything, and checking up front keeps a failed shared + // prefix from paying for a metadata lookup once per path below it. + if (val.empty()) { return {}; } + device_span sub_val = val; - for (size_type i = 0; i < path.size(); ++i) { + for (size_type i = step_begin; i < step_end; ++i) { auto const step = path.element(i); if (step.size_bytes() >= 1 && step.data()[0] == '[') { @@ -537,6 +547,14 @@ __device__ device_span resolve_path(device_span me return sub_val; } +// Walk every step of `path` starting at `val`. +__device__ device_span resolve_path(device_span meta, + device_span val, + column_device_view path) +{ + return resolve_steps(meta, val, path, 0, path.size()); +} + __device__ cuda::std::optional> decode_string( device_span enc) { @@ -625,6 +643,131 @@ CUDF_KERNEL __launch_bounds__(block_size) void locate_variant_fields_kernel( } } +/** + * @brief Where one trie slot's value was found within a row's value blob. + * + * Validity lives entirely in `size`: a slot whose steps did not resolve for this row has + * `invalid_slot_size`, and resolving anything below it fails immediately because its span is empty. + */ +struct slot_result { + size_type src_offset; + size_type size; +}; + +constexpr size_type invalid_slot_size = -1; + +__device__ slot_result make_slot_result(device_span field, uint8_t const* val_base) +{ + if (field.empty()) { return {0, invalid_slot_size}; } + return {static_cast(field.data() - val_base), static_cast(field.size())}; +} + +__device__ bool slot_is_valid(slot_result const& result) +{ + return result.size != invalid_slot_size; +} + +// The span a slot located, or an empty span if it did not resolve. +__device__ device_span slot_span(device_span val, + slot_result const& result) +{ + if (!slot_is_valid(result)) { return {}; } + return val.subspan(result.src_offset, result.size); +} + +// Tries up to this deep are walked with a per-thread stack the compiler can keep in registers or +// local memory; a deeper one uses a global scratch allocation instead. +constexpr size_type max_local_trie_depth = 16; + +// Global scratch is allocated per thread, so the grid has to be capped for the allocation to stay +// independent of the row count. This many blocks still saturates the walk. +constexpr int max_global_scratch_blocks = 256; + +/** + * @brief Resolves a whole trie of VARIANT paths in each row, recording each path's result. + * + * Slots are visited in index order, which is depth-first pre-order, so the walk only has to + * remember one located value per depth: a slot's parent is the entry one level up, still untouched + * from when the walk descended. A shared prefix is therefore resolved once per row and reused by + * every path below it, and a prefix that fails leaves an empty span behind that makes its whole + * subtree fail at its first step. + * + * For each path `p` and row, the located field's byte length is written to + * `d_sizes[p * num_rows + row]` and its offset within the row's value blob to `d_src_offsets`, so + * each path's outputs are contiguous. Rows that are null in `d_row_valid`, or whose path does not + * resolve, get a size of 0 and are marked null in that path's mask in `d_null_masks`. + * + * @tparam UseLocalScratch Keep the per-thread depth stack in local memory rather than `d_scratch` + */ +template +CUDF_KERNEL __launch_bounds__(block_size) void locate_variant_field_trie_kernel( + cudf::lists_column_device_view metadata, + cudf::lists_column_device_view values, + column_device_view steps, + device_span slot_steps, + device_span slot_depth, + device_span output_offsets, + device_span output_paths, + bitmask_type const* d_row_valid, + size_type num_rows, + size_type trie_depth, + device_span d_sizes, + device_span d_src_offsets, + device_span d_null_masks, + device_span d_scratch) +{ + auto const num_slots = static_cast(slot_depth.size()); + auto const num_paths = static_cast(d_null_masks.size()); + auto const tid = cudf::detail::grid_1d::global_thread_id(); + auto const stride = cudf::detail::grid_1d::grid_stride(); + + [[maybe_unused]] cuda::std::array + local_stack; + auto* const located = [&]() -> slot_result* { + if constexpr (UseLocalScratch) { + return local_stack.data(); + } else { + return d_scratch.data() + tid * trie_depth; + } + }(); + + for (auto row = tid; row < num_rows; row += stride) { + bool const row_valid = d_row_valid == nullptr || cudf::bit_is_set(d_row_valid, row); + + if (!row_valid) { + for (size_type path = 0; path < num_paths; ++path) { + auto const out = path * num_rows + static_cast(row); + d_sizes[out] = 0; + d_src_offsets[out] = 0; + cudf::clear_bit(d_null_masks[path], row); + } + continue; + } + + auto const [meta, val] = metadata_and_value_at(metadata, values, row); + + for (size_type slot = 0; slot < num_slots; ++slot) { + auto const depth = slot_depth[slot]; + auto const parent = depth == 0 ? val : slot_span(val, located[depth - 1]); + auto const field = resolve_steps(meta, parent, steps, slot_steps[slot], slot_steps[slot + 1]); + located[depth] = make_slot_result(field, val.data()); + + for (auto out_idx = output_offsets[slot]; out_idx < output_offsets[slot + 1]; ++out_idx) { + auto const path = output_paths[out_idx]; + auto const out = path * num_rows + static_cast(row); + if (slot_is_valid(located[depth])) { + d_sizes[out] = located[depth].size; + d_src_offsets[out] = located[depth].src_offset; + } else { + d_sizes[out] = 0; + d_src_offsets[out] = 0; + cudf::clear_bit(d_null_masks[path], row); + } + } + } + } +} + /** * @brief Per-row kernel: decode each VARIANT value blob into a fixed-width primitive of type `T`. * @@ -699,6 +842,13 @@ struct cast_variant_string_fn { } }; +// An empty `list` column: the shape a VARIANT field extraction produces for an empty input. +std::unique_ptr make_empty_variant_value_column() +{ + return cudf::make_lists_column( + 0, make_empty_column(type_id::INT32), make_empty_column(type_id::UINT8), 0, {}); +} + void validate_variant_child(column_view const& child) { CUDF_EXPECTS(child.type().id() == type_id::LIST, @@ -880,10 +1030,7 @@ std::unique_ptr get_variant_field(column_view const& variant_column, auto const steps = parse_variant_path(path); auto const num_rows = variant_column.size(); - if (num_rows == 0) { - return cudf::make_lists_column( - 0, make_empty_column(type_id::INT32), make_empty_column(type_id::UINT8), 0, {}); - } + if (num_rows == 0) { return make_empty_variant_value_column(); } auto const temp_mr = cudf::get_current_device_resource_ref(); @@ -958,6 +1105,195 @@ std::unique_ptr get_variant_field(column_view const& variant_column, null_count > 0 ? std::move(null_mask) : rmm::device_buffer{}); } +std::unique_ptr
get_variant_fields(column_view const& variant_column, + host_span paths, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + // Validate the variant column + CUDF_EXPECTS(variant_column.type().id() == type_id::STRUCT, + "VARIANT column must be struct type", + std::invalid_argument); + CUDF_EXPECTS(variant_column.num_children() >= 2, + "VARIANT struct must have at least two children", + std::invalid_argument); + validate_variant_child(variant_column.child(0)); + validate_variant_child(variant_column.child(1)); + + auto const num_paths = static_cast(paths.size()); + auto const num_rows = variant_column.size(); + + std::vector> output; + output.reserve(num_paths); + if (num_paths == 0) { return std::make_unique
(std::move(output)); } + + // A single path has no prefixes to share, so it is exactly the single-path entry point; the + // batched setup would only add fixed overhead. + if (num_paths == 1) { + output.push_back(get_variant_field(variant_column, paths.front(), stream, mr)); + return std::make_unique
(std::move(output)); + } + + // Validate and merge the paths even for empty input columns + auto const trie = build_variant_path_trie(paths); + + if (num_rows == 0) { + std::generate_n( + std::back_inserter(output), num_paths, [] { return make_empty_variant_value_column(); }); + return std::make_unique
(std::move(output)); + } + + auto const temp_mr = cudf::get_current_device_resource_ref(); + + auto steps_column = build_path_column(trie.steps, stream, temp_mr); + auto steps_device_view = column_device_view::create(steps_column->view(), stream); + auto const d_slot_steps = + cudf::detail::make_device_uvector_async(trie.slot_steps, stream, temp_mr); + auto const d_slot_depth = + cudf::detail::make_device_uvector_async(trie.slot_depth, stream, temp_mr); + auto const d_output_offsets = + cudf::detail::make_device_uvector_async(trie.output_offsets, stream, temp_mr); + auto const d_output_paths = + cudf::detail::make_device_uvector_async(trie.output_paths, stream, temp_mr); + + // Resolve children with respect to any slice/offset on the parent struct + structs_column_view const variant_struct{variant_column}; + auto const meta_view = variant_struct.get_sliced_child(0, stream); + auto const val_view = variant_struct.get_sliced_child(1, stream); + + auto meta_device_view = column_device_view::create(meta_view, stream); + auto val_device_view = column_device_view::create(val_view, stream); + cudf::lists_column_device_view meta_lists_device_view(*meta_device_view); + cudf::lists_column_device_view val_lists_device_view(*val_device_view); + + // Input row validity, copied so that it is indexable by row regardless of any slice offset + auto const row_mask = variant_column.nullable() + ? cudf::detail::copy_bitmask(variant_column, stream, temp_mr) + : rmm::device_buffer{}; + auto const* d_row_valid = static_cast(row_mask.data()); + + // Per-path outputs are contiguous, so each path's sizes can be scanned on their own + CUDF_EXPECTS(static_cast(num_paths) * num_rows <= std::numeric_limits::max(), + "VARIANT paths times rows exceeds cudf size_type limit", + std::overflow_error); + auto const num_outputs = num_paths * num_rows; + rmm::device_uvector d_sizes(num_outputs, stream, temp_mr); + rmm::device_uvector d_src_offsets(num_outputs, stream, temp_mr); + + // One null mask per output column, narrowed from all-valid by the walk + std::vector null_masks; + null_masks.reserve(num_paths); + std::vector h_null_masks(num_paths); + for (size_type p = 0; p < num_paths; ++p) { + null_masks.push_back(cudf::create_null_mask(num_rows, mask_state::ALL_VALID, stream, mr)); + h_null_masks[p] = static_cast(null_masks.back().data()); + } + auto const d_null_masks = cudf::detail::make_device_uvector_async(h_null_masks, stream, temp_mr); + + // Resolve the whole trie per row and compute the output sizes. The walk keeps one located value + // per trie level, so only a pathologically deep trie needs scratch outside the thread. + auto const trie_depth = 1 + *std::max_element(trie.slot_depth.begin(), trie.slot_depth.end()); + bool const use_local_scratch = trie_depth <= max_local_trie_depth; + auto const grid = cudf::detail::grid_1d{num_rows, block_size}; + auto const num_blocks = + use_local_scratch + ? grid.num_blocks + : std::min(grid.num_blocks, static_cast(max_global_scratch_blocks)); + rmm::device_uvector d_scratch( + use_local_scratch ? 0 : static_cast(num_blocks) * block_size * trie_depth, + stream, + temp_mr); + + auto const launch = [&](auto use_local) { + locate_variant_field_trie_kernel + <<>>(meta_lists_device_view, + val_lists_device_view, + *steps_device_view, + d_slot_steps, + d_slot_depth, + d_output_offsets, + d_output_paths, + d_row_valid, + num_rows, + trie_depth, + d_sizes, + d_src_offsets, + d_null_masks, + d_scratch); + }; + if (use_local_scratch) { + launch(cuda::std::true_type{}); + } else { + launch(cuda::std::false_type{}); + } + CUDF_CUDA_TRY(cudaGetLastError()); + + // Convert each path's sizes to offsets and allocate its output bytes + std::vector> offsets_columns; + std::vector> value_children; + offsets_columns.reserve(num_paths); + value_children.reserve(num_paths); + std::vector h_offsets(num_paths); + std::vector h_out(num_paths); + int64_t all_paths_bytes = 0; + for (size_type p = 0; p < num_paths; ++p) { + device_span const path_sizes{ + d_sizes.data() + static_cast(p) * num_rows, static_cast(num_rows)}; + auto [offsets_column, total_bytes] = + cudf::strings::detail::make_offsets_child_column(path_sizes, stream, mr); + CUDF_EXPECTS(total_bytes <= std::numeric_limits::max(), + "VARIANT extracted bytes exceed cudf size_type limit", + std::overflow_error); + + auto value_child = make_numeric_column(data_type{type_id::UINT8}, + static_cast(total_bytes), + mask_state::UNALLOCATED, + stream, + mr); + h_offsets[p] = offsets_column->view().data(); + h_out[p] = value_child->mutable_view().data(); + all_paths_bytes += total_bytes; + offsets_columns.push_back(std::move(offsets_column)); + value_children.push_back(std::move(value_child)); + } + + // Copy the located values of every (path, row) pair in one pass + if (all_paths_bytes > 0) { + auto const d_offsets = cudf::detail::make_device_uvector_async(h_offsets, stream, temp_mr); + auto const d_out = cudf::detail::make_device_uvector_async(h_out, stream, temp_mr); + + auto src_iter = cudf::detail::make_counting_transform_iterator( + size_type{0}, + cuda::proclaim_return_type( + [vlv = val_lists_device_view, d_src = d_src_offsets.data(), num_rows] __device__( + size_type i) -> uint8_t const* { + auto const row = i % num_rows; + return vlv.child().template data() + vlv.offset_at(row) + d_src[i]; + })); + auto dst_iter = cudf::detail::make_counting_transform_iterator( + size_type{0}, + cuda::proclaim_return_type([d_off = d_offsets.data(), + d_dst = d_out.data(), + num_rows] __device__(size_type i) -> uint8_t* { + return d_dst[i / num_rows] + d_off[i / num_rows][i % num_rows]; + })); + cudf::detail::batched_memcpy_async(src_iter, dst_iter, d_sizes.begin(), num_outputs, stream); + } + + for (size_type p = 0; p < num_paths; ++p) { + auto const null_count = + num_rows - cudf::detail::count_set_bits(h_null_masks[p], 0, num_rows, stream); + output.push_back( + make_lists_column(num_rows, + std::move(offsets_columns[p]), + std::move(value_children[p]), + null_count, + null_count > 0 ? std::move(null_masks[p]) : rmm::device_buffer{})); + } + + return std::make_unique
(std::move(output)); +} + std::unique_ptr cast_variant(column_view const& values, data_type desired_type, cuda::stream_ref stream, @@ -999,6 +1335,27 @@ std::unique_ptr cast_variant(column_view const& values, mr}); } +std::unique_ptr
extract_variant_fields(column_view const& variant_column, + host_span paths, + host_span desired_types, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + CUDF_EXPECTS(paths.size() == desired_types.size(), + "VARIANT paths and desired types must have the same size", + std::invalid_argument); + + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const values = get_variant_fields(variant_column, paths, stream, temp_mr); + + std::vector> output; + output.reserve(paths.size()); + for (size_type p = 0; p < values->num_columns(); ++p) { + output.push_back(cast_variant(values->get_column(p).view(), desired_types[p], stream, mr)); + } + return std::make_unique
(std::move(output)); +} + std::unique_ptr get_variant_type_id(column_view const& values, cuda::stream_ref stream, rmm::device_async_resource_ref mr) @@ -1078,5 +1435,24 @@ std::unique_ptr extract_variant_field(column_view const& variant_column, return detail::cast_variant(value->view(), desired_type, stream, mr); } +std::unique_ptr
get_variant_fields(column_view const& variant_column, + host_span paths, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + return detail::get_variant_fields(variant_column, paths, stream, mr); +} + +std::unique_ptr
extract_variant_fields(column_view const& variant_column, + host_span paths, + host_span desired_types, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + return detail::extract_variant_fields(variant_column, paths, desired_types, stream, mr); +} + } // namespace io::parquet::experimental } // namespace cudf diff --git a/cpp/src/io/parquet/experimental/variant_path.cpp b/cpp/src/io/parquet/experimental/variant_path.cpp index 58dc519e7965..7f888e7a384b 100644 --- a/cpp/src/io/parquet/experimental/variant_path.cpp +++ b/cpp/src/io/parquet/experimental/variant_path.cpp @@ -10,10 +10,13 @@ #include #include +#include +#include #include #include #include #include +#include #include namespace cudf::io::parquet::experimental::detail { @@ -61,6 +64,21 @@ namespace { return std::string{tail.substr(0, n + 1)}; // include the closing ']' } +// A prefix-tree node built while merging paths. Children are keyed by step token, which is also how +// the device walk tells the steps apart, so an index step and a same-spelled name cannot collide. +// The map keeps the child order deterministic across runs. +struct trie_builder_node { + std::map children; + bool ends_a_path = false; +}; + +// A prefix is worth caching only if more than one path continues past it, or a path ends there. +// Any other prefix is folded into its single descendant's step range. +[[nodiscard]] bool needs_slot(trie_builder_node const& node) +{ + return node.ends_a_path || node.children.size() > 1; +} + } // namespace std::vector parse_variant_path(std::string_view path) @@ -99,4 +117,87 @@ std::vector parse_variant_path(std::string_view path) return steps; } +variant_path_trie build_variant_path_trie(host_span paths) +{ + // Node 0 is the root: the value blob itself, before any step is applied. + std::vector nodes(1); + std::vector path_end_node(paths.size()); + + for (std::size_t p = 0; p < paths.size(); ++p) { + std::size_t node = 0; + for (auto const& step : parse_variant_path(paths[p])) { + auto const child = nodes[node].children.find(step); + if (child != nodes[node].children.end()) { + node = child->second; + } else { + auto const new_node = nodes.size(); + nodes.emplace_back(); + nodes[node].children.emplace(step, new_node); + node = new_node; + } + } + nodes[node].ends_a_path = true; + path_end_node[p] = node; + } + + variant_path_trie trie; + trie.slot_steps.push_back(0); + std::vector slot_of_node(nodes.size(), -1); + + // Depth-first descent from the root, emitting slots in pre-order. `pending` carries the steps of + // the collapsed single-child chain walked since the last slot, and becomes that slot's step + // range. + struct descent_state { + std::size_t node; + size_type depth; + std::vector pending; + }; + std::vector stack; + stack.push_back({0, 0, {}}); + + while (!stack.empty()) { + auto state = std::move(stack.back()); + stack.pop_back(); + + // The root is the value blob itself and never becomes a slot. + auto child_depth = state.depth; + std::vector child_pending; + if (state.node != 0 && needs_slot(nodes[state.node])) { + slot_of_node[state.node] = static_cast(trie.slot_depth.size()); + trie.steps.insert(trie.steps.end(), state.pending.begin(), state.pending.end()); + trie.slot_steps.push_back(static_cast(trie.steps.size())); + trie.slot_depth.push_back(state.depth); + child_depth = state.depth + 1; + } else { + child_pending = std::move(state.pending); + } + + // Reverse order, so that popping visits the first child first and keeps each subtree + // contiguous. + for (auto const& [step, child] : std::ranges::reverse_view(nodes[state.node].children)) { + auto pending = child_pending; + pending.push_back(step); + stack.push_back({child, child_depth, std::move(pending)}); + } + } + + // Invert the path-to-slot mapping into CSR form, counting then filling. + auto const num_slots = trie.slot_depth.size(); + trie.output_offsets.assign(num_slots + 1, 0); + for (auto const node : path_end_node) { + ++trie.output_offsets[slot_of_node[node] + 1]; + } + for (std::size_t slot = 0; slot < num_slots; ++slot) { + trie.output_offsets[slot + 1] += trie.output_offsets[slot]; + } + trie.output_paths.resize(paths.size()); + auto fill_position = trie.output_offsets; + size_type path_index = 0; + for (auto const node : path_end_node) { + trie.output_paths[fill_position[slot_of_node[node]]++] = path_index++; + } + + return trie; +} + } // namespace cudf::io::parquet::experimental::detail diff --git a/cpp/src/io/parquet/experimental/variant_path.hpp b/cpp/src/io/parquet/experimental/variant_path.hpp index 760b36a7962c..db6e25045eaf 100644 --- a/cpp/src/io/parquet/experimental/variant_path.hpp +++ b/cpp/src/io/parquet/experimental/variant_path.hpp @@ -5,6 +5,9 @@ #pragma once +#include +#include + #include #include #include @@ -30,4 +33,47 @@ namespace cudf::io::parquet::experimental::detail { */ [[nodiscard]] std::vector parse_variant_path(std::string_view path); +/** + * @brief A set of VARIANT paths merged into a prefix tree, flattened for a device walk. + * + * Slots are the units of work of the device walk: slot `k` starts from the value located by its + * parent slot (or from the whole value blob when `slot_depth[k]` is 0) and applies the steps + * `steps[slot_steps[k] : slot_steps[k + 1]]`. + * + * Only prefixes worth caching get a slot: a prefix shared by more than one path (a branch point) or + * one that ends a path. A prefix with a single continuation and no path ending on it is folded into + * its descendant's step range, so `$.user.addr.zip` is one slot holding three steps rather than + * three slots. + * + * Slots are in depth-first pre-order, so a slot's subtree is contiguous and immediately follows it. + * A walk can therefore keep one located value per depth: when it reaches slot `k`, the entry at + * `slot_depth[k] - 1` still holds its parent's value. That bounds the walk's scratch by the depth + * of the trie rather than by its slot count. + * + * `output_offsets` and `output_paths` are the reverse of the path-to-slot mapping in CSR form: the + * value of slot `k` is the output of the input paths `output_paths[output_offsets[k] : + * output_offsets[k + 1]]`, which is empty for a slot that only exists as a shared prefix. Duplicate + * paths land on one slot with several outputs. + */ +struct variant_path_trie { + std::vector steps; ///< Step tokens of every slot, concatenated in slot order + std::vector slot_steps; ///< Size `num_slots + 1`; each slot's range in `steps` + std::vector slot_depth; ///< Depth of each slot; 0 starts from the value blob + std::vector output_offsets; ///< Size `num_slots + 1`; CSR offsets into `output_paths` + std::vector output_paths; ///< Input paths that each slot's value is the output of +}; + +/** + * @brief Merge VARIANT paths into a `variant_path_trie` so shared prefixes resolve once per row. + * + * Each path is parsed with `parse_variant_path`, so every path is validated before any of them is + * used. + * + * @param paths JSONPath-like path strings, in output order + * @return The flattened trie + * + * @throws std::invalid_argument if any path is empty or malformed + */ +[[nodiscard]] variant_path_trie build_variant_path_trie(host_span paths); + } // namespace cudf::io::parquet::experimental::detail diff --git a/cpp/tests/io/experimental/variant_extract_test.cpp b/cpp/tests/io/experimental/variant_extract_test.cpp index 31b3f34a0897..ff1e70a77c44 100644 --- a/cpp/tests/io/experimental/variant_extract_test.cpp +++ b/cpp/tests/io/experimental/variant_extract_test.cpp @@ -1039,6 +1039,204 @@ TEST_F(GetVariantFieldTest, EmptyInput) EXPECT_EQ(cudf::lists_column_view{got->view()}.child().type().id(), cudf::type_id::UINT8); } +namespace { + +// The batched APIs must agree with the single-path API applied to each path in turn, which is what +// most of the tests below check: the single-path behavior is the specification. +void expect_matches_looped_get(cudf::column_view const& variant, + std::vector const& paths) +{ + auto const stream = cudf::test::get_default_stream(); + std::vector const path_views(paths.begin(), paths.end()); + + auto const got = cudf::io::parquet::experimental::get_variant_fields(variant, path_views, stream); + ASSERT_EQ(got->num_columns(), static_cast(paths.size())); + + for (std::size_t p = 0; p < paths.size(); ++p) { + SCOPED_TRACE(std::string{"path: "} + paths[p]); + auto const expected = + cudf::io::parquet::experimental::get_variant_field(variant, paths[p], stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(got->get_column(static_cast(p)), *expected); + } +} + +} // namespace + +struct GetVariantFieldsTest : public cudf::test::BaseFixture {}; + +TEST_F(GetVariantFieldsTest, MultiRowMatchesLoopedSingleField) +{ + auto const col = make_xyz_three_row_variant(); + expect_matches_looped_get(col, {"x", "y", "z", "no_such_field"}); +} + +TEST_F(GetVariantFieldsTest, SharedPrefixes) +{ + auto const col = make_apache_variant(avf::object_nested); + expect_matches_looped_get(col, + {"$.observation.location", + "$.observation.time", + "$.observation.value.temperature", + "$.species.name", + "$.species.population", + "$.id"}); +} + +TEST_F(GetVariantFieldsTest, DuplicateAndPrefixPaths) +{ + // "$.observation" is both a requested output and the prefix of another path, and appears twice. + auto const col = make_apache_variant(avf::object_nested); + expect_matches_looped_get( + col, {"$.observation", "$.observation.time", "$.observation", "$.observation.value"}); +} + +TEST_F(GetVariantFieldsTest, ArrayIndexSteps) +{ + // array_primitive encodes the int8 array [2, 1, 5, 9]; "[9]" is out of bounds. + auto const col = make_apache_variant(avf::array_primitive); + expect_matches_looped_get(col, {"[0]", "[3]", "[9]", "[1]"}); +} + +TEST_F(GetVariantFieldsTest, NullRowsAndSlicedInput) +{ + std::vector const m = {0x01, 0x01, 0x00, 0x01, 'x'}; + std::vector const v = {0x02, 0x01, 0x00, 0x00, 0x05, 0x14, 0x07, 0x00, 0x00, 0x00}; + cudf::test::lists_column_wrapper meta{ + {m.begin(), m.end()}, {0x00}, {m.begin(), m.end()}}; + cudf::test::lists_column_wrapper val{{v.begin(), v.end()}, {0x00}, {v.begin(), v.end()}}; + cudf::test::structs_column_wrapper const col{{meta, val}, std::vector{true, false, true}}; + + expect_matches_looped_get(col, {"x", "y"}); + expect_matches_looped_get(cudf::slice(col, {1, 3}).front(), {"x", "y"}); +} + +TEST_F(GetVariantFieldsTest, ManyPathsUseGlobalScratch) +{ + // More leaves than `max_local_trie_slots`, so the walk falls back to global slot scratch. + auto paths = std::vector{"x", "y", "z"}; + for (int i = 0; i < 40; ++i) { + paths.push_back(std::format("$.field_{}", i)); + } + expect_matches_looped_get(make_xyz_three_row_variant(), paths); +} + +TEST_F(GetVariantFieldsTest, DeepTrieUsesGlobalScratch) +{ + // Every path extends the previous one, so the trie is a chain deeper than the per-thread stack + // the walk normally uses, which pushes it onto global scratch. Only "x" resolves. + std::vector paths{"x"}; + for (int i = 0; i < 20; ++i) { + paths.push_back(std::format("{}.s{}", paths.back(), i)); + } + expect_matches_looped_get(make_xyz_three_row_variant(), paths); +} + +TEST_F(GetVariantFieldsTest, NoPathsYieldsNoColumns) +{ + auto const col = make_xyz_three_row_variant(); + auto const got = cudf::io::parquet::experimental::get_variant_fields( + col, std::vector{}, cudf::test::get_default_stream()); + EXPECT_EQ(got->num_columns(), 0); +} + +TEST_F(GetVariantFieldsTest, EmptyInput) +{ + auto const stream = cudf::test::get_default_stream(); + auto const variant = cudf::empty_like(make_xyz_three_row_variant()); + + std::vector const paths{"x", "$.y.z"}; + auto const got = cudf::io::parquet::experimental::get_variant_fields(*variant, paths, stream); + + ASSERT_EQ(got->num_columns(), 2); + for (auto const& column : got->view()) { + EXPECT_EQ(column.type().id(), cudf::type_id::LIST); + EXPECT_EQ(column.size(), 0); + EXPECT_EQ(cudf::lists_column_view{column}.child().type().id(), cudf::type_id::UINT8); + } +} + +TEST_F(GetVariantFieldsTest, MalformedPathThrows) +{ + auto const col = make_xyz_three_row_variant(); + auto const bad = std::vector{"x", "$.a[", "y"}; + EXPECT_THROW(static_cast(cudf::io::parquet::experimental::get_variant_fields( + col, bad, cudf::test::get_default_stream())), + std::invalid_argument); +} + +struct ExtractVariantFieldsTest : public cudf::test::BaseFixture {}; + +TEST_F(ExtractVariantFieldsTest, MatchesLoopedExtract) +{ + auto const col = make_xyz_three_row_variant(); + auto const stream = cudf::test::get_default_stream(); + + std::vector const paths{"x", "y", "z"}; + std::vector const types{cudf::data_type{cudf::type_id::INT32}, + cudf::data_type{cudf::type_id::STRING}, + cudf::data_type{cudf::type_id::INT32}}; + + auto const got = + cudf::io::parquet::experimental::extract_variant_fields(col, paths, types, stream); + + ASSERT_EQ(got->num_columns(), 3); + for (std::size_t p = 0; p < paths.size(); ++p) { + SCOPED_TRACE(std::string{"path: "} + std::string{paths[p]}); + auto const expected = + cudf::io::parquet::experimental::extract_variant_field(col, paths[p], types[p], stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(got->get_column(static_cast(p)), *expected); + } +} + +TEST_F(ExtractVariantFieldsTest, SharedPrefixesTypedValues) +{ + auto const col = make_apache_variant(avf::object_nested); + std::vector const paths{"$.observation.location", + "$.observation.value.temperature", + "$.species.population", + "$.species.nope"}; + std::vector const types{cudf::data_type{cudf::type_id::STRING}, + cudf::data_type{cudf::type_id::INT8}, + cudf::data_type{cudf::type_id::INT16}, + cudf::data_type{cudf::type_id::STRING}}; + + auto const got = cudf::io::parquet::experimental::extract_variant_fields( + col, paths, types, cudf::test::get_default_stream()); + + ASSERT_EQ(got->num_columns(), 4); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(got->get_column(0), + cudf::test::strings_column_wrapper({"In the Volcano"})); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(got->get_column(1), + cudf::test::fixed_width_column_wrapper{int8_t{123}}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(got->get_column(2), + cudf::test::fixed_width_column_wrapper{int16_t{6789}}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(got->get_column(3), + cudf::test::strings_column_wrapper({"donotread"}, {false})); +} + +TEST_F(ExtractVariantFieldsTest, MismatchedDesiredTypesThrows) +{ + auto const col = make_xyz_three_row_variant(); + std::vector const paths{"x", "y"}; + std::vector const types{cudf::data_type{cudf::type_id::INT32}}; + + EXPECT_THROW(static_cast(cudf::io::parquet::experimental::extract_variant_fields( + col, paths, types, cudf::test::get_default_stream())), + std::invalid_argument); +} + +TEST_F(ExtractVariantFieldsTest, UnsupportedDesiredTypeThrows) +{ + auto const col = make_xyz_three_row_variant(); + std::vector const paths{"x", "y"}; + std::vector const types{cudf::data_type{cudf::type_id::INT32}, + cudf::data_type{cudf::type_id::TIMESTAMP_DAYS}}; + + EXPECT_THROW(static_cast(cudf::io::parquet::experimental::extract_variant_fields( + col, paths, types, cudf::test::get_default_stream())), + std::invalid_argument); +} + template std::unique_ptr cast_apache_primitive(avf::fixture const& fixture) { From dfad7f78ac9edf88d047868584d79b85a8937dd5 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 20 Aug 2026 15:48:16 +0000 Subject: [PATCH 2/2] Benchmark batched variant extraction on a realistic path set The existing multi-field benchmark shares a prefix that sorts first in the dictionary, so resolving it costs one comparison and sharing it saves almost nothing. Add a workload whose shape and 50 paths come from the variant_workload example: an 85-key dictionary, a root object that nests most of its data under a mid-dictionary key, and a fan-out four to five steps deep below it. --- .../parquet/experimental/variant/extract.cpp | 220 ++++++++++++++++++ 1 file changed, 220 insertions(+) diff --git a/cpp/benchmarks/io/parquet/experimental/variant/extract.cpp b/cpp/benchmarks/io/parquet/experimental/variant/extract.cpp index abe59646c29a..614eb3387fbf 100644 --- a/cpp/benchmarks/io/parquet/experimental/variant/extract.cpp +++ b/cpp/benchmarks/io/parquet/experimental/variant/extract.cpp @@ -355,6 +355,168 @@ std::vector build_object( return out; } +// Dictionary key "item%03d" of the workload below. Zero padding makes the numeric order the +// lexicographic order, so key `n` sits at dictionary index `n - 1`. +std::string item_key(int n) +{ + auto const digits = std::to_string(n); + return "item" + std::string(3 - digits.size(), '0') + digits; +} + +// A bare VARIANT short string value. +std::vector build_short_string(std::string_view s) +{ + std::vector out{make_variant_short_string_header(s.size())}; + out.insert(out.end(), s.begin(), s.end()); + return out; +} + +// A VARIANT array value holding `elements` in order. +std::vector build_array(std::vector> const& elements) +{ + constexpr std::size_t max_single_byte_offset = 255; + auto const values_bytes = + std::accumulate(elements.begin(), elements.end(), std::size_t{0}, [](auto acc, auto const& e) { + return acc + e.size(); + }); + int const offset_size = values_bytes > max_single_byte_offset ? 2 : 1; + + // array value_header: | unused (3) | is_large (1) | offset_size-1 (2) | + std::vector out{ + make_variant_header(variant_basic_type::ARRAY, static_cast(offset_size - 1)), + static_cast(elements.size())}; + std::size_t running = 0; + for (auto const& element : elements) { + append_le(out, running, offset_size); + running += element.size(); + } + append_le(out, running, offset_size); + for (auto const& element : elements) { + out.insert(out.end(), element.begin(), element.end()); + } + return out; +} + +// A VARIANT object whose fields are named rather than pre-assigned ids. Field ids are the keys' +// positions in the sorted dictionary, and the spec wants them in name order, which for a sorted +// dictionary is id order. +std::vector build_named_object( + std::vector const& dict, + std::vector>> fields) +{ + std::ranges::sort(fields, {}, &std::pair>::first); + + std::vector>> by_id; + by_id.reserve(fields.size()); + for (auto& [key, value] : fields) { + auto const entry = std::ranges::lower_bound(dict, key); + CUDF_EXPECTS(entry != dict.end() && *entry == key, "Key missing from the VARIANT dictionary"); + by_id.emplace_back(static_cast(std::distance(dict.begin(), entry)), std::move(value)); + } + return build_object(by_id); +} + +// Build the value blob of the multi-path workload: a root object of six fields whose `item016` +// child fans out into ~40 sibling sub-trees, most of them `{item085: [{item018: "..."}]}`. The +// paths in `workload_paths` below match this shape. +std::vector build_workload_value(std::vector const& dict) +{ + auto const fanout_leaf = [&](int n) { + auto inner = build_named_object(dict, {{item_key(18), build_short_string("C_" + item_key(n))}}); + return build_named_object(dict, {{item_key(85), build_array({std::move(inner)})}}); + }; + + std::vector>> item016_fields; + for (int n : {19, 20, 21, 22, 23, 24, 25, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, + 41, 42, 43, 44, 45, 46, 47, 48, 49, 53, 54, 57, 58, 59, 61, 62}) { + item016_fields.emplace_back(item_key(n), fanout_leaf(n)); + } + // Sub-trees that several paths descend into past `item085` or `item027`. + for (int n : {26, 55}) { + auto element = build_named_object( + dict, + {{item_key(18), build_short_string("C_18")}, {item_key(30), build_short_string("C_30")}}); + auto item027 = build_named_object( + dict, + {{item_key(28), build_short_string("C_28")}, {item_key(29), build_short_string("C_29")}}); + item016_fields.emplace_back( + item_key(n), + build_named_object( + dict, + {{item_key(85), build_array({std::move(element)})}, {item_key(27), std::move(item027)}})); + } + for (int n : {50, 60}) { + auto element = build_named_object( + dict, + {{item_key(51), build_short_string("C_51")}, {item_key(52), build_short_string("C_52")}}); + item016_fields.emplace_back( + item_key(n), build_named_object(dict, {{item_key(85), build_array({std::move(element)})}})); + } + // `item056` is an object where one path expects an array, so that path misses. + item016_fields.emplace_back( + item_key(56), + build_named_object(dict, + {{item_key(27), + build_named_object(dict, + {{item_key(28), build_short_string("C_28")}, + {item_key(29), build_short_string("C_29")}})}})); + + auto item009 = build_named_object( + dict, + {{item_key(10), + build_named_object(dict, + {{item_key(84), + build_array({build_named_object( + dict, {{item_key(11), build_short_string("C_011")}})})}})}}); + + return build_named_object( + dict, + {{item_key(6), build_short_string("C_006")}, + {item_key(7), build_named_object(dict, {{item_key(8), build_short_string("C_008")}})}, + {item_key(9), std::move(item009)}, + {item_key(12), + build_named_object(dict, + {{item_key(13), build_short_string("C_013")}, + {item_key(14), build_short_string("C_014")}, + {item_key(15), build_short_string("C_015")}})}, + {item_key(16), build_named_object(dict, std::move(item016_fields))}, + {item_key(63), build_short_string("C_063")}}); +} + +// The paths of the multi-path workload: a few shallow ones plus a wide fan-out that all shares the +// `$.item016` prefix. +std::vector workload_paths() +{ + std::vector paths{ + "$." + item_key(6), + "$." + item_key(7) + "." + item_key(8), + "$." + item_key(9) + "." + item_key(10) + "." + item_key(84) + "[0]." + item_key(11), + "$." + item_key(12) + "." + item_key(13), + "$." + item_key(12) + "." + item_key(14), + "$." + item_key(12) + "." + item_key(15), + "$." + item_key(63)}; + + auto const under_016 = [](int n) { return "$." + item_key(16) + "." + item_key(n); }; + for (int n : {19, 20, 21, 22, 23, 24, 25, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, + 41, 42, 43, 44, 45, 46, 47, 48, 49, 53, 54, 57, 58, 59, 61, 62}) { + paths.push_back(under_016(n) + "." + item_key(85) + "[0]." + item_key(18)); + } + for (int n : {26, 55}) { + paths.push_back(under_016(n) + "." + item_key(85) + "[0]." + item_key(18)); + paths.push_back(under_016(n) + "." + item_key(85) + "[0]." + item_key(30)); + paths.push_back(under_016(n) + "." + item_key(27) + "." + item_key(28)); + paths.push_back(under_016(n) + "." + item_key(27) + "." + item_key(29)); + } + for (int n : {50, 60}) { + paths.push_back(under_016(n) + "." + item_key(85) + "[0]." + item_key(51)); + paths.push_back(under_016(n) + "." + item_key(85) + "[0]." + item_key(52)); + } + paths.push_back(under_016(56) + "." + item_key(27) + "." + item_key(28)); + paths.push_back(under_016(56) + "." + item_key(27) + "." + item_key(29)); + paths.push_back(under_016(56) + "[0]." + item_key(18)); + return paths; +} + // Build the JSONPath-like extraction path. // For nesting=2, type=array: "a.b[1]" // For nesting=3, type=string: "a.b.c" @@ -624,3 +786,61 @@ NVBENCH_BENCH(bench_variant_extract_multi_field) .add_string_axis("prefix", {"shared", "disjoint"}) .add_string_axis("api", {"batched", "looped"}) .add_int64_axis("hit_rate", {80}); + +// Compares batched against looped extraction on a workload shaped like a real one: an 85-key +// dictionary, a root object that nests most of its data under `item016`, and 50 paths that fan out +// below that shared prefix at a depth of four to five steps. +static void bench_variant_extract_workload(nvbench::state& state) +{ + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + + auto const num_rows = static_cast(state.get_int64("num_rows")); + bool const batched = state.get_string("api") == "batched"; + + std::vector dict; + dict.reserve(85); + for (int n = 1; n <= 85; ++n) { + dict.push_back(item_key(n)); + } + + auto const meta_blob = build_metadata(dict); + auto const val_blob = build_workload_value(dict); + + std::vector> meta_spans(num_rows, std::span{meta_blob}); + std::vector> val_spans(num_rows, std::span{val_blob}); + auto col = build_variant_column(meta_spans, val_spans, stream, mr); + CUDF_CUDA_TRY(cudaStreamSynchronize(stream.value())); + + auto const path_strings = workload_paths(); + std::vector const paths(path_strings.begin(), path_strings.end()); + auto const target_type = cudf::data_type{cudf::type_id::STRING}; + std::vector const target_types(paths.size(), target_type); + + auto const data_size = static_cast(num_rows) * (meta_blob.size() + val_blob.size()); + + auto mem_stats_logger = cudf::memory_stats_logger(); + mr = cudf::get_current_device_resource_ref(); + state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value())); + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) { + if (batched) { + std::ignore = cudf::io::parquet::experimental::extract_variant_fields( + col->view(), paths, target_types, stream, mr); + } else { + for (auto const& path : path_strings) { + std::ignore = cudf::io::parquet::experimental::extract_variant_field( + col->view(), path, target_type, stream, mr); + } + } + }); + + auto const time = state.get_summary("nv/cold/time/gpu/mean").get_float64("value"); + state.add_element_count(static_cast(data_size) / time, "bytes_per_second"); + state.add_buffer_size( + mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); +} + +NVBENCH_BENCH(bench_variant_extract_workload) + .set_name("bench_variant_extract_workload") + .add_int64_axis("num_rows", {262144, 1048576}) + .add_string_axis("api", {"batched", "looped"});