Fix test flakiness

Author: cevian

pgvectorscale/src/access_method/build.rs

@@ -469,15 +469,17 @@ pub mod tests {
     #[cfg(any(test, feature = "pg_test"))]
     pub unsafe fn test_index_creation_and_accuracy_scaffold(
         index_options: &str,
+        name: &str,
     ) -> spi::Result<()> {
+        let table_name = format!("test_data_icaa_{}", name);
         Spi::run(&format!(
-            "CREATE TABLE test_data (
+            "CREATE TABLE {table_name} (
                 embedding vector (1536)
             );
 
             select setseed(0.5);
            -- generate 300 vectors
-            INSERT INTO test_data (embedding)
+            INSERT INTO {table_name} (embedding)
             SELECT
                 *
             FROM (
@@ -488,15 +490,15 @@ pub mod tests {
                 GROUP BY
                     i % 300) g;
 
-            CREATE INDEX idx_diskann_bq ON test_data USING diskann (embedding) WITH ({index_options});
+            CREATE INDEX ON {table_name} USING diskann (embedding) WITH ({index_options});
 
 
             SET enable_seqscan = 0;
             -- perform index scans on the vectors
             SELECT
                 *
             FROM
-                test_data
+                {table_name}
             ORDER BY
                 embedding <=> (
                     SELECT
@@ -514,7 +516,7 @@ pub mod tests {
             SET enable_seqscan = 0;
             SET enable_indexscan = 1;
             SET diskann.query_search_list_size = 2;
-            WITH cte as (select * from test_data order by embedding <=> $1::vector) SELECT count(*) from cte;
+            WITH cte as (select * from {table_name} order by embedding <=> $1::vector) SELECT count(*) from cte;
             ",
                 ),
                 vec![(
@@ -530,7 +532,7 @@ pub mod tests {
 
         Spi::run(&format!("
             -- test insert 2 vectors
-            INSERT INTO test_data (embedding)
+            INSERT INTO {table_name} (embedding)
             SELECT
                 *
             FROM (
@@ -546,15 +548,15 @@ pub mod tests {
             SELECT
                 *
             FROM
-                test_data
+                {table_name}
             ORDER BY
                 embedding <=> (
                     SELECT
                         ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding
             FROM generate_series(1, 1536));
 
             -- test insert 10 vectors to search for that aren't random
-            INSERT INTO test_data (embedding)
+            INSERT INTO {table_name} (embedding)
             SELECT
                 *
             FROM (
@@ -578,7 +580,7 @@ pub mod tests {
             SELECT
                 ctid
             FROM
-                test_data
+                {table_name}
             ORDER BY
                 embedding <=> $1::vector
             LIMIT 10
@@ -601,7 +603,7 @@ pub mod tests {
             SELECT
                 ctid
             FROM
-                test_data
+                {table_name}
             ORDER BY
                 embedding <=> $1::vector
             LIMIT 10
@@ -625,7 +627,7 @@ pub mod tests {
             SELECT
                 ctid
             FROM
-                test_data
+                {table_name}
             ORDER BY
                 embedding <=> $1::vector
             LIMIT 10
@@ -661,7 +663,7 @@ pub mod tests {
         SET enable_seqscan = 0;
         SET enable_indexscan = 1;
         SET diskann.query_search_list_size = 2;
-        WITH cte as (select * from test_data order by embedding <=> $1::vector) SELECT count(*) from cte;
+        WITH cte as (select * from {table_name} order by embedding <=> $1::vector) SELECT count(*) from cte;
         ",
             ),
             vec![(
@@ -736,16 +738,21 @@ pub mod tests {
     }
 
     #[cfg(any(test, feature = "pg_test"))]
-    pub unsafe fn test_index_updates(index_options: &str, expected_cnt: i64) -> spi::Result<()> {
+    pub unsafe fn test_index_updates(
+        index_options: &str,
+        expected_cnt: i64,
+        name: &str,
+    ) -> spi::Result<()> {
+        let table_name = format!("test_data_index_updates_{}", name);
         Spi::run(&format!(
-            "CREATE TABLE test_data (
+            "CREATE TABLE {table_name} (
                 id int,
                 embedding vector (1536)
             );
 
             select setseed(0.5);
            -- generate 300 vectors
-            INSERT INTO test_data (id, embedding)
+            INSERT INTO {table_name} (id, embedding)
             SELECT
                 *
             FROM (
@@ -757,15 +764,15 @@ pub mod tests {
                 GROUP BY
                     i % {expected_cnt}) g;
 
-            CREATE INDEX idx_diskann_bq ON test_data USING diskann (embedding) WITH ({index_options});
+            CREATE INDEX ON {table_name} USING diskann (embedding) WITH ({index_options});
 
 
             SET enable_seqscan = 0;
             -- perform index scans on the vectors
             SELECT
                 *
             FROM
-                test_data
+                {table_name}
             ORDER BY
                 embedding <=> (
                     SELECT
@@ -783,7 +790,7 @@ pub mod tests {
             SET enable_seqscan = 0;
             SET enable_indexscan = 1;
             SET diskann.query_search_list_size = 2;
-            WITH cte as (select * from test_data order by embedding <=> $1::vector) SELECT count(*) from cte;
+            WITH cte as (select * from {table_name} order by embedding <=> $1::vector) SELECT count(*) from cte;
             ",
                 ),
                 vec![(
@@ -797,7 +804,7 @@ pub mod tests {
         Spi::run(&format!(
             "
 
-        --CREATE INDEX idx_id ON test_data(id);
+        --CREATE INDEX idx_id ON {table_name}(id);
 
         WITH CTE as (
             SELECT
@@ -808,9 +815,9 @@ pub mod tests {
             GROUP BY
             i % {expected_cnt}
         )
-        UPDATE test_data SET embedding = cte.embedding
+        UPDATE {table_name} SET embedding = cte.embedding
         FROM cte
-        WHERE test_data.id = cte.id;
+        WHERE {table_name}.id = cte.id;
 
         --DROP INDEX idx_id;
             ",
@@ -822,7 +829,7 @@ pub mod tests {
         SET enable_seqscan = 0;
         SET enable_indexscan = 1;
         SET diskann.query_search_list_size = 2;
-        WITH cte as (select * from test_data order by embedding <=> $1::vector) SELECT count(*) from cte;
+        WITH cte as (select * from {table_name} order by embedding <=> $1::vector) SELECT count(*) from cte;
         ",
             ),
             vec![(

pgvectorscale/src/access_method/plain_storage.rs

@@ -360,6 +360,7 @@ mod tests {
     unsafe fn test_plain_storage_index_creation_many_neighbors() -> spi::Result<()> {
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "num_neighbors=38, storage_layout = plain",
+            "plain_many_neighbors",
         )?;
         Ok(())
     }
@@ -369,6 +370,7 @@ mod tests {
         //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past.
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "num_neighbors=10, storage_layout = plain",
+            "plain_few_neighbors",
         )?;
         Ok(())
     }
@@ -412,6 +414,7 @@ mod tests {
     unsafe fn test_plain_storage_num_dimensions() -> spi::Result<()> {
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "num_neighbors=38, storage_layout = plain, num_dimensions=768",
+            "plain_num_dimensions",
         )?;
         Ok(())
     }
@@ -421,6 +424,7 @@ mod tests {
         crate::access_method::build::tests::test_index_updates(
             "storage_layout = plain, num_neighbors=30",
             50,
+            "plain",
         )?;
         Ok(())
     }

pgvectorscale/src/access_method/sbq.rs

@@ -969,6 +969,7 @@ mod tests {
     unsafe fn test_bq_speedup_storage_index_creation_default_neighbors() -> spi::Result<()> {
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "storage_layout = io_optimized",
+            "bq_speedup_default_neighbors",
         )?;
         Ok(())
     }
@@ -978,6 +979,7 @@ mod tests {
         //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past.
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "num_neighbors=10, storage_layout = io_optimized",
+            "bq_speedup_few_neighbors",
         )?;
         Ok(())
     }
@@ -1014,6 +1016,7 @@ mod tests {
     unsafe fn test_bq_speedup_storage_index_creation_num_dimensions() -> spi::Result<()> {
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "storage_layout = io_optimized, num_dimensions=768",
+            "bq_speedup_num_dimensions",
         )?;
         Ok(())
     }
@@ -1023,14 +1026,16 @@ mod tests {
         crate::access_method::build::tests::test_index_updates(
             "storage_layout = io_optimized, num_neighbors=10",
             300,
+            "bq_speedup",
         )?;
         Ok(())
     }
 
     #[pg_test]
-    unsafe fn test_bq_speedup_compressed_index_creation_default_neighbors() -> spi::Result<()> {
+    unsafe fn test_bq_compressed_index_creation_default_neighbors() -> spi::Result<()> {
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "storage_layout = memory_optimized",
+            "bq_compressed_default_neighbors",
         )?;
         Ok(())
     }
@@ -1040,6 +1045,7 @@ mod tests {
         //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past.
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "num_neighbors=10, storage_layout = memory_optimized",
+            "bq_compressed_few_neighbors",
         )?;
         Ok(())
     }
@@ -1082,6 +1088,7 @@ mod tests {
     unsafe fn test_bq_compressed_storage_index_creation_num_dimensions() -> spi::Result<()> {
         crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold(
             "storage_layout = memory_optimized, num_dimensions=768",
+            "bq_compressed_num_dimensions",
         )?;
         Ok(())
     }
@@ -1091,6 +1098,7 @@ mod tests {
         crate::access_method::build::tests::test_index_updates(
             "storage_layout = memory_optimized, num_neighbors=10",
             300,
+            "bq_compressed",
         )?;
         Ok(())
     }