com.fnklabs.draenei.analytics.search.SearchTask Maven / Gradle / Ivy

Go to download
package com.fnklabs.draenei.analytics.search;

import com.codahale.metrics.Timer;
import com.fnklabs.draenei.MetricsFactoryImpl;
import org.apache.ignite.Ignite;
import org.apache.ignite.IgniteException;
import org.apache.ignite.cluster.ClusterNode;
import org.apache.ignite.compute.ComputeJobResult;
import org.apache.ignite.compute.ComputeTaskAdapter;
import org.apache.ignite.resources.IgniteInstanceResource;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

class SearchTask extends ComputeTaskAdapter> {
    /**
     *
     */
    @NotNull
    private final Set searchFacets;

    private Ignite ignite;


    public SearchTask(@NotNull Set searchFacets) {
        this.searchFacets = searchFacets;
    }

    @Nullable
    @Override
    public Map map(List subgrid, @Nullable Object arg) throws IgniteException {
        Map map = new HashMap<>(subgrid.size());

        subgrid.forEach(clusterNode -> {
            map.put(new SearchJob(searchFacets), clusterNode);
        });

        return map;
    }

    @Nullable
    @Override
    public List reduce(List results) throws IgniteException {

        Timer.Context totalTimer = MetricsFactoryImpl.getTimer("search_task.reduce.total").time();

        Set documentIndexes = results.stream()
                                                    .flatMap(computeJobResult -> {
                                                        return computeJobResult.>getData().stream();
                                                    })
                                                    .collect(Collectors.toSet());

        Map> facetKeyListMap = documentIndexes.stream()
                                                                    .flatMap(documentIndex -> {
                                                                        Set facets = documentIndex.getFacets()
                                                                                                         .stream()
                                                                                                         .collect(Collectors.toSet());

                                                                        return facets.stream();
                                                                    })
                                                                    .collect(Collectors.groupingBy(new Function() {
                                                                        @Override
                                                                        public FacetKey apply(Facet facet) {
                                                                            return facet.getKey();
                                                                        }
                                                                    }));


        int totalDocuments = ignite.getOrCreateCache(SearchServiceImpl.getDocumentsCacheConfiguration()).size();

        Map facetIdfMap = new HashMap<>();

        facetKeyListMap.entrySet()
                       .forEach(entry -> {
                           double idf = TfIdfUtils.calculateIdf(entry.getValue().size(), totalDocuments);

                           facetIdfMap.put(entry.getKey(), idf);
                       });

        SimilarityCosineAlgorithm algorithm = new SimilarityCosineAlgorithm(MetricsFactoryImpl.getInstance());

        List searchFacetTfIdf = searchFacets.stream()
                                                   .map(facet -> {
                                                       FacetKey facetKey = facet.getKey();
                                                       double tfIdf = TfIdfUtils.calculateTfIdf(facet.getRank(), facetIdfMap.getOrDefault(facetKey, 0d));
                                                       return new Facet(facetKey, tfIdf, facet.getDocument());
                                                   })
                                                   .collect(Collectors.toList());
        Timer.Context mapDocumentsTimer = MetricsFactoryImpl.getTimer("search_task.reduce.documents_map").time();

        List searchResults = documentIndexes.stream()
                                                          .map(documentIndex -> {
                                                              Set tfIdfFacets = documentIndex.getFacets()
                                                                                                    .stream()
                                                                                                    .map(facet -> {
                                                                                                        FacetKey facetKey = facet.getKey();
                                                                                                        double tfIdf = TfIdfUtils.calculateTfIdf(facet.getRank(), facetIdfMap.getOrDefault(facetKey, 0d));
                                                                                                        return new Facet(facetKey, tfIdf, facet.getDocument());
                                                                                                    })
                                                                                                    .collect(Collectors.toSet());


                                                              double similarity = algorithm.getSimilarity(tfIdfFacets, searchFacetTfIdf);

                                                              return new SearchResult(documentIndex.getDocument(), similarity);
                                                          })
                                                          .sorted()
                                                          .collect(Collectors.toList());

        mapDocumentsTimer.stop();
        totalTimer.stop();

        return searchResults;
    }

    @IgniteInstanceResource
    protected void setIgnite(Ignite ignite) {
        this.ignite = ignite;
    }
}