[FLINK-40233][table] Correctly share the parsed JSON across JSON_VALUE/JSON_QUERY calls when the owning call is skipped because of NULL or not taken branch

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala

@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.table.planner.codegen.calls
+
+import org.apache.flink.table.planner.codegen.{CodeGeneratorContext, CodeGenUtils, GeneratedExpression}
+import org.apache.flink.table.planner.codegen.CodeGenUtils.qualifyMethod
+import org.apache.flink.table.runtime.functions.SqlJsonUtils
+
+/**
+ * Shares the parsed JSON input between the JSON function calls of a single generated expression,
+ * see [[JsonValueCallGen]] and [[JsonQueryCallGen]].
+ *
+ * The input is parsed at most once per record, by whichever call runs first, and the result is held
+ * in a member variable that the following calls on the same input read instead of parsing again.
+ */
+object JsonParseReuse {
+
+  /**
+   * Returns the expression holding the parsed input. The caller generates its own call with the
+   * original operands and reads the parsed input from the returned expression.
+   *
+   * The parse is lazy: the result term is a call to a member method that parses on its first
+   * invocation for the record. `generateCallWithStmtIfArgsNotNull` places a call under a guard that
+   * requires *all* of its arguments to be non-null, so no single call can be made the owner of the
+   * parse - in
+   * {{{
+   * SELECT JSON_VALUE(v, CAST(NULL AS STRING)), JSON_QUERY(v, '$.a')
+   * }}}
+   * the NULL path makes the first call short-circuit while the second one still needs the parse.
+   * Conversely, when no call runs, no parse happens at all.
+   */
+  def parseSharedInput(
+      ctx: CodeGeneratorContext,
+      operands: Seq[GeneratedExpression]): GeneratedExpression = {
+    val input = operands.head
+    val inputTerm = s"${input.resultTerm}.toString()"
+
+    ctx.getReusableInputUnboxingExprs(inputTerm, Int.MinValue) match {
+      case Some(expr) => expr
+      case None =>
+        val varName = CodeGenUtils.newName(ctx, "jsonParsed")
+        val lastInputName = CodeGenUtils.newName(ctx, "jsonParsedInput")
+        val methodName = CodeGenUtils.newName(ctx, "parseJson")
+        val typeName = classOf[SqlJsonUtils.JsonValueContext].getName
+        val inputType = CodeGenUtils.BINARY_STRING
+        ctx.addReusableMember(s"$typeName $varName;")
+        ctx.addReusableMember(s"$inputType $lastInputName;")
+
+        // keyed on the immutable input so it re-parses on change; no per-record reset needed
+        ctx.addReusableMember(
+          s"""
+             |private $typeName $methodName($inputType in) {
+             |  if (in != $lastInputName) {
+             |    $lastInputName = in;
+             |    $varName = ${qualifyMethod(BuiltInMethods.JSON_PARSE)}(in.toString());
+             |  }
+             |  return $varName;
+             |}
+             |""".stripMargin)
+
+        val parsed =
+          GeneratedExpression(s"$methodName(${input.resultTerm})", "false", "", null)
+        ctx.addReusableInputUnboxingExprs(inputTerm, Int.MinValue, parsed)
+
+        parsed
+    }
+  }
+}

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonQueryCallGen.scala

@@ -21,7 +21,6 @@ import org.apache.flink.table.api.{JsonQueryOnEmptyOrError, JsonQueryWrapper, Js
 import org.apache.flink.table.planner.codegen.{CodeGeneratorContext, CodeGenException, CodeGenUtils, GeneratedExpression}
 import org.apache.flink.table.planner.codegen.CodeGenUtils.{qualifyEnum, qualifyMethod, BINARY_STRING, GENERIC_ARRAY}
 import org.apache.flink.table.planner.codegen.GenerateUtils.generateCallWithStmtIfArgsNotNull
-import org.apache.flink.table.runtime.functions.SqlJsonUtils
 import org.apache.flink.table.runtime.functions.SqlJsonUtils.JsonQueryReturnType
 import org.apache.flink.table.types.logical.{ArrayType, LogicalType, LogicalTypeRoot}
 
@@ -41,14 +40,22 @@ import org.apache.calcite.sql.SqlJsonEmptyOrError
  * }}}
  * generates code similar to:
  * {{{
- * // member variable (declared once)
+ * // members (declared once)
  * SqlJsonUtils.JsonValueContext jsonParsed$0;
+ * BinaryStringData jsonParsedInput$1;
  *
- * // in processElement (parse emitted only by the first function)
- * jsonParsed$0 = SqlJsonUtils.jsonParse(field$0.toString());
- * Object rawResult$1 = SqlJsonUtils.jsonValue(jsonParsed$0, "$.type", ...);
- * // second call reuses jsonParsed$123 without re-parsing
- * Object rawResult$2 = SqlJsonUtils.jsonQuery(jsonParsed$0, "$.address", ...);
+ * // parses once per input value and reuses the result for the same input
+ * private SqlJsonUtils.JsonValueContext parseJson$2(BinaryStringData in) {
+ *   if (in != jsonParsedInput$1) {
+ *     jsonParsedInput$1 = in;
+ *     jsonParsed$0 = SqlJsonUtils.jsonParse(in.toString());
+ *   }
+ *   return jsonParsed$0;
+ * }
+ *
+ * // whichever call runs first parses, the other ones reuse the result
+ * Object rawResult$3 = SqlJsonUtils.jsonValue(parseJson$2(field$0), "$.type", ...);
+ * Object rawResult$4 = SqlJsonUtils.jsonQuery(parseJson$2(field$0), "$.address", ...);
  * }}}
  */
 class JsonQueryCallGen extends CallGenerator {
@@ -57,6 +64,8 @@ class JsonQueryCallGen extends CallGenerator {
       operands: Seq[GeneratedExpression],
       returnType: LogicalType): GeneratedExpression = {
 
+    val parsed = JsonParseReuse.parseSharedInput(ctx, operands)
+
     generateCallWithStmtIfArgsNotNull(ctx, returnType, operands, resultNullable = true) {
       argTerms =>
         {
@@ -68,26 +77,8 @@ class JsonQueryCallGen extends CallGenerator {
           } else {
             JsonQueryReturnType.STRING
           }
-          val inputTerm = s"${argTerms.head}.toString()"
-
-          val (varName, parseCode) =
-            ctx.getReusableInputUnboxingExprs(inputTerm, Int.MinValue) match {
-              case Some(expr) => (expr.resultTerm, "")
-              case None =>
-                val newVarName = CodeGenUtils.newName(ctx, "jsonParsed")
-                val typeName = classOf[SqlJsonUtils.JsonValueContext].getName
-                ctx.addReusableMember(s"$typeName $newVarName;")
-                ctx.addReusableInputUnboxingExprs(
-                  inputTerm,
-                  Int.MinValue,
-                  GeneratedExpression(newVarName, "false", "", null))
-                val assign =
-                  s"$newVarName = ${qualifyMethod(BuiltInMethods.JSON_PARSE)}($inputTerm);"
-                (newVarName, assign)
-            }
-
           val terms = Seq(
-            varName,
+            parsed.resultTerm,
             s"${argTerms(1)}.toString()",
             qualifyEnum(jsonQueryReturnType),
             qualifyEnum(wrapperBehavior),
@@ -97,7 +88,6 @@ class JsonQueryCallGen extends CallGenerator {
 
           val rawResultTerm = CodeGenUtils.newName(ctx, "rawResult")
           val call = s"""
-                        |$parseCode
                         |Object $rawResultTerm =
                         |    ${qualifyMethod(BuiltInMethods.JSON_QUERY_PARSED)}(${terms
                          .mkString(", ")});

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonValueCallGen.scala

@@ -21,7 +21,6 @@ import org.apache.flink.table.api.JsonValueOnEmptyOrError
 import org.apache.flink.table.planner.codegen.{CodeGeneratorContext, CodeGenException, CodeGenUtils, GeneratedExpression}
 import org.apache.flink.table.planner.codegen.CodeGenUtils.{qualifyEnum, qualifyMethod, BINARY_STRING}
 import org.apache.flink.table.planner.codegen.GenerateUtils.generateCallWithStmtIfArgsNotNull
-import org.apache.flink.table.runtime.functions.SqlJsonUtils
 import org.apache.flink.table.types.logical.{LogicalType, LogicalTypeRoot}
 
 import org.apache.calcite.sql.SqlJsonEmptyOrError
@@ -41,14 +40,22 @@ import org.apache.calcite.sql.SqlJsonEmptyOrError
  * }}}
  * generates code similar to:
  * {{{
- * // member variable (declared once)
+ * // members (declared once)
  * SqlJsonUtils.JsonValueContext jsonParsed$0;
+ * BinaryStringData jsonParsedInput$1;
  *
- * // in processElement (parse emitted only by the first function)
- * jsonParsed$0 = SqlJsonUtils.jsonParse(field$0.toString());
- * Object rawResult$1 = SqlJsonUtils.jsonValue(jsonParsed$0, "$.type", ...);
- * // second call reuses jsonParsed$0 without re-parsing
- * Object rawResult$2 = SqlJsonUtils.jsonValue(jsonParsed$0, "$.age", ...);
+ * // parses once per input value and reuses the result for the same input
+ * private SqlJsonUtils.JsonValueContext parseJson$2(BinaryStringData in) {
+ *   if (in != jsonParsedInput$1) {
+ *     jsonParsedInput$1 = in;
+ *     jsonParsed$0 = SqlJsonUtils.jsonParse(in.toString());
+ *   }
+ *   return jsonParsed$0;
+ * }
+ *
+ * // whichever call runs first parses, the other ones reuse the result
+ * Object rawResult$3 = SqlJsonUtils.jsonValue(parseJson$2(field$0), "$.type", ...);
+ * Object rawResult$4 = SqlJsonUtils.jsonValue(parseJson$2(field$0), "$.age", ...);
  * }}}
  */
 class JsonValueCallGen extends CallGenerator {
@@ -57,31 +64,16 @@ class JsonValueCallGen extends CallGenerator {
       operands: Seq[GeneratedExpression],
       returnType: LogicalType): GeneratedExpression = {
 
+    val parsed = JsonParseReuse.parseSharedInput(ctx, operands)
+
     generateCallWithStmtIfArgsNotNull(ctx, returnType, operands, resultNullable = true) {
       argTerms =>
         {
           val emptyBehavior = getBehavior(operands, SqlJsonEmptyOrError.EMPTY)
           val errorBehavior = getBehavior(operands, SqlJsonEmptyOrError.ERROR)
-          val inputTerm = s"${argTerms.head}.toString()"
-
-          val (varName, parseCode) =
-            ctx.getReusableInputUnboxingExprs(inputTerm, Int.MinValue) match {
-              case Some(expr) => (expr.resultTerm, "")
-              case None =>
-                val newVarName = CodeGenUtils.newName(ctx, "jsonParsed")
-                val typeName = classOf[SqlJsonUtils.JsonValueContext].getName
-                ctx.addReusableMember(s"$typeName $newVarName;")
-                ctx.addReusableInputUnboxingExprs(
-                  inputTerm,
-                  Int.MinValue,
-                  GeneratedExpression(newVarName, "false", "", null))
-                val assign =
-                  s"$newVarName = ${qualifyMethod(BuiltInMethods.JSON_PARSE)}($inputTerm);"
-                (newVarName, assign)
-            }
 
           val terms = Seq(
-            varName,
+            parsed.resultTerm,
             s"${argTerms(1)}.toString()",
             qualifyEnum(emptyBehavior._1),
             emptyBehavior._2,
@@ -91,7 +83,6 @@ class JsonValueCallGen extends CallGenerator {
 
           val rawResultTerm = CodeGenUtils.newName(ctx, "rawResult")
           val call = s"""
-                        |$parseCode
                         |Object $rawResultTerm =
                         |    ${qualifyMethod(BuiltInMethods.JSON_VALUE)}(${terms
                          .mkString(", ")});