def METHOD_NAME(self, input_id=0):
    # key is mapper chunk key, index is mapper chunk index.
    input_chunk = self.inputs[input_id]
    if isinstance(input_chunk.op, ShuffleProxy):
        keys = [inp.key for inp in input_chunk.inputs]
    else:
        assert isinstance(input_chunk.op, FetchShuffle), input_chunk.op
        if input_chunk.op.shuffle_fetch_type == ShuffleFetchType.FETCH_BY_INDEX:
            # For fetch shuffle by index, all shuffle block of same reducers are
            # identified by their index. chunk key are not needed any more.
            # so just mock key here.
            # keep this in sync with ray executor `execute_subtask`.
            return list(range(input_chunk.op.n_mappers))
        keys = input_chunk.op.source_keys
    return keys