k-Nearest Neighbor (kNN) Search

The k-nearest neighbor (kNN) algorithm performs a similarity search on fields of dense_vector type. This type of search, which is more appropriately called "approximate kNN", accepts a vector or embedding as a search term, and finds entries in the index that are close.

In this section you are going to learn how to run a kNN search using the document embeddings created in the previous section.

The knn Query

In the full-text search section of the tutorial you learned about the query option passed to the search() method of the Elasticsearch client. When searching vectors, the knn option is used instead.

Below you can see a new version of the handle_search() function in app.py that runs a kNN search for the query entered by the user in the search form.

@app.post('/')
def handle_search():
    query = request.form.get('query', '')
    filters, parsed_query = extract_filters(query)
    from_ = request.form.get('from_', type=int, default=0)

    results = es.search(
        knn={
            'field': 'embedding',
            'query_vector': es.get_embedding(parsed_query),
            'num_candidates': 50,
            'k': 10,
        },
        size=5,
        from_=from_
    )
    return render_template('index.html', results=results['hits']['hits'],
                           query=query, from_=from_,
                           total=results['hits']['total']['value'])

In this version of the function, the query option was replaced with knn. The size and from_ options for pagination remain the same, and everything else in the function and the index.html template are also the same as before.

The knn search option accepts a number of parameters that configure the search:

• field: the field in the index to search. The field must have a dense_vector type.
• query_vector: the embedding to search for. This should be an embedding generated from the search text.
• num_candidates: the number of candidate documents to consider from each shard. Elasticsearch retrieves this many candidates from each shard, combines them into a single list and then finds the closest "k" to return as results.
• k: the number of results to return. This number has a direct effect on performance, so it should be kept as small as possible. The value passed in this option must be less than num_candidates.

With the settings used in the code above, the 10 best matching results will be returned.

You are welcome to experiment with this new version of the application. Here are a pair of good examples to appreciate how useful this type of search:

• Searching for "holiday", which is the British English equivalent to "vacation" in American English, kNN search returns the document "Vacation Policy" as top result, even though the word holiday itself does not appear in the document.
• Searching for "cats and dogs" or any other term related to pets brings the "Office Pet Policy" document as top result, even though the document summary does not mention any specific pets.

Using Filters in kNN Queries

The search query, as defined in the full-text section of this tutorial, allowed the user to request a specific category to be used, using the syntax category:<category-name> in any place of the search text. The extract_filters() function in app.py is in charge of finding and separating these filter expressions from the search query. In the version of the handle_search() function from the previous section the filters variable is not used, so the category filters are ignored.

Luckily, the knn option also supports filtering. The filter option actually accepts the same type of filters, so the filters can be inserted directly into the knn query, exactly as they are returned by the extract_filters() function:

@app.post('/')
def handle_search():
    query = request.form.get('query', '')
    filters, parsed_query = extract_filters(query)
    from_ = request.form.get('from_', type=int, default=0)

    results = es.search(
        knn={
            'field': 'embedding',
            'query_vector': es.get_embedding(parsed_query),
            'k': 10,
            'num_candidates': 50,
            **filters,
        },
        size=5,
        from_=from_
    )
    return render_template('index.html', results=results['hits']['hits'],
                           query=query, from_=from_,
                           total=results['hits']['total']['value'])

Aggregations also work well in kNN queries, so they can also be added back:

@app.post('/')
def handle_search():
    query = request.form.get('query', '')
    filters, parsed_query = extract_filters(query)
    from_ = request.form.get('from_', type=int, default=0)

    results = es.search(
        knn={
            'field': 'embedding',
            'query_vector': es.get_embedding(parsed_query),
            'k': 10,
            'num_candidates': 50,
            **filters,
        },
        aggs={
            'category-agg': {
                'terms': {
                    'field': 'category.keyword',
                }
            },
            'year-agg': {
                'date_histogram': {
                    'field': 'updated_at',
                    'calendar_interval': 'year',
                    'format': 'yyyy',
                },
            },
        },
        size=5,
        from_=from_
    )
    aggs = {
        'Category': {
            bucket['key']: bucket['doc_count']
            for bucket in results['aggregations']['category-agg']['buckets']
        },
        'Year': {
            bucket['key_as_string']: bucket['doc_count']
            for bucket in results['aggregations']['year-agg']['buckets']
            if bucket['doc_count'] > 0
        },
    }
    return render_template('index.html', results=results['hits']['hits'],
                           query=query, from_=from_,
                           total=results['hits']['total']['value'], aggs=aggs)

This version of the handle_search() function has the same functionality as the full-text search version, implemented using vector search instead of keyword-based search.

In the next section, you'll learn how to combine results from these two different search methods.