From ab131498cdabfe3daa7ec8cd7cbfb04c2fda3633 Mon Sep 17 00:00:00 2001
From: vibhatha <vibhatha@gmail.com>
Date: Tue, 21 Nov 2023 17:15:23 +0530
Subject: [PATCH] fix: adding a method to get codecs

---
 .../apache/arrow/vector/ipc/ArrowWriter.java  | 23 +++++++++----------
 1 file changed, 11 insertions(+), 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
index 9953428031372..1cd26b64de116 100644
--- 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,12 +92,6 @@ 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;
@@ -108,6 +99,8 @@ protected ArrowWriter(VectorSchemaRoot root, DictionaryProvider provider, Writab
     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;