Variable-length jsonb KV map entries

Author: Nikita Glukhov

src/backend/utils/adt/jsonb_util.c

@@ -66,6 +66,27 @@ typedef struct CompressedJsonb
 	int			offset;
 } CompressedJsonb;
 
+typedef struct JsonbKVMap
+{
+	union
+	{
+		const uint8 *entries1;
+		const uint16 *entries2;
+		const int32 *entries4;
+		const void *entries;
+	}			map;
+	int			entry_size;
+} JsonbKVMap;
+
+#define JSONB_KVMAP_ENTRY_SIZE(nPairs) \
+	((nPairs) < 256 ? 1 : (nPairs) < 65536 ? 2 : 4)
+
+#define JSONB_KVMAP_ENTRY(kvmap, index) \
+	(!(kvmap)->entry_size ? (index) : \
+	 (kvmap)->entry_size == 1 ? (int32) (kvmap)->map.entries1[index] : \
+	 (kvmap)->entry_size == 2 ? (int32) (kvmap)->map.entries2[index] : \
+	 (kvmap)->map.entries4[index])
+
 struct JsonbIterator
 {
 	JsonIterator	ji;
@@ -81,7 +102,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;
+	JsonbKVMap	kvmap;
 
 	/* Current item in buffer (up to nElems) */
 	int			curIndex;
@@ -549,6 +570,24 @@ JsonFindValueInContainer(JsonContainer *json, uint32 flags, JsonValue *key)
 	return NULL;
 }
 
+static void *
+initKVMap(JsonbKVMap *kvmap, void *pentries, int field_count, bool sorted)
+{
+	if (sorted)
+	{
+		kvmap->map.entries = pentries;
+		kvmap->entry_size = JSONB_KVMAP_ENTRY_SIZE(field_count);
+
+		return (char *) pentries + INTALIGN(field_count * kvmap->entry_size);
+	}
+	else
+	{
+		kvmap->entry_size = 0;
+
+		return pentries;
+	}
+}
+
 /*
  * Find value by key in Jsonb object and fetch it into 'res', which is also
  * returned.
@@ -562,9 +601,9 @@ getKeyJsonValueFromContainer(JsonContainer *jsc,
 	const JsonbContainer *container = JsonContainerDataPtr(jsc);
 	const JEntry *children = container->children;
 	int			count = JsonContainerSize(jsc);
-	char	   *baseAddr;
+	char	   *baseAddr = (char *) (children + count * 2);
 	bool		sorted_values = (container->header & JB_TMASK) == JB_TOBJECT_SORTED;
-	const uint32 *kvmap;
+	JsonbKVMap	kvmap;
 	uint32		stopLow,
 				stopHigh;
 
@@ -578,16 +617,8 @@ getKeyJsonValueFromContainer(JsonContainer *jsc,
 	 * Binary search the container. Since we know this is an object, account
 	 * for *Pairs* of Jentrys
 	 */
-	if (sorted_values)
-	{
-		kvmap = &children[count * 2];
-		baseAddr = (char *) &kvmap[count];
-	}
-	else
-	{
-		kvmap = NULL;
-		baseAddr = (char *) (children + count * 2);
-	}
+	baseAddr = initKVMap(&kvmap, baseAddr, count, sorted_values);
+
 	stopLow = 0;
 	stopHigh = count;
 	while (stopLow < stopHigh)
@@ -608,7 +639,7 @@ getKeyJsonValueFromContainer(JsonContainer *jsc,
 		if (difference == 0)
 		{
 			/* Found our key, return corresponding value */
-			int			index = (sorted_values ? kvmap[stopMiddle] : stopMiddle) + count;
+			int			index = JSONB_KVMAP_ENTRY(&kvmap, stopMiddle) + count;
 
 			if (!res)
 				res = palloc(sizeof(JsonbValue));
@@ -1097,6 +1128,7 @@ JsonbIteratorToken
 JsonbIteratorNext(JsonIterator **jsit, JsonbValue *val, bool skipNested)
 {
 	JsonbIterator **it = (JsonbIterator **) jsit;
+	int			entry_index;
 
 	if (*it == NULL)
 		return WJB_DONE;
@@ -1209,17 +1241,19 @@ JsonbIteratorNext(JsonIterator **jsit, JsonbValue *val, bool skipNested)
 			/* Set state for next call */
 			(*it)->state = JBI_OBJECT_KEY;
 
+			entry_index = JSONB_KVMAP_ENTRY(&(*it)->kvmap, (*it)->curIndex) + (*it)->nElems;
+
 			fillCompressedJsonbValue((*it)->compressed, (*it)->container,
-									 ((*it)->kvMap ? (*it)->kvMap[(*it)->curIndex] : (*it)->curIndex) + (*it)->nElems,
+									 entry_index,
 									 (*it)->dataProper,
-									 (*it)->kvMap ?
-									 getJsonbOffset((*it)->container, (*it)->kvMap[(*it)->curIndex] + (*it)->nElems) :
+									 (*it)->kvmap.entry_size ?
+									 getJsonbOffset((*it)->container, entry_index) :
 									 (*it)->curValueOffset,
 									 val);
 
 			JBE_ADVANCE_OFFSET((*it)->curDataOffset,
 							   (*it)->children[(*it)->curIndex]);
-			if (!(*it)->kvMap)
+			if (!(*it)->kvmap.entry_size)
 				JBE_ADVANCE_OFFSET((*it)->curValueOffset,
 								   (*it)->children[(*it)->curIndex + (*it)->nElems]);
 			(*it)->curIndex++;
@@ -1261,6 +1295,7 @@ jsonbIteratorInit(JsonContainer *cont, const JsonbContainer *container,
 				  struct CompressedJsonb *cjb)
 {
 	JsonbIterator *it;
+	int			type = container->header & JB_TMASK;
 
 	it = palloc0(sizeof(JsonbIterator));
 	it->ji.container = cont;
@@ -1273,7 +1308,7 @@ jsonbIteratorInit(JsonContainer *cont, const JsonbContainer *container,
 	/* Array starts just after header */
 	it->children = container->children;
 
-	switch (container->header & JB_TMASK)
+	switch (type)
 	{
 		case JB_TSCALAR:
 			it->isScalar = true;
@@ -1288,16 +1323,12 @@ jsonbIteratorInit(JsonContainer *cont, const JsonbContainer *container,
 			break;
 
 		case JB_TOBJECT:
-			it->kvMap = NULL;
+		case JB_TOBJECT_SORTED:
 			it->dataProper =
 				(char *) it->children + it->nElems * sizeof(JEntry) * 2;
-			it->state = JBI_OBJECT_START;
-			break;
+			it->dataProper = initKVMap(&it->kvmap, it->dataProper, it->nElems,
+									   type == JB_TOBJECT_SORTED);
 
-		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;
 
@@ -2021,6 +2052,7 @@ convertJsonbObject(StringInfo buffer, JEntry *pheader, const JsonbValue *val, in
 	uint32		header;
 	int			nPairs = val->val.object.nPairs;
 	int			reserved_size;
+	int			kvmap_entry_size;
 	bool		sorted_values = jsonb_sort_field_values && nPairs > 1;
 	struct
 	{
@@ -2047,6 +2079,7 @@ convertJsonbObject(StringInfo buffer, JEntry *pheader, const JsonbValue *val, in
 		{
 			if (values[i].index != i)
 			{
+				kvmap_entry_size = JSONB_KVMAP_ENTRY_SIZE(nPairs);
 				sorted_values = true;
 				break;
 			}
@@ -2069,16 +2102,45 @@ convertJsonbObject(StringInfo buffer, JEntry *pheader, const JsonbValue *val, in
 	/* Reserve space for the JEntries of the keys and values. */
 	reserved_size = sizeof(JEntry) * nPairs * 2;
 	if (sorted_values)
-		reserved_size += sizeof(int32) * nPairs;
+		reserved_size += INTALIGN(kvmap_entry_size * nPairs);
 
 	jentry_offset = reserveFromBuffer(buffer, reserved_size);
 
 	/* Write key-value map */
 	if (sorted_values)
 	{
+		int			kvmap_offset = jentry_offset + sizeof(JEntry) * nPairs * 2;
+
 		for (i = 0; i < nPairs; i++)
-			copyToBuffer(buffer, jentry_offset + sizeof(JEntry) * nPairs * 2 + values[i].index * sizeof(int32),
-						 &i, sizeof(int32));
+		{
+			uint8		entry1;
+			uint16		entry2;
+			uint32		entry4;
+			void	   *pentry;
+
+			if (kvmap_entry_size == 1)
+			{
+				entry1 = (uint8) i;
+				pentry = &entry1;
+			}
+			else if (kvmap_entry_size == 2)
+			{
+				entry2 = (uint16) i;
+				pentry = &entry2;
+			}
+			else
+			{
+				entry4 = (int32) i;
+				pentry = &entry4;
+			}
+
+			copyToBuffer(buffer, kvmap_offset + values[i].index * kvmap_entry_size,
+						 pentry, kvmap_entry_size);
+		}
+
+		if ((kvmap_entry_size * nPairs) % ALIGNOF_INT)
+			memset(buffer->data + kvmap_offset + kvmap_entry_size * nPairs, 0,
+				   ALIGNOF_INT - (kvmap_entry_size * nPairs) % ALIGNOF_INT);
 	}
 
 	/*
@@ -2670,9 +2732,9 @@ findValueInCompressedJsonbObject(CompressedJsonb *cjb, const char *keystr, int k
 	int			count = container->header & JB_CMASK;
 	/* Since this is an object, account for *Pairs* of Jentrys */
 	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;
+	char	   *base_addr = (char *) (children + count * 2);
+	JsonbKVMap	kvmap;
+	Size		base_offset;
 	uint32		stopLow = 0,
 				stopHigh = count;
 
@@ -2682,6 +2744,9 @@ findValueInCompressedJsonbObject(CompressedJsonb *cjb, const char *keystr, int k
 	if (count <= 0)
 		return NULL;
 
+	base_addr = initKVMap(&kvmap, base_addr, count, sorted_values);
+	base_offset = base_addr - (char *) jb;
+
 	key.type = jbvString;
 	key.val.string.val = keystr;
 	key.val.string.len = keylen;
@@ -2713,7 +2778,7 @@ findValueInCompressedJsonbObject(CompressedJsonb *cjb, const char *keystr, int k
 		if (difference == 0)
 		{
 			/* Found our key, return corresponding value */
-			int			index = (sorted_values ? kvmap[stopMiddle] : stopMiddle) + count;
+			int			index = JSONB_KVMAP_ENTRY(&kvmap, stopMiddle) + count;
 
 			return fillCompressedJsonbValue(cjb, container, index, base_addr,
 											getJsonbOffset(container, index),