add quantization

Author: b41sh

Cargo.lock

@@ -261,6 +261,7 @@ base64 = "0.22"
 bincode = { version = "2.0.0-rc.3", features = ["serde", "std", "alloc"] }
 bincode_v1 = { package = "bincode", version = "1.3.3" }
 bitpacking = "0.8.0"
+bitvec = "1.0.1"
 blake3 = "1.3.1"
 bollard = { version = "0.17" }
 borsh = { version = "1.2.1", features = ["derive"] }
@@ -269,8 +270,7 @@ bstr = "1"
 buf-list = "1.0.3"
 bumpalo = "3.12.0"
 byte-unit = "5.1.6"
-#bytemuck = { version = "1", features = ["derive"] }
-bytemuck = { version = "1", features = ["derive", "extern_crate_alloc", "must_cast", "transparentwrapper_extra"] }
+bytemuck = { version = "1", features = ["derive"] }
 byteorder = "1.4.3"
 bytes = "1.5.0"
 bytesize = "1.1.0"
@@ -310,6 +310,7 @@ enumflags2 = { version = "0.7.7", features = ["serde"] }
 ethnum = { version = "1.5.1" }
 faststr = "0.2"
 feature-set = { version = "0.1.1" }
+feistel-permutation-rs = "0.1.1"
 flatbuffers = "25" # Must use the same version with arrow-ipc
 foreign_vec = "0.1.0"
 form_urlencoded = { version = "1" }
@@ -391,7 +392,7 @@ num = "0.4.0"
 num-bigint = "0.4.6"
 num-derive = "0.4.2"
 num-traits = "0.2.19"
-num_cpus = "1.13.1"
+num_cpus = "1.17"
 object = "0.36.5"
 object_store_opendal = { version = "0.52.0" }
 once_cell = "1.15.0"
@@ -470,6 +471,7 @@ rustls-pemfile = "2"
 rustls-pki-types = "1"
 rustyline = "14"
 scroll = "0.12.0"
+self_cell = "1.2.0"
 semver = "1.0.14"
 serde = { version = "1.0.164", features = ["derive", "rc"] }
 serde_derive = "1"
@@ -540,6 +542,7 @@ wiremock = "0.6"
 wkt = "0.11.1"
 xorf = { version = "0.11.0", default-features = false, features = ["binary-fuse"] }
 xorfilter-rs = "0.5"
+zerocopy = "0.8.26"
 zstd = "0.12.3"
 
 # AST needed

Cargo.toml

@@ -177,6 +177,24 @@ static BLOCK_INVERTED_INDEX_READ_MILLISECONDS: LazyLock<Histogram> = LazyLock::n
 static BLOCK_INVERTED_INDEX_SEARCH_MILLISECONDS: LazyLock<Histogram> = LazyLock::new(|| {
     register_histogram_in_milliseconds("fuse_block_inverted_index_search_milliseconds")
 });
+static BLOCK_VECTOR_INDEX_WRITE_NUMS: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_block_vector_index_write_nums"));
+static BLOCK_VECTOR_INDEX_WRITE_BYTES: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_block_vector_index_write_bytes"));
+static BLOCK_VECTOR_INDEX_WRITE_MILLISECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+    register_histogram_in_milliseconds("fuse_block_vector_index_write_milliseconds")
+});
+static BLOCK_VECTOR_INDEX_GENERATE_MILLISECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+    register_histogram_in_milliseconds("fuse_block_vector_index_generate_milliseconds")
+});
+static BLOCK_VECTOR_INDEX_READ_MILLISECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+    register_histogram_in_milliseconds("fuse_block_vector_index_read_milliseconds")
+});
+static BLOCK_VECTOR_INDEX_SEARCH_MILLISECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+    register_histogram_in_milliseconds("fuse_block_vector_index_search_milliseconds")
+});
+static BLOCK_VECTOR_INDEX_READ_BYTES: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_block_vector_index_read_bytes"));
 static COMPACT_BLOCK_READ_NUMS: LazyLock<Counter> =
     LazyLock::new(|| register_counter("fuse_compact_block_read_nums"));
 static COMPACT_BLOCK_READ_BYTES: LazyLock<Counter> =
@@ -548,6 +566,10 @@ pub fn metrics_inc_block_index_write_bytes(c: u64) {
     BLOCK_INDEX_WRITE_BYTES.inc_by(c);
 }
 
+pub fn metrics_inc_block_index_read_bytes(c: u64) {
+    BLOCK_INDEX_READ_BYTES.inc_by(c);
+}
+
 pub fn metrics_inc_block_index_write_milliseconds(c: u64) {
     BLOCK_INDEX_WRITE_MILLISECONDS.observe(c as f64);
 }
@@ -576,8 +598,32 @@ pub fn metrics_inc_block_inverted_index_search_milliseconds(c: u64) {
     BLOCK_INVERTED_INDEX_SEARCH_MILLISECONDS.observe(c as f64);
 }
 
-pub fn metrics_inc_block_index_read_bytes(c: u64) {
-    BLOCK_INDEX_READ_BYTES.inc_by(c);
+pub fn metrics_inc_block_vector_index_write_nums(c: u64) {
+    BLOCK_VECTOR_INDEX_WRITE_NUMS.inc_by(c);
+}
+
+pub fn metrics_inc_block_vector_index_write_bytes(c: u64) {
+    BLOCK_VECTOR_INDEX_WRITE_BYTES.inc_by(c);
+}
+
+pub fn metrics_inc_block_vector_index_write_milliseconds(c: u64) {
+    BLOCK_VECTOR_INDEX_WRITE_MILLISECONDS.observe(c as f64);
+}
+
+pub fn metrics_inc_block_vector_index_generate_milliseconds(c: u64) {
+    BLOCK_VECTOR_INDEX_GENERATE_MILLISECONDS.observe(c as f64);
+}
+
+pub fn metrics_inc_block_vector_index_read_milliseconds(c: u64) {
+    BLOCK_VECTOR_INDEX_READ_MILLISECONDS.observe(c as f64);
+}
+
+pub fn metrics_inc_block_vector_index_search_milliseconds(c: u64) {
+    BLOCK_VECTOR_INDEX_SEARCH_MILLISECONDS.observe(c as f64);
+}
+
+pub fn metrics_inc_block_vector_index_read_bytes(c: u64) {
+    BLOCK_VECTOR_INDEX_READ_BYTES.inc_by(c);
 }
 
 /// Compact metrics.

src/common/metrics/src/metrics/storage.rs

@@ -16,69 +16,101 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use ndarray::ArrayView;
 
-pub fn cosine_distance(from: &[f32], to: &[f32]) -> Result<f32> {
-    if from.len() != to.len() {
+pub fn cosine_distance(lhs: &[f32], rhs: &[f32]) -> Result<f32> {
+    if lhs.len() != rhs.len() {
         return Err(ErrorCode::InvalidArgument(format!(
             "Vector length not equal: {:} != {:}",
-            from.len(),
-            to.len(),
+            lhs.len(),
+            rhs.len(),
         )));
     }
 
-    let a = ArrayView::from(from);
-    let b = ArrayView::from(to);
+    let a = ArrayView::from(lhs);
+    let b = ArrayView::from(rhs);
     let aa_sum = (&a * &a).sum();
     let bb_sum = (&b * &b).sum();
 
     Ok(1.0 - (&a * &b).sum() / ((aa_sum).sqrt() * (bb_sum).sqrt()))
 }
 
-pub fn l2_distance(from: &[f32], to: &[f32]) -> Result<f32> {
-    if from.len() != to.len() {
+pub fn l1_distance(lhs: &[f32], rhs: &[f32]) -> Result<f32> {
+    if lhs.len() != rhs.len() {
         return Err(ErrorCode::InvalidArgument(format!(
             "Vector length not equal: {:} != {:}",
-            from.len(),
-            to.len(),
+            lhs.len(),
+            rhs.len(),
         )));
     }
 
-    Ok(from
+    Ok(lhs
         .iter()
-        .zip(to.iter())
+        .zip(rhs.iter())
+        .map(|(a, b)| (a - b).abs())
+        .sum::<f32>())
+}
+
+pub fn l2_distance(lhs: &[f32], rhs: &[f32]) -> Result<f32> {
+    if lhs.len() != rhs.len() {
+        return Err(ErrorCode::InvalidArgument(format!(
+            "Vector length not equal: {:} != {:}",
+            lhs.len(),
+            rhs.len(),
+        )));
+    }
+
+    Ok(lhs
+        .iter()
+        .zip(rhs.iter())
         .map(|(a, b)| (a - b).powi(2))
         .sum::<f32>()
         .sqrt())
 }
 
-pub fn cosine_distance_64(from: &[f64], to: &[f64]) -> Result<f64> {
-    if from.len() != to.len() {
+pub fn cosine_distance_64(lhs: &[f64], rhs: &[f64]) -> Result<f64> {
+    if lhs.len() != rhs.len() {
         return Err(ErrorCode::InvalidArgument(format!(
             "Vector length not equal: {:} != {:}",
-            from.len(),
-            to.len(),
+            lhs.len(),
+            rhs.len(),
         )));
     }
 
-    let a = ArrayView::from(from);
-    let b = ArrayView::from(to);
+    let a = ArrayView::from(lhs);
+    let b = ArrayView::from(rhs);
     let aa_sum = (&a * &a).sum();
     let bb_sum = (&b * &b).sum();
 
     Ok(1.0 - (&a * &b).sum() / ((aa_sum).sqrt() * (bb_sum).sqrt()))
 }
 
-pub fn l2_distance_64(from: &[f64], to: &[f64]) -> Result<f64> {
-    if from.len() != to.len() {
+pub fn l1_distance_64(lhs: &[f64], rhs: &[f64]) -> Result<f64> {
+    if lhs.len() != rhs.len() {
+        return Err(ErrorCode::InvalidArgument(format!(
+            "Vector length not equal: {:} != {:}",
+            lhs.len(),
+            rhs.len(),
+        )));
+    }
+
+    Ok(lhs
+        .iter()
+        .zip(rhs.iter())
+        .map(|(a, b)| (a - b).abs())
+        .sum::<f64>())
+}
+
+pub fn l2_distance_64(lhs: &[f64], rhs: &[f64]) -> Result<f64> {
+    if lhs.len() != rhs.len() {
         return Err(ErrorCode::InvalidArgument(format!(
             "Vector length not equal: {:} != {:}",
-            from.len(),
-            to.len(),
+            lhs.len(),
+            rhs.len(),
         )));
     }
 
-    Ok(from
+    Ok(lhs
         .iter()
-        .zip(to.iter())
+        .zip(rhs.iter())
         .map(|(a, b)| (a - b).powi(2))
         .sum::<f64>()
         .sqrt())

src/common/vector/src/distance.rs

@@ -16,5 +16,7 @@ mod distance;
 
 pub use distance::cosine_distance;
 pub use distance::cosine_distance_64;
+pub use distance::l1_distance;
+pub use distance::l1_distance_64;
 pub use distance::l2_distance;
 pub use distance::l2_distance_64;

src/common/vector/src/lib.rs

@@ -42,6 +42,7 @@ use databend_common_expression::SEARCH_MATCHED_COLUMN_ID;
 use databend_common_expression::SEARCH_SCORE_COLUMN_ID;
 use databend_common_expression::SEGMENT_NAME_COLUMN_ID;
 use databend_common_expression::SNAPSHOT_NAME_COLUMN_ID;
+use databend_common_expression::VECTOR_SCORE_COLUMN_ID;
 use databend_storages_common_table_meta::meta::try_extract_uuid_str_from_path;
 use databend_storages_common_table_meta::meta::NUM_BLOCK_ID_BITS;
 
@@ -106,6 +107,8 @@ pub struct InternalColumnMeta {
     pub inner: Option<BlockMetaInfoPtr>,
     // The search matched rows and optional scores in the block.
     pub matched_rows: Option<Vec<(usize, Option<F32>)>>,
+    // The vector topn rows and scores in the block.
+    pub vector_scores: Option<Vec<(usize, F32)>>,
 }
 
 #[typetag::serde(name = "internal_column_meta")]
@@ -142,6 +145,9 @@ pub enum InternalColumnType {
     SearchMatched,
     SearchScore,
 
+    // vector columns
+    VectorScore,
+
     FileName,
     FileRowNumber,
 }
@@ -176,6 +182,7 @@ impl InternalColumn {
             )),
             InternalColumnType::SearchMatched => TableDataType::Boolean,
             InternalColumnType::SearchScore => TableDataType::Number(NumberDataType::Float32),
+            InternalColumnType::VectorScore => TableDataType::Number(NumberDataType::Float32),
             InternalColumnType::FileName => TableDataType::String,
             InternalColumnType::FileRowNumber => TableDataType::Number(NumberDataType::UInt64),
         }
@@ -200,6 +207,7 @@ impl InternalColumn {
             InternalColumnType::BaseBlockIds => BASE_BLOCK_IDS_COLUMN_ID,
             InternalColumnType::SearchMatched => SEARCH_MATCHED_COLUMN_ID,
             InternalColumnType::SearchScore => SEARCH_SCORE_COLUMN_ID,
+            InternalColumnType::VectorScore => VECTOR_SCORE_COLUMN_ID,
             InternalColumnType::FileName => FILENAME_COLUMN_ID,
             InternalColumnType::FileRowNumber => FILE_ROW_NUMBER_COLUMN_ID,
         }
@@ -291,6 +299,20 @@ impl InternalColumn {
                 }
                 Float32Type::from_data(scores).into()
             }
+            InternalColumnType::VectorScore => {
+                assert!(meta.vector_scores.is_some());
+                let vector_scores = meta.vector_scores.as_ref().unwrap();
+
+                // The smaller the score, the closer the distance.
+                // Fill other rows with the maximum value and they will be filtered out.
+                let mut scores = vec![F32::from(f32::MAX); num_rows];
+                for (idx, score) in vector_scores.iter() {
+                    if let Some(val) = scores.get_mut(*idx) {
+                        *val = *score;
+                    }
+                }
+                Float32Type::from_data(scores).into()
+            }
             InternalColumnType::FileName | InternalColumnType::FileRowNumber => {
                 todo!("generate_column_values not support for file related")
             }