package edu.washington.escience.myria.operator.network.partition;

import java.io.Serializable;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonSubTypes.Type;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
import com.google.common.base.Objects;
import com.google.common.base.Preconditions;

import edu.washington.escience.myria.storage.TupleBatch;

/**
 * The ShuffleProducer class uses an instance of the PartitionFunction class to decide which worker a tuple should be
 * routed to. Typically, the ShuffleProducer class invokes {@link partition(Tuple, Schema) partition} on every tuple it
 * generates.
 */
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
@JsonSubTypes({
    @Type(value = RoundRobinPartitionFunction.class, name = "RoundRobin"),
    @Type(value = SingleFieldHashPartitionFunction.class, name = "SingleFieldHash"),
    @Type(value = MultiFieldHashPartitionFunction.class, name = "MultiFieldHash"),
    @Type(value = WholeTupleHashPartitionFunction.class, name = "WholeTupleHash") })
public abstract class PartitionFunction implements Serializable {

  /** Required for Java serialization. */
  private static final long serialVersionUID = 1L;

  /** Invalid number of partitions. */
  private static final int INVALID_NUM_PARTITIONS = -1;

  /**
   * The number of partitions into which input tuples can be divided.
   */
  @JsonProperty
  private int numPartitions = INVALID_NUM_PARTITIONS;

  /**
   * @param numPartitions the number of partitions into which input tuples can be divided. Note that this is a
   *          {@link Integer} not an {@link int} so that it can properly handle <code>null</code> values, e.g., in JSON
   *          deserialization.
   */
  public PartitionFunction(@Nullable final Integer numPartitions) {
    Preconditions.checkArgument((numPartitions == null) || (numPartitions > 0),
        "numPartitions argument must be null or > 0");
    this.numPartitions = Objects.firstNonNull(numPartitions, INVALID_NUM_PARTITIONS);
  }

  /**
   * @return the number of partitions.
   */
  public final int numPartition() {
    if (numPartitions == INVALID_NUM_PARTITIONS) {
      throw new IllegalStateException("numPartitions has not been set");
    }
    return numPartitions;
  }

  /**
   * Given that the TupleBatches expose only the valid tuples, partition functions using TB.get** methods should be of
   * little overhead comparing with direct Column access.
   * 
   * @param data the data to be partitioned.
   * 
   * @return an int[] of length specified by <code>data.{@link TupleBatch#numTuples}</code>, specifying which partition
   *         every tuple should be sent to.
   * 
   */
  public abstract int[] partition(@Nonnull final TupleBatch data);

  /**
   * Set the number of output partitions.
   * 
   * @param numPartitions the number of output partitions. Must be greater than 0.
   */
  public final void setNumPartitions(final int numPartitions) {
    Preconditions.checkArgument(numPartitions > 0, "numPartitions must be > 0");
    this.numPartitions = numPartitions;
  }
}

@@ -383,47 +383,8 @@
 fo;%0A
-import com.google.common.base.Objects;%0A
 impo
@@ -424,16 +424,16 @@
 itions;%0A
+
 %0Aimport 
@@ -1349,104 +1349,8 @@
 L;%0A%0A
-  /** Invalid number of partitions. */%0A  private static final int INVALID_NUM_PARTITIONS = -1;%0A%0A
   /*
@@ -1453,19 +1453,23 @@
 private 
-i
+I
 nt
+eger
  numPart
@@ -1477,38 +1477,20 @@
 tions = 
-INVALID_NUM_PARTITIONS
+null
 ;%0A%0A  /**
@@ -1990,67 +1990,21 @@
 s = 
-Objects.firstNonNull(numPartitions, INVALID_NUM_PARTITIONS)
+numPartitions
 ;%0A  
@@ -2101,11 +2101,32 @@
     
-if 
+Preconditions.checkState
 (num
@@ -2140,75 +2140,17 @@
 ons 
-== INVALID_NUM_PARTITIONS) %7B%0A      throw new IllegalStateException(
+!= null, 
 %22num
@@ -2184,14 +2184,8 @@
 %22);%0A
-    %7D%0A
     

fix deserializing uninstantiated PFs