Vamsi-klu commented on code in PR #18979:
URL: https://github.com/apache/pinot/pull/18979#discussion_r3725704223
##########
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:
Good catch, thanks. The transform parses BYTES columns via
ParseContext.parseUtf8 (JsonExtractScalarTransformFunction line 804), while my
helpers only special-case String, so a byte[] falls into parse(Object), which
Jayway treats as an already parsed document instead of raw JSON. Because of the
catch block, that would silently resolve to the default or throw instead of
extracting. I will add an instanceof byte[] branch in both readJsonPathValue
and readJsonPathArray that routes through parseContext.parseUtf8, and push it
with the next update.
##########
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:
Agreed. The current byte[] input path is broken (it falls through to
parse(Object)), so this gap is real. Along with the parseUtf8 fix in
readJsonPathValue and readJsonPathArray, I will add tests that feed the same
sample JSON as UTF-8 bytes (json.getBytes(StandardCharsets.UTF_8)) and assert
parity with the String-input results, covering single-value extraction, array
extraction, and the default-value path.
##########
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:
Fair point. The removed case only asserted the old failure mode, and the
remaining case at line 333 exercises the leaf stage. I will add a positive case
to the row-count data provider that forces intermediate-stage evaluation across
a JOIN, for example SELECT json_extract_scalar(a.col1, '$.x', 'INT', 0) FROM a
JOIN b ON a.col1 = b.col1 (with a default value since col1 test data is not
JSON), asserting the expected row count. That verifies the function resolves
and runs in the v2 intermediate stage and guards against registration
regressions.
--
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]
