From 720bbd794540edf8dfe4377a5e0d7181eb167b1d Mon Sep 17 00:00:00 2001
From: yah01 <yang.cen@zilliz.com>
Date: Tue, 6 Jun 2023 16:28:37 +0800
Subject: [PATCH] Add array_at method to get array from JSON (#24671) (#24676)

Signed-off-by: yah01 <yang.cen@zilliz.com>
---
 internal/core/src/common/Json.h | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/internal/core/src/common/Json.h b/internal/core/src/common/Json.h
index 03ded9a259..a422d4ee0c 100644
--- a/internal/core/src/common/Json.h
+++ b/internal/core/src/common/Json.h
@@ -29,6 +29,12 @@
 
 #include "exceptions/EasyAssert.h"
 #include "simdjson.h"
+#include "simdjson/common_defs.h"
+#include "simdjson/dom/array.h"
+#include "simdjson/dom/document.h"
+#include "simdjson/dom/element.h"
+#include "simdjson/error.h"
+#include "simdjson/padded_string.h"
 
 namespace milvus {
 using document = simdjson::ondemand::document;
@@ -99,6 +105,18 @@ class Json {
         return doc;
     }
 
+    value_result<simdjson::dom::element>
+    dom_doc() const {
+        thread_local simdjson::dom::parser parser;
+
+        // it's always safe to add the padding,
+        // as we have allocated the memory with this padding
+        auto doc = parser.parse(data_);
+        AssertInfo(doc.error() == simdjson::SUCCESS, std::string("failed to parse the json ") + data_.data() + ": " +
+                                                         simdjson::error_message(doc.error()));
+        return doc;
+    }
+
     bool
     exist(std::string_view pointer) const {
         return doc().at_pointer(pointer).error() == simdjson::SUCCESS;
@@ -120,6 +138,15 @@ class Json {
         return doc().at_pointer(pointer).get<T>();
     }
 
+    // get dom array by JSON pointer,
+    // call `size()` to get array size,
+    // call `at()` to get array element by index,
+    // iterate through array elements by iterator.
+    value_result<simdjson::dom::array>
+    array_at(std::string_view pointer) const {
+        return dom_doc().at_pointer(pointer).get_array();
+    }
+
     std::string_view
     data() const {
         return data_;