package org.apache.hadoop.hive.cassandra.input;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import org.apache.cassandra.exceptions.ConfigurationException;
import org.apache.cassandra.db.marshal.AbstractType;
import org.apache.cassandra.db.marshal.BytesType;
import org.apache.cassandra.db.marshal.TypeParser;
import org.apache.cassandra.hadoop.ColumnFamilyInputFormat;
import org.apache.cassandra.hadoop.ColumnFamilyRecordReader;
import org.apache.cassandra.hadoop.ColumnFamilySplit;
import org.apache.cassandra.hadoop.ConfigHelper;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.thrift.SliceRange;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.cassandra.serde.AbstractColumnSerDe;
import org.apache.hadoop.hive.serde2.ColumnProjectionUtils;
import org.apache.hadoop.io.BytesWritable;
import org.apache.hadoop.io.MapWritable;
import org.apache.hadoop.mapred.InputSplit;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.mapred.RecordReader;
import org.apache.hadoop.mapred.Reporter;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.TaskAttemptID;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
@SuppressWarnings("deprecation")
public class HiveCassandraStandardColumnInputFormat extends InputFormat<BytesWritable, MapWritable>
implements org.apache.hadoop.mapred.InputFormat<BytesWritable, MapWritable> {
static final Log LOG = LogFactory.getLog(HiveCassandraStandardColumnInputFormat.class);
private boolean isTransposed;
private final ColumnFamilyInputFormat cfif = new ColumnFamilyInputFormat();
@Override
public RecordReader<BytesWritable, MapWritable> getRecordReader(InputSplit split,
JobConf jobConf, final Reporter reporter) throws IOException {
HiveCassandraStandardSplit cassandraSplit = (HiveCassandraStandardSplit) split;
List<String> columns = AbstractColumnSerDe.parseColumnMapping(cassandraSplit.getColumnMapping());
isTransposed = AbstractColumnSerDe.isTransposed(columns);
List<Integer> readColIDs = ColumnProjectionUtils.getReadColumnIDs(jobConf);
if (columns.size() < readColIDs.size()) {
throw new IOException("Cannot read more columns than the given table contains.");
}
org.apache.cassandra.hadoop.ColumnFamilySplit cfSplit = cassandraSplit.getSplit();
Job job = new Job(jobConf);
TaskAttemptContext tac = new TaskAttemptContext(job.getConfiguration(), new TaskAttemptID()) {
@Override
public void progress() {
reporter.progress();
}
};
SlicePredicate predicate = new SlicePredicate();
if (isTransposed || readColIDs.size() == columns.size() || readColIDs.size() == 0) {
SliceRange range = new SliceRange();
AbstractType comparator = BytesType.instance;
String comparatorType = jobConf.get(AbstractColumnSerDe.CASSANDRA_SLICE_PREDICATE_RANGE_COMPARATOR);
if (comparatorType != null && !comparatorType.equals("")) {
try {
comparator = TypeParser.parse(comparatorType);
} catch (Exception ex) {
throw new IOException("Comparator class not found.");
}
}
String sliceStart = jobConf.get(AbstractColumnSerDe.CASSANDRA_SLICE_PREDICATE_RANGE_START);
String sliceEnd = jobConf.get(AbstractColumnSerDe.CASSANDRA_SLICE_PREDICATE_RANGE_FINISH);
String reversed = jobConf.get(AbstractColumnSerDe.CASSANDRA_SLICE_PREDICATE_RANGE_REVERSED);
range.setStart(comparator.fromString(sliceStart == null ? "" : sliceStart));
range.setFinish(comparator.fromString(sliceEnd == null ? "" : sliceEnd));
range.setReversed(reversed == null ? false : reversed.equals("true"));
range.setCount(cassandraSplit.getSlicePredicateSize());
predicate.setSlice_range(range);
} else {
int iKey = columns.indexOf(AbstractColumnSerDe.CASSANDRA_KEY_COLUMN);
predicate.setColumn_names(getColumnNames(iKey, columns, readColIDs));
}
try {
ConfigHelper.setInputColumnFamily(tac.getConfiguration(),
cassandraSplit.getKeyspace(), cassandraSplit.getColumnFamily());
ConfigHelper.setInputSlicePredicate(tac.getConfiguration(), predicate);
ConfigHelper.setRangeBatchSize(tac.getConfiguration(), cassandraSplit.getRangeBatchSize());
ConfigHelper.setInputRpcPort(tac.getConfiguration(), cassandraSplit.getPort() + "");
ConfigHelper.setInputInitialAddress(tac.getConfiguration(), cassandraSplit.getHost());
ConfigHelper.setInputPartitioner(tac.getConfiguration(), cassandraSplit.getPartitioner());
// Set Split Size
ConfigHelper.setInputSplitSize(tac.getConfiguration(), cassandraSplit.getSplitSize());
CassandraHiveRecordReader rr = null;
if(isTransposed && tac.getConfiguration().getBoolean(AbstractColumnSerDe.CASSANDRA_ENABLE_WIDEROW_ITERATOR, true)) {
rr = new CassandraHiveRecordReader(new ColumnFamilyWideRowRecordReader(), isTransposed);
} else {
rr = new CassandraHiveRecordReader(new ColumnFamilyRecordReader(), isTransposed);
}
rr.initialize(cfSplit, tac);
return rr;
} catch (Exception ie) {
throw new IOException(ie);
}
}
@Override
public InputSplit[] getSplits(JobConf jobConf, int numSplits) throws IOException {
String ks = jobConf.get(AbstractColumnSerDe.CASSANDRA_KEYSPACE_NAME);
String cf = jobConf.get(AbstractColumnSerDe.CASSANDRA_CF_NAME);
int slicePredicateSize = jobConf.getInt(AbstractColumnSerDe.CASSANDRA_SLICE_PREDICATE_SIZE,
AbstractColumnSerDe.DEFAULT_SLICE_PREDICATE_SIZE);
int sliceRangeSize = jobConf.getInt(
AbstractColumnSerDe.CASSANDRA_RANGE_BATCH_SIZE,
AbstractColumnSerDe.DEFAULT_RANGE_BATCH_SIZE);
int splitSize = jobConf.getInt(
AbstractColumnSerDe.CASSANDRA_SPLIT_SIZE,
AbstractColumnSerDe.DEFAULT_SPLIT_SIZE);
String cassandraColumnMapping = jobConf.get(AbstractColumnSerDe.CASSANDRA_COL_MAPPING);
int rpcPort = jobConf.getInt(AbstractColumnSerDe.CASSANDRA_PORT, 9160);
String host = jobConf.get(AbstractColumnSerDe.CASSANDRA_HOST);
String partitioner = jobConf.get(AbstractColumnSerDe.CASSANDRA_PARTITIONER);
if (cassandraColumnMapping == null) {
throw new IOException("cassandra.columns.mapping required for Cassandra Table.");
}
SliceRange range = new SliceRange();
range.setStart(new byte[0]);
range.setFinish(new byte[0]);
range.setReversed(false);
range.setCount(slicePredicateSize);
SlicePredicate predicate = new SlicePredicate();
predicate.setSlice_range(range);
ConfigHelper.setInputRpcPort(jobConf, "" + rpcPort);
ConfigHelper.setInputInitialAddress(jobConf, host);
ConfigHelper.setInputPartitioner(jobConf, partitioner);
ConfigHelper.setInputSlicePredicate(jobConf, predicate);
ConfigHelper.setInputColumnFamily(jobConf, ks, cf);
ConfigHelper.setRangeBatchSize(jobConf, sliceRangeSize);
ConfigHelper.setInputSplitSize(jobConf, splitSize);
Job job = new Job(jobConf);
JobContext jobContext = new JobContext(job.getConfiguration(), job.getJobID());
Path[] tablePaths = FileInputFormat.getInputPaths(jobContext);
List<org.apache.hadoop.mapreduce.InputSplit> splits = getSplits(jobContext);
InputSplit[] results = new InputSplit[splits.size()];
for (int i = 0; i < splits.size(); ++i) {
HiveCassandraStandardSplit csplit = new HiveCassandraStandardSplit(
(ColumnFamilySplit) splits.get(i), cassandraColumnMapping, tablePaths[0]);
csplit.setKeyspace(ks);
csplit.setColumnFamily(cf);
csplit.setRangeBatchSize(sliceRangeSize);
csplit.setSplitSize(splitSize);
csplit.setHost(host);
csplit.setPort(rpcPort);
csplit.setSlicePredicateSize(slicePredicateSize);
csplit.setPartitioner(partitioner);
csplit.setColumnMapping(cassandraColumnMapping);
results[i] = csplit;
}
return results;
}
/**
* Return a list of columns names to read from cassandra. The column defined as the key in the
* column mapping
* should be skipped.
*
* @param iKey
* the index of the key defined in the column mappping
* @param columns
* column mapping
* @param readColIDs
* column names to read from cassandra
*/
private List<ByteBuffer> getColumnNames(int iKey, List<String> columns, List<Integer> readColIDs) {
List<ByteBuffer> results = new ArrayList();
int maxSize = columns.size();
for (Integer i : readColIDs) {
assert (i < maxSize);
if (i != iKey) {
results.add(ByteBufferUtil.bytes(columns.get(i.intValue())));
}
}
return results;
}
@Override
public List<org.apache.hadoop.mapreduce.InputSplit> getSplits(JobContext context)
throws IOException {
return cfif.getSplits(context);
}
@Override
public org.apache.hadoop.mapreduce.RecordReader<BytesWritable, MapWritable> createRecordReader(
org.apache.hadoop.mapreduce.InputSplit arg0, TaskAttemptContext tac) throws IOException,
InterruptedException {
if(isTransposed && tac.getConfiguration().getBoolean(AbstractColumnSerDe.CASSANDRA_ENABLE_WIDEROW_ITERATOR, true)) {
return new CassandraHiveRecordReader(new ColumnFamilyWideRowRecordReader(), isTransposed);
} else {
return new CassandraHiveRecordReader(new ColumnFamilyRecordReader(), isTransposed);
}
}
}
