diff --git a/server/src/main/java/org/elasticsearch/index/mapper/BlockLoaderStoredFieldsFromLeafLoader.java b/server/src/main/java/org/elasticsearch/index/mapper/BlockLoaderStoredFieldsFromLeafLoader.java index cf2cc1331b0ae..a197db1c44d1d 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/BlockLoaderStoredFieldsFromLeafLoader.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/BlockLoaderStoredFieldsFromLeafLoader.java @@ -35,7 +35,9 @@ public void advanceTo(int docId) { private void advanceIfNeeded() throws IOException { if (loaderDocId != docId) { - loader.advanceTo(docId); + if (loader != null) { + loader.advanceTo(docId); + } loaderDocId = docId; } } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BlockScriptReader.java b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BlockScriptReader.java new file mode 100644 index 0000000000000..c37f33c39d54d --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BlockScriptReader.java @@ -0,0 +1,97 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.mapper.blockloader.script; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.SortedSetDocValues; +import org.apache.lucene.util.IOFunction; +import org.elasticsearch.common.breaker.CircuitBreaker; +import org.elasticsearch.common.unit.ByteSizeValue; +import org.elasticsearch.index.mapper.BlockLoader; +import org.elasticsearch.search.fetch.StoredFieldsSpec; + +import java.io.IOException; + +/** + * A reader that supports reading doc-values from a Lucene segment in Block fashion. + */ +public abstract class BlockScriptReader implements BlockLoader.RowStrideReader { + protected final CircuitBreaker breaker; + private final long byteSize; + private final Thread creationThread; + + public BlockScriptReader(CircuitBreaker breaker, long byteSize) { + this.breaker = breaker; + this.byteSize = byteSize; + this.creationThread = Thread.currentThread(); + } + + protected abstract int docId(); + + /** + * Checks if the reader can be used to read a range documents starting with the given docID by the current thread. + */ + @Override + public final boolean canReuse(int startingDocID) { + return creationThread == Thread.currentThread() && docId() <= startingDocID; + } + + @Override + public final void close() { + breaker.addWithoutBreaking(-byteSize); + } + + @Override + public abstract String toString(); + + public abstract static class ScriptBlockLoader implements BlockLoader { + private final long byteSize; + + protected ScriptBlockLoader(ByteSizeValue byteSize) { + this.byteSize = byteSize.getBytes(); + } + + public abstract BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException; + + @Override + public final IOFunction columnAtATimeReader(LeafReaderContext context) { + return null; + } + + @Override + public final RowStrideReader rowStrideReader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { + breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks"); + RowStrideReader reader = null; + try { + reader = reader(breaker, context); + return reader; + } finally { + if (reader == null) { + breaker.addWithoutBreaking(-byteSize); + } + } + } + + @Override + public final StoredFieldsSpec rowStrideStoredFieldSpec() { + return StoredFieldsSpec.NO_REQUIREMENTS; + } + + @Override + public boolean supportsOrdinals() { + return false; + } + + @Override + public SortedSetDocValues ordinals(LeafReaderContext context) throws IOException { + throw new UnsupportedOperationException(); + } + } +} diff --git a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BooleanScriptBlockDocValuesReader.java b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BooleanScriptBlockDocValuesReader.java index 499df2d829454..b31f8a736800f 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BooleanScriptBlockDocValuesReader.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BooleanScriptBlockDocValuesReader.java @@ -21,12 +21,13 @@ /** * {@link BlockDocValuesReader} implementation for {@code boolean} scripts. */ -public class BooleanScriptBlockDocValuesReader extends BlockDocValuesReader { - public static class BooleanScriptBlockLoader extends DocValuesBlockLoader { +public class BooleanScriptBlockDocValuesReader extends BlockScriptReader { + public static class BooleanScriptBlockLoader extends ScriptBlockLoader { private final BooleanFieldScript.LeafFactory factory; private final long byteSize; public BooleanScriptBlockLoader(BooleanFieldScript.LeafFactory factory, ByteSizeValue byteSize) { + super(byteSize); this.factory = factory; this.byteSize = byteSize.getBytes(); } @@ -37,28 +38,17 @@ public Builder builder(BlockFactory factory, int expectedCount) { } @Override - public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { - breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks"); - BooleanFieldScript script = null; - try { - script = factory.newInstance(context); - return new BooleanScriptBlockDocValuesReader(breaker, script, byteSize); - } finally { - if (script == null) { - breaker.addWithoutBreaking(-byteSize); - } - } + public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { + return new BooleanScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize); } } private final BooleanFieldScript script; - private final long byteSize; private int docId; BooleanScriptBlockDocValuesReader(CircuitBreaker breaker, BooleanFieldScript script, long byteSize) { - super(breaker); + super(breaker, byteSize); this.script = script; - this.byteSize = byteSize; } @Override @@ -67,24 +57,9 @@ public int docId() { } @Override - public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered) - throws IOException { - // Note that we don't emit falses before trues so we conform to the doc values contract and can use booleansFromDocValues - try (BlockLoader.BooleanBuilder builder = factory.booleans(docs.count() - offset)) { - for (int i = offset; i < docs.count(); i++) { - read(docs.get(i), builder); - } - return builder.build(); - } - } - - @Override - public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException { + public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException { + BlockLoader.BooleanBuilder builder = (BlockLoader.BooleanBuilder) b; this.docId = docId; - read(docId, (BlockLoader.BooleanBuilder) builder); - } - - private void read(int docId, BlockLoader.BooleanBuilder builder) { script.runForDoc(docId); int total = script.falses() + script.trues(); switch (total) { @@ -107,9 +82,4 @@ private void read(int docId, BlockLoader.BooleanBuilder builder) { public String toString() { return "ScriptBooleans"; } - - @Override - public void close() { - breaker.addWithoutBreaking(-byteSize); - } } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DateScriptBlockDocValuesReader.java b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DateScriptBlockDocValuesReader.java index 7199dfc4eb2da..042b1cfde99d4 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DateScriptBlockDocValuesReader.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DateScriptBlockDocValuesReader.java @@ -21,12 +21,13 @@ /** * {@link BlockDocValuesReader} implementation for date scripts. */ -public class DateScriptBlockDocValuesReader extends BlockDocValuesReader { - public static class DateScriptBlockLoader extends DocValuesBlockLoader { +public class DateScriptBlockDocValuesReader extends BlockScriptReader { + public static class DateScriptBlockLoader extends ScriptBlockLoader { private final DateFieldScript.LeafFactory factory; private final long byteSize; public DateScriptBlockLoader(DateFieldScript.LeafFactory factory, ByteSizeValue byteSize) { + super(byteSize); this.factory = factory; this.byteSize = byteSize.getBytes(); } @@ -37,28 +38,17 @@ public Builder builder(BlockFactory factory, int expectedCount) { } @Override - public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { - breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks"); - DateFieldScript script = null; - try { - script = factory.newInstance(context); - return new DateScriptBlockDocValuesReader(breaker, script, byteSize); - } finally { - if (script == null) { - breaker.addWithoutBreaking(-byteSize); - } - } + public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { + return new DateScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize); } } private final DateFieldScript script; - private final long byteSize; private int docId; DateScriptBlockDocValuesReader(CircuitBreaker breaker, DateFieldScript script, long byteSize) { - super(breaker); + super(breaker, byteSize); this.script = script; - this.byteSize = byteSize; } @Override @@ -67,24 +57,9 @@ public int docId() { } @Override - public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered) - throws IOException { - // Note that we don't sort the values sort, so we can't use factory.longsFromDocValues - try (BlockLoader.LongBuilder builder = factory.longs(docs.count() - offset)) { - for (int i = offset; i < docs.count(); i++) { - read(docs.get(i), builder); - } - return builder.build(); - } - } - - @Override - public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException { + public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException { + BlockLoader.LongBuilder builder = (BlockLoader.LongBuilder) b; this.docId = docId; - read(docId, (BlockLoader.LongBuilder) builder); - } - - private void read(int docId, BlockLoader.LongBuilder builder) { script.runForDoc(docId); switch (script.count()) { case 0 -> builder.appendNull(); @@ -103,9 +78,4 @@ private void read(int docId, BlockLoader.LongBuilder builder) { public String toString() { return "ScriptDates"; } - - @Override - public void close() { - breaker.addWithoutBreaking(-byteSize); - } } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DoubleScriptBlockDocValuesReader.java b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DoubleScriptBlockDocValuesReader.java index 63bef89f5e2cf..23e737e5a7228 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DoubleScriptBlockDocValuesReader.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DoubleScriptBlockDocValuesReader.java @@ -21,12 +21,13 @@ /** * {@link BlockDocValuesReader} implementation for {@code double} scripts. */ -public class DoubleScriptBlockDocValuesReader extends BlockDocValuesReader { - public static class DoubleScriptBlockLoader extends DocValuesBlockLoader { +public class DoubleScriptBlockDocValuesReader extends BlockScriptReader { + public static class DoubleScriptBlockLoader extends ScriptBlockLoader { private final DoubleFieldScript.LeafFactory factory; private final long byteSize; public DoubleScriptBlockLoader(DoubleFieldScript.LeafFactory factory, ByteSizeValue byteSize) { + super(byteSize); this.factory = factory; this.byteSize = byteSize.getBytes(); } @@ -37,29 +38,17 @@ public Builder builder(BlockFactory factory, int expectedCount) { } @Override - public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { - breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks"); - DoubleFieldScript script = null; - try { - script = factory.newInstance(context); - return new DoubleScriptBlockDocValuesReader(breaker, script, byteSize); - } finally { - if (script == null) { - breaker.addWithoutBreaking(-byteSize); - } - } + public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { + return new DoubleScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize); } } private final DoubleFieldScript script; - private final long byteSize; private int docId; DoubleScriptBlockDocValuesReader(CircuitBreaker breaker, DoubleFieldScript script, long byteSize) { - super(breaker); + super(breaker, byteSize); this.script = script; - this.byteSize = byteSize; - } @Override @@ -68,24 +57,9 @@ public int docId() { } @Override - public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered) - throws IOException { - // Note that we don't sort the values sort, so we can't use factory.doublesFromDocValues - try (BlockLoader.DoubleBuilder builder = factory.doubles(docs.count() - offset)) { - for (int i = offset; i < docs.count(); i++) { - read(docs.get(i), builder); - } - return builder.build(); - } - } - - @Override - public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException { + public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException { + BlockLoader.DoubleBuilder builder = (BlockLoader.DoubleBuilder) b; this.docId = docId; - read(docId, (BlockLoader.DoubleBuilder) builder); - } - - private void read(int docId, BlockLoader.DoubleBuilder builder) { script.runForDoc(docId); switch (script.count()) { case 0 -> builder.appendNull(); @@ -104,9 +78,4 @@ private void read(int docId, BlockLoader.DoubleBuilder builder) { public String toString() { return "ScriptDoubles"; } - - @Override - public void close() { - breaker.addWithoutBreaking(-byteSize); - } } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/IpScriptBlockDocValuesReader.java b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/IpScriptBlockDocValuesReader.java index b90cc94a0c707..669d5be50c136 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/IpScriptBlockDocValuesReader.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/IpScriptBlockDocValuesReader.java @@ -21,12 +21,13 @@ /** * {@link BlockDocValuesReader} implementation for keyword scripts. */ -public class IpScriptBlockDocValuesReader extends BlockDocValuesReader { - public static class IpScriptBlockLoader extends DocValuesBlockLoader { +public class IpScriptBlockDocValuesReader extends BlockScriptReader { + public static class IpScriptBlockLoader extends ScriptBlockLoader { private final IpFieldScript.LeafFactory factory; private final long byteSize; public IpScriptBlockLoader(IpFieldScript.LeafFactory factory, ByteSizeValue byteSize) { + super(byteSize); this.factory = factory; this.byteSize = byteSize.getBytes(); } @@ -37,28 +38,17 @@ public Builder builder(BlockFactory factory, int expectedCount) { } @Override - public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { - breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks"); - IpFieldScript script = null; - try { - script = factory.newInstance(context); - return new IpScriptBlockDocValuesReader(breaker, script, byteSize); - } finally { - if (script == null) { - breaker.addWithoutBreaking(-byteSize); - } - } + public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { + return new IpScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize); } } private final IpFieldScript script; - private final long byteSize; private int docId; IpScriptBlockDocValuesReader(CircuitBreaker breaker, IpFieldScript script, long byteSize) { - super(breaker); + super(breaker, byteSize); this.script = script; - this.byteSize = byteSize; } @Override @@ -67,24 +57,9 @@ public int docId() { } @Override - public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered) - throws IOException { - // Note that we don't pre-sort our output so we can't use bytesRefsFromDocValues - try (BlockLoader.BytesRefBuilder builder = factory.bytesRefs(docs.count() - offset)) { - for (int i = offset; i < docs.count(); i++) { - read(docs.get(i), builder); - } - return builder.build(); - } - } - - @Override - public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException { + public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException { + BlockLoader.BytesRefBuilder builder = (BlockLoader.BytesRefBuilder) b; this.docId = docId; - read(docId, (BlockLoader.BytesRefBuilder) builder); - } - - private void read(int docId, BlockLoader.BytesRefBuilder builder) { script.runForDoc(docId); switch (script.count()) { case 0 -> builder.appendNull(); @@ -105,9 +80,4 @@ private void read(int docId, BlockLoader.BytesRefBuilder builder) { public String toString() { return "ScriptIps"; } - - @Override - public void close() { - breaker.addWithoutBreaking(-byteSize); - } } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/KeywordScriptBlockDocValuesReader.java b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/KeywordScriptBlockDocValuesReader.java index 7758977a4c462..99e5cee27ffe9 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/KeywordScriptBlockDocValuesReader.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/KeywordScriptBlockDocValuesReader.java @@ -22,12 +22,13 @@ /** * {@link BlockDocValuesReader} implementation for keyword scripts. */ -public class KeywordScriptBlockDocValuesReader extends BlockDocValuesReader { - public static class KeywordScriptBlockLoader extends DocValuesBlockLoader { +public class KeywordScriptBlockDocValuesReader extends BlockScriptReader { + public static class KeywordScriptBlockLoader extends ScriptBlockLoader { private final StringFieldScript.LeafFactory factory; private final long byteSize; public KeywordScriptBlockLoader(StringFieldScript.LeafFactory factory, ByteSizeValue byteSize) { + super(byteSize); this.factory = factory; this.byteSize = byteSize.getBytes(); } @@ -38,29 +39,18 @@ public Builder builder(BlockFactory factory, int expectedCount) { } @Override - public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { - breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks"); - StringFieldScript script = null; - try { - script = factory.newInstance(context); - return new KeywordScriptBlockDocValuesReader(breaker, script, byteSize); - } finally { - if (script == null) { - breaker.addWithoutBreaking(-byteSize); - } - } + public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { + return new KeywordScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize); } } private final BytesRefBuilder bytesBuild = new BytesRefBuilder(); // TODO breaking builder private final StringFieldScript script; - private final long byteSize; private int docId; KeywordScriptBlockDocValuesReader(CircuitBreaker breaker, StringFieldScript script, long byteSize) { - super(breaker); + super(breaker, byteSize); this.script = script; - this.byteSize = byteSize; } @Override @@ -69,24 +59,9 @@ public int docId() { } @Override - public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered) - throws IOException { - // Note that we don't pre-sort our output so we can't use bytesRefsFromDocValues - try (BlockLoader.BytesRefBuilder builder = factory.bytesRefs(docs.count() - offset)) { - for (int i = offset; i < docs.count(); i++) { - read(docs.get(i), builder); - } - return builder.build(); - } - } - - @Override - public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException { + public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException { + BlockLoader.BytesRefBuilder builder = (BlockLoader.BytesRefBuilder) b; this.docId = docId; - read(docId, (BlockLoader.BytesRefBuilder) builder); - } - - private void read(int docId, BlockLoader.BytesRefBuilder builder) { script.runForDoc(docId); switch (script.getValues().size()) { case 0 -> builder.appendNull(); @@ -109,9 +84,4 @@ private void read(int docId, BlockLoader.BytesRefBuilder builder) { public String toString() { return "ScriptKeywords"; } - - @Override - public void close() { - breaker.addWithoutBreaking(-byteSize); - } } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/LongScriptBlockDocValuesReader.java b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/LongScriptBlockDocValuesReader.java index aaf97a85b4aab..4fce0b403319c 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/LongScriptBlockDocValuesReader.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/LongScriptBlockDocValuesReader.java @@ -21,12 +21,13 @@ /** * {@link BlockDocValuesReader} implementation for {@code long} scripts. */ -public class LongScriptBlockDocValuesReader extends BlockDocValuesReader { - public static class LongScriptBlockLoader extends DocValuesBlockLoader { +public class LongScriptBlockDocValuesReader extends BlockScriptReader { + public static class LongScriptBlockLoader extends ScriptBlockLoader { private final LongFieldScript.LeafFactory factory; private final long byteSize; public LongScriptBlockLoader(LongFieldScript.LeafFactory factory, ByteSizeValue byteSize) { + super(byteSize); this.factory = factory; this.byteSize = byteSize.getBytes(); } @@ -37,28 +38,17 @@ public Builder builder(BlockFactory factory, int expectedCount) { } @Override - public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { - breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks"); - LongFieldScript script = null; - try { - script = factory.newInstance(context); - return new LongScriptBlockDocValuesReader(breaker, script, byteSize); - } finally { - if (script == null) { - breaker.addWithoutBreaking(-byteSize); - } - } + public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException { + return new LongScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize); } } private final LongFieldScript script; - private final long byteSize; private int docId; LongScriptBlockDocValuesReader(CircuitBreaker breaker, LongFieldScript script, long byteSize) { - super(breaker); + super(breaker, byteSize); this.script = script; - this.byteSize = byteSize; } @Override @@ -67,24 +57,9 @@ public int docId() { } @Override - public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered) - throws IOException { - // Note that we don't pre-sort our output so we can't use longsFromDocValues - try (BlockLoader.LongBuilder builder = factory.longs(docs.count() - offset)) { - for (int i = offset; i < docs.count(); i++) { - read(docs.get(i), builder); - } - return builder.build(); - } - } - - @Override - public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException { + public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException { + BlockLoader.LongBuilder builder = (BlockLoader.LongBuilder) b; this.docId = docId; - read(docId, (BlockLoader.LongBuilder) builder); - } - - private void read(int docId, BlockLoader.LongBuilder builder) { script.runForDoc(docId); switch (script.count()) { case 0 -> builder.appendNull(); @@ -103,9 +78,4 @@ private void read(int docId, BlockLoader.LongBuilder builder) { public String toString() { return "ScriptLongs"; } - - @Override - public void close() { - breaker.addWithoutBreaking(-byteSize); - } } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/BooleanScriptFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/BooleanScriptFieldTypeTests.java index 0d783bd33ed7a..7fc4890b96fe6 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/BooleanScriptFieldTypeTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/BooleanScriptFieldTypeTests.java @@ -499,9 +499,8 @@ private void testBlockLoader( ); try (DirectoryReader reader = iw.getReader()) { BooleanScriptFieldType fieldType = buildWrapped("xor_param", Map.of("param", false), leafFactoryWrapper); + assertColumnAtATimeReaderNotSupported(reader, fieldType); List expected = List.of(false, true); - assertThat(blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 0), equalTo(expected)); - assertThat(blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 1), equalTo(expected.subList(1, 2))); assertThat(blockLoaderReadValuesFromRowStrideReader(breaker, reader, fieldType), equalTo(expected)); } } @@ -535,23 +534,16 @@ public void testBlockLoaderSourceOnlyRuntimeField() throws IOException { // then // assert loader is of expected instance type + assertColumnAtATimeReaderNotSupported(reader, fieldType); assertThat(loader, instanceOf(BooleanScriptBlockDocValuesReader.BooleanScriptBlockLoader.class)); CircuitBreaker breaker = newLimitedBreaker(ByteSizeValue.ofMb(1)); - // ignored source doesn't support column at a time loading: - try ( - BlockLoader.ColumnAtATimeReader columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst()) - .apply(breaker) - ) { - assertThat(columnAtATimeLoader, instanceOf(BooleanScriptBlockDocValuesReader.class)); - } try (BlockLoader.RowStrideReader rowStrideReader = loader.rowStrideReader(breaker, reader.leaves().getFirst())) { assertThat(rowStrideReader, instanceOf(BooleanScriptBlockDocValuesReader.class)); } // assert values - assertThat(blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 0), equalTo(expected)); assertThat(blockLoaderReadValuesFromRowStrideReader(breaker, reader, fieldType), equalTo(expected)); } } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/DateScriptFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/DateScriptFieldTypeTests.java index a2299eb8ee2ca..33897714b5e90 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/DateScriptFieldTypeTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/DateScriptFieldTypeTests.java @@ -543,11 +543,7 @@ private void testBlockLoader(CircuitBreaker breaker, Function expected = List.of( new BytesRef(InetAddressPoint.encode(InetAddresses.forString("192.168.0.1"))), new BytesRef(InetAddressPoint.encode(InetAddresses.forString("192.168.1.1"))) ); - assertThat(blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 0), equalTo(expected)); - assertThat(blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 1), equalTo(expected.subList(1, 2))); assertThat(blockLoaderReadValuesFromRowStrideReader(breaker, reader, fieldType), equalTo(expected)); } } @@ -358,26 +357,18 @@ public void testBlockLoaderSourceOnlyRuntimeField() throws IOException { ); try (DirectoryReader reader = iw.getReader()) { - // when BlockLoader loader = fieldType.blockLoader(blContext(Settings.EMPTY, true)); - - // then + assertColumnAtATimeReaderNotSupported(reader, fieldType); // assert loader is of expected instance type assertThat(loader, instanceOf(IpScriptBlockDocValuesReader.IpScriptBlockLoader.class)); - // ignored source doesn't support column at a time loading: CircuitBreaker breaker = newLimitedBreaker(ByteSizeValue.ofMb(1)); - try (var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst()).apply(breaker)) { - assertThat(columnAtATimeLoader, instanceOf(IpScriptBlockDocValuesReader.class)); - } - try (var rowStrideReader = loader.rowStrideReader(breaker, reader.leaves().getFirst())) { assertThat(rowStrideReader, instanceOf(IpScriptBlockDocValuesReader.class)); } // assert values - assertThat(blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 0), equalTo(expected)); assertThat(blockLoaderReadValuesFromRowStrideReader(breaker, reader, fieldType), equalTo(expected)); } } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/KeywordScriptFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/KeywordScriptFieldTypeTests.java index 10bd3f76af256..e6478868feeea 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/KeywordScriptFieldTypeTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/KeywordScriptFieldTypeTests.java @@ -459,14 +459,7 @@ private void testBlockLoader( ); try (DirectoryReader reader = iw.getReader()) { KeywordScriptFieldType fieldType = buildWrapped("append_param", Map.of("param", "-Suffix"), factoryWrapper); - assertThat( - blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 0), - equalTo(List.of(new BytesRef("1-Suffix"), new BytesRef("2-Suffix"))) - ); - assertThat( - blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 1), - equalTo(List.of(new BytesRef("2-Suffix"))) - ); + assertColumnAtATimeReaderNotSupported(reader, fieldType); assertThat( blockLoaderReadValuesFromRowStrideReader(breaker, reader, fieldType), equalTo(List.of(new BytesRef("1-Suffix"), new BytesRef("2-Suffix"))) @@ -488,15 +481,12 @@ public void testBlockLoaderSourceOnlyRuntimeField() throws IOException { ); try (DirectoryReader reader = iw.getReader()) { KeywordScriptFieldType fieldType = simpleSourceOnlyMappedFieldType(); + assertColumnAtATimeReaderNotSupported(reader, fieldType); // Assert implementations: BlockLoader loader = fieldType.blockLoader(blContext(Settings.EMPTY, true)); assertThat(loader, instanceOf(KeywordScriptBlockDocValuesReader.KeywordScriptBlockLoader.class)); CircuitBreaker breaker = newLimitedBreaker(ByteSizeValue.ofMb(1)); - // ignored source doesn't support column at a time loading: - try (var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst()).apply(breaker)) { - assertThat(columnAtATimeLoader, instanceOf(KeywordScriptBlockDocValuesReader.class)); - } try (var rowStrideReader = loader.rowStrideReader(breaker, reader.leaves().getFirst())) { assertThat(rowStrideReader, instanceOf(KeywordScriptBlockDocValuesReader.class)); } @@ -505,11 +495,6 @@ public void testBlockLoaderSourceOnlyRuntimeField() throws IOException { var dogBytes = new BytesRef("dog"); // Assert values: - assertThat( - blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 0), - equalTo(List.of(catBytes, dogBytes)) - ); - assertThat(blockLoaderReadValuesFromColumnAtATimeReader(breaker, reader, fieldType, 1), equalTo(List.of(dogBytes))); assertThat(blockLoaderReadValuesFromRowStrideReader(breaker, reader, fieldType), equalTo(List.of(catBytes, dogBytes))); } } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java index 46ada704ac211..48a66427996f5 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java @@ -344,8 +344,7 @@ private void testBlockLoader(CircuitBreaker breaker, Function blockLoaderReadValuesFromColumnAtATimeReader( - CircuitBreaker breaker, - DirectoryReader reader, - MappedFieldType fieldType, - int offset - ) throws IOException { - return blockLoaderReadValuesFromColumnAtATimeReader(breaker, Settings.EMPTY, reader, fieldType, offset); - } - - protected final List blockLoaderReadValuesFromColumnAtATimeReader( - CircuitBreaker breaker, - Settings settings, - DirectoryReader reader, - MappedFieldType fieldType, - int offset - ) throws IOException { - BlockLoader loader = fieldType.blockLoader(blContext(settings, true)); - List all = new ArrayList<>(); + protected final void assertColumnAtATimeReaderNotSupported(DirectoryReader reader, MappedFieldType fieldType) throws IOException { + BlockLoader loader = fieldType.blockLoader(blContext(Settings.EMPTY, true)); for (LeafReaderContext ctx : reader.leaves()) { - try (BlockLoader.ColumnAtATimeReader columnReader = loader.columnAtATimeReader(ctx).apply(breaker)) { - TestBlock block = (TestBlock) columnReader.read(TestBlock.factory(), TestBlock.docs(ctx), offset, false); - for (int i = 0; i < block.size(); i++) { - all.add(block.get(i)); - } - } + assertThat(loader.columnAtATimeReader(ctx), nullValue()); } - return all; } protected final List blockLoaderReadValuesFromRowStrideReader( diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/ValuesFromSingleReader.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/ValuesFromSingleReader.java index 26db1d577c7c9..43db38df2b4ea 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/ValuesFromSingleReader.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/ValuesFromSingleReader.java @@ -146,19 +146,7 @@ private void loadFromRowStrideReaders( sourceLoader = shardContext.newSourceLoader().apply(storedFieldsSpec.sourcePaths()); storedFieldsSpec = storedFieldsSpec.merge(new StoredFieldsSpec(true, false, sourceLoader.requiredStoredFields())); } - if (storedFieldsSpec.equals(StoredFieldsSpec.NO_REQUIREMENTS)) { - throw new IllegalStateException( - "found row stride readers [" + rowStrideReaders + "] without stored fields [" + storedFieldsSpec + "]" - ); - } - StoredFieldLoader storedFieldLoader; - if (useSequentialStoredFieldsReader(docs, shardContext.storedFieldsSequentialProportion())) { - storedFieldLoader = StoredFieldLoader.fromSpecSequential(storedFieldsSpec); - operator.trackStoredFields(storedFieldsSpec, true); - } else { - storedFieldLoader = StoredFieldLoader.fromSpec(storedFieldsSpec); - operator.trackStoredFields(storedFieldsSpec, false); - } + StoredFieldLoader storedFieldLoader = storedFieldLoader(storedFieldsSpec, shardContext, docs); BlockLoaderStoredFieldsFromLeafLoader storedFields = new BlockLoaderStoredFieldsFromLeafLoader( storedFieldLoader.getLoader(ctx, null), sourceLoader != null ? sourceLoader.leaf(ctx.reader(), null) : null @@ -188,6 +176,22 @@ private void loadFromRowStrideReaders( docs.setCount(p); } + private StoredFieldLoader storedFieldLoader( + StoredFieldsSpec storedFieldsSpec, + ValuesSourceReaderOperator.ShardContext shardContext, + ValuesReaderDocs docs + ) { + if (storedFieldsSpec.equals(StoredFieldsSpec.NO_REQUIREMENTS)) { + return StoredFieldLoader.empty(); + } + if (useSequentialStoredFieldsReader(docs, shardContext.storedFieldsSequentialProportion())) { + operator.trackStoredFields(storedFieldsSpec, true); + return StoredFieldLoader.fromSpecSequential(storedFieldsSpec); + } + operator.trackStoredFields(storedFieldsSpec, false); + return StoredFieldLoader.fromSpec(storedFieldsSpec); + } + /** * Is it more efficient to use a sequential stored field reader * when reading stored fields for the documents contained in {@code docIds}? diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java index 1a0f30572f979..9df1129654588 100644 --- a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java +++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java @@ -160,7 +160,8 @@ public void testTaskContents() throws Exception { ValuesSourceReaderOperatorStatus oStatus = (ValuesSourceReaderOperatorStatus) o.status(); assertMap( oStatus.readersBuilt(), - matchesMap().entry("pause_me:column_at_a_time:ScriptLongs", greaterThanOrEqualTo(1)) + matchesMap().entry("pause_me:column_at_a_time:null", greaterThanOrEqualTo(1)) + .entry("pause_me:row_stride:ScriptLongs", greaterThanOrEqualTo(1)) ); assertThat(oStatus.pagesReceived(), greaterThanOrEqualTo(1)); assertThat(oStatus.pagesEmitted(), greaterThanOrEqualTo(1));