Skip to content

Commit 10ee388

Browse files
lawofcycleslawofcycles
authored
Enable parallel reads for OpenSearch Serverless via PIT + Slice (#783)
* Enable parallel reads for OpenSearch Serverless via PIT + Slice Remove the validation that rejected opensearch.input.max.docs.per.partition in serverless mode. When this setting is provided, the connector now creates multiple sliced partitions per index using PIT + Slice, enabling parallel reads across Spark tasks. This leverages the sliced search support introduced in next-generation OpenSearch Serverless. On Classic Serverless (which does not support the Slice API), using this setting will result in a server-side error. Changes: - RestService.findServerlessPartitions(): count docs and create sliced partitions when maxDocsPerPartition is set - SearchRequestBuilder.assembleSearchAfterBody(): include slice parameter in search_after request body - ServerlessModeTest: replace rejection test with parallel partition tests - USER_GUIDE.md: add Serverless documentation section Signed-off-by: Sotaro Hikita <bering1814@gmail.com> * Default to parallel reads in serverless mode (50,000 docs/partition) Instead of requiring explicit opensearch.input.max.docs.per.partition to enable parallel reads, serverless mode now automatically splits reads using a default of 50,000 documents per partition. Users can still override this value. Also removes incorrect documentation about Classic Serverless not supporting the Slice API. Signed-off-by: Sotaro Hikita <bering1814@gmail.com> * update document Signed-off-by: Sotaro Hikita <bering1814@gmail.com> --------- Signed-off-by: Sotaro Hikita <bering1814@gmail.com>
1 parent 402a2e8 commit 10ee388

6 files changed

Lines changed: 113 additions & 14 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
2121
- Add Apache Spark 4.0 support ([#684](https://github.com/opensearch-project/opensearch-hadoop/pull/684))
2222
- Add dedicated `opensearch.search_after.size` setting for serverless mode page size ([#695](https://github.com/opensearch-project/opensearch-hadoop/pull/695))
2323
- Add Apache Spark 3.5 support with dedicated opensearch-spark-35 module ([#717](https://github.com/opensearch-project/opensearch-hadoop/pull/717))
24+
- Add parallel read support for next-generation OpenSearch Serverless via PIT + Slice ([#783](https://github.com/opensearch-project/opensearch-hadoop/pull/783))
2425

2526
### Changed
2627
- Switched to more reliable OpenSearch Lucene snapshot location ([#597](https://github.com/opensearch-project/opensearch-hadoop/pull/597))

‎USER_GUIDE.md‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -392,6 +392,28 @@ df = spark.read.format("opensearch") \
392392
.load("my-index")
393393
```
394394

395+
### Parallel Reads
396+
397+
Serverless mode automatically splits reads into parallel partitions using PIT + Slice. By default, the connector creates one partition per 50,000 documents in each index.
398+
399+
To customize the partition size, set `opensearch.input.max.docs.per.partition`:
400+
401+
```python
402+
df = spark.read.format("opensearch") \
403+
.option("opensearch.nodes", "https://xxx.us-east-1.aoss.amazonaws.com") \
404+
.option("opensearch.port", "443") \
405+
.option("opensearch.net.ssl", "true") \
406+
.option("opensearch.nodes.wan.only", "true") \
407+
.option("opensearch.aws.sigv4.enabled", "true") \
408+
.option("opensearch.aws.sigv4.region", "us-east-1") \
409+
.option("opensearch.aws.sigv4.service.name", "aoss") \
410+
.option("opensearch.serverless", "true") \
411+
.option("opensearch.input.max.docs.per.partition", "100000") \
412+
.load("my-index")
413+
```
414+
415+
The connector will count the documents in each index, divide by this value to determine the number of slices, and create one Spark partition per slice. For example, an index with 1,000,000 documents and `max.docs.per.partition=100000` will produce 10 parallel read tasks.
416+
395417
## Map/Reduce
396418

397419
For low-level Hadoop Map/Reduce jobs, opensearch-hadoop provides `OpenSearchInputFormat` and `OpenSearchOutputFormat`. Add `opensearch-hadoop-mr-2.0.0.jar` to your job classpath.
@@ -471,3 +493,4 @@ TBLPROPERTIES(
471493
'opensearch.aws.sigv4.enabled' = 'true',
472494
'opensearch.aws.sigv4.region' = 'us-east-1');
473495
```
496+

‎mr/src/main/java/org/opensearch/hadoop/cfg/ConfigurationOptions.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,7 @@ public interface ConfigurationOptions {
164164

165165
/** Input options **/
166166
String OPENSEARCH_MAX_DOCS_PER_PARTITION = "opensearch.input.max.docs.per.partition";
167+
String OPENSEARCH_MAX_DOCS_PER_PARTITION_SERVERLESS_DEFAULT = "50000";
167168

168169
String OPENSEARCH_INPUT_JSON = "opensearch.input.json";
169170
String OPENSEARCH_INPUT_JSON_DEFAULT = "no";

‎mr/src/main/java/org/opensearch/hadoop/rest/RestService.java‎

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -229,15 +229,10 @@ public static List<PartitionDefinition> findPartitions(Settings settings, Log lo
229229

230230
/**
231231
* Create partitions for OpenSearch Serverless mode, which doesn't support shard APIs.
232-
* This function creates a single partition for each index as serverless doesn't expose shard information.
232+
* When maxDocsPerPartition is set, creates multiple sliced partitions per index for parallel reads
233+
* using PIT + Slice.
233234
*/
234235
static List<PartitionDefinition> findServerlessPartitions(RestRepository client, Settings settings, MappingSet mappingSet, Log log) {
235-
if (settings.getMaxDocsPerPartition() != null) {
236-
throw new OpenSearchHadoopIllegalArgumentException(
237-
"maxDocsPerPartition setting is not supported in OpenSearch Serverless mode. " +
238-
"Serverless does not support Slice API which is required for parallel partition reads.");
239-
}
240-
241236
Resource readResource = new Resource(settings, true);
242237
Mapping resolvedMapping = mappingSet == null ? null : mappingSet.getResolvedView();
243238
PartitionDefinition.PartitionDefinitionBuilder partitionBuilder = PartitionDefinition.builder(settings, resolvedMapping);
@@ -247,9 +242,25 @@ static List<PartitionDefinition> findServerlessPartitions(RestRepository client,
247242
List<PartitionDefinition> partitions = new ArrayList<PartitionDefinition>();
248243
String[] indices = readResource.index().split(",");
249244

245+
Integer maxDocsPerPartition = settings.getMaxDocsPerPartition();
246+
if (maxDocsPerPartition == null) {
247+
maxDocsPerPartition = Integer.parseInt(ConfigurationOptions.OPENSEARCH_MAX_DOCS_PER_PARTITION_SERVERLESS_DEFAULT);
248+
}
249+
250250
for (String indexName : indices) {
251251
indexName = indexName.trim();
252-
partitions.add(partitionBuilder.build(indexName, 0, new String[0]));
252+
QueryBuilder query = QueryUtils.parseQueryAndFilters(settings);
253+
long numDocs = client.getRestClient().count(indexName, query);
254+
int numPartitions = (int) Math.max(1, numDocs / maxDocsPerPartition);
255+
if (numPartitions > 1) {
256+
log.info(String.format("Serverless parallel read: index [%s] has [%d] docs, creating [%d] sliced partitions", indexName, numDocs, numPartitions));
257+
for (int i = 0; i < numPartitions; i++) {
258+
PartitionDefinition.Slice slice = new PartitionDefinition.Slice(i, numPartitions);
259+
partitions.add(partitionBuilder.build(indexName, 0, slice, new String[0]));
260+
}
261+
} else {
262+
partitions.add(partitionBuilder.build(indexName, 0, new String[0]));
263+
}
253264
}
254265

255266
return partitions;

‎mr/src/main/java/org/opensearch/hadoop/rest/SearchRequestBuilder.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,6 +377,15 @@ private BytesArray assembleSearchAfterBody() {
377377
JacksonJsonGenerator generator = new JacksonJsonGenerator(out);
378378
try {
379379
generator.writeBeginObject();
380+
if (slice != null && slice.max > 1) {
381+
generator.writeFieldName("slice");
382+
generator.writeBeginObject();
383+
generator.writeFieldName("id");
384+
generator.writeNumber(slice.id);
385+
generator.writeFieldName("max");
386+
generator.writeNumber(slice.max);
387+
generator.writeEndObject();
388+
}
380389
generator.writeFieldName("query");
381390
generator.writeBeginObject();
382391
root.toJson(generator);

‎mr/src/test/java/org/opensearch/hadoop/rest/ServerlessModeTest.java‎

Lines changed: 60 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
import org.apache.commons.logging.impl.NoOpLog;
1515
import org.junit.Test;
1616
import org.mockito.Mockito;
17-
import org.opensearch.hadoop.OpenSearchHadoopIllegalArgumentException;
1817
import org.opensearch.hadoop.cfg.ConfigurationOptions;
1918
import org.opensearch.hadoop.cfg.PropertiesSettings;
2019
import org.opensearch.hadoop.cfg.Settings;
@@ -58,8 +57,16 @@ public void testFindServerlessPartitionsSingleIndex() {
5857
settings.setProperty(ConfigurationOptions.OPENSEARCH_RESOURCE_READ, "test-index");
5958
settings.setServerlessMode(true);
6059

61-
List<PartitionDefinition> partitions = RestService.findServerlessPartitions(null, settings, null, LOGGER);
60+
RestClient restClient = Mockito.mock(RestClient.class);
61+
Mockito.when(restClient.count(Mockito.eq("test-index"), Mockito.any()))
62+
.thenReturn(100L);
6263

64+
RestRepository repository = Mockito.mock(RestRepository.class);
65+
Mockito.when(repository.getRestClient()).thenReturn(restClient);
66+
67+
List<PartitionDefinition> partitions = RestService.findServerlessPartitions(repository, settings, null, LOGGER);
68+
69+
// 100 docs / 50000 default = 1 partition (no slicing needed)
6370
assertEquals(1, partitions.size());
6471
assertEquals("test-index", partitions.get(0).getIndex());
6572
assertEquals(0, partitions.get(0).getShardId());
@@ -71,22 +78,69 @@ public void testFindServerlessPartitionsMultipleIndices() {
7178
settings.setProperty(ConfigurationOptions.OPENSEARCH_RESOURCE_READ, "index1,index2,index3");
7279
settings.setServerlessMode(true);
7380

74-
List<PartitionDefinition> partitions = RestService.findServerlessPartitions(null, settings, null, LOGGER);
81+
RestClient restClient = Mockito.mock(RestClient.class);
82+
Mockito.when(restClient.count(Mockito.anyString(), Mockito.any()))
83+
.thenReturn(100L);
84+
85+
RestRepository repository = Mockito.mock(RestRepository.class);
86+
Mockito.when(repository.getRestClient()).thenReturn(restClient);
87+
88+
List<PartitionDefinition> partitions = RestService.findServerlessPartitions(repository, settings, null, LOGGER);
7589

7690
assertEquals(3, partitions.size());
7791
assertEquals("index1", partitions.get(0).getIndex());
7892
assertEquals("index2", partitions.get(1).getIndex());
7993
assertEquals("index3", partitions.get(2).getIndex());
8094
}
8195

82-
@Test(expected = OpenSearchHadoopIllegalArgumentException.class)
83-
public void testFindServerlessPartitionsRejectsMaxDocsPerPartition() {
96+
@Test
97+
public void testFindServerlessPartitionsWithMaxDocsPerPartition() {
8498
Settings settings = new PropertiesSettings();
8599
settings.setProperty(ConfigurationOptions.OPENSEARCH_RESOURCE_READ, "test-index");
86100
settings.setServerlessMode(true);
101+
settings.setMaxDocsPerPartition(500);
102+
103+
RestClient restClient = Mockito.mock(RestClient.class);
104+
Mockito.when(restClient.count(Mockito.eq("test-index"), Mockito.any()))
105+
.thenReturn(2000L);
106+
107+
RestRepository repository = Mockito.mock(RestRepository.class);
108+
Mockito.when(repository.getRestClient()).thenReturn(restClient);
109+
110+
List<PartitionDefinition> partitions = RestService.findServerlessPartitions(repository, settings, null, LOGGER);
111+
112+
assertEquals(4, partitions.size());
113+
for (int i = 0; i < 4; i++) {
114+
assertEquals("test-index", partitions.get(i).getIndex());
115+
assertNotNull(partitions.get(i).getSlice());
116+
assertEquals(i, partitions.get(i).getSlice().id);
117+
assertEquals(4, partitions.get(i).getSlice().max);
118+
}
119+
}
120+
121+
@Test
122+
public void testFindServerlessPartitionsWithMaxDocsPerPartitionMultipleIndices() {
123+
Settings settings = new PropertiesSettings();
124+
settings.setProperty(ConfigurationOptions.OPENSEARCH_RESOURCE_READ, "index1,index2");
125+
settings.setServerlessMode(true);
87126
settings.setMaxDocsPerPartition(1000);
88127

89-
RestService.findServerlessPartitions(null, settings, null, LOGGER);
128+
RestClient restClient = Mockito.mock(RestClient.class);
129+
Mockito.when(restClient.count(Mockito.eq("index1"), Mockito.any()))
130+
.thenReturn(3000L);
131+
Mockito.when(restClient.count(Mockito.eq("index2"), Mockito.any()))
132+
.thenReturn(1500L);
133+
134+
RestRepository repository = Mockito.mock(RestRepository.class);
135+
Mockito.when(repository.getRestClient()).thenReturn(restClient);
136+
137+
List<PartitionDefinition> partitions = RestService.findServerlessPartitions(repository, settings, null, LOGGER);
138+
139+
// index1: 3000/1000 = 3 partitions, index2: 1500/1000 = 1 partition
140+
assertEquals(4, partitions.size());
141+
assertEquals("index1", partitions.get(0).getIndex());
142+
assertEquals("index1", partitions.get(2).getIndex());
143+
assertEquals("index2", partitions.get(3).getIndex());
90144
}
91145

92146
@Test

0 commit comments

Comments
 (0)