[FEA] Support multi-field extraction from variant columns
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
Part of https://github.com/rapidsai/cudf/issues/22312
Add a batched entry point that extracts **many** fields from a Parquet variant column in one call, returning a `cudf::table` (one column per path) instead of calling `get_variant_field` / `extract_variant_field` once per field. When paths share prefixes (e.g. `$.user.id`, `$.user.name`, `$.user.addr.zip`), the shared prefix is resolved once per row instead of once per field.
## Motivation
`get_variant_field` walks a path step by step, and each name step re-parses the per-row metadata dictionary (`find_key_in_metadata`) and re-scans the object header (`locate_object_field`). So per-row latency grows linearly with path depth, and calling it `M` times redoes every shared prefix `M` times — plus `M` kernel launches and `M` copy passes. Pulling 20 fields out of one variant blob re-descends the same top-level objects 20× per row.
## Proposed API
```cpp
// One list column per path.
[[nodiscard]] std::unique_ptr get_variant_fields(
column_view const& variant_column,
host_span paths,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
// One decoded column per path (desired_types parallels paths).
[[nodiscard]] std::unique_ptr extract_variant_fields(
column_view const& variant_column,
host_span paths,
host_span desired_types,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
```
Output columns follow input path order; duplicate/overlapping paths are allowed.
Supported target types unchanged: `STRING` / `INT{8,16,32,64}`.
## Implementation overview
Couple of key points to ensure high peformance:
1. **Host trie.** Parse paths with the existing `parse_variant_path`, insert into a trie so each distinct prefix is one node, and flatten in topological order: `node_parent[k]`, `node_step[k]`, and `leaf_of_path[p]` mapping each input path to its leaf node.
2. **Trie-walk sizing kernel.** Per row, walk nodes parent-before-child, applying each step on the parent's span via the existing `locate_object_field` / `locate_array_element`. Only persist a node result
`{src_offset, size, valid}` at branch points (>1 child) and leaves - single-child chains are collapsed (path-compressed) and carried in a register, so scratch `K` is bounded by branch points + leaves, not total
nodes. Shared prefixes resolve once; an invalid parent short-circuits its whole subtree. When `K` is small (tens) — thread-local array, with a global scratch fallback when `K` is large (e.g. >32).
3. **Assembly.** Reuse `make_offsets_child_column` per output column and a single `batched_memcpy_async` over all `(leaf, row)` pairs. For the typed API, run each `list` column through the existing `cast_variant`.
Contributor guide
Assessment
This issue has not been assessed yet.