feat: add a field to meta tuple

tensorchord · Nov 14, 2024 · 1e33fec · 1e33fec
1 parent 3a30e21
commit 1e33fec
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs
@@ -48,7 +48,7 @@ pub fn build<T: HeapRelation, R: Reporter>(
             let mut tuples_total = 0_u64;
             let samples = {
                 let mut rand = rand::thread_rng();
-                let max_number_of_samples = internal_build.lists.saturating_mul(256);
+                let max_number_of_samples = internal_build.lists[0].saturating_mul(256);
                 let mut samples = Vec::new();
                 let mut number_of_samples = 0_u32;
                 heap_relation.traverse(false, |(_, vector)| {
@@ -167,6 +167,7 @@ pub fn build<T: HeapRelation, R: Reporter>(
     forwards.head.get_mut().get_opaque_mut().fast_forward = vectors.first();
     meta.push(&MetaTuple {
         dims,
+        height_of_root: 2,
         is_residual,
         vectors_first: vectors.first(),
         forwards_first: forwards.first(),
@@ -196,7 +197,7 @@ impl Structure {
             |parallelism| {
                 k_means::k_means(
                     parallelism,
-                    internal_build.lists as usize,
+                    internal_build.lists[0] as usize,
                     vector_options.dims as usize,
                     samples,
                     internal_build.spherical_centroids,

diff --git a/src/algorithm/tuples.rs b/src/algorithm/tuples.rs
@@ -5,6 +5,7 @@ use rkyv::{Archive, Deserialize, Serialize};
 #[archive(check_bytes)]
 pub struct MetaTuple {
     pub dims: u32,
+    pub height_of_root: u32,
     pub is_residual: bool,
     pub vectors_first: u32,
     pub forwards_first: u32,

diff --git a/src/types.rs b/src/types.rs
@@ -5,8 +5,7 @@ use validator::Validate;
 #[serde(deny_unknown_fields)]
 pub struct RabbitholeInternalBuildOptions {
     #[serde(default = "RabbitholeInternalBuildOptions::default_lists")]
-    #[validate(range(min = 1, max = 1_000_000))]
-    pub lists: u32,
+    pub lists: [u32; 1],
     #[serde(default = "RabbitholeInternalBuildOptions::default_spherical_centroids")]
     pub spherical_centroids: bool,
     #[serde(default = "RabbitholeInternalBuildOptions::default_build_threads")]
@@ -15,8 +14,8 @@ pub struct RabbitholeInternalBuildOptions {
 }
 
 impl RabbitholeInternalBuildOptions {
-    fn default_lists() -> u32 {
-        1000
+    fn default_lists() -> [u32; 1] {
+        [1000]
     }
     fn default_spherical_centroids() -> bool {
         false