<commit_before>package com.etisalat.log.query;

import com.etisalat.log.config.LogConfFactory;
import com.etisalat.log.parser.QueryCondition;
import com.etisalat.log.sort.SortUtils;
import com.google.gson.JsonObject;
import org.apache.hadoop.hbase.client.Get;
import org.apache.http.client.HttpClient;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;

public class SolrQueryTask implements Runnable {
    protected static Logger logger = LoggerFactory.getLogger(SolrQueryTask.class);
    private String cacheKey;
    private String shardId;
    private String reqUrl;
    private String qString;
    private String collection;
    private String rows;
    private HttpClient httpClient;
    private QueryCondition qCondition;
    private boolean lazyFetch = false;
    private int startRows = 0;
    private int fetchRows = 0;

    private HBaseQueryHandlerFactory queryHBaseHandlerFactory = null;

    public SolrQueryTask(String qString, String reqUrl) {
        this.qString = qString;
        this.reqUrl = reqUrl;
    }

    public String getCacheKey() {
        return cacheKey;
    }

    public void setCacheKey(String cacheKey) {
        this.cacheKey = cacheKey;
    }

    public String getCollection() {
        return collection;
    }

    public void setCollection(String collection) {
        this.collection = collection;
    }

    public HttpClient getHttpClient() {
        return httpClient;
    }

    public void setHttpClient(HttpClient httpClient) {
        this.httpClient = httpClient;
    }

    public QueryCondition getqCondition() {
        return qCondition;
    }

    public void setqCondition(QueryCondition qCondition) {
        this.qCondition = qCondition;
    }

    public String getqString() {
        return qString;
    }

    public void setqString(String qString) {
        this.qString = qString;
    }

    public HBaseQueryHandlerFactory getQueryHBaseHandlerFactory() {
        return queryHBaseHandlerFactory;
    }

    public void setQueryHBaseHandlerFactory(HBaseQueryHandlerFactory queryHBaseHandlerFactory) {
        this.queryHBaseHandlerFactory = queryHBaseHandlerFactory;
    }

    public String getReqUrl() {
        return reqUrl;
    }

    public void setReqUrl(String reqUrl) {
        this.reqUrl = reqUrl;
    }

    public String getRows() {
        return rows;
    }

    public void setRows(String rows) {
        this.rows = rows;
    }

    public String getShardId() {
        return shardId;
    }

    public void setShardId(String shardId) {
        this.shardId = shardId;
    }

    public boolean isLazyFetch() {
        return lazyFetch;
    }

    public void setLazyFetch(boolean lazyFetch) {
        this.lazyFetch = lazyFetch;
    }

    public int getStartRows() {
        return startRows;
    }

    public void setStartRows(int startRows) {
        this.startRows = startRows;
    }

    public int getFetchRows() {
        return fetchRows;
    }

    public void setFetchRows(int fetchRows) {
        this.fetchRows = fetchRows;
    }

    @Override
    public void run() {
        long start = System.currentTimeMillis();
        final SolrQuery myParameters = new SolrQuery();
        myParameters.add("q", qString);
        myParameters.add("rows", rows);
        myParameters.add("distrib", "false");
        myParameters.add("fl", LogConfFactory.rowkeyName);
        myParameters.set("collection", collection);

        final HttpSolrClient httpSolrClient = new HttpSolrClient(reqUrl, httpClient);
        try {
            logger.debug("Query session {}, query {}, QUERY request url is {}.", cacheKey, shardId,
                    httpSolrClient.getBaseURL());
            logger.debug("Query session {}, query {}, QUERY myParameters is {}.", cacheKey, shardId,
                    myParameters.toQueryString());

            QueryResponse response = httpSolrClient.query(myParameters);

            if (response.getStatus() != 0) {
                logger.error(
                        "Query session {}, Error(SolrServerException) sending live Query command, url is {}, query is {} with status return code {}",
                        cacheKey, reqUrl, qString, response.getStatus());
                
                checkAndAddResultToCache();
                return;
            }

            logger.info("Query session {}, {} with shardId {}, rows {}, solr cost {} ms", cacheKey, reqUrl,
                    shardId, rows, (System.currentTimeMillis() - start));
            fetchResultFromHBase(queryHBaseHandlerFactory, response);

        } catch (Exception e) {
            logger.error("Query session {}, query {}, Error(SolrServerException) sending live Query command, url is {}",
                    cacheKey, shardId, httpSolrClient.getBaseURL(), e);
        }
        
        checkAndAddResultToCache();
    }

    private void fetchResultFromHBase(HBaseQueryHandlerFactory queryHBaseHandlerFactory, QueryResponse queryResponse) {
        SolrDocumentList results = queryResponse.getResults();

        int hbaseSize = (results.size() < LogConfFactory.hbaseBatchSize * 10 ?
                LogConfFactory.hbaseBatchMinSize :
                LogConfFactory.hbaseBatchSize);

        logger.debug("Query session {}, query {}, get hbase row key with hbase batch size {}, result size {}", cacheKey, shardId, hbaseSize, 
        		results.size());

        CompletionService<HBaseQueryCursorRsp> completionService = queryHBaseHandlerFactory.newCompletionService();
        Set<Future<HBaseQueryCursorRsp>> pending = new HashSet<Future<HBaseQueryCursorRsp>>();

        List<Get> gets = new ArrayList<Get>();
        long start = System.currentTimeMillis();
        logger.debug("Query session {}, query {}, query DEBUG results size {} ", cacheKey, shardId, results.size());
        List<String> rowKeyList = new ArrayList<String>();

        int cnt = 0;
        for (int i = 0; i < results.size(); i++) {
            String solrKey = ((SolrDocument) results.get(i)).getFieldValue(LogConfFactory.rowkeyName).toString();
            rowKeyList.add(solrKey);
            if(i < startRows || (lazyFetch && cnt >= fetchRows)) {
                continue;
            }
            byte[] rowkey = SolrUtils.exchangeKey(solrKey);
            Get get = new Get(rowkey);
            gets.add(get);

            if (gets.size() == hbaseSize) {
                HBaseQueryCursorHandler hBaseQueryCursorHandler = new HBaseQueryCursorHandler(collection, cacheKey,
                        shardId);
                hBaseQueryCursorHandler.setGets(gets);
                hBaseQueryCursorHandler.setCompletionService(completionService);
                hBaseQueryCursorHandler.setPending(pending);
                hBaseQueryCursorHandler.submit();
                gets = new ArrayList<Get>();
            }

            cnt ++;
        }

        if (gets.size() != 0) {
            HBaseQueryCursorHandler hBaseQueryCursorHandler = new HBaseQueryCursorHandler(collection, cacheKey,
                    shardId);
            hBaseQueryCursorHandler.setGets(gets);
            hBaseQueryCursorHandler.setCompletionService(completionService);
            hBaseQueryCursorHandler.setPending(pending);
            hBaseQueryCursorHandler.submit();
        }

        Map<String, String> jsonObjectsMap = new HashMap<String, String>();
        while (pending.size() > 0) {
            try {
                Future<HBaseQueryCursorRsp> future = completionService.take();
                HBaseQueryCursorRsp rsp = future.get();
                pending.remove(future);
                if (rsp != null && rsp.getResultJsonObjMap() != null && rsp.getResultJsonObjMap().size() > 0) {
                    jsonObjectsMap.putAll(rsp.getResultJsonObjMap());
                }
            } catch (Exception e) {
                logger.error("Query session {}, query {}, hbase data fetch task failed and IOException arised",
                        cacheKey, shardId, e);
                continue;
            }
        }

        if (lazyFetch) {
            ConcurrentHashMap<String, List<String>> map = QueryBatch.RESULTS_FOR_SHARDS.get(cacheKey);
            if (map == null) {
                map = new ConcurrentHashMap<String, List<String>>();
                QueryBatch.RESULTS_FOR_SHARDS.put(cacheKey, map);
            }
            map.put(shardId, rowKeyList);

            addResultToCache2(jsonObjectsMap, rowKeyList);
        } else {
            addResultToCache(jsonObjectsMap, rowKeyList);
        }

        logger.info("Query session {}, query {}, hbase cost: {} ms", cacheKey, shardId, (System.currentTimeMillis() - start));
    }

    private void addResultToCache(Map<String, String> jsonObjectsMap, List<String> rowKeyList) {
        ConcurrentHashMap<String, List<String>> map = QueryBatch.RESULTS_FOR_SHARDS.get(cacheKey);
        if (map == null) {
            map = new ConcurrentHashMap<String, List<String>>();
            QueryBatch.RESULTS_FOR_SHARDS.put(cacheKey, map);
        }

        JsonObject jsonObject = new JsonObject();
        List<String> jsonObjects = new ArrayList<String>();
        String jsonStr = null;
        if (jsonObjectsMap.size() != 0) {
            for (String rowKey : rowKeyList) {
                jsonStr = jsonObjectsMap.get(rowKey);
                if (jsonStr == null) {
                    logger.debug("Query session {}, query {}, rowKey: {} does not have the related record.", cacheKey,
                            shardId, rowKey);
                    jsonObject.addProperty(LogConfFactory.rowkeyName, rowKey);
                    jsonStr = jsonObject.toString();
                }

                jsonObjects.add(jsonStr);
            }
        }

        map.put(shardId, jsonObjects);

        if (qCondition.isSort() && qCondition.isExportOp() && jsonObjects.size() != 0) {
            SortUtils.sortSingleShardRsp(qCondition.getSortedFields(), jsonObjects);
        }

        if (Integer.valueOf(rows) != jsonObjects.size()) {
        	logger.warn("Query session {}, query {}, solr result size {}, but hbase result size is {} ", cacheKey, shardId, rows, jsonObjects.size());
        }

        logger.info("Query session {}, query {}, add {} record to cache.", cacheKey, shardId, jsonObjects.size() );
    }

    private void addResultToCache2(Map<String, String> jsonObjectsMap, List<String> rowKeyList) {
        List<String> jsonObjects = QueryBatch.RESULTS_FOR_SHARDS.get(cacheKey).get(shardId);
        String jsonStr = null;
        int resSize = jsonObjectsMap.size();
        JsonObject jsonObject = new JsonObject();
        if (resSize != 0) {
            int i = startRows;
            String rowKey = null;
            for (; i < jsonObjects.size(); i++) {
                rowKey = rowKeyList.get(i);
                jsonStr = jsonObjectsMap.get(rowKey);
                if (jsonStr == null) {
                    logger.debug("Query session {}, query {}, rowKey: {} does not have the related record.", cacheKey,
                            shardId, rowKey);
                    jsonObject.addProperty(LogConfFactory.rowkeyName, rowKey);
                    jsonStr = jsonObject.toString();
                }

                jsonObjects.set(i, jsonStr);
            }
        }

        if (qCondition.isSort() && qCondition.isExportOp() && jsonObjects.size() != 0) {
            SortUtils.sortSingleShardRsp(qCondition.getSortedFields(), jsonObjects);
        }<commit_msg>Update SolrQueryTask.java<commit_after>package com.etisalat.log.query;

import com.etisalat.log.config.LogConfFactory;
import com.etisalat.log.parser.QueryCondition;
import com.etisalat.log.sort.SortUtils;
import com.google.gson.JsonObject;
import org.apache.hadoop.hbase.client.Get;
import org.apache.http.client.HttpClient;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;

public class SolrQueryTask implements Runnable {
    protected static Logger logger = LoggerFactory.getLogger(SolrQueryTask.class);
    private String cacheKey;
    private String shardId;
    private String reqUrl;
    private String qString;
    private String collection;
    private String rows;
    private HttpClient httpClient;
    private QueryCondition qCondition;
    private boolean lazyFetch = false;
    private int startRows = 0;
    private int fetchRows = 0;

    private HBaseQueryHandlerFactory queryHBaseHandlerFactory = null;

    public SolrQueryTask(String qString, String reqUrl) {
        this.qString = qString;
        this.reqUrl = reqUrl;
    }

    public String getCacheKey() {
        return cacheKey;
    }

    public void setCacheKey(String cacheKey) {
        this.cacheKey = cacheKey;
    }

    public String getCollection() {
        return collection;
    }

    public void setCollection(String collection) {
        this.collection = collection;
    }

    public HttpClient getHttpClient() {
        return httpClient;
    }

    public void setHttpClient(HttpClient httpClient) {
        this.httpClient = httpClient;
    }

    public QueryCondition getqCondition() {
        return qCondition;
    }

    public void setqCondition(QueryCondition qCondition) {
        this.qCondition = qCondition;
    }

    public String getqString() {
        return qString;
    }

    public void setqString(String qString) {
        this.qString = qString;
    }

    public HBaseQueryHandlerFactory getQueryHBaseHandlerFactory() {
        return queryHBaseHandlerFactory;
    }

    public void setQueryHBaseHandlerFactory(HBaseQueryHandlerFactory queryHBaseHandlerFactory) {
        this.queryHBaseHandlerFactory = queryHBaseHandlerFactory;
    }

    public String getReqUrl() {
        return reqUrl;
    }

    public void setReqUrl(String reqUrl) {
        this.reqUrl = reqUrl;
    }

    public String getRows() {
        return rows;
    }

    public void setRows(String rows) {
        this.rows = rows;
    }

    public String getShardId() {
        return shardId;
    }

    public void setShardId(String shardId) {
        this.shardId = shardId;
    }

    public boolean isLazyFetch() {
        return lazyFetch;
    }

    public void setLazyFetch(boolean lazyFetch) {
        this.lazyFetch = lazyFetch;
    }

    public int getStartRows() {
        return startRows;
    }

    public void setStartRows(int startRows) {
        this.startRows = startRows;
    }

    public int getFetchRows() {
        return fetchRows;
    }

    public void setFetchRows(int fetchRows) {
        this.fetchRows = fetchRows;
    }

    @Override
    public void run() {
        long start = System.currentTimeMillis();
        final SolrQuery myParameters = new SolrQuery();
        myParameters.add("q", qString);
        myParameters.add("rows", rows);
        myParameters.add("distrib", "false");
        myParameters.add("fl", LogConfFactory.rowkeyName);
        myParameters.set("collection", collection);

        final HttpSolrClient httpSolrClient = new HttpSolrClient(reqUrl, httpClient);
        try {
            logger.debug("Query session {}, query {}, QUERY request url is {}.", cacheKey, shardId,
                    httpSolrClient.getBaseURL());
            logger.debug("Query session {}, query {}, QUERY myParameters is {}.", cacheKey, shardId,
                    myParameters.toQueryString());

            QueryResponse response = httpSolrClient.query(myParameters);

            if (response.getStatus() != 0) {
                logger.error(
                        "Query session {}, Error(SolrServerException) sending live Query command, url is {}, query is {} with status return code {}",
                        cacheKey, reqUrl, qString, response.getStatus());
                
                checkAndAddResultToCache();
                return;
            }

            logger.info("Query session {}, {} with shardId {}, rows {}, solr cost {} ms", cacheKey, reqUrl,
                    shardId, rows, (System.currentTimeMillis() - start));
            fetchResultFromHBase(queryHBaseHandlerFactory, response);

        } catch (Exception e) {
            logger.error("Query session {}, query {}, Error(SolrServerException) sending live Query command, url is {}",
                    cacheKey, shardId, httpSolrClient.getBaseURL(), e);
        }
        
        checkAndAddResultToCache();
    }

    private void fetchResultFromHBase(HBaseQueryHandlerFactory queryHBaseHandlerFactory, QueryResponse queryResponse) {
        SolrDocumentList results = queryResponse.getResults();

        int hbaseSize = qCondition.isExportOp() ? LogConfFactory.hbaseBatchSize : ((results.size() < LogConfFactory.hbaseBatchSize * 10 ?
                LogConfFactory.hbaseBatchMinSize :
                LogConfFactory.hbaseBatchSize));

        logger.debug("Query session {}, query {}, get hbase row key with hbase batch size {}, result size {}", cacheKey, shardId, hbaseSize, 
        		results.size());

        CompletionService<HBaseQueryCursorRsp> completionService = queryHBaseHandlerFactory.newCompletionService();
        Set<Future<HBaseQueryCursorRsp>> pending = new HashSet<Future<HBaseQueryCursorRsp>>();

        List<Get> gets = new ArrayList<Get>();
        long start = System.currentTimeMillis();
        logger.debug("Query session {}, query {}, query DEBUG results size {} ", cacheKey, shardId, results.size());
        List<String> rowKeyList = new ArrayList<String>();

        int cnt = 0;
        for (int i = 0; i < results.size(); i++) {
            String solrKey = ((SolrDocument) results.get(i)).getFieldValue(LogConfFactory.rowkeyName).toString();
            rowKeyList.add(solrKey);
            if(i < startRows || (lazyFetch && cnt >= fetchRows)) {
                continue;
            }
            byte[] rowkey = SolrUtils.exchangeKey(solrKey);
            Get get = new Get(rowkey);
            gets.add(get);

            if (gets.size() == hbaseSize) {
                HBaseQueryCursorHandler hBaseQueryCursorHandler = new HBaseQueryCursorHandler(collection, cacheKey,
                        shardId);
                hBaseQueryCursorHandler.setGets(gets);
                hBaseQueryCursorHandler.setCompletionService(completionService);
                hBaseQueryCursorHandler.setPending(pending);
                hBaseQueryCursorHandler.submit();
                gets = new ArrayList<Get>();
            }

            cnt ++;
        }

        if (gets.size() != 0) {
            HBaseQueryCursorHandler hBaseQueryCursorHandler = new HBaseQueryCursorHandler(collection, cacheKey,
                    shardId);
            hBaseQueryCursorHandler.setGets(gets);
            hBaseQueryCursorHandler.setCompletionService(completionService);
            hBaseQueryCursorHandler.setPending(pending);
            hBaseQueryCursorHandler.submit();
        }

        Map<String, String> jsonObjectsMap = new HashMap<String, String>();
        while (pending.size() > 0) {
            try {
                Future<HBaseQueryCursorRsp> future = completionService.take();
                HBaseQueryCursorRsp rsp = future.get();
                pending.remove(future);
                if (rsp != null && rsp.getResultJsonObjMap() != null && rsp.getResultJsonObjMap().size() > 0) {
                    jsonObjectsMap.putAll(rsp.getResultJsonObjMap());
                }
            } catch (Exception e) {
                logger.error("Query session {}, query {}, hbase data fetch task failed and IOException arised",
                        cacheKey, shardId, e);
                continue;
            }
        }

        if (lazyFetch) {
            ConcurrentHashMap<String, List<String>> map = QueryBatch.RESULTS_FOR_SHARDS.get(cacheKey);
            if (map == null) {
                map = new ConcurrentHashMap<String, List<String>>();
                QueryBatch.RESULTS_FOR_SHARDS.put(cacheKey, map);
            }
            map.put(shardId, rowKeyList);

            addResultToCache2(jsonObjectsMap, rowKeyList);
        } else {
            addResultToCache(jsonObjectsMap, rowKeyList);
        }

        logger.info("Query session {}, query {}, hbase cost: {} ms", cacheKey, shardId, (System.currentTimeMillis() - start));
    }

    private void addResultToCache(Map<String, String> jsonObjectsMap, List<String> rowKeyList) {
        ConcurrentHashMap<String, List<String>> map = QueryBatch.RESULTS_FOR_SHARDS.get(cacheKey);
        if (map == null) {
            map = new ConcurrentHashMap<String, List<String>>();
            QueryBatch.RESULTS_FOR_SHARDS.put(cacheKey, map);
        }

        JsonObject jsonObject = new JsonObject();
        List<String> jsonObjects = new ArrayList<String>();
        String jsonStr = null;
        if (jsonObjectsMap.size() != 0) {
            for (String rowKey : rowKeyList) {
                jsonStr = jsonObjectsMap.get(rowKey);
                if (jsonStr == null) {
                    logger.debug("Query session {}, query {}, rowKey: {} does not have the related record.", cacheKey,
                            shardId, rowKey);
                    jsonObject.addProperty(LogConfFactory.rowkeyName, rowKey);
                    jsonStr = jsonObject.toString();
                }

                jsonObjects.add(jsonStr);
            }
        }

        map.put(shardId, jsonObjects);

        if (qCondition.isSort() && qCondition.isExportOp() && jsonObjects.size() != 0) {
            SortUtils.sortSingleShardRsp(qCondition.getSortedFields(), jsonObjects);
        }

        if (Integer.valueOf(rows) != jsonObjects.size()) {
        	logger.warn("Query session {}, query {}, solr result size {}, but hbase result size is {} ", cacheKey, shardId, rows, jsonObjects.size());
        }

        logger.info("Query session {}, query {}, add {} record to cache.", cacheKey, shardId, jsonObjects.size() );
    }

    private void addResultToCache2(Map<String, String> jsonObjectsMap, List<String> rowKeyList) {
        List<String> jsonObjects = QueryBatch.RESULTS_FOR_SHARDS.get(cacheKey).get(shardId);
        String jsonStr = null;
        int resSize = jsonObjectsMap.size();
        JsonObject jsonObject = new JsonObject();
        if (rowKeyList.size() > 0 || resSize != 0) {
            int i = startRows;
            String rowKey = null;
            for (; i < jsonObjects.size(); i++) {
                rowKey = rowKeyList.get(i);
                jsonStr = jsonObjectsMap.get(rowKey);
                if (jsonStr == null) {
                    logger.debug("Query session {}, query {}, rowKey: {} does not have the related record.", cacheKey,
                            shardId, rowKey);
                    jsonObject.addProperty(LogConfFactory.rowkeyName, rowKey);
                    jsonStr = jsonObject.toString();
                }

                jsonObjects.set(i, jsonStr);
            }
        }

        if (qCondition.isSort() && qCondition.isExportOp() && jsonObjects.size() != 0) {
            SortUtils.sortSingleShardRsp(qCondition.getSortedFields(), jsonObjects);
        }
{'fraction_non_alphanumeric': 0.10271814477166744, 'fraction_numerical': 0.00155442591270008, 'mean_word_length': 2.3460781557492267, 'pattern_counts': {'":': 0, '<': 67, '<?xml version=': 0, '>': 66, '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': '5324154', 'n_tokens_mistral': 6572, 'n_tokens_neox': 6315, 'n_words': 1397}