Faiss(17)：IndexIVFPQ文件格式分析

1. 说明

Faiss版本：1.6.3
nb = 1M, nlist = 4000, d = 64, m = 64, nbits = 8

2. 应用层接口

python接口

#将生成的index实例保存到index_name.bin命名的文件中
faiss.write_index(index, "index_name.bin")

c++接口

在C++代码中，write_index()函数是一个在faiss命名空间下的独立的函数，即它不依赖于某个特定的类实例，而在代码中可直接调用。如：

//demo.cpp
...
faiss::write_index(&index, "index_name.bin");
...

write_index()是一个重载类型函数，有多种参数可供调用，函数声明在index_io.h文件中，如下：

namespace faiss {
...
void write_index (const Index *idx, const char *fname);
void write_index (const Index *idx, FILE *f);
void write_index (const Index *idx, IOWriter *writer);
...
}

其中常用的是前两种接口，要么使用一个字符串作为参数，要么直接传递一个流文件描述符。这两种最终都会调用到第三种接口中去。
下面以示例中使用的第一种函数进行代码分析。

3. C++源代码 write_index()函数

该函数定义在impl/index_write.cpp文件中

void write_index (const Index *idx, const char *fname) {
    FileIOWriter writer(fname);      //FileIOWrite根据文件名生成一个打开的流文件，用于保存后续从index中提取的数据流
    write_index (idx, &writer);
}


void write_index (const Index *idx, IOWriter *f) {
    ...
    // 这里仅保留IndexIVFPQ类型相关的操作
   else if(const IndexIVFPQ * ivpq =
              dynamic_cast<const IndexIVFPQ *> (idx)) {
        const IndexIVFPQR * ivfpqr = dynamic_cast<const IndexIVFPQR *> (idx);
        
        // h = 0x51507749, 其实就是Index的类型string转换为以小端方式表示的十六进制数
        uint32_t h = fourcc (ivfpqr ? "IwQR" : "IwPQ");
        WRITE1 (h);                                                 // WRITE1宏将括号内参数写入到f流中, 4bytes
        write_ivf_header (ivpq, f);
        WRITE1 (ivpq->by_residual);
        WRITE1 (ivpq->code_size);
        write_ProductQuantizer (&ivpq->pq, f);
        write_InvertedLists (ivpq->invlists, f);
        if (ivfpqr) {
            write_ProductQuantizer (&ivfpqr->refine_pq, f);
            WRITEVECTOR (ivfpqr->refine_codes);
            WRITE1 (ivfpqr->k_factor);
        }

    } 
    ...
}

WRITE1宏是写入操作的基本调用，其定义如下

//WRITE1宏定义如下：
#define WRITE1(x) WRITEANDCHECK(&(x), 1)

#define WRITEANDCHECK(ptr, n) {                                 \
        size_t ret = (*f)(ptr, sizeof(*(ptr)), n);              \
        FAISS_THROW_IF_NOT_FMT(ret == (n),                      \
            "write error in %s: %ld != %ld (%s)",               \
            f->name.c_str(), ret, size_t(n), strerror(errno));  \
    }

进入此函数后，首先会判断idx的具体index类型执行相应的操作，支持的类型有：IndexFlat, IndexLSH, IndexPQ, Index2Layer, IndexScalarQuantizer, IndexLattice, IndexIVFFlatDedup, IndexIVFFlat, IndexIVFScalarQuantizer, IndexIVFSpectralHash, IndexIVFPQ, IndexIVFPQR, IndexPreTransform, MultiIndexQuantizer, IndexRefineFlat, IndexDMap, IndexHNSW。

IndexIVFPQ与IndexIVFPQR在这里被认为是同一种类型进行处理，不同之处在于后者有额外的写入ProductQuantizer的过程。

写入IndexIVFPQ文件的总体步骤为：

1. 写入magic number：IwPQ (0x51507749)
2. 写入倒排信息 write_ivf_header()
3. 写入by_residual，该变量用于表明数据是编码残差(1)或纯矢量(0)，我这里数据为1，占1byte
4. 写入code_size，该变量用于表示每个vector的代码大小，即我们传入的参数M，我的值为64，占8bytes
5. 写入ProductQuantizer信息 write_ProductQuantizer()
6. 写入倒排序列write_InvertedLists()

对于IndexBinary相关类型的Index会有另一套存储格式，不从这个接口进入。

3.1 write_ivf_header ()函数

该函数用于写入倒序列表的表头，包含index的头部，nlist, nprobe等

static void write_ivf_header (const IndexIVF *ivf, IOWriter *f) {
    write_index_header (ivf, f);            // 写入index文件头部
    WRITE1 (ivf->nlist);                    // 写入nlist, 8bytes
    WRITE1 (ivf->nprobe);                   // 写入nprobe， 8bytes
    write_index (ivf->quantizer, f);
    write_direct_map (&ivf->direct_map, f);
}

此函数写入内容包括：

1. index文件头部，这部分通常为所有index通用部分， write_index_header()；
2. 写入nlist值， 4000， 占8bytes；
3. 写入nprobe值, 如果之前没有设置过，那么为1，占8bytes；
4. 写入粗聚类中心write_index (ivf->quantizer, f)；
5. 写入direct_map

3.1.1 写入索引头部信息 write_index_header()

static void write_index_header (const Index *idx, IOWriter *f) {
    WRITE1 (idx->d);                    // 写入数据维度，4bytes
    WRITE1 (idx->ntotal);               // 写入nb， 8bytes
    Index::idx_t dummy = 1 << 20;
    WRITE1 (dummy);                     // 写入dummy, 0x100000, 8bytes 
    WRITE1 (dummy);                     // 写入dummy, 0x100000, 8bytes 
    WRITE1 (idx->is_trained);           // 是否已训练， 1byte
    WRITE1 (idx->metric_type);          // metric_type, METRIC_L2=1, 4bytes
    if (idx->metric_type > 1) {
        WRITE1 (idx->metric_arg);      // 默认不写入metric_arg, 4bytes
    }
}

Index文件头部包含：数据维度(d), 添加的向量总数(ntotal)， 两次dummy值(0x100000)，训练标志位(is_trained)， metric_type， metric_arg（默认情况下IndexIVFPQ使用的metric_type=METRIC_L2=1，所以不需要写入metric_arg）

3.1.2 写入粗聚类中心

写入粗聚类中心仍然是调用了write_index()函数，只是这次传入的参数是IndexIVFPQ->quantizer，该quantizer是用于将向量映射到反向列表(Inverted list)的量化器。

void write_index (const Index *idx, IOWriter *f) {
    // 这里仅保留粗聚类中心相关的操作
    if (const IndexFlat * idxf = dynamic_cast<const IndexFlat *> (idx)) {
        uint32_t h = fourcc (
              idxf->metric_type == METRIC_INNER_PRODUCT ? "IxFI" :
              idxf->metric_type == METRIC_L2 ? "IxF2" : "IxFl");
        WRITE1 (h);                      // 这里写入IxF2,值为0x32467849
        write_index_header (idx, f);     // 写入quantizer的头部信息
        WRITEVECTOR (idxf->xb);          // 写入粗聚类中心的中心向量，先写size(nlist*d)，8bytes，再写如内容
    }
    ...
}

// WRITEVECTOR 宏定义如下：
#define WRITEVECTOR(vec) {                      \
        size_t size = (vec).size ();            \
        WRITEANDCHECK (&size, 1);               \
        WRITEANDCHECK ((vec).data (), size);    \
}

步骤如下：

1. 写入magic number： IxF2(0x32467849)，占4bytes
2. 写入quantizer的头部信息（注：quantizer内的ntotal其实是index->nlist）
3. 写入粗聚类中心的中心向量，先写入向量大小(nlist*d)，占8bytes，再写入中心向量内容，所占空间为：(nlist * d * 4) bytes

3.1.3 写入direct_map write_direct_map()

Direct map是一种从ID映射回反向列表的方法，支持数组和哈希列表方式，默认为NoMap。

statci void write_direct_map (const DirectMap *dm, IOWrite *f) {
    char maintain_direct_map = (char)dm->type; // for backwards compatibility with bool
    WRITE1 (maintain_direct_map);              // 写入maintain_direct_map值
    WRITEVECTOR (dm->array);                   // 写入dm->array,默认为空，但是会先写入size，占8bytes，全为0
    if (dm->type == DirectMap::Hashtable) {    // 默认不符合条件，不写入
        using idx_t = Index::idx_t;
        std::vector<std::pair<idx_t, idx_t>> v;
        const std::unordered_map<idx_t, idx_t> & map = dm->hashtable;
        v.resize (map.size());
        std::copy(map.bengin(), map.end(), v.begin());
        WRITEVECTOR (v);
    }
}

1. 写入direct map的类型，默认为NoMap，即：0，占1byte
2. 先写入direct map array的size，再写如vector内容。默认情况下size=0，此时没有写入的内容，占8bytes
3. 当Direct Map 类型为Hashtable时还要写入哈希表中的内容，默认情况下不写入。

3.2 写入pq信息 write_ProductQuantizer()

void write_ProductQuantizer (const ProductQuantizer *pq, IOWriter *f) {
    WRITE1 (pq->d);                   // dimension
    WRITE1 (pq->M);                 
    WRITE1 (pq->nbits);
    WRITEVECTOR (pq->centroids);
}

1. 写入dimension，我这里值为64， 占8bytes;
2. 写入M值，我这里值为64，占8bytes，
3. 写入nbits，我这里值为8， 占8bytes
4. 写入centroids，先写size，占8bytes，再写内容，占 (M * ksub *dsub * 4)bytes

3.3 写入倒排序列write_InvertedLists()

void write_InvertedLists (const InvertedLists *ils, IOWriter *f) {
    ...
    // 这里仅保留使用的InvertedLists类型相关的操作
    else if (const auto & ails =
               dynamic_cast<const ArrayInvertedLists *>(ils)) {
        uint32_t h = fourcc ("ilar");                   
        WRITE1 (h);                        // maigc number： 0x72616c69
        WRITE1 (ails->nlist); 
        WRITE1 (ails->code_size);
        // here we store either as a full or a sparse data buffer
        size_t n_non0 = 0;
        
        // 统计有向量的聚类个数
        for (size_t i = 0; i < ails->nlist; i++) {
            if (ails->ids[i].size() > 0)
                n_non0++;
        }
        if (n_non0 > ails->nlist / 2) {                
            uint32_t list_type = fourcc("full");     // 如果有向量的聚类个数超过nlist的一半，那么list_type="full"
            WRITE1 (list_type);
            std::vector<size_t> sizes;               // sizes中依次保存了每个cluster中向量的个数
            for (size_t i = 0; i < ails->nlist; i++) {
                sizes.push_back (ails->ids[i].size());
            }
            WRITEVECTOR (sizes);                     // 写入 sizes，注：先写入sizes大小，在依次写入内容
        } else {
            int list_type = fourcc("sprs"); // sparse   // 如果有向量的聚类个数不足nlist的一半，那么list_type="sprs"
            WRITE1 (list_type);
            std::vector<size_t> sizes;                 // sizes中分别保存了cluster编号以及该cluster中向量个数
            for (size_t i = 0; i < ails->nlist; i++) {
                size_t n = ails->ids[i].size();
                if (n > 0) {
                    sizes.push_back (i);
                    sizes.push_back (n);
                }
            }
            WRITEVECTOR (sizes);                      // 写入 sizes，注：先写入sizes大小，在依次写入内容
        }
        // make a single contiguous data buffer (useful for mmapping)
        // 写入倒排code
        for (size_t i = 0; i < ails->nlist; i++) {
            size_t n = ails->ids[i].size();
            if (n > 0) {
                WRITEANDCHECK (ails->codes[i].data(), n * ails->code_size);
                WRITEANDCHECK (ails->ids[i].data(), n);
            }
        }
    }
    ...
}

1. 写入magic number：“ilar”，4bytes
2. 写入nlist：4000，8bytes
3. 写入code_size：64，8bytes
4. 写入列表类型，这里类型分为两种，如果超过一半的聚类有向量，那么该列表类型为"full"，否则为"sprs"。我用的index例子中为full类型。占4bytes
5. 先写入cluster的个数，即nlist：4000，占8bytes。再依次写入每个cluster中向量的个数，每个数据占8bytes，总共 nlist * 8个。
6. 写入倒排code，遍历每个cluster，先写入cluster内的codes，占 (n * code_size)个bytes，再写入该cluster内每个向量的id，占 (n * 8)bytes。
   总计codes占(nb * code_size)bytes， ids占(nb * 8)bytes。
   
   

4. Index数据结构

综上代码加上我使用的Index实例数据分析，index的数据结构如下图所示：

注：右侧Data部分要加8bytes是因为在保存data之前，要先保存data size，该值占8bytes。

4.1 IndexIVFPQ文件内容

综上，IndexIVFPQ文件依次包括7类信息：

1. Index头部信息
   包括：Index type, dimension, ntotal, dummy *2, is_train标志， metric_type, nlist, nprobe.

2. Quantizer 头部信息
   包括： Quantizer type, dimension, nlist, dummy * 2, id_train标志， metric_type, quantizer size

3. Quantizer Data
   依次存放nlist个粗聚类中心的d维向量
   

注：红色方框内为Index头部信息，绿色方框为Quantizer头部信息，后面接着Quantizer Data

4. ProductorQuantizer头部信息
   包括：Direct map类型，DirectMap array size， code_size，pq->d，pq-M，pq->nbits，centroids size

5. ProductorQuantizer Data
   

注：红色方框内为ProductorQuantizer头部信息，其上为Quantizer Data，下接ProductorQuantizer Data

6. 倒排表头部信息
   包括：倒排表类型，nlist，code_size，list_type，nlist，每个cluster内向量个数
   

7. 倒排表Data
   按照cluster顺序依次存放每个cluster内的codes和对应的id号