Improve BenchmarkPartitionedOutputOperator

core/trino-main/src/test/java/io/trino/block/BlockAssertions.java

@@ -355,6 +355,16 @@ public static <T> List<T> generateListWithNulls(int positionCount, float nullRat
         return unmodifiableList(result);
     }
 
+    public static Set<Integer> chooseNullPositions(int positionCount, float nullRate)
+    {
+        int nullCount = (int) (positionCount * nullRate);
+        if (nullCount == 0) {
+            verify(nullRate == 0 || positionCount == 0, "position count %s too small to have at least one null with rate %s", (Object) positionCount, nullRate);
+            return ImmutableSet.of();
+        }
+        return chooseRandomUnique(positionCount, nullCount);
+    }
+
     public static Block createStringsBlock(String... values)
     {
         requireNonNull(values, "values is null");
@@ -871,16 +881,6 @@ private interface ValueWriter<T>
         void write(BlockBuilder builder, T value);
     }
 
-    private static Set<Integer> chooseNullPositions(int positionCount, float nullRate)
-    {
-        int nullCount = (int) (positionCount * nullRate);
-        if (nullCount == 0) {
-            verify(nullRate == 0, "position count %s too small to have at least one null with rate %s", (Object) positionCount, nullRate);
-            return ImmutableSet.of();
-        }
-        return chooseRandomUnique(positionCount, nullCount);
-    }
-
     private static Set<Integer> chooseRandomUnique(int bound, int count)
     {
         if (count < bound / 10) {

core/trino-main/src/test/java/io/trino/operator/PageTestUtils.java

@@ -75,24 +75,32 @@ public static Page createRandomPage(
             float nullRate,
             Optional<Wrapping> wrapping)
     {
-        int channelCount = types.size();
-        ImmutableList.Builder<Block> blocks = ImmutableList.builder();
+        List<Block> blocks = types.stream()
+                .map(type -> {
+                    Block block = createRandomBlockForType(type, positionCount, nullRate);
+                    return wrapping.map(w -> w.wrap(block, positionCount)).orElse(block);
+                })
+                .collect(toImmutableList());
 
-        for (int i = 0; i < channelCount; i++) {
-            Block block = createRandomBlockForType(types.get(i), positionCount, nullRate);
-            blocks.add(wrapping.map(w -> w.wrap(block, positionCount)).orElse(block));
-        }
+        return createPage(types, positionCount, hashChannels, blocks);
+    }
 
-        hashChannels.ifPresent(channels -> {
-            ImmutableList<Block> blocksWithoutHash = blocks.build();
+    public static Page createPage(
+            List<Type> types,
+            int positionCount,
+            Optional<List<Integer>> hashChannels,
+            List<Block> blocks)
+    {
+        ImmutableList.Builder<Block> finalBlocks = ImmutableList.<Block>builder().addAll(blocks);
 
-            blocks.add(getHashBlock(
+        hashChannels.ifPresent(channels -> {
+            finalBlocks.add(getHashBlock(
                     channels.stream()
                             .map(types::get)
                             .collect(toImmutableList()),
-                    channels.stream().map(blocksWithoutHash::get).toArray(Block[]::new)));
+                    channels.stream().map(blocks::get).toArray(Block[]::new)));
         });
 
-        return new Page(positionCount, blocks.build().toArray(Block[]::new));
+        return new Page(positionCount, finalBlocks.build().toArray(Block[]::new));
     }
 }