Jackie-Jiang commented on code in PR #18979:
URL: https://github.com/apache/pinot/pull/18979#discussion_r3873872971
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunction.java:
##########
@@ -831,7 +813,9 @@ private <T> IntFunction<T> getResultExtractor(ValueBlock
valueBlock, ParseContex
boolean useBigDecimal) {
if (_jsonFieldTransformFunction.getResultMetadata().getDataType() ==
DataType.BYTES) {
byte[][] jsonBytes =
_jsonFieldTransformFunction.transformToBytesValuesSV(valueBlock);
- IntFunction<T> jaywayExtractor = i ->
parseContext.parseUtf8(jsonBytes[i]).read(_jsonPath);
+ IntFunction<T> jaywayExtractor =
Review Comment:
(minor) Suggest reverting these 2 changes as the type is known, and this
introduces small overhead without gain
##########
pinot-common/src/main/java/org/apache/pinot/common/function/scalar/JsonFunctions.java:
##########
@@ -611,6 +630,321 @@ public static Object jsonExtractObject(@Nullable Object
object) {
return null;
}
+ /// Extract a scalar (or scalar-array) value from a JSON document and coerce
it to `resultsType`.
+ ///
+ /// Scalar-function counterpart of the `jsonExtractScalar` transform
(`JsonExtractScalarTransformFunction` in
+ /// pinot-core), so that `json_extract_scalar(...)` resolves in the
multi-stage engine and in ad-hoc scalar
+ /// contexts. `resultsType` is a Pinot [DataType] name, optionally suffixed
with `_ARRAY` for a multi-value
+ /// result. Supported types are
`INT/LONG/FLOAT/DOUBLE/BIG_DECIMAL/BOOLEAN/TIMESTAMP/STRING/JSON/BYTES` and
+ /// the `INT/LONG/FLOAT/DOUBLE/BIG_DECIMAL/STRING` array variants.
+ ///
+ /// The document may be a `String`, a UTF-8 encoded `byte[]` (BYTES columns)
or an already-parsed container.
+ /// Coercion mirrors the transform exactly: `BOOLEAN` is returned as its
stored `INT` (0/1), `TIMESTAMP` as
+ /// epoch millis (numeric values as-is, strings via ISO-8601), `BIG_DECIMAL`
/ `STRING` / `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 `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 [#jsonExtractScalar(Object, String, String)]. `defaultValue` is
returned (coerced to `resultsType`)
+ /// when the path resolves to `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);
+ }
+
+ /// Reads `jsonPath` from a JSON `String`, UTF-8 `byte[]`, or already-parsed
document.
+ /// Shared with `JsonExtractScalarTransformFunction` so both stay on the
same input dispatch.
+ /// A missing path returns `null` (`Option.SUPPRESS_EXCEPTIONS`). Malformed
input throws.
+ @Nullable
+ public static <T> T readJsonPathInternal(Object jsonInput, String jsonPath,
ParseContext parseContext) {
+ return parseJsonDocument(jsonInput, parseContext).read(jsonPath,
NO_PREDICATES);
+ }
+
+ /// Compiled-path counterpart of [#readJsonPathInternal(Object, String,
ParseContext)].
+ @Nullable
+ public static <T> T readJsonPathInternal(Object jsonInput, JsonPath
jsonPath, ParseContext parseContext) {
+ return parseJsonDocument(jsonInput, parseContext).read(jsonPath);
+ }
+
+ private static DocumentContext parseJsonDocument(Object jsonInput,
ParseContext parseContext) {
+ if (jsonInput instanceof String) {
+ return parseContext.parse((String) jsonInput);
+ }
+ if (jsonInput instanceof byte[]) {
+ // BYTES columns carry the raw UTF-8 document; parse(Object) would treat
the array as already parsed.
+ return parseContext.parseUtf8((byte[]) jsonInput);
+ }
+ return parseContext.parse(jsonInput);
+ }
+
+ @Nullable
+ private static Object readJsonPathValue(@Nullable Object jsonInput, String
jsonPath, boolean useBigDecimal) {
+ if (jsonInput == null) {
+ return null;
+ }
+ try {
+ return readJsonPathInternal(jsonInput, jsonPath,
+ useBigDecimal ? PARSE_CONTEXT_WITH_BIG_DECIMAL : PARSE_CONTEXT);
+ } catch (Exception e) {
+ // Malformed JSON (e.g. a plain-text row) is treated as unresolved,
mirroring the transform which swallows
+ // per-row extraction errors; the caller then applies the default or
throws.
+ return null;
+ }
+ }
+
+ @Nullable
+ private static Object[] readJsonPathArray(@Nullable Object jsonInput, String
jsonPath, boolean useBigDecimal) {
+ if (jsonInput == null) {
+ return null;
+ }
+ try {
+ return convertObjectToArray(readJsonPathInternal(jsonInput, jsonPath,
+ useBigDecimal ? PARSE_CONTEXT_WITH_BIG_DECIMAL : PARSE_CONTEXT));
+ } catch (Exception e) {
+ return null;
+ }
+ }
+
+ private static Object coerceScalar(Object value, DataType dataType) {
+ switch (dataType) {
+ case INT:
+ return toInt(value, false);
+ case BOOLEAN:
+ return toInt(value, true);
+ case LONG:
+ return toLong(value, false);
+ case TIMESTAMP:
+ return toLong(value, true);
+ case FLOAT:
+ return toFloat(value);
+ case DOUBLE:
+ return toDouble(value);
+ case BIG_DECIMAL:
+ return toBigDecimal(value);
+ case STRING:
+ case JSON:
+ return toStringValue(value);
+ case BYTES:
+ return BytesUtils.toBytes(value.toString());
+ default:
+ throw new
IllegalArgumentException(unsupportedResultsTypeMessage(dataType.name()));
+ }
+ }
+
+ private static Object coerceScalarArray(@Nullable Object[] array, DataType
dataType, @Nullable Object defaultValue) {
+ switch (dataType) {
Review Comment:
BOOLEAN_ARRAY and TIMESTAMP_ARRAY is missing
##########
pinot-common/src/main/java/org/apache/pinot/common/function/scalar/JsonFunctions.java:
##########
@@ -611,6 +630,321 @@ public static Object jsonExtractObject(@Nullable Object
object) {
return null;
}
+ /// Extract a scalar (or scalar-array) value from a JSON document and coerce
it to `resultsType`.
+ ///
+ /// Scalar-function counterpart of the `jsonExtractScalar` transform
(`JsonExtractScalarTransformFunction` in
+ /// pinot-core), so that `json_extract_scalar(...)` resolves in the
multi-stage engine and in ad-hoc scalar
+ /// contexts. `resultsType` is a Pinot [DataType] name, optionally suffixed
with `_ARRAY` for a multi-value
+ /// result. Supported types are
`INT/LONG/FLOAT/DOUBLE/BIG_DECIMAL/BOOLEAN/TIMESTAMP/STRING/JSON/BYTES` and
+ /// the `INT/LONG/FLOAT/DOUBLE/BIG_DECIMAL/STRING` array variants.
+ ///
+ /// The document may be a `String`, a UTF-8 encoded `byte[]` (BYTES columns)
or an already-parsed container.
+ /// Coercion mirrors the transform exactly: `BOOLEAN` is returned as its
stored `INT` (0/1), `TIMESTAMP` as
+ /// epoch millis (numeric values as-is, strings via ISO-8601), `BIG_DECIMAL`
/ `STRING` / `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 `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 [#jsonExtractScalar(Object, String, String)]. `defaultValue` is
returned (coerced to `resultsType`)
+ /// when the path resolves to `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);
+ }
+
+ /// Reads `jsonPath` from a JSON `String`, UTF-8 `byte[]`, or already-parsed
document.
+ /// Shared with `JsonExtractScalarTransformFunction` so both stay on the
same input dispatch.
+ /// A missing path returns `null` (`Option.SUPPRESS_EXCEPTIONS`). Malformed
input throws.
+ @Nullable
+ public static <T> T readJsonPathInternal(Object jsonInput, String jsonPath,
ParseContext parseContext) {
Review Comment:
(minor) This is no longer an internal (usually private) method.
Take a look at the comment in `JsonExtractScalarTransformFunction`, maybe we
can make it private/package-private
--
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]
