forked from lancedb/lancedb
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_distance_range.py
More file actions
62 lines (48 loc) · 1.83 KB
/
test_distance_range.py
File metadata and controls
62 lines (48 loc) · 1.83 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import shutil
import pytest
# --8<-- [start:imports]
import lancedb
import numpy as np
# --8<-- [end:imports]
shutil.rmtree("data/distance_range_demo", ignore_errors=True)
def test_binary_vector():
# --8<-- [start:sync_distance_range]
db = lancedb.connect("data/distance_range_demo")
data = [
{
"id": i,
"vector": np.random.random(256),
}
for i in range(1024)
]
tbl = db.create_table("my_table", data=data)
query = np.random.random(256)
# Search for the vectors within the range of [0.1, 0.5)
tbl.search(query).distance_range(0.1, 0.5).to_arrow()
# Search for the vectors with the distance less than 0.5
tbl.search(query).distance_range(upper_bound=0.5).to_arrow()
# Search for the vectors with the distance greater or equal to 0.1
tbl.search(query).distance_range(lower_bound=0.1).to_arrow()
# --8<-- [end:sync_distance_range]
db.drop_table("my_table")
@pytest.mark.asyncio
async def test_binary_vector_async():
# --8<-- [start:async_distance_range]
db = await lancedb.connect_async("data/distance_range_demo")
data = [
{
"id": i,
"vector": np.random.random(256),
}
for i in range(1024)
]
tbl = await db.create_table("my_table", data=data)
query = np.random.random(256)
# Search for the vectors within the range of [0.1, 0.5)
await tbl.query().nearest_to(query).distance_range(0.1, 0.5).to_arrow()
# Search for the vectors with the distance less than 0.5
await tbl.query().nearest_to(query).distance_range(upper_bound=0.5).to_arrow()
# Search for the vectors with the distance greater or equal to 0.1
await tbl.query().nearest_to(query).distance_range(lower_bound=0.1).to_arrow()
# --8<-- [end:async_distance_range]
await db.drop_table("my_table")