prefix search base redis

TestRedisPrefix.java

package com.xxx.xxx;

import org.junit.Test;
import redis.clients.jedis.Jedis;

import java.util.*;

/**
 * User: evanhuang
 */
public class TestRedisPrefix {

    private static final String statusString = "*";
    private Jedis baseYZedis = new Jedis("127.0.0.1", 6379);
    class ValueComparator implements Comparator<String> {
        Map<String, Integer> base;
        public ValueComparator(Map<String, Integer> base) {
            this.base = base;
        }
        public int compare(String a, String b) {
            if (base.get(a) >= base.get(b)) {
                return -1;
            } else {
                return 1;
            }
        }
    }

    @Test
    public void testPrefixIndexAndSearch(){
        String[] wordArrays = {"数据库", "数据挖掘", "数据", "数据分析师", "数据可视化", "电影", "科幻电影", "数学", "数学家"};
        for(String word : wordArrays){
            prefixIndexWord(word);
        }
        System.out.println("prefixSearch, more " + prefixSearch("数", 10));
        System.out.println("prefixSearch, less " + prefixSearch("数", 3));
        System.out.println("hit, no size " + prefixSearch("电影", 0));
        System.out.println("hit " + prefixSearch("电影", 3));
        System.out.println("miss " + prefixSearch("艺", 5));
    }

    private void prefixIndexWord(String word){
        int defaultScore = 0;
        String prefixKey = "prefixKey";
        String detailKey = "detailKey";
        baseYZedis.zadd(prefixKey, defaultScore, word + statusString);
        for(int i=1,len=word.length(); i<len; i++){
            baseYZedis.zadd(prefixKey, 0, word.substring(0, i));
        }
        baseYZedis.hset(detailKey, word, (int)(Math.random()*100)+"");
    }

    private Map<String, Integer> prefixSearch(String word, int size){
        if(size <= 0){
            return null;
        }
        int step = 10;
        String sCount;
        String prefixKey = "prefixKey";
        String detailKey = "detailKey";
        Map<String, Integer> hitsMap = new HashMap<String, Integer>();
        sCount = baseYZedis.hget(detailKey, word);
        if(sCount != null){
            hitsMap.put(word, Integer.parseInt(sCount));
        }
        Long index = baseYZedis.zrank(prefixKey, word);
        if(index != null){
            while (size>0){
                Set<String> prefixSet =  baseYZedis.zrange(prefixKey, index, index + step -1);
                if(prefixSet.isEmpty()){
                    break;
                }
                index += step;
                for(String prefixWord : prefixSet){
                    if(prefixWord.contains(statusString) && prefixWord.contains(word)){
                        prefixWord = prefixWord.substring(0, prefixWord.length()-1);
                        sCount = baseYZedis.hget(detailKey, prefixWord);
                        hitsMap.put(prefixWord, Integer.parseInt(sCount));
                        size--;
                        if(size == 0){
                            break;
                        }
                    }
                }
            }
        }
        if(hitsMap.isEmpty()){
            return null;
        }else{
            ValueComparator bvc =  new ValueComparator(hitsMap);
            TreeMap<String,Integer> sortedMap = new TreeMap<String,Integer>(bvc);
            sortedMap.putAll(hitsMap);
            return sortedMap;
        }
    }
}