This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 7df4748319 GH-44353: [Java] Implement `map()` for `UnionMapWriter`
(#44390)
7df4748319 is described below
commit 7df47483197186b564a9882ac6d1b5f32b2e3d51
Author: 0xderek <[email protected]>
AuthorDate: Tue Oct 15 08:41:45 2024 +0800
GH-44353: [Java] Implement `map()` for `UnionMapWriter` (#44390)
### Rationale for this change
See #44353
`UnionMapWriter` does not override `map()` and falls back to implementation
in `UnionListWriter` which causes exception when used.
### What changes are included in this PR?
Implement `map()` for `UnionMapWriter`.
### Are these changes tested?
Yes.
### Are there any user-facing changes?
Changed behavior of the related function.
* GitHub Issue: #44353
Authored-by: 0xderek <[email protected]>
Signed-off-by: David Li <[email protected]>
---
.../src/main/codegen/templates/UnionMapWriter.java | 12 ++++++
.../org/apache/arrow/vector/TestMapVector.java | 49 +++++++++++++++++++---
2 files changed, 55 insertions(+), 6 deletions(-)
diff --git a/java/vector/src/main/codegen/templates/UnionMapWriter.java
b/java/vector/src/main/codegen/templates/UnionMapWriter.java
index 606f880377..90b55cb65e 100644
--- a/java/vector/src/main/codegen/templates/UnionMapWriter.java
+++ b/java/vector/src/main/codegen/templates/UnionMapWriter.java
@@ -219,4 +219,16 @@ public class UnionMapWriter extends UnionListWriter {
return super.map();
}
}
+
+ @Override
+ public MapWriter map() {
+ switch (mode) {
+ case KEY:
+ return entryWriter.map(MapVector.KEY_NAME);
+ case VALUE:
+ return entryWriter.map(MapVector.VALUE_NAME);
+ default:
+ return super.map();
+ }
+ }
}
diff --git
a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
index 213ffced27..a4197c50b5 100644
--- a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
+++ b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
@@ -640,11 +640,12 @@ public class TestMapVector {
MapWriter valueWriter;
// we are essentially writing Map<Long, Map<Long, Long>>
- // populate map vector with the following three records
+ // populate map vector with the following four records
// [
// null,
// [1:[50: 100, 200:400], 2:[75: 175, 150: 250]],
- // [3:[10: 20], 4:[15: 20], 5:[25: 30, 35: null]]
+ // [3:[10: 20], 4:[15: 20], 5:[25: 30, 35: null]],
+ // [8:[15: 30, 10: 20]]
// ]
/* write null at index 0 */
@@ -706,11 +707,26 @@ public class TestMapVector {
mapWriter.endMap();
- assertEquals(2, mapVector.getLastSet());
+ /* write one or more maps at index 3 */
+ mapWriter.setPosition(3);
+ mapWriter.startMap();
+
+ mapWriter.startEntry();
+ mapWriter.key().bigInt().writeBigInt(8);
+ valueWriter = mapWriter.value().map();
+ valueWriter.startMap();
+ writeEntry(valueWriter, 15, 30L);
+ writeEntry(valueWriter, 10, 20L);
+ valueWriter.endMap();
+ mapWriter.endEntry();
+
+ mapWriter.endMap();
+
+ assertEquals(3, mapVector.getLastSet());
- mapWriter.setValueCount(3);
+ mapWriter.setValueCount(4);
- assertEquals(3, mapVector.getValueCount());
+ assertEquals(4, mapVector.getValueCount());
// Get mapVector element at index 0
Object result = mapVector.getObject(0);
@@ -784,19 +800,40 @@ public class TestMapVector {
assertEquals(35L, getResultKey(innerMap));
assertNull(innerMap.get(MapVector.VALUE_NAME));
+ // Get mapVector element at index 3
+ result = mapVector.getObject(3);
+ resultSet = (ArrayList<?>) result;
+
+ // only 1 map entry at index 3
+ assertEquals(1, resultSet.size());
+
+ resultStruct = (Map<?, ?>) resultSet.get(0);
+ assertEquals(8L, getResultKey(resultStruct));
+ list = (ArrayList<Map<?, ?>>) getResultValue(resultStruct);
+ assertEquals(2, list.size()); // value is a list of 2 maps
+ innerMap = list.get(0);
+ assertEquals(15L, getResultKey(innerMap));
+ assertEquals(30L, getResultValue(innerMap));
+ innerMap = list.get(1);
+ assertEquals(10L, getResultKey(innerMap));
+ assertEquals(20L, getResultValue(innerMap));
+
/* check underlying bitVector */
assertTrue(mapVector.isNull(0));
assertFalse(mapVector.isNull(1));
assertFalse(mapVector.isNull(2));
+ assertFalse(mapVector.isNull(3));
/* check underlying offsets */
final ArrowBuf offsetBuffer = mapVector.getOffsetBuffer();
- /* mapVector has 0 entries at index 0, 2 entries at index 1, and 3
entries at index 2 */
+ // mapVector has 0 entries at index 0, 2 entries at index 1, 3 entries
at index 2,
+ // and 1 entry at index 3
assertEquals(0, offsetBuffer.getInt(0 * MapVector.OFFSET_WIDTH));
assertEquals(0, offsetBuffer.getInt(1 * MapVector.OFFSET_WIDTH));
assertEquals(2, offsetBuffer.getInt(2 * MapVector.OFFSET_WIDTH));
assertEquals(5, offsetBuffer.getInt(3 * MapVector.OFFSET_WIDTH));
+ assertEquals(6, offsetBuffer.getInt(4 * MapVector.OFFSET_WIDTH));
}
}
