Hey Alexander,
> But I think we can support strict inequalities too (e.g.
less and greater without equals). Could you please make it a bool
argument equal_matches?
Sure, an argument is a good idea to keep the code shorter.
> Could you please extract this change as a separate patch.
Done!
On Thu, 30 Jun 2022 at 14:34, Alexander Korotkov <aekorot...@gmail.com>
wrote:
> Hi, Steve!
>
> Thank you for working on this.
>
> On Thu, Jun 30, 2022 at 7:51 PM Steve Chavez <st...@supabase.io> wrote:
> > Currently the red-black tree implementation only has an equality search.
> Other extensions might need other comparison searches, like less-or-equal
> or greater-or-equal. For example OrioleDB defines a greater-or-equal search
> on its postgres fork:
> >
> >
> https://github.com/orioledb/postgres/blob/4c18ae94c20e3e95c374b9947f1ace7d1d6497a1/src/backend/lib/rbtree.c#L164-L186
> >
> > So I thought this might be valuable to have in core. I've added
> less-or-equal and greater-or-equal searches functions plus tests in the
> test_rbtree module. I can add the remaining less/great searches if this is
> deemed worth it.
>
> Looks good. But I think we can support strict inequalities too (e.g.
> less and greater without equals). Could you please make it a bool
> argument equal_matches?
>
> > Also I refactored the sentinel used in the rbtree.c to use C99
> designators.
>
> Could you please extract this change as a separate patch.
>
> ------
> Regards,
> Alexander Korotkov
>
From 8046463e2c98fb4c51ec05740cee130bdf2ad533 Mon Sep 17 00:00:00 2001
From: steve-chavez <stevechavez...@gmail.com>
Date: Tue, 28 Jun 2022 23:46:38 -0500
Subject: [PATCH 1/2] Change rbtree sentinel to C99 designator
---
src/backend/lib/rbtree.c | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/src/backend/lib/rbtree.c b/src/backend/lib/rbtree.c
index a9981dbada..e1cc4110bd 100644
--- a/src/backend/lib/rbtree.c
+++ b/src/backend/lib/rbtree.c
@@ -62,7 +62,7 @@ struct RBTree
static RBTNode sentinel =
{
- RBTBLACK, RBTNIL, RBTNIL, NULL
+ .color = RBTBLACK,.left = RBTNIL,.right = RBTNIL,.parent = NULL
};
--
2.34.1
From 5d49ddd764b05083e1c276df0b50d43ec6896931 Mon Sep 17 00:00:00 2001
From: steve-chavez <stevechavez...@gmail.com>
Date: Wed, 29 Jun 2022 16:28:02 -0500
Subject: [PATCH 2/2] Add rbtree missing comparison searches
* Add find great/less or equal by adding an argument
---
src/backend/lib/rbtree.c | 61 +++++++++++++
src/include/lib/rbtree.h | 2 +
src/test/modules/test_rbtree/test_rbtree.c | 101 +++++++++++++++++++++
3 files changed, 164 insertions(+)
diff --git a/src/backend/lib/rbtree.c b/src/backend/lib/rbtree.c
index e1cc4110bd..92f6b99164 100644
--- a/src/backend/lib/rbtree.c
+++ b/src/backend/lib/rbtree.c
@@ -161,6 +161,67 @@ rbt_find(RBTree *rbt, const RBTNode *data)
return NULL;
}
+/*
+ * rbt_find_great: search for a greater value in an RBTree
+ *
+ * If equal_match is true, this will be a great or equal search.
+ *
+ * Returns the matching tree entry, or NULL if no match is found.
+ */
+RBTNode *
+rbt_find_great(RBTree *rbt, const RBTNode *data, bool equal_match)
+{
+ RBTNode *node = rbt->root;
+ RBTNode *greater = NULL;
+
+ while (node != RBTNIL)
+ {
+ int cmp = rbt->comparator(data, node, rbt->arg);
+
+ if (equal_match && cmp == 0)
+ return node;
+ else if (cmp < 0)
+ {
+ greater = node;
+ node = node->left;
+ }
+ else
+ node = node->right;
+ }
+
+ return greater;
+}
+
+/*
+ * rbt_find_less: search for a lesser value in an RBTree
+ *
+ * If equal_match is true, this will be a less or equal search.
+ *
+ * Returns the matching tree entry, or NULL if no match is found.
+ */
+RBTNode *
+rbt_find_less(RBTree *rbt, const RBTNode *data, bool equal_match)
+{
+ RBTNode *node = rbt->root;
+ RBTNode *lesser = NULL;
+
+ while (node != RBTNIL)
+ {
+ int cmp = rbt->comparator(data, node, rbt->arg);
+
+ if (equal_match && cmp == 0)
+ return node;
+ else if (cmp > 0){
+ lesser = node;
+ node = node->right;
+ }
+ else
+ node = node->left;
+ }
+
+ return lesser;
+}
+
/*
* rbt_leftmost: fetch the leftmost (smallest-valued) tree node.
* Returns NULL if tree is empty.
diff --git a/src/include/lib/rbtree.h b/src/include/lib/rbtree.h
index 580a3e3439..13cf68415e 100644
--- a/src/include/lib/rbtree.h
+++ b/src/include/lib/rbtree.h
@@ -67,6 +67,8 @@ extern RBTree *rbt_create(Size node_size,
void *arg);
extern RBTNode *rbt_find(RBTree *rbt, const RBTNode *data);
+extern RBTNode *rbt_find_great(RBTree *rbt, const RBTNode *data, bool equal_match);
+extern RBTNode *rbt_find_less(RBTree *rbt, const RBTNode *data, bool equal_match);
extern RBTNode *rbt_leftmost(RBTree *rbt);
extern RBTNode *rbt_insert(RBTree *rbt, const RBTNode *data, bool *isNew);
diff --git a/src/test/modules/test_rbtree/test_rbtree.c b/src/test/modules/test_rbtree/test_rbtree.c
index 7cb38759a2..c17d524e01 100644
--- a/src/test/modules/test_rbtree/test_rbtree.c
+++ b/src/test/modules/test_rbtree/test_rbtree.c
@@ -278,6 +278,105 @@ testfind(int size)
}
}
+/*
+ * Check the correctness of the rbt_find_great function by searching for
+ * an equal key and all of the greater keys.
+ */
+static void
+testfindgte(int size)
+{
+ RBTree *tree = create_int_rbtree();
+
+ /*
+ * Using the size as the random key to search wouldn't allow us to get at
+ * least one greater match, so we do size - 1
+ */
+ int randomKey = pg_prng_uint64_range(&pg_global_prng_state, 0, size - 1);
+ IntRBTreeNode searchNode = {.key = randomKey};
+ IntRBTreeNode *eqNode;
+ IntRBTreeNode *gtNode;
+
+ /* Insert natural numbers */
+ rbt_populate(tree, size, 1);
+
+ /*
+ * Since the search key is included in the naturals of the tree, we're
+ * sure to find an equal match
+ */
+ eqNode = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, true);
+
+ if (eqNode == NULL)
+ elog(ERROR, "key was not found");
+
+ /* ensure we find an equal match */
+ if (!(eqNode->key == searchNode.key))
+ elog(ERROR, "rbt_find_great operation in rbtree didn't find an equal key");
+
+ /* Find the rest of the naturals greater than the search key */
+ for (int i = 1; i < size - searchNode.key; i++)
+ {
+ // switch equal_match to false so we only find greater matches now
+ gtNode = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, false);
+
+ /* ensure we find a greater match */
+ if (!(gtNode->key > searchNode.key))
+ elog(ERROR, "rbt_find_great operation in rbtree didn't find a greater key");
+
+ /* delete the previous match to find a greater one */
+ rbt_delete(tree, (RBTNode *) gtNode);
+ }
+}
+
+/*
+ * Check the correctness of the rbt_find_less operation by searching for
+ * an equal key and all of the lesser keys.
+ */
+static void
+testfindlte(int size)
+{
+ RBTree *tree = create_int_rbtree();
+
+ /*
+ * Using 0 as the random key to search wouldn't allow us to get at least
+ * one lesser match, so we start from 1
+ */
+ int randomKey = pg_prng_uint64_range(&pg_global_prng_state, 1, size);
+ IntRBTreeNode searchNode = {.key = randomKey};
+ IntRBTreeNode *eqNode;
+ IntRBTreeNode *ltNode;
+
+ /* Insert natural numbers */
+ rbt_populate(tree, size, 1);
+
+ /*
+ * Since the search key is included in the naturals of the tree, we're
+ * sure to find an equal match
+ */
+ eqNode = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, true);
+
+ if (eqNode == NULL)
+ elog(ERROR, "key was not found");
+
+ /* ensure we find an equal match */
+ if (!(eqNode->key == searchNode.key)){
+ elog(ERROR, "rbt_find_less operation in rbtree didn't find an equal key");
+ }
+
+ /* Find the rest of the naturals lesser than the search key */
+ for (int i = searchNode.key; i > 0; i--)
+ {
+ // switch equal_match to false so we only find lesser matches now
+ ltNode = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, false);
+
+ /* ensure we find a lesser match */
+ if (!(ltNode->key < searchNode.key))
+ elog(ERROR, "rbt_find_less operation in rbtree didn't find a lesser key");
+
+ /* delete the previous match to find a lesser one */
+ rbt_delete(tree, (RBTNode *) ltNode);
+ }
+}
+
/*
* Check the correctness of the rbt_leftmost operation.
* This operation should always return the smallest element of the tree.
@@ -408,6 +507,8 @@ test_rb_tree(PG_FUNCTION_ARGS)
testleftright(size);
testrightleft(size);
testfind(size);
+ testfindgte(size);
+ testfindlte(size);
testleftmost(size);
testdelete(size, Max(size / 10, 1));
PG_RETURN_VOID();
--
2.34.1
