Parse a JSON column once per record for repeated jsonPath* transforms

pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java

@@ -19,13 +19,15 @@
 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;
 import org.apache.commons.lang3.StringUtils;
 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, ',') + ')';

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
+   * <b>unchanged</b> for anything else (a non-container or invalid string, a scalar, an already-parsed value, null).
+   * <p>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 <b>without throwing</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)");

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<String, Object> alreadyParsed = Map.of("a", "b");
+    assertSame(JsonFunctions.parseDocOrSelf(alreadyParsed), alreadyParsed);
+  }
+
   @Test
   public void testJsonFunctionExtractingArray()
       throws JsonProcessingException {

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<String> _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<Object, Object> _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<String> primaryKeyColumns) {
     int numPrimaryKeyColumns = primaryKeyColumns.size();
@@ -263,6 +287,9 @@ public void clear() {
     _nullValueFields.clear();
     _incomplete = false;
     _sanitized = false;
+    if (_scratch != null) {
+      _scratch.clear();
+    }
   }
 
   @Override