Jackie-Jiang commented on code in PR #16306:
URL: https://github.com/apache/pinot/pull/16306#discussion_r2216667172


##########
pinot-common/src/main/java/org/apache/pinot/sql/parsers/ParserUtils.java:
##########
@@ -77,15 +77,25 @@ private static void 
validateJsonExtractScalarFunction(List<Expression> operands)
   }
 
   private static void validateJsonExtractKeyFunction(List<Expression> 
operands) {
-    // Check that there are exactly 2 arguments
-    if (operands.size() != 2) {
+    // Check that there are 2, 3, or 4 arguments

Review Comment:
   Update comment



##########
pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java:
##########
@@ -113,7 +113,9 @@ public enum TransformFunctionType {
           List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, 
SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER,
               SqlTypeFamily.CHARACTER), i -> i > 2)),
   JSON_EXTRACT_KEY("jsonExtractKey", ReturnTypes.TO_ARRAY,
-      OperandTypes.family(List.of(SqlTypeFamily.CHARACTER, 
SqlTypeFamily.CHARACTER))),
+      OperandTypes.family(
+          List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, 
SqlTypeFamily.CHARACTER),
+          i -> i >= 2 && i <= 3)),

Review Comment:
   (minor) No need to set upper bound
   ```suggestion
             i -> i > 1)),
   ```



##########
pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/JsonExtractKeyTransformFunction.java:
##########
@@ -90,8 +102,24 @@ public String[][] transformToStringValuesMV(ValueBlock 
valueBlock) {
     initStringValuesMV(length);
     String[] jsonStrings = 
_jsonFieldTransformFunction.transformToStringValuesSV(valueBlock);
     for (int i = 0; i < length; i++) {
-      List<String> values = 
JSON_PARSER_CONTEXT.parse(jsonStrings[i]).read(_jsonPath);
-      _stringValuesMV[i] = values.toArray(new String[0]);
+      // Call the appropriate JsonFunctions method based on available 
parameters
+      List values;
+      try {
+        if (_maxDepth != Integer.MAX_VALUE || _dotNotation) {
+          // Call 4-parameter method

Review Comment:
   Update comment



##########
pinot-common/src/main/java/org/apache/pinot/common/function/scalar/JsonFunctions.java:
##########
@@ -321,17 +329,226 @@ private static void setValuesToMap(String keyColumnName, 
String valueColumnName,
       Map<String, String> objMap = (Map) obj;
       result.put(objMap.get(keyColumnName), objMap.get(valueColumnName));
     } else {
-      ObjectMapper mapper = new ObjectMapper();
       JsonNode mapNode;
       try {
-        mapNode = mapper.readTree(obj.toString());
-      } catch (JsonProcessingException e) {
+        mapNode = JsonUtils.stringToJsonNode(obj.toString());
+      } catch (IOException e) {
         throw new RuntimeException(e);
       }
       result.put(mapNode.get(keyColumnName).asText(), 
mapNode.get(valueColumnName).asText());
     }
   }
 
+  /**
+   * Extract keys from JSON object using JsonPath.
+   * <p>
+   * Examples:
+   * - jsonExtractKey('{"a": 1, "b": {"c": 2}}', '$.*') returns ["$['a']", 
"$['b']"]
+   * - jsonExtractKey('{"a": 1, "b": {"c": 2}}', '$.*') returns ["$['a']", 
"$['b']"]
+   * - jsonExtractKey('{"a": 1, "b": {"c": 2}}', '$..**') returns ["$['a']", 
"$['b']", "$['b']['c']"]
+   * - jsonExtractKey('{"a": [1, 2]}', '$.*') returns ["$['a']"]
+   *
+   * @param jsonObj  JSON object or string
+   * @param jsonPath JsonPath expression to extract keys
+   * @return List of key paths matching the JsonPath, empty list if input is 
null or invalid
+   */
+  @ScalarFunction
+  public static List jsonExtractKey(Object jsonObj, String jsonPath) {

Review Comment:
   Does this always return `List<String>`?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to