diff --git a/src/main/java/org/elasticsearch/search/highlight/vectorhighlight/SourceScoreOrderFragmentsBuilder.java b/src/main/java/org/elasticsearch/search/highlight/vectorhighlight/SourceScoreOrderFragmentsBuilder.java
index 84ec390b2e6..1ea8deb9069 100644
--- a/src/main/java/org/elasticsearch/search/highlight/vectorhighlight/SourceScoreOrderFragmentsBuilder.java
+++ b/src/main/java/org/elasticsearch/search/highlight/vectorhighlight/SourceScoreOrderFragmentsBuilder.java
@@ -20,6 +20,7 @@
 package org.elasticsearch.search.highlight.vectorhighlight;
 
 import org.apache.lucene.document.Field;
+import org.apache.lucene.index.AtomicReaderContext;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.vectorhighlight.BoundaryScanner;
 import org.apache.lucene.search.vectorhighlight.XScoreOrderFragmentsBuilder;
@@ -50,7 +51,7 @@ public class SourceScoreOrderFragmentsBuilder extends XScoreOrderFragmentsBuilde
     protected Field[] getFields(IndexReader reader, int docId, String fieldName) throws IOException {
         // we know its low level reader, and matching docId, since that's how we call the highlighter with
         SearchLookup lookup = searchContext.lookup();
-        lookup.setNextReader(reader);
+        lookup.setNextReader((AtomicReaderContext) reader.getContext());
         lookup.setNextDocId(docId);
 
         List<Object> values = lookup.source().extractRawValues(mapper.names().sourcePath());
diff --git a/src/main/java/org/elasticsearch/search/highlight/vectorhighlight/SourceSimpleFragmentsBuilder.java b/src/main/java/org/elasticsearch/search/highlight/vectorhighlight/SourceSimpleFragmentsBuilder.java
index f3fd38ff1b3..1bf3cc2ba6d 100644
--- a/src/main/java/org/elasticsearch/search/highlight/vectorhighlight/SourceSimpleFragmentsBuilder.java
+++ b/src/main/java/org/elasticsearch/search/highlight/vectorhighlight/SourceSimpleFragmentsBuilder.java
@@ -20,6 +20,7 @@
 package org.elasticsearch.search.highlight.vectorhighlight;
 
 import org.apache.lucene.document.Field;
+import org.apache.lucene.index.AtomicReaderContext;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.vectorhighlight.BoundaryScanner;
 import org.apache.lucene.search.vectorhighlight.XSimpleFragmentsBuilder;
@@ -52,7 +53,7 @@ public class SourceSimpleFragmentsBuilder extends XSimpleFragmentsBuilder {
     protected Field[] getFields(IndexReader reader, int docId, String fieldName) throws IOException {
         // we know its low level reader, and matching docId, since that's how we call the highlighter with
         SearchLookup lookup = searchContext.lookup();
-        lookup.setNextReader(reader);
+        lookup.setNextReader((AtomicReaderContext) reader.getContext());
         lookup.setNextDocId(docId);
 
         List<Object> values = lookup.source().extractRawValues(mapper.names().sourcePath());
diff --git a/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java b/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java
index dbe9e03811d..efa151b5cd5 100644
--- a/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java
+++ b/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java
@@ -20,11 +20,8 @@
 package org.elasticsearch.search.lookup;
 
 import com.google.common.collect.ImmutableMap;
-import org.apache.lucene.index.AtomicReader;
 import org.apache.lucene.index.AtomicReaderContext;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.Scorer;
-import org.elasticsearch.ElasticSearchIllegalArgumentException;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.cache.field.data.FieldDataCache;
 import org.elasticsearch.index.mapper.MapperService;
@@ -80,14 +77,4 @@ public class SearchLookup {
         sourceLookup.setNextDocId(docId);
         fieldsLookup.setNextDocId(docId);
     }
-
-    // Sneaky way fvh to get the ARC. See SourceScoreOrderFragmentsBuilder. Passed reader must be atomic.
-    // but that is the case in SourceScoreOrderFragmentsBuilder.
-    public void setNextReader(IndexReader reader) {
-        if (AtomicReader.class.isInstance(reader)) {
-            throw new ElasticSearchIllegalArgumentException("reader not instance of AtomicReader, but " + reader.getClass());
-        }
-        AtomicReaderContext context = (AtomicReaderContext) reader.getContext();
-        setNextReader(context);
-    }
 }