feat: Add cosine_similarity function for two arrays (facebookincubator#13311)

Summary:

Similar to prestodb/presto#25056

Reviewed By: yuandagits

Differential Revision: D74549208

Author: amitkdutta

velox/docs/functions/presto/math.rst

@@ -38,6 +38,19 @@ Mathematical Functions
 
         SELECT cosine_similarity(MAP(ARRAY[], ARRAY[]), MAP(ARRAY['a', 'b'], ARRAY[2, 3])); -- NaN
 
+.. function:: cosine_similarity(array(double), array(double)) -> double
+
+    Returns the `cosine similarity <https://en.wikipedia.org/wiki/Cosine_similarity>`_ between the vectors represented as array(double).
+    If any input array is empty, the function returns NaN. If the input arrays have different sizes, the function throws VeloxUserError.
+
+        SELECT cosine_similarity(ARRAY[1], ARRAY[2]); -- 1.0
+
+        SELECT cosine_similarity(ARRAY[1.0, 2.0], ARRAY[NULL, 3.0]); -- NULL
+
+        SELECT cosine_similarity(ARRAY[], ARRAY[2, 3]); -- Throws VeloxUserError
+
+        SELECT cosine_similarity(ARRAY[], ARRAY[]); -- NaN
+
 .. function:: degrees(x) -> double
 
     Converts angle x in radians to degrees.

velox/functions/prestosql/Arithmetic.h

@@ -656,18 +656,18 @@ struct WilsonIntervalLowerFunction {
 };
 
 template <typename T>
-struct CosineSimilarityFunction {
+struct CosineSimilarityFunctionMap {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
-  double normalizeMap(const null_free_arg_type<Map<Varchar, double>>& map) {
+  double normalize(const null_free_arg_type<Map<Varchar, double>>& map) {
     double norm = 0.0;
     for (const auto& [key, value] : map) {
       norm += (value * value);
     }
     return std::sqrt(norm);
   }
 
-  double mapDotProduct(
+  double dotProduct(
       const null_free_arg_type<Map<Varchar, double>>& leftMap,
       const null_free_arg_type<Map<Varchar, double>>& rightMap) {
     double result = 0.0;
@@ -689,20 +689,67 @@ struct CosineSimilarityFunction {
       return;
     }
 
-    double normLeftMap = normalizeMap(leftMap);
+    double normLeftMap = normalize(leftMap);
     if (normLeftMap == 0.0) {
       result = std::numeric_limits<double>::quiet_NaN();
       return;
     }
 
-    double normRightMap = normalizeMap(rightMap);
+    double normRightMap = normalize(rightMap);
     if (normRightMap == 0.0) {
       result = std::numeric_limits<double>::quiet_NaN();
       return;
     }
 
-    double dotProduct = mapDotProduct(leftMap, rightMap);
-    result = dotProduct / (normLeftMap * normRightMap);
+    double product = dotProduct(leftMap, rightMap);
+    result = product / (normLeftMap * normRightMap);
+  }
+};
+
+template <typename T>
+struct CosineSimilarityFunctionArray {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  double normalize(const null_free_arg_type<Array<double>>& map) {
+    double norm = 0.0;
+    for (const auto value : map) {
+      norm += (value * value);
+    }
+    return std::sqrt(norm);
+  }
+
+  double dotProduct(
+      const null_free_arg_type<Array<double>>& leftArray,
+      const null_free_arg_type<Array<double>>& rightArray) {
+    double result = 0.0;
+    for (size_t i = 0; i < leftArray.size(); i++) {
+      result += leftArray[i] * rightArray[i];
+    }
+    return result;
+  }
+
+  void callNullFree(
+      out_type<double>& result,
+      const null_free_arg_type<Array<double>>& leftArray,
+      const null_free_arg_type<Array<double>>& rightArray) {
+    VELOX_USER_CHECK(
+        leftArray.size() == rightArray.size(),
+        "Both arrays need to have identical size");
+
+    double normLeftArray = normalize(leftArray);
+    if (normLeftArray == 0.0) {
+      result = std::numeric_limits<double>::quiet_NaN();
+      return;
+    }
+
+    double normRightArray = normalize(rightArray);
+    if (normRightArray == 0.0) {
+      result = std::numeric_limits<double>::quiet_NaN();
+      return;
+    }
+
+    double product = dotProduct(leftArray, rightArray);
+    result = product / (normLeftArray * normRightArray);
   }
 };
 

velox/functions/prestosql/registration/MathematicalFunctionsRegistration.cpp

@@ -114,10 +114,15 @@ void registerMathFunctions(const std::string& prefix) {
   registerTruncate({prefix + "truncate"});
 
   registerFunction<
-      CosineSimilarityFunction,
+      CosineSimilarityFunctionMap,
       double,
       Map<Varchar, double>,
       Map<Varchar, double>>({prefix + "cosine_similarity"});
+  registerFunction<
+      CosineSimilarityFunctionArray,
+      double,
+      Array<double>,
+      Array<double>>({prefix + "cosine_similarity"});
 }
 
 } // namespace

velox/functions/prestosql/tests/ArithmeticTest.cpp

@@ -1088,5 +1088,32 @@ TEST_F(ArithmeticTest, cosineSimilarity) {
                    .has_value());
 }
 
+TEST_F(ArithmeticTest, cosineSimilarityArray) {
+  const auto cosineSimilarity = [&](const std::vector<double>& left,
+                                    const std::vector<double>& right) {
+    auto leftMap = makeArrayVector<double>({left});
+    auto rightMap = makeArrayVector<double>({right});
+    return evaluateOnce<double>(
+               "cosine_similarity(c0,c1)", makeRowVector({leftMap, rightMap}))
+        .value();
+  };
+
+  EXPECT_DOUBLE_EQ(
+      (1 * 1 * 1 + 2 * 3) / (std::sqrt(5.0) * std::sqrt(10.0)),
+      cosineSimilarity({{1, 2}}, {{1, 3}}));
+
+  EXPECT_DOUBLE_EQ(
+      (1 * 1 + 2 * 3 + (-1) * 5) /
+          (std::sqrt(1 + 4 + 1) * std::sqrt(1 + 9 + 25)),
+      cosineSimilarity({{1, 2, -1}}, {{1, 3, 5}}));
+
+  EXPECT_TRUE(std::isnan(cosineSimilarity({}, {})));
+  VELOX_ASSERT_THROW(
+      cosineSimilarity({1, 3}, {}), "Both arrays need to have identical size");
+  EXPECT_TRUE(std::isnan(cosineSimilarity({1, 3}, {0, 0})));
+  EXPECT_TRUE(std::isnan(cosineSimilarity({1, 3}, {kNan, 1})));
+  EXPECT_TRUE(std::isnan(cosineSimilarity({1, 3}, {kInf, 1})));
+}
+
 } // namespace
 } // namespace facebook::velox