/**
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with this
* work for additional information regarding copyright ownership. The ASF
* licenses this file to you 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 org.apache.pig.backend.hadoop.accumulo;
import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Map.Entry;
import java.util.SortedMap;
import org.apache.accumulo.core.client.IteratorSetting;
import org.apache.accumulo.core.client.mapreduce.AccumuloInputFormat;
import org.apache.accumulo.core.data.Key;
import org.apache.accumulo.core.data.Mutation;
import org.apache.accumulo.core.data.PartialKey;
import org.apache.accumulo.core.data.Range;
import org.apache.accumulo.core.data.Value;
import org.apache.accumulo.core.iterators.user.WholeRowIterator;
import org.apache.commons.cli.ParseException;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.log4j.Logger;
import org.apache.pig.ResourceSchema.ResourceFieldSchema;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataByteArray;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
/**
* Basic PigStorage implementation that uses Accumulo as the backing store.
*
* <p>
* When writing data, the first entry in the {@link Tuple} is treated as the row
* in the Accumulo key, while subsequent entries in the tuple are handled as
* columns in that row. {@link Map}s are expanded, placing the map key in the
* column family and the map value in the Accumulo value. Scalars are placed
* directly into the value with an empty column qualifier. If the columns
* argument on the constructor is omitted, null or the empty String, no column
* family is provided on the Keys created for Accumulo
* </p>
*
* <p>
* When reading data, if aggregateColfams is true, elements in the same row and
* column family are aggregated into a single {@link Map}. This will result in a
* {@link Tuple} of length (unique_column_families + 1) for the given row. If
* aggregateColfams is false, column family and column qualifier are
* concatenated (separated by a colon), and placed into a {@link Map}. This will
* result in a {@link Tuple} with two entries, where the latter element has a
* number of elements equal to the number of columns in the given row.
* </p>
*/
public class AccumuloStorage extends AbstractAccumuloStorage {
private static final Logger log = Logger.getLogger(AccumuloStorage.class);
private static final String COLON = ":", EMPTY = "";
private static final Text EMPTY_TEXT = new Text(new byte[0]);
private static final DataByteArray EMPTY_DATA_BYTE_ARRAY = new DataByteArray(
new byte[0]);
// Not sure if AccumuloStorage instances need to be thread-safe or not
final Text _cfHolder = new Text(), _cqHolder = new Text();
/**
* Creates an AccumuloStorage which writes all values in a {@link Tuple}
* with an empty column family and doesn't group column families together on
* read (creates on {@link Map} for all columns)
*/
public AccumuloStorage() throws ParseException, IOException {
this(EMPTY, EMPTY);
}
/**
* Create an AccumuloStorage with a CSV of columns-families to use on write
* and whether columns in a row should be grouped by family on read.
*
* @param columns
* A comma-separated list of column families to use when writing
* data, aligned to the n'th entry in the tuple
* @param aggregateColfams
* Should unique column qualifier and value pairs be grouped
* together by column family when reading data
*/
public AccumuloStorage(String columns) throws ParseException, IOException {
this(columns, EMPTY);
}
public AccumuloStorage(String columnStr, String args)
throws ParseException, IOException {
super(columnStr, args);
}
@Override
protected Tuple getTuple(Key key, Value value) throws IOException {
SortedMap<Key, Value> rowKVs = WholeRowIterator.decodeRow(key, value);
Tuple tuple = TupleFactory.getInstance().newTuple(columns.size() + 1);
final Text cfHolder = new Text();
final Text cqHolder = new Text();
final Text row = key.getRow();
int tupleOffset = 0;
tuple.set(
tupleOffset,
new DataByteArray(Text.decode(row.getBytes(), 0,
row.getLength())));
for (Column column : this.columns) {
tupleOffset++;
switch (column.getType()) {
case LITERAL:
cfHolder.set(column.getColumnFamily());
if (null != column.getColumnQualifier()) {
cqHolder.set(column.getColumnQualifier());
} else {
cqHolder.set(EMPTY_TEXT);
}
// Get the key where our literal would exist (accounting for
// "colf:colq" or "colf:" empty colq)
Key literalStartKey = new Key(row, cfHolder, cqHolder);
SortedMap<Key, Value> tailMap = rowKVs.tailMap(literalStartKey);
// Find the element
if (tailMap.isEmpty()) {
tuple.set(tupleOffset, EMPTY_DATA_BYTE_ARRAY);
} else {
Key actualKey = tailMap.firstKey();
// Only place it in the tuple if it matches the user
// request, avoid using a value from a
// key with the wrong colqual
if (0 == literalStartKey.compareTo(actualKey,
PartialKey.ROW_COLFAM_COLQUAL)) {
tuple.set(tupleOffset,
new DataByteArray(tailMap.get(actualKey).get()));
} else {
// This row doesn't have the column we were looking for
tuple.set(tupleOffset, EMPTY_DATA_BYTE_ARRAY);
}
}
break;
case COLFAM_PREFIX:
cfHolder.set(column.getColumnFamily());
Range colfamPrefixRange = Range.prefix(row, cfHolder);
Key colfamPrefixStartKey = new Key(row, cfHolder);
SortedMap<Key, Value> cfTailMap = rowKVs
.tailMap(colfamPrefixStartKey);
// Find the element
if (cfTailMap.isEmpty()) {
tuple.set(tupleOffset, EMPTY_DATA_BYTE_ARRAY);
} else {
HashMap<String, DataByteArray> tupleMap = new HashMap<String, DataByteArray>();
// Build up a map for all the entries in this row that match
// the colfam prefix
for (Entry<Key, Value> entry : cfTailMap.entrySet()) {
if (colfamPrefixRange.contains(entry.getKey())) {
entry.getKey().getColumnFamily(cfHolder);
entry.getKey().getColumnQualifier(cqHolder);
DataByteArray val = new DataByteArray(entry
.getValue().get());
// Avoid adding an extra ':' when colqual is empty
if (0 == cqHolder.getLength()) {
tupleMap.put(cfHolder.toString(), val);
} else {
tupleMap.put(cfHolder.toString() + COLON
+ cqHolder.toString(), val);
}
} else {
break;
}
}
if (!tupleMap.isEmpty()) {
tuple.set(tupleOffset, tupleMap);
}
}
break;
case COLQUAL_PREFIX:
cfHolder.set(column.getColumnFamily());
cqHolder.set(column.getColumnQualifier());
Range colqualPrefixRange = Range
.prefix(row, cfHolder, cqHolder);
Key colqualPrefixStartKey = new Key(row, cfHolder, cqHolder);
SortedMap<Key, Value> cqTailMap = rowKVs
.tailMap(colqualPrefixStartKey);
if (cqTailMap.isEmpty()) {
tuple.set(tupleOffset, EMPTY_DATA_BYTE_ARRAY);
} else {
HashMap<String, DataByteArray> tupleMap = new HashMap<String, DataByteArray>();
// Build up a map for all the entries in this row that match
// the colqual prefix
for (Entry<Key, Value> entry : cqTailMap.entrySet()) {
if (colqualPrefixRange.contains(entry.getKey())) {
entry.getKey().getColumnFamily(cfHolder);
entry.getKey().getColumnQualifier(cqHolder);
DataByteArray val = new DataByteArray(entry
.getValue().get());
// Avoid the extra ':' on empty colqual
if (0 == cqHolder.getLength()) {
tupleMap.put(cfHolder.toString(), val);
} else {
tupleMap.put(cfHolder.toString() + COLON
+ cqHolder.toString(), val);
}
} else {
break;
}
}
if (!tupleMap.isEmpty()) {
tuple.set(tupleOffset, tupleMap);
}
}
break;
default:
break;
}
}
return tuple;
}
@Override
protected void configureInputFormat(Job job) {
AccumuloInputFormat.addIterator(job, new IteratorSetting(100,
WholeRowIterator.class));
}
@Override
protected Collection<Mutation> getMutations(Tuple tuple)
throws ExecException, IOException {
final ResourceFieldSchema[] fieldSchemas = (schema == null) ? null
: schema.getFields();
Iterator<Object> tupleIter = tuple.iterator();
if (1 >= tuple.size()) {
log.debug("Ignoring tuple of size " + tuple.size());
return Collections.emptyList();
}
Mutation mutation = new Mutation(objectToText(tupleIter.next(),
(null == fieldSchemas) ? null : fieldSchemas[0]));
int tupleOffset = 1;
Iterator<Column> columnIter = columns.iterator();
while (tupleIter.hasNext() && columnIter.hasNext()) {
Object o = tupleIter.next();
Column column = columnIter.next();
// Grab the type for this field
final byte type = schemaToType(o, (null == fieldSchemas) ? null
: fieldSchemas[tupleOffset]);
switch (column.getType()) {
case LITERAL:
byte[] bytes = objToBytes(o, type);
if (null != bytes) {
Value value = new Value(bytes);
// We don't have any column name from non-Maps
addColumn(mutation, column.getColumnFamily(),
column.getColumnQualifier(), value);
}
break;
case COLFAM_PREFIX:
case COLQUAL_PREFIX:
Map<String, Object> map;
try {
map = (Map<String, Object>) o;
} catch (ClassCastException e) {
log.error("Expected Map at tuple offset " + tupleOffset
+ " but was " + o.getClass().getSimpleName());
throw e;
}
for (Entry<String, Object> entry : map.entrySet()) {
String key = entry.getKey();
Object objValue = entry.getValue();
byte valueType = DataType.findType(objValue);
byte[] mapValue = objToBytes(objValue, valueType);
if (Column.Type.COLFAM_PREFIX == column.getType()) {
addColumn(mutation, column.getColumnFamily() + key,
null, new Value(mapValue));
} else if (Column.Type.COLQUAL_PREFIX == column.getType()) {
addColumn(mutation, column.getColumnFamily(),
column.getColumnQualifier() + key, new Value(
mapValue));
} else {
throw new IOException("Unknown column type");
}
}
break;
default:
log.info("Ignoring unhandled column type");
continue;
}
tupleOffset++;
}
if (0 == mutation.size()) {
return Collections.emptyList();
}
return Collections.singletonList(mutation);
}
/**
* Adds the given column family, column qualifier and value to the given
* mutation
*
* @param mutation
* @param colfam
* @param colqual
* @param columnValue
*/
protected void addColumn(Mutation mutation, String colfam, String colqual,
Value columnValue) {
if (null != colfam) {
_cfHolder.set(colfam);
} else {
_cfHolder.clear();
}
if (null != colqual) {
_cqHolder.set(colqual);
} else {
_cqHolder.clear();
}
mutation.put(_cfHolder, _cqHolder, columnValue);
}
}