Support multi-field extraction from variant columns - #23737
Draft
vuule wants to merge 2 commits into
Draft
Conversation
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds
get_variant_fieldsandextract_variant_fields, which take a set of paths and return one output column per path. Extracting many fields previously meant callingget_variant_fieldonce per path, which re-resolved every shared prefix and re-launched the locate and copy passes per field.The paths are merged host-side into a prefix trie, flattened so that the device walk resolves each shared prefix once per row, locates every path in a single kernel, and copies every (path, row) pair in one batched memcpy. Single-child chains are collapsed, and slots are emitted in depth-first pre-order so the walk only keeps one located value per trie level rather than one per path.
On an A100 at 2M rows and an 80% hit rate, the batched API is 1.2-1.4x faster than the equivalent loop over
extract_variant_fieldfor 4 to 64 fields, with or without shared prefixes.Closes #22897
Checklist