ARROW-8646: [Java] Allow UnionListWriter to write null values

Closes apache#7070 from tvamsikalyan/master Authored-by: Vamsi <vamsi@dremio.com> Signed-off-by: Pindikura Ravindra <ravindra@dremio.com>
dremio · May 1, 2020 · daf2e01 · daf2e01
1 parent 1f175e1
commit daf2e01
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 0 deletions.
diff --git a/java/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java b/java/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java
@@ -181,6 +181,11 @@ public void write(DecimalHolder holder) {
     writer.setPosition(writer.idx() + 1);
   }
 
+  @Override
+  public void writeNull() {
+    writer.setPosition(writer.idx()+1);
+  }
+
   public void writeDecimal(int start, ArrowBuf buffer, ArrowType arrowType) {
     if (writer.idx() >= (idx() + 1) * listSize) {
       throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));

diff --git a/java/vector/src/main/codegen/templates/UnionListWriter.java b/java/vector/src/main/codegen/templates/UnionListWriter.java
@@ -178,6 +178,11 @@ public void write(DecimalHolder holder) {
     writer.setPosition(writer.idx()+1);
   }
 
+  @Override
+  public void writeNull() {
+    writer.setPosition(writer.idx()+1);
+  }
+
   public void writeDecimal(int start, ArrowBuf buffer, ArrowType arrowType) {
     writer.writeDecimal(start, buffer, arrowType);
     writer.setPosition(writer.idx()+1);

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java
@@ -240,6 +240,38 @@ public void testConsistentChildName() throws Exception {
     }
   }
 
+  @Test
+  public void testUnionFixedSizeListWriterWithNulls() throws Exception {
+    /* Write to a decimal list vector
+     * each list of size 3 and having its data values alternating between null and a non-null.
+     * Read and verify
+     */
+    try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /*listSize=*/3, allocator)) {
+
+      UnionFixedSizeListWriter writer = vector.getWriter();
+      writer.allocate();
+
+      final int valueCount = 100;
+
+      for (int i = 0; i < valueCount; i++) {
+        writer.startList();
+        writer.decimal().writeDecimal(new BigDecimal(i));
+        writer.writeNull();
+        writer.decimal().writeDecimal(new BigDecimal(i * 3));
+        writer.endList();
+      }
+      vector.setValueCount(valueCount);
+
+      for (int i = 0; i < valueCount; i++) {
+        List<BigDecimal> values = (List<BigDecimal>) vector.getObject(i);
+        assertEquals(3, values.size());
+        assertEquals(new BigDecimal(i), values.get(0));
+        assertEquals(null, values.get(1));
+        assertEquals(new BigDecimal(i * 3), values.get(2));
+      }
+    }
+  }
+
   @Test
   public void testUnionFixedSizeListWriter() throws Exception {
     try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", 3, allocator)) {

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java
@@ -266,6 +266,45 @@ public void listScalarType() {
     }
   }
 
+  @Test
+  public void testListScalarNull() {
+    /* Write to a integer list vector
+     * each list of size 8 and having it's data values alternating between null and a non-null.
+     * Read and verify
+     */
+    try (ListVector listVector = ListVector.empty("list", allocator)) {
+      listVector.allocateNew();
+      UnionListWriter listWriter = new UnionListWriter(listVector);
+      for (int i = 0; i < COUNT; i++) {
+        listWriter.startList();
+        for (int j = 0; j < i % 7; j++) {
+          if (j % 2 == 0) {
+            listWriter.writeNull();
+          } else {
+            IntHolder holder = new IntHolder();
+            holder.value = j;
+            listWriter.write(holder);
+          }
+        }
+        listWriter.endList();
+      }
+      listWriter.setValueCount(COUNT);
+      UnionListReader listReader = new UnionListReader(listVector);
+      for (int i = 0; i < COUNT; i++) {
+        listReader.setPosition(i);
+        for (int j = 0; j < i % 7; j++) {
+          listReader.next();
+          if (j % 2 == 0) {
+            assertFalse("index is set: " + j, listReader.reader().isSet());
+          } else {
+            assertTrue("index is not set: " + j, listReader.reader().isSet());
+            assertEquals(j, listReader.reader().readInteger().intValue());
+          }
+        }
+      }
+    }
+  }
+
   @Test
   public void listDecimalType() {
     try (ListVector listVector = ListVector.empty("list", allocator)) {