lidavidm commented on code in PR #34873:
URL: https://github.com/apache/arrow/pull/34873#discussion_r1173148616
##########
java/compression/src/test/java/org/apache/arrow/compression/TestArrowReaderWriterWithCompression.java:
##########
@@ -85,4 +85,46 @@ public void testArrowFileZstdRoundTrip() throws Exception {
}
}
+ @Test
+ public void testArrowFileZstdWithCompressionLevelRoundTrip() throws
Exception {
+ // Prepare sample data
+ final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+ List<Field> fields = new ArrayList<>();
+ fields.add(new Field("col", FieldType.notNullable(new ArrowType.Utf8()),
new ArrayList<>()));
+ VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(fields),
allocator);
+ final int rowCount = 10;
+ GenerateSampleData.generateTestData(root.getVector(0), rowCount);
+ root.setRowCount(rowCount);
+
+ // Write an in-memory compressed arrow file
+ ByteArrayOutputStream out = new ByteArrayOutputStream();
+ try (final ArrowFileWriter writer =
+ new ArrowFileWriter(root, null, Channels.newChannel(out), new
HashMap<>(),
+ IpcOption.DEFAULT, CommonsCompressionFactory.INSTANCE,
CompressionUtil.CodecType.ZSTD, 7)) {
+ writer.start();
+ writer.writeBatch();
+ writer.end();
+ }
+
+ // Read the in-memory compressed arrow file with CommonsCompressionFactory
provided
+ try (ArrowFileReader reader =
+ new ArrowFileReader(new
ByteArrayReadableSeekableByteChannel(out.toByteArray()),
+ allocator, CommonsCompressionFactory.INSTANCE)) {
+ Assert.assertEquals(1, reader.getRecordBlocks().size());
+ Assert.assertTrue(reader.loadNextBatch());
+ Assert.assertTrue(root.equals(reader.getVectorSchemaRoot()));
+ Assert.assertFalse(reader.loadNextBatch());
+ }
+
+ // Read the in-memory compressed arrow file without CompressionFactory
provided
Review Comment:
Isn't this tested elsewhere already?
##########
java/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowWriter.java:
##########
@@ -119,6 +119,53 @@ protected ArrowWriter(VectorSchemaRoot root,
DictionaryProvider provider, Writab
this.schema = new Schema(fields, root.getSchema().getCustomMetadata());
}
+ /**
+ * Note: fields are not closed when the writer is closed.
+ *
+ * @param root the vectors to write to the output
+ * @param provider where to find the dictionaries
+ * @param out the output where to write
+ * @param option IPC write options
+ * @param compressionFactory Compression codec factory
+ * @param codecType Compression codec
+ * @param compressionLevel Compression level
+ */
+ protected ArrowWriter(VectorSchemaRoot root, DictionaryProvider provider,
WritableByteChannel out, IpcOption option,
Review Comment:
Agreed, can we consolidate?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
