siddharthteotia commented on a change in pull request #9151:
URL: https://github.com/apache/arrow/pull/9151#discussion_r566552494
##########
File path: java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
##########
@@ -624,6 +625,423 @@ public void testMapWithListValue() throws Exception {
}
}
+ @Test
+ public void testMapWithMapValue() throws Exception {
+ try (MapVector mapVector = MapVector.empty("sourceVector", allocator,
false)) {
+
+ UnionMapWriter mapWriter = mapVector.getWriter();
+ MapWriter valueWriter;
+
+ // we are essentially writing Map<Long, Map<Long, Long>>
+
+ mapWriter.setPosition(0);
+ mapWriter.startMap();
+
+ mapWriter.startEntry();
+ mapWriter.key().bigInt().writeBigInt(1);
+ valueWriter = mapWriter.value().map(false);
+ valueWriter.startMap();
+ valueWriter.startEntry();
+ valueWriter.key().bigInt().writeBigInt(50);
+ valueWriter.value().bigInt().writeBigInt(100);
+ valueWriter.endEntry();
+ valueWriter.startEntry();
+ valueWriter.key().bigInt().writeBigInt(200);
+ valueWriter.value().bigInt().writeBigInt(400);
+ valueWriter.endEntry();
+ valueWriter.endMap();
+ mapWriter.endEntry();
+
+ mapWriter.startEntry();
+ mapWriter.key().bigInt().writeBigInt(2);
+ valueWriter = mapWriter.value().map(false);
+ valueWriter.startMap();
+ valueWriter.startEntry();
+ valueWriter.key().bigInt().writeBigInt(75);
+ valueWriter.value().bigInt().writeBigInt(175);
+ valueWriter.endEntry();
+ valueWriter.startEntry();
+ valueWriter.key().bigInt().writeBigInt(150);
+ valueWriter.value().bigInt().writeBigInt(250);
+ valueWriter.endEntry();
+ valueWriter.endMap();
+ mapWriter.endEntry();
+
+ mapWriter.endMap();
+
+ /* write one or more maps at index 1 */
+ mapWriter.setPosition(1);
+ mapWriter.startMap();
+
+ mapWriter.startEntry();
+ mapWriter.key().bigInt().writeBigInt(3);
+ valueWriter = mapWriter.value().map(true);
+ valueWriter.startMap();
+ valueWriter.startEntry();
+ valueWriter.key().bigInt().writeBigInt(10);
+ valueWriter.value().bigInt().writeBigInt(20);
+ valueWriter.endEntry();
+ valueWriter.endMap();
+ mapWriter.endEntry();
+
+ mapWriter.startEntry();
+ mapWriter.key().bigInt().writeBigInt(4);
+ valueWriter = mapWriter.value().map(false);
+ valueWriter.startMap();
+ valueWriter.startEntry();
+ valueWriter.key().bigInt().writeBigInt(15);
+ valueWriter.value().bigInt().writeBigInt(20);
+ valueWriter.endEntry();
+ valueWriter.endMap();
+ mapWriter.endEntry();
+
+ mapWriter.startEntry();
+ mapWriter.key().bigInt().writeBigInt(5);
+ valueWriter = mapWriter.value().map(false);
+ valueWriter.startMap();
+ valueWriter.startEntry();
+ valueWriter.key().bigInt().writeBigInt(25);
+ valueWriter.value().bigInt().writeBigInt(30);
+ valueWriter.endEntry();
+ valueWriter.startEntry();
+ valueWriter.key().bigInt().writeBigInt(35);
+ valueWriter.endEntry();
+ valueWriter.endMap();
+ mapWriter.endEntry();
+
+ mapWriter.endMap();
+
+ assertEquals(1, mapVector.getLastSet());
+
+ mapWriter.setValueCount(2);
+
+ assertEquals(2, mapVector.getValueCount());
+
+ // Get mapVector element at index 0
+ Object result = mapVector.getObject(0);
+ ArrayList<?> resultSet = (ArrayList<?>) result;
+
+ // 2 map entries at index 0
+ assertEquals(2, resultSet.size());
+
+ // First Map entry
+ Map<?, ?> resultStruct = (Map<?, ?>) resultSet.get(0);
+ assertEquals(1L, getResultKey(resultStruct));
+ ArrayList<Map<?, ?>> list = (ArrayList<Map<?, ?>>)
getResultValue(resultStruct);
+ assertEquals(2, list.size()); // value is a list of 2 two maps
+ Map<?, ?> innerMap = list.get(0);
+ assertEquals(50L, getResultKey(innerMap));
+ assertEquals(100L, getResultValue(innerMap));
+ innerMap = list.get(1);
+ assertEquals(200L, getResultKey(innerMap));
+ assertEquals(400L, getResultValue(innerMap));
+
+ // Second Map entry
+ resultStruct = (Map<?, ?>) resultSet.get(1);
+ assertEquals(2L, getResultKey(resultStruct));
+ list = (ArrayList<Map<?, ?>>) getResultValue(resultStruct);
+ assertEquals(2, list.size()); // value is a list of two maps
+ innerMap = list.get(0);
+ assertEquals(75L, getResultKey(innerMap));
+ assertEquals(175L, getResultValue(innerMap));
+ innerMap = list.get(1);
+ assertEquals(150L, getResultKey(innerMap));
+ assertEquals(250L, getResultValue(innerMap));
+
+ // Get mapVector element at index 1
+ result = mapVector.getObject(1);
+ resultSet = (ArrayList<?>) result;
+
+ // 3 map entries at index 1
+ assertEquals(3, resultSet.size());
+
+ // First Map entry
+ resultStruct = (Map<?, ?>) resultSet.get(0);
+ assertEquals(3L, getResultKey(resultStruct));
+ list = (ArrayList<Map<?, ?>>) getResultValue(resultStruct);
+ assertEquals(1, list.size()); // value is a list of maps with 1 element
+ innerMap = list.get(0);
+ assertEquals(10L, getResultKey(innerMap));
+ assertEquals(20L, getResultValue(innerMap));
+
+ // Second Map entry
+ resultStruct = (Map<?, ?>) resultSet.get(1);
+ assertEquals(4L, getResultKey(resultStruct));
+ list = (ArrayList<Map<?, ?>>) getResultValue(resultStruct);
+ assertEquals(1, list.size()); // value is a list of maps with 1 element
+ innerMap = list.get(0);
+ assertEquals(15L, getResultKey(innerMap));
+ assertEquals(20L, getResultValue(innerMap));
+
+ // Third Map entry
+ resultStruct = (Map<?, ?>) resultSet.get(2);
+ assertEquals(5L, getResultKey(resultStruct));
+ list = (ArrayList<Map<?, ?>>) getResultValue(resultStruct);
+ assertEquals(2, list.size()); // value is a list of maps with 2 elements
+ innerMap = list.get(0);
+ assertEquals(25L, getResultKey(innerMap));
+ assertEquals(30L, getResultValue(innerMap));
+ innerMap = list.get(1);
+ assertEquals(35L, getResultKey(innerMap));
+ assertNull(innerMap.get(MapVector.VALUE_NAME));
+
+ /* check underlying bitVector */
+ assertFalse(mapVector.isNull(0));
+ assertFalse(mapVector.isNull(1));
+
+ /* check underlying offsets */
+ final ArrowBuf offsetBuffer = mapVector.getOffsetBuffer();
+
+ /* mapVector has 2 entries at index 0 and 3 entries at index 1 */
+ assertEquals(0, offsetBuffer.getInt(0 * MapVector.OFFSET_WIDTH));
+ assertEquals(2, offsetBuffer.getInt(1 * MapVector.OFFSET_WIDTH));
+ assertEquals(5, offsetBuffer.getInt(2 * MapVector.OFFSET_WIDTH));
+ }
+ }
+
+ @Test
+ public void testMapWithMapKeyAndMapValue() throws Exception {
+ try (MapVector mapVector = MapVector.empty("sourceVector", allocator,
false)) {
+
+ UnionMapWriter mapWriter = mapVector.getWriter();
+ MapWriter keyWriter;
+ MapWriter valueWriter;
+
+ // we are essentially writing Map<Map<Integer, Integer>, Map<Long, Long>>
+
+ mapWriter.setPosition(0);
Review comment:
I think is going to be hard to debug without actually having a
description of what the map will look like finally. So would be great to add
that in comments. Also, you should consider extracting this into a function and
pass all the values and writers.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
