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

Go to download
package org.elasticsearch.plugin.nlpcn;

import com.google.common.collect.ImmutableMap;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.client.Client;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.common.xcontent.XContentType;

import org.elasticsearch.rest.BytesRestResponse;
import org.elasticsearch.rest.RestChannel;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.internal.InternalSearchHit;
import org.elasticsearch.search.internal.InternalSearchHits;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.nlpcn.es4sql.domain.Field;
import org.nlpcn.es4sql.exception.SqlParseException;
import org.nlpcn.es4sql.query.SqlElasticRequestBuilder;
import org.nlpcn.es4sql.query.join.HashJoinElasticRequestBuilder;
import org.nlpcn.es4sql.query.join.JoinRequestBuilder;
import org.nlpcn.es4sql.query.join.NestedLoopsElasticRequestBuilder;
import org.nlpcn.es4sql.query.join.TableInJoinRequestBuilder;

import java.io.IOException;
import java.util.*;

/**
 * Created by Eliran on 15/9/2015.
 */
public abstract class ElasticJoinExecutor implements ElasticHitsExecutor {
    protected SearchHits results ;
    protected MetaSearchResult metaResults;
    protected final int MAX_RESULTS_ON_ONE_FETCH = 10000;
    private Set aliasesOnReturn;
    private boolean allFieldsReturn;

    protected ElasticJoinExecutor(JoinRequestBuilder requestBuilder) {
        metaResults = new MetaSearchResult();
        aliasesOnReturn = new HashSet<>();
        List firstTableReturnedField = requestBuilder.getFirstTable().getReturnedFields();
        List secondTableReturnedField = requestBuilder.getSecondTable().getReturnedFields();
        allFieldsReturn = (firstTableReturnedField == null || firstTableReturnedField.size() == 0)
                            && (secondTableReturnedField == null || secondTableReturnedField.size() == 0);
    }

    public void  sendResponse(RestChannel channel){
        try {
            String json = ElasticUtils.hitsAsStringResult(results,metaResults);
            BytesRestResponse bytesRestResponse = new BytesRestResponse(RestStatus.OK, json);
            channel.sendResponse(bytesRestResponse);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public void run() throws IOException, SqlParseException {
        long timeBefore = System.currentTimeMillis();
        List combinedSearchHits =  innerRun();
        int resultsSize = combinedSearchHits.size();
        InternalSearchHit[] hits = combinedSearchHits.toArray(new InternalSearchHit[resultsSize]);
        this.results = new InternalSearchHits(hits, resultsSize,1.0f);
        long joinTimeInMilli = System.currentTimeMillis() - timeBefore;
        this.metaResults.setTookImMilli(joinTimeInMilli);
    }


    protected abstract List innerRun() throws IOException, SqlParseException ;

    public SearchHits getHits(){
        return results;
    }

    public static ElasticJoinExecutor createJoinExecutor(Client client, SqlElasticRequestBuilder requestBuilder){
        if(requestBuilder instanceof HashJoinElasticRequestBuilder) {
            HashJoinElasticRequestBuilder hashJoin = (HashJoinElasticRequestBuilder) requestBuilder;
            return new HashJoinElasticExecutor(client, hashJoin);
        }
        else if (requestBuilder instanceof NestedLoopsElasticRequestBuilder){
            NestedLoopsElasticRequestBuilder nestedLoops = (NestedLoopsElasticRequestBuilder) requestBuilder;
            return  new NestedLoopsElasticExecutor(client,nestedLoops);
        }
        else {
            throw new RuntimeException("Unsuported requestBuilder of type: " + requestBuilder.getClass());
        }
    }

    protected void mergeSourceAndAddAliases(Map secondTableHitSource, InternalSearchHit searchHit,String t1Alias,String t2Alias) {
        Map results = mapWithAliases(searchHit.getSource(), t1Alias);
        results.putAll(mapWithAliases(secondTableHitSource, t2Alias));
        searchHit.getSource().clear();
        searchHit.getSource().putAll(results);
    }

    protected Map mapWithAliases(Map source, String alias) {
        Map mapWithAliases = new HashMap<>();
        for(Map.Entry fieldNameToValue : source.entrySet()) {
            if(!aliasesOnReturn.contains(fieldNameToValue.getKey()))
                mapWithAliases.put(alias + "." + fieldNameToValue.getKey(), fieldNameToValue.getValue());
            else mapWithAliases.put(fieldNameToValue.getKey(),fieldNameToValue.getValue());
        }
        return mapWithAliases;
    }

    protected void  onlyReturnedFields(Map fieldsMap, List required,boolean allRequired) {
        HashMap filteredMap = new HashMap<>();
        if(allFieldsReturn || allRequired) {
            filteredMap.putAll(fieldsMap);
            return;
        }
        for(Field field: required){
            String name = field.getName();
            String returnName = name;
            String alias = field.getAlias();
            if(alias !=null && alias !=""){
                returnName = alias;
                aliasesOnReturn.add(alias);
            }
            filteredMap.put(returnName, deepSearchInMap(fieldsMap, name));
        }
        fieldsMap.clear();
        fieldsMap.putAll(filteredMap);

    }

    protected Object deepSearchInMap(Map fieldsMap, String name) {
        if(name.contains(".")){
            String[] path = name.split("\\.");
            Map currentObject = fieldsMap;
            for(int i=0;i) valueFromCurrentMap;
            }
            return currentObject.get(path[path.length-1]);
        }

        return fieldsMap.get(name);
    }


    protected void addUnmatchedResults(List combinedResults, Collection firstTableSearchHits, List secondTableReturnedFields,int currentNumOfIds, int totalLimit,String t1Alias,String t2Alias) {
        boolean limitReached = false;
        for(SearchHitsResult hitsResult : firstTableSearchHits){
            if(!hitsResult.isMatchedWithOtherTable())
                for (InternalSearchHit hit : hitsResult.getSearchHits()) {

                    //todo: decide which id to put or type. or maby its ok this way. just need to doc.
                    InternalSearchHit unmachedResult = createUnmachedResult(secondTableReturnedFields, hit.docId(), t1Alias, t2Alias, hit);
                    combinedResults.add(unmachedResult);
                    currentNumOfIds++;
                    if (currentNumOfIds >= totalLimit) {
                        limitReached = true;
                        break;
                    }

                }
            if(limitReached) break;
        }
    }

    protected InternalSearchHit createUnmachedResult( List secondTableReturnedFields, int docId, String t1Alias, String t2Alias, SearchHit hit) {
        String unmatchedId = hit.id() + "|0";
        Text unamatchedType = new Text(hit.getType() + "|null");

        InternalSearchHit searchHit = new InternalSearchHit(docId, unmatchedId, unamatchedType, hit.getFields());

        searchHit.sourceRef(hit.getSourceRef());
        searchHit.sourceAsMap().clear();
        searchHit.sourceAsMap().putAll(hit.sourceAsMap());
        Map emptySecondTableHitSource = createNullsSource(secondTableReturnedFields);

        mergeSourceAndAddAliases(emptySecondTableHitSource, searchHit,t1Alias,t2Alias);

        return searchHit;
    }

    protected Map createNullsSource(List secondTableReturnedFields) {
        Map nulledSource = new HashMap<>();
        for(Field field : secondTableReturnedFields){
            if(!field.getName().equals("*")){
                nulledSource.put(field.getName(),null);
            }
        }
        return nulledSource;
    }

    protected void updateMetaSearchResults( SearchResponse searchResponse) {
        this.metaResults.addSuccessfulShards(searchResponse.getSuccessfulShards());
        this.metaResults.addFailedShards(searchResponse.getFailedShards());
        this.metaResults.addTotalNumOfShards(searchResponse.getTotalShards());
        this.metaResults.updateTimeOut(searchResponse.isTimedOut());
    }

    protected SearchResponse scrollOneTimeWithMax(Client client,TableInJoinRequestBuilder tableRequest) {
        SearchResponse responseWithHits;SearchRequestBuilder scrollRequest = tableRequest.getRequestBuilder()
                .setScroll(new TimeValue(60000))
                .setSize(MAX_RESULTS_ON_ONE_FETCH);
        boolean ordered = tableRequest.getOriginalSelect().isOrderdSelect();
        if(!ordered) scrollRequest.addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC);
        responseWithHits = scrollRequest.get();
        //on ordered select - not using SCAN , elastic returns hits on first scroll
        //es5.0 elastic always return docs on scan
//        if(!ordered)
//            responseWithHits = client.prepareSearchScroll(responseWithHits.getScrollId()).setScroll(new TimeValue(600000)).get();
        return responseWithHits;
    }


}