Add segment sorter for data streams (#75195)

It is beneficial to sort segments within a datastream's index
by desc order of their max timestamp field, so
that the most recent (in terms of timestamp) segments
will be first.

This allows to speed up sort query on @timestamp desc field,
which is the most common type of query for datastreams,
as we are mostly concerned with the recent data.
This patch addressed this for writable indices.

Segments' sorter is different from index sorting.
An index sorter by itself is  concerned about the order of docs
within an individual segment (and not how the segments are organized),
while the segment sorter is only used during search and allows
to start docs collection with the "right" segment,
so we can terminate the collection faster.

This PR adds a property to IndexShard `isDataStreamIndex` that
shows if a shard is a part of datastream.

Author: mayya-sharipova

server/src/internalClusterTest/java/org/elasticsearch/indices/IndexingMemoryControllerIT.java

@@ -57,7 +57,7 @@ EngineConfig engineConfigWithLargerIndexingMemory(EngineConfig config) {
                 config.getQueryCachingPolicy(), config.getTranslogConfig(), config.getFlushMergesAfter(),
                 config.getExternalRefreshListener(), config.getInternalRefreshListener(), config.getIndexSort(),
                 config.getCircuitBreakerService(), config.getGlobalCheckpointSupplier(), config.retentionLeasesSupplier(),
-                config.getPrimaryTermSupplier(), config.getSnapshotCommitSupplier());
+                config.getPrimaryTermSupplier(), config.getSnapshotCommitSupplier(), config.getLeafSorter());
         }
 
         @Override

server/src/main/java/org/elasticsearch/cluster/metadata/DataStream.java

@@ -7,6 +7,9 @@
  */
 package org.elasticsearch.cluster.metadata;
 
+import org.apache.lucene.document.LongPoint;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.PointValues;
 import org.elasticsearch.cluster.AbstractDiffable;
 import org.elasticsearch.cluster.Diff;
 import org.elasticsearch.core.Nullable;
@@ -25,6 +28,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
@@ -36,6 +40,25 @@ public final class DataStream extends AbstractDiffable<DataStream> implements To
 
     public static final String BACKING_INDEX_PREFIX = ".ds-";
     public static final DateFormatter DATE_FORMATTER = DateFormatter.forPattern("uuuu.MM.dd");
+    // Datastreams' leaf readers should be sorted by desc order of their timestamp field, as it allows search time optimizations
+    public static Comparator<LeafReader> DATASTREAM_LEAF_READERS_SORTER =
+        Comparator.comparingLong(
+            (LeafReader r) -> {
+                try {
+                    PointValues points = r.getPointValues(DataStream.TimestampField.FIXED_TIMESTAMP_FIELD);
+                    if (points != null) {
+                        byte[] sortValue = points.getMaxPackedValue();
+                        return LongPoint.decodeDimension(sortValue, 0);
+                    } else if (r.numDocs() == 0) {
+                        // points can be null if the segment contains only deleted documents
+                        return Long.MIN_VALUE;
+                    }
+                } catch (IOException e) {
+                }
+                throw new IllegalStateException("Can't access [" +
+                    DataStream.TimestampField.FIXED_TIMESTAMP_FIELD + "] field for the data stream!");
+            })
+        .reversed();
 
     private final LongSupplier timeProvider;
     private final String name;

server/src/main/java/org/elasticsearch/index/engine/EngineConfig.java

@@ -9,6 +9,7 @@
 
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.search.QueryCache;
 import org.apache.lucene.search.QueryCachingPolicy;
@@ -32,6 +33,7 @@
 import org.elasticsearch.plugins.IndexStorePlugin;
 import org.elasticsearch.threadpool.ThreadPool;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.LongSupplier;
@@ -70,6 +72,7 @@ public final class EngineConfig {
     private final CircuitBreakerService circuitBreakerService;
     private final LongSupplier globalCheckpointSupplier;
     private final Supplier<RetentionLeases> retentionLeasesSupplier;
+    private final Comparator<LeafReader> leafSorter;
 
     /**
      * A supplier of the outstanding retention leases. This is used during merged operations to determine which operations that have been
@@ -131,7 +134,8 @@ public EngineConfig(
             LongSupplier globalCheckpointSupplier,
             Supplier<RetentionLeases> retentionLeasesSupplier,
             LongSupplier primaryTermSupplier,
-            IndexStorePlugin.SnapshotCommitSupplier snapshotCommitSupplier) {
+            IndexStorePlugin.SnapshotCommitSupplier snapshotCommitSupplier,
+            Comparator<LeafReader> leafSorter) {
         this.shardId = shardId;
         this.indexSettings = indexSettings;
         this.threadPool = threadPool;
@@ -169,6 +173,7 @@ public EngineConfig(
         this.retentionLeasesSupplier = Objects.requireNonNull(retentionLeasesSupplier);
         this.primaryTermSupplier = primaryTermSupplier;
         this.snapshotCommitSupplier = snapshotCommitSupplier;
+        this.leafSorter = leafSorter;
     }
 
     /**
@@ -353,4 +358,12 @@ public LongSupplier getPrimaryTermSupplier() {
     public IndexStorePlugin.SnapshotCommitSupplier getSnapshotCommitSupplier() {
         return snapshotCommitSupplier;
     }
+
+    /**
+     * Returns how segments should be sorted for reading or @null if no sorting should be applied.
+     */
+    @Nullable
+    public Comparator<LeafReader> getLeafSorter() {
+        return leafSorter;
+    }
 }

server/src/main/java/org/elasticsearch/index/engine/InternalEngine.java

@@ -2193,6 +2193,11 @@ private IndexWriterConfig getIndexWriterConfig() {
         if (config().getIndexSort() != null) {
             iwc.setIndexSort(config().getIndexSort());
         }
+        // Provide a custom leaf sorter, so that index readers opened from this writer
+        // will have its leaves sorted according the given leaf sorter.
+        if (engineConfig.getLeafSorter() != null) {
+            iwc.setLeafSorter(engineConfig.getLeafSorter());
+        }
         return iwc;
     }
 

server/src/main/java/org/elasticsearch/index/engine/ReadOnlyEngine.java

@@ -199,6 +199,8 @@ protected final ElasticsearchDirectoryReader wrapReader(DirectoryReader reader,
     }
 
     protected DirectoryReader open(IndexCommit commit) throws IOException {
+        // TODO: provide engineConfig.getLeafSorter() when opening a DirectoryReader from a commit
+        // should be available from Lucene v 8.10
         assert Transports.assertNotTransportThread("opening index commit of a read-only engine");
         if (lazilyLoadSoftDeletes) {
             return new LazySoftDeletesDirectoryReaderWrapper(DirectoryReader.open(commit), Lucene.SOFT_DELETES_FIELD);

server/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.index.shard;
 
 import com.carrotsearch.hppc.ObjectLongMap;
+
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.apache.lucene.analysis.Analyzer;
@@ -185,6 +186,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 
+import static org.elasticsearch.cluster.metadata.DataStream.DATASTREAM_LEAF_READERS_SORTER;
 import static org.elasticsearch.index.seqno.RetentionLeaseActions.RETAIN_ALL;
 import static org.elasticsearch.index.seqno.SequenceNumbers.UNASSIGNED_SEQ_NO;
 
@@ -283,6 +285,7 @@ Runnable getGlobalCheckpointSyncer() {
     private final AtomicReference<Translog.Location> pendingRefreshLocation = new AtomicReference<>();
     private final RefreshPendingLocationListener refreshPendingLocationListener;
     private volatile boolean useRetentionLeasesInPeerRecovery;
+    private final boolean isDataStreamIndex; // if a shard is a part of data stream
 
     public IndexShard(
             final ShardRouting shardRouting,
@@ -387,6 +390,7 @@ public boolean shouldCache(Query query) {
         persistMetadata(path, indexSettings, shardRouting, null, logger);
         this.useRetentionLeasesInPeerRecovery = replicationTracker.hasAllPeerRecoveryRetentionLeases();
         this.refreshPendingLocationListener = new RefreshPendingLocationListener();
+        this.isDataStreamIndex = mapperService == null ? false : mapperService.mappingLookup().isDataStreamTimestampFieldEnabled();
     }
 
     public ThreadPool getThreadPool() {
@@ -2912,7 +2916,8 @@ private EngineConfig newEngineConfig(LongSupplier globalCheckpointSupplier) {
                 globalCheckpointSupplier,
                 replicationTracker::getRetentionLeases,
                 this::getOperationPrimaryTerm,
-                snapshotCommitSupplier);
+                snapshotCommitSupplier,
+                isDataStreamIndex ? DATASTREAM_LEAF_READERS_SORTER : null);
     }
 
     /**

server/src/test/java/org/elasticsearch/index/engine/InternalEngineTests.java

@@ -2948,7 +2948,8 @@ public void testRecoverFromForeignTranslog() throws IOException {
                 () -> UNASSIGNED_SEQ_NO,
                 () -> RetentionLeases.EMPTY,
                 primaryTerm::get,
-                IndexModule.DEFAULT_SNAPSHOT_COMMIT_SUPPLIER);
+                IndexModule.DEFAULT_SNAPSHOT_COMMIT_SUPPLIER,
+                null);
         expectThrows(EngineCreationFailureException.class, () -> new InternalEngine(brokenConfig));
 
         engine = createEngine(store, primaryTranslogDir); // and recover again!
@@ -6021,7 +6022,7 @@ public void testNotWarmUpSearcherInEngineCtor() throws Exception {
                 config.getQueryCachingPolicy(), translogConfig, config.getFlushMergesAfter(),
                 config.getExternalRefreshListener(), config.getInternalRefreshListener(), config.getIndexSort(),
                 config.getCircuitBreakerService(), config.getGlobalCheckpointSupplier(), config.retentionLeasesSupplier(),
-                config.getPrimaryTermSupplier(), config.getSnapshotCommitSupplier());
+                config.getPrimaryTermSupplier(), config.getSnapshotCommitSupplier(), config.getLeafSorter());
             try (InternalEngine engine = createEngine(configWithWarmer)) {
                 assertThat(warmedUpReaders, empty());
                 assertThat(expectThrows(Throwable.class, () -> engine.acquireSearcher("test")).getMessage(),

server/src/test/java/org/elasticsearch/index/shard/IndexShardTests.java

@@ -3928,7 +3928,7 @@ public void testCloseShardWhileEngineIsWarming() throws Exception {
                 config.getQueryCachingPolicy(), config.getTranslogConfig(), config.getFlushMergesAfter(),
                 config.getExternalRefreshListener(), config.getInternalRefreshListener(), config.getIndexSort(),
                 config.getCircuitBreakerService(), config.getGlobalCheckpointSupplier(), config.retentionLeasesSupplier(),
-                config.getPrimaryTermSupplier(), IndexModule.DEFAULT_SNAPSHOT_COMMIT_SUPPLIER);
+                config.getPrimaryTermSupplier(), IndexModule.DEFAULT_SNAPSHOT_COMMIT_SUPPLIER, config.getLeafSorter());
             return new InternalEngine(configWithWarmer);
         });
         Thread recoveryThread = new Thread(() -> expectThrows(AlreadyClosedException.class, () -> recoverShardFromStore(shard)));

server/src/test/java/org/elasticsearch/index/shard/RefreshListenersTests.java

@@ -139,7 +139,8 @@ public void onFailedEngine(String reason, @Nullable Exception e) {
                 () -> SequenceNumbers.NO_OPS_PERFORMED,
                 () -> RetentionLeases.EMPTY,
                 () -> primaryTerm,
-                IndexModule.DEFAULT_SNAPSHOT_COMMIT_SUPPLIER);
+                IndexModule.DEFAULT_SNAPSHOT_COMMIT_SUPPLIER,
+                null);
         engine = new InternalEngine(config);
         engine.recoverFromTranslog((e, s) -> 0, Long.MAX_VALUE);
         listeners.setCurrentRefreshLocationSupplier(engine::getTranslogLastWriteLocation);

server/src/test/java/org/elasticsearch/indices/IndexingMemoryControllerTests.java

@@ -371,7 +371,7 @@ EngineConfig configWithRefreshListener(EngineConfig config, ReferenceManager.Ref
             config.getQueryCachingPolicy(), config.getTranslogConfig(), config.getFlushMergesAfter(),
             config.getExternalRefreshListener(), internalRefreshListener, config.getIndexSort(),
             config.getCircuitBreakerService(), config.getGlobalCheckpointSupplier(), config.retentionLeasesSupplier(),
-            config.getPrimaryTermSupplier(), config.getSnapshotCommitSupplier());
+            config.getPrimaryTermSupplier(), config.getSnapshotCommitSupplier(), config.getLeafSorter());
     }
 
     ThreadPoolStats.Stats getRefreshThreadPoolStats() {