Streaming PQ

protos/index.proto

@@ -72,6 +72,23 @@ message PQ {
   repeated float codebook = 4;
 }
 
+// Transform type
+enum TransformType {
+  OPQ = 0;
+}
+
+// A transform matrix to apply to a vector or vectors.
+message Transform {
+  // The file offset the matrix is stored
+  uint64 position = 1;
+
+  // Data shape of the matrix, [rows, cols].
+  repeated uint32 shape = 2;
+
+  // Transform type.
+  TransformType type = 3;
+}
+
 // Flat Index
 message Flat {}
 
@@ -84,6 +101,8 @@ message VectorIndexStage {
     IVF ivf = 2;
     // Product Quantization
     PQ pq = 3;
+    // Transformer
+    Transform transform = 4;
   }
 }
 

rust/src/arrow.rs

@@ -138,7 +138,7 @@ pub trait GenericListArrayExt<Offset: ArrowNumericType>
 where
     Offset::Native: OffsetSizeTrait,
 {
-    /// Create an [`ListArray`] from values and offsets.
+    /// Create an [`GenericListArray`] from values and offsets.
     ///
     /// ```
     /// use arrow_array::{Int32Array, Int64Array, ListArray};

rust/src/arrow/linalg.rs

@@ -23,7 +23,8 @@ use arrow::{
     array::{as_primitive_array, Float32Builder},
     datatypes::Float32Type,
 };
-use arrow_array::{Array, Float32Array};
+use arrow_array::{Array, FixedSizeListArray, Float32Array};
+use arrow_schema::DataType;
 use rand::{distributions::Standard, rngs::SmallRng, seq::IteratorRandom, Rng, SeedableRng};
 
 #[allow(unused_imports)]
@@ -130,6 +131,15 @@ impl MatrixView {
         }
     }
 
+    pub fn row(&self, i: usize) -> Option<Float32Array> {
+        if i >= self.num_rows() {
+            None
+        } else {
+            let slice_arr = self.data.slice(i * self.num_columns(), self.num_columns());
+            Some(as_primitive_array(slice_arr.as_ref()).clone())
+        }
+    }
+
     /// (Lazy) transpose of the matrix.
     ///
     pub fn transpose(&self) -> Self {
@@ -225,6 +235,25 @@ impl MatrixView {
     }
 }
 
+impl TryFrom<&FixedSizeListArray> for MatrixView {
+    type Error = Error;
+
+    fn try_from(fsl: &FixedSizeListArray) -> Result<Self> {
+        if !matches!(fsl.value_type(), DataType::Float32) {
+            return Err(Error::Arrow(format!(
+                "Only support convert f32 FixedSizeListArray to MatrixView, got {}",
+                fsl.data_type()
+            )));
+        }
+        let values = fsl.values();
+        Ok(Self {
+            data: Arc::new(as_primitive_array(values.as_ref()).clone()),
+            num_columns: fsl.value_length() as usize,
+            transpose: false,
+        })
+    }
+}
+
 /// Single Value Decomposition.
 ///
 /// <https://en.wikipedia.org/wiki/Singular_value_decomposition>

rust/src/dataset.rs

@@ -36,10 +36,8 @@ use self::scanner::Scanner;
 use crate::arrow::*;
 use crate::datatypes::Schema;
 use crate::format::{pb, Fragment, Index, Manifest};
-use crate::index::{
-    vector::{ivf::IvfPqIndexBuilder, VectorIndexParams},
-    IndexBuilder, IndexParams, IndexType,
-};
+use crate::index::vector::ivf::{build_ivf_pq_index, IvfBuildParams, PQBuildParams};
+use crate::index::{vector::VectorIndexParams, IndexParams, IndexType};
 use crate::io::{
     object_reader::{read_message, read_struct},
     read_manifest, read_metadata_offset, write_manifest, FileReader, FileWriter, ObjectStore,
@@ -277,7 +275,7 @@ impl Dataset {
         let base = object_store.base_path().clone();
         Ok(Self {
             object_store,
-            base,
+            base: base.into(),
             manifest: Arc::new(manifest.clone()),
         })
     }
@@ -381,16 +379,27 @@ impl Dataset {
                     }
                 }
 
-                let builder = IvfPqIndexBuilder::try_new(
+                let ivf_params = IvfBuildParams {
+                    num_partitions: vec_params.num_partitions as usize,
+                    metric_type: vec_params.metric_type,
+                    max_iters: 50,
+                };
+                let pq_params = PQBuildParams {
+                    num_sub_vectors: vec_params.num_sub_vectors as usize,
+                    num_bits: 8,
+                    metric_type: vec_params.metric_type,
+                    use_opq: true,
+                    max_iters: 100,
+                };
+                build_ivf_pq_index(
                     self,
-                    index_id,
-                    &index_name,
                     column,
-                    vec_params.num_partitions,
-                    vec_params.num_sub_vectors,
-                    vec_params.metric_type,
-                )?;
-                builder.build().await?
+                    &index_name,
+                    &index_id,
+                    &ivf_params,
+                    &pq_params,
+                )
+                .await?
             }
         }
 
@@ -521,6 +530,14 @@ impl Dataset {
         Ok(as_struct_array(&reordered).into())
     }
 
+    /// Sample `n` rows from the dataset.
+    pub(crate) async fn sample(&self, n: usize, projection: &Schema) -> Result<RecordBatch> {
+        use rand::seq::IteratorRandom;
+        let num_rows = self.count_rows().await?;
+        let ids = (0..num_rows).choose_multiple(&mut rand::thread_rng(), n);
+        Ok(self.take(&ids[..], &projection).await?)
+    }
+
     pub(crate) fn object_store(&self) -> &ObjectStore {
         &self.object_store
     }

rust/src/index/vector.rs

@@ -30,8 +30,9 @@ mod kmeans;
 mod opq;
 mod pq;
 
-use super::IndexParams;
+use super::{pb, IndexParams};
 use crate::{
+    arrow::linalg::MatrixView,
     utils::distance::{cosine::cosine_distance, l2::l2_distance},
     Error, Result,
 };
@@ -79,6 +80,26 @@ pub trait VectorIndex {
     async fn search(&self, query: &Query) -> Result<RecordBatch>;
 }
 
+/// Transformer on vectors.
+#[async_trait]
+pub trait Transformer: std::fmt::Debug + Sync + Send {
+    /// Train the transformer.
+    ///
+    /// Parameters:
+    /// - *data*: training vectors.
+    async fn train(&mut self, data: &MatrixView) -> Result<()>;
+
+    /// Apply transform on the matrix `data`.
+    ///
+    /// Returns a new Matrix instead.
+    async fn transform(&self, data: &MatrixView) -> Result<MatrixView>;
+
+    /// Try to convert into protobuf.
+    ///
+    /// TODO: can we use TryFrom/TryInto as trait constrats?
+    fn try_into_pb(&self) -> Result<pb::Transform>;
+}
+
 /// Distance metrics type.
 #[derive(Debug, Copy, Clone, PartialEq)]
 pub enum MetricType {