ARROW-274: Add NullableMapVector to support nullable maps

Author: Julien Le Dem <[email protected]>

Closes apache#128 from julienledem/nullable_map and squashes the following commits:

d98580a [Julien Le Dem] review feedback
ee1dd45 [Julien Le Dem] Fix complex writers/readers
8780f48 [Julien Le Dem] ARROW-274: Add NullableMapVector to support nullable maps

Author: julienledem

vector/src/main/codegen/templates/MapWriters.java

@@ -17,14 +17,13 @@
  */
 
 <@pp.dropOutputFile />
-<#list ["Single"] as mode>
+<#list ["Nullable", "Single"] as mode>
 <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/${mode}MapWriter.java" />
+<#assign index = "idx()">
 <#if mode == "Single">
 <#assign containerClass = "MapVector" />
-<#assign index = "idx()">
 <#else>
-<#assign containerClass = "RepeatedMapVector" />
-<#assign index = "currentChildIndex">
+<#assign containerClass = "NullableMapVector" />
 </#if>
 
 <#include "/@includes/license.ftl" />
@@ -49,9 +48,13 @@ public class ${mode}MapWriter extends AbstractFieldWriter {
 
   protected final ${containerClass} container;
   private final Map<String, FieldWriter> fields = Maps.newHashMap();
-  <#if mode == "Repeated">private int currentChildIndex = 0;</#if>
 
   public ${mode}MapWriter(${containerClass} container) {
+    <#if mode == "Single">
+    if (container instanceof NullableMapVector) {
+      throw new IllegalArgumentException("Invalid container: " + container);
+    }
+    </#if>
     this.container = container;
   }
 
@@ -75,12 +78,12 @@ public MapWriter map(String name) {
       FieldWriter writer = fields.get(name.toLowerCase());
     if(writer == null){
       int vectorCount=container.size();
-      MapVector vector = container.addOrGet(name, MinorType.MAP, MapVector.class);
+      NullableMapVector vector = container.addOrGet(name, MinorType.MAP, NullableMapVector.class);
       writer = new PromotableWriter(vector, container);
       if(vectorCount != container.size()) {
         writer.allocate();
       }
-      writer.setPosition(${index});
+      writer.setPosition(idx());
       fields.put(name.toLowerCase(), writer);
     }
     return writer;
@@ -117,40 +120,12 @@ public ListWriter list(String name) {
       if (container.size() > vectorCount) {
         writer.allocate();
       }
-      writer.setPosition(${index});
+      writer.setPosition(idx());
       fields.put(name.toLowerCase(), writer);
     }
     return writer;
   }
 
-  <#if mode == "Repeated">
-  public void start() {
-      // update the repeated vector to state that there is current+1 objects.
-    final RepeatedMapHolder h = new RepeatedMapHolder();
-    final RepeatedMapVector map = (RepeatedMapVector) container;
-    final RepeatedMapVector.Mutator mutator = map.getMutator();
-
-    // Make sure that the current vector can support the end position of this list.
-    if(container.getValueCapacity() <= idx()) {
-      mutator.setValueCount(idx()+1);
-    }
-
-    map.getAccessor().get(idx(), h);
-    if (h.start >= h.end) {
-      container.getMutator().startNewValue(idx());
-    }
-    currentChildIndex = container.getMutator().add(idx());
-    for(final FieldWriter w : fields.values()) {
-      w.setPosition(currentChildIndex);
-    }
-  }
-
-
-  public void end() {
-    // noop
-  }
-  <#else>
-
   public void setValueCount(int count) {
     container.getMutator().setValueCount(count);
   }
@@ -165,14 +140,16 @@ public void setPosition(int index) {
 
   @Override
   public void start() {
+    <#if mode == "Single">
+    <#else>
+    container.getMutator().setIndexDefined(idx());
+    </#if>
   }
 
   @Override
   public void end() {
   }
 
-  </#if>
-
   <#list vv.types as type><#list type.minor as minor>
   <#assign lowerName = minor.class?uncap_first />
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>
@@ -204,7 +181,7 @@ public void end() {
       if (currentVector == null || currentVector != vector) {
         vector.allocateNewSafe();
       } 
-      writer.setPosition(${index});
+      writer.setPosition(idx());
       fields.put(name.toLowerCase(), writer);
     }
     return writer;

vector/src/main/codegen/templates/UnionListWriter.java

@@ -160,11 +160,13 @@ public void start() {
     vector.getMutator().setNotNull(idx());
     offsets.getMutator().setSafe(idx() + 1, nextOffset);
     writer.setPosition(nextOffset);
+    writer.start();
   }
 
   @Override
   public void end() {
 //    if (inMap) {
+      writer.end();
       inMap = false;
       final int nextOffset = offsets.getAccessor().get(idx() + 1);
       offsets.getMutator().setSafe(idx() + 1, nextOffset + 1);

vector/src/main/codegen/templates/UnionVector.java

@@ -72,7 +72,7 @@ public class UnionVector implements FieldVector {
   MapVector internalMap;
   UInt1Vector typeVector;
 
-  private MapVector mapVector;
+  private NullableMapVector mapVector;
   private ListVector listVector;
 
   private FieldReader reader;
@@ -127,10 +127,10 @@ public List<BufferBacked> getFieldInnerVectors() {
     throw new UnsupportedOperationException();
   }
 
-  public MapVector getMap() {
+  public NullableMapVector getMap() {
     if (mapVector == null) {
       int vectorCount = internalMap.size();
-      mapVector = internalMap.addOrGet("map", MinorType.MAP, MapVector.class);
+      mapVector = internalMap.addOrGet("map", MinorType.MAP, NullableMapVector.class);
       if (internalMap.size() > vectorCount) {
         mapVector.allocateNew();
         if (callBack != null) {

vector/src/main/codegen/templates/UnionWriter.java

@@ -74,7 +74,7 @@ public void endList() {
 
   private MapWriter getMapWriter() {
     if (mapWriter == null) {
-      mapWriter = new SingleMapWriter(data.getMap());
+      mapWriter = new NullableMapWriter(data.getMap());
       mapWriter.setPosition(idx());
       writers.add(mapWriter);
     }

vector/src/main/java/org/apache/arrow/vector/NullableVector.java

@@ -17,7 +17,7 @@
  */
 package org.apache.arrow.vector;
 
-public interface NullableVector extends ValueVector{
+public interface NullableVector extends ValueVector {
 
   ValueVector getValuesVector();
 }

vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java

@@ -68,7 +68,9 @@ private void appendNodes(FieldVector vector, List<ArrowFieldNode> nodes, List<Ar
     List<ArrowBuf> fieldBuffers = vector.getFieldBuffers();
     List<ArrowVectorType> expectedBuffers = vector.getField().getTypeLayout().getVectorTypes();
     if (fieldBuffers.size() != expectedBuffers.size()) {
-      throw new IllegalArgumentException("wrong number of buffers for field " + vector.getField() + ". found: " + fieldBuffers);
+      throw new IllegalArgumentException(String.format(
+          "wrong number of buffers for field %s in vector %s. found: %s",
+          vector.getField(), vector.getClass().getSimpleName(), fieldBuffers));
     }
     buffers.addAll(fieldBuffers);
     for (FieldVector child : vector.getChildrenFromFields()) {

vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java

@@ -18,25 +18,20 @@
 package org.apache.arrow.vector.complex;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 
 import javax.annotation.Nullable;
 
 import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.vector.BaseDataValueVector;
 import org.apache.arrow.vector.BaseValueVector;
-import org.apache.arrow.vector.BufferBacked;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.complex.impl.SingleMapReaderImpl;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.holders.ComplexHolder;
-import org.apache.arrow.vector.schema.ArrowFieldNode;
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.ArrowType.Tuple;
@@ -49,26 +44,20 @@
 import com.google.common.collect.Ordering;
 import com.google.common.primitives.Ints;
 
-import io.netty.buffer.ArrowBuf;
-
-public class MapVector extends AbstractMapVector implements FieldVector {
+public class MapVector extends AbstractMapVector {
   //private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(MapVector.class);
 
-  private final SingleMapReaderImpl reader = new SingleMapReaderImpl(MapVector.this);
+  private final SingleMapReaderImpl reader = new SingleMapReaderImpl(this);
   private final Accessor accessor = new Accessor();
   private final Mutator mutator = new Mutator();
   int valueCount;
 
-  // TODO: validity vector
-  private final List<BufferBacked> innerVectors = Collections.unmodifiableList(Arrays.<BufferBacked>asList());
-
-  public MapVector(String name, BufferAllocator allocator, CallBack callBack){
+  public MapVector(String name, BufferAllocator allocator, CallBack callBack) {
     super(name, allocator, callBack);
   }
 
   @Override
   public FieldReader getReader() {
-    //return new SingleMapReaderImpl(MapVector.this);
     return reader;
   }
 
@@ -124,18 +113,9 @@ public int getBufferSizeFor(final int valueCount) {
     return (int) bufferSize;
   }
 
-  @Override
-  public ArrowBuf[] getBuffers(boolean clear) {
-    int expectedSize = getBufferSize();
-    int actualSize   = super.getBufferSize();
-
-    Preconditions.checkArgument(expectedSize == actualSize, expectedSize + " != " + actualSize);
-    return super.getBuffers(clear);
-  }
-
   @Override
   public TransferPair getTransferPair(BufferAllocator allocator) {
-    return new MapTransferPair(this, name, allocator);
+    return new MapTransferPair(this, new MapVector(name, allocator, callBack), false);
   }
 
   @Override
@@ -145,18 +125,14 @@ public TransferPair makeTransferPair(ValueVector to) {
 
   @Override
   public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
-    return new MapTransferPair(this, ref, allocator);
+    return new MapTransferPair(this, new MapVector(ref, allocator, callBack), false);
   }
 
   protected static class MapTransferPair implements TransferPair{
     private final TransferPair[] pairs;
     private final MapVector from;
     private final MapVector to;
 
-    public MapTransferPair(MapVector from, String name, BufferAllocator allocator) {
-      this(from, new MapVector(name, allocator, from.callBack), false);
-    }
-
     public MapTransferPair(MapVector from, MapVector to) {
       this(from, to, true);
     }
@@ -335,7 +311,6 @@ public void close() {
     super.close();
  }
 
-  @Override
   public void initializeChildrenFromFields(List<Field> children) {
     for (Field field : children) {
       MinorType minorType = Types.getMinorTypeForArrowType(field.getType());
@@ -344,25 +319,9 @@ public void initializeChildrenFromFields(List<Field> children) {
     }
   }
 
-  @Override
+
   public List<FieldVector> getChildrenFromFields() {
     return getChildren();
   }
 
-  @Override
-  public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> ownBuffers) {
-    BaseDataValueVector.load(getFieldInnerVectors(), ownBuffers);
-    // TODO: something with fieldNode?
-  }
-
-  @Override
-  public List<ArrowBuf> getFieldBuffers() {
-    return BaseDataValueVector.unload(getFieldInnerVectors());
-  }
-
-  @Override
-  public List<BufferBacked> getFieldInnerVectors() {
-    return innerVectors;
-  }
-
 }