2121
2222import org .elasticsearch .action .OriginalIndices ;
2323import org .elasticsearch .cluster .routing .PlainShardIterator ;
24- import org .elasticsearch .cluster .routing .ShardIterator ;
2524import org .elasticsearch .cluster .routing .ShardRouting ;
2625import org .elasticsearch .common .Nullable ;
2726import org .elasticsearch .common .unit .TimeValue ;
27+ import org .elasticsearch .common .util .Countable ;
28+ import org .elasticsearch .common .util .PlainIterator ;
2829import org .elasticsearch .index .shard .ShardId ;
2930import org .elasticsearch .search .SearchShardTarget ;
3031import org .elasticsearch .search .internal .ShardSearchContextId ;
3132
33+ import java .util .Comparator ;
3234import java .util .List ;
3335import java .util .Objects ;
36+ import java .util .stream .Collectors ;
3437
3538/**
3639 * Extension of {@link PlainShardIterator} used in the search api, which also holds the {@link OriginalIndices}
3740 * of the search request (useful especially with cross-cluster search, as each cluster has its own set of original indices) as well as
3841 * the cluster alias.
3942 * @see OriginalIndices
4043 */
41- public final class SearchShardIterator extends PlainShardIterator {
44+ public final class SearchShardIterator implements Comparable < SearchShardIterator >, Countable {
4245
4346 private final OriginalIndices originalIndices ;
4447 private final String clusterAlias ;
48+ private final ShardId shardId ;
4549 private boolean skip = false ;
4650
4751 private final ShardSearchContextId searchContextId ;
4852 private final TimeValue searchContextKeepAlive ;
53+ private final PlainIterator <String > targetNodesIterator ;
4954
5055 /**
5156 * Creates a {@link PlainShardIterator} instance that iterates over a subset of the given shards
5257 * this the a given <code>shardId</code>.
5358 *
54- * @param clusterAlias the alias of the cluster where the shard is located
55- * @param shardId shard id of the group
56- * @param shards shards to iterate
59+ * @param clusterAlias the alias of the cluster where the shard is located
60+ * @param shardId shard id of the group
61+ * @param shards shards to iterate
5762 * @param originalIndices the indices that the search request originally related to (before any rewriting happened)
5863 */
5964 public SearchShardIterator (@ Nullable String clusterAlias , ShardId shardId , List <ShardRouting > shards , OriginalIndices originalIndices ) {
60- this (clusterAlias , shardId , shards , originalIndices , null , null );
65+ this (clusterAlias , shardId , shards .stream ().map (ShardRouting ::currentNodeId ).collect (Collectors .toList ()),
66+ originalIndices , null , null );
6167 }
6268
6369 public SearchShardIterator (@ Nullable String clusterAlias , ShardId shardId ,
64- List <ShardRouting > shards , OriginalIndices originalIndices ,
70+ List <String > targetNodeIds , OriginalIndices originalIndices ,
6571 ShardSearchContextId searchContextId , TimeValue searchContextKeepAlive ) {
66- super (shardId , shards );
72+ this .shardId = shardId ;
73+ this .targetNodesIterator = new PlainIterator <>(targetNodeIds );
6774 this .originalIndices = originalIndices ;
6875 this .clusterAlias = clusterAlias ;
6976 this .searchContextId = searchContextId ;
@@ -86,12 +93,16 @@ public String getClusterAlias() {
8693 return clusterAlias ;
8794 }
8895
89- /**
90- * Creates a new shard target from this iterator, pointing at the node identified by the provided identifier.
91- * @see SearchShardTarget
92- */
93- SearchShardTarget newSearchShardTarget (String nodeId ) {
94- return new SearchShardTarget (nodeId , shardId (), clusterAlias , originalIndices );
96+ SearchShardTarget nextOrNull () {
97+ final String nodeId = targetNodesIterator .nextOrNull ();
98+ if (nodeId != null ) {
99+ return new SearchShardTarget (nodeId , shardId , clusterAlias , originalIndices );
100+ }
101+ return null ;
102+ }
103+
104+ int remaining () {
105+ return targetNodesIterator .remaining ();
95106 }
96107
97108 /**
@@ -105,6 +116,10 @@ TimeValue getSearchContextKeepAlive() {
105116 return searchContextKeepAlive ;
106117 }
107118
119+ List <String > getTargetNodeIds () {
120+ return targetNodesIterator .asList ();
121+ }
122+
108123 /**
109124 * Reset the iterator and mark it as skippable
110125 * @see #skip()
@@ -114,49 +129,44 @@ void resetAndSkip() {
114129 skip = true ;
115130 }
116131
132+ void reset () {
133+ targetNodesIterator .reset ();
134+ }
135+
117136 /**
118137 * Returns <code>true</code> if the search execution should skip this shard since it can not match any documents given the query.
119138 */
120139 boolean skip () {
121140 return skip ;
122141 }
123142
143+
144+ @ Override
145+ public int size () {
146+ return targetNodesIterator .size ();
147+ }
148+
149+ ShardId shardId () {
150+ return shardId ;
151+ }
152+
124153 @ Override
125154 public boolean equals (Object o ) {
126- if (this == o ) {
127- return true ;
128- }
129- if (o == null || getClass () != o .getClass ()) {
130- return false ;
131- }
132- if (super .equals (o ) == false ) {
133- return false ;
134- }
155+ if (this == o ) return true ;
156+ if (o == null || getClass () != o .getClass ()) return false ;
135157 SearchShardIterator that = (SearchShardIterator ) o ;
136- return Objects .equals (clusterAlias , that .clusterAlias );
158+ return shardId . equals ( that . shardId ) && Objects .equals (clusterAlias , that .clusterAlias );
137159 }
138160
139161 @ Override
140162 public int hashCode () {
141- return Objects .hash (super . hashCode (), clusterAlias );
163+ return Objects .hash (clusterAlias , shardId );
142164 }
143165
144166 @ Override
145- public int compareTo (ShardIterator o ) {
146- int superCompareTo = super .compareTo (o );
147- if (superCompareTo != 0 || (o instanceof SearchShardIterator == false )) {
148- return superCompareTo ;
149- }
150- SearchShardIterator searchShardIterator = (SearchShardIterator )o ;
151- if (clusterAlias == null && searchShardIterator .getClusterAlias () == null ) {
152- return 0 ;
153- }
154- if (clusterAlias == null ) {
155- return -1 ;
156- }
157- if (searchShardIterator .getClusterAlias () == null ) {
158- return 1 ;
159- }
160- return clusterAlias .compareTo (searchShardIterator .getClusterAlias ());
167+ public int compareTo (SearchShardIterator o ) {
168+ return Comparator .comparing (SearchShardIterator ::shardId )
169+ .thenComparing (SearchShardIterator ::getClusterAlias , Comparator .nullsFirst (String ::compareTo ))
170+ .compare (this , o );
161171 }
162172}
0 commit comments