/**
* Copyright 2011 LiveRamp
*
* Licensed 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 com.liveramp.hank.hadoop;
import com.liveramp.hank.config.CoordinatorConfigurator;
import com.liveramp.hank.config.InvalidConfigurationException;
import com.liveramp.hank.config.yaml.YamlClientConfigurator;
import com.liveramp.hank.coordinator.Coordinator;
import com.liveramp.hank.coordinator.CoordinatorFactory;
import com.liveramp.hank.coordinator.Domain;
import com.liveramp.hank.coordinator.DomainVersion;
import com.liveramp.hank.coordinator.mock.MockCoordinator;
import com.liveramp.hank.coordinator.mock.MockDomain;
import com.liveramp.hank.coordinator.mock.MockDomainVersion;
import com.liveramp.hank.partitioner.Partitioner;
import com.liveramp.hank.storage.PartitionRemoteFileOps;
import com.liveramp.hank.storage.Writer;
import com.liveramp.hank.storage.incremental.IncrementalDomainVersionProperties;
import com.liveramp.hank.storage.mock.MockStorageEngine;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.util.Map;
// Integer String key storage engine.
// Store records (key, value) where key is an Integer's String representation
// The partition is key % numPartitions
// Format is "key value\n"
public class IntStringKeyStorageEngineCoordinator extends MockCoordinator {
private final int numPartitions;
public IntStringKeyStorageEngineCoordinator(int numPartitions) {
this.numPartitions = numPartitions;
}
private static class IntStringKeyWriter implements Writer {
protected final OutputStream outputStream;
IntStringKeyWriter(DomainVersion domainVersion,
PartitionRemoteFileOps partitionRemoteFileOps,
int partNum) throws IOException {
IncrementalDomainVersionProperties domainVersionProperties =
(IncrementalDomainVersionProperties) domainVersion.getProperties();
if (domainVersionProperties == null) {
throw new RuntimeException("IntStringKeyWriter needs a non null DomainVersionProperties");
}
this.outputStream = partitionRemoteFileOps.getOutputStream(Integer.toString(domainVersion.getVersionNumber()) + "." + (domainVersionProperties.isBase() ? "base" : "nobase"));
}
public void write(ByteBuffer key, ByteBuffer value) throws IOException {
outputStream.write(key.array(), key.position(), key.remaining());
outputStream.write(" ".getBytes());
outputStream.write(value.array(), value.position(), value.remaining());
outputStream.write("\n".getBytes());
}
public void close() throws IOException {
outputStream.close();
}
@Override
public long getNumBytesWritten() {
return 0;
}
@Override
public long getNumRecordsWritten() {
return 0;
}
}
private static class IntStringKeyModPartitioner implements Partitioner {
public int partition(ByteBuffer key, int numPartitions) {
String keyString = new String(key.array(), key.position(), key.remaining());
Integer keyInteger = Integer.valueOf(keyString);
return keyInteger % numPartitions;
}
}
private static class IntStringKeyStorageEngine extends MockStorageEngine {
@Override
public Writer getWriter(DomainVersion domainVersion,
PartitionRemoteFileOps partitionRemoteFileOps,
int partitionNumber) throws IOException {
return new IntStringKeyWriter(domainVersion, partitionRemoteFileOps, partitionNumber);
}
@Override
public ByteBuffer getComparableKey(ByteBuffer key) {
return key;
}
}
public static class Factory implements CoordinatorFactory {
public Coordinator getCoordinator(Map<String, Object> options) {
Integer numPartitions = (Integer) options.get("num_partitions");
if (numPartitions != null) {
return new IntStringKeyStorageEngineCoordinator(numPartitions);
} else {
return new IntStringKeyStorageEngineCoordinator(1);
}
}
}
@Override
public Domain getDomain(String domainName) {
return new MockDomain(domainName,
1,
numPartitions,
new IntStringKeyModPartitioner(),
new IntStringKeyStorageEngine(),
null,
new MockDomainVersion(0, null, new IncrementalDomainVersionProperties.Base()));
}
static public CoordinatorConfigurator getConfigurator(int numPartitions) {
YamlClientConfigurator configurator = new YamlClientConfigurator();
try {
configurator.loadFromYaml("coordinator:\n factory: com.liveramp.hank.hadoop.IntStringKeyStorageEngineCoordinator$Factory\n options:\n num_partitions: " + numPartitions + "\n");
} catch (InvalidConfigurationException e) {
throw new RuntimeException(e);
}
return configurator;
}
}