NVIDIA / NVIDIA/cudf

[FEA] Support multi-field extraction from variant columns

Open
#22,897 0 comments 0 reactions 0 assignees View on GitHub
cuIO feature request
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.