diff --git a/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java b/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java index 5db855b7327a..f34fa6779a07 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java @@ -19,6 +19,7 @@ package org.apache.pinot.common.evaluator; import com.google.common.base.Preconditions; +import java.lang.reflect.Method; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -26,6 +27,7 @@ import org.apache.pinot.common.function.FunctionInfo; import org.apache.pinot.common.function.FunctionInvoker; import org.apache.pinot.common.function.FunctionRegistry; +import org.apache.pinot.common.function.scalar.JsonFunctions; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.FunctionContext; import org.apache.pinot.common.request.context.RequestContextUtils; @@ -255,11 +257,18 @@ private static class FunctionExecutionNode implements ExecutableNode { final ExecutableNode[] _argumentNodes; final Object[] _arguments; + // True when this is a JsonFunctions.jsonPath* function whose first argument is parsed as a JSON document. For + // these, a JSON-string first argument is parsed once per record and reused via the record's scratch cache, so + // multiple extractions from the same column (e.g. several JSONPATHSTRING(message, ...) transforms) do not + // re-parse the document once per field. + final boolean _parseFirstArgAsJson; + FunctionExecutionNode(FunctionInfo functionInfo, ExecutableNode[] argumentNodes) { _functionInvoker = new FunctionInvoker(functionInfo); _functionInfo = functionInfo; _argumentNodes = argumentNodes; _arguments = new Object[_argumentNodes.length]; + _parseFirstArgAsJson = isJsonDocumentFunction(_functionInvoker.getMethod()); } @Override @@ -269,6 +278,17 @@ public Object execute(GenericRow row) { for (int i = 0; i < numArguments; i++) { _arguments[i] = _argumentNodes[i].execute(row); } + if (_parseFirstArgAsJson && _arguments[0] instanceof String) { + // Parse the JSON document once per record and stash the parsed form on the record, so the other jsonPath* + // extractions on the same column reuse it instead of re-parsing. Non-JSON / scalar input is returned + // unchanged by parseDocOrSelf, so behavior is preserved. + Object parsed = row.getScratchValue(_arguments[0]); + if (parsed == null) { + parsed = JsonFunctions.parseDocOrSelf(_arguments[0]); + row.putScratchValue(_arguments[0], parsed); + } + _arguments[0] = parsed; + } if (!_functionInfo.hasNullableParameters()) { // Preserve null values during ingestion transformation if function is an inbuilt // scalar function that cannot handle nulls, and invoked with null parameter(s). @@ -314,6 +334,14 @@ public Object execute(Object[] values) { } } + // The JsonFunctions.jsonPath* family (jsonPath, jsonPathString, jsonPathLong, jsonPathDouble, jsonPathArray, + // jsonPathArrayDefaultEmpty, jsonPathExists) all take the JSON document as their first argument. Recognizing them + // lets the evaluator parse that document once per record (see execute(GenericRow)). + private static boolean isJsonDocumentFunction(Method method) { + return method.getParameterCount() >= 1 && method.getDeclaringClass() == JsonFunctions.class + && method.getName().startsWith("jsonPath"); + } + @Override public String toString() { return _functionInvoker.getMethod().getName() + '(' + StringUtils.join(_argumentNodes, ',') + ')'; diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/JsonFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/JsonFunctions.java index 2a5b35166304..12a9f5061174 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/JsonFunctions.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/JsonFunctions.java @@ -367,6 +367,25 @@ public static Object jsonStringToListOrMap(String jsonString) { return null; } + /** + * Parse a JSON object/array string into a reusable {@code Map}/{@code List} once, or return the input + * unchanged for anything else (a non-container or invalid string, a scalar, an already-parsed value, null). + *

Used by the ingestion evaluator to parse a JSON column once per record and feed the parsed form to the many + * {@code jsonPath*} extractions on that record, instead of re-parsing the document once per extracted field. + * Because non-container input is returned unchanged, the downstream {@code jsonPath*} call behaves exactly as if it + * received the raw value (it does its own scalar parse / non-JSON fast-path / exception), so this is + * behavior-preserving. The returned container is a live model that {@code jsonPath} navigates without re-parsing; + * callers must treat it as read-only (same contract as {@code jsonExtractObject}). + */ + @Nullable + public static Object parseDocOrSelf(@Nullable Object object) { + if (object instanceof String) { + Object parsed = jsonStringToListOrMap((String) object); + return parsed != null ? parsed : object; + } + return object; + } + /** * Parse a JSON document into a reusable {@code Map}/{@code List}, or pass through an already-parsed container, * returning {@code null} for null, scalar, or non-JSON input without throwing. diff --git a/pinot-common/src/test/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluatorTest.java b/pinot-common/src/test/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluatorTest.java index 116517a6c790..572373ca7695 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluatorTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluatorTest.java @@ -19,6 +19,7 @@ package org.apache.pinot.common.evaluator; import java.util.Collections; +import java.util.Map; import org.apache.pinot.common.function.FunctionUtils; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.utils.PinotDataType; @@ -28,6 +29,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; @@ -43,6 +45,31 @@ public void booleanLiteralTest() { checkBooleanLiteralExpression("0", 0); } + @Test + public void testJsonPathParsesDocumentOncePerRecord() { + // A JSON-string column read by several jsonPath* transforms is parsed once and the parsed document is stashed on + // the record (GenericRow scratch), so subsequent extractions reuse it instead of re-parsing. + String message = "{\"a\": \"x\", \"b\": {\"c\": 2}}"; + GenericRow row = new GenericRow(); + row.putValue("message", message); + assertNull(row.getScratchValue(message)); + + assertEquals(new InbuiltFunctionEvaluator("jsonPathString(message, '$.a')").evaluate(row), "x"); + Object cached = row.getScratchValue(message); + assertTrue(cached instanceof Map); + + // subsequent extractions reuse the same parsed document and remain correct + assertEquals(new InbuiltFunctionEvaluator("jsonPathString(message, '$.b.c')").evaluate(row), "2"); + assertEquals(new InbuiltFunctionEvaluator("jsonPathLong(message, '$.b.c')").evaluate(row), 2L); + assertSame(row.getScratchValue(message), cached); + + // a non-JSON string column is passed through unchanged (no parse), and cached as itself + String plain = "INFO plain text log line"; + row.putValue("msg", plain); + assertEquals(new InbuiltFunctionEvaluator("jsonPathString(msg, '$.level', 'def')").evaluate(row), "def"); + assertSame(row.getScratchValue(plain), plain); + } + @Test public void testOrWithNulls() { InbuiltFunctionEvaluator evaluator = new InbuiltFunctionEvaluator("or(null, false, true)"); diff --git a/pinot-common/src/test/java/org/apache/pinot/common/function/JsonFunctionsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/function/JsonFunctionsTest.java index 96d83b951190..bf5df784d03d 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/function/JsonFunctionsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/function/JsonFunctionsTest.java @@ -243,6 +243,31 @@ public void testJsonExtractObject() assertEquals(JsonFunctions.jsonPathString(parsed, "$.id", "def"), "abc"); } + /** + * {@code parseDocOrSelf} parses a JSON object/array string into a reusable Map/List once, and returns the input + * unchanged for everything else (non-container/invalid string, scalar, already-parsed value, null) so it is + * behavior-preserving when fed to a downstream {@code jsonPath*} call. + */ + @Test + public void testParseDocOrSelf() { + // Object/array strings -> reusable parsed container, navigable by jsonPath without re-parsing. + Object obj = JsonFunctions.parseDocOrSelf("{\"level\": \"info\", \"n\": 1}"); + assertTrue(obj instanceof Map); + assertEquals(JsonFunctions.jsonPathString(obj, "$.level", "def"), "info"); + assertTrue(JsonFunctions.parseDocOrSelf("[1, 2, 3]") instanceof List); + + // Anything that is not a JSON container is returned unchanged (same instance), so behavior is preserved. + String plain = "INFO plain text"; + assertSame(JsonFunctions.parseDocOrSelf(plain), plain); + String scalar = "12345"; + assertSame(JsonFunctions.parseDocOrSelf(scalar), scalar); + String invalid = "{not valid"; + assertSame(JsonFunctions.parseDocOrSelf(invalid), invalid); + assertNull(JsonFunctions.parseDocOrSelf(null)); + Map alreadyParsed = Map.of("a", "b"); + assertSame(JsonFunctions.parseDocOrSelf(alreadyParsed), alreadyParsed); + } + @Test public void testJsonFunctionExtractingArray() throws JsonProcessingException { diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/GenericRow.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/GenericRow.java index 0b1862ca4dc1..e255f5b71f08 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/GenericRow.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/GenericRow.java @@ -24,6 +24,7 @@ import java.util.Collection; import java.util.HashMap; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -68,6 +69,12 @@ public class GenericRow implements Serializable { private final Set _nullValueFields = new HashSet<>(); private boolean _incomplete; private boolean _sanitized; + // Transient, per-record scratch used to memoize an expensive per-record derivation that several transform + // functions share within one record - currently parsing a JSON-string column once that is then read by multiple + // jsonPath* extractions. It is NOT record data: not serialized, not indexed, not copied by init(), and reset by + // clear(). Keyed by reference identity (the same column resolves to the same value instance within a record), so + // it is created lazily only when used and freed with the record. + private transient IdentityHashMap _scratch; /** * Initializes the generic row from the given generic row (shallow copy). The row should be new created or cleared @@ -107,6 +114,23 @@ public Object getValue(String fieldName) { return _fieldToValueMap.get(fieldName); } + /// Returns a previously memoized derivation for the given key from the transient per-record scratch, or `null` if + /// absent. See [#_scratch]. Keyed by reference identity. + @Nullable + public Object getScratchValue(Object key) { + return _scratch != null ? _scratch.get(key) : null; + } + + /// Memoizes a derivation in the transient per-record scratch under the given key (reference identity). The scratch + /// is created lazily on first use and reset by [#clear()]; it is never serialized, indexed, or copied. See + /// [#_scratch]. + public void putScratchValue(Object key, Object value) { + if (_scratch == null) { + _scratch = new IdentityHashMap<>(); + } + _scratch.put(key, value); + } + /// Constructs a [PrimaryKey] from the given list of primary key columns. public PrimaryKey getPrimaryKey(List primaryKeyColumns) { int numPrimaryKeyColumns = primaryKeyColumns.size(); @@ -263,6 +287,9 @@ public void clear() { _nullValueFields.clear(); _incomplete = false; _sanitized = false; + if (_scratch != null) { + _scratch.clear(); + } } @Override