org.elasticsearch.plugin.nlpcn.UnionExecutor Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of elasticsearch-sql Show documentation
Query elasticsearch using SQL
The newest version!
package org.elasticsearch.plugin.nlpcn;

import org.apache.lucene.search.TotalHits;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.nlpcn.es4sql.Util;
import org.nlpcn.es4sql.exception.SqlParseException;
import org.nlpcn.es4sql.query.multi.MultiQueryRequestBuilder;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * Created by Eliran on 21/8/2016.
 */
public class UnionExecutor implements ElasticHitsExecutor {

    private MultiQueryRequestBuilder multiQueryBuilder;
    private SearchHits results;
    private Client client;
    private int currentId;

    public UnionExecutor(Client client,MultiQueryRequestBuilder builder) {
        multiQueryBuilder = builder;
        this.client = client;
        currentId = 0;
    }

    @Override
    public void run() throws IOException, SqlParseException {
        SearchResponse firstResponse = this.multiQueryBuilder.getFirstSearchRequest().get();
        SearchHit[] hits = firstResponse.getHits().getHits();
        List unionHits = new ArrayList<>(hits.length);
        fillInternalSearchHits(unionHits,hits,this.multiQueryBuilder.getFirstTableFieldToAlias());
        SearchResponse secondResponse = this.multiQueryBuilder.getSecondSearchRequest().get();
        fillInternalSearchHits(unionHits,secondResponse.getHits().getHits(),this.multiQueryBuilder.getSecondTableFieldToAlias());
        int totalSize = unionHits.size();
        SearchHit[] unionHitsArr = unionHits.toArray(new SearchHit[totalSize]);
        this.results = new SearchHits(unionHitsArr, new TotalHits(totalSize, TotalHits.Relation.EQUAL_TO), 1.0f);
    }

    private void fillInternalSearchHits(List unionHits, SearchHit[] hits, Map fieldNameToAlias) {
        for(SearchHit hit : hits){
            SearchHit searchHit = new SearchHit(currentId, hit.getId(), new Text(hit.getType()), hit.getFields(), null);
            searchHit.sourceRef(hit.getSourceRef());
            searchHit.getSourceAsMap().clear();
            Map sourceAsMap = hit.getSourceAsMap();
            if(!fieldNameToAlias.isEmpty()){
                updateFieldNamesToAlias(sourceAsMap, fieldNameToAlias);
            }
            searchHit.getSourceAsMap().putAll(sourceAsMap);
            currentId++;
            unionHits.add(searchHit);
        }
    }


    private void updateFieldNamesToAlias(Map sourceAsMap, Map fieldNameToAlias) {
        for(Map.Entry fieldToAlias : fieldNameToAlias.entrySet()){
            String fieldName = fieldToAlias.getKey();
            Object value = null;
            Map deleteFrom = null;
            if(fieldName.contains(".")){
                String[] split = fieldName.split("\\.");
                String[] path = Arrays.copyOf(split, split.length - 1);
                Object placeInMap = Util.searchPathInMap(sourceAsMap, path);
                if(placeInMap != null){
                    if(!Map.class.isAssignableFrom(placeInMap.getClass())){
                        continue;
                    }
                }
                deleteFrom = (Map) placeInMap;
                value = deleteFrom.get(split[split.length-1]);
            }
            else if(sourceAsMap.containsKey(fieldName)){
                value = sourceAsMap.get(fieldName);
                deleteFrom = sourceAsMap;
            }
            if(value!=null){
                sourceAsMap.put(fieldToAlias.getValue(),value);
                deleteFrom.remove(fieldName);
            }
        }
        Util.clearEmptyPaths(sourceAsMap);
    }

    @Override
    public SearchHits getHits() {
        return results;
    }
}