Similarity / vector search

The Xata vector type can be combined with the vectorSearch API to perform similarity search on your data based on embeddings.

An embedding is a vector of floating point numbers that represents the original data (text, image, audio, etc.). Embeddings are typically produced by a machine learning model. You can store the embeddings in Xata in a column of type vector and use the vectorSearch API call to find the nearest neighbors of a given embedding. This is useful in a number of use cases related to machine learning, such as: semantic search, recommendations, clustering, classification, and others.

Vector search runs in the Search store which is eventually consistent. The Search store is enabled by default and can be disabled in the Database Settings from the Web UI. The vectorSearch API and features described in this page require the Search store to be enabled. It works at the table level and has the following format:

const results = await xata.db.Docs.vectorSearch(
  "column name",
  [<query array of floats>],
  {
    similarityFunction: "<space function>",
    size: <number of results to return>,
    filter: "<filter expression>",
  }
);

results = xata.data().vector_search("Docs", {
  "queryVector": [0.1, 0.2],  # array of floats
  "column": "embeddings",     # column name,
  "similarityFunction": "l1", # space function
  "size": 5,                  # number of results to return
  "filter": {},               # filter expression
})

searchClient, _ := xata.NewSearchAndFilterClient()
results, _ := searchClient.VectorSearch(context.TODO(), xata.VectorSearchTableRequest{
  TableName: "Docs",
  Payload: xata.VectorSearchTableRequestPayload{
    QueryVector:        []float64{0.1, 0.2},      // array of floats
    Column:             "embeddings",             // column name
    SimilarityFunction: xata.String("l1"),        // space function
    Size:               xata.Int(5),              // number of results to return
    Filter:             &xata.FilterExpression{}, // filter expression
  },
})

// POST https://{workspace}.{region}.xata.sh/db/{db}:{branch}/tables/{table}/vectorSearch
{
    "queryVector": [/*<array of floats>*/],
    "column": /*<column name>*/,
    "similarityFunction": /*<space function>*/,
    "size": /*<number of results to return>*/,
    "filter": /*<filter expression>*/
}

With the following required parameters:

• column - the name of the column in which to search. It must be of type vector.
• queryVector - the vector to search for. It needs to have the same dimension as the vector column.

Additionally you can customize the results with these optional parameters:

• similarityFunction - the function used to measure the distance between two points. It can be one of the following: cosineSimilarity, l1, l2. The default is cosineSimilarity for which you can expect a score between 0 and 2. Because cosineSimilarity returns a value between -1 and 1, but relevance scores must not be below 0, we add 1 for the final score.
• size - the number of results to return. The default is 10.
• filter - a filter expression that can be used to filter before applying the search (pre-filter). See Filtering for more details.

Here is a simple example:

const results = await xata.db.Docs.vectorSearch('embeddings', [0.1, 0.2, 0.3], { size: 5 });

results = xata.data().vector_search("Docs", {
  "queryVector": [0.1, 0.2, 0.3],
  "column": "embeddings",
  "size": 5
})

searchClient, _ := xata.NewSearchAndFilterClient()
results, _ := searchClient.VectorSearch(context.TODO(), xata.VectorSearchTableRequest{
  TableName: "Docs",
  Payload: xata.VectorSearchTableRequestPayload{
    QueryVector:        []float64{0.1, 0.2, 0.3}, // array of floats
    Column:             "embeddings",             // column name
    Size:               xata.Int(5),              // number of results to return
  },
})

// POST https://{workspace}.{region}.xata.sh/db/{db}:{branch}/tables/{table}/vectorSearch
{
  "queryVector": [0.1, 0.2, 0.3],
  "column": "embeddings",
  "size": 5
}

The results of a vector search query return a set of records from the specified table, that are most similar to the provided query vector. The response includes the following:

• totalCount: The total number of records matching the query vector.
• records: An array of the most relevant records. Each item in records includes the id, and score within its xata column.
• id: A unique identifier for the record.
• score: A similarity score (usually between 0 and 1) indicating how closely the record's vector matches the query vector.

Note: Currently this endpoint executes exact vector search. We are planning to add approximate vector search in the future, which scales better to large datasets.

In this example, we will use the OpenAI embeddings API to generate embeddings for a few sentences and then use the /vectorSearch endpoint to find the nearest neighbors of a given embedding.

We will assume the schema that we use has a content column and an embeddings column of type vector. The embeddings column needs to have a dimension of 1536, because this is what the OpenAI embeddings API returns.

The following code snippet loads the senteces into Xata together with the embeddings generated by the OpenAI API:

import { getXataClient } from './xata.ts';
import { Configuration, OpenAIApi } from 'openai';

const docs: string[] = [
  'This is a story about a quick brown fox that jumps over the lazy dog.',
  `Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut
  labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris
  nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit
   esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident,
    sunt in culpa qui officia deserunt mollit anim id est laborum.`,
  'The lazy dog was not happy about the quick brown fox jumping over him.'
];

const openAIConfig = new Configuration({
  apiKey: `YOUR OPENAI API KEY`
});
const openAI = new OpenAIApi(openAIConfig);

const xata = getXataClient();

for (const doc of docs) {
  const resp = await openAI.createEmbedding({
    input: doc,
    model: 'text-embedding-ada-002'
  });
  const [{ embedding }] = resp.data.data;

  xata.db.docs.create({
    contents: doc,
    embedding
  });
}

The following code snippet searches for the nearest neighbors of a given phrase. It first generates the embedding for the input phrase and then uses vector search to find the nearest neighbors:

import { getXataClient } from './xata.ts';
import { Configuration, OpenAIApi } from 'openai';

const question = 'The quick brown fox jumps over the lazy dog.';

const openAIConfig = new Configuration({
  apiKey: `YOUR OPENAI API KEY`
});
const openAI = new OpenAIApi(openAIConfig);
const resp = await openAI.createEmbedding({
  input: question,
  model: 'text-embedding-ada-002'
});
const [{ embedding }] = resp.data.data;

const xata = getXataClient();

const records = await xata.db.docs.vectorSearch('embedding', embedding);

for (const record of records) {
  console.log(record.contents, record.xata);
}