@Override
  public void prepare() throws Exception {
    final TaskConfig config = this.taskContext.getTaskConfig();
    if (config.getDriverStrategy() != DriverStrategy.ALL_REDUCE) {
      throw new Exception(
          "Unrecognized driver strategy for AllReduce driver: "
              + config.getDriverStrategy().name());
    }

    TypeSerializerFactory<T> serializerFactory = this.taskContext.getInputSerializer(0);
    this.serializer = serializerFactory.getSerializer();
    this.input = this.taskContext.getInput(0);
  }
  @Override
  public void prepare() throws Exception {
    final TaskConfig config = this.taskContext.getTaskConfig();
    if (config.getDriverStrategy() != DriverStrategy.CO_GROUP) {
      throw new Exception(
          "Unrecognized driver strategy for CoGoup driver: " + config.getDriverStrategy().name());
    }

    final MutableObjectIterator<IT1> in1 = this.taskContext.getInput(0);
    final MutableObjectIterator<IT2> in2 = this.taskContext.getInput(1);

    // get the key positions and types
    final TypeSerializer<IT1> serializer1 =
        this.taskContext.<IT1>getInputSerializer(0).getSerializer();
    final TypeSerializer<IT2> serializer2 =
        this.taskContext.<IT2>getInputSerializer(1).getSerializer();
    final TypeComparator<IT1> groupComparator1 = this.taskContext.getInputComparator(0);
    final TypeComparator<IT2> groupComparator2 = this.taskContext.getInputComparator(1);

    final TypePairComparatorFactory<IT1, IT2> pairComparatorFactory =
        config.getPairComparatorFactory(this.taskContext.getUserCodeClassLoader());
    if (pairComparatorFactory == null) {
      throw new Exception("Missing pair comparator factory for CoGroup driver");
    }

    // create CoGropuTaskIterator according to provided local strategy.
    this.coGroupIterator =
        new SortMergeCoGroupIterator<IT1, IT2>(
            in1,
            in2,
            serializer1,
            groupComparator1,
            serializer2,
            groupComparator2,
            pairComparatorFactory.createComparator12(groupComparator1, groupComparator2));

    // open CoGroupTaskIterator - this triggers the sorting and blocks until the iterator is ready
    this.coGroupIterator.open();

    if (LOG.isDebugEnabled()) {
      LOG.debug(this.taskContext.formatLogString("CoGroup task iterator ready."));
    }
  }