/*

  * Licensed to the Apache Software Foundation (ASF) under one

  * or more contributor license agreements.  See the NOTICE file

  * distributed with this work for additional information

  * regarding copyright ownership.  The ASF licenses this file

  * to you under the Apache License, Version 2.0 (the

  * "License"); you may not use this file except in compliance

  * with the License.  You may obtain a copy of the License at

  *

  *     http://www.apache.org/licenses/LICENSE-2.0

  *

  * Unless required by applicable law or agreed to in writing, software

  * distributed under the License is distributed on an "AS IS" BASIS,

  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  * See the License for the specific language governing permissions and

  * limitations under the License.

  */

 package org.apache.giraph.partition;

 import com.google.common.collect.MapMaker;

 import com.google.common.primitives.Ints;

 import org.apache.giraph.edge.Edge;

 import org.apache.giraph.graph.Vertex;

 import org.apache.giraph.utils.UnsafeByteArrayInputStream;

 import org.apache.giraph.utils.WritableUtils;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import org.apache.hadoop.util.Progressable;

 import javax.annotation.concurrent.NotThreadSafe;

 import java.io.DataInput;

 import java.io.DataOutput;

 import java.io.IOException;

 import java.util.Iterator;

 import java.util.Map;

 import java.util.concurrent.ConcurrentMap;

 /**

  * Byte array based partition.  Should reduce the amount of memory used since

  * the entire graph is compressed into byte arrays.  Must guarantee, however,

  * that only one thread at a time will call getVertex since it is a singleton.

  *

  * @param <I> Vertex index value

  * @param <V> Vertex value

  * @param <E> Edge value

  */

 @NotThreadSafe

 public class ByteArrayPartition<I extends WritableComparable,

     V extends Writable, E extends Writable>

     extends BasicPartition<I, V, E>

     implements ReusesObjectsPartition<I, V, E> {

   /**

    * Vertex map for this range (keyed by index).  Note that the byte[] is a

    * serialized vertex with the first four bytes as the length of the vertex

    * to read.

    */

   private ConcurrentMap<I, byte[]> vertexMap;

   /** Representative vertex */

   private Vertex<I, V, E> representativeVertex;

   /** Representative combiner vertex */

   private Vertex<I, V, E> representativeCombinerVertex;

   /** Use unsafe serialization */

   private boolean useUnsafeSerialization;

   /**

    * Constructor for reflection.

    */

   public ByteArrayPartition() { }

   @Override

   public void initialize(int partitionId, Progressable progressable) {

     super.initialize(partitionId, progressable);

     vertexMap = new MapMaker().concurrencyLevel(

         getConf().getNettyServerExecutionConcurrency()).makeMap();

     representativeVertex = getConf().createVertex();

     representativeVertex.initialize(

         getConf().createVertexId(),

         getConf().createVertexValue(),

         getConf().createOutEdges());

     representativeCombinerVertex = getConf().createVertex();

     representativeCombinerVertex.initialize(

         getConf().createVertexId(),

         getConf().createVertexValue(),

         getConf().createOutEdges());

     useUnsafeSerialization = getConf().useUnsafeSerialization();

   }

   @Override

   public Vertex<I, V, E> getVertex(I vertexIndex) {

     byte[] vertexData = vertexMap.get(vertexIndex);

     if (vertexData == null) {

       return null;

     }

     WritableUtils.reinitializeVertexFromByteArray(

         vertexData, representativeVertex, useUnsafeSerialization, getConf());

     return representativeVertex;

   }

   @Override

   public Vertex<I, V, E> putVertex(Vertex<I, V, E> vertex) {

     byte[] vertexData =

         WritableUtils.writeVertexToByteArray(

             vertex, useUnsafeSerialization, getConf());

     byte[] oldVertexBytes = vertexMap.put(vertex.getId(), vertexData);

     if (oldVertexBytes == null) {

       return null;

     } else {

       WritableUtils.reinitializeVertexFromByteArray(oldVertexBytes,

           representativeVertex, useUnsafeSerialization, getConf());

       return representativeVertex;

     }

   }

   @Override

   public Vertex<I, V, E> removeVertex(I vertexIndex) {

     byte[] vertexBytes = vertexMap.remove(vertexIndex);

     if (vertexBytes == null) {

       return null;

     }

     WritableUtils.reinitializeVertexFromByteArray(vertexBytes,

         representativeVertex, useUnsafeSerialization, getConf());

     return representativeVertex;

   }

   @Override

   public void addPartition(Partition<I, V, E> partition) {

     // Only work with other ByteArrayPartition instances

     if (!(partition instanceof ByteArrayPartition)) {

       throw new IllegalStateException("addPartition: Cannot add partition " +

           "of type " + partition.getClass());

     }

     ByteArrayPartition<I, V, E> byteArrayPartition =

         (ByteArrayPartition<I, V, E>) partition;

     for (Map.Entry<I, byte[]> entry :

         byteArrayPartition.vertexMap.entrySet()) {

       byte[] oldVertexBytes =

           vertexMap.putIfAbsent(entry.getKey(), entry.getValue());

       if (oldVertexBytes == null) {

         continue;

       }

       // Note that vertex combining is going to be expensive compared to

       // SimplePartition since here we have to deserialize the vertices,

       // combine them, and then reserialize them.  If the vertex doesn't exist,

       // just add the new vertex as a byte[]

       synchronized (this) {

         // Combine the vertex values

         WritableUtils.reinitializeVertexFromByteArray(oldVertexBytes,

             representativeVertex, useUnsafeSerialization, getConf());

         WritableUtils.reinitializeVertexFromByteArray(entry.getValue(),

             representativeCombinerVertex, useUnsafeSerialization, getConf());

         combine(representativeVertex, representativeCombinerVertex);

       }

     }

   }

   @Override

   public synchronized boolean putOrCombine(Vertex<I, V, E> vertex) {

     // Optimistically try to first put and then combine if this fails

     byte[] vertexData =

         WritableUtils.writeVertexToByteArray(

             vertex, useUnsafeSerialization, getConf());

     byte[] oldVertexBytes = vertexMap.putIfAbsent(vertex.getId(), vertexData);

     if (oldVertexBytes == null) {

       return true;

     }

     WritableUtils.reinitializeVertexFromByteArray(oldVertexBytes,

         representativeVertex, useUnsafeSerialization, getConf());

     combine(representativeVertex, vertex);

     return false;

   }

   /**

    * Combine two vertices together and store the serialized bytes

    * in the vertex map.

    *

    * @param representativeVertex existing vertex

    * @param representativeCombinerVertex new vertex to combine

    */

   private void combine(Vertex<I, V, E> representativeVertex,

       Vertex<I, V, E> representativeCombinerVertex) {

     getVertexValueCombiner().combine(representativeVertex.getValue(),

         representativeCombinerVertex.getValue());

     // Add the edges to the representative vertex

     for (Edge<I, E> edge : representativeCombinerVertex.getEdges()) {

       representativeVertex.addEdge(edge);

     }

     vertexMap.put(representativeCombinerVertex.getId(),

         WritableUtils.writeVertexToByteArray(

             representativeVertex, useUnsafeSerialization, getConf()));

   }

   @Override

   public long getVertexCount() {

     return vertexMap.size();

   }

   @Override

   public long getEdgeCount() {

     long edges = 0;

     for (byte[] vertexBytes : vertexMap.values()) {

       WritableUtils.reinitializeVertexFromByteArray(vertexBytes,

           representativeVertex, useUnsafeSerialization, getConf());

       edges += representativeVertex.getNumEdges();

     }

     return edges;

   }

   @Override

   public void saveVertex(Vertex<I, V, E> vertex) {

     // Reuse the old buffer whenever possible

     byte[] oldVertexData = vertexMap.get(vertex.getId());

     if (oldVertexData != null) {

       vertexMap.put(vertex.getId(),

           WritableUtils.writeVertexToByteArray(

               vertex, oldVertexData, useUnsafeSerialization, getConf()));

     } else {

       vertexMap.put(vertex.getId(),

           WritableUtils.writeVertexToByteArray(

               vertex, useUnsafeSerialization, getConf()));

     }

   }

   @Override

   public void write(DataOutput output) throws IOException {

     super.write(output);

     output.writeInt(vertexMap.size());

     for (Map.Entry<I, byte[]> entry : vertexMap.entrySet()) {

       progress();

       entry.getKey().write(output);

       // Note here that we are writing the size of the vertex data first

       // as it is encoded in the first four bytes of the byte[]

       int vertexDataSize;

       if (useUnsafeSerialization) {

         vertexDataSize = UnsafeByteArrayInputStream.getInt(entry.getValue(),

             0);

       } else {

         vertexDataSize = Ints.fromByteArray(entry.getValue());

       }

       output.writeInt(vertexDataSize);

       output.write(entry.getValue(), 0, vertexDataSize);

     }

   }

   @Override

   public void readFields(DataInput input) throws IOException {

     super.readFields(input);

     int size = input.readInt();

     vertexMap = new MapMaker().concurrencyLevel(

         getConf().getNettyServerExecutionConcurrency()).initialCapacity(

         size).makeMap();

     representativeVertex = getConf().createVertex();

     representativeVertex.initialize(

         getConf().createVertexId(),

         getConf().createVertexValue(),

         getConf().createOutEdges());

     useUnsafeSerialization = getConf().useUnsafeSerialization();

     for (int i = 0; i < size; ++i) {

       progress();

       I vertexId = getConf().createVertexId();

       vertexId.readFields(input);

       int vertexDataSize = input.readInt();

       byte[] vertexData = new byte[vertexDataSize];

       input.readFully(vertexData);

       if (vertexMap.put(vertexId, vertexData) != null) {

         throw new IllegalStateException("readFields: Already saw vertex " +

             vertexId);

       }

     }

   }

   @Override

   public Iterator<Vertex<I, V, E>> iterator() {

     return new RepresentativeVertexIterator();

   }

   /**

    * Iterator that deserializes a vertex from a byte array on the fly, using

    * the same representative vertex object.

    */

   private class RepresentativeVertexIterator implements

       Iterator<Vertex<I, V, E>> {

     /** Iterator to the vertex values */

     private Iterator<byte[]> vertexDataIterator =

         vertexMap.values().iterator();

     @Override

     public boolean hasNext() {

       return vertexDataIterator.hasNext();

     }

     @Override

     public Vertex<I, V, E> next() {

       WritableUtils.reinitializeVertexFromByteArray(

           vertexDataIterator.next(), representativeVertex,

           useUnsafeSerialization, getConf());

       return representativeVertex;

     }

     @Override

     public void remove() {

       throw new IllegalAccessError("remove: This method is not supported.");

     }

   }

 }