Fix HOT update handling

Previously we weren't handling HOT chain correctly because we used
fetch_row_version function instead of the index_fetch_tuple function
which correctly handles HOT. This commit changes that.

We also fix the snapshot that is used to be aligned with the snapshot
of the scan, which is more correct.

Finally we handle the case that no valid heap tuple is found. While
we don't quite know how that can be hit, it's better to handle that
case.

Author: cevian

pgvectorscale/src/access_method/plain_storage.rs

@@ -11,7 +11,7 @@ use super::{
     storage_common::get_attribute_number_from_index,
 };
 
-use pgrx::PgRelation;
+use pgrx::{PgBox, PgRelation};
 
 use crate::util::{
     page::PageType, table_slot::TableSlot, tape::Tape, HeapPointer, IndexPointer, ItemPointer,
@@ -213,21 +213,31 @@ impl<'a> Storage for PlainStorage<'a> {
     }
     fn get_full_distance_for_resort<S: StatsHeapNodeRead + StatsDistanceComparison>(
         &self,
+        scan: &PgBox<pgrx::pg_sys::IndexScanDescData>,
         qdm: &Self::QueryDistanceMeasure,
         _index_pointer: IndexPointer,
         heap_pointer: HeapPointer,
         meta_page: &MetaPage,
         stats: &mut S,
-    ) -> f32 {
+    ) -> Option<f32> {
         /* Plain storage only needs to resort when the index is using less dimensions than the underlying data. */
         assert!(meta_page.get_num_dimensions() > meta_page.get_num_dimensions_to_index());
 
-        let slot = unsafe { TableSlot::new(self.heap_rel, heap_pointer, stats) };
+        let slot_opt = unsafe {
+            TableSlot::from_index_heap_pointer(self.heap_rel, heap_pointer, scan.xs_snapshot, stats)
+        };
+        let slot = slot_opt?;
         match qdm {
             PlainDistanceMeasure::Full(query) => {
-                let datum = unsafe { slot.get_attribute(self.heap_attr).unwrap() };
+                let datum = unsafe {
+                    slot.get_attribute(self.heap_attr)
+                        .expect("vector attribute should exist in the heap")
+                };
                 let vec = unsafe { PgVector::from_datum(datum, meta_page, false, true) };
-                self.get_distance_function()(vec.to_full_slice(), query.to_full_slice())
+                Some(self.get_distance_function()(
+                    vec.to_full_slice(),
+                    query.to_full_slice(),
+                ))
             }
         }
     }
@@ -378,6 +388,13 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_plain_storage_update_with_null() {
+        crate::access_method::vacuum::tests::test_update_with_null_scaffold(
+            "num_neighbors = 38, storage_layout = plain",
+        );
+    }
+
     #[pg_test]
     unsafe fn test_plain_storage_empty_table_insert() -> spi::Result<()> {
         crate::access_method::build::tests::test_empty_table_insert_scaffold(

pgvectorscale/src/access_method/sbq.rs

@@ -14,7 +14,7 @@ use std::{cell::RefCell, collections::HashMap, iter::once, marker::PhantomData,
 
 use pgrx::{
     pg_sys::{InvalidBlockNumber, InvalidOffsetNumber, BLCKSZ},
-    PgRelation,
+    PgBox, PgRelation,
 };
 use rkyv::{vec::ArchivedVec, Archive, Deserialize, Serialize};
 
@@ -276,12 +276,22 @@ impl SbqSearchDistanceMeasure {
                         &bq_vector[..self.quantized_dimensions],
                     )
                 } else {
-                    debug_assert!(self.quantized_vector.len() == bq_vector.len());
+                    debug_assert!(
+                        self.quantized_vector.len() == bq_vector.len(),
+                        "self.quantized_vector.len()={} bq_vector.len()={}",
+                        self.quantized_vector.len(),
+                        bq_vector.len()
+                    );
                     (self.quantized_vector.as_slice(), bq_vector)
                 }
             }
             GraphNeighborStore::Builder(_b) => {
-                debug_assert!(self.quantized_vector.len() == bq_vector.len());
+                debug_assert!(
+                    self.quantized_vector.len() == bq_vector.len(),
+                    "self.quantized_vector.len()={} bq_vector.len()={}",
+                    self.quantized_vector.len(),
+                    bq_vector.len()
+                );
                 (self.quantized_vector.as_slice(), bq_vector)
             }
         };
@@ -538,14 +548,6 @@ impl<'a> SbqSpeedupStorage<'a> {
             }
         }
     }
-
-    unsafe fn get_heap_table_slot_from_heap_pointer<T: StatsHeapNodeRead>(
-        &self,
-        heap_pointer: HeapPointer,
-        stats: &mut T,
-    ) -> TableSlot {
-        TableSlot::new(self.heap_rel, heap_pointer, stats)
-    }
 }
 
 pub type SbqSpeedupStorageLsnPrivateData = PhantomData<bool>; //no data stored
@@ -635,17 +637,28 @@ impl<'a> Storage for SbqSpeedupStorage<'a> {
 
     fn get_full_distance_for_resort<S: StatsHeapNodeRead + StatsDistanceComparison>(
         &self,
+        scan: &PgBox<pgrx::pg_sys::IndexScanDescData>,
         qdm: &Self::QueryDistanceMeasure,
         _index_pointer: IndexPointer,
         heap_pointer: HeapPointer,
         meta_page: &MetaPage,
         stats: &mut S,
-    ) -> f32 {
-        let slot = unsafe { self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) };
+    ) -> Option<f32> {
+        let slot_opt = unsafe {
+            TableSlot::from_index_heap_pointer(self.heap_rel, heap_pointer, scan.xs_snapshot, stats)
+        };
+
+        let slot = slot_opt?;
 
-        let datum = unsafe { slot.get_attribute(self.heap_attr).unwrap() };
+        let datum = unsafe {
+            slot.get_attribute(self.heap_attr)
+                .expect("vector attribute should exist in the heap")
+        };
         let vec = unsafe { PgVector::from_datum(datum, meta_page, false, true) };
-        self.get_distance_function()(vec.to_full_slice(), qdm.query.to_full_slice())
+        Some(self.get_distance_function()(
+            vec.to_full_slice(),
+            qdm.query.to_full_slice(),
+        ))
     }
 
     fn get_neighbors_with_distances_from_disk<S: StatsNodeRead + StatsDistanceComparison>(
@@ -1045,6 +1058,12 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_bq_compressed_storage_update_with_null() {
+        crate::access_method::vacuum::tests::test_update_with_null_scaffold(
+            "num_neighbors = 38, storage_layout = memory_optimized",
+        );
+    }
     #[pg_test]
     unsafe fn test_bq_compressed_storage_empty_table_insert() -> spi::Result<()> {
         crate::access_method::build::tests::test_empty_table_insert_scaffold(

pgvectorscale/src/access_method/scan.rs

@@ -226,6 +226,7 @@ impl<QDM, PD> TSVResponseIterator<QDM, PD> {
 
     fn next_with_resort<S: Storage<QueryDistanceMeasure = QDM, LSNPrivateData = PD>>(
         &mut self,
+        scan: &PgBox<pg_sys::IndexScanDescData>,
         _index: &PgRelation,
         storage: &S,
     ) -> Option<(HeapPointer, IndexPointer)> {
@@ -243,23 +244,32 @@ impl<QDM, PD> TSVResponseIterator<QDM, PD> {
                 Some((heap_pointer, index_pointer)) => {
                     self.full_distance_comparisons += 1;
                     let distance = storage.get_full_distance_for_resort(
+                        scan,
                         self.lsr.sdm.as_ref().unwrap(),
                         index_pointer,
                         heap_pointer,
                         &self.meta_page,
                         &mut self.lsr.stats,
                     );
 
-                    if self.resort_buffer.len() > 1 {
-                        self.streaming_stats
-                            .update(distance, distance - self.streaming_stats.max_distance);
+                    match distance {
+                        None => {
+                            /* No entry found in heap */
+                            continue;
+                        }
+                        Some(distance) => {
+                            if self.resort_buffer.len() > 1 {
+                                self.streaming_stats
+                                    .update(distance, distance - self.streaming_stats.max_distance);
+                            }
+
+                            self.resort_buffer.push(ResortData {
+                                heap_pointer,
+                                index_pointer,
+                                distance,
+                            });
+                        }
                     }
-
-                    self.resort_buffer.push(ResortData {
-                        heap_pointer,
-                        index_pointer,
-                        distance,
-                    });
                 }
                 None => {
                     break;
@@ -373,7 +383,7 @@ pub extern "C" fn amgettuple(
                 quantizer,
                 &state.meta_page,
             );
-            let next = iter.next_with_resort(&indexrel, &bq);
+            let next = iter.next_with_resort(&scan, &indexrel, &bq);
             get_tuple(state, next, scan)
         }
         StorageState::Plain(iter) => {
@@ -385,7 +395,7 @@ pub extern "C" fn amgettuple(
                 /* no need to resort */
                 iter.next(&storage)
             } else {
-                iter.next_with_resort(&indexrel, &storage)
+                iter.next_with_resort(&scan, &indexrel, &storage)
             };
             get_tuple(state, next, scan)
         }

pgvectorscale/src/access_method/storage.rs

@@ -1,5 +1,7 @@
 use std::pin::Pin;
 
+use pgrx::{pg_sys, PgBox};
+
 use crate::util::{page::PageType, tape::Tape, HeapPointer, IndexPointer, ItemPointer};
 
 use super::{
@@ -74,12 +76,13 @@ pub trait Storage {
 
     fn get_full_distance_for_resort<S: StatsHeapNodeRead + StatsDistanceComparison>(
         &self,
+        scan: &PgBox<pg_sys::IndexScanDescData>,
         query: &Self::QueryDistanceMeasure,
         index_pointer: IndexPointer,
         heap_pointer: HeapPointer,
         meta_page: &MetaPage,
         stats: &mut S,
-    ) -> f32;
+    ) -> Option<f32>;
 
     fn visit_lsn(
         &self,

pgvectorscale/src/access_method/vacuum.rs

@@ -355,6 +355,92 @@ pub mod tests {
         client.execute("DROP TABLE test_vac_full", &[]).unwrap();
     }
 
+    /* This test inserts data with a NULL embedding and then updates the row with the actual embedding.
+      We do this in several transactions to trigger Heap-only-tuples (HOT) updates. This test is specifically structured to ensure to hit HOT code paths.
+    */
+    #[cfg(test)]
+    pub fn test_update_with_null_scaffold(index_options: &str) {
+        //do not run this test in parallel
+        let _lock = VAC_FULL_MUTEX.lock().unwrap();
+        let expected_cnt = 1000;
+
+        //we need to run a few txn in this test which cannot be run from SPI.
+        //so we cannot use the pg_test framework here. Thus we do a bit of
+        //hackery to bring up the test db and then use a client to run queries against it.
+
+        //bring up the test db by running a fake test on a fake fn
+        let mut opts = crate::pg_test::postgresql_conf_options();
+        /* this is a special key that causes all messages to be printed out by the pgrx test framework */
+        opts.push("log_line_prefix='TMSG: '");
+        pgrx_tests::run_test("test_delete_mock_fn", None, opts).unwrap();
+
+        let (mut client, _) = pgrx_tests::client().unwrap();
+
+        client
+            .batch_execute(&format!(
+                "CREATE TABLE test_data_hot_test_1 (
+                id bigint not null primary key generated by default as identity,
+                article text,
+                embedding vector(1024)
+            );
+            
+            insert into test_data_hot_test_1(article) SELECT 'abc' FROM generate_series(1,{expected_cnt});
+            "
+            ))
+            .unwrap();
+
+        client
+            .execute(
+                &format!("do $$
+            declare
+                _id bigint;
+                _article text;
+                _embedding vector(1024);
+            begin
+                loop
+                    select id, article into _id, _article
+                    from test_data_hot_test_1
+                    where embedding is null
+                    for update skip locked
+                    limit 1;
+
+                    if not found then
+                        exit;
+                    end if;
+
+
+                    select  ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector into _embedding from generate_series(1,1024);
+                    update test_data_hot_test_1 set embedding = _embedding where id = _id;
+                    commit;
+                end loop;
+            end;
+            $$;"),
+                &[],
+            )
+            .unwrap();
+
+        client
+            .execute(
+                &format!("CREATE INDEX idx_diskann_bq ON test_data_hot_test_1 USING diskann (embedding) WITH({index_options});"),
+                &[],
+            )
+            .unwrap();
+
+        client.execute(&format!(
+                "with q as
+                (
+                    select  ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector q from generate_series(1,1024)
+                )
+                select article
+                from test_data_hot_test_1
+                order by embedding <=> (select q from q limit 1)
+                limit 15;"), &[]).unwrap();
+
+        client
+            .execute("DROP TABLE test_data_hot_test_1", &[])
+            .unwrap();
+    }
+
     #[pg_test]
     ///This function is only a mock to bring up the test framewokr in test_delete_vacuum
     fn test_delete_mock_fn() -> spi::Result<()> {

pgvectorscale/src/util/table_slot.rs

@@ -1,6 +1,4 @@
-use std::ptr::addr_of_mut;
-
-use pgrx::pg_sys::{Datum, TupleTableSlot};
+use pgrx::pg_sys::{Datum, SnapshotData, TupleTableSlot};
 use pgrx::{pg_sys, PgBox, PgRelation};
 
 use crate::access_method::stats::StatsHeapNodeRead;
@@ -12,31 +10,46 @@ pub struct TableSlot {
 }
 
 impl TableSlot {
-    pub unsafe fn new<S: StatsHeapNodeRead>(
+    pub unsafe fn from_index_heap_pointer<S: StatsHeapNodeRead>(
         heap_rel: &PgRelation,
         heap_pointer: HeapPointer,
+        snapshot: *mut SnapshotData,
         stats: &mut S,
-    ) -> Self {
+    ) -> Option<Self> {
         let slot = PgBox::from_pg(pg_sys::table_slot_create(
             heap_rel.as_ptr(),
             std::ptr::null_mut(),
         ));
 
         let table_am = heap_rel.rd_tableam;
-        let fetch_row_version = (*table_am).tuple_fetch_row_version.unwrap();
         let mut ctid: pg_sys::ItemPointerData = pg_sys::ItemPointerData {
             ..Default::default()
         };
         heap_pointer.to_item_pointer_data(&mut ctid);
-        fetch_row_version(
-            heap_rel.as_ptr(),
+
+        let scan = (*table_am).index_fetch_begin.unwrap()(heap_rel.as_ptr());
+        let mut call_again = false;
+        /* all_dead can be ignored, only used in optimizations we don't implement */
+        let mut all_dead = false;
+        let valid = (*table_am).index_fetch_tuple.unwrap()(
+            scan,
             &mut ctid,
-            addr_of_mut!(pg_sys::SnapshotAnyData),
+            snapshot,
             slot.as_ptr(),
+            &mut call_again,
+            &mut all_dead,
         );
+        (*table_am).index_fetch_end.unwrap()(scan);
+
+        assert!(!call_again, "MVCC snapshots should not require call_again");
         stats.record_heap_read();
 
-        Self { slot }
+        if !valid {
+            /* no valid tuples found in HOT-chain */
+            return None;
+        }
+
+        Some(Self { slot })
     }
 
     pub unsafe fn get_attribute(&self, attribute_number: pg_sys::AttrNumber) -> Option<Datum> {