|
@@ -9,6 +9,7 @@ package org.elasticsearch.compute;
|
|
|
|
|
|
import org.apache.lucene.document.Document;
|
|
|
import org.apache.lucene.document.Field;
|
|
|
+import org.apache.lucene.document.LongField;
|
|
|
import org.apache.lucene.document.LongPoint;
|
|
|
import org.apache.lucene.document.SortedSetDocValuesField;
|
|
|
import org.apache.lucene.index.DirectoryReader;
|
|
@@ -43,9 +44,11 @@ import org.elasticsearch.compute.data.DocVector;
|
|
|
import org.elasticsearch.compute.data.ElementType;
|
|
|
import org.elasticsearch.compute.data.IntVector;
|
|
|
import org.elasticsearch.compute.data.LongBlock;
|
|
|
+import org.elasticsearch.compute.data.LongVector;
|
|
|
import org.elasticsearch.compute.data.Page;
|
|
|
import org.elasticsearch.compute.lucene.DataPartitioning;
|
|
|
import org.elasticsearch.compute.lucene.LuceneOperator;
|
|
|
+import org.elasticsearch.compute.lucene.LuceneSliceQueue;
|
|
|
import org.elasticsearch.compute.lucene.LuceneSourceOperator;
|
|
|
import org.elasticsearch.compute.lucene.LuceneSourceOperatorTests;
|
|
|
import org.elasticsearch.compute.lucene.ShardContext;
|
|
@@ -54,7 +57,6 @@ import org.elasticsearch.compute.operator.AbstractPageMappingOperator;
|
|
|
import org.elasticsearch.compute.operator.Driver;
|
|
|
import org.elasticsearch.compute.operator.DriverContext;
|
|
|
import org.elasticsearch.compute.operator.HashAggregationOperator;
|
|
|
-import org.elasticsearch.compute.operator.LimitOperator;
|
|
|
import org.elasticsearch.compute.operator.Operator;
|
|
|
import org.elasticsearch.compute.operator.OrdinalsGroupingOperator;
|
|
|
import org.elasticsearch.compute.operator.PageConsumerOperator;
|
|
@@ -64,9 +66,11 @@ import org.elasticsearch.compute.test.BlockTestUtils;
|
|
|
import org.elasticsearch.compute.test.OperatorTestCase;
|
|
|
import org.elasticsearch.compute.test.SequenceLongBlockSourceOperator;
|
|
|
import org.elasticsearch.compute.test.TestDriverFactory;
|
|
|
+import org.elasticsearch.compute.test.TestResultPageSinkOperator;
|
|
|
import org.elasticsearch.core.CheckedConsumer;
|
|
|
import org.elasticsearch.core.Releasables;
|
|
|
import org.elasticsearch.index.IndexSettings;
|
|
|
+import org.elasticsearch.index.mapper.BlockDocValuesReader;
|
|
|
import org.elasticsearch.index.mapper.FieldNamesFieldMapper;
|
|
|
import org.elasticsearch.index.mapper.KeywordFieldMapper;
|
|
|
import org.elasticsearch.index.mapper.MappedFieldType;
|
|
@@ -75,7 +79,6 @@ import org.elasticsearch.index.mapper.SourceLoader;
|
|
|
import org.elasticsearch.index.mapper.Uid;
|
|
|
import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
|
|
|
import org.elasticsearch.search.lookup.SearchLookup;
|
|
|
-import org.elasticsearch.test.ESTestCase;
|
|
|
|
|
|
import java.io.IOException;
|
|
|
import java.util.ArrayList;
|
|
@@ -90,9 +93,9 @@ import java.util.TreeMap;
|
|
|
import static org.elasticsearch.compute.aggregation.AggregatorMode.FINAL;
|
|
|
import static org.elasticsearch.compute.aggregation.AggregatorMode.INITIAL;
|
|
|
import static org.elasticsearch.compute.test.OperatorTestCase.randomPageSize;
|
|
|
-import static org.hamcrest.Matchers.contains;
|
|
|
import static org.hamcrest.Matchers.empty;
|
|
|
import static org.hamcrest.Matchers.equalTo;
|
|
|
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
|
|
|
|
|
|
/**
|
|
|
* This venerable test builds {@link Driver}s by hand and runs them together, simulating
|
|
@@ -112,7 +115,11 @@ public class OperatorTests extends MapperServiceTestCase {
|
|
|
final long from = randomBoolean() ? Long.MIN_VALUE : randomLongBetween(0, 10000);
|
|
|
final long to = randomBoolean() ? Long.MAX_VALUE : randomLongBetween(from, from + 10000);
|
|
|
final Query query = LongPoint.newRangeQuery("pt", from, to);
|
|
|
- LuceneOperator.Factory factory = luceneOperatorFactory(reader, query, LuceneOperator.NO_LIMIT);
|
|
|
+ LuceneOperator.Factory factory = luceneOperatorFactory(
|
|
|
+ reader,
|
|
|
+ List.of(new LuceneSliceQueue.QueryAndTags(query, List.of())),
|
|
|
+ LuceneOperator.NO_LIMIT
|
|
|
+ );
|
|
|
List<Driver> drivers = new ArrayList<>();
|
|
|
try {
|
|
|
Set<Integer> actualDocIds = ConcurrentCollections.newConcurrentSet();
|
|
@@ -221,7 +228,11 @@ public class OperatorTests extends MapperServiceTestCase {
|
|
|
);
|
|
|
Driver driver = TestDriverFactory.create(
|
|
|
driverContext,
|
|
|
- luceneOperatorFactory(reader, new MatchAllDocsQuery(), LuceneOperator.NO_LIMIT).get(driverContext),
|
|
|
+ luceneOperatorFactory(
|
|
|
+ reader,
|
|
|
+ List.of(new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of())),
|
|
|
+ LuceneOperator.NO_LIMIT
|
|
|
+ ).get(driverContext),
|
|
|
operators,
|
|
|
new PageConsumerOperator(page -> {
|
|
|
BytesRefBlock keys = page.getBlock(0);
|
|
@@ -243,31 +254,109 @@ public class OperatorTests extends MapperServiceTestCase {
|
|
|
assertThat(blockFactory.breaker().getUsed(), equalTo(0L));
|
|
|
}
|
|
|
|
|
|
- public void testLimitOperator() {
|
|
|
- var positions = 100;
|
|
|
- var limit = randomIntBetween(90, 101);
|
|
|
- var values = randomList(positions, positions, ESTestCase::randomLong);
|
|
|
+ public void testPushRoundToToQuery() throws IOException {
|
|
|
+ long firstGroupMax = randomLong();
|
|
|
+ long secondGroupMax = randomLong();
|
|
|
+ long thirdGroupMax = randomLong();
|
|
|
|
|
|
- var results = new ArrayList<Long>();
|
|
|
- DriverContext driverContext = driverContext();
|
|
|
- try (
|
|
|
- var driver = TestDriverFactory.create(
|
|
|
- driverContext,
|
|
|
- new SequenceLongBlockSourceOperator(driverContext.blockFactory(), values, 100),
|
|
|
- List.of((new LimitOperator.Factory(limit)).get(driverContext)),
|
|
|
- new PageConsumerOperator(page -> {
|
|
|
- LongBlock block = page.getBlock(0);
|
|
|
- for (int i = 0; i < page.getPositionCount(); i++) {
|
|
|
- results.add(block.getLong(i));
|
|
|
+ CheckedConsumer<DirectoryReader, IOException> verifier = reader -> {
|
|
|
+ Query firstGroupQuery = LongPoint.newRangeQuery("g", Long.MIN_VALUE, 99);
|
|
|
+ Query secondGroupQuery = LongPoint.newRangeQuery("g", 100, 9999);
|
|
|
+ Query thirdGroupQuery = LongPoint.newRangeQuery("g", 10000, Long.MAX_VALUE);
|
|
|
+
|
|
|
+ LuceneSliceQueue.QueryAndTags firstGroupQueryAndTags = new LuceneSliceQueue.QueryAndTags(firstGroupQuery, List.of(0L));
|
|
|
+ LuceneSliceQueue.QueryAndTags secondGroupQueryAndTags = new LuceneSliceQueue.QueryAndTags(secondGroupQuery, List.of(100L));
|
|
|
+ LuceneSliceQueue.QueryAndTags thirdGroupQueryAndTags = new LuceneSliceQueue.QueryAndTags(thirdGroupQuery, List.of(10000L));
|
|
|
+
|
|
|
+ LuceneOperator.Factory factory = luceneOperatorFactory(
|
|
|
+ reader,
|
|
|
+ List.of(firstGroupQueryAndTags, secondGroupQueryAndTags, thirdGroupQueryAndTags),
|
|
|
+ LuceneOperator.NO_LIMIT
|
|
|
+ );
|
|
|
+ ValuesSourceReaderOperator.Factory load = new ValuesSourceReaderOperator.Factory(
|
|
|
+ List.of(
|
|
|
+ new ValuesSourceReaderOperator.FieldInfo("v", ElementType.LONG, f -> new BlockDocValuesReader.LongsBlockLoader("v"))
|
|
|
+ ),
|
|
|
+ List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> {
|
|
|
+ throw new UnsupportedOperationException();
|
|
|
+ }, 0.8)),
|
|
|
+ 0
|
|
|
+ );
|
|
|
+ List<Page> pages = new ArrayList<>();
|
|
|
+ DriverContext driverContext = driverContext();
|
|
|
+ try (
|
|
|
+ Driver driver = TestDriverFactory.create(
|
|
|
+ driverContext,
|
|
|
+ factory.get(driverContext),
|
|
|
+ List.of(load.get(driverContext)),
|
|
|
+ new TestResultPageSinkOperator(pages::add)
|
|
|
+ )
|
|
|
+ ) {
|
|
|
+ OperatorTestCase.runDriver(driver);
|
|
|
+ }
|
|
|
+ assertDriverContext(driverContext);
|
|
|
+
|
|
|
+ boolean sawFirstMax = false;
|
|
|
+ boolean sawSecondMax = false;
|
|
|
+ boolean sawThirdMax = false;
|
|
|
+ for (Page page : pages) {
|
|
|
+ logger.error("ADFA {}", page);
|
|
|
+ LongVector group = page.<LongBlock>getBlock(1).asVector();
|
|
|
+ LongVector value = page.<LongBlock>getBlock(2).asVector();
|
|
|
+ for (int p = 0; p < page.getPositionCount(); p++) {
|
|
|
+ long g = group.getLong(p);
|
|
|
+ long v = value.getLong(p);
|
|
|
+ switch ((int) g) {
|
|
|
+ case 0 -> {
|
|
|
+ assertThat(v, lessThanOrEqualTo(firstGroupMax));
|
|
|
+ sawFirstMax |= v == firstGroupMax;
|
|
|
+ }
|
|
|
+ case 100 -> {
|
|
|
+ assertThat(v, lessThanOrEqualTo(secondGroupMax));
|
|
|
+ sawSecondMax |= v == secondGroupMax;
|
|
|
+ }
|
|
|
+ case 10000 -> {
|
|
|
+ assertThat(v, lessThanOrEqualTo(thirdGroupMax));
|
|
|
+ sawThirdMax |= v == thirdGroupMax;
|
|
|
+ }
|
|
|
+ default -> throw new IllegalArgumentException("Unknown group [" + g + "]");
|
|
|
}
|
|
|
- })
|
|
|
- )
|
|
|
- ) {
|
|
|
- OperatorTestCase.runDriver(driver);
|
|
|
- }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ assertTrue(sawFirstMax);
|
|
|
+ assertTrue(sawSecondMax);
|
|
|
+ assertTrue(sawThirdMax);
|
|
|
+ };
|
|
|
|
|
|
- assertThat(results, contains(values.stream().limit(limit).toArray()));
|
|
|
- assertDriverContext(driverContext);
|
|
|
+ try (Directory dir = newDirectory(); RandomIndexWriter w = new RandomIndexWriter(random(), dir)) {
|
|
|
+ int numDocs = randomIntBetween(0, 10_000);
|
|
|
+ for (int i = 0; i < numDocs; i++) {
|
|
|
+ long g, v;
|
|
|
+ switch (between(0, 2)) {
|
|
|
+ case 0 -> {
|
|
|
+ g = randomLongBetween(Long.MIN_VALUE, 99);
|
|
|
+ v = randomLongBetween(Long.MIN_VALUE, firstGroupMax);
|
|
|
+ }
|
|
|
+ case 1 -> {
|
|
|
+ g = randomLongBetween(100, 9999);
|
|
|
+ v = randomLongBetween(Long.MIN_VALUE, secondGroupMax);
|
|
|
+ }
|
|
|
+ case 2 -> {
|
|
|
+ g = randomLongBetween(10000, Long.MAX_VALUE);
|
|
|
+ v = randomLongBetween(Long.MIN_VALUE, thirdGroupMax);
|
|
|
+ }
|
|
|
+ default -> throw new IllegalArgumentException();
|
|
|
+ }
|
|
|
+ w.addDocument(List.of(new LongField("g", g, Field.Store.NO), new LongField("v", v, Field.Store.NO)));
|
|
|
+ }
|
|
|
+ w.addDocument(List.of(new LongField("g", 0, Field.Store.NO), new LongField("v", firstGroupMax, Field.Store.NO)));
|
|
|
+ w.addDocument(List.of(new LongField("g", 200, Field.Store.NO), new LongField("v", secondGroupMax, Field.Store.NO)));
|
|
|
+ w.addDocument(List.of(new LongField("g", 20000, Field.Store.NO), new LongField("v", thirdGroupMax, Field.Store.NO)));
|
|
|
+
|
|
|
+ try (DirectoryReader reader = w.getReader()) {
|
|
|
+ verifier.accept(reader);
|
|
|
+ }
|
|
|
+ }
|
|
|
}
|
|
|
|
|
|
private static Set<Integer> searchForDocIds(IndexReader reader, Query query) throws IOException {
|
|
@@ -388,11 +477,11 @@ public class OperatorTests extends MapperServiceTestCase {
|
|
|
assertThat(driverContext.getSnapshot().releasables(), empty());
|
|
|
}
|
|
|
|
|
|
- static LuceneOperator.Factory luceneOperatorFactory(IndexReader reader, Query query, int limit) {
|
|
|
+ static LuceneOperator.Factory luceneOperatorFactory(IndexReader reader, List<LuceneSliceQueue.QueryAndTags> queryAndTags, int limit) {
|
|
|
final ShardContext searchContext = new LuceneSourceOperatorTests.MockShardContext(reader, 0);
|
|
|
return new LuceneSourceOperator.Factory(
|
|
|
List.of(searchContext),
|
|
|
- ctx -> query,
|
|
|
+ ctx -> queryAndTags,
|
|
|
randomFrom(DataPartitioning.values()),
|
|
|
randomIntBetween(1, 10),
|
|
|
randomPageSize(),
|