diff --git a/cpp/include/cudf/io/experimental/variant.hpp b/cpp/include/cudf/io/experimental/variant.hpp index 62763c8588d4..159fbaa4498a 100644 --- a/cpp/include/cudf/io/experimental/variant.hpp +++ b/cpp/include/cudf/io/experimental/variant.hpp @@ -26,34 +26,38 @@ namespace io::parquet::experimental { */ /** - * @brief Extract the raw VARIANT-encoded bytes of a nested object field by JSONPath-like path. + * @brief Extract the raw VARIANT-encoded bytes of a nested field by JSONPath-like path. * - * Walks `path` step by step, descending into object values (`basic_type == 2`) at each name step. - * Returns a `list` column containing the raw encoded bytes of the value at the end of - * the path for each row. + * Walks `path` step by step, descending into object values (`basic_type == 2`) at each name step + * and into array values (`basic_type == 3`) at each `[N]` index step. Returns a `list` + * column containing the raw encoded bytes of the value at the end of the path for each row. * * Null is produced when the struct row is null, a name step's key is absent from the dictionary, - * or the current value is not an object (`basic_type != 2`). + * an index step is out of bounds, or the current value's basic type does not match the step kind. * * Path grammar: - * path := "$"? first_step ("." name)* - * first := name | "." name - * name := [^.\[]+ // any byte except '.' (step separator) and '[' (reserved) + * path := "$"? first_step step* + * first := name | "." name | "[" index "]" + * step := "." name | "[" index "]" + * name := any sequence of bytes other than '.' or '[' + * index := non-negative base-10 integer * * Examples: - * "x" -> top-level field "x" (leading $ optional) - * "$.foo" -> top-level field "foo" - * "$.foo.bar" -> object descent foo -> bar + * "x" -> top-level field "x" (leading $ optional) + * "$.foo" -> top-level field "foo" + * "$.foo.bar" -> object descent foo -> bar + * "$[0]" -> first element of a top-level array + * "$.a[0].b" -> object key "a" -> first array element -> object key "b" * * @param variant_column Struct column (VARIANT materialization) with `list` children * (`metadata`, `value`), plus optional shredded siblings - * @param path JSONPath-like path string identifying the target object field + * @param path JSONPath-like path string identifying the target field * @param stream CUDA stream * @param mr Device memory resource - * @return `list` column with the extracted field's encoded bytes + * @return `list` column with the extracted value's encoded bytes * - * @throws std::invalid_argument on empty path or malformed syntax (including bracket steps, - * which require array-indexing support that is not yet implemented) + * @throws std::invalid_argument on empty path or malformed syntax (`[*]` wildcards, negative + * indices, out-of-range indices, and quoted names inside `[...]` are not supported) */ [[nodiscard]] std::unique_ptr get_variant_field( column_view const& variant_column, diff --git a/cpp/src/io/parquet/experimental/variant_extract.cu b/cpp/src/io/parquet/experimental/variant_extract.cu index 79a3361bbf83..87f2a8a69686 100644 --- a/cpp/src/io/parquet/experimental/variant_extract.cu +++ b/cpp/src/io/parquet/experimental/variant_extract.cu @@ -369,6 +369,57 @@ __device__ device_span locate_object_field(device_span locate_array_element(device_span val, + size_type index) +{ + auto const val_len = static_cast(val.size()); + if (val_len < 1) { return {}; } + uint8_t const value_metadata = val[0]; + if (variant_basic_type(value_metadata) != basic_type::array) { return {}; } + + int const value_header = variant_value_header(value_metadata); + [[maybe_unused]] auto const [offset_size, id_size, num_elements_size] = + decode_object_array_header(value_header, false); + + size_type pos = 1; + auto const num_elements = narrow_cast(read_uint64(val, pos, num_elements_size)); + if (!num_elements.has_value()) { return {}; } + auto const n = num_elements.value(); + pos += num_elements_size; + + if (index < 0 || index >= n) { return {}; } + + size_type const offsets_start = pos; + auto const offsets_bytes = (static_cast(n) + 1) * offset_size; + if (offsets_bytes > static_cast(val_len - offsets_start)) { return {}; } + size_type const values_base = offsets_start + static_cast(offsets_bytes); + auto const values_extent = val_len - values_base; + + auto const o0 = read_uint64(val, offsets_start + index * offset_size, offset_size); + auto const o1 = read_uint64(val, offsets_start + (index + 1) * offset_size, offset_size); + if (!o0 || !o1) { return {}; } + auto const start_u = *o0; + auto const end_u = *o1; + if (end_u < start_u || end_u > static_cast(values_extent)) { return {}; } + return val.subspan(values_base + start_u, end_u - start_u); +} + // The fixed-width signed integers a VARIANT value can be cast to: INT{8,16,32,64}. Matches the // exact width types (not e.g. __int128) since those are the only variant primitive int headers. template @@ -401,18 +452,50 @@ __device__ inline cuda::std::optional decode_int(device_span e return cudf::io::unaligned_load(enc.data() + 1); } +// 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. +// +// 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_span sub_val = val; for (size_type i = 0; i < path.size(); ++i) { - auto const name = path.element(i); - - auto const field_id = find_key_in_metadata(meta, name); - if (!field_id.has_value()) { return {}; } - - sub_val = locate_object_field(sub_val, field_id.value()); + auto const step = path.element(i); + auto const slen = step.size_bytes(); + auto const* sd = step.data(); + + if (slen >= 2 && sd[0] == '[') { + // Accumulate the index in an unsigned 64-bit value so a long digit run cannot overflow the + // signed size_type accumulator (which would be UB). An index that does not fit in size_type + // is out of range for any array and resolves to a missing element (empty span). + uint64_t index = 0; + bool ok = (sd[slen - 1] == ']'); + size_type const lo = 1; + size_type const hi = slen - 1; + for (size_type k = lo; ok && k < hi; ++k) { + char const c = sd[k]; + if (c < '0' || c > '9') { + ok = false; + break; + } + index = index * 10 + static_cast(c - '0'); + if (index > static_cast(cuda::std::numeric_limits::max())) { + ok = false; + break; + } + } + if (!ok || lo == hi) { return {}; } + sub_val = locate_array_element(sub_val, static_cast(index)); + } 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()); + } if (sub_val.empty()) { return {}; } } return sub_val; diff --git a/cpp/src/io/parquet/experimental/variant_path.cpp b/cpp/src/io/parquet/experimental/variant_path.cpp index 3d846c75bf3b..faf3791c93d3 100644 --- a/cpp/src/io/parquet/experimental/variant_path.cpp +++ b/cpp/src/io/parquet/experimental/variant_path.cpp @@ -5,13 +5,16 @@ #include "variant_path.hpp" +#include #include +#include #include #include #include #include #include +#include #include namespace cudf::io::parquet::experimental::detail { @@ -56,13 +59,53 @@ std::vector parse_variant_path(std::string_view path) if (pos >= len || !is_name_char(path[pos])) { throw_parse_error(path, pos - 1, "trailing '.' with no field name"); } - } else if (!(first && is_name_char(c))) { - // Not a '.' step or a bare leading name (e.g. "x", "$foo") - invalid + steps.emplace_back(read_unquoted_name(path.substr(pos))); + pos += steps.back().size(); + } else if (c == '[') { + // Array-index step: "[]". + // The literal token (e.g. "[42]") is preserved so the GPU-side path walker can + // distinguish array-index steps from object-key steps by inspecting the first byte. + auto const tok_start = pos; + ++pos; // consume '[' + if (pos >= len) { throw_parse_error(path, tok_start, "unterminated '[' in variant path"); } + char const nc = path[pos]; + if (nc == '*') { + throw_parse_error(path, pos, "variant path wildcard '[*]' is not supported"); + } + if (nc == '-') { + throw_parse_error(path, pos, "negative variant path index is not supported"); + } + if (nc == '\'' || nc == '"') { + throw_parse_error(path, pos, "quoted names in '[...]' are not supported"); + } + if (!(nc >= '0' && nc <= '9')) { + throw_parse_error(path, pos, "expected non-negative integer after '['"); + } + auto const digits_start = pos; + while (pos < len && path[pos] >= '0' && path[pos] <= '9') { + ++pos; + } + // Reject indices that cannot be a valid array position (don't fit in cudf::size_type), so the + // GPU-side path walker never has to parse an out-of-range value. + cudf::size_type index_value = 0; + [[maybe_unused]] auto const [ptr, ec] = + std::from_chars(path.data() + digits_start, path.data() + pos, index_value); + if (ec != std::errc{}) { + throw_parse_error(path, digits_start, "variant path index is out of range"); + } + if (pos >= len || path[pos] != ']') { + throw_parse_error(path, pos, "expected ']' after index"); + } + ++pos; // consume ']' + steps.emplace_back(path.data() + tok_start, pos - tok_start); + } else if (first && is_name_char(c)) { + // Allow a bare leading name (e.g. "x" or "foo" with no leading '$'). + steps.emplace_back(read_unquoted_name(path.substr(pos))); + pos += steps.back().size(); + } else { throw_parse_error(path, pos, "unexpected character in variant path"); } - steps.emplace_back(read_unquoted_name(path.substr(pos))); - pos += steps.back().size(); first = false; } diff --git a/cpp/tests/io/experimental/variant_extract_test.cpp b/cpp/tests/io/experimental/variant_extract_test.cpp index 051c3d94e240..a84fd5ce3013 100644 --- a/cpp/tests/io/experimental/variant_extract_test.cpp +++ b/cpp/tests/io/experimental/variant_extract_test.cpp @@ -490,9 +490,17 @@ TEST_F(ExtractVariantFieldTest, SyntaxErrors) { auto col = wrap_single_variant(build_metadata({}), enc_int32(1)); auto stream = cudf::test::get_default_stream(); - // Only object-key descent is supported — array indexing, bracket steps, and quoted keys should - // throw, alongside malformed paths. - for (auto const* bad : {"$..a", "$.a[0]", "$.a[", "$.a[]", "$.", "$['x']", "$.a[*]"}) { + // Object-key descent and array-index steps are supported; wildcards, quoted keys, negative + // indices, out-of-range indices, and other malformed bracket forms must throw. + for (auto const* bad : {"$..a", + "$.a[", + "$.a[]", + "$.", + "$['x']", + "$.a[*]", + "$.a[-1]", + "$.a[1", + "$.a[99999999999999999999]"}) { EXPECT_THROW( static_cast(cudf::io::parquet::experimental::get_variant_field(col, bad, stream)), std::invalid_argument) @@ -500,6 +508,48 @@ TEST_F(ExtractVariantFieldTest, SyntaxErrors) } } +TEST_F(ExtractVariantFieldTest, ApacheArrayPrimitiveIndexing) +{ + // array_primitive encodes the int8 array [2, 1, 5, 9]; index into it via "[N]" steps. + auto col = make_apache_variant(avf::array_primitive); + auto stream = cudf::test::get_default_stream(); + auto const i8 = cudf::data_type{cudf::type_id::INT8}; + auto const get = [&](char const* path) { + return cudf::io::parquet::experimental::extract_variant_field(col, path, i8, stream); + }; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[0]"), + cudf::test::fixed_width_column_wrapper{int8_t{2}}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[2]"), + cudf::test::fixed_width_column_wrapper{int8_t{5}}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[3]"), + cudf::test::fixed_width_column_wrapper{int8_t{9}}); + + // Out-of-bounds index resolves to null. + cudf::test::fixed_width_column_wrapper const null_expected({0}, {false}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[4]"), null_expected); +} + +TEST_F(ExtractVariantFieldTest, ArrayIndexingTypeMismatchAndBounds) +{ + // array_primitive is the int8 array [2, 1, 5, 9]. An object-key step against an array, an + // out-of-bounds index, and an index step against a non-array element all resolve to null. + auto col = make_apache_variant(avf::array_primitive); + auto stream = cudf::test::get_default_stream(); + auto const i8 = cudf::data_type{cudf::type_id::INT8}; + cudf::test::fixed_width_column_wrapper const null_expected({0}, {false}); + + // Object-key descent into an array value: no such key -> null. + auto key_on_array = + cudf::io::parquet::experimental::extract_variant_field(col, "$.foo", i8, stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*key_on_array, null_expected); + + // Index step against a primitive element (after first descending into it): non-array -> null. + auto index_on_primitive = + cudf::io::parquet::experimental::extract_variant_field(col, "$[0][0]", i8, stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*index_on_primitive, null_expected); +} + TEST_F(ExtractVariantFieldTest, LargeDictionaryAndObjectScan) { auto const keys = make_numeric_keys(50);