forked from pgvector/pgvector-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbit.py
More file actions
38 lines (28 loc) · 1.09 KB
/
bit.py
File metadata and controls
38 lines (28 loc) · 1.09 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
from sqlalchemy.dialects.postgresql.base import ischema_names
from sqlalchemy.types import UserDefinedType, Float
class BIT(UserDefinedType):
cache_ok = True
def __init__(self, length=None):
super(UserDefinedType, self).__init__()
self.length = length
def get_col_spec(self, **kw):
if self.length is None:
return 'BIT'
return 'BIT(%d)' % self.length
def bind_processor(self, dialect):
if dialect.__class__.__name__ == 'PGDialect_asyncpg':
import asyncpg
def process(value):
if isinstance(value, str):
return asyncpg.BitString(value)
return value
return process
else:
return super().bind_processor(dialect)
class comparator_factory(UserDefinedType.Comparator):
def hamming_distance(self, other):
return self.op('<~>', return_type=Float)(other)
def jaccard_distance(self, other):
return self.op('<%>', return_type=Float)(other)
# for reflection
ischema_names['bit'] = BIT # type: ignore