vlambda博客
学习文章列表

  // 持久化的相关文件的输出流
  // meta：向量元信息文件输出流
  // vectorData: 向量数据文件输出流
  // vectorIndex: 向量索引文件输出流（存储邻居信息）
  private final IndexOutput meta, vectorData, vectorIndex;

  Lucene91HnswVectorsWriter(SegmentWriteState state, int maxConn, int beamWidth)
      throws IOException {
    this.maxConn = maxConn;
    this.beamWidth = beamWidth;

    assert state.fieldInfos.hasVectorValues();
    segmentWriteState = state;

    // 文件名是 _段编号_段后缀名.文件扩展名
    // 段后缀名三个文件是统一的：Lucene91HnswVectorsFormat_0
    String metaFileName =
        IndexFileNames.segmentFileName(
            state.segmentInfo.name, state.segmentSuffix, Lucene91HnswVectorsFormat.META_EXTENSION);

    String vectorDataFileName =
        IndexFileNames.segmentFileName(
            state.segmentInfo.name,
            state.segmentSuffix,
            Lucene91HnswVectorsFormat.VECTOR_DATA_EXTENSION);

    String indexDataFileName =
        IndexFileNames.segmentFileName(
            state.segmentInfo.name,
            state.segmentSuffix,
            Lucene91HnswVectorsFormat.VECTOR_INDEX_EXTENSION);

    boolean success = false;
    try {
      // 初始化话三个文件的输出流 
      meta = state.directory.createOutput(metaFileName, state.context);
      vectorData = state.directory.createOutput(vectorDataFileName, state.context);
      vectorIndex = state.directory.createOutput(indexDataFileName, state.context);

      // 分别写入文件头
      // 文件头信息有：
      // 1.文件头魔数（同一lucene版本所有文件相同）
      // 2.该文件使用的codec名称：Lucene91HnswVectorsFormatData
      // 3.codec版本：0
      // 4.segment id: 段唯一标识符，标识这个文件属于哪个段的
      // 5.segment后缀名：Lucene91HnswVectorsFormat_0
      CodecUtil.writeIndexHeader(
          meta,
          Lucene91HnswVectorsFormat.META_CODEC_NAME,
          Lucene91HnswVectorsFormat.VERSION_CURRENT,
          state.segmentInfo.getId(),
          state.segmentSuffix);
      CodecUtil.writeIndexHeader(
          vectorData,
          Lucene91HnswVectorsFormat.VECTOR_DATA_CODEC_NAME,
          Lucene91HnswVectorsFormat.VERSION_CURRENT,
          state.segmentInfo.getId(),
          state.segmentSuffix);
      CodecUtil.writeIndexHeader(
          vectorIndex,
          Lucene91HnswVectorsFormat.VECTOR_INDEX_CODEC_NAME,
          Lucene91HnswVectorsFormat.VERSION_CURRENT,
          state.segmentInfo.getId(),
          state.segmentSuffix);
      maxDoc = state.segmentInfo.maxDoc();
      success = true;
    } finally {
      if (success == false) {
        IOUtils.closeWhileHandlingException(this);
      }
    }
  }

  public void finish() throws IOException {
    if (finished) {
      throw new IllegalStateException("already finished");
    }
    finished = true;

    if (meta != null) {
      // -1是表示向量元信息的主内容已经结束。
      // 多个向量字段的向量元信息是写在一起的，因为没有单独记录有多少个向量字段，所以用-1作为字段结束标记
      meta.writeInt(-1);
      // 写文件注脚，注脚内容：
      // 1.文件尾魔数（同一个lucene版本所有文件一样）
      // 2.0
      // 3.校验码
      CodecUtil.writeFooter(meta);
    }
    if (vectorData != null) {
      CodecUtil.writeFooter(vectorData);
      CodecUtil.writeFooter(vectorIndex);
    }
  }

  public void writeField(FieldInfo fieldInfo, KnnVectorsReader knnVectorsReader)
      throws IOException {
    // 计算当前字段中所有向量数据在向量数据文件的起始位置，该信息会写入元信息文件
    long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES);
    VectorValues vectors = knnVectorsReader.getVectorValues(fieldInfo.name);

    IndexOutput tempVectorData =
        segmentWriteState.directory.createTempOutput(
            vectorData.getName(), "temp", segmentWriteState.context);
    IndexInput vectorDataInput = null;
    boolean success = false;
    try {
      // 向量的数据先存临时文件，构建HNSW失败不会删除临时文件，构建成功则会删除临时文件
      DocsWithFieldSet docsWithField = writeVectorData(tempVectorData, vectors);
      CodecUtil.writeFooter(tempVectorData);
      IOUtils.close(tempVectorData);

      // 从临时文件中读取向量数据
      vectorDataInput =
          segmentWriteState.directory.openInput(
              tempVectorData.getName(), segmentWriteState.context);
      // 重点来了，向量的数据都拷贝到最后落盘文件的输出流中 
      vectorData.copyBytes(vectorDataInput, vectorDataInput.length() - CodecUtil.footerLength());
      CodecUtil.retrieveChecksum(vectorDataInput);
      // 计算当前字段所有向量数据的总长度，该信息会写入元信息文件
      long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset;

     。。。（省略元信息文件和索引文件持久化逻辑，后面单独讨论）
    } finally {
      。。。（省略资源关闭）
    }
  }

  @Override
  public void writeField(FieldInfo fieldInfo, KnnVectorsReader knnVectorsReader)
      throws IOException {
    。。。（忽略和索引文件落盘无关逻辑）
    try {
      。。。（忽略和索引文件落盘无关逻辑）
      // 计算当前处理的字段的索引文件信息的起始位置，该信息会写入元信息文件
      long vectorIndexOffset = vectorIndex.getFilePointer();
      // 
      Lucene91HnswVectorsReader.OffHeapVectorValues offHeapVectors =
          new Lucene91HnswVectorsReader.OffHeapVectorValues(
              vectors.dimension(), docsWithField.cardinality(), null, vectorDataInput);
      OnHeapHnswGraph graph =
          offHeapVectors.size() == 0
              ? null
              // 索引信息存储的是邻居信息，所以必须先构建HNSW 
              : writeGraph(offHeapVectors, fieldInfo.getVectorSimilarityFunction());
      // 计算当前处理的字段的索引文件信息的总长度，该信息会写入元信息文件 
      long vectorIndexLength = vectorIndex.getFilePointer() - vectorIndexOffset;

      。。。（忽略和索引文件落盘无关逻辑）
    } finally {
      。。。（忽略和索引文件落盘无关逻辑）
  }

  private OnHeapHnswGraph writeGraph(
      RandomAccessVectorValuesProducer vectorValues, VectorSimilarityFunction similarityFunction)
      throws IOException {

    // 构建HNSW（详细逻辑请看《lucene 9.0.1 HNSW的源码解析》）
    HnswGraphBuilder hnswGraphBuilder =
        new HnswGraphBuilder(
            vectorValues, similarityFunction, maxConn, beamWidth, HnswGraphBuilder.randSeed);
    hnswGraphBuilder.setInfoStream(segmentWriteState.infoStream);
    OnHeapHnswGraph graph = hnswGraphBuilder.build(vectorValues.randomAccess());

    // 邻居信息落盘
    // HNSW总共有多少个节点，也就是一共有多少个向量。
    int countOnLevel0 = graph.size();
    // 从对底层开始处理每一层的邻居信息
    for (int level = 0; level < graph.numLevels(); level++) {
      // 获取当前层的节点迭代器 
      NodesIterator nodesOnLevel = graph.getNodesOnLevel(level);
      while (nodesOnLevel.hasNext()) {
        int node = nodesOnLevel.nextInt();
        // 获取当前层当前节点的邻居信息 
        NeighborArray neighbors = graph.getNeighbors(level, node);
        // 邻居的总数，因为不是每个节点都是maxConn个邻居 
        int size = neighbors.size();
        // 先写入邻居的个数 
        vectorIndex.writeInt(size);
        // 先对所有邻居按照节点编号从小到大排序，以后如果需要确认某个节点是否是其邻居，就可以通过二分查找了
        int[] nnodes = neighbors.node();
        Arrays.sort(nnodes, 0, size);
        // 写入所有的邻居节点编号 
        for (int i = 0; i < size; i++) {
          int nnode = nnodes[i];
          vectorIndex.writeInt(nnode);
        }
        // 对于邻居没有达到maxConn个用0补齐，这是为了方便计算每个节点邻居存储的起始offset，后面读取的时候会介绍。
        for (int i = size; i < maxConn; i++) {
          vectorIndex.writeInt(0);
        }
      }
    }
    return graph;
  }

  private void writeMeta(
      FieldInfo field,        // 当前处理的向量字段信息
      long vectorDataOffset,  // 当前字段在向量数据文件中的起始位置
      long vectorDataLength,  // 当前字段的向量数据总长度
      long vectorIndexOffset, // 当前字段在向量索引文件中的起始位置
      long vectorIndexLength, // 当前字段在向量索引数据总长度
      DocsWithFieldSet docsWithField, // 包含当前字段的的文档位图
      OnHeapHnswGraph graph) // HNSW，用来获取每层的节点
      throws IOException {
    // 写入当前处理的字段的编号
    meta.writeInt(field.number);
    // 写入向量字段的距离度量方式
    meta.writeInt(field.getVectorSimilarityFunction().ordinal());
    // 写入当前字段在向量数据文件中的起始位置
    meta.writeVLong(vectorDataOffset);
    // 写入当前字段的向量数据总长度
    meta.writeVLong(vectorDataLength);
    // 写入当前字段在向量索引文件中的起始位置
    meta.writeVLong(vectorIndexOffset);
    // 写入当前字段在向量索引数据总长度
    meta.writeVLong(vectorIndexLength);
    // 写入当前字段的向量的维度
    meta.writeInt(field.getVectorDimension());

    int count = docsWithField.cardinality();
    // 写入所有包含此字段的文档总数
    meta.writeInt(count);
    // 如果所有的文档都包含此字段，则写入-1作为标记
    if (count == maxDoc) {
      meta.writeByte((byte) -1);
    } else {
      // 否则写0，表示只有部分文档包含此字段
      meta.writeByte((byte) 0);
      DocIdSetIterator iter = docsWithField.iterator();
      // 遍历所有的文档，写入文档编号
      for (int doc = iter.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = iter.nextDoc()) {
        meta.writeInt(doc);
      }
    }

    // 写入最大邻居个数
    meta.writeInt(maxConn);
    // 如果没有节点，则写入0作为标记
    if (graph == null) {
      meta.writeInt(0);
    } else {
      // 否则写入hnsw的总层数
      meta.writeInt(graph.numLevels());
      // 从低到高遍历所有的层
      for (int level = 0; level < graph.numLevels(); level++) {
        NodesIterator nodesOnLevel = graph.getNodesOnLevel(level);
        // 写入当前层的节点个数
        meta.writeInt(nodesOnLevel.size()); 
        // 因为第0层默认包含所有的节点，因此从第一层开始写入所有的节点
        if (level > 0) {
          while (nodesOnLevel.hasNext()) {
            int node = nodesOnLevel.nextInt();
            meta.writeInt(node); 
          }
        }
      }
    }
  }

// 所有字段的信息 
private final FieldInfos fieldInfos;

// key是向量字段名，向量字段元信息从元信息文件中解析存在在FieldEntry中
private final Map<String, FieldEntry> fields = new HashMap<>();

// 向量数据文件，可以根据FieldEntry中的信息获取到特定字段的所有向量数据
private final IndexInput vectorData;

// 向量索引文件（邻居信息），可以根据FieldEntry中的信息获取到特定字段的所有邻居数据
private final IndexInput vectorIndex;

  Lucene91HnswVectorsReader(SegmentReadState state) throws IOException {
    this.fieldInfos = state.fieldInfos;
    // 解析向量元信息文件，是最为核心的逻辑 
    int versionMeta = readMetadata(state);
    boolean success = false;
    try {
      // 读取向量数据 
      vectorData =
          openDataInput(
              state,
              versionMeta,
              Lucene91HnswVectorsFormat.VECTOR_DATA_EXTENSION,
              Lucene91HnswVectorsFormat.VECTOR_DATA_CODEC_NAME);
      // 读取向量索引数据 
      vectorIndex =
          openDataInput(
              state,
              versionMeta,
              Lucene91HnswVectorsFormat.VECTOR_INDEX_EXTENSION,
              Lucene91HnswVectorsFormat.VECTOR_INDEX_CODEC_NAME);
      success = true;
    } finally {
      if (success == false) {
        IOUtils.closeWhileHandlingException(this);
      }
    }
  }

  private int readMetadata(SegmentReadState state) throws IOException {
    String metaFileName =
        IndexFileNames.segmentFileName(
            state.segmentInfo.name, state.segmentSuffix, Lucene91HnswVectorsFormat.META_EXTENSION);
    int versionMeta = -1;
    try (ChecksumIndexInput meta = state.directory.openChecksumInput(metaFileName, state.context)) {
      Throwable priorE = null;
      try {
        versionMeta =
            CodecUtil.checkIndexHeader(
                meta,
                Lucene91HnswVectorsFormat.META_CODEC_NAME,
                Lucene91HnswVectorsFormat.VERSION_START,
                Lucene91HnswVectorsFormat.VERSION_CURRENT,
                state.segmentInfo.getId(),
                state.segmentSuffix);
        // 解析所有向量字段的元信息
        readFields(meta, state.fieldInfos);
      } catch (Throwable exception) {
        priorE = exception;
      } finally {
        CodecUtil.checkFooter(meta, priorE);
      }
    }
    return versionMeta;
  }

  private void readFields(ChecksumIndexInput meta, FieldInfos infos) throws IOException {
    // 遍历所有的字段元信息，指导发现了-1结束标记（还记得落盘中说向量元信息的结束标记吗） 
    for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) {
      FieldInfo info = infos.fieldInfo(fieldNumber);
      if (info == null) {
        throw new CorruptIndexException("Invalid field number: " + fieldNumber, meta);
      }
      // 解析元信息数据，封装成FieldEntry对象
      FieldEntry fieldEntry = readField(meta);
      // 向量维度和向量数据长度的校验 
      validateFieldEntry(info, fieldEntry);
      // 缓存字段的FieldEntry信息
      fields.put(info.name, fieldEntry);
    }
  }

  private FieldEntry readField(DataInput input) throws IOException {
    // 读取距离度量信息
    VectorSimilarityFunction similarityFunction = readSimilarityFunction(input);
    // 构建字段的FieldEntry 
    return new FieldEntry(input, similarityFunction);
  }

  private VectorSimilarityFunction readSimilarityFunction(DataInput input) throws IOException {
    int similarityFunctionId = input.readInt();
    if (similarityFunctionId < 0
        || similarityFunctionId >= VectorSimilarityFunction.values().length) {
      throw new CorruptIndexException(
          "Invalid similarity function id: " + similarityFunctionId, input);
    }
    return VectorSimilarityFunction.values()[similarityFunctionId];
  }

    // 这几个就不再说了
    final VectorSimilarityFunction similarityFunction;
    final long vectorDataOffset;
    final long vectorDataLength;
    final long vectorIndexOffset;
    final long vectorIndexLength;
    final int maxConn;
    final int numLevels;
    final int dimension;
    private final int size;

    // 存储的是节点id对应的docId
    final int[] ordToDoc;
    // 根据节点id获取docId操作器
    private final IntUnaryOperator ordToDocOperator;
    // 每一层中的所有节点，注意构建的时候已经决定了每层的节点编号是从小到大的
    final int[][] nodesByLevel;
    // 每一层中邻居信息在向量文件中索引文件的起始位置，这个是提前算出来的，怎么算后面有解释
    final long[] graphOffsetsByLevel;

    FieldEntry(DataInput input, VectorSimilarityFunction similarityFunction) throws IOException {
      this.similarityFunction = similarityFunction;
      vectorDataOffset = input.readVLong();
      vectorDataLength = input.readVLong();
      vectorIndexOffset = input.readVLong();
      vectorIndexLength = input.readVLong();
      dimension = input.readInt();
      size = input.readInt();

      。。。（解析文档编号）
      。。。（解析每一层的node）
      。。。（初始化graphOffsetsByLevel）
    }

    FieldEntry(DataInput input, VectorSimilarityFunction similarityFunction) throws IOException {
      。。。（解析部分元信息）
      // 读取是否是稀有数据的标记
      int denseSparseMarker = input.readByte();
      if (denseSparseMarker == -1) {
        // -1代表所有的文档都包含该字段 
        ordToDoc = null; 
      } else {
        // 一次解析所有的文档id，放入ordToDoc数组中，数组的小标是节点编号，数组值就是docid
        ordToDoc = new int[size];
        for (int i = 0; i < size; i++) {
          int doc = input.readInt();
          ordToDoc[i] = doc;
        }
      }
      ordToDocOperator = ordToDoc == null ? IntUnaryOperator.identity() : (ord) -> ordToDoc[ord];

     。。。（解析每一层的节点）
      。。。（初始化graphOffsetsByLevel）
    }

FieldEntry(DataInput input, VectorSimilarityFunction similarityFunction) throws IOException {
  。。。（解析部分元信息）

  。。。（解析文档编号）



  // 最大邻居个数
  maxConn = input.readInt();
  // HNSW的层数 
  numLevels = input.readInt();
  nodesByLevel = new int[numLevels][];
  // 遍历所有层 
  for (int level = 0; level < numLevels; level++) {
    // 当前层的节点个数 
    int numNodesOnLevel = input.readInt();
    if (level == 0) {
      // 第0层默认存储了所有的节点
      nodesByLevel[0] = null;
    } else {
      nodesByLevel[level] = new int[numNodesOnLevel];
      for (int i = 0; i < numNodesOnLevel; i++) {
        // 获取节点编号 
        nodesByLevel[level][i] = input.readInt();
      }
    }
  }

  。。。（初始化graphOffsetsByLevel）
}

    FieldEntry(DataInput input, VectorSimilarityFunction similarityFunction) throws IOException {
      。。。（解析部分元信息）
      。。。（解析文档编号）
      。。。（解析每一层的节点）

      // 每一层的邻居数据的起始位置
      graphOffsetsByLevel = new long[numLevels];
      for (int level = 0; level < numLevels; level++) {
        if (level == 0) {
          graphOffsetsByLevel[level] = 0;
        } else {
          int numNodesOnPrevLevel = level == 1 ? size : nodesByLevel[level - 1].length;
          // 1 + maxConn：1个int记录真实的邻居个数 + maxConn个int记录所有的邻居编号
          // (1 + maxConn) * Integer.BYTES * numNodesOnPrevLevel：前一层的邻居数据长度
          // graphOffsetsByLevel[level - 1]：前一层邻居数据的起始位置
          // 从而得到graphOffsetsByLevel[level]是当前层邻居数据的起始位置 
          graphOffsetsByLevel[level] =
              graphOffsetsByLevel[level - 1] + (1 + maxConn) * Integer.BYTES * numNodesOnPrevLevel;
        }
      }
    }

public HnswGraph getGraph(String field) throws IOException {
  FieldInfo info = fieldInfos.fieldInfo(field);
  if (info == null) {
    throw new IllegalArgumentException("No such field '" + field + "'");
  }
  FieldEntry entry = fields.get(field);
  if (entry != null && entry.vectorIndexLength > 0) {
      // 通过FieldEntry构建HNSW
    return getGraph(entry);
  } else {
    return HnswGraph.EMPTY;
  }
}

  private HnswGraph getGraph(FieldEntry entry) throws IOException {
    // 向量索引数据(邻居信息数据)
    IndexInput bytesSlice =
        vectorIndex.slice("graph-data", entry.vectorIndexOffset, entry.vectorIndexLength);
    // 构建HNSW 
    return new OffHeapHnswGraph(entry, bytesSlice);
  }

    // 向量索引文件数据
    final IndexInput dataIn;

    // 每层的所有节点
    final int[][] nodesByLevel;

    // 在向量索引文件中的每层的邻居信息起始位置
    final long[] graphOffsetsByLevel;

    // HNSW的层数
    final int numLevels;

    // 搜索的起始节点
    final int entryNode;

    // hnsw中的向量总数
    final int size;

    // 每个节点的邻居信息在向量索引文件中占用的字节，
    // 因为所有节点邻居都是maxConn个（不足的文件写入会补0），所以长度是固定的，可以提前算出来
    final long bytesForConns;

    // 邻居的个数
    int arcCount;

    // 下一个要遍历的邻居编号
    int arcUpTo;

    // 邻居编号
    int arc;

    OffHeapHnswGraph(FieldEntry entry, IndexInput dataIn) {
      // 只要有数据，就把搜索的起始遍历节点设置为最顶层的第一个节点
      this.entryNode = numLevels > 1 ? nodesByLevel[numLevels - 1][0] : 0;

      // 加1是因为存储所有邻居信息之前会额外存储真实邻居个数（并不是所有的节点邻居都是maxConn，有些是补0的） 
      this.bytesForConns = ((long) entry.maxConn + 1) * Integer.BYTES;
    }

  // 定位邻居信息的起始位置思路：
  // nodesByLevel[level]存放的是第level层中所有的节点，节点是按编号从到到大排列。
  // 对于第0层，所有节点都存在，所以节点编号和数组下标一一对应。
  // 对于其他层，可以通过二分法找到对应节点在数组中的下标，数组的下标i就可以用来确定第i个节点邻居信息的起始位置
  public void seek(int level, int targetOrd) throws IOException {
      // 先找到节点编号为targetOrd在数组中的下标
      int targetIndex =
          level == 0
                // 第0层包含所有的节点，数组下标就是节点的编号，可以直接获取
              ? targetOrd
                // 二分法查找指定值的数组下标
              : Arrays.binarySearch(nodesByLevel[level], 0, nodesByLevel[level].length, targetOrd);

      // graphOffsetsByLevel[level]是第level层的所有节点的邻居信息的起始位置。
      // bytesForConns是一个节点的邻居信息占用的大小
      // graphOffsetsByLevel[level] + targetIndex * bytesForConns就是第targetIndex节点邻居信息的起始位置
      long graphDataOffset = graphOffsetsByLevel[level] + targetIndex * bytesForConns;
      // 定位到节点邻居信息的起始位置
      dataIn.seek(graphDataOffset);
      // 获取真实邻居总数
      arcCount = dataIn.readInt();
      // 当前邻居编号初始化为-1，表示还没开始遍历
      arc = -1;
      // 下一个要遍历的编号
      arcUpTo = 0;
    }

    @Override
    public int nextNeighbor() throws IOException {
      // 判断是否遍历结束，注意编号是从0开始的 
      if (arcUpTo >= arcCount) {
        return NO_MORE_DOCS;
      }
      ++arcUpTo;
      // 读取邻居编号 
      arc = dataIn.readInt();
      return arc;
    }