@SuppressWarnings("unchecked")
public void insertOrReplace(BT record) {
int hashCode = hash(buildComparator.hash(record));
int index = hashIndex(hashCode, data.length);
buildComparator.setReference(record);
HashEntry entry = data[index];
while (entry != null) {
if (entryHashCode(entry) == hashCode
&& buildComparator.equalToReference((BT) entry.getValue())) {
entry.setValue(record);
return;
}
entry = entryNext(entry);
}
addMapping(index, hashCode, null, record);
}
@SuppressWarnings("unchecked")
public BT lookupMatch(PT record) {
int hashCode = hash(probeComparator.hash(record));
int index = hashIndex(hashCode, data.length);
pairComparator.setReference(record);
HashEntry entry = data[index];
while (entry != null) {
if (entryHashCode(entry) == hashCode
&& pairComparator.equalToReference((BT) entry.getValue())) {
return (BT) entry.getValue();
}
entry = entryNext(entry);
}
return null;
}
@SuppressWarnings("unchecked")
public OutputEmitter(
ShipStrategyType strategy,
int indexInSubtaskGroup,
TypeComparator<T> comparator,
Partitioner<?> partitioner,
DataDistribution distribution) {
if (strategy == null) {
throw new NullPointerException();
}
this.strategy = strategy;
this.nextChannelToSendTo = indexInSubtaskGroup;
this.comparator = comparator;
this.partitioner = (Partitioner<Object>) partitioner;
this.distribution = distribution;
switch (strategy) {
case PARTITION_CUSTOM:
extractedKeys = new Object[1];
case FORWARD:
case PARTITION_HASH:
case PARTITION_RANDOM:
case PARTITION_FORCED_REBALANCE:
channels = new int[1];
break;
case PARTITION_RANGE:
channels = new int[1];
if (comparator != null) {
this.flatComparators = comparator.getFlatComparators();
this.keys = new Object[flatComparators.length];
}
break;
case BROADCAST:
break;
default:
throw new IllegalArgumentException(
"Invalid shipping strategy for OutputEmitter: " + strategy.name());
}
if (strategy == ShipStrategyType.PARTITION_CUSTOM && partitioner == null) {
throw new NullPointerException(
"Partitioner must not be null when the ship strategy is set to custom partitioning.");
}
}
private int[] customPartition(T record, int numberOfChannels) {
if (channels == null) {
channels = new int[1];
extractedKeys = new Object[1];
}
try {
if (comparator.extractKeys(record, extractedKeys, 0) == 1) {
final Object key = extractedKeys[0];
channels[0] = partitioner.partition(key, numberOfChannels);
return channels;
} else {
throw new RuntimeException(
"Inconsistency in the key comparator - comparator extracted more than one field.");
}
} catch (Throwable t) {
throw new RuntimeException("Error while calling custom partitioner.", t);
}
}
@SuppressWarnings("unchecked")
public long addRecord(T record) throws IOException {
if (recordsOutFile == null) {
if (closed) {
throw new IllegalStateException("The large record handler has been closed.");
}
if (recordsReader != null) {
throw new IllegalStateException("The handler has already switched to sorting.");
}
LOG.debug("Initializing the large record spilling...");
// initialize the utilities
{
final TypeComparator<?>[] keyComps = comparator.getFlatComparators();
numKeyFields = keyComps.length;
Object[] keyHolder = new Object[numKeyFields];
comparator.extractKeys(record, keyHolder, 0);
TypeSerializer<?>[] keySers = new TypeSerializer<?>[numKeyFields];
TypeSerializer<?>[] tupleSers = new TypeSerializer<?>[numKeyFields + 1];
int[] keyPos = new int[numKeyFields];
for (int i = 0; i < numKeyFields; i++) {
keyPos[i] = i;
keySers[i] = createSerializer(keyHolder[i], i);
tupleSers[i] = keySers[i];
}
// add the long serializer for the offset
tupleSers[numKeyFields] = LongSerializer.INSTANCE;
keySerializer =
new TupleSerializer<Tuple>(
(Class<Tuple>) Tuple.getTupleClass(numKeyFields + 1), tupleSers);
keyComparator = new TupleComparator<Tuple>(keyPos, keyComps, keySers);
keySerializerFactory =
new RuntimeSerializerFactory<Tuple>(keySerializer, keySerializer.getTupleClass());
keyTuple = keySerializer.createInstance();
}
// initialize the spilling
final int totalNumSegments = memory.size();
final int segmentsForKeys =
(totalNumSegments >= 2 * MAX_SEGMENTS_FOR_KEY_SPILLING)
? MAX_SEGMENTS_FOR_KEY_SPILLING
: Math.max(
MIN_SEGMENTS_FOR_KEY_SPILLING, totalNumSegments - MAX_SEGMENTS_FOR_KEY_SPILLING);
List<MemorySegment> recordsMemory = new ArrayList<MemorySegment>();
List<MemorySegment> keysMemory = new ArrayList<MemorySegment>();
for (int i = 0; i < segmentsForKeys; i++) {
keysMemory.add(memory.get(i));
}
for (int i = segmentsForKeys; i < totalNumSegments; i++) {
recordsMemory.add(memory.get(i));
}
recordsChannel = ioManager.createChannel();
keysChannel = ioManager.createChannel();
recordsOutFile =
new FileChannelOutputView(
ioManager.createBlockChannelWriter(recordsChannel),
memManager,
recordsMemory,
memManager.getPageSize());
keysOutFile =
new FileChannelOutputView(
ioManager.createBlockChannelWriter(keysChannel),
memManager,
keysMemory,
memManager.getPageSize());
}
final long offset = recordsOutFile.getWriteOffset();
if (offset < 0) {
throw new RuntimeException("wrong offset");
}
Object[] keyHolder = new Object[numKeyFields];
comparator.extractKeys(record, keyHolder, 0);
for (int i = 0; i < numKeyFields; i++) {
keyTuple.setField(keyHolder[i], i);
}
keyTuple.setField(offset, numKeyFields);
keySerializer.serialize(keyTuple, keysOutFile);
serializer.serialize(record, recordsOutFile);
recordCounter++;
return offset;
}
