[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

@@ -68,10 +68,13 @@ object JsonPathParser extends RegexParsers {
       Subscript :: operand :: Nil
     }
 
+  private val legacyNamedPartRegexp = "[^\\'\\?]+"
+  private val fixedNamedPartRegexp = "[^\\']+"
+
   // parse `.name` or `['name']` child expressions
-  def named: Parser[List[PathInstruction]] =
+  private def named(partRegexpInNamed: String): Parser[List[PathInstruction]] =
     for {
-      name <- '.' ~> "[^\\.\\[]+".r | "['" ~> GetJsonObjectShim.partRegexpInNamed.r <~ "']"
+      name <- '.' ~> "[^\\.\\[]+".r | "['" ~> partRegexpInNamed.r <~ "']"
     } yield {
       Key :: Named(name) :: Nil
     }
@@ -80,16 +83,25 @@ object JsonPathParser extends RegexParsers {
   def wildcard: Parser[List[PathInstruction]] =
     (".*" | "['*']") ^^^ List(Wildcard)
 
-  def node: Parser[List[PathInstruction]] =
+  private def node(partRegexpInNamed: String): Parser[List[PathInstruction]] =
     wildcard |
-      named |
+      named(partRegexpInNamed) |
       subscript
 
-  val expression: Parser[List[PathInstruction]] = {
-    phrase(root ~> rep(node) ^^ (x => x.flatten))
+  private def pathExpression(partRegexpInNamed: String): Parser[List[PathInstruction]] = {
+    phrase(root ~> rep(node(partRegexpInNamed)) ^^ (x => x.flatten))
   }
 
-  def parse(str: String): Option[List[PathInstruction]] = {
+  private lazy val legacyPathExpression = pathExpression(legacyNamedPartRegexp)
+  private lazy val fixedPathExpression = pathExpression(fixedNamedPartRegexp)
+
+  def parse(str: String, allowQuestionMarkInQuotedName: Boolean):
+      Option[List[PathInstruction]] = {
+    val expression = if (allowQuestionMarkInQuotedName) {
+      fixedPathExpression
+    } else {
+      legacyPathExpression
+    }
     this.parseAll(expression, str) match {
       case Success(result, _) =>
         Some(result)
@@ -156,29 +168,61 @@ object JsonPathParser extends RegexParsers {
   }
 }
 
+object GpuGetJsonObjectMeta {
+  private[rapids] val UNKNOWN_QUESTION_MARK_SUPPORT_REASON =
+    "Could not determine whether this Spark runtime accepts question marks in quoted " +
+      "get_json_object paths"
+
+  private[rapids] def parseLiteralPath(
+      value: Any,
+      allowQuestionMarkInQuotedName: Boolean): Option[List[PathInstruction]] = {
+    Option(value).map(_.asInstanceOf[UTF8String].toString).flatMap { path =>
+      JsonPathParser.parse(path, allowQuestionMarkInQuotedName)
+    }
+  }
+
+  private[rapids] def unsupportedReason(
+      quotedQuestionMarkSupport: Option[Boolean]): Option[String] = {
+    if (quotedQuestionMarkSupport.isDefined) {
+      None
+    } else {
+      Some(UNKNOWN_QUESTION_MARK_SUPPORT_REASON)
+    }
+  }
+}
+
 class GpuGetJsonObjectMeta(
     expr: GetJsonObject,
     conf: RapidsConf,
     parent: Option[RapidsMeta[_, _, _]],
     rule: DataFromReplacementRule
   ) extends BinaryExprMeta[GetJsonObject](expr, conf, parent, rule) {
 
+  private val quotedQuestionMarkSupport = GetJsonObjectShim.quotedQuestionMarkSupport
+
   override def tagExprForGpu(): Unit = {
-    val lit = GpuOverrides.extractLit(expr.right)
-    lit.foreach { l =>
-      val instructions = JsonPathParser.parse(l.value.asInstanceOf[UTF8String].toString)
-      val updated = instructions.map(JsonPathParser.filterInstructionsForJni)
-      if (updated.exists(JsonPathParser.fallbackCheck)) {
-        willNotWorkOnGpu(s"get_json_object on GPU does not support more " +
-          s"than ${JsonPathParser.MAX_PATH_DEPTH} nested paths." +
-          instructions.map(i => s" (Found ${i.length})").getOrElse(""))
+    GpuGetJsonObjectMeta.unsupportedReason(quotedQuestionMarkSupport).foreach(willNotWorkOnGpu)
+    quotedQuestionMarkSupport.foreach { allowQuestionMark =>
+      val lit = GpuOverrides.extractLit(expr.right)
+      lit.foreach { l =>
+        val instructions =
+          GpuGetJsonObjectMeta.parseLiteralPath(l.value, allowQuestionMark)
+        val updated = instructions.map(JsonPathParser.filterInstructionsForJni)
+        if (updated.exists(JsonPathParser.fallbackCheck)) {
+          willNotWorkOnGpu(s"get_json_object on GPU does not support more " +
+            s"than ${JsonPathParser.MAX_PATH_DEPTH} nested paths." +
+            instructions.map(i => s" (Found ${i.length})").getOrElse(""))
+        }
       }
     }
   }
 
   override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
+    val allowQuestionMark = quotedQuestionMarkSupport.getOrElse {
+      throw new IllegalStateException(GpuGetJsonObjectMeta.UNKNOWN_QUESTION_MARK_SUPPORT_REASON)
+    }
     GpuGetJsonObject(lhs, rhs)(
-      conf.testGetJsonObjectSavePath, conf.testGetJsonObjectSaveRows)
+      conf.testGetJsonObjectSavePath, conf.testGetJsonObjectSaveRows, allowQuestionMark)
   }
 }
 
@@ -216,13 +260,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
+            }
           }
         }
 
@@ -311,7 +357,8 @@ class GetJsonObjectCombiner(private val exp: GpuGetJsonObject) extends GpuExpres
             case u: UTF8String => u.toString
             case _ => null.asInstanceOf[String]
           }
-          val pathInstructions = parseJsonPath(str)
+          val pathInstructions =
+            parseJsonPath(str, e.allowQuestionMarkInQuotedName)
           if (hasSeparateWildcard(pathInstructions)) {
             // If has separate wildcard path, should return all nulls
             None
@@ -336,18 +383,22 @@ class GetJsonObjectCombiner(private val exp: GpuGetJsonObject) extends GpuExpres
 }
 
 object GpuGetJsonObject {
-  def parseJsonPath(path: GpuScalar): Option[List[PathInstruction]] = {
+  def parseJsonPath(
+      path: GpuScalar,
+      allowQuestionMarkInQuotedName: Boolean): Option[List[PathInstruction]] = {
     if (path.isValid) {
       val pathStr = path.getValue.toString
-      JsonPathParser.parse(pathStr)
+      JsonPathParser.parse(pathStr, allowQuestionMarkInQuotedName)
     } else {
       None
     }
   }
 
-  def parseJsonPath(pathStr: String): Option[List[PathInstruction]] = {
+  def parseJsonPath(
+      pathStr: String,
+      allowQuestionMarkInQuotedName: Boolean): Option[List[PathInstruction]] = {
     if (pathStr != null) {
-      JsonPathParser.parse(pathStr)
+      JsonPathParser.parse(pathStr, allowQuestionMarkInQuotedName)
     } else {
       None
     }
@@ -392,7 +443,8 @@ case class GpuGetJsonObject(
      json: Expression,
      path: Expression)(
   val savePathForVerify: Option[String],
-  val saveRowsForVerify: Int)
+  val saveRowsForVerify: Int,
+  val allowQuestionMarkInQuotedName: Boolean)
     extends GpuBinaryExpressionArgsAnyScalar
         with ExpectsInputTypes
         with GpuCombinable {
@@ -405,8 +457,10 @@ case class GpuGetJsonObject(
   }
   val seed = System.nanoTime()
 
-  override def otherCopyArgs: Seq[AnyRef] = Seq(savePathForVerify,
-    saveRowsForVerify.asInstanceOf[java.lang.Integer])
+  override def otherCopyArgs: Seq[AnyRef] = Seq(
+    savePathForVerify,
+    saveRowsForVerify.asInstanceOf[java.lang.Integer],
+    allowQuestionMarkInQuotedName.asInstanceOf[java.lang.Boolean])
 
   override def left: Expression = json
   override def right: Expression = path
@@ -420,7 +474,7 @@ case class GpuGetJsonObject(
 
   override def doColumnar(lhs: GpuColumnVector, rhs: GpuScalar): ColumnVector = {
     val fromGpu = cachedInstructions.getOrElse {
-      val pathInstructions = parseJsonPath(rhs)
+      val pathInstructions = parseJsonPath(rhs, allowQuestionMarkInQuotedName)
       val checkedPathInstructions = if (hasSeparateWildcard(pathInstructions)) {
         // If has separate wildcard path, should return all nulls
         None

sql-plugin/src/main/scala/com/nvidia/spark/rapids/shims/GetJsonObjectRuntimeSemantics.scala

@@ -0,0 +1,33 @@
+/*
+ * 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.shims
+
+import scala.util.Try
+
+import org.apache.spark.unsafe.types.UTF8String
+
+private[rapids] object GetJsonObjectRuntimeSemantics {
+  private val ExpectedQuestionMarkValue = "QUESTION"
+
+  def classifyQuotedQuestionMarkResult(result: => Any): Option[Boolean] = {
+    Try(result).toOption match {
+      case Some(null) => Some(false)
+      case Some(value: UTF8String) if value.toString == ExpectedQuestionMarkValue => Some(true)
+      case _ => None
+    }
+  }
+}

sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala

@@ -41,17 +41,22 @@
 spark-rapids-shim-json-lines ***/
 package com.nvidia.spark.rapids.shims
 
+import org.apache.spark.sql.catalyst.expressions.{GetJsonObject, Literal}
+import org.apache.spark.sql.types.StringType
+import org.apache.spark.unsafe.types.UTF8String
+
 object GetJsonObjectShim {
+  private lazy val runtimeQuotedQuestionMarkSupport: Option[Boolean] = {
+    val json = Literal.create(UTF8String.fromString("""{"?":"QUESTION"}"""), StringType)
+    val path = Literal.create(UTF8String.fromString("$['?']"), StringType)
+    GetJsonObjectRuntimeSemantics.classifyQuotedQuestionMarkResult {
+      GetJsonObject(json, path).eval(null)
+    }
+  }
+
   /**
-   * 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: '
+   * Detect whether this Spark runtime includes SPARK-46761 semantics. Some vendors backported the
+   * fix without changing the upstream Spark version, so a version check is not sufficient.
    */
-  def partRegexpInNamed: String = "[^\\'\\?]+"
+  def quotedQuestionMarkSupport: Option[Boolean] = runtimeQuotedQuestionMarkSupport
 }

sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/GetJsonObjectShim.scala

@@ -32,15 +32,7 @@ package com.nvidia.spark.rapids.shims
 
 object GetJsonObjectShim {
   /**
-   * 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 quotedQuestionMarkSupport: Option[Boolean] = Some(true)
 }