Add array-index support to Parquet variant field extraction

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<uint8>` 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<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`).
+ * 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<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
  *
- * @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

@@ -369,6 +369,57 @@ __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
+//
+// Unlike object field offsets (which are ordered by field name and so are not necessarily
+// monotonic, see locate_object_field), array element offsets are monotonically increasing per the
+// Variant spec, 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> val,
+                                                           size_type index)
+{
+  auto const val_len = static_cast<size_type>(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<uint64_t>(n) + 1) * offset_size;
+  if (offsets_bytes > static_cast<uint64_t>(val_len - offsets_start)) { return {}; }
+  size_type const values_base = offsets_start + static_cast<size_type>(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<uint64_t>(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 <typename T>
@@ -401,18 +452,50 @@ __device__ inline cuda::std::optional<T> decode_int(device_span<uint8_t const> e
   return cudf::io::unaligned_load<T>(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:
+//   - "<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 {}; }
-
-    sub_val = locate_object_field(sub_val, field_id.value());
+    auto const step = path.element<cudf::string_view>(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<uint64_t>(c - '0');
+        if (index > static_cast<uint64_t>(cuda::std::numeric_limits<size_type>::max())) {
+          ok = false;
+          break;
+        }
+      }
+      if (!ok || lo == hi) { return {}; }
+      sub_val = locate_array_element(sub_val, static_cast<size_type>(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;

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

@@ -5,13 +5,16 @@
 
 #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 {
@@ -56,13 +59,53 @@ std::vector<std::string> 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: "[<non-negative integer>]".
+      // 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;
   }
 

cpp/tests/io/experimental/variant_extract_test.cpp

@@ -490,16 +490,66 @@ 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<void>(cudf::io::parquet::experimental::get_variant_field(col, bad, stream)),
       std::invalid_argument)
       << "path that should have thrown: " << bad;
   }
 }
 
+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>{int8_t{2}});
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[2]"),
+                                 cudf::test::fixed_width_column_wrapper<int8_t>{int8_t{5}});
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*get("$[3]"),
+                                 cudf::test::fixed_width_column_wrapper<int8_t>{int8_t{9}});
+
+  // Out-of-bounds index resolves to null.
+  cudf::test::fixed_width_column_wrapper<int8_t> 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<int8_t> 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);