添加是否聚合匹配的选项

xmh
Commit fda29510 authored Jun 23, 2021 by xmh
Showing 2 changed files with 53 additions and 44 deletions
src/main/java/com/viontech/match/entity/vo/RequestVo.java
src/main/java/com/viontech/match/service/PersonService.java
--- a/src/main/java/com/viontech/match/entity/vo/RequestVo.java
+++ b/src/main/java/com/viontech/match/entity/vo/RequestVo.java
@@ -43,6 +43,7 @@ public class RequestVo {
    private String personPoolId;
    private List<String> unionPersonPoolId;
    private Integer size;
+    private Boolean agg = false;
    private FaceFeature newFaceFeature;

--- a/src/main/java/com/viontech/match/service/PersonService.java
+++ b/src/main/java/com/viontech/match/service/PersonService.java
@@ -29,6 +29,8 @@ import org.elasticsearch.index.reindex.DeleteByQueryRequest;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptType;
+import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.SearchHits;
 import org.elasticsearch.search.aggregations.AggregationBuilders;
 import org.elasticsearch.search.aggregations.Aggregations;
 import org.elasticsearch.search.aggregations.BucketOrder;
@@ -41,6 +43,7 @@ import org.springframework.stereotype.Service;
 import javax.annotation.Resource;
 import java.io.IOException;
+import java.text.ParseException;
 import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -76,6 +79,7 @@ public class PersonService {
        String poolId = requestVo.getPersonPoolId();
        List<String> unionPersonPoolId = requestVo.getUnionPersonPoolId();
        Integer size = requestVo.getSize();
+        Boolean agg = requestVo.getAgg();
        List<String> poolIds = new ArrayList<>();
        List<Person> matchFaces = new ArrayList<>();
        List<Person> matchBodies = new ArrayList<>();
@@ -94,9 +98,9 @@ public class PersonService {
                pool.setPersonPoolId(id);
                if (poolService.existPool(id)) {
                    pool.setStatus(0);
-                    List<Person> face = matchPerson(id, requestVo.getPerson(), 0, size);
+                    List<Person> face = matchPerson(id, requestVo.getPerson(), 0, size, agg);
                    matchFaces.addAll(face);
-                    List<Person> body = matchPerson(id, requestVo.getPerson(), 1, size);
+                    List<Person> body = matchPerson(id, requestVo.getPerson(), 1, size, agg);
                    matchBodies.addAll(body);
                } else {
                    pool.setStatus(1);
@@ -250,14 +254,14 @@ public class PersonService {
     * @return 匹配结果
     * @throws Exception --
     */
-    public List<Person> matchPerson(String poolId, Person person, int type, Integer size) throws Exception {
+    public List<Person> matchPerson(String poolId, Person person, int type, Integer size, Boolean agg) throws Exception {
        List<Person> matchResult = new ArrayList<>();
        int matchResultSize;
        if (type == 0) {
-            matchFace(poolId, matchResult, person, size);
+            matchFace(poolId, matchResult, person, size, agg);
            matchResultSize = Constant.FACE_MATCH_RESULT_SIZE;
        } else {
-            matchBody(poolId, matchResult, person, size);
+            matchBody(poolId, matchResult, person, size, agg);
            matchResultSize = Constant.BODY_MATCH_RESULT_SIZE;
        }
        if (size != null) {
@@ -272,7 +276,7 @@ public class PersonService {
        return matchResult;
    }
-    private void matchFace(String poolId, List<Person> matchResult, Person person, Integer size) throws Exception {
+    private void matchFace(String poolId, List<Person> matchResult, Person person, Integer size, Boolean agg) throws Exception {
        List<FaceFeature> faceFeatures = person.getFaceFeatures();
        if (faceFeatures != null && faceFeatures.size() > 0) {
            for (FaceFeature faceFeature : faceFeatures) {
@@ -282,15 +286,15 @@ public class PersonService {
                    continue;
                }
-                SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.FACE_MATCH_RESULT_SIZE : size, feature, person, 0);
+                SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.FACE_MATCH_RESULT_SIZE : size, feature, person, 0, agg);
-                matchResult.addAll(match0(searchRequest));
+                matchResult.addAll(match0(searchRequest, agg));
            }
        } else {
            log.info("no face feature");
        }
    }
-    private void matchBody(String poolId, List<Person> matchResult, Person person, Integer size) throws Exception {
+    private void matchBody(String poolId, List<Person> matchResult, Person person, Integer size, Boolean agg) throws Exception {
        List<BodyFeature> bodyFeatures = person.getBodyFeatures();
        if (bodyFeatures != null && bodyFeatures.size() > 0) {
            for (BodyFeature faceFeature : bodyFeatures) {
@@ -306,14 +310,14 @@ public class PersonService {
                    feature = Arrays.copyOfRange(feature, 3, Constant.BODY_FEATURE_DIMS_2048 + 3);
                }
-                SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.BODY_MATCH_RESULT_SIZE : size, feature, person, 1);
+                SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.BODY_MATCH_RESULT_SIZE : size, feature, person, 1, agg);
-                matchResult.addAll(match0(searchRequest));
+                matchResult.addAll(match0(searchRequest, agg));
            }
        }
    }
-    private SearchRequest getSearchRequest(String poolId, Integer matchResultSize, Double[] feature, Person person, int type) {
+    private SearchRequest getSearchRequest(String poolId, Integer matchResultSize, Double[] feature, Person person, int type, Boolean agg) {
        Script script;
        BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
        if (type == 0) {
@@ -353,23 +357,26 @@ public class PersonService {
        ScriptScoreQueryBuilder queryBuilder = QueryBuilders.scriptScoreQuery(boolQuery, script);
+        SearchSourceBuilder builder = new SearchSourceBuilder().query(queryBuilder);
+        if (agg) {
            MaxAggregationBuilder maxScoreAgg = AggregationBuilders.max("max_score").script(new Script("_score"));
            TermsAggregationBuilder personIdAgg = AggregationBuilders.terms("by_personId").field("personId");
            personIdAgg.subAggregation(maxScoreAgg);
            personIdAgg.order(BucketOrder.aggregation("max_score", false));
+            personIdAgg.size(matchResultSize);
-        SearchSourceBuilder builder = new SearchSourceBuilder()
+            builder.aggregation(personIdAgg);
-                .size(matchResultSize)
+        } else {
-                .query(queryBuilder)
+            builder.fetchSource(FETCH_SOURCE, null)
-//                .fetchSource(FETCH_SOURCE, null)
+                    .size(matchResultSize);
-                .aggregation(personIdAgg);
+        }
        return new SearchRequest(poolId).source(builder);
    }
-    private List<Person> match0(SearchRequest searchRequest) throws Exception {
+    private List<Person> match0(SearchRequest searchRequest, Boolean agg) throws Exception {
        List<Person> persons = new ArrayList<>();
        SearchResponse search = client.search(searchRequest, RequestOptions.DEFAULT);
+        if (agg) {
            Aggregations aggregations = search.getAggregations();
            Terms byPersonId = aggregations.get("by_personId");
            List<? extends Terms.Bucket> buckets = byPersonId.getBuckets();
@@ -380,32 +387,33 @@ public class PersonService {
                Person person = new Person().setPersonId(personId).setScore((float) value);
                persons.add(person);
            }
+        } else {
+            SearchHits hits = search.getHits();
+            SearchHit[] hits1 = hits.getHits();
+            for (SearchHit item : hits1) {
+                Map<String, Object> source = item.getSourceAsMap();
+                Person p = new Person();
+                p.setPersonId((String) source.get("personId"));
+                p.setAge((Integer) source.get("age"));
+                p.setGender((String) source.get("gender"));
+                p.setChannelSerialNum((String) source.get("channelSerialNum"));
+                p.setBodyType((Integer) source.get("body_type"));
+                p.setCounttime(Optional.ofNullable((String) source.get("counttime"))
+                        .map(x -> {
+                            try {
+                                return Constant.DATE_FORMAT.get().parse(x);
+                            } catch (ParseException ignore) {
+                            }
+                            return null;
+                        })
+                        .orElse(null));
-//        SearchHits hits = search.getHits();
+                p.setScore(item.getScore());
-//        SearchHit[] hits1 = hits.getHits();
+                p.setPersonPoolId(item.getIndex());
-//        for (SearchHit item : hits1) {
+                persons.add(p);
-//            Map<String, Object> source = item.getSourceAsMap();
+            }
-//            Person p = new Person();
+        }
-//            p.setPersonId((String) source.get("personId"));
-//            p.setAge((Integer) source.get("age"));
-//            p.setGender((String) source.get("gender"));
-//            p.setChannelSerialNum((String) source.get("channelSerialNum"));
-//            p.setBodyType((Integer) source.get("body_type"));
-//
-//            p.setCounttime(Optional.ofNullable((String) source.get("counttime"))
-//                    .map(x -> {
-//                        try {
-//                            return Constant.DATE_FORMAT.get().parse(x);
-//                        } catch (ParseException ignore) {
-//                        }
-//                        return null;
-//                    })
-//                    .orElse(null));
-//
-//            p.setScore(item.getScore());
-//            p.setPersonPoolId(item.getIndex());
-//            persons.add(p);
-//        }
        return persons;
    }