Skip to content

Commit 61e0a57

Browse files
authored
fix: fix offset factory (#208)
1 parent f6fdc9d commit 61e0a57

File tree

3 files changed

+56
-6
lines changed

3 files changed

+56
-6
lines changed

annlite/storage/table.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -306,9 +306,14 @@ def _offset_factory(_, record):
306306
self._conn.row_factory = _offset_factory
307307

308308
cursor = self._conn.cursor()
309-
offsets = cursor.execute(sql, params).fetchall()
310-
self._conn.row_factory = None
311-
return offsets if offsets else []
309+
310+
try:
311+
offsets = cursor.execute(sql, params).fetchall()
312+
self._conn.row_factory = None
313+
return offsets if offsets else []
314+
except Exception as e:
315+
self._conn.row_factory = None
316+
raise e
312317

313318
def delete(self, doc_ids: List[str]):
314319
"""Delete the docs
@@ -361,7 +366,6 @@ def count(self, where_clause: str = '', where_params: Tuple = ()):
361366
# # EXPLAIN SQL query
362367
# for row in self._conn.execute('EXPLAIN QUERY PLAN ' + sql, params):
363368
# print(row)
364-
365369
return self._conn.execute(sql, params).fetchone()[0]
366370
else:
367371
sql = f'SELECT MAX(_id) from {self.name} LIMIT 1;'

tests/test_filter.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,3 +209,49 @@ def test_filter_with_limit_offset(tmpfile, limit, offset, order_by, ascending):
209209
assert m.tags[order_by] <= matches[i + 1].tags[order_by]
210210
else:
211211
assert m.tags[order_by] >= matches[i + 1].tags[order_by]
212+
213+
214+
@pytest.mark.parametrize('limit', [1, 5])
215+
def test_filter_with_wrong_columns(tmpfile, limit):
216+
N = 100
217+
D = 128
218+
219+
index = AnnLite(
220+
D,
221+
columns=[('price', float)],
222+
data_path=tmpfile,
223+
)
224+
X = np.random.random((N, D)).astype(np.float32)
225+
226+
docs = DocumentArray(
227+
[
228+
Document(id=f'{i}', embedding=X[i], tags={'price': random.random()})
229+
for i in range(N)
230+
]
231+
)
232+
233+
index.index(docs)
234+
235+
matches = index.filter(
236+
filter={'price': {'$lte': 50}},
237+
limit=limit,
238+
include_metadata=True,
239+
)
240+
241+
assert len(matches) == limit
242+
243+
import sqlite3
244+
245+
with pytest.raises(sqlite3.OperationalError):
246+
matches = index.filter(
247+
filter={'price_': {'$lte': 50}},
248+
include_metadata=True,
249+
)
250+
251+
matches = index.filter(
252+
filter={'price': {'$lte': 50}},
253+
limit=limit,
254+
include_metadata=True,
255+
)
256+
257+
assert len(matches) == limit

tests/test_index.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -306,10 +306,10 @@ def test_remote_backup_restore(tmpdir):
306306
index.index(docs)
307307

308308
tmpname = uuid.uuid4().hex
309-
index.backup(target_name='test', token=token)
309+
index.backup(target_name='test_remote_backup_restore', token=token)
310310

311311
index = AnnLite(n_dim=D, data_path=tmpdir / 'workspace' / '0')
312-
index.restore(source_name='test', token=token)
312+
index.restore(source_name='test_remote_backup_restore', token=token)
313313

314314
delete_artifact(tmpname)
315315
status = index.stat

0 commit comments

Comments
 (0)