ceph/src/arrow/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java

   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.vector.complex.writer;
  19
  20 import static org.junit.Assert.*;
  21
  22 import java.math.BigDecimal;
  23 import java.time.LocalDateTime;
  24 import java.util.HashSet;
  25 import java.util.List;
  26 import java.util.Set;
  27
  28 import org.apache.arrow.memory.ArrowBuf;
  29 import org.apache.arrow.memory.BufferAllocator;
  30 import org.apache.arrow.memory.RootAllocator;
  31 import org.apache.arrow.util.AutoCloseables;
  32 import org.apache.arrow.vector.BigIntVector;
  33 import org.apache.arrow.vector.DecimalVector;
  34 import org.apache.arrow.vector.Float4Vector;
  35 import org.apache.arrow.vector.Float8Vector;
  36 import org.apache.arrow.vector.IntVector;
  37 import org.apache.arrow.vector.SchemaChangeCallBack;
  38 import org.apache.arrow.vector.complex.ListVector;
  39 import org.apache.arrow.vector.complex.MapVector;
  40 import org.apache.arrow.vector.complex.NonNullableStructVector;
  41 import org.apache.arrow.vector.complex.StructVector;
  42 import org.apache.arrow.vector.complex.UnionVector;
  43 import org.apache.arrow.vector.complex.impl.ComplexWriterImpl;
  44 import org.apache.arrow.vector.complex.impl.SingleStructReaderImpl;
  45 import org.apache.arrow.vector.complex.impl.SingleStructWriter;
  46 import org.apache.arrow.vector.complex.impl.UnionListReader;
  47 import org.apache.arrow.vector.complex.impl.UnionListWriter;
  48 import org.apache.arrow.vector.complex.impl.UnionMapReader;
  49 import org.apache.arrow.vector.complex.impl.UnionReader;
  50 import org.apache.arrow.vector.complex.impl.UnionWriter;
  51 import org.apache.arrow.vector.complex.reader.BaseReader.StructReader;
  52 import org.apache.arrow.vector.complex.reader.BigIntReader;
  53 import org.apache.arrow.vector.complex.reader.FieldReader;
  54 import org.apache.arrow.vector.complex.reader.Float4Reader;
  55 import org.apache.arrow.vector.complex.reader.Float8Reader;
  56 import org.apache.arrow.vector.complex.reader.IntReader;
  57 import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter;
  58 import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter;
  59 import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;
  60 import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
  61 import org.apache.arrow.vector.holders.DecimalHolder;
  62 import org.apache.arrow.vector.holders.IntHolder;
  63 import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder;
  64 import org.apache.arrow.vector.types.pojo.ArrowType;
  65 import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID;
  66 import org.apache.arrow.vector.types.pojo.ArrowType.Int;
  67 import org.apache.arrow.vector.types.pojo.ArrowType.Struct;
  68 import org.apache.arrow.vector.types.pojo.ArrowType.Timestamp;
  69 import org.apache.arrow.vector.types.pojo.ArrowType.Union;
  70 import org.apache.arrow.vector.types.pojo.ArrowType.Utf8;
  71 import org.apache.arrow.vector.types.pojo.Field;
  72 import org.apache.arrow.vector.types.pojo.FieldType;
  73 import org.apache.arrow.vector.util.CallBack;
  74 import org.apache.arrow.vector.util.DecimalUtility;
  75 import org.apache.arrow.vector.util.JsonStringArrayList;
  76 import org.apache.arrow.vector.util.JsonStringHashMap;
  77 import org.apache.arrow.vector.util.Text;
  78 import org.apache.arrow.vector.util.TransferPair;
  79 import org.junit.After;
  80 import org.junit.Assert;
  81 import org.junit.Before;
  82 import org.junit.Test;
  83
  84 public class TestComplexWriter {
  85
  86   private BufferAllocator allocator;
  87
  88   private static final int COUNT = 100;
  89
  90   @Before
  91   public void init() {
  92     allocator = new RootAllocator(Integer.MAX_VALUE);
  93   }
  94
  95   @After
  96   public void terminate() throws Exception {
  97     allocator.close();
  98   }
  99
 100   @Test
 101   public void simpleNestedTypes() {
 102     NonNullableStructVector parent = populateStructVector(null);
 103     StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
 104     for (int i = 0; i < COUNT; i++) {
 105       rootReader.setPosition(i);
 106       Assert.assertEquals(i, rootReader.reader("int").readInteger().intValue());
 107       Assert.assertEquals(i, rootReader.reader("bigInt").readLong().longValue());
 108     }
 109
 110     parent.close();
 111   }
 112
 113   @Test
 114   public void transferPairSchemaChange() {
 115     SchemaChangeCallBack callBack1 = new SchemaChangeCallBack();
 116     SchemaChangeCallBack callBack2 = new SchemaChangeCallBack();
 117     try (NonNullableStructVector parent = populateStructVector(callBack1)) {
 118       TransferPair tp = parent.getTransferPair("newVector", allocator, callBack2);
 119
 120       ComplexWriter writer = new ComplexWriterImpl("newWriter", parent);
 121       StructWriter rootWriter = writer.rootAsStruct();
 122       IntWriter intWriter = rootWriter.integer("newInt");
 123       intWriter.writeInt(1);
 124       writer.setValueCount(1);
 125
 126       assertTrue(callBack1.getSchemaChangedAndReset());
 127       // The second vector should not have registered a schema change
 128       assertFalse(callBack1.getSchemaChangedAndReset());
 129     }
 130   }
 131
 132   private NonNullableStructVector populateStructVector(CallBack callBack) {
 133     NonNullableStructVector parent =
 134         new NonNullableStructVector("parent", allocator, new FieldType(false, Struct.INSTANCE, null, null), callBack);
 135     ComplexWriter writer = new ComplexWriterImpl("root", parent);
 136     StructWriter rootWriter = writer.rootAsStruct();
 137     IntWriter intWriter = rootWriter.integer("int");
 138     BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt");
 139     for (int i = 0; i < COUNT; i++) {
 140       rootWriter.start();
 141       intWriter.writeInt(i);
 142       bigIntWriter.writeBigInt(i);
 143       rootWriter.end();
 144     }
 145     writer.setValueCount(COUNT);
 146     return parent;
 147   }
 148
 149   @Test
 150   public void nullableStruct() {
 151     try (NonNullableStructVector structVector = NonNullableStructVector.empty("parent", allocator)) {
 152       ComplexWriter writer = new ComplexWriterImpl("root", structVector);
 153       StructWriter rootWriter = writer.rootAsStruct();
 154       for (int i = 0; i < COUNT; i++) {
 155         rootWriter.start();
 156         if (i % 2 == 0) {
 157           StructWriter structWriter = rootWriter.struct("struct");
 158           structWriter.setPosition(i);
 159           structWriter.start();
 160           structWriter.bigInt("nested").writeBigInt(i);
 161           structWriter.end();
 162         }
 163         rootWriter.end();
 164       }
 165       writer.setValueCount(COUNT);
 166       checkNullableStruct(structVector);
 167     }
 168   }
 169
 170   /**
 171    * This test is similar to {@link #nullableStruct()} ()} but we get the inner struct writer once at the beginning.
 172    */
 173   @Test
 174   public void nullableStruct2() {
 175     try (NonNullableStructVector structVector = NonNullableStructVector.empty("parent", allocator)) {
 176       ComplexWriter writer = new ComplexWriterImpl("root", structVector);
 177       StructWriter rootWriter = writer.rootAsStruct();
 178       StructWriter structWriter = rootWriter.struct("struct");
 179
 180       for (int i = 0; i < COUNT; i++) {
 181         rootWriter.start();
 182         if (i % 2 == 0) {
 183           structWriter.setPosition(i);
 184           structWriter.start();
 185           structWriter.bigInt("nested").writeBigInt(i);
 186           structWriter.end();
 187         }
 188         rootWriter.end();
 189       }
 190       writer.setValueCount(COUNT);
 191       checkNullableStruct(structVector);
 192     }
 193   }
 194
 195   private void checkNullableStruct(NonNullableStructVector structVector) {
 196     StructReader rootReader = new SingleStructReaderImpl(structVector).reader("root");
 197     for (int i = 0; i < COUNT; i++) {
 198       rootReader.setPosition(i);
 199       assertTrue("index is set: " + i, rootReader.isSet());
 200       FieldReader struct = rootReader.reader("struct");
 201       if (i % 2 == 0) {
 202         assertTrue("index is set: " + i, struct.isSet());
 203         assertNotNull("index is set: " + i, struct.readObject());
 204         assertEquals(i, struct.reader("nested").readLong().longValue());
 205       } else {
 206         assertFalse("index is not set: " + i, struct.isSet());
 207         assertNull("index is not set: " + i, struct.readObject());
 208       }
 209     }
 210   }
 211
 212   @Test
 213   public void testList() {
 214     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
 215       ComplexWriter writer = new ComplexWriterImpl("root", parent);
 216       StructWriter rootWriter = writer.rootAsStruct();
 217
 218       rootWriter.start();
 219       rootWriter.bigInt("int").writeBigInt(0);
 220       rootWriter.list("list").startList();
 221       rootWriter.list("list").bigInt().writeBigInt(0);
 222       rootWriter.list("list").endList();
 223       rootWriter.end();
 224
 225       rootWriter.start();
 226       rootWriter.bigInt("int").writeBigInt(1);
 227       rootWriter.end();
 228
 229       writer.setValueCount(2);
 230
 231       StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
 232
 233       rootReader.setPosition(0);
 234       assertTrue("row 0 list is not set", rootReader.reader("list").isSet());
 235       assertEquals(Long.valueOf(0), rootReader.reader("list").reader().readLong());
 236
 237       rootReader.setPosition(1);
 238       assertFalse("row 1 list is set", rootReader.reader("list").isSet());
 239     }
 240   }
 241
 242   @Test
 243   public void listScalarType() {
 244     try (ListVector listVector = ListVector.empty("list", allocator)) {
 245       listVector.allocateNew();
 246       UnionListWriter listWriter = new UnionListWriter(listVector);
 247       for (int i = 0; i < COUNT; i++) {
 248         listWriter.startList();
 249         for (int j = 0; j < i % 7; j++) {
 250           if (j % 2 == 0) {
 251             listWriter.writeInt(j);
 252           } else {
 253             IntHolder holder = new IntHolder();
 254             holder.value = j;
 255             listWriter.write(holder);
 256           }
 257         }
 258         listWriter.endList();
 259       }
 260       listWriter.setValueCount(COUNT);
 261       UnionListReader listReader = new UnionListReader(listVector);
 262       for (int i = 0; i < COUNT; i++) {
 263         listReader.setPosition(i);
 264         for (int j = 0; j < i % 7; j++) {
 265           listReader.next();
 266           assertEquals(j, listReader.reader().readInteger().intValue());
 267         }
 268       }
 269     }
 270   }
 271
 272   @Test
 273   public void testListScalarNull() {
 274     /* Write to a integer list vector
 275      * each list of size 8 and having it's data values alternating between null and a non-null.
 276      * Read and verify
 277      */
 278     try (ListVector listVector = ListVector.empty("list", allocator)) {
 279       listVector.allocateNew();
 280       UnionListWriter listWriter = new UnionListWriter(listVector);
 281       for (int i = 0; i < COUNT; i++) {
 282         listWriter.startList();
 283         for (int j = 0; j < i % 7; j++) {
 284           if (j % 2 == 0) {
 285             listWriter.writeNull();
 286           } else {
 287             IntHolder holder = new IntHolder();
 288             holder.value = j;
 289             listWriter.write(holder);
 290           }
 291         }
 292         listWriter.endList();
 293       }
 294       listWriter.setValueCount(COUNT);
 295       UnionListReader listReader = new UnionListReader(listVector);
 296       for (int i = 0; i < COUNT; i++) {
 297         listReader.setPosition(i);
 298         for (int j = 0; j < i % 7; j++) {
 299           listReader.next();
 300           if (j % 2 == 0) {
 301             assertFalse("index is set: " + j, listReader.reader().isSet());
 302           } else {
 303             assertTrue("index is not set: " + j, listReader.reader().isSet());
 304             assertEquals(j, listReader.reader().readInteger().intValue());
 305           }
 306         }
 307       }
 308     }
 309   }
 310
 311   @Test
 312   public void listDecimalType() {
 313     try (ListVector listVector = ListVector.empty("list", allocator)) {
 314       listVector.allocateNew();
 315       UnionListWriter listWriter = new UnionListWriter(listVector);
 316       DecimalHolder holder = new DecimalHolder();
 317       holder.buffer = allocator.buffer(DecimalVector.TYPE_WIDTH);
 318       ArrowType arrowType = new ArrowType.Decimal(10, 0, 128);
 319       for (int i = 0; i < COUNT; i++) {
 320         listWriter.startList();
 321         for (int j = 0; j < i % 7; j++) {
 322           if (j % 4 == 0) {
 323             listWriter.writeDecimal(new BigDecimal(j));
 324           } else if (j % 4 == 1) {
 325             DecimalUtility.writeBigDecimalToArrowBuf(new BigDecimal(j), holder.buffer, 0, DecimalVector.TYPE_WIDTH);
 326             holder.start = 0;
 327             holder.scale = 0;
 328             holder.precision = 10;
 329             listWriter.write(holder);
 330           } else if (j % 4 == 2) {
 331             DecimalUtility.writeBigDecimalToArrowBuf(new BigDecimal(j), holder.buffer, 0, DecimalVector.TYPE_WIDTH);
 332             listWriter.writeDecimal(0, holder.buffer, arrowType);
 333           } else {
 334             byte[] value = BigDecimal.valueOf(j).unscaledValue().toByteArray();
 335             listWriter.writeBigEndianBytesToDecimal(value, arrowType);
 336           }
 337         }
 338         listWriter.endList();
 339       }
 340       listWriter.setValueCount(COUNT);
 341       UnionListReader listReader = new UnionListReader(listVector);
 342       for (int i = 0; i < COUNT; i++) {
 343         listReader.setPosition(i);
 344         for (int j = 0; j < i % 7; j++) {
 345           listReader.next();
 346           Object expected = new BigDecimal(j);
 347           Object actual = listReader.reader().readBigDecimal();
 348           assertEquals(expected, actual);
 349         }
 350       }
 351       holder.buffer.close();
 352     }
 353   }
 354
 355   @Test
 356   public void listScalarTypeNullable() {
 357     try (ListVector listVector = ListVector.empty("list", allocator)) {
 358       listVector.allocateNew();
 359       UnionListWriter listWriter = new UnionListWriter(listVector);
 360       for (int i = 0; i < COUNT; i++) {
 361         if (i % 2 == 0) {
 362           listWriter.setPosition(i);
 363           listWriter.startList();
 364           for (int j = 0; j < i % 7; j++) {
 365             listWriter.writeInt(j);
 366           }
 367           listWriter.endList();
 368         }
 369       }
 370       listWriter.setValueCount(COUNT);
 371       UnionListReader listReader = new UnionListReader(listVector);
 372       for (int i = 0; i < COUNT; i++) {
 373         listReader.setPosition(i);
 374         if (i % 2 == 0) {
 375           assertTrue("index is set: " + i, listReader.isSet());
 376           assertEquals("correct length at: " + i, i % 7, ((List<?>) listReader.readObject()).size());
 377         } else {
 378           assertFalse("index is not set: " + i, listReader.isSet());
 379           assertNull("index is not set: " + i, listReader.readObject());
 380         }
 381       }
 382     }
 383   }
 384
 385   @Test
 386   public void listStructType() {
 387     try (ListVector listVector = ListVector.empty("list", allocator)) {
 388       listVector.allocateNew();
 389       UnionListWriter listWriter = new UnionListWriter(listVector);
 390       StructWriter structWriter = listWriter.struct();
 391       for (int i = 0; i < COUNT; i++) {
 392         listWriter.startList();
 393         for (int j = 0; j < i % 7; j++) {
 394           structWriter.start();
 395           structWriter.integer("int").writeInt(j);
 396           structWriter.bigInt("bigInt").writeBigInt(j);
 397           structWriter.end();
 398         }
 399         listWriter.endList();
 400       }
 401       listWriter.setValueCount(COUNT);
 402       UnionListReader listReader = new UnionListReader(listVector);
 403       for (int i = 0; i < COUNT; i++) {
 404         listReader.setPosition(i);
 405         for (int j = 0; j < i % 7; j++) {
 406           listReader.next();
 407           Assert.assertEquals("record: " + i, j, listReader.reader().reader("int").readInteger().intValue());
 408           Assert.assertEquals(j, listReader.reader().reader("bigInt").readLong().longValue());
 409         }
 410       }
 411     }
 412   }
 413
 414   @Test
 415   public void listListType() {
 416     try (ListVector listVector = ListVector.empty("list", allocator)) {
 417       listVector.allocateNew();
 418       UnionListWriter listWriter = new UnionListWriter(listVector);
 419       for (int i = 0; i < COUNT; i++) {
 420         listWriter.startList();
 421         for (int j = 0; j < i % 7; j++) {
 422           ListWriter innerListWriter = listWriter.list();
 423           innerListWriter.startList();
 424           for (int k = 0; k < i % 13; k++) {
 425             innerListWriter.integer().writeInt(k);
 426           }
 427           innerListWriter.endList();
 428         }
 429         listWriter.endList();
 430       }
 431       listWriter.setValueCount(COUNT);
 432       checkListOfLists(listVector);
 433     }
 434   }
 435
 436   /**
 437    * This test is similar to {@link #listListType()} but we get the inner list writer once at the beginning.
 438    */
 439   @Test
 440   public void listListType2() {
 441     try (ListVector listVector = ListVector.empty("list", allocator)) {
 442       listVector.allocateNew();
 443       UnionListWriter listWriter = new UnionListWriter(listVector);
 444       ListWriter innerListWriter = listWriter.list();
 445
 446       for (int i = 0; i < COUNT; i++) {
 447         listWriter.startList();
 448         for (int j = 0; j < i % 7; j++) {
 449           innerListWriter.startList();
 450           for (int k = 0; k < i % 13; k++) {
 451             innerListWriter.integer().writeInt(k);
 452           }
 453           innerListWriter.endList();
 454         }
 455         listWriter.endList();
 456       }
 457       listWriter.setValueCount(COUNT);
 458       checkListOfLists(listVector);
 459     }
 460   }
 461
 462   private void checkListOfLists(final ListVector listVector) {
 463     UnionListReader listReader = new UnionListReader(listVector);
 464     for (int i = 0; i < COUNT; i++) {
 465       listReader.setPosition(i);
 466       for (int j = 0; j < i % 7; j++) {
 467         listReader.next();
 468         FieldReader innerListReader = listReader.reader();
 469         for (int k = 0; k < i % 13; k++) {
 470           innerListReader.next();
 471           Assert.assertEquals("record: " + i, k, innerListReader.reader().readInteger().intValue());
 472         }
 473       }
 474     }
 475   }
 476
 477   @Test
 478   public void unionListListType() {
 479     try (ListVector listVector = ListVector.empty("list", allocator)) {
 480       listVector.allocateNew();
 481       UnionListWriter listWriter = new UnionListWriter(listVector);
 482       for (int i = 0; i < COUNT; i++) {
 483         listWriter.startList();
 484         for (int j = 0; j < i % 7; j++) {
 485           ListWriter innerListWriter = listWriter.list();
 486           innerListWriter.startList();
 487           for (int k = 0; k < i % 13; k++) {
 488             if (k % 2 == 0) {
 489               innerListWriter.integer().writeInt(k);
 490             } else {
 491               innerListWriter.bigInt().writeBigInt(k);
 492             }
 493           }
 494           innerListWriter.endList();
 495         }
 496         listWriter.endList();
 497       }
 498       listWriter.setValueCount(COUNT);
 499       checkUnionList(listVector);
 500     }
 501   }
 502
 503   /**
 504    * This test is similar to {@link #unionListListType()} but we get the inner list writer once at the beginning.
 505    */
 506   @Test
 507   public void unionListListType2() {
 508     try (ListVector listVector = ListVector.empty("list", allocator)) {
 509       listVector.allocateNew();
 510       UnionListWriter listWriter = new UnionListWriter(listVector);
 511       ListWriter innerListWriter = listWriter.list();
 512
 513       for (int i = 0; i < COUNT; i++) {
 514         listWriter.startList();
 515         for (int j = 0; j < i % 7; j++) {
 516           innerListWriter.startList();
 517           for (int k = 0; k < i % 13; k++) {
 518             if (k % 2 == 0) {
 519               innerListWriter.integer().writeInt(k);
 520             } else {
 521               innerListWriter.bigInt().writeBigInt(k);
 522             }
 523           }
 524           innerListWriter.endList();
 525         }
 526         listWriter.endList();
 527       }
 528       listWriter.setValueCount(COUNT);
 529       checkUnionList(listVector);
 530     }
 531   }
 532
 533   private void checkUnionList(ListVector listVector) {
 534     UnionListReader listReader = new UnionListReader(listVector);
 535     for (int i = 0; i < COUNT; i++) {
 536       listReader.setPosition(i);
 537       for (int j = 0; j < i % 7; j++) {
 538         listReader.next();
 539         FieldReader innerListReader = listReader.reader();
 540         for (int k = 0; k < i % 13; k++) {
 541           innerListReader.next();
 542           if (k % 2 == 0) {
 543             Assert.assertEquals("record: " + i, k, innerListReader.reader().readInteger().intValue());
 544           } else {
 545             Assert.assertEquals("record: " + i, k, innerListReader.reader().readLong().longValue());
 546           }
 547         }
 548       }
 549     }
 550   }
 551
 552   @Test
 553   public void testListMapType() {
 554     try (ListVector listVector = ListVector.empty("list", allocator)) {
 555       listVector.allocateNew();
 556       UnionListWriter listWriter = new UnionListWriter(listVector);
 557       MapWriter innerMapWriter = listWriter.map(true);
 558
 559       for (int i = 0; i < COUNT; i++) {
 560         listWriter.startList();
 561         for (int j = 0; j < i % 7; j++) {
 562           innerMapWriter.startMap();
 563           for (int k = 0; k < i % 13; k++) {
 564             innerMapWriter.startEntry();
 565             innerMapWriter.key().integer().writeInt(k);
 566             if (k % 2 == 0) {
 567               innerMapWriter.value().bigInt().writeBigInt(k);
 568             }
 569             innerMapWriter.endEntry();
 570           }
 571           innerMapWriter.endMap();
 572         }
 573         listWriter.endList();
 574       }
 575       listWriter.setValueCount(COUNT);
 576       checkListMap(listVector);
 577
 578       // Verify that the map vector has keysSorted = true
 579       MapVector mapVector = (MapVector) listVector.getDataVector();
 580       ArrowType arrowType = mapVector.getField().getFieldType().getType();
 581       assertTrue(((ArrowType.Map) arrowType).getKeysSorted());
 582     }
 583   }
 584
 585   private void checkListMap(ListVector listVector) {
 586     UnionListReader listReader = new UnionListReader(listVector);
 587     for (int i = 0; i < COUNT; i++) {
 588       listReader.setPosition(i);
 589       for (int j = 0; j < i % 7; j++) {
 590         listReader.next();
 591         UnionMapReader mapReader = (UnionMapReader) listReader.reader();
 592         for (int k = 0; k < i % 13; k++) {
 593           mapReader.next();
 594           Assert.assertEquals("record key: " + i, k, mapReader.key().readInteger().intValue());
 595           if (k % 2 == 0) {
 596             Assert.assertEquals("record value: " + i, k, mapReader.value().readLong().longValue());
 597           } else {
 598             Assert.assertNull("record value: " + i, mapReader.value().readLong());
 599           }
 600         }
 601       }
 602     }
 603   }
 604
 605   @Test
 606   public void simpleUnion() {
 607     UnionVector vector = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null);
 608     UnionWriter unionWriter = new UnionWriter(vector);
 609     unionWriter.allocate();
 610     for (int i = 0; i < COUNT; i++) {
 611       unionWriter.setPosition(i);
 612       if (i % 2 == 0) {
 613         unionWriter.writeInt(i);
 614       } else {
 615         unionWriter.writeFloat4((float) i);
 616       }
 617     }
 618     vector.setValueCount(COUNT);
 619     UnionReader unionReader = new UnionReader(vector);
 620     for (int i = 0; i < COUNT; i++) {
 621       unionReader.setPosition(i);
 622       if (i % 2 == 0) {
 623         Assert.assertEquals(i, i, unionReader.readInteger());
 624       } else {
 625         Assert.assertEquals((float) i, unionReader.readFloat(), 1e-12);
 626       }
 627     }
 628     vector.close();
 629   }
 630
 631   @Test
 632   public void promotableWriter() {
 633     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
 634
 635       ComplexWriter writer = new ComplexWriterImpl("root", parent);
 636       StructWriter rootWriter = writer.rootAsStruct();
 637       for (int i = 0; i < 100; i++) {
 638         BigIntWriter bigIntWriter = rootWriter.bigInt("a");
 639         bigIntWriter.setPosition(i);
 640         bigIntWriter.writeBigInt(i);
 641       }
 642       Field field = parent.getField().getChildren().get(0).getChildren().get(0);
 643       Assert.assertEquals("a", field.getName());
 644       Assert.assertEquals(Int.TYPE_TYPE, field.getType().getTypeID());
 645       Int intType = (Int) field.getType();
 646
 647       Assert.assertEquals(64, intType.getBitWidth());
 648       Assert.assertTrue(intType.getIsSigned());
 649       for (int i = 100; i < 200; i++) {
 650         VarCharWriter varCharWriter = rootWriter.varChar("a");
 651         varCharWriter.setPosition(i);
 652         byte[] bytes = Integer.toString(i).getBytes();
 653         ArrowBuf tempBuf = allocator.buffer(bytes.length);
 654         tempBuf.setBytes(0, bytes);
 655         varCharWriter.writeVarChar(0, bytes.length, tempBuf);
 656         tempBuf.close();
 657       }
 658       field = parent.getField().getChildren().get(0).getChildren().get(0);
 659       Assert.assertEquals("a", field.getName());
 660       Assert.assertEquals(Union.TYPE_TYPE, field.getType().getTypeID());
 661       Assert.assertEquals(Int.TYPE_TYPE, field.getChildren().get(0).getType().getTypeID());
 662       Assert.assertEquals(Utf8.TYPE_TYPE, field.getChildren().get(1).getType().getTypeID());
 663       StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
 664       for (int i = 0; i < 100; i++) {
 665         rootReader.setPosition(i);
 666         FieldReader reader = rootReader.reader("a");
 667         Long value = reader.readLong();
 668         Assert.assertNotNull("index: " + i, value);
 669         Assert.assertEquals(i, value.intValue());
 670       }
 671       for (int i = 100; i < 200; i++) {
 672         rootReader.setPosition(i);
 673         FieldReader reader = rootReader.reader("a");
 674         Text value = reader.readText();
 675         Assert.assertEquals(Integer.toString(i), value.toString());
 676       }
 677     }
 678   }
 679
 680   /**
 681    * Even without writing to the writer, the union schema is created correctly.
 682    */
 683   @Test
 684   public void promotableWriterSchema() {
 685     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
 686       ComplexWriter writer = new ComplexWriterImpl("root", parent);
 687       StructWriter rootWriter = writer.rootAsStruct();
 688       rootWriter.bigInt("a");
 689       rootWriter.varChar("a");
 690
 691       Field field = parent.getField().getChildren().get(0).getChildren().get(0);
 692       Assert.assertEquals("a", field.getName());
 693       Assert.assertEquals(ArrowTypeID.Union, field.getType().getTypeID());
 694
 695       Assert.assertEquals(ArrowTypeID.Int, field.getChildren().get(0).getType().getTypeID());
 696       Int intType = (Int) field.getChildren().get(0).getType();
 697       Assert.assertEquals(64, intType.getBitWidth());
 698       Assert.assertTrue(intType.getIsSigned());
 699       Assert.assertEquals(ArrowTypeID.Utf8, field.getChildren().get(1).getType().getTypeID());
 700     }
 701   }
 702
 703   private Set<String> getFieldNames(List<Field> fields) {
 704     Set<String> fieldNames = new HashSet<>();
 705     for (Field field : fields) {
 706       fieldNames.add(field.getName());
 707       if (!field.getChildren().isEmpty()) {
 708         for (String name : getFieldNames(field.getChildren())) {
 709           fieldNames.add(field.getName() + "::" + name);
 710         }
 711       }
 712     }
 713     return fieldNames;
 714   }
 715
 716   @Test
 717   public void structWriterMixedCaseFieldNames() {
 718     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
 719       // test case-sensitive StructWriter
 720       ComplexWriter writer = new ComplexWriterImpl("rootCaseSensitive", parent, false, true);
 721       StructWriter rootWriterCaseSensitive = writer.rootAsStruct();
 722       rootWriterCaseSensitive.bigInt("int_field");
 723       rootWriterCaseSensitive.bigInt("Int_Field");
 724       rootWriterCaseSensitive.float4("float_field");
 725       rootWriterCaseSensitive.float4("Float_Field");
 726       StructWriter structFieldWriterCaseSensitive = rootWriterCaseSensitive.struct("struct_field");
 727       structFieldWriterCaseSensitive.varChar("char_field");
 728       structFieldWriterCaseSensitive.varChar("Char_Field");
 729       ListWriter listFieldWriterCaseSensitive = rootWriterCaseSensitive.list("list_field");
 730       StructWriter listStructFieldWriterCaseSensitive = listFieldWriterCaseSensitive.struct();
 731       listStructFieldWriterCaseSensitive.bit("bit_field");
 732       listStructFieldWriterCaseSensitive.bit("Bit_Field");
 733
 734       List<Field> fieldsCaseSensitive = parent.getField().getChildren().get(0).getChildren();
 735       Set<String> fieldNamesCaseSensitive = getFieldNames(fieldsCaseSensitive);
 736       Assert.assertEquals(11, fieldNamesCaseSensitive.size());
 737       Assert.assertTrue(fieldNamesCaseSensitive.contains("int_field"));
 738       Assert.assertTrue(fieldNamesCaseSensitive.contains("Int_Field"));
 739       Assert.assertTrue(fieldNamesCaseSensitive.contains("float_field"));
 740       Assert.assertTrue(fieldNamesCaseSensitive.contains("Float_Field"));
 741       Assert.assertTrue(fieldNamesCaseSensitive.contains("struct_field"));
 742       Assert.assertTrue(fieldNamesCaseSensitive.contains("struct_field::char_field"));
 743       Assert.assertTrue(fieldNamesCaseSensitive.contains("struct_field::Char_Field"));
 744       Assert.assertTrue(fieldNamesCaseSensitive.contains("list_field"));
 745       Assert.assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$"));
 746       Assert.assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$::bit_field"));
 747       Assert.assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$::Bit_Field"));
 748
 749       // test case-insensitive StructWriter
 750       ComplexWriter writerCaseInsensitive = new ComplexWriterImpl("rootCaseInsensitive", parent, false, false);
 751       StructWriter rootWriterCaseInsensitive = writerCaseInsensitive.rootAsStruct();
 752
 753       rootWriterCaseInsensitive.bigInt("int_field");
 754       rootWriterCaseInsensitive.bigInt("Int_Field");
 755       rootWriterCaseInsensitive.float4("float_field");
 756       rootWriterCaseInsensitive.float4("Float_Field");
 757       StructWriter structFieldWriterCaseInsensitive = rootWriterCaseInsensitive.struct("struct_field");
 758       structFieldWriterCaseInsensitive.varChar("char_field");
 759       structFieldWriterCaseInsensitive.varChar("Char_Field");
 760       ListWriter listFieldWriterCaseInsensitive = rootWriterCaseInsensitive.list("list_field");
 761       StructWriter listStructFieldWriterCaseInsensitive = listFieldWriterCaseInsensitive.struct();
 762       listStructFieldWriterCaseInsensitive.bit("bit_field");
 763       listStructFieldWriterCaseInsensitive.bit("Bit_Field");
 764
 765       List<Field> fieldsCaseInsensitive = parent.getField().getChildren().get(1).getChildren();
 766       Set<String> fieldNamesCaseInsensitive = getFieldNames(fieldsCaseInsensitive);
 767       Assert.assertEquals(7, fieldNamesCaseInsensitive.size());
 768       Assert.assertTrue(fieldNamesCaseInsensitive.contains("int_field"));
 769       Assert.assertTrue(fieldNamesCaseInsensitive.contains("float_field"));
 770       Assert.assertTrue(fieldNamesCaseInsensitive.contains("struct_field"));
 771       Assert.assertTrue(fieldNamesCaseInsensitive.contains("struct_field::char_field"));
 772       Assert.assertTrue(fieldNamesCaseSensitive.contains("list_field"));
 773       Assert.assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$"));
 774       Assert.assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$::bit_field"));
 775     }
 776   }
 777
 778   @Test
 779   public void timeStampSecWriter() throws Exception {
 780     // test values
 781     final long expectedSecs = 981173106L;
 782     final LocalDateTime expectedSecDateTime = LocalDateTime.of(2001, 2, 3, 4, 5, 6, 0);
 783
 784     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
 785       // write
 786
 787       ComplexWriter writer = new ComplexWriterImpl("root", parent);
 788       StructWriter rootWriter = writer.rootAsStruct();
 789
 790       {
 791         TimeStampSecWriter timeStampSecWriter = rootWriter.timeStampSec("sec");
 792         timeStampSecWriter.setPosition(0);
 793         timeStampSecWriter.writeTimeStampSec(expectedSecs);
 794       }
 795       {
 796         TimeStampSecTZWriter timeStampSecTZWriter = rootWriter.timeStampSecTZ("secTZ", "UTC");
 797         timeStampSecTZWriter.setPosition(1);
 798         timeStampSecTZWriter.writeTimeStampSecTZ(expectedSecs);
 799       }
 800       // schema
 801       List<Field> children = parent.getField().getChildren().get(0).getChildren();
 802       checkTimestampField(children.get(0), "sec");
 803       checkTimestampTZField(children.get(1), "secTZ", "UTC");
 804
 805       // read
 806       StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
 807       {
 808         FieldReader secReader = rootReader.reader("sec");
 809         secReader.setPosition(0);
 810         LocalDateTime secDateTime = secReader.readLocalDateTime();
 811         Assert.assertEquals(expectedSecDateTime, secDateTime);
 812         long secLong = secReader.readLong();
 813         Assert.assertEquals(expectedSecs, secLong);
 814       }
 815       {
 816         FieldReader secTZReader = rootReader.reader("secTZ");
 817         secTZReader.setPosition(1);
 818         long secTZLong = secTZReader.readLong();
 819         Assert.assertEquals(expectedSecs, secTZLong);
 820       }
 821     }
 822   }
 823
 824   @Test
 825   public void timeStampMilliWriters() throws Exception {
 826     // test values
 827     final long expectedMillis = 981173106123L;
 828     final LocalDateTime expectedMilliDateTime = LocalDateTime.of(2001, 2, 3, 4, 5, 6, 123 * 1_000_000);
 829
 830     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator);) {
 831       // write
 832       ComplexWriter writer = new ComplexWriterImpl("root", parent);
 833       StructWriter rootWriter = writer.rootAsStruct();
 834       {
 835         TimeStampMilliWriter timeStampWriter = rootWriter.timeStampMilli("milli");
 836         timeStampWriter.setPosition(0);
 837         timeStampWriter.writeTimeStampMilli(expectedMillis);
 838       }
 839       String tz = "UTC";
 840       {
 841         TimeStampMilliTZWriter timeStampTZWriter = rootWriter.timeStampMilliTZ("milliTZ", tz);
 842         timeStampTZWriter.setPosition(0);
 843         timeStampTZWriter.writeTimeStampMilliTZ(expectedMillis);
 844       }
 845       // schema
 846       List<Field> children = parent.getField().getChildren().get(0).getChildren();
 847       checkTimestampField(children.get(0), "milli");
 848       checkTimestampTZField(children.get(1), "milliTZ", tz);
 849
 850       // read
 851       StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
 852
 853       {
 854         FieldReader milliReader = rootReader.reader("milli");
 855         milliReader.setPosition(0);
 856         LocalDateTime milliDateTime = milliReader.readLocalDateTime();
 857         Assert.assertEquals(expectedMilliDateTime, milliDateTime);
 858         long milliLong = milliReader.readLong();
 859         Assert.assertEquals(expectedMillis, milliLong);
 860       }
 861       {
 862         FieldReader milliTZReader = rootReader.reader("milliTZ");
 863         milliTZReader.setPosition(0);
 864         long milliTZLong = milliTZReader.readLong();
 865         Assert.assertEquals(expectedMillis, milliTZLong);
 866       }
 867     }
 868   }
 869
 870   private void checkTimestampField(Field field, String name) {
 871     Assert.assertEquals(name, field.getName());
 872     Assert.assertEquals(ArrowType.Timestamp.TYPE_TYPE, field.getType().getTypeID());
 873   }
 874
 875   private void checkTimestampTZField(Field field, String name, String tz) {
 876     checkTimestampField(field, name);
 877     Assert.assertEquals(tz, ((Timestamp) field.getType()).getTimezone());
 878   }
 879
 880   @Test
 881   public void timeStampMicroWriters() throws Exception {
 882     // test values
 883     final long expectedMicros = 981173106123456L;
 884     final LocalDateTime expectedMicroDateTime = LocalDateTime.of(2001, 2, 3, 4, 5, 6, 123456 * 1000);
 885
 886     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
 887       // write
 888       ComplexWriter writer = new ComplexWriterImpl("root", parent);
 889       StructWriter rootWriter = writer.rootAsStruct();
 890
 891       {
 892         TimeStampMicroWriter timeStampMicroWriter = rootWriter.timeStampMicro("micro");
 893         timeStampMicroWriter.setPosition(0);
 894         timeStampMicroWriter.writeTimeStampMicro(expectedMicros);
 895       }
 896       String tz = "UTC";
 897       {
 898         TimeStampMicroTZWriter timeStampMicroWriter = rootWriter.timeStampMicroTZ("microTZ", tz);
 899         timeStampMicroWriter.setPosition(1);
 900         timeStampMicroWriter.writeTimeStampMicroTZ(expectedMicros);
 901       }
 902
 903       // schema
 904       List<Field> children = parent.getField().getChildren().get(0).getChildren();
 905       checkTimestampField(children.get(0), "micro");
 906       checkTimestampTZField(children.get(1), "microTZ", tz);
 907
 908       // read
 909       StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
 910       {
 911         FieldReader microReader = rootReader.reader("micro");
 912         microReader.setPosition(0);
 913         LocalDateTime microDateTime = microReader.readLocalDateTime();
 914         Assert.assertEquals(expectedMicroDateTime, microDateTime);
 915         long microLong = microReader.readLong();
 916         Assert.assertEquals(expectedMicros, microLong);
 917       }
 918       {
 919         FieldReader microReader = rootReader.reader("microTZ");
 920         microReader.setPosition(1);
 921         long microLong = microReader.readLong();
 922         Assert.assertEquals(expectedMicros, microLong);
 923       }
 924     }
 925   }
 926
 927   @Test
 928   public void timeStampNanoWriters() throws Exception {
 929     // test values
 930     final long expectedNanos = 981173106123456789L;
 931     final LocalDateTime expectedNanoDateTime = LocalDateTime.of(2001, 2, 3, 4, 5, 6, 123456789);
 932
 933     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
 934       // write
 935       ComplexWriter writer = new ComplexWriterImpl("root", parent);
 936       StructWriter rootWriter = writer.rootAsStruct();
 937
 938       {
 939         TimeStampNanoWriter timeStampNanoWriter = rootWriter.timeStampNano("nano");
 940         timeStampNanoWriter.setPosition(0);
 941         timeStampNanoWriter.writeTimeStampNano(expectedNanos);
 942       }
 943       String tz = "UTC";
 944       {
 945         TimeStampNanoTZWriter timeStampNanoWriter = rootWriter.timeStampNanoTZ("nanoTZ", tz);
 946         timeStampNanoWriter.setPosition(0);
 947         timeStampNanoWriter.writeTimeStampNanoTZ(expectedNanos);
 948       }
 949       // schema
 950       List<Field> children = parent.getField().getChildren().get(0).getChildren();
 951       checkTimestampField(children.get(0), "nano");
 952       checkTimestampTZField(children.get(1), "nanoTZ", tz);
 953       // read
 954       StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
 955
 956       {
 957         FieldReader nanoReader = rootReader.reader("nano");
 958         nanoReader.setPosition(0);
 959         LocalDateTime nanoDateTime = nanoReader.readLocalDateTime();
 960         Assert.assertEquals(expectedNanoDateTime, nanoDateTime);
 961         long nanoLong = nanoReader.readLong();
 962         Assert.assertEquals(expectedNanos, nanoLong);
 963       }
 964       {
 965         FieldReader nanoReader = rootReader.reader("nanoTZ");
 966         nanoReader.setPosition(0);
 967         long nanoLong = nanoReader.readLong();
 968         Assert.assertEquals(expectedNanos, nanoLong);
 969         NullableTimeStampNanoTZHolder h = new NullableTimeStampNanoTZHolder();
 970         nanoReader.read(h);
 971         Assert.assertEquals(expectedNanos, h.value);
 972       }
 973     }
 974
 975   }
 976
 977   @Test
 978   public void fixedSizeBinaryWriters() throws Exception {
 979     // test values
 980     int numValues = 10;
 981     int byteWidth = 9;
 982     byte[][] values = new byte[numValues][byteWidth];
 983     for (int i = 0; i < numValues; i++) {
 984       for (int j = 0; j < byteWidth; j++) {
 985         values[i][j] = ((byte) i);
 986       }
 987     }
 988     ArrowBuf[] bufs = new ArrowBuf[numValues];
 989     for (int i = 0; i < numValues; i++) {
 990       bufs[i] = allocator.buffer(byteWidth);
 991       bufs[i].setBytes(0, values[i]);
 992     }
 993
 994     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
 995       // write
 996       ComplexWriter writer = new ComplexWriterImpl("root", parent);
 997       StructWriter rootWriter = writer.rootAsStruct();
 998
 999       String fieldName = "fixedSizeBinary";
1000       FixedSizeBinaryWriter fixedSizeBinaryWriter = rootWriter.fixedSizeBinary(fieldName, byteWidth);
1001       for (int i = 0; i < numValues; i++) {
1002         fixedSizeBinaryWriter.setPosition(i);
1003         fixedSizeBinaryWriter.writeFixedSizeBinary(bufs[i]);
1004       }
1005
1006       // schema
1007       List<Field> children = parent.getField().getChildren().get(0).getChildren();
1008       Assert.assertEquals(fieldName, children.get(0).getName());
1009       Assert.assertEquals(ArrowType.FixedSizeBinary.TYPE_TYPE, children.get(0).getType().getTypeID());
1010
1011       // read
1012       StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
1013
1014       FieldReader fixedSizeBinaryReader = rootReader.reader(fieldName);
1015       for (int i = 0; i < numValues; i++) {
1016         fixedSizeBinaryReader.setPosition(i);
1017         byte[] readValues = fixedSizeBinaryReader.readByteArray();
1018         Assert.assertArrayEquals(values[i], readValues);
1019       }
1020     }
1021
1022     AutoCloseables.close(bufs);
1023   }
1024
1025   @Test
1026   public void complexCopierWithList() {
1027     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
1028       ComplexWriter writer = new ComplexWriterImpl("root", parent);
1029       StructWriter rootWriter = writer.rootAsStruct();
1030       ListWriter listWriter = rootWriter.list("list");
1031       StructWriter innerStructWriter = listWriter.struct();
1032       IntWriter outerIntWriter = listWriter.integer();
1033       rootWriter.start();
1034       listWriter.startList();
1035       outerIntWriter.writeInt(1);
1036       outerIntWriter.writeInt(2);
1037       innerStructWriter.start();
1038       IntWriter intWriter = innerStructWriter.integer("a");
1039       intWriter.writeInt(1);
1040       innerStructWriter.end();
1041       innerStructWriter.start();
1042       intWriter = innerStructWriter.integer("a");
1043       intWriter.writeInt(2);
1044       innerStructWriter.end();
1045       listWriter.endList();
1046       rootWriter.end();
1047       writer.setValueCount(1);
1048
1049       StructVector structVector = (StructVector) parent.getChild("root");
1050       TransferPair tp = structVector.getTransferPair(allocator);
1051       tp.splitAndTransfer(0, 1);
1052       NonNullableStructVector toStructVector = (NonNullableStructVector) tp.getTo();
1053       JsonStringHashMap<?, ?> toMapValue = (JsonStringHashMap<?, ?>) toStructVector.getObject(0);
1054       JsonStringArrayList<?> object = (JsonStringArrayList<?>) toMapValue.get("list");
1055       assertEquals(1, object.get(0));
1056       assertEquals(2, object.get(1));
1057       JsonStringHashMap<?, ?> innerStruct = (JsonStringHashMap<?, ?>) object.get(2);
1058       assertEquals(1, innerStruct.get("a"));
1059       innerStruct = (JsonStringHashMap<?, ?>) object.get(3);
1060       assertEquals(2, innerStruct.get("a"));
1061       toStructVector.close();
1062     }
1063   }
1064
1065   @Test
1066   public void testSingleStructWriter1() {
1067     /* initialize a SingleStructWriter with empty StructVector and then lazily
1068      * create all vectors with expected initialCapacity.
1069      */
1070     try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
1071       SingleStructWriter singleStructWriter = new SingleStructWriter(parent);
1072
1073       int initialCapacity = 1024;
1074       singleStructWriter.setInitialCapacity(initialCapacity);
1075
1076       IntWriter intWriter = singleStructWriter.integer("intField");
1077       BigIntWriter bigIntWriter = singleStructWriter.bigInt("bigIntField");
1078       Float4Writer float4Writer = singleStructWriter.float4("float4Field");
1079       Float8Writer float8Writer = singleStructWriter.float8("float8Field");
1080       ListWriter listWriter = singleStructWriter.list("listField");
1081       MapWriter mapWriter = singleStructWriter.map("mapField", false);
1082
1083       int intValue = 100;
1084       long bigIntValue = 10000;
1085       float float4Value = 100.5f;
1086       double float8Value = 100.375;
1087
1088       for (int i = 0; i < initialCapacity; i++) {
1089         singleStructWriter.start();
1090
1091         intWriter.writeInt(intValue + i);
1092         bigIntWriter.writeBigInt(bigIntValue + (long) i);
1093         float4Writer.writeFloat4(float4Value + (float) i);
1094         float8Writer.writeFloat8(float8Value + (double) i);
1095
1096         listWriter.setPosition(i);
1097         listWriter.startList();
1098         listWriter.integer().writeInt(intValue + i);
1099         listWriter.integer().writeInt(intValue + i + 1);
1100         listWriter.integer().writeInt(intValue + i + 2);
1101         listWriter.integer().writeInt(intValue + i + 3);
1102         listWriter.endList();
1103
1104         mapWriter.setPosition(i);
1105         mapWriter.startMap();
1106         mapWriter.startEntry();
1107         mapWriter.key().integer().writeInt(intValue + i);
1108         mapWriter.value().integer().writeInt(intValue + i + 1);
1109         mapWriter.endEntry();
1110         mapWriter.startEntry();
1111         mapWriter.key().integer().writeInt(intValue + i + 2);
1112         mapWriter.value().integer().writeInt(intValue + i + 3);
1113         mapWriter.endEntry();
1114         mapWriter.endMap();
1115
1116         singleStructWriter.end();
1117       }
1118
1119       IntVector intVector = (IntVector) parent.getChild("intField");
1120       BigIntVector bigIntVector = (BigIntVector) parent.getChild("bigIntField");
1121       Float4Vector float4Vector = (Float4Vector) parent.getChild("float4Field");
1122       Float8Vector float8Vector = (Float8Vector) parent.getChild("float8Field");
1123
1124       int capacity = singleStructWriter.getValueCapacity();
1125       assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2);
1126       capacity = intVector.getValueCapacity();
1127       assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2);
1128       capacity = bigIntVector.getValueCapacity();
1129       assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2);
1130       capacity = float4Vector.getValueCapacity();
1131       assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2);
1132       capacity = float8Vector.getValueCapacity();
1133       assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2);
1134
1135       StructReader singleStructReader = new SingleStructReaderImpl(parent);
1136
1137       IntReader intReader = singleStructReader.reader("intField");
1138       BigIntReader bigIntReader = singleStructReader.reader("bigIntField");
1139       Float4Reader float4Reader = singleStructReader.reader("float4Field");
1140       Float8Reader float8Reader = singleStructReader.reader("float8Field");
1141       UnionListReader listReader = (UnionListReader) singleStructReader.reader("listField");
1142       UnionMapReader mapReader = (UnionMapReader) singleStructReader.reader("mapField");
1143
1144       for (int i = 0; i < initialCapacity; i++) {
1145         intReader.setPosition(i);
1146         bigIntReader.setPosition(i);
1147         float4Reader.setPosition(i);
1148         float8Reader.setPosition(i);
1149         listReader.setPosition(i);
1150         mapReader.setPosition(i);
1151
1152         assertEquals(intValue + i, intReader.readInteger().intValue());
1153         assertEquals(bigIntValue + (long) i, bigIntReader.readLong().longValue());
1154         assertEquals(float4Value + (float) i, float4Reader.readFloat().floatValue(), 0);
1155         assertEquals(float8Value + (double) i, float8Reader.readDouble().doubleValue(), 0);
1156
1157         for (int j = 0; j < 4; j++) {
1158           listReader.next();
1159           assertEquals(intValue + i + j, listReader.reader().readInteger().intValue());
1160         }
1161
1162         for (int k = 0; k < 4; k += 2) {
1163           mapReader.next();
1164           assertEquals(intValue + k + i, mapReader.key().readInteger().intValue());
1165           assertEquals(intValue + k + i + 1, mapReader.value().readInteger().intValue());
1166         }
1167       }
1168     }
1169
1170
1171   }
1172
1173   @Test
1174   public void testListWriterWithNulls() {
1175     try (ListVector listVector = ListVector.empty("list", allocator)) {
1176       listVector.setInitialCapacity(COUNT);
1177       listVector.allocateNew();
1178       listVector.getValidityBuffer().setOne(0, (int) listVector.getValidityBuffer().capacity());
1179
1180       UnionListWriter listWriter = listVector.getWriter();
1181
1182       // expected listVector :  [[null], null, [2, 4], null, [null], null, [6, 12], ...]
1183       for (int i = 0; i < COUNT; i++) {
1184         listWriter.setPosition(i);
1185         if (i % 2 == 0) {
1186           listWriter.startList();
1187           if (i % 4 == 0) {
1188             listWriter.integer().writeNull();
1189           } else {
1190             listWriter.integer().writeInt(i);
1191             listWriter.integer().writeInt(i * 2);
1192           }
1193           listWriter.endList();
1194         } else {
1195           listWriter.writeNull();
1196         }
1197       }
1198       listVector.setValueCount(COUNT);
1199
1200       UnionListReader listReader = new UnionListReader(listVector);
1201       for (int i = 0; i < COUNT; i++) {
1202         listReader.setPosition(i);
1203         if (i % 2 == 0) {
1204           Assert.assertTrue(listReader.isSet());
1205           listReader.next();
1206           if (i % 4 == 0) {
1207             Assert.assertNull(listReader.reader().readInteger());
1208           } else {
1209             Assert.assertEquals(i, listReader.reader().readInteger().intValue());
1210             listReader.next();
1211             Assert.assertEquals(i * 2, listReader.reader().readInteger().intValue());
1212           }
1213         } else {
1214           Assert.assertFalse(listReader.isSet());
1215         }
1216       }
1217     }
1218   }
1219
1220   @Test
1221   public void testListOfListWriterWithNulls() {
1222     try (ListVector listVector = ListVector.empty("listoflist", allocator)) {
1223       listVector.setInitialCapacity(COUNT);
1224       listVector.allocateNew();
1225       listVector.getValidityBuffer().setOne(0, (int) listVector.getValidityBuffer().capacity());
1226
1227       UnionListWriter listWriter = listVector.getWriter();
1228
1229       // create list : [ [null], null, [[null, 2, 4]], null, [null], null, [[null, 6, 12]], ... ]
1230       for (int i = 0; i < COUNT; i++) {
1231         listWriter.setPosition(i);
1232         if (i % 2 == 0) {
1233           listWriter.startList();
1234           if (i % 4 == 0) {
1235             listWriter.list().writeNull();
1236           } else {
1237             listWriter.list().startList();
1238             listWriter.list().integer().writeNull();
1239             listWriter.list().integer().writeInt(i);
1240             listWriter.list().integer().writeInt(i * 2);
1241             listWriter.list().endList();
1242           }
1243           listWriter.endList();
1244         } else {
1245           listWriter.writeNull();
1246         }
1247       }
1248       listVector.setValueCount(COUNT);
1249
1250       UnionListReader listReader = new UnionListReader(listVector);
1251       for (int i = 0; i < COUNT; i++) {
1252         listReader.setPosition(i);
1253         if (i % 2 == 0) {
1254           Assert.assertTrue(listReader.isSet());
1255           listReader.next();
1256           if (i % 4 == 0) {
1257             Assert.assertFalse(listReader.reader().isSet());
1258           } else {
1259             listReader.reader().next();
1260             Assert.assertFalse(listReader.reader().reader().isSet());
1261             listReader.reader().next();
1262             Assert.assertEquals(i, listReader.reader().reader().readInteger().intValue());
1263             listReader.reader().next();
1264             Assert.assertEquals(i * 2, listReader.reader().reader().readInteger().intValue());
1265           }
1266         } else {
1267           Assert.assertFalse(listReader.isSet());
1268         }
1269       }
1270     }
1271   }
1272
1273   @Test
1274   public void testListOfListOfListWriterWithNulls() {
1275     try (ListVector listVector = ListVector.empty("listoflistoflist", allocator)) {
1276       listVector.setInitialCapacity(COUNT);
1277       listVector.allocateNew();
1278       listVector.getValidityBuffer().setOne(0, (int) listVector.getValidityBuffer().capacity());
1279
1280       UnionListWriter listWriter = listVector.getWriter();
1281
1282       // create list : [ null, [null], [[null]], [[[null, 1, 2]]], null, [null], ...
1283       for (int i = 0; i < COUNT; i++) {
1284         listWriter.setPosition(i);
1285         if (i % 4 == 0) {
1286           listWriter.writeNull();
1287         } else {
1288           listWriter.startList();
1289           if (i % 4 == 1) {
1290             listWriter.list().writeNull();
1291           } else if (i % 4 == 2) {
1292             listWriter.list().startList();
1293             listWriter.list().list().writeNull();
1294             listWriter.list().endList();
1295           } else {
1296             listWriter.list().startList();
1297             listWriter.list().list().startList();
1298             listWriter.list().list().integer().writeNull();
1299             listWriter.list().list().integer().writeInt(i);
1300             listWriter.list().list().integer().writeInt(i * 2);
1301             listWriter.list().list().endList();
1302             listWriter.list().endList();
1303           }
1304           listWriter.endList();
1305         }
1306       }
1307       listVector.setValueCount(COUNT);
1308
1309       UnionListReader listReader = new UnionListReader(listVector);
1310       for (int i = 0; i < COUNT; i++) {
1311         listReader.setPosition(i);
1312         if (i % 4 == 0) {
1313           Assert.assertFalse(listReader.isSet());
1314         } else {
1315           Assert.assertTrue(listReader.isSet());
1316           listReader.next();
1317           if (i % 4 == 1) {
1318             Assert.assertFalse(listReader.reader().isSet());
1319           } else if (i % 4 == 2) {
1320             listReader.reader().next();
1321             Assert.assertFalse(listReader.reader().reader().isSet());
1322           } else {
1323             listReader.reader().next();
1324             listReader.reader().reader().next();
1325             Assert.assertFalse(listReader.reader().reader().reader().isSet());
1326             listReader.reader().reader().next();
1327             Assert.assertEquals(i, listReader.reader().reader().reader().readInteger().intValue());
1328             listReader.reader().reader().next();
1329             Assert.assertEquals(i * 2, listReader.reader().reader().reader().readInteger().intValue());
1330           }
1331         }
1332       }
1333     }
1334   }
1335 }