giraph-block-app/src/main/java/org/apache/giraph/block_app/library/Pieces.java

/*
 * 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.block_app.library;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import org.apache.giraph.block_app.framework.api.BlockMasterApi;
import org.apache.giraph.block_app.framework.api.BlockWorkerReceiveApi;
import org.apache.giraph.block_app.framework.api.BlockWorkerSendApi;
import org.apache.giraph.block_app.framework.api.CreateReducersApi;
import org.apache.giraph.block_app.framework.piece.Piece;
import org.apache.giraph.block_app.framework.piece.global_comm.ReducerAndBroadcastWrapperHandle;
import org.apache.giraph.block_app.framework.piece.global_comm.ReducerHandle;
import org.apache.giraph.block_app.framework.piece.global_comm.array.BroadcastArrayHandle;
import org.apache.giraph.block_app.framework.piece.interfaces.VertexReceiver;
import org.apache.giraph.block_app.framework.piece.interfaces.VertexSender;
import org.apache.giraph.block_app.library.internal.SendMessagePiece;
import org.apache.giraph.block_app.library.internal.SendMessageWithCombinerPiece;
import org.apache.giraph.block_app.reducers.array.ArrayOfHandles;
import org.apache.giraph.combiner.MessageCombiner;
import org.apache.giraph.function.Consumer;
import org.apache.giraph.function.PairConsumer;
import org.apache.giraph.function.Supplier;
import org.apache.giraph.function.vertex.ConsumerWithVertex;
import org.apache.giraph.function.vertex.SupplierFromVertex;
import org.apache.giraph.graph.Vertex;
import org.apache.giraph.reducers.ReduceOperation;
import org.apache.giraph.reducers.impl.SumReduce;
import org.apache.giraph.types.NoMessage;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableComparable;
import org.apache.log4j.Logger;

/**
 * Utility class for creating common Pieces and computations for processing
 * graphs.
 */
public class Pieces {
  private static final Logger LOG = Logger.getLogger(Pieces.class);

  private Pieces() { }

  /**
   * For each vertex execute given process function.
   * Computation is happening in send phase of the returned Piece.
   */
  public static
  <I extends WritableComparable, V extends Writable, E extends Writable>
  Piece<I, V, E, NoMessage, Object> forAllVertices(
      final String pieceName, final Consumer<Vertex<I, V, E>> process) {
    return new Piece<I, V, E, NoMessage, Object>() {
      @Override
      public VertexSender<I, V, E> getVertexSender(
          BlockWorkerSendApi<I, V, E, NoMessage> workerApi,
          Object executionStage) {
        return new InnerVertexSender() {
          @Override
          public void vertexSend(Vertex<I, V, E> vertex) {
            process.apply(vertex);
          }
        };
      }

      @Override
      public String toString() {
        return pieceName;
      }
    };
  }

  /**
   * Execute given function on master.
   */
  public static
  Piece<WritableComparable, Writable,  Writable, NoMessage,
    Object> masterCompute(
      final String pieceName, final Consumer<BlockMasterApi> process) {
    return new Piece<WritableComparable, Writable,  Writable, NoMessage,
        Object>() {
      @Override
      public void masterCompute(
          BlockMasterApi masterApi, Object executionStage) {
        process.apply(masterApi);
      }
    };
  }

  /**
   * For each vertex execute given process function.
   * Computation is happening in the receive phase of the returned Piece.
   * This function should be used if you need returned Piece to interact with
   * subsequent Piece, as that requires passed function to be executed
   * during receive phase,
   */
  public static
  <I extends WritableComparable, V extends Writable, E extends Writable>
  Piece<I, V, E, NoMessage, Object> forAllVerticesOnReceive(
      final String pieceName, final Consumer<Vertex<I, V, E>> process) {
    return new Piece<I, V, E, NoMessage, Object>() {
      @Override
      public VertexReceiver<I, V, E, NoMessage> getVertexReceiver(
          BlockWorkerReceiveApi<I> workerApi, Object executionStage) {
        return new InnerVertexReceiver() {
          @Override
          public void vertexReceive(
              Vertex<I, V, E> vertex, Iterable<NoMessage> messages) {
            process.apply(vertex);
          }
        };
      }

      @Override
      public String toString() {
        return pieceName;
      }
    };
  }

  /**
   * Creates Piece which removes vertices for which supplier returns true.
   */
  public static
  <I extends WritableComparable, V extends Writable, E extends Writable>
  Piece<I, V, E, NoMessage, Object> removeVertices(
      final String pieceName,
      final SupplierFromVertex<I, V, E, Boolean> shouldRemoveVertex) {
    return new Piece<I, V, E, NoMessage, Object>() {
      private ReducerHandle<LongWritable, LongWritable> countRemovedAgg;

      @Override
      public void registerReducers(
          CreateReducersApi reduceApi, Object executionStage) {
        countRemovedAgg = reduceApi.createLocalReducer(SumReduce.LONG);
      }

      @Override
      public VertexSender<I, V, E> getVertexSender(
          final BlockWorkerSendApi<I, V, E, NoMessage> workerApi,
          Object executionStage) {
        return new InnerVertexSender() {
          @Override
          public void vertexSend(Vertex<I, V, E> vertex) {
            if (shouldRemoveVertex.get(vertex)) {
              workerApi.removeVertexRequest(vertex.getId());
              reduceLong(countRemovedAgg, 1);
            }
          }
        };
      }

      @Override
      public void masterCompute(BlockMasterApi master, Object executionStage) {
        LOG.info("Removed " + countRemovedAgg.getReducedValue(master) +
            " vertices from the graph, during stage " + executionStage);
      }

      @Override
      public String toString() {
        return pieceName;
      }
    };
  }

  /**
   * Creates single reducer piece - given reduce class, supplier of values on
   * worker, reduces and passes the result to given consumer on master.
   *
   * @param <S> Single value type, objects passed on workers
   * @param <R> Reduced value type
   * @param <I> Vertex id type
   * @param <V> Vertex value type
   * @param <E> Edge value type
   */
  public static
  <S, R extends Writable, I extends WritableComparable, V extends Writable,
  E extends Writable>
  Piece<I, V, E, NoMessage, Object> reduce(
      String name,
      ReduceOperation<S, R> reduceOp,
      SupplierFromVertex<I, V, E, S> valueSupplier,
      final Consumer<R> reducedValueConsumer) {
    return reduceWithMaster(
        name, reduceOp, valueSupplier,
        new PairConsumer<R, BlockMasterApi>() {
          @Override
          public void apply(R input, BlockMasterApi master) {
            reducedValueConsumer.apply(input);
          }
        });
  }

  /**
   * Creates single reducer piece - given reduce class, supplier of values on
   * worker, reduces and passes the result to given consumer on master.
   *
   * @param <S> Single value type, objects passed on workers
   * @param <R> Reduced value type
   * @param <I> Vertex id type
   * @param <V> Vertex value type
   * @param <E> Edge value type
   */
  public static
  <S, R extends Writable, I extends WritableComparable, V extends Writable,
  E extends Writable>
  Piece<I, V, E, NoMessage, Object> reduceWithMaster(
      final String name,
      final ReduceOperation<S, R> reduceOp,
      final SupplierFromVertex<I, V, E, S> valueSupplier,
      final PairConsumer<R, BlockMasterApi> reducedValueConsumer) {
    return new Piece<I, V, E, NoMessage, Object>() {
      private ReducerHandle<S, R> handle;

      @Override
      public void registerReducers(
          CreateReducersApi reduceApi, Object executionStage) {
        handle = reduceApi.createLocalReducer(reduceOp);
      }

      @Override
      public VertexSender<I, V, E> getVertexSender(
          BlockWorkerSendApi<I, V, E, NoMessage> workerApi,
          Object executionStage) {
        return new InnerVertexSender() {
          @Override
          public void vertexSend(Vertex<I, V, E> vertex) {
            handle.reduce(valueSupplier.get(vertex));
          }
        };
      }

      @Override
      public void masterCompute(BlockMasterApi master, Object executionStage) {
        reducedValueConsumer.apply(handle.getReducedValue(master), master);
      }

      @Override
      public String toString() {
        return name;
      }
    };
  }

  /**
   * Creates single reducer and broadcast piece - given reduce class, supplier
   * of values on worker, reduces and broadcasts the value, passing it to the
   * consumer on worker for each vertex.
   *
   * @param <S> Single value type, objects passed on workers
   * @param <R> Reduced value type
   * @param <I> Vertex id type
   * @param <V> Vertex value type
   * @param <E> Edge value type
   */
  public static
  <S, R extends Writable, I extends WritableComparable, V extends Writable,
  E extends Writable>
  Piece<I, V, E, NoMessage, Object> reduceAndBroadcast(
      final String name,
      final ReduceOperation<S, R> reduceOp,
      final SupplierFromVertex<I, V, E, S> valueSupplier,
      final ConsumerWithVertex<I, V, E, R> reducedValueConsumer) {
    return new Piece<I, V, E, NoMessage, Object>() {
      private final ReducerAndBroadcastWrapperHandle<S, R> handle =
          new ReducerAndBroadcastWrapperHandle<>();

      @Override
      public void registerReducers(
          CreateReducersApi reduceApi, Object executionStage) {
        handle.registeredReducer(reduceApi.createLocalReducer(reduceOp));
      }

      @Override
      public VertexSender<I, V, E> getVertexSender(
          BlockWorkerSendApi<I, V, E, NoMessage> workerApi,
          Object executionStage) {
        return new InnerVertexSender() {
          @Override
          public void vertexSend(Vertex<I, V, E> vertex) {
            handle.reduce(valueSupplier.get(vertex));
          }
        };
      }

      @Override
      public void masterCompute(BlockMasterApi master, Object executionStage) {
        handle.broadcastValue(master);
      }

      @Override
      public VertexReceiver<I, V, E, NoMessage> getVertexReceiver(
          BlockWorkerReceiveApi<I> workerApi, Object executionStage) {
        final R value = handle.getBroadcast(workerApi);
        return new InnerVertexReceiver() {
          @Override
          public void vertexReceive(
              Vertex<I, V, E> vertex, Iterable<NoMessage> messages) {
            reducedValueConsumer.apply(vertex, value);
          }
        };
      }

      @Override
      public String toString() {
        return name;
      }
    };
  }

  /**
   * Like reduceAndBroadcast, but uses array of handles for reducers and
   * broadcasts, to make it feasible and performant when values are large.
   * Each supplied value to reduce will be reduced in the handle defined by
   * handleHashSupplier%numHandles
   *
   * @param <S> Single value type, objects passed on workers
   * @param <R> Reduced value type
   * @param <I> Vertex id type
   * @param <V> Vertex value type
   * @param <E> Edge value type
   */
  public static
  <S, R extends Writable, I extends WritableComparable, V extends Writable,
      E extends Writable>
  Piece<I, V, E, NoMessage, Object> reduceAndBroadcastWithArrayOfHandles(
      final String name,
      final int numHandles,
      final Supplier<ReduceOperation<S, R>> reduceOp,
      final SupplierFromVertex<I, V, E, Long> handleHashSupplier,
      final SupplierFromVertex<I, V, E, S> valueSupplier,
      final ConsumerWithVertex<I, V, E, R> reducedValueConsumer) {
    return new Piece<I, V, E, NoMessage, Object>() {
      protected ArrayOfHandles.ArrayOfReducers<S, R> reducers;
      protected BroadcastArrayHandle<R> broadcasts;

      private int getHandleIndex(Vertex<I, V, E> vertex) {
        return (int) Math.abs(handleHashSupplier.get(vertex) % numHandles);
      }

      @Override
      public void registerReducers(
          final CreateReducersApi reduceApi, Object executionStage) {
        reducers = new ArrayOfHandles.ArrayOfReducers<>(
            numHandles,
            new Supplier<ReducerHandle<S, R>>() {
              @Override
              public ReducerHandle<S, R> get() {
                return reduceApi.createLocalReducer(reduceOp.get());
              }
            });
      }

      @Override
      public VertexSender<I, V, E> getVertexSender(
          BlockWorkerSendApi<I, V, E, NoMessage> workerApi,
          Object executionStage) {
        return new InnerVertexSender() {
          @Override
          public void vertexSend(Vertex<I, V, E> vertex) {
            reducers.get(getHandleIndex(vertex)).reduce(
                valueSupplier.get(vertex));
          }
        };
      }

      @Override
      public void masterCompute(BlockMasterApi master, Object executionStage) {
        broadcasts = reducers.broadcastValue(master);
      }

      @Override
      public VertexReceiver<I, V, E, NoMessage> getVertexReceiver(
          BlockWorkerReceiveApi<I> workerApi, Object executionStage) {
        final List<R> values = new ArrayList<>();
        for (int i = 0; i < numHandles; i++) {
          values.add(broadcasts.get(i).getBroadcast(workerApi));
        }
        return new InnerVertexReceiver() {
          @Override
          public void vertexReceive(
              Vertex<I, V, E> vertex, Iterable<NoMessage> messages) {
            reducedValueConsumer.apply(
                vertex, values.get(getHandleIndex(vertex)));
          }
        };
      }

      @Override
      public String toString() {
        return name;
      }
    };
  }

  /**
   * Creates Piece that for each vertex, sends message provided by
   * messageSupplier to all targets provided by targetsSupplier.
   * Received messages are then passed to and processed by provided
   * messagesConsumer.
   *
   * If messageSupplier or targetsSupplier returns null, current vertex
   * is not going to send any messages.
   */
  public static
  <I extends WritableComparable, V extends Writable, E extends Writable,
  M extends Writable>
  SendMessagePiece<I, V, E, M> sendMessage(
      String name,
      Class<M> messageClass,
      SupplierFromVertex<I, V, E, M> messageSupplier,
      SupplierFromVertex<I, V, E, Iterator<I>> targetsSupplier,
      ConsumerWithVertex<I, V, E, Iterable<M>> messagesConsumer) {
    return new SendMessagePiece<>(
        name, messageClass, messageSupplier, targetsSupplier, messagesConsumer);
  }

  /**
   * Creates Piece that for each vertex, sends message provided by
   * messageSupplier to all neighbors of current vertex.
   * Received messages are then passed to and processed by provided
   * messagesConsumer.
   *
   * If messageSupplier returns null, current vertex
   * is not going to send any messages.
   */
  public static
  <I extends WritableComparable, V extends Writable, E extends Writable,
  M extends Writable>
  SendMessagePiece<I, V, E, M> sendMessageToNeighbors(
      String name,
      Class<M> messageClass,
      SupplierFromVertex<I, V, E, M> messageSupplier,
      ConsumerWithVertex<I, V, E, Iterable<M>> messagesConsumer) {
    return sendMessage(
        name, messageClass, messageSupplier,
        VertexSuppliers.<I, V, E>vertexNeighborsSupplier(),
        messagesConsumer);
  }

  /**
   * Creates Piece that for each vertex, sends message provided by
   * messageSupplier to all targets provided by targetsSupplier,
   * and uses given messageCombiner to combine messages together.
   * Received combined message is then passed to and processed by provided
   * messageConsumer. (null is passed to it, if vertex received no messages)
   *
   * If messageSupplier or targetsSupplier returns null, current vertex
   * is not going to send any messages.
   */
  public static
  <I extends WritableComparable, V extends Writable, E extends Writable,
  M extends Writable>
  SendMessageWithCombinerPiece<I, V, E, M> sendMessage(
      String name,
      MessageCombiner<? super I, M> messageCombiner,
      SupplierFromVertex<I, V, E, M> messageSupplier,
      SupplierFromVertex<I, V, E, Iterator<I>> targetsSupplier,
      ConsumerWithVertex<I, V, E, M> messagesConsumer) {
    return new SendMessageWithCombinerPiece<>(
        name, messageCombiner,
        messageSupplier, targetsSupplier, messagesConsumer);
  }

  /**
   * Creates Piece that for each vertex, sends message provided by
   * messageSupplier to all neighbors of current vertex,
   * and uses given messageCombiner to combine messages together.
   * Received combined message is then passed to and processed by provided
   * messageConsumer. (null is passed to it, if vertex received no messages)
   *
   * If messageSupplier returns null, current vertex
   * is not going to send any messages.
   */
  public static
  <I extends WritableComparable, V extends Writable, E extends Writable,
  M extends Writable>
  SendMessageWithCombinerPiece<I, V, E, M> sendMessageToNeighbors(
      String name,
      MessageCombiner<? super I, M> messageCombiner,
      SupplierFromVertex<I, V, E, M> messageSupplier,
      ConsumerWithVertex<I, V, E, M> messagesConsumer) {
    return sendMessage(
        name, messageCombiner, messageSupplier,
        VertexSuppliers.<I, V, E>vertexNeighborsSupplier(),
        messagesConsumer);
  }
}