From b3a594e01c71299572a5775ad367f703e3cde99b Mon Sep 17 00:00:00 2001
From: Clint Wylie <cwylie@apache.org>
Date: Wed, 27 Jul 2022 22:20:54 -0700
Subject: [PATCH] fix bugs with nested column jsonpath parser

---
 .../druid/segment/nested/NestedPathFinder.java  | 17 ++++++++++-------
 .../frame/segment/FrameStorageAdapterTest.java  |  1 +
 .../segment/nested/NestedPathFinderTest.java    |  7 ++++++-
 .../sql/calcite/CalciteNestedDataQueryTest.java |  2 +-
 4 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/segment/nested/NestedPathFinder.java b/processing/src/main/java/org/apache/druid/segment/nested/NestedPathFinder.java
index 7dc60a027fed..e77a629bc23f 100644
--- a/processing/src/main/java/org/apache/druid/segment/nested/NestedPathFinder.java
+++ b/processing/src/main/java/org/apache/druid/segment/nested/NestedPathFinder.java
@@ -33,21 +33,20 @@ public class NestedPathFinder
   public static String toNormalizedJsonPath(List<NestedPathPart> paths)
   {
     if (paths.isEmpty()) {
-      return "$.";
+      return "$";
     }
     StringBuilder bob = new StringBuilder();
     boolean first = true;
     for (NestedPathPart partFinder : paths) {
       if (partFinder instanceof NestedPathField) {
         if (first) {
-          bob.append("$.");
-        } else {
-          bob.append(".");
+          bob.append("$");
         }
         final String id = partFinder.getPartIdentifier();
         if (id.contains(".") || id.contains("'") || id.contains("\"") || id.contains("[") || id.contains("]")) {
           bob.append("['").append(id).append("']");
         } else {
+          bob.append(".");
           bob.append(id);
         }
       } else if (partFinder instanceof NestedPathArrayElement) {
@@ -124,6 +123,13 @@ public static List<NestedPathPart> parseJsonPath(@Nullable String path)
         quoteMark = i;
         partMark = i + 1;
       } else if (current == '\'' && quoteMark >= 0 && path.charAt(i - 1) != '\\') {
+        if (path.charAt(i + 1) != ']') {
+          if (arrayMark >= 0) {
+            continue;
+          }
+          badFormatJsonPath(path, "closing ' must immediately precede ']'");
+        }
+
         parts.add(new NestedPathField(getPathSubstring(path, partMark, i)));
         dotMark = -1;
         quoteMark = -1;
@@ -131,9 +137,6 @@ public static List<NestedPathPart> parseJsonPath(@Nullable String path)
         if (++i == path.length()) {
           break;
         }
-        if (path.charAt(i) != ']') {
-          badFormatJsonPath(path, "closing ' must immediately precede ']'");
-        }
         partMark = i + 1;
         arrayMark = -1;
       }
diff --git a/processing/src/test/java/org/apache/druid/frame/segment/FrameStorageAdapterTest.java b/processing/src/test/java/org/apache/druid/frame/segment/FrameStorageAdapterTest.java
index 94cee3ed1ce7..1f4b0e45aee8 100644
--- a/processing/src/test/java/org/apache/druid/frame/segment/FrameStorageAdapterTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/segment/FrameStorageAdapterTest.java
@@ -100,6 +100,7 @@ public static Iterable<Object[]> constructorFeeder()
     @Before
     public void setUp()
     {
+
       queryableAdapter = new QueryableIndexStorageAdapter(TestIndex.getMMappedTestIndex());
       frameSegment = FrameTestUtil.adapterToFrameSegment(queryableAdapter, frameType);
       frameAdapter = frameSegment.asStorageAdapter();
diff --git a/processing/src/test/java/org/apache/druid/segment/nested/NestedPathFinderTest.java b/processing/src/test/java/org/apache/druid/segment/nested/NestedPathFinderTest.java
index d8c771cf3a4c..b60a30c8a605 100644
--- a/processing/src/test/java/org/apache/druid/segment/nested/NestedPathFinderTest.java
+++ b/processing/src/test/java/org/apache/druid/segment/nested/NestedPathFinderTest.java
@@ -309,9 +309,14 @@ public void testParseJsonPath()
         NestedPathFinder.toNormalizedJqPath(pathParts)
     );
     Assert.assertEquals(
-        "$.['x.y.z][\\']][]'].['13234.12[]][23'].['fo.o'].['.b.a.r.']",
+        "$['x.y.z][\\']][]']['13234.12[]][23']['fo.o']['.b.a.r.']",
         NestedPathFinder.toNormalizedJsonPath(pathParts)
     );
+
+    pathParts = NestedPathFinder.parseJsonPath("$['hell'o']");
+    Assert.assertEquals(1, pathParts.size());
+    Assert.assertEquals("hell'o", pathParts.get(0).getPartIdentifier());
+    Assert.assertEquals("$['hell'o']", NestedPathFinder.toNormalizedJsonPath(pathParts));
   }
 
   @Test
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
index ea3c66d11da9..e66919e009c4 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
@@ -1955,7 +1955,7 @@ public void testGroupByAllPaths() throws Exception
                         .build()
         ),
         ImmutableList.of(
-            new Object[]{"[\"$.\"]", 5L},
+            new Object[]{"[\"$\"]", 5L},
             new Object[]{"[\"$.n.x\",\"$.array[0]\",\"$.array[1]\"]", 2L}
         )
     );