chore: type check and external test

scripts/train.py

@@ -170,7 +170,7 @@ def kmeans_cluster(
             verbose=True,
             niter=niter,
             seed=SEED,
-            spherical=metric == "cos",
+            spherical=metric != "l2",
         )
         child_kmeans.train(child_train)
         centroids.append(child_kmeans.centroids)

src/vchordrq/algorithm/build.rs

@@ -203,28 +203,40 @@ impl Structure {
     ) -> Vec<Self> {
         use std::collections::BTreeMap;
         let VchordrqExternalBuildOptions { table } = external_build;
-        let query = format!("SELECT id, parent, vector FROM {table};");
         let mut parents = BTreeMap::new();
         let mut vectors = BTreeMap::new();
         pgrx::spi::Spi::connect(|client| {
             use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput;
             use base::vector::VectorBorrowed;
             use pgrx::pg_sys::panic::ErrorReportable;
-            let table = client.select(&query, None, None).unwrap_or_report();
-            for row in table {
+            let schema_query = "SELECT n.nspname::TEXT 
+                FROM pg_catalog.pg_extension e
+                LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace
+                WHERE e.extname = 'vector';";
+            let pgvector_schema: String = client
+                .select(schema_query, None, None)
+                .unwrap_or_report()
+                .first()
+                .get_by_name("nspname")
+                .expect("external build: cannot get schema of pgvector")
+                .expect("external build: cannot get schema of pgvector");
+            let dump_query =
+                format!("SELECT id, parent, vector::{pgvector_schema}.vector FROM {table};");
+            let centroids = client.select(&dump_query, None, None).unwrap_or_report();
+            for row in centroids {
                 let id: Option<i32> = row.get_by_name("id").unwrap();
                 let parent: Option<i32> = row.get_by_name("parent").unwrap();
                 let vector: Option<PgvectorVectorOutput> = row.get_by_name("vector").unwrap();
-                let id = id.expect("extern build: id could not be NULL");
-                let vector = vector.expect("extern build: vector could not be NULL");
+                let id = id.expect("external build: id could not be NULL");
+                let vector = vector.expect("external build: vector could not be NULL");
                 let pop = parents.insert(id, parent);
                 if pop.is_some() {
                     pgrx::error!(
                         "external build: there are at least two lines have same id, id = {id}"
                     );
                 }
                 if vector_options.dims != vector.as_borrowed().dims() {
-                    pgrx::error!("extern build: incorrect dimension, id = {id}");
+                    pgrx::error!("external build: incorrect dimension, id = {id}");
                 }
                 vectors.insert(id, crate::projection::project(vector.as_borrowed().slice()));
             }
@@ -275,7 +287,7 @@ impl Structure {
             }
         }
         let Some(root) = root else {
-            pgrx::error!("extern build: there are no root");
+            pgrx::error!("external build: there are no root");
         };
         let mut heights = BTreeMap::<_, _>::new();
         fn dfs_for_heights(
@@ -284,7 +296,7 @@ impl Structure {
             u: i32,
         ) {
             if heights.contains_key(&u) {
-                pgrx::error!("extern build: detect a cycle, id = {u}");
+                pgrx::error!("external build: detect a cycle, id = {u}");
             }
             heights.insert(u, None);
             let mut height = None;
@@ -293,7 +305,7 @@ impl Structure {
                 let new = heights[&v].unwrap() + 1;
                 if let Some(height) = height {
                     if height != new {
-                        pgrx::error!("extern build: two heights, id = {u}");
+                        pgrx::error!("external build: two heights, id = {u}");
                     }
                 } else {
                     height = Some(new);
@@ -311,7 +323,7 @@ impl Structure {
             .collect::<BTreeMap<_, _>>();
         if !(1..=8).contains(&(heights[&root] - 1)) {
             pgrx::error!(
-                "extern build: unexpected tree height, height = {}",
+                "external build: unexpected tree height, height = {}",
                 heights[&root]
             );
         }

tests/logic/external_build.slt

@@ -0,0 +1,117 @@
+statement ok
+CREATE TABLE t (val0 vector(3), val1 halfvec(3));
+
+statement ok
+INSERT INTO t (val0, val1) 
+SELECT 
+    ARRAY[random(), random(), random()]::real[]::vector, 
+    ARRAY[random(), random(), random()]::real[]::halfvec
+FROM generate_series(1, 100);
+
+statement ok
+CREATE TABLE vector_centroid (id integer, parent integer, vector vector(3));
+
+statement ok
+INSERT INTO vector_centroid (id, vector) VALUES 
+    (0, '[1.0, 0.0, 0.0]'),
+    (1, '[0.0, 1.0, 0.0]'),
+    (2, '[0.0, 0.0, 1.0]');
+
+statement ok
+CREATE TABLE halfvec_centroid (id integer, parent integer, vector halfvec(3));
+
+statement ok
+INSERT INTO halfvec_centroid (id, vector) VALUES 
+    (0, '[1.0, 0.0, 0.0]'),
+    (1, '[0.0, 1.0, 0.0]'),
+    (2, '[0.0, 0.0, 1.0]');
+
+statement ok
+CREATE TABLE real_centroid (id integer, parent integer, vector real[]);
+
+statement ok
+INSERT INTO real_centroid (id, vector) VALUES 
+    (0, '{1.0, 0.0, 0.0}'),
+    (1, '{0.0, 1.0, 0.0}'),
+    (2, '{0.0, 0.0, 1.0}');
+
+statement ok
+CREATE TABLE bad_type_centroid (id integer, parent integer, vector integer);
+
+statement ok
+INSERT INTO bad_type_centroid (id, vector) VALUES 
+    (0, 0),
+    (1, 0),
+    (2, 0);
+
+statement ok
+CREATE TABLE bad_duplicate_id (id integer, parent integer, vector vector(3));
+
+statement ok
+INSERT INTO bad_duplicate_id (id, vector) VALUES 
+    (1, '[1.0, 0.0, 0.0]'),
+    (1, '[0.0, 1.0, 0.0]'),
+    (2, '[0.0, 0.0, 1.0]');
+
+# external build for vector column
+
+statement ok
+CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops)
+WITH (options = $$
+residual_quantization = true
+[build.external]
+table = 'public.vector_centroid'
+$$);
+
+# external build for halfvec column
+
+statement ok
+CREATE INDEX ON t USING vchordrq (val1 halfvec_l2_ops)
+WITH (options = $$
+residual_quantization = true
+[build.external]
+table = 'public.vector_centroid'
+$$);
+
+# external build for halfvec column by a halfvec table
+
+statement ok
+CREATE INDEX ON t USING vchordrq (val1 halfvec_l2_ops)
+WITH (options = $$
+residual_quantization = true
+[build.external]
+table = 'public.halfvec_centroid'
+$$);
+
+# external build for halfvec column by a real[] table
+
+statement ok
+CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops)
+WITH (options = $$
+residual_quantization = true
+[build.external]
+table = 'public.real_centroid'
+$$);
+
+# failed: bad vector data type
+
+statement error cannot cast type integer to (.*)vector
+CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops)
+WITH (options = $$
+residual_quantization = true
+[build.external]
+table = 'public.bad_type_centroid'
+$$);
+
+# failed: duplicate id
+
+statement error external build: there are at least two lines have same id, id = 1
+CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops)
+WITH (options = $$
+residual_quantization = true
+[build.external]
+table = 'public.bad_duplicate_id'
+$$);
+
+statement ok
+DROP TABLE t, vector_centroid, halfvec_centroid, real_centroid, bad_type_centroid, bad_duplicate_id;