Merge pull request #102 from psqlpy-python/feature/support_pg_vector

chandr-andr · web-flow · commit e831c5860dbe · 2024-11-19T19:57:05.000+01:00
Added PgVector integration
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -56,3 +56,6 @@ itertools = "0.12.1"
 openssl-src = "300.2.2"
 openssl-sys = "0.9.102"
 pg_interval = { git = "https://github.com/chandr-andr/rust-postgres-interval.git", branch = "psqlpy" }
+pgvector = { git = "https://github.com/chandr-andr/pgvector-rust.git", branch = "psqlpy", features = [
+    "postgres",
+] }
diff --git a/python/psqlpy/_internal/extra_types.pyi b/python/psqlpy/_internal/extra_types.pyi
@@ -774,3 +774,16 @@ class IntervalArray:
         ### Parameters:
         - `inner`: inner value, sequence of timedelta values.
         """
+
+class PgVector:
+    """Represent VECTOR in PostgreSQL."""
+
+    def __init__(
+        self: Self,
+        inner: typing.Sequence[float | int],
+    ) -> None:
+        """Create new instance of PgVector.
+
+        ### Parameters:
+        - `inner`: inner value, sequence of float or int values.
+        """
diff --git a/python/psqlpy/extra_types.py b/python/psqlpy/extra_types.py
@@ -26,6 +26,7 @@
     MoneyArray,
     NumericArray,
     PathArray,
+    PgVector,
     PointArray,
     PyBox,
     PyCircle,
@@ -98,4 +99,5 @@
     "LsegArray",
     "CircleArray",
     "IntervalArray",
+    "PgVector",
 ]
diff --git a/python/tests/conftest.py b/python/tests/conftest.py
@@ -117,7 +117,7 @@ async def psql_pool_with_cert_file(
 
 
 @pytest.fixture(autouse=True)
-async def create_deafult_data_for_tests(
+async def create_default_data_for_tests(
     psql_pool: ConnectionPool,
     table_name: str,
     number_database_records: int,
diff --git a/src/extra_types.rs b/src/extra_types.rs
@@ -18,6 +18,25 @@ use crate::{
     },
 };
 
+#[pyclass]
+#[derive(Clone)]
+pub struct PgVector(Vec<f32>);
+
+#[pymethods]
+impl PgVector {
+    #[new]
+    fn new(vector: Vec<f32>) -> Self {
+        Self(vector)
+    }
+}
+
+impl PgVector {
+    #[must_use]
+    pub fn inner_value(self) -> Vec<f32> {
+        self.0
+    }
+}
+
 macro_rules! build_python_type {
     ($st_name:ident, $rust_type:ty) => {
         #[pyclass]
@@ -412,5 +431,6 @@ pub fn extra_types_module(_py: Python<'_>, pymod: &Bound<'_, PyModule>) -> PyRes
     pymod.add_class::<LsegArray>()?;
     pymod.add_class::<CircleArray>()?;
     pymod.add_class::<IntervalArray>()?;
+    pymod.add_class::<PgVector>()?;
     Ok(())
 }
diff --git a/src/value_converter.rs b/src/value_converter.rs
