Skip to content

Commit 12a6d76

Browse files
xxlaykxxxxlaykxx
committed
GH-836: more tests
1 parent 8835842 commit 12a6d76

File tree

2 files changed

+139
-0
lines changed

2 files changed

+139
-0
lines changed

‎vector/src/test/java/org/apache/arrow/vector/TestListVector.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1271,6 +1271,49 @@ public void testListVectorReaderForExtensionType() throws Exception {
12711271
}
12721272
}
12731273

1274+
@Test
1275+
public void testCopyFromForExtensionType() throws Exception {
1276+
try (ListVector inVector = ListVector.empty("input", allocator);
1277+
ListVector outVector = ListVector.empty("output", allocator)) {
1278+
UnionListWriter writer = inVector.getWriter();
1279+
writer.allocate();
1280+
writer.setPosition(0);
1281+
UUID u1 = UUID.randomUUID();
1282+
UUID u2 = UUID.randomUUID();
1283+
writer.startList();
1284+
ExtensionWriter extensionWriter = writer.extension(new UuidType());
1285+
extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory());
1286+
extensionWriter.writeExtension(u1);
1287+
extensionWriter.writeExtension(u2);
1288+
extensionWriter.writeNull();
1289+
writer.endList();
1290+
1291+
writer.setValueCount(1);
1292+
1293+
// copy values from input to output
1294+
outVector.allocateNew();
1295+
outVector.copyFrom(0, 0, inVector, new UuidWriterFactory());
1296+
outVector.setValueCount(1);
1297+
1298+
UnionListReader reader = outVector.getReader();
1299+
assertTrue(reader.isSet(), "shouldn't be null");
1300+
reader.setPosition(0);
1301+
reader.next();
1302+
FieldReader uuidReader = reader.reader();
1303+
UuidHolder holder = new UuidHolder();
1304+
uuidReader.read(holder);
1305+
ByteBuffer bb = ByteBuffer.wrap(holder.value);
1306+
UUID actualUuid = new UUID(bb.getLong(), bb.getLong());
1307+
assertEquals(u1, actualUuid);
1308+
reader.next();
1309+
uuidReader = reader.reader();
1310+
uuidReader.read(holder);
1311+
bb = ByteBuffer.wrap(holder.value);
1312+
actualUuid = new UUID(bb.getLong(), bb.getLong());
1313+
assertEquals(u2, actualUuid);
1314+
}
1315+
}
1316+
12741317
private void writeIntValues(UnionListWriter writer, int[] values) {
12751318
writer.startList();
12761319
for (int v : values) {

‎vector/src/test/java/org/apache/arrow/vector/TestMapVector.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,24 +22,30 @@
2222
import static org.junit.jupiter.api.Assertions.assertSame;
2323
import static org.junit.jupiter.api.Assertions.assertTrue;
2424

25+
import java.nio.ByteBuffer;
2526
import java.util.ArrayList;
2627
import java.util.Collections;
2728
import java.util.List;
2829
import java.util.Map;
30+
import java.util.UUID;
2931
import org.apache.arrow.memory.ArrowBuf;
3032
import org.apache.arrow.memory.BufferAllocator;
3133
import org.apache.arrow.vector.complex.MapVector;
3234
import org.apache.arrow.vector.complex.StructVector;
3335
import org.apache.arrow.vector.complex.impl.UnionMapReader;
3436
import org.apache.arrow.vector.complex.impl.UnionMapWriter;
37+
import org.apache.arrow.vector.complex.impl.UuidWriterFactory;
3538
import org.apache.arrow.vector.complex.reader.FieldReader;
39+
import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
3640
import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter;
3741
import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;
3842
import org.apache.arrow.vector.complex.writer.FieldWriter;
43+
import org.apache.arrow.vector.holder.UuidHolder;
3944
import org.apache.arrow.vector.types.Types.MinorType;
4045
import org.apache.arrow.vector.types.pojo.ArrowType;
4146
import org.apache.arrow.vector.types.pojo.Field;
4247
import org.apache.arrow.vector.types.pojo.FieldType;
48+
import org.apache.arrow.vector.types.pojo.UuidType;
4349
import org.apache.arrow.vector.util.JsonStringArrayList;
4450
import org.apache.arrow.vector.util.TransferPair;
4551
import org.junit.jupiter.api.AfterEach;
@@ -1263,4 +1269,94 @@ public void testMapTypeReturnsSupportedMapWriter() {
12631269
assertEquals(11, getResultValue(resultStruct));
12641270
}
12651271
}
1272+
1273+
@Test
1274+
public void testMapVectorWithExtensionType() throws Exception {
1275+
try (final MapVector inVector = MapVector.empty("map", allocator, false)) {
1276+
inVector.allocateNew();
1277+
UnionMapWriter writer = inVector.getWriter();
1278+
writer.setPosition(0);
1279+
UUID u1 = UUID.randomUUID();
1280+
UUID u2 = UUID.randomUUID();
1281+
writer.startMap();
1282+
writer.startEntry();
1283+
writer.key().bigInt().writeBigInt(0);
1284+
ExtensionWriter extensionWriter = writer.value().extension(new UuidType());
1285+
extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory());
1286+
extensionWriter.writeExtension(u1);
1287+
writer.endEntry();
1288+
writer.startEntry();
1289+
writer.key().bigInt().writeBigInt(1);
1290+
extensionWriter = writer.value().extension(new UuidType());
1291+
extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory());
1292+
extensionWriter.writeExtension(u2);
1293+
writer.endEntry();
1294+
writer.endMap();
1295+
1296+
writer.setValueCount(1);
1297+
1298+
UnionMapReader mapReader = inVector.getReader();
1299+
mapReader.setPosition(0);
1300+
mapReader.next();
1301+
FieldReader uuidReader = mapReader.value();
1302+
UuidHolder holder = new UuidHolder();
1303+
uuidReader.read(holder);
1304+
ByteBuffer bb = ByteBuffer.wrap(holder.value);
1305+
UUID actualUuid = new UUID(bb.getLong(), bb.getLong());
1306+
assertEquals(u1, actualUuid);
1307+
mapReader.next();
1308+
uuidReader = mapReader.value();
1309+
uuidReader.read(holder);
1310+
bb = ByteBuffer.wrap(holder.value);
1311+
actualUuid = new UUID(bb.getLong(), bb.getLong());
1312+
assertEquals(u2, actualUuid);
1313+
}
1314+
}
1315+
1316+
@Test
1317+
public void testCopyFromForExtensionType() throws Exception {
1318+
try (final MapVector inVector = MapVector.empty("in", allocator, false);
1319+
final MapVector outVector = MapVector.empty("out", allocator, false)) {
1320+
inVector.allocateNew();
1321+
UnionMapWriter writer = inVector.getWriter();
1322+
writer.setPosition(0);
1323+
UUID u1 = UUID.randomUUID();
1324+
UUID u2 = UUID.randomUUID();
1325+
writer.startMap();
1326+
writer.startEntry();
1327+
writer.key().bigInt().writeBigInt(0);
1328+
ExtensionWriter extensionWriter = writer.value().extension(new UuidType());
1329+
extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory());
1330+
extensionWriter.writeExtension(u1);
1331+
writer.endEntry();
1332+
writer.startEntry();
1333+
writer.key().bigInt().writeBigInt(1);
1334+
extensionWriter = writer.value().extension(new UuidType());
1335+
extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory());
1336+
extensionWriter.writeExtension(u2);
1337+
writer.endEntry();
1338+
writer.endMap();
1339+
1340+
writer.setValueCount(1);
1341+
outVector.allocateNew();
1342+
outVector.copyFrom(0, 0, inVector, new UuidWriterFactory());
1343+
outVector.setValueCount(1);
1344+
1345+
UnionMapReader mapReader = outVector.getReader();
1346+
mapReader.setPosition(0);
1347+
mapReader.next();
1348+
FieldReader uuidReader = mapReader.value();
1349+
UuidHolder holder = new UuidHolder();
1350+
uuidReader.read(holder);
1351+
ByteBuffer bb = ByteBuffer.wrap(holder.value);
1352+
UUID actualUuid = new UUID(bb.getLong(), bb.getLong());
1353+
assertEquals(u1, actualUuid);
1354+
mapReader.next();
1355+
uuidReader = mapReader.value();
1356+
uuidReader.read(holder);
1357+
bb = ByteBuffer.wrap(holder.value);
1358+
actualUuid = new UUID(bb.getLong(), bb.getLong());
1359+
assertEquals(u2, actualUuid);
1360+
}
1361+
}
12661362
}

0 commit comments

Comments
 (0)