Add array indexing support to Parquet variant field extraction

cpp/include/cudf/io/experimental/variant.hpp

@@ -29,34 +29,33 @@ namespace io::parquet::experimental {
  */
 
 /**
- * @brief Extract the raw VARIANT-encoded bytes of a nested object field by JSONPath-like path.
- *
- * Walks `path` step by step, descending into object values (`basic_type == 2`) at each name step.
- * Returns a `list<uint8>` 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`).
+ * @brief Extract the raw VARIANT-encoded bytes of a nested field by JSONPath-like path.
  *
  * 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 (leading zeros are allowed, e.g. "[01]" == "[1]")
  *
  * 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<uint8>` 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<uint8>` column with the extracted field's encoded bytes
+ * @return `list<uint8>` column with the extracted value's encoded bytes. A row is null when the
+ *         input row is null, a name is absent, an index is out of bounds, or a step does not match
+ *         the current value.
  *
- * @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<column> get_variant_field(
   column_view const& variant_column,

cpp/src/io/parquet/experimental/variant_extract.cu

@@ -266,7 +266,7 @@ __device__ cuda::std::optional<size_type> find_key_in_metadata(device_span<uint8
 
   auto const offsets_start = pos;
   auto const offsets_bytes = (static_cast<uint64_t>(num_entries.value()) + 1) * offset_size;
-  if (offsets_bytes > static_cast<uint64_t>(meta_len - offsets_start)) {
+  if (cuda::std::cmp_greater(offsets_bytes, meta_len - offsets_start)) {
     return cuda::std::nullopt;
   }
 
@@ -369,6 +369,62 @@ __device__ device_span<uint8_t const> locate_object_field(device_span<uint8_t co
   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.
+//
+// Array layout per the Variant spec:
+//   byte 0: header (basic_type=array in low 2 bits; value_header in high 6 bits)
+//     value_header bits: (offset_size - 1) in bits 0-1, is_large in bit 2, bits 3-5 unused
+//   num_elements: 1 byte if !is_large else 4 bytes (little-endian)
+//   offsets:      (num_elements + 1) entries, each `offset_size` bytes, relative to the end of
+//                 offsets
+//   values:       concatenated element blobs
+//
+// Array element offsets are monotonically increasing, so the element length is taken directly from
+// the offset delta (o1 - o0) rather than from the element's own header.
+__device__ device_span<uint8_t const> locate_array_element(device_span<uint8_t const> value,
+                                                           size_type index)
+{
+  if (index < 0) { return {}; }
+
+  auto const value_size = static_cast<size_type>(value.size());
+  if (value_size < 1) { return {}; }
+  uint8_t const value_metadata = value[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, _, num_elements_size] =
+    decode_object_array_header(value_header, false);
+
+  size_type position            = 1;
+  auto const num_elements_value = narrow_cast(read_uint64(value, position, num_elements_size));
+  if (!num_elements_value.has_value()) { return {}; }
+  auto const num_elements = num_elements_value.value();
+  if (index >= num_elements) { return {}; }
+  position += num_elements_size;
+
+  size_type const offsets_start = position;
+  auto const offsets_bytes      = (static_cast<uint64_t>(num_elements) + 1) * offset_size;
+  if (cuda::std::cmp_greater(offsets_bytes, value_size - offsets_start)) { return {}; }
+  size_type const values_base = offsets_start + static_cast<size_type>(offsets_bytes);
+  auto const values_extent    = value_size - values_base;
+
+  auto const start_offset_pos = offsets_start + static_cast<uint64_t>(index) * offset_size;
+  auto const end_offset_pos   = offsets_start + (static_cast<uint64_t>(index) + 1) * offset_size;
+  if (cuda::std::cmp_greater(end_offset_pos + offset_size, value_size)) { return {}; }
+
+  auto const start_offset = read_uint64(value, start_offset_pos, offset_size);
+  auto const end_offset   = read_uint64(value, end_offset_pos, offset_size);
+  if (!start_offset.has_value() || !end_offset.has_value()) { return {}; }
+  auto const element_start = *start_offset;
+  auto const element_end   = *end_offset;
+  if (element_end < element_start || cuda::std::cmp_greater(element_end, values_extent)) {
+    return {};
+  }
+  return value.subspan(values_base + element_start, element_end - element_start);
+}
+
 // 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 <typename T>
@@ -401,18 +457,55 @@ __device__ inline cuda::std::optional<T> decode_int(device_span<uint8_t const> e
   return cudf::io::unaligned_load<T>(enc.data() + 1);
 }
 
+// Parse an array-index step token of the form "[<N>]" into its zero-based index. Returns nullopt
+// for any malformed token or an index that does not fit in `size_type` (such an index is out of
+// range for any array, so the caller treats it as a missing element).
+__device__ cuda::std::optional<size_type> parse_index_step(cudf::string_view step)
+{
+  auto const step_size  = step.size_bytes();
+  auto const* step_data = step.data();
+  if (step_size < 3 || step_data[0] != '[' || step_data[step_size - 1] != ']') {
+    return cuda::std::nullopt;
+  }
+
+  // The index is accumulated in an unsigned 64-bit value so a long digit run cannot overflow the
+  // signed `size_type` accumulator (which would be UB) before the range check rejects it.
+  uint64_t index = 0;
+  for (size_type k = 1; k < step_size - 1; ++k) {
+    char const c = step_data[k];
+    if (c < '0' || c > '9') { return cuda::std::nullopt; }
+    index = index * 10 + static_cast<uint64_t>(c - '0');
+    if (cuda::std::cmp_greater(index, cuda::std::numeric_limits<size_type>::max())) {
+      return cuda::std::nullopt;
+    }
+  }
+  return static_cast<size_type>(index);
+}
+
+// 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:
+//   - "<name>"  -> descend into an object by dictionary key, or
+//   - "[<N>]"   -> descend into an array by zero-based integer index.
+// The step kind is inferred from the first byte (`'['` means index).
 __device__ device_span<uint8_t const> resolve_path(device_span<uint8_t const> meta,
                                                    device_span<uint8_t const> val,
                                                    column_device_view path)
 {
   device_span<uint8_t const> sub_val = val;
   for (size_type i = 0; i < path.size(); ++i) {
-    auto const name = path.element<cudf::string_view>(i);
-
-    auto const field_id = find_key_in_metadata(meta, name);
-    if (!field_id.has_value()) { return {}; }
+    auto const step = path.element<cudf::string_view>(i);
 
-    sub_val = locate_object_field(sub_val, field_id.value());
+    if (step.size_bytes() >= 1 && step.data()[0] == '[') {
+      auto const index = parse_index_step(step);
+      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());
+    }
     if (sub_val.empty()) { return {}; }
   }
   return sub_val;

cpp/src/io/parquet/experimental/variant_path.cpp

@@ -1,17 +1,19 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include "variant_path.hpp"
 
+#include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
 
+#include <charconv>
 #include <cstddef>
-#include <format>
 #include <stdexcept>
 #include <string>
 #include <string_view>
+#include <system_error>
 #include <vector>
 
 namespace cudf::io::parquet::experimental::detail {
@@ -21,12 +23,6 @@ namespace {
 // Dot-notation field names accept any byte except the structural characters '.' and '['.
 [[nodiscard]] constexpr bool is_name_char(char c) { return c != '.' && c != '['; }
 
-[[noreturn]] void throw_parse_error(std::string_view path, std::size_t pos, std::string_view msg)
-{
-  CUDF_FAIL(std::format("invalid variant path \"{}\" at position {}: {}", path, pos, msg),
-            std::invalid_argument);
-}
-
 // Reads a maximal run of name characters from the front of `tail`.
 [[nodiscard]] std::string read_unquoted_name(std::string_view tail)
 {
@@ -37,6 +33,34 @@ namespace {
   return std::string{tail.substr(0, n)};
 }
 
+// Reads a bracket step "[<non-negative integer>]" from the front of `tail`.
+// The returned token keeps its brackets (e.g. "[42]").
+[[nodiscard]] std::string read_bracket_step(std::string_view tail)
+{
+  CUDF_EXPECTS(!tail.empty() && tail.front() == '[',
+               "expected '[' to open variant path index",
+               std::invalid_argument);
+
+  // Consume the maximal run of decimal digits.
+  std::size_t n = 1;
+  while (n < tail.size() && tail[n] >= '0' && tail[n] <= '9') {
+    ++n;
+  }
+  CUDF_EXPECTS(
+    n != 1, "expected non-negative integer after '[' in variant path", std::invalid_argument);
+
+  // Reject indices that cannot be a valid array position (don't fit in cudf::size_type)
+  cudf::size_type index = 0;
+  auto const result     = std::from_chars(tail.data() + 1, tail.data() + n, index);
+  CUDF_EXPECTS(
+    result.ec == std::errc{}, "variant path index is out of range", std::invalid_argument);
+
+  CUDF_EXPECTS(n < tail.size() && tail[n] == ']',
+               "expected ']' to close variant path index",
+               std::invalid_argument);
+  return std::string{tail.substr(0, n + 1)};  // include the closing ']'
+}
+
 }  // namespace
 
 std::vector<std::string> parse_variant_path(std::string_view path)
@@ -51,17 +75,21 @@ std::vector<std::string> parse_variant_path(std::string_view path)
   bool first = true;
   while (pos < len) {
     char const c = path[pos];
-    if (c == '.') {
-      ++pos;
-      if (pos >= len || !is_name_char(path[pos])) {
-        throw_parse_error(path, pos - 1, "trailing '.' with no field name");
+    if (c == '[') {
+      steps.emplace_back(read_bracket_step(path.substr(pos)));
+    } else {
+      if (c == '.') {
+        ++pos;
+        CUDF_EXPECTS(pos < len && is_name_char(path[pos]),
+                     "trailing '.' with no field name",
+                     std::invalid_argument);
+      } else {
+        // Neither a '.'/'[' step nor a valid leading name (e.g. a stray ']' or a name after a step)
+        CUDF_EXPECTS(
+          first && is_name_char(c), "unexpected character in variant path", std::invalid_argument);
       }
-    } else if (!(first && is_name_char(c))) {
-      // Neither a '.' step nor a valid leading name (e.g. a bracket step like "[0]" or "foo[1]")
-      throw_parse_error(path, pos, "unexpected character in variant path");
+      steps.emplace_back(read_unquoted_name(path.substr(pos)));
     }
-
-    steps.emplace_back(read_unquoted_name(path.substr(pos)));
     pos += steps.back().size();
     first = false;
   }

cpp/src/io/parquet/experimental/variant_path.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -12,18 +12,21 @@
 namespace cudf::io::parquet::experimental::detail {
 
 /**
- * @brief Parse a JSONPath-like VARIANT path string into an ordered sequence of object-key steps.
+ * @brief Parse a JSONPath-like VARIANT path string into an ordered sequence of steps.
  *
- * Grammar — object descent only:
- *   path  := "$"? first_step ("." name)*
- *   first := name | "." name
+ * Grammar — object descent and array indexing:
+ *   path  := "$"? first_step (("." name) | index)*
+ *   first := name | "." name | index
  *   name  := [^.\[]+
+ *   index := "[" [0-9]+ "]"
  *
- * Names accept any byte except '.' (step separator) and '[' (start of a bracket step,
- * reserved for future array indexing and quoted-name syntax).
+ * A step is either an object-key name or an array index. Names accept any byte except '.' (step
+ * separator) and '[' (start of an index step). Index steps hold a non-negative integer and are
+ * returned with their brackets kept (e.g. "[42]"), which is how downstream consumers tell an index
+ * step apart from an object key.
  *
- * @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 an empty path or malformed syntax (e.g. a non-integer, negative,
+ *         or out-of-range array index, an unterminated '[', or a trailing '.')
  */
 [[nodiscard]] std::vector<std::string> parse_variant_path(std::string_view path);