diff --git a/src/backend/access/nbtree/nbtree.c b/src/backend/access/nbtree/nbtree.c index 0fd595a2d4..b055eaff5b 100644 --- a/src/backend/access/nbtree/nbtree.c +++ b/src/backend/access/nbtree/nbtree.c @@ -584,6 +584,10 @@ btmarkpos(PG_FUNCTION_ARGS) else so->markItemIndex = -1; + /* Also record the current positions of any array keys */ + if (so->numArrayKeys) + _bt_mark_array_keys(scan); + PG_RETURN_VOID(); } @@ -596,6 +600,10 @@ btrestrpos(PG_FUNCTION_ARGS) IndexScanDesc scan = (IndexScanDesc) PG_GETARG_POINTER(0); BTScanOpaque so = (BTScanOpaque) scan->opaque; + /* Restore the marked positions of any array keys */ + if (so->numArrayKeys) + _bt_restore_array_keys(scan); + if (so->markItemIndex >= 0) { /* diff --git a/src/backend/access/nbtree/nbtutils.c b/src/backend/access/nbtree/nbtutils.c index 33ad8915f5..50df5e0e5b 100644 --- a/src/backend/access/nbtree/nbtutils.c +++ b/src/backend/access/nbtree/nbtutils.c @@ -595,6 +595,65 @@ _bt_advance_array_keys(IndexScanDesc scan, ScanDirection dir) return found; } +/* + * _bt_mark_array_keys() -- Handle array keys during btmarkpos + * + * Save the current state of the array keys as the "mark" position. + */ +void +_bt_mark_array_keys(IndexScanDesc scan) +{ + BTScanOpaque so = (BTScanOpaque) scan->opaque; + int i; + + for (i = 0; i < so->numArrayKeys; i++) + { + BTArrayKeyInfo *curArrayKey = &so->arrayKeys[i]; + + curArrayKey->mark_elem = curArrayKey->cur_elem; + } +} + +/* + * _bt_restore_array_keys() -- Handle array keys during btrestrpos + * + * Restore the array keys to where they were when the mark was set. + */ +void +_bt_restore_array_keys(IndexScanDesc scan) +{ + BTScanOpaque so = (BTScanOpaque) scan->opaque; + bool changed = false; + int i; + + /* Restore each array key to its position when the mark was set */ + for (i = 0; i < so->numArrayKeys; i++) + { + BTArrayKeyInfo *curArrayKey = &so->arrayKeys[i]; + ScanKey skey = &so->arrayKeyData[curArrayKey->scan_key]; + int mark_elem = curArrayKey->mark_elem; + + if (curArrayKey->cur_elem != mark_elem) + { + curArrayKey->cur_elem = mark_elem; + skey->sk_argument = curArrayKey->elem_values[mark_elem]; + changed = true; + } + } + + /* + * If we changed any keys, we must redo _bt_preprocess_keys. That might + * sound like overkill, but in cases with multiple keys per index column + * it seems necessary to do the full set of pushups. + */ + if (changed) + { + _bt_preprocess_keys(scan); + /* The mark should have been set on a consistent set of keys... */ + Assert(so->qual_ok); + } +} + /* * _bt_preprocess_keys() -- Preprocess scan keys diff --git a/src/include/access/nbtree.h b/src/include/access/nbtree.h index f23ac3559a..d4941e0ff8 100644 --- a/src/include/access/nbtree.h +++ b/src/include/access/nbtree.h @@ -535,6 +535,7 @@ typedef struct BTArrayKeyInfo { int scan_key; /* index of associated key in arrayKeyData */ int cur_elem; /* index of current element in elem_values */ + int mark_elem; /* index of marked element in elem_values */ int num_elems; /* number of elems in current array value */ Datum *elem_values; /* array of num_elems Datums */ } BTArrayKeyInfo; @@ -665,6 +666,8 @@ extern void _bt_freestack(BTStack stack); extern void _bt_preprocess_array_keys(IndexScanDesc scan); extern void _bt_start_array_keys(IndexScanDesc scan, ScanDirection dir); extern bool _bt_advance_array_keys(IndexScanDesc scan, ScanDirection dir); +extern void _bt_mark_array_keys(IndexScanDesc scan); +extern void _bt_restore_array_keys(IndexScanDesc scan); extern void _bt_preprocess_keys(IndexScanDesc scan); extern IndexTuple _bt_checkkeys(IndexScanDesc scan, Page page, OffsetNumber offnum,