Vamsi-klu commented on code in PR #18979:
URL: https://github.com/apache/pinot/pull/18979#discussion_r3755547877
##########
pinot-common/src/main/java/org/apache/pinot/common/function/scalar/JsonFunctions.java:
##########
@@ -394,6 +408,306 @@ public static Object jsonExtractObject(@Nullable Object
object) {
return null;
}
+ /**
+ * Extract a scalar (or scalar-array) value from a JSON document and coerce
it to {@code resultsType}.
+ * <p>Scalar-function counterpart of the {@code jsonExtractScalar} transform
(see
+ * {@link
org.apache.pinot.core.operator.transform.function.JsonExtractScalarTransformFunction}),
so that
+ * {@code json_extract_scalar(...)} resolves in the multi-stage engine and
in ad-hoc scalar contexts.
+ * {@code resultsType} is a Pinot {@link DataType} name, optionally suffixed
with {@code _ARRAY} for a
+ * multi-value result. Supported types are
+ * {@code
INT/LONG/FLOAT/DOUBLE/BIG_DECIMAL/BOOLEAN/TIMESTAMP/STRING/JSON/BYTES} and the
+ * {@code INT/LONG/FLOAT/DOUBLE/BIG_DECIMAL/STRING} array variants.
+ * <p>Coercion mirrors the transform exactly: {@code BOOLEAN} is returned as
its stored {@code INT} (0/1),
+ * {@code TIMESTAMP} as epoch millis (numeric values as-is, strings via
ISO-8601), {@code BIG_DECIMAL} /
+ * {@code STRING} / {@code JSON} use a BigDecimal-preserving parser. Without
a default value an unresolved
+ * single-value path throws; a multi-value path yields an empty array, but a
{@code null} element inside a
+ * resolved array still throws. A malformed JSON document is treated as
unresolved.
+ */
+ @ScalarFunction
+ public static Object jsonExtractScalar(Object jsonInput, String jsonPath,
String resultsType) {
+ return jsonExtractScalar(jsonInput, jsonPath, resultsType, null);
+ }
+
+ /**
+ * See {@link #jsonExtractScalar(Object, String, String)}. {@code
defaultValue} is returned (coerced to
+ * {@code resultsType}) when the path resolves to {@code null} or the
document is malformed.
+ */
+ @ScalarFunction(nullableParameters = true)
+ public static Object jsonExtractScalar(@Nullable Object jsonInput, String
jsonPath, String resultsType,
+ @Nullable Object defaultValue) {
+ String type = resultsType.toUpperCase();
+ boolean isSingleValue = !type.endsWith("_ARRAY");
+ DataType dataType;
+ try {
+ dataType = DataType.valueOf(isSingleValue ? type : type.substring(0,
type.length() - 6));
+ } catch (IllegalArgumentException e) {
+ throw new
IllegalArgumentException(unsupportedResultsTypeMessage(resultsType));
+ }
+ // BIG_DECIMAL / STRING / JSON must read floats as BigDecimal to preserve
precision, matching the transform.
+ boolean useBigDecimal =
+ dataType == DataType.BIG_DECIMAL || dataType == DataType.STRING ||
dataType == DataType.JSON;
+ if (isSingleValue) {
+ Object value = readJsonPathValue(jsonInput, jsonPath, useBigDecimal);
+ if (value == null) {
+ if (defaultValue != null) {
+ return coerceScalar(defaultValue, dataType);
+ }
+ throw new IllegalArgumentException(
+ "Cannot resolve JSON path on some records. Consider setting a
default value.");
+ }
+ return coerceScalar(value, dataType);
+ }
+ return coerceScalarArray(readJsonPathArray(jsonInput, jsonPath,
useBigDecimal), dataType, defaultValue);
+ }
+
+ @Nullable
+ private static Object readJsonPathValue(@Nullable Object jsonInput, String
jsonPath, boolean useBigDecimal) {
+ if (jsonInput == null) {
+ return null;
+ }
+ try {
+ ParseContext parseContext = useBigDecimal ?
PARSE_CONTEXT_WITH_BIG_DECIMAL : PARSE_CONTEXT;
+ if (jsonInput instanceof String) {
+ return parseContext.parse((String) jsonInput).read(jsonPath,
NO_PREDICATES);
+ }
+ return parseContext.parse(jsonInput).read(jsonPath, NO_PREDICATES);
Review Comment:
Fixed in `5ea20f0`: `byte[]` input routes through `parseContext.parseUtf8`
in both `readJsonPathValue` and `readJsonPathArray`, and for both the plain
parse context and the BigDecimal one. Missing the BigDecimal context would have
left BIG_DECIMAL, STRING and JSON result types broken on BYTES input while INT
and LONG worked. This now matches what `JsonExtractScalarTransformFunction`
already does for BYTES columns.
##########
pinot-common/src/test/java/org/apache/pinot/common/function/JsonFunctionsTest.java:
##########
@@ -794,4 +796,119 @@ public void testJsonExtractKeySpecialCharacters()
Assert.assertTrue(jsonPathResult.contains("$['field_with_underscores']"));
Assert.assertTrue(jsonPathResult.contains("$['field with spaces']"));
}
+
+ private static String scalarSampleJson() {
+ return "{"
+ +
"\"i\":42,\"l\":9999999999,\"f\":1.5,\"d\":2.5,\"s\":\"hi\",\"b\":true,"
+ + "\"num5\":5,\"zero\":0,\"bstr1\":\"1\","
+ + "\"tnum\":1514805173000,\"tiso\":\"2018-01-01T11:12:53Z\","
+ + "\"hp\":0.1234567890123456789,\"obj\":{\"k\":\"v\"},"
+ +
"\"arr\":[1,2,3],\"arrnull\":[1,null,3],\"lstr\":\"1.234\",\"lexp\":\"1E1\"}";
+ }
+
+ @Test
+ public void testJsonExtractScalarSingleValueCoercions() {
+ String json = scalarSampleJson();
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.i", "INT"), 42);
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.l", "LONG"),
9999999999L);
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.f", "FLOAT"), 1.5f);
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.d", "DOUBLE"), 2.5d);
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.s", "STRING"), "hi");
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.obj", "STRING"),
"{\"k\":\"v\"}");
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.obj", "JSON"),
"{\"k\":\"v\"}");
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.i", "BIG_DECIMAL"),
new BigDecimal("42"));
+ // BOOLEAN follows Pinot's numeric convention (non-zero number is true)
and is returned as stored INT (0/1).
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.b", "BOOLEAN"), 1);
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.num5", "BOOLEAN"),
1);
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.zero", "BOOLEAN"),
0);
+ // The string "1" must go through BooleanUtils.toInt;
Boolean.parseBoolean("1") would wrongly yield 0.
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.bstr1", "BOOLEAN"),
1);
+ // TIMESTAMP: numeric epoch as-is; ISO-8601 string via TimestampUtils
(Long.parseLong would throw).
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.tnum", "TIMESTAMP"),
1514805173000L);
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.tiso", "TIMESTAMP"),
1514805173000L);
+ // LONG from numeric strings uses BigDecimal truncation/exponent handling;
Long.parseLong would throw.
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.lstr", "LONG"), 1L);
+ assertEquals(JsonFunctions.jsonExtractScalar(json, "$.lexp", "LONG"), 10L);
+ }
+
+ @Test
+ public void testJsonExtractScalarBigDecimalPreservesPrecision() {
+ // Default float parsing collapses to a double (0.12345678901234568); the
BigDecimal parser must not.
+ Object result = JsonFunctions.jsonExtractScalar(scalarSampleJson(),
"$.hp", "BIG_DECIMAL");
+ assertEquals(((BigDecimal) result).compareTo(new
BigDecimal("0.1234567890123456789")), 0);
+ }
+
+ @Test
+ public void testJsonExtractScalarBytes() {
+ assertEquals((byte[]) JsonFunctions.jsonExtractScalar("{\"h\":\"0a0b\"}",
"$.h", "BYTES"),
+ new byte[]{0x0a, 0x0b});
+ }
Review Comment:
Added in `5ea20f0`: BYTES-input parity tests asserting the same results as
the String input for each result type, plus array extraction, the default-value
path, and empty and invalid UTF-8 input, neither of which may NPE or leak a
Jayway internal exception. `JsonFunctionsTest` is 76 tests green locally.
##########
pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTest.java:
##########
@@ -331,11 +331,6 @@ protected Iterator<Object[]>
provideTestSqlWithExecutionException() {
// - checked "Illegal Json Path" as col1 is not actually a json string,
but the call is correctly triggered.
testCases.add(
new Object[]{"SELECT CAST(jsonExtractScalar(col1, 'path', 'INT') AS
INT) FROM a", "Cannot resolve JSON path"});
- // - checked function cannot be found b/c there's no intermediate stage
impl for json_extract_scalar
- testCases.add(new Object[]{
- "SELECT CAST(json_extract_scalar(a.col1, b.col2, 'INT') AS INT) FROM a
JOIN b ON a.col1 = b.col1",
- "Unsupported function: JSONEXTRACTSCALAR"
- });
Review Comment:
Added in `5ea20f0`: a positive intermediate-stage case in the row-count
provider that resolves `json_extract_scalar` across a JOIN with a default
value, since the `col1` test data is not JSON. That restores multi-stage
coverage in the positive direction rather than just removing the negative test.
--
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]
