add ut

marin-ma · Apr 9, 2024 · f450417 · f450417
1 parent d945e1d
commit f450417
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 6 deletions.
diff --git a/velox/functions/sparksql/Hash.cpp b/velox/functions/sparksql/Hash.cpp
@@ -638,6 +638,8 @@ void checkArgTypes(const std::vector<exec::VectorFunctionArg>& args) {
       case TypeKind::HUGEINT:
       case TypeKind::TIMESTAMP:
       case TypeKind::ARRAY:
+      case TypeKind::MAP:
+      case TypeKind::ROW:
         break;
       default:
         VELOX_USER_FAIL("Unsupported type for hash: {}", arg.type->toString())

diff --git a/velox/functions/sparksql/tests/XxHash64Test.cpp b/velox/functions/sparksql/tests/XxHash64Test.cpp
@@ -29,8 +29,8 @@ class XxHash64Test : public SparkFunctionBaseTest {
     return evaluateOnce<int64_t>("xxhash64(c0)", arg);
   }
 
-  VectorPtr xxhash64(std::shared_ptr<ArrayVector> array) {
-    return evaluate("xxhash64(c0)", makeRowVector({array}));
+  VectorPtr xxhash64(VectorPtr vector) {
+    return evaluate("xxhash64(c0)", makeRowVector({vector}));
   }
 };
 
@@ -145,15 +145,18 @@ TEST_F(XxHash64Test, float) {
 }
 
 TEST_F(XxHash64Test, array) {
-  assertEqualVectors(
-      makeFlatVector<int64_t>(std::vector<int64_t>{4904562767517797033}),
-      xxhash64(makeArrayVector<int64_t>({{1, 2, 3}})));
-
   assertEqualVectors(
       makeFlatVector<int64_t>(
           std::vector<int64_t>{-6041664978295882827, 42, 4904562767517797033}),
       xxhash64(makeArrayVector<int64_t>({{1, 2, 3, 4, 5}, {}, {1, 2, 3}})));
 
+  assertEqualVectors(
+      makeFlatVector<int64_t>(
+          std::vector<int64_t>{-6698625589789238999, 8420071140774656230, 42}),
+      xxhash64(makeNullableArrayVector<int32_t>(
+          {{1, std::nullopt}, {std::nullopt, 2}, {std::nullopt}})));
+
+  // Nested array.
   {
     using innerArrayType = std::vector<std::optional<int64_t>>;
     using outerArrayType =
@@ -174,6 +177,7 @@ TEST_F(XxHash64Test, array) {
         xxhash64(arrayVector));
   }
 
+  // Array of map.
   {
     using S = StringView;
     using P = std::pair<int64_t, std::optional<S>>;
@@ -186,6 +190,7 @@ TEST_F(XxHash64Test, array) {
         xxhash64(arrayVector));
   }
 
+  // Array of row.
   {
     std::vector<std::vector<std::optional<std::tuple<int32_t, std::string>>>>
         data = {
@@ -201,6 +206,57 @@ TEST_F(XxHash64Test, array) {
   }
 }
 
+TEST_F(XxHash64Test, map) {
+  auto mapVector = makeMapVector<int64_t, double>(
+      {{{1, 17.0}, {2, 36.0}, {3, 8.0}, {4, 28.0}, {5, 24.0}, {6, 32.0}}});
+  assertEqualVectors(
+      makeFlatVector<int64_t>(std::vector<int64_t>{-6303587702533348160}),
+      xxhash64(mapVector));
+
+  auto mapOfArrays = createMapOfArraysVector<int32_t, int32_t>(
+      {{{1, {{1, 2, 3}}}}, {{2, {{4, 5, 6}}}}, {{3, {{7, 8, 9}}}}});
+  assertEqualVectors(
+      makeFlatVector<int64_t>(std::vector<int64_t>{
+          -2103781794412908874, 1112887818746642853, 5787852566364222439}),
+      xxhash64(mapOfArrays));
+
+  auto mapWithNullArrays = createMapOfArraysVector<int64_t, int64_t>(
+      {{{1, std::nullopt}}, {{2, {{4, 5, std::nullopt}}}}, {{3, {{7, 8, 9}}}}});
+  assertEqualVectors(
+      makeFlatVector<int64_t>(std::vector<int64_t>{
+          -7001672635703045582, 7217681953522744649, 5785528104873330081}),
+      xxhash64(mapWithNullArrays));
+}
+
+TEST_F(XxHash64Test, row) {
+  auto row = makeRowVector({
+      makeFlatVector<int64_t>({1, 3}),
+      makeFlatVector<int64_t>({2, 4}),
+  });
+  assertEqualVectors(
+      makeFlatVector<int64_t>(
+          std::vector<int64_t>{-8198029865082835910, 351067884137457704}),
+      xxhash64(row));
+
+  row = makeRowVector({
+      makeNullableFlatVector<int64_t>({1, std::nullopt}),
+      makeNullableFlatVector<int64_t>({std::nullopt, 4}),
+  });
+  assertEqualVectors(
+      makeFlatVector<int64_t>(
+          std::vector<int64_t>{-7001672635703045582, 404280023041566627}),
+      xxhash64(row));
+
+  row->setNull(0, true);
+  assertEqualVectors(
+      makeFlatVector<int64_t>(std::vector<int64_t>{42, 404280023041566627}),
+      xxhash64(row));
+
+  row->setNull(1, true);
+  assertEqualVectors(
+      makeFlatVector<int64_t>(std::vector<int64_t>{42, 42}), xxhash64(row));
+}
+
 TEST_F(XxHash64Test, hashSeed) {
   auto xxhash64WithSeed = [&](int64_t seed, const std::optional<int64_t>& arg) {
     return evaluateOnce<int64_t>(