From 153a23282a7427c33018b7ce54a3bf9a09aa7dc8 Mon Sep 17 00:00:00 2001
From: Dmitrii Dolgov <9erthalion6@gmail.com>
Date: Tue, 4 Aug 2020 17:41:42 +0200
Subject: [PATCH v33 5/5] Filling gaps in jsonb arrays

Appending or prepending array elements on the specified position, gaps
filled with nulls (similar to JavaScript behavior). Originally proposed
by Nikita Glukhov based on polymorphic subscripting patch, but
transformed into an independent change.
---
 src/backend/utils/adt/jsonfuncs.c   | 43 +++++++++++++++++++++++++----
 src/test/regress/expected/jsonb.out | 24 ++++++++++++++++
 src/test/regress/sql/jsonb.sql      | 13 +++++++++
 3 files changed, 74 insertions(+), 6 deletions(-)

diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
index f46a2828b3..6bb81f4c40 100644
--- a/src/backend/utils/adt/jsonfuncs.c
+++ b/src/backend/utils/adt/jsonfuncs.c
@@ -47,6 +47,7 @@
 #define JB_PATH_INSERT_AFTER			0x0010
 #define JB_PATH_CREATE_OR_INSERT \
 	(JB_PATH_INSERT_BEFORE | JB_PATH_INSERT_AFTER | JB_PATH_CREATE)
+#define JB_PATH_FILL_GAPS				0x0020
 
 /* state for json_object_keys */
 typedef struct OkeysState
@@ -1492,10 +1493,8 @@ get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
 static Datum
 jsonb_get_element(Jsonb *jb, Datum *path, int npath, bool *isnull, bool as_text)
 {
-	Jsonb		   *res;
 	JsonbContainer *container = &jb->root;
 	JsonbValue	   *jbvp = NULL;
-	JsonbValue		tv;
 	int				i;
 	bool			have_object = false,
 					have_array = false;
@@ -1657,13 +1656,24 @@ jsonb_set_element(Datum jsonbdatum, Datum *path, int path_len,
 	it = JsonbIteratorInit(&jb->root);
 
 	res = setPath(&it, path, path_nulls, path_len, &state, 0,
-				  newval, JB_PATH_CREATE);
+				  newval, JB_PATH_CREATE | JB_PATH_FILL_GAPS);
 
 	pfree(path_nulls);
 
 	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
 }
 
+static void
+push_null_elements(JsonbParseState **ps, int num)
+{
+		JsonbValue	null;
+
+		null.type = jbvNull;
+
+		while (num-- > 0)
+				pushJsonbValue(ps, WJB_ELEM, &null);
+}
+
 /*
  * Return the text representation of the given JsonbValue.
  */
@@ -4811,6 +4821,13 @@ IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
  * Bits JB_PATH_INSERT_BEFORE and JB_PATH_INSERT_AFTER in op_type
  * behave as JB_PATH_CREATE if new value is inserted in JsonbObject.
  *
+ * If JB_PATH_FILL_GAPS bit is set, this will change an assignment logic in
+ * case if target is an array. The assignment index will not be restricted by
+ * number of elements in the array, and if there are any empty slots between
+ * last element of the array and a new one they will be filled with nulls. If
+ * the index is negative, it still will be considered an an index from the end
+ * of the array.
+ *
  * All path elements before the last must already exist
  * whatever bits in op_type are set, or nothing is done.
  */
@@ -5012,15 +5029,21 @@ setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 			idx = nelems + idx;
 	}
 
-	if (idx > 0 && idx > nelems)
-		idx = nelems;
+	/*
+	 * Filling the gaps means there are no limits on the positive index are
+	 * imposed, we can set any element. Otherwise limit the index by nelems.
+	 */
+	if (!(op_type & JB_PATH_FILL_GAPS))
+	{
+		if (idx > 0 && idx > nelems)
+			idx = nelems;
+	}
 
 	/*
 	 * if we're creating, and idx == INT_MIN, we prepend the new value to the
 	 * array also if the array is empty - in which case we don't really care
 	 * what the idx value is
 	 */
-
 	if ((idx == INT_MIN || nelems == 0) && (level == path_len - 1) &&
 		(op_type & JB_PATH_CREATE_OR_INSERT))
 	{
@@ -5086,10 +5109,18 @@ setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 			if ((op_type & JB_PATH_CREATE_OR_INSERT) && !done &&
 				level == path_len - 1 && i == nelems - 1)
 			{
+				/*
+				 * If asked to fill the gaps, idx could be bigger than nelems,
+				 * so prepend the new element with nulls if that's the case.
+				 */
+				if (op_type & JB_PATH_FILL_GAPS && idx > nelems)
+					push_null_elements(st, idx - nelems);
+
 				(void) pushJsonbValue(st, WJB_ELEM, newval);
 			}
 		}
 	}
+
 }
 
 /*
diff --git a/src/test/regress/expected/jsonb.out b/src/test/regress/expected/jsonb.out
index 04a146a7d0..4d52652688 100644
--- a/src/test/regress/expected/jsonb.out
+++ b/src/test/regress/expected/jsonb.out
@@ -4928,6 +4928,30 @@ select * from test_jsonb_subscript;
   2 | {"a": [1, 2, 3], "key": "value", "another_key": null}
 (2 rows)
 
+-- Fill the gaps logic
+delete from test_jsonb_subscript;
+insert into test_jsonb_subscript values (1, '[0]');
+update test_jsonb_subscript set test_json[5] = 1;
+select * from test_jsonb_subscript;
+ id |           test_json            
+----+--------------------------------
+  1 | [0, null, null, null, null, 1]
+(1 row)
+
+update test_jsonb_subscript set test_json[-4] = 1;
+select * from test_jsonb_subscript;
+ id |          test_json          
+----+-----------------------------
+  1 | [0, null, 1, null, null, 1]
+(1 row)
+
+update test_jsonb_subscript set test_json[-8] = 1;
+select * from test_jsonb_subscript;
+ id |           test_json            
+----+--------------------------------
+  1 | [1, 0, null, 1, null, null, 1]
+(1 row)
+
 -- jsonb to tsvector
 select to_tsvector('{"a": "aaa bbb ddd ccc", "b": ["eee fff ggg"], "c": {"d": "hhh iii"}}'::jsonb);
                                 to_tsvector                                
diff --git a/src/test/regress/sql/jsonb.sql b/src/test/regress/sql/jsonb.sql
index 12541e7e50..584f145bab 100644
--- a/src/test/regress/sql/jsonb.sql
+++ b/src/test/regress/sql/jsonb.sql
@@ -1269,6 +1269,19 @@ update test_jsonb_subscript set test_json[NULL] = 1;
 update test_jsonb_subscript set test_json['another_key'] = NULL;
 select * from test_jsonb_subscript;
 
+-- Fill the gaps logic
+delete from test_jsonb_subscript;
+insert into test_jsonb_subscript values (1, '[0]');
+
+update test_jsonb_subscript set test_json[5] = 1;
+select * from test_jsonb_subscript;
+
+update test_jsonb_subscript set test_json[-4] = 1;
+select * from test_jsonb_subscript;
+
+update test_jsonb_subscript set test_json[-8] = 1;
+select * from test_jsonb_subscript;
+
 -- jsonb to tsvector
 select to_tsvector('{"a": "aaa bbb ddd ccc", "b": ["eee fff ggg"], "c": {"d": "hhh iii"}}'::jsonb);
 
-- 
2.21.0

