diff --git a/integration_tests/src/main/python/get_json_test.py b/integration_tests/src/main/python/get_json_test.py index 2b4218692a6..8fc37df0955 100644 --- a/integration_tests/src/main/python/get_json_test.py +++ b/integration_tests/src/main/python/get_json_test.py @@ -14,12 +14,13 @@ import pytest -from asserts import assert_gpu_and_cpu_are_equal_collect, assert_gpu_fallback_collect, with_gpu_session +from asserts import (assert_cpu_and_gpu_are_equal_collect_with_capture, + assert_gpu_and_cpu_are_equal_collect, assert_gpu_fallback_collect, + with_gpu_session) from data_gen import * from pyspark.sql.types import * from marks import * from spark_init_internal import spark_version -from conftest import is_dataproc_runtime, is_dataproc_serverless_runtime from spark_session import is_before_spark_400, is_databricks113_or_later, is_databricks_runtime def mk_json_str_gen(pattern): @@ -123,15 +124,39 @@ def test_get_json_object_normalize_non_string_output(): f.col('jsonStr'), f.get_json_object('jsonStr', '$'))) -@pytest.mark.xfail(condition=is_dataproc_runtime() or is_dataproc_serverless_runtime(), - reason="https://github.com/NVIDIA/spark-rapids/issues/14290") def test_get_json_object_quoted_question(): schema = StructType([StructField("jsonStr", StringType())]) data = [[r'{"?":"QUESTION"}']] - assert_gpu_and_cpu_are_equal_collect( + assert_cpu_and_gpu_are_equal_collect_with_capture( + lambda spark: spark.createDataFrame(data,schema=schema).select( + f.get_json_object('jsonStr',r'''$['?']''').alias('question')), + exist_classes='GpuGetJsonObject') + + +def test_multi_get_json_object_quoted_question(): + schema = StructType([StructField("jsonStr", StringType())]) + data = [[r'{"?":"QUESTION","a?b":"EMBEDDED","outer":{"?":"NESTED"}}']] + + assert_cpu_and_gpu_are_equal_collect_with_capture( lambda spark: spark.createDataFrame(data,schema=schema).select( - f.get_json_object('jsonStr',r'''$['?']''').alias('question'))) + f.get_json_object('jsonStr',r'''$['?']''').alias('question'), + f.get_json_object('jsonStr',r'''$['a?b']''').alias('embedded'), + f.get_json_object('jsonStr',r'''$.outer['?']''').alias('nested')), + exist_classes='GpuProjectExec,GpuGetJsonObject') + + +def test_multi_get_json_object_all_invalid_paths(): + schema = StructType([StructField("jsonStr", StringType())]) + data = [['{"a":"A"}']] + + assert_cpu_and_gpu_are_equal_collect_with_capture( + lambda spark: spark.createDataFrame(data,schema=schema).selectExpr( + 'get_json_object(jsonStr, CAST(NULL AS STRING)) AS null_path', + 'get_json_object(jsonStr, "$[") AS malformed_path', + 'get_json_object(jsonStr, "not_a_path") AS missing_root'), + exist_classes='GpuProjectExec,GpuGetJsonObject') + def test_get_json_object_escaped_string_data(): schema = StructType([StructField("jsonStr", StringType())]) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuGetJsonObject.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuGetJsonObject.scala index bd6760f6765..d08bb6e5957 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuGetJsonObject.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuGetJsonObject.scala @@ -17,7 +17,6 @@ package com.nvidia.spark.rapids import scala.collection.mutable -import scala.util.parsing.combinator.RegexParsers import ai.rapids.cudf.ColumnVector import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} @@ -46,7 +45,7 @@ object PathInstruction { case class Named(name: String) extends PathInstruction } -object JsonPathParser extends RegexParsers { +object JsonPathParser { // Mirrors JSONUtils.MAX_PATH_DEPTH from spark-rapids-jni (get_json_object.hpp). // Duplicated here to avoid triggering JNI native library loading during // Driver-side plan conversion (see github.com/NVIDIA/spark-rapids/issues/14184). @@ -54,49 +53,8 @@ object JsonPathParser extends RegexParsers { import PathInstruction._ - def root: Parser[Char] = '$' - - def long: Parser[Long] = "\\d+".r ^? { - case x => x.toLong - } - - // parse `[*]` and `[123]` subscripts - def subscript: Parser[List[PathInstruction]] = - for { - operand <- '[' ~> ('*' ^^^ Wildcard | long ^^ Index) <~ ']' - } yield { - Subscript :: operand :: Nil - } - - // parse `.name` or `['name']` child expressions - def named: Parser[List[PathInstruction]] = - for { - name <- '.' ~> "[^\\.\\[]+".r | "['" ~> GetJsonObjectShim.partRegexpInNamed.r <~ "']" - } yield { - Key :: Named(name) :: Nil - } - - // child wildcards: `..`, `.*` or `['*']` - def wildcard: Parser[List[PathInstruction]] = - (".*" | "['*']") ^^^ List(Wildcard) - - def node: Parser[List[PathInstruction]] = - wildcard | - named | - subscript - - val expression: Parser[List[PathInstruction]] = { - phrase(root ~> rep(node) ^^ (x => x.flatten)) - } - def parse(str: String): Option[List[PathInstruction]] = { - this.parseAll(expression, str) match { - case Success(result, _) => - Some(result) - - case _ => - None - } + GetJsonObjectShim.parse(str) } def filterInstructionsForJni(instructions: List[PathInstruction]): List[PathInstruction] = @@ -156,6 +114,14 @@ object JsonPathParser extends RegexParsers { } } +object GpuGetJsonObjectMeta { + private[rapids] def parseLiteralPath(value: Any): Option[List[PathInstruction]] = { + Option(value).map(_.asInstanceOf[UTF8String].toString).flatMap { path => + JsonPathParser.parse(path) + } + } +} + class GpuGetJsonObjectMeta( expr: GetJsonObject, conf: RapidsConf, @@ -166,7 +132,7 @@ class GpuGetJsonObjectMeta( override def tagExprForGpu(): Unit = { val lit = GpuOverrides.extractLit(expr.right) lit.foreach { l => - val instructions = JsonPathParser.parse(l.value.asInstanceOf[UTF8String].toString) + val instructions = GpuGetJsonObjectMeta.parseLiteralPath(l.value) val updated = instructions.map(JsonPathParser.filterInstructionsForJni) if (updated.exists(JsonPathParser.fallbackCheck)) { willNotWorkOnGpu(s"get_json_object on GPU does not support more " + @@ -216,13 +182,15 @@ case class GpuMultiGetJsonObject(json: Expression, val validPaths = validPathsWithIndexes.map(_._1) withResource(new Array[ColumnVector](validPaths.length)) { validPathColumns => withResource(json.columnarEval(batch)) { input => - // Last argument -1 indicates to use automatically calculated parallelism - withResource(JSONUtils.getJsonObjectMultiplePaths(input.getBase, - java.util.Arrays.asList(validPaths: _*), 4 * targetBatchSize, - -1)) { chunkedResult => - chunkedResult.foreach { cr => - validPathColumns(validPathsIndex) = cr.incRefCount() - validPathsIndex += 1 + if (validPaths.nonEmpty) { + // Last argument -1 indicates to use automatically calculated parallelism + withResource(JSONUtils.getJsonObjectMultiplePaths(input.getBase, + java.util.Arrays.asList(validPaths: _*), 4 * targetBatchSize, + -1)) { chunkedResult => + chunkedResult.foreach { cr => + validPathColumns(validPathsIndex) = cr.incRefCount() + validPathsIndex += 1 + } } } diff --git a/sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala b/sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala index c8ab49e07f0..5cc2fb84c32 100644 --- a/sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala +++ b/sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala @@ -31,7 +31,6 @@ {"spark": "350db143"} {"spark": "351"} {"spark": "352"} -{"spark": "353"} {"spark": "354"} {"spark": "355"} {"spark": "356"} @@ -41,17 +40,70 @@ spark-rapids-shim-json-lines ***/ package com.nvidia.spark.rapids.shims +import scala.util.parsing.combinator.RegexParsers + +import com.nvidia.spark.rapids.PathInstruction + +import org.apache.spark.SparkConf + object GetJsonObjectShim { + // Copied from Apache Spark 3.5.3 JsonPathParser in jsonExpressions.scala. + private object JsonPathParser extends RegexParsers { + import com.nvidia.spark.rapids.PathInstruction._ + + def root: Parser[Char] = '$' + + def long: Parser[Long] = "\\d+".r ^? { + case x => x.toLong + } + + // parse `[*]` and `[123]` subscripts + def subscript: Parser[List[PathInstruction]] = + for { + operand <- '[' ~> ('*' ^^^ Wildcard | long ^^ Index) <~ ']' + } yield { + Subscript :: operand :: Nil + } + + // parse `.name` or `['name']` child expressions + def named: Parser[List[PathInstruction]] = + for { + name <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\'\\?]+".r <~ "']" + } yield { + Key :: Named(name) :: Nil + } + + // child wildcards: `..`, `.*` or `['*']` + def wildcard: Parser[List[PathInstruction]] = + (".*" | "['*']") ^^^ List(Wildcard) + + def node: Parser[List[PathInstruction]] = + wildcard | + named | + subscript + + val expression: Parser[List[PathInstruction]] = { + phrase(root ~> rep(node) ^^ (x => x.flatten)) + } + + def parse(str: String): Option[List[PathInstruction]] = { + this.parseAll(expression, str) match { + case Success(result, _) => + Some(result) + + case _ => + None + } + } + } + + private[rapids] def parse( + str: String, + _conf: SparkConf): Option[List[PathInstruction]] = JsonPathParser.parse(str) + /** - * Return a shim string for a part in named Regexp. - * For Spark versions before 400, named Regexp is: - * name <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\'\\?]+".r <~ "']" - * For Spark versions 400 and 400+, named Regexp is: - * name <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\']+".r <~ "']" - * This is the shim to distinct "[^\\'\\?]+" and "[^\\']+" - * - * "[^\\'\\?]+" : One or more chars which are not: ' or ? - * "[^\\']+" : One or more chars which are not: ' + * Spark 3.x uses the legacy quoted-name parser. The Dataproc runtimes that backport + * SPARK-46761 use the Spark 3.5.3-specific shim instead. */ - def partRegexpInNamed: String = "[^\\'\\?]+" + def parse(str: String): Option[List[PathInstruction]] = JsonPathParser.parse(str) } diff --git a/sql-plugin/src/main/spark353/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala b/sql-plugin/src/main/spark353/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala new file mode 100644 index 00000000000..8d02a4767d6 --- /dev/null +++ b/sql-plugin/src/main/spark353/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala @@ -0,0 +1,155 @@ +/* + * Copyright (c) 2025-2026, NVIDIA CORPORATION. + * + * Licensed 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. + */ + +/*** spark-rapids-shim-json-lines +{"spark": "353"} +spark-rapids-shim-json-lines ***/ +package com.nvidia.spark.rapids.shims + +import scala.util.parsing.combinator.RegexParsers + +import com.nvidia.spark.rapids.PathInstruction + +import org.apache.spark.{SparkConf, SparkEnv} + +object GetJsonObjectShim { + private val DATAPROC_ENGINE_KEY = "spark.dataproc.engine" + + // Copied from Apache Spark 3.5.3 JsonPathParser in jsonExpressions.scala. + private object LegacyJsonPathParser extends RegexParsers { + import com.nvidia.spark.rapids.PathInstruction._ + + def root: Parser[Char] = '$' + + def long: Parser[Long] = "\\d+".r ^? { + case x => x.toLong + } + + // parse `[*]` and `[123]` subscripts + def subscript: Parser[List[PathInstruction]] = + for { + operand <- '[' ~> ('*' ^^^ Wildcard | long ^^ Index) <~ ']' + } yield { + Subscript :: operand :: Nil + } + + // parse `.name` or `['name']` child expressions + def named: Parser[List[PathInstruction]] = + for { + name <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\'\\?]+".r <~ "']" + } yield { + Key :: Named(name) :: Nil + } + + // child wildcards: `..`, `.*` or `['*']` + def wildcard: Parser[List[PathInstruction]] = + (".*" | "['*']") ^^^ List(Wildcard) + + def node: Parser[List[PathInstruction]] = + wildcard | + named | + subscript + + val expression: Parser[List[PathInstruction]] = { + phrase(root ~> rep(node) ^^ (x => x.flatten)) + } + + def parse(str: String): Option[List[PathInstruction]] = { + this.parseAll(expression, str) match { + case Success(result, _) => + Some(result) + + case _ => + None + } + } + } + + // Copied from Apache Spark 4.0.0 JsonPathParser after SPARK-46761. + private object FixedJsonPathParser extends RegexParsers { + import com.nvidia.spark.rapids.PathInstruction._ + + def root: Parser[Char] = '$' + + def long: Parser[Long] = "\\d+".r ^? { + case x => x.toLong + } + + // parse `[*]` and `[123]` subscripts + def subscript: Parser[List[PathInstruction]] = + for { + operand <- '[' ~> ('*' ^^^ Wildcard | long ^^ Index) <~ ']' + } yield { + Subscript :: operand :: Nil + } + + // parse `.name` or `['name']` child expressions + def named: Parser[List[PathInstruction]] = + for { + name <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\']+".r <~ "']" + } yield { + Key :: Named(name) :: Nil + } + + // child wildcards: `..`, `.*` or `['*']` + def wildcard: Parser[List[PathInstruction]] = + (".*" | "['*']") ^^^ List(Wildcard) + + def node: Parser[List[PathInstruction]] = + wildcard | + named | + subscript + + val expression: Parser[List[PathInstruction]] = { + phrase(root ~> rep(node) ^^ (x => x.flatten)) + } + + def parse(str: String): Option[List[PathInstruction]] = { + this.parseAll(expression, str) match { + case Success(result, _) => + Some(result) + + case _ => + None + } + } + } + + private def useFixedParser(conf: SparkConf): Boolean = conf.contains(DATAPROC_ENGINE_KEY) + + private lazy val activeParserUsesFixedSemantics = + Option(SparkEnv.get).exists(env => useFixedParser(env.conf)) + + private[rapids] def parse(str: String, conf: SparkConf): Option[List[PathInstruction]] = { + if (useFixedParser(conf)) { + FixedJsonPathParser.parse(str) + } else { + LegacyJsonPathParser.parse(str) + } + } + + /** + * Dataproc classic 2.2/2.3 and Serverless 2.2/2.3 use Spark 3.5.3 builds with + * SPARK-46761 backported. Vanilla Spark 3.5.3 keeps the legacy parser. + */ + def parse(str: String): Option[List[PathInstruction]] = { + if (activeParserUsesFixedSemantics) { + FixedJsonPathParser.parse(str) + } else { + LegacyJsonPathParser.parse(str) + } + } +} diff --git a/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala b/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala index efb91ce6c2c..0f33441b674 100644 --- a/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala +++ b/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala @@ -30,17 +30,71 @@ spark-rapids-shim-json-lines ***/ package com.nvidia.spark.rapids.shims +import scala.util.parsing.combinator.RegexParsers + +import com.nvidia.spark.rapids.PathInstruction + +import org.apache.spark.SparkConf + object GetJsonObjectShim { + // Copied from Apache Spark 4.0.0 JsonPathParser after SPARK-46761. + private object JsonPathParser extends RegexParsers { + import com.nvidia.spark.rapids.PathInstruction._ + + def root: Parser[Char] = '$' + + def long: Parser[Long] = "\\d+".r ^? { + case x => x.toLong + } + + // parse `[*]` and `[123]` subscripts + def subscript: Parser[List[PathInstruction]] = + for { + operand <- '[' ~> ('*' ^^^ Wildcard | long ^^ Index) <~ ']' + } yield { + Subscript :: operand :: Nil + } + + // parse `.name` or `['name']` child expressions + def named: Parser[List[PathInstruction]] = + for { + name <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\']+".r <~ "']" + } yield { + Key :: Named(name) :: Nil + } + + // child wildcards: `..`, `.*` or `['*']` + def wildcard: Parser[List[PathInstruction]] = + (".*" | "['*']") ^^^ List(Wildcard) + + def node: Parser[List[PathInstruction]] = + wildcard | + named | + subscript + + val expression: Parser[List[PathInstruction]] = { + phrase(root ~> rep(node) ^^ (x => x.flatten)) + } + + def parse(str: String): Option[List[PathInstruction]] = { + this.parseAll(expression, str) match { + case Success(result, _) => + Some(result) + + case _ => + None + } + } + } + + private[rapids] def parse( + str: String, + _conf: SparkConf): Option[List[PathInstruction]] = { + JsonPathParser.parse(str) + } + /** - * Return a shim string for a part in named Regexp. - * For Spark versions before 400, named Regexp is: - * name <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\'\\?]+".r <~ "']" - * For Spark versions 400 and 400+, named Regexp is: - * name <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\']+".r <~ "']" - * This is the shim to distinct "[^\\'\\?]+" and "[^\\']+" - * - * "[^\\'\\?]+" : One or more chars which are not: ' or ? - * "[^\\']+" : One or more chars which are not: ' + * Spark 4 includes SPARK-46761, which accepts question marks in quoted path names. */ - def partRegexpInNamed: String = "[^\\']+" + def parse(str: String): Option[List[PathInstruction]] = JsonPathParser.parse(str) } diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/JsonPathParserSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/JsonPathParserSuite.scala new file mode 100644 index 00000000000..d43507de162 --- /dev/null +++ b/tests/src/test/scala/com/nvidia/spark/rapids/JsonPathParserSuite.scala @@ -0,0 +1,86 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed 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 com.nvidia.spark.rapids + +import com.nvidia.spark.rapids.PathInstruction.{Key, Named} +import com.nvidia.spark.rapids.shims.GetJsonObjectShim +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.SparkConf +import org.apache.spark.sql.catalyst.expressions.{GetJsonObject, Literal} +import org.apache.spark.sql.types.StringType +import org.apache.spark.unsafe.types.UTF8String + +class JsonPathParserSuite extends AnyFunSuite { + private val questionMarkPath = List(Key, Named("?")) + + // Classic 2.1 and Serverless 1.2 use the legacy parser through Spark 3.3.2/3.5.1. + // Classic 2.2/2.3 and Serverless 2.2/2.3 use patched Spark 3.5.3; Spark 4 is fixed. + private val dataprocUsesFixedParser = { + val sparkVersion = org.apache.spark.SPARK_VERSION + sparkVersion.startsWith("3.5.3") || sparkVersion.split('.').head.toInt >= 4 + } + + test("supported Dataproc shims select measured quoted question mark semantics") { + val dataprocConf = new SparkConf(false).set("spark.dataproc.engine", "default") + val questionMarkCases = Seq( + "$['?']" -> List(Key, Named("?")), + "$['a?b']" -> List(Key, Named("a?b")), + "$.outer['?']" -> List(Key, Named("outer"), Key, Named("?"))) + + questionMarkCases.foreach { case (path, expected) => + val expectedResult = if (dataprocUsesFixedParser) Some(expected) else None + assert(GetJsonObjectShim.parse(path, dataprocConf) === expectedResult) + } + } + + test("unquoted and malformed paths are independent of the configured platform") { + val vanillaConf = new SparkConf(false) + val dataprocConf = new SparkConf(false).set("spark.dataproc.engine", "default") + + Seq(vanillaConf, dataprocConf).foreach { conf => + assert(GetJsonObjectShim.parse("$.?", conf) === + Some(questionMarkPath)) + assert(GetJsonObjectShim.parse("$['ordinary']", conf) === + Some(List(Key, Named("ordinary")))) + assert(GetJsonObjectShim.parse("$['']", conf).isEmpty) + assert(GetJsonObjectShim.parse("$['unterminated]", conf).isEmpty) + } + } + + test("literal path parsing handles null") { + assert(GpuGetJsonObjectMeta.parseLiteralPath(null).isEmpty) + } + + test("vanilla shim parser matches the active Spark CPU expression") { + val expectedValue = "QUESTION" + val json = Literal.create( + UTF8String.fromString(s"""{"?":"$expectedValue"}"""), StringType) + val path = Literal.create(UTF8String.fromString("$['?']"), StringType) + val cpuResult = Option(GetJsonObject(json, path).eval(null)).map(_.toString) + val expectedInstructions = cpuResult match { + case Some(`expectedValue`) => Some(questionMarkPath) + case None => None + case other => fail(s"Unexpected CPU get_json_object result: $other") + } + val vanillaConf = new SparkConf(false) + + assert(GetJsonObjectShim.parse("$['?']", vanillaConf) === + expectedInstructions) + assert(JsonPathParser.parse("$['?']") === expectedInstructions) + } +}