pgvector 0.4.2__tar.gz → 0.5.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (95) hide show
  1. {pgvector-0.4.2 → pgvector-0.5.1}/LICENSE.txt +1 -1
  2. {pgvector-0.4.2/pgvector.egg-info → pgvector-0.5.1}/PKG-INFO +50 -17
  3. pgvector-0.4.2/PKG-INFO → pgvector-0.5.1/README.md +48 -27
  4. pgvector-0.5.1/pgvector/_utils.py +23 -0
  5. pgvector-0.5.1/pgvector/asyncpg/register.py +32 -0
  6. pgvector-0.5.1/pgvector/bit.py +94 -0
  7. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/django/__init__.py +1 -6
  8. pgvector-0.5.1/pgvector/django/bit.py +37 -0
  9. pgvector-0.5.1/pgvector/django/extensions.py +11 -0
  10. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/django/functions.py +9 -10
  11. pgvector-0.5.1/pgvector/django/halfvec.py +54 -0
  12. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/django/indexes.py +7 -6
  13. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/django/sparsevec.py +19 -13
  14. pgvector-0.5.1/pgvector/django/vector.py +54 -0
  15. pgvector-0.5.1/pgvector/halfvec.py +114 -0
  16. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/peewee/__init__.py +1 -6
  17. pgvector-0.5.1/pgvector/peewee/bit.py +22 -0
  18. pgvector-0.5.1/pgvector/peewee/halfvec.py +35 -0
  19. pgvector-0.5.1/pgvector/peewee/sparsevec.py +35 -0
  20. pgvector-0.5.1/pgvector/peewee/vector.py +35 -0
  21. pgvector-0.5.1/pgvector/pg8000/__init__.py +5 -0
  22. pgvector-0.5.1/pgvector/pg8000/register.py +29 -0
  23. pgvector-0.5.1/pgvector/psycopg/__init__.py +6 -0
  24. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/psycopg/bit.py +10 -7
  25. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/psycopg/halfvec.py +14 -16
  26. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/psycopg/register.py +4 -2
  27. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/psycopg/sparsevec.py +14 -16
  28. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/psycopg/vector.py +22 -18
  29. pgvector-0.5.1/pgvector/psycopg2/__init__.py +5 -0
  30. pgvector-0.5.1/pgvector/psycopg2/halfvec.py +30 -0
  31. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/psycopg2/register.py +3 -3
  32. pgvector-0.5.1/pgvector/psycopg2/sparsevec.py +30 -0
  33. pgvector-0.5.1/pgvector/psycopg2/vector.py +40 -0
  34. pgvector-0.5.1/pgvector/py.typed +0 -0
  35. pgvector-0.5.1/pgvector/sparsevec.py +191 -0
  36. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/sqlalchemy/__init__.py +0 -5
  37. pgvector-0.5.1/pgvector/sqlalchemy/bit.py +24 -0
  38. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/sqlalchemy/functions.py +2 -1
  39. pgvector-0.5.1/pgvector/sqlalchemy/halfvec.py +55 -0
  40. pgvector-0.5.1/pgvector/sqlalchemy/sparsevec.py +55 -0
  41. pgvector-0.5.1/pgvector/sqlalchemy/vector.py +55 -0
  42. pgvector-0.5.1/pgvector/vector.py +112 -0
  43. pgvector-0.4.2/README.md → pgvector-0.5.1/pgvector.egg-info/PKG-INFO +60 -14
  44. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector.egg-info/SOURCES.txt +2 -2
  45. pgvector-0.5.1/pyproject.toml +87 -0
  46. pgvector-0.5.1/tests/test_asyncpg.py +147 -0
  47. pgvector-0.5.1/tests/test_bit.py +103 -0
  48. {pgvector-0.4.2 → pgvector-0.5.1}/tests/test_django.py +88 -77
  49. pgvector-0.5.1/tests/test_half_vector.py +75 -0
  50. {pgvector-0.4.2 → pgvector-0.5.1}/tests/test_peewee.py +44 -44
  51. {pgvector-0.4.2 → pgvector-0.5.1}/tests/test_pg8000.py +21 -30
  52. pgvector-0.5.1/tests/test_psycopg.py +236 -0
  53. {pgvector-0.4.2 → pgvector-0.5.1}/tests/test_psycopg2.py +31 -38
  54. pgvector-0.5.1/tests/test_sparse_vector.py +166 -0
  55. {pgvector-0.4.2 → pgvector-0.5.1}/tests/test_sqlalchemy.py +212 -182
  56. {pgvector-0.4.2 → pgvector-0.5.1}/tests/test_sqlmodel.py +74 -70
  57. pgvector-0.5.1/tests/test_vector.py +75 -0
  58. pgvector-0.4.2/pgvector/asyncpg/__init__.py +0 -11
  59. pgvector-0.4.2/pgvector/asyncpg/register.py +0 -31
  60. pgvector-0.4.2/pgvector/bit.py +0 -75
  61. pgvector-0.4.2/pgvector/django/bit.py +0 -32
  62. pgvector-0.4.2/pgvector/django/extensions.py +0 -11
  63. pgvector-0.4.2/pgvector/django/halfvec.py +0 -60
  64. pgvector-0.4.2/pgvector/django/vector.py +0 -73
  65. pgvector-0.4.2/pgvector/halfvec.py +0 -83
  66. pgvector-0.4.2/pgvector/peewee/bit.py +0 -21
  67. pgvector-0.4.2/pgvector/peewee/halfvec.py +0 -34
  68. pgvector-0.4.2/pgvector/peewee/sparsevec.py +0 -34
  69. pgvector-0.4.2/pgvector/peewee/vector.py +0 -34
  70. pgvector-0.4.2/pgvector/pg8000/register.py +0 -23
  71. pgvector-0.4.2/pgvector/psycopg/__init__.py +0 -13
  72. pgvector-0.4.2/pgvector/psycopg2/__init__.py +0 -10
  73. pgvector-0.4.2/pgvector/psycopg2/halfvec.py +0 -25
  74. pgvector-0.4.2/pgvector/psycopg2/sparsevec.py +0 -25
  75. pgvector-0.4.2/pgvector/psycopg2/vector.py +0 -27
  76. pgvector-0.4.2/pgvector/sparsevec.py +0 -161
  77. pgvector-0.4.2/pgvector/sqlalchemy/bit.py +0 -38
  78. pgvector-0.4.2/pgvector/sqlalchemy/halfvec.py +0 -51
  79. pgvector-0.4.2/pgvector/sqlalchemy/sparsevec.py +0 -51
  80. pgvector-0.4.2/pgvector/sqlalchemy/vector.py +0 -51
  81. pgvector-0.4.2/pgvector/utils/__init__.py +0 -9
  82. pgvector-0.4.2/pgvector/vector.py +0 -83
  83. pgvector-0.4.2/pgvector.egg-info/requires.txt +0 -1
  84. pgvector-0.4.2/pyproject.toml +0 -24
  85. pgvector-0.4.2/tests/test_asyncpg.py +0 -146
  86. pgvector-0.4.2/tests/test_bit.py +0 -63
  87. pgvector-0.4.2/tests/test_half_vector.py +0 -59
  88. pgvector-0.4.2/tests/test_psycopg.py +0 -223
  89. pgvector-0.4.2/tests/test_sparse_vector.py +0 -112
  90. pgvector-0.4.2/tests/test_vector.py +0 -59
  91. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector/__init__.py +0 -0
  92. {pgvector-0.4.2/pgvector/pg8000 → pgvector-0.5.1/pgvector/asyncpg}/__init__.py +0 -0
  93. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector.egg-info/dependency_links.txt +0 -0
  94. {pgvector-0.4.2 → pgvector-0.5.1}/pgvector.egg-info/top_level.txt +0 -0
  95. {pgvector-0.4.2 → pgvector-0.5.1}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  The MIT License (MIT)
2
2
 
3
- Copyright (c) 2021-2025 Andrew Kane
3
+ Copyright (c) 2021-2026 Andrew Kane
4
4
 
5
5
  Permission is hereby granted, free of charge, to any person obtaining a copy
6
6
  of this software and associated documentation files (the "Software"), to deal
@@ -1,21 +1,20 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pgvector
3
- Version: 0.4.2
3
+ Version: 0.5.1
4
4
  Summary: pgvector support for Python
5
5
  Author-email: Andrew Kane <andrew@ankane.org>
6
6
  License-Expression: MIT
7
7
  Project-URL: Homepage, https://github.com/pgvector/pgvector-python
8
- Requires-Python: >=3.9
8
+ Requires-Python: >=3.10
9
9
  Description-Content-Type: text/markdown
10
10
  License-File: LICENSE.txt
11
- Requires-Dist: numpy
12
11
  Dynamic: license-file
13
12
 
14
13
  # pgvector-python
15
14
 
16
15
  [pgvector](https://github.com/pgvector/pgvector) support for Python
17
16
 
18
- Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://github.com/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee)
17
+ Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://codeberg.org/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee)
19
18
 
20
19
  [![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions)
21
20
 
@@ -190,10 +189,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
190
189
  Add a vector column
191
190
 
192
191
  ```python
193
- from pgvector.sqlalchemy import Vector
192
+ from pgvector.sqlalchemy import VECTOR
194
193
 
195
194
  class Item(Base):
196
- embedding = mapped_column(Vector(3))
195
+ embedding: Mapped[list[float]] = mapped_column(VECTOR(3))
197
196
  ```
198
197
 
199
198
  Also supports `HALFVEC`, `BIT`, and `SPARSEVEC`
@@ -319,11 +318,11 @@ session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([
319
318
  Add an array column
320
319
 
321
320
  ```python
322
- from pgvector.sqlalchemy import Vector
321
+ from pgvector.sqlalchemy import VECTOR
323
322
  from sqlalchemy import ARRAY
324
323
 
325
324
  class Item(Base):
326
- embeddings = mapped_column(ARRAY(Vector(3)))
325
+ embeddings: Mapped[list[list[float]]] = mapped_column(ARRAY(VECTOR(3), dimensions=1))
327
326
  ```
328
327
 
329
328
  And register the types with the underlying driver
@@ -358,7 +357,7 @@ from sqlalchemy import event
358
357
 
359
358
  @event.listens_for(engine, "connect")
360
359
  def connect(dbapi_connection, connection_record):
361
- register_vector(dbapi_connection, arrays=True)
360
+ register_vector(dbapi_connection)
362
361
  ```
363
362
 
364
363
  ## SQLModel
@@ -372,10 +371,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector'))
372
371
  Add a vector column
373
372
 
374
373
  ```python
375
- from pgvector.sqlalchemy import Vector
374
+ from pgvector.sqlalchemy import VECTOR
376
375
 
377
376
  class Item(SQLModel, table=True):
378
- embedding: Any = Field(sa_type=Vector(3))
377
+ embedding: list[float] = Field(sa_type=VECTOR(3))
379
378
  ```
380
379
 
381
380
  Also supports `HALFVEC`, `BIT`, and `SPARSEVEC`
@@ -486,7 +485,9 @@ conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))
486
485
  Insert a vector
487
486
 
488
487
  ```python
489
- embedding = np.array([1, 2, 3])
488
+ from pgvector import Vector
489
+
490
+ embedding = Vector([1, 2, 3])
490
491
  conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
491
492
  ```
492
493
 
@@ -532,7 +533,9 @@ cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))'
532
533
  Insert a vector
533
534
 
534
535
  ```python
535
- embedding = np.array([1, 2, 3])
536
+ from pgvector import Vector
537
+
538
+ embedding = Vector([1, 2, 3])
536
539
  cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
537
540
  ```
538
541
 
@@ -587,7 +590,9 @@ await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vect
587
590
  Insert a vector
588
591
 
589
592
  ```python
590
- embedding = np.array([1, 2, 3])
593
+ from pgvector import Vector
594
+
595
+ embedding = Vector([1, 2, 3])
591
596
  await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding)
592
597
  ```
593
598
 
@@ -632,7 +637,9 @@ conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
632
637
  Insert a vector
633
638
 
634
639
  ```python
635
- embedding = np.array([1, 2, 3])
640
+ from pgvector import Vector
641
+
642
+ embedding = Vector([1, 2, 3])
636
643
  conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding)
637
644
  ```
638
645
 
@@ -711,6 +718,32 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc
711
718
 
712
719
  ## Reference
713
720
 
721
+ ### Vectors
722
+
723
+ Create a vector from a list
724
+
725
+ ```python
726
+ vec = Vector([1, 2, 3])
727
+ ```
728
+
729
+ Or a NumPy array
730
+
731
+ ```python
732
+ vec = Vector(np.array([1, 2, 3]))
733
+ ```
734
+
735
+ Get a list
736
+
737
+ ```python
738
+ lst = vec.to_list()
739
+ ```
740
+
741
+ Get a NumPy array
742
+
743
+ ```python
744
+ arr = vec.to_numpy()
745
+ ```
746
+
714
747
  ### Half Vectors
715
748
 
716
749
  Create a half vector from a list
@@ -820,7 +853,7 @@ To get started with development:
820
853
  ```sh
821
854
  git clone https://github.com/pgvector/pgvector-python.git
822
855
  cd pgvector-python
823
- pip install -r requirements.txt
856
+ pip install --group dev
824
857
  createdb pgvector_python_test
825
858
  pytest
826
859
  ```
@@ -829,7 +862,7 @@ To run an example:
829
862
 
830
863
  ```sh
831
864
  cd examples/loading
832
- pip install -r requirements.txt
865
+ pip install --group dev
833
866
  createdb pgvector_example
834
867
  python3 example.py
835
868
  ```
@@ -1,21 +1,8 @@
1
- Metadata-Version: 2.4
2
- Name: pgvector
3
- Version: 0.4.2
4
- Summary: pgvector support for Python
5
- Author-email: Andrew Kane <andrew@ankane.org>
6
- License-Expression: MIT
7
- Project-URL: Homepage, https://github.com/pgvector/pgvector-python
8
- Requires-Python: >=3.9
9
- Description-Content-Type: text/markdown
10
- License-File: LICENSE.txt
11
- Requires-Dist: numpy
12
- Dynamic: license-file
13
-
14
1
  # pgvector-python
15
2
 
16
3
  [pgvector](https://github.com/pgvector/pgvector) support for Python
17
4
 
18
- Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://github.com/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee)
5
+ Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://codeberg.org/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee)
19
6
 
20
7
  [![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions)
21
8
 
@@ -190,10 +177,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
190
177
  Add a vector column
191
178
 
192
179
  ```python
193
- from pgvector.sqlalchemy import Vector
180
+ from pgvector.sqlalchemy import VECTOR
194
181
 
195
182
  class Item(Base):
196
- embedding = mapped_column(Vector(3))
183
+ embedding: Mapped[list[float]] = mapped_column(VECTOR(3))
197
184
  ```
198
185
 
199
186
  Also supports `HALFVEC`, `BIT`, and `SPARSEVEC`
@@ -319,11 +306,11 @@ session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([
319
306
  Add an array column
320
307
 
321
308
  ```python
322
- from pgvector.sqlalchemy import Vector
309
+ from pgvector.sqlalchemy import VECTOR
323
310
  from sqlalchemy import ARRAY
324
311
 
325
312
  class Item(Base):
326
- embeddings = mapped_column(ARRAY(Vector(3)))
313
+ embeddings: Mapped[list[list[float]]] = mapped_column(ARRAY(VECTOR(3), dimensions=1))
327
314
  ```
328
315
 
329
316
  And register the types with the underlying driver
@@ -358,7 +345,7 @@ from sqlalchemy import event
358
345
 
359
346
  @event.listens_for(engine, "connect")
360
347
  def connect(dbapi_connection, connection_record):
361
- register_vector(dbapi_connection, arrays=True)
348
+ register_vector(dbapi_connection)
362
349
  ```
363
350
 
364
351
  ## SQLModel
@@ -372,10 +359,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector'))
372
359
  Add a vector column
373
360
 
374
361
  ```python
375
- from pgvector.sqlalchemy import Vector
362
+ from pgvector.sqlalchemy import VECTOR
376
363
 
377
364
  class Item(SQLModel, table=True):
378
- embedding: Any = Field(sa_type=Vector(3))
365
+ embedding: list[float] = Field(sa_type=VECTOR(3))
379
366
  ```
380
367
 
381
368
  Also supports `HALFVEC`, `BIT`, and `SPARSEVEC`
@@ -486,7 +473,9 @@ conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))
486
473
  Insert a vector
487
474
 
488
475
  ```python
489
- embedding = np.array([1, 2, 3])
476
+ from pgvector import Vector
477
+
478
+ embedding = Vector([1, 2, 3])
490
479
  conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
491
480
  ```
492
481
 
@@ -532,7 +521,9 @@ cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))'
532
521
  Insert a vector
533
522
 
534
523
  ```python
535
- embedding = np.array([1, 2, 3])
524
+ from pgvector import Vector
525
+
526
+ embedding = Vector([1, 2, 3])
536
527
  cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
537
528
  ```
538
529
 
@@ -587,7 +578,9 @@ await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vect
587
578
  Insert a vector
588
579
 
589
580
  ```python
590
- embedding = np.array([1, 2, 3])
581
+ from pgvector import Vector
582
+
583
+ embedding = Vector([1, 2, 3])
591
584
  await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding)
592
585
  ```
593
586
 
@@ -632,7 +625,9 @@ conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
632
625
  Insert a vector
633
626
 
634
627
  ```python
635
- embedding = np.array([1, 2, 3])
628
+ from pgvector import Vector
629
+
630
+ embedding = Vector([1, 2, 3])
636
631
  conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding)
637
632
  ```
638
633
 
@@ -711,6 +706,32 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc
711
706
 
712
707
  ## Reference
713
708
 
709
+ ### Vectors
710
+
711
+ Create a vector from a list
712
+
713
+ ```python
714
+ vec = Vector([1, 2, 3])
715
+ ```
716
+
717
+ Or a NumPy array
718
+
719
+ ```python
720
+ vec = Vector(np.array([1, 2, 3]))
721
+ ```
722
+
723
+ Get a list
724
+
725
+ ```python
726
+ lst = vec.to_list()
727
+ ```
728
+
729
+ Get a NumPy array
730
+
731
+ ```python
732
+ arr = vec.to_numpy()
733
+ ```
734
+
714
735
  ### Half Vectors
715
736
 
716
737
  Create a half vector from a list
@@ -820,7 +841,7 @@ To get started with development:
820
841
  ```sh
821
842
  git clone https://github.com/pgvector/pgvector-python.git
822
843
  cd pgvector-python
823
- pip install -r requirements.txt
844
+ pip install --group dev
824
845
  createdb pgvector_python_test
825
846
  pytest
826
847
  ```
@@ -829,7 +850,7 @@ To run an example:
829
850
 
830
851
  ```sh
831
852
  cd examples/loading
832
- pip install -r requirements.txt
853
+ pip install --group dev
833
854
  createdb pgvector_example
834
855
  python3 example.py
835
856
  ```
@@ -0,0 +1,23 @@
1
+ import sys
2
+ from typing import TYPE_CHECKING, TypeAlias
3
+
4
+ if TYPE_CHECKING:
5
+ import numpy as np
6
+
7
+ ndarray: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.floating]]
8
+ else:
9
+ # any value works since not type checking
10
+ # TODO use Never when Python 3.10 no longer supported
11
+ ndarray = None
12
+
13
+
14
+ def is_ndarray(value: object, /) -> bool:
15
+ if (numpy := sys.modules.get('numpy')):
16
+ return isinstance(value, numpy.ndarray)
17
+ return False
18
+
19
+
20
+ def is_sparse_array(value: object, /) -> bool:
21
+ if (sparse := sys.modules.get('scipy.sparse')):
22
+ return isinstance(value, (sparse.sparray, sparse.spmatrix))
23
+ return False
@@ -0,0 +1,32 @@
1
+ from asyncpg import Connection
2
+ from .. import Vector, HalfVector, SparseVector
3
+
4
+
5
+ async def register_vector(conn: Connection, /, *, schema: str = 'public') -> None:
6
+ await conn.set_type_codec(
7
+ 'vector',
8
+ schema=schema,
9
+ encoder=lambda v: (v if isinstance(v, Vector) else Vector(v)).to_binary(),
10
+ decoder=Vector.from_binary,
11
+ format='binary'
12
+ )
13
+
14
+ try:
15
+ await conn.set_type_codec(
16
+ 'halfvec',
17
+ schema=schema,
18
+ encoder=lambda v: (v if isinstance(v, HalfVector) else HalfVector(v)).to_binary(),
19
+ decoder=HalfVector.from_binary,
20
+ format='binary'
21
+ )
22
+
23
+ await conn.set_type_codec(
24
+ 'sparsevec',
25
+ schema=schema,
26
+ encoder=lambda v: (v if isinstance(v, SparseVector) else SparseVector(v)).to_binary(),
27
+ decoder=SparseVector.from_binary,
28
+ format='binary'
29
+ )
30
+ except ValueError as e:
31
+ if not str(e).startswith('unknown type:'):
32
+ raise e
@@ -0,0 +1,94 @@
1
+ from __future__ import annotations
2
+ from struct import pack, unpack_from
3
+ from typing import TYPE_CHECKING
4
+ from ._utils import is_ndarray
5
+
6
+ if TYPE_CHECKING:
7
+ import numpy as np
8
+
9
+
10
+ class Bit:
11
+ _length: int
12
+ _data: bytes
13
+
14
+ def __init__(
15
+ self,
16
+ value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool_ | np.uint8]],
17
+ /
18
+ ) -> None:
19
+ if isinstance(value, bytes):
20
+ self._length = 8 * len(value)
21
+ self._data = value
22
+ elif isinstance(value, (list, str)):
23
+ if isinstance(value, list):
24
+ bits = {True: '1', False: '0'}
25
+ try:
26
+ value = ''.join([bits[v] for v in value])
27
+ except (KeyError, TypeError):
28
+ raise ValueError('expected list[bool]')
29
+ elif not set(value).issubset({'0', '1'}):
30
+ raise ValueError('expected bit string')
31
+
32
+ length = len(value)
33
+ if length % 8 != 0:
34
+ value += '0' * (8 - (length % 8))
35
+
36
+ self._length = length
37
+ self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big')
38
+ elif is_ndarray(value):
39
+ import numpy as np
40
+
41
+ if value.dtype != np.bool_:
42
+ # skip error for result of np.unpackbits
43
+ if value.dtype != np.uint8 or np.any(value > 1):
44
+ raise ValueError('expected elements to be boolean')
45
+ value = value.astype(bool)
46
+
47
+ if value.ndim != 1:
48
+ raise ValueError('expected ndim to be 1')
49
+
50
+ self._length = len(value)
51
+ self._data = np.packbits(value).tobytes() # type: ignore
52
+ else:
53
+ raise ValueError('expected bytes, str, list, or ndarray')
54
+
55
+ def __repr__(self) -> str:
56
+ return f'Bit({self.to_text()})'
57
+
58
+ def __eq__(self, other: object, /) -> bool:
59
+ if not isinstance(other, self.__class__):
60
+ return NotImplemented
61
+ return self._length == other._length and self._data == other._data
62
+
63
+ def to_list(self) -> list[bool]:
64
+ # TODO improve
65
+ return [v != '0' for v in self.to_text()]
66
+
67
+ def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool_]]:
68
+ import numpy as np
69
+
70
+ return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool)
71
+
72
+ def to_text(self) -> str:
73
+ return ''.join(format(v, '08b') for v in self._data)[:self._length]
74
+
75
+ def to_binary(self) -> bytes:
76
+ return pack('>i', self._length) + self._data
77
+
78
+ @classmethod
79
+ def from_text(cls, value: str, /) -> Bit:
80
+ # cast to ensure always uses str constructor
81
+ return cls(str(value))
82
+
83
+ @classmethod
84
+ def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Bit:
85
+ length, = unpack_from('>i', value)
86
+ data = memoryview(value)[4:].tobytes()
87
+
88
+ if len(data) != (length + 7) // 8:
89
+ raise ValueError('invalid length')
90
+
91
+ bit = cls.__new__(cls)
92
+ bit._length = length
93
+ bit._data = data
94
+ return bit
@@ -6,9 +6,6 @@ from .indexes import IvfflatIndex, HnswIndex
6
6
  from .sparsevec import SparseVectorField
7
7
  from .vector import VectorField
8
8
 
9
- # TODO remove
10
- from .. import HalfVector, SparseVector
11
-
12
9
  __all__ = [
13
10
  'VectorExtension',
14
11
  'VectorField',
@@ -22,7 +19,5 @@ __all__ = [
22
19
  'CosineDistance',
23
20
  'L1Distance',
24
21
  'HammingDistance',
25
- 'JaccardDistance',
26
- 'HalfVector',
27
- 'SparseVector'
22
+ 'JaccardDistance'
28
23
  ]
@@ -0,0 +1,37 @@
1
+ from django import forms
2
+ from django.db.models import Field
3
+ from typing import Any
4
+
5
+
6
+ # https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/
7
+ class BitField(Field):
8
+ description = 'Bit string'
9
+
10
+ def __init__(self, *args: Any, length: int | None = None, **kwargs: Any) -> None:
11
+ self.length = length
12
+ super().__init__(*args, **kwargs)
13
+
14
+ def deconstruct(self) -> tuple[Any, Any, Any, Any]:
15
+ name, path, args, kwargs = super().deconstruct()
16
+ if self.length is not None:
17
+ kwargs['length'] = self.length
18
+ return name, path, args, kwargs
19
+
20
+ def db_type(self, connection: Any) -> str:
21
+ if self.length is None:
22
+ return 'bit'
23
+ return 'bit(%d)' % self.length
24
+
25
+ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field:
26
+ return super().formfield(
27
+ form_class=BitFormField if form_class is None else form_class,
28
+ choices_form_class=choices_form_class,
29
+ **kwargs
30
+ )
31
+
32
+
33
+ class BitFormField(forms.CharField):
34
+ def to_python(self, value: Any) -> Any:
35
+ if isinstance(value, str) and value == '':
36
+ return None
37
+ return super().to_python(value)
@@ -0,0 +1,11 @@
1
+ from django import VERSION
2
+ from django.contrib.postgres.operations import CreateExtension
3
+ from typing import Any
4
+
5
+
6
+ class VectorExtension(CreateExtension):
7
+ def __init__(self, hints: Any = None) -> None:
8
+ if VERSION[0] >= 6:
9
+ super().__init__('vector', hints=hints) # type: ignore
10
+ else:
11
+ self.name = 'vector'
@@ -1,18 +1,17 @@
1
1
  from django.db.models import FloatField, Func, Value
2
2
  from .. import Vector, HalfVector, SparseVector
3
+ from typing import Any
3
4
 
4
5
 
5
6
  class DistanceBase(Func):
6
- output_field = FloatField()
7
+ output_field = FloatField() # type: ignore
7
8
 
8
- def __init__(self, expression, vector, **extra):
9
+ def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None:
9
10
  if not hasattr(vector, 'resolve_expression'):
10
- if isinstance(vector, HalfVector):
11
- vector = Value(HalfVector._to_db(vector))
12
- elif isinstance(vector, SparseVector):
13
- vector = Value(SparseVector._to_db(vector))
14
- else:
15
- vector = Value(Vector._to_db(vector))
11
+ if isinstance(vector, (Vector, HalfVector, SparseVector)):
12
+ vector = Value(vector.to_text())
13
+ elif vector is not None:
14
+ vector = Value(Vector(vector).to_text())
16
15
 
17
16
  # prevent error with unhashable types
18
17
  self._constructor_args = ((expression, vector), extra)
@@ -21,9 +20,9 @@ class DistanceBase(Func):
21
20
 
22
21
 
23
22
  class BitDistanceBase(Func):
24
- output_field = FloatField()
23
+ output_field = FloatField() # type: ignore
25
24
 
26
- def __init__(self, expression, vector, **extra):
25
+ def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None:
27
26
  if not hasattr(vector, 'resolve_expression'):
28
27
  vector = Value(vector)
29
28
  super().__init__(expression, vector, **extra)
@@ -0,0 +1,54 @@
1
+ from django import forms
2
+ from django.db.models import Field
3
+ from typing import Any
4
+ from .. import HalfVector
5
+
6
+
7
+ # https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/
8
+ class HalfVectorField(Field):
9
+ description = 'Half vector'
10
+ empty_strings_allowed = False
11
+
12
+ def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None:
13
+ self.dimensions = dimensions
14
+ super().__init__(*args, **kwargs)
15
+
16
+ def deconstruct(self) -> tuple[Any, Any, Any, Any]:
17
+ name, path, args, kwargs = super().deconstruct()
18
+ if self.dimensions is not None:
19
+ kwargs['dimensions'] = self.dimensions
20
+ return name, path, args, kwargs
21
+
22
+ def db_type(self, connection: Any) -> str:
23
+ if self.dimensions is None:
24
+ return 'halfvec'
25
+ return 'halfvec(%d)' % self.dimensions
26
+
27
+ def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None:
28
+ return HalfVector._from_db(value)
29
+
30
+ def to_python(self, value: Any) -> list[float] | None:
31
+ if isinstance(value, list):
32
+ return value
33
+ return HalfVector._from_db(value)
34
+
35
+ def get_prep_value(self, value: Any) -> str | None:
36
+ return HalfVector._to_db(value)
37
+
38
+ def value_to_string(self, obj: Any) -> str:
39
+ value = self.get_prep_value(self.value_from_object(obj))
40
+ return '' if value is None else value
41
+
42
+ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field:
43
+ return super().formfield(
44
+ form_class=HalfVectorFormField if form_class is None else form_class,
45
+ choices_form_class=choices_form_class,
46
+ **kwargs
47
+ )
48
+
49
+
50
+ class HalfVectorFormField(forms.CharField):
51
+ def to_python(self, value: Any) -> Any:
52
+ if isinstance(value, str) and value == '':
53
+ return None
54
+ return super().to_python(value)