search-parent/src/main/java/com/uas/search/util/SearchUtils.java

package com.uas.search.util;

import com.uas.search.LuceneProperties;
import com.uas.search.analyzer.IKAnalyzer;
import com.uas.search.constant.SearchConstants;
import com.uas.search.constant.model.SPage;
import com.uas.search.exception.SearchException;
import com.uas.search.model.*;
import com.uas.search.support.IndexSearcherManager;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 搜索相关的工具类
 * 
 * @author sunyj
 * @since 2016年10月17日 下午1:56:16
 */
public class SearchUtils {

	/**
	 * 默认的页码
	 */
	private static final int PAGE_INDEX = 1;

	/**
	 * 默认每页的大小
	 */
	private static final int PAGE_SIZE = 20;

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

	/**
	 * IndexSearcher的管理器
	 */
	private static IndexSearcherManager searcherManager = IndexSearcherManager.getInstance();

	private static LuceneProperties luceneProperties = ContextUtils.getBean(LuceneProperties.class);

	/**
	 * 判断搜索词是否为无效的（比如只包含特殊字符，不含有任何字母、数字、汉字等有意义的字符）
	 * 
	 * @param keyword
	 * @return 搜索词是否无效
	 */
	public static boolean isKeywordInvalid(String keyword) {
		if (!StringUtils.isEmpty(keyword)) {
			// 将特殊字符剔除
			keyword = keyword.replaceAll("[`~!@#$^&*()=|{}':;',\\[\\].<>/?~！@#￥……&*（）——|{}【】‘；：”“'。，、？ ]+", "");
			return StringUtils.isEmpty(keyword);
		}
		return true;
	}

	/**
	 * 对搜索词进行分词后组合得到BooleanQuery
	 * 
	 * @param field
	 *            搜索的域名
	 * @param keyword
	 *            搜索关键词
	 * @return
	 */
	public static BooleanQuery getBooleanQuery(String field, String keyword) {
		return getBooleanQuery(field, keyword, true, true, Occur.MUST);
	}

	/**
	 * 对搜索词进行分词后组合得到BooleanQuery
	 * 
	 * @param field
	 *            搜索的域名
	 * @param keyword
	 *            搜索关键词
     * @param useSmart
     *            是否以最大粒度进行分词
	 * @param loadExtDic
	 *            是否加载自定义词典
	 * @param occur
	 *            多个Query之间的关系
	 * @return
	 */
	public static BooleanQuery getBooleanQuery(String field, String keyword, boolean useSmart, boolean loadExtDic, Occur occur) {
		if (StringUtils.isEmpty(field) || StringUtils.isEmpty(keyword)) {
			return null;
		}
		BooleanQuery booleanQuery = new BooleanQuery();
		Analyzer analyzer = new IKAnalyzer(useSmart, loadExtDic);
		try {
			TokenStream tokenStream = analyzer.tokenStream(field, keyword);
			tokenStream.reset();
			CharTermAttribute cta = tokenStream.addAttribute(CharTermAttribute.class);
			while (tokenStream.incrementToken()) {
				booleanQuery.add(new PrefixQuery(new Term(field, cta.toString())), occur);
			}
			tokenStream.close();
			analyzer.close();
		} catch (IOException e) {
			logger.error("", e);
		}
		return booleanQuery;
	}

	/**
	 * 对搜索词进行分词后组合得到RegexpQuery
	 * 
	 * @param field
	 *            搜索的域名
	 * @param keyword
	 *            搜索关键词
	 * @return
	 */
	// 如果用正则表达式进行搜索，就不需要再分词（分词的字段可能会导致边界问题）
	public static RegexpQuery getRegexpQuery(String field, String keyword) {
		if (StringUtils.isEmpty(field)) {
			return null;
		}
		// TODO 针对的是分词了的字段（未分词的字段可能搜索不准确）
		return new RegexpQuery(new Term(field, ".*" + keyword.toLowerCase() + ".*"));
	}

	/**
	 * 获取指定表的IndexSearcher对象，若为空，抛出异常
	 * 
	 * @param tableName
	 *            表名
	 * @return IndexSearcher对象
	 */
	public static IndexSearcher getIndexSearcher(String tableName) {
		IndexSearcher indexSearcher = searcherManager.get(tableName);
		if (indexSearcher == null) {
			throw new SearchException("获取索引文件失败");
		}
		return indexSearcher;
	}

	/**
	 * 释放indexSearcher
	 * 
	 * @param indexSearcher
	 */
	public static void releaseIndexSearcher(IndexSearcher indexSearcher) {
		searcherManager.release(indexSearcher);
	}

	/**
	 * 根据域和搜索词获取指定表的Document（每个id最多只能对应一条数据）
	 * 
	 * @param tableName
	 *            表名
	 * @param field
	 *            搜索域
	 * @param id
	 *            搜索词
	 * @return Document
	 */
	public static Document getDocumentById(String tableName, String field, Long id) {
		if (id == null) {
			throw new SearchException("输入无效：" + id);
		}
		return getDocumentById(tableName, field, Long.toString(id));
	}

	/**
	 * 根据域和搜索词获取指定表的Document（每个id最多只能对应一条数据）
	 * 
	 * @param tableName
	 *            表名
	 * @param field
	 *            搜索域
	 * @param id
	 *            搜索词
	 * @return Document
	 */
	public static Document getDocumentById(String tableName, String field, String id) {
		if (StringUtils.isEmpty(field) || StringUtils.isEmpty(id)) {
			throw new SearchException("搜索的域和搜索词不能为空");
		}
		TermQuery termQuery = new TermQuery(new Term(field, id));
		List<Document> documents = getDocuments(tableName, termQuery).getContent();
		if (CollectionUtils.isEmpty(documents)) {
			return null;
		} else if (documents.size() > 1) {
			throw new SearchException("索引中存在不止一个对象：" + field + "=" + id);
		}
		return documents.get(0);
	}

	/**
	 * 根据查询条件获取指定表的Document列表
	 * 
	 * @param tableName
	 *            表名
	 * @param query
	 *            查询条件
	 * @return
	 */
	public static SPage<Document> getDocuments(String tableName, Query query) {
		return getDocuments(tableName, query, null, null);
	}

	/**
	 * 根据查询条件获取指定表的Document列表
	 * 
	 * @param tableName
	 *            表名
	 * @param query
	 *            查询条件
	 * @param page
	 *            页码
	 * @param size
	 *            页大小
	 * @return
	 */
	public static SPage<Document> getDocuments(String tableName, Query query, Integer page, Integer size) {
		return getDocuments(tableName, query, null, page, size);
	}

	/**
	 * 根据查询条件获取指定表的Document列表
	 *
	 * @param tableName
	 *            表名
	 * @param query
	 *            查询条件
	 * @param sort 排序规则
	 * @param page
	 *            页码
	 * @param size
	 *            页大小
	 * @return
	 */
	public static SPage<Document> getDocuments(String tableName, Query query, Sort sort, Integer page, Integer size) {
		if (query == null) {
			throw new SearchException("query不能为null");
		}
		SPage<Document> sPage = new SPage<>();
		if (page != null && page > 0) {
			sPage.setPage(page);
		} else {
			sPage.setPage(PAGE_INDEX);
			sPage.setFirst(true);
		}
		if (size != null && size > 0) {
			sPage.setSize(size);
		} else {
			sPage.setSize(PAGE_SIZE);
		}
		IndexSearcher indexSearcher = getIndexSearcher(tableName);

		TopDocs topDocs;
		try {
			// 如果页码不为1
			if (sPage.getPage() > 1) {
				TopDocs previousTopDocs;
				previousTopDocs = sort == null ? indexSearcher.search(query, (sPage.getPage() - 1) * sPage.getSize()) :
						indexSearcher.search(query, (sPage.getPage() - 1) * sPage.getSize(), sort);
				int totalHits = previousTopDocs.totalHits;
				ScoreDoc[] previousScoreDocs = previousTopDocs.scoreDocs;
				if ((sPage.getPage() - 1) * sPage.getSize() >= totalHits) {
					throw new SearchException("页码过大：元素总数量为" + totalHits);
				}
				topDocs = sort == null ? indexSearcher.searchAfter(previousScoreDocs[previousScoreDocs.length - 1], query, sPage.getSize()) :
						indexSearcher.searchAfter(previousScoreDocs[previousScoreDocs.length - 1], query, sPage.getSize(), sort);
			} else {
				sPage.setFirst(true);
				topDocs = sort == null ? indexSearcher.search(query, sPage.getSize()) : indexSearcher.search(query, sPage.getSize(), sort);
			}

			int totalHits = topDocs.totalHits;
			// 设置总元素个数、页数等信息
			sPage.setTotalElement(totalHits);
			int totalPage = (int) Math.ceil(totalHits / (1.0 * sPage.getSize()));
			sPage.setTotalPage(totalPage);
			if (totalPage == sPage.getPage()) {
				sPage.setLast(true);
			}

			List<Document> documents = new ArrayList<>();
			for (ScoreDoc scoreDoc : topDocs.scoreDocs) {
				documents.add(indexSearcher.doc(scoreDoc.doc));
			}

			sPage.setContent(documents);
		} catch (IOException e) {
			throw new SearchException(e).setDetailedMessage(e);
		} finally {
			releaseIndexSearcher(indexSearcher);
		}
		return sPage;
	}

	/**
	 * 获取各表索引的id字段名
	 * 
	 * @param tableName
	 * @return
	 */
	public static String getIdField(String tableName) {
		if (tableName == null) {
			return null;
		}
		if (tableName.equals(SearchConstants.KIND_TABLE_NAME)) {
			return SearchConstants.KIND_ID_FIELD;
		} else if (tableName.equals(SearchConstants.BRAND_TABLE_NAME)) {
			return SearchConstants.BRAND_ID_FIELD;
		} else if (tableName.equals(SearchConstants.COMPONENT_TABLE_NAME)) {
			return SearchConstants.COMPONENT_ID_FIELD;
		} else if (tableName.equals(SearchConstants.ORDER_TABLE_NAME)) {
			return SearchConstants.ORDER_ID_FIELD;
		} else if (tableName.equals(SearchConstants.ORDER_INVOICE_TABLE_NAME)) {
			return SearchConstants.ORDER_INVOICE_ID_FIELD;
		} else if (tableName.equals(SearchConstants.PURCHASE_TABLE_NAME)) {
			return SearchConstants.PURCHASE_ID_FIELD;
		} else if (tableName.equals(SearchConstants.PURCHASE_INVOICE_TABLE_NAME)) {
			return SearchConstants.PURCHASE_INVOICE_ID_FIELD;
		}
		throw new SearchException("不支持的表：" + tableName);
	}

	/**
	 * 获取所有需要建立索引的表名
	 * 
	 * @return
	 */
	public static List<String> getTableNames() {
		List<String> tableNames = new ArrayList<>();
		tableNames.add(SearchConstants.KIND_TABLE_NAME);
		tableNames.add(SearchConstants.BRAND_TABLE_NAME);
		tableNames.add(SearchConstants.COMPONENT_TABLE_NAME);
		tableNames.add(SearchConstants.GOODS_TABLE_NAME);
		tableNames.add(SearchConstants.ORDER_TABLE_NAME);
		tableNames.add(SearchConstants.ORDER_INVOICE_TABLE_NAME);
		tableNames.add(SearchConstants.PURCHASE_TABLE_NAME);
		tableNames.add(SearchConstants.PURCHASE_INVOICE_TABLE_NAME);
		return tableNames;
	}

	/**
	 * 获取指定的实体类对应的需要建立索引的表名
	 * 
	 * @param clazz
	 * @return
	 */
	public static String getTableName(Class<?> clazz) {
		if (clazz == null) {
			return null;
		}
		if (clazz == Kind.class) {
			return SearchConstants.KIND_TABLE_NAME;
		} else if (clazz == Brand.class) {
			return SearchConstants.BRAND_TABLE_NAME;
		} else if (clazz == Component.class) {
			return SearchConstants.COMPONENT_TABLE_NAME;
		} else if (clazz == Goods.class) {
			return SearchConstants.GOODS_TABLE_NAME;
		} else if (clazz == Order.class) {
			return SearchConstants.ORDER_TABLE_NAME;
		} else if (clazz == OrderInvoice.class) {
			return SearchConstants.ORDER_INVOICE_TABLE_NAME;
		} else if (clazz == Purchase.class) {
			return SearchConstants.PURCHASE_TABLE_NAME;
		} else if (clazz == PurchaseInvoice.class) {
			return SearchConstants.PURCHASE_INVOICE_TABLE_NAME;
		} else {
			throw new SearchException("该实体类没有对应的需要建立索引的表：" + clazz);
		}
	}

	/**
	 * 获取指定表的索引路径
	 * 
	 * @param tableName
	 *            指定表
	 * @return 索引路径
	 */
	public static String getIndexPath(String tableName) {
		if (tableName == null) {
			return null;
		}
		return luceneProperties.getIndexesDir() + "/" + tableName;
	}

	/**
	 * 获取指定表的数据路径
	 * 
	 * @param tableName
	 *            指定表
	 * @return 数据路径
	 */
	public static String getDataPath(String tableName) {
		if (tableName == null) {
			return null;
		}
		return luceneProperties.getDataDir() + "/" + tableName;
	}
}