Support DWRF stripes above MAX_INT rows in reader

OrcReader fails when a stripe has more than INT32_MAX rows.
Both ORC/DWRF format uses long for stripe number of rows.
We ran into this edge case on a file with only string as a column
and file is produced by a different writer compatible with
DWRF spec.

Author: Arunachalam Thirupathi

presto-orc/src/main/java/com/facebook/presto/orc/AbstractOrcRecordReader.java

@@ -689,7 +689,7 @@ private void validateWrite(Predicate<OrcWriteValidation> test, String messageFor
         }
     }
 
-    private void validateWriteStripe(int rowCount)
+    private void validateWriteStripe(long rowCount)
     {
         if (writeChecksumBuilder.isPresent()) {
             writeChecksumBuilder.get().addStripe(rowCount);

presto-orc/src/main/java/com/facebook/presto/orc/OrcWriteValidation.java

@@ -481,10 +481,10 @@ public static WriteChecksumBuilder createWriteChecksumBuilder(Map<Integer, Type>
             return new WriteChecksumBuilder(types.build());
         }
 
-        public void addStripe(int rowCount)
+        public void addStripe(long rowCount)
         {
-            longSlice.setInt(0, rowCount);
-            stripeHash.update(longBuffer, 0, Integer.BYTES);
+            longSlice.setLong(0, rowCount);
+            stripeHash.update(longBuffer, 0, Long.BYTES);
         }
 
         public void addPage(Page page)
@@ -888,7 +888,7 @@ public OrcWriteValidationBuilder addMetadataProperty(String key, Slice value)
             return this;
         }
 
-        public OrcWriteValidationBuilder addStripe(int rowCount)
+        public OrcWriteValidationBuilder addStripe(long rowCount)
         {
             checksum.addStripe(rowCount);
             return this;

presto-orc/src/main/java/com/facebook/presto/orc/OrcWriterFlushStats.java

@@ -57,7 +57,7 @@ public DistributionStat getDictionaryBytes()
         return dictionaryBytes;
     }
 
-    public void recordStripeWritten(long stripeBytes, int stripeRows, int dictionaryBytes)
+    public void recordStripeWritten(long stripeBytes, long stripeRows, int dictionaryBytes)
     {
         this.stripeBytes.add(stripeBytes);
         this.stripeRows.add(stripeRows);

presto-orc/src/main/java/com/facebook/presto/orc/OrcWriterStats.java

@@ -45,7 +45,7 @@ public void recordStripeWritten(
             StripeInformation stripeInformation)
     {
         long stripeBytes = stripeInformation.getTotalLength();
-        int stripeRows = stripeInformation.getNumberOfRows();
+        long stripeRows = stripeInformation.getNumberOfRows();
         getFlushStats(flushReason).recordStripeWritten(stripeBytes, stripeRows, dictionaryBytes);
         allFlush.recordStripeWritten(stripeBytes, stripeRows, dictionaryBytes);
     }

presto-orc/src/main/java/com/facebook/presto/orc/StripeReader.java

@@ -420,7 +420,7 @@ public InputStreamSources createDictionaryStreamSources(Map<StreamId, Stream> st
     }
 
     private List<RowGroup> createRowGroups(
-            int rowsInStripe,
+            long rowsInStripe,
             Map<StreamId, Stream> streams,
             Map<StreamId, ValueInputStream<?>> valueStreams,
             Map<StreamId, List<RowGroupIndex>> columnIndexes,
@@ -433,7 +433,7 @@ private List<RowGroup> createRowGroups(
         for (int rowGroupId : selectedRowGroups) {
             Map<StreamId, StreamCheckpoint> checkpoints = getStreamCheckpoints(includedOrcColumns, types, decompressor.isPresent(), rowGroupId, encodings, streams, columnIndexes);
             int rowOffset = rowGroupId * rowsInRowGroup;
-            int rowsInGroup = Math.min(rowsInStripe - rowOffset, rowsInRowGroup);
+            int rowsInGroup = toIntExact(Math.min(rowsInStripe - rowOffset, rowsInRowGroup));
             long minAverageRowBytes = columnIndexes
                     .entrySet()
                     .stream()
@@ -536,13 +536,13 @@ private Map<StreamId, List<RowGroupIndex>> readColumnIndexes(Map<StreamId, Strea
 
     private Set<Integer> selectRowGroups(StripeInformation stripe, Map<StreamId, List<RowGroupIndex>> columnIndexes)
     {
-        int rowsInStripe = toIntExact(stripe.getNumberOfRows());
+        long rowsInStripe = stripe.getNumberOfRows();
         int groupsInStripe = ceil(rowsInStripe, rowsInRowGroup);
 
         ImmutableSet.Builder<Integer> selectedRowGroups = ImmutableSet.builder();
-        int remainingRows = rowsInStripe;
+        long remainingRows = rowsInStripe;
         for (int rowGroup = 0; rowGroup < groupsInStripe; ++rowGroup) {
-            int rows = Math.min(remainingRows, rowsInRowGroup);
+            int rows = toIntExact(Math.min(remainingRows, rowsInRowGroup));
             Map<Integer, ColumnStatistics> statistics = getRowGroupStatistics(types.get(0), columnIndexes, rowGroup);
             if (predicate.matches(rows, statistics)) {
                 selectedRowGroups.add(rowGroup);
@@ -613,9 +613,10 @@ public static Map<StreamId, DiskRange> getDiskRanges(List<List<Stream>> streams)
     /**
      * Ceiling of integer division
      */
-    private static int ceil(int dividend, int divisor)
+    private static int ceil(long dividend, int divisor)
     {
-        return ((dividend + divisor) - 1) / divisor;
+        long ceil = ((dividend + divisor) - 1) / divisor;
+        return toIntExact(ceil);
     }
 
     public static class StripeId

presto-orc/src/main/java/com/facebook/presto/orc/metadata/DwrfMetadataReader.java

@@ -286,7 +286,7 @@ private static StripeInformation toStripeInformation(DwrfProto.StripeInformation
             keyMetadata = previousKeyMetadata;
         }
         return new StripeInformation(
-                toIntExact(stripeInformation.getNumberOfRows()),
+                stripeInformation.getNumberOfRows(),
                 stripeInformation.getOffset(),
                 stripeInformation.getIndexLength(),
                 stripeInformation.getDataLength(),

presto-orc/src/main/java/com/facebook/presto/orc/metadata/DwrfMetadataWriter.java

@@ -24,6 +24,7 @@
 import com.facebook.presto.orc.proto.DwrfProto.UserMetadataItem;
 import com.facebook.presto.orc.protobuf.ByteString;
 import com.facebook.presto.orc.protobuf.MessageLite;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.io.CountingOutputStream;
@@ -136,7 +137,8 @@ public int writeFooter(SliceOutput output, Footer footer)
         return writeProtobufObject(output, footerProtobuf.build());
     }
 
-    private static DwrfProto.StripeInformation toStripeInformation(StripeInformation stripe)
+    @VisibleForTesting
+    static DwrfProto.StripeInformation toStripeInformation(StripeInformation stripe)
     {
         return DwrfProto.StripeInformation.newBuilder()
                 .setNumberOfRows(stripe.getNumberOfRows())

presto-orc/src/main/java/com/facebook/presto/orc/metadata/OrcMetadataReader.java

@@ -163,7 +163,7 @@ private static List<StripeInformation> toStripeInformation(List<OrcProto.StripeI
     private static StripeInformation toStripeInformation(OrcProto.StripeInformation stripeInformation)
     {
         return new StripeInformation(
-                toIntExact(stripeInformation.getNumberOfRows()),
+                stripeInformation.getNumberOfRows(),
                 stripeInformation.getOffset(),
                 stripeInformation.getIndexLength(),
                 stripeInformation.getDataLength(),

presto-orc/src/main/java/com/facebook/presto/orc/metadata/StripeInformation.java

@@ -23,7 +23,7 @@
 
 public class StripeInformation
 {
-    private final int numberOfRows;
+    private final long numberOfRows;
     private final long offset;
     private final long indexLength;
     private final long dataLength;
@@ -34,7 +34,7 @@ public class StripeInformation
     // only set for run start, and reuse until next run
     private final List<byte[]> keyMetadata;
 
-    public StripeInformation(int numberOfRows, long offset, long indexLength, long dataLength, long footerLength, List<byte[]> keyMetadata)
+    public StripeInformation(long numberOfRows, long offset, long indexLength, long dataLength, long footerLength, List<byte[]> keyMetadata)
     {
         // dataLength can be zero when the stripe only contains empty flat maps.
         checkArgument(numberOfRows > 0, "Stripe must have at least one row");
@@ -48,7 +48,7 @@ public StripeInformation(int numberOfRows, long offset, long indexLength, long d
         this.keyMetadata = ImmutableList.copyOf(requireNonNull(keyMetadata, "keyMetadata is null"));
     }
 
-    public int getNumberOfRows()
+    public long getNumberOfRows()
     {
         return numberOfRows;
     }

presto-orc/src/test/java/com/facebook/presto/orc/metadata/TestDwrfMetadataReader.java

@@ -16,7 +16,6 @@
 import com.facebook.presto.orc.DwrfEncryptionProvider;
 import com.facebook.presto.orc.DwrfKeyProvider;
 import com.facebook.presto.orc.OrcCorruptionException;
-import com.facebook.presto.orc.OrcDataSource;
 import com.facebook.presto.orc.OrcDataSourceId;
 import com.facebook.presto.orc.metadata.PostScript.HiveWriterVersion;
 import com.facebook.presto.orc.metadata.statistics.StringStatistics;
@@ -30,6 +29,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 
@@ -126,6 +126,32 @@ public void testToStripeCacheMode()
         assertEquals(DwrfMetadataReader.toStripeCacheMode(DwrfProto.StripeCacheMode.NA), DwrfStripeCacheMode.NONE);
     }
 
+    @Test
+    public void testStripeInformationRows()
+            throws IOException
+    {
+        long aLongNumber = Integer.MAX_VALUE + 1000L;
+        StripeInformation expectedStripeInformation = new StripeInformation(aLongNumber, aLongNumber, aLongNumber, aLongNumber, aLongNumber, ImmutableList.of());
+        DwrfProto.StripeInformation dwrfStripeInformation = DwrfMetadataWriter.toStripeInformation(expectedStripeInformation);
+
+        DwrfProto.Footer protoFooter = DwrfProto.Footer.newBuilder()
+                .setNumberOfRows(aLongNumber)
+                .setRowIndexStride(10_000)
+                .addStripes(dwrfStripeInformation)
+                .build();
+
+        Footer footer = convertToFooter(protoFooter);
+        assertEquals(footer.getNumberOfRows(), aLongNumber);
+        assertEquals(footer.getStripes().size(), 1);
+        StripeInformation actualStripeInformation = footer.getStripes().get(0);
+        assertEquals(actualStripeInformation.getOffset(), expectedStripeInformation.getOffset());
+        assertEquals(actualStripeInformation.getNumberOfRows(), expectedStripeInformation.getNumberOfRows());
+        assertEquals(actualStripeInformation.getIndexLength(), expectedStripeInformation.getIndexLength());
+        assertEquals(actualStripeInformation.getDataLength(), expectedStripeInformation.getDataLength());
+        assertEquals(actualStripeInformation.getFooterLength(), expectedStripeInformation.getFooterLength());
+        assertEquals(actualStripeInformation.getTotalLength(), expectedStripeInformation.getTotalLength());
+    }
+
     @Test
     public void testReadFooter()
             throws IOException
@@ -139,20 +165,26 @@ public void testReadFooter()
                 .setRowIndexStride(rowIndexStride)
                 .addAllStripeCacheOffsets(stripeCacheOffsets)
                 .build();
+        Footer footer = convertToFooter(protoFooter);
+
+        assertEquals(footer.getNumberOfRows(), numberOfRows);
+        assertEquals(footer.getRowsInRowGroup(), rowIndexStride);
+        assertEquals(footer.getDwrfStripeCacheOffsets().get(), stripeCacheOffsets);
+        assertEquals(footer.getStripes(), Collections.emptyList());
+    }
+
+    private Footer convertToFooter(DwrfProto.Footer protoFooter)
+            throws IOException
+    {
         byte[] data = protoFooter.toByteArray();
         InputStream inputStream = new ByteArrayInputStream(data);
-        OrcDataSource orcDataSource = null; // orcDataSource only needed for encrypted files
 
-        Footer footer = dwrfMetadataReader.readFooter(HiveWriterVersion.ORC_HIVE_8732,
+        return dwrfMetadataReader.readFooter(HiveWriterVersion.ORC_HIVE_8732,
                 inputStream,
                 DwrfEncryptionProvider.NO_ENCRYPTION,
                 DwrfKeyProvider.EMPTY,
-                orcDataSource,
+                null, // orcDataSource only needed for encrypted files
                 Optional.empty());
-
-        assertEquals(footer.getNumberOfRows(), numberOfRows);
-        assertEquals(footer.getRowsInRowGroup(), rowIndexStride);
-        assertEquals(footer.getDwrfStripeCacheOffsets().get(), stripeCacheOffsets);
     }
 
     @Test