Add test helper for creating a StoragePageSink

prestodb · Oct 23, 2015 · 1b93363 · 1b93363
1 parent f975b50
commit 1b93363
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 8 deletions.
diff --git a/presto-raptor/src/test/java/com/facebook/presto/raptor/storage/TestOrcStorageManager.java b/presto-raptor/src/test/java/com/facebook/presto/raptor/storage/TestOrcStorageManager.java
@@ -177,7 +177,7 @@ public void testWriter()
         List<Long> columnIds = ImmutableList.of(3L, 7L);
         List<Type> columnTypes = ImmutableList.<Type>of(BIGINT, VARCHAR);
 
-        StoragePageSink sink = manager.createStoragePageSink(columnIds, columnTypes);
+        StoragePageSink sink = createStoragePageSink(manager, columnIds, columnTypes);
         List<Page> pages = rowPagesBuilder(columnTypes)
                 .row(123, "hello")
                 .row(456, "bye")
@@ -239,7 +239,7 @@ public void testReader()
         byte[] bytes1 = octets(0x00, 0xFE, 0xFF);
         byte[] bytes3 = octets(0x01, 0x02, 0x19, 0x80);
 
-        StoragePageSink sink = manager.createStoragePageSink(columnIds, columnTypes);
+        StoragePageSink sink = createStoragePageSink(manager, columnIds, columnTypes);
 
         Object[][] doubles = {
                 {881, "-inf", null, null, null, Double.NEGATIVE_INFINITY},
@@ -311,7 +311,7 @@ public void testRewriter()
         List<Type> columnTypes = ImmutableList.<Type>of(BIGINT, VARCHAR);
 
         // create file with 2 rows
-        StoragePageSink sink = manager.createStoragePageSink(columnIds, columnTypes);
+        StoragePageSink sink = createStoragePageSink(manager, columnIds, columnTypes);
         List<Page> pages = rowPagesBuilder(columnTypes)
                 .row(123, "hello")
                 .row(456, "bye")
@@ -450,7 +450,7 @@ public void testMaxShardRows()
         List<Long> columnIds = ImmutableList.of(3L, 7L);
         List<Type> columnTypes = ImmutableList.<Type>of(BIGINT, VARCHAR);
 
-        StoragePageSink sink = manager.createStoragePageSink(columnIds, columnTypes);
+        StoragePageSink sink = createStoragePageSink(manager, columnIds, columnTypes);
         List<Page> pages = rowPagesBuilder(columnTypes)
                 .row(123, "hello")
                 .row(456, "bye")
@@ -473,11 +473,16 @@ public void testMaxFileSize()
 
         // Set maxFileSize to 1 byte, so adding any page makes the StoragePageSink full
         OrcStorageManager manager = createOrcStorageManager(storageService, backupStore, recoveryManager, 20, new DataSize(1, BYTE));
-        StoragePageSink sink = manager.createStoragePageSink(columnIds, columnTypes);
+        StoragePageSink sink = createStoragePageSink(manager, columnIds, columnTypes);
         sink.appendPages(pages);
         assertTrue(sink.isFull());
     }
 
+    private static StoragePageSink createStoragePageSink(StorageManager manager, List<Long> columnIds, List<Type> columnTypes)
+    {
+        return manager.createStoragePageSink(columnIds, columnTypes);
+    }
+
     private OrcStorageManager createOrcStorageManager()
     {
         return createOrcStorageManager(storageService, backupStore, recoveryManager);
@@ -581,7 +586,7 @@ private List<ColumnStats> columnStats(List<Type> columnTypes, Object[]... rows)
         List<Long> columnIds = list.build();
 
         OrcStorageManager manager = createOrcStorageManager();
-        StoragePageSink sink = manager.createStoragePageSink(columnIds, columnTypes);
+        StoragePageSink sink = createStoragePageSink(manager, columnIds, columnTypes);
         sink.appendPages(rowPagesBuilder(columnTypes).rows(rows).build());
         List<ShardInfo> shards = sink.commit();
 

diff --git a/presto-raptor/src/test/java/com/facebook/presto/raptor/storage/TestShardCompactor.java b/presto-raptor/src/test/java/com/facebook/presto/raptor/storage/TestShardCompactor.java
@@ -213,7 +213,7 @@ private MaterializedResult getMaterializedRows(Set<UUID> uuids, List<Long> colum
 
     private static List<ShardInfo> createShardsSorted(StorageManager storageManager, List<Long> columnIds, List<Type> columnTypes, List<Integer> sortChannels, List<SortOrder> sortOrders, int count, int length)
     {
-        StoragePageSink sink = storageManager.createStoragePageSink(columnIds, columnTypes);
+        StoragePageSink sink = createStoragePageSink(storageManager, columnIds, columnTypes);
         for (int numShards = 0; numShards < count; numShards++) {
             List<Page> pages = createPages(columnTypes, 1, length);
             long[] orderedAddresses = PAGE_SORTER.sort(columnTypes, pages, sortChannels, sortOrders, 10_000);
@@ -232,14 +232,19 @@ private static List<ShardInfo> createShardsSorted(StorageManager storageManager,
     private static List<ShardInfo> createShards(StorageManager storageManager, List<Long> columnIds, List<Type> columnTypes, int count, int length)
     {
         List<Page> pages = createPages(columnTypes, count, length);
-        StoragePageSink sink = storageManager.createStoragePageSink(columnIds, columnTypes);
+        StoragePageSink sink = createStoragePageSink(storageManager, columnIds, columnTypes);
         for (Page page : pages) {
             sink.appendPages(ImmutableList.of(page));
             sink.flush();
         }
         return sink.commit();
     }
 
+    private static StoragePageSink createStoragePageSink(StorageManager manager, List<Long> columnIds, List<Type> columnTypes)
+    {
+        return manager.createStoragePageSink(columnIds, columnTypes);
+    }
+
     private static List<Page> createPages(List<Type> columnTypes, int count, int length)
     {
         ImmutableList.Builder<Page> pages = ImmutableList.builder();