From 2afc18db559a4d266c2c3754974e5b4f6f1509f2 Mon Sep 17 00:00:00 2001
From: Nikita Glukhov <n.gluhov@postgrespro.ru>
Date: Thu, 21 Feb 2019 19:19:22 +0300
Subject: [PATCH v3 5/5] Optimize resulting JsonbValue allocation in
 findJsonbValueFromContainer() and getIthJsonbValueFromContainer()

---
 src/backend/utils/adt/jsonb_op.c   |  9 ++--
 src/backend/utils/adt/jsonb_util.c | 72 ++++++++++++++++++------------
 src/backend/utils/adt/jsonfuncs.c  | 11 +++--
 src/include/utils/jsonb.h          |  5 ++-
 4 files changed, 59 insertions(+), 38 deletions(-)

diff --git a/src/backend/utils/adt/jsonb_op.c b/src/backend/utils/adt/jsonb_op.c
index a64206eeb1..82c4b0b2cb 100644
--- a/src/backend/utils/adt/jsonb_op.c
+++ b/src/backend/utils/adt/jsonb_op.c
@@ -24,6 +24,7 @@ jsonb_exists(PG_FUNCTION_ARGS)
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	text	   *key = PG_GETARG_TEXT_PP(1);
 	JsonbValue	kval;
+	JsonbValue	vval;
 	JsonbValue *v = NULL;
 
 	/*
@@ -38,7 +39,7 @@ jsonb_exists(PG_FUNCTION_ARGS)
 
 	v = findJsonbValueFromContainer(&jb->root,
 									JB_FOBJECT | JB_FARRAY,
-									&kval);
+									&kval, &vval);
 
 	PG_RETURN_BOOL(v != NULL);
 }
@@ -59,6 +60,7 @@ jsonb_exists_any(PG_FUNCTION_ARGS)
 	for (i = 0; i < elem_count; i++)
 	{
 		JsonbValue	strVal;
+		JsonbValue	valVal;
 
 		if (key_nulls[i])
 			continue;
@@ -69,7 +71,7 @@ jsonb_exists_any(PG_FUNCTION_ARGS)
 
 		if (findJsonbValueFromContainer(&jb->root,
 										JB_FOBJECT | JB_FARRAY,
-										&strVal) != NULL)
+										&strVal, &valVal) != NULL)
 			PG_RETURN_BOOL(true);
 	}
 
@@ -92,6 +94,7 @@ jsonb_exists_all(PG_FUNCTION_ARGS)
 	for (i = 0; i < elem_count; i++)
 	{
 		JsonbValue	strVal;
+		JsonbValue	valVal;
 
 		if (key_nulls[i])
 			continue;
@@ -102,7 +105,7 @@ jsonb_exists_all(PG_FUNCTION_ARGS)
 
 		if (findJsonbValueFromContainer(&jb->root,
 										JB_FOBJECT | JB_FARRAY,
-										&strVal) == NULL)
+										&strVal, &valVal) == NULL)
 			PG_RETURN_BOOL(false);
 	}
 
diff --git a/src/backend/utils/adt/jsonb_util.c b/src/backend/utils/adt/jsonb_util.c
index afde5148aa..95c23bbb81 100644
--- a/src/backend/utils/adt/jsonb_util.c
+++ b/src/backend/utils/adt/jsonb_util.c
@@ -299,6 +299,37 @@ compareJsonbContainers(JsonbContainer *a, JsonbContainer *b)
 	return res;
 }
 
+/* Find scalar element in Jsonb array and return it. */
+static JsonbValue *
+findJsonbElementInArray(JsonbContainer *container, JsonbValue *elem,
+						JsonbValue *res)
+{
+	JsonbValue *result;
+	JEntry	   *children = container->children;
+	int			count = JsonContainerSize(container);
+	char	   *baseAddr = (char *) (children + count);
+	uint32		offset = 0;
+	int			i;
+
+	result = res ? res : palloc(sizeof(*result));
+
+	for (i = 0; i < count; i++)
+	{
+		fillJsonbValue(container, i, baseAddr, offset, result);
+
+		if (elem->type == result->type &&
+			equalsJsonbScalarValue(elem, result))
+			return result;
+
+		JBE_ADVANCE_OFFSET(offset, children[i]);
+	}
+
+	if (!res)
+		pfree(result);
+
+	return NULL;
+}
+
 /* Find value by key in Jsonb object and fetch it into 'res'. */
 JsonbValue *
 findJsonbKeyInObject(JsonbContainer *container, const char *keyVal, int keyLen,
@@ -392,38 +423,19 @@ findJsonbKeyInObject(JsonbContainer *container, const char *keyVal, int keyLen,
  */
 JsonbValue *
 findJsonbValueFromContainer(JsonbContainer *container, uint32 flags,
-							JsonbValue *key)
+							JsonbValue *key, JsonbValue *res)
 {
-	JEntry	   *children = container->children;
 	int			count = JsonContainerSize(container);
 
 	Assert((flags & ~(JB_FARRAY | JB_FOBJECT)) == 0);
 
-	/* Quick out without a palloc cycle if object/array is empty */
+	/* Quick out if object/array is empty */
 	if (count <= 0)
 		return NULL;
 
 	if ((flags & JB_FARRAY) && JsonContainerIsArray(container))
 	{
-		JsonbValue *result = palloc(sizeof(JsonbValue));
-		char	   *base_addr = (char *) (children + count);
-		uint32		offset = 0;
-		int			i;
-
-		for (i = 0; i < count; i++)
-		{
-			fillJsonbValue(container, i, base_addr, offset, result);
-
-			if (key->type == result->type)
-			{
-				if (equalsJsonbScalarValue(key, result))
-					return result;
-			}
-
-			JBE_ADVANCE_OFFSET(offset, children[i]);
-		}
-
-		pfree(result);
+		return findJsonbElementInArray(container, key, res);
 	}
 	else if ((flags & JB_FOBJECT) && JsonContainerIsObject(container))
 	{
@@ -431,7 +443,7 @@ findJsonbValueFromContainer(JsonbContainer *container, uint32 flags,
 		Assert(key->type == jbvString);
 
 		return findJsonbKeyInObject(container, key->val.string.val,
-									key->val.string.len, NULL);
+									key->val.string.len, res);
 	}
 
 	/* Not found */
@@ -444,9 +456,9 @@ findJsonbValueFromContainer(JsonbContainer *container, uint32 flags,
  * Returns palloc()'d copy of the value, or NULL if it does not exist.
  */
 JsonbValue *
-getIthJsonbValueFromContainer(JsonbContainer *container, uint32 i)
+getIthJsonbValueFromContainer(JsonbContainer *container, uint32 i,
+							  JsonbValue *result)
 {
-	JsonbValue *result;
 	char	   *base_addr;
 	uint32		nelements;
 
@@ -459,7 +471,8 @@ getIthJsonbValueFromContainer(JsonbContainer *container, uint32 i)
 	if (i >= nelements)
 		return NULL;
 
-	result = palloc(sizeof(JsonbValue));
+	if (!result)
+		result = palloc(sizeof(JsonbValue));
 
 	fillJsonbValue(container, i, base_addr,
 				   getJsonbOffset(container, i),
@@ -1157,9 +1170,10 @@ JsonbDeepContains(JsonbIterator **val, JsonbIterator **mContained)
 
 			if (IsAJsonbScalar(&vcontained))
 			{
-				if (!findJsonbValueFromContainer((*val)->container,
-												 JB_FARRAY,
-												 &vcontained))
+				JsonbValue	elemBuf;
+
+				if (!findJsonbElementInArray((*val)->container, &vcontained,
+											 &elemBuf))
 					return false;
 			}
 			else
diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
index 3732fb5e6b..71f27a5d16 100644
--- a/src/backend/utils/adt/jsonfuncs.c
+++ b/src/backend/utils/adt/jsonfuncs.c
@@ -827,6 +827,7 @@ jsonb_array_element(PG_FUNCTION_ARGS)
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	int			element = PG_GETARG_INT32(1);
 	JsonbValue *v;
+	JsonbValue	vbuf;
 
 	if (!JB_ROOT_IS_ARRAY(jb))
 		PG_RETURN_NULL();
@@ -842,7 +843,7 @@ jsonb_array_element(PG_FUNCTION_ARGS)
 			element += nelements;
 	}
 
-	v = getIthJsonbValueFromContainer(&jb->root, element);
+	v = getIthJsonbValueFromContainer(&jb->root, element, &vbuf);
 	if (v != NULL)
 		PG_RETURN_JSONB_P(JsonbValueToJsonb(v));
 
@@ -870,6 +871,7 @@ jsonb_array_element_text(PG_FUNCTION_ARGS)
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	int			element = PG_GETARG_INT32(1);
 	JsonbValue *v;
+	JsonbValue	vbuf;
 
 	if (!JB_ROOT_IS_ARRAY(jb))
 		PG_RETURN_NULL();
@@ -885,7 +887,7 @@ jsonb_array_element_text(PG_FUNCTION_ARGS)
 			element += nelements;
 	}
 
-	v = getIthJsonbValueFromContainer(&jb->root, element);
+	v = getIthJsonbValueFromContainer(&jb->root, element, &vbuf);
 
 	if (v != NULL && v->type != jbvNull)
 		PG_RETURN_TEXT_P(JsonbValueAsText(v));
@@ -1373,6 +1375,7 @@ get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
 	bool		have_object = false,
 				have_array = false;
 	JsonbValue *jbvp = NULL;
+	JsonbValue	jbvbuf;
 	JsonbContainer *container;
 
 	/*
@@ -1400,7 +1403,7 @@ get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
 		Assert(JB_ROOT_IS_ARRAY(jb) && JB_ROOT_IS_SCALAR(jb));
 		/* Extract the scalar value, if it is what we'll return */
 		if (npath <= 0)
-			jbvp = getIthJsonbValueFromContainer(container, 0);
+			jbvp = getIthJsonbValueFromContainer(container, 0, &jbvbuf);
 	}
 
 	/*
@@ -1469,7 +1472,7 @@ get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
 					index = nelements + lindex;
 			}
 
-			jbvp = getIthJsonbValueFromContainer(container, index);
+			jbvp = getIthJsonbValueFromContainer(container, index, &jbvbuf);
 		}
 		else
 		{
diff --git a/src/include/utils/jsonb.h b/src/include/utils/jsonb.h
index bddf272b53..096e766927 100644
--- a/src/include/utils/jsonb.h
+++ b/src/include/utils/jsonb.h
@@ -363,12 +363,13 @@ extern uint32 getJsonbLength(const JsonbContainer *jc, int index);
 extern int	compareJsonbContainers(JsonbContainer *a, JsonbContainer *b);
 extern JsonbValue *findJsonbValueFromContainer(JsonbContainer *sheader,
 											   uint32 flags,
-											   JsonbValue *key);
+											   JsonbValue *key,
+											   JsonbValue *res);
 extern JsonbValue *findJsonbKeyInObject(JsonbContainer *container,
 				  						const char *keyVal, int keyLen,
 										JsonbValue *res);
 extern JsonbValue *getIthJsonbValueFromContainer(JsonbContainer *sheader,
-												 uint32 i);
+												 uint32 i, JsonbValue *result);
 extern JsonbValue *pushJsonbValue(JsonbParseState **pstate,
 								  JsonbIteratorToken seq, JsonbValue *jbval);
 extern JsonbIterator *JsonbIteratorInit(JsonbContainer *container);
-- 
2.17.1

