[Distance metrics M1] support for L2 distance (#156)

This PR adds support for distance metrics besides cosine, beginning with L2. The primitives for this metric were already present (earlier versions of pgvectorscale used L2 rather than cosine), so this PR just adds the operator class-related plumbing. As discussed in the design doc (see references) we follow pgvector syntax.

---------

Signed-off-by: tjgreen42 <tj@timescale.com>
Co-authored-by: Matvey Arye <cevian@gmail.com>

Author: tjgreen42

pgvectorscale/Cargo.toml

@@ -22,9 +22,9 @@ pg_test = []
 [dependencies]
 memoffset = "0.9.0"
 pgrx = "=0.12.5"
-rkyv = { version="0.7.42", features=["validation"]}
-simdeez = {version = "1.0.8"}
-rand = { version = "0.8", features = [ "small_rng" ] }
+rkyv = { version = "0.7.42", features = ["validation"] }
+simdeez = { version = "1.0.8" }
+rand = { version = "0.8", features = ["small_rng"] }
 pgvectorscale_derive = { path = "pgvectorscale_derive" }
 semver = "1.0.22"
 once_cell = "1.20.1"

pgvectorscale/src/access_method/build.rs

@@ -1,14 +1,17 @@
 use std::time::Instant;
 
-use pgrx::pg_sys::{pgstat_progress_update_param, AsPgCStr};
+use pg_sys::{FunctionCall0Coll, InvalidOid};
+use pgrx::pg_sys::{index_getprocinfo, pgstat_progress_update_param, AsPgCStr};
 use pgrx::*;
 
+use crate::access_method::distance::DistanceType;
 use crate::access_method::graph::Graph;
 use crate::access_method::graph_neighbor_store::GraphNeighborStore;
 use crate::access_method::options::TSVIndexOptions;
 use crate::access_method::pg_vector::PgVector;
 use crate::access_method::stats::{InsertStats, WriteStats};
 
+use crate::access_method::DISKANN_DISTANCE_TYPE_PROC;
 use crate::util::page::PageType;
 use crate::util::tape::Tape;
 use crate::util::*;
@@ -79,7 +82,18 @@ pub extern "C" fn ambuild(
 
     let dimensions = index_relation.tuple_desc().get(0).unwrap().atttypmod;
     assert!(dimensions > 0 && dimensions <= 2000);
-    let meta_page = unsafe { MetaPage::create(&index_relation, dimensions as _, opt) };
+
+    let distance_type = unsafe {
+        let fmgr_info = index_getprocinfo(indexrel, 1, DISKANN_DISTANCE_TYPE_PROC);
+        if fmgr_info == std::ptr::null_mut() {
+            error!("No distance type function found for index");
+        }
+        let result = FunctionCall0Coll(fmgr_info, InvalidOid).value() as u16;
+        DistanceType::from_u16(result)
+    };
+
+    let meta_page =
+        unsafe { MetaPage::create(&index_relation, dimensions as _, distance_type, opt) };
 
     let ntuples = do_heap_scan(index_info, &heap_relation, &index_relation, meta_page);
 
@@ -487,15 +501,19 @@ pub unsafe extern "C" fn ambuildphasename(phasenum: i64) -> *mut ffi::c_char {
 pub mod tests {
     use std::collections::HashSet;
 
+    use crate::access_method::distance::DistanceType;
     use pgrx::*;
 
     //TODO: add test where inserting and querying with vectors that are all the same.
 
     #[cfg(any(test, feature = "pg_test"))]
     pub unsafe fn test_index_creation_and_accuracy_scaffold(
+        distance_type: DistanceType,
         index_options: &str,
         name: &str,
     ) -> spi::Result<()> {
+        let operator = distance_type.get_operator();
+        let operator_class = distance_type.get_operator_class();
         let table_name = format!("test_data_icaa_{}", name);
         Spi::run(&format!(
             "CREATE TABLE {table_name} (
@@ -515,7 +533,7 @@ pub mod tests {
                 GROUP BY
                     i % 300) g;
 
-            CREATE INDEX ON {table_name} USING diskann (embedding) WITH ({index_options});
+            CREATE INDEX ON {table_name} USING diskann (embedding {operator_class}) WITH ({index_options});
 
 
             SET enable_seqscan = 0;
@@ -525,7 +543,7 @@ pub mod tests {
             FROM
                 {table_name}
             ORDER BY
-                embedding <=> (
+                embedding {operator} (
                     SELECT
                         ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding
             FROM generate_series(1, 1536));"))?;
@@ -542,7 +560,7 @@ pub mod tests {
             SET enable_seqscan = 0;
             SET enable_indexscan = 1;
             SET diskann.query_search_list_size = 2;
-            WITH cte as (select * from {table_name} order by embedding <=> $1::vector) SELECT count(*) from cte;
+            WITH cte as (select * from {table_name} order by embedding {operator} $1::vector) SELECT count(*) from cte;
             ",
                 ),
                 vec![(
@@ -576,7 +594,7 @@ pub mod tests {
             FROM
                 {table_name}
             ORDER BY
-                embedding <=> (
+                embedding {operator} (
                     SELECT
                         ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding
             FROM generate_series(1, 1536));
@@ -608,7 +626,7 @@ pub mod tests {
             FROM
                 {table_name}
             ORDER BY
-                embedding <=> $1::vector
+                embedding {operator} $1::vector
             LIMIT 10
         )
         SELECT array_agg(ctid) from cte;"
@@ -631,7 +649,7 @@ pub mod tests {
             FROM
                 {table_name}
             ORDER BY
-                embedding <=> $1::vector
+                embedding {operator} $1::vector
             LIMIT 10
         )
         SELECT array_agg(ctid) from cte;"
@@ -655,7 +673,7 @@ pub mod tests {
             FROM
                 {table_name}
             ORDER BY
-                embedding <=> $1::vector
+                embedding {operator} $1::vector
             LIMIT 10
         )
         SELECT array_agg(ctid) from cte;"
@@ -686,7 +704,7 @@ pub mod tests {
         SET enable_seqscan = 0;
         SET enable_indexscan = 1;
         SET diskann.query_search_list_size = 2;
-        WITH cte as (select * from {table_name} order by embedding <=> $1::vector) SELECT count(*) from cte;
+        WITH cte as (select * from {table_name} order by embedding {operator} $1::vector) SELECT count(*) from cte;
         ",
             ),
             vec![(
@@ -759,10 +777,14 @@ pub mod tests {
 
     #[cfg(any(test, feature = "pg_test"))]
     pub unsafe fn test_index_updates(
+        distance_type: DistanceType,
         index_options: &str,
         expected_cnt: i64,
         name: &str,
     ) -> spi::Result<()> {
+        let operator_class = distance_type.get_operator_class();
+        let operator = distance_type.get_operator();
+
         let table_name = format!("test_data_index_updates_{}", name);
         Spi::run(&format!(
             "CREATE TABLE {table_name} (
@@ -784,7 +806,7 @@ pub mod tests {
                 GROUP BY
                     i % {expected_cnt}) g;
 
-            CREATE INDEX ON {table_name} USING diskann (embedding) WITH ({index_options});
+            CREATE INDEX ON {table_name} USING diskann (embedding {operator_class}) WITH ({index_options});
 
 
             SET enable_seqscan = 0;
@@ -794,7 +816,7 @@ pub mod tests {
             FROM
                 {table_name}
             ORDER BY
-                embedding <=> (
+                embedding {operator} (
                     SELECT
                         ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding
             FROM generate_series(1, 1536));"))?;
@@ -811,7 +833,7 @@ pub mod tests {
             SET enable_seqscan = 0;
             SET enable_indexscan = 1;
             SET diskann.query_search_list_size = 2;
-            WITH cte as (select * from {table_name} order by embedding <=> $1::vector) SELECT count(*) from cte;
+            WITH cte as (select * from {table_name} order by embedding {operator} $1::vector) SELECT count(*) from cte;
             ",
                 ),
                 vec![(
@@ -850,7 +872,7 @@ pub mod tests {
         SET enable_seqscan = 0;
         SET enable_indexscan = 1;
         SET diskann.query_search_list_size = 2;
-        WITH cte as (select * from {table_name} order by embedding <=> $1::vector) SELECT count(*) from cte;
+        WITH cte as (select * from {table_name} order by embedding {operator} $1::vector) SELECT count(*) from cte;
         ",
             ),
             vec![(

pgvectorscale/src/access_method/distance.rs

@@ -1,5 +1,54 @@
+use pgrx::pg_extern;
+
 pub type DistanceFn = fn(&[f32], &[f32]) -> f32;
 
+#[derive(Debug, PartialEq)]
+pub enum DistanceType {
+    Cosine = 0,
+    L2 = 1,
+}
+
+impl DistanceType {
+    pub fn from_u16(value: u16) -> Self {
+        match value {
+            0 => DistanceType::Cosine,
+            1 => DistanceType::L2,
+            _ => panic!("Unknown DistanceType number {}", value),
+        }
+    }
+
+    pub fn get_operator(&self) -> &str {
+        match self {
+            DistanceType::Cosine => "<=>",
+            DistanceType::L2 => "<->",
+        }
+    }
+
+    pub fn get_operator_class(&self) -> &str {
+        match self {
+            DistanceType::Cosine => "vector_cosine_ops",
+            DistanceType::L2 => "vector_l2_ops",
+        }
+    }
+
+    pub fn get_distance_function(&self) -> DistanceFn {
+        match self {
+            DistanceType::Cosine => distance_cosine,
+            DistanceType::L2 => distance_l2,
+        }
+    }
+}
+
+#[pg_extern(immutable, parallel_safe)]
+pub fn distance_type_cosine() -> i16 {
+    DistanceType::Cosine as i16
+}
+
+#[pg_extern(immutable, parallel_safe)]
+pub fn distance_type_l2() -> i16 {
+    DistanceType::L2 as i16
+}
+
 /* we use the avx2 version of x86 functions. This verifies that's kosher */
 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
 #[cfg(not(target_feature = "avx2"))]

pgvectorscale/src/access_method/meta_page.rs

@@ -8,7 +8,7 @@ use crate::access_method::options::TSVIndexOptions;
 use crate::util::page;
 use crate::util::*;
 
-use super::distance::{self, DistanceFn};
+use super::distance::{DistanceFn, DistanceType};
 use super::options::{
     NUM_DIMENSIONS_DEFAULT_SENTINEL, NUM_NEIGHBORS_DEFAULT_SENTINEL,
     SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL,
@@ -93,22 +93,6 @@ pub struct MetaPageHeader {
     version: u32,
 }
 
-/// TODO: move to distance.rs
-enum DistanceType {
-    Cosine = 0,
-    L2 = 1,
-}
-
-impl DistanceType {
-    fn from_u16(value: u16) -> Self {
-        match value {
-            0 => DistanceType::Cosine,
-            1 => DistanceType::L2,
-            _ => panic!("Unknown DistanceType number {}", value),
-        }
-    }
-}
-
 /// This is metadata about the entire index.
 /// Stored as the first page (offset 2) in the index relation.
 #[derive(Clone, PartialEq, Archive, Deserialize, Serialize, Readable, Writeable)]
@@ -175,10 +159,7 @@ impl MetaPage {
     }
 
     pub fn get_distance_function(&self) -> DistanceFn {
-        match DistanceType::from_u16(self.distance_type) {
-            DistanceType::Cosine => distance::distance_cosine,
-            DistanceType::L2 => distance::distance_l2,
-        }
+        DistanceType::from_u16(self.distance_type).get_distance_function()
     }
 
     pub fn get_storage_type(&self) -> StorageType {
@@ -234,6 +215,7 @@ impl MetaPage {
     pub unsafe fn create(
         index: &PgRelation,
         num_dimensions: u32,
+        distance_type: DistanceType,
         opt: PgBox<TSVIndexOptions>,
     ) -> MetaPage {
         let version = Version::parse(env!("CARGO_PKG_VERSION")).unwrap();
@@ -272,7 +254,7 @@ impl MetaPage {
             magic_number: TSV_MAGIC_NUMBER,
             version: TSV_VERSION,
             extension_version_when_built: version.to_string(),
-            distance_type: DistanceType::Cosine as u16,
+            distance_type: distance_type as u16,
             num_dimensions,
             num_dimensions_to_index,
             storage_type: (*opt).get_storage_type() as u8,