Add custom toaster encoders to jsonb

Author: Nikita Glukhov

src/backend/utils/adt/json_generic.c

@@ -81,6 +81,7 @@ JsonExpand(Json *tmp, Datum value, bool freeValue, JsonContainerOps *ops)
 	json->root.ops = ops;
 	json->root.size = -1;
 	json->root.type = jbvBinary;
+	json->root.toasterid = InvalidOid;
 	json->is_json = false;
 
 	memset(json->root._data, 0, ops->data_size);
@@ -159,6 +160,21 @@ JsonValueToJson(JsonValue *val)
 	}
 }
 
+Datum
+JsonbValueToOrigJsonbDatum(JsonValue *val, Json *orig_json)
+{
+	if (val->type != jbvBinary &&
+		JsonRoot(orig_json)->ops->encode)
+	{
+		void	   *res = JsonRoot(orig_json)->ops->encode(val, &jsonbContainerOps);
+
+		if (res)
+			return PointerGetDatum(res);
+	}
+
+	return JsonValueToJsonbDatum(val);
+}
+
 JsonContainer *
 JsonCopyFlat(JsonContainer *jc)
 {

src/backend/utils/adt/jsonb_util.c

@@ -84,7 +84,6 @@ static void fillJsonbValue(const JsonbContainerHeader *container, int index,
 						   char *base_addr, uint32 offset,
 						   JsonbValue *result);
 static int	compareJsonbScalarValue(const JsonbValue *a, const JsonbValue *b);
-static void *convertToJsonb(const JsonbValue *val, JsonValueEncoder encoder);
 static void convertJsonbValue(StringInfo buffer, JEntry *header, const JsonbValue *val, int level);
 static void convertJsonbArray(StringInfo buffer, JEntry *header, const JsonbValue *val, int level);
 static void convertJsonbObject(StringInfo buffer, JEntry *header, const JsonbValue *val, int level);
@@ -121,6 +120,18 @@ JsonContainerFlatten(JsonContainer *jc, JsonValueEncoder encoder,
 {
 	JsonbValue jbv;
 
+	if (jc->ops->encode)
+	{
+		JsonbValue bin;
+		void	   *res;
+
+		JsonValueInitBinary(&bin, jc);
+		res = jc->ops->encode(&bin, ops);
+
+		if (res)
+			return res;
+	}
+
 	if (jc->ops == ops)
 	{
 		int			size = jc->len;
@@ -137,7 +148,7 @@ JsonContainerFlatten(JsonContainer *jc, JsonValueEncoder encoder,
 	else
 		binary = JsonValueInitBinary(&jbv, jc);
 
-	return convertToJsonb(binary, encoder);
+	return JsonEncode(binary, encoder, NULL);
 }
 
 /*
@@ -169,7 +180,7 @@ JsonValueFlatten(const JsonValue *val, JsonValueEncoder encoder,
 		Assert(val->type == jbvObject || val->type == jbvArray);
 	}
 
-	return convertToJsonb(val, encoder);
+	return JsonEncode(val, encoder, NULL);
 }
 
 /*
@@ -1714,41 +1725,36 @@ padBufferToInt(StringInfo buffer)
 }
 
 void
-JsonbEncode(StringInfoData *buffer, const JsonbValue *val)
+JsonbEncode(StringInfoData *buffer, const JsonbValue *val, void *cxt)
 {
 	JEntry	jentry;
 
+	/* Make room for the varlena header */
+	reserveFromBuffer(buffer, VARHDRSZ);
 	convertJsonbValue(buffer, &jentry, val, 0);
+	SET_VARSIZE(buffer->data, buffer->len);
 }
 
 /*
  * Given a JsonbValue, convert to Jsonb. The result is palloc'd.
  */
-static void *
-convertToJsonb(const JsonbValue *val, JsonValueEncoder encoder)
+void *
+JsonEncode(const JsonbValue *val, JsonValueEncoder encoder, void *cxt)
 {
 	StringInfoData	buffer;
-	void		   *res;
 
 	/* Allocate an output buffer. It will be enlarged as needed */
 	initStringInfo(&buffer);
 
-	/* Make room for the varlena header */
-	reserveFromBuffer(&buffer, VARHDRSZ);
-
-	(*encoder)(&buffer, val);
+	(*encoder)(&buffer, val, cxt);
 
 	/*
 	 * Note: the JEntry of the root is discarded. Therefore the root
 	 * JsonbContainer struct must contain enough information to tell what kind
 	 * of value it is.
 	 */
 
-	res = (void *) buffer.data;
-
-	SET_VARSIZE(res, buffer.len);
-
-	return res;
+	return buffer.data;
 }
 
 /*
@@ -2204,6 +2210,7 @@ jsonbInitContainer(JsonContainerData *jc, JsonbContainerHeader *jbc, int len)
 	jc->ops = &jsonbContainerOps;
 	JsonContainerDataPtr(jc) = jbc;
 	jc->len = len;
+	jc->toasterid = InvalidOid;
 	jc->size = jbc->header & JBC_CMASK;
 	jc->type = jbc->header & JBC_FOBJECT ? jbvObject :
 			   jbc->header & JBC_FSCALAR ? jbvArray | jbvScalar :
@@ -2230,4 +2237,5 @@ jsonbContainerOps =
 	JsonbToCStringRaw,
 	JsonCopyFlat,
 	NULL,
+	NULL
 };

src/backend/utils/adt/jsonfuncs.c

@@ -1631,7 +1631,7 @@ jsonb_set_element(Jsonb *jb, Datum *path, int path_len,
 
 	pfree(path_nulls);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	PG_RETURN_DATUM(JsonbValueToOrigJsonbDatum(res, jb));
 }
 
 static void
@@ -4209,7 +4209,7 @@ jsonb_strip_nulls(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	PG_RETURN_DATUM(JsonbValueToOrigJsonbDatum(res, jb));
 }
 
 /*
@@ -4264,7 +4264,7 @@ jsonb_concat(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	PG_RETURN_DATUM(JsonbValueToOrigJsonbDatum(res, jb1));
 }
 
 
@@ -4318,7 +4318,7 @@ jsonb_delete(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	PG_RETURN_DATUM(JsonbValueToOrigJsonbDatum(res, in));
 }
 
 /*
@@ -4403,7 +4403,7 @@ jsonb_delete_array(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	PG_RETURN_DATUM(JsonbValueToOrigJsonbDatum(res, in));
 }
 
 /*
@@ -4471,7 +4471,16 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	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_DATUM(JsonbValueToOrigJsonbDatum(res, in));
 }
 
 /*
@@ -4519,7 +4528,7 @@ jsonb_set(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	PG_RETURN_DATUM(JsonbValueToOrigJsonbDatum(res, in));
 }
 
 
@@ -4630,7 +4639,7 @@ jsonb_delete_path(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	PG_RETURN_DATUM(JsonbValueToOrigJsonbDatum(res, in));
 }
 
 /*
@@ -4675,7 +4684,7 @@ jsonb_insert(PG_FUNCTION_ARGS)
 
 	Assert(res != NULL);
 
-	PG_RETURN_JSONB_VALUE_P(res);
+	PG_RETURN_DATUM(JsonbValueToOrigJsonbDatum(res, in));
 }
 
 /*

src/include/utils/json_generic.h

@@ -29,6 +29,7 @@ typedef struct JsonContainerData
 	int					len;
 	int					size;
 	JsonbValueType		type;
+	Oid					toasterid;
 	void			   *_data[FLEXIBLE_ARRAY_MEMBER];
 } JsonContainerData;
 
@@ -63,6 +64,7 @@ struct JsonContainerOps
 								int estimated_len);
 	JsonContainer  *(*copy)(JsonContainer *jc);
 	void			(*free)(JsonContainer *jc);
+	void		   *(*encode)(JsonValue *jc, JsonContainerOps *ops);
 };
 
 typedef struct CompressedObject
@@ -248,6 +250,7 @@ JsonValueInitBinary(JsonValue *val, JsonContainer *cont)
 }
 
 extern Json *JsonValueToJson(JsonValue *val);
+extern Datum JsonbValueToOrigJsonbDatum(JsonValue *val, Json *origjs);
 extern JsonValue *JsonToJsonValue(Json *json, JsonValue *jv);
 extern JsonValue *JsonValueUnpackBinary(const JsonValue *jbv);
 extern JsonValue *JsonValueCopy(JsonValue *res, const JsonValue *val);
@@ -285,21 +288,22 @@ extern char *JsonbToCStringIndent(StringInfo out, JsonContainer *in,
 extern bool JsonValueScalarEquals(const JsonValue *aScalar,
 								  const JsonValue *bScalar);
 
-typedef void (*JsonValueEncoder)(StringInfo, const JsonValue *);
+typedef void (*JsonValueEncoder)(StringInfo, const JsonValue *, void *cxt);
 
 extern void *JsonContainerFlatten(JsonContainer *jc, JsonValueEncoder encoder,
 								  JsonContainerOps *ops, const JsonValue *binary);
 
 extern void *JsonValueFlatten(const JsonValue *val, JsonValueEncoder encoder,
 							  JsonContainerOps *ops);
+extern void *JsonEncode(const JsonbValue *val, JsonValueEncoder encoder, void *cxt);
 
 static inline void *
 JsonFlatten(Json *json, JsonValueEncoder encoder, JsonContainerOps *ops)
 {
 	return JsonContainerFlatten(JsonRoot(json), encoder, ops, NULL);
 }
 
-extern void JsonbEncode(StringInfo, const JsonValue *);
+extern void JsonbEncode(StringInfo, const JsonValue *, void *cxt);
 
 #define JsonValueToJsonb(val) \
 		JsonValueFlatten(val, JsonbEncode, &jsonbContainerOps)