@Override
public int getPartition(PigNullableWritable wrappedKey, Writable value, int numPartitions) {
// for streaming tables, return the partition index blindly
if (wrappedKey instanceof NullablePartitionWritable
&& (((NullablePartitionWritable) wrappedKey).getPartition()) != -1) {
return ((NullablePartitionWritable) wrappedKey).getPartition();
}
// for partition table, compute the index based on the sampler output
Pair<Integer, Integer> indexes;
Integer curIndex = -1;
Tuple keyTuple = TupleFactory.getInstance().newTuple(1);
// extract the key from nullablepartitionwritable
PigNullableWritable key = ((NullablePartitionWritable) wrappedKey).getKey();
try {
keyTuple.set(0, key.getValueAsPigType());
} catch (ExecException e) {
return -1;
}
// if the key is not null and key
if (key instanceof NullableTuple && key.getValueAsPigType() != null) {
keyTuple = (Tuple) key.getValueAsPigType();
}
// if the partition file is empty, use numPartitions
totalReducers = (totalReducers > 0) ? totalReducers : numPartitions;
indexes = reducerMap.get(keyTuple);
// if the reducerMap does not contain the key, do the default hash based partitioning
if (indexes == null) {
return (Math.abs(keyTuple.hashCode() % totalReducers));
}
if (currentIndexMap.containsKey(keyTuple)) {
curIndex = currentIndexMap.get(keyTuple);
}
if (curIndex >= (indexes.first + indexes.second) || curIndex == -1) {
curIndex = indexes.first;
} else {
curIndex++;
}
// set it in the map
currentIndexMap.put(keyTuple, curIndex);
return (curIndex % totalReducers);
}
