diff --git a/java/src/main/java/ai/rapids/cudf/Table.java b/java/src/main/java/ai/rapids/cudf/Table.java index cbd212c13d5..e9033c91480 100644 --- a/java/src/main/java/ai/rapids/cudf/Table.java +++ b/java/src/main/java/ai/rapids/cudf/Table.java @@ -1588,30 +1588,60 @@ public static Table readORC(ORCOptions opts, DataSource ds) { * Get a table writer to write parquet data to a file. * @param options the parquet writer options. * @param outputFile where to write the file. + * @return a table writer to use for writing out multiple tables. + */ + public static TableWriter writeParquetChunked(ParquetWriterOptions options, File outputFile) { + return writeParquetChunkedWithFooter(options, outputFile); + } + + /** + * Get a table writer to write parquet data and handle each chunk with a callback. + * @param options the parquet writer options. + * @param consumer a class that will be called when host buffers are ready with parquet + * formatted data in them. + * @param hostMemoryAllocator allocator for host memory buffers + * @return a table writer to use for writing out multiple tables. + */ + public static TableWriter writeParquetChunked(ParquetWriterOptions options, + HostBufferConsumer consumer, + HostMemoryAllocator hostMemoryAllocator) { + return writeParquetChunkedWithFooter(options, consumer, hostMemoryAllocator); + } + + public static TableWriter writeParquetChunked(ParquetWriterOptions options, + HostBufferConsumer consumer) { + return writeParquetChunked(options, consumer, DefaultHostMemoryAllocator.get()); + } + + /** + * Get a Parquet table writer that can return footer metadata when writing to a file. + * @param options the parquet writer options. + * @param outputFile where to write the file. * @return a Parquet table writer to use for writing out multiple tables. */ - public static ParquetTableWriter writeParquetChunked(ParquetWriterOptions options, - File outputFile) { + public static ParquetTableWriter writeParquetChunkedWithFooter(ParquetWriterOptions options, + File outputFile) { return new ParquetTableWriter(options, outputFile); } /** - * Get a table writer to write parquet data and handle each chunk with a callback. + * Get a Parquet table writer that can return footer metadata and handle each chunk with a + * callback. * @param options the parquet writer options. * @param consumer a class that will be called when host buffers are ready with parquet * formatted data in them. * @param hostMemoryAllocator allocator for host memory buffers * @return a Parquet table writer to use for writing out multiple tables. */ - public static ParquetTableWriter writeParquetChunked(ParquetWriterOptions options, - HostBufferConsumer consumer, - HostMemoryAllocator hostMemoryAllocator) { + public static ParquetTableWriter writeParquetChunkedWithFooter(ParquetWriterOptions options, + HostBufferConsumer consumer, + HostMemoryAllocator hostMemoryAllocator) { return new ParquetTableWriter(options, consumer, hostMemoryAllocator); } - public static ParquetTableWriter writeParquetChunked(ParquetWriterOptions options, - HostBufferConsumer consumer) { - return writeParquetChunked(options, consumer, DefaultHostMemoryAllocator.get()); + public static ParquetTableWriter writeParquetChunkedWithFooter(ParquetWriterOptions options, + HostBufferConsumer consumer) { + return writeParquetChunkedWithFooter(options, consumer, DefaultHostMemoryAllocator.get()); } /** diff --git a/java/src/test/java/ai/rapids/cudf/TableTest.java b/java/src/test/java/ai/rapids/cudf/TableTest.java index 3090ce59978..45013ffdfea 100644 --- a/java/src/test/java/ai/rapids/cudf/TableTest.java +++ b/java/src/test/java/ai/rapids/cudf/TableTest.java @@ -10182,7 +10182,7 @@ void testParquetWriterCloseAndGetFooterToFile() throws IOException { .column("a", "b", "c") .build(); ParquetTableWriter writer = - Table.writeParquetChunked(options, tempFile.getFile())) { + Table.writeParquetChunkedWithFooter(options, tempFile.getFile())) { writer.write(table); writer.write(table); try (HostMemoryBuffer footer = writer.closeAndGetFooter()) { @@ -10210,7 +10210,7 @@ void testParquetWriterCloseAndGetFooterToBuffer() { .build(); MyBufferConsumer consumer = new MyBufferConsumer(); ParquetTableWriter writer = - Table.writeParquetChunked(options, consumer, allocator)) { + Table.writeParquetChunkedWithFooter(options, consumer, allocator)) { writer.write(table); try (HostMemoryBuffer footer = writer.closeAndGetFooter()) { byte[] parquetData = new byte[(int) consumer.offset]; @@ -10231,7 +10231,8 @@ void testParquetWriterCloseAndGetFooterAfterClose() throws IOException { .build(); try (TempFile tempFile = TempFile.create("discarded-footer", ".parquet"); Table table = new Table.TestBuilder().column(1, 2, 3).build()) { - ParquetTableWriter writer = Table.writeParquetChunked(options, tempFile.getFile()); + ParquetTableWriter writer = + Table.writeParquetChunkedWithFooter(options, tempFile.getFile()); writer.write(table); writer.close(); assertThrows(IllegalStateException.class, writer::closeAndGetFooter); @@ -10239,6 +10240,29 @@ void testParquetWriterCloseAndGetFooterAfterClose() throws IOException { } } + @Test + void testParquetWriterFactoryReturnTypes() throws NoSuchMethodException { + assertEquals(TableWriter.class, + Table.class.getMethod("writeParquetChunked", ParquetWriterOptions.class, File.class) + .getReturnType()); + assertEquals(TableWriter.class, + Table.class.getMethod("writeParquetChunked", ParquetWriterOptions.class, + HostBufferConsumer.class, HostMemoryAllocator.class).getReturnType()); + assertEquals(TableWriter.class, + Table.class.getMethod("writeParquetChunked", ParquetWriterOptions.class, + HostBufferConsumer.class).getReturnType()); + + assertEquals(ParquetTableWriter.class, + Table.class.getMethod("writeParquetChunkedWithFooter", ParquetWriterOptions.class, + File.class).getReturnType()); + assertEquals(ParquetTableWriter.class, + Table.class.getMethod("writeParquetChunkedWithFooter", ParquetWriterOptions.class, + HostBufferConsumer.class, HostMemoryAllocator.class).getReturnType()); + assertEquals(ParquetTableWriter.class, + Table.class.getMethod("writeParquetChunkedWithFooter", ParquetWriterOptions.class, + HostBufferConsumer.class).getReturnType()); + } + private static void assertReturnedFooterMatches(byte[] parquetData, HostMemoryBuffer returnedFooter) { int fileLength = parquetData.length;