From 7e0b75ac034f4a46d827c04af477bc0a61593b42 Mon Sep 17 00:00:00 2001 From: John Mazanec Date: Wed, 29 Jan 2025 21:25:08 -0800 Subject: [PATCH] Switch away from getMergeInstance Signed-off-by: John Mazanec --- .../DerivedSourceStoredFieldsReader.java | 20 ++++++++++++++++--- .../DerivedSourceStoredFieldsWriter.java | 4 ++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsReader.java b/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsReader.java index 6c1ade140..24900eb19 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsReader.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsReader.java @@ -107,12 +107,12 @@ public void close() throws IOException { /** * For merging, we need to tell the derived source stored fields reader to skip injecting the source. Otherwise, - * on merge we will end up just writing the source to disk + * on merge we will end up just writing the source to disk. We cant override + * {@link StoredFieldsReader#getMergeInstance()} because it is used elsewhere than just merging. * * @return Merged instance that wont inject by default */ - @Override - public StoredFieldsReader getMergeInstance() { + private StoredFieldsReader cloneForMerge() { try { return new DerivedSourceStoredFieldsReader( delegate.getMergeInstance(), @@ -125,4 +125,18 @@ public StoredFieldsReader getMergeInstance() { throw new RuntimeException(e); } } + + /** + * For merging, we need to tell the derived source stored fields reader to skip injecting the source. Otherwise, + * on merge we will end up just writing the source to disk + * + * @param storedFieldsReader stored fields reader to wrap + * @return wrapped stored fields reader + */ + public static StoredFieldsReader wrapForMerge(StoredFieldsReader storedFieldsReader) { + if (storedFieldsReader instanceof DerivedSourceStoredFieldsReader) { + return ((DerivedSourceStoredFieldsReader) storedFieldsReader).cloneForMerge(); + } + return storedFieldsReader; + } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsWriter.java b/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsWriter.java index 576bc9e98..c585b09f7 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsWriter.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsWriter.java @@ -65,6 +65,10 @@ public void writeField(FieldInfo info, DataInput value, int length) throws IOExc @Override public int merge(MergeState mergeState) throws IOException { + // We have to wrap these here to avoid storing the vectors during merge + for (int i = 0; i < mergeState.storedFieldsReaders.length; i++) { + mergeState.storedFieldsReaders[i] = DerivedSourceStoredFieldsReader.wrapForMerge(mergeState.storedFieldsReaders[i]); + } return delegate.merge(mergeState); }