fix: adding new constructors

vibhatha · May 22, 2024 · 67fc7ce · 67fc7ce
1 parent 704767d
commit 67fc7ce
Showing 1 changed file with 56 additions and 0 deletions.
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowRecordBatch.java b/java/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowRecordBatch.java
@@ -71,6 +71,21 @@ public ArrowRecordBatch(
     this(length, nodes, buffers, bodyCompression, null, true);
   }
 
+  /**
+   * Construct a record batch from nodes.
+   *
+   * @param length  how many rows in this batch
+   * @param nodes   field level info
+   * @param buffers will be retained until this recordBatch is closed
+   * @param bodyCompression compression info.
+   * @param alignBuffers Whether to align buffers to an 8 byte boundary.
+   */
+  public ArrowRecordBatch(
+      int length, List<ArrowFieldNode> nodes, List<ArrowBuf> buffers,
+      ArrowBodyCompression bodyCompression, boolean alignBuffers) {
+    this(length, nodes, buffers, bodyCompression, alignBuffers, /*retainBuffers*/ true);
+  }
+
   /**
    * Construct a record batch from nodes.
    *
@@ -87,6 +102,47 @@ public ArrowRecordBatch(
     this(length, nodes, buffers, bodyCompression, variadicBufferCounts, alignBuffers, /*retainBuffers*/ true);
   }
 
+  /**
+   * Construct a record batch from nodes.
+   *
+   * @param length  how many rows in this batch
+   * @param nodes   field level info
+   * @param buffers will be retained until this recordBatch is closed
+   * @param bodyCompression compression info.
+   * @param alignBuffers Whether to align buffers to an 8 byte boundary.
+   * @param retainBuffers Whether to retain() each source buffer in the constructor. If false, the caller is
+   *                      responsible for retaining the buffers beforehand.
+   */
+  public ArrowRecordBatch(
+      int length, List<ArrowFieldNode> nodes, List<ArrowBuf> buffers,
+      ArrowBodyCompression bodyCompression, boolean alignBuffers,
+      boolean retainBuffers) {
+    super();
+    this.length = length;
+    this.nodes = nodes;
+    this.buffers = buffers;
+    Preconditions.checkArgument(bodyCompression != null, "body compression cannot be null");
+    this.bodyCompression = bodyCompression;
+    List<ArrowBuffer> arrowBuffers = new ArrayList<>(buffers.size());
+    long offset = 0;
+    for (ArrowBuf arrowBuf : buffers) {
+      if (retainBuffers) {
+        arrowBuf.getReferenceManager().retain();
+      }
+      long size = arrowBuf.readableBytes();
+      arrowBuffers.add(new ArrowBuffer(offset, size));
+      if (LOGGER.isTraceEnabled()) {
+        LOGGER.trace("Buffer in RecordBatch at {}, length: {}", offset, size);
+      }
+      offset += size;
+      if (alignBuffers) { // align on 8 byte boundaries
+        offset = DataSizeRoundingUtil.roundUpTo8Multiple(offset);
+      }
+    }
+    this.buffersLayout = Collections.unmodifiableList(arrowBuffers);
+    this.variadicBufferCounts = null;
+  }
+
   /**
    * Construct a record batch from nodes.
    *