vectors: avoid expensive comparisons between numpy ints and Python ints (#10992)

* vectors: avoid expensive comparisons between numpy ints and Python ints

* vectors: avoid failure on lists of ints

* Convert another numpy int to Python
This commit is contained in:
Daniël de Kok 2022-06-29 12:58:31 +02:00 committed by GitHub
parent dd038b536c
commit 0ff14aabce
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -336,10 +336,10 @@ cdef class Vectors:
xp = get_array_module(self.data) xp = get_array_module(self.data)
if key is not None: if key is not None:
key = get_string_id(key) key = get_string_id(key)
return self.key2row.get(key, -1) return self.key2row.get(int(key), -1)
elif keys is not None: elif keys is not None:
keys = [get_string_id(key) for key in keys] keys = [get_string_id(key) for key in keys]
rows = [self.key2row.get(key, -1) for key in keys] rows = [self.key2row.get(int(key), -1) for key in keys]
return xp.asarray(rows, dtype="i") return xp.asarray(rows, dtype="i")
else: else:
row2key = {row: key for key, row in self.key2row.items()} row2key = {row: key for key, row in self.key2row.items()}