Add functions to binaryheap for efficient key removal and update.

MasahikoSawada · MasahikoSawada · commit b84050864415 · 2024-04-03T10:44:21.000+09:00
Previously, binaryheap didn't support updating a key and removing a node in an efficient way. For example, in order to remove a node from the binaryheap, the caller had to pass the node's position within the array that the binaryheap internally has. Removing a node from the binaryheap is done in O(log n) but searching for the key's position is done in O(n). This commit adds a hash table to binaryheap in order to track the position of each nodes in the binaryheap. That way, by using newly added functions such as binaryheap_update_up() etc., both updating a key and removing a node can be done in O(1) on an average and O(log n) in worst case. This is known as the indexed binary heap. The caller can specify to use the indexed binaryheap by passing indexed = true. The current code does not use the new indexing logic, but it will be used by an upcoming patch. Reviewed-by: Vignesh C, Peter Smith, Hayato Kuroda, Ajin Cherian, Tomas Vondra, Shubham Khanna Discussion: https://postgr.es/m/CAD21AoDffo37RC-eUuyHJKVEr017V2YYDLyn1xF_00ofptWbkg%40mail.gmail.com
diff --git a/src/backend/executor/nodeGatherMerge.c b/src/backend/executor/nodeGatherMerge.c
@@ -422,6 +422,7 @@ gather_merge_setup(GatherMergeState *gm_state)
 	/* Allocate the resources for the merge */
 	gm_state->gm_heap = binaryheap_allocate(nreaders + 1,
 											heap_compare_slots,
+											false,
 											gm_state);
 }
 
diff --git a/src/backend/executor/nodeMergeAppend.c b/src/backend/executor/nodeMergeAppend.c
@@ -125,7 +125,7 @@ ExecInitMergeAppend(MergeAppend *node, EState *estate, int eflags)
 	mergestate->ms_nplans = nplans;
 
 	mergestate->ms_slots = (TupleTableSlot **) palloc0(sizeof(TupleTableSlot *) * nplans);
-	mergestate->ms_heap = binaryheap_allocate(nplans, heap_compare_slots,
+	mergestate->ms_heap = binaryheap_allocate(nplans, heap_compare_slots, false,
 											  mergestate);
 
 	/*
diff --git a/src/backend/postmaster/pgarch.c b/src/backend/postmaster/pgarch.c
@@ -254,7 +254,8 @@ PgArchiverMain(char *startup_data, size_t startup_data_len)
 
 	/* Initialize our max-heap for prioritizing files to archive. */
 	arch_files->arch_heap = binaryheap_allocate(NUM_FILES_PER_DIRECTORY_SCAN,
-												ready_file_comparator, NULL);
+												ready_file_comparator, false,
+												NULL);
 
 	/* Load the archive_library. */
 	LoadArchiveLibrary();
diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
@@ -1294,6 +1294,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	/* allocate heap */
 	state->heap = binaryheap_allocate(state->nr_txns,
 									  ReorderBufferIterCompare,
+									  false,
 									  state);
 
 	/* Now that the state fields are initialized, it is safe to return it. */
diff --git a/src/backend/storage/buffer/bufmgr.c b/src/backend/storage/buffer/bufmgr.c
@@ -3014,6 +3014,7 @@ BufferSync(int flags)
 	 */
 	ts_heap = binaryheap_allocate(num_spaces,
 								  ts_ckpt_progress_comparator,
+								  false,
 								  NULL);
 
 	for (i = 0; i < num_spaces; i++)
diff --git a/src/bin/pg_dump/pg_backup_archiver.c b/src/bin/pg_dump/pg_backup_archiver.c
@@ -4200,6 +4200,7 @@ restore_toc_entries_parallel(ArchiveHandle *AH, ParallelState *pstate,
 	/* Set up ready_heap with enough room for all known TocEntrys */
 	ready_heap = binaryheap_allocate(AH->tocCount,
 									 TocEntrySizeCompareBinaryheap,
+									 false,
 									 NULL);
 
 	/*
diff --git a/src/bin/pg_dump/pg_dump_sort.c b/src/bin/pg_dump/pg_dump_sort.c
@@ -405,7 +405,7 @@ TopoSort(DumpableObject **objs,
 		return true;
 
 	/* Create workspace for the above-described heap */
-	pendingHeap = binaryheap_allocate(numObjs, int_cmp, NULL);
+	pendingHeap = binaryheap_allocate(numObjs, int_cmp, false, NULL);
 
 	/*
 	 * Scan the constraints, and for each item in the input, generate a count
diff --git a/src/common/binaryheap.c b/src/common/binaryheap.c
@@ -22,8 +22,30 @@
 #ifdef FRONTEND
 #include "common/logging.h"
 #endif
+#include "common/hashfn.h"
 #include "lib/binaryheap.h"
 
+/*
+ * Define parameters for hash table code generation. The interface is *also*
+ * declared in binaryheap.h (to generate the types, which are externally
+ * visible).
+ */
+#define SH_PREFIX bh_nodeidx
+#define SH_ELEMENT_TYPE bh_nodeidx_entry
+#define SH_KEY_TYPE bh_node_type
+#define SH_KEY key
+#define SH_HASH_KEY(tb, key) \
+	hash_bytes((const unsigned char *) &key, sizeof(bh_node_type))
+#define SH_EQUAL(tb, a, b) (memcmp(&a, &b, sizeof(bh_node_type)) == 0)
+#define SH_SCOPE extern
+#ifdef FRONTEND
+#define SH_RAW_ALLOCATOR pg_malloc0
+#endif
+#define SH_STORE_HASH
+#define SH_GET_HASH(tb, a) a->hash
+#define SH_DEFINE
+#include "lib/simplehash.h"
+
 static void sift_down(binaryheap *heap, int node_off);
 static void sift_up(binaryheap *heap, int node_off);
 
@@ -34,9 +56,14 @@ static void sift_up(binaryheap *heap, int node_off);
  * of nodes, and with the heap property defined by the given comparator
  * function, which will be invoked with the additional argument specified by
  * 'arg'.
+ *
+ * If 'indexed' is true, we create a hash table to track each node's
+ * index in the heap, enabling to perform some operations such as
+ * binaryheap_remove_node_ptr() etc.
  */
 binaryheap *
-binaryheap_allocate(int num_nodes, binaryheap_comparator compare, void *arg)
+binaryheap_allocate(int num_nodes, binaryheap_comparator compare,
+					bool indexed, void *arg)
 {
 	binaryheap *heap;
 
@@ -48,6 +75,17 @@ binaryheap_allocate(int num_nodes, binaryheap_comparator compare, void *arg)
 	heap->bh_size = 0;
 	heap->bh_has_heap_property = true;
 	heap->bh_nodes = (bh_node_type *) palloc(sizeof(bh_node_type) * num_nodes);
+	heap->bh_nodeidx = NULL;
+
+	if (indexed)
+	{
+#ifdef FRONTEND
+		heap->bh_nodeidx = bh_nodeidx_create(num_nodes, NULL);
+#else
+		heap->bh_nodeidx = bh_nodeidx_create(CurrentMemoryContext, num_nodes,
+											 NULL);
+#endif
+	}
 
 	return heap;
 }
@@ -63,6 +101,9 @@ binaryheap_reset(binaryheap *heap)
 {
 	heap->bh_size = 0;
 	heap->bh_has_heap_property = true;
+
+	if (binaryheap_indexed(heap))
+		bh_nodeidx_reset(heap->bh_nodeidx);
 }
 
 /*
@@ -73,6 +114,9 @@ binaryheap_reset(binaryheap *heap)
 void
 binaryheap_free(binaryheap *heap)
 {
+	if (binaryheap_indexed(heap))
+		bh_nodeidx_destroy(heap->bh_nodeidx);
+
 	pfree(heap->bh_nodes);
 	pfree(heap);
 }
@@ -115,6 +159,67 @@ enlarge_node_array(binaryheap *heap)
 							  sizeof(bh_node_type) * heap->bh_space);
 }
 
+/*
+ * Set the given node at the 'index' and track it if required.
+ *
+ * Return true if the node's index is already tracked.
+ */
+static bool
+set_node(binaryheap *heap, bh_node_type node, int index)
+{
+	bool		found = false;
+
+	/* Set the node to the nodes array */
+	heap->bh_nodes[index] = node;
+
+	if (binaryheap_indexed(heap))
+	{
+		bh_nodeidx_entry *ent;
+
+		/* Keep track of the node index */
+		ent = bh_nodeidx_insert(heap->bh_nodeidx, node, &found);
+		ent->index = index;
+	}
+
+	return found;
+}
+
+/*
+ * Remove the node's index from the hash table if the heap is indexed.
+ */
+static inline void
+delete_nodeidx(binaryheap *heap, bh_node_type node)
+{
+	if (binaryheap_indexed(heap))
+		bh_nodeidx_delete(heap->bh_nodeidx, node);
+}
+
+/*
+ * Replace the existing node at 'idx' with the given 'new_node'. Also
+ * update their positions accordingly. Note that we assume the new_node's
+ * position is already tracked if enabled, i.e. the new_node is already
+ * present in the heap.
+ */
+static void
+replace_node(binaryheap *heap, int index, bh_node_type new_node)
+{
+	bool		found PG_USED_FOR_ASSERTS_ONLY;
+
+	/* Quick return if not necessary to move */
+	if (heap->bh_nodes[index] == new_node)
+		return;
+
+	/* Remove the overwritten node's index */
+	delete_nodeidx(heap, heap->bh_nodes[index]);
+
+	/*
+	 * Replace it with the given new node. This node's position must also be
+	 * tracked as we assume to replace the node with the existing node.
+	 */
+	found = set_node(heap, new_node, index);
+	Assert(!binaryheap_indexed(heap) || found);
+}
+
 /*
  * binaryheap_add_unordered
  *
@@ -131,7 +236,7 @@ binaryheap_add_unordered(binaryheap *heap, bh_node_type d)
 		enlarge_node_array(heap);
 
 	heap->bh_has_heap_property = false;
-	heap->bh_nodes[heap->bh_size] = d;
+	set_node(heap, d, heap->bh_size);
 	heap->bh_size++;
 }
 
@@ -164,7 +269,7 @@ binaryheap_add(binaryheap *heap, bh_node_type d)
 	if (heap->bh_size >= heap->bh_space)
 		enlarge_node_array(heap);
 
-	heap->bh_nodes[heap->bh_size] = d;
+	set_node(heap, d, heap->bh_size);
 	heap->bh_size++;
 	sift_up(heap, heap->bh_size - 1);
 }
@@ -205,14 +310,16 @@ binaryheap_remove_first(binaryheap *heap)
 	if (heap->bh_size == 1)
 	{
 		heap->bh_size--;
+		delete_nodeidx(heap, result);
+
 		return result;
 	}
 
 	/*
 	 * Remove the last node, placing it in the vacated root entry, and sift
 	 * the new root node down to its correct position.
 	 */
-	heap->bh_nodes[0] = heap->bh_nodes[--heap->bh_size];
+	replace_node(heap, 0, heap->bh_nodes[--heap->bh_size]);
 	sift_down(heap, 0);
 
 	return result;
@@ -238,7 +345,7 @@ binaryheap_remove_node(binaryheap *heap, int n)
 						   heap->bh_arg);
 
 	/* remove the last node, placing it in the vacated entry */
-	heap->bh_nodes[n] = heap->bh_nodes[heap->bh_size];
+	replace_node(heap, n, heap->bh_nodes[heap->bh_size]);
 
 	/* sift as needed to preserve the heap property */
 	if (cmp > 0)
@@ -247,6 +354,77 @@ binaryheap_remove_node(binaryheap *heap, int n)
 		sift_down(heap, n);
 }
 
+/*
+ * binaryheap_remove_node_ptr
+ *
+ * Similar to binaryheap_remove_node() but removes the given node. The caller
+ * must ensure that the given node is in the heap. O(log n) worst case.
+ *
+ * This function can be used only if the heap is indexed.
+ */
+void
+binaryheap_remove_node_ptr(binaryheap *heap, bh_node_type d)
+{
+	bh_nodeidx_entry *ent;
+
+	Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
+	Assert(binaryheap_indexed(heap));
+
+	ent = bh_nodeidx_lookup(heap->bh_nodeidx, d);
+	Assert(ent);
+
+	binaryheap_remove_node(heap, ent->index);
+}
+
+/*
+ * Workhorse for binaryheap_update_up and binaryheap_update_down.
+ */
+static void
+resift_node(binaryheap *heap, bh_node_type node, bool sift_dir_up)
+{
+	bh_nodeidx_entry *ent;
+
+	Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
+	Assert(binaryheap_indexed(heap));
+
+	ent = bh_nodeidx_lookup(heap->bh_nodeidx, node);
+	Assert(ent);
+	Assert(ent->index >= 0 && ent->index < heap->bh_size);
+
+	if (sift_dir_up)
+		sift_up(heap, ent->index);
+	else
+		sift_down(heap, ent->index);
+}
+
+/*
+ * binaryheap_update_up
+ *
+ * Sift the given node up after the node's key is updated. The caller must
+ * ensure that the given node is in the heap. O(log n) worst case.
+ *
+ * This function can be used only if the heap is indexed.
+ */
+void
+binaryheap_update_up(binaryheap *heap, bh_node_type d)
+{
+	resift_node(heap, d, true);
+}
+
+/*
+ * binaryheap_update_down
+ *
+ * Sift the given node down after the node's key is updated. The caller must
+ * ensure that the given node is in the heap. O(log n) worst case.
+ *
+ * This function can be used only if the heap is indexed.
+ */
+void
+binaryheap_update_down(binaryheap *heap, bh_node_type d)
+{
+	resift_node(heap, d, false);
+}
+
 /*
  * binaryheap_replace_first
  *
@@ -259,7 +437,7 @@ binaryheap_replace_first(binaryheap *heap, bh_node_type d)
 {
 	Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
 
-	heap->bh_nodes[0] = d;
+	replace_node(heap, 0, d);
 
 	if (heap->bh_size > 1)
 		sift_down(heap, 0);
@@ -301,11 +479,11 @@ sift_up(binaryheap *heap, int node_off)
 		 * Otherwise, swap the parent value with the hole, and go on to check
 		 * the node's new parent.
 		 */
-		heap->bh_nodes[node_off] = parent_val;
+		set_node(heap, parent_val, node_off);
 		node_off = parent_off;
 	}
 	/* Re-fill the hole */
-	heap->bh_nodes[node_off] = node_val;
+	set_node(heap, node_val, node_off);
 }
 
 /*
@@ -360,9 +538,9 @@ sift_down(binaryheap *heap, int node_off)
 		 * Otherwise, swap the hole with the child that violates the heap
 		 * property; then go on to check its children.
 		 */
-		heap->bh_nodes[node_off] = heap->bh_nodes[swap_off];
+		set_node(heap, heap->bh_nodes[swap_off], node_off);
 		node_off = swap_off;
 	}
 	/* Re-fill the hole */
-	heap->bh_nodes[node_off] = node_val;
+	set_node(heap, node_val, node_off);
 }
diff --git a/src/include/lib/binaryheap.h b/src/include/lib/binaryheap.h
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list

Original file line number	Diff line number	Diff line change
`@@ -422,6 +422,7 @@ gather_merge_setup(GatherMergeState *gm_state)`
`422`	`422`	`/* Allocate the resources for the merge */`
`423`	`423`	`gm_state->gm_heap = binaryheap_allocate(nreaders + 1,`
`424`	`424`	`heap_compare_slots,`
	`425`	`+ false,`
`425`	`426`	`gm_state);`
`426`	`427`	`}`
`427`	`428`