sunchao commented on code in PR #57069:
URL: https://github.com/apache/spark/pull/57069#discussion_r3537375436


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala:
##########
@@ -145,25 +145,37 @@ case class GetJsonObject(json: Expression, path: 
Expression)
 object GetJsonObject {
   import PathInstruction._
 
-  private[sql] def simpleNamedPath(path: UTF8String): Option[Seq[String]] = {
+  private[sql] sealed trait SimpleJsonPathSegment
+  private[sql] case class NamedPathSegment(name: String) extends 
SimpleJsonPathSegment
+  private[sql] case class IndexedPathSegment(index: Long) extends 
SimpleJsonPathSegment
+
+  private[sql] def simplePath(path: UTF8String): 
Option[Seq[SimpleJsonPathSegment]] = {
     try {
       Option(path).flatMap(value => 
JsonPathParser.parse(value.toString)).flatMap { instructions =>
-        val names = instructions.grouped(2).map {
-          case List(Key, Named(fieldName)) => Some(fieldName)
+        val segments = instructions.grouped(2).map {
+          case List(Key, Named(fieldName)) => Some(NamedPathSegment(fieldName))
+          case List(Subscript, Index(index)) if index >= 0 => 
Some(IndexedPathSegment(index))
           case _ => None
         }.toSeq
-        if (names.nonEmpty && names.forall(_.isDefined)) Some(names.flatten) 
else None
+        if (segments.nonEmpty && segments.forall(_.isDefined)) 
Some(segments.flatten) else None
       }
     } catch {
       // Numeric subscripts are parsed as Long and can overflow before the 
parser returns None.
       case _: NumberFormatException => None
     }
   }
+
+  private[sql] def simpleNamedPath(path: UTF8String): Option[Seq[String]] = {

Review Comment:
   Good catch — this was left behind when all consumers moved to `simplePath`; 
it was not retained for compatibility. Removed in d9ab71ed4b4.



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala:
##########
@@ -44,76 +44,127 @@ import org.apache.spark.unsafe.types.UTF8String
 object OptimizeCsvJsonExprs extends Rule[LogicalPlan] {
   private def nameOfCorruptRecord = conf.columnNameOfCorruptRecord
 
+  private type SimpleJsonPath = Seq[GetJsonObject.SimpleJsonPathSegment]
+  private type SharedJsonCandidate = (GetJsonObject, SimpleJsonPath, String)
+  private type SharedJsonCandidateUnit = Seq[SharedJsonCandidate]
+  private type RequestedJsonPath = (SimpleJsonPath, String)
+  private type RequestedJsonPathUnit = Seq[RequestedJsonPath]
+
   private case class SharedJsonFields(
       json: Expression,
-      paths: Seq[Seq[String]],
+      paths: Seq[SimpleJsonPath],
       alias: Alias) {
-    val ordinalMapping: Map[Seq[String], Int] = paths.zipWithIndex.toMap
+    val ordinalMapping: Map[SimpleJsonPath, Int] = paths.zipWithIndex.toMap
   }
 
   private final class SelectedJsonPathTrieNode {
     var isTerminal: Boolean = false
     var hasSelectedPathInSubtree: Boolean = false
-    val children: mutable.HashMap[String, SelectedJsonPathTrieNode] = 
mutable.HashMap.empty
+    val children: mutable.HashMap[GetJsonObject.SimpleJsonPathSegment, 
SelectedJsonPathTrieNode] =
+      mutable.HashMap.empty
   }
 
   private final class SelectedJsonPathGroup {
     val root = new SelectedJsonPathTrieNode
-    val paths = mutable.ArrayBuffer.empty[(Seq[String], String)]
+    val paths = mutable.ArrayBuffer.empty[RequestedJsonPath]
 
-    def tryAdd(candidate: (Seq[String], String)): Boolean = {
-      val added = addIfNoPrefixConflict(root, candidate._1)
-      if (added) {
-        paths += candidate
+    def tryAdd(unit: RequestedJsonPathUnit): Boolean = {
+      val uniquePaths = mutable.LinkedHashMap.empty[SimpleJsonPath, String]
+      unit.foreach { case (path, fallbackPath) =>
+        uniquePaths.getOrElseUpdate(path, fallbackPath)
+      }
+      val newPaths = uniquePaths.iterator.filterNot { case (path, _) =>
+        containsExactPath(root, path)
+      }.toSeq
+      val conflictsWithGroup = newPaths.exists { case (path, _) =>
+        hasPrefixConflict(root, path)
+      }
+      val conflictsWithinUnit = newPaths.indices.exists { leftIndex =>

Review Comment:
   Yes, the current units are prefix-free by construction: either a singleton 
or a named-root/indexed-root pair. I removed `conflictsWithinUnit` and its sole 
helper `isStrictPrefix` in d9ab71ed4b4, and clarified that invariant on the 
grouping helper.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to