Skip to content

Commit b2f48bc

Browse files
davemoore-davemoore-
committed
Resolution input now supports an optional "ids" field to start the job 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.
1 parent 9ea4b95 commit b2f48bc

6 files changed

Lines changed: 375 additions & 42 deletions

File tree

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,29 @@
11
package io.zentity.common;
22

3+
import com.fasterxml.jackson.core.io.JsonStringEncoder;
34
import com.fasterxml.jackson.databind.ObjectMapper;
45
import com.fasterxml.jackson.databind.SerializationFeature;
56

67
public class Json {
78

89
public static final ObjectMapper MAPPER = new ObjectMapper();
910
public static final ObjectMapper ORDERED_MAPPER = new ObjectMapper().configure(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS, true);
11+
private static final JsonStringEncoder STRING_ENCODER = new JsonStringEncoder();
12+
13+
public static String quoteString(String value) {
14+
return jsonStringFormat(value);
15+
}
16+
17+
private static String jsonStringEscape(String value) {
18+
return new String(STRING_ENCODER.quoteAsString(value));
19+
}
20+
21+
private static String jsonStringQuote(String value) {
22+
return "\"" + value + "\"";
23+
}
24+
25+
private static String jsonStringFormat(String value) {
26+
return jsonStringQuote(jsonStringEscape(value));
27+
}
1028

1129
}

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

Lines changed: 32 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package io.zentity.resolution;
22

3-
import com.fasterxml.jackson.core.io.JsonStringEncoder;
43
import com.fasterxml.jackson.databind.JsonNode;
54
import com.fasterxml.jackson.databind.node.ObjectNode;
65
import io.zentity.common.Json;
@@ -46,7 +45,6 @@ public class Job {
4645
public static final int DEFAULT_MAX_HOPS = 100;
4746
public static final boolean DEFAULT_PRETTY = false;
4847
public static final boolean DEFAULT_PROFILE = false;
49-
private final JsonStringEncoder encoder = new JsonStringEncoder();
5048

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

502-
private String jsonStringEscape(String value) {
503-
return new String(encoder.quoteAsString(value));
504-
}
505-
506-
private String jsonStringQuote(String value) {
507-
return "\"" + value + "\"";
508-
}
509-
510-
private String jsonStringFormat(String value) {
511-
return jsonStringQuote(jsonStringEscape(value));
512-
}
513-
514500
/**
515501
* Submit a search query to Elasticsearch.
516502
*
@@ -549,12 +535,14 @@ private void traverse() throws IOException, ValidationException {
549535
if (!this.docIds.containsKey(indexName))
550536
this.docIds.put(indexName, new TreeSet<>());
551537

538+
boolean filterIds = this.hop == 0 && this.input().ids().containsKey(indexName) && !this.input().ids().get(indexName).isEmpty();
539+
552540
// Determine which resolvers can be queried for this index.
553541
List<String> resolvers = new ArrayList<>();
554542
for (String resolverName : this.input.model().resolvers().keySet())
555543
if (canQuery(this.input.model(), indexName, resolverName, this.attributes))
556544
resolvers.add(resolverName);
557-
if (resolvers.size() == 0)
545+
if (resolvers.size() == 0 && !filterIds)
558546
continue;
559547

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

569558
// Exclude docs by _id
570-
Set<String> ids = this.docIds.get(indexName);
571-
if (!ids.isEmpty())
572-
queryMustNotClauses.add("{\"ids\":{\"values\":[" + String.join(",", ids) + "]}}");
559+
Set<String> docIds = this.docIds.get(indexName);
560+
if (!docIds.isEmpty())
561+
queryMustNotClauses.add("{\"ids\":{\"values\":[" + String.join(",", docIds) + "]}}");
573562

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

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

587+
// Construct the "ids" clause if this is the first hop and if any ids are specified for this index.
588+
String idsClause = "{}";
589+
if (filterIds) {
590+
Set<String> ids = this.input().ids().get(indexName);
591+
idsClause = "{\"bool\":{\"filter\":[{\"ids\":{\"values\":[" + String.join(",", ids) + "]}}]}}";
592+
}
593+
598594
// Construct the resolvers clause.
599-
Map<String, Integer> counts = countAttributesAcrossResolvers(this.input.model(), resolvers);
600-
List<List<String>> resolversSorted = sortResolverAttributes(this.input.model(), resolvers, counts);
601-
TreeMap<String, TreeMap> resolversFilterTree = makeResolversFilterTree(resolversSorted);
602-
String resolversClause = populateResolversFilterTree(this.input.model(), indexName, resolversFilterTree, this.attributes);
603-
queryFilterClauses.add(resolversClause);
595+
String resolversClause = "{}";
596+
TreeMap<String, TreeMap> resolversFilterTree = new TreeMap<>();
597+
if (!this.attributes.isEmpty()) {
598+
Map<String, Integer> counts = countAttributesAcrossResolvers(this.input.model(), resolvers);
599+
List<List<String>> resolversSorted = sortResolverAttributes(this.input.model(), resolvers, counts);
600+
resolversFilterTree = makeResolversFilterTree(resolversSorted);
601+
resolversClause = populateResolversFilterTree(this.input.model(), indexName, resolversFilterTree, this.attributes);
602+
}
603+
604+
// Combine the ids clause and resolvers clause in a "should" clause if necessary.
605+
if (!idsClause.equals("{}") && !resolversClause.equals("{}"))
606+
queryFilterClauses.add("{\"bool\":{\"should\":[" + idsClause + "," + resolversClause + "]}}");
607+
else if (!idsClause.equals("{}"))
608+
queryFilterClauses.add(idsClause);
609+
else
610+
queryFilterClauses.add(resolversClause);
604611

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

664671
// Skip doc if already fetched. Otherwise mark doc as fetched and then proceed.
665-
String _id = jsonStringFormat(doc.get("_id").textValue());
672+
String _id = Json.quoteString(doc.get("_id").textValue());
666673
if (this.docIds.get(indexName).contains(_id))
667674
continue;
668675
this.docIds.get(indexName).add(_id);

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

Lines changed: 56 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
public class Input {
2020

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

131+
/**
132+
* Parse and validate the "ids" field of the request body.
133+
*
134+
* @param requestBody The request body.
135+
* @param model The entity model.
136+
* @return The parsed "ids" field from the request body.
137+
* @throws ValidationException
138+
*/
139+
public static Map<String, Set<String>> parseIds(JsonNode requestBody, Model model) throws ValidationException {
140+
Map<String, Set<String>> idsObj = new TreeMap<>();
141+
if (!requestBody.has("ids") || requestBody.get("ids").size() == 0)
142+
return idsObj;
143+
JsonNode ids = requestBody.get("ids");
144+
Iterator<Map.Entry<String, JsonNode>> indices = ids.fields();
145+
while (indices.hasNext()) {
146+
Map.Entry<String, JsonNode> index = indices.next();
147+
String indexName = index.getKey();
148+
JsonNode idsValues = index.getValue();
149+
150+
// Validate that the index exists in the entity model.
151+
if (!model.indices().containsKey(indexName))
152+
throw new ValidationException("'ids." + indexName + "' is not defined in the entity model.");
153+
154+
// Parse the id values.
155+
idsObj.put(indexName, new TreeSet<>());
156+
if (!idsValues.isNull() && !idsValues.isArray())
157+
throw new ValidationException("'ids." + indexName + "' must be an array.");
158+
Iterator<JsonNode> idsNode = idsValues.elements();
159+
while (idsNode.hasNext()) {
160+
JsonNode idNode = idsNode.next();
161+
if (!idNode.isTextual())
162+
throw new ValidationException("'ids." + indexName + "' must be an array of strings.");
163+
String id = idNode.asText();
164+
if (Patterns.EMPTY_STRING.matcher(id).matches())
165+
throw new ValidationException("'ids." + indexName + "' must be an array of non-empty strings.");
166+
idsObj.get(indexName).add(Json.quoteString(id));
167+
}
168+
}
169+
return idsObj;
170+
}
171+
130172
/**
131173
* Parse and validate the "attributes" field of the request body.
132174
*
@@ -137,12 +179,10 @@ public static Model includeResolvers(Model model, Set<String> resolvers) throws
137179
* @throws JsonProcessingException
138180
*/
139181
public static Map<String, Attribute> parseAttributes(JsonNode requestBody, Model model) throws ValidationException, JsonProcessingException {
140-
if (!requestBody.has("attributes"))
141-
throw new ValidationException("The 'attributes' field is missing from the request body.");
142-
if (requestBody.get("attributes").size() == 0)
143-
throw new ValidationException("The 'attributes' field of the request body must not be empty.");
144-
JsonNode attributes = requestBody.get("attributes");
145182
Map<String, Attribute> attributesObj = new TreeMap<>();
183+
if (!requestBody.has("attributes") || requestBody.get("attributes").size() == 0)
184+
return attributesObj;
185+
JsonNode attributes = requestBody.get("attributes");
146186
Iterator<String> attributeFields = attributes.fieldNames();
147187
while (attributeFields.hasNext()) {
148188
String attributeName = attributeFields.next();
@@ -219,6 +259,10 @@ public Map<String, Attribute> attributes() {
219259
return this.attributes;
220260
}
221261

262+
public Map<String, Set<String>> ids() {
263+
return this.ids;
264+
}
265+
222266
public Model model() {
223267
return this.model;
224268
}
@@ -243,6 +287,13 @@ public void deserialize(JsonNode json) throws ValidationException, IOException {
243287
// Parse and validate the "attributes" field of the request body.
244288
this.attributes = parseAttributes(json, this.model);
245289

290+
// Parse and validate the "ids" field of the request body.
291+
this.ids = parseIds(json, this.model);
292+
293+
// Ensure that either the "attributes" or "ids" field exists and is valid.
294+
if (this.attributes().isEmpty() && this.ids.isEmpty())
295+
throw new ValidationException("The 'attributes' and 'ids' fields are missing from the request body. At least one must exist.");
296+
246297
// Parse and validate the "scope" field of the request body.
247298
if (json.has("scope")) {
248299
this.scope.deserialize(json.get("scope"), this.model);

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,8 @@ protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient
6565
String model = getResponse.getSourceAsString();
6666
input = new Input(body, new Model(model));
6767
}
68-
if (input.attributes().isEmpty())
69-
throw new ValidationException("'attributes' is missing.");
7068
if (input.model() == null)
71-
throw new ValidationException("'model' is missing.");
69+
throw new ValidationException("An entity model was not given for this request.");
7270

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

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

Lines changed: 98 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ public class JobIT extends AbstractITCase {
1919

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

22-
private final StringEntity TEST_PAYLOAD_JOB = new StringEntity("{\n" +
22+
private final StringEntity TEST_PAYLOAD_JOB_ATTRIBUTES = new StringEntity("{\n" +
2323
" \"attributes\": {\n" +
2424
" \"attribute_a\": [ \"a_00\" ]\n" +
2525
" },\n" +
@@ -31,6 +31,33 @@ public class JobIT extends AbstractITCase {
3131
" }\n" +
3232
"}", ContentType.APPLICATION_JSON);
3333

34+
private final StringEntity TEST_PAYLOAD_JOB_IDS = new StringEntity("{\n" +
35+
" \"ids\": {\n" +
36+
" \".zentity_test_index_a\": [ \"a0\" ]\n" +
37+
" },\n" +
38+
" \"scope\": {\n" +
39+
" \"include\": {\n" +
40+
" \"indices\": [ \".zentity_test_index_a\", \".zentity_test_index_b\", \".zentity_test_index_c\" ],\n" +
41+
" \"resolvers\": [ \"resolver_a\", \"resolver_b\" ]\n" +
42+
" }\n" +
43+
" }\n" +
44+
"}", ContentType.APPLICATION_JSON);
45+
46+
private final StringEntity TEST_PAYLOAD_JOB_ATTRIBUTES_IDS = new StringEntity("{\n" +
47+
" \"attributes\": {\n" +
48+
" \"attribute_a\": [ \"a_00\" ]\n" +
49+
" },\n" +
50+
" \"ids\": {\n" +
51+
" \".zentity_test_index_a\": [ \"a6\" ]\n" +
52+
" },\n" +
53+
" \"scope\": {\n" +
54+
" \"include\": {\n" +
55+
" \"indices\": [ \".zentity_test_index_a\", \".zentity_test_index_b\", \".zentity_test_index_c\" ],\n" +
56+
" \"resolvers\": [ \"resolver_a\", \"resolver_b\" ]\n" +
57+
" }\n" +
58+
" }\n" +
59+
"}", ContentType.APPLICATION_JSON);
60+
3461
private final StringEntity TEST_PAYLOAD_JOB_MAX_HOPS_AND_DOCS = new StringEntity("{\n" +
3562
" \"attributes\": {\n" +
3663
" \"attribute_d\": { \"values\": [ \"d_00\" ] }\n" +
@@ -265,11 +292,11 @@ private Set<String> getActual(JsonNode json) {
265292
return docsActual;
266293
}
267294

268-
public void testJob() throws Exception {
295+
public void testJobAttributes() throws Exception {
269296
try {
270297
prepareTestResources();
271298
String endpoint = "_zentity/resolution/zentity_test_entity_a";
272-
Response response = client.performRequest("POST", endpoint, params, TEST_PAYLOAD_JOB);
299+
Response response = client.performRequest("POST", endpoint, params, TEST_PAYLOAD_JOB_ATTRIBUTES);
273300
JsonNode json = Json.MAPPER.readTree(response.getEntity().getContent());
274301
assertEquals(json.get("hits").get("total").asInt(), 6);
275302

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

317+
public void testJobIds() throws Exception {
318+
try {
319+
prepareTestResources();
320+
String endpoint = "_zentity/resolution/zentity_test_entity_a";
321+
Response response = client.performRequest("POST", endpoint, params, TEST_PAYLOAD_JOB_IDS);
322+
JsonNode json = Json.MAPPER.readTree(response.getEntity().getContent());
323+
assertEquals(json.get("hits").get("total").asInt(), 6);
324+
325+
Set<String> docsExpected = new TreeSet<>();
326+
docsExpected.add("a0,0");
327+
docsExpected.add("b0,1");
328+
docsExpected.add("c0,2");
329+
docsExpected.add("a1,3");
330+
docsExpected.add("b1,4");
331+
docsExpected.add("c1,5");
332+
333+
assertEquals(docsExpected, getActual(json));
334+
} finally {
335+
destroyTestResources();
336+
}
337+
}
338+
339+
public void testJobAttributesIds() throws Exception {
340+
try {
341+
prepareTestResources();
342+
String endpoint = "_zentity/resolution/zentity_test_entity_a";
343+
Response response = client.performRequest("POST", endpoint, params, TEST_PAYLOAD_JOB_ATTRIBUTES_IDS);
344+
JsonNode json = Json.MAPPER.readTree(response.getEntity().getContent());
345+
assertEquals(json.get("hits").get("total").asInt(), 30);
346+
347+
Set<String> docsExpected = new TreeSet<>();
348+
docsExpected.add("a0,0");
349+
docsExpected.add("a6,0");
350+
docsExpected.add("b0,0");
351+
docsExpected.add("a2,1");
352+
docsExpected.add("a7,1");
353+
docsExpected.add("a8,1");
354+
docsExpected.add("a9,1");
355+
docsExpected.add("b2,1");
356+
docsExpected.add("b6,1");
357+
docsExpected.add("b7,1");
358+
docsExpected.add("b8,1");
359+
docsExpected.add("b9,1");
360+
docsExpected.add("c0,1");
361+
docsExpected.add("c2,1");
362+
docsExpected.add("c6,1");
363+
docsExpected.add("c7,1");
364+
docsExpected.add("c8,1");
365+
docsExpected.add("c9,1");
366+
docsExpected.add("a1,2");
367+
docsExpected.add("a3,2");
368+
docsExpected.add("a4,2");
369+
docsExpected.add("a5,2");
370+
docsExpected.add("b3,2");
371+
docsExpected.add("b4,2");
372+
docsExpected.add("b5,2");
373+
docsExpected.add("c3,2");
374+
docsExpected.add("c4,2");
375+
docsExpected.add("c5,2");
376+
docsExpected.add("b1,3");
377+
docsExpected.add("c1,4");
378+
379+
assertEquals(docsExpected, getActual(json));
380+
} finally {
381+
destroyTestResources();
382+
}
383+
}
384+
290385
public void testJobMaxHopsAndDocs() throws Exception {
291386
try {
292387
prepareTestResources();

0 commit comments

Comments
 (0)