This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new e40d93100bc Add fast jsonExtractScalar transform variants (#19012)
e40d93100bc is described below
commit e40d93100bce828d11ba6e167dea8c04b6559ac3
Author: Xiang Fu <[email protected]>
AuthorDate: Sun Jul 26 13:28:40 2026 -0700
Add fast jsonExtractScalar transform variants (#19012)
* Add fast jsonExtractScalar transform functions
* Fix jsonExtractScalar operand checker to accept constant-foldable jsonPath
The shared operand type checker required jsonPath and resultsType to be
literal SqlNodes. Operand checking runs before PinotEvaluateLiteralRule
folds constants, so this rejected queries such as
jsonExtractScalar(col, CONCAT('$.', 'foo'), 'INT') that fold to a literal
and plan/execute correctly on master.
Relax the jsonPath operand to accept any CHARACTER expression. Keep the
literal requirement on resultsType (read during validation for return-type
inference; a foldable value there would silently infer VARCHAR) and on
defaultValue. Genuinely non-literal operands are still rejected on the leaf
stage by ParserUtils#validateFunction.
Restores the QueryRunnerTest case the original checker had forced to change,
and adds planner + runtime regression coverage for foldable and non-literal
jsonPath across all three variants.
---
.../common/function/FastJsonPathExtractor.java | 11 +
.../common/function/TransformFunctionType.java | 56 ++++-
.../org/apache/pinot/sql/parsers/ParserUtils.java | 19 +-
.../common/function/FastJsonPathExtractorTest.java | 3 +
.../JsonExtractScalarTransformFunction.java | 173 ++++++++++++--
.../function/TransformFunctionFactory.java | 4 +
.../pinot/core/function/FunctionRegistryTest.java | 3 +-
.../JsonExtractScalarTransformFunctionTest.java | 260 +++++++++++++--------
.../integration/tests/custom/JsonPathTest.java | 39 +++-
.../resources/udf-test-results/all-functions.yaml | 32 +++
.../apache/pinot/query/QueryCompilationTest.java | 68 ++++++
.../pinot/query/QueryEnvironmentTestBase.java | 7 +
.../query/runtime/queries/QueryRunnerTest.java | 17 ++
13 files changed, 559 insertions(+), 133 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/function/FastJsonPathExtractor.java
b/pinot-common/src/main/java/org/apache/pinot/common/function/FastJsonPathExtractor.java
index 2979d8b4bcc..c2dd133e462 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/function/FastJsonPathExtractor.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/function/FastJsonPathExtractor.java
@@ -164,6 +164,17 @@ public final class FastJsonPathExtractor {
}
}
+ /// UTF-8 overload of [#extract(String, SimpleJsonPath[], Object[], boolean,
boolean)] with the same output and
+ /// exception contract.
+ public static void extract(byte[] json, SimpleJsonPath[] paths, Object[]
out, boolean useBigDecimal,
+ boolean earlyExit) {
+ try (JsonParser parser = FACTORY.createParser(json)) {
+ extract(parser, paths, out, useBigDecimal, earlyExit);
+ } catch (IOException e) {
+ throw new InvalidJsonException(e);
+ }
+ }
+
private static void extract(JsonParser parser, SimpleJsonPath[] paths,
Object[] out, boolean useBigDecimal,
boolean earlyExit)
throws IOException {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java
b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java
index e06918a49e2..f4584173ebc 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java
@@ -30,6 +30,7 @@ import org.apache.calcite.sql.type.ReturnTypes;
import org.apache.calcite.sql.type.SqlOperandCountRanges;
import org.apache.calcite.sql.type.SqlOperandTypeChecker;
import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.apache.calcite.sql.type.SqlSingleOperandTypeChecker;
import org.apache.calcite.sql.type.SqlTypeFamily;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.calcite.sql.type.SqlTypeTransforms;
@@ -103,10 +104,11 @@ public enum TransformFunctionType {
// JSON extract functions
JSON_EXTRACT_SCALAR("jsonExtractScalar",
- opBinding -> positionalReturnTypeInferenceFromStringLiteral(opBinding,
2, SqlTypeName.VARCHAR),
- OperandTypes.family(
- List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER,
SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER),
- i -> i == 3)),
+ TransformFunctionType::jsonExtractScalarReturnTypeInference,
jsonExtractScalarOperandTypeChecker()),
+ JSON_EXTRACT_SCALAR_FAST("jsonExtractScalarFast",
+ TransformFunctionType::jsonExtractScalarReturnTypeInference,
jsonExtractScalarOperandTypeChecker()),
+ JSON_EXTRACT_SCALAR_FIRST_MATCH("jsonExtractScalarFirstMatch",
+ TransformFunctionType::jsonExtractScalarReturnTypeInference,
jsonExtractScalarOperandTypeChecker()),
JSON_EXTRACT_INDEX("jsonExtractIndex",
opBinding -> positionalReturnTypeInferenceFromStringLiteral(opBinding,
2, SqlTypeName.VARCHAR),
OperandTypes.family(
@@ -317,6 +319,50 @@ public enum TransformFunctionType {
return opBinding.getTypeFactory().createSqlType(defaultSqlType);
}
+ private static RelDataType
jsonExtractScalarReturnTypeInference(SqlOperatorBinding opBinding) {
+ if (opBinding.getOperandCount() > 2 && opBinding.isOperandLiteral(2,
false)) {
+ String resultsType = opBinding.getOperandLiteralValue(2,
String.class).toUpperCase();
+ RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
+ switch (resultsType) {
+ case "JSON":
+ return typeFactory.createSqlType(SqlTypeName.VARCHAR);
+ case "BOOLEAN_ARRAY":
+ return
typeFactory.createArrayType(typeFactory.createSqlType(SqlTypeName.BOOLEAN), -1);
+ case "TIMESTAMP_ARRAY":
+ return
typeFactory.createArrayType(typeFactory.createSqlType(SqlTypeName.TIMESTAMP),
-1);
+ default:
+ break;
+ }
+ }
+ return positionalReturnTypeInferenceFromStringLiteral(opBinding, 2,
SqlTypeName.VARCHAR);
+ }
+
+ /// Operand checker shared by `jsonExtractScalar` and its `Fast` /
`FirstMatch` variants.
+ ///
+ /// `jsonPath` deliberately does **not** require [OperandTypes#LITERAL].
Operand checking runs on the raw
+ /// `SqlNode` tree, before `PinotEvaluateLiteralRule` folds constant
expressions, so demanding a literal here
+ /// would reject `jsonExtractScalar(col, CONCAT('$.', 'foo'), 'INT')` and
other constant-foldable paths that fold
+ /// to a literal and run correctly. A genuinely non-literal `jsonPath` is
rejected later, when the leaf stage
+ /// converts the folded call back into a Pinot expression.
+ ///
+ /// `resultsType` **does** require a literal, because
[#jsonExtractScalarReturnTypeInference] reads it during
+ /// validation — also before folding — to derive the return type. Accepting
a foldable `resultsType` there would
+ /// silently infer `VARCHAR` while the leaf stage extracts the real type,
producing a plan whose schema disagrees
+ /// with its data.
+ private static SqlOperandTypeChecker jsonExtractScalarOperandTypeChecker() {
+ SqlSingleOperandTypeChecker jsonInputTypeChecker =
OperandTypes.or(OperandTypes.CHARACTER, OperandTypes.BINARY);
+ SqlSingleOperandTypeChecker resultsTypeChecker =
OperandTypes.and(OperandTypes.CHARACTER, OperandTypes.LITERAL);
+ return OperandTypes.or(
+ OperandTypes.sequence(
+ (operator, ignored) -> "'" + operator.getName()
+ + "(<CHARACTER_OR_BINARY>, <CHARACTER>, <CHARACTER_LITERAL>)'",
+ jsonInputTypeChecker, OperandTypes.CHARACTER, resultsTypeChecker),
+ OperandTypes.sequence(
+ (operator, ignored) -> "'" + operator.getName()
+ + "(<CHARACTER_OR_BINARY>, <CHARACTER>, <CHARACTER_LITERAL>,
<LITERAL>)'",
+ jsonInputTypeChecker, OperandTypes.CHARACTER, resultsTypeChecker,
OperandTypes.NULLABLE_LITERAL));
+ }
+
private static RelDataType componentType(SqlOperatorBinding opBinding) {
return opBinding.getOperandType(0).getComponentType();
}
@@ -362,6 +408,8 @@ public enum TransformFunctionType {
return typeFactory.createSqlType(SqlTypeName.VARBINARY);
case "BIG_DECIMAL":
return typeFactory.createSqlType(SqlTypeName.DECIMAL);
+ case "BIG_DECIMAL_ARRAY":
+ return
typeFactory.createArrayType(typeFactory.createSqlType(SqlTypeName.DECIMAL), -1);
default:
SqlTypeName sqlTypeName = SqlTypeName.get(operandTypeStr);
if (sqlTypeName == null) {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/ParserUtils.java
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/ParserUtils.java
index 23b8a70237d..4406a44707d 100644
--- a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/ParserUtils.java
+++ b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/ParserUtils.java
@@ -29,7 +29,13 @@ public class ParserUtils {
public static void validateFunction(String canonicalName, List<Expression>
operands) {
switch (canonicalName) {
case "jsonextractscalar":
- validateJsonExtractScalarFunction(operands);
+ validateJsonExtractScalarFunction("jsonExtractScalar", operands);
+ break;
+ case "jsonextractscalarfast":
+ validateJsonExtractScalarFunction("jsonExtractScalarFast", operands);
+ break;
+ case "jsonextractscalarfirstmatch":
+ validateJsonExtractScalarFunction("jsonExtractScalarFirstMatch",
operands);
break;
case "jsonextractkey":
validateJsonExtractKeyFunction(operands);
@@ -133,19 +139,20 @@ public class ParserUtils {
return result.toString();
}
- private static void validateJsonExtractScalarFunction(List<Expression>
operands) {
+ private static void validateJsonExtractScalarFunction(String functionName,
List<Expression> operands) {
// Check that there are 3 or 4 arguments
int numOperands = operands.size();
if (numOperands != 3 && numOperands != 4) {
throw new SqlCompilationException(
- "Expect 3 or 4 arguments for transform function:
jsonExtractScalar(jsonFieldName, 'jsonPath', "
- + "'resultsType', ['defaultValue'])");
+ "Expect 3 or 4 arguments for transform function: " + functionName
+ + "(jsonFieldName, 'jsonPath', 'resultsType',
['defaultValue'])");
}
if (!operands.get(1).isSetLiteral() || !operands.get(2).isSetLiteral() ||
(numOperands == 4 && !operands.get(3)
.isSetLiteral())) {
throw new SqlCompilationException(
- "Expect the 2nd/3rd/4th argument of transform function:
jsonExtractScalar(jsonFieldName, 'jsonPath', "
- + "'resultsType', ['defaultValue']) to be a single-quoted
literal value.");
+ "Expect the 2nd and 3rd arguments of transform function: " +
functionName
+ + "(jsonFieldName, 'jsonPath', 'resultsType', ['defaultValue'])
to be single-quoted literal values, "
+ + "and the optional 4th argument to be a literal value.");
}
}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/function/FastJsonPathExtractorTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/function/FastJsonPathExtractorTest.java
index a29ac274f84..41c398d961d 100644
---
a/pinot-common/src/test/java/org/apache/pinot/common/function/FastJsonPathExtractorTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/common/function/FastJsonPathExtractorTest.java
@@ -620,7 +620,10 @@ public class FastJsonPathExtractorTest {
}
Object[] batch = new Object[paths.length];
FastJsonPathExtractor.extract(json, compiled, batch, false, false);
+ Object[] bytesBatch = new Object[paths.length];
+ FastJsonPathExtractor.extract(json.getBytes(StandardCharsets.UTF_8),
compiled, bytesBatch, false, false);
for (int i = 0; i < paths.length; i++) {
+ assertEquals(bytesBatch[i], batch[i], paths[i]);
assertEquals(batch[i], FastJsonPathExtractor.extract(json, compiled[i],
false, false), paths[i]);
assertEquals(batch[i], PLAIN_CONTEXT.parse(json).read(paths[i],
NO_PREDICATES), paths[i]);
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunction.java
index dee336e2dc9..21e0f3e0cde 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunction.java
@@ -32,7 +32,9 @@ import java.util.List;
import java.util.Map;
import java.util.function.IntFunction;
import javax.annotation.Nullable;
+import org.apache.pinot.common.function.FastJsonPathExtractor;
import org.apache.pinot.common.function.JsonPathCache;
+import org.apache.pinot.common.function.SimpleJsonPath;
import org.apache.pinot.core.operator.ColumnContext;
import org.apache.pinot.core.operator.blocks.ValueBlock;
import org.apache.pinot.core.operator.transform.TransformResultMetadata;
@@ -41,15 +43,36 @@ import org.apache.pinot.core.util.NumericException;
import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.utils.BooleanUtils;
import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.PinotDataType;
import org.apache.pinot.spi.utils.TimestampUtils;
import org.roaringbitmap.RoaringBitmap;
-/// Implements the `jsonExtractScalar(jsonField, jsonPath, resultsType[,
defaultValue])` transform.
-/// Reads a JSON document from `jsonField` for each row, resolves the
+/// Implements the typed JSON extraction transforms:
+///
+/// | Function | Simple-path behavior |
+/// | --- | --- |
+/// | `jsonExtractScalar` | Builds the existing Jayway DOM; unchanged for
backward compatibility. |
+/// | `jsonExtractScalarFast` | Uses [FastJsonPathExtractor] and scans the
full root value, preserving Jayway results. |
+/// | `jsonExtractScalarFirstMatch` | Uses [FastJsonPathExtractor] and stops
when the addressed value is found. |
+///
+/// Each function reads a JSON document from `jsonField` for each row,
resolves the
/// [Stefan Goessner JsonPath](https://goessner.net/articles/JsonPath/)
expression against it, and
/// converts the resolved value to `resultsType`.
///
+/// ```sql
+/// SELECT jsonExtractScalarFast(payload, '$.user.id', 'LONG', '-1'),
+/// jsonExtractScalarFirstMatch(payload, '$.service.name', 'STRING',
'unknown')
+/// FROM events
+/// ```
+///
+/// The fast variants optimize simple linear paths (`$` followed by `.key`,
`['key']`, or `[index]`) and fall back
+/// to the existing Jayway implementation for complex paths, unsupported input
roots, and any fast-extractor
+/// exception. `FirstMatch` is faster for fields near the start of a document,
but deliberately resolves duplicate
+/// keys to the first non-null occurrence and does not validate malformed
content after the resolved value. Use it
+/// only for well-formed, duplicate-free JSON. `Fast` scans the full root
value and retains Jayway's last-key-wins
+/// and malformed-document behavior; see [FastJsonPathExtractor] for one
documented unaddressed-value edge case.
+///
/// **Arguments:**
/// - `jsonField` — single-value `STRING` or `BYTES` column / transform
expression containing JSON.
/// - `jsonPath` — JsonPath expression used to read the value.
@@ -60,7 +83,10 @@ import org.roaringbitmap.RoaringBitmap;
///
/// **Supported `resultsType`:** `INT`, `LONG`, `FLOAT`, `DOUBLE`,
`BIG_DECIMAL`, `BOOLEAN`, `TIMESTAMP`,
/// `STRING`, `JSON`, `BYTES`, plus `_ARRAY` variants of `INT` / `LONG` /
`FLOAT` / `DOUBLE` /
-/// `BIG_DECIMAL` / `STRING`.
+/// `BIG_DECIMAL` / `BOOLEAN` / `TIMESTAMP` / `STRING`.
+///
+/// Multi-stage planning accepts all of these result types, but currently
lowers `BIG_DECIMAL_ARRAY` to
+/// `DOUBLE_ARRAY`. Use the single-stage engine when decimal-array precision
must be preserved.
///
/// **Per-row coercion** of the JsonPath result to `resultsType`:
/// - `BOOLEAN` (stored as `INT`) follows Pinot's numeric convention — any
non-zero `Number` is true;
@@ -69,6 +95,7 @@ import org.roaringbitmap.RoaringBitmap;
/// [TimestampUtils#toMillisSinceEpoch] (ISO-8601 and numeric millis
strings).
/// - `STRING` returns `String` values as-is; other JSON values are serialized
via
/// [JsonUtils#objectToString].
+/// - `BYTES` decodes a Base64-encoded JSON string, matching
[PinotDataType#JSON].
/// - `BIG_DECIMAL` and `STRING` paths use a BigDecimal-preserving JSON parser
/// (`JSON_PARSER_CONTEXT_WITH_BIG_DECIMAL`) to avoid precision loss on
numeric values; other paths use
/// the default parser.
@@ -76,6 +103,14 @@ import org.roaringbitmap.RoaringBitmap;
/// `parse*(toString())`.
public class JsonExtractScalarTransformFunction extends BaseTransformFunction {
public static final String FUNCTION_NAME = "jsonExtractScalar";
+ public static final String FAST_FUNCTION_NAME = "jsonExtractScalarFast";
+ public static final String FIRST_MATCH_FUNCTION_NAME =
"jsonExtractScalarFirstMatch";
+
+ private enum ExtractionMode {
+ JAYWAY,
+ FAST,
+ FIRST_MATCH
+ }
// This ObjectMapper requires special configurations, hence we can't use
pinot JsonUtils here.
private static final ObjectMapper OBJECT_MAPPER_WITH_BIG_DECIMAL =
@@ -89,17 +124,44 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
new Configuration.ConfigurationBuilder().jsonProvider(new
JacksonJsonProvider())
.mappingProvider(new
JacksonMappingProvider()).options(Option.SUPPRESS_EXCEPTIONS).build());
+ private final String _functionName;
+ private final ExtractionMode _extractionMode;
private TransformFunction _jsonFieldTransformFunction;
private JsonPath _jsonPath;
+ @Nullable
+ private SimpleJsonPath _simpleJsonPath;
private DataType _dataType;
private DataType _storedType;
private Object _defaultValue;
private boolean _defaultIsNull;
private TransformResultMetadata _resultMetadata;
+ public JsonExtractScalarTransformFunction() {
+ this(FUNCTION_NAME, ExtractionMode.JAYWAY);
+ }
+
+ private JsonExtractScalarTransformFunction(String functionName,
ExtractionMode extractionMode) {
+ _functionName = functionName;
+ _extractionMode = extractionMode;
+ }
+
+ /// Full-scan fast variant of [JsonExtractScalarTransformFunction].
+ public static final class Fast extends JsonExtractScalarTransformFunction {
+ public Fast() {
+ super(FAST_FUNCTION_NAME, ExtractionMode.FAST);
+ }
+ }
+
+ /// First-match fast variant of [JsonExtractScalarTransformFunction].
+ public static final class FirstMatch extends
JsonExtractScalarTransformFunction {
+ public FirstMatch() {
+ super(FIRST_MATCH_FUNCTION_NAME, ExtractionMode.FIRST_MATCH);
+ }
+ }
+
@Override
public String getName() {
- return FUNCTION_NAME;
+ return _functionName;
}
@Override
@@ -109,28 +171,30 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
// Check that there are exactly 3 or 4 arguments
if (arguments.size() < 3 || arguments.size() > 4) {
throw new IllegalArgumentException(
- "Expected 3/4 arguments for transform function:
jsonExtractScalar(jsonFieldName, 'jsonPath', 'resultsType',"
- + " ['defaultValue'])");
+ "Expected 3/4 arguments for transform function: " + _functionName
+ + "(jsonFieldName, 'jsonPath', 'resultsType',
['defaultValue'])");
}
TransformFunction firstArgument = arguments.get(0);
if (firstArgument instanceof LiteralTransformFunction ||
!firstArgument.getResultMetadata().isSingleValue()) {
throw new IllegalArgumentException(
- "The first argument of jsonExtractScalar transform function must be
a single-valued column or a transform "
- + "function");
+ "The first argument of " + _functionName
+ + " transform function must be a single-valued column or a
transform function");
}
_jsonFieldTransformFunction = firstArgument;
String jsonPathString = ((LiteralTransformFunction)
arguments.get(1)).getStringLiteral();
_jsonPath = JsonPathCache.INSTANCE.getOrCompute(jsonPathString);
+ _simpleJsonPath = _extractionMode == ExtractionMode.JAYWAY ? null :
SimpleJsonPath.compile(jsonPathString);
String resultsType = ((LiteralTransformFunction)
arguments.get(2)).getStringLiteral().toUpperCase();
boolean isSingleValue = !resultsType.endsWith("_ARRAY");
try {
_dataType = DataType.valueOf(isSingleValue ? resultsType :
resultsType.substring(0, resultsType.length() - 6));
} catch (Exception e) {
throw new IllegalArgumentException(String.format(
- "Unsupported results type: %s for jsonExtractScalar function.
Supported types are: "
+ "Unsupported results type: %s for %s function. Supported types are: "
+
"INT/LONG/FLOAT/DOUBLE/BIG_DECIMAL/BOOLEAN/TIMESTAMP/STRING/JSON/BYTES/"
- +
"INT_ARRAY/LONG_ARRAY/FLOAT_ARRAY/DOUBLE_ARRAY/BIG_DECIMAL_ARRAY/STRING_ARRAY",
resultsType));
+ +
"INT_ARRAY/LONG_ARRAY/FLOAT_ARRAY/DOUBLE_ARRAY/BIG_DECIMAL_ARRAY/BOOLEAN_ARRAY/TIMESTAMP_ARRAY/"
+ + "STRING_ARRAY", resultsType, _functionName));
}
_storedType = _dataType.getStoredType();
if (arguments.size() == 4) {
@@ -170,7 +234,8 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
break;
default:
throw new IllegalArgumentException(
- "Unsupported results type: " + _dataType + " for
jsonExtractScalar function. Supported types are: "
+ "Unsupported results type: " + _dataType + " for " +
_functionName
+ + " function. Supported types are: "
+
"INT/LONG/FLOAT/DOUBLE/BIG_DECIMAL/BOOLEAN/TIMESTAMP/STRING/JSON/BYTES"
);
}
@@ -388,6 +453,34 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
return _stringValuesSV;
}
+ @Override
+ public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) {
+ if (_storedType != DataType.BYTES) {
+ return super.transformToBytesValuesSV(valueBlock);
+ }
+ initBytesValuesSV(valueBlock.getNumDocs());
+ IntFunction<Object> resultExtractor = getResultExtractor(valueBlock);
+ byte[] defaultValue = (byte[]) _defaultValue;
+ int numDocs = valueBlock.getNumDocs();
+ for (int i = 0; i < numDocs; i++) {
+ Object result = null;
+ try {
+ result = resultExtractor.apply(i);
+ } catch (Exception ignored) {
+ }
+ if (result == null) {
+ if (_defaultValue != null) {
+ _bytesValuesSV[i] = defaultValue;
+ continue;
+ }
+ throw new IllegalArgumentException(
+ "Cannot resolve JSON path on some records. Consider setting a
default value.");
+ }
+ _bytesValuesSV[i] = PinotDataType.JSON.toBytes(result);
+ }
+ return _bytesValuesSV;
+ }
+
@Override
public int[][] transformToIntValuesMV(ValueBlock valueBlock) {
if (_storedType != DataType.INT) {
@@ -419,7 +512,7 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
}
throw new IllegalArgumentException(
"At least one of the resolved JSON arrays include nulls, which
is not supported in Pinot. "
- + "Consider setting a default value as the fourth argument
of json_extract_scalar.");
+ + "Consider setting a default value as the fourth argument
of " + _functionName + ".");
}
values[j] = toInt(element, isBoolean);
}
@@ -459,7 +552,7 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
}
throw new IllegalArgumentException(
"At least one of the resolved JSON arrays include nulls, which
is not supported in Pinot. "
- + "Consider setting a default value as the fourth argument
of json_extract_scalar.");
+ + "Consider setting a default value as the fourth argument
of " + _functionName + ".");
}
values[j] = toLong(element, isTimestamp);
}
@@ -498,7 +591,7 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
}
throw new IllegalArgumentException(
"At least one of the resolved JSON arrays include nulls, which
is not supported in Pinot. "
- + "Consider setting a default value as the fourth argument
of json_extract_scalar.");
+ + "Consider setting a default value as the fourth argument
of " + _functionName + ".");
}
values[j] = toFloat(element);
}
@@ -537,7 +630,7 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
}
throw new IllegalArgumentException(
"At least one of the resolved JSON arrays include nulls, which
is not supported in Pinot. "
- + "Consider setting a default value as the fourth argument
of json_extract_scalar.");
+ + "Consider setting a default value as the fourth argument
of " + _functionName + ".");
}
values[j] = toDouble(element);
}
@@ -576,7 +669,7 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
}
throw new IllegalArgumentException(
"At least one of the resolved JSON arrays include nulls, which
is not supported in Pinot. "
- + "Consider setting a default value as the fourth argument
of json_extract_scalar.");
+ + "Consider setting a default value as the fourth argument
of " + _functionName + ".");
}
values[j] = toBigDecimal(element);
}
@@ -615,7 +708,7 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
}
throw new IllegalArgumentException(
"At least one of the resolved JSON arrays include nulls, which
is not supported in Pinot. "
- + "Consider setting a default value as the fourth argument
of json_extract_scalar.");
+ + "Consider setting a default value as the fourth argument
of " + _functionName + ".");
}
values[j] = toString(element);
}
@@ -703,21 +796,57 @@ public class JsonExtractScalarTransformFunction extends
BaseTransformFunction {
}
}
- private <T> IntFunction<T> getResultExtractor(ValueBlock valueBlock,
ParseContext parseContext) {
+ @SuppressWarnings("unchecked")
+ private <T> IntFunction<T> getResultExtractor(ValueBlock valueBlock,
ParseContext parseContext,
+ boolean useBigDecimal) {
if (_jsonFieldTransformFunction.getResultMetadata().getDataType() ==
DataType.BYTES) {
byte[][] jsonBytes =
_jsonFieldTransformFunction.transformToBytesValuesSV(valueBlock);
- return i -> parseContext.parseUtf8(jsonBytes[i]).read(_jsonPath);
+ IntFunction<T> jaywayExtractor = i ->
parseContext.parseUtf8(jsonBytes[i]).read(_jsonPath);
+ if (_simpleJsonPath == null) {
+ return jaywayExtractor;
+ }
+ SimpleJsonPath[] paths = {_simpleJsonPath};
+ Object[] result = new Object[1];
+ boolean earlyExit = _extractionMode == ExtractionMode.FIRST_MATCH;
+ return i -> {
+ if (FastJsonPathExtractor.canExtract(jsonBytes[i])) {
+ try {
+ FastJsonPathExtractor.extract(jsonBytes[i], paths, result,
useBigDecimal, earlyExit);
+ return (T) result[0];
+ } catch (Exception ignored) {
+ // Retry with Jayway so a fast-extractor failure cannot change the
existing result.
+ }
+ }
+ return jaywayExtractor.apply(i);
+ };
} else {
String[] jsonStrings =
_jsonFieldTransformFunction.transformToStringValuesSV(valueBlock);
- return i -> parseContext.parse(jsonStrings[i]).read(_jsonPath);
+ IntFunction<T> jaywayExtractor = i ->
parseContext.parse(jsonStrings[i]).read(_jsonPath);
+ if (_simpleJsonPath == null) {
+ return jaywayExtractor;
+ }
+ SimpleJsonPath[] paths = {_simpleJsonPath};
+ Object[] result = new Object[1];
+ boolean earlyExit = _extractionMode == ExtractionMode.FIRST_MATCH;
+ return i -> {
+ if (FastJsonPathExtractor.canExtract(jsonStrings[i])) {
+ try {
+ FastJsonPathExtractor.extract(jsonStrings[i], paths, result,
useBigDecimal, earlyExit);
+ return (T) result[0];
+ } catch (Exception ignored) {
+ // Retry with Jayway so a fast-extractor failure cannot change the
existing result.
+ }
+ }
+ return jaywayExtractor.apply(i);
+ };
}
}
private <T> IntFunction<T> getResultExtractor(ValueBlock valueBlock) {
- return getResultExtractor(valueBlock, JSON_PARSER_CONTEXT);
+ return getResultExtractor(valueBlock, JSON_PARSER_CONTEXT, false);
}
private <T> IntFunction<T> getResultExtractorWithBigDecimal(ValueBlock
valueBlock) {
- return getResultExtractor(valueBlock,
JSON_PARSER_CONTEXT_WITH_BIG_DECIMAL);
+ return getResultExtractor(valueBlock,
JSON_PARSER_CONTEXT_WITH_BIG_DECIMAL, true);
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
index 4c050f61b73..ca8e18b937c 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
@@ -130,6 +130,10 @@ public class TransformFunctionFactory {
typeToImplementation.put(TransformFunctionType.CAST,
CastTransformFunction.class);
typeToImplementation.put(TransformFunctionType.JSON_EXTRACT_SCALAR,
JsonExtractScalarTransformFunction.class);
+ typeToImplementation.put(TransformFunctionType.JSON_EXTRACT_SCALAR_FAST,
+ JsonExtractScalarTransformFunction.Fast.class);
+
typeToImplementation.put(TransformFunctionType.JSON_EXTRACT_SCALAR_FIRST_MATCH,
+ JsonExtractScalarTransformFunction.FirstMatch.class);
typeToImplementation.put(TransformFunctionType.JSON_EXTRACT_KEY,
JsonExtractKeyTransformFunction.class);
typeToImplementation.put(TransformFunctionType.TIME_CONVERT,
TimeConversionTransformFunction.class);
typeToImplementation.put(TransformFunctionType.DATE_TIME_CONVERT,
DateTimeConversionTransformFunction.class);
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/function/FunctionRegistryTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/function/FunctionRegistryTest.java
index abc996d2447..37efd3c5b10 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/function/FunctionRegistryTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/function/FunctionRegistryTest.java
@@ -46,7 +46,8 @@ public class FunctionRegistryTest {
// TODO: Support these functions
TransformFunctionType.IN, TransformFunctionType.NOT_IN,
TransformFunctionType.IS_TRUE,
TransformFunctionType.IS_NOT_TRUE, TransformFunctionType.IS_FALSE,
TransformFunctionType.IS_NOT_FALSE,
- TransformFunctionType.JSON_EXTRACT_SCALAR,
+ TransformFunctionType.JSON_EXTRACT_SCALAR,
TransformFunctionType.JSON_EXTRACT_SCALAR_FAST,
+ TransformFunctionType.JSON_EXTRACT_SCALAR_FIRST_MATCH,
TransformFunctionType.JSON_EXTRACT_KEY,
TransformFunctionType.TIME_CONVERT,
TransformFunctionType.DATE_TIME_CONVERT_WINDOW_HOP,
TransformFunctionType.ARRAY_LENGTH,
TransformFunctionType.ARRAY_AVERAGE, TransformFunctionType.ARRAY_MIN,
TransformFunctionType.ARRAY_MAX,
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunctionTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunctionTest.java
index 458d549a941..19ae0283a8f 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunctionTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/JsonExtractScalarTransformFunctionTest.java
@@ -22,12 +22,14 @@ import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.math.BigDecimal;
+import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Random;
+import javax.annotation.Nullable;
import org.apache.commons.io.FileUtils;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.RequestContextUtils;
@@ -37,6 +39,7 @@ import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.exception.BadQueryRequestException;
+import org.apache.pinot.spi.utils.BytesUtils;
import org.apache.pinot.spi.utils.JsonUtils;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.apache.pinot.sql.parsers.SqlCompilationException;
@@ -49,16 +52,22 @@ import org.testng.annotations.Test;
public class JsonExtractScalarTransformFunctionTest extends
BaseTransformFunctionTest {
+ private static final String[] JSON_EXTRACT_SCALAR_FUNCTIONS = {
+ JsonExtractScalarTransformFunction.FUNCTION_NAME,
+ JsonExtractScalarTransformFunction.FAST_FUNCTION_NAME,
+ JsonExtractScalarTransformFunction.FIRST_MATCH_FUNCTION_NAME
+ };
protected File _baseDir;
@Test(dataProvider = "testJsonPathTransformFunction")
- public void testJsonPathTransformFunction(String expressionStr, DataType
resultsDataType, boolean isSingleValue) {
+ public void testJsonPathTransformFunction(String functionName, String
expressionStr, DataType resultsDataType,
+ boolean isSingleValue) {
ExpressionContext expression =
RequestContextUtils.getExpression(expressionStr);
TransformFunction transformFunction =
TransformFunctionFactory.get(expression, _dataSourceMap);
Assert.assertTrue(transformFunction instanceof
JsonExtractScalarTransformFunction);
- Assert.assertEquals(transformFunction.getName(),
JsonExtractScalarTransformFunction.FUNCTION_NAME);
+ Assert.assertEquals(transformFunction.getName(), functionName);
Assert.assertEquals(transformFunction.getResultMetadata().getDataType(),
resultsDataType);
Assert.assertEquals(transformFunction.getResultMetadata().isSingleValue(),
isSingleValue);
@@ -170,52 +179,53 @@ public class JsonExtractScalarTransformFunctionTest
extends BaseTransformFunctio
// Test operating on both column and output of another transform (trim) to
avoid passing the evaluation down to the
// storage in order to test transformTransformedValuesToXXXValuesSV()
methods.
- for (String input : new String[]{JSON_COLUMN, String.format("trim(%s)",
JSON_COLUMN)}) {
- // Without default value
- testArguments.add(
- new Object[]{String.format("jsonExtractScalar(%s,'$.intSV','INT')",
input), DataType.INT, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.intMV','INT_ARRAY')", input),
DataType.INT, false});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.longSV','LONG')", input),
DataType.LONG, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.floatSV','FLOAT')", input),
DataType.FLOAT, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.doubleSV','DOUBLE')", input),
DataType.DOUBLE, true});
- testArguments.add(new Object[]{
-
String.format("jsonExtractScalar(%s,'$.bigDecimalSV','BIG_DECIMAL')", input),
DataType.BIG_DECIMAL, true
- });
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.stringSV','STRING')", input),
DataType.STRING, true});
-
- // With default value
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.intSV','INT',0)", input),
DataType.INT, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.intMV','INT_ARRAY',0)", input),
DataType.INT, false});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.longSV','LONG',0)", input),
DataType.LONG, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.floatSV','FLOAT',0)", input),
DataType.FLOAT, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.doubleSV','DOUBLE',0)", input),
DataType.DOUBLE, true});
- testArguments.add(new Object[]{
-
String.format("jsonExtractScalar(%s,'$.bigDecimalSV','BIG_DECIMAL',0)", input),
DataType.BIG_DECIMAL, true
- });
- testArguments.add(new Object[]{
- String.format("jsonExtractScalar(%s,'$.stringSV','STRING','null')",
input), DataType.STRING, true
- });
+ for (String functionName : JSON_EXTRACT_SCALAR_FUNCTIONS) {
+ for (String input : new String[]{JSON_COLUMN, String.format("trim(%s)",
JSON_COLUMN)}) {
+ // Without default value
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.intSV','INT')", functionName, input),
DataType.INT, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.intMV','INT_ARRAY')", functionName,
input), DataType.INT, false});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.longSV','LONG')", functionName, input),
DataType.LONG, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.floatSV','FLOAT')", functionName, input),
DataType.FLOAT, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.doubleSV','DOUBLE')", functionName,
input), DataType.DOUBLE, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.bigDecimalSV','BIG_DECIMAL')",
functionName, input), DataType.BIG_DECIMAL, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.stringSV','STRING')", functionName,
input), DataType.STRING, true});
+
+ // With default value
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.intSV','INT',0)", functionName, input),
DataType.INT, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.intMV','INT_ARRAY',0)", functionName,
input), DataType.INT, false});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.longSV','LONG',0)", functionName, input),
DataType.LONG, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.floatSV','FLOAT',0)", functionName,
input), DataType.FLOAT, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.doubleSV','DOUBLE',0)", functionName,
input), DataType.DOUBLE, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.bigDecimalSV','BIG_DECIMAL',0)",
functionName, input), DataType.BIG_DECIMAL,
+ true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.stringSV','STRING','null')", functionName,
input), DataType.STRING, true});
+ }
}
return testArguments.toArray(new Object[0][]);
}
@Test(dataProvider = "testDefaultValue")
- public void testDefaultValue(String expressionStr, DataType resultsDataType,
boolean isSingleValue) {
+ public void testDefaultValue(String functionName, String expressionStr,
DataType resultsDataType,
+ boolean isSingleValue) {
ExpressionContext expression =
RequestContextUtils.getExpression(expressionStr);
TransformFunction transformFunction =
TransformFunctionFactory.get(expression, _dataSourceMap);
Assert.assertTrue(transformFunction instanceof
JsonExtractScalarTransformFunction);
- Assert.assertEquals(transformFunction.getName(),
JsonExtractScalarTransformFunction.FUNCTION_NAME);
+ Assert.assertEquals(transformFunction.getName(), functionName);
Assert.assertEquals(transformFunction.getResultMetadata().getDataType(),
resultsDataType);
Assert.assertEquals(transformFunction.getResultMetadata().isSingleValue(),
isSingleValue);
@@ -280,23 +290,24 @@ public class JsonExtractScalarTransformFunctionTest
extends BaseTransformFunctio
// Test operating on both column and output of another transform (trim) to
avoid passing the evaluation down to the
// storage in order to test transformTransformedValuesToXXXValuesSV()
methods.
- for (String input : new String[]{DEFAULT_JSON_COLUMN,
String.format("trim(%s)", DEFAULT_JSON_COLUMN)}) {
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.intSV','INT',0)", input),
DataType.INT, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.intMV','INT_ARRAY',0)", input),
DataType.INT, false});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.longSV','LONG',0)", input),
DataType.LONG, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.floatSV','FLOAT',0)", input),
DataType.FLOAT, true});
- testArguments.add(
- new
Object[]{String.format("jsonExtractScalar(%s,'$.doubleSV','DOUBLE',0)", input),
DataType.DOUBLE, true});
- testArguments.add(new Object[]{
-
String.format("jsonExtractScalar(%s,'$.bigDecimalSV','BIG_DECIMAL',0)", input),
DataType.BIG_DECIMAL, true
- });
- testArguments.add(new Object[]{
- String.format("jsonExtractScalar(%s,'$.stringSV','STRING','null')",
input), DataType.STRING, true
- });
+ for (String functionName : JSON_EXTRACT_SCALAR_FUNCTIONS) {
+ for (String input : new String[]{DEFAULT_JSON_COLUMN,
String.format("trim(%s)", DEFAULT_JSON_COLUMN)}) {
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.intSV','INT',0)", functionName, input),
DataType.INT, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.intMV','INT_ARRAY',0)", functionName,
input), DataType.INT, false});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.longSV','LONG',0)", functionName, input),
DataType.LONG, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.floatSV','FLOAT',0)", functionName,
input), DataType.FLOAT, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.doubleSV','DOUBLE',0)", functionName,
input), DataType.DOUBLE, true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.bigDecimalSV','BIG_DECIMAL',0)",
functionName, input), DataType.BIG_DECIMAL,
+ true});
+ testArguments.add(new Object[]{functionName,
+ String.format("%s(%s,'$.stringSV','STRING','null')", functionName,
input), DataType.STRING, true});
+ }
}
return testArguments.toArray(new Object[0][]);
@@ -304,38 +315,40 @@ public class JsonExtractScalarTransformFunctionTest
extends BaseTransformFunctio
@Test
public void testJsonPathTransformFunctionWithPredicate() {
- String jsonPathExpressionStr =
-
String.format("jsonExtractScalar(json,'[?($.stringSV==''%s'')]','STRING')",
_stringSVValues[0]);
- ExpressionContext expression =
RequestContextUtils.getExpression(jsonPathExpressionStr);
- TransformFunction transformFunction =
TransformFunctionFactory.get(expression, _dataSourceMap);
- Assert.assertTrue(transformFunction instanceof
JsonExtractScalarTransformFunction);
- Assert.assertEquals(transformFunction.getName(),
JsonExtractScalarTransformFunction.FUNCTION_NAME);
- // Note: transformToStringValuesSV() calls
IdentifierTransformFunction.transformToStringValuesSV() which in turns
- // call DataFetcher.readStringValues() which calls
DefaultJsonPathEvaluator.evaluateBlock() that parses String w/o
- // support for exact BigDecimal. Therefore, testing string parsing of
BigDecimal is disabled.
- String[] resultValues =
transformFunction.transformToStringValuesSV(_projectionBlock);
- for (int i = 0; i < NUM_ROWS; i++) {
- if (_stringSVValues[i].equals(_stringSVValues[0])) {
- try {
- final List<HashMap<String, Object>> resultMap =
JsonUtils.stringToObject(resultValues[i], List.class);
- Assert.assertEquals(_intSVValues[i], resultMap.get(0).get("intSV"));
- for (int j = 0; j < _intMVValues[i].length; j++) {
- Assert.assertEquals(_intMVValues[i][j], ((List)
resultMap.get(0).get("intMV")).get(j));
+ for (String functionName : JSON_EXTRACT_SCALAR_FUNCTIONS) {
+ String jsonPathExpressionStr =
+ String.format("%s(json,'[?($.stringSV==''%s'')]','STRING')",
functionName, _stringSVValues[0]);
+ ExpressionContext expression =
RequestContextUtils.getExpression(jsonPathExpressionStr);
+ TransformFunction transformFunction =
TransformFunctionFactory.get(expression, _dataSourceMap);
+ Assert.assertTrue(transformFunction instanceof
JsonExtractScalarTransformFunction);
+ Assert.assertEquals(transformFunction.getName(), functionName);
+ // Note: transformToStringValuesSV() calls
IdentifierTransformFunction.transformToStringValuesSV(), which calls
+ // DataFetcher.readStringValues() and
DefaultJsonPathEvaluator.evaluateBlock(). The evaluator parses String values
+ // without exact BigDecimal support, so testing string parsing of
BigDecimal is disabled.
+ String[] resultValues =
transformFunction.transformToStringValuesSV(_projectionBlock);
+ for (int i = 0; i < NUM_ROWS; i++) {
+ if (_stringSVValues[i].equals(_stringSVValues[0])) {
+ try {
+ final List<HashMap<String, Object>> resultMap =
JsonUtils.stringToObject(resultValues[i], List.class);
+ Assert.assertEquals(_intSVValues[i],
resultMap.get(0).get("intSV"));
+ for (int j = 0; j < _intMVValues[i].length; j++) {
+ Assert.assertEquals(_intMVValues[i][j], ((List)
resultMap.get(0).get("intMV")).get(j));
+ }
+ Assert.assertEquals(_longSVValues[i],
resultMap.get(0).get("longSV"));
+ // Notes: since we use currently a mapper that parses exact big
decimals, doubles may get parsed as
+ // big decimals. Confirm this is a backward compatible change?
+ Assert.assertEquals(
+ Float.compare(_floatSVValues[i], ((Number)
resultMap.get(0).get("floatSV")).floatValue()), 0);
+ Assert.assertEquals(_doubleSVValues[i],
resultMap.get(0).get("doubleSV"));
+ // Disabled:
+ // Assert.assertEquals(_bigDecimalSVValues[i], (BigDecimal)
resultMap.get(0).get("bigDecimalSV"));
+ Assert.assertEquals(_stringSVValues[i],
resultMap.get(0).get("stringSV"));
+ } catch (IOException e) {
+ throw new RuntimeException();
}
- Assert.assertEquals(_longSVValues[i],
resultMap.get(0).get("longSV"));
- // Notes: since we use currently a mapper that parses exact big
decimals, doubles may get parsed as
- // big decimals. Confirm this is a backward compatible change?
- Assert.assertEquals(Float.compare(_floatSVValues[i], ((Number)
resultMap.get(0).get("floatSV")).floatValue()),
- 0);
- Assert.assertEquals(_doubleSVValues[i],
resultMap.get(0).get("doubleSV"));
- // Disabled:
- // Assert.assertEquals(_bigDecimalSVValues[i], (BigDecimal)
resultMap.get(0).get("bigDecimalSV"));
- Assert.assertEquals(_stringSVValues[i],
resultMap.get(0).get("stringSV"));
- } catch (IOException e) {
- throw new RuntimeException();
+ } else {
+ Assert.assertEquals(resultValues[i], "[]");
}
- } else {
- Assert.assertEquals(resultValues[i], "[]");
}
}
}
@@ -527,6 +540,10 @@ public class JsonExtractScalarTransformFunctionTest
extends BaseTransformFunctio
new Object[]{String.format("jsonExtractScalar(%s,
\"$.store.book[0].author\", 'String')", JSON_COLUMN)},
new Object[]{String.format("jsonExtractScalar(%s,
'$.store.book[0].author', \"String\")", JSON_COLUMN)},
new Object[]{String.format("json_extract_scalar(%s,
\"$.store.book[0].author\", 'String','abc')", JSON_COLUMN)},
+ new Object[]{String.format("jsonExtractScalarFast(%s)", JSON_COLUMN)},
+ new Object[]{String.format("jsonExtractScalarFast(%s,
\"$.store.book[0].author\", 'String')", JSON_COLUMN)},
+ new Object[]{String.format("jsonExtractScalarFirstMatch(%s,
'$.store.book[0].author', \"String\")",
+ JSON_COLUMN)},
new Object[]{String.format("jsonExtractKey(%s, \"$.*\")",
JSON_COLUMN)},
new Object[]{String.format("json_extract_key(%s, \"$.*\")",
JSON_COLUMN)}};
//@formatter:on
@@ -660,14 +677,20 @@ public class JsonExtractScalarTransformFunctionTest
extends BaseTransformFunctio
// -- Single-value (SV) tests --
- /// Runs `SELECT jsonExtractScalar(json, '$.v', resultsType) FROM testTable`
against a single-row table
- /// containing the given JSON document, and asserts the result for the
(always-duplicated) two
- /// expected rows.
+ /// Runs each JSON scalar extraction variant against a single-row table
containing the given JSON document, and
+ /// asserts the result for the (always-duplicated) two expected rows.
private void assertJsonExtractScalar(String json, String resultsType, Object
expectedValue) {
+ for (String functionName : JSON_EXTRACT_SCALAR_FUNCTIONS) {
+ assertJsonExtractScalar(functionName, json, DataType.JSON, resultsType,
null, expectedValue);
+ }
+ }
+
+ private void assertJsonExtractScalar(String functionName, Object json,
DataType inputType, String resultsType,
+ @Nullable String defaultValueSql, Object expectedValue) {
Schema schema = new Schema.SchemaBuilder()
.setSchemaName("testTable")
.setEnableColumnBasedNullHandling(true)
- .addDimensionField("json", DataType.JSON)
+ .addDimensionField("json", inputType)
.build();
TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE)
.setTableName("testTable")
@@ -677,10 +700,42 @@ public class JsonExtractScalarTransformFunctionTest
extends BaseTransformFunctio
.withNullHandling(false)
.givenTable(schema, tableConfig)
.onFirstInstance(new Object[]{json})
- .whenQuery("SELECT jsonExtractScalar(json, '$.v', '" + resultsType +
"') FROM testTable")
+ .whenQuery("SELECT " + functionName + "(json, '$.v', '" + resultsType
+ "'"
+ + (defaultValueSql != null ? ", " + defaultValueSql : "") + ")
FROM testTable")
.thenResultIs(expectedRow, expectedRow);
}
+ @Test
+ public void testFastExtractionModes() {
+ String duplicateKeys = "{\"v\": 1, \"v\": 2}";
+ assertJsonExtractScalar(JsonExtractScalarTransformFunction.FUNCTION_NAME,
duplicateKeys, DataType.STRING, "INT",
+ null, 2);
+
assertJsonExtractScalar(JsonExtractScalarTransformFunction.FAST_FUNCTION_NAME,
duplicateKeys, DataType.STRING,
+ "INT", null, 2);
+
assertJsonExtractScalar(JsonExtractScalarTransformFunction.FIRST_MATCH_FUNCTION_NAME,
duplicateKeys,
+ DataType.STRING, "INT", null, 1);
+
+ String malformedTail = "{\"v\": 1, \"broken\": [}";
+ assertJsonExtractScalar(JsonExtractScalarTransformFunction.FUNCTION_NAME,
malformedTail, DataType.STRING, "INT",
+ "-1", -1);
+
assertJsonExtractScalar(JsonExtractScalarTransformFunction.FAST_FUNCTION_NAME,
malformedTail, DataType.STRING,
+ "INT", "-1", -1);
+
assertJsonExtractScalar(JsonExtractScalarTransformFunction.FIRST_MATCH_FUNCTION_NAME,
malformedTail,
+ DataType.STRING, "INT", "-1", 1);
+ }
+
+ @Test
+ public void testFastExtractionFromBytesWithBigDecimal() {
+ byte[] json = "{\"label\":\"crème
brûlée\",\"v\":12345678901234567890.123456789}"
+ .getBytes(StandardCharsets.UTF_8);
+ String expected = "12345678901234567890.123456789";
+ String hexEncodedJson = BytesUtils.toHexString(json);
+
assertJsonExtractScalar(JsonExtractScalarTransformFunction.FAST_FUNCTION_NAME,
hexEncodedJson, DataType.BYTES,
+ "BIG_DECIMAL", null, expected);
+
assertJsonExtractScalar(JsonExtractScalarTransformFunction.FIRST_MATCH_FUNCTION_NAME,
hexEncodedJson,
+ DataType.BYTES, "BIG_DECIMAL", null, expected);
+ }
+
@Test
public void testExtractBooleanAsNumeric() {
// JSON true / false coerces to 1 / 0 across all numeric result types —
matches PinotDataType's
@@ -807,10 +862,15 @@ public class JsonExtractScalarTransformFunctionTest
extends BaseTransformFunctio
"12345678901234567890.123456789");
}
+ @Test
+ public void testExtractBytes() {
+ assertJsonExtractScalar("{\"v\": \"AAGl/w==\"}", "BYTES", "0001a5ff");
+ }
+
// -- Multi-value (MV) tests --
- /// Asserts that `SELECT jsonExtractScalar(json, '$.v', resultsType)` over a
single-row table with the
- /// given JSON document produces the given primitive-array result.
+ /// Asserts that each JSON scalar extraction variant produces the given
primitive-array result for the supplied
+ /// single-row JSON document.
private void assertJsonExtractMv(String json, String resultsType, Object
expectedArray) {
Schema schema = new Schema.SchemaBuilder()
.setSchemaName("testTable")
@@ -821,12 +881,14 @@ public class JsonExtractScalarTransformFunctionTest
extends BaseTransformFunctio
.setTableName("testTable")
.build();
Object[] expectedRow = new Object[]{expectedArray};
- FluentQueryTest.withBaseDir(_baseDir)
- .withNullHandling(false)
- .givenTable(schema, tableConfig)
- .onFirstInstance(new Object[]{json})
- .whenQuery("SELECT jsonExtractScalar(json, '$.v', '" + resultsType +
"') FROM testTable")
- .thenResultIs(expectedRow, expectedRow);
+ for (String functionName : JSON_EXTRACT_SCALAR_FUNCTIONS) {
+ FluentQueryTest.withBaseDir(_baseDir)
+ .withNullHandling(false)
+ .givenTable(schema, tableConfig)
+ .onFirstInstance(new Object[]{json})
+ .whenQuery("SELECT " + functionName + "(json, '$.v', '" +
resultsType + "') FROM testTable")
+ .thenResultIs(expectedRow, expectedRow);
+ }
}
@Test
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonPathTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonPathTest.java
index 95a9cceafae..6d77b3c69a4 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonPathTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonPathTest.java
@@ -23,6 +23,8 @@ import com.fasterxml.jackson.databind.node.ArrayNode;
import com.jayway.jsonpath.spi.cache.Cache;
import com.jayway.jsonpath.spi.cache.CacheProvider;
import java.io.File;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
@@ -58,6 +60,7 @@ public class JsonPathTest extends
CustomDataQueryClusterIntegrationTest {
// path (which visits every doc) return the same result set but follow
visibly different code paths.
private static final int NUM_DISTINCT_K1 = 100;
private static final String MY_MAP_STR_FIELD_NAME = "myMapStr";
+ private static final String MY_MAP_BYTES_FIELD_NAME = "myMapBytes";
private static final String MY_MAP_STR_K1_FIELD_NAME = "myMapStr_k1";
private static final String MY_MAP_STR_K2_FIELD_NAME = "myMapStr_k2";
/// Derived columns that exercise the opt-in fast scalar functions through
the ingestion transform path.
@@ -89,6 +92,7 @@ public class JsonPathTest extends
CustomDataQueryClusterIntegrationTest {
.setSchemaName(getTableName())
.addSingleValueDimension("myMap", DataType.STRING)
.addSingleValueDimension(MY_MAP_STR_FIELD_NAME, DataType.STRING)
+ .addSingleValueDimension(MY_MAP_BYTES_FIELD_NAME, DataType.BYTES)
.addSingleValueDimension(MY_MAP_STR_K1_FIELD_NAME, DataType.STRING)
.addSingleValueDimension(MY_MAP_STR_K2_FIELD_NAME, DataType.STRING)
.addSingleValueDimension(MY_MAP_STR_K1_FAST_FIELD_NAME,
DataType.STRING)
@@ -125,6 +129,8 @@ public class JsonPathTest extends
CustomDataQueryClusterIntegrationTest {
List<org.apache.avro.Schema.Field> fields = List.of(
new org.apache.avro.Schema.Field(MY_MAP_STR_FIELD_NAME,
org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING),
null, null),
+ new org.apache.avro.Schema.Field(MY_MAP_BYTES_FIELD_NAME,
+ org.apache.avro.Schema.create(org.apache.avro.Schema.Type.BYTES),
null, null),
new org.apache.avro.Schema.Field(COMPLEX_MAP_STR_FIELD_NAME,
org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING),
null, null)
);
@@ -136,7 +142,9 @@ public class JsonPathTest extends
CustomDataQueryClusterIntegrationTest {
map.put("k1", "value-k1-" + (i % NUM_DISTINCT_K1));
map.put("k2", "value-k2-" + i);
GenericData.Record record = new GenericData.Record(avroSchema);
- record.put(MY_MAP_STR_FIELD_NAME, JsonUtils.objectToString(map));
+ String myMapJson = JsonUtils.objectToString(map);
+ record.put(MY_MAP_STR_FIELD_NAME, myMapJson);
+ record.put(MY_MAP_BYTES_FIELD_NAME,
ByteBuffer.wrap(myMapJson.getBytes(StandardCharsets.UTF_8)));
Map<String, Object> complexMap = new HashMap<>();
complexMap.put("k1", "value-k1-" + i);
@@ -410,6 +418,35 @@ public class JsonPathTest extends
CustomDataQueryClusterIntegrationTest {
}
}
+ /// Query-time coverage for the typed transforms backed by the same fast
extractor. Both modes must match the
+ /// existing `jsonExtractScalar` transform on this clean, duplicate-free
data set.
+ @Test(dataProvider = "useBothQueryEngines")
+ void testFastJsonExtractScalarTransforms(boolean useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ String query = "SELECT jsonExtractScalar(myMapStr, '$.k1', 'STRING'), "
+ + "jsonExtractScalarFast(myMapStr, '$.k1', 'STRING'), "
+ + "jsonExtractScalarFirstMatch(myMapStr, '$.k1', 'STRING') FROM " +
getTableName() + " LIMIT 1000";
+ JsonNode rows = postQuery(query).get("resultTable").get("rows");
+ assertTrue(rows.size() > 0, "expected non-empty result set");
+ for (JsonNode row : rows) {
+ String jayway = row.get(0).asText();
+ assertEquals(row.get(1).asText(), jayway, "jsonExtractScalarFast must
equal jsonExtractScalar");
+ assertEquals(row.get(2).asText(), jayway, "jsonExtractScalarFirstMatch
must equal Jayway on clean data");
+ }
+
+ query = "SELECT jsonExtractScalar(myMapBytes, '$.k1', 'STRING'), "
+ + "jsonExtractScalarFast(myMapBytes, '$.k1', 'STRING'), "
+ + "jsonExtractScalarFirstMatch(myMapBytes, '$.k1', 'STRING') FROM " +
getTableName() + " LIMIT 1000";
+ rows = postQuery(query).get("resultTable").get("rows");
+ assertTrue(rows.size() > 0, "expected non-empty BYTES result set");
+ for (JsonNode row : rows) {
+ String jayway = row.get(0).asText();
+ assertEquals(row.get(1).asText(), jayway, "BYTES fast extraction must
equal jsonExtractScalar");
+ assertEquals(row.get(2).asText(), jayway, "BYTES first-match extraction
must equal Jayway on clean data");
+ }
+ }
+
@Test
public void testJsonPathCache() {
Cache cache = CacheProvider.getCache();
diff --git
a/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml
b/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml
index 606f79f6d12..0fe80bc4ceb 100644
---
a/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml
+++
b/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml
@@ -981,6 +981,14 @@ jsonextractscalar:
scalar: null
transform:
"org.apache.pinot.core.operator.transform.function.JsonExtractScalarTransformFunction"
udf: null
+jsonextractscalarfast:
+ scalar: null
+ transform:
"org.apache.pinot.core.operator.transform.function.JsonExtractScalarTransformFunction.Fast"
+ udf: null
+jsonextractscalarfirstmatch:
+ scalar: null
+ transform:
"org.apache.pinot.core.operator.transform.function.JsonExtractScalarTransformFunction.FirstMatch"
+ udf: null
jsonformat:
scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonFormat}"
transform: null
@@ -1007,6 +1015,14 @@ jsonpathdouble:
\ 3:
org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathDouble]}"
transform: null
udf: null
+jsonpathdoublefast:
+ scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathDoubleFast}"
+ transform: null
+ udf: null
+jsonpathdoublefirstmatch:
+ scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathDoubleFirstMatch}"
+ transform: null
+ udf: null
jsonpathexists:
scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathExists}"
transform: null
@@ -1016,11 +1032,27 @@ jsonpathlong:
\ 3: org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathLong]}"
transform: null
udf: null
+jsonpathlongfast:
+ scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathLongFast}"
+ transform: null
+ udf: null
+jsonpathlongfirstmatch:
+ scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathLongFirstMatch}"
+ transform: null
+ udf: null
jsonpathstring:
scalar: "ArgumentCountBasedScalarFunction{[2:
org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathString,\
\ 3:
org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathString]}"
transform: null
udf: null
+jsonpathstringfast:
+ scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathStringFast}"
+ transform: null
+ udf: null
+jsonpathstringfirstmatch:
+ scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathStringFirstMatch}"
+ transform: null
+ udf: null
jsonstringtoarray:
scalar:
"ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonStringToArray}"
transform: null
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
index e3b87b197fa..6f22a19ae17 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
@@ -30,6 +30,7 @@ import java.util.concurrent.locks.ReentrantLock;
import java.util.stream.Collectors;
import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.pinot.query.planner.PlannerUtils;
import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
@@ -70,6 +71,73 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
assertNotNull(dispatchableSubPlan);
}
+ @Test
+ public void testFastJsonExtractScalarTypeInference() {
+ RelDataType rowType = _queryEnvironment.compile(
+ "SELECT JSON_EXTRACT_SCALAR_FAST(hexToBytes(col1), '$.foo',
'BIG_DECIMAL'), "
+ + "JSON_EXTRACT_SCALAR(col1, '$.foo', 'LONG', -1), "
+ + "JSON_EXTRACT_SCALAR_FAST(col1, '$.foo', 'BOOLEAN', FALSE), "
+ + "JSON_EXTRACT_SCALAR_FIRST_MATCH(col1, '$.foo', 'LONG', -1),
"
+ + "JSON_EXTRACT_SCALAR_FAST(col1, '$.foo', 'DOUBLE_ARRAY'), "
+ + "JSON_EXTRACT_SCALAR_FIRST_MATCH(col1, '$.foo',
'BIG_DECIMAL_ARRAY'), "
+ + "JSON_EXTRACT_SCALAR_FAST(col1, '$.foo', 'BOOLEAN_ARRAY'), "
+ + "JSON_EXTRACT_SCALAR_FIRST_MATCH(col1, '$.foo',
'TIMESTAMP_ARRAY'), "
+ + "JSON_EXTRACT_SCALAR_FAST(col1, '$.foo', 'JSON') FROM a")
+ .getRelRoot().validatedRowType;
+ assertEquals(rowType.getFieldList().get(0).getType().getSqlTypeName(),
SqlTypeName.DECIMAL);
+ assertEquals(rowType.getFieldList().get(1).getType().getSqlTypeName(),
SqlTypeName.BIGINT);
+ assertEquals(rowType.getFieldList().get(2).getType().getSqlTypeName(),
SqlTypeName.BOOLEAN);
+ assertEquals(rowType.getFieldList().get(3).getType().getSqlTypeName(),
SqlTypeName.BIGINT);
+ RelDataType arrayType = rowType.getFieldList().get(4).getType();
+ assertEquals(arrayType.getSqlTypeName(), SqlTypeName.ARRAY);
+ assertEquals(arrayType.getComponentType().getSqlTypeName(),
SqlTypeName.DOUBLE);
+ arrayType = rowType.getFieldList().get(5).getType();
+ assertEquals(arrayType.getSqlTypeName(), SqlTypeName.ARRAY);
+ assertEquals(arrayType.getComponentType().getSqlTypeName(),
SqlTypeName.DECIMAL);
+ arrayType = rowType.getFieldList().get(6).getType();
+ assertEquals(arrayType.getSqlTypeName(), SqlTypeName.ARRAY);
+ assertEquals(arrayType.getComponentType().getSqlTypeName(),
SqlTypeName.BOOLEAN);
+ arrayType = rowType.getFieldList().get(7).getType();
+ assertEquals(arrayType.getSqlTypeName(), SqlTypeName.ARRAY);
+ assertEquals(arrayType.getComponentType().getSqlTypeName(),
SqlTypeName.TIMESTAMP);
+ assertEquals(rowType.getFieldList().get(8).getType().getSqlTypeName(),
SqlTypeName.VARCHAR);
+
+ // A non-literal resultsType or defaultValue is rejected during
validation. jsonPath is deliberately not in this
+ // list -- see testJsonExtractScalarAcceptsFoldableJsonPath.
+ List<String> invalidQueries = List.of(
+ "SELECT JSON_EXTRACT_SCALAR_FAST(col1, '$.foo', 'LONG', col3) FROM a",
+ "SELECT JSON_EXTRACT_SCALAR_FIRST_MATCH(col1, '$.foo', col2, -1) FROM
a");
+ for (String invalidQuery : invalidQueries) {
+ Throwable invalidOperand = expectThrows(RuntimeException.class, () ->
_queryEnvironment.compile(invalidQuery));
+
assertTrue(Throwables.getStackTraceAsString(invalidOperand).contains("Cannot
apply 'JSONEXTRACTSCALAR"),
+ "Unexpected failure for " + invalidQuery + ": " +
Throwables.getStackTraceAsString(invalidOperand));
+ }
+ }
+
+ /// `jsonPath` must resolve to a literal, but the operand type checker
deliberately does not demand a literal
+ /// `SqlNode` in that position: operand checking runs before
`PinotEvaluateLiteralRule` folds constant
+ /// expressions, so an argument such as `CONCAT('$.', 'foo')` folds to a
literal and plans and executes
+ /// correctly. Requiring [org.apache.calcite.sql.type.OperandTypes#LITERAL]
there would reject these queries,
+ /// which plan and execute successfully on master. Regression guard for all
three JSON scalar transforms, which
+ /// share one operand checker;
`QueryRunnerTest#provideTestSqlWithExecutionException` covers the end-to-end
half,
+ /// asserting that a folded path is actually applied on the leaf stage.
+ ///
+ /// `resultsType` is checked separately in
[#testFastJsonExtractScalarTypeInference], since it must stay a literal
+ /// for return-type inference to see it.
+ @Test
+ public void testJsonExtractScalarAcceptsFoldableJsonPath() {
+ List<String> functions =
+ List.of("JSON_EXTRACT_SCALAR", "JSON_EXTRACT_SCALAR_FAST",
"JSON_EXTRACT_SCALAR_FIRST_MATCH");
+ for (String function : functions) {
+ for (String path : List.of("CONCAT('$.', 'foo')", "CAST('$.foo' AS
VARCHAR)", "UPPER('$.foo')")) {
+ String query = "SELECT " + function + "(col1, " + path + ", 'INT')
FROM a";
+ // The return type must still come from the literal resultsType rather
than falling back to VARCHAR.
+
assertEquals(_queryEnvironment.compile(query).getRelRoot().validatedRowType.getFieldList().get(0).getType()
+ .getSqlTypeName(), SqlTypeName.INTEGER, query);
+ }
+ }
+ }
+
@Test
public void testPolymorphicArithmeticScalarFunctionsPlanQuery() {
DispatchableSubPlan dispatchableSubPlan = _queryEnvironment.planQuery(
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java
index c82c255339b..254a9359567 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java
@@ -251,6 +251,13 @@ public class QueryEnvironmentTestBase {
new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'FLOAT_ARRAY')
FROM a"},
new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo',
'DOUBLE_ARRAY') FROM a"},
new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo',
'STRING_ARRAY') FROM a"},
+ new Object[]{"SELECT JSON_EXTRACT_SCALAR_FAST(col1, '$.foo',
'BIG_DECIMAL') FROM a"},
+ new Object[]{"SELECT JSON_EXTRACT_SCALAR_FAST(col1, '$.foo',
'DOUBLE_ARRAY') FROM a"},
+ new Object[]{"SELECT JSON_EXTRACT_SCALAR_FAST(col1, '$.foo',
'BOOLEAN_ARRAY') FROM a"},
+ new Object[]{"SELECT JSON_EXTRACT_SCALAR_FAST(col1, '$.foo', 'JSON')
FROM a"},
+ new Object[]{"SELECT JSON_EXTRACT_SCALAR_FIRST_MATCH(col1, '$.foo',
'LONG', '0') FROM a"},
+ new Object[]{"SELECT JSON_EXTRACT_SCALAR_FIRST_MATCH(col1, '$.foo',
'STRING_ARRAY') FROM a"},
+ new Object[]{"SELECT JSON_EXTRACT_SCALAR_FIRST_MATCH(col1, '$.foo',
'TIMESTAMP_ARRAY') FROM a"},
new Object[]{"SELECT ts_timestamp FROM a WHERE ts_timestamp BETWEEN
TIMESTAMP '2016-01-01 00:00:00' AND "
+ "TIMESTAMP '2016-01-01 10:00:00'"},
new Object[]{"SELECT ts_timestamp FROM a WHERE ts_timestamp >=
CAST(1454284798000 AS TIMESTAMP)"},
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTest.java
index 6059b8031aa..bf4a7095926 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTest.java
@@ -331,6 +331,23 @@ public class QueryRunnerTest extends QueryRunnerTestBase {
// - 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"});
+ // - a constant-foldable jsonPath must be folded by
PinotEvaluateLiteralRule and then applied like a literal.
+ // Reaching "Cannot resolve JSON path" (rather than ParserUtils'
"single-quoted literal values") is what
+ // proves the fold happened, so this pins the reason jsonPath does
not require a literal in
+ // TransformFunctionType#jsonExtractScalarOperandTypeChecker.
+ testCases.add(new Object[]{
+ "SELECT CAST(jsonExtractScalar(col1, CONCAT('pa', 'th'), 'INT') AS
INT) FROM a", "Cannot resolve JSON path"});
+ // - the flip side: a jsonPath that cannot fold to a literal is still
rejected, on the leaf stage rather than
+ // during validation. Covers all three variants, which share the
operand type checker.
+ for (String jsonExtractScalar : new String[]{
+ "jsonExtractScalar", "jsonExtractScalarFast",
"jsonExtractScalarFirstMatch"
+ }) {
+ testCases.add(new Object[]{
+ "SELECT " + jsonExtractScalar + "(col1, col2, 'INT') FROM a",
+ "Expect the 2nd and 3rd arguments of transform function: " +
jsonExtractScalar
+ + "(jsonFieldName, 'jsonPath', 'resultsType', ['defaultValue'])
to be single-quoted literal values"
+ });
+ }
// - 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",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
