feat: add cosine_distance for vector similarity compute (#10737)

* feat: add cosine_distance for vector similarity compute

* fix cosine_distance return value

* cosine_distance support empty array

* change the distance from vec to one

* fix function list.txt

* change to ndarray

* change to ndarray

* change to array view avoid mem copy

* add logic test

* fix unit test

---------

Co-authored-by: baishen <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 3 people

Cargo.lock

@@ -26,6 +26,7 @@ members = [
     "src/common/tracing",
     "src/common/storage",
     "src/common/profile",
+    "src/common/vector",
     # Query
     "src/query/ast",
     "src/query/codegen",

Cargo.toml

@@ -28,5 +28,6 @@ pub mod containers;
 pub mod mem_allocator;
 pub mod rangemap;
 pub mod runtime;
+
 pub use runtime::match_join_handle;
 pub use runtime::set_alloc_error_hook;

src/common/base/src/lib.rs

@@ -0,0 +1,23 @@
+[package]
+name = "common-vector"
+version = { workspace = true }
+authors = { workspace = true }
+license = { workspace = true }
+publish = { workspace = true }
+edition = { workspace = true }
+
+[lib]
+doctest = false
+test = false
+
+[dependencies] # In alphabetical order
+common-exception = { path = "../exception" }
+
+ndarray = "0.15.6"
+
+[build-dependencies]
+
+[features]
+
+[dev-dependencies]
+approx = "0.5.1"

src/common/vector/Cargo.toml

@@ -0,0 +1,34 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_exception::ErrorCode;
+use common_exception::Result;
+use ndarray::ArrayView;
+
+pub fn cosine_distance(from: &[f32], to: &[f32]) -> Result<f32> {
+    if from.len() != to.len() {
+        return Err(ErrorCode::InvalidArgument(format!(
+            "Vector length not equal: {:} != {:}",
+            from.len(),
+            to.len(),
+        )));
+    }
+
+    let a = ArrayView::from(from);
+    let b = ArrayView::from(to);
+    let aa_sum = (&a * &a).sum();
+    let bb_sum = (&b * &b).sum();
+
+    Ok((&a * &b).sum() / ((aa_sum).sqrt() * (bb_sum).sqrt()))
+}

src/common/vector/src/distance.rs

@@ -0,0 +1,17 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+mod distance;
+
+pub use distance::cosine_distance;

src/common/vector/src/lib.rs

@@ -0,0 +1,41 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_vector::cosine_distance;
+
+#[test]
+fn test_cosine() {
+    {
+        let x: Vec<f32> = (1..9).map(|v| v as f32).collect();
+        let y: Vec<f32> = (100..108).map(|v| v as f32).collect();
+        let d = cosine_distance(&x, &y).unwrap();
+        // from scipy.spatial.distance.cosine
+        approx::assert_relative_eq!(d, 0.900_957);
+    }
+
+    {
+        let x = vec![3.0, 45.0, 7.0, 2.0, 5.0, 20.0, 13.0, 12.0];
+        let y = vec![2.0, 54.0, 13.0, 15.0, 22.0, 34.0, 50.0, 1.0];
+        let d = cosine_distance(&x, &y).unwrap();
+        // from sklearn.metrics.pairwise import cosine_similarity
+        approx::assert_relative_eq!(d, 0.873_580_6);
+    }
+
+    {
+        let x = vec![3.0, 45.0, 7.0, 2.0, 5.0, 20.0, 13.0, 12.0];
+        let y = vec![2.0, 54.0];
+        let d = cosine_distance(&x, &y);
+        assert!(d.is_err());
+    }
+}

src/common/vector/tests/it/distance.rs

@@ -0,0 +1,15 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+mod distance;

src/common/vector/tests/it/main.rs

@@ -17,6 +17,7 @@ common-exception = { path = "../../common/exception" }
 common-expression = { path = "../expression" }
 common-hashtable = { path = "../../common/hashtable" }
 common-io = { path = "../../common/io" }
+common-vector = { path = "../../common/vector" }
 jsonb = { workspace = true }
 
 # Crates.io dependencies

src/query/functions/Cargo.toml

@@ -25,6 +25,7 @@ mod map;
 mod math;
 mod tuple;
 mod variant;
+mod vector;
 
 mod comparison;
 mod decimal;
@@ -55,4 +56,5 @@ pub fn register(registry: &mut FunctionRegistry) {
     hash::register(registry);
     other::register(registry);
     decimal::register(registry);
+    vector::register(registry);
 }