fix: adding a method to get codecs

vibhatha · Nov 21, 2023 · ab13149 · ab13149
1 parent b2ade6e
commit ab13149
Showing 1 changed file with 11 additions and 12 deletions.
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowWriter.java b/java/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowWriter.java
@@ -61,17 +61,14 @@ public abstract class ArrowWriter implements AutoCloseable {
   private final DictionaryProvider dictionaryProvider;
   private final Set<Long> dictionaryIdsUsed = new HashSet<>();
 
+  private final CompressionCodec.Factory compressionFactory;
+  private final CompressionUtil.CodecType codecType;
+  private final Optional<Integer> compressionLevel;
   private boolean started = false;
   private boolean ended = false;
 
   protected IpcOption option;
 
-  private CompressionCodec.Factory compressionFactory;
-
-  private CompressionUtil.CodecType codecType;
-
-  private Optional<Integer> compressionLevel;
-
   protected ArrowWriter(VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out) {
     this(root, provider, out, IpcOption.DEFAULT);
   }
@@ -95,19 +92,15 @@ protected ArrowWriter(VectorSchemaRoot root, DictionaryProvider provider, Writab
   protected ArrowWriter(VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, IpcOption option,
                         CompressionCodec.Factory compressionFactory, CompressionUtil.CodecType codecType,
                         Optional<Integer> compressionLevel) {
-    this.unloader = new VectorUnloader(
-        root, /*includeNullCount*/ true,
-        compressionLevel.isPresent() ?
-            compressionFactory.createCodec(codecType, compressionLevel.get()) :
-            compressionFactory.createCodec(codecType),
-        /*alignBuffers*/ true);
     this.out = new WriteChannel(out);
     this.option = option;
     this.dictionaryProvider = provider;
 
     this.compressionFactory = compressionFactory;
     this.codecType = codecType;
     this.compressionLevel = compressionLevel;
+    this.unloader = new VectorUnloader(root, /*includeNullCount*/ true,
+        getCodec(), /*alignBuffers*/ true);
 
     List<Field> fields = new ArrayList<>(root.getSchema().getFields().size());
 
@@ -188,6 +181,12 @@ public long bytesWritten() {
     return out.getCurrentPosition();
   }
 
+  private CompressionCodec getCodec() {
+    return this.compressionLevel.isPresent() ?
+        this.compressionFactory.createCodec(this.codecType, this.compressionLevel.get()) :
+        this.compressionFactory.createCodec(this.codecType);
+  }
+
   private void ensureStarted() throws IOException {
     if (!started) {
       started = true;