Skip to content

Commit 7bcfa92

Browse files
karenyrxkarenx
authored and
jainankitk
committed
[GRPC] Optimize source conversion in gRPC search hits using zero-copy BytesRef (opensearch-project#19280)
* [GRPC] Optimize source conversion in gRPC search hits using zero-copy BytesRef Signed-off-by: karenx <karenx@uber.com> * code cov Signed-off-by: karenx <karenx@uber.com> * resolve merge conflicts Signed-off-by: karenx <karenx@uber.com> --------- Signed-off-by: karenx <karenx@uber.com> Signed-off-by: Karen X <karenxyr@gmail.com> Co-authored-by: karenx <karenx@uber.com>
1 parent e5c9c90 commit 7bcfa92

File tree

3 files changed

+240
-1
lines changed

3 files changed

+240
-1
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
3636
- Make all methods in Engine.Result public ([#19276](https://github.com/opensearch-project/OpenSearch/pull/19275))
3737
- Create and attach interclusterTest and yamlRestTest code coverage reports to gradle check task([#19165](https://github.com/opensearch-project/OpenSearch/pull/19165))
3838
- Optimized date histogram aggregations by preventing unnecessary object allocations in date rounding utils ([19088](https://github.com/opensearch-project/OpenSearch/pull/19088))
39+
- Optimize source conversion in gRPC search hits using zero-copy BytesRef ([#19280](https://github.com/opensearch-project/OpenSearch/pull/19280))
3940

4041
### Fixed
4142
- Fix unnecessary refreshes on update preparation failures ([#15261](https://github.com/opensearch-project/OpenSearch/issues/15261))

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/response/search/SearchHitProtoUtils.java‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,11 @@
88
package org.opensearch.transport.grpc.proto.response.search;
99

1010
import com.google.protobuf.ByteString;
11+
import com.google.protobuf.UnsafeByteOperations;
1112
import org.apache.lucene.search.Explanation;
13+
import org.apache.lucene.util.BytesRef;
1214
import org.opensearch.common.document.DocumentField;
15+
import org.opensearch.core.common.bytes.BytesArray;
1316
import org.opensearch.core.common.bytes.BytesReference;
1417
import org.opensearch.core.xcontent.ToXContent;
1518
import org.opensearch.core.xcontent.XContentBuilder;
@@ -194,7 +197,18 @@ private static void processMetadataFields(SearchHit hit, org.opensearch.protobuf
194197
*/
195198
private static void processSource(SearchHit hit, org.opensearch.protobufs.HitsMetadataHitsInner.Builder hitBuilder) {
196199
if (hit.getSourceRef() != null) {
197-
hitBuilder.setXSource(ByteString.copyFrom(BytesReference.toBytes(hit.getSourceRef())));
200+
BytesReference sourceRef = hit.getSourceRef();
201+
BytesRef bytesRef = sourceRef.toBytesRef();
202+
203+
if (sourceRef instanceof BytesArray) {
204+
if (bytesRef.offset == 0 && bytesRef.length == bytesRef.bytes.length) {
205+
hitBuilder.setXSource(UnsafeByteOperations.unsafeWrap(bytesRef.bytes));
206+
} else {
207+
hitBuilder.setXSource(UnsafeByteOperations.unsafeWrap(bytesRef.bytes, bytesRef.offset, bytesRef.length));
208+
}
209+
} else {
210+
hitBuilder.setXSource(ByteString.copyFrom(bytesRef.bytes, bytesRef.offset, bytesRef.length));
211+
}
198212
}
199213
}
200214

‎modules/transport-grpc/src/test/java/org/opensearch/transport/grpc/proto/response/search/SearchHitProtoUtilsTests.java‎

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.apache.lucene.search.TotalHits;
1313
import org.opensearch.common.document.DocumentField;
1414
import org.opensearch.core.common.bytes.BytesArray;
15+
import org.opensearch.core.common.bytes.CompositeBytesReference;
1516
import org.opensearch.core.common.text.Text;
1617
import org.opensearch.core.index.shard.ShardId;
1718
import org.opensearch.index.seqno.SequenceNumbers;
@@ -296,4 +297,227 @@ public void testToProtoWithScoreStructure() throws IOException {
296297
assertEquals("Negative score value should match", -1.5, hitWithNegativeScore.getXScore().getDouble(), 0.0);
297298
assertFalse("Negative score should not have null value", hitWithNegativeScore.getXScore().hasNullValue());
298299
}
300+
301+
public void testToProtoWithBytesArrayZeroCopyOptimization() throws IOException {
302+
// Create a SearchHit with BytesArray source (should use zero-copy optimization)
303+
SearchHit searchHit = new SearchHit(1);
304+
byte[] sourceBytes = "{\"field\":\"value\",\"number\":42}".getBytes(StandardCharsets.UTF_8);
305+
BytesArray bytesArray = new BytesArray(sourceBytes);
306+
searchHit.sourceRef(bytesArray);
307+
308+
// Call the method under test
309+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
310+
311+
// Verify the result
312+
assertNotNull("Hit should not be null", hit);
313+
assertTrue("Source should be set", hit.hasXSource());
314+
assertArrayEquals("Source bytes should match exactly", sourceBytes, hit.getXSource().toByteArray());
315+
316+
// Verify that the ByteString was created using UnsafeByteOperations.unsafeWrap
317+
// This is an indirect test - we verify the content is correct and assume the optimization was used
318+
assertEquals("Source size should match", sourceBytes.length, hit.getXSource().size());
319+
}
320+
321+
public void testToProtoWithBytesArrayWithOffsetZeroCopyOptimization() throws IOException {
322+
// Create a SearchHit with BytesArray source that has offset/length (should use zero-copy with offset)
323+
SearchHit searchHit = new SearchHit(1);
324+
byte[] fullBytes = "prefix{\"field\":\"value\"}suffix".getBytes(StandardCharsets.UTF_8);
325+
byte[] expectedBytes = "{\"field\":\"value\"}".getBytes(StandardCharsets.UTF_8);
326+
int offset = 6; // "prefix".length()
327+
int length = expectedBytes.length;
328+
BytesArray bytesArray = new BytesArray(fullBytes, offset, length);
329+
searchHit.sourceRef(bytesArray);
330+
331+
// Call the method under test
332+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
333+
334+
// Verify the result
335+
assertNotNull("Hit should not be null", hit);
336+
assertTrue("Source should be set", hit.hasXSource());
337+
assertArrayEquals("Source bytes should match the sliced portion", expectedBytes, hit.getXSource().toByteArray());
338+
assertEquals("Source size should match expected length", length, hit.getXSource().size());
339+
}
340+
341+
public void testToProtoWithCompositeBytesReferenceUsesDeepCopy() throws IOException {
342+
// Create a SearchHit with CompositeBytesReference source (should use ByteString.copyFrom)
343+
SearchHit searchHit = new SearchHit(1);
344+
byte[] bytes1 = "{\"field1\":".getBytes(StandardCharsets.UTF_8);
345+
byte[] bytes2 = "\"value1\"}".getBytes(StandardCharsets.UTF_8);
346+
BytesArray part1 = new BytesArray(bytes1);
347+
BytesArray part2 = new BytesArray(bytes2);
348+
CompositeBytesReference compositeBytesRef = (CompositeBytesReference) CompositeBytesReference.of(part1, part2);
349+
searchHit.sourceRef(compositeBytesRef);
350+
351+
// Call the method under test
352+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
353+
354+
// Verify the result
355+
assertNotNull("Hit should not be null", hit);
356+
assertTrue("Source should be set", hit.hasXSource());
357+
358+
// Verify the combined content is correct
359+
String expectedJson = "{\"field1\":\"value1\"}";
360+
byte[] expectedBytes = expectedJson.getBytes(StandardCharsets.UTF_8);
361+
assertArrayEquals("Source bytes should match the combined content", expectedBytes, hit.getXSource().toByteArray());
362+
assertEquals("Source size should match combined length", expectedBytes.length, hit.getXSource().size());
363+
}
364+
365+
public void testToProtoWithEmptyBytesArraySource() throws IOException {
366+
// Create a SearchHit with minimal valid JSON as source (empty object)
367+
SearchHit searchHit = new SearchHit(1);
368+
byte[] emptyJsonBytes = "{}".getBytes(StandardCharsets.UTF_8);
369+
BytesArray emptyJsonBytesArray = new BytesArray(emptyJsonBytes);
370+
searchHit.sourceRef(emptyJsonBytesArray);
371+
372+
// Call the method under test
373+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
374+
375+
// Verify the result
376+
assertNotNull("Hit should not be null", hit);
377+
assertTrue("Source should be set", hit.hasXSource());
378+
assertEquals("Source should contain empty JSON object", emptyJsonBytes.length, hit.getXSource().size());
379+
assertArrayEquals("Source bytes should match empty JSON object", emptyJsonBytes, hit.getXSource().toByteArray());
380+
}
381+
382+
public void testToProtoWithLargeBytesArrayZeroCopyOptimization() throws IOException {
383+
// Create a SearchHit with large BytesArray source to test performance benefit
384+
SearchHit searchHit = new SearchHit(1);
385+
StringBuilder largeJsonBuilder = new StringBuilder("{\"data\":[");
386+
for (int i = 0; i < 1000; i++) {
387+
if (i > 0) largeJsonBuilder.append(",");
388+
largeJsonBuilder.append("{\"id\":").append(i).append(",\"value\":\"item").append(i).append("\"}");
389+
}
390+
largeJsonBuilder.append("]}");
391+
392+
byte[] largeSourceBytes = largeJsonBuilder.toString().getBytes(StandardCharsets.UTF_8);
393+
BytesArray largeBytesArray = new BytesArray(largeSourceBytes);
394+
searchHit.sourceRef(largeBytesArray);
395+
396+
// Call the method under test
397+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
398+
399+
// Verify the result
400+
assertNotNull("Hit should not be null", hit);
401+
assertTrue("Source should be set", hit.hasXSource());
402+
assertEquals("Source size should match large content", largeSourceBytes.length, hit.getXSource().size());
403+
assertArrayEquals("Source bytes should match exactly", largeSourceBytes, hit.getXSource().toByteArray());
404+
}
405+
406+
public void testToProtoWithBytesArraySliceZeroCopyOptimization() throws IOException {
407+
// Test the optimization with a BytesArray that represents a slice of a larger array
408+
SearchHit searchHit = new SearchHit(1);
409+
410+
// Create a larger byte array
411+
String fullContent = "HEADER{\"important\":\"data\",\"field\":\"value\"}FOOTER";
412+
byte[] fullBytes = fullContent.getBytes(StandardCharsets.UTF_8);
413+
414+
// Create a BytesArray that represents just the JSON part
415+
int jsonStart = 6; // "HEADER".length()
416+
String jsonContent = "{\"important\":\"data\",\"field\":\"value\"}";
417+
int jsonLength = jsonContent.length();
418+
BytesArray slicedBytesArray = new BytesArray(fullBytes, jsonStart, jsonLength);
419+
searchHit.sourceRef(slicedBytesArray);
420+
421+
// Call the method under test
422+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
423+
424+
// Verify the result
425+
assertNotNull("Hit should not be null", hit);
426+
assertTrue("Source should be set", hit.hasXSource());
427+
assertEquals("Source size should match JSON length", jsonLength, hit.getXSource().size());
428+
429+
byte[] expectedJsonBytes = jsonContent.getBytes(StandardCharsets.UTF_8);
430+
assertArrayEquals("Source bytes should match only the JSON portion", expectedJsonBytes, hit.getXSource().toByteArray());
431+
}
432+
433+
public void testToProtoSourceOptimizationBehaviorComparison() throws IOException {
434+
// Test to demonstrate the difference in behavior between BytesArray and other BytesReference types
435+
String jsonContent = "{\"test\":\"optimization\"}";
436+
byte[] jsonBytes = jsonContent.getBytes(StandardCharsets.UTF_8);
437+
438+
// Test with BytesArray (should use zero-copy optimization)
439+
SearchHit searchHitWithBytesArray = new SearchHit(1);
440+
BytesArray bytesArray = new BytesArray(jsonBytes);
441+
searchHitWithBytesArray.sourceRef(bytesArray);
442+
HitsMetadataHitsInner hitWithBytesArray = SearchHitProtoUtils.toProto(searchHitWithBytesArray);
443+
444+
// Test with CompositeBytesReference (should use deep copy)
445+
SearchHit searchHitWithComposite = new SearchHit(2);
446+
BytesArray part1 = new BytesArray(jsonBytes, 0, jsonBytes.length / 2);
447+
BytesArray part2 = new BytesArray(jsonBytes, jsonBytes.length / 2, jsonBytes.length - jsonBytes.length / 2);
448+
CompositeBytesReference composite = (CompositeBytesReference) CompositeBytesReference.of(part1, part2);
449+
searchHitWithComposite.sourceRef(composite);
450+
HitsMetadataHitsInner hitWithComposite = SearchHitProtoUtils.toProto(searchHitWithComposite);
451+
452+
// Both should produce the same result
453+
assertArrayEquals(
454+
"Both approaches should produce identical byte content",
455+
hitWithBytesArray.getXSource().toByteArray(),
456+
hitWithComposite.getXSource().toByteArray()
457+
);
458+
assertEquals(
459+
"Both approaches should produce same size",
460+
hitWithBytesArray.getXSource().size(),
461+
hitWithComposite.getXSource().size()
462+
);
463+
}
464+
465+
public void testToProtoWithNullSourceRef() throws IOException {
466+
// Test behavior when source reference is null
467+
SearchHit searchHit = new SearchHit(1);
468+
// Don't set any source reference (sourceRef remains null)
469+
470+
// Call the method under test
471+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
472+
473+
// Verify the result
474+
assertNotNull("Hit should not be null", hit);
475+
assertFalse("Source should not be set when sourceRef is null", hit.hasXSource());
476+
}
477+
478+
public void testToProtoWithActuallyEmptyBytesArray() throws IOException {
479+
// Test the edge case of truly empty bytes - this should be handled gracefully
480+
// by checking if the source is null before processing
481+
SearchHit searchHit = new SearchHit(1);
482+
// Explicitly set source to null to test null handling
483+
searchHit.sourceRef(null);
484+
485+
// Call the method under test
486+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
487+
488+
// Verify the result
489+
assertNotNull("Hit should not be null", hit);
490+
assertFalse("Source should not be set when explicitly null", hit.hasXSource());
491+
}
492+
493+
public void testToProtoWithBytesArrayOffsetConditionCoverage() throws IOException {
494+
// Test the specific condition coverage for BytesArray when offset != 0 OR length != bytes.length
495+
// This covers the else branch in the optimization logic (lines 207-208)
496+
SearchHit searchHit = new SearchHit(1);
497+
498+
// Create a larger byte array with prefix and suffix
499+
String prefix = "PREFIX";
500+
String jsonContent = "{\"field\":\"value\"}";
501+
String suffix = "SUFFIX";
502+
String fullContent = prefix + jsonContent + suffix;
503+
byte[] fullBytes = fullContent.getBytes(StandardCharsets.UTF_8);
504+
505+
// Create BytesArray with offset > 0 to extract just the JSON part
506+
// This will trigger the condition: bytesRef.offset != 0 || bytesRef.length != bytesRef.bytes.length
507+
int offset = prefix.length();
508+
int length = jsonContent.length();
509+
BytesArray slicedBytesArray = new BytesArray(fullBytes, offset, length);
510+
searchHit.sourceRef(slicedBytesArray);
511+
512+
// Call the method under test
513+
HitsMetadataHitsInner hit = SearchHitProtoUtils.toProto(searchHit);
514+
515+
// Verify the result - should use the offset/length version of unsafeWrap
516+
assertNotNull("Hit should not be null", hit);
517+
assertTrue("Source should be set", hit.hasXSource());
518+
assertEquals("Source size should match JSON length", length, hit.getXSource().size());
519+
520+
byte[] expectedBytes = jsonContent.getBytes(StandardCharsets.UTF_8);
521+
assertArrayEquals("Source bytes should match JSON content", expectedBytes, hit.getXSource().toByteArray());
522+
}
299523
}

0 commit comments

Comments
 (0)