Skip to content

Commit

Permalink
vectorsearch-cities example: use explicit imports and fix formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
greenrobot committed May 28, 2024
1 parent ca79c96 commit cf8afb8
Showing 1 changed file with 24 additions and 22 deletions.
46 changes: 24 additions & 22 deletions example/vectorsearch-cities/main.py
Original file line number Diff line number Diff line change
@@ -1,52 +1,54 @@
from cmd import Cmd
from objectbox import *
from objectbox import Entity, Float32Vector, HnswIndex, Id, Store, String
import time
import csv
import os


@Entity()
class City:
id = Id()
name = String()
location = Float32Vector(index=HnswIndex(
dimensions=2,
distance_type=VectorDistanceType.EUCLIDEAN
))
location = Float32Vector(index=HnswIndex(dimensions=2))


def list_cities(cities):
print("{:3s} {:25s} {:>9s} {:>9s}".format("ID", "Name", "Latitude", "Longitude"))
for city in cities:
print("{:3d} {:25s} {:>9.2f} {:>9.2f}".format(
city.id, city.name, city.location[0], city.location[1]))
city.id, city.name, city.location[0], city.location[1]))


def list_cities_with_scores(city_score_tuples):
print("{:3s} {:25s} {:>9s} {:>9s} {:>5s}".format("ID", "Name", "Latitude", "Longitude", "Score"))
for (city,score) in city_score_tuples:
for (city, score) in city_score_tuples:
print("{:3d} {:25s} {:>9.2f} {:>9.2f} {:>5.2f}".format(
city.id, city.name, city.location[0], city.location[1], score))
city.id, city.name, city.location[0], city.location[1], score))


class VectorSearchCitiesCmd(Cmd):
prompt = "> "

def __init__(self, *args):
Cmd.__init__(self, *args)
dbdir = "cities-db"
new_db = not os.path.exists(dbdir)
self._store = Store(directory=dbdir)
self._box = self._store.box(City)
if new_db:
if new_db:
with open(os.path.join(os.path.dirname(__file__), 'cities.csv')) as f:
r = csv.reader(f)
r = csv.reader(f)
cities = []
for row in r:
city = City()
city.name = row[0]
city.location = [ row[1], row[2] ]
city.location = [row[1], row[2]]
cities.append(city)
self._box.put(*cities)

def do_ls(self, name: str = ""):
"""list all cities or starting with <prefix>\nusage: ls [<prefix>]"""
qb = self._box.query( City.name.starts_with(name) )
qb = self._box.query(City.name.starts_with(name))
query = qb.build()
list_cities(query.find())

Expand All @@ -62,38 +64,38 @@ def do_city_neighbors(self, args: str):
num = 5
if len(args) == 2:
num = int(args[1])
qb = self._box.query( City.name.equals(city) )
qb = self._box.query(City.name.equals(city))
query = qb.build()
cities = query.find()
if len(cities) == 1:
location = cities[0].location
# +1 for the city
qb = self._box.query(
City.location.nearest_neighbor(location, num+1) & City.name.not_equals(city)
City.location.nearest_neighbor(location, num + 1) & City.name.not_equals(city)
)
neighbors = qb.build().find_with_scores()
list_cities_with_scores(neighbors)
else:
print(f"no city found named '{city}'")
except ValueError:
except ValueError:
print("usage: city_neighbors <name>[,<num: default 5>]")

def do_neighbors(self, args):
"""find <num> neighbors next to geo-coord <lat> <long>.\nusage: neighbors <num>,<latitude>,<longitude>"""
try:
args = args.split(',')
if len(args) != 3:
raise ValueError()
num = int(args[0])
geocoord = [ float(args[1]), float(args[2]) ]
geocoord = [float(args[1]), float(args[2])]
qb = self._box.query(
City.location.nearest_neighbor(geocoord, num)
)
neighbors = qb.build().find_with_scores()
list_cities_with_scores(neighbors)
except ValueError:
except ValueError:
print("usage: neighbors <num>,<latitude>,<longitude>")

def do_add(self, args: str):
"""add new location\nusage: add <name>,<lat>,<long>"""
try:
Expand All @@ -105,11 +107,11 @@ def do_add(self, args: str):
long = float(args[2])
city = City()
city.name = name
city.location = [lat,long]
city.location = [lat, long]
self._box.put(city)
except ValueError:
except ValueError:
print("usage: add <name>,<latitude>,<longitude>")

def do_exit(self, _):
"""close the program"""
raise SystemExit()
Expand Down

0 comments on commit cf8afb8

Please sign in to comment.