Added offset index based parquet reading support

extensions/parquet/base/src/main/java/io/deephaven/parquet/base/ColumnChunkReader.java

@@ -4,18 +4,15 @@
 package io.deephaven.parquet.base;
 
 import org.apache.parquet.column.Dictionary;
+import org.apache.parquet.internal.column.columnindex.OffsetIndex;
 import org.apache.parquet.schema.PrimitiveType;
+import org.jetbrains.annotations.Nullable;
 
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.function.Supplier;
 
 public interface ColumnChunkReader {
-    /**
-     * @return -1 if the current column doesn't guarantee fixed page size, otherwise the fixed page size
-     */
-    int getPageFixedSize();
-
     /**
      * @return The number of rows in this ColumnChunk, or -1 if it's unknown.
      */
@@ -32,16 +29,34 @@ public interface ColumnChunkReader {
      */
     int getMaxRl();
 
+    /**
+     * @return The offset index for this column chunk, or null if it not found in the metadata.
+     */
+    @Nullable
+    OffsetIndex getOffsetIndex();
+
     interface ColumnPageReaderIterator extends Iterator<ColumnPageReader>, AutoCloseable {
         @Override
-        void close() throws IOException;
+        void close();
     }
 
     /**
      * @return An iterator over individual parquet pages
      */
     ColumnPageReaderIterator getPageIterator() throws IOException;
 
+    interface ColumnPageDirectAccessor {
+        /**
+         * Directly access a page reader for a given page number.
+         */
+        ColumnPageReader getPageReader(final int pageNum);
+    }
+
+    /**
+     * @return An accessor for individual parquet pages
+     */
+    ColumnPageDirectAccessor getPageAccessor();
+
     /**
      * @return Whether this column chunk uses a dictionary-based encoding on every page
      */
@@ -69,4 +84,10 @@ public int getMaxId() {
     }
 
     PrimitiveType getType();
+
+    /**
+     * @return The "version" string from deephaven specific parquet metadata, or null if it's not present.
+     */
+    @Nullable
+    String getVersion();
 }

extensions/parquet/base/src/main/java/io/deephaven/parquet/base/ColumnChunkReaderImpl.java

@@ -28,6 +28,7 @@
 import java.nio.channels.SeekableByteChannel;
 import java.nio.file.Path;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.function.Supplier;
 
 import static org.apache.parquet.format.Encoding.PLAIN_DICTIONARY;
@@ -46,10 +47,18 @@ public class ColumnChunkReaderImpl implements ColumnChunkReader {
     private final PageMaterializer.Factory nullMaterializerFactory;
 
     private Path filePath;
-
-    ColumnChunkReaderImpl(
-            ColumnChunk columnChunk, SeekableChannelsProvider channelsProvider,
-            Path rootPath, MessageType type, OffsetIndex offsetIndex, List<Type> fieldTypes) {
+    /**
+     * Number of rows in the row group of this column chunk.
+     */
+    private final long numRows;
+    /**
+     * Version string from deephaven specific parquet metadata, or null if it's not present.
+     */
+    private final String version;
+
+    ColumnChunkReaderImpl(ColumnChunk columnChunk, SeekableChannelsProvider channelsProvider, Path rootPath,
+            MessageType type, OffsetIndex offsetIndex, List<Type> fieldTypes, final long numRows,
+            final String version) {
         this.channelsProvider = channelsProvider;
         this.columnChunk = columnChunk;
         this.rootPath = rootPath;
@@ -65,16 +74,13 @@ public class ColumnChunkReaderImpl implements ColumnChunkReader {
         this.fieldTypes = fieldTypes;
         this.dictionarySupplier = new LazyCachingSupplier<>(this::getDictionary);
         this.nullMaterializerFactory = PageMaterializer.factoryForType(path.getPrimitiveType().getPrimitiveTypeName());
-    }
-
-    @Override
-    public int getPageFixedSize() {
-        return -1;
+        this.numRows = numRows;
+        this.version = version;
     }
 
     @Override
     public long numRows() {
-        return numValues();
+        return numRows;
     }
 
     @Override
@@ -87,17 +93,28 @@ public int getMaxRl() {
         return path.getMaxRepetitionLevel();
     }
 
+    public final OffsetIndex getOffsetIndex() {
+        return offsetIndex;
+    }
+
     @Override
     public ColumnPageReaderIterator getPageIterator() {
         final long dataPageOffset = columnChunk.meta_data.getData_page_offset();
         if (offsetIndex == null) {
-            return new ColumnPageReaderIteratorImpl(dataPageOffset, columnChunk.getMeta_data().getNum_values(),
-                    path, channelsProvider);
+            return new ColumnPageReaderIteratorImpl(dataPageOffset, columnChunk.getMeta_data().getNum_values());
         } else {
-            return new ColumnPageReaderIteratorIndexImpl(path, channelsProvider);
+            return new ColumnPageReaderIteratorIndexImpl();
         }
     }
 
+    @Override
+    public final ColumnPageDirectAccessor getPageAccessor() {
+        if (offsetIndex == null) {
+            throw new UnsupportedOperationException("Cannot use direct accessor without offset index");
+        }
+        return new ColumnPageDirectAccessorImpl();
+    }
+
     private Path getFilePath() {
         if (filePath != null) {
             return filePath;
@@ -166,6 +183,11 @@ public PrimitiveType getType() {
         return path.getPrimitiveType();
     }
 
+    @Override
+    public String getVersion() {
+        return version;
+    }
+
     @NotNull
     private Dictionary readDictionary(ReadableByteChannel file) throws IOException {
         // explicitly not closing this, caller is responsible
@@ -192,21 +214,13 @@ private Dictionary readDictionary(ReadableByteChannel file) throws IOException {
         return dictionaryPage.getEncoding().initDictionary(path, dictionaryPage);
     }
 
-    class ColumnPageReaderIteratorImpl implements ColumnPageReaderIterator {
-        private final SeekableChannelsProvider channelsProvider;
+    private final class ColumnPageReaderIteratorImpl implements ColumnPageReaderIterator {
         private long currentOffset;
-        private final ColumnDescriptor path;
-
         private long remainingValues;
 
-        ColumnPageReaderIteratorImpl(final long startOffset,
-                final long numValues,
-                @NotNull final ColumnDescriptor path,
-                @NotNull final SeekableChannelsProvider channelsProvider) {
+        ColumnPageReaderIteratorImpl(final long startOffset, final long numValues) {
             this.remainingValues = numValues;
             this.currentOffset = startOffset;
-            this.path = path;
-            this.channelsProvider = channelsProvider;
         }
 
         @Override
@@ -267,15 +281,10 @@ nullMaterializerFactory, path, getFilePath(), fieldTypes, readChannel.position()
         public void close() {}
     }
 
-    class ColumnPageReaderIteratorIndexImpl implements ColumnPageReaderIterator {
-        private final SeekableChannelsProvider channelsProvider;
+    private final class ColumnPageReaderIteratorIndexImpl implements ColumnPageReaderIterator {
         private int pos;
-        private final ColumnDescriptor path;
 
-        ColumnPageReaderIteratorIndexImpl(ColumnDescriptor path,
-                SeekableChannelsProvider channelsProvider) {
-            this.path = path;
-            this.channelsProvider = channelsProvider;
+        ColumnPageReaderIteratorIndexImpl() {
             pos = 0;
         }
 
@@ -287,20 +296,39 @@ public boolean hasNext() {
         @Override
         public ColumnPageReader next() {
             if (!hasNext()) {
-                throw new RuntimeException("No next element");
+                throw new NoSuchElementException("No next element");
             }
-            int rowCount =
-                    (int) (offsetIndex.getLastRowIndex(pos, columnChunk.getMeta_data().getNum_values())
-                            - offsetIndex.getFirstRowIndex(pos) + 1);
+            // Following logic assumes that offsetIndex will store the number of values for a page instead of number
+            // of rows (which can be different for array and vector columns). This behavior is because of a bug on
+            // parquet writing side which got fixed in deephaven-core/pull/4844 and is only kept to support reading
+            // parquet files written before deephaven-core/pull/4844.
+            final int numValues = (int) (offsetIndex.getLastRowIndex(pos, columnChunk.getMeta_data().getNum_values())
+                    - offsetIndex.getFirstRowIndex(pos) + 1);
             ColumnPageReaderImpl columnPageReader =
                     new ColumnPageReaderImpl(channelsProvider, decompressor, dictionarySupplier,
                             nullMaterializerFactory, path, getFilePath(), fieldTypes, offsetIndex.getOffset(pos), null,
-                            rowCount);
+                            numValues);
             pos++;
             return columnPageReader;
         }
 
         @Override
-        public void close() throws IOException {}
+        public void close() {}
+    }
+
+    private final class ColumnPageDirectAccessorImpl implements ColumnPageDirectAccessor {
+
+        ColumnPageDirectAccessorImpl() {}
+
+        @Override
+        public ColumnPageReader getPageReader(final int pageNum) {
+            if (pageNum < 0 || pageNum >= offsetIndex.getPageCount()) {
+                throw new IndexOutOfBoundsException(
+                        "pageNum=" + pageNum + ", offsetIndex.getPageCount()=" + offsetIndex.getPageCount());
+            }
+            final int numValues = -1; // Will be populated properly when we read the page header
+            return new ColumnPageReaderImpl(channelsProvider, decompressor, dictionarySupplier, nullMaterializerFactory,
+                    path, getFilePath(), fieldTypes, offsetIndex.getOffset(pageNum), null, numValues);
+        }
     }
 }

extensions/parquet/base/src/main/java/io/deephaven/parquet/base/ColumnPageReaderImpl.java

@@ -54,11 +54,19 @@ public class ColumnPageReaderImpl implements ColumnPageReader {
     private final Path filePath;
     private final List<Type> fieldTypes;
 
+    /**
+     * Stores the offset from where the next byte should be read. Can be the offset of page header if
+     * {@link #pageHeader} is {@code null}, else will be the offset of data.
+     */
     private long offset;
     private PageHeader pageHeader;
     private int numValues;
     private int rowCount = -1;
 
+    /**
+     * @param offset The offset for page header if supplied {@code pageHeader} is {@code null}. Else, the offset of data
+     *        in the page.
+     */
     ColumnPageReaderImpl(SeekableChannelsProvider channelsProvider,
             CompressorAdapter compressorAdapter,
             Supplier<Dictionary> dictionarySupplier,
@@ -108,9 +116,17 @@ public IntBuffer readKeyValues(IntBuffer keyDest, int nullPlaceholder) throws IO
         }
     }
 
+    /**
+     * If {@link #pageHeader} is {@code null}, read it from the file and increment the {@link #offset} and file position
+     * by the length of page header. This method assumes that file position is set to {@link #offset} before calling.
+     * This method also read the number of values in the page from the header.
+     */
     private synchronized void ensurePageHeader(SeekableByteChannel file) throws IOException {
         if (pageHeader == null) {
-            offset = file.position();
+            if (file.position() != offset) {
+                throw new IllegalStateException("File position = " + file.position() + " not equal to expected offset ="
+                        + offset);
+            }
             int maxHeader = START_HEADER;
             boolean success;
             do {
@@ -132,6 +148,13 @@ private synchronized void ensurePageHeader(SeekableByteChannel file) throws IOEx
                 }
             } while (!success);
             file.position(offset);
+            if (numValues >= 0) {
+                // Make sure the number of values are same as those in the header
+                if (numValues != readNumValuesFromPageHeader(pageHeader)) {
+                    throw new IllegalStateException("numValues = " + numValues + " different from number of values " +
+                            "read from the page header");
+                }
+            }
         }
         ensureNumValues();
     }
@@ -569,25 +592,34 @@ private ValuesReader getDataReader(Encoding dataEncoding, ByteBuffer in, int val
 
     @Override
     public int numValues() throws IOException {
-        ensureNumValues();
-        return numValues;
+        if (numValues >= 0) {
+            return numValues;
+        }
+        try (final SeekableByteChannel readChannel = channelsProvider.getReadChannel(filePath)) {
+            readChannel.position(offset);
+            ensurePageHeader(readChannel);
+            // Above will automatically populate numValues
+            Assert.geq(numValues, "numValues", 0);
+            return numValues;
+        }
     }
 
     private void ensureNumValues() throws IOException {
         if (numValues >= 0) {
             return;
         }
         Assert.neqNull(pageHeader, "pageHeader");
-        switch (pageHeader.type) {
+        numValues = readNumValuesFromPageHeader(pageHeader);
+    }
+
+    private static int readNumValuesFromPageHeader(@NotNull final PageHeader header) throws IOException {
+        switch (header.type) {
             case DATA_PAGE:
-                numValues = pageHeader.getData_page_header().getNum_values();
-                break;
+                return header.getData_page_header().getNum_values();
             case DATA_PAGE_V2:
-                numValues = pageHeader.getData_page_header_v2().getNum_values();
-                break;
+                return header.getData_page_header_v2().getNum_values();
             default:
-                throw new IOException(
-                        String.format("Unexpected page of type {%s}", pageHeader.getType()));
+                throw new IOException(String.format("Unexpected page of type {%s}", header.getType()));
         }
     }