diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestListVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestListVector.java index a1762c466ce..eb30fdd46bf 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestListVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestListVector.java @@ -38,8 +38,10 @@ import org.junit.Before; import org.junit.Test; +import java.util.ArrayList; import java.util.List; + public class TestListVector { private BufferAllocator allocator; @@ -92,6 +94,26 @@ public void testCopyFrom() throws Exception { Assert.assertFalse("should be null", reader.isSet()); reader.setPosition(2); Assert.assertTrue("shouldn't be null", reader.isSet()); + + /* check the exact contents of vector */ + final ListVector.Accessor accessor = outVector.getAccessor(); + + /* index 0 */ + Object result = accessor.getObject(0); + ArrayList resultSet = (ArrayList)result; + assertEquals(3, resultSet.size()); + assertEquals(new Long(1), (Long)resultSet.get(0)); + assertEquals(new Long(2), (Long)resultSet.get(1)); + assertEquals(new Long(3), (Long)resultSet.get(2)); + + /* index 1 */ + result = accessor.getObject(1); + assertNull(result); + + /* index 2 */ + result = accessor.getObject(2); + resultSet = (ArrayList)result; + assertEquals(0, resultSet.size()); } } @@ -127,7 +149,7 @@ public void testSetLastSetUsage() throws Exception { dataVector.getMutator().setSafe(2, 1, 12); offsetVector.getMutator().setSafe(index + 1, 3); - index += 1; + index += 1; /* write [13, 14] to the list vector at index 1 */ bitVector.getMutator().setSafe(index, 1); @@ -193,7 +215,7 @@ public void testSetLastSetUsage() throws Exception { final UInt4Vector.Accessor offsetAccessor = offsetVector.getAccessor(); final ValueVector.Accessor valueAccessor = dataVector.getAccessor(); - index = 0; + index = 0; offset = offsetAccessor.get(index); assertEquals(Integer.toString(0), Integer.toString(offset)); @@ -300,7 +322,7 @@ public void testSplitAndTransfer() throws Exception { final UInt4Vector.Accessor offsetAccessor = offsetVector.getAccessor(); final ValueVector.Accessor valueAccessor = dataVector.getAccessor(); - int index = 0; + int index = 0; int offset = 0; Object actual = null; @@ -386,10 +408,7 @@ public void testSplitAndTransfer() throws Exception { TransferPair transferPair = listVector.makeTransferPair(toVector); - int[][] transferLengths = { {0, 2}, - {3, 1}, - {4, 1} - }; + int[][] transferLengths = {{0, 2}, {3, 1}, {4, 1}}; for (final int[] transferLength : transferLengths) { int start = transferLength[0]; @@ -423,7 +442,7 @@ public void testSplitAndTransfer() throws Exception { for(int j = 0; j < dataLength1; j++) { assertEquals("Different data at indexes: " + offset1 + " and " + offset2, - valueAccessor.getObject(offset1), valueAccessor1.getObject(offset2)); + valueAccessor.getObject(offset1), valueAccessor1.getObject(offset2)); offset1++; offset2++; @@ -433,4 +452,120 @@ public void testSplitAndTransfer() throws Exception { } } } + + @Test + public void testNestedListVector() throws Exception { + try (ListVector listVector = ListVector.empty("sourceVector", allocator)) { + + UnionListWriter listWriter = listVector.getWriter(); + + /* allocate memory */ + listWriter.allocate(); + + /* the dataVector that backs a listVector will also be a + * listVector for this test. + */ + + /* write one or more inner lists at index 0 */ + listWriter.setPosition(0); + listWriter.startList(); + + listWriter.list().startList(); + listWriter.list().bigInt().writeBigInt(50); + listWriter.list().bigInt().writeBigInt(100); + listWriter.list().bigInt().writeBigInt(200); + listWriter.list().endList(); + + listWriter.list().startList(); + listWriter.list().bigInt().writeBigInt(75); + listWriter.list().bigInt().writeBigInt(125); + listWriter.list().bigInt().writeBigInt(150); + listWriter.list().bigInt().writeBigInt(175); + listWriter.list().endList(); + + listWriter.endList(); + + /* write one or more inner lists at index 1 */ + listWriter.setPosition(1); + listWriter.startList(); + + listWriter.list().startList(); + listWriter.list().bigInt().writeBigInt(10); + listWriter.list().endList(); + + listWriter.list().startList(); + listWriter.list().bigInt().writeBigInt(15); + listWriter.list().bigInt().writeBigInt(20); + listWriter.list().endList(); + + listWriter.list().startList(); + listWriter.list().bigInt().writeBigInt(25); + listWriter.list().bigInt().writeBigInt(30); + listWriter.list().bigInt().writeBigInt(35); + listWriter.list().endList(); + + listWriter.endList(); + + assertEquals(2, listVector.getMutator().getLastSet()); + + listVector.getMutator().setValueCount(2); + + final ListVector.Accessor accessor = listVector.getAccessor(); + assertEquals(2, accessor.getValueCount()); + + /* get listVector value at index 0 -- the value itself is a listvector */ + Object result = accessor.getObject(0); + ArrayList> resultSet = (ArrayList>)result; + ArrayList list; + + assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ + assertEquals(3, resultSet.get(0).size()); /* size of first inner list */ + assertEquals(4, resultSet.get(1).size()); /* size of second inner list */ + + list = resultSet.get(0); + assertEquals(new Long(50), list.get(0)); + assertEquals(new Long(100), list.get(1)); + assertEquals(new Long(200), list.get(2)); + + list = resultSet.get(1); + assertEquals(new Long(75), list.get(0)); + assertEquals(new Long(125), list.get(1)); + assertEquals(new Long(150), list.get(2)); + assertEquals(new Long(175), list.get(3)); + + /* get listVector value at index 1 -- the value itself is a listvector */ + result = accessor.getObject(1); + resultSet = (ArrayList>)result; + + assertEquals(3, resultSet.size()); /* 3 inner lists at index 1 */ + assertEquals(1, resultSet.get(0).size()); /* size of first inner list */ + assertEquals(2, resultSet.get(1).size()); /* size of second inner list */ + assertEquals(3, resultSet.get(2).size()); /* size of third inner list */ + + list = resultSet.get(0); + assertEquals(new Long(10), list.get(0)); + + list = resultSet.get(1); + assertEquals(new Long(15), list.get(0)); + assertEquals(new Long(20), list.get(1)); + + list = resultSet.get(2); + assertEquals(new Long(25), list.get(0)); + assertEquals(new Long(30), list.get(1)); + assertEquals(new Long(35), list.get(2)); + + /* check underlying bitVector */ + assertFalse(accessor.isNull(0)); + assertFalse(accessor.isNull(1)); + + /* check underlying offsetVector */ + UInt4Vector offsetVector = listVector.getOffsetVector(); + final UInt4Vector.Accessor offsetAccessor = offsetVector.getAccessor(); + + /* listVector has 2 lists at index 0 and 3 lists at index 1 */ + assertEquals(0, offsetAccessor.get(0)); + assertEquals(2, offsetAccessor.get(1)); + assertEquals(5, offsetAccessor.get(2)); + } + } }