Java Code Examples for org.elasticsearch.cluster.routing.ShardRouting#currentNodeId()

private void processShardRouting(Map<String, Map<String, List<Integer>>> routing, ShardRouting shardRouting, ShardId shardId) {
    String node;
    int id;
    String index = shardId.getIndex();

    if (shardRouting == null) {
        node = service.localNode().id();
        id = UnassignedShard.markUnassigned(shardId.id());
    } else {
        node = shardRouting.currentNodeId();
        id = shardRouting.id();
    }
    Map<String, List<Integer>> nodeMap = routing.get(node);
    if (nodeMap == null) {
        nodeMap = new TreeMap<>();
        routing.put(node, nodeMap);
    }

    List<Integer> shards = nodeMap.get(index);
    if (shards == null) {
        shards = new ArrayList<>();
        nodeMap.put(index, shards);
    }
    shards.add(id);
}
 

private void processShardRouting(Map<String, Map<String, List<Integer>>> locations, ShardRouting shardRouting, ShardId shardId) {
    String node;
    if (shardRouting == null) {
        throw new NoShardAvailableActionException(shardId);
    }
    node = shardRouting.currentNodeId();
    Map<String, List<Integer>> nodeMap = locations.get(node);
    if (nodeMap == null) {
        nodeMap = new TreeMap<>();
        locations.put(shardRouting.currentNodeId(), nodeMap);
    }

    List<Integer> shards = nodeMap.get(shardRouting.getIndex());
    if (shards == null) {
        shards = new ArrayList<>();
        nodeMap.put(shardRouting.getIndex(), shards);
    }
    shards.add(shardRouting.id());
}
 

private int findReplicaLevel(ShardIterator shardIt) {
    int replicaLevel = 0;
    shardIt.reset();
    ShardRouting shard;
    while ((shard = shardIt.nextOrNull()) != null) {
        if (shard.unassigned()) {
            continue;
        }
        boolean doOnlyOnRelocating = false;
        if (shard.primary()) {
            if (shard.relocating()) {
                doOnlyOnRelocating = true;
            } else {
                continue;
            }
        }
        String nodeId = !doOnlyOnRelocating ? shard.currentNodeId() : shard.relocating() ? shard.relocatingNodeId() : null;
        if (nodeId == null) {
            continue;
        }
        replicaLevel++;
    }
    return replicaLevel;
}
 

@Override
public void performOn(
        final ShardRouting replica,
        final ReplicaRequest request,
        final long globalCheckpoint,
        final long maxSeqNoOfUpdatesOrDeletes,
        final ActionListener<ReplicationOperation.ReplicaResponse> listener) {
    String nodeId = replica.currentNodeId();
    final DiscoveryNode node = clusterService.state().nodes().get(nodeId);
    if (node == null) {
        listener.onFailure(new NoNodeAvailableException("unknown node [" + nodeId + "]"));
        return;
    }
    final ConcreteReplicaRequest<ReplicaRequest> replicaRequest = new ConcreteReplicaRequest<>(
        request, replica.allocationId().getId(), primaryTerm, globalCheckpoint, maxSeqNoOfUpdatesOrDeletes);
    sendReplicaRequest(replicaRequest, node, listener);
}
 

private static ShardLocation getShardLocation(String indexName,
                                              String id,
                                              @Nullable String routing,
                                              ClusterService clusterService) {
    ShardIterator shardIterator = clusterService.operationRouting().indexShards(
        clusterService.state(),
        indexName,
        id,
        routing);

    final String nodeId;
    ShardRouting shardRouting = shardIterator.nextOrNull();
    if (shardRouting == null) {
        nodeId = null;
    } else if (shardRouting.active() == false) {
        nodeId = shardRouting.relocatingNodeId();
    } else {
        nodeId = shardRouting.currentNodeId();
    }
    return new ShardLocation(shardIterator.shardId(), nodeId);
}
 

private static ClusterAllocationExplanation explainShard(ShardRouting shardRouting,
                                                         RoutingAllocation allocation,
                                                         GatewayAllocator gatewayAllocator,
                                                         ShardsAllocator shardAllocator) {
    allocation.setDebugMode(RoutingAllocation.DebugMode.EXCLUDE_YES_DECISIONS);

    ShardAllocationDecision shardDecision;
    if (shardRouting.initializing() || shardRouting.relocating()) {
        shardDecision = ShardAllocationDecision.NOT_TAKEN;
    } else {
        AllocateUnassignedDecision allocateDecision = shardRouting.unassigned() ?
            gatewayAllocator.decideUnassignedShardAllocation(shardRouting, allocation) : AllocateUnassignedDecision.NOT_TAKEN;
        if (allocateDecision.isDecisionTaken() == false) {
            shardDecision = shardAllocator.decideShardAllocation(shardRouting, allocation);
        } else {
            shardDecision = new ShardAllocationDecision(allocateDecision, MoveDecision.NOT_TAKEN);
        }
    }
    return new ClusterAllocationExplanation(
        shardRouting,
        shardRouting.currentNodeId() != null ? allocation.nodes().get(shardRouting.currentNodeId()) : null,
        shardRouting.relocatingNodeId() != null ? allocation.nodes().get(shardRouting.relocatingNodeId()) : null,
        null,
        shardDecision
    );
}
 

/**
 * Asserts that all shards are allocated on nodes matching the given node pattern.
 */
public Set<String> assertAllShardsOnNodes(String index, String... pattern) {
    Set<String> nodes = new HashSet<>();
    ClusterState clusterState = client().admin().cluster().prepareState().execute().actionGet().getState();
    for (IndexRoutingTable indexRoutingTable : clusterState.routingTable()) {
        for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable) {
            for (ShardRouting shardRouting : indexShardRoutingTable) {
                if (shardRouting.currentNodeId() != null && index.equals(shardRouting.getIndexName())) {
                    String name = clusterState.nodes().get(shardRouting.currentNodeId()).getName();
                    nodes.add(name);
                    assertThat("Allocated on new node: " + name, Regex.simpleMatch(pattern, name), is(true));
                }
            }
        }
    }
    return nodes;
}
 

private void processShardRouting(Map<String, Map<String, List<Integer>>> locations, ShardRouting shardRouting) {
    String node = shardRouting.currentNodeId();
    Map<String, List<Integer>> nodeMap = locations.get(node);
    if (nodeMap == null) {
        nodeMap = new TreeMap<>();
        locations.put(shardRouting.currentNodeId(), nodeMap);
    }

    List<Integer> shards = nodeMap.get(shardRouting.getIndex());
    if (shards == null) {
        shards = new ArrayList<>();
        nodeMap.put(shardRouting.getIndex(), shards);
    }
    shards.add(shardRouting.id());
}
 

private CompletableFuture<List<Row>> retrieveRows(ShardRouting activePrimaryRouting,
                                                  RoutedCollectPhase collectPhase,
                                                  CollectTask collectTask,
                                                  ShardCollectorProviderFactory shardCollectorProviderFactory) {
    Collector<Row, ?, List<Object[]>> listCollector = Collectors.mapping(Row::materialize, Collectors.toList());
    CollectingRowConsumer<?, List<Object[]>> consumer = new CollectingRowConsumer<>(listCollector);
    String nodeId = activePrimaryRouting.currentNodeId();
    String localNodeId = clusterService.localNode().getId();
    if (localNodeId.equalsIgnoreCase(nodeId)) {
        var indexShard = indicesService.indexServiceSafe(activePrimaryRouting.index())
            .getShard(activePrimaryRouting.shardId().id());
        var collectorProvider = shardCollectorProviderFactory.create(indexShard);
        BatchIterator<Row> it;
        try {
            it = collectorProvider.getIterator(collectPhase, consumer.requiresScroll(), collectTask);
        } catch (Exception e) {
            return Exceptions.rethrowRuntimeException(e);
        }
        consumer.accept(it, null);
    } else {
        UUID childJobId = UUID.randomUUID();
        RemoteCollector remoteCollector = new RemoteCollector(
            childJobId,
            collectTask.txnCtx().sessionSettings(),
            localNodeId,
            nodeId,
            transportActionProvider.transportJobInitAction(),
            transportActionProvider.transportKillJobsNodeAction(),
            searchTp,
            tasksService,
            collectTask.getRamAccounting(),
            consumer,
            createRemoteCollectPhase(childJobId, collectPhase, activePrimaryRouting.shardId(), nodeId)
        );
        remoteCollector.doCollect();
    }
    return consumer
        .completionFuture()
        .thenApply(rows -> LazyMapList.of(rows, Buckets.arrayToSharedRow()));
}
 

@Nullable
private ShardLocation getShardLocation(String indexName, String id, @Nullable String routing) {
    try {
        ShardIterator shardIterator = clusterService.operationRouting().indexShards(
            clusterService.state(),
            indexName,
            id,
            routing
        );

        final String nodeId;
        ShardRouting shardRouting = shardIterator.nextOrNull();
        if (shardRouting == null) {
            nodeId = null;
        } else if (shardRouting.active() == false) {
            nodeId = shardRouting.relocatingNodeId();
        } else {
            nodeId = shardRouting.currentNodeId();
        }
        if (nodeId == null && LOGGER.isDebugEnabled()) {
            LOGGER.debug("Unable to get the node id for index {} and shard {}", indexName, id);
        }
        return new ShardLocation(shardIterator.shardId(), nodeId);
    } catch (IndexNotFoundException e) {
        if (!autoCreateIndices) {
            throw e;
        }
        return null;
    }
}
 

private Table buildTable(RestRequest request, final ClusterStateResponse state, final NodesStatsResponse stats) {
    final ObjectIntScatterMap<String> allocs = new ObjectIntScatterMap<>();

    for (ShardRouting shard : state.getState().routingTable().allShards()) {
        String nodeId = "UNASSIGNED";

        if (shard.assignedToNode()) {
            nodeId = shard.currentNodeId();
        }

        allocs.addTo(nodeId, 1);
    }

    Table table = getTableWithHeader(request);

    for (NodeStats nodeStats : stats.getNodes()) {
        DiscoveryNode node = nodeStats.getNode();

        int shardCount = allocs.getOrDefault(node.id(), 0);

        ByteSizeValue total = nodeStats.getFs().getTotal().getTotal();
        ByteSizeValue avail = nodeStats.getFs().getTotal().getAvailable();
        //if we don't know how much we use (non data nodes), it means 0
        long used = 0;
        short diskPercent = -1;
        if (total.bytes() > 0) {
            used = total.bytes() - avail.bytes();
            if (used >= 0 && avail.bytes() >= 0) {
                diskPercent = (short) (used * 100 / (used + avail.bytes()));
            }
        }

        table.startRow();
        table.addCell(shardCount);
        table.addCell(nodeStats.getIndices().getStore().getSize());
        table.addCell(used < 0 ? null : new ByteSizeValue(used));
        table.addCell(avail.bytes() < 0 ? null : avail);
        table.addCell(total.bytes() < 0 ? null : total);
        table.addCell(diskPercent < 0 ? null : diskPercent);
        table.addCell(node.getHostName());
        table.addCell(node.getHostAddress());
        table.addCell(node.name());
        table.endRow();
    }

    final String UNASSIGNED = "UNASSIGNED";
    if (allocs.containsKey(UNASSIGNED)) {
        table.startRow();
        table.addCell(allocs.get(UNASSIGNED));
        table.addCell(null);
        table.addCell(null);
        table.addCell(null);
        table.addCell(null);
        table.addCell(null);
        table.addCell(null);
        table.addCell(null);
        table.addCell(UNASSIGNED);
        table.endRow();
    }

    return table;
}
 

private Decision underCapacity(ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation, boolean moveToNode) {
    if (awarenessAttributes.length == 0) {
        return allocation.decision(Decision.YES, NAME, "no allocation awareness enabled");
    }

    IndexMetaData indexMetaData = allocation.metaData().index(shardRouting.index());
    int shardCount = indexMetaData.getNumberOfReplicas() + 1; // 1 for primary
    for (String awarenessAttribute : awarenessAttributes) {
        // the node the shard exists on must be associated with an awareness attribute
        if (!node.node().attributes().containsKey(awarenessAttribute)) {
            return allocation.decision(Decision.NO, NAME, "node does not contain awareness attribute: [%s]", awarenessAttribute);
        }

        // build attr_value -> nodes map
        ObjectIntHashMap<String> nodesPerAttribute = allocation.routingNodes().nodesPerAttributesCounts(awarenessAttribute);

        // build the count of shards per attribute value
        ObjectIntHashMap<String> shardPerAttribute = new ObjectIntHashMap<>();
        for (ShardRouting assignedShard : allocation.routingNodes().assignedShards(shardRouting)) {
            if (assignedShard.started() || assignedShard.initializing()) {
                // Note: this also counts relocation targets as that will be the new location of the shard.
                // Relocation sources should not be counted as the shard is moving away
                RoutingNode routingNode = allocation.routingNodes().node(assignedShard.currentNodeId());
                shardPerAttribute.addTo(routingNode.node().attributes().get(awarenessAttribute), 1);
            }
        }

        if (moveToNode) {
            if (shardRouting.assignedToNode()) {
                String nodeId = shardRouting.relocating() ? shardRouting.relocatingNodeId() : shardRouting.currentNodeId();
                if (!node.nodeId().equals(nodeId)) {
                    // we work on different nodes, move counts around
                    shardPerAttribute.putOrAdd(allocation.routingNodes().node(nodeId).node().attributes().get(awarenessAttribute), 0, -1);
                    shardPerAttribute.addTo(node.node().attributes().get(awarenessAttribute), 1);
                }
            } else {
                shardPerAttribute.addTo(node.node().attributes().get(awarenessAttribute), 1);
            }
        }

        int numberOfAttributes = nodesPerAttribute.size();
        String[] fullValues = forcedAwarenessAttributes.get(awarenessAttribute);
        if (fullValues != null) {
            for (String fullValue : fullValues) {
                if (!shardPerAttribute.containsKey(fullValue)) {
                    numberOfAttributes++;
                }
            }
        }
        // TODO should we remove ones that are not part of full list?

        int averagePerAttribute = shardCount / numberOfAttributes;
        int totalLeftover = shardCount % numberOfAttributes;
        int requiredCountPerAttribute;
        if (averagePerAttribute == 0) {
            // if we have more attributes values than shard count, no leftover
            totalLeftover = 0;
            requiredCountPerAttribute = 1;
        } else {
            requiredCountPerAttribute = averagePerAttribute;
        }
        int leftoverPerAttribute = totalLeftover == 0 ? 0 : 1;

        int currentNodeCount = shardPerAttribute.get(node.node().attributes().get(awarenessAttribute));
        // if we are above with leftover, then we know we are not good, even with mod
        if (currentNodeCount > (requiredCountPerAttribute + leftoverPerAttribute)) {
            return allocation.decision(Decision.NO, NAME,
                    "too many shards on node for attribute: [%s], required per attribute: [%d], node count: [%d], leftover: [%d]",
                    awarenessAttribute, requiredCountPerAttribute, currentNodeCount, leftoverPerAttribute);
        }
        // all is well, we are below or same as average
        if (currentNodeCount <= requiredCountPerAttribute) {
            continue;
        }
    }

    return allocation.decision(Decision.YES, NAME, "node meets awareness requirements");
}
 

protected AsyncAction(Task task, Request request, ActionListener<Response> listener) {
    this.task = task;
    this.request = request;
    this.listener = listener;

    ClusterState state = clusterService.state();
    if (request.getHeader(LoginUserContext.TENANT_FILTER) != null) {
        state = AuthService.filterState(state, state.metaData(), (Long) request.getHeader(LoginUserContext.TENANT_FILTER));
    }
    clusterState = state;
    nodes = clusterState.nodes();

    ClusterBlockException globalBlockException = checkGlobalBlock(clusterState, request);
    if (globalBlockException != null) {
        throw globalBlockException;
    }

    String[] concreteIndices = indexNameExpressionResolver.concreteIndices(clusterState, request);
    ClusterBlockException requestBlockException = checkRequestBlock(clusterState, request, concreteIndices);
    if (requestBlockException != null) {
        throw requestBlockException;
    }

    if (logger.isTraceEnabled()) {
        logger.trace("resolving shards for [{}] based on cluster state version [{}]", actionName, clusterState.version());
    }
    ShardsIterator shardIt = shards(clusterState, request, concreteIndices);
    nodeIds = Maps.newHashMap();

    for (ShardRouting shard : shardIt.asUnordered()) {
        // send a request to the shard only if it is assigned to a node that is in the local node's cluster state
        // a scenario in which a shard can be assigned but to a node that is not in the local node's cluster state
        // is when the shard is assigned to the master node, the local node has detected the master as failed
        // and a new master has not yet been elected; in this situation the local node will have removed the
        // master node from the local cluster state, but the shards assigned to the master will still be in the
        // routing table as such
        if (shard.assignedToNode() && nodes.get(shard.currentNodeId()) != null) {
            String nodeId = shard.currentNodeId();
            if (!nodeIds.containsKey(nodeId)) {
                nodeIds.put(nodeId, new ArrayList<ShardRouting>());
            }
            nodeIds.get(nodeId).add(shard);
        } else {
            unavailableShardExceptions.add(
                    new NoShardAvailableActionException(
                            shard.shardId(),
                            " no shards available for shard " + shard.toString() + " while executing " + actionName
                    )
            );
        }
    }

    responses = new AtomicReferenceArray<>(nodeIds.size());
}
 

private Decision underCapacity(ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation, boolean moveToNode) {
    if (awarenessAttributes.isEmpty()) {
        return allocation.decision(Decision.YES, NAME,
            "allocation awareness is not enabled, set cluster setting [%s] to enable it",
            CLUSTER_ROUTING_ALLOCATION_AWARENESS_ATTRIBUTE_SETTING.getKey());
    }

    IndexMetaData indexMetaData = allocation.metaData().getIndexSafe(shardRouting.index());
    int shardCount = indexMetaData.getNumberOfReplicas() + 1; // 1 for primary
    for (String awarenessAttribute : awarenessAttributes) {
        // the node the shard exists on must be associated with an awareness attribute
        if (!node.node().getAttributes().containsKey(awarenessAttribute)) {
            return allocation.decision(Decision.NO, NAME,
                "node does not contain the awareness attribute [%s]; required attributes cluster setting [%s=%s]",
                awarenessAttribute, CLUSTER_ROUTING_ALLOCATION_AWARENESS_ATTRIBUTE_SETTING.getKey(),
                allocation.debugDecision() ? Strings.collectionToCommaDelimitedString(awarenessAttributes) : null);
        }

        // build attr_value -> nodes map
        ObjectIntHashMap<String> nodesPerAttribute = allocation.routingNodes().nodesPerAttributesCounts(awarenessAttribute);

        // build the count of shards per attribute value
        ObjectIntHashMap<String> shardPerAttribute = new ObjectIntHashMap<>();
        for (ShardRouting assignedShard : allocation.routingNodes().assignedShards(shardRouting.shardId())) {
            if (assignedShard.started() || assignedShard.initializing()) {
                // Note: this also counts relocation targets as that will be the new location of the shard.
                // Relocation sources should not be counted as the shard is moving away
                RoutingNode routingNode = allocation.routingNodes().node(assignedShard.currentNodeId());
                shardPerAttribute.addTo(routingNode.node().getAttributes().get(awarenessAttribute), 1);
            }
        }

        if (moveToNode) {
            if (shardRouting.assignedToNode()) {
                String nodeId = shardRouting.relocating() ? shardRouting.relocatingNodeId() : shardRouting.currentNodeId();
                if (!node.nodeId().equals(nodeId)) {
                    // we work on different nodes, move counts around
                    shardPerAttribute.putOrAdd(allocation.routingNodes().node(nodeId).node().getAttributes().get(awarenessAttribute),
                            0, -1);
                    shardPerAttribute.addTo(node.node().getAttributes().get(awarenessAttribute), 1);
                }
            } else {
                shardPerAttribute.addTo(node.node().getAttributes().get(awarenessAttribute), 1);
            }
        }

        int numberOfAttributes = nodesPerAttribute.size();
        List<String> fullValues = forcedAwarenessAttributes.get(awarenessAttribute);
        if (fullValues != null) {
            for (String fullValue : fullValues) {
                if (!shardPerAttribute.containsKey(fullValue)) {
                    numberOfAttributes++;
                }
            }
        }
        // TODO should we remove ones that are not part of full list?

        int averagePerAttribute = shardCount / numberOfAttributes;
        int totalLeftover = shardCount % numberOfAttributes;
        int requiredCountPerAttribute;
        if (averagePerAttribute == 0) {
            // if we have more attributes values than shard count, no leftover
            totalLeftover = 0;
            requiredCountPerAttribute = 1;
        } else {
            requiredCountPerAttribute = averagePerAttribute;
        }
        int leftoverPerAttribute = totalLeftover == 0 ? 0 : 1;

        int currentNodeCount = shardPerAttribute.get(node.node().getAttributes().get(awarenessAttribute));
        // if we are above with leftover, then we know we are not good, even with mod
        if (currentNodeCount > (requiredCountPerAttribute + leftoverPerAttribute)) {
            return allocation.decision(Decision.NO, NAME,
                    "there are too many copies of the shard allocated to nodes with attribute [%s], there are [%d] total configured " +
                    "shard copies for this shard id and [%d] total attribute values, expected the allocated shard count per " +
                    "attribute [%d] to be less than or equal to the upper bound of the required number of shards per attribute [%d]",
                    awarenessAttribute,
                    shardCount,
                    numberOfAttributes,
                    currentNodeCount,
                    requiredCountPerAttribute + leftoverPerAttribute);
        }
        // all is well, we are below or same as average
        if (currentNodeCount <= requiredCountPerAttribute) {
            continue;
        }
    }

    return allocation.decision(Decision.YES, NAME, "node meets all awareness attribute requirements");
}
 

protected AsyncAction(Task task, Request request, ActionListener<Response> listener) {
    this.task = task;
    this.request = request;
    this.listener = listener;

    clusterState = clusterService.state();
    nodes = clusterState.nodes();

    ClusterBlockException globalBlockException = checkGlobalBlock(clusterState, request);
    if (globalBlockException != null) {
        throw globalBlockException;
    }

    String[] concreteIndices = indexNameExpressionResolver.concreteIndexNames(clusterState, request);
    ClusterBlockException requestBlockException = checkRequestBlock(clusterState, request, concreteIndices);
    if (requestBlockException != null) {
        throw requestBlockException;
    }

    if (logger.isTraceEnabled()) {
        logger.trace("resolving shards for [{}] based on cluster state version [{}]", actionName, clusterState.version());
    }
    ShardsIterator shardIt = shards(clusterState, request, concreteIndices);
    nodeIds = new HashMap<>();

    for (ShardRouting shard : shardIt) {
        // send a request to the shard only if it is assigned to a node that is in the local node's cluster state
        // a scenario in which a shard can be assigned but to a node that is not in the local node's cluster state
        // is when the shard is assigned to the master node, the local node has detected the master as failed
        // and a new master has not yet been elected; in this situation the local node will have removed the
        // master node from the local cluster state, but the shards assigned to the master will still be in the
        // routing table as such
        if (shard.assignedToNode() && nodes.get(shard.currentNodeId()) != null) {
            String nodeId = shard.currentNodeId();
            if (!nodeIds.containsKey(nodeId)) {
                nodeIds.put(nodeId, new ArrayList<>());
            }
            nodeIds.get(nodeId).add(shard);
        } else {
            unavailableShardExceptions.add(
                    new NoShardAvailableActionException(
                            shard.shardId(),
                            " no shards available for shard " + shard.toString() + " while executing " + actionName
                    )
            );
        }
    }

    responses = new AtomicReferenceArray<>(nodeIds.size());
}
 

@Override
public ExecutionPlan build(PlannerContext plannerContext,
                           ProjectionBuilder projectionBuilder,
                           int limitHint,
                           int offsetHint,
                           @Nullable OrderBy order,
                           @Nullable Integer pageSizeHint,
                           Row params,
                           SubQueryResults subQueryResults) {
    HashMap<String, Map<ShardId, List<PKAndVersion>>> idsByShardByNode = new HashMap<>();
    DocTableInfo docTableInfo = tableRelation.tableInfo();
    List<Symbol> boundOutputs = Lists2.map(
        outputs, s -> SubQueryAndParamBinder.convert(s, params, subQueryResults));
    for (DocKeys.DocKey docKey : docKeys) {
        String id = docKey.getId(plannerContext.transactionContext(), plannerContext.functions(), params, subQueryResults);
        if (id == null) {
            continue;
        }
        List<String> partitionValues = docKey.getPartitionValues(
            plannerContext.transactionContext(), plannerContext.functions(), params, subQueryResults);
        String indexName = indexName(docTableInfo, partitionValues);

        String routing = docKey.getRouting(
            plannerContext.transactionContext(), plannerContext.functions(), params, subQueryResults);
        ShardRouting shardRouting;
        try {
            shardRouting = plannerContext.resolveShard(indexName, id, routing);
        } catch (IndexNotFoundException e) {
            if (docTableInfo.isPartitioned()) {
                continue;
            }
            throw e;
        }
        String currentNodeId = shardRouting.currentNodeId();
        if (currentNodeId == null) {
            // If relocating is fast enough this will work, otherwise it will result in a shard failure which
            // will cause a statement retry
            currentNodeId = shardRouting.relocatingNodeId();
            if (currentNodeId == null) {
                throw new ShardNotFoundException(shardRouting.shardId());
            }
        }
        Map<ShardId, List<PKAndVersion>> idsByShard = idsByShardByNode.get(currentNodeId);
        if (idsByShard == null) {
            idsByShard = new HashMap<>();
            idsByShardByNode.put(currentNodeId, idsByShard);
        }
        List<PKAndVersion> pkAndVersions = idsByShard.get(shardRouting.shardId());
        if (pkAndVersions == null) {
            pkAndVersions = new ArrayList<>();
            idsByShard.put(shardRouting.shardId(), pkAndVersions);
        }
        long version = docKey
            .version(plannerContext.transactionContext(), plannerContext.functions(), params, subQueryResults)
            .orElse(Versions.MATCH_ANY);
        long sequenceNumber = docKey.sequenceNo(plannerContext.transactionContext(), plannerContext.functions(), params, subQueryResults)
            .orElse(SequenceNumbers.UNASSIGNED_SEQ_NO);
        long primaryTerm = docKey.primaryTerm(plannerContext.transactionContext(), plannerContext.functions(), params, subQueryResults)
            .orElse(SequenceNumbers.UNASSIGNED_PRIMARY_TERM);
        pkAndVersions.add(new PKAndVersion(id, version, sequenceNumber, primaryTerm));
    }
    return new Collect(
        new PKLookupPhase(
            plannerContext.jobId(),
            plannerContext.nextExecutionPhaseId(),
            docTableInfo.partitionedBy(),
            boundOutputs,
            idsByShardByNode
        ),
        TopN.NO_LIMIT,
        0,
        boundOutputs.size(),
        docKeys.size(),
        null
    );
}