前言
搜索框的提示词功能是个很常见的功能，第一反应是用字典树（Trie Tree）来实现，这里不赘述 Trie Tree 的原理，这个经典的数据结构有很多书本和网络的资料可以参考。
本文以 Redis 的 zset 为基本工具，来实现一个直白简单，但可能并非最有效（空间复杂度高）的前缀匹配功能。
语言/框架/数据库
Spring Boot 3.x
Redis 7.x
Java 21
Python 3.8
实现功能
实现一个接口，get 请求，query 参数是一个字符串（例如：城市名称），返回该字符串作为前缀，匹配到的在数据库中的所有相关城市信息。
例如，数据库有以下省市名称：
- shanghai
- shanxi
- beijing
- tianjin
- fujian
- zhejiang
query = fu 时，返回 ['fujian']，query = shan，返回 ['shanghai', 'shanxi']。
本项目的接口实现结果如下：
数据集
我爬取了以下网站的数据：
UN/LOCODE Code List by Country and Territory | UNECE
这里的数据是 联合国口岸及相关地点代码 ，用于运输，贸易行业，里面包含了各个国家的口岸地址信息。
比如，query = shanghai 匹配到的，可以是如下口岸信息列表：
{
"success": true,
"message": "ok",
"data": [
{
"countryName": "China",
"countryCode": "CN",
"cityName": "shanghai railway station",
"unCode": "CNSHZ"
},
{
"countryName": "China",
"countryCode": "CN",
"cityName": "shanghai hongqiao international apt",
"unCode": "CNSHA"
},
{
"countryName": "China",
"countryCode": "CN",
"cityName": "shanghai pudong international apt",
"unCode": "CNPVG"
},
{
"countryName": "China",
"countryCode": "CN",
"cityName": "shanghai",
"unCode": "CNSGH"
},
{
"countryName": "China",
"countryCode": "CN",
"cityName": "shanghai pt",
"unCode": "CNSHG"
}
]
}
实现
zset、zrank、zrange
set 是集合，包含不重复的元素，而 zset 则是 set 的扩展，它存储了有序的，不重复的元素。
zset 的顺序是依靠评分（score）来排的，比如，插入几个学生的分数：
zadd stu 67 bob 88 alice 54 tom 72 jack 91 mike
然后查询：
zrange stu 0 -1
结果：
1 tom
2 bob
3 jack
4 alice
5 mike
集合是通过哈希表实现的，所以添加，删除，查找的复杂度都是 O(1)。
排序是根据 score 从小到大排，有序集合的成员是唯一的，但分数值可以重复。
当分数相同时，按照字典排序，例如，所有学生都以 score = 0 插入 zset：
zadd stu 0 bob 0 alice 0 tom 0 jack 0 mike
结果：
1 alice
2 bob
3 jack
4 mike
5 tom
我们可以利用这一特性，来存储城市名称（中文也可以，因为是基于 Unicode 码值的字典序进行排序的）。
Redis 的 Zrank 返回有序集中指定成员的排名，下标从 0 开始，比如：
zadd stu 0 bob 0 alice 0 tom 0 jack 0 mike
zrank stu alice # 0
zrank stu tom # 4
Redis Zrange 返回有序集中，指定区间内的成员，比如，返回第零至第二个成员：
zrange stu 0 2
# 返回
alice
bob
jack
zadd，zrank，zrange 就是本文实现前缀匹配所需要的所有命令了。
插入单个名称
我们以 cat，bed，bag，foo，bar 这几个单词为例子，描述一下插入一个新的名称进 zset 的方法。
按照前缀树的逻辑，我们需要构造以下的树：
从数据库存储的角度来说，我们需要把单词的所有前缀都存储进去，这是一种空间换时间的方法，以存储 cat 为例：
zadd word 0 c 0 ca 0 cat 0 cat*
zrange word 0 -1
我把 cat（包括 cat 本身）的前缀都存储进 zset，并且用 * 结尾来标识一个有效单词，没有 * 结尾的都是作为前缀存储。
zrange 需要一个范围，除了开始区间值（也就是用户输入的查询前缀），我们还需要一个结尾区间值。
c 的结尾值是 d，ca 的结尾值是 cb，cat 的结尾值是 cau，这样有了结尾值，我们就可以用 zrange 来查询区间的有效单词了。
以下是插入 cat 这一个单词的所有内容：
zadd word 0 c 0 ca 0 cat 0 d 0 cb 0 cau 0 cat*
插入一个单词，所需要的所有成员数是 s = 2 * (len(word)) + 1。
用 python 代码来表示这个过程：
def save_word_in_sorted_set(set_name, word):
"""
将单个 word 存储 redis sorted set
:param set_name: 有序集合的名称
:param word: 单词
"""
# 有效单词以 * 结尾
# 分数值都为 0
r.zadd(name=set_name, mapping={word + "*": 0})
for i in range(len(word)):
tmp = word[:i + 1]
# 前缀
r.zadd(name=set_name, mapping={tmp: 0})
# 前缀最后一个字符码值 + 1, 用于 zrange 的结束区间值
r.zadd(name=set_name, mapping={tmp[:len(tmp) - 1] + chr(ord(tmp[-1]) + 1): 0})
插入所有的城市口岸名称
我把所有名称，写到了一个 json 文件中，格式如下：
["Ali Shirzayi", "Bagram", "Bamian", "Bazar-E-Panjwai", "Bost", "Camp Bastion", "Camp Dwyer", "Camp Leatherneck", "Camp Salerno", "Camp Wolverine", ... ]
总数是 106843 个。
然后循环插入到 redis 中：
def save_cities_in_sorted_set(set_name):
"""
存储城市名称列表到 redis sorted set
:return:
"""
key = set_name
# 城市数量: 106843
with open('./cities.json', mode='r', encoding='utf-8') as f:
cities = json.load(f)
total = len(cities)
for index, city_name in enumerate(cities, start=1):
city_name = city_name.lower()
print(f'开始处理第 {index} 个, 总数: {total}')
save_word_in_sorted_set(key, city_name)
数量以及占用空间
总共有 10 万个城市口岸名称，用 zcard 和 memory usage 来查看下最终的总数和占用空间大小：
zcard "quote:location" # 1171525
memory usage "quote:location" # 110500016
前缀占用的成员数量就有 100万+，空间占用了约 110 M，所以这种方法的缺点就是空间复杂度高。
搜索
当用户输入 ca 的时候，只需要四步就可以查出结果：
- 在 zset 中找到 ca 的 rank 值
- 在 zset 中找到 cb 的 rank 值
- zrange zset_name ca_rank cb_rank，就可以找到具体的区间（包含了前缀和有效单词）
- 通过后缀 *，过滤出有效单词
转换成 Python 代码：
def auto_complete(set_name, query):
"""
根据 query 查出相关的自动补全的选项
:param set_name: 有序集合名称
:param query: 查询前缀
"""
# query 作为 start
start_query = query
# 最后一个字符 unicode + 1 作为 end
end_query = query[:len(query) - 1] + chr(ord(query[-1]) + 1)
start_index = r.zrank(name=set_name, value=start_query)
end_index = r.zrank(name=set_name, value=end_query)
tmp_res = r.zrange(set_name, start=start_index, end=end_index)
res = []
for item in tmp_res:
if item[-1] == ord('*'):
word = item.decode('utf-8')
word = word[:len(word) - 1]
res.append(word)
return res
Spring Boot 接口
下面用 Spring Boot 来编写一个后端 http 接口，来实现自动补全提示的功能。
模型类，也就是最终返回结果：
@Data
public class QuoteLocation {
private String countryName;
private String countryCode;
private String cityName;
private String unCode;
}
controller 类：
@RestController
@RequestMapping("/open")
public class LocationController {
private final LocationService locationService;
public LocationController(LocationService locationService) {
this.locationService = locationService;
}
@GetMapping("/locations")
public JsonResult<List<QuoteLocation>> locations(@Validated LocationQuery query) {
return JsonResult.success(locationService.getLocations(query));
}
}
service 类：
public interface LocationService {
List<QuoteLocation> getLocations(LocationQuery query);
}
service impl 类：
@Slf4j
@Service
public class LocationServiceImpl implements LocationService {
private final static String LOCATION_SET = "location";
private final static String LOCATION_DETAIL_SET = "location_detail";
private final RedisTemplate<String, Object> redisTemplate;
private final ObjectMapper objectMapper;
public LocationServiceImpl(RedisTemplate<String, Object> redisTemplate, ObjectMapper objectMapper) {
this.redisTemplate = redisTemplate;
this.objectMapper = objectMapper;
}
@Override
public List<QuoteLocation> getLocations(LocationQuery query) {
ZSetOperations<String, Object> zSetOps = redisTemplate.opsForZSet();
HashOperations<String, Object, Object> hashOps = redisTemplate.opsForHash();
// 前缀树索引信息 sorted set
String setKey = RedisUtil.generateKey(LOCATION_SET);
// 城市国家代码详细信息 hash map
String hashKey = RedisUtil.generateKey(LOCATION_DETAIL_SET);
String queryName = query.getCityName().toLowerCase();
String endQuery = queryName.substring(0, queryName.length() - 1) +
(char) (queryName.charAt(queryName.length() - 1) + 1);
// 获取 start 和 end 索引
Long startIndex = zSetOps.rank(setKey, queryName);
Long endIndex = zSetOps.rank(setKey, endQuery);
if (startIndex == null || endIndex == null) {
return new ArrayList<>();
}
// 获取前缀匹配结果（包含有效单词和前缀）
Set<Object> tmpRes = zSetOps.range(setKey, startIndex, endIndex);
List<String> res = new ArrayList<>();
// 有效单词以 * 结尾
// 过滤掉前缀，保留有效单词，并去掉结尾的 *
if (!CollectionUtils.isEmpty(tmpRes)) {
for (Object item : tmpRes) {
String cityName = (String) item;
if (cityName.endsWith("*")) {
res.add(cityName.substring(0, cityName.length() - 1));
}
}
}
List<QuoteLocation> locations = new ArrayList<>();
// 根据城市名获取所有相关的城市国家代码信息
List<Object> details = hashOps.multiGet(hashKey, new HashSet<>(res));
for (Object detail : details) {
String cityDetail = (String) detail;
QuoteLocation location;
try {
location = objectMapper.readValue(cityDetail, QuoteLocation.class);
}
catch (JsonProcessingException e) {
throw new AppException("解析城市详情异常");
}
locations.add(location);
}
return locations;
}
}
我用一个 zset 存储所有城市口岸名称，再用一个 hash 来存储每个城市口岸对应的具体信息（所属国家，代码等）。
先在 zset 找到可能的匹配项，然后再去 hash 里检索具体信息。
RedisTemplate 配置以及其他工具类，这些细节就不在本文赘述了。
总结
这并不是一个很好的方案，因为占用空间大，但是搜索的效果还不错（本地测试，3-5毫秒），也很稳定，是一个简单直白的思路。
如果接口访问量极大，不想用关系型数据库的 Like 语句做匹配，也不想引入 ES 等其他较重的工具，可以考虑这个 Redis 的解决方案。
参考
- Can Redis do prefix matching? - Stack Overflow
- http://antirez.com/post/autocomplete-with-redis.html
- Two ways of using Redis to build a NoSQL autocomplete search index - Pat Shaughnessy