From 626ebb6d8b977b7a126984d73d3079d068b3d261 Mon Sep 17 00:00:00 2001
From: Bikramjeet Vig <bikramjeet@fb.com>
Date: Tue, 5 Nov 2024 17:28:48 -0800
Subject: [PATCH] Add support for using negative array indices in json path

This allows the use of negative array indices in json path to access elements from the end.

Test Plan:
Added unit tests
---
 velox/functions/prestosql/json/JsonExtractor.cpp  |  1 +
 .../prestosql/json/tests/JsonExtractorTest.cpp    | 15 +++++++++++++++
 2 files changed, 16 insertions(+)

diff --git a/velox/functions/prestosql/json/JsonExtractor.cpp b/velox/functions/prestosql/json/JsonExtractor.cpp
index 1c1127d062c1a..243af8a8dc9f7 100644
--- a/velox/functions/prestosql/json/JsonExtractor.cpp
+++ b/velox/functions/prestosql/json/JsonExtractor.cpp
@@ -122,6 +122,7 @@ void extractArray(
     auto rv = folly::tryTo<int32_t>(key);
     if (rv.hasValue()) {
       auto idx = rv.value();
+      idx = idx >= 0 ? idx : arrayLen + idx;
       if (idx >= 0 && idx < arrayLen) {
         ret.push_back(jsonArray->get_ptr(idx));
       }
diff --git a/velox/functions/prestosql/json/tests/JsonExtractorTest.cpp b/velox/functions/prestosql/json/tests/JsonExtractorTest.cpp
index 900a903a0db23..e8fa0baadbac7 100644
--- a/velox/functions/prestosql/json/tests/JsonExtractorTest.cpp
+++ b/velox/functions/prestosql/json/tests/JsonExtractorTest.cpp
@@ -243,6 +243,8 @@ TEST(JsonExtractorTest, arrayJsonValueTest) {
   EXPECT_JSON_VALUE_EQ("[1, 2, 3]"s, "$[0]"s, "1"s);
   EXPECT_JSON_VALUE_EQ("[1, 2]"s, "$[1]"s, "2"s);
   EXPECT_JSON_VALUE_EQ("[1, null]"s, "$[1]"s, "null"s);
+  // Negative Index
+  EXPECT_JSON_VALUE_EQ("[1, 2, 3]"s, "$[-1]"s, "3"s);
   // Out of bounds
   EXPECT_JSON_VALUE_NULL("[1]"s, "$[1]"s);
   // Check skipping complex structures
@@ -412,6 +414,19 @@ TEST(JsonExtractorTest, fullJsonValueTest) {
       "{\"15day\" : 0, \"30day\" : 1, \"90day\" : 2 }"s, "$[\"30day\"]"s, "1"s);
   EXPECT_JSON_VALUE_EQ("{\"a\\\\b\": 4}"s, "$[\"a\\\\b\"]"s, "4"s);
   EXPECT_JSON_VALUE_NULL("{\"fuu\" : null}"s, "$.a.b"s);
+
+ // Negative array index
+  EXPECT_JSON_VALUE_EQ(
+     "[{\"id\": 1, \"text\": \"First\"}, {\"id\": 2, \"text\": \"Second\"},"
+     "{\"id\": 3, \"text\": \"Last\"}]"s, "[-1][\"text\"]"s, "\"Last\""s);
+ EXPECT_JSON_VALUE_EQ(
+     "{\"id\": 1, \"entries\": [\"First\", \"Second\", \"Last\"]}"s,
+     "$.entries[-2]"s,
+     "\"Second\""s);
+ EXPECT_JSON_VALUE_EQ(
+     "{\"id\": 1, \"entries\": [\"First\", \"Second\", \"Last\"]}"s,
+     "$.entries.-3"s,
+     "\"First\""s);
 }
 
 TEST(JsonExtractorTest, invalidJsonPathTest) {