From e8491fdf166ec6f97a98622c7ec561dbca46bd23 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 |  5 +++--
 .../prestosql/json/tests/JsonExtractorTest.cpp   | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/velox/functions/prestosql/json/JsonExtractor.cpp b/velox/functions/prestosql/json/JsonExtractor.cpp
index 1c1127d062c1a..abc9c55d05fd6 100644
--- a/velox/functions/prestosql/json/JsonExtractor.cpp
+++ b/velox/functions/prestosql/json/JsonExtractor.cpp
@@ -113,15 +113,16 @@ void extractArray(
     const folly::dynamic* jsonArray,
     const std::string& key,
     JsonVector& ret) {
-  auto arrayLen = jsonArray->size();
+  int64_t arrayLen = jsonArray->size();
   if (key == "*") {
     for (size_t i = 0; i < arrayLen; ++i) {
       ret.push_back(jsonArray->get_ptr(i));
     }
   } else {
-    auto rv = folly::tryTo<int32_t>(key);
+    auto rv = folly::tryTo<int64_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..e7e296e941eaa 100644
--- a/velox/functions/prestosql/json/tests/JsonExtractorTest.cpp
+++ b/velox/functions/prestosql/json/tests/JsonExtractorTest.cpp
@@ -243,6 +243,9 @@ 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);
+  EXPECT_JSON_VALUE_NULL("[1, 2, 3]"s, "$[-4]"s);
   // Out of bounds
   EXPECT_JSON_VALUE_NULL("[1]"s, "$[1]"s);
   // Check skipping complex structures
@@ -412,6 +415,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) {