Skip to content

Commit d2e7a0c

Browse files
l-wangl-wang
committed
WIP: jsonb dot notation
1 parent 5261f02 commit d2e7a0c

File tree

6 files changed

+282
-40
lines changed

6 files changed

+282
-40
lines changed

‎src/backend/parser/parse_expr.c

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -500,11 +500,16 @@ transformIndirection(ParseState *pstate, A_Indirection *ind)
500500
result_basetypid = (result_typid == JSONOID || result_typid == JSONBOID) ?
501501
result_typid : getBaseType(result_typid);
502502

503-
if (result_basetypid == JSONBOID) // TODO
503+
if (result_basetypid == JSONBOID)
504+
{
505+
json_accessor_chain_first = (i == list_head(ind->indirection));
506+
if (lnext(ind->indirection, i) == NULL)
507+
json_accessor_chain_last = true;
504508
newresult = ParseJsonbSimplifiedAccessorObjectField(pstate,
505-
strVal(n),
506-
result,
507-
location, result_basetypid);
509+
strVal(n),
510+
result,
511+
location, result_basetypid, json_accessor_chain_first, json_accessor_chain_last);
512+
}
508513
else if (result_basetypid == JSONOID)
509514
{
510515
json_accessor_chain_first = (i == list_head(ind->indirection));

‎src/backend/parser/parse_func.c

Lines changed: 22 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1952,36 +1952,32 @@ ParseJsonSimplifiedAccessorArrayElement(ParseState *pstate, A_Indices *subscript
19521952
* transformed expression tree. If not, return NULL.
19531953
*/
19541954
Node *
1955-
ParseJsonbSimplifiedAccessorObjectField(ParseState *pstate, const char *funcname,
1956-
Node *first_arg, int location, Oid basetypid)
1955+
ParseJsonbSimplifiedAccessorObjectField(ParseState *pstate, const char *funcname, Node *first_arg, int location,
1956+
Oid basetypid, bool first_op, bool last_op)
19571957
{
1958-
OpExpr *result;
1958+
FuncExpr *result;
19591959
Node *rexpr;
19601960

1961-
result = makeNode(OpExpr);
1962-
result->opresulttype = basetypid;
1963-
switch (basetypid)
1964-
{
1965-
case JSONOID:
1966-
result->opno = OID_JSON_OBJECT_FIELD_OP;
1967-
break;
1968-
case JSONBOID:
1969-
result->opno = OID_JSONB_OBJECT_FIELD_OP;
1970-
break;
1971-
default:
1972-
elog(ERROR, "unsupported type OID: %u", basetypid);
1973-
}
1974-
result->opfuncid = get_opcode(result->opno);
1961+
if (basetypid != JSONBOID)
1962+
elog(ERROR, "unsupported type OID: %u", basetypid);
1963+
19751964
rexpr = (Node *) makeConst(
1976-
TEXTOID,
1977-
-1,
1978-
InvalidOid,
1979-
-1,
1980-
CStringGetTextDatum(funcname),
1981-
false,
1982-
false);
1983-
result->args = list_make2(first_arg, rexpr);
1984-
result->location = location;
1965+
TEXTOID,
1966+
-1,
1967+
InvalidOid,
1968+
-1,
1969+
CStringGetTextDatum(funcname),
1970+
false,
1971+
false);
1972+
result = makeFuncExpr(4100,
1973+
JSONBOID,
1974+
list_make4(first_arg, rexpr,
1975+
makeBoolConst(first_op, false),
1976+
makeBoolConst(last_op, false)),
1977+
InvalidOid,
1978+
InvalidOid,
1979+
COERCE_EXPLICIT_CALL);
1980+
19851981
return (Node *) result;
19861982
}
19871983

‎src/backend/utils/adt/jsonfuncs.c

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -898,6 +898,242 @@ jsonb_object_field(PG_FUNCTION_ARGS)
898898
PG_RETURN_NULL();
899899
}
900900

901+
static void
902+
expand_jbvBinary_in_memory(const JsonbValue *binaryVal, JsonbValue *expandedVal)
903+
{
904+
Assert(binaryVal->type == jbvBinary);
905+
JsonbContainer *container = (JsonbContainer *) binaryVal->val.binary.data;
906+
JsonbIterator *it;
907+
JsonbValue v;
908+
JsonbIteratorToken token;
909+
910+
it = JsonbIteratorInit(container);
911+
token = JsonbIteratorNext(&it, &v, false);
912+
913+
if (token == WJB_BEGIN_OBJECT)
914+
{
915+
/*
916+
* We'll read all keys/values until WJB_END_OBJECT
917+
* and build expandedVal->type = jbvObject.
918+
*/
919+
List *keyvals = NIL;
920+
921+
while ((token = JsonbIteratorNext(&it, &v, false)) != WJB_END_OBJECT)
922+
{
923+
if (token == WJB_KEY)
924+
{
925+
JsonbValue key;
926+
927+
key= v;
928+
token = JsonbIteratorNext(&it, &v, true);
929+
if (token == WJB_VALUE)
930+
{
931+
JsonbValue val;
932+
JsonbValue *objPair;
933+
934+
if (v.type == jbvBinary)
935+
{
936+
/* Recursively expand nested objects/arrays. */
937+
expand_jbvBinary_in_memory(&v, &val);
938+
}
939+
else
940+
{
941+
/* Scalar or already jbvObject/jbvArray. Copy as-is. */
942+
val = v;
943+
}
944+
945+
/* Build a small pair (key, value). We'll store them in a list. */
946+
objPair = palloc(sizeof(JsonbValue) * 2);
947+
objPair[0] = key;
948+
objPair[1] = val;
949+
keyvals = lappend(keyvals, objPair);
950+
}
951+
}
952+
}
953+
954+
/* Now convert our keyvals list into a jbvObject. */
955+
int nPairs = list_length(keyvals);
956+
expandedVal->type = jbvObject;
957+
expandedVal->val.object.nPairs = nPairs;
958+
expandedVal->val.object.pairs = palloc(sizeof(JsonbPair) * nPairs);
959+
960+
{
961+
int i = 0;
962+
ListCell *lc;
963+
foreach(lc, keyvals)
964+
{
965+
JsonbValue *kv = (JsonbValue *) lfirst(lc);
966+
/* kv[0] = key, kv[1] = value */
967+
968+
expandedVal->val.object.pairs[i].key = kv[0];
969+
expandedVal->val.object.pairs[i].value = kv[1];
970+
expandedVal->val.object.pairs[i].order = i;
971+
i++;
972+
}
973+
}
974+
}
975+
else if (token == WJB_BEGIN_ARRAY)
976+
{
977+
/*
978+
* We'll read array elems until WJB_END_ARRAY
979+
* and build expandedVal->type = jbvArray.
980+
*/
981+
List *elems = NIL;
982+
983+
while ((token = JsonbIteratorNext(&it, &v, true)) != WJB_END_ARRAY)
984+
{
985+
if (token == WJB_ELEM)
986+
{
987+
/* If it's jbvBinary, recursively expand again. */
988+
JsonbValue val;
989+
if (v.type == jbvBinary)
990+
{
991+
expand_jbvBinary_in_memory(&v, &val);
992+
}
993+
else
994+
{
995+
val = v;
996+
}
997+
JsonbValue *elemCopy = palloc(sizeof(JsonbValue));
998+
*elemCopy = val;
999+
elems = lappend(elems, elemCopy);
1000+
}
1001+
}
1002+
1003+
expandedVal->type = jbvArray;
1004+
expandedVal->val.array.nElems = list_length(elems);
1005+
expandedVal->val.array.rawScalar = false;
1006+
expandedVal->val.array.elems = palloc(sizeof(JsonbValue) * expandedVal->val.array.nElems);
1007+
1008+
{
1009+
int i = 0;
1010+
ListCell *lc;
1011+
foreach(lc, elems)
1012+
{
1013+
JsonbValue *vptr = (JsonbValue *) lfirst(lc);
1014+
expandedVal->val.array.elems[i++] = *vptr;
1015+
}
1016+
}
1017+
}
1018+
else
1019+
{
1020+
/*
1021+
* Possibly a scalar container (WJB_ELEM or WJB_VALUE with jbvNumeric, jbvString, etc.).
1022+
* If this container truly only has one scalar, token might be WJB_ELEM or similar.
1023+
* For simplicity, let's check tmp.type. If it's jbvString/jbvNumeric, copy it directly.
1024+
*/
1025+
expandedVal->type = v.type;
1026+
expandedVal->val = v.val;
1027+
}
1028+
}
1029+
1030+
static List *
1031+
jsonb_object_field_unwrap_array(JsonbContainer *jb, text *key, bool unwrap_nested)
1032+
{
1033+
JsonbIterator *it;
1034+
JsonbValue v;
1035+
JsonbIteratorToken token;
1036+
List *resultList = NIL;
1037+
int arraySize;
1038+
1039+
it = JsonbIteratorInit(jb);
1040+
token = JsonbIteratorNext(&it, &v, false);
1041+
1042+
Assert(token == WJB_BEGIN_ARRAY);
1043+
arraySize = v.val.array.nElems;
1044+
1045+
/* Unwrap out-most array elements and extract the key value */
1046+
for (int i = 0; i < arraySize; i++) {
1047+
token = JsonbIteratorNext(&it, &v, true);
1048+
if (token == WJB_ELEM && v.type == jbvBinary) {
1049+
JsonbContainer *elemContainer = (JsonbContainer *) v.val.binary.data;
1050+
if (JsonContainerIsObject(elemContainer)) {
1051+
JsonbValue *extractedValue;
1052+
JsonbValue vbuf;
1053+
extractedValue = getKeyJsonValueFromContainer(elemContainer,
1054+
VARDATA_ANY(key),
1055+
VARSIZE_ANY_EXHDR(key),
1056+
&vbuf);
1057+
if (extractedValue != NULL) {
1058+
JsonbValue *copy;
1059+
if (extractedValue->type == jbvBinary) {
1060+
JsonbValue expanded;
1061+
expand_jbvBinary_in_memory(extractedValue, &expanded);
1062+
1063+
copy = palloc(sizeof(expanded));
1064+
*copy = expanded;
1065+
resultList = lappend(resultList, copy);
1066+
} else {
1067+
copy = palloc(sizeof(*extractedValue));
1068+
*copy = *extractedValue;
1069+
resultList = lappend(resultList, copy);
1070+
}
1071+
}
1072+
} else if (unwrap_nested && JsonContainerIsArray(elemContainer)) {
1073+
resultList = jsonb_object_field_unwrap_array(elemContainer, key, false);
1074+
}
1075+
}
1076+
}
1077+
token = JsonbIteratorNext(&it, &v, true);
1078+
1079+
return resultList;
1080+
}
1081+
1082+
Datum
1083+
jsonb_object_field_dot(PG_FUNCTION_ARGS)
1084+
{
1085+
Jsonb *jb = PG_GETARG_JSONB_P(0);
1086+
text *key = PG_GETARG_TEXT_PP(1);
1087+
bool first_op = PG_GETARG_BOOL(2);
1088+
bool last_op = PG_GETARG_BOOL(3);
1089+
1090+
if (JB_ROOT_IS_OBJECT(jb))
1091+
{
1092+
JsonbValue *v;
1093+
JsonbValue vbuf;
1094+
v = getKeyJsonValueFromContainer(&jb->root,
1095+
VARDATA_ANY(key),
1096+
VARSIZE_ANY_EXHDR(key),
1097+
&vbuf);
1098+
1099+
if (v != NULL)
1100+
PG_RETURN_JSONB_P(JsonbValueToJsonb(v));
1101+
}
1102+
else if (JB_ROOT_IS_ARRAY(jb))
1103+
{
1104+
List *resultList;
1105+
JsonbValue resultVal;
1106+
1107+
resultList = jsonb_object_field_unwrap_array(&jb->root, key, !first_op);
1108+
1109+
if (list_length(resultList) == 0)
1110+
PG_RETURN_NULL();
1111+
else if (!last_op || list_length(resultList) > 1)
1112+
{
1113+
/* Conditional wrap result */
1114+
resultVal.type = jbvArray;
1115+
resultVal.val.array.rawScalar = false;
1116+
resultVal.val.array.nElems = list_length(resultList);
1117+
resultVal.val.array.elems = (JsonbValue *)palloc(sizeof(JsonbValue) * list_length(resultList));
1118+
1119+
int idx = 0;
1120+
ListCell *lc;
1121+
JsonbValue *elem;
1122+
foreach(lc, resultList)
1123+
{
1124+
elem = (JsonbValue *)lfirst(lc);
1125+
resultVal.val.array.elems[idx++] = *elem;
1126+
}
1127+
}
1128+
else
1129+
resultVal = *(JsonbValue *)linitial(resultList);
1130+
1131+
PG_RETURN_JSONB_P(JsonbValueToJsonb(&resultVal));
1132+
}
1133+
1134+
PG_RETURN_NULL();
1135+
}
1136+
9011137
Datum
9021138
json_object_field_text(PG_FUNCTION_ARGS)
9031139
{

‎src/include/catalog/pg_proc.dat

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9317,6 +9317,10 @@
93179317
proname => 'json_object_field_dot', prorettype => 'json',
93189318
proargtypes => 'json text bool bool', proargnames => '{from_json, field_name, first_op, last_op}',
93199319
prosrc => 'json_object_field_dot' },
9320+
{ oid => '4100',
9321+
proname => 'jsonb_object_field_dot', prorettype => 'jsonb',
9322+
proargtypes => 'jsonb text bool bool', proargnames => '{from_json, field_name, first_op, last_op}',
9323+
prosrc => 'jsonb_object_field_dot' },
93209324

93219325
# uuid
93229326
{ oid => '2952', descr => 'I/O',

‎src/include/parser/parse_func.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,9 @@ typedef enum
3434
extern Node *ParseFuncOrColumn(ParseState *pstate, List *funcname, List *fargs,
3535
Node *last_srf, FuncCall *fn, bool proc_call,
3636
int location);
37-
extern Node *ParseJsonbSimplifiedAccessorObjectField(ParseState *pstate, const char *funcname,
38-
Node *first_arg, int location, Oid basetypid);
37+
extern Node *
38+
ParseJsonbSimplifiedAccessorObjectField(ParseState *pstate, const char *funcname, Node *first_arg, int location,
39+
Oid basetypid, bool first_op, bool last_op);
3940
extern Node *
4041
ParseJsonSimplifiedAccessorObjectField(ParseState *pstate, const char *funcname, Node *first_arg, int location,
4142
Oid basetypid, bool first_op, bool last_op);

‎src/test/regress/expected/jsonb.out

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5784,7 +5784,7 @@ select * from test_jsonb_dot_notation('{"a": 4, "b": {"c": {"d": [11, 12]}}}'::j
57845784
select * from test_jsonb_dot_notation('[{"x": 42}]'::jsonb, 'x');
57855785
dot_access | expected
57865786
------------+----------
5787-
| 42
5787+
42 | 42
57885788
(1 row)
57895789

57905790
select * from test_jsonb_dot_notation('["x"]'::jsonb, 'x');
@@ -5801,20 +5801,20 @@ select * from test_jsonb_dot_notation('[[{"x": 42}]]'::jsonb, 'x');
58015801

58025802
-- wrap the result into an array on the conditional of more than one matched object keys
58035803
select * from test_jsonb_dot_notation('[{"x": 42}, {"x": {"y": {"z": 12}}}]'::jsonb, 'x');
5804-
dot_access | expected
5805-
------------+------------------------
5806-
| [42, {"y": {"z": 12}}]
5804+
dot_access | expected
5805+
------------------------+------------------------
5806+
[42, {"y": {"z": 12}}] | [42, {"y": {"z": 12}}]
58075807
(1 row)
58085808

58095809
select * from test_jsonb_dot_notation('[{"x": 42}, [{"x": {"y": {"z": 12}}}]]'::jsonb, 'x');
58105810
dot_access | expected
58115811
------------+----------
5812-
| 42
5812+
42 | 42
58135813
(1 row)
58145814

58155815
select * from test_jsonb_dot_notation('[{"x": 42}, {"x": [{"y": 12}, {"y": {"z": 12}}]}]'::jsonb, 'x.y');
5816-
dot_access | expected
5817-
------------+-----------------
5818-
| [12, {"z": 12}]
5816+
dot_access | expected
5817+
-----------------+-----------------
5818+
[12, {"z": 12}] | [12, {"z": 12}]
58195819
(1 row)
58205820

0 commit comments

Comments
 (0)