{"id":111190,"date":"2024-05-29T15:54:32","date_gmt":"2024-05-29T15:54:32","guid":{"rendered":"https:\/\/foojay.io\/?p=111190"},"modified":"2024-12-30T10:55:24","modified_gmt":"2024-12-30T10:55:24","slug":"indexing-all-of-wikipedia-on-a-laptop","status":"publish","type":"post","link":"https:\/\/foojay.io\/today\/indexing-all-of-wikipedia-on-a-laptop\/","title":{"rendered":"Indexing all of Wikipedia, on a laptop"},"content":{"rendered":"\n
\n
\n \n \n \n \n \n \n \n <\/svg>\n Table of Contents\n \n \n <\/svg>\n <\/div>\n
Compression parameters<\/a><\/span>GraphIndexBuilder<\/a><\/span>Chronicle Map and RowData<\/a><\/span>Ingesting the data<\/a><\/span>Loading the index (after construction)<\/a><\/span>Performing a search<\/a><\/span><\/div><\/div>\n

\n
\n
Sponsored Content<\/div>

AI4J - The Intelligent Java Conference<\/h2>\n

\n This exclusive virtual event brings together leading AI innovators and renowned Java Champions to unpack what\u2019s changing today, what\u2019s coming next, and how enterprise Java teams can stay ahead. April 14, 2026 @ 9am PDT | 12pm PDT. <\/p>\n\t \n Register Now \n \n \n <\/svg>\n <\/a>\n <\/div>\n

\n \"\" <\/div>\n<\/div>\n<\/p>\n\n\n\n

In November, Cohere released a dataset containing all of Wikipedia<\/a>, chunked and embedded to vectors with their multilingual-v3 model<\/a>. <\/p>\n\n\n\n

Computing this many embeddings yourself would cost in the neighborhood of $5000, so the public release of this dataset makes creating a semantic, vector-based index<\/a> of Wikipedia practical for an individual for the first time.<\/p>\n\n\n\n

Here’s what we’re building: <\/p>\n\n\n\n

\"\"<\/figure>\n\n\n\n

You can try searching the completed index on a public demo instance here<\/a>.<\/p>\n\n\n\n

Why this is hard<\/h1>\n\n\n\n

Sure, the dataset is big (180GB for the English corpus), but that’s not the obstacle per se.  We’ve been able to build full-text indexes on larger datasets for a long time.<\/p>\n\n\n\n

The obstacle is that until now, off-the-shelf vector databases could not index a dataset larger than memory, because both the full-resolution vectors and the index (edge list) needed to be kept in memory during index construction.  Larger datasets could be split into segments<\/a>, but this means that at query time they need to search each segment separately, then combine the results, turning an O(log N) search per segment into O(N) overall.  (In their latest release, Lucene attempts to mitigate this by processing segments in parallel with multiple threads<\/a>, but obviously (1) this only gives you a constant factor of improvement before you run out of CPU cores and (2) this does not improve throughput.)<\/p>\n\n\n\n

Specifically, if you’re indexing 1536-dimension vectors (the size of ada002 or openai-v3-small), then you can fit about 5M vectors and their associated edge lists in a 32GB index construction RAM budget.<\/p>\n\n\n\n

\"\"<\/figure>\n\n\n\n

JVector<\/a>, the library that powers DataStax Astra<\/a> vector search, now supports indexing larger-than-memory datasets by performing construction-related searches with compressed vectors.  This means that the edge lists need to fit in memory, but the uncompressed vectors do not, which gives us enough headroom to index Wikipedia-en on a laptop.<\/p>\n\n\n\n

Requirements<\/h1>\n\n\n\n
    \n
  1. Linux or MacOS.  It will not work on Windows because ChronicleMap, which we are going to use for the non-vector data, is limited to a 4GB size there.  (If you are interested enough, you could shard the Map by vector id to keep each shard under 4GB and still have O(1) lookup times.)<\/li>\n\n\n\n
  2. About 180GB of free space for the dataset, and 90GB for the completed index.<\/li>\n\n\n\n
  3. Enough RAM to run a JVM with 36GB of heap space during construction (~28GB for the index, 8GB for GC headroom).<\/li>\n\n\n\n
  4. Disable swap before building the index.  Linux will aggressively try to cache the index being constructed to the point of swapping out parts of the JVM heap, which is obviously counterproductive.  In my test, building with swap enabled was almost twice as slow as with it off.<\/li>\n<\/ol>\n\n\n\n

    Building and searching the index<\/h1>\n\n\n\n
      \n
    1. Check out the project:
      $ git clone       https:\/\/github.com\/jbellis\/coherepedia-jvector
      <\/a>$ cd coherepedia-jvector<\/li>\n\n\n\n
    2. Edit config.properties<\/em> to set the locations for the dataset and the index. <\/li>\n\n\n\n
    3. Run pip install datasets<\/em>.  (Setting up a venv<\/a> or conda environment first is recommended but not strictly necessary.)<\/li>\n\n\n\n
    4. Run python download.py.  <\/em>This downloads the 180 GB dataset to the location you configured.  For me that took about half an hour.<\/li>\n\n\n\n
    5. Run .\/mvnw compile exec:exec@buildindex.<\/em>  This took about 5 and a half hours on my machine (with an i9-12900 CPU).<\/li>\n\n\n\n
    6. Run .\/mvnw compile exec:exec@serve <\/em>and open a browser to http:\/\/localhost:4567<\/a>.  Search away!<\/li>\n<\/ol>\n\n\n\n

      How it works<\/h1>\n\n\n\n

      We’re using JVector<\/a> for the vector index and Chronicle Map<\/a> for the article data.  There are several<\/a> things<\/a> I don’t love about Chronicle Map, but nothing else touches it for simple disk-based key\/value performance.<\/p>\n\n\n\n

      The full source of the index construction class is here<\/a>.  I’ll explain it next in pieces.<\/p>\n\n\n\n

      Compression parameters<\/h2>\n\n\n\n

      JVector is based on the DiskANN<\/a> vector index design, which performs an initial search using vectors compressed lossily with product quantization (PQ)<\/a> in memory, then reranks the results using high-resolution vectors from disk.  However, while DiskANN stores full, uncompressed vectors to perform reranking, JVector is able to improve on that using Locally-Adaptive Quantization (LVQ)<\/a> compression.<\/p>\n\n\n\n

      To set this up, we’ll first load some vectors into a RandomAccessVectorValues (RAVV) instance.  RAVV is a JVector interface for a vector container; it could be List or Map based, in-memory or on-disk.  In this case we’ll use a simple List-backed RAVV.  We’ll compute the parameters for both compressions (kmeans clustering for PQ, global mean for LVQ) from a single shard of the dataset.  At about 110k rows, this is enough data to have a statistically valid sample.<\/p>\n\n\n\n

      \"\"<\/figure>\n\n\n\n

      Next, we compute the PQ compression codebook; we’re compressing the vectors by a factor of 64, because the Cohere v3 embeddings can be PQ-compressed that much without losing accuracy, after reranking.  Binary Quantization only gives us 32x compression and is less accurate<\/a>.<\/p>\n\n\n\n

      \"\"<\/figure>\n\n\n\n

      Finally, we need to set up LVQ.  LVQ gives us 4x compression while losing no measurable accuracy over the full uncompressed vectors, resulting in both a smaller footprint on disk and faster searches.  (I thank the vector search team at Intel Research for pointing this out to us.)<\/p>\n\n\n\n

      \"\"<\/figure>\n\n\n\n

      GraphIndexBuilder<\/h2>\n\n\n\n

      Next, we need to instantiate and configure our GraphIndexBuilder.<\/p>\n\n\n\n

      \"\"<\/figure>\n\n\n\n

      This instantiates a JVector GraphIndexBuilder and connects it to an OnDiskGraphIndexWriter, and tells it to use the PQ-compressed vectors list (which starts empty and will grow as we add vectors to the index) during construction (in the BuildScoreProvider).<\/p>\n\n\n\n

      Chronicle Map and RowData<\/h2>\n\n\n\n

      We’ll store article contents in RowData records.  This content is what has been encoded as the corresponding vector in the dataset, and is what we want to return to the user in our search results.<\/p>\n\n\n\n

      \"\"<\/figure>\n\n\n\n

      To turn the vector index’s search results (a list of integer vector ids) into RowData, we store the RowData in a Map keyed by the vector id.  This will be a lot of data, so we use ChronicleMap<\/a> to store this on disk with a minimal in-memory footprint.<\/p>\n\n\n\n

      \"\"<\/figure>\n\n\n\n

      We need to tell ChronicleMap how large it’s going to be, both in entry count and entry size.  Undersizing these will cause it to crash (my primary complaint<\/a> about ChronicleMap), so we deliberately use a high estimate.<\/p>\n\n\n\n

      We do not<\/em> need to explicitly tell ChronicleMap how to read and write RowData objects, instead we just have RowData implement Serializable.  While ChronicleMap supports custom de\/serialize code, it’s perfectly happy to use simple out-of-the-box serialization and since profiling shows that’s not a bottleneck for us we’ll just leave it at that.<\/p>\n\n\n\n

      Ingesting the data<\/h2>\n\n\n\n

      We use Java’s parallel Streams to process the shards in parallel.  For each row in each shard, we<\/p>\n\n\n\n

        \n
      1. Add it to pqVectorsList<\/em><\/li>\n\n\n\n
      2. Call writer.writeInline<\/em> to add the LVQ-compressed vector to disk<\/li>\n\n\n\n
      3. Call builder.addGraphNode <\/em>– order is important because both (1) and (2) are used when we call addGraphNode<\/li>\n\n\n\n
      4. Call contentMap.put<\/em> with the article chunk data.<\/li>\n<\/ol>\n\n\n\n
        \"\"<\/figure>\n\n\n\n
        \"\"<\/figure>\n\n\n\n

        You can look at the full source<\/a> if you’re curious about forEachRow<\/em>, it’s just standard “pull data out of Arrow” stuff.<\/p>\n\n\n\n

        When the build completes, you should see files like this:<\/p>\n\n\n\n

        $ ls -lh ~\/coherepedia<\/p>\n\n\n\n

        -rw-rw-r-- 1 jonathan jonathan  48G May 20 15:53 coherepedia.ann<\/p>\n\n\n\n

        -rw-rw-r-- 1 jonathan jonathan  36G May 20 18:05 coherepedia.map<\/p>\n\n\n\n

        -rw-rw-r-- 1 jonathan jonathan 2.5G May 20 15:53 coherepedia.pqv<\/p>\n\n\n\n

        -rw-rw-r-- 1 jonathan jonathan 4.1K May 17 23:04 coherepedia.lvq<\/p>\n\n\n\n

        -rw-rw-r-- 1 jonathan jonathan 1.1M May 17 23:04 coherepedia.pq<\/p>\n\n\n\n

        These are respectively<\/p>\n\n\n\n