Added support to write metadata files in parquet

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

@@ -31,10 +31,17 @@ public interface ColumnChunkReader {
     int getMaxRl();
 
     /**
-     * @return The offset index for this column chunk, or null if it not found in the metadata.
+     * @return Whether the column chunk has offset index information set in the metadata or not.
      */
-    @Nullable
-    OffsetIndex getOffsetIndex();
+    boolean hasOffsetIndex();
+
+    /**
+     * Get the offset index for a column chunk.
+     *
+     * @param context The channel context to use for reading the offset index.
+     *
+     */
+    OffsetIndex getOffsetIndex(final SeekableChannelContext context);
 
     /**
      * Used to iterate over column page readers for each page with the capability to set channel context to for reading
@@ -69,9 +76,9 @@ interface ColumnPageDirectAccessor {
     }
 
     /**
-     * @return An accessor for individual parquet pages.
+     * @return An accessor for individual parquet pages which uses the provided offset index.
      */
-    ColumnPageDirectAccessor getPageAccessor();
+    ColumnPageDirectAccessor getPageAccessor(OffsetIndex offsetIndex);
 
     /**
      * @return Whether this column chunk uses a dictionary-based encoding on every page.

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

@@ -27,31 +27,24 @@
 import java.io.UncheckedIOException;
 import java.net.URI;
 import java.nio.channels.SeekableByteChannel;
-import java.nio.file.Path;
 import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.function.Function;
 
-import static io.deephaven.parquet.base.ParquetFileReader.FILE_URI_SCHEME;
 import static org.apache.parquet.format.Encoding.PLAIN_DICTIONARY;
 import static org.apache.parquet.format.Encoding.RLE_DICTIONARY;
 
 final class ColumnChunkReaderImpl implements ColumnChunkReader {
 
     private final ColumnChunk columnChunk;
     private final SeekableChannelsProvider channelsProvider;
-    /**
-     * If reading a single parquet file, root URI is the URI of the file, else the parent directory for a metadata file
-     */
-    private final URI rootURI;
     private final CompressorAdapter decompressor;
     private final ColumnDescriptor path;
-    private final OffsetIndex offsetIndex;
+    private OffsetIndexReader offsetIndexReader;
     private final List<Type> fieldTypes;
     private final Function<SeekableChannelContext, Dictionary> dictionarySupplier;
     private final PageMaterializerFactory nullMaterializerFactory;
-
-    private URI uri;
+    private final URI columnChunkURI;
     /**
      * Number of rows in the row group of this column chunk.
      */
@@ -61,12 +54,11 @@ final class ColumnChunkReaderImpl implements ColumnChunkReader {
      */
     private final String version;
 
-    ColumnChunkReaderImpl(ColumnChunk columnChunk, SeekableChannelsProvider channelsProvider, URI rootURI,
-            MessageType type, OffsetIndex offsetIndex, List<Type> fieldTypes, final long numRows,
-            final String version) {
+    ColumnChunkReaderImpl(ColumnChunk columnChunk, SeekableChannelsProvider channelsProvider, URI columnChunkURI,
+            MessageType type, List<Type> fieldTypes, final long numRows, final String version) {
         this.channelsProvider = channelsProvider;
         this.columnChunk = columnChunk;
-        this.rootURI = rootURI;
+        this.columnChunkURI = columnChunkURI;
         this.path = type
                 .getColumnDescription(columnChunk.meta_data.getPath_in_schema().toArray(new String[0]));
         if (columnChunk.getMeta_data().isSetCodec()) {
@@ -75,12 +67,15 @@ final class ColumnChunkReaderImpl implements ColumnChunkReader {
         } else {
             decompressor = CompressorAdapter.PASSTHRU;
         }
-        this.offsetIndex = offsetIndex;
         this.fieldTypes = fieldTypes;
         this.dictionarySupplier = new LazyCachingFunction<>(this::getDictionary);
         this.nullMaterializerFactory = PageMaterializer.factoryForType(path.getPrimitiveType().getPrimitiveTypeName());
         this.numRows = numRows;
         this.version = version;
+        // Construct the reader object but don't read the offset index yet
+        this.offsetIndexReader = (columnChunk.isSetOffset_index_offset())
+                ? new OffsetIndexReaderImpl(channelsProvider, columnChunk, columnChunkURI)
+                : OffsetIndexReader.NULL;
     }
 
     @Override
@@ -98,8 +93,15 @@ public int getMaxRl() {
         return path.getMaxRepetitionLevel();
     }
 
-    public OffsetIndex getOffsetIndex() {
-        return offsetIndex;
+    @Override
+    public boolean hasOffsetIndex() {
+        return columnChunk.isSetOffset_index_offset();
+    }
+
+    @Override
+    public OffsetIndex getOffsetIndex(final SeekableChannelContext context) {
+        // Read the offset index if it hasn't been read yet
+        return offsetIndexReader.getOffsetIndex(context);
     }
 
     @Override
@@ -108,23 +110,15 @@ public ColumnPageReaderIterator getPageIterator() {
     }
 
     @Override
-    public final ColumnPageDirectAccessor getPageAccessor() {
+    public ColumnPageDirectAccessor getPageAccessor(final OffsetIndex offsetIndex) {
         if (offsetIndex == null) {
             throw new UnsupportedOperationException("Cannot use direct accessor without offset index");
         }
-        return new ColumnPageDirectAccessorImpl();
+        return new ColumnPageDirectAccessorImpl(offsetIndex);
     }
 
     private URI getURI() {
-        if (uri != null) {
-            return uri;
-        }
-        if (columnChunk.isSetFile_path() && FILE_URI_SCHEME.equals(rootURI.getScheme())) {
-            return uri = Path.of(rootURI).resolve(columnChunk.getFile_path()).toUri();
-        } else {
-            // TODO(deephaven-core#5066): Add support for reading metadata files from non-file URIs
-            return uri = rootURI;
-        }
+        return columnChunkURI;
     }
 
     @Override
@@ -307,7 +301,11 @@ private static int getNumValues(PageHeader pageHeader) {
 
     private final class ColumnPageDirectAccessorImpl implements ColumnPageDirectAccessor {
 
-        ColumnPageDirectAccessorImpl() {}
+        private final OffsetIndex offsetIndex;
+
+        ColumnPageDirectAccessorImpl(final OffsetIndex offsetIndex) {
+            this.offsetIndex = offsetIndex;
+        }
 
         @Override
         public ColumnPageReader getPageReader(final int pageNum, final SeekableChannelContext channelContext) {

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

@@ -0,0 +1,25 @@
+package io.deephaven.parquet.base;
+
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+
+import java.io.File;
+
+/**
+ * A no-op implementation of MetadataFileWriterBase when we don't want to write metadata files for Parquet files.
+ */
+public final class NullParquetMetadataFileWriter implements ParquetMetadataFileWriter {
+
+    public static final NullParquetMetadataFileWriter INSTANCE = new NullParquetMetadataFileWriter();
+
+    private NullParquetMetadataFileWriter() {}
+
+    @Override
+    public void addParquetFileMetadata(final File parquetFile, final ParquetMetadata metadata) {}
+
+    @Override
+    public void writeMetadataFiles(final File metadataFile, final File commonMetadataFile) {}
+
+    @Override
+    public void clear() {}
+}
+

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

@@ -0,0 +1,22 @@
+package io.deephaven.parquet.base;
+
+import io.deephaven.util.channel.SeekableChannelContext;
+import org.apache.parquet.internal.column.columnindex.OffsetIndex;
+
+/**
+ * Interface for reading the offset index for a column chunk.
+ */
+public interface OffsetIndexReader {
+
+    /**
+     * Get the offset index for a column chunk.
+     *
+     * @param context The channel context to use for reading the offset index.
+     */
+    OffsetIndex getOffsetIndex(SeekableChannelContext context);
+
+    /**
+     * A null implementation of the offset index reader.
+     */
+    OffsetIndexReader NULL = context -> null;
+}

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

@@ -0,0 +1,55 @@
+package io.deephaven.parquet.base;
+
+import io.deephaven.util.channel.SeekableChannelContext;
+import io.deephaven.util.channel.SeekableChannelsProvider;
+import org.apache.parquet.format.ColumnChunk;
+import org.apache.parquet.format.converter.ParquetMetadataConverter;
+import org.apache.parquet.internal.column.columnindex.OffsetIndex;
+import org.apache.parquet.format.Util;
+import org.jetbrains.annotations.NotNull;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UncheckedIOException;
+import java.net.URI;
+import java.nio.channels.SeekableByteChannel;
+
+/**
+ * Implementation of {@link OffsetIndexReader}, which reads the offset index for a column chunk on demand.
+ */
+final class OffsetIndexReaderImpl implements OffsetIndexReader {
+
+    private final SeekableChannelsProvider channelsProvider;
+    private final ColumnChunk chunk;
+    private final URI columnChunkURI;
+    private OffsetIndex offsetIndex;
+
+    OffsetIndexReaderImpl(final SeekableChannelsProvider channelsProvider, final ColumnChunk chunk,
+            final URI columnChunkURI) {
+        this.channelsProvider = channelsProvider;
+        this.chunk = chunk;
+        this.columnChunkURI = columnChunkURI;
+        this.offsetIndex = null;
+    }
+
+    @Override
+    public OffsetIndex getOffsetIndex(@NotNull final SeekableChannelContext context) {
+        if (offsetIndex != null) {
+            return offsetIndex;
+        }
+        return readOffsetIndex(context);
+    }
+
+    private OffsetIndex readOffsetIndex(@NotNull final SeekableChannelContext channelContext) {
+        try (
+                final SeekableChannelContext.ContextHolder holder =
+                        SeekableChannelContext.ensureContext(channelsProvider, channelContext);
+                final SeekableByteChannel readChannel = channelsProvider.getReadChannel(holder.get(), columnChunkURI);
+                final InputStream in =
+                        channelsProvider.getInputStream(readChannel.position(chunk.getOffset_index_offset()))) {
+            return (offsetIndex = ParquetMetadataConverter.fromParquetOffsetIndex(Util.readOffsetIndex(in)));
+        } catch (final IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+}

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

@@ -27,7 +27,7 @@
 public class ParquetFileReader {
     private static final int FOOTER_LENGTH_SIZE = 4;
     private static final String MAGIC_STR = "PAR1";
-    static final byte[] MAGIC = MAGIC_STR.getBytes(StandardCharsets.US_ASCII);
+    public static final byte[] MAGIC = MAGIC_STR.getBytes(StandardCharsets.US_ASCII);
     public static final String FILE_URI_SCHEME = "file";
 
     public final FileMetaData fileMetaData;

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

@@ -16,7 +16,9 @@
 import org.apache.parquet.internal.column.columnindex.OffsetIndex;
 import org.apache.parquet.internal.hadoop.metadata.IndexReference;
 import org.apache.parquet.schema.MessageType;
+import org.jetbrains.annotations.NotNull;
 
+import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -38,23 +40,29 @@ public final class ParquetFileWriter {
     private final Map<String, String> extraMetaData;
     private final List<BlockMetaData> blocks = new ArrayList<>();
     private final List<List<OffsetIndex>> offsetIndexes = new ArrayList<>();
+    private final File metadataFilePath;
+    private final ParquetMetadataFileWriter metadataFileWriter;
 
     public ParquetFileWriter(
-            final String filePath,
+            final File destFile,
+            final File metadataFilePath,
             final SeekableChannelsProvider channelsProvider,
             final int targetPageSize,
             final ByteBufferAllocator allocator,
             final MessageType type,
             final String codecName,
-            final Map<String, String> extraMetaData) throws IOException {
+            final Map<String, String> extraMetaData,
+            @NotNull final ParquetMetadataFileWriter metadataFileWriter) throws IOException {
         this.targetPageSize = targetPageSize;
         this.allocator = allocator;
         this.extraMetaData = new HashMap<>(extraMetaData);
-        bufferedOutput = new PositionedBufferedOutputStream(channelsProvider.getWriteChannel(filePath, false),
+        bufferedOutput = new PositionedBufferedOutputStream(channelsProvider.getWriteChannel(destFile.getPath(), false),
                 OUTPUT_BUFFER_SIZE);
         bufferedOutput.write(ParquetFileReader.MAGIC);
         this.type = type;
         this.compressorAdapter = DeephavenCompressorAdapterFactory.getInstance().getByName(codecName);
+        this.metadataFilePath = metadataFilePath;
+        this.metadataFileWriter = metadataFileWriter;
     }
 
     public RowGroupWriter addRowGroup(final long size) {
@@ -70,13 +78,16 @@ public void close() throws IOException {
         serializeOffsetIndexes();
         final ParquetMetadata footer =
                 new ParquetMetadata(new FileMetaData(type, extraMetaData, Version.FULL_VERSION), blocks);
-        serializeFooter(footer);
+        serializeFooter(footer, bufferedOutput);
+        metadataFileWriter.addParquetFileMetadata(metadataFilePath, footer);
         // Flush any buffered data and close the channel
         bufferedOutput.close();
         compressorAdapter.close();
     }
 
-    private void serializeFooter(final ParquetMetadata footer) throws IOException {
+    public static void serializeFooter(final ParquetMetadata footer,
+            final PositionedBufferedOutputStream bufferedOutput)
+            throws IOException {
         final long footerIndex = bufferedOutput.position();
         final org.apache.parquet.format.FileMetaData parquetMetadata =
                 metadataConverter.toParquetMetadata(VERSION, footer);

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

@@ -0,0 +1,34 @@
+package io.deephaven.parquet.base;
+
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+
+import java.io.File;
+import java.io.IOException;
+
+/**
+ * Used to write _metadata and _common_metadata files for Parquet.
+ */
+public interface ParquetMetadataFileWriter {
+
+    /**
+     * Add the parquet metadata for the provided parquet file to the list of metadata to be written to combined metadata
+     * files.
+     *
+     * @param parquetFile The parquet file destination path
+     * @param metadata The parquet metadata corresponding to the parquet file
+     */
+    void addParquetFileMetadata(File parquetFile, ParquetMetadata metadata);
+
+    /**
+     * Write the combined metadata files for all metadata accumulated so far and clear the list.
+     *
+     * @param metadataFile The destination file for the _metadata file
+     * @param commonMetadataFile The destination file for the _common_metadata file
+     */
+    void writeMetadataFiles(File metadataFile, File commonMetadataFile) throws IOException;
+
+    /**
+     * Clear the list of metadata accumulated so far.
+     */
+    void clear();
+}

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

@@ -5,17 +5,22 @@
 import java.nio.channels.Channels;
 import java.nio.channels.SeekableByteChannel;
 
-final class PositionedBufferedOutputStream extends BufferedOutputStream {
+public final class PositionedBufferedOutputStream extends BufferedOutputStream {
 
     private final SeekableByteChannel writeChannel;
 
+    public PositionedBufferedOutputStream(final SeekableByteChannel writeChannel) {
+        super(Channels.newOutputStream(writeChannel));
+        this.writeChannel = writeChannel;
+    }
+
     PositionedBufferedOutputStream(final SeekableByteChannel writeChannel, final int size) {
         super(Channels.newOutputStream(writeChannel), size);
         this.writeChannel = writeChannel;
     }
 
     /**
-     * Get total number of bytes written to this stream
+     * Get the total number of bytes written to this stream
      */
     long position() throws IOException {
         // Number of bytes buffered in the stream + bytes written to the underlying channel