package org.vertexium.elasticsearch.score;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders;
import org.vertexium.*;
import org.vertexium.elasticsearch.ElasticsearchSingleDocumentSearchIndex;
import org.vertexium.elasticsearch.IndexInfo;
import org.vertexium.elasticsearch.utils.GetResponseUtil;
import org.vertexium.search.SearchIndex;
import java.io.IOException;
import java.util.List;
@Deprecated
public class EdgeCountScoringStrategy extends ScoringStrategy {
private final EdgeCountScoringStrategyConfiguration config;
public EdgeCountScoringStrategy(GraphConfiguration config) {
super(config);
this.config = new EdgeCountScoringStrategyConfiguration(config);
}
public EdgeCountScoringStrategyConfiguration getConfig() {
return config;
}
@Override
public void addElement(SearchIndex searchIndex, Graph graph, Element element, Authorizations authorizations) {
if (!getConfig().isUpdateEdgeBoost()) {
return;
}
if (!(element instanceof Edge)) {
return;
}
Element vOut = ((Edge) element).getVertex(Direction.OUT, authorizations);
if (vOut != null) {
searchIndex.addElement(graph, vOut, authorizations);
}
Element vIn = ((Edge) element).getVertex(Direction.IN, authorizations);
if (vIn != null) {
searchIndex.addElement(graph, vIn, authorizations);
}
}
@Override
public int addElement(ElasticsearchSingleDocumentSearchIndex searchIndex, Graph graph, BulkRequest bulkRequest, IndexInfo indexInfo, Element element, Authorizations authorizations) {
int totalCount = 0;
if (!getConfig().isUpdateEdgeBoost()) {
return totalCount;
}
if (!(element instanceof Edge)) {
return totalCount;
}
Element vOut = ((Edge) element).getVertex(Direction.OUT, authorizations);
if (vOut != null) {
searchIndex.addElementToBulkRequest(graph, bulkRequest, indexInfo, vOut, authorizations);
totalCount++;
}
Element vIn = ((Edge) element).getVertex(Direction.IN, authorizations);
if (vIn != null) {
searchIndex.addElementToBulkRequest(graph, bulkRequest, indexInfo, vIn, authorizations);
totalCount++;
}
return totalCount;
}
@Override
public void addFieldsToElementType(XContentBuilder builder) throws IOException {
builder
.startObject(EdgeCountScoringStrategyConfiguration.IN_EDGE_COUNT_FIELD_NAME).field("type", "integer").field("store", "true").endObject()
.startObject(EdgeCountScoringStrategyConfiguration.OUT_EDGE_COUNT_FIELD_NAME).field("type", "integer").field("store", "true").endObject()
;
}
@Override
public List<String> getFieldNames() {
List<String> fieldNames = super.getFieldNames();
fieldNames.add(EdgeCountScoringStrategyConfiguration.IN_EDGE_COUNT_FIELD_NAME);
fieldNames.add(EdgeCountScoringStrategyConfiguration.OUT_EDGE_COUNT_FIELD_NAME);
return fieldNames;
}
@Override
public QueryBuilder updateQuery(QueryBuilder query) {
if (!getConfig().isUseEdgeBoost()) {
return query;
}
ScoreFunctionBuilder scoreFunction = ScoreFunctionBuilders
.scriptFunction(getConfig().getScoreFormula(), "groovy")
.param("inEdgeMultiplier", getConfig().getInEdgeBoost())
.param("outEdgeMultiplier", getConfig().getOutEdgeBoost());
return QueryBuilders.functionScoreQuery(query, scoreFunction);
}
@Override
public boolean addFieldsToVertexDocument(SearchIndex searchIndex, XContentBuilder jsonBuilder, Vertex vertex, GetResponse existingDocument, Authorizations authorizations) throws IOException {
if (existingDocument != null && !getConfig().isUpdateEdgeBoost()) {
return false;
}
boolean changed = false;
int inEdgeCount = vertex.getEdgeCount(Direction.IN, authorizations);
Long existingInEdgeCount = existingDocument == null ? null : GetResponseUtil.getFieldValueLong(existingDocument, EdgeCountScoringStrategyConfiguration.IN_EDGE_COUNT_FIELD_NAME);
if (existingInEdgeCount == null || existingInEdgeCount.intValue() != inEdgeCount) {
jsonBuilder.field(EdgeCountScoringStrategyConfiguration.IN_EDGE_COUNT_FIELD_NAME, inEdgeCount);
changed = true;
} else {
jsonBuilder.field(EdgeCountScoringStrategyConfiguration.IN_EDGE_COUNT_FIELD_NAME, existingInEdgeCount.intValue());
}
int outEdgeCount = vertex.getEdgeCount(Direction.OUT, authorizations);
Long existingOutEdgeCount = existingDocument == null ? null : GetResponseUtil.getFieldValueLong(existingDocument, EdgeCountScoringStrategyConfiguration.OUT_EDGE_COUNT_FIELD_NAME);
if (existingOutEdgeCount == null || existingOutEdgeCount.intValue() != outEdgeCount) {
jsonBuilder.field(EdgeCountScoringStrategyConfiguration.OUT_EDGE_COUNT_FIELD_NAME, outEdgeCount);
changed = true;
} else {
jsonBuilder.field(EdgeCountScoringStrategyConfiguration.OUT_EDGE_COUNT_FIELD_NAME, existingOutEdgeCount.intValue());
}
return changed;
}
@Override
public boolean addFieldsToEdgeDocument(SearchIndex searchIndex, XContentBuilder jsonBuilder, Edge edge, GetResponse existingDocument, Authorizations authorizations) throws IOException {
return false;
}
}