package org.vertexium.accumulo.iterator;
import org.apache.accumulo.core.client.IteratorSetting;
import org.apache.accumulo.core.data.Key;
import org.apache.accumulo.core.data.Value;
import org.apache.accumulo.core.iterators.IteratorEnvironment;
import org.apache.accumulo.core.iterators.SortedKeyValueIterator;
import org.apache.accumulo.core.iterators.user.RowEncodingIterator;
import org.apache.hadoop.io.Text;
import org.vertexium.accumulo.iterator.model.EdgeInfo;
import org.vertexium.accumulo.iterator.model.VertexiumAccumuloIteratorException;
import org.vertexium.accumulo.iterator.util.DataInputStreamUtils;
import org.vertexium.accumulo.iterator.util.DataOutputStreamUtils;
import java.io.*;
import java.util.*;
public class ConnectedVertexIdsIterator extends RowEncodingIterator {
public static final String SETTING_LABEL_PREFIX = "label:";
public static final String SETTING_EXCLUDED_LABEL_PREFIX = "excludedLabel:";
private Set<String> labels;
private Set<String> excludedLabels;
public static void setLabels(IteratorSetting settings, String[] labels) {
addLabelsToSettings(settings, SETTING_LABEL_PREFIX, labels);
}
public static void setExcludedLabels(IteratorSetting settings, String[] labels) {
addLabelsToSettings(settings, SETTING_EXCLUDED_LABEL_PREFIX, labels);
}
private static void addLabelsToSettings(IteratorSetting settings, String settingPrefix, String[] labels) {
if (labels == null) {
return;
}
for (int i = 0; i < labels.length; i++) {
settings.addOption(settingPrefix + i, labels[i]);
}
}
@Override
public void init(SortedKeyValueIterator<Key, Value> source, Map<String, String> options, IteratorEnvironment env) throws IOException {
super.init(source, options, env);
Set<String> labels = new HashSet<>();
Set<String> excludedLabels = new HashSet<>();
for (Map.Entry<String, String> option : options.entrySet()) {
if (option.getKey().startsWith(SETTING_LABEL_PREFIX)) {
labels.add(option.getValue());
} else if (option.getKey().startsWith(SETTING_EXCLUDED_LABEL_PREFIX)) {
excludedLabels.add(option.getValue());
}
}
this.labels = labels.size() == 0 ? null : labels;
this.excludedLabels = excludedLabels.size() == 0 ? null : excludedLabels;
}
public static Set<String> decodeValue(Value value) throws IOException {
ByteArrayInputStream bais = new ByteArrayInputStream(value.get());
DataInputStream in = new DataInputStream(bais);
return DataInputStreamUtils.decodeSetOfStrings(in);
}
@Override
public SortedMap<Key, Value> rowDecoder(Key rowKey, Value rowValue) throws IOException {
throw new VertexiumAccumuloIteratorException("not implemented");
}
@Override
public Value rowEncoder(List<Key> keys, List<Value> values) throws IOException {
Map<Text, String> inVertexIds = new HashMap<>();
Map<Text, String> outVertexIds = new HashMap<>();
for (int i = 0; i < keys.size(); i++) {
Key key = keys.get(i);
Value value = values.get(i);
if (key.getColumnFamily().equals(VertexIterator.CF_OUT_EDGE)) {
EdgeInfo edgeInfo = new EdgeInfo(value.get(), key.getTimestamp());
if (isMatch(edgeInfo)) {
outVertexIds.put(key.getColumnQualifier(), edgeInfo.getVertexId());
}
} else if (key.getColumnFamily().equals(VertexIterator.CF_OUT_EDGE_HIDDEN)
|| key.getColumnFamily().equals(VertexIterator.CF_OUT_EDGE_SOFT_DELETE)) {
outVertexIds.remove(key.getColumnQualifier());
} else if (key.getColumnFamily().equals(VertexIterator.CF_IN_EDGE)) {
EdgeInfo edgeInfo = new EdgeInfo(value.get(), key.getTimestamp());
if (isMatch(edgeInfo)) {
inVertexIds.put(key.getColumnQualifier(), edgeInfo.getVertexId());
}
} else if (key.getColumnFamily().equals(VertexIterator.CF_IN_EDGE_HIDDEN)
|| key.getColumnFamily().equals(VertexIterator.CF_IN_EDGE_SOFT_DELETE)) {
inVertexIds.remove(key.getColumnQualifier());
}
}
ByteArrayOutputStream baos = new ByteArrayOutputStream();
DataOutputStream out = new DataOutputStream(baos);
Set<String> vertexIds = new HashSet<>();
vertexIds.addAll(inVertexIds.values());
vertexIds.addAll(outVertexIds.values());
DataOutputStreamUtils.encodeSetOfStrings(out, vertexIds);
return new Value(baos.toByteArray());
}
private boolean isMatch(EdgeInfo edgeInfo) {
if (excludedLabels != null && excludedLabels.contains(edgeInfo.getLabel())) {
return false;
}
return labels == null || labels.contains(edgeInfo.getLabel());
}
@Override
public SortedKeyValueIterator<Key, Value> deepCopy(IteratorEnvironment env) {
return new ConnectedVertexIdsIterator();
}
}