Add jsonpath operator @# for returning conditionally wrapped in array sequences

Author: Nikita Glukhov

src/backend/utils/adt/jsonpath_exec.c

@@ -2480,6 +2480,44 @@ jsonb_jsonpath_query3(PG_FUNCTION_ARGS)
 	return jsonb_jsonpath_query(fcinfo);
 }
 
+static Datum
+jsonb_jsonpath_query_wrapped(FunctionCallInfo fcinfo, List *vars)
+{
+	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
+	JsonPath   *jp = PG_GETARG_JSONPATH_P(1);
+	JsonValueList found = { 0 };
+	JsonPathExecResult	res;
+	int			size;
+
+	res = executeJsonPath(jp, vars, jb, &found);
+
+	if (jperIsError(res))
+		throwJsonPathError(res);
+
+	size = JsonValueListLength(&found);
+
+	if (size == 0)
+		PG_RETURN_NULL();
+
+	if (size == 1)
+		PG_RETURN_JSONB_P(JsonbValueToJsonb(JsonValueListHead(&found)));
+
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(wrapItemsInArray(&found)));
+}
+
+Datum
+jsonb_jsonpath_query_wrapped2(PG_FUNCTION_ARGS)
+{
+	return jsonb_jsonpath_query_wrapped(fcinfo, NIL);
+}
+
+Datum
+jsonb_jsonpath_query_wrapped3(PG_FUNCTION_ARGS)
+{
+	return jsonb_jsonpath_query_wrapped(fcinfo,
+										makePassingVars(PG_GETARG_JSONB_P(2)));
+}
+
 /* Construct a JSON array from the item list */
 static inline JsonbValue *
 wrapItemsInArray(const JsonValueList *items)

src/include/catalog/pg_operator.dat

@@ -3206,5 +3206,8 @@
   oprname => '@~', oprleft => 'jsonb', oprright => 'jsonpath',
   oprresult => 'bool', oprcode => 'jsonpath_predicate(jsonb,jsonpath)',
   oprrest => 'contsel', oprjoin => 'contjoinsel' },
+{ oid => '6122', descr => 'jsonpath items wrapped',
+  oprname => '@#', oprleft => 'jsonb', oprright => 'jsonpath',
+  oprresult => 'jsonb', oprcode => 'jsonpath_query_wrapped(jsonb,jsonpath)' },
 
 ]

src/include/catalog/pg_proc.dat

@@ -9117,13 +9117,20 @@
   proname => 'jsonpath_query', prorows => '1000', proretset => 't',
   prorettype => 'jsonb', proargtypes => 'jsonb jsonpath',
   prosrc => 'jsonb_jsonpath_query2' },
+{ oid => '6124', descr => 'implementation of @# operator',
+  proname => 'jsonpath_query_wrapped', prorettype => 'jsonb',
+  proargtypes => 'jsonb jsonpath', prosrc => 'jsonb_jsonpath_query_wrapped2' },
 { oid => '6056', descr => 'jsonpath exists test',
   proname => 'jsonpath_exists', prorettype => 'bool',
   proargtypes => 'jsonb jsonpath jsonb', prosrc => 'jsonb_jsonpath_exists3' },
 { oid => '6057', descr => 'jsonpath query',
   proname => 'jsonpath_query', prorows => '1000', proretset => 't',
   prorettype => 'jsonb', proargtypes => 'jsonb jsonpath jsonb',
   prosrc => 'jsonb_jsonpath_query3' },
+{ oid => '6125', descr => 'jsonpath query with conditional wrapper',
+  proname => 'jsonpath_query_wrapped', prorettype => 'jsonb',
+  proargtypes => 'jsonb jsonpath jsonb',
+  prosrc => 'jsonb_jsonpath_query_wrapped3' },
 { oid => '6073', descr => 'implementation of @~ operator',
   proname => 'jsonpath_predicate', prorettype => 'bool',
   proargtypes => 'jsonb jsonpath', prosrc => 'jsonb_jsonpath_predicate2' },

src/test/regress/expected/jsonb_jsonpath.out

@@ -1620,6 +1620,24 @@ SELECT jsonb '[{"a": 1}, {"a": 2}]' @* '$[*] ? (@.a > 10)';
 ----------
 (0 rows)
 
+SELECT jsonb '[{"a": 1}, {"a": 2}]' @# '$[*].a';
+ ?column? 
+----------
+ [1, 2]
+(1 row)
+
+SELECT jsonb '[{"a": 1}, {"a": 2}]' @# '$[*].a ? (@ == 1)';
+ ?column? 
+----------
+ 1
+(1 row)
+
+SELECT jsonb '[{"a": 1}, {"a": 2}]' @# '$[*].a ? (@ > 10)';
+ ?column? 
+----------
+ 
+(1 row)
+
 SELECT jsonb '[{"a": 1}, {"a": 2}]' @? '$[*].a ? (@ > 1)';
  ?column? 
 ----------

src/test/regress/sql/jsonb_jsonpath.sql

@@ -361,6 +361,10 @@ set time zone default;
 SELECT jsonb '[{"a": 1}, {"a": 2}]' @* '$[*]';
 SELECT jsonb '[{"a": 1}, {"a": 2}]' @* '$[*] ? (@.a > 10)';
 
+SELECT jsonb '[{"a": 1}, {"a": 2}]' @# '$[*].a';
+SELECT jsonb '[{"a": 1}, {"a": 2}]' @# '$[*].a ? (@ == 1)';
+SELECT jsonb '[{"a": 1}, {"a": 2}]' @# '$[*].a ? (@ > 10)';
+
 SELECT jsonb '[{"a": 1}, {"a": 2}]' @? '$[*].a ? (@ > 1)';
 SELECT jsonb '[{"a": 1}, {"a": 2}]' @? '$[*] ? (@.a > 2)';