test(fuzzer): add supporting map functions (#11536)

Summary:

This diff adds some supporting functions for my diff for new benchmark tests for map_concat.

Reviewed By: darrenfu

Differential Revision: D65755011

velox/vector/fuzzer/VectorFuzzer.cpp

@@ -537,16 +537,8 @@ VectorPtr VectorFuzzer::fuzzFlat(const TypePtr& type, vector_size_t size) {
     // Do not initialize keys and values inline in the fuzzMap call as C++ does
     // not specify the order they'll be called in, leading to inconsistent
     // results across platforms.
-    const auto& keyType = type->asMap().keyType();
-    const auto& valueType = type->asMap().valueType();
-    auto length = getElementsVectorLength(opts_, size);
-
-    auto keys = opts_.normalizeMapKeys || !opts_.containerHasNulls
-        ? fuzzFlatNotNull(keyType, length)
-        : fuzzFlat(keyType, length);
-    auto values = opts_.containerHasNulls ? fuzzFlat(valueType, length)
-                                          : fuzzFlatNotNull(valueType, length);
-    return fuzzMap(keys, values, size);
+    auto keys = fuzzKeys(type->asMap().keyType(), size);
+    return fuzzMap(keys, type->asMap().valueType(), size);
   }
   // Rows.
   else if (type->isRow()) {
@@ -568,6 +560,27 @@ VectorPtr VectorFuzzer::fuzzFlat(const TypePtr& type, vector_size_t size) {
   }
 }
 
+VectorPtr VectorFuzzer::fuzzKeys(const TypePtr& keyType) {
+  return fuzzKeys(keyType, opts_.vectorSize);
+}
+
+VectorPtr VectorFuzzer::fuzzKeys(const TypePtr& keyType, vector_size_t size) {
+  auto length = getElementsVectorLength(opts_, size);
+  return opts_.normalizeMapKeys || !opts_.containerHasNulls
+      ? fuzzFlatNotNull(keyType, length)
+      : fuzzFlat(keyType, length);
+}
+
+VectorPtr VectorFuzzer::fuzzMap(
+    const VectorPtr keys,
+    const TypePtr valueType,
+    vector_size_t size) {
+  auto length = getElementsVectorLength(opts_, size);
+  auto values = opts_.containerHasNulls ? fuzzFlat(valueType, length)
+                                        : fuzzFlatNotNull(valueType, length);
+  return fuzzMap(keys, values, size);
+}
+
 VectorPtr VectorFuzzer::fuzzFlatPrimitive(
     const TypePtr& type,
     vector_size_t size) {
@@ -912,6 +925,10 @@ TypePtr VectorFuzzer::randType(
   return velox::randType(rng_, scalarTypes, maxDepth);
 }
 
+TypePtr VectorFuzzer::randMapType(int maxDepth) {
+  return velox::randMapType(rng_, defaultScalarTypes(), maxDepth);
+}
+
 RowTypePtr VectorFuzzer::randRowType(int maxDepth) {
   return velox::randRowType(rng_, maxDepth);
 }
@@ -922,6 +939,14 @@ RowTypePtr VectorFuzzer::randRowType(
   return velox::randRowType(rng_, scalarTypes, maxDepth);
 }
 
+RowTypePtr VectorFuzzer::randRowType(size_t numFields, int maxDepth) {
+  return velox::randRowType(rng_, numFields, maxDepth);
+}
+
+size_t VectorFuzzer::randNumColumns(size_t min, size_t max) {
+  return min + rand<size_t>(rng_) % (max - min + 1);
+}
+
 VectorPtr VectorFuzzer::wrapInLazyVector(VectorPtr baseVector) {
   if (hasNestedDictionaryLayers(baseVector)) {
     auto indices = baseVector->wrapInfo();
@@ -1139,16 +1164,22 @@ TypePtr randType(
   }
   switch (rand<uint32_t>(rng) % 3) {
     case 0:
-      return MAP(
-          randType(rng, scalarTypes, 0),
-          randType(rng, scalarTypes, maxDepth - 1));
+      return randMapType(rng, scalarTypes, maxDepth);
     case 1:
       return ARRAY(randType(rng, scalarTypes, maxDepth - 1));
     default:
       return randRowType(rng, scalarTypes, maxDepth - 1);
   }
 }
 
+TypePtr randMapType(
+    FuzzerGenerator& rng,
+    const std::vector<TypePtr>& scalarTypes,
+    int maxDepth) {
+  return MAP(
+      randType(rng, scalarTypes, 0), randType(rng, scalarTypes, maxDepth - 1));
+}
+
 TypePtr randOrderableType(FuzzerGenerator& rng, int maxDepth) {
   return randOrderableType(rng, defaultScalarTypes(), maxDepth);
 }
@@ -1185,10 +1216,22 @@ RowTypePtr randRowType(
     FuzzerGenerator& rng,
     const std::vector<TypePtr>& scalarTypes,
     int maxDepth) {
-  int numFields = 1 + rand<uint32_t>(rng) % 7;
+  size_t numFields = 1 + rand<uint32_t>(rng) % 7;
+  return randRowType(rng, scalarTypes, numFields, maxDepth);
+}
+
+RowTypePtr randRowType(FuzzerGenerator& rng, size_t numFields, int maxDepth) {
+  return randRowType(rng, defaultScalarTypes(), numFields, maxDepth);
+}
+
+RowTypePtr randRowType(
+    FuzzerGenerator& rng,
+    const std::vector<TypePtr>& scalarTypes,
+    size_t numFields,
+    int maxDepth) {
   std::vector<std::string> names;
   std::vector<TypePtr> fields;
-  for (int i = 0; i < numFields; ++i) {
+  for (size_t i = 0; i < numFields; ++i) {
     names.push_back(fmt::format("f{}", i));
     fields.push_back(randType(rng, scalarTypes, maxDepth));
   }

velox/vector/fuzzer/VectorFuzzer.h

@@ -202,6 +202,13 @@ class VectorFuzzer {
   VectorPtr fuzzFlatNotNull(const TypePtr& type);
   VectorPtr fuzzFlatNotNull(const TypePtr& type, vector_size_t size);
 
+  /// Returns a map vector with randomized values and nulls.Returns a vector
+  /// containing `opts_.vectorSize` or `size` elements.
+  VectorPtr
+  fuzzMap(const VectorPtr keys, const TypePtr valueType, vector_size_t size);
+  VectorPtr fuzzKeys(const TypePtr& keyType);
+  VectorPtr fuzzKeys(const TypePtr& keyType, vector_size_t size);
+
   /// Returns a random constant vector (which could be a null constant). Returns
   /// a vector with size set to `opts_.vectorSize` or 'size'.
   VectorPtr fuzzConstant(const TypePtr& type);
@@ -284,6 +291,11 @@ class VectorFuzzer {
   RowTypePtr randRowType(
       const std::vector<TypePtr>& scalarTypes,
       int maxDepth = 5);
+  RowTypePtr randRowType(size_t numFields, int maxDepth = 5);
+
+  size_t randNumColumns(size_t min, size_t max);
+
+  TypePtr randMapType(int maxDepth = 5);
 
   /// Generates short decimal TypePtr with random precision and scale.
   inline TypePtr randShortDecimalType() {
@@ -406,6 +418,11 @@ TypePtr randType(
     const std::vector<TypePtr>& scalarTypes,
     int maxDepth = 5);
 
+TypePtr randMapType(
+    FuzzerGenerator& rng,
+    const std::vector<TypePtr>& scalarTypes,
+    int maxDepth = 5);
+
 /// Same as the function above, but only generate orderable types.
 /// MAP types are not generated as they are not orderable.
 TypePtr randOrderableType(FuzzerGenerator& rng, int maxDepth = 5);
@@ -423,6 +440,14 @@ RowTypePtr randRowType(
     const std::vector<TypePtr>& scalarTypes,
     int maxDepth = 5);
 
+RowTypePtr randRowType(FuzzerGenerator& rng, size_t numFields, int maxDepth);
+
+RowTypePtr randRowType(
+    FuzzerGenerator& rng,
+    const std::vector<TypePtr>& scalarTypes,
+    size_t numFields,
+    int maxDepth = 5);
+
 /// Default set of scalar types to be chosen from when generating random types.
 const std::vector<TypePtr>& defaultScalarTypes();