Remove io_optimized storage type, which does not perform well and is …

…unused
timescale · Jan 21, 2025 · f7b4a86 · f7b4a86
1 parent 6f8c9d1
commit f7b4a86
Show file tree

Hide file tree

Showing 7 changed files with 296 additions and 37 deletions.
diff --git a/pgvectorscale/src/access_method/build.rs b/pgvectorscale/src/access_method/build.rs
@@ -198,7 +198,7 @@ unsafe fn aminsert_internal(
                 &mut stats,
             );
         }
-        StorageType::SbqCompression => {
+        StorageType::SbqSpeedup | StorageType::SbqCompression => {
             let bq = SbqSpeedupStorage::load_for_insert(
                 &heap_relation,
                 &index_relation,
@@ -282,7 +282,7 @@ fn do_heap_scan<'a>(
 
             finalize_index_build(&mut plain, &mut bs, write_stats)
         }
-        StorageType::SbqCompression => {
+        StorageType::SbqSpeedup | StorageType::SbqCompression => {
             let mut bq =
                 SbqSpeedupStorage::new_for_build(index_relation, heap_relation, &meta_page);
 

diff --git a/pgvectorscale/src/access_method/meta_page.rs b/pgvectorscale/src/access_method/meta_page.rs
@@ -13,6 +13,7 @@ use super::options::{
     NUM_DIMENSIONS_DEFAULT_SENTINEL, NUM_NEIGHBORS_DEFAULT_SENTINEL,
     SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL,
 };
+use super::sbq::SbqNode;
 use super::stats::StatsNodeModify;
 use super::storage::StorageType;
 
@@ -133,6 +134,16 @@ impl MetaPage {
         self.bq_num_bits_per_dimension
     }
 
+    pub fn get_num_dimensions_for_neighbors(&self) -> u32 {
+        match StorageType::from_u8(self.storage_type) {
+            StorageType::Plain => {
+                error!("get_num_dimensions_for_neighbors should not be called for Plain storage")
+            }
+            StorageType::SbqSpeedup => self.num_dimensions_to_index,
+            StorageType::SbqCompression => 0,
+        }
+    }
+
     /// Maximum number of neigbors per node. Given that we pre-allocate
     /// these many slots for each node, this cannot change after the graph is built.
     pub fn get_num_neighbors(&self) -> u32 {
@@ -178,15 +189,24 @@ impl MetaPage {
 
         match self.get_storage_type() {
             StorageType::Plain => None,
-            StorageType::SbqCompression => Some(self.quantizer_metadata),
+            StorageType::SbqSpeedup | StorageType::SbqCompression => Some(self.quantizer_metadata),
         }
     }
 
-    fn calculate_num_neighbors(opt: &PgBox<TSVIndexOptions>) -> u32 {
+    fn calculate_num_neighbors(
+        num_dimensions: u32,
+        num_bits_per_dimension: u8,
+        opt: &PgBox<TSVIndexOptions>,
+    ) -> u32 {
         let num_neighbors = (*opt).get_num_neighbors();
         if num_neighbors == NUM_NEIGHBORS_DEFAULT_SENTINEL {
             match (*opt).get_storage_type() {
                 StorageType::Plain => 50,
+                StorageType::SbqSpeedup => SbqNode::get_default_num_neighbors(
+                    num_dimensions as usize,
+                    num_dimensions as usize,
+                    num_bits_per_dimension,
+                ) as u32,
                 StorageType::SbqCompression => 50,
             }
         } else {
@@ -242,7 +262,11 @@ impl MetaPage {
             num_dimensions,
             num_dimensions_to_index,
             storage_type: (*opt).get_storage_type() as u8,
-            num_neighbors: Self::calculate_num_neighbors(&opt),
+            num_neighbors: Self::calculate_num_neighbors(
+                num_dimensions,
+                bq_num_bits_per_dimension,
+                &opt,
+            ),
             bq_num_bits_per_dimension,
             search_list_size: opt.search_list_size,
             max_alpha: opt.max_alpha,

diff --git a/pgvectorscale/src/access_method/options.rs b/pgvectorscale/src/access_method/options.rs
@@ -298,6 +298,28 @@ mod tests {
         Ok(())
     }
 
+    #[pg_test]
+    unsafe fn test_index_options_bq() -> spi::Result<()> {
+        Spi::run(
+            "CREATE TABLE test(encoding vector(3));
+        CREATE INDEX idxtest
+                  ON test
+               USING diskann(encoding)
+               WITH (storage_layout = io_optimized);",
+        )?;
+
+        let index_oid =
+            Spi::get_one::<pg_sys::Oid>("SELECT 'idxtest'::regclass::oid")?.expect("oid was null");
+        let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid));
+        let options = TSVIndexOptions::from_relation(&indexrel);
+        assert_eq!(options.get_num_neighbors(), NUM_NEIGHBORS_DEFAULT_SENTINEL);
+        assert_eq!(options.search_list_size, 100);
+        assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA);
+        assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL);
+        assert_eq!(options.get_storage_type(), StorageType::SbqSpeedup);
+        Ok(())
+    }
+
     #[pg_test]
     unsafe fn test_index_options_plain() -> spi::Result<()> {
         Spi::run(