Add the dot_product UDF (facebookincubator#15971)

Summary:
Pull Request resolved: facebookincubator#15971

Add a DOT_PRODUCT function to compute the dot product between two arrays or maps, enabling efficient vector operations for ML and analytics workloads.

## Function Signatures
- `dot_product(array(T), array(T)) -> bigint/double`
- `dot_product(map(K, V), map(K, V)) -> bigint/double`

## Key Features
- **Array dot product**: Computes sum of element-wise products of corresponding elements
- **Map dot product**: Multiplies values with matching keys and sums the results
- **Equal length requirement**: Arrays must have the same length (throws error otherwise)
- **Null handling**: Null elements are treated as zero, null arguments return null
- **Type support**: int8, int16, int32, int64, float, double for arrays; integer/varchar keys with integer/double values for maps
- **Overflow protection**: Uses checked arithmetic for integer operations

## Behavior Examples
```sql
SELECT dot_product(ARRAY[1, 2, 3], ARRAY[4, 5, 6]); -- 32 (1*4 + 2*5 + 3*6)
SELECT dot_product(ARRAY[1.0, 2.0], ARRAY[3.0, 4.0]); -- 11.0
SELECT dot_product(ARRAY[1, NULL, 3], ARRAY[4, 5, 6]); -- 22 (nulls treated as 0)
SELECT dot_product(MAP(ARRAY[1, 2], ARRAY[10, 20]), MAP(ARRAY[1, 2], ARRAY[3, 4])); -- 110
```

Reviewed By: zacw7

Differential Revision: D90476184

fbshipit-source-id: 8bffe8a70c52748d62d9fb36e039ad89291006db

Author: abhinavmuk04

velox/docs/functions/presto/array.rst

@@ -272,6 +272,30 @@ Array Functions
           SELECT l2_norm(MAP(ARRAY[1, 2], ARRAY[3.0, 4.0])); -- 5.0
           SELECT l2_norm(MAP(ARRAY[], ARRAY[])); -- 0.0
 
+.. function:: dot_product(array(T), array(T)) -> bigint/double
+
+    Computes the dot product of two arrays. The dot product is the sum of element-wise
+    products of corresponding elements. Both arrays must have the same length.
+    If either array is null, returns null. If arrays have different lengths, throws an error.
+    Null elements in arrays are treated as zero.
+    Returns bigint for integer arrays, double for floating-point arrays. ::
+
+          SELECT dot_product(ARRAY[1, 2, 3], ARRAY[4, 5, 6]); -- 32 (1*4 + 2*5 + 3*6)
+          SELECT dot_product(ARRAY[1.0, 2.0], ARRAY[3.0, 4.0]); -- 11.0 (1.0*3.0 + 2.0*4.0)
+          SELECT dot_product(ARRAY[1, NULL, 3], ARRAY[4, 5, 6]); -- 22 (1*4 + 0*5 + 3*6)
+          SELECT dot_product(ARRAY[], ARRAY[]); -- 0
+
+.. function:: dot_product(map(K, V), map(K, V)) -> bigint/double
+
+    Computes the dot product of two maps. For maps, the dot product is computed by
+    multiplying values with matching keys and summing the results. Keys present in only
+    one map contribute zero to the result. If either map is null, returns null.
+    Null values in maps are treated as zero.
+    Returns bigint for integer value maps, double for floating-point value maps. ::
+
+          SELECT dot_product(MAP(ARRAY[1, 2], ARRAY[10, 20]), MAP(ARRAY[1, 2], ARRAY[3, 4])); -- 110 (10*3 + 20*4)
+          SELECT dot_product(MAP(ARRAY['a', 'b'], ARRAY[1.0, 2.0]), MAP(ARRAY['a', 'c'], ARRAY[3.0, 4.0])); -- 3.0 (only 'a' matches)
+
 .. function:: array_sum(array(T)) -> bigint/double
 
     Returns the sum of all non-null elements of the array. If there is no non-null elements, returns 0. The behaviour is similar to aggregation function sum().

velox/expression/fuzzer/ExpressionFuzzerTest.cpp

@@ -302,6 +302,7 @@ std::unordered_set<std::string> skipFunctionsSOT = {
     "array_subset", // Velox-only function, not available in Presto
     "map_values_in_range", // Velox-only function, not available in Presto
     "transform_with_index", // Velox-only function, not available in Presto
+    "dot_product", // Velox-only function, not available in Presto
     "remap_keys", // Velox-only function, not available in Presto
     "map_intersect", // Velox-only function, not available in Presto
     "map_keys_overlap", // Velox-only function, not available in Presto

velox/functions/prestosql/DotProduct.h

@@ -0,0 +1,127 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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.
+ */
+#pragma once
+
+#include <cstdint>
+
+#include <folly/CPortability.h>
+#include <folly/container/F14Map.h>
+
+#include "velox/common/base/Exceptions.h"
+#include "velox/functions/Macros.h"
+#include "velox/functions/lib/CheckedArithmetic.h"
+#include "velox/type/SimpleFunctionApi.h"
+
+namespace facebook::velox::functions {
+
+/// Computes the dot product of two arrays.
+/// The dot product is the sum of element-wise products of corresponding
+/// elements. Both arrays must have the same length. If either array is null,
+/// returns null. If arrays have different lengths, throws an error.
+/// Null elements in arrays are treated as zero.
+template <typename TExec, typename T>
+struct DotProductFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(TExec);
+
+  template <typename TOutput>
+  FOLLY_ALWAYS_INLINE bool call(
+      TOutput& out,
+      const arg_type<Array<T>>& array1,
+      const arg_type<Array<T>>& array2) {
+    const auto size1 = array1.size();
+    const auto size2 = array2.size();
+
+    VELOX_USER_CHECK_EQ(
+        size1,
+        size2,
+        "dot_product requires arrays of equal length, but got {} and {}",
+        size1,
+        size2);
+
+    TOutput sum = 0;
+    for (vector_size_t i = 0; i < size1; ++i) {
+      const auto& val1 = array1[i];
+      const auto& val2 = array2[i];
+
+      if (val1.has_value() && val2.has_value()) {
+        if constexpr (std::is_same_v<TOutput, int64_t>) {
+          auto product = checkedMultiply<TOutput>(
+              static_cast<TOutput>(val1.value()),
+              static_cast<TOutput>(val2.value()));
+          sum = checkedPlus<TOutput>(sum, product);
+        } else {
+          sum += static_cast<TOutput>(val1.value()) *
+              static_cast<TOutput>(val2.value());
+        }
+      }
+    }
+    out = sum;
+    return true;
+  }
+};
+
+/// Computes the dot product of two maps.
+/// For maps, the dot product is computed by multiplying values with matching
+/// keys and summing the results. Keys present in only one map contribute zero
+/// to the result. If either map is null, returns null.
+/// Null values in maps are treated as zero.
+template <typename TExec, typename K, typename V>
+struct MapDotProductFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(TExec);
+
+  template <typename TOutput>
+  FOLLY_ALWAYS_INLINE bool call(
+      TOutput& out,
+      const arg_type<Map<K, V>>& map1,
+      const arg_type<Map<K, V>>& map2) {
+    TOutput sum = 0;
+
+    // Build a lookup map from map2 for O(1) key lookup.
+    // This reduces complexity from O(n*m) to O(n+m).
+    folly::F14FastMap<arg_type<K>, arg_type<V>> map2Lookup;
+    map2Lookup.reserve(map2.size());
+    for (const auto& [key, val] : map2) {
+      if (val.has_value()) {
+        map2Lookup.emplace(key, val.value());
+      }
+    }
+
+    // Iterate through map1 and look up matching keys in map2.
+    for (const auto& [key1, val1] : map1) {
+      if (!val1.has_value()) {
+        continue;
+      }
+
+      auto it = map2Lookup.find(key1);
+      if (it != map2Lookup.end()) {
+        if constexpr (std::is_same_v<TOutput, int64_t>) {
+          auto product = checkedMultiply<TOutput>(
+              static_cast<TOutput>(val1.value()),
+              static_cast<TOutput>(it->second));
+          sum = checkedPlus<TOutput>(sum, product);
+        } else {
+          sum += static_cast<TOutput>(val1.value()) *
+              static_cast<TOutput>(it->second);
+        }
+      }
+    }
+
+    out = sum;
+    return true;
+  }
+};
+
+} // namespace facebook::velox::functions

velox/functions/prestosql/registration/ArrayFunctionsRegistration.cpp

@@ -26,6 +26,7 @@
 #include "velox/functions/prestosql/ArrayFunctions.h"
 #include "velox/functions/prestosql/ArraySort.h"
 #include "velox/functions/prestosql/ArraySubset.h"
+#include "velox/functions/prestosql/DotProduct.h"
 #include "velox/functions/prestosql/L2Norm.h"
 #include "velox/functions/prestosql/WidthBucketArray.h"
 #include "velox/functions/prestosql/types/JsonRegistration.h"
@@ -448,5 +449,66 @@ void registerArrayFunctions(const std::string& prefix) {
       ParameterBinder<MapL2NormFunction, int64_t, double>,
       double,
       Map<int64_t, double>>({prefix + "l2_norm"});
+
+  // Register dot_product for integer arrays only.
+  // Float and double array versions already exist in
+  // MathematicalFunctionsRegistration.cpp (DotProductArray,
+  // DotProductFloatArray) with different semantics: they return NaN for empty
+  // arrays to maintain compatibility with cosine_similarity and other distance
+  // functions there. Integer versions here return 0 for empty arrays.
+  registerFunction<
+      ParameterBinder<DotProductFunction, int8_t>,
+      int64_t,
+      Array<int8_t>,
+      Array<int8_t>>({prefix + "dot_product"});
+  registerFunction<
+      ParameterBinder<DotProductFunction, int16_t>,
+      int64_t,
+      Array<int16_t>,
+      Array<int16_t>>({prefix + "dot_product"});
+  registerFunction<
+      ParameterBinder<DotProductFunction, int32_t>,
+      int64_t,
+      Array<int32_t>,
+      Array<int32_t>>({prefix + "dot_product"});
+  registerFunction<
+      ParameterBinder<DotProductFunction, int64_t>,
+      int64_t,
+      Array<int64_t>,
+      Array<int64_t>>({prefix + "dot_product"});
+
+  // Register dot_product for maps with integer keys
+  registerFunction<
+      ParameterBinder<MapDotProductFunction, int32_t, int64_t>,
+      int64_t,
+      Map<int32_t, int64_t>,
+      Map<int32_t, int64_t>>({prefix + "dot_product"});
+  registerFunction<
+      ParameterBinder<MapDotProductFunction, int64_t, int64_t>,
+      int64_t,
+      Map<int64_t, int64_t>,
+      Map<int64_t, int64_t>>({prefix + "dot_product"});
+  registerFunction<
+      ParameterBinder<MapDotProductFunction, int32_t, double>,
+      double,
+      Map<int32_t, double>,
+      Map<int32_t, double>>({prefix + "dot_product"});
+  registerFunction<
+      ParameterBinder<MapDotProductFunction, int64_t, double>,
+      double,
+      Map<int64_t, double>,
+      Map<int64_t, double>>({prefix + "dot_product"});
+
+  // Register dot_product for maps with varchar keys
+  registerFunction<
+      ParameterBinder<MapDotProductFunction, Varchar, int64_t>,
+      int64_t,
+      Map<Varchar, int64_t>,
+      Map<Varchar, int64_t>>({prefix + "dot_product"});
+  registerFunction<
+      ParameterBinder<MapDotProductFunction, Varchar, double>,
+      double,
+      Map<Varchar, double>,
+      Map<Varchar, double>>({prefix + "dot_product"});
 }
 } // namespace facebook::velox::functions

velox/functions/prestosql/tests/CMakeLists.txt

@@ -51,6 +51,7 @@ add_executable(
   ArraySumTest.cpp
   ArrayTrimTest.cpp
   ArrayUnionTest.cpp
+  DotProductTest.cpp
   ArgTypesGeneratorTest.cpp
   BinaryFunctionsTest.cpp
   BingTileCastTest.cpp