Introduction
We were on our second search backend in a year but it was time to move again.
All the main features of Legora are based on document search under the hood. Our user document collection was at 2B+ search chunks and growing like crazy. Performance and reliability was struggling to keep up.
This is about how we found a scalable solution: the winning move was to exploit the specific properties of our dataset.
Our Search Problem
Work in Legora happens in projects. These often map 1:1 with legal matters. Think of a matter as one case for one client. The project is where we store all legal documents that a user has uploaded or connected to the matter.
Within a project, the Legora agent needs powerful search tools to find its way around the matter, to answer user questions, to run batch analysis, and to perform drafting tasks.
We use a combination of keyword search and semantic search. With this, we can already see a couple of specifics about our search problem:
We primarily need to search within a project. Organization level search is a thing too, but less important and we can solve that separately.
Projects are never huge. Organizations (i.e. law firms) can be huge, but individual matters are not. 10k-100k documents is a large matter.
Projects go out of date. A matter will go through an active phase and then become mostly dormant. On any given day, we can expect some small fraction of projects seeing most of the read+write traffic, and most others not accessed at all.
Next, since we are working on legal documents here, we have some key requirements regarding data protection.
Data can only travel to a small number of 3rd parties mentioned in customer agreements.
In our highest isolation tiers, data for a firm must use customer managed encryption keys and separate tenants: isolated storage and database deployments.
Act 1: Elasticsearch
We started on self-hosted Elasticsearch. This served us well for some time. But as we needed to support stricter data protection tiers, we had to decide whether we wanted to run an ES cluster per tenant. We went a different route, namely:
Act 2: PostgreSQL + pgvector + DiskANN
Moving the search index to postgreSQL was a perfect solution for data protection. We already had a fully isolated tenant setup for postgres. PG has extensions for keyword and vector search.
Performance was excellent initially, life was good.
But at the time of this story, we found ourselves struggling with latency on document inserts and reads, and the index was growing quickly, turning this database into our main operations challenge.
Time for a deep dive to diagnose..
Anatomy of a Search Index
Our main search needs are keywords (BM25) and semantic (vector). Keyword search is based on an inverted index, think of it as word→[documents] lists. Searching for a word starts with looking up the documents containing it, and then a hundred additional tricks to filter and rank.
Search vectors are embeddings of text chunks. A chunk can be a text paragraph or section. So we split up all or docs and each chunk goes into an embedding model which returns a list of floats, aka a vector. Crucially, the model is trained so that chunks that are semantically "similar" end up with similar vectors: vectors that are close in space.
This means we can take an agent query, embed that as a vector, and use vector math to decide which other chunks in our project are close semantically. The catch is how to do this nearest neighbors search in a large set of vectors. Fortunately, approximate nearest neighbor (ANN) is a well researched problem. And this is the foundation of semantic search: vector embeds and ANN.
One final thing to note about this index: it's surprisingly large! The vectors are high dimensional, we use 1536-float vectors, So each chunk ends up with a 6kB vector. This is typically more than the actual chunk, i.e. the text content.
The situation is similar for the keyword index: it's on the same order as the source text.
Neighbors
Back to the issues at hand. The latency problems turned out to be all about the vector index. Both search and inserts were starting to see latency spikes, up to 20s for inserts. And significant memory pressure.
We were using pgvector and the DiskANN extension. Aptly named, DiskANN is a family of algorithms to run ANN efficiently with the index stored on disk. Very useful since the index is so large. DiskANN maintains a graph structure to find nearest neighbors with only a few hops, i.e. disk accesses.
Recall that our projects are completely separate from a search point of view. And that projects are something like 1-10M chunks tops. This should be light work for DiskANN. Either using per project indexes, or the Filtered variation of the algorithm.
But, in the postgres extension, these options are not available. PG needs, in practice, a fixed partitioning, where each partition ends up storing a number of projects.
For a nearest neighbor search, this means first finding neighbors, and then filtering on project. We may need to search through a large number of candidates to find enough that will pass the filter. This way, unrelated projects stack up to increase latency for a partition.
Growing Pains
Zooming out to the overall scaling problem. Besides the vector search, there was also the looming issue of storing an index the size of our user document collection in postgres. A collection that was growing very quickly.
We were not using the key structural advantages of our data: that 1) search is project scoped, and 2) that most projects become dormant.
Instead, unrelated projects interfered in ANN search to impact performance. And old projects were staying ready in the index when they might as well be moved out to colder storage, growing index size and memory pressure.
At the current growth rate, even if we could tune e.g. the vector search partitioning, we reckoned we only had a few months of runway before the database would become very challenging to operate. We had about 2B vectors at the time, and we wanted something that could get us to 10B but also onward towards 100B.
So we went searching for a new solution.
Act 3: Turbopuffer
We arrived at a shortlist of candidates, and one standout option was Turbopuffer
Turbopuffer uses namespaces as search partitions. Namespaces are low overhead, it's cheap to have a large number of them. This was a perfect fit for our projects in terms of search partitioning.
Turbopuffer separates compute from storage. Namespaces sit on blob until actually queried, when they are quickly loaded onto an NVMe powered search instance. This means our dormant project would only consume blob storage.
Turbopuffer supports BYOC: we can run search instances in a k8s namespace we own, and of course store the index data in our blob containers. Namespaces are stored separately on blob, and are collected in Turbopuffer organizations which we map 1:1 to our Legora tenants. Each such org uses a separate storage account and key config. This creates the separation we need to meet our data protection requirements.
We made a decision fairly quickly. After that we worked with the Turbopuffer team to iron out some last details around the temporary storage on search instances, and then of course the data migration with double writes, dark-loading, result monitoring. But that’s for another post.
Results
We saw p99 latency drop from 20s to 2s, and it’s held steady over time.
Overall we have had very few scaling issues with Turbopuffer. We do get hotspots on reads and writes: for large batch analysis jobs or for certain import workloads, we can get rate limited on reads or index updates in that project.
But these are localized issues and we can deal with them by simply slowing down.
The Lesson
Our main takeaway from these three search architectures is an old one: understand the design decisions and tradeoffs of the system you pick or build, and make sure they can make use of the specific properties of your dataset. That’s when you will be able to get maximum scale and performance.
You can also read turbopuffer‘s write-up of this migration.
