Sort jsonb object values by length

Author: Nikita Glukhov

src/backend/utils/adt/jsonb_op.c

@@ -324,10 +324,10 @@ json_hash_internal(FunctionCallInfo fcinfo, bool is_jsonb)
 		{
 				/* Rotation is left to JsonbHashScalarValue() */
 			case WJB_BEGIN_ARRAY:
-				hash ^= JB_FARRAY;
+				hash ^= JB_TARRAY;
 				break;
 			case WJB_BEGIN_OBJECT:
-				hash ^= JB_FOBJECT;
+				hash ^= JB_TOBJECT;
 				break;
 			case WJB_KEY:
 			case WJB_VALUE:
@@ -382,10 +382,10 @@ json_hash_extended_internal(FunctionCallInfo fcinfo, bool is_jsonb)
 		{
 				/* Rotation is left to JsonbHashScalarValueExtended() */
 			case WJB_BEGIN_ARRAY:
-				hash ^= ((uint64) JB_FARRAY) << 32 | JB_FARRAY;
+				hash ^= ((uint64) JB_TARRAY) << 32 | JB_TARRAY;
 				break;
 			case WJB_BEGIN_OBJECT:
-				hash ^= ((uint64) JB_FOBJECT) << 32 | JB_FOBJECT;
+				hash ^= ((uint64) JB_TOBJECT) << 32 | JB_TOBJECT;
 				break;
 			case WJB_KEY:
 			case WJB_VALUE:

src/backend/utils/adt/jsonb_util.c

@@ -31,6 +31,8 @@
 #include "utils/memutils.h"
 #include "utils/varlena.h"
 
+#define JSONB_SORTED_VALUES 1
+
 /*
  * Maximum number of elements in an array (or key/value pairs in an object).
  * This is limited by two things: the size of the JEntry array must fit
@@ -81,6 +83,7 @@ struct JsonbIterator
 	const JEntry *children;		/* JEntrys for child nodes */
 	/* Data proper.  This points to the beginning of the variable-length data */
 	char	   *dataProper;
+	uint32	   *kvMap;
 
 	/* Current item in buffer (up to nElems) */
 	int			curIndex;
@@ -561,6 +564,8 @@ getKeyJsonValueFromContainer(JsonContainer *jsc,
 	const JEntry *children = container->children;
 	int			count = JsonContainerSize(jsc);
 	char	   *baseAddr;
+	bool		sorted_values = (container->header & JB_TMASK) == JB_TOBJECT_SORTED;
+	const uint32 *kvmap;
 	uint32		stopLow,
 				stopHigh;
 
@@ -574,7 +579,16 @@ getKeyJsonValueFromContainer(JsonContainer *jsc,
 	 * Binary search the container. Since we know this is an object, account
 	 * for *Pairs* of Jentrys
 	 */
-	baseAddr = (char *) (children + count * 2);
+	if (sorted_values)
+	{
+		kvmap = &children[count * 2];
+		baseAddr = (char *) &kvmap[count];
+	}
+	else
+	{
+		kvmap = NULL;
+		baseAddr = (char *) (children + count * 2);
+	}
 	stopLow = 0;
 	stopHigh = count;
 	while (stopLow < stopHigh)
@@ -595,7 +609,7 @@ getKeyJsonValueFromContainer(JsonContainer *jsc,
 		if (difference == 0)
 		{
 			/* Found our key, return corresponding value */
-			int			index = stopMiddle + count;
+			int			index = (sorted_values ? kvmap[stopMiddle] : stopMiddle) + count;
 
 			if (!res)
 				res = palloc(sizeof(JsonbValue));
@@ -1197,14 +1211,18 @@ JsonbIteratorNext(JsonIterator **jsit, JsonbValue *val, bool skipNested)
 			(*it)->state = JBI_OBJECT_KEY;
 
 			fillCompressedJsonbValue((*it)->compressed, (*it)->container,
-									 (*it)->curIndex + (*it)->nElems,
-									 (*it)->dataProper, (*it)->curValueOffset,
+									 ((*it)->kvMap ? (*it)->kvMap[(*it)->curIndex] : (*it)->curIndex) + (*it)->nElems,
+									 (*it)->dataProper,
+									 (*it)->kvMap ?
+									 getJsonbOffset((*it)->container, (*it)->kvMap[(*it)->curIndex] + (*it)->nElems) :
+									 (*it)->curValueOffset,
 									 val);
 
 			JBE_ADVANCE_OFFSET((*it)->curDataOffset,
 							   (*it)->children[(*it)->curIndex]);
-			JBE_ADVANCE_OFFSET((*it)->curValueOffset,
-							   (*it)->children[(*it)->curIndex + (*it)->nElems]);
+			if (!(*it)->kvMap)
+				JBE_ADVANCE_OFFSET((*it)->curValueOffset,
+								   (*it)->children[(*it)->curIndex + (*it)->nElems]);
 			(*it)->curIndex++;
 
 			/*
@@ -1256,24 +1274,34 @@ jsonbIteratorInit(JsonContainer *cont, const JsonbContainer *container,
 	/* Array starts just after header */
 	it->children = container->children;
 
-	switch (container->header & (JB_FARRAY | JB_FOBJECT))
+	switch (container->header & JB_TMASK)
 	{
-		case JB_FARRAY:
+		case JB_TSCALAR:
+			it->isScalar = true;
+			/* FALLTHROUGH */
+		case JB_TARRAY:
 			it->dataProper =
 				(char *) it->children + it->nElems * sizeof(JEntry);
-			it->isScalar = (container->header & JB_FSCALAR) != 0;
 			/* This is either a "raw scalar", or an array */
 			Assert(!it->isScalar || it->nElems == 1);
 
 			it->state = JBI_ARRAY_START;
 			break;
 
-		case JB_FOBJECT:
+		case JB_TOBJECT:
+			it->kvMap = NULL;
 			it->dataProper =
 				(char *) it->children + it->nElems * sizeof(JEntry) * 2;
 			it->state = JBI_OBJECT_START;
 			break;
 
+		case JB_TOBJECT_SORTED:
+			it->kvMap = (uint32 *)
+				((char *) it->children + it->nElems * sizeof(JEntry) * 2);
+			it->dataProper = (char *) &it->kvMap[it->nElems];
+			it->state = JBI_OBJECT_START;
+			break;
+
 		default:
 			elog(ERROR, "unknown type of jsonb container");
 	}
@@ -1877,13 +1905,14 @@ convertJsonbArray(StringInfo buffer, JEntry *pheader, const JsonbValue *val, int
 	 * Construct the header Jentry and store it in the beginning of the
 	 * variable-length payload.
 	 */
-	header = nElems | JB_FARRAY;
 	if (val->val.array.rawScalar)
 	{
 		Assert(nElems == 1);
 		Assert(level == 0);
-		header |= JB_FSCALAR;
+		header = nElems | JB_TSCALAR;
 	}
+	else
+		header = nElems | JB_TARRAY;
 
 	appendToBuffer(buffer, (char *) &header, sizeof(uint32));
 
@@ -1941,6 +1970,48 @@ convertJsonbArray(StringInfo buffer, JEntry *pheader, const JsonbValue *val, int
 	*pheader = JENTRY_ISCONTAINER | totallen;
 }
 
+static int
+int_cmp(const void *a, const void *b)
+{
+	int			x = *(const int *) a;
+	int			y = *(const int *) b;
+
+	return x == y ? 0 : x > y ? 1 : -1;
+}
+
+static int
+estimateJsonbValueSize(const JsonbValue *jbv)
+{
+	int			size;
+
+	switch (jbv->type)
+	{
+		case jbvNull:
+		case jbvBool:
+			return 0;
+		case jbvString:
+			return jbv->val.string.len;
+		case jbvNumeric:
+			return VARSIZE_ANY(jbv->val.numeric);
+		case jbvArray:
+			size = offsetof(JsonbContainer, children[jbv->val.array.nElems]);
+			for (int i = 0; i < jbv->val.array.nElems; i++)
+				size += estimateJsonbValueSize(&jbv->val.array.elems[i]);
+			return size;
+		case jbvObject:
+			size = offsetof(JsonbContainer, children[jbv->val.object.nPairs * 2]);
+			for (int i = 0; i < jbv->val.object.nPairs; i++)
+			{
+				size += estimateJsonbValueSize(&jbv->val.object.pairs[i].key);
+				size += estimateJsonbValueSize(&jbv->val.object.pairs[i].value);
+			}
+			return size;
+		default:
+			elog(ERROR, "invalid jsonb value type: %d", jbv->type);
+			return 0;
+	}
+}
+
 static void
 convertJsonbObject(StringInfo buffer, JEntry *pheader, const JsonbValue *val, int level)
 {
@@ -1950,9 +2021,39 @@ convertJsonbObject(StringInfo buffer, JEntry *pheader, const JsonbValue *val, in
 	int			totallen;
 	uint32		header;
 	int			nPairs = val->val.object.nPairs;
+	int			reserved_size;
+	bool		sorted_values = JSONB_SORTED_VALUES && nPairs > 1;
+	struct
+	{
+		int			size;
+		int32		index;
+	}		   *values = sorted_values ? palloc(sizeof(*values) * nPairs) : NULL;
 
 	Assert(nPairs >= 0);
 
+	if (sorted_values)
+	{
+		for (i = 0; i < nPairs; i++)
+		{
+			values[i].index = i;
+			values[i].size = estimateJsonbValueSize(&val->val.object.pairs[i].value);
+		}
+
+		qsort(values, nPairs, sizeof(*values), int_cmp);
+
+		/* check if keys were really moved */
+		sorted_values = false;
+
+		for (i = 0; i < nPairs; i++)
+		{
+			if (values[i].index != i)
+			{
+				sorted_values = true;
+				break;
+			}
+		}
+	}
+
 	/* Remember where in the buffer this object starts. */
 	base_offset = buffer->len;
 
@@ -1963,17 +2064,30 @@ convertJsonbObject(StringInfo buffer, JEntry *pheader, const JsonbValue *val, in
 	 * Construct the header Jentry and store it in the beginning of the
 	 * variable-length payload.
 	 */
-	header = nPairs | JB_FOBJECT;
+	header = nPairs | (sorted_values ? JB_TOBJECT_SORTED : JB_TOBJECT);
 	appendToBuffer(buffer, (char *) &header, sizeof(uint32));
 
 	/* Reserve space for the JEntries of the keys and values. */
-	jentry_offset = reserveFromBuffer(buffer, sizeof(JEntry) * nPairs * 2);
+	reserved_size = sizeof(JEntry) * nPairs * 2;
+	if (sorted_values)
+		reserved_size += sizeof(int32) * nPairs;
+
+	jentry_offset = reserveFromBuffer(buffer, reserved_size);
+
+	/* Write key-value map */
+	if (sorted_values)
+	{
+		for (i = 0; i < nPairs; i++)
+			copyToBuffer(buffer, jentry_offset + sizeof(JEntry) * nPairs * 2 + values[i].index * sizeof(int32),
+						 &i, sizeof(int32));
+	}
 
 	/*
 	 * Iterate over the keys, then over the values, since that is the ordering
 	 * we want in the on-disk representation.
 	 */
 	totallen = 0;
+
 	for (i = 0; i < nPairs; i++)
 	{
 		JsonbPair  *pair = &val->val.object.pairs[i];
@@ -2009,9 +2123,11 @@ convertJsonbObject(StringInfo buffer, JEntry *pheader, const JsonbValue *val, in
 		copyToBuffer(buffer, jentry_offset, (char *) &meta, sizeof(JEntry));
 		jentry_offset += sizeof(JEntry);
 	}
+
 	for (i = 0; i < nPairs; i++)
 	{
-		JsonbPair  *pair = &val->val.object.pairs[i];
+		int			val_index = sorted_values ? values[i].index : i;
+		JsonbPair  *pair = &val->val.object.pairs[val_index];
 		int			len;
 		JEntry		meta;
 
@@ -2045,6 +2161,9 @@ convertJsonbObject(StringInfo buffer, JEntry *pheader, const JsonbValue *val, in
 		jentry_offset += sizeof(JEntry);
 	}
 
+	if (values)
+		pfree(values);
+
 	/* Total data size is everything we've appended to buffer */
 	totallen = buffer->len - base_offset;
 
@@ -2339,16 +2458,35 @@ JsonUniquify(Json *json)
 	return json;
 }
 
+static void
+jsonbInitContainerFromHeader(JsonContainerData *jc, JsonbContainer *jbc)
+{
+	jc->size = jbc->header & JB_CMASK;
+	switch (jbc->header & JB_TMASK)
+	{
+		case JB_TOBJECT:
+		case JB_TOBJECT_SORTED:
+			jc->type = jbvObject;
+			break;
+		case JB_TARRAY:
+			jc->type = jbvArray;
+			break;
+		case JB_TSCALAR:
+			jc->type = jbvArray | jbvScalar;
+			break;
+		default:
+			elog(ERROR, "invalid jsonb container type: %d",
+				 jbc->header & JB_TMASK);
+	}
+}
+
 static void
 jsonbInitContainer(JsonContainerData *jc, JsonbContainer *jbc, int len)
 {
 	jc->ops = &jsonbContainerOps;
 	JsonContainerDataPtr(jc) = jbc;
 	jc->len = len;
-	jc->size = jbc->header & JB_CMASK;
-	jc->type = jbc->header & JB_FOBJECT ? jbvObject :
-			   jbc->header & JB_FSCALAR ? jbvArray | jbvScalar :
-										  jbvArray;
+	jsonbInitContainerFromHeader(jc, jbc);
 }
 
 static void
@@ -2462,10 +2600,7 @@ jsonbzInitContainer(JsonContainerData *jc, CompressedJsonb *cjb, int len)
 
 	jc->ops = &jsonbzContainerOps;
 	jc->len = len;
-	jc->size = jbc->header & JB_CMASK;
-	jc->type = jbc->header & JB_FOBJECT ? jbvObject :
-			   jbc->header & JB_FSCALAR ? jbvArray | jbvScalar :
-										  jbvArray;
+	jsonbInitContainerFromHeader(jc, jbc);
 }
 
 static JsonbContainer *
@@ -2533,7 +2668,9 @@ findValueInCompressedJsonbObject(CompressedJsonb *cjb, const char *keystr, int k
 	JEntry	   *children = container->children;
 	int			count = container->header & JB_CMASK;
 	/* Since this is an object, account for *Pairs* of Jentrys */
-	char	   *base_addr = (char *) (children + count * 2);
+	bool		sorted_values = (container->header & JB_TMASK) == JB_TOBJECT_SORTED;
+	char	   *base_addr = (char *) (children + count * 2) + (sorted_values ? sizeof(uint32) * count : 0);
+	uint32	   *kvmap = sorted_values ? &container->children[count * 2] : NULL;
 	Size		base_offset = base_addr - (char *) jb;
 	uint32		stopLow = 0,
 				stopHigh = count;
@@ -2575,7 +2712,7 @@ findValueInCompressedJsonbObject(CompressedJsonb *cjb, const char *keystr, int k
 		if (difference == 0)
 		{
 			/* Found our key, return corresponding value */
-			int			index = stopMiddle + count;
+			int			index = (sorted_values ? kvmap[stopMiddle] : stopMiddle) + count;
 
 			return fillCompressedJsonbValue(cjb, container, index, base_addr,
 											getJsonbOffset(container, index),