From bd385c6d1553c189022326a812d4e028a10a5301 Mon Sep 17 00:00:00 2001 From: Sergey Nuyanzin Date: Sun, 26 Jul 2026 13:59:55 +0200 Subject: [PATCH 1/2] [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 --- .../codegen/calls/JsonParseReuse.scala | 82 +++++++++++ .../codegen/calls/JsonQueryCallGen.scala | 43 +++--- .../codegen/calls/JsonValueCallGen.scala | 42 +++--- .../planner/codegen/JsonParseReuseTest.java | 127 ++++++++++++++++-- 4 files changed, 233 insertions(+), 61 deletions(-) create mode 100644 flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala new file mode 100644 index 00000000000000..19c2d55752a991 --- /dev/null +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala @@ -0,0 +1,82 @@ +/* + * 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 methodName = CodeGenUtils.newName(ctx, "parseJson") + val typeName = classOf[SqlJsonUtils.JsonValueContext].getName + ctx.addReusableMember(s"$typeName $varName;") + + // null means not parsed yet - jsonParse never returns null for a non-null input + ctx.addReusableMember( + s""" + |private $typeName $methodName(Object in) { + | if ($varName == null) { + | $varName = ${qualifyMethod(BuiltInMethods.JSON_PARSE)}(in.toString()); + | } + | return $varName; + |} + |""".stripMargin) + + // reset once per record so a call outside a not-taken branch never reuses a stale parse + ctx.addReusablePerRecordStatement(s"$varName = null;") + + val parsed = + GeneratedExpression(s"$methodName(${input.resultTerm})", "false", "", null) + ctx.addReusableInputUnboxingExprs(inputTerm, Int.MinValue, parsed) + + parsed + } + } +} diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonQueryCallGen.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonQueryCallGen.scala index 2b386ac1fb74ef..f5ab6ebc37c7a7 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonQueryCallGen.scala +++ b/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,21 @@ import org.apache.calcite.sql.SqlJsonEmptyOrError * }}} * generates code similar to: * {{{ - * // member variable (declared once) + * // members (declared once) * SqlJsonUtils.JsonValueContext jsonParsed$0; * - * // 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", ...); + * private SqlJsonUtils.JsonValueContext parseJson$1(Object in) { + * if (jsonParsed$0 == null) { + * jsonParsed$0 = SqlJsonUtils.jsonParse(in.toString()); + * } + * return jsonParsed$0; + * } + * + * // in processElement, reset once per record + * jsonParsed$0 = null; + * // whichever call runs first parses, the other ones reuse the result + * Object rawResult$2 = SqlJsonUtils.jsonValue(parseJson$1(field$0), "$.type", ...); + * Object rawResult$3 = SqlJsonUtils.jsonQuery(parseJson$1(field$0), "$.address", ...); * }}} */ class JsonQueryCallGen extends CallGenerator { @@ -57,6 +63,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 +76,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 +87,6 @@ class JsonQueryCallGen extends CallGenerator { val rawResultTerm = CodeGenUtils.newName(ctx, "rawResult") val call = s""" - |$parseCode |Object $rawResultTerm = | ${qualifyMethod(BuiltInMethods.JSON_QUERY_PARSED)}(${terms .mkString(", ")}); diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonValueCallGen.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonValueCallGen.scala index f28b60e85e1523..fcc223b5ab571b 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonValueCallGen.scala +++ b/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,21 @@ import org.apache.calcite.sql.SqlJsonEmptyOrError * }}} * generates code similar to: * {{{ - * // member variable (declared once) + * // members (declared once) * SqlJsonUtils.JsonValueContext jsonParsed$0; * - * // 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", ...); + * private SqlJsonUtils.JsonValueContext parseJson$1(Object in) { + * if (jsonParsed$0 == null) { + * jsonParsed$0 = SqlJsonUtils.jsonParse(in.toString()); + * } + * return jsonParsed$0; + * } + * + * // in processElement, reset once per record + * jsonParsed$0 = null; + * // whichever call runs first parses, the other ones reuse the result + * Object rawResult$2 = SqlJsonUtils.jsonValue(parseJson$1(field$0), "$.type", ...); + * Object rawResult$3 = SqlJsonUtils.jsonValue(parseJson$1(field$0), "$.age", ...); * }}} */ class JsonValueCallGen extends CallGenerator { @@ -57,31 +63,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 +82,6 @@ class JsonValueCallGen extends CallGenerator { val rawResultTerm = CodeGenUtils.newName(ctx, "rawResult") val call = s""" - |$parseCode |Object $rawResultTerm = | ${qualifyMethod(BuiltInMethods.JSON_VALUE)}(${terms .mkString(", ")}); diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/codegen/JsonParseReuseTest.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/codegen/JsonParseReuseTest.java index a1591841b5b9e4..7c5c980ff67108 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/codegen/JsonParseReuseTest.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/codegen/JsonParseReuseTest.java @@ -23,8 +23,10 @@ import org.apache.flink.streaming.api.transformations.OneInputTransformation; import org.apache.flink.table.api.EnvironmentSettings; import org.apache.flink.table.api.Table; +import org.apache.flink.table.api.TableEnvironment; import org.apache.flink.table.api.TableResult; import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; +import org.apache.flink.table.api.config.TableConfigOptions; import org.apache.flink.table.planner.codegen.calls.BuiltInMethods; import org.apache.flink.table.runtime.operators.CodeGenOperatorFactory; import org.apache.flink.types.Row; @@ -103,9 +105,10 @@ void testTwoJsonValueCalls() { "SELECT JSON_VALUE(json_data, '$.type'), JSON_VALUE(json_data, '$.age') FROM json_src"; final List rows = collect(sql); assertThat(rows).containsExactlyInAnyOrder(Row.of("account", "42"), Row.of("admin", "30")); - assertThat(countJsonParse(extractGeneratedCode(sql))) + final String code = extractGeneratedCode(sql); + assertThat(countJsonParse(code)) .as("Two JSON_VALUE calls on the same input should parse once") - .isEqualTo(1); + .isOne(); } @Test @@ -118,9 +121,10 @@ void testTwoJsonQueryCalls() { .containsExactlyInAnyOrder( Row.of("{\"city\":\"Munich\"}", "[[\"user\",\"viewer\"]]"), Row.of("{\"city\":\"Berlin\"}", "[[\"admin\"]]")); - assertThat(countJsonParse(extractGeneratedCode(sql))) + final String code = extractGeneratedCode(sql); + assertThat(countJsonParse(code)) .as("Two JSON_QUERY calls on the same input should parse once") - .isEqualTo(1); + .isOne(); } @Test @@ -133,9 +137,10 @@ void testJsonValueAndJsonQueryMixed() { .containsExactlyInAnyOrder( Row.of("account", "{\"city\":\"Munich\"}"), Row.of("admin", "{\"city\":\"Berlin\"}")); - assertThat(countJsonParse(extractGeneratedCode(sql))) + final String code = extractGeneratedCode(sql); + assertThat(countJsonParse(code)) .as("JSON_VALUE + JSON_QUERY on the same input should parse once") - .isEqualTo(1); + .isOdd(); } @Test @@ -149,9 +154,81 @@ void testThreeJsonFunctionCalls() { .containsExactlyInAnyOrder( Row.of("account", "42", "{\"city\":\"Munich\"}"), Row.of("admin", "30", "{\"city\":\"Berlin\"}")); - assertThat(countJsonParse(extractGeneratedCode(sql))) + final String code = extractGeneratedCode(sql); + assertThat(countJsonParse(code)) .as("Three JSON function calls on the same input should parse once") - .isEqualTo(1); + .isOne(); + } + + @Test + void testReuseSurvivesCodeSplitting() { + // an aggressive split must still route the input local into each parseJson call site + tEnv.getConfig().set(TableConfigOptions.MAX_LENGTH_GENERATED_CODE, 1); + final String sql = + "SELECT JSON_VALUE(json_data, '$.type'), " + + "JSON_VALUE(json_data, '$.age'), " + + "JSON_QUERY(json_data, '$.address'), " + + "JSON_QUERY(json_data, '$.roles' WITH WRAPPER) FROM json_src"; + final List rows = collect(sql); + assertThat(rows) + .containsExactlyInAnyOrder( + Row.of("account", "42", "{\"city\":\"Munich\"}", "[[\"user\",\"viewer\"]]"), + Row.of("admin", "30", "{\"city\":\"Berlin\"}", "[[\"admin\"]]")); + final String code = extractGeneratedCode(sql); + assertThat(countJsonParse(code)) + .as("Even with code splitting the input is parsed once") + .isOne(); + } + + @Test + void testFirstCallWithNullArgumentStillParses() { + // The first JSON call owns the parse. Its arguments are null, so its own result is NULL, + // but the parse must still happen for the following calls on the same input. + final String sql = + "SELECT JSON_VALUE(json_data, CAST(NULL AS STRING)), " + + "JSON_QUERY(json_data, '$.address') FROM json_src"; + final List rows = collect(sql); + assertThat(rows) + .containsExactlyInAnyOrder( + Row.of(null, "{\"city\":\"Munich\"}"), + Row.of(null, "{\"city\":\"Berlin\"}")); + final String code = extractGeneratedCode(sql); + assertThat(countJsonParse(code)) + .as("JSON_VALUE + JSON_QUERY on the same input should parse once") + .isOne(); + } + + @Test + void testFirstCallInsideNotTakenBranchStillParses() { + // The first JSON call owns the parse but sits in a CASE branch that is never taken. + final String sql = + "SELECT CASE WHEN json_data IS NULL " + + "THEN JSON_VALUE(json_data, '$.type') ELSE 'fallback' END, " + + "JSON_QUERY(json_data, '$.address') FROM json_src"; + final List rows = collect(sql); + assertThat(rows) + .containsExactlyInAnyOrder( + Row.of("fallback", "{\"city\":\"Munich\"}"), + Row.of("fallback", "{\"city\":\"Berlin\"}")); + } + + @Test + void testCallInsideSurvivingBranchSharesWithCallOutside() { + // Branch guarded by an unrelated column survives the optimizer; the per-record reset lets + // the '{}' row (branch not taken) still parse correctly while sharing a single parse. + final String sql = + "SELECT CASE WHEN CHARACTER_LENGTH(other_json) > 3 " + + "THEN JSON_VALUE(json_data, '$.type') ELSE 'fb' END, " + + "JSON_QUERY(json_data, '$.address') FROM json_src"; + final List rows = collect(sql); + assertThat(rows) + .containsExactlyInAnyOrder( + Row.of("fb", "{\"city\":\"Munich\"}"), + Row.of("admin", "{\"city\":\"Berlin\"}")); + final String code = extractGeneratedCode(sql); + assertThat(countJsonParse(code)) + .as("Calls on the same input share one parse even across a CASE boundary") + .isOne(); } @Test @@ -165,4 +242,38 @@ void testDifferentJsonInputs() { .as("JSON_VALUE calls on different inputs should parse separately") .isEqualTo(2); } + + @Test + void testReuseInOverWindowQueryIsResetPerRow() { + // JSON scalars run in a Calc ahead of the OverAggregate; each row must get its own parse + final TableEnvironment bEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + bEnv.createTemporaryView( + "over_src", + bEnv.fromValues(Row.of(1, JSON_ROW1), Row.of(2, JSON_ROW2)).as("id", "j")); + final String sql = + "SELECT id, " + + "MAX(JSON_VALUE(j, '$.type')) OVER (ORDER BY id ROWS BETWEEN CURRENT ROW AND CURRENT ROW), " + + "MAX(JSON_QUERY(j, '$.address')) OVER (ORDER BY id ROWS BETWEEN CURRENT ROW AND CURRENT ROW) " + + "FROM over_src"; + final List rows = new ArrayList<>(); + bEnv.executeSql(sql).collect().forEachRemaining(rows::add); + assertThat(rows) + .containsExactlyInAnyOrder( + Row.of(1, "account", "{\"city\":\"Munich\"}"), + Row.of(2, "admin", "{\"city\":\"Berlin\"}")); + } + + @Test + void testFilterAndProjectionShareParse() { + // JSON_VALUE in WHERE and JSON_QUERY in SELECT on the same input share one parse + final String sql = + "SELECT JSON_QUERY(json_data, '$.address') FROM json_src " + + "WHERE JSON_VALUE(json_data, '$.type') = 'admin'"; + final List rows = collect(sql); + assertThat(rows).containsExactly(Row.of("{\"city\":\"Berlin\"}")); + final String code = extractGeneratedCode(sql); + assertThat(countJsonParse(code)) + .as("Filter and projection on the same input should parse once") + .isOne(); + } } From d19a41c861fa77b9d5741904d88ad3fc7a2dcea4 Mon Sep 17 00:00:00 2001 From: Sergey Nuyanzin Date: Thu, 30 Jul 2026 14:12:00 +0200 Subject: [PATCH 2/2] Address feedback --- .../codegen/calls/JsonParseReuse.scala | 13 +- .../codegen/calls/JsonQueryCallGen.scala | 13 +- .../codegen/calls/JsonValueCallGen.scala | 13 +- .../planner/codegen/JsonParseReuseTest.java | 116 ++++++++++++++++-- 4 files changed, 126 insertions(+), 29 deletions(-) diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala index 19c2d55752a991..7a340dee0492c7 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonParseReuse.scala @@ -54,24 +54,25 @@ object JsonParseReuse { 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;") - // null means not parsed yet - jsonParse never returns null for a non-null input + // keyed on the immutable input so it re-parses on change; no per-record reset needed ctx.addReusableMember( s""" - |private $typeName $methodName(Object in) { - | if ($varName == null) { + |private $typeName $methodName($inputType in) { + | if (in != $lastInputName) { + | $lastInputName = in; | $varName = ${qualifyMethod(BuiltInMethods.JSON_PARSE)}(in.toString()); | } | return $varName; |} |""".stripMargin) - // reset once per record so a call outside a not-taken branch never reuses a stale parse - ctx.addReusablePerRecordStatement(s"$varName = null;") - val parsed = GeneratedExpression(s"$methodName(${input.resultTerm})", "false", "", null) ctx.addReusableInputUnboxingExprs(inputTerm, Int.MinValue, parsed) diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonQueryCallGen.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonQueryCallGen.scala index f5ab6ebc37c7a7..5eb269ff4f8d0d 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonQueryCallGen.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonQueryCallGen.scala @@ -42,19 +42,20 @@ import org.apache.calcite.sql.SqlJsonEmptyOrError * {{{ * // members (declared once) * SqlJsonUtils.JsonValueContext jsonParsed$0; + * BinaryStringData jsonParsedInput$1; * - * private SqlJsonUtils.JsonValueContext parseJson$1(Object in) { - * if (jsonParsed$0 == null) { + * // 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; * } * - * // in processElement, reset once per record - * jsonParsed$0 = null; * // whichever call runs first parses, the other ones reuse the result - * Object rawResult$2 = SqlJsonUtils.jsonValue(parseJson$1(field$0), "$.type", ...); - * Object rawResult$3 = SqlJsonUtils.jsonQuery(parseJson$1(field$0), "$.address", ...); + * Object rawResult$3 = SqlJsonUtils.jsonValue(parseJson$2(field$0), "$.type", ...); + * Object rawResult$4 = SqlJsonUtils.jsonQuery(parseJson$2(field$0), "$.address", ...); * }}} */ class JsonQueryCallGen extends CallGenerator { diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonValueCallGen.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonValueCallGen.scala index fcc223b5ab571b..321a6d072c362a 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonValueCallGen.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/JsonValueCallGen.scala @@ -42,19 +42,20 @@ import org.apache.calcite.sql.SqlJsonEmptyOrError * {{{ * // members (declared once) * SqlJsonUtils.JsonValueContext jsonParsed$0; + * BinaryStringData jsonParsedInput$1; * - * private SqlJsonUtils.JsonValueContext parseJson$1(Object in) { - * if (jsonParsed$0 == null) { + * // 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; * } * - * // in processElement, reset once per record - * jsonParsed$0 = null; * // whichever call runs first parses, the other ones reuse the result - * Object rawResult$2 = SqlJsonUtils.jsonValue(parseJson$1(field$0), "$.type", ...); - * Object rawResult$3 = SqlJsonUtils.jsonValue(parseJson$1(field$0), "$.age", ...); + * Object rawResult$3 = SqlJsonUtils.jsonValue(parseJson$2(field$0), "$.type", ...); + * Object rawResult$4 = SqlJsonUtils.jsonValue(parseJson$2(field$0), "$.age", ...); * }}} */ class JsonValueCallGen extends CallGenerator { diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/codegen/JsonParseReuseTest.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/codegen/JsonParseReuseTest.java index 7c5c980ff67108..61acdd86eea20a 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/codegen/JsonParseReuseTest.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/codegen/JsonParseReuseTest.java @@ -26,18 +26,25 @@ import org.apache.flink.table.api.TableEnvironment; import org.apache.flink.table.api.TableResult; import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; +import org.apache.flink.table.api.config.ExecutionConfigOptions; +import org.apache.flink.table.api.config.OptimizerConfigOptions; import org.apache.flink.table.api.config.TableConfigOptions; +import org.apache.flink.table.codesplit.JavaCodeSplitter; import org.apache.flink.table.planner.codegen.calls.BuiltInMethods; +import org.apache.flink.table.planner.factories.TestValuesTableFactory; import org.apache.flink.table.runtime.operators.CodeGenOperatorFactory; import org.apache.flink.types.Row; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.time.Instant; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.regex.Matcher; import java.util.regex.Pattern; +import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; @@ -82,10 +89,10 @@ private static int countJsonParse(final String code) { return count; } - private String extractGeneratedCode(final String sql) { + private List generatedClassCodes(final String sql) { final Table table = tEnv.sqlQuery(sql); final Transformation root = tEnv.toChangelogStream(table).getTransformation(); - final StringBuilder allCode = new StringBuilder(); + final List codes = new ArrayList<>(); for (final Transformation t : root.getTransitivePredecessors()) { if (t instanceof OneInputTransformation && ((OneInputTransformation) t).getOperatorFactory() @@ -93,10 +100,14 @@ private String extractGeneratedCode(final String sql) { final CodeGenOperatorFactory factory = (CodeGenOperatorFactory) ((OneInputTransformation) t).getOperatorFactory(); - allCode.append(factory.getGeneratedClass().getCode()); + codes.add(factory.getGeneratedClass().getCode()); } } - return allCode.toString(); + return codes; + } + + private String extractGeneratedCode(final String sql) { + return String.join("", generatedClassCodes(sql)); } @Test @@ -140,7 +151,7 @@ void testJsonValueAndJsonQueryMixed() { final String code = extractGeneratedCode(sql); assertThat(countJsonParse(code)) .as("JSON_VALUE + JSON_QUERY on the same input should parse once") - .isOdd(); + .isOne(); } @Test @@ -162,28 +173,58 @@ void testThreeJsonFunctionCalls() { @Test void testReuseSurvivesCodeSplitting() { - // an aggressive split must still route the input local into each parseJson call site + // an aggressive split must still route the input into each parseJson call site tEnv.getConfig().set(TableConfigOptions.MAX_LENGTH_GENERATED_CODE, 1); final String sql = "SELECT JSON_VALUE(json_data, '$.type'), " + "JSON_VALUE(json_data, '$.age'), " + "JSON_QUERY(json_data, '$.address'), " + "JSON_QUERY(json_data, '$.roles' WITH WRAPPER) FROM json_src"; + // GeneratedClass compiles splitCode, so correct results already prove reuse survives it final List rows = collect(sql); assertThat(rows) .containsExactlyInAnyOrder( Row.of("account", "42", "{\"city\":\"Munich\"}", "[[\"user\",\"viewer\"]]"), Row.of("admin", "30", "{\"city\":\"Berlin\"}", "[[\"admin\"]]")); + + // assert on the split code, not getCode() (unsplit), since only splitCode is compiled + final int maxLength = tEnv.getConfig().get(TableConfigOptions.MAX_LENGTH_GENERATED_CODE); + final int maxMembers = tEnv.getConfig().get(TableConfigOptions.MAX_MEMBERS_GENERATED_CODE); + final List splitCodes = + generatedClassCodes(sql).stream() + .map(code -> JavaCodeSplitter.split(code, maxLength, maxMembers)) + .collect(Collectors.toList()); + assertThat(generatedClassCodes(sql)) + .as("the aggressive limit must actually split some generated class") + .anySatisfy( + code -> + assertThat(JavaCodeSplitter.split(code, maxLength, maxMembers)) + .isNotEqualTo(code)); + assertThat(splitCodes.stream().mapToInt(JsonParseReuseTest::countJsonParse).sum()) + .as("Even in the split code the input is parsed once") + .isOne(); + } + + @Test + void testComputedColumnInputSharesParse() { + // calls on the same computed input (TRIM) must still share a parse + final String sql = + "SELECT JSON_VALUE(TRIM(json_data), '$.type'), " + + "JSON_QUERY(TRIM(json_data), '$.address') FROM json_src"; + final List rows = collect(sql); + assertThat(rows) + .containsExactlyInAnyOrder( + Row.of("account", "{\"city\":\"Munich\"}"), + Row.of("admin", "{\"city\":\"Berlin\"}")); final String code = extractGeneratedCode(sql); assertThat(countJsonParse(code)) - .as("Even with code splitting the input is parsed once") + .as("Calls on the same computed input should parse once") .isOne(); } @Test void testFirstCallWithNullArgumentStillParses() { - // The first JSON call owns the parse. Its arguments are null, so its own result is NULL, - // but the parse must still happen for the following calls on the same input. + // first call's args are null so its result is NULL, but the parse must still happen final String sql = "SELECT JSON_VALUE(json_data, CAST(NULL AS STRING)), " + "JSON_QUERY(json_data, '$.address') FROM json_src"; @@ -214,8 +255,7 @@ void testFirstCallInsideNotTakenBranchStillParses() { @Test void testCallInsideSurvivingBranchSharesWithCallOutside() { - // Branch guarded by an unrelated column survives the optimizer; the per-record reset lets - // the '{}' row (branch not taken) still parse correctly while sharing a single parse. + // branch guarded by an unrelated column survives the optimizer, yet both rows share a parse final String sql = "SELECT CASE WHEN CHARACTER_LENGTH(other_json) > 3 " + "THEN JSON_VALUE(json_data, '$.type') ELSE 'fb' END, " @@ -276,4 +316,58 @@ void testFilterAndProjectionShareParse() { .as("Filter and projection on the same input should parse once") .isOne(); } + + @Test + void testReuseIsResetPerRowInMatchRecognize() { + // A matches the first three rows; the per-row parse must give SUM 1+2+3, not 1+1+1 + final List data = + Arrays.asList( + Row.of(1000, "{\"n\":1}", Instant.ofEpochMilli(1000L)), + Row.of(2000, "{\"n\":2}", Instant.ofEpochMilli(2000L)), + Row.of(3000, "{\"n\":3}", Instant.ofEpochMilli(3000L)), + Row.of(9000, "{\"n\":9}", Instant.ofEpochMilli(9000L))); + final String dataId = TestValuesTableFactory.registerData(data); + tEnv.executeSql( + "CREATE TABLE events (" + + " f0 INT," + + " f1 STRING," + + " ts TIMESTAMP_LTZ(3)," + + " WATERMARK FOR ts AS ts" + + ") WITH (" + + " 'connector' = 'values'," + + " 'data-id' = '" + + dataId + + "'," + + " 'bounded' = 'true')"); + final String sql = + "SELECT total FROM events MATCH_RECOGNIZE (" + + " ORDER BY ts" + + " MEASURES SUM(CAST(JSON_VALUE(A.f1, '$.n') AS INT)) AS total" + + " AFTER MATCH SKIP PAST LAST ROW" + + " PATTERN (A+ B)" + + " DEFINE A AS A.f0 < 9000, B AS B.f0 >= 9000)"; + final List rows = collect(sql); + assertThat(rows).containsExactly(Row.of(6)); + } + + @Test + void testReuseIsResetPerRowInBatchFusion() { + // a projection Calc is only fused on top of a HashJoin, so force one (disable the rest) + final TableEnvironment bEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + bEnv.getConfig() + .set(ExecutionConfigOptions.TABLE_EXEC_OPERATOR_FUSION_CODEGEN_ENABLED, true) + .set( + ExecutionConfigOptions.TABLE_EXEC_DISABLED_OPERATORS, + "NestedLoopJoin,SortMergeJoin") + .set(OptimizerConfigOptions.TABLE_OPTIMIZER_BROADCAST_JOIN_THRESHOLD, -1L); + bEnv.createTemporaryView( + "src", bEnv.fromValues(Row.of(1, JSON_ROW1), Row.of(2, JSON_ROW2)).as("id", "j")); + bEnv.createTemporaryView("dim", bEnv.fromValues(Row.of(1), Row.of(2)).as("id")); + final String sql = + "SELECT JSON_VALUE(src.j, '$.type'), JSON_VALUE(src.j, '$.age') " + + "FROM src JOIN dim ON src.id = dim.id"; + final List rows = new ArrayList<>(); + bEnv.executeSql(sql).collect().forEachRemaining(rows::add); + assertThat(rows).containsExactlyInAnyOrder(Row.of("account", "42"), Row.of("admin", "30")); + } }