Add more test cases for json_query function

Author: XavieLee

core/trino-main/src/test/java/io/trino/sql/query/TestJsonQueryFunction.java

@@ -425,7 +425,7 @@ public void testNullInput()
     }
 
     @Test
-    public void testDescendantMemberAccessor()
+    public void testDescendantMemberWithArrayWrapperAccessor()
     {
         assertThat(assertions.query(
                 """
@@ -438,5 +438,64 @@ SELECT json_query(
                         """
                         VALUES cast('[[true,{"c":{"c":null}}],{"c":null},null]'AS varchar)
                         """);
+
+        // `json_query` must return a single JSON item.
+        // `WITH ARRAY WRAPPER` wraps `Diana` and `John` into a single JSON array, which can then be successfully returned.
+        assertThat(assertions.query(
+                """
+                SELECT json_query(
+                                '{"author":"Diana","id":{"value":1, "author":"John"},"notes":[{"type":1,"comment":"foo"},{"type":2,"comment":null}],"comment":["bar","baz"]}',
+                                'lax $..author'
+                                WITH ARRAY WRAPPER)
+                """))
+                .matches(
+                        """
+                        VALUES cast('["Diana","John"]'AS varchar)
+                        """);
+    }
+
+    @Test
+    public void testDescendantMemberWithoutArrayWrapperAccessor()
+    {
+        // Test `json_query` to find a non-descendant member in the JSON.
+        assertThat(assertions.query(
+                """
+                select json_query(
+                                '{"id":{"value":1},"notes":[{"type":1,"comment":"foo"},{"type":2,"comment":null}],"comment":["bar","baz"]}',
+                                'lax $..author'
+                                omit quotes)
+                """))
+                .matches(
+                        """
+                        VALUES cast(NULL as varchar)
+                    """);
+
+        // Test `json_query` to find a single descendant member in the JSON.
+        assertThat(assertions.query(
+                """
+                select json_query(
+                                '{"author":"Diana","id":{"value":1},"notes":[{"type":1,"comment":"foo"},{"type":2,"comment":null}],"comment":["bar","baz"]}',
+                                'lax $..author'
+                                omit quotes)
+                """))
+                .matches(
+                        """
+                        VALUES cast('Diana' as varchar)
+                    """);
+
+        // Test `json_query` to find multiple descendant members in the JSON.
+        // `json_query` must return a single JSON item.
+        // `WITHOUT ARRAY WRAPPER` returns `NULL`, which is the default result in case of an error.
+        assertThat(assertions.query(
+                """
+                select json_query(
+                                '{"author":"Diana","id":{"value":1, "author":"John"},"notes":[{"type":1,"comment":"foo"},{"type":2,"comment":null}],"comment":["bar","baz"]}',
+                                'lax $..author'
+                                omit quotes)
+                """))
+                .matches(
+                        """
+                        VALUES cast(NULL as varchar)
+                    """);
     }
 }