[SkipRecovery] Fix get_json_object runtime path semantics [databricks]

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())])

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,57 +45,16 @@ 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).
   val MAX_PATH_DEPTH: Int = 16
 
   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
+            }
           }
         }
 

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)
 }