diff --git a/core/src/main/java/org/elasticsearch/search/internal/ScrollContext.java b/core/src/main/java/org/elasticsearch/search/internal/ScrollContext.java index 1b7bcfb93c7..163dbcc73d9 100644 --- a/core/src/main/java/org/elasticsearch/search/internal/ScrollContext.java +++ b/core/src/main/java/org/elasticsearch/search/internal/ScrollContext.java @@ -22,11 +22,35 @@ package org.elasticsearch.search.internal; import org.apache.lucene.search.ScoreDoc; import org.elasticsearch.search.Scroll; +import java.util.HashMap; +import java.util.Map; + /** Wrapper around information that needs to stay around when scrolling. */ -public class ScrollContext { +public final class ScrollContext { + + private Map context = null; public int totalHits = -1; public float maxScore; public ScoreDoc lastEmittedDoc; public Scroll scroll; + + /** + * Returns the object or null if the given key does not have a + * value in the context + */ + @SuppressWarnings("unchecked") // (T)object + public T getFromContext(String key) { + return context != null ? (T) context.get(key) : null; + } + + /** + * Puts the object into the context + */ + public void putInContext(String key, Object value) { + if (context == null) { + context = new HashMap<>(); + } + context.put(key, value); + } } diff --git a/core/src/test/java/org/elasticsearch/search/internal/ScrollContextTests.java b/core/src/test/java/org/elasticsearch/search/internal/ScrollContextTests.java new file mode 100644 index 00000000000..de4863dd92a --- /dev/null +++ b/core/src/test/java/org/elasticsearch/search/internal/ScrollContextTests.java @@ -0,0 +1,36 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.internal; + +import org.elasticsearch.test.ESTestCase; + +public class ScrollContextTests extends ESTestCase { + + public void testStoringObjectsInScrollContext() { + final ScrollContext scrollContext = new ScrollContext(); + final String key = randomAlphaOfLengthBetween(1, 16); + assertNull(scrollContext.getFromContext(key)); + + final String value = randomAlphaOfLength(6); + scrollContext.putInContext(key, value); + + assertEquals(value, scrollContext.getFromContext(key)); + } +}