opensearch-pyd/samples/knn/knn-efficient-filter.py

#!/usr/bin/env python

# -*- coding: utf-8 -*-
# SPDX-License-Identifier: Apache-2.0
#
# The OpenSearch Contributors require contributions made to
# this file be licensed under the Apache-2.0 license or a
# compatible open source license.
#
# Modifications Copyright OpenSearch Contributors. See
# GitHub history for details.


import os

from opensearchpy import OpenSearch, helpers

# connect to an instance of OpenSearch

host = os.getenv("HOST", default="localhost")
port = int(os.getenv("PORT", 9200))
auth = (os.getenv("USERNAME", "admin"), os.getenv("PASSWORD", "admin"))

client = OpenSearch(
    hosts=[{"host": host, "port": port}],
    http_auth=auth,
    use_ssl=True,
    verify_certs=False,
    ssl_show_warn=False,
)

# check whether an index exists
index_name = "hotels-index"

if not client.indices.exists(index_name):
    client.indices.create(
        index_name,
        body={
            "settings": {
                "index.knn": True,
                "knn.algo_param.ef_search": 100,
                "number_of_shards": 1,
                "number_of_replicas": 0,
            },
            "mappings": {
                "properties": {
                    "location": {
                        "type": "knn_vector",
                        "dimension": 2,
                        "method": {
                            "name": "hnsw",
                            "space_type": "l2",
                            "engine": "lucene",
                            "parameters": {"ef_construction": 100, "m": 16},
                        },
                    },
                }
            },
        },
    )

# index data
vectors = [
    {
        "_index": "hotels-index",
        "_id": "1",
        "location": [5.2, 4.4],
        "parking": "true",
        "rating": 5,
    },
    {
        "_index": "hotels-index",
        "_id": "2",
        "location": [5.2, 3.9],
        "parking": "false",
        "rating": 4,
    },
    {
        "_index": "hotels-index",
        "_id": "3",
        "location": [4.9, 3.4],
        "parking": "true",
        "rating": 9,
    },
    {
        "_index": "hotels-index",
        "_id": "4",
        "location": [4.2, 4.6],
        "parking": "false",
        "rating": 6,
    },
    {
        "_index": "hotels-index",
        "_id": "5",
        "location": [3.3, 4.5],
        "parking": "true",
        "rating": 8,
    },
    {
        "_index": "hotels-index",
        "_id": "6",
        "location": [6.4, 3.4],
        "parking": "true",
        "rating": 9,
    },
    {
        "_index": "hotels-index",
        "_id": "7",
        "location": [4.2, 6.2],
        "parking": "true",
        "rating": 5,
    },
    {
        "_index": "hotels-index",
        "_id": "8",
        "location": [2.4, 4.0],
        "parking": "true",
        "rating": 8,
    },
    {
        "_index": "hotels-index",
        "_id": "9",
        "location": [1.4, 3.2],
        "parking": "false",
        "rating": 5,
    },
    {
        "_index": "hotels-index",
        "_id": "10",
        "location": [7.0, 9.9],
        "parking": "true",
        "rating": 9,
    },
    {
        "_index": "hotels-index",
        "_id": "11",
        "location": [3.0, 2.3],
        "parking": "false",
        "rating": 6,
    },
    {
        "_index": "hotels-index",
        "_id": "12",
        "location": [5.0, 1.0],
        "parking": "true",
        "rating": 3,
    },
]

helpers.bulk(client, vectors)

client.indices.refresh(index=index_name)

# search
search_query = {
    "size": 3,
    "query": {
        "knn": {
            "location": {
                "vector": [5, 4],
                "k": 3,
                "filter": {
                    "bool": {
                        "must": [
                            {"range": {"rating": {"gte": 8, "lte": 10}}},
                            {"term": {"parking": "true"}},
                        ]
                    }
                },
            }
        }
    },
}

results = client.search(index=index_name, body=search_query)
for hit in results["hits"]["hits"]:
    print(hit)

# delete index
client.indices.delete(index=index_name)