Elasticsearch scroll java

package com.ygsoft.jvm.read;

import java.net.InetAddress;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequestBuilder;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.InetSocketTransportAddress;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.PipelineAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramInterval;
import org.elasticsearch.search.aggregations.bucket.histogram.Histogram.Order;
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregatorBuilders;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
import org.elasticsearch.transport.client.PreBuiltTransportClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.ygsoft.jvm.dto.AggregatePoorFormulaDTO;
import com.ygsoft.jvm.dto.EsAggsVO;
import com.ygsoft.jvm.dto.EsQueryVO;
import com.ygsoft.jvm.type.AggregationType;
import com.ygsoft.jvm.type.EsAggregationType;
import com.ygsoft.jvm.type.EsQueryType;
import com.ygsoft.jvm.utils.StringUtils;

/**
 * ES Client
 * 
 * @author duhai
 * @date 2019年5月15日
 */
public class ElasticsearchUtils {

	private static final Logger logger = LoggerFactory.getLogger(ElasticsearchUtils.class);

	private static TransportClient client;

	private final static String TIMESTAMP = "timestamp";

	private final static String DEVICEID = "deviceid";

	private final static String METRIC = "metric";

	private final static String VALUE = "value";

	public static final int BULK_SIZE = 5000;

	// 獲取線程池
	public static TransportClient getClient() {
		return client;
	}

	/**
	 * 
	 * @param hostName
	 *            elk集羣地址
	 * @param port
	 *            端口
	 * @param clusterName
	 *            集羣名稱
	 * @param poolSize
	 *            連接池
	 * @return
	 */
	@SuppressWarnings("unchecked")
	public TransportClient init(final String hostName, final String port) {
		try {
			if (client == null) {
				// 配置信息
				Settings esSetting = Settings.builder()// 初始化
						.put("cluster.name", "elsearch")// 名字
						.put("client.transport.sniff", true)// 增加嗅探機制，找到ES集羣
						.put("thread_pool.search.size", 5)// 增加線程池個數，暫時設爲5
						.build();

				client = new PreBuiltTransportClient(esSetting);
				InetSocketTransportAddress inetSocketTransportAddress = new InetSocketTransportAddress(
						InetAddress.getByName(hostName), Integer.valueOf(port));
				client.addTransportAddresses(inetSocketTransportAddress);
			}
		} catch (Exception e) {
			logger.error("elasticsearch 客戶端 創建 error!!!", e);
		}
		return client;
	}

	/**
	 * 條件查詢出數據
	 * 
	 * @param index
	 *            es的查詢index
	 * @param type
	 *            VirtualMetricType表類型
	 * @param dto
	 *            聚合查詢參數對象
	 * @return
	 */
	public JSONArray caculateVirtualMetricQuery(final String index, final AggregatePoorFormulaDTO dto) {
		final JSONArray result = new JSONArray();
		/*
		 * 1.1、設置查詢
		 */
		final BoolQueryBuilder mainBoolQueryBuild = QueryBuilders.boolQuery();

		// 開始時間
		final String startTime = dto.getStartTime();
		// 結束時間
		final String endTime = dto.getEndTime();
		final EsQueryVO rangeVO = new EsQueryVO();
		rangeVO.setType(EsQueryType.RANGEQUERY);
		rangeVO.setStartTime(startTime);
		rangeVO.setEndTime(endTime);
		rangeVO.setCondition(TIMESTAMP);
		final QueryBuilder query1 = buildQueryBuild(rangeVO);
		mainBoolQueryBuild.filter(query1);// 設置查詢條件

		// 設備code
		final EsQueryVO deviceIdTermVO = new EsQueryVO();
		deviceIdTermVO.setType(EsQueryType.TERM);
		deviceIdTermVO.setCondition(DEVICEID);
		deviceIdTermVO.setValue(dto.getDeviceCode());
		final QueryBuilder deviceIdTermQuery = buildQueryBuild(deviceIdTermVO);
		mainBoolQueryBuild.filter(deviceIdTermQuery);

		// 測點列表
		final BoolQueryBuilder orBoolQueryBuild = QueryBuilders.boolQuery();
		// {"metricA":"3080"}
		final JSONObject params = dto.getParams();
		if (null == params || params.size() <= 0) {
			return result;
		}
		for (Entry<String, Object> entry : params.entrySet()) {
			// 測點
			final EsQueryVO metricTermVO = new EsQueryVO();
			metricTermVO.setType(EsQueryType.TERM);
			metricTermVO.setCondition(METRIC);
			metricTermVO.setValue(entry.getValue().toString());
			final QueryBuilder metricTermQuery = buildQueryBuild(metricTermVO);
			// 滿足一個要求的就可以
			orBoolQueryBuild.should(metricTermQuery);
		}

		mainBoolQueryBuild.filter(orBoolQueryBuild);
		final QueryBuilder builder1 = QueryBuilders.rangeQuery(VALUE).lt(-10000);
		final QueryBuilder builder2 = QueryBuilders.rangeQuery(VALUE).gt(10000);
		mainBoolQueryBuild.mustNot(builder1);
		mainBoolQueryBuild.mustNot(builder2);

		// 設置查詢 TODO
		SearchRequestBuilder requestBuilder = client.prepareSearch(index).setTypes(index)
				.setSearchType(SearchType.DFS_QUERY_THEN_FETCH);
		requestBuilder.setQuery(mainBoolQueryBuild);
		requestBuilder.addSort(SortBuilders.fieldSort("timestamp").order(SortOrder.DESC));
		requestBuilder.setFrom(0);
		requestBuilder.setScroll(TimeValue.timeValueMinutes(2));
		requestBuilder.setSize(10000);
		// requestBuilder.setExplain(true);

		logger.info("requestBuilder查詢:");
		logger.info(requestBuilder.toString());

		SearchResponse searchResponse = requestBuilder.execute().actionGet();
		do {
			Long time1 = System.currentTimeMillis();
			// 解析結果
			final SearchHit[] hits = searchResponse.getHits().getHits();
			if (hits.length > 0) {
				for (SearchHit hit : hits) {
					final Map<String, Object> map = hit.getSource();
					final JSONObject queryResultDTO = JSONObject.parseObject(JSONObject.toJSONString(map));
					result.add(queryResultDTO);
				}
			}
			Long time2 = System.currentTimeMillis();
			logger.info("1萬條添加完成時間" + (time2 - time1));
			SearchScrollRequestBuilder searchScrollRequestBuilder = client
					.prepareSearchScroll(searchResponse.getScrollId()).setScroll(TimeValue.timeValueMinutes(4));
			logger.info("searchScrollRequestBuilder查詢:");
			logger.info(requestBuilder.toString());
			searchResponse = searchScrollRequestBuilder.execute().actionGet();
			logger.info("1萬查詢完成時間" + (System.currentTimeMillis() - time2));
		} while (searchResponse.getHits().getHits().length != 0);
		logger.info("-----------------------全部查詢" + result.size() + "條");
		return result;
	}

	/**
	 * 構建查詢語句
	 * 
	 * @param vo
	 *            查詢參數
	 * @return QueryBuilder 查詢結構
	 */
	public QueryBuilder buildQueryBuild(final EsQueryVO vo) {
		QueryBuilder builder = null;
		if (vo.getType() == EsQueryType.RANGEQUERY) {
			builder = QueryBuilders.rangeQuery(vo.getCondition()).gte(vo.getStartTime()).lte(vo.getEndTime());
		} else if (vo.getType() == EsQueryType.TERM) {
			builder = QueryBuilders.termQuery(vo.getCondition(), vo.getValue());
		} else if (vo.getType() == EsQueryType.TERMS) {
			builder = QueryBuilders.termsQuery(vo.getCondition(), vo.getMultiValue());
		}

		return builder;
	}

	/**
	 * 設置聚合
	 * 
	 * @param vo
	 * @return
	 */
	public AggregationBuilder buildAggregationModel(final EsAggsVO vo) {
		AggregationBuilder aggregation = null;
		if (vo.getType() == EsAggregationType.DATEHISTOGRAM) {
			final DateHistogramInterval interval = new DateHistogramInterval(vo.getInterval());
			aggregation = AggregationBuilders.dateHistogram(AggregationType.DATE_HISTOGRAM.getName())
					.field(vo.getValue()).format(vo.getFormat()).dateHistogramInterval(interval)
					.minDocCount(vo.getMinDocCount()).order(vo.getOrder() == 0 ? Order.KEY_ASC : Order.KEY_DESC);
		} else if (vo.getType() == EsAggregationType.MAX) {
			aggregation = AggregationBuilders.max(AggregationType.MAX.getName()).field(vo.getValue());
		} else if (vo.getType() == EsAggregationType.MIN) {
			aggregation = AggregationBuilders.min(AggregationType.MIN.getName()).field(vo.getValue());
		} else if (vo.getType() == EsAggregationType.SUM) {
			aggregation = AggregationBuilders
					.sum(StringUtils.isEmpty(vo.getBucketStr()) ? AggregationType.SUM.getName() : vo.getBucketStr())
					.field(vo.getValue());
		} else if (vo.getType() == EsAggregationType.AVG) {
			aggregation = AggregationBuilders.avg(AggregationType.AVG.getName()).field(vo.getValue());
		} else if (vo.getType() == EsAggregationType.FILTER) {
			final QueryBuilder build = QueryBuilders.termQuery(vo.getKey(), vo.getValue());
			aggregation = AggregationBuilders.filter(vo.getValue(), build);
		} else if (vo.getType() == EsAggregationType.TERMS) {
			aggregation = AggregationBuilders.terms(AggregationType.TERMS.getName()).field(vo.getKey()).size(100);
		}

		return aggregation;
	}

	/**
	 * 設置聚合
	 * 
	 * @param vo
	 * @return
	 */
	public PipelineAggregationBuilder buildPipelineAggregationModel(final EsAggsVO vo) {
		PipelineAggregationBuilder aggregation = null;
		String bucketName = vo.getBucketStr();
		if (vo.getType() == EsAggregationType.MAX_BUCKET) {
			aggregation = PipelineAggregatorBuilders.maxBucket(bucketName, vo.getBucketName());
		} else if (vo.getType() == EsAggregationType.MIN_BUCKET) {
			aggregation = PipelineAggregatorBuilders.minBucket(bucketName, vo.getBucketName());
		} else if (vo.getType() == EsAggregationType.AVG_BUCKET) {
			aggregation = PipelineAggregatorBuilders.avgBucket(bucketName, vo.getBucketName());
		} else if (vo.getType() == EsAggregationType.SUM_BUCKET) {
			aggregation = PipelineAggregatorBuilders.sumBucket(bucketName, vo.getBucketName());
		} else if (vo.getType() == EsAggregationType.POOR_FORMULE) {
			Map<String, String> bucketsPathsMap = new HashMap<String, String>();
			if (null == vo.getBucketsPathsMap()) {
				bucketsPathsMap.put(AggregationType.MAX.getName(), AggregationType.MAX.getName());
				bucketsPathsMap.put(AggregationType.MIN.getName(), AggregationType.MIN.getName());
			} else {
				bucketsPathsMap = vo.getBucketsPathsMap();
			}

			String model = vo.getBucketModel();
			if (StringUtils.isEmpty(model)) {
				model = String.format("params.%s - params.%s", AggregationType.MAX.getName(),
						AggregationType.MIN.getName());
			}

			String name = vo.getBucketStr();
			if (null == name || name.isEmpty()) {
				name = AggregationType.POOR_FORMULA.getName();
			}

			final Script script = new Script(model);
			aggregation = PipelineAggregatorBuilders.bucketScript(name, bucketsPathsMap, script);
		}

		return aggregation;
	}

}