HedgehogCode commented on a change in pull request #8363:
URL: https://github.com/apache/arrow/pull/8363#discussion_r503140409
##########
File path:
java/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowReaderWriter.java
##########
@@ -305,6 +325,77 @@ public void testWriteReadWithDictionaries() throws
IOException {
}
}
+ @Test
+ public void testWriteReadWithStructDictionaries() throws IOException {
+ DictionaryProvider.MapDictionaryProvider provider = new
DictionaryProvider.MapDictionaryProvider();
+ provider.put(dictionary4);
+
+ final StructVector vector = newVector(StructVector.class, "D4",
MinorType.STRUCT, allocator);
+ final Map<String, List<Integer>> values = new HashMap<>();
+ // Index: 0, 2, 1, 2, 1, 0, 0
+ values.put("a", Arrays.asList(1, 3, 2, 3, 2, 1, 1));
+ values.put("b", Arrays.asList(4, 6, 5, 6, 5, 4, 4));
+ setVector(vector, values);
+ FieldVector encodedVector = (FieldVector) DictionaryEncoder.encode(vector,
dictionary4);
+
+ List<Field> fields = Arrays.asList(encodedVector.getField());
+ List<FieldVector> vectors = Collections2.asImmutableList(encodedVector);
+ try (VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors,
encodedVector.getValueCount());
+ ByteArrayOutputStream out = new ByteArrayOutputStream();
+ ArrowFileWriter writer = new ArrowFileWriter(root, provider,
newChannel(out));) {
+
+ writer.start();
+ writer.writeBatch();
+ writer.end();
+
+ try (SeekableReadChannel channel = new SeekableReadChannel(
+ new ByteArrayReadableSeekableByteChannel(out.toByteArray()));
+ ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
+ final VectorSchemaRoot readRoot = reader.getVectorSchemaRoot();
+ final Schema readSchema = readRoot.getSchema();
+ assertEquals(root.getSchema(), readSchema);
+ assertEquals(1, reader.getDictionaryBlocks().size());
+ assertEquals(1, reader.getRecordBlocks().size());
+
+ reader.loadNextBatch();
+ assertEquals(1, readRoot.getFieldVectors().size());
+ assertEquals(1, reader.getDictionaryVectors().size());
+
+ // Read the encoded vector and check it
+ final FieldVector readEncoded = readRoot.getVector(0);
+ assertEquals(encodedVector.getValueCount(),
readEncoded.getValueCount());
+ assertTrue(new RangeEqualsVisitor(encodedVector, readEncoded)
+ .rangeEquals(new Range(0, 0, encodedVector.getValueCount())));
+
+ // Read the dictionary
+ final Map<Long, Dictionary> readDictionaryMap =
reader.getDictionaryVectors();
+ final Dictionary readDictionary =
+
readDictionaryMap.get(readEncoded.getField().getDictionary().getId());
+ assertNotNull(readDictionary);
+
+ // Assert the dictionary vector is correct
+ final FieldVector readDictionaryVector = readDictionary.getVector();
+ assertEquals(dictionaryVector4.getValueCount(),
readDictionaryVector.getValueCount());
+ final BiFunction<ValueVector, ValueVector, Boolean>
typeComparatorIgnoreName =
+ (v1, v2) -> new TypeEqualsVisitor(v1, false, true).equals(v2);
+ assertTrue("Dictionary vectors are not equal",
+ new RangeEqualsVisitor(dictionaryVector4, readDictionaryVector,
+ typeComparatorIgnoreName)
+ .rangeEquals(new Range(0, 0,
dictionaryVector4.getValueCount())));
+
+ // Assert the decoded vector is correct
+ final ValueVector readVector = DictionaryEncoder.decode(readEncoded,
readDictionary);
+ assertEquals(vector.getValueCount(), readVector.getValueCount());
+ assertTrue("Decoded vectors are not equal",
+ new RangeEqualsVisitor(vector, readVector,
typeComparatorIgnoreName)
+ .rangeEquals(new Range(0, 0, vector.getValueCount())));
+
+ vector.close();
+ readVector.close();
Review comment:
Yes, I changed it. I did not do it before because I did not want the
deep nesting.
----------------------------------------------------------------
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:
[email protected]
