<commit_before>package org.unipop.elastic.helpers;

import org.apache.tinkerpop.gremlin.process.traversal.P;
import org.apache.tinkerpop.gremlin.process.traversal.step.util.HasContainer;
import org.apache.tinkerpop.gremlin.structure.T;
import org.elasticsearch.action.get.*;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.index.query.QueryBuilders;
import org.unipop.controller.Predicates;
import org.unipop.helpers.LazzyGetter;
import org.unipop.structure.BaseVertex;

import java.util.*;

public class ElasticLazyGetter implements LazzyGetter{

    private static final int MAX_LAZY_GET = 50;
    private Client client;
    private TimingAccessor timing;
    private boolean executed = false;
    private HashMap<GetKey, List<BaseVertex>> keyToVertices = new HashMap();

    public ElasticLazyGetter(Client client, TimingAccessor timing) {
        this.client = client;
        this.timing = timing;
    }

    @Override
    public Boolean canRegister() {
        return !executed && keyToVertices.keySet().size() < MAX_LAZY_GET;
    }

    @Override
    public void register(BaseVertex v, String label, String indexName) {
        if(executed) System.out.println("This ElasticLazyGetter has already been executed.");

        GetKey key = new GetKey(v.id(), label, indexName);

        List<BaseVertex> vertices = keyToVertices.get(key);
        if (vertices == null) {
            vertices = new ArrayList();
            keyToVertices.put(key, vertices);
        }
        vertices.add(v);
    }

    @Override
    public void execute() {
        if (executed) return;
        executed = true;

        timing.start("lazyMultiGet");


        Set<String> types = new HashSet<>();
        Set<String> indices = new HashSet<>();
        Set<Object> ids = new HashSet<>();

        keyToVertices.keySet().forEach(getKey -> {
            types.add(getKey.type);
            ids.add(getKey.id);
            indices.add(getKey.indexName);
        });

        Predicates p = new Predicates();
        p.hasContainers.add(new HasContainer(T.id.getAccessor(), P.within(ids)));

        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(indices.toArray(new String[indices.size()]))
                .setQuery(QueryBuilders.filteredQuery(QueryBuilders.matchAllQuery(), ElasticHelper.createFilterBuilder(p.hasContainers))).setSize(ids.size());

        SearchResponse response = searchRequestBuilder.execute().actionGet();
        response.getHits().forEach(hit -> {
            keyToVertices.get(new GetKey(hit.id(), hit.type(), hit.getIndex())).forEach(baseVertex ->
            {
                Map<String,Object> source = hit.getSource();

                baseVertex.applyLazyFields(hit.type(), source);
            });
        });
        timing.stop("lazyMultiGet");

        keyToVertices = null;
        client = null;
    }

    private class GetKey {
        private final String id;
        private final String type;
        private final String indexName;<commit_msg>ElasticLazyGetter should use MultiGet API instead of the Search API in order to return the correct results, according to the id+type+index combination<commit_after>package org.unipop.elastic.helpers;

import org.apache.tinkerpop.gremlin.process.traversal.P;
import org.apache.tinkerpop.gremlin.process.traversal.step.util.HasContainer;
import org.apache.tinkerpop.gremlin.structure.T;
import org.elasticsearch.action.get.*;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.index.query.QueryBuilders;
import org.unipop.controller.Predicates;
import org.unipop.helpers.LazzyGetter;
import org.unipop.structure.BaseVertex;

import java.util.*;

public class ElasticLazyGetter implements LazzyGetter{

    private static final int MAX_LAZY_GET = 50;
    private Client client;
    private TimingAccessor timing;
    private boolean executed = false;
    private HashMap<GetKey, List<BaseVertex>> keyToVertices = new HashMap();

    public ElasticLazyGetter(Client client, TimingAccessor timing) {
        this.client = client;
        this.timing = timing;
    }

    @Override
    public Boolean canRegister() {
        return !executed && keyToVertices.keySet().size() < MAX_LAZY_GET;
    }

    @Override
    public void register(BaseVertex v, String label, String indexName) {
        if(executed) System.out.println("This ElasticLazyGetter has already been executed.");

        GetKey key = new GetKey(v.id(), label, indexName);

        List<BaseVertex> vertices = keyToVertices.get(key);
        if (vertices == null) {
            vertices = new ArrayList();
            keyToVertices.put(key, vertices);
        }
        vertices.add(v);
    }

    @Override
    public void execute() {
        if (executed) return;
        executed = true;

        timing.start("lazyMultiGet");

        MultiGetRequestBuilder multiGetRequestBuilder = client.prepareMultiGet();
        keyToVertices.keySet().forEach(key-> multiGetRequestBuilder.add(key.indexName, key.type, key.id));

        MultiGetResponse response = multiGetRequestBuilder.execute().actionGet();
        response.forEach(hit -> {
            List<BaseVertex> baseVertices = keyToVertices.get(new GetKey(hit.getId(), hit.getType(), hit.getIndex()));
            Map<String, Object> hitSource = hit.getResponse().getSource();
            baseVertices.forEach(baseVertex -> baseVertex.applyLazyFields(hit.getType(), hitSource));
        });
        timing.stop("lazyMultiGet");

        keyToVertices = null;
        client = null;
    }

    private class GetKey {
        private final String id;
        private final String type;
        private final String indexName;
{'fraction_non_alphanumeric': 0.10799796230259806, 'fraction_numerical': 0.0006792324673119375, 'mean_word_length': 3.1802696948190206, 'pattern_counts': {'":': 0, '<': 20, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3751527', 'n_tokens_mistral': 1672, 'n_tokens_neox': 1655, 'n_words': 354}