feat(postgres): integrate pgvector support in Postgres transport

- Added support for pgvector types (Vector, HalfVector, Bit, SparseVector) in the Postgres transport.
- Updated Cargo.toml and Cargo.lock to include the pgvector dependency.
- Enhanced type conversion for pgvector types in PostgresArrowTransport.
- Created SQL scripts to define and populate a new table for vector types.
- Added tests to verify the handling of pgvector types in the Postgres integration.

Author: longshan.lu

Cargo.lock

@@ -60,6 +60,7 @@ datafusion = {version = "46", optional = true}
 prusto = {version = "0.5", optional = true}
 serde = {version = "1", optional = true}
 cidr-02 = { version = "0.2", package = "cidr", optional = true }
+pgvector = { version = "0.4", features = [ "postgres", "halfvec"], optional = true }
 
 [lib]
 crate-type = ["cdylib", "rlib"]
@@ -99,6 +100,7 @@ src_postgres = [
   "openssl",
   "postgres-openssl",
   "cidr-02",
+  "pgvector",
 ]
 src_sqlite = ["rusqlite", "r2d2_sqlite", "fallible-streaming-iterator", "r2d2", "urlencoding"]
 src_trino = ["prusto", "uuid", "urlencoding", "rust_decimal", "tokio", "num-traits", "serde"]

connectorx/Cargo.toml

@@ -7,6 +7,7 @@ mod typesystem;
 pub use self::errors::PostgresSourceError;
 use cidr_02::IpInet;
 pub use connection::rewrite_tls_args;
+use pgvector::{Bit, HalfVector, SparseVector, Vector};
 pub use typesystem::{PostgresTypePairs, PostgresTypeSystem};
 
 use crate::constants::DB_BUFFER_SIZE;
@@ -52,6 +53,47 @@ pub enum SimpleProtocol {}
 type PgManager<C> = PostgresConnectionManager<C>;
 type PgConn<C> = PooledConnection<PgManager<C>>;
 
+macro_rules! impl_produce_unimplemented {
+    ($(($protocol: ty, $t: ty, $msg: expr),)+) => {
+        $(
+            impl<'r> Produce<'r, $t> for $protocol {
+                type Error = PostgresSourceError;
+
+                #[throws(PostgresSourceError)]
+                fn produce(&'r mut self) -> $t {
+                   unimplemented!($msg);
+                }
+            }
+
+            impl<'r> Produce<'r, Option<$t>> for $protocol {
+                type Error = PostgresSourceError;
+
+                #[throws(PostgresSourceError)]
+                fn produce(&'r mut self) -> Option<$t> {
+                   unimplemented!($msg);
+                }
+            }
+        )+
+    };
+}
+
+impl_produce_unimplemented!(
+    (PostgresCSVSourceParser<'_>, HashMap<String, Option<String>>, "Please use `cursor` protocol for hstore type"),
+    (PostgresCSVSourceParser<'_>, Vector, "Please use `binary` protocol for vector type"),
+    (PostgresCSVSourceParser<'_>, HalfVector, "Please use `binary` protocol for halfvector type"),
+    (PostgresCSVSourceParser<'_>, Bit, "Please use `binary` protocol for bit type"),
+    (PostgresCSVSourceParser<'_>, SparseVector, "Please use `binary` protocol for sparsevector type"),
+
+
+    (PostgresSimpleSourceParser,HashMap<String, Option<String>>, "unimplemented"),
+    (PostgresSimpleSourceParser,Value, "unimplemented"),
+    (PostgresSimpleSourceParser, Vector, "Please use `binary` protocol for vector type"),
+    (PostgresSimpleSourceParser, HalfVector, "Please use `binary` protocol for halfvector type"),
+    (PostgresSimpleSourceParser, Bit, "Please use `binary` protocol for bit type"),
+    (PostgresSimpleSourceParser, SparseVector, "Please use `binary` protocol for sparsevector type"),
+
+);
+
 // take a row and unwrap the interior field from column 0
 fn convert_row<'b, R: TryFrom<usize> + postgres::types::FromSql<'b> + Clone>(row: &'b Row) -> R {
     let nrows: Option<R> = row.get(0);
@@ -482,6 +524,10 @@ impl_produce!(
     Uuid,
     Value,
     IpInet,
+    Vector,
+    HalfVector,
+    Bit,
+    SparseVector,
     Vec<Option<bool>>,
     Vec<Option<i16>>,
     Vec<Option<i32>>,
@@ -773,22 +819,6 @@ macro_rules! impl_csv_vec_produce {
 
 impl_csv_vec_produce!(i8, i16, i32, i64, f32, f64, Decimal, String,);
 
-impl Produce<'_, HashMap<String, Option<String>>> for PostgresCSVSourceParser<'_> {
-    type Error = PostgresSourceError;
-    #[throws(PostgresSourceError)]
-    fn produce(&mut self) -> HashMap<String, Option<String>> {
-        unimplemented!("Please use `cursor` protocol for hstore type");
-    }
-}
-
-impl Produce<'_, Option<HashMap<String, Option<String>>>> for PostgresCSVSourceParser<'_> {
-    type Error = PostgresSourceError;
-    #[throws(PostgresSourceError)]
-    fn produce(&mut self) -> Option<HashMap<String, Option<String>>> {
-        unimplemented!("Please use `cursor` protocol for hstore type");
-    }
-}
-
 impl Produce<'_, bool> for PostgresCSVSourceParser<'_> {
     type Error = PostgresSourceError;
 
@@ -1219,6 +1249,10 @@ impl_produce!(
     Uuid,
     Value,
     IpInet,
+    Vector,
+    HalfVector,
+    Bit,
+    SparseVector,
     HashMap<String, Option<String>>,
     Vec<Option<bool>>,
     Vec<Option<String>>,
@@ -1403,30 +1437,6 @@ impl PartitionParser<'_> for PostgresSimpleSourceParser {
     }
 }
 
-macro_rules! impl_simple_produce_unimplemented {
-    ($($t: ty,)+) => {
-        $(
-            impl<'r, 'a> Produce<'r, $t> for PostgresSimpleSourceParser {
-                type Error = PostgresSourceError;
-
-                #[throws(PostgresSourceError)]
-                fn produce(&'r mut self) -> $t {
-                   unimplemented!("not implemented!");
-                }
-            }
-
-            impl<'r, 'a> Produce<'r, Option<$t>> for PostgresSimpleSourceParser {
-                type Error = PostgresSourceError;
-
-                #[throws(PostgresSourceError)]
-                fn produce(&'r mut self) -> Option<$t> {
-                   unimplemented!("not implemented!");
-                }
-            }
-        )+
-    };
-}
-
 macro_rules! impl_simple_produce {
     ($($t: ty,)+) => {
         $(
@@ -1591,10 +1601,6 @@ impl<'r> Produce<'r, Option<Decimal>> for PostgresSimpleSourceParser {
     }
 }
 
-impl_simple_produce_unimplemented!(
-    Value,
-    HashMap<String, Option<String>>,);
-
 impl<'r> Produce<'r, &'r str> for PostgresSimpleSourceParser {
     type Error = PostgresSourceError;
 

connectorx/src/sources/postgres/mod.rs

@@ -6,6 +6,8 @@ use serde_json::Value;
 use std::collections::HashMap;
 use uuid::Uuid;
 
+use pgvector::{Bit, HalfVector, SparseVector, Vector};
+
 #[derive(Copy, Clone, Debug)]
 pub enum PostgresTypeSystem {
     Bool(bool),
@@ -40,6 +42,10 @@ pub enum PostgresTypeSystem {
     HSTORE(bool),
     Name(bool),
     Inet(bool),
+    Vector(bool),
+    HalfVec(bool),
+    Bit(bool),
+    SparseVec(bool),
 }
 
 impl_typesystem! {
@@ -71,6 +77,10 @@ impl_typesystem! {
         { JSON | JSONB => Value }
         { HSTORE => HashMap<String, Option<String>> }
         { Inet => IpInet }
+        { Vector => Vector }
+        { HalfVec => HalfVector }
+        { Bit => Bit }
+        { SparseVec => SparseVector }
     }
 }
 
@@ -108,6 +118,10 @@ impl<'a> From<&'a Type> for PostgresTypeSystem {
             "jsonb" => JSONB(true),
             "hstore" => HSTORE(true),
             "inet" => Inet(true),
+            "vector" => Vector(true),
+            "halfvec" => HalfVec(true),
+            "bit" => Bit(true),
+            "sparsevec" => SparseVec(true),
             _ => match ty.kind() {
                 postgres::types::Kind::Enum(_) => Enum(true),
                 _ => unimplemented!("{}", ty.name()),

connectorx/src/sources/postgres/typesystem.rs

@@ -14,6 +14,7 @@ use crate::typesystem::TypeConversion;
 use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
 use cidr_02::IpInet;
 use num_traits::ToPrimitive;
+use pgvector::{Bit, HalfVector, SparseVector, Vector};
 use postgres::NoTls;
 use postgres_openssl::MakeTlsConnector;
 use rust_decimal::Decimal;
@@ -76,6 +77,10 @@ macro_rules! impl_postgres_transport {
                 { Float4Array[Vec<Option<f32>>]      => Float32Array[Vec<Option<f32>>]         | conversion auto   }
                 { Float8Array[Vec<Option<f64>>]      => Float64Array[Vec<Option<f64>>]         | conversion auto   }
                 { NumericArray[Vec<Option<Decimal>>] => DecimalArray[Vec<Option<Decimal>>]     | conversion auto   }
+                { Vector[Vector]                     => Float32Array[Vec<Option<f32>>]         | conversion none   }
+                { HalfVec[HalfVector]                => Float32Array[Vec<Option<f32>>]         | conversion none   }
+                { Bit[Bit]                           => LargeBinary[Vec<u8>]                   | conversion none   }
+                { SparseVec[SparseVector]            => Float32Array[Vec<Option<f32>>]         | conversion none   }
             }
         );
     }
@@ -140,3 +145,57 @@ impl<P, C> TypeConversion<Value, String> for PostgresArrowTransport<P, C> {
         val.to_string()
     }
 }
+
+impl<P, C> TypeConversion<Vector, Vec<Option<f32>>> for PostgresArrowTransport<P, C> {
+    fn convert(val: Vector) -> Vec<Option<f32>> {
+        val.to_vec().into_iter().map(Some).collect()
+    }
+}
+
+impl<P, C> TypeConversion<Option<Vector>, Option<Vec<Option<f32>>>>
+    for PostgresArrowTransport<P, C>
+{
+    fn convert(val: Option<Vector>) -> Option<Vec<Option<f32>>> {
+        val.map(|val| val.to_vec().into_iter().map(Some).collect())
+    }
+}
+
+impl<P, C> TypeConversion<HalfVector, Vec<Option<f32>>> for PostgresArrowTransport<P, C> {
+    fn convert(val: HalfVector) -> Vec<Option<f32>> {
+        val.to_vec().into_iter().map(|v| Some(v.to_f32())).collect()
+    }
+}
+
+impl<P, C> TypeConversion<Option<HalfVector>, Option<Vec<Option<f32>>>>
+    for PostgresArrowTransport<P, C>
+{
+    fn convert(val: Option<HalfVector>) -> Option<Vec<Option<f32>>> {
+        val.map(|val| val.to_vec().into_iter().map(|v| Some(v.to_f32())).collect())
+    }
+}
+
+impl<P, C> TypeConversion<Bit, Vec<u8>> for PostgresArrowTransport<P, C> {
+    fn convert(val: Bit) -> Vec<u8> {
+        val.as_bytes().into()
+    }
+}
+
+impl<P, C> TypeConversion<Option<Bit>, Option<Vec<u8>>> for PostgresArrowTransport<P, C> {
+    fn convert(val: Option<Bit>) -> Option<Vec<u8>> {
+        val.map(|val| val.as_bytes().into())
+    }
+}
+
+impl<P, C> TypeConversion<SparseVector, Vec<Option<f32>>> for PostgresArrowTransport<P, C> {
+    fn convert(val: SparseVector) -> Vec<Option<f32>> {
+        val.to_vec().into_iter().map(Some).collect()
+    }
+}
+
+impl<P, C> TypeConversion<Option<SparseVector>, Option<Vec<Option<f32>>>>
+    for PostgresArrowTransport<P, C>
+{
+    fn convert(val: Option<SparseVector>) -> Option<Vec<Option<f32>>> {
+        val.map(|val| val.to_vec().into_iter().map(Some).collect())
+    }
+}