Replace Scroll API with Search After API for read pagination in serverless mode

Signed-off-by: Sotaro Hikita <bering1814@gmail.com>

Author: lawofcycles

mr/src/main/java/org/opensearch/hadoop/rest/RestClient.java

@@ -641,6 +641,35 @@ static BytesArray searchRequest(QueryBuilder query) {
         return out.bytes();
     }
 
+    /**
+     * Build a search_after request body by appending search_after to the original query body.
+     */
+    public BytesArray buildSearchAfterBody(Object[] searchAfter) {
+        FastByteArrayOutputStream out = new FastByteArrayOutputStream(256);
+        JacksonJsonGenerator generator = new JacksonJsonGenerator(out);
+        try {
+            generator.writeBeginObject();
+            generator.writeFieldName("search_after");
+            generator.writeBeginArray();
+            for (Object value : searchAfter) {
+                if (value instanceof Long) {
+                    generator.writeNumber((Long) value);
+                } else if (value instanceof Double) {
+                    generator.writeNumber((Double) value);
+                } else if (value instanceof Integer) {
+                    generator.writeNumber((Integer) value);
+                } else {
+                    generator.writeString(String.valueOf(value));
+                }
+            }
+            generator.writeEndArray();
+            generator.writeEndObject();
+        } finally {
+            generator.close();
+        }
+        return out.bytes();
+    }
+
     public boolean isAlias(String query) {
         Map<String, Object> aliases = (Map<String, Object>) get(query, null);
         return (aliases.size() > 1);

mr/src/main/java/org/opensearch/hadoop/rest/RestRepository.java

@@ -348,6 +348,39 @@ Scroll scroll(String scrollId, ScrollReader reader) throws IOException {
         }
     }
 
+    Scroll searchAfter(String queryUri, BytesArray baseBody, Object[] searchAfter, ScrollReader reader) throws IOException {
+        BytesArray body = mergeSearchAfterIntoBody(baseBody, searchAfter);
+        InputStream response = client.execute(Request.Method.POST, queryUri, body).body();
+        try {
+            return reader.read(response);
+        } finally {
+            if (response instanceof StatsAware) {
+                stats.aggregate(((StatsAware) response).stats());
+            }
+        }
+    }
+
+    /**
+     * Merge search_after values into the base request body.
+     * Inserts "search_after":[...] before the closing brace of the JSON body.
+     */
+    private BytesArray mergeSearchAfterIntoBody(BytesArray baseBody, Object[] searchAfter) {
+        BytesArray searchAfterJson = client.buildSearchAfterBody(searchAfter);
+        // baseBody: {...}  searchAfterJson: {"search_after":[...]}
+        // Result: base body with search_after inserted
+        String base = baseBody.toString().trim();
+        String saFragment = searchAfterJson.toString().trim();
+        // Remove outer braces from search_after fragment: {"search_after":[...]} -> "search_after":[...]
+        String saContent = saFragment.substring(1, saFragment.length() - 1);
+        // Insert before closing brace of base body
+        String merged = base.substring(0, base.length() - 1) + "," + saContent + "}";
+        return new BytesArray(merged);
+    }
+
+    ScrollQuery scanLimitSearchAfter(String query, BytesArray body, long limit, ScrollReader reader) {
+        return new ScrollQuery(this, query, body, limit, reader, true);
+    }
+
     public boolean resourceExists(boolean read) {
         Resource res = (read ? resources.getResourceRead() : resources.getResourceWrite());
         // cheap hit - works for exact index names, index patterns, the `_all` resource, and alias names

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

@@ -78,15 +78,21 @@ public static class PartitionReader implements Closeable {
         public final ScrollReader scrollReader;
         public final RestRepository client;
         public final SearchRequestBuilder queryBuilder;
+        private final boolean serverlessMode;
 
         private ScrollQuery scrollQuery;
 
         private boolean closed = false;
 
         PartitionReader(ScrollReader scrollReader, RestRepository client, SearchRequestBuilder queryBuilder) {
+            this(scrollReader, client, queryBuilder, false);
+        }
+
+        PartitionReader(ScrollReader scrollReader, RestRepository client, SearchRequestBuilder queryBuilder, boolean serverlessMode) {
             this.scrollReader = scrollReader;
             this.client = client;
             this.queryBuilder = queryBuilder;
+            this.serverlessMode = serverlessMode;
         }
 
         @Override
@@ -102,7 +108,11 @@ public void close() {
 
         public ScrollQuery scrollQuery() {
             if (scrollQuery == null) {
-                scrollQuery = queryBuilder.build(client, scrollReader);
+                if (serverlessMode) {
+                    scrollQuery = queryBuilder.buildSearchAfter(client, scrollReader);
+                } else {
+                    scrollQuery = queryBuilder.build(client, scrollReader);
+                }
             }
 
             return scrollQuery;
@@ -222,9 +232,10 @@ public static List<PartitionDefinition> findPartitions(Settings settings, Log lo
      * This function creates a single partition for each index as serverless doesn't expose shard information.
      */
     static List<PartitionDefinition> findServerlessPartitions(RestRepository client, Settings settings, MappingSet mappingSet, Log log) {
-        // Check if maxDocsPerPartition is set and throw exception if it is
         if (settings.getMaxDocsPerPartition() != null) {
-            throw new OpenSearchHadoopIllegalArgumentException("maxDocsPerPartition setting is not supported in OpenSearch Serverless mode");
+            throw new OpenSearchHadoopIllegalArgumentException(
+                "maxDocsPerPartition setting is not supported in OpenSearch Serverless mode. " +
+                "Serverless does not support Slice API which is required for parallel partition reads.");
         }
 
         Resource readResource = new Resource(settings, true);
@@ -234,25 +245,13 @@ static List<PartitionDefinition> findServerlessPartitions(RestRepository client,
         log.info(String.format("Reading from [%s] in serverless mode", settings.getResourceRead()));
 
         List<PartitionDefinition> partitions = new ArrayList<PartitionDefinition>();
-
-        // For serverless mode, we create partitions based on the index names in the resource
-        // Since we cannot get shard information, we treat each index as a single partition
         String[] indices = readResource.index().split(",");
-        
+
         for (String indexName : indices) {
             indexName = indexName.trim();
-            // Create a single partition per index with dummy shard ID 0 (serverless mode detection via settings)
             partitions.add(partitionBuilder.build(indexName, 0, new String[0]));
         }
 
-        if (partitions.isEmpty()) {
-            if (log.isDebugEnabled()) {
-                log.debug("No partitions created for serverless mode - creating default partition");
-            }
-            // Fallback: create at least one partition for the primary index
-            partitions.add(partitionBuilder.build(readResource.index(), 0, new String[0]));
-        }
-
         return partitions;
     }
 
@@ -462,7 +461,7 @@ public static PartitionReader createReader(Settings settings, PartitionDefinitio
                 requestBuilder = applyAliasMetadata(clusterInfo.getMajorVersion(), aliases, requestBuilder, partition.getIndex(), indices);
             }
         }
-        return new PartitionReader(scrollReader, repository, requestBuilder);
+        return new PartitionReader(scrollReader, repository, requestBuilder, settings.getServerlessMode());
     }
 
     /**

mr/src/main/java/org/opensearch/hadoop/rest/ScrollQuery.java

@@ -66,13 +66,23 @@ public class ScrollQuery implements Iterator<Object>, Closeable, StatsAware {
     private boolean initialized = false;
     private String query;
     private BytesArray body;
+    private final boolean serverlessMode;
+    private Object[] searchAfter;
+    // Retained for search_after continuation in serverless mode
+    private String searchAfterUri;
+    private BytesArray searchAfterBaseBody;
 
     ScrollQuery(RestRepository client, String query, BytesArray body, long size, ScrollReader reader) {
+        this(client, query, body, size, reader, false);
+    }
+
+    ScrollQuery(RestRepository client, String query, BytesArray body, long size, ScrollReader reader, boolean serverlessMode) {
         this.repository = client;
         this.size = size;
         this.reader = reader;
         this.query = query;
         this.body = body;
+        this.serverlessMode = serverlessMode;
     }
 
     @Override
@@ -84,7 +94,7 @@ public void close() {
             reader.close();
             // typically the scroll is closed after it is consumed so this will trigger a 404
             // however we're closing it either way
-            if (StringUtils.hasText(scrollId)) {
+            if (!serverlessMode && StringUtils.hasText(scrollId)) {
                 repository.getRestClient().deleteScroll(scrollId);
             }
             repository.close();
@@ -111,12 +121,18 @@ public boolean hasNext() {
                 scrollId = scroll.getScrollId();
                 batch = scroll.getHits();
                 finished = scroll.isConcluded();
+                searchAfter = scroll.getSearchAfter();
             } catch (IOException ex) {
                 throw new OpenSearchHadoopIllegalStateException(String.format("Cannot create scroll for query [%s/%s]", query, body), ex);
             }
             read += batch.size();
             stats.docsReceived += batch.size();
-            // no longer needed
+            if (serverlessMode) {
+                // Retain query URI and body for search_after continuation
+                searchAfterUri = query;
+                searchAfterBaseBody = body;
+            }
+            // no longer needed for scroll mode
             body = null;
             query = null;
         }
@@ -128,14 +144,25 @@ public boolean hasNext() {
             }
 
             try {
-                Scroll scroll = repository.scroll(scrollId, reader);
+                Scroll scroll;
+                if (serverlessMode) {
+                    scroll = repository.searchAfter(searchAfterUri, searchAfterBaseBody, searchAfter, reader);
+                } else {
+                    scroll = repository.scroll(scrollId, reader);
+                }
                 if (scroll == null) {
                     finished = true;
                     return false;
                 }
                 scrollId = scroll.getScrollId();
                 batch = scroll.getHits();
                 finished = scroll.isConcluded();
+                if (serverlessMode) {
+                    searchAfter = scroll.getSearchAfter();
+                    if (searchAfter == null) {
+                        finished = true;
+                    }
+                }
             } catch (IOException ex) {
                 throw new OpenSearchHadoopIllegalStateException("Cannot retrieve scroll [" + scrollId + "]", ex);
             }

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

@@ -314,6 +314,95 @@ public ScrollQuery build(RestRepository client, ScrollReader reader) {
         return client.scanLimit(scrollUri, requestBody, limit, reader);
     }
 
+    /**
+     * Build a ScrollQuery using Search After API for serverless mode.
+     * Differs from build(): no scroll parameter in URI, sort added to body.
+     */
+    public ScrollQuery buildSearchAfter(RestRepository client, ScrollReader reader) {
+        String searchUri = assembleSearchAfterUri();
+        BytesArray requestBody = assembleSearchAfterBody();
+        return client.scanLimitSearchAfter(searchUri, requestBody, limit, reader);
+    }
+
+    private String assembleSearchAfterUri() {
+        if (limit > 0 && size > limit) {
+            size = limit;
+        }
+        Map<String, String> uriParams = new LinkedHashMap<String, String>();
+        StringBuilder sb = new StringBuilder();
+        sb.append(indices);
+        if (StringUtils.hasLength(types)) {
+            sb.append("/");
+            sb.append(types);
+        }
+        sb.append("/_search?");
+
+        uriParams.put("size", String.valueOf(size));
+        if (includeVersion) {
+            uriParams.put("version", "true");
+        }
+        if (routing != null) {
+            uriParams.put("routing", HttpEncodingTools.encode(routing));
+        }
+        uriParams.put("track_total_hits", "true");
+        if (readMetadata) {
+            uriParams.put("track_scores", "true");
+        }
+
+        for (Iterator<Entry<String, String>> it = uriParams.entrySet().iterator(); it.hasNext();) {
+            Entry<String, String> entry = it.next();
+            sb.append(entry.getKey());
+            if (StringUtils.hasText(entry.getValue())) {
+                sb.append("=");
+                sb.append(entry.getValue());
+            }
+            if (it.hasNext()) {
+                sb.append("&");
+            }
+        }
+        return sb.toString();
+    }
+
+    private BytesArray assembleSearchAfterBody() {
+        QueryBuilder root = query;
+        if (root == null) {
+            root = MatchAllQueryBuilder.MATCH_ALL;
+        }
+        if (filters.isEmpty() == false) {
+            root = new BoolQueryBuilder().must(root).filters(filters);
+        }
+        FastByteArrayOutputStream out = new FastByteArrayOutputStream(256);
+        JacksonJsonGenerator generator = new JacksonJsonGenerator(out);
+        try {
+            generator.writeBeginObject();
+            generator.writeFieldName("query");
+            generator.writeBeginObject();
+            root.toJson(generator);
+            generator.writeEndObject();
+            // sort by _doc for search_after pagination
+            generator.writeFieldName("sort");
+            generator.writeBeginArray();
+            generator.writeString("_doc");
+            generator.writeEndArray();
+            if (StringUtils.hasText(fields)) {
+                generator.writeFieldName("_source");
+                generator.writeBeginArray();
+                final List<String> fieldsArray = StringUtils.tokenize(fields, StringUtils.DEFAULT_DELIMITER);
+                for (String field : fieldsArray) {
+                    generator.writeString(field);
+                }
+                generator.writeEndArray();
+            } else if (excludeSource) {
+                generator.writeFieldName("_source");
+                generator.writeBoolean(false);
+            }
+            generator.writeEndObject();
+        } finally {
+            generator.close();
+        }
+        return out.bytes();
+    }
+
     @Override
     public String toString() {
         return "QueryBuilder [" + assemble() + "][" + assembleBody() + "]";