]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one or more | |
3 | * contributor license agreements. See the NOTICE file distributed with | |
4 | * this work for additional information regarding copyright ownership. | |
5 | * The ASF licenses this file to You under the Apache License, Version 2.0 | |
6 | * (the "License"); you may not use this file except in compliance with | |
7 | * the License. You may obtain a copy of the License at | |
8 | * | |
9 | * http://www.apache.org/licenses/LICENSE-2.0 | |
10 | * | |
11 | * Unless required by applicable law or agreed to in writing, software | |
12 | * distributed under the License is distributed on an "AS IS" BASIS, | |
13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
14 | * See the License for the specific language governing permissions and | |
15 | * limitations under the License. | |
16 | */ | |
17 | ||
18 | package org.apache.arrow.c; | |
19 | ||
20 | import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; | |
21 | import static org.junit.jupiter.api.Assertions.assertEquals; | |
22 | import static org.junit.jupiter.api.Assertions.assertThrows; | |
23 | import static org.junit.jupiter.api.Assertions.assertTrue; | |
24 | ||
25 | import java.nio.ByteBuffer; | |
26 | import java.nio.charset.StandardCharsets; | |
27 | import java.util.ArrayList; | |
28 | import java.util.Arrays; | |
29 | import java.util.Collections; | |
30 | import java.util.HashMap; | |
31 | import java.util.List; | |
32 | import java.util.Map; | |
33 | import java.util.UUID; | |
34 | import java.util.stream.Collectors; | |
35 | ||
36 | import org.apache.arrow.c.ArrowArray; | |
37 | import org.apache.arrow.c.ArrowSchema; | |
38 | import org.apache.arrow.c.Data; | |
39 | import org.apache.arrow.memory.ArrowBuf; | |
40 | import org.apache.arrow.memory.BufferAllocator; | |
41 | import org.apache.arrow.memory.RootAllocator; | |
42 | import org.apache.arrow.memory.util.hash.ArrowBufHasher; | |
43 | import org.apache.arrow.vector.BigIntVector; | |
44 | import org.apache.arrow.vector.BitVector; | |
45 | import org.apache.arrow.vector.DateDayVector; | |
46 | import org.apache.arrow.vector.DateMilliVector; | |
47 | import org.apache.arrow.vector.DecimalVector; | |
48 | import org.apache.arrow.vector.DurationVector; | |
49 | import org.apache.arrow.vector.ExtensionTypeVector; | |
50 | import org.apache.arrow.vector.FieldVector; | |
51 | import org.apache.arrow.vector.FixedSizeBinaryVector; | |
52 | import org.apache.arrow.vector.Float4Vector; | |
53 | import org.apache.arrow.vector.Float8Vector; | |
54 | import org.apache.arrow.vector.IntVector; | |
55 | import org.apache.arrow.vector.IntervalDayVector; | |
56 | import org.apache.arrow.vector.IntervalYearVector; | |
57 | import org.apache.arrow.vector.LargeVarBinaryVector; | |
58 | import org.apache.arrow.vector.LargeVarCharVector; | |
59 | import org.apache.arrow.vector.NullVector; | |
60 | import org.apache.arrow.vector.SmallIntVector; | |
61 | import org.apache.arrow.vector.TimeMicroVector; | |
62 | import org.apache.arrow.vector.TimeMilliVector; | |
63 | import org.apache.arrow.vector.TimeNanoVector; | |
64 | import org.apache.arrow.vector.TimeSecVector; | |
65 | import org.apache.arrow.vector.TimeStampMicroTZVector; | |
66 | import org.apache.arrow.vector.TimeStampMicroVector; | |
67 | import org.apache.arrow.vector.TimeStampMilliTZVector; | |
68 | import org.apache.arrow.vector.TimeStampMilliVector; | |
69 | import org.apache.arrow.vector.TimeStampNanoTZVector; | |
70 | import org.apache.arrow.vector.TimeStampNanoVector; | |
71 | import org.apache.arrow.vector.TimeStampSecTZVector; | |
72 | import org.apache.arrow.vector.TimeStampSecVector; | |
73 | import org.apache.arrow.vector.TinyIntVector; | |
74 | import org.apache.arrow.vector.UInt1Vector; | |
75 | import org.apache.arrow.vector.UInt2Vector; | |
76 | import org.apache.arrow.vector.UInt4Vector; | |
77 | import org.apache.arrow.vector.UInt8Vector; | |
78 | import org.apache.arrow.vector.ValueVector; | |
79 | import org.apache.arrow.vector.VarBinaryVector; | |
80 | import org.apache.arrow.vector.VarCharVector; | |
81 | import org.apache.arrow.vector.VectorSchemaRoot; | |
82 | import org.apache.arrow.vector.ZeroVector; | |
83 | import org.apache.arrow.vector.compare.VectorEqualsVisitor; | |
84 | import org.apache.arrow.vector.complex.FixedSizeListVector; | |
85 | import org.apache.arrow.vector.complex.LargeListVector; | |
86 | import org.apache.arrow.vector.complex.ListVector; | |
87 | import org.apache.arrow.vector.complex.MapVector; | |
88 | import org.apache.arrow.vector.complex.StructVector; | |
89 | import org.apache.arrow.vector.complex.UnionVector; | |
90 | import org.apache.arrow.vector.complex.impl.UnionMapWriter; | |
91 | import org.apache.arrow.vector.holders.IntervalDayHolder; | |
92 | import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder; | |
93 | import org.apache.arrow.vector.holders.NullableUInt4Holder; | |
94 | import org.apache.arrow.vector.types.TimeUnit; | |
95 | import org.apache.arrow.vector.types.Types.MinorType; | |
96 | import org.apache.arrow.vector.types.pojo.ArrowType; | |
97 | import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; | |
98 | import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; | |
99 | import org.apache.arrow.vector.types.pojo.Field; | |
100 | import org.apache.arrow.vector.types.pojo.FieldType; | |
101 | import org.apache.arrow.vector.types.pojo.Schema; | |
102 | import org.junit.jupiter.api.AfterEach; | |
103 | import org.junit.jupiter.api.BeforeEach; | |
104 | import org.junit.jupiter.api.Test; | |
105 | ||
106 | public class RoundtripTest { | |
107 | private static final String EMPTY_SCHEMA_PATH = ""; | |
108 | private RootAllocator allocator = null; | |
109 | ||
110 | @BeforeEach | |
111 | public void setUp() { | |
112 | allocator = new RootAllocator(Long.MAX_VALUE); | |
113 | } | |
114 | ||
115 | @AfterEach | |
116 | public void tearDown() { | |
117 | allocator.close(); | |
118 | } | |
119 | ||
120 | FieldVector vectorRoundtrip(FieldVector vector) { | |
121 | // Consumer allocates empty structures | |
122 | try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); | |
123 | ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { | |
124 | ||
125 | // Producer creates structures from existing memory pointers | |
126 | try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress()); | |
127 | ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) { | |
128 | // Producer exports vector into the C Data Interface structures | |
129 | Data.exportVector(allocator, vector, null, arrowArray, arrowSchema); | |
130 | } | |
131 | ||
132 | // Consumer imports vector | |
133 | return Data.importVector(allocator, consumerArrowArray, consumerArrowSchema, null); | |
134 | } | |
135 | } | |
136 | ||
137 | VectorSchemaRoot vectorSchemaRootRoundtrip(VectorSchemaRoot root) { | |
138 | // Consumer allocates empty structures | |
139 | try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); | |
140 | ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { | |
141 | ||
142 | // Producer creates structures from existing memory pointers | |
143 | try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress()); | |
144 | ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) { | |
145 | // Producer exports vector into the C Data Interface structures | |
146 | Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchema); | |
147 | } | |
148 | ||
149 | // Consumer imports vector | |
150 | return Data.importVectorSchemaRoot(allocator, consumerArrowArray, consumerArrowSchema, null); | |
151 | } | |
152 | } | |
153 | ||
154 | boolean roundtrip(FieldVector vector, Class<?> clazz) { | |
155 | try (ValueVector imported = vectorRoundtrip(vector)) { | |
156 | assertTrue(clazz.isInstance(imported), String.format("expected %s but was %s", clazz, imported.getClass())); | |
157 | return VectorEqualsVisitor.vectorEquals(vector, imported); | |
158 | } | |
159 | } | |
160 | ||
161 | @Test | |
162 | public void testBitVector() { | |
163 | BitVector imported; | |
164 | ||
165 | try (final BitVector vector = new BitVector(EMPTY_SCHEMA_PATH, allocator)) { | |
166 | vector.allocateNew(1024); | |
167 | vector.setValueCount(1024); | |
168 | ||
169 | // Put and set a few values | |
170 | vector.set(0, 1); | |
171 | vector.set(1, 0); | |
172 | vector.set(100, 0); | |
173 | vector.set(1022, 1); | |
174 | ||
175 | vector.setValueCount(1024); | |
176 | ||
177 | imported = (BitVector) vectorRoundtrip(vector); | |
178 | assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); | |
179 | } | |
180 | ||
181 | assertEquals(1, imported.get(0)); | |
182 | assertEquals(0, imported.get(1)); | |
183 | assertEquals(0, imported.get(100)); | |
184 | assertEquals(1, imported.get(1022)); | |
185 | assertEquals(1020, imported.getNullCount()); | |
186 | imported.close(); | |
187 | } | |
188 | ||
189 | @Test | |
190 | public void testIntVector() { | |
191 | IntVector imported; | |
192 | try (final IntVector vector = new IntVector("v", allocator)) { | |
193 | setVector(vector, 1, 2, 3, null); | |
194 | imported = (IntVector) vectorRoundtrip(vector); | |
195 | assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); | |
196 | } | |
197 | assertEquals(1, imported.get(0)); | |
198 | assertEquals(2, imported.get(1)); | |
199 | assertEquals(3, imported.get(2)); | |
200 | assertEquals(4, imported.getValueCount()); | |
201 | assertEquals(1, imported.getNullCount()); | |
202 | imported.close(); | |
203 | } | |
204 | ||
205 | @Test | |
206 | public void testBigIntVector() { | |
207 | BigIntVector imported; | |
208 | try (final BigIntVector vector = new BigIntVector("v", allocator)) { | |
209 | setVector(vector, 1L, 2L, 3L, null); | |
210 | imported = (BigIntVector) vectorRoundtrip(vector); | |
211 | assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); | |
212 | } | |
213 | assertEquals(1, imported.get(0)); | |
214 | assertEquals(2, imported.get(1)); | |
215 | assertEquals(3, imported.get(2)); | |
216 | assertEquals(4, imported.getValueCount()); | |
217 | assertEquals(1, imported.getNullCount()); | |
218 | imported.close(); | |
219 | } | |
220 | ||
221 | @Test | |
222 | public void testDateDayVector() { | |
223 | DateDayVector imported; | |
224 | try (final DateDayVector vector = new DateDayVector("v", allocator)) { | |
225 | setVector(vector, 1, 2, 3, null); | |
226 | imported = (DateDayVector) vectorRoundtrip(vector); | |
227 | assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); | |
228 | } | |
229 | assertEquals(1, imported.get(0)); | |
230 | assertEquals(2, imported.get(1)); | |
231 | assertEquals(3, imported.get(2)); | |
232 | assertEquals(4, imported.getValueCount()); | |
233 | assertEquals(1, imported.getNullCount()); | |
234 | imported.close(); | |
235 | } | |
236 | ||
237 | @Test | |
238 | public void testDateMilliVector() { | |
239 | DateMilliVector imported; | |
240 | try (final DateMilliVector vector = new DateMilliVector("v", allocator)) { | |
241 | setVector(vector, 1L, 2L, 3L, null); | |
242 | imported = (DateMilliVector) vectorRoundtrip(vector); | |
243 | assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); | |
244 | } | |
245 | assertEquals(1, imported.get(0)); | |
246 | assertEquals(2, imported.get(1)); | |
247 | assertEquals(3, imported.get(2)); | |
248 | assertEquals(4, imported.getValueCount()); | |
249 | assertEquals(1, imported.getNullCount()); | |
250 | imported.close(); | |
251 | } | |
252 | ||
253 | @Test | |
254 | public void testDecimalVector() { | |
255 | try (final DecimalVector vector = new DecimalVector("v", allocator, 1, 1)) { | |
256 | setVector(vector, 1L, 2L, 3L, null); | |
257 | assertTrue(roundtrip(vector, DecimalVector.class)); | |
258 | } | |
259 | } | |
260 | ||
261 | @Test | |
262 | public void testDurationVector() { | |
263 | for (TimeUnit unit : TimeUnit.values()) { | |
264 | final FieldType fieldType = FieldType.nullable(new ArrowType.Duration(unit)); | |
265 | try (final DurationVector vector = new DurationVector("v", fieldType, allocator)) { | |
266 | setVector(vector, 1L, 2L, 3L, null); | |
267 | assertTrue(roundtrip(vector, DurationVector.class)); | |
268 | } | |
269 | } | |
270 | } | |
271 | ||
272 | @Test | |
273 | public void testZeroVectorEquals() { | |
274 | try (final ZeroVector vector = new ZeroVector()) { | |
275 | // A ZeroVector is imported as a NullVector | |
276 | assertTrue(roundtrip(vector, NullVector.class)); | |
277 | } | |
278 | } | |
279 | ||
280 | @Test | |
281 | public void testFixedSizeBinaryVector() { | |
282 | try (final FixedSizeBinaryVector vector = new FixedSizeBinaryVector("v", allocator, 2)) { | |
283 | setVector(vector, new byte[] { 0b0000, 0b0001 }, new byte[] { 0b0010, 0b0011 }); | |
284 | assertTrue(roundtrip(vector, FixedSizeBinaryVector.class)); | |
285 | } | |
286 | } | |
287 | ||
288 | @Test | |
289 | public void testFloat4Vector() { | |
290 | try (final Float4Vector vector = new Float4Vector("v", allocator)) { | |
291 | setVector(vector, 0.1f, 0.2f, 0.3f, null); | |
292 | assertTrue(roundtrip(vector, Float4Vector.class)); | |
293 | } | |
294 | } | |
295 | ||
296 | @Test | |
297 | public void testFloat8Vector() { | |
298 | try (final Float8Vector vector = new Float8Vector("v", allocator)) { | |
299 | setVector(vector, 0.1d, 0.2d, 0.3d, null); | |
300 | assertTrue(roundtrip(vector, Float8Vector.class)); | |
301 | } | |
302 | } | |
303 | ||
304 | @Test | |
305 | public void testIntervalDayVector() { | |
306 | try (final IntervalDayVector vector = new IntervalDayVector("v", allocator)) { | |
307 | IntervalDayHolder value = new IntervalDayHolder(); | |
308 | value.days = 5; | |
309 | value.milliseconds = 100; | |
310 | setVector(vector, value, null); | |
311 | assertTrue(roundtrip(vector, IntervalDayVector.class)); | |
312 | } | |
313 | } | |
314 | ||
315 | @Test | |
316 | public void testIntervalYearVector() { | |
317 | try (final IntervalYearVector vector = new IntervalYearVector("v", allocator)) { | |
318 | setVector(vector, 1990, 2000, 2010, 2020, null); | |
319 | assertTrue(roundtrip(vector, IntervalYearVector.class)); | |
320 | } | |
321 | } | |
322 | ||
323 | @Test | |
324 | public void testSmallIntVector() { | |
325 | try (final SmallIntVector vector = new SmallIntVector("v", allocator)) { | |
326 | setVector(vector, (short) 0, (short) 256, null); | |
327 | assertTrue(roundtrip(vector, SmallIntVector.class)); | |
328 | } | |
329 | } | |
330 | ||
331 | @Test | |
332 | public void testTimeMicroVector() { | |
333 | try (final TimeMicroVector vector = new TimeMicroVector("v", allocator)) { | |
334 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
335 | assertTrue(roundtrip(vector, TimeMicroVector.class)); | |
336 | } | |
337 | } | |
338 | ||
339 | @Test | |
340 | public void testTimeMilliVector() { | |
341 | try (final TimeMilliVector vector = new TimeMilliVector("v", allocator)) { | |
342 | setVector(vector, 0, 1, 2, 3, null); | |
343 | assertTrue(roundtrip(vector, TimeMilliVector.class)); | |
344 | } | |
345 | } | |
346 | ||
347 | @Test | |
348 | public void testTimeNanoVector() { | |
349 | try (final TimeNanoVector vector = new TimeNanoVector("v", allocator)) { | |
350 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
351 | assertTrue(roundtrip(vector, TimeNanoVector.class)); | |
352 | } | |
353 | } | |
354 | ||
355 | @Test | |
356 | public void testTimeSecVector() { | |
357 | try (final TimeSecVector vector = new TimeSecVector("v", allocator)) { | |
358 | setVector(vector, 0, 1, 2, 3, null); | |
359 | assertTrue(roundtrip(vector, TimeSecVector.class)); | |
360 | } | |
361 | } | |
362 | ||
363 | @Test | |
364 | public void testTimeStampMicroTZVector() { | |
365 | try (final TimeStampMicroTZVector vector = new TimeStampMicroTZVector("v", allocator, "UTC")) { | |
366 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
367 | assertTrue(roundtrip(vector, TimeStampMicroTZVector.class)); | |
368 | } | |
369 | } | |
370 | ||
371 | @Test | |
372 | public void testTimeStampMicroVector() { | |
373 | try (final TimeStampMicroVector vector = new TimeStampMicroVector("v", allocator)) { | |
374 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
375 | assertTrue(roundtrip(vector, TimeStampMicroVector.class)); | |
376 | } | |
377 | } | |
378 | ||
379 | @Test | |
380 | public void testTimeStampMilliTZVector() { | |
381 | try (final TimeStampMilliTZVector vector = new TimeStampMilliTZVector("v", allocator, "UTC")) { | |
382 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
383 | assertTrue(roundtrip(vector, TimeStampMilliTZVector.class)); | |
384 | } | |
385 | } | |
386 | ||
387 | @Test | |
388 | public void testTimeStampMilliVector() { | |
389 | try (final TimeStampMilliVector vector = new TimeStampMilliVector("v", allocator)) { | |
390 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
391 | assertTrue(roundtrip(vector, TimeStampMilliVector.class)); | |
392 | } | |
393 | } | |
394 | ||
395 | @Test | |
396 | public void testTimeTimeStampNanoTZVector() { | |
397 | try (final TimeStampNanoTZVector vector = new TimeStampNanoTZVector("v", allocator, "UTC")) { | |
398 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
399 | assertTrue(roundtrip(vector, TimeStampNanoTZVector.class)); | |
400 | } | |
401 | } | |
402 | ||
403 | @Test | |
404 | public void testTimeStampNanoVector() { | |
405 | try (final TimeStampNanoVector vector = new TimeStampNanoVector("v", allocator)) { | |
406 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
407 | assertTrue(roundtrip(vector, TimeStampNanoVector.class)); | |
408 | } | |
409 | } | |
410 | ||
411 | @Test | |
412 | public void testTimeStampSecTZVector() { | |
413 | try (final TimeStampSecTZVector vector = new TimeStampSecTZVector("v", allocator, "UTC")) { | |
414 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
415 | assertTrue(roundtrip(vector, TimeStampSecTZVector.class)); | |
416 | } | |
417 | } | |
418 | ||
419 | @Test | |
420 | public void testTimeStampSecVector() { | |
421 | try (final TimeStampSecVector vector = new TimeStampSecVector("v", allocator)) { | |
422 | setVector(vector, 0L, 1L, 2L, 3L, null); | |
423 | assertTrue(roundtrip(vector, TimeStampSecVector.class)); | |
424 | } | |
425 | } | |
426 | ||
427 | @Test | |
428 | public void testTinyIntVector() { | |
429 | try (final TinyIntVector vector = new TinyIntVector("v", allocator)) { | |
430 | setVector(vector, (byte) 0, (byte) 1, null); | |
431 | assertTrue(roundtrip(vector, TinyIntVector.class)); | |
432 | } | |
433 | } | |
434 | ||
435 | @Test | |
436 | public void testUInt1Vector() { | |
437 | try (final UInt1Vector vector = new UInt1Vector("v", allocator)) { | |
438 | setVector(vector, (byte) 0, (byte) 1, null); | |
439 | assertTrue(roundtrip(vector, UInt1Vector.class)); | |
440 | } | |
441 | } | |
442 | ||
443 | @Test | |
444 | public void testUInt2Vector() { | |
445 | try (final UInt2Vector vector = new UInt2Vector("v", allocator)) { | |
446 | setVector(vector, '0', '1', null); | |
447 | assertTrue(roundtrip(vector, UInt2Vector.class)); | |
448 | } | |
449 | } | |
450 | ||
451 | @Test | |
452 | public void testUInt4Vector() { | |
453 | try (final UInt4Vector vector = new UInt4Vector("v", allocator)) { | |
454 | setVector(vector, 0, 1, null); | |
455 | assertTrue(roundtrip(vector, UInt4Vector.class)); | |
456 | } | |
457 | } | |
458 | ||
459 | @Test | |
460 | public void testUInt8Vector() { | |
461 | try (final UInt8Vector vector = new UInt8Vector("v", allocator)) { | |
462 | setVector(vector, 0L, 1L, null); | |
463 | assertTrue(roundtrip(vector, UInt8Vector.class)); | |
464 | } | |
465 | } | |
466 | ||
467 | @Test | |
468 | public void testVarBinaryVector() { | |
469 | try (final VarBinaryVector vector = new VarBinaryVector("v", allocator)) { | |
470 | setVector(vector, "abc".getBytes(), "def".getBytes(), null); | |
471 | assertTrue(roundtrip(vector, VarBinaryVector.class)); | |
472 | } | |
473 | } | |
474 | ||
475 | @Test | |
476 | public void testVarCharVector() { | |
477 | try (final VarCharVector vector = new VarCharVector("v", allocator)) { | |
478 | setVector(vector, "abc", "def", null); | |
479 | assertTrue(roundtrip(vector, VarCharVector.class)); | |
480 | } | |
481 | } | |
482 | ||
483 | @Test | |
484 | public void testLargeVarBinaryVector() { | |
485 | try (final LargeVarBinaryVector vector = new LargeVarBinaryVector("", allocator)) { | |
486 | vector.allocateNew(5, 1); | |
487 | ||
488 | NullableLargeVarBinaryHolder nullHolder = new NullableLargeVarBinaryHolder(); | |
489 | nullHolder.isSet = 0; | |
490 | ||
491 | NullableLargeVarBinaryHolder binHolder = new NullableLargeVarBinaryHolder(); | |
492 | binHolder.isSet = 1; | |
493 | ||
494 | String str = "hello world"; | |
495 | try (ArrowBuf buf = allocator.buffer(16)) { | |
496 | buf.setBytes(0, str.getBytes()); | |
497 | binHolder.start = 0; | |
498 | binHolder.end = str.length(); | |
499 | binHolder.buffer = buf; | |
500 | vector.setSafe(0, binHolder); | |
501 | vector.setSafe(1, nullHolder); | |
502 | ||
503 | assertTrue(roundtrip(vector, LargeVarBinaryVector.class)); | |
504 | } | |
505 | } | |
506 | } | |
507 | ||
508 | @Test | |
509 | public void testLargeVarCharVector() { | |
510 | try (final LargeVarCharVector vector = new LargeVarCharVector("v", allocator)) { | |
511 | setVector(vector, "abc", "def", null); | |
512 | assertTrue(roundtrip(vector, LargeVarCharVector.class)); | |
513 | } | |
514 | } | |
515 | ||
516 | @Test | |
517 | public void testListVector() { | |
518 | try (final ListVector vector = ListVector.empty("v", allocator)) { | |
519 | setVector(vector, Arrays.stream(new int[] { 1, 2 }).boxed().collect(Collectors.toList()), | |
520 | Arrays.stream(new int[] { 3, 4 }).boxed().collect(Collectors.toList()), new ArrayList<Integer>()); | |
521 | assertTrue(roundtrip(vector, ListVector.class)); | |
522 | } | |
523 | } | |
524 | ||
525 | @Test | |
526 | public void testLargeListVector() { | |
527 | try (final LargeListVector vector = LargeListVector.empty("v", allocator)) { | |
528 | setVector(vector, Arrays.stream(new int[] { 1, 2 }).boxed().collect(Collectors.toList()), | |
529 | Arrays.stream(new int[] { 3, 4 }).boxed().collect(Collectors.toList()), new ArrayList<Integer>()); | |
530 | assertTrue(roundtrip(vector, LargeListVector.class)); | |
531 | } | |
532 | } | |
533 | ||
534 | @Test | |
535 | public void testFixedSizeListVector() { | |
536 | try (final FixedSizeListVector vector = FixedSizeListVector.empty("v", 2, allocator)) { | |
537 | setVector(vector, Arrays.stream(new int[] { 1, 2 }).boxed().collect(Collectors.toList()), | |
538 | Arrays.stream(new int[] { 3, 4 }).boxed().collect(Collectors.toList())); | |
539 | assertTrue(roundtrip(vector, FixedSizeListVector.class)); | |
540 | } | |
541 | } | |
542 | ||
543 | @Test | |
544 | public void testMapVector() { | |
545 | int count = 5; | |
546 | try (final MapVector vector = MapVector.empty("v", allocator, false)) { | |
547 | vector.allocateNew(); | |
548 | UnionMapWriter mapWriter = vector.getWriter(); | |
549 | for (int i = 0; i < count; i++) { | |
550 | mapWriter.startMap(); | |
551 | for (int j = 0; j < i + 1; j++) { | |
552 | mapWriter.startEntry(); | |
553 | mapWriter.key().bigInt().writeBigInt(j); | |
554 | mapWriter.value().integer().writeInt(j); | |
555 | mapWriter.endEntry(); | |
556 | } | |
557 | mapWriter.endMap(); | |
558 | } | |
559 | mapWriter.setValueCount(count); | |
560 | ||
561 | assertTrue(roundtrip(vector, MapVector.class)); | |
562 | } | |
563 | } | |
564 | ||
565 | @Test | |
566 | public void testUnionVector() { | |
567 | final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); | |
568 | uInt4Holder.value = 100; | |
569 | uInt4Holder.isSet = 1; | |
570 | ||
571 | try (UnionVector vector = UnionVector.empty("v", allocator)) { | |
572 | vector.allocateNew(); | |
573 | ||
574 | // write some data | |
575 | vector.setType(0, MinorType.UINT4); | |
576 | vector.setSafe(0, uInt4Holder); | |
577 | vector.setType(2, MinorType.UINT4); | |
578 | vector.setSafe(2, uInt4Holder); | |
579 | vector.setValueCount(4); | |
580 | ||
581 | assertTrue(roundtrip(vector, UnionVector.class)); | |
582 | } | |
583 | } | |
584 | ||
585 | @Test | |
586 | public void testStructVector() { | |
587 | try (final StructVector vector = StructVector.empty("v", allocator)) { | |
588 | Map<String, List<Integer>> data = new HashMap<>(); | |
589 | data.put("col_1", Arrays.stream(new int[] { 1, 2 }).boxed().collect(Collectors.toList())); | |
590 | data.put("col_2", Arrays.stream(new int[] { 3, 4 }).boxed().collect(Collectors.toList())); | |
591 | setVector(vector, data); | |
592 | assertTrue(roundtrip(vector, StructVector.class)); | |
593 | } | |
594 | } | |
595 | ||
596 | @Test | |
597 | public void testExtensionTypeVector() { | |
598 | ExtensionTypeRegistry.register(new UuidType()); | |
599 | final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new UuidType()))); | |
600 | try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { | |
601 | // Fill with data | |
602 | UUID u1 = UUID.randomUUID(); | |
603 | UUID u2 = UUID.randomUUID(); | |
604 | UuidVector vector = (UuidVector) root.getVector("a"); | |
605 | vector.setValueCount(2); | |
606 | vector.set(0, u1); | |
607 | vector.set(1, u2); | |
608 | root.setRowCount(2); | |
609 | ||
610 | // Roundtrip (export + import) | |
611 | VectorSchemaRoot importedRoot = vectorSchemaRootRoundtrip(root); | |
612 | ||
613 | // Verify correctness | |
614 | assertEquals(root.getSchema(), importedRoot.getSchema()); | |
615 | ||
616 | final Field field = importedRoot.getSchema().getFields().get(0); | |
617 | final UuidType expectedType = new UuidType(); | |
618 | assertEquals(field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME), expectedType.extensionName()); | |
619 | assertEquals(field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA), expectedType.serialize()); | |
620 | ||
621 | final UuidVector deserialized = (UuidVector) importedRoot.getFieldVectors().get(0); | |
622 | assertEquals(vector.getValueCount(), deserialized.getValueCount()); | |
623 | for (int i = 0; i < vector.getValueCount(); i++) { | |
624 | assertEquals(vector.isNull(i), deserialized.isNull(i)); | |
625 | if (!vector.isNull(i)) { | |
626 | assertEquals(vector.getObject(i), deserialized.getObject(i)); | |
627 | } | |
628 | } | |
629 | ||
630 | importedRoot.close(); | |
631 | } | |
632 | } | |
633 | ||
634 | @Test | |
635 | public void testVectorSchemaRoot() { | |
636 | VectorSchemaRoot imported; | |
637 | ||
638 | // Consumer allocates empty structures | |
639 | try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); | |
640 | ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { | |
641 | try (VectorSchemaRoot vsr = createTestVSR()) { | |
642 | // Producer creates structures from existing memory pointers | |
643 | try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress()); | |
644 | ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) { | |
645 | // Producer exports vector into the C Data Interface structures | |
646 | Data.exportVectorSchemaRoot(allocator, vsr, null, arrowArray, arrowSchema); | |
647 | } | |
648 | } | |
649 | // Consumer imports vector | |
650 | imported = Data.importVectorSchemaRoot(allocator, consumerArrowArray, consumerArrowSchema, null); | |
651 | } | |
652 | ||
653 | // Ensure that imported VectorSchemaRoot is valid even after C Data Interface | |
654 | // structures are closed | |
655 | try (VectorSchemaRoot original = createTestVSR()) { | |
656 | assertTrue(imported.equals(original)); | |
657 | } | |
658 | imported.close(); | |
659 | } | |
660 | ||
661 | @Test | |
662 | public void testSchema() { | |
663 | Field decimalField = new Field("inner1", FieldType.nullable(new ArrowType.Decimal(19, 4, 128)), null); | |
664 | Field strField = new Field("inner2", FieldType.nullable(new ArrowType.Utf8()), null); | |
665 | Field itemField = new Field("col1", FieldType.nullable(new ArrowType.Struct()), | |
666 | Arrays.asList(decimalField, strField)); | |
667 | Field intField = new Field("col2", FieldType.nullable(new ArrowType.Int(32, true)), null); | |
668 | Schema schema = new Schema(Arrays.asList(itemField, intField)); | |
669 | // Consumer allocates empty ArrowSchema | |
670 | try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) { | |
671 | // Producer fills the schema with data | |
672 | try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress())) { | |
673 | Data.exportSchema(allocator, schema, null, arrowSchema); | |
674 | } | |
675 | // Consumer imports schema | |
676 | Schema importedSchema = Data.importSchema(allocator, consumerArrowSchema, null); | |
677 | assertEquals(schema.toJson(), importedSchema.toJson()); | |
678 | } | |
679 | } | |
680 | ||
681 | @Test | |
682 | public void testImportReleasedArray() { | |
683 | // Consumer allocates empty structures | |
684 | try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); | |
685 | ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { | |
686 | // Producer creates structures from existing memory pointers | |
687 | try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress()); | |
688 | ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) { | |
689 | // Producer exports vector into the C Data Interface structures | |
690 | try (final NullVector vector = new NullVector()) { | |
691 | Data.exportVector(allocator, vector, null, arrowArray, arrowSchema); | |
692 | } | |
693 | } | |
694 | ||
695 | // Release array structure | |
696 | consumerArrowArray.markReleased(); | |
697 | ||
698 | // Consumer tried to imports vector but fails | |
699 | Exception e = assertThrows(IllegalStateException.class, () -> { | |
700 | Data.importVector(allocator, consumerArrowArray, consumerArrowSchema, null); | |
701 | }); | |
702 | ||
703 | assertEquals("Cannot import released ArrowArray", e.getMessage()); | |
704 | } | |
705 | } | |
706 | ||
707 | private VectorSchemaRoot createTestVSR() { | |
708 | BitVector bitVector = new BitVector("boolean", allocator); | |
709 | ||
710 | Map<String, String> metadata = new HashMap<>(); | |
711 | metadata.put("key", "value"); | |
712 | FieldType fieldType = new FieldType(true, ArrowType.Utf8.INSTANCE, null, metadata); | |
713 | VarCharVector varCharVector = new VarCharVector("varchar", fieldType, allocator); | |
714 | ||
715 | bitVector.allocateNew(); | |
716 | varCharVector.allocateNew(); | |
717 | for (int i = 0; i < 10; i++) { | |
718 | bitVector.setSafe(i, i % 2 == 0 ? 0 : 1); | |
719 | varCharVector.setSafe(i, ("test" + i).getBytes(StandardCharsets.UTF_8)); | |
720 | } | |
721 | bitVector.setValueCount(10); | |
722 | varCharVector.setValueCount(10); | |
723 | ||
724 | List<Field> fields = Arrays.asList(bitVector.getField(), varCharVector.getField()); | |
725 | List<FieldVector> vectors = Arrays.asList(bitVector, varCharVector); | |
726 | ||
727 | return new VectorSchemaRoot(fields, vectors); | |
728 | } | |
729 | ||
730 | static class UuidType extends ExtensionType { | |
731 | ||
732 | @Override | |
733 | public ArrowType storageType() { | |
734 | return new ArrowType.FixedSizeBinary(16); | |
735 | } | |
736 | ||
737 | @Override | |
738 | public String extensionName() { | |
739 | return "uuid"; | |
740 | } | |
741 | ||
742 | @Override | |
743 | public boolean extensionEquals(ExtensionType other) { | |
744 | return other instanceof UuidType; | |
745 | } | |
746 | ||
747 | @Override | |
748 | public ArrowType deserialize(ArrowType storageType, String serializedData) { | |
749 | if (!storageType.equals(storageType())) { | |
750 | throw new UnsupportedOperationException("Cannot construct UuidType from underlying type " + storageType); | |
751 | } | |
752 | return new UuidType(); | |
753 | } | |
754 | ||
755 | @Override | |
756 | public String serialize() { | |
757 | return ""; | |
758 | } | |
759 | ||
760 | @Override | |
761 | public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { | |
762 | return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16)); | |
763 | } | |
764 | } | |
765 | ||
766 | static class UuidVector extends ExtensionTypeVector<FixedSizeBinaryVector> { | |
767 | ||
768 | public UuidVector(String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) { | |
769 | super(name, allocator, underlyingVector); | |
770 | } | |
771 | ||
772 | @Override | |
773 | public UUID getObject(int index) { | |
774 | final ByteBuffer bb = ByteBuffer.wrap(getUnderlyingVector().getObject(index)); | |
775 | return new UUID(bb.getLong(), bb.getLong()); | |
776 | } | |
777 | ||
778 | @Override | |
779 | public int hashCode(int index) { | |
780 | return hashCode(index, null); | |
781 | } | |
782 | ||
783 | @Override | |
784 | public int hashCode(int index, ArrowBufHasher hasher) { | |
785 | return getUnderlyingVector().hashCode(index, hasher); | |
786 | } | |
787 | ||
788 | public void set(int index, UUID uuid) { | |
789 | ByteBuffer bb = ByteBuffer.allocate(16); | |
790 | bb.putLong(uuid.getMostSignificantBits()); | |
791 | bb.putLong(uuid.getLeastSignificantBits()); | |
792 | getUnderlyingVector().set(index, bb.array()); | |
793 | } | |
794 | } | |
795 | } |