Add jsonpath .map() item method

Author: Nikita Glukhov

src/backend/utils/adt/jsonpath.c

@@ -137,6 +137,7 @@ flattenJsonPathParseItem(StringInfo buf, JsonPathParseItem *item,
 		case jpiPlus:
 		case jpiMinus:
 		case jpiExists:
+		case jpiMap:
 			{
 				int32 arg;
 
@@ -145,6 +146,7 @@ flattenJsonPathParseItem(StringInfo buf, JsonPathParseItem *item,
 
 				chld = flattenJsonPathParseItem(buf, item->value.arg,
 												item->type == jpiFilter ||
+												item->type == jpiMap ||
 												allowCurrent,
 												insideArraySubscript);
 				*(int32*)(buf->data + arg) = chld;
@@ -551,6 +553,12 @@ printJsonPathItem(StringInfo buf, JsonPathItem *v, bool inKey, bool printBracket
 		case jpiKeyValue:
 			appendBinaryStringInfo(buf, ".keyvalue()", 11);
 			break;
+		case jpiMap:
+			appendBinaryStringInfo(buf, ".map(", 5);
+			jspGetArg(v, &elem);
+			printJsonPathItem(buf, &elem, false, false);
+			appendStringInfoChar(buf, ')');
+			break;
 		default:
 			elog(ERROR, "Unknown jsonpath item type: %d", v->type);
 	}
@@ -684,6 +692,7 @@ jspInitByBuffer(JsonPathItem *v, char *base, int32 pos)
 		case jpiMinus:
 		case jpiFilter:
 		case jpiDatetime:
+		case jpiMap:
 			read_int32(v->content.arg, base, pos);
 			break;
 		case jpiIndexArray:
@@ -710,7 +719,8 @@ jspGetArg(JsonPathItem *v, JsonPathItem *a)
 		v->type == jpiExists ||
 		v->type == jpiPlus ||
 		v->type == jpiMinus ||
-		v->type == jpiDatetime
+		v->type == jpiDatetime ||
+		v->type == jpiMap
 	);
 
 	jspInitByBuffer(a, v->base, v->content.arg);
@@ -762,7 +772,8 @@ jspGetNext(JsonPathItem *v, JsonPathItem *a)
 			v->type == jpiDouble ||
 			v->type == jpiDatetime ||
 			v->type == jpiKeyValue ||
-			v->type == jpiStartsWith
+			v->type == jpiStartsWith ||
+			v->type == jpiMap
 		);
 
 		if (a)

src/backend/utils/adt/jsonpath_exec.c

@@ -2057,6 +2057,85 @@ recursiveExecuteNoUnwrap(JsonPathExecContext *cxt, JsonPathItem *jsp,
 			res = executeLikeRegexPredicate(cxt, jsp, jb);
 			res = appendBoolResult(cxt, jsp, found, res, needBool);
 			break;
+		case jpiMap:
+			if (JsonbType(jb) != jbvArray)
+			{
+				if (cxt->lax)
+				{
+					JsonValueList reslist = { 0 };
+
+					jspGetArg(jsp, &elem);
+					res = recursiveExecute(cxt, &elem, jb, &reslist);
+
+					if (jperIsError(res))
+						return res;
+
+					if (JsonValueListLength(&reslist) != 1)
+						return jperMakeError(ERRCODE_SINGLETON_JSON_ITEM_REQUIRED);
+
+					res = recursiveExecuteNext(cxt, jsp, NULL,
+											   JsonValueListHead(&reslist),
+											   found, true);
+				}
+				else
+					res = jperMakeError(ERRCODE_JSON_ARRAY_NOT_FOUND);
+			}
+			else
+			{
+				JsonbValue	element_buf;
+				JsonbValue *element;
+				JsonbIterator *it = NULL;
+				JsonbIteratorToken tok;
+				JsonValueList result = { 0 };
+				int			size = JsonbArraySize(jb);
+				int			i;
+
+				jspGetArg(jsp, &elem);
+
+				if (jb->type == jbvBinary && size > 0)
+				{
+					element = &element_buf;
+					it = JsonbIteratorInit(jb->val.binary.data);
+					tok = JsonbIteratorNext(&it, &element_buf, false);
+					if (tok != WJB_BEGIN_ARRAY)
+						elog(ERROR, "unexpected jsonb token at the array start");
+				}
+
+				for (i = 0; i < size; i++)
+				{
+					JsonValueList reslist = { 0 };
+
+					if (it)
+					{
+						tok = JsonbIteratorNext(&it, element, true);
+						if (tok != WJB_ELEM)
+							break;
+					}
+					else
+						element = &jb->val.array.elems[i];
+
+					res = recursiveExecute(cxt, &elem, element, &reslist);
+
+					if (jperIsError(res))
+						break;
+
+					if (JsonValueListLength(&reslist) != 1)
+					{
+						res = jperMakeError(ERRCODE_SINGLETON_JSON_ITEM_REQUIRED);
+						break;
+					}
+
+					JsonValueListConcat(&result, reslist);
+				}
+
+				if (jperIsError(res))
+					break;
+
+				res = recursiveExecuteNext(cxt, jsp, NULL,
+										   wrapItemsInArray(&result),
+										   found, false);
+			}
+			break;
 		default:
 			elog(ERROR,"2Wrong state: %d", jsp->type);
 	}

src/backend/utils/adt/jsonpath_gram.y

@@ -266,7 +266,7 @@ makeItemLikeRegex(JsonPathParseItem *expr, string *pattern, string *flags)
 %token	<str>		LESS_P LESSEQUAL_P EQUAL_P NOTEQUAL_P GREATEREQUAL_P GREATER_P
 %token	<str>		ANY_P STRICT_P LAX_P LAST_P STARTS_P WITH_P LIKE_REGEX_P FLAG_P
 %token	<str>		ABS_P SIZE_P TYPE_P FLOOR_P DOUBLE_P CEILING_P DATETIME_P
-%token	<str>		KEYVALUE_P
+%token	<str>		KEYVALUE_P MAP_P
 
 %type	<result>	result
 
@@ -425,6 +425,8 @@ accessor_op:
 	| '.' method '(' ')'			{ $$ = makeItemType($2); }
 	| '.' DATETIME_P '(' opt_datetime_template ')'
 									{ $$ = makeItemUnary(jpiDatetime, $4); }
+	| '.' MAP_P '(' expr_or_predicate ')'
+									{ $$ = makeItemUnary(jpiMap, $4); }
 	| '?' '(' predicate ')'			{ $$ = makeItemUnary(jpiFilter, $3); }
 	;
 
@@ -462,6 +464,7 @@ key_name:
 	| WITH_P
 	| LIKE_REGEX_P
 	| FLAG_P
+	| MAP_P
 	;
 
 method:

src/backend/utils/adt/jsonpath_scan.l

@@ -273,6 +273,7 @@ static keyword keywords[] = {
 	{ 2, false,	TO_P,		"to"},
 	{ 3, false,	ABS_P,		"abs"},
 	{ 3, false,	LAX_P,		"lax"},
+	{ 3, false,	MAP_P,		"map"},
 	{ 4, false,	FLAG_P,		"flag"},
 	{ 4, false,	LAST_P,		"last"},
 	{ 4, true,	NULL_P,		"null"},

src/include/utils/jsonpath.h

@@ -84,6 +84,7 @@ typedef enum JsonPathItemType {
 		jpiLast,
 		jpiStartsWith,
 		jpiLikeRegex,
+		jpiMap,
 } JsonPathItemType;
 
 /* XQuery regex mode flags for LIKE_REGEX predicate */

src/test/regress/expected/jsonb_jsonpath.out

@@ -1651,3 +1651,24 @@ SELECT jsonb '[{"a": 1}, {"a": 2}]' @~ '$[*].a > 2';
  f
 (1 row)
 
+-- extension: map item method
+select jsonb '1' @* 'strict $.map(@ + 10)';
+ERROR:  SQL/JSON array not found
+select jsonb '1' @* 'lax $.map(@ + 10)';
+ ?column? 
+----------
+ 11
+(1 row)
+
+select jsonb '[1, 2, 3]' @* '$.map(@ + 10)';
+   ?column?   
+--------------
+ [11, 12, 13]
+(1 row)
+
+select jsonb '[[1, 2], [3, 4, 5], [], [6, 7]]' @* '$.map(@.map(@ + 10))';
+                ?column?                
+----------------------------------------
+ [[11, 12], [13, 14, 15], [], [16, 17]]
+(1 row)
+

src/test/regress/sql/jsonb_jsonpath.sql

@@ -369,3 +369,9 @@ SELECT jsonb '[{"a": 1}, {"a": 2}]' @? '$[*] ? (@.a > 2)';
 
 SELECT jsonb '[{"a": 1}, {"a": 2}]' @~ '$[*].a > 1';
 SELECT jsonb '[{"a": 1}, {"a": 2}]' @~ '$[*].a > 2';
+
+-- extension: map item method
+select jsonb '1' @* 'strict $.map(@ + 10)';
+select jsonb '1' @* 'lax $.map(@ + 10)';
+select jsonb '[1, 2, 3]' @* '$.map(@ + 10)';
+select jsonb '[[1, 2], [3, 4, 5], [], [6, 7]]' @* '$.map(@.map(@ + 10))';