fix: 修复知识库索引并发与向量化边界
- 复用 Lucene 和 Elasticsearch 客户端并支持有界批量写入与删除 - 记录脱敏后的 Embedding 失败请求与完整响应 - 为 BGE-M3 分块统一增加上下文硬上限
This commit is contained in:
@@ -2,8 +2,6 @@ package com.easyagents.engine.es;
|
||||
|
||||
import co.elastic.clients.elasticsearch.ElasticsearchClient;
|
||||
import co.elastic.clients.elasticsearch.core.*;
|
||||
import co.elastic.clients.elasticsearch.core.bulk.BulkOperation;
|
||||
import co.elastic.clients.elasticsearch.core.bulk.IndexOperation;
|
||||
import co.elastic.clients.elasticsearch.core.search.SourceConfig;
|
||||
import co.elastic.clients.json.JsonData;
|
||||
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
|
||||
@@ -25,19 +23,46 @@ import org.slf4j.LoggerFactory;
|
||||
import javax.net.ssl.SSLContext;
|
||||
import javax.net.ssl.TrustManager;
|
||||
import javax.net.ssl.X509TrustManager;
|
||||
import java.io.IOException;
|
||||
import java.security.KeyManagementException;
|
||||
import java.security.NoSuchAlgorithmException;
|
||||
import java.security.cert.X509Certificate;
|
||||
import java.util.*;
|
||||
|
||||
public class ElasticSearcher implements DocumentSearcher {
|
||||
/**
|
||||
* 基于 Elasticsearch 的关键词搜索器。
|
||||
*
|
||||
* <p>底层客户端线程安全并与搜索器实例共享生命周期，避免逐文档重复创建网络连接。</p>
|
||||
*/
|
||||
public class ElasticSearcher implements DocumentSearcher, AutoCloseable {
|
||||
|
||||
private static final Logger LOG = LoggerFactory.getLogger(ElasticSearcher.class);
|
||||
static final int MAX_BULK_OPERATIONS = 200;
|
||||
static final long MAX_BULK_ESTIMATED_BYTES = 5L * 1024L * 1024L;
|
||||
private static final long BULK_OPERATION_METADATA_BYTES = 128L;
|
||||
|
||||
private final ESConfig esConfig;
|
||||
private final JacksonJsonpMapper jsonpMapper = new JacksonJsonpMapper();
|
||||
private final ElasticsearchTransport transport;
|
||||
private final ElasticsearchClient client;
|
||||
|
||||
/**
|
||||
* 创建 Elasticsearch 搜索器及其共享客户端。
|
||||
*
|
||||
* @param esConfig Elasticsearch 配置
|
||||
* @throws IllegalStateException 客户端初始化失败时抛出
|
||||
*/
|
||||
public ElasticSearcher(ESConfig esConfig) {
|
||||
this.esConfig = esConfig;
|
||||
this.esConfig = Objects.requireNonNull(esConfig, "ESConfig 不能为 null");
|
||||
RestClient openedRestClient = null;
|
||||
try {
|
||||
openedRestClient = buildRestClient();
|
||||
this.transport = new RestClientTransport(openedRestClient, jsonpMapper);
|
||||
this.client = new ElasticsearchClient(transport);
|
||||
} catch (Exception e) {
|
||||
closeQuietly(openedRestClient);
|
||||
throw new IllegalStateException("初始化 Elasticsearch 客户端失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
// 忽略 SSL 的 client 构建逻辑
|
||||
@@ -76,52 +101,128 @@ public class ElasticSearcher implements DocumentSearcher {
|
||||
|
||||
|
||||
/**
|
||||
* 添加文档到Elasticsearch
|
||||
* 添加文档到 Elasticsearch。
|
||||
*
|
||||
* @param document 待写入文档
|
||||
* @return 写入成功时返回 {@code true}
|
||||
*/
|
||||
@Override
|
||||
public boolean addDocument(Document document) {
|
||||
if (document == null || document.getContent() == null) {
|
||||
return addDocuments(Collections.singletonList(document));
|
||||
}
|
||||
|
||||
/**
|
||||
* 使用有界 Bulk 请求批量写入文档。
|
||||
*
|
||||
* @param documents 待写入文档
|
||||
* @return 全部文档写入成功时返回 {@code true}
|
||||
*/
|
||||
@Override
|
||||
public boolean addDocuments(List<Document> documents) {
|
||||
if (!areValidDocuments(documents)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
RestClient restClient = null;
|
||||
ElasticsearchTransport transport = null;
|
||||
try {
|
||||
restClient = buildRestClient();
|
||||
transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
|
||||
ElasticsearchClient client = new ElasticsearchClient(transport);
|
||||
|
||||
Map<String, Object> source = buildSource(document);
|
||||
String documentId = document.getId().toString();
|
||||
IndexOperation<?> indexOp = IndexOperation.of(i -> i
|
||||
.index(esConfig.getIndexName())
|
||||
.id(documentId)
|
||||
.document(JsonData.of(source))
|
||||
);
|
||||
|
||||
BulkOperation bulkOp = BulkOperation.of(b -> b.index(indexOp));
|
||||
BulkRequest request = BulkRequest.of(b -> b.operations(Collections.singletonList(bulkOp)));
|
||||
BulkResponse response = client.bulk(request);
|
||||
return !response.errors();
|
||||
|
||||
List<List<Document>> batches = partitionDocuments(documents);
|
||||
for (int batchIndex = 0; batchIndex < batches.size(); batchIndex++) {
|
||||
List<Document> batch = batches.get(batchIndex);
|
||||
BulkResponse response = client.bulk(buildBulkRequest(batch));
|
||||
if (!isBulkSuccessful(response, "写入", batchIndex, batch.size())) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
} catch (Exception e) {
|
||||
LOG.error(e.getMessage(), e);
|
||||
LOG.error("Elasticsearch 批量写入异常: count={}", documents.size(), e);
|
||||
return false;
|
||||
} finally {
|
||||
closeResources(transport, restClient);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 按操作数量和序列化后的估算字节数拆分 Bulk 批次。
|
||||
*
|
||||
* @param documents 已完成校验的文档
|
||||
* @return 有界文档批次
|
||||
* @throws IOException 文档序列化失败时抛出
|
||||
* @throws IllegalArgumentException 单个文档超过批次字节限制时抛出
|
||||
*/
|
||||
List<List<Document>> partitionDocuments(List<Document> documents) throws IOException {
|
||||
List<List<Document>> batches = new ArrayList<>();
|
||||
List<Document> currentBatch = new ArrayList<>(Math.min(documents.size(), MAX_BULK_OPERATIONS));
|
||||
long currentBatchBytes = 0L;
|
||||
for (Document document : documents) {
|
||||
long documentBytes = estimateBulkOperationBytes(document);
|
||||
if (documentBytes > MAX_BULK_ESTIMATED_BYTES) {
|
||||
throw new IllegalArgumentException(
|
||||
"单个 Elasticsearch 索引文档超过 Bulk 字节限制: id=" + document.getId());
|
||||
}
|
||||
if (!currentBatch.isEmpty()
|
||||
&& (currentBatch.size() >= MAX_BULK_OPERATIONS
|
||||
|| currentBatchBytes + documentBytes > MAX_BULK_ESTIMATED_BYTES)) {
|
||||
batches.add(currentBatch);
|
||||
currentBatch = new ArrayList<>(Math.min(documents.size(), MAX_BULK_OPERATIONS));
|
||||
currentBatchBytes = 0L;
|
||||
}
|
||||
currentBatch.add(document);
|
||||
currentBatchBytes += documentBytes;
|
||||
}
|
||||
if (!currentBatch.isEmpty()) {
|
||||
batches.add(currentBatch);
|
||||
}
|
||||
return batches;
|
||||
}
|
||||
|
||||
private long estimateBulkOperationBytes(Document document) throws IOException {
|
||||
long sourceBytes = jsonpMapper.objectMapper().writeValueAsBytes(buildSource(document)).length;
|
||||
return sourceBytes
|
||||
+ utf8Length(esConfig.getIndexName())
|
||||
+ utf8Length(document.getId().toString())
|
||||
+ BULK_OPERATION_METADATA_BYTES;
|
||||
}
|
||||
|
||||
private int utf8Length(String value) {
|
||||
return value == null ? 0 : value.getBytes(java.nio.charset.StandardCharsets.UTF_8).length;
|
||||
}
|
||||
|
||||
/**
|
||||
* 构建包含全部文档的 Bulk 请求。
|
||||
*
|
||||
* @param documents 已完成校验的文档
|
||||
* @return Bulk 请求
|
||||
*/
|
||||
BulkRequest buildBulkRequest(List<Document> documents) {
|
||||
BulkRequest.Builder builder = new BulkRequest.Builder();
|
||||
for (Document document : documents) {
|
||||
builder.operations(operation -> operation.index(index -> index
|
||||
.index(esConfig.getIndexName())
|
||||
.id(document.getId().toString())
|
||||
.document(JsonData.of(buildSource(document)))
|
||||
));
|
||||
}
|
||||
return builder.build();
|
||||
}
|
||||
|
||||
private boolean areValidDocuments(List<Document> documents) {
|
||||
if (documents == null || documents.isEmpty()) {
|
||||
return false;
|
||||
}
|
||||
for (Document document : documents) {
|
||||
if (document == null || document.getId() == null || document.getContent() == null) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* 按请求条件搜索文档。
|
||||
*
|
||||
* @param request 搜索请求
|
||||
* @return 命中文档
|
||||
*/
|
||||
@Override
|
||||
public List<Document> searchDocuments(KeywordSearchRequest request) {
|
||||
RestClient restClient = null;
|
||||
ElasticsearchTransport transport = null;
|
||||
|
||||
try {
|
||||
restClient = buildRestClient();
|
||||
transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
|
||||
ElasticsearchClient client = new ElasticsearchClient(transport);
|
||||
|
||||
SearchResponse<Map> response = client.search(buildSearchRequest(request), Map.class);
|
||||
List<Document> results = new ArrayList<>();
|
||||
response.hits().hits().forEach(hit -> {
|
||||
@@ -137,54 +238,114 @@ public class ElasticSearcher implements DocumentSearcher {
|
||||
} catch (Exception e) {
|
||||
LOG.error(e.getMessage(), e);
|
||||
return Collections.emptyList();
|
||||
} finally {
|
||||
closeResources(transport, restClient);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 删除指定文档。
|
||||
*
|
||||
* @param id 文档 ID
|
||||
* @return 删除成功时返回 {@code true}
|
||||
*/
|
||||
@Override
|
||||
public boolean deleteDocument(Object id) {
|
||||
if (id == null) {
|
||||
return deleteDocuments(Collections.singletonList(id));
|
||||
}
|
||||
|
||||
/**
|
||||
* 使用有界 Bulk 请求批量删除文档。
|
||||
*
|
||||
* <p>删除不存在的文档按幂等成功处理；仅 Bulk 条目包含实际错误时返回失败。</p>
|
||||
*
|
||||
* @param ids 文档 ID 集合
|
||||
* @return 全部删除操作成功时返回 {@code true}
|
||||
*/
|
||||
@Override
|
||||
public boolean deleteDocuments(Collection<?> ids) {
|
||||
if (ids == null || ids.isEmpty()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
RestClient restClient = null;
|
||||
ElasticsearchTransport transport = null;
|
||||
List<String> documentIds = new ArrayList<>(ids.size());
|
||||
for (Object id : ids) {
|
||||
if (id == null) {
|
||||
return false;
|
||||
}
|
||||
documentIds.add(id.toString());
|
||||
}
|
||||
try {
|
||||
restClient = buildRestClient();
|
||||
transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
|
||||
ElasticsearchClient client = new ElasticsearchClient(transport);
|
||||
|
||||
DeleteRequest request = DeleteRequest.of(d -> d
|
||||
.index(esConfig.getIndexName())
|
||||
.id(id.toString())
|
||||
);
|
||||
|
||||
DeleteResponse response = client.delete(request);
|
||||
return response.result() == co.elastic.clients.elasticsearch._types.Result.Deleted;
|
||||
|
||||
for (int start = 0, batchIndex = 0;
|
||||
start < documentIds.size();
|
||||
start += MAX_BULK_OPERATIONS, batchIndex++) {
|
||||
int end = Math.min(start + MAX_BULK_OPERATIONS, documentIds.size());
|
||||
List<String> batch = documentIds.subList(start, end);
|
||||
BulkResponse response = client.bulk(buildDeleteBulkRequest(batch));
|
||||
if (!isBulkSuccessful(response, "删除", batchIndex, batch.size())) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
} catch (Exception e) {
|
||||
LOG.error("Error deleting document with id: " + id, e);
|
||||
LOG.error("Elasticsearch 批量删除异常: count={}", documentIds.size(), e);
|
||||
return false;
|
||||
} finally {
|
||||
closeResources(transport, restClient);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 构建批量删除请求。
|
||||
*
|
||||
* @param ids 文档 ID
|
||||
* @return Bulk 删除请求
|
||||
*/
|
||||
BulkRequest buildDeleteBulkRequest(List<String> ids) {
|
||||
BulkRequest.Builder builder = new BulkRequest.Builder();
|
||||
for (String id : ids) {
|
||||
builder.operations(operation -> operation.delete(delete -> delete
|
||||
.index(esConfig.getIndexName())
|
||||
.id(id)
|
||||
));
|
||||
}
|
||||
return builder.build();
|
||||
}
|
||||
|
||||
private boolean isBulkSuccessful(BulkResponse response,
|
||||
String operation,
|
||||
int batchIndex,
|
||||
int batchSize) {
|
||||
if (response != null && !response.errors()) {
|
||||
return true;
|
||||
}
|
||||
if (response == null) {
|
||||
LOG.error("Elasticsearch 批量{}未返回结果: batchIndex={}, batchSize={}",
|
||||
operation, batchIndex, batchSize);
|
||||
return false;
|
||||
}
|
||||
response.items().stream()
|
||||
.filter(item -> item.error() != null)
|
||||
.forEach(item -> LOG.error(
|
||||
"Elasticsearch 批量{}失败: batchIndex={}, index={}, id={}, status={}, reason={}",
|
||||
operation,
|
||||
batchIndex,
|
||||
item.index(),
|
||||
item.id(),
|
||||
item.status(),
|
||||
item.error().reason()
|
||||
));
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* 更新指定文档。
|
||||
*
|
||||
* @param document 待更新文档
|
||||
* @return 更新成功时返回 {@code true}
|
||||
*/
|
||||
@Override
|
||||
public boolean updateDocument(Document document) {
|
||||
if (document == null || document.getId() == null) {
|
||||
return false;
|
||||
}
|
||||
|
||||
RestClient restClient = null;
|
||||
ElasticsearchTransport transport = null;
|
||||
|
||||
try {
|
||||
restClient = buildRestClient();
|
||||
transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
|
||||
ElasticsearchClient client = new ElasticsearchClient(transport);
|
||||
|
||||
UpdateRequest<Map<String, Object>, Map<String, Object>> request = UpdateRequest.of(u -> u
|
||||
.index(esConfig.getIndexName())
|
||||
.id(document.getId().toString())
|
||||
@@ -199,20 +360,17 @@ public class ElasticSearcher implements DocumentSearcher {
|
||||
} catch (Exception e) {
|
||||
LOG.error("Error updating document with id: " + document.getId(), e);
|
||||
return false;
|
||||
} finally {
|
||||
closeResources(transport, restClient);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
private void closeResources(AutoCloseable... closeables) {
|
||||
for (AutoCloseable closeable : closeables) {
|
||||
try {
|
||||
if (closeable != null)
|
||||
closeable.close();
|
||||
} catch (Exception e) {
|
||||
LOG.error("Error closing resource", e);
|
||||
}
|
||||
private static void closeQuietly(AutoCloseable closeable) {
|
||||
if (closeable == null) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
closeable.close();
|
||||
} catch (Exception ignored) {
|
||||
// 初始化失败时仅执行尽力清理，原始异常由构造方法继续抛出。
|
||||
}
|
||||
}
|
||||
|
||||
@@ -284,19 +442,27 @@ public class ElasticSearcher implements DocumentSearcher {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 检查 Elasticsearch 服务是否可用。
|
||||
*
|
||||
* @return 服务可访问时返回 {@code true}
|
||||
*/
|
||||
public boolean checkAvailable() {
|
||||
RestClient restClient = null;
|
||||
ElasticsearchTransport transport = null;
|
||||
try {
|
||||
restClient = buildRestClient();
|
||||
transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
|
||||
ElasticsearchClient client = new ElasticsearchClient(transport);
|
||||
return client.info() != null;
|
||||
} catch (Exception e) {
|
||||
LOG.error("Elasticsearch availability check failed", e);
|
||||
return false;
|
||||
} finally {
|
||||
closeResources(transport, restClient);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 关闭共享传输层及其底层 RestClient。
|
||||
*
|
||||
* @throws IOException 关闭客户端资源失败时抛出
|
||||
*/
|
||||
@Override
|
||||
public void close() throws IOException {
|
||||
transport.close();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,48 +1,185 @@
|
||||
package com.easyagents.engine.es;
|
||||
|
||||
import co.elastic.clients.elasticsearch.core.SearchRequest;
|
||||
import co.elastic.clients.elasticsearch.core.BulkRequest;
|
||||
import com.easyagents.core.document.Document;
|
||||
import com.easyagents.search.engine.service.KeywordSearchMetadataKeys;
|
||||
import com.easyagents.search.engine.service.KeywordSearchRequest;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* {@link ElasticSearcher} 请求构建回归测试。
|
||||
*/
|
||||
public class ElasticSearcherQueryBuilderTest {
|
||||
|
||||
/**
|
||||
* 验证搜索请求同时包含多字段检索和知识库过滤。
|
||||
*
|
||||
* @throws Exception Elasticsearch 客户端关闭失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void shouldBuildSearchRequestWithMultiMatchAndKnowledgeFilter() {
|
||||
ElasticSearcher searcher = new ElasticSearcher(config());
|
||||
KeywordSearchRequest request = KeywordSearchRequest.of("客服", 5);
|
||||
request.setKnowledgeId("100");
|
||||
public void shouldBuildSearchRequestWithMultiMatchAndKnowledgeFilter() throws Exception {
|
||||
try (ElasticSearcher searcher = new ElasticSearcher(config())) {
|
||||
KeywordSearchRequest request = KeywordSearchRequest.of("客服", 5);
|
||||
request.setKnowledgeId("100");
|
||||
|
||||
SearchRequest searchRequest = searcher.buildSearchRequest(request);
|
||||
SearchRequest searchRequest = searcher.buildSearchRequest(request);
|
||||
|
||||
Assert.assertEquals(5, searchRequest.size().intValue());
|
||||
Assert.assertNotNull(searchRequest.query().bool());
|
||||
Assert.assertEquals(1, searchRequest.query().bool().must().size());
|
||||
Assert.assertNotNull(searchRequest.query().bool().must().get(0).multiMatch());
|
||||
Assert.assertEquals(2, searchRequest.query().bool().must().get(0).multiMatch().fields().size());
|
||||
Assert.assertEquals(1, searchRequest.query().bool().filter().size());
|
||||
Assert.assertEquals("knowledgeId", searchRequest.query().bool().filter().get(0).term().field());
|
||||
Assert.assertEquals(5, searchRequest.size().intValue());
|
||||
Assert.assertNotNull(searchRequest.query().bool());
|
||||
Assert.assertEquals(1, searchRequest.query().bool().must().size());
|
||||
Assert.assertNotNull(searchRequest.query().bool().must().get(0).multiMatch());
|
||||
Assert.assertEquals(2, searchRequest.query().bool().must().get(0).multiMatch().fields().size());
|
||||
Assert.assertEquals(1, searchRequest.query().bool().filter().size());
|
||||
Assert.assertEquals("knowledgeId", searchRequest.query().bool().filter().get(0).term().field());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证知识库 ID 会写入顶层字段。
|
||||
*
|
||||
* @throws Exception Elasticsearch 客户端关闭失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void shouldExtractKnowledgeIdToTopLevelSource() {
|
||||
ElasticSearcher searcher = new ElasticSearcher(config());
|
||||
public void shouldExtractKnowledgeIdToTopLevelSource() throws Exception {
|
||||
try (ElasticSearcher searcher = new ElasticSearcher(config())) {
|
||||
Document document = new Document();
|
||||
document.setId("1");
|
||||
document.setTitle("title");
|
||||
document.setContent("content");
|
||||
document.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "100");
|
||||
|
||||
Map<String, Object> source = searcher.buildSource(document);
|
||||
|
||||
Assert.assertEquals("100", source.get(KeywordSearchMetadataKeys.KNOWLEDGE_ID));
|
||||
Assert.assertTrue(source.get("metadataMap") instanceof Map);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证多个文档会进入同一个 Bulk 请求并保留稳定文档 ID。
|
||||
*
|
||||
* @throws Exception Elasticsearch 客户端关闭失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void shouldBuildSingleBulkRequestForMultipleDocuments() throws Exception {
|
||||
try (ElasticSearcher searcher = new ElasticSearcher(config())) {
|
||||
Document first = new Document();
|
||||
first.setId("first");
|
||||
first.setContent("first content");
|
||||
Document second = new Document();
|
||||
second.setId("second");
|
||||
second.setContent("second content");
|
||||
|
||||
BulkRequest request = searcher.buildBulkRequest(List.of(first, second));
|
||||
|
||||
Assert.assertEquals(2, request.operations().size());
|
||||
Assert.assertEquals("first", request.operations().get(0).index().id());
|
||||
Assert.assertEquals("second", request.operations().get(1).index().id());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 Bulk 写入会按最大操作数拆分，避免单次请求无界增长。
|
||||
*
|
||||
* @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void shouldPartitionBulkRequestsByOperationCount() throws Exception {
|
||||
try (ElasticSearcher searcher = new ElasticSearcher(config())) {
|
||||
List<Document> documents = new ArrayList<>();
|
||||
for (int index = 0; index <= ElasticSearcher.MAX_BULK_OPERATIONS; index++) {
|
||||
documents.add(document("count-" + index, "content"));
|
||||
}
|
||||
|
||||
List<List<Document>> batches = searcher.partitionDocuments(documents);
|
||||
|
||||
Assert.assertEquals(2, batches.size());
|
||||
Assert.assertEquals(ElasticSearcher.MAX_BULK_OPERATIONS, batches.get(0).size());
|
||||
Assert.assertEquals(1, batches.get(1).size());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 Bulk 写入会按序列化字节数拆分。
|
||||
*
|
||||
* @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void shouldPartitionBulkRequestsByEstimatedBytes() throws Exception {
|
||||
try (ElasticSearcher searcher = new ElasticSearcher(config())) {
|
||||
int contentLength = (int) (ElasticSearcher.MAX_BULK_ESTIMATED_BYTES / 2L);
|
||||
Document first = document("bytes-first", "a".repeat(contentLength));
|
||||
Document second = document("bytes-second", "b".repeat(contentLength));
|
||||
|
||||
List<List<Document>> batches = searcher.partitionDocuments(List.of(first, second));
|
||||
|
||||
Assert.assertEquals(2, batches.size());
|
||||
Assert.assertEquals(1, batches.get(0).size());
|
||||
Assert.assertEquals(1, batches.get(1).size());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证单个超限文档会在发送请求前失败。
|
||||
*
|
||||
* @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void shouldRejectSingleDocumentOverByteLimit() throws Exception {
|
||||
try (ElasticSearcher searcher = new ElasticSearcher(config())) {
|
||||
int contentLength = (int) ElasticSearcher.MAX_BULK_ESTIMATED_BYTES;
|
||||
Document oversized = document("oversized", "a".repeat(contentLength));
|
||||
|
||||
try {
|
||||
searcher.partitionDocuments(List.of(oversized));
|
||||
Assert.fail("单个超限文档应拒绝构建 Bulk 请求");
|
||||
} catch (IllegalArgumentException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("oversized"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证批量删除请求包含全部稳定文档 ID。
|
||||
*
|
||||
* @throws Exception Elasticsearch 客户端关闭失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void shouldBuildBulkDeleteRequest() throws Exception {
|
||||
try (ElasticSearcher searcher = new ElasticSearcher(config())) {
|
||||
BulkRequest request = searcher.buildDeleteBulkRequest(List.of("first", "second"));
|
||||
|
||||
Assert.assertEquals(2, request.operations().size());
|
||||
Assert.assertEquals("first", request.operations().get(0).delete().id());
|
||||
Assert.assertEquals("second", request.operations().get(1).delete().id());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建测试文档。
|
||||
*
|
||||
* @param id 文档 ID
|
||||
* @param content 文档内容
|
||||
* @return 文档
|
||||
*/
|
||||
private Document document(String id, String content) {
|
||||
Document document = new Document();
|
||||
document.setId("1");
|
||||
document.setTitle("title");
|
||||
document.setContent("content");
|
||||
document.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "100");
|
||||
|
||||
Map<String, Object> source = searcher.buildSource(document);
|
||||
|
||||
Assert.assertEquals("100", source.get(KeywordSearchMetadataKeys.KNOWLEDGE_ID));
|
||||
Assert.assertTrue(source.get("metadataMap") instanceof Map);
|
||||
document.setId(id);
|
||||
document.setContent(content);
|
||||
return document;
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建测试用 Elasticsearch 配置。
|
||||
*
|
||||
* @return Elasticsearch 配置
|
||||
*/
|
||||
private ESConfig config() {
|
||||
ESConfig config = new ESConfig();
|
||||
config.setHost("http://127.0.0.1:9200");
|
||||
|
||||
Reference in New Issue
Block a user