fix: addressing reviews v2

vibhatha · Jun 4, 2024 · 6003d3e · 6003d3e
1 parent 038ed93
commit 6003d3e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 12 deletions.
diff --git a/java/c/src/main/java/org/apache/arrow/c/StructVectorLoader.java b/java/c/src/main/java/org/apache/arrow/c/StructVectorLoader.java
@@ -104,9 +104,10 @@ public StructVector load(BufferAllocator allocator, ArrowRecordBatch recordBatch
       loadBuffers(fieldVector, fieldVector.getField(), buffers, nodes, codec, variadicBufferCounts);
     }
     result.loadFieldBuffers(new ArrowFieldNode(recordBatch.getLength(), 0), Collections.singletonList(null));
-    if (nodes.hasNext() || buffers.hasNext()) {
-      throw new IllegalArgumentException("not all nodes and buffers were consumed. nodes: " + 
-        Collections2.toList(nodes).toString() + " buffers: " + Collections2.toList(buffers).toString());
+    if (nodes.hasNext() || buffers.hasNext() || variadicBufferCounts.hasNext()) {
+      throw new IllegalArgumentException("not all nodes, buffers and variadicBufferCounts were consumed. nodes: " +
+        Collections2.toString(nodes) + " buffers: " + Collections2.toString(buffers) + " variadicBufferCounts: " +
+          Collections2.toString(variadicBufferCounts));
     }
     return result;
   }
@@ -115,10 +116,14 @@ private void loadBuffers(FieldVector vector, Field field, Iterator<ArrowBuf> buf
       CompressionCodec codec, Iterator<Long> variadicBufferCounts) {
     checkArgument(nodes.hasNext(), "no more field nodes for field %s and vector %s", field, vector);
     ArrowFieldNode fieldNode = nodes.next();
-    // variadicBufferLayoutCount will be 0 for vectors of type except BaseVariableWidthViewVector
+    // variadicBufferLayoutCount will be 0 for vectors of a type except BaseVariableWidthViewVector
     long variadicBufferLayoutCount = 0;
-    if (vector instanceof BaseVariableWidthViewVector && variadicBufferCounts.hasNext()) {
-      variadicBufferLayoutCount = variadicBufferCounts.next();
+    if (vector instanceof BaseVariableWidthViewVector) {
+      if (variadicBufferCounts.hasNext()) {
+        variadicBufferLayoutCount = variadicBufferCounts.next();
+      } else {
+        throw new IllegalStateException("No variadicBufferCounts available for BaseVariableWidthViewVector");
+      }
     }
     int bufferLayoutCount = (int) (variadicBufferLayoutCount + TypeLayout.getTypeBufferCount(field.getType()));
     List<ArrowBuf> ownBuffers = new ArrayList<>(bufferLayoutCount);

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java b/java/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java
@@ -90,9 +90,10 @@ public void load(ArrowRecordBatch recordBatch) {
       loadBuffers(fieldVector, fieldVector.getField(), buffers, nodes, codec, variadicBufferCounts);
     }
     root.setRowCount(recordBatch.getLength());
-    if (nodes.hasNext() || buffers.hasNext()) {
-      throw new IllegalArgumentException("not all nodes and buffers were consumed. nodes: " +
-          Collections2.toString(nodes) + " buffers: " + Collections2.toString(buffers));
+    if (nodes.hasNext() || buffers.hasNext() || variadicBufferCounts.hasNext()) {
+      throw new IllegalArgumentException("not all nodes, buffers and variadicBufferCounts were consumed. nodes: " +
+          Collections2.toString(nodes) + " buffers: " + Collections2.toString(buffers) + " variadicBufferCounts: " +
+          Collections2.toString(variadicBufferCounts));
     }
   }
 
@@ -105,10 +106,14 @@ private void loadBuffers(
       Iterator<Long> variadicBufferCounts) {
     checkArgument(nodes.hasNext(), "no more field nodes for field %s and vector %s", field, vector);
     ArrowFieldNode fieldNode = nodes.next();
-    // variadicBufferLayoutCount will be 0 for vectors of type except BaseVariableWidthViewVector
+    // variadicBufferLayoutCount will be 0 for vectors of a type except BaseVariableWidthViewVector
     long variadicBufferLayoutCount = 0;
-    if (vector instanceof BaseVariableWidthViewVector && variadicBufferCounts.hasNext()) {
-      variadicBufferLayoutCount = variadicBufferCounts.next();
+    if (vector instanceof BaseVariableWidthViewVector) {
+      if (variadicBufferCounts.hasNext()) {
+        variadicBufferLayoutCount = variadicBufferCounts.next();
+      } else {
+        throw new IllegalStateException("No variadicBufferCounts available for BaseVariableWidthViewVector");
+      }
     }
     int bufferLayoutCount = (int) (variadicBufferLayoutCount + TypeLayout.getTypeBufferCount(field.getType()));
     List<ArrowBuf> ownBuffers = new ArrayList<>(bufferLayoutCount);