Added support to write metadata files in parquet

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

@@ -40,18 +40,13 @@ 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 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 +56,12 @@ final class ColumnChunkReaderImpl implements ColumnChunkReader {
      */
     private final String version;
 
-    ColumnChunkReaderImpl(ColumnChunk columnChunk, SeekableChannelsProvider channelsProvider, URI rootURI,
+    ColumnChunkReaderImpl(ColumnChunk columnChunk, SeekableChannelsProvider channelsProvider, URI columnChunkURI,
             MessageType type, OffsetIndex offsetIndex, 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()) {
@@ -121,15 +116,7 @@ public final ColumnPageDirectAccessor getPageAccessor() {
     }
 
     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

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/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

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

@@ -21,14 +21,15 @@
 import java.io.UncheckedIOException;
 import java.net.URI;
 import java.nio.channels.SeekableByteChannel;
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
-final class RowGroupReaderImpl implements RowGroupReader {
+import static io.deephaven.parquet.base.ParquetFileReader.FILE_URI_SCHEME;
 
-    private static final int BUFFER_SIZE = 65536;
+final class RowGroupReaderImpl implements RowGroupReader {
     private final RowGroup rowGroup;
     private final SeekableChannelsProvider channelsProvider;
     private final MessageType type;
@@ -78,23 +79,32 @@ public ColumnChunkReaderImpl getColumnChunk(@NotNull final List<String> path,
         if (columnChunk == null) {
             return null;
         }
-        final OffsetIndex offsetIndex = offsetIndex(columnChunk, channelContext);
-        return new ColumnChunkReaderImpl(columnChunk, channelsProvider, rootURI, type, offsetIndex, fieldTypes,
+
+        final URI columnChunkURI;
+        if (columnChunk.isSetFile_path() && FILE_URI_SCHEME.equals(rootURI.getScheme())) {
+            columnChunkURI = Path.of(rootURI).resolve(columnChunk.getFile_path()).toUri();
+        } else {
+            // TODO(deephaven-core#5066): Add support for reading metadata files from non-file URIs
+            columnChunkURI = rootURI;
+        }
+
+        final OffsetIndex offsetIndex = offsetIndex(columnChunk, columnChunkURI, channelContext);
+        return new ColumnChunkReaderImpl(columnChunk, channelsProvider, columnChunkURI, type, offsetIndex, fieldTypes,
                 numRows(), version);
     }
 
-    private OffsetIndex offsetIndex(ColumnChunk chunk, @NotNull SeekableChannelContext context) {
+    private OffsetIndex offsetIndex(ColumnChunk chunk, URI columnChunkURI, @NotNull SeekableChannelContext context) {
         if (!chunk.isSetOffset_index_offset()) {
             return null;
         }
-        return ParquetMetadataConverter.fromParquetOffsetIndex(readOffsetIndex(chunk, context));
+        return ParquetMetadataConverter.fromParquetOffsetIndex(readOffsetIndex(chunk, columnChunkURI, context));
     }
 
-    private org.apache.parquet.format.OffsetIndex readOffsetIndex(ColumnChunk chunk,
+    private org.apache.parquet.format.OffsetIndex readOffsetIndex(ColumnChunk chunk, URI columnChunkURI,
             @NotNull SeekableChannelContext channelContext) {
         try (
                 final ContextHolder holder = SeekableChannelContext.ensureContext(channelsProvider, channelContext);
-                final SeekableByteChannel readChannel = channelsProvider.getReadChannel(holder.get(), rootURI);
+                final SeekableByteChannel readChannel = channelsProvider.getReadChannel(holder.get(), columnChunkURI);
                 final InputStream in =
                         channelsProvider.getInputStream(readChannel.position(chunk.getOffset_index_offset()))) {
             return Util.readOffsetIndex(in);

extensions/parquet/table/src/main/java/io/deephaven/parquet/table/ParquetInstructions.java

@@ -115,6 +115,8 @@ public static int getDefaultTargetPageSize() {
         return defaultTargetPageSize;
     }
 
+    static final String DEFAULT_METADATA_ROOT_DIR = ""; // Empty = No metadata files written
+
     public ParquetInstructions() {}
 
     public final String getColumnNameFromParquetColumnNameOrDefault(final String parquetColumnName) {
@@ -164,6 +166,11 @@ public final String getColumnNameFromParquetColumnNameOrDefault(final String par
      */
     public abstract boolean isRefreshing();
 
+    /**
+     * @return the directory in which metadata files should be stored.
+     */
+    public abstract String getMetadataRootDir();
+
     @VisibleForTesting
     public static boolean sameColumnNamesAndCodecMappings(final ParquetInstructions i1, final ParquetInstructions i2) {
         if (i1 == EMPTY) {
@@ -238,6 +245,11 @@ public int getTargetPageSize() {
         public boolean isRefreshing() {
             return DEFAULT_IS_REFRESHING;
         }
+
+        @Override
+        public String getMetadataRootDir() {
+            return DEFAULT_METADATA_ROOT_DIR;
+        }
     };
 
     private static class ColumnInstructions {
@@ -306,6 +318,7 @@ private static final class ReadOnly extends ParquetInstructions {
         private final int targetPageSize;
         private final boolean isRefreshing;
         private final Object specialInstructions;
+        private final String metadataRootDir;
 
         private ReadOnly(
                 final KeyedObjectHashMap<String, ColumnInstructions> columnNameToInstructions,
@@ -316,7 +329,8 @@ private ReadOnly(
                 final boolean isLegacyParquet,
                 final int targetPageSize,
                 final boolean isRefreshing,
-                final Object specialInstructions) {
+                final Object specialInstructions,
+                final String metadataRootDir) {
             this.columnNameToInstructions = columnNameToInstructions;
             this.parquetColumnNameToInstructions = parquetColumnNameToColumnName;
             this.compressionCodecName = compressionCodecName;
@@ -326,6 +340,7 @@ private ReadOnly(
             this.targetPageSize = targetPageSize;
             this.isRefreshing = isRefreshing;
             this.specialInstructions = specialInstructions;
+            this.metadataRootDir = metadataRootDir;
         }
 
         private String getOrDefault(final String columnName, final String defaultValue,
@@ -419,6 +434,10 @@ public boolean isRefreshing() {
             return specialInstructions;
         }
 
+        @Override
+        public String getMetadataRootDir() {
+            return metadataRootDir;
+        }
 
         KeyedObjectHashMap<String, ColumnInstructions> copyColumnNameToInstructions() {
             // noinspection unchecked
@@ -471,6 +490,7 @@ public static class Builder {
         private int targetPageSize = defaultTargetPageSize;
         private boolean isRefreshing = DEFAULT_IS_REFRESHING;
         private Object specialInstructions;
+        private String metadataRootDir = DEFAULT_METADATA_ROOT_DIR;
 
         public Builder() {}
 
@@ -647,6 +667,17 @@ public Builder setSpecialInstructions(final Object specialInstructions) {
             return this;
         }
 
+        /**
+         * Set the default metadata root directory.
+         *
+         * @param metadataRootDir the root directory to store metadata files in. All the parquet destinations should be
+         *        inside this directory.
+         */
+        public Builder setMetadataRootDir(final String metadataRootDir) {
+            this.metadataRootDir = metadataRootDir;
+            return this;
+        }
+
         public ParquetInstructions build() {
             final KeyedObjectHashMap<String, ColumnInstructions> columnNameToInstructionsOut = columnNameToInstructions;
             columnNameToInstructions = null;
@@ -655,7 +686,7 @@ public ParquetInstructions build() {
             parquetColumnNameToInstructions = null;
             return new ReadOnly(columnNameToInstructionsOut, parquetColumnNameToColumnNameOut, compressionCodecName,
                     maximumDictionaryKeys, maximumDictionarySize, isLegacyParquet, targetPageSize, isRefreshing,
-                    specialInstructions);
+                    specialInstructions, metadataRootDir);
         }
     }