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]