Add support of json object/arrays with unknown size

Nikita Glukhov · Nikita Glukhov · commit 1e2c67342c90 · 2021-03-29T20:40:08.000+03:00
diff --git a/src/backend/utils/adt/jsonb_gin.c b/src/backend/utils/adt/jsonb_gin.c
@@ -244,6 +244,9 @@ gin_extract_jsonb(PG_FUNCTION_ARGS)
 		PG_RETURN_POINTER(NULL);
 	}
 
+	if (total < 0)
+		total = 8;
+
 	/* Otherwise, use 2 * root count as initial estimate of result size */
 	init_gin_entries(&entries, 2 * total);
 
@@ -1108,6 +1111,9 @@ gin_extract_jsonb_path(PG_FUNCTION_ARGS)
 		PG_RETURN_POINTER(NULL);
 	}
 
+	if (total < 0)
+		total = 8;
+
 	/* Otherwise, use 2 * root count as initial estimate of result size */
 	init_gin_entries(&entries, 2 * total);
 
diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
@@ -567,6 +567,8 @@ jsonb_object_keys(PG_FUNCTION_ARGS)
 		state = palloc(sizeof(OkeysState));
 
 		state->result_size = JB_ROOT_COUNT(jb);
+		if (state->result_size < 0)
+			state->result_size = 8;
 		state->result_count = 0;
 		state->sent_count = 0;
 		state->result = palloc(state->result_size * sizeof(char *));
@@ -584,6 +586,12 @@ jsonb_object_keys(PG_FUNCTION_ARGS)
 				cstr = palloc(v.val.string.len + 1 * sizeof(char));
 				memcpy(cstr, v.val.string.val, v.val.string.len);
 				cstr[v.val.string.len] = '\0';
+				if (state->result_count >= state->result_size)
+				{
+					state->result_size *= 2;
+					state->result = repalloc(state->result, state->result_size *
+															sizeof(char *));
+				}
 				state->result[state->result_count++] = cstr;
 			}
 		}
@@ -900,7 +908,9 @@ jsonb_array_element(PG_FUNCTION_ARGS)
 	/* Handle negative subscript */
 	if (element < 0)
 	{
-		uint32		nelements = JB_ROOT_COUNT(jb);
+		int		nelements = JB_ROOT_COUNT(jb);
+		if (nelements < 0)
+			nelements = JsonGetArraySize(JsonRoot(jb));
 
 		if (-element > nelements)
 			PG_RETURN_NULL();
@@ -944,6 +954,8 @@ jsonb_array_element_text(PG_FUNCTION_ARGS)
 	if (element < 0)
 	{
 		uint32		nelements = JB_ROOT_COUNT(jb);
+		if (nelements < 0)
+			nelements = JsonGetArraySize(JsonRoot(jb));
 
 		if (-element > nelements)
 			PG_RETURN_NULL();
@@ -1827,7 +1839,8 @@ jsonb_array_length(PG_FUNCTION_ARGS)
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 				 errmsg("cannot get array length of a non-array")));
 
-	PG_RETURN_INT32(JB_ROOT_COUNT(jb));
+	PG_RETURN_INT32(JB_ROOT_COUNT(jb) >= 0 ? JB_ROOT_COUNT(jb)
+										   : JsonGetArraySize(JsonRoot(jb)));
 }
 
 /*
@@ -4536,16 +4549,19 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 	Assert(r == WJB_BEGIN_ARRAY);
 	n = v.val.array.nElems;
 
-	if (idx < 0)
+	if (v.val.array.nElems >= 0)
 	{
-		if (-idx > n)
-			idx = n;
-		else
-			idx = n + idx;
-	}
+		if (idx < 0)
+		{
+			if (-idx > n)
+				idx = n;
+			else
+				idx = n + idx;
+		}
 
-	if (idx >= n)
-		PG_RETURN_JSONB_P(in);
+		if (idx >= n)
+			PG_RETURN_JSONB_P(in);
+	}
 
 	pushJsonbValue(&state, r, NULL);
 
@@ -4562,6 +4578,15 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
+	if (idx < 0 && -idx <= res->val.array.nElems)
+	{
+		idx = res->val.array.nElems + idx;
+		res->val.array.nElems--;
+		memmove(&res->val.array.elems[idx],
+				&res->val.array.elems[idx + 1],
+				sizeof(JsonValue) * (res->val.array.nElems - idx));
+	}
+
 	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
 }
 
@@ -4946,7 +4971,8 @@ setPath(JsonbIterator **it, Datum *path_elems,
 
 			(void) pushJsonbValue(st, r, NULL);
 			setPathArray(it, path_elems, path_nulls, path_len, st, level,
-						 newval, v.val.array.nElems, op_type);
+						 newval, v.val.array.nElems >= 0 ? v.val.array.nElems :
+						 JsonGetArraySize((*it)->container), op_type);
 			r = JsonbIteratorNext(it, &v, false);
 			Assert(r == WJB_END_ARRAY);
 			res = pushJsonbValue(st, r, NULL);
diff --git a/src/backend/utils/adt/jsonpath_exec.c b/src/backend/utils/adt/jsonpath_exec.c
@@ -2158,7 +2158,11 @@ JsonbArraySize(JsonbValue *jb)
 		JsonbContainer *jbc = jb->val.binary.data;
 
 		if (JsonContainerIsArray(jbc) && !JsonContainerIsScalar(jbc))
-			return JsonContainerSize(jbc);
+		{
+			int			size = JsonContainerSize(jbc);
+
+			return size >= 0 ? size : JsonGetArraySize(jbc);
+		}
 	}
 
 	return -1;
