Skip to content

Commit d7068db

Browse files
chishuireta
authored and
finnegancarroll
committed
Support batch ingestion in bulk API (opensearch-project#12457) (opensearch-project#13306)
* [PoC][issues-12457] Support Batch Ingestion Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Rewrite batch interface and handle error and metrics Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Remove unnecessary change Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Revert some unnecessary test change Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Keep executeBulkRequest main logic untouched Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Add UT Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Add UT & yamlRest test, fix BulkRequest se/deserialization Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Add missing java docs Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Remove Writable from BatchIngestionOption Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Add more UTs Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Fix spotlesscheck Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Rename parameter name to batch_size Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Add more rest yaml tests & update rest spec Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Remove batch_ingestion_option and only use batch_size Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Throw invalid request exception for invalid batch_size Signed-off-by: Liyun Xiu <xiliyun@amazon.com> * Update server/src/main/java/org/opensearch/action/bulk/BulkRequest.java Co-authored-by: Andriy Redko <drreta@gmail.com> Signed-off-by: Liyun Xiu <chishui2@gmail.com> * Remove version constant Signed-off-by: Liyun Xiu <xiliyun@amazon.com> --------- Signed-off-by: Liyun Xiu <xiliyun@amazon.com> Signed-off-by: Liyun Xiu <chishui2@gmail.com> Co-authored-by: Andriy Redko <drreta@gmail.com>
1 parent baabb33 commit d7068db

File tree

18 files changed

+1486
-50
lines changed

18 files changed

+1486
-50
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
2828
- [Tiered Caching] Gate new stats logic behind FeatureFlags.PLUGGABLE_CACHE ([#13238](https://github.com/opensearch-project/OpenSearch/pull/13238))
2929
- [Tiered Caching] Add a dynamic setting to disable/enable disk cache. ([#13373](https://github.com/opensearch-project/OpenSearch/pull/13373))
3030
- [Remote Store] Add capability of doing refresh as determined by the translog ([#12992](https://github.com/opensearch-project/OpenSearch/pull/12992))
31+
- [Batch Ingestion] Add `batch_size` to `_bulk` API. ([#12457](https://github.com/opensearch-project/OpenSearch/issues/12457))
3132
- [Tiered caching] Make Indices Request Cache Stale Key Mgmt Threshold setting dynamic ([#12941](https://github.com/opensearch-project/OpenSearch/pull/12941))
3233
- Batch mode for async fetching shard information in GatewayAllocator for unassigned shards ([#8746](https://github.com/opensearch-project/OpenSearch/pull/8746))
3334
- [Remote Store] Add settings for remote path type and hash algorithm ([#13225](https://github.com/opensearch-project/OpenSearch/pull/13225))

‎modules/ingest-common/src/yamlRestTest/resources/rest-api-spec/test/ingest/70_bulk.yml‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,3 +167,90 @@ teardown:
167167
index: test_index
168168
id: test_id3
169169
- match: { _source: {"f1": "v2", "f2": 47, "field1": "value1"}}
170+
171+
---
172+
"Test bulk API with batch enabled happy case":
173+
- skip:
174+
version: " - 2.13.99"
175+
reason: "Added in 2.14.0"
176+
177+
- do:
178+
bulk:
179+
refresh: true
180+
batch_size: 2
181+
pipeline: "pipeline1"
182+
body:
183+
- '{"index": {"_index": "test_index", "_id": "test_id1"}}'
184+
- '{"text": "text1"}'
185+
- '{"index": {"_index": "test_index", "_id": "test_id2"}}'
186+
- '{"text": "text2"}'
187+
- '{"index": {"_index": "test_index", "_id": "test_id3"}}'
188+
- '{"text": "text3"}'
189+
- '{"index": {"_index": "test_index", "_id": "test_id4"}}'
190+
- '{"text": "text4"}'
191+
- '{"index": {"_index": "test_index", "_id": "test_id5", "pipeline": "pipeline2"}}'
192+
- '{"text": "text5"}'
193+
- '{"index": {"_index": "test_index", "_id": "test_id6", "pipeline": "pipeline2"}}'
194+
- '{"text": "text6"}'
195+
196+
- match: { errors: false }
197+
198+
- do:
199+
get:
200+
index: test_index
201+
id: test_id5
202+
- match: { _source: {"text": "text5", "field2": "value2"}}
203+
204+
- do:
205+
get:
206+
index: test_index
207+
id: test_id3
208+
- match: { _source: { "text": "text3", "field1": "value1" } }
209+
210+
---
211+
"Test bulk API with batch_size missing":
212+
- skip:
213+
version: " - 2.13.99"
214+
reason: "Added in 2.14.0"
215+
216+
- do:
217+
bulk:
218+
refresh: true
219+
pipeline: "pipeline1"
220+
body:
221+
- '{"index": {"_index": "test_index", "_id": "test_id1"}}'
222+
- '{"text": "text1"}'
223+
- '{"index": {"_index": "test_index", "_id": "test_id2"}}'
224+
- '{"text": "text2"}'
225+
226+
- match: { errors: false }
227+
228+
- do:
229+
get:
230+
index: test_index
231+
id: test_id1
232+
- match: { _source: { "text": "text1", "field1": "value1" } }
233+
234+
- do:
235+
get:
236+
index: test_index
237+
id: test_id2
238+
- match: { _source: { "text": "text2", "field1": "value1" } }
239+
240+
---
241+
"Test bulk API with invalid batch_size":
242+
- skip:
243+
version: " - 2.13.99"
244+
reason: "Added in 2.14.0"
245+
246+
- do:
247+
catch: bad_request
248+
bulk:
249+
refresh: true
250+
batch_size: -1
251+
pipeline: "pipeline1"
252+
body:
253+
- '{"index": {"_index": "test_index", "_id": "test_id1"}}'
254+
- '{"text": "text1"}'
255+
- '{"index": {"_index": "test_index", "_id": "test_id2"}}'
256+
- '{"text": "text2"}'

‎rest-api-spec/src/main/resources/rest-api-spec/api/bulk.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,10 @@
7474
"require_alias": {
7575
"type": "boolean",
7676
"description": "Sets require_alias for all incoming documents. Defaults to unset (false)"
77+
},
78+
"batch_size": {
79+
"type": "int",
80+
"description": "Sets the batch size"
7781
}
7882
},
7983
"body":{

‎server/src/main/java/org/opensearch/action/bulk/BulkRequest.java‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434

3535
import org.apache.lucene.util.Accountable;
3636
import org.apache.lucene.util.RamUsageEstimator;
37+
import org.opensearch.Version;
3738
import org.opensearch.action.ActionRequest;
3839
import org.opensearch.action.ActionRequestValidationException;
3940
import org.opensearch.action.CompositeIndicesRequest;
@@ -80,7 +81,6 @@ public class BulkRequest extends ActionRequest implements CompositeIndicesReques
8081
private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(BulkRequest.class);
8182

8283
private static final int REQUEST_OVERHEAD = 50;
83-
8484
/**
8585
* Requests that are part of this request. It is only possible to add things that are both {@link ActionRequest}s and
8686
* {@link WriteRequest}s to this but java doesn't support syntax to declare that everything in the array has both types so we declare
@@ -96,6 +96,7 @@ public class BulkRequest extends ActionRequest implements CompositeIndicesReques
9696
private String globalRouting;
9797
private String globalIndex;
9898
private Boolean globalRequireAlias;
99+
private int batchSize = 1;
99100

100101
private long sizeInBytes = 0;
101102

@@ -107,6 +108,9 @@ public BulkRequest(StreamInput in) throws IOException {
107108
requests.addAll(in.readList(i -> DocWriteRequest.readDocumentRequest(null, i)));
108109
refreshPolicy = RefreshPolicy.readFrom(in);
109110
timeout = in.readTimeValue();
111+
if (in.getVersion().onOrAfter(Version.V_2_14_0)) {
112+
batchSize = in.readInt();
113+
}
110114
}
111115

112116
public BulkRequest(@Nullable String globalIndex) {
@@ -346,6 +350,27 @@ public final BulkRequest timeout(TimeValue timeout) {
346350
return this;
347351
}
348352

353+
/**
354+
* Set batch size
355+
* @param size batch size from input
356+
* @return {@link BulkRequest}
357+
*/
358+
public BulkRequest batchSize(int size) {
359+
if (size < 1) {
360+
throw new IllegalArgumentException("batch_size must be greater than 0");
361+
}
362+
this.batchSize = size;
363+
return this;
364+
}
365+
366+
/**
367+
* Get batch size
368+
* @return batch size
369+
*/
370+
public int batchSize() {
371+
return this.batchSize;
372+
}
373+
349374
/**
350375
* Note for internal callers (NOT high level rest client),
351376
* the global parameter setting is ignored when used with:
@@ -453,6 +478,9 @@ public void writeTo(StreamOutput out) throws IOException {
453478
out.writeCollection(requests, DocWriteRequest::writeDocumentRequest);
454479
refreshPolicy.writeTo(out);
455480
out.writeTimeValue(timeout);
481+
if (out.getVersion().onOrAfter(Version.V_2_14_0)) {
482+
out.writeInt(batchSize);
483+
}
456484
}
457485

458486
@Override

‎server/src/main/java/org/opensearch/action/bulk/TransportBulkAction.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -923,7 +923,8 @@ public boolean isForceExecution() {
923923
}
924924
},
925925
bulkRequestModifier::markItemAsDropped,
926-
executorName
926+
executorName,
927+
original
927928
);
928929
}
929930

‎server/src/main/java/org/opensearch/common/metrics/OperationMetrics.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,14 @@ public void before() {
3737
current.incrementAndGet();
3838
}
3939

40+
/**
41+
* Invoke before the given operation begins in multiple items at the same time.
42+
* @param n number of items
43+
*/
44+
public void beforeN(int n) {
45+
current.addAndGet(n);
46+
}
47+
4048
/**
4149
* Invoked upon completion (success or failure) of the given operation
4250
* @param currentTime elapsed time of the operation
@@ -46,13 +54,35 @@ public void after(long currentTime) {
4654
time.inc(currentTime);
4755
}
4856

57+
/**
58+
* Invoked upon completion (success or failure) of the given operation for multiple items.
59+
* @param n number of items completed
60+
* @param currentTime elapsed time of the operation
61+
*/
62+
public void afterN(int n, long currentTime) {
63+
current.addAndGet(-n);
64+
for (int i = 0; i < n; ++i) {
65+
time.inc(currentTime);
66+
}
67+
}
68+
4969
/**
5070
* Invoked upon failure of the operation.
5171
*/
5272
public void failed() {
5373
failed.inc();
5474
}
5575

76+
/**
77+
* Invoked upon failure of the operation on multiple items.
78+
* @param n number of items on operation.
79+
*/
80+
public void failedN(int n) {
81+
for (int i = 0; i < n; ++i) {
82+
failed.inc();
83+
}
84+
}
85+
5686
public void add(OperationMetrics other) {
5787
// Don't try copying over current, since in-flight requests will be linked to the existing metrics instance.
5888
failed.inc(other.failed.count());

‎server/src/main/java/org/opensearch/ingest/CompoundProcessor.java‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,13 @@
3939
import java.util.ArrayList;
4040
import java.util.Arrays;
4141
import java.util.Collections;
42+
import java.util.HashMap;
4243
import java.util.List;
4344
import java.util.Map;
4445
import java.util.concurrent.TimeUnit;
46+
import java.util.concurrent.atomic.AtomicInteger;
4547
import java.util.function.BiConsumer;
48+
import java.util.function.Consumer;
4649
import java.util.function.LongSupplier;
4750
import java.util.stream.Collectors;
4851

@@ -150,6 +153,108 @@ public void execute(IngestDocument ingestDocument, BiConsumer<IngestDocument, Ex
150153
innerExecute(0, ingestDocument, handler);
151154
}
152155

156+
@Override
157+
public void batchExecute(List<IngestDocumentWrapper> ingestDocumentWrappers, Consumer<List<IngestDocumentWrapper>> handler) {
158+
innerBatchExecute(0, ingestDocumentWrappers, handler);
159+
}
160+
161+
/**
162+
* Internal logic to process documents with current processor.
163+
*
164+
* @param currentProcessor index of processor to process batched documents
165+
* @param ingestDocumentWrappers batched documents to be processed
166+
* @param handler callback function
167+
*/
168+
void innerBatchExecute(
169+
int currentProcessor,
170+
List<IngestDocumentWrapper> ingestDocumentWrappers,
171+
Consumer<List<IngestDocumentWrapper>> handler
172+
) {
173+
if (currentProcessor == processorsWithMetrics.size()) {
174+
handler.accept(ingestDocumentWrappers);
175+
return;
176+
}
177+
Tuple<Processor, OperationMetrics> processorWithMetric = processorsWithMetrics.get(currentProcessor);
178+
final Processor processor = processorWithMetric.v1();
179+
final OperationMetrics metric = processorWithMetric.v2();
180+
final long startTimeInNanos = relativeTimeProvider.getAsLong();
181+
int size = ingestDocumentWrappers.size();
182+
metric.beforeN(size);
183+
// Use synchronization to ensure batches are processed by processors in sequential order
184+
AtomicInteger counter = new AtomicInteger(size);
185+
List<IngestDocumentWrapper> allResults = Collections.synchronizedList(new ArrayList<>());
186+
Map<Integer, IngestDocumentWrapper> slotToWrapperMap = createSlotIngestDocumentWrapperMap(ingestDocumentWrappers);
187+
processor.batchExecute(ingestDocumentWrappers, results -> {
188+
if (results.isEmpty()) return;
189+
allResults.addAll(results);
190+
// counter equals to 0 means all documents are processed and called back.
191+
if (counter.addAndGet(-results.size()) == 0) {
192+
long ingestTimeInMillis = TimeUnit.NANOSECONDS.toMillis(relativeTimeProvider.getAsLong() - startTimeInNanos);
193+
metric.afterN(allResults.size(), ingestTimeInMillis);
194+
195+
List<IngestDocumentWrapper> documentsDropped = new ArrayList<>();
196+
List<IngestDocumentWrapper> documentsWithException = new ArrayList<>();
197+
List<IngestDocumentWrapper> documentsToContinue = new ArrayList<>();
198+
int totalFailed = 0;
199+
// iterate all results to categorize them to: to continue, to drop, with exception
200+
for (IngestDocumentWrapper resultDocumentWrapper : allResults) {
201+
IngestDocumentWrapper originalDocumentWrapper = slotToWrapperMap.get(resultDocumentWrapper.getSlot());
202+
if (resultDocumentWrapper.getException() != null) {
203+
++totalFailed;
204+
if (ignoreFailure) {
205+
documentsToContinue.add(originalDocumentWrapper);
206+
} else {
207+
IngestProcessorException compoundProcessorException = newCompoundProcessorException(
208+
resultDocumentWrapper.getException(),
209+
processor,
210+
originalDocumentWrapper.getIngestDocument()
211+
);
212+
documentsWithException.add(
213+
new IngestDocumentWrapper(
214+
resultDocumentWrapper.getSlot(),
215+
originalDocumentWrapper.getIngestDocument(),
216+
compoundProcessorException
217+
)
218+
);
219+
}
220+
} else {
221+
if (resultDocumentWrapper.getIngestDocument() == null) {
222+
documentsDropped.add(resultDocumentWrapper);
223+
} else {
224+
documentsToContinue.add(resultDocumentWrapper);
225+
}
226+
}
227+
}
228+
if (totalFailed > 0) {
229+
metric.failedN(totalFailed);
230+
}
231+
if (!documentsDropped.isEmpty()) {
232+
handler.accept(documentsDropped);
233+
}
234+
if (!documentsToContinue.isEmpty()) {
235+
innerBatchExecute(currentProcessor + 1, documentsToContinue, handler);
236+
}
237+
if (!documentsWithException.isEmpty()) {
238+
if (onFailureProcessors.isEmpty()) {
239+
handler.accept(documentsWithException);
240+
} else {
241+
documentsWithException.forEach(
242+
doc -> executeOnFailureAsync(
243+
0,
244+
doc.getIngestDocument(),
245+
(IngestProcessorException) doc.getException(),
246+
(result, ex) -> {
247+
handler.accept(Collections.singletonList(new IngestDocumentWrapper(doc.getSlot(), result, ex)));
248+
}
249+
)
250+
);
251+
}
252+
}
253+
}
254+
assert counter.get() >= 0;
255+
});
256+
}
257+
153258
void innerExecute(int currentProcessor, IngestDocument ingestDocument, BiConsumer<IngestDocument, Exception> handler) {
154259
if (currentProcessor == processorsWithMetrics.size()) {
155260
handler.accept(ingestDocument, null);
@@ -266,4 +371,12 @@ static IngestProcessorException newCompoundProcessorException(Exception e, Proce
266371
return exception;
267372
}
268373

374+
private Map<Integer, IngestDocumentWrapper> createSlotIngestDocumentWrapperMap(List<IngestDocumentWrapper> ingestDocumentWrappers) {
375+
Map<Integer, IngestDocumentWrapper> slotIngestDocumentWrapperMap = new HashMap<>();
376+
for (IngestDocumentWrapper ingestDocumentWrapper : ingestDocumentWrappers) {
377+
slotIngestDocumentWrapperMap.put(ingestDocumentWrapper.getSlot(), ingestDocumentWrapper);
378+
}
379+
return slotIngestDocumentWrapperMap;
380+
}
381+
269382
}

0 commit comments

Comments
 (0)