From f1b1dae2ce0f4083007c58e6479898e622075c56 Mon Sep 17 00:00:00 2001 From: tianchen Date: Fri, 10 Jan 2020 12:23:53 +0800 Subject: [PATCH] ARROW-7539: [Java] FieldVector getFieldBuffers API should not set reader/writer indices --- .../main/codegen/templates/UnionVector.java | 7 ++--- .../arrow/vector/BaseFixedWidthVector.java | 11 +++----- .../arrow/vector/BaseVariableWidthVector.java | 14 ++++------ .../apache/arrow/vector/VectorUnloader.java | 27 ++++++++++--------- .../complex/BaseRepeatedValueVector.java | 12 +++------ .../vector/complex/FixedSizeListVector.java | 13 +++------ .../arrow/vector/complex/ListVector.java | 15 ++++------- .../arrow/vector/complex/StructVector.java | 13 +++------ .../apache/arrow/vector/TestValueVector.java | 8 +++--- 9 files changed, 46 insertions(+), 74 deletions(-) diff --git a/java/vector/src/main/codegen/templates/UnionVector.java b/java/vector/src/main/codegen/templates/UnionVector.java index 8bf5d8a9b20..679fbe4fd50 100644 --- a/java/vector/src/main/codegen/templates/UnionVector.java +++ b/java/vector/src/main/codegen/templates/UnionVector.java @@ -159,7 +159,6 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers @Override public List getFieldBuffers() { List result = new ArrayList<>(1); - setReaderAndWriterIndex(); result.add(typeBuffer); return result; @@ -534,10 +533,8 @@ public int getBufferSizeFor(final int valueCount) { public ArrowBuf[] getBuffers(boolean clear) { List list = new java.util.ArrayList<>(); setReaderAndWriterIndex(); - if (getBufferSize() != 0) { - list.add(typeBuffer); - list.addAll(java.util.Arrays.asList(internalStruct.getBuffers(clear))); - } + list.add(typeBuffer); + list.addAll(java.util.Arrays.asList(internalStruct.getBuffers(clear))); if (clear) { valueCount = 0; typeBuffer.getReferenceManager().retain(); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java b/java/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java index a4e94bcac09..ea6da19ca9d 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java @@ -396,13 +396,9 @@ public Field getField() { public ArrowBuf[] getBuffers(boolean clear) { final ArrowBuf[] buffers; setReaderAndWriterIndex(); - if (getBufferSize() == 0) { - buffers = new ArrowBuf[0]; - } else { - buffers = new ArrowBuf[2]; - buffers[0] = validityBuffer; - buffers[1] = valueBuffer; - } + buffers = new ArrowBuf[2]; + buffers[0] = validityBuffer; + buffers[1] = valueBuffer; if (clear) { for (final ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(1); @@ -503,7 +499,6 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers */ public List getFieldBuffers() { List result = new ArrayList<>(2); - setReaderAndWriterIndex(); result.add(validityBuffer); result.add(valueBuffer); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java b/java/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java index 89395ef7e22..ba2d71affcf 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java @@ -337,7 +337,6 @@ public List getFieldBuffers() { fillHoles(valueCount); List result = new ArrayList<>(3); - setReaderAndWriterIndex(); result.add(validityBuffer); result.add(offsetBuffer); result.add(valueBuffer); @@ -638,15 +637,12 @@ public Field getField() { @Override public ArrowBuf[] getBuffers(boolean clear) { final ArrowBuf[] buffers; + fillHoles(valueCount); setReaderAndWriterIndex(); - if (getBufferSize() == 0) { - buffers = new ArrowBuf[0]; - } else { - buffers = new ArrowBuf[3]; - buffers[0] = validityBuffer; - buffers[1] = offsetBuffer; - buffers[2] = valueBuffer; - } + buffers = new ArrowBuf[3]; + buffers[0] = validityBuffer; + buffers[1] = offsetBuffer; + buffers[2] = valueBuffer; if (clear) { for (final ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java b/java/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java index 8c936aa0ba9..cf3d1e50507 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java @@ -18,6 +18,7 @@ package org.apache.arrow.vector; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; @@ -63,24 +64,26 @@ public ArrowRecordBatch getRecordBatch() { List nodes = new ArrayList<>(); List buffers = new ArrayList<>(); for (FieldVector vector : root.getFieldVectors()) { - appendNodes(vector, nodes, buffers); + ArrowBuf[] fieldBuffers = vector.getBuffers(false); + buffers.addAll(Arrays.asList(fieldBuffers)); + int expectedBufferCount = appendNodes(vector, nodes); + if (fieldBuffers.length != expectedBufferCount) { + throw new IllegalArgumentException(String.format( + "wrong number of buffers for field %s in vector %s. found: %s", + vector.getField(), vector.getClass().getSimpleName(), fieldBuffers)); + } } return new ArrowRecordBatch(root.getRowCount(), nodes, buffers, alignBuffers); } - private void appendNodes(FieldVector vector, List nodes, List buffers) { + private int appendNodes(FieldVector vector, List nodes) { nodes.add(new ArrowFieldNode(vector.getValueCount(), includeNullCount ? vector.getNullCount() : -1)); - List fieldBuffers = vector.getFieldBuffers(); - int expectedBufferCount = TypeLayout.getTypeBufferCount(vector.getField().getType()); - if (fieldBuffers.size() != expectedBufferCount) { - throw new IllegalArgumentException(String.format( - "wrong number of buffers for field %s in vector %s. found: %s", - vector.getField(), vector.getClass().getSimpleName(), fieldBuffers)); - } - buffers.addAll(fieldBuffers); + + int expected = TypeLayout.getTypeBufferCount(vector.getField().getType()); + for (FieldVector child : vector.getChildrenFromFields()) { - appendNodes(child, nodes, buffers); + expected += appendNodes(child, nodes); } + return expected; } - } diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java index 9f6a5ee61e7..d3767e5fe9c 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java @@ -253,14 +253,10 @@ public void reset() { @Override public ArrowBuf[] getBuffers(boolean clear) { final ArrowBuf[] buffers; - if (getBufferSize() == 0) { - buffers = new ArrowBuf[0]; - } else { - List list = new ArrayList<>(); - list.add(offsetBuffer); - list.addAll(Arrays.asList(vector.getBuffers(false))); - buffers = list.toArray(new ArrowBuf[list.size()]); - } + List list = new ArrayList<>(); + list.add(offsetBuffer); + list.addAll(Arrays.asList(vector.getBuffers(false))); + buffers = list.toArray(new ArrowBuf[list.size()]); if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java index 8fa43fb06b7..9c6cc8bc010 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java @@ -173,7 +173,6 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers @Override public List getFieldBuffers() { List result = new ArrayList<>(1); - setReaderAndWriterIndex(); result.add(validityBuffer); return result; @@ -348,14 +347,10 @@ public void reset() { public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; - if (getBufferSize() == 0) { - buffers = new ArrowBuf[0]; - } else { - List list = new ArrayList<>(); - list.add(validityBuffer); - list.addAll(Arrays.asList(vector.getBuffers(false))); - buffers = list.toArray(new ArrowBuf[list.size()]); - } + List list = new ArrayList<>(); + list.add(validityBuffer); + list.addAll(Arrays.asList(vector.getBuffers(false))); + buffers = list.toArray(new ArrowBuf[list.size()]); if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java index 66c8a3f4323..c9b81bb2e93 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java @@ -225,7 +225,6 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers @Override public List getFieldBuffers() { List result = new ArrayList<>(2); - setReaderAndWriterIndex(); result.add(validityBuffer); result.add(offsetBuffer); @@ -662,15 +661,11 @@ public void reset() { public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; - if (getBufferSize() == 0) { - buffers = new ArrowBuf[0]; - } else { - List list = new ArrayList<>(); - list.add(offsetBuffer); - list.add(validityBuffer); - list.addAll(Arrays.asList(vector.getBuffers(false))); - buffers = list.toArray(new ArrowBuf[list.size()]); - } + List list = new ArrayList<>(); + list.add(validityBuffer); + list.add(offsetBuffer); + list.addAll(Arrays.asList(vector.getBuffers(false))); + buffers = list.toArray(new ArrowBuf[list.size()]); if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java index 7b22835963a..5aff252baf8 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java @@ -118,7 +118,6 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers @Override public List getFieldBuffers() { List result = new ArrayList<>(1); - setReaderAndWriterIndex(); result.add(validityBuffer); return result; @@ -296,14 +295,10 @@ public int getValueCapacity() { public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; - if (getBufferSize() == 0) { - buffers = new ArrowBuf[0]; - } else { - List list = new ArrayList<>(); - list.add(validityBuffer); - list.addAll(Arrays.asList(super.getBuffers(false))); - buffers = list.toArray(new ArrowBuf[list.size()]); - } + List list = new ArrayList<>(); + list.add(validityBuffer); + list.addAll(Arrays.asList(super.getBuffers(false))); + buffers = list.toArray(new ArrowBuf[list.size()]); if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java index 65bc963acb5..c7c2a5b4c8f 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java @@ -2673,11 +2673,11 @@ public void testUnloadVariableWidthVector() { varCharVector.set(0, "abcd".getBytes()); - List bufs = varCharVector.getFieldBuffers(); - assertEquals(3, bufs.size()); + ArrowBuf[] bufs = varCharVector.getBuffers(false); + assertEquals(3, bufs.length); - ArrowBuf offsetBuf = bufs.get(1); - ArrowBuf dataBuf = bufs.get(2); + ArrowBuf offsetBuf = bufs[1]; + ArrowBuf dataBuf = bufs[2]; assertEquals(12, offsetBuf.writerIndex()); assertEquals(4, offsetBuf.getInt(4));