Resolution input now supports an optional "ids" field to start the jo…

…b by document _id(s) by index. In the first iteration of a job, documents are fetched by _id(s) by index and the attributes of those documents are acquired. Any _id supplied to "ids" is assumed to be representative of the entity. Jobs must have either or both of "ids" and "attributes" to run.

src/main/java/io/zentity/common/Json.java

@@ -1,11 +1,29 @@
package io.zentity.common;

import com.fasterxml.jackson.core.io.JsonStringEncoder;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;

public class Json {

    public static final ObjectMapper MAPPER = new ObjectMapper();
    public static final ObjectMapper ORDERED_MAPPER = new ObjectMapper().configure(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS, true);
    private static final JsonStringEncoder STRING_ENCODER = new JsonStringEncoder();

    public static String quoteString(String value) {
        return jsonStringFormat(value);
    }

    private static String jsonStringEscape(String value) {
        return new String(STRING_ENCODER.quoteAsString(value));
    }

    private static String jsonStringQuote(String value) {
        return "\"" + value + "\"";
    }

    private static String jsonStringFormat(String value) {
        return jsonStringQuote(jsonStringEscape(value));
    }

}

src/main/java/io/zentity/resolution/Job.java

@@ -1,6 +1,5 @@
package io.zentity.resolution;

import com.fasterxml.jackson.core.io.JsonStringEncoder;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import io.zentity.common.Json;
@@ -46,7 +45,6 @@
    public static final int DEFAULT_MAX_HOPS = 100;
    public static final boolean DEFAULT_PRETTY = false;
    public static final boolean DEFAULT_PROFILE = false;
    private final JsonStringEncoder encoder = new JsonStringEncoder();

    // Job configuration
    private Input input;
@@ -499,18 +497,6 @@ public void input(Input input) {
        this.input = input;
    }

    private String jsonStringEscape(String value) {
        return new String(encoder.quoteAsString(value));
    }

    private String jsonStringQuote(String value) {
        return "\"" + value + "\"";
    }

    private String jsonStringFormat(String value) {
        return jsonStringQuote(jsonStringEscape(value));
    }

    /**
     * Submit a search query to Elasticsearch.
     *
@@ -549,12 +535,14 @@ private void traverse() throws IOException, ValidationException {
            if (!this.docIds.containsKey(indexName))
                this.docIds.put(indexName, new TreeSet<>());

            boolean filterIds = this.hop == 0 && this.input().ids().containsKey(indexName) && !this.input().ids().get(indexName).isEmpty();

            // Determine which resolvers can be queried for this index.
            List<String> resolvers = new ArrayList<>();
            for (String resolverName : this.input.model().resolvers().keySet())
                if (canQuery(this.input.model(), indexName, resolverName, this.attributes))
                    resolvers.add(resolverName);
            if (resolvers.size() == 0)
            if (resolvers.size() == 0 && !filterIds)
                continue;

            // Construct query for this index.
@@ -563,13 +551,14 @@ private void traverse() throws IOException, ValidationException {
            List<String> queryClauses = new ArrayList<>();
            List<String> queryMustNotClauses = new ArrayList<>();
            List<String> queryFilterClauses = new ArrayList<>();
            List<String> queryResolverClauses = new ArrayList<>();
            List<String> topLevelClauses = new ArrayList<>();
            topLevelClauses.add("\"_source\":true");

            // Exclude docs by _id
            Set<String> ids = this.docIds.get(indexName);
            if (!ids.isEmpty())
                queryMustNotClauses.add("{\"ids\":{\"values\":[" + String.join(",", ids) + "]}}");
            Set<String> docIds = this.docIds.get(indexName);
            if (!docIds.isEmpty())
                queryMustNotClauses.add("{\"ids\":{\"values\":[" + String.join(",", docIds) + "]}}");

            // Create "scope.exclude.attributes" clauses. Combine them into a single "should" clause.
            if (!this.input.scope().exclude().attributes().isEmpty()) {
@@ -585,7 +574,7 @@ else if (size == 1)
            if (!queryMustNotClauses.isEmpty())
                queryClauses.add("\"must_not\":[" + String.join(",", queryMustNotClauses) + "]");

            // Create "scope.include.attributes" clauses. Combine them into a single "filter" clause.
            // Construct "scope.include.attributes" clauses. Combine them into a single "filter" clause.
            if (!this.input.scope().include().attributes().isEmpty()) {
                List<String> attributeClauses = makeAttributeClauses(this.input.model(), indexName, this.input.scope().include().attributes(), "filter");
                int size = attributeClauses.size();
@@ -595,12 +584,30 @@ else if (size == 1)
                    queryFilterClauses.add(attributeClauses.get(0));
            }

            // Construct the "ids" clause if this is the first hop and if any ids are specified for this index.
            String idsClause = "{}";
            if (filterIds) {
                Set<String> ids = this.input().ids().get(indexName);
                idsClause = "{\"bool\":{\"filter\":[{\"ids\":{\"values\":[" + String.join(",", ids) + "]}}]}}";
            }

            // Construct the resolvers clause.
            Map<String, Integer> counts = countAttributesAcrossResolvers(this.input.model(), resolvers);
            List<List<String>> resolversSorted = sortResolverAttributes(this.input.model(), resolvers, counts);
            TreeMap<String, TreeMap> resolversFilterTree = makeResolversFilterTree(resolversSorted);
            String resolversClause = populateResolversFilterTree(this.input.model(), indexName, resolversFilterTree, this.attributes);
            queryFilterClauses.add(resolversClause);
            String resolversClause = "{}";
            TreeMap<String, TreeMap> resolversFilterTree = new TreeMap<>();
            if (!this.attributes.isEmpty()) {
                Map<String, Integer> counts = countAttributesAcrossResolvers(this.input.model(), resolvers);
                List<List<String>> resolversSorted = sortResolverAttributes(this.input.model(), resolvers, counts);
                resolversFilterTree = makeResolversFilterTree(resolversSorted);
                resolversClause = populateResolversFilterTree(this.input.model(), indexName, resolversFilterTree, this.attributes);
            }

            // Combine the ids clause and resolvers clause in a "should" clause if necessary.
            if (!idsClause.equals("{}") && !resolversClause.equals("{}"))
                queryFilterClauses.add("{\"bool\":{\"should\":[" + idsClause + "," + resolversClause + "]}}");
            else if (!idsClause.equals("{}"))
                queryFilterClauses.add(idsClause);
            else
                queryFilterClauses.add(resolversClause);

            // Construct the top-level "filter" clause.
            if (!queryFilterClauses.isEmpty()) {
@@ -662,7 +669,7 @@ else if (size == 1)
            for (JsonNode doc : responseData.get("hits").get("hits")) {

                // Skip doc if already fetched. Otherwise mark doc as fetched and then proceed.
                String _id = jsonStringFormat(doc.get("_id").textValue());
                String _id = Json.quoteString(doc.get("_id").textValue());
                if (this.docIds.get(indexName).contains(_id))
                    continue;
                this.docIds.get(indexName).add(_id);

src/main/java/io/zentity/resolution/input/Input.java

@@ -19,6 +19,7 @@
public class Input {

    private Map<String, Attribute> attributes = new TreeMap<>();
    private Map<String, Set<String>> ids = new TreeMap<>();
    private String entityType;
    private Model model;
    private Scope scope = new Scope();
@@ -127,6 +128,47 @@ public static Model includeResolvers(Model model, Set<String> resolvers) throws
        return model;
    }

    /**
     * Parse and validate the "ids" field of the request body.
     *
     * @param requestBody The request body.
     * @param model       The entity model.
     * @return The parsed "ids" field from the request body.
     * @throws ValidationException
     */
    public static Map<String, Set<String>> parseIds(JsonNode requestBody, Model model) throws ValidationException {
        Map<String, Set<String>> idsObj = new TreeMap<>();
        if (!requestBody.has("ids") || requestBody.get("ids").size() == 0)
            return idsObj;
        JsonNode ids = requestBody.get("ids");
        Iterator<Map.Entry<String, JsonNode>> indices = ids.fields();
        while (indices.hasNext()) {
            Map.Entry<String, JsonNode> index = indices.next();
            String indexName = index.getKey();
            JsonNode idsValues = index.getValue();

            // Validate that the index exists in the entity model.
            if (!model.indices().containsKey(indexName))
                throw new ValidationException("'ids." + indexName + "' is not defined in the entity model.");

            // Parse the id values.
            idsObj.put(indexName, new TreeSet<>());
            if (!idsValues.isNull() && !idsValues.isArray())
                throw new ValidationException("'ids." + indexName + "' must be an array.");
            Iterator<JsonNode> idsNode = idsValues.elements();
            while (idsNode.hasNext()) {
                JsonNode idNode = idsNode.next();
                if (!idNode.isTextual())
                    throw new ValidationException("'ids." + indexName + "' must be an array of strings.");
                String id = idNode.asText();
                if (Patterns.EMPTY_STRING.matcher(id).matches())
                    throw new ValidationException("'ids." + indexName + "' must be an array of non-empty strings.");
                idsObj.get(indexName).add(Json.quoteString(id));
            }
        }
        return idsObj;
    }

    /**
     * Parse and validate the "attributes" field of the request body.
     *
@@ -137,12 +179,10 @@ public static Model includeResolvers(Model model, Set<String> resolvers) throws
     * @throws JsonProcessingException
     */
    public static Map<String, Attribute> parseAttributes(JsonNode requestBody, Model model) throws ValidationException, JsonProcessingException {
        if (!requestBody.has("attributes"))
            throw new ValidationException("The 'attributes' field is missing from the request body.");
        if (requestBody.get("attributes").size() == 0)
            throw new ValidationException("The 'attributes' field of the request body must not be empty.");
        JsonNode attributes = requestBody.get("attributes");
        Map<String, Attribute> attributesObj = new TreeMap<>();
        if (!requestBody.has("attributes") || requestBody.get("attributes").size() == 0)
            return attributesObj;
        JsonNode attributes = requestBody.get("attributes");
        Iterator<String> attributeFields = attributes.fieldNames();
        while (attributeFields.hasNext()) {
            String attributeName = attributeFields.next();
@@ -219,6 +259,10 @@ private void validateObject(String field, JsonNode object) throws ValidationExce
        return this.attributes;
    }

    public Map<String, Set<String>> ids() {
        return this.ids;
    }

    public Model model() {
        return this.model;
    }
@@ -243,6 +287,13 @@ public void deserialize(JsonNode json) throws ValidationException, IOException {
        // Parse and validate the "attributes" field of the request body.
        this.attributes = parseAttributes(json, this.model);

        // Parse and validate the "ids" field of the request body.
        this.ids = parseIds(json, this.model);

        // Ensure that either the "attributes" or "ids" field exists and is valid.
        if (this.attributes().isEmpty() && this.ids.isEmpty())
            throw new ValidationException("The 'attributes' and 'ids' fields are missing from the request body. At least one must exist.");

        // Parse and validate the "scope" field of the request body.
        if (json.has("scope")) {
            this.scope.deserialize(json.get("scope"), this.model);

src/main/java/org/elasticsearch/plugin/zentity/ResolutionAction.java

@@ -65,10 +65,8 @@ protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient
                    String model = getResponse.getSourceAsString();
                    input = new Input(body, new Model(model));
                }
                if (input.attributes().isEmpty())
                    throw new ValidationException("'attributes' is missing.");
                if (input.model() == null)
                    throw new ValidationException("'model' is missing.");
                    throw new ValidationException("An entity model was not given for this request.");

                // Prepare the entity resolution job.
                Job job = new Job(client);

src/test/java/io/zentity/resolution/JobIT.java

@@ -19,7 +19,7 @@

    private final Map<String, String> params = Collections.emptyMap();

    private final StringEntity TEST_PAYLOAD_JOB = new StringEntity("{\n" +
    private final StringEntity TEST_PAYLOAD_JOB_ATTRIBUTES = new StringEntity("{\n" +
            "  \"attributes\": {\n" +
            "    \"attribute_a\": [ \"a_00\" ]\n" +
            "  },\n" +
@@ -31,6 +31,33 @@
            "  }\n" +
            "}", ContentType.APPLICATION_JSON);

    private final StringEntity TEST_PAYLOAD_JOB_IDS = new StringEntity("{\n" +
            "  \"ids\": {\n" +
            "    \".zentity_test_index_a\": [ \"a0\" ]\n" +
            "  },\n" +
            "  \"scope\": {\n" +
            "    \"include\": {\n" +
            "      \"indices\": [ \".zentity_test_index_a\", \".zentity_test_index_b\", \".zentity_test_index_c\" ],\n" +
            "      \"resolvers\": [ \"resolver_a\", \"resolver_b\" ]\n" +
            "    }\n" +
            "  }\n" +
            "}", ContentType.APPLICATION_JSON);

    private final StringEntity TEST_PAYLOAD_JOB_ATTRIBUTES_IDS = new StringEntity("{\n" +
            "  \"attributes\": {\n" +
            "    \"attribute_a\": [ \"a_00\" ]\n" +
            "  },\n" +
            "  \"ids\": {\n" +
            "    \".zentity_test_index_a\": [ \"a6\" ]\n" +
            "  },\n" +
            "  \"scope\": {\n" +
            "    \"include\": {\n" +
            "      \"indices\": [ \".zentity_test_index_a\", \".zentity_test_index_b\", \".zentity_test_index_c\" ],\n" +
            "      \"resolvers\": [ \"resolver_a\", \"resolver_b\" ]\n" +
            "    }\n" +
            "  }\n" +
            "}", ContentType.APPLICATION_JSON);

    private final StringEntity TEST_PAYLOAD_JOB_MAX_HOPS_AND_DOCS = new StringEntity("{\n" +
            "  \"attributes\": {\n" +
            "    \"attribute_d\": { \"values\": [ \"d_00\" ] }\n" +
@@ -265,11 +292,11 @@ private void prepareTestResources() throws IOException {
        return docsActual;
    }

    public void testJob() throws Exception {
    public void testJobAttributes() throws Exception {
        try {
            prepareTestResources();
            String endpoint = "_zentity/resolution/zentity_test_entity_a";
            Response response = client.performRequest("POST", endpoint, params, TEST_PAYLOAD_JOB);
            Response response = client.performRequest("POST", endpoint, params, TEST_PAYLOAD_JOB_ATTRIBUTES);
            JsonNode json = Json.MAPPER.readTree(response.getEntity().getContent());
            assertEquals(json.get("hits").get("total").asInt(), 6);

@@ -287,6 +314,74 @@ public void testJob() throws Exception {
        }
    }

    public void testJobIds() throws Exception {
        try {
            prepareTestResources();
            String endpoint = "_zentity/resolution/zentity_test_entity_a";
            Response response = client.performRequest("POST", endpoint, params, TEST_PAYLOAD_JOB_IDS);
            JsonNode json = Json.MAPPER.readTree(response.getEntity().getContent());
            assertEquals(json.get("hits").get("total").asInt(), 6);

            Set<String> docsExpected = new TreeSet<>();
            docsExpected.add("a0,0");
            docsExpected.add("b0,1");
            docsExpected.add("c0,2");
            docsExpected.add("a1,3");
            docsExpected.add("b1,4");
            docsExpected.add("c1,5");

            assertEquals(docsExpected, getActual(json));
        } finally {
            destroyTestResources();
        }
    }

    public void testJobAttributesIds() throws Exception {
        try {
            prepareTestResources();
            String endpoint = "_zentity/resolution/zentity_test_entity_a";
            Response response = client.performRequest("POST", endpoint, params, TEST_PAYLOAD_JOB_ATTRIBUTES_IDS);
            JsonNode json = Json.MAPPER.readTree(response.getEntity().getContent());
            assertEquals(json.get("hits").get("total").asInt(), 30);

            Set<String> docsExpected = new TreeSet<>();
            docsExpected.add("a0,0");
            docsExpected.add("a6,0");
            docsExpected.add("b0,0");
            docsExpected.add("a2,1");
            docsExpected.add("a7,1");
            docsExpected.add("a8,1");
            docsExpected.add("a9,1");
            docsExpected.add("b2,1");
            docsExpected.add("b6,1");
            docsExpected.add("b7,1");
            docsExpected.add("b8,1");
            docsExpected.add("b9,1");
            docsExpected.add("c0,1");
            docsExpected.add("c2,1");
            docsExpected.add("c6,1");
            docsExpected.add("c7,1");
            docsExpected.add("c8,1");
            docsExpected.add("c9,1");
            docsExpected.add("a1,2");
            docsExpected.add("a3,2");
            docsExpected.add("a4,2");
            docsExpected.add("a5,2");
            docsExpected.add("b3,2");
            docsExpected.add("b4,2");
            docsExpected.add("b5,2");
            docsExpected.add("c3,2");
            docsExpected.add("c4,2");
            docsExpected.add("c5,2");
            docsExpected.add("b1,3");
            docsExpected.add("c1,4");

            assertEquals(docsExpected, getActual(json));
        } finally {
            destroyTestResources();
        }
    }

    public void testJobMaxHopsAndDocs() throws Exception {
        try {
            prepareTestResources();