From b74b54f22545aff74c41d84efbe19992d39faa8f Mon Sep 17 00:00:00 2001 From: Abigale Kim Date: Tue, 11 Aug 2026 23:43:04 +0000 Subject: [PATCH 1/2] binsearch changes --- .../parquet/experimental/variant_extract.cu | 263 ++++++++---------- 1 file changed, 119 insertions(+), 144 deletions(-) diff --git a/cpp/src/io/parquet/experimental/variant_extract.cu b/cpp/src/io/parquet/experimental/variant_extract.cu index 8327e6c5cde..bf697de71ac 100644 --- a/cpp/src/io/parquet/experimental/variant_extract.cu +++ b/cpp/src/io/parquet/experimental/variant_extract.cu @@ -136,7 +136,8 @@ __device__ object_array_header decode_object_array_header(uint8_t value_header, * - short_string (1): header + `value_header` payload bytes (the header is the string length). * - object/array (2/3): header + num_elements + field-id list + field-offset list + values; the * total values-region size is read from the trailing field_offset (the "sentinel" at index - * num_elements). See decode_object_array_header / locate_object_field for the sub-layout. + * num_elements). See decode_object_array_header / locate_object_field_by_name for the + * sub-layout. * * @param enc The encoded value bytes (must begin at the value metadata byte) * @return The total length in bytes of the value, or nullopt if `enc` is empty/malformed or the @@ -210,146 +211,6 @@ __device__ cuda::std::optional variant_value_length(device_span find_key_in_metadata(device_span meta, - cudf::string_view key) -{ - auto const meta_len = static_cast(meta.size()); - if (meta_len < 1) { return cuda::std::nullopt; } - - auto const header = meta[0]; - int const version = header & 0x0F; - if (version != variant_version_v1) { return cuda::std::nullopt; } - int const offset_size = ((header >> 6) & 0x03) + 1; - - size_type pos = 1; - auto const num_entries = narrow_cast(read_uint64(meta, pos, offset_size)); - if (!num_entries.has_value()) { return cuda::std::nullopt; } - pos += offset_size; - - auto const offsets_start = pos; - auto const offsets_bytes = (static_cast(num_entries.value()) + 1) * offset_size; - if (cuda::std::cmp_greater(offsets_bytes, meta_len - offsets_start)) { - return cuda::std::nullopt; - } - - auto start_off = read_uint64(meta, offsets_start, offset_size); - if (!start_off.has_value()) { return cuda::std::nullopt; } - auto const strings_base = offsets_start + static_cast(offsets_bytes); - // Bytes available for dictionary string payloads - auto const strings_extent = meta_len - strings_base; - for (size_type i = 0; i < num_entries.value(); ++i) { - auto const end_off = read_uint64(meta, offsets_start + (i + 1) * offset_size, offset_size); - if (!end_off.has_value()) { return cuda::std::nullopt; } - if (end_off.value() < start_off.value() || end_off.value() > strings_extent) { - return cuda::std::nullopt; - } - cudf::string_view const entry{ - reinterpret_cast(meta.data() + strings_base + start_off.value()), - static_cast(end_off.value() - start_off.value())}; - if (entry == key) { return i; } - start_off = end_off; - } - return cuda::std::nullopt; -} - -/** - * @brief Locate the encoded bytes of a single field within an object value by field id. - * - * An object value is laid out as: - * - * byte 0: value metadata | basic_type=object (2) | value_header (6) | - * bytes 1..: num_elements (num_elements_size bytes) = number of fields N - * next N*field_id_size bytes: field_ids[0..N-1] (sorted by field name) - * next (N+1)*field_offset_size bytes: field_offsets[0..N] (relative to values_base) - * remaining bytes (values_base..): the concatenated field values - * - * `num_elements_size`, `field_id_size`, and `field_offset_size` come from the value header (see - * decode_object_array_header). The trailing offset `field_offsets[N]` is the total size of the - * values region. This scans `field_ids` for `id`, then uses the matching `field_offsets[i]` to - * slice out the field's value, whose length is derived from its own header via - * variant_value_length. - * - * Per the spec, field ids/offsets are ordered by the corresponding field names (lexicographically), - * but the values themselves may be in any order, so `field_offsets` are not necessarily monotonic — - * hence the value length is taken from each field's own header rather than from offset deltas. - * - * @param val The object value bytes - * @param id The dictionary index of the field to locate - * @return The encoded bytes of the field value, or an empty span if `val` is not an object, the - * field is absent, or the blob is malformed - */ -__device__ device_span locate_object_field(device_span val, int id) -{ - auto const val_len = static_cast(val.size()); - if (val_len < 1) { return {}; } - auto const value_metadata = val[0]; - if (decode_basic_type(value_metadata) != basic_type::OBJECT) { return {}; } - - auto const [offset_size, id_size, num_elements_size] = - decode_object_array_header(variant_value_header(value_metadata), true); - - size_type pos = 1; - auto const num_fields = narrow_cast(read_uint64(val, pos, num_elements_size)); - if (!num_fields.has_value()) { return {}; } - pos += num_elements_size; - - auto const ids_start = pos; - auto const ids_bytes = static_cast(num_fields.value()) * id_size; - if (ids_bytes > val_len - ids_start) { return {}; } - - auto const offsets_start = ids_start + static_cast(ids_bytes); - auto const offsets_bytes = (static_cast(num_fields.value()) + 1) * offset_size; - if (offsets_bytes > val_len - offsets_start) { return {}; } - - auto const values_base = offsets_start + static_cast(offsets_bytes); - // Maximum legitimate field-offset value: bytes available after values_base - auto const values_extent = val_len - values_base; - - // Find the matching field ID and its start offset - bool found = false; - uint64_t match_start = 0; - for (size_type i = 0; i < num_fields.value(); ++i) { - auto const current_id = read_uint64(val, ids_start + i * id_size, id_size); - if (!current_id.has_value()) { return {}; } - if (cuda::std::cmp_not_equal(current_id.value(), id)) { continue; } - - auto const match_offset = read_uint64(val, offsets_start + i * offset_size, offset_size); - if (!match_offset.has_value()) { return {}; } - if (match_offset.value() > values_extent) { return {}; } - match_start = match_offset.value(); - found = true; - break; - } - if (!found) { return {}; } - - // Derive field's value length from its header - auto const value = val.subspan(values_base + match_start); - auto const value_len = variant_value_length(value); - if (!value_len.has_value()) { return {}; } - auto const match_end = match_start + value_len.value(); - if (match_end > values_extent) { return {}; } - return val.subspan(values_base + match_start, value_len.value()); -} - // Parse an array value header and return the sub-span of the element at `index` (0-based) within // `val`. Returns an empty span if `val` is not an array (`basic_type != array`), if `index` is out // of bounds, or if the encoded data is truncated. @@ -508,6 +369,122 @@ __device__ cuda::std::optional parse_index_step(cudf::string_view ste return index; } +/** + * @brief Locate a field within an object value by key name, using the metadata blob. + * + * Combines metadata key lookup and object field traversal into a single O(log N) binary search. + * The Variant spec guarantees that `field_ids[0..N-1]` in every object value are ordered by the + * corresponding field name (lexicographically), regardless of the metadata `sorted` bit. This + * function exploits that invariant unconditionally: at each binary search probe it reads + * `field_ids[mid]` and uses the metadata offset table (random-access by id) to retrieve the + * name for that id in O(1), then compares it with `key`. + * + * This replaces the two-step chain `find_key_in_metadata` → `locate_object_field`, giving + * O(log N) even when the metadata `sorted` bit is 0 (where that chain falls back to O(N)). + * + * Not using thrust::lower_bound since it does not propagate entry read failures. + * + * @param meta The metadata blob for this row + * @param val The object value bytes + * @param key The field name to locate + * @return The encoded bytes of the matching field value, or an empty span if the field is absent + * or either blob is malformed + */ +__device__ device_span locate_object_field_by_name(device_span meta, + device_span val, + cudf::string_view key) +{ + // --- Parse metadata header --- + auto const meta_len = static_cast(meta.size()); + if (meta_len < 1) { return {}; } + auto const meta_header = meta[0]; + if ((meta_header & 0x0F) != variant_version_v1) { return {}; } + int const meta_offset_size = ((meta_header >> 6) & 0x03) + 1; + + size_type meta_pos = 1; + auto const num_meta_entries = narrow_cast(read_uint64(meta, meta_pos, meta_offset_size)); + if (!num_meta_entries.has_value()) { return {}; } + meta_pos += meta_offset_size; + + auto const meta_offsets_start = meta_pos; + auto const meta_offsets_bytes = + (static_cast(num_meta_entries.value()) + 1) * meta_offset_size; + if (cuda::std::cmp_greater(meta_offsets_bytes, meta_len - meta_offsets_start)) { return {}; } + auto const meta_strings_base = meta_offsets_start + static_cast(meta_offsets_bytes); + auto const meta_strings_extent = meta_len - meta_strings_base; + + // O(1) name lookup by id: two offset reads into the metadata table. + auto name_for_id = [&](size_type id) -> cuda::std::optional { + auto const s = read_uint64(meta, meta_offsets_start + id * meta_offset_size, meta_offset_size); + auto const e = + read_uint64(meta, meta_offsets_start + (id + 1) * meta_offset_size, meta_offset_size); + if (!s.has_value() || !e.has_value()) { return cuda::std::nullopt; } + if (e.value() < s.value() || cuda::std::cmp_greater(e.value(), meta_strings_extent)) { + return cuda::std::nullopt; + } + return cudf::string_view{ + reinterpret_cast(meta.data() + meta_strings_base + s.value()), + static_cast(e.value() - s.value())}; + }; + + // --- Parse object value header --- + auto const val_len = static_cast(val.size()); + if (val_len < 1) { return {}; } + auto const value_metadata = val[0]; + if (decode_basic_type(value_metadata) != basic_type::OBJECT) { return {}; } + + auto const [offset_size, id_size, num_elements_size] = + decode_object_array_header(variant_value_header(value_metadata), true); + + size_type pos = 1; + auto const num_fields = narrow_cast(read_uint64(val, pos, num_elements_size)); + if (!num_fields.has_value()) { return {}; } + pos += num_elements_size; + + auto const ids_start = pos; + auto const ids_bytes = static_cast(num_fields.value()) * id_size; + if (cuda::std::cmp_greater(ids_bytes, val_len - ids_start)) { return {}; } + + auto const offsets_start = ids_start + static_cast(ids_bytes); + auto const offsets_bytes = (static_cast(num_fields.value()) + 1) * offset_size; + if (cuda::std::cmp_greater(offsets_bytes, val_len - offsets_start)) { return {}; } + + auto const values_base = offsets_start + static_cast(offsets_bytes); + auto const values_extent = val_len - values_base; + + // --- Binary search field_ids by name --- + // field_ids are always name-sorted per spec; name_for_id provides O(1) lookup per probe. + size_type lo = 0; + size_type hi = num_fields.value(); + while (lo < hi) { + size_type const mid = lo + (hi - lo) / 2; + auto const probe_id = narrow_cast(read_uint64(val, ids_start + mid * id_size, id_size)); + if (!probe_id.has_value()) { return {}; } + auto const probe_name = name_for_id(probe_id.value()); + if (!probe_name.has_value()) { return {}; } + auto const cmp = probe_name.value().compare(key); + if (cmp == 0) { + auto const match_offset = read_uint64(val, offsets_start + mid * offset_size, offset_size); + if (!match_offset.has_value() || + cuda::std::cmp_greater(match_offset.value(), values_extent)) { + return {}; + } + auto const field_val = val.subspan(values_base + match_offset.value()); + auto const field_len = variant_value_length(field_val); + if (!field_len.has_value() || match_offset.value() + field_len.value() > values_extent) { + return {}; + } + return val.subspan(values_base + match_offset.value(), field_len.value()); + } + if (cmp < 0) { + lo = mid + 1; + } else { + hi = mid; + } + } + return {}; +} + // 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. // @@ -528,9 +505,7 @@ __device__ device_span resolve_path(device_span me if (!index.has_value()) { return {}; } sub_val = locate_array_element(sub_val, index.value()); } else { - auto const field_id = find_key_in_metadata(meta, step); - if (!field_id.has_value()) { return {}; } - sub_val = locate_object_field(sub_val, field_id.value()); + sub_val = locate_object_field_by_name(meta, sub_val, step); } if (sub_val.empty()) { return {}; } } From 002058460b4632c8d8aa9bf6df7562e584986830 Mon Sep 17 00:00:00 2001 From: Abigale Kim Date: Fri, 14 Aug 2026 02:27:08 +0000 Subject: [PATCH 2/2] changes --- .../parquet/experimental/variant_extract.cu | 164 +++++++++++++----- 1 file changed, 124 insertions(+), 40 deletions(-) diff --git a/cpp/src/io/parquet/experimental/variant_extract.cu b/cpp/src/io/parquet/experimental/variant_extract.cu index bf697de71ac..0edd1b0c18c 100644 --- a/cpp/src/io/parquet/experimental/variant_extract.cu +++ b/cpp/src/io/parquet/experimental/variant_extract.cu @@ -136,8 +136,7 @@ __device__ object_array_header decode_object_array_header(uint8_t value_header, * - short_string (1): header + `value_header` payload bytes (the header is the string length). * - object/array (2/3): header + num_elements + field-id list + field-offset list + values; the * total values-region size is read from the trailing field_offset (the "sentinel" at index - * num_elements). See decode_object_array_header / locate_object_field_by_name for the - * sub-layout. + * num_elements). See decode_object_array_header / locate_object_field for the sub-layout. * * @param enc The encoded value bytes (must begin at the value metadata byte) * @return The total length in bytes of the value, or nullopt if `enc` is empty/malformed or the @@ -370,31 +369,105 @@ __device__ cuda::std::optional parse_index_step(cudf::string_view ste } /** - * @brief Locate a field within an object value by key name, using the metadata blob. + * @brief Find the dictionary index of a key in a VARIANT metadata blob. * - * Combines metadata key lookup and object field traversal into a single O(log N) binary search. - * The Variant spec guarantees that `field_ids[0..N-1]` in every object value are ordered by the - * corresponding field name (lexicographically), regardless of the metadata `sorted` bit. This - * function exploits that invariant unconditionally: at each binary search probe it reads - * `field_ids[mid]` and uses the metadata offset table (random-access by id) to retrieve the - * name for that id in O(1), then compares it with `key`. + * The metadata blob is the per-row string dictionary, laid out as: * - * This replaces the two-step chain `find_key_in_metadata` → `locate_object_field`, giving - * O(log N) even when the metadata `sorted` bit is 0 (where that chain falls back to O(N)). + * byte 0: header | version (4) | sorted (1) | unused (1) | offset_size-1 (2) | + * bytes 1..: dictionary_size (offset_size bytes, little-endian) = number of keys N + * next (N+1)*offset_size bytes: offsets[0..N] (offset_size bytes each, little-endian) + * remaining bytes: string_data (concatenated UTF-8 key bytes) + * + * `offset_size` (1..4 bytes, from the header) is the width of every dictionary_size/offset entry. + * Key `i` occupies `string_data[offsets[i] : offsets[i+1]]`; the trailing offset `offsets[N]` is + * the total length of `string_data`. Offsets are relative to the start of `string_data`, i.e. to `1 + * + offset_size + (N+1)*offset_size`. + * + * @param meta The metadata blob bytes for a single row + * @param key The key to search for + * @return The dictionary index of `key`, or nullopt if absent or the blob is malformed + */ +__device__ cuda::std::optional find_key_in_metadata(device_span meta, + cudf::string_view key) +{ + auto const meta_len = static_cast(meta.size()); + if (meta_len < 1) { return cuda::std::nullopt; } + + auto const header = meta[0]; + int const version = header & 0x0F; + if (version != variant_version_v1) { return cuda::std::nullopt; } + int const offset_size = ((header >> 6) & 0x03) + 1; + + size_type pos = 1; + auto const num_entries = narrow_cast(read_uint64(meta, pos, offset_size)); + if (!num_entries.has_value()) { return cuda::std::nullopt; } + pos += offset_size; + + auto const offsets_start = pos; + auto const offsets_bytes = (static_cast(num_entries.value()) + 1) * offset_size; + if (cuda::std::cmp_greater(offsets_bytes, meta_len - offsets_start)) { + return cuda::std::nullopt; + } + + auto start_off = read_uint64(meta, offsets_start, offset_size); + if (!start_off.has_value()) { return cuda::std::nullopt; } + auto const strings_base = offsets_start + static_cast(offsets_bytes); + // Bytes available for dictionary string payloads + auto const strings_extent = meta_len - strings_base; + for (size_type i = 0; i < num_entries.value(); ++i) { + auto const end_off = read_uint64(meta, offsets_start + (i + 1) * offset_size, offset_size); + if (!end_off.has_value()) { return cuda::std::nullopt; } + if (end_off.value() < start_off.value() || end_off.value() > strings_extent) { + return cuda::std::nullopt; + } + cudf::string_view const entry{ + reinterpret_cast(meta.data() + strings_base + start_off.value()), + static_cast(end_off.value() - start_off.value())}; + if (entry == key) { return i; } + start_off = end_off; + } + return cuda::std::nullopt; +} + +/** + * @brief Locate the encoded bytes of a single field within an object value by field id. + * + * Object value layout, following the 1-byte value metadata header (basic_type=object in the low 2 + * bits; value_header in the high 6 bits, see decode_object_array_header): + * + * bytes 1..: num_elements (num_elements_size bytes) = number of fields N + * next N*field_id_size bytes: field_ids[0..N-1] (sorted by field name) + * next (N+1)*field_offset_size bytes: field_offsets[0..N] (relative to values_base) + * remaining bytes (values_base..): the concatenated field values + * + * `num_elements_size`, `field_id_size`, and `field_offset_size` come from the value header (see + * decode_object_array_header). The trailing offset `field_offsets[N]` is the total size of the + * values region. + * + * Per the spec, `field_ids[0..N-1]` are ordered by the corresponding field name + * (lexicographically), not by the numeric id value, and the values themselves may be in any order, + * so `field_offsets` are not necessarily monotonic -- hence the value length is taken from each + * field's own header rather than from offset deltas. + * + * Because field_ids are name-ordered rather than id-ordered, `id` cannot be binary searched + * directly. Instead, at each probe this turns `field_ids[mid]` into its dictionary name via an + * O(1) lookup in `meta` (the metadata offset table is indexed directly by id, independent of + * whether the dictionary strings themselves are sorted) and compares that name against the target + * field's name -- resolved the same way, once, up front -- giving O(log N) unconditionally. * * Not using thrust::lower_bound since it does not propagate entry read failures. * - * @param meta The metadata blob for this row - * @param val The object value bytes - * @param key The field name to locate - * @return The encoded bytes of the matching field value, or an empty span if the field is absent - * or either blob is malformed + * @param meta The metadata blob for this row, used to resolve field ids to names + * @param val The object value bytes + * @param id The dictionary index of the field to locate + * @return The encoded bytes of the field value, or an empty span if `val` is not an object, the + * field is absent, or either blob is malformed */ -__device__ device_span locate_object_field_by_name(device_span meta, - device_span val, - cudf::string_view key) +__device__ device_span locate_object_field(device_span meta, + device_span val, + int id) { - // --- Parse metadata header --- + // --- Parse metadata header (for O(1) id -> name lookups) --- auto const meta_len = static_cast(meta.size()); if (meta_len < 1) { return {}; } auto const meta_header = meta[0]; @@ -414,10 +487,11 @@ __device__ device_span locate_object_field_by_name(device_span cuda::std::optional { - auto const s = read_uint64(meta, meta_offsets_start + id * meta_offset_size, meta_offset_size); + auto name_for_id = [&](size_type field_id) -> cuda::std::optional { + auto const s = + read_uint64(meta, meta_offsets_start + field_id * meta_offset_size, meta_offset_size); auto const e = - read_uint64(meta, meta_offsets_start + (id + 1) * meta_offset_size, meta_offset_size); + read_uint64(meta, meta_offsets_start + (field_id + 1) * meta_offset_size, meta_offset_size); if (!s.has_value() || !e.has_value()) { return cuda::std::nullopt; } if (e.value() < s.value() || cuda::std::cmp_greater(e.value(), meta_strings_extent)) { return cuda::std::nullopt; @@ -427,6 +501,9 @@ __device__ device_span locate_object_field_by_name(device_span(e.value() - s.value())}; }; + auto const key = name_for_id(id); + if (!key.has_value()) { return {}; } + // --- Parse object value header --- auto const val_len = static_cast(val.size()); if (val_len < 1) { return {}; } @@ -449,32 +526,29 @@ __device__ device_span locate_object_field_by_name(device_span(num_fields.value()) + 1) * offset_size; if (cuda::std::cmp_greater(offsets_bytes, val_len - offsets_start)) { return {}; } - auto const values_base = offsets_start + static_cast(offsets_bytes); + auto const values_base = offsets_start + static_cast(offsets_bytes); + // Maximum legitimate field-offset value: bytes available after values_base auto const values_extent = val_len - values_base; // --- Binary search field_ids by name --- - // field_ids are always name-sorted per spec; name_for_id provides O(1) lookup per probe. - size_type lo = 0; - size_type hi = num_fields.value(); + bool found = false; + uint64_t match_start = 0; + size_type lo = 0; + size_type hi = num_fields.value(); while (lo < hi) { size_type const mid = lo + (hi - lo) / 2; auto const probe_id = narrow_cast(read_uint64(val, ids_start + mid * id_size, id_size)); if (!probe_id.has_value()) { return {}; } auto const probe_name = name_for_id(probe_id.value()); if (!probe_name.has_value()) { return {}; } - auto const cmp = probe_name.value().compare(key); + auto const cmp = probe_name.value().compare(key.value()); if (cmp == 0) { auto const match_offset = read_uint64(val, offsets_start + mid * offset_size, offset_size); - if (!match_offset.has_value() || - cuda::std::cmp_greater(match_offset.value(), values_extent)) { - return {}; - } - auto const field_val = val.subspan(values_base + match_offset.value()); - auto const field_len = variant_value_length(field_val); - if (!field_len.has_value() || match_offset.value() + field_len.value() > values_extent) { - return {}; - } - return val.subspan(values_base + match_offset.value(), field_len.value()); + if (!match_offset.has_value()) { return {}; } + if (match_offset.value() > values_extent) { return {}; } + match_start = match_offset.value(); + found = true; + break; } if (cmp < 0) { lo = mid + 1; @@ -482,7 +556,15 @@ __device__ device_span locate_object_field_by_name(device_span values_extent) { return {}; } + return val.subspan(values_base + match_start, value_len.value()); } // Walk a path of object-key or array-index steps level by level starting at `val` and return @@ -505,7 +587,9 @@ __device__ device_span resolve_path(device_span me if (!index.has_value()) { return {}; } sub_val = locate_array_element(sub_val, index.value()); } else { - sub_val = locate_object_field_by_name(meta, sub_val, step); + auto const field_id = find_key_in_metadata(meta, step); + if (!field_id.has_value()) { return {}; } + sub_val = locate_object_field(meta, sub_val, field_id.value()); } if (sub_val.empty()) { return {}; } }