Skip to content

Commit 37b1b78

Browse files
chandr-andrchandr-andr
committed
Added PgVector integration
Signed-off-by: chandr-andr (Kiselev Aleksandr) <[email protected]>
1 parent 6c903e5 commit 37b1b78

File tree

3 files changed

+144
-86
lines changed

3 files changed

+144
-86
lines changed

‎python/tests/conftest.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,16 @@ async def psql_pool_with_cert_file(
117117

118118

119119
@pytest.fixture(autouse=True)
120-
async def create_deafult_data_for_tests(
120+
async def setup_postgresql_for_tests(
121+
psql_pool: ConnectionPool,
122+
) -> None:
123+
await psql_pool.execute(
124+
"CREATE EXTENSION IF NOT EXISTS vector;",
125+
)
126+
127+
128+
@pytest.fixture(autouse=True)
129+
async def create_default_data_for_tests(
121130
psql_pool: ConnectionPool,
122131
table_name: str,
123132
number_database_records: int,

‎python/tests/test_value_converter.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
MoneyArray,
3838
NumericArray,
3939
PathArray,
40+
PgVector,
4041
PointArray,
4142
PyBox,
4243
PyCircle,
@@ -609,6 +610,11 @@ async def test_as_class(
609610
[datetime.timedelta(days=100, microseconds=100), datetime.timedelta(days=100, microseconds=100)],
610611
[datetime.timedelta(days=100, microseconds=100), datetime.timedelta(days=100, microseconds=100)],
611612
),
613+
(
614+
"VECTOR",
615+
PgVector([123123, 1222211, 9999112, 12999873]),
616+
[123123, 1222211, 9999112, 12999873],
617+
),
612618
),
613619
)
614620
async def test_deserialization_simple_into_python(

‎src/value_converter.rs

Lines changed: 128 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -1325,19 +1325,18 @@ fn postgres_bytes_to_py(
13251325
// ---------- Bytes Types ----------
13261326
// Convert BYTEA type into Vector<u8>, then into PyBytes
13271327
Type::BYTEA => {
1328-
let vec_of_bytes = _composite_field_postgres_to_py::<Option<Vec<u8>>>(
1329-
type_, buf, is_simple,
1330-
)?;
1328+
let vec_of_bytes =
1329+
_composite_field_postgres_to_py::<Option<Vec<u8>>>(type_, buf, is_simple)?;
13311330
if let Some(vec_of_bytes) = vec_of_bytes {
13321331
return Ok(PyBytes::new_bound(py, &vec_of_bytes).to_object(py));
13331332
}
13341333
Ok(py.None())
1335-
},
1334+
}
13361335
// // ---------- String Types ----------
13371336
// // Convert TEXT and VARCHAR type into String, then into str
1338-
Type::TEXT | Type::VARCHAR | Type::XML => Ok(_composite_field_postgres_to_py::<Option<String>>(
1339-
type_, buf, is_simple,
1340-
)?
1337+
Type::TEXT | Type::VARCHAR | Type::XML => Ok(_composite_field_postgres_to_py::<
1338+
Option<String>,
1339+
>(type_, buf, is_simple)?
13411340
.to_object(py)),
13421341
// ---------- Boolean Types ----------
13431342
// Convert BOOL type into bool
@@ -1354,9 +1353,10 @@ fn postgres_bytes_to_py(
13541353
_composite_field_postgres_to_py::<Option<i32>>(type_, buf, is_simple)?.to_object(py),
13551354
),
13561355
// Convert BigInt into i64, then into int
1357-
Type::INT8 | Type::MONEY => Ok(
1358-
_composite_field_postgres_to_py::<Option<i64>>(type_, buf, is_simple)?.to_object(py),
1359-
),
1356+
Type::INT8 | Type::MONEY => Ok(_composite_field_postgres_to_py::<Option<i64>>(
1357+
type_, buf, is_simple,
1358+
)?
1359+
.to_object(py)),
13601360
// Convert REAL into f32, then into float
13611361
Type::FLOAT4 => Ok(
13621362
_composite_field_postgres_to_py::<Option<f32>>(type_, buf, is_simple)?.to_object(py),
@@ -1432,16 +1432,17 @@ fn postgres_bytes_to_py(
14321432
}
14331433
}
14341434
Type::NUMERIC => {
1435-
if let Some(numeric_) = _composite_field_postgres_to_py::<Option<Decimal>>(
1436-
type_, buf, is_simple,
1437-
)? {
1435+
if let Some(numeric_) =
1436+
_composite_field_postgres_to_py::<Option<Decimal>>(type_, buf, is_simple)?
1437+
{
14381438
return Ok(InnerDecimal(numeric_).to_object(py));
14391439
}
14401440
Ok(py.None().to_object(py))
14411441
}
14421442
// ---------- Geo Types ----------
14431443
Type::POINT => {
1444-
let point_ = _composite_field_postgres_to_py::<Option<RustPoint>>(type_, buf, is_simple)?;
1444+
let point_ =
1445+
_composite_field_postgres_to_py::<Option<RustPoint>>(type_, buf, is_simple)?;
14451446

14461447
match point_ {
14471448
Some(point_) => Ok(point_.into_py(py)),
@@ -1457,7 +1458,8 @@ fn postgres_bytes_to_py(
14571458
}
14581459
}
14591460
Type::PATH => {
1460-
let path_ = _composite_field_postgres_to_py::<Option<RustLineString>>(type_, buf, is_simple)?;
1461+
let path_ =
1462+
_composite_field_postgres_to_py::<Option<RustLineString>>(type_, buf, is_simple)?;
14611463

14621464
match path_ {
14631465
Some(path_) => Ok(path_.into_py(py)),
@@ -1473,7 +1475,8 @@ fn postgres_bytes_to_py(
14731475
}
14741476
}
14751477
Type::LSEG => {
1476-
let lseg_ = _composite_field_postgres_to_py::<Option<RustLineSegment>>(type_, buf, is_simple)?;
1478+
let lseg_ =
1479+
_composite_field_postgres_to_py::<Option<RustLineSegment>>(type_, buf, is_simple)?;
14771480

14781481
match lseg_ {
14791482
Some(lseg_) => Ok(lseg_.into_py(py)),
@@ -1489,139 +1492,179 @@ fn postgres_bytes_to_py(
14891492
}
14901493
}
14911494
Type::INTERVAL => {
1492-
let interval = _composite_field_postgres_to_py::<Option<Interval>>(type_, buf, is_simple)?;
1495+
let interval =
1496+
_composite_field_postgres_to_py::<Option<Interval>>(type_, buf, is_simple)?;
14931497
if let Some(interval) = interval {
14941498
return Ok(InnerInterval(interval).to_object(py));
14951499
}
14961500
Ok(py.None())
14971501
}
14981502
// ---------- Array Text Types ----------
1499-
Type::BOOL_ARRAY => Ok(postgres_array_to_py(py, _composite_field_postgres_to_py::<Option<Array<bool>>>(
1500-
type_, buf, is_simple,
1501-
)?)
1503+
Type::BOOL_ARRAY => Ok(postgres_array_to_py(
1504+
py,
1505+
_composite_field_postgres_to_py::<Option<Array<bool>>>(type_, buf, is_simple)?,
1506+
)
15021507
.to_object(py)),
15031508
// Convert ARRAY of TEXT or VARCHAR into Vec<String>, then into list[str]
1504-
Type::TEXT_ARRAY | Type::VARCHAR_ARRAY | Type::XML_ARRAY => Ok(
1505-
postgres_array_to_py(
1506-
py,
1507-
_composite_field_postgres_to_py::<Option<Array<String>>>(type_, buf, is_simple)?,
1508-
).to_object(py)),
1509+
Type::TEXT_ARRAY | Type::VARCHAR_ARRAY | Type::XML_ARRAY => Ok(postgres_array_to_py(
1510+
py,
1511+
_composite_field_postgres_to_py::<Option<Array<String>>>(type_, buf, is_simple)?,
1512+
)
1513+
.to_object(py)),
15091514
// ---------- Array Integer Types ----------
15101515
// Convert ARRAY of SmallInt into Vec<i16>, then into list[int]
1511-
Type::INT2_ARRAY => Ok(
1512-
postgres_array_to_py(
1513-
py,
1514-
_composite_field_postgres_to_py::<Option<Array<i16>>>(
1515-
type_, buf, is_simple,
1516-
)?,
1516+
Type::INT2_ARRAY => Ok(postgres_array_to_py(
1517+
py,
1518+
_composite_field_postgres_to_py::<Option<Array<i16>>>(type_, buf, is_simple)?,
15171519
)
15181520
.to_object(py)),
15191521
// Convert ARRAY of Integer into Vec<i32>, then into list[int]
1520-
Type::INT4_ARRAY => {
1521-
Ok(postgres_array_to_py(
1522-
py,
1523-
_composite_field_postgres_to_py::<Option<Array<i32>>>(
1524-
type_,
1525-
buf,
1526-
is_simple,
1527-
)?
1528-
).to_object(py))
1529-
},
1522+
Type::INT4_ARRAY => Ok(postgres_array_to_py(
1523+
py,
1524+
_composite_field_postgres_to_py::<Option<Array<i32>>>(type_, buf, is_simple)?,
1525+
)
1526+
.to_object(py)),
15301527
// Convert ARRAY of BigInt into Vec<i64>, then into list[int]
1531-
Type::INT8_ARRAY | Type::MONEY_ARRAY => Ok(postgres_array_to_py(py, _composite_field_postgres_to_py::<Option<Array<i64>>>(
1532-
type_, buf, is_simple,
1533-
)?).to_object(py)),
1528+
Type::INT8_ARRAY | Type::MONEY_ARRAY => Ok(postgres_array_to_py(
1529+
py,
1530+
_composite_field_postgres_to_py::<Option<Array<i64>>>(type_, buf, is_simple)?,
1531+
)
1532+
.to_object(py)),
15341533
// Convert ARRAY of Float4 into Vec<f32>, then into list[float]
1535-
Type::FLOAT4_ARRAY => Ok(postgres_array_to_py(py,_composite_field_postgres_to_py::<Option<Array<f32>>>(
1536-
type_, buf, is_simple,
1537-
)?)
1534+
Type::FLOAT4_ARRAY => Ok(postgres_array_to_py(
1535+
py,
1536+
_composite_field_postgres_to_py::<Option<Array<f32>>>(type_, buf, is_simple)?,
1537+
)
15381538
.to_object(py)),
15391539
// Convert ARRAY of Float8 into Vec<f64>, then into list[float]
1540-
Type::FLOAT8_ARRAY => Ok(postgres_array_to_py(py,_composite_field_postgres_to_py::<Option<Array<f64>>>(
1541-
type_, buf, is_simple,
1542-
)?)
1540+
Type::FLOAT8_ARRAY => Ok(postgres_array_to_py(
1541+
py,
1542+
_composite_field_postgres_to_py::<Option<Array<f64>>>(type_, buf, is_simple)?,
1543+
)
15431544
.to_object(py)),
15441545
// Convert ARRAY of Date into Vec<NaiveDate>, then into list[datetime.date]
1545-
Type::DATE_ARRAY => Ok(postgres_array_to_py(py,_composite_field_postgres_to_py::<Option<Array<NaiveDate>>>(
1546-
type_, buf, is_simple,
1547-
)?)
1546+
Type::DATE_ARRAY => Ok(postgres_array_to_py(
1547+
py,
1548+
_composite_field_postgres_to_py::<Option<Array<NaiveDate>>>(type_, buf, is_simple)?,
1549+
)
15481550
.to_object(py)),
15491551
// Convert ARRAY of Time into Vec<NaiveTime>, then into list[datetime.date]
1550-
Type::TIME_ARRAY => Ok(postgres_array_to_py(py, _composite_field_postgres_to_py::<Option<Array<NaiveTime>>>(
1551-
type_, buf, is_simple,
1552-
)?)
1552+
Type::TIME_ARRAY => Ok(postgres_array_to_py(
1553+
py,
1554+
_composite_field_postgres_to_py::<Option<Array<NaiveTime>>>(type_, buf, is_simple)?,
1555+
)
15531556
.to_object(py)),
15541557
// Convert ARRAY of TIMESTAMP into Vec<NaiveDateTime>, then into list[datetime.date]
1555-
Type::TIMESTAMP_ARRAY => Ok(
1556-
postgres_array_to_py(py, _composite_field_postgres_to_py::<Option<Array<NaiveDateTime>>>(type_, buf, is_simple)?)
1557-
.to_object(py),
1558-
),
1558+
Type::TIMESTAMP_ARRAY => Ok(postgres_array_to_py(
1559+
py,
1560+
_composite_field_postgres_to_py::<Option<Array<NaiveDateTime>>>(type_, buf, is_simple)?,
1561+
)
1562+
.to_object(py)),
15591563
// Convert ARRAY of TIMESTAMPTZ into Vec<DateTime<FixedOffset>>, then into list[datetime.date]
1560-
Type::TIMESTAMPTZ_ARRAY => Ok(postgres_array_to_py(py, _composite_field_postgres_to_py::<
1561-
Option<Array<DateTime<FixedOffset>>>,
1562-
>(type_, buf, is_simple)?)
1564+
Type::TIMESTAMPTZ_ARRAY => Ok(postgres_array_to_py(
1565+
py,
1566+
_composite_field_postgres_to_py::<Option<Array<DateTime<FixedOffset>>>>(
1567+
type_, buf, is_simple,
1568+
)?,
1569+
)
15631570
.to_object(py)),
15641571
// Convert ARRAY of UUID into Vec<Array<InternalUuid>>, then into list[UUID]
15651572
Type::UUID_ARRAY => {
1566-
let uuid_array =
1567-
_composite_field_postgres_to_py::<Option<Array<InternalUuid>>>(type_, buf, is_simple)?;
1573+
let uuid_array = _composite_field_postgres_to_py::<Option<Array<InternalUuid>>>(
1574+
type_, buf, is_simple,
1575+
)?;
15681576
Ok(postgres_array_to_py(py, uuid_array).to_object(py))
15691577
}
15701578
// Convert ARRAY of INET into Vec<INET>, then into list[IPv4Address | IPv6Address]
1571-
Type::INET_ARRAY => Ok(postgres_array_to_py(py, _composite_field_postgres_to_py::<Option<Array<IpAddr>>>(
1572-
type_, buf, is_simple,
1573-
)?)
1579+
Type::INET_ARRAY => Ok(postgres_array_to_py(
1580+
py,
1581+
_composite_field_postgres_to_py::<Option<Array<IpAddr>>>(type_, buf, is_simple)?,
1582+
)
15741583
.to_object(py)),
15751584
Type::JSONB_ARRAY | Type::JSON_ARRAY => {
1576-
let db_json_array =
1577-
_composite_field_postgres_to_py::<Option<Array<InternalSerdeValue>>>(type_, buf, is_simple)?;
1585+
let db_json_array = _composite_field_postgres_to_py::<Option<Array<InternalSerdeValue>>>(
1586+
type_, buf, is_simple,
1587+
)?;
15781588
Ok(postgres_array_to_py(py, db_json_array).to_object(py))
15791589
}
1580-
Type::NUMERIC_ARRAY => {
1581-
Ok(postgres_array_to_py(py, _composite_field_postgres_to_py::<Option<Array<InnerDecimal>>>(
1582-
type_, buf, is_simple,
1583-
)?).to_object(py))
1584-
},
1590+
Type::NUMERIC_ARRAY => Ok(postgres_array_to_py(
1591+
py,
1592+
_composite_field_postgres_to_py::<Option<Array<InnerDecimal>>>(type_, buf, is_simple)?,
1593+
)
1594+
.to_object(py)),
15851595
// ---------- Array Geo Types ----------
15861596
Type::POINT_ARRAY => {
1587-
let point_array_ = _composite_field_postgres_to_py::<Option<Array<RustPoint>>>(type_, buf, is_simple)?;
1597+
let point_array_ =
1598+
_composite_field_postgres_to_py::<Option<Array<RustPoint>>>(type_, buf, is_simple)?;
15881599

15891600
Ok(postgres_array_to_py(py, point_array_).to_object(py))
15901601
}
15911602
Type::BOX_ARRAY => {
1592-
let box_array_ = _composite_field_postgres_to_py::<Option<Array<RustRect>>>(type_, buf, is_simple)?;
1603+
let box_array_ =
1604+
_composite_field_postgres_to_py::<Option<Array<RustRect>>>(type_, buf, is_simple)?;
15931605

15941606
Ok(postgres_array_to_py(py, box_array_).to_object(py))
15951607
}
15961608
Type::PATH_ARRAY => {
1597-
let path_array_ = _composite_field_postgres_to_py::<Option<Array<RustLineString>>>(type_, buf, is_simple)?;
1609+
let path_array_ = _composite_field_postgres_to_py::<Option<Array<RustLineString>>>(
1610+
type_, buf, is_simple,
1611+
)?;
15981612

15991613
Ok(postgres_array_to_py(py, path_array_).to_object(py))
16001614
}
16011615
Type::LINE_ARRAY => {
1602-
let line_array_ = _composite_field_postgres_to_py::<Option<Array<Line>>>(type_, buf, is_simple)?;
1616+
let line_array_ =
1617+
_composite_field_postgres_to_py::<Option<Array<Line>>>(type_, buf, is_simple)?;
16031618

16041619
Ok(postgres_array_to_py(py, line_array_).to_object(py))
16051620
}
16061621
Type::LSEG_ARRAY => {
1607-
let lseg_array_ = _composite_field_postgres_to_py::<Option<Array<RustLineSegment>>>(type_, buf, is_simple)?;
1622+
let lseg_array_ = _composite_field_postgres_to_py::<Option<Array<RustLineSegment>>>(
1623+
type_, buf, is_simple,
1624+
)?;
16081625

16091626
Ok(postgres_array_to_py(py, lseg_array_).to_object(py))
16101627
}
16111628
Type::CIRCLE_ARRAY => {
1612-
let circle_array_ = _composite_field_postgres_to_py::<Option<Array<Circle>>>(type_, buf, is_simple)?;
1629+
let circle_array_ =
1630+
_composite_field_postgres_to_py::<Option<Array<Circle>>>(type_, buf, is_simple)?;
16131631

16141632
Ok(postgres_array_to_py(py, circle_array_).to_object(py))
16151633
}
16161634
Type::INTERVAL_ARRAY => {
1617-
let interval_array_ = _composite_field_postgres_to_py::<Option<Array<InnerInterval>>>(type_, buf, is_simple)?;
1635+
let interval_array_ = _composite_field_postgres_to_py::<Option<Array<InnerInterval>>>(
1636+
type_, buf, is_simple,
1637+
)?;
16181638

16191639
Ok(postgres_array_to_py(py, interval_array_).to_object(py))
16201640
}
1621-
_ => Err(RustPSQLDriverError::RustToPyValueConversionError(
1622-
format!("Cannot convert {type_} into Python type, please look at the custom_decoders functionality.")
1623-
)),
1641+
_ => other_postgres_bytes_to_py(py, type_, buf, is_simple),
1642+
}
1643+
}
1644+
1645+
/// Convert OTHER type to python.
1646+
///
1647+
/// # Errors
1648+
/// May return result if type is unknown.
1649+
pub fn other_postgres_bytes_to_py(
1650+
py: Python<'_>,
1651+
type_: &Type,
1652+
buf: &mut &[u8],
1653+
is_simple: bool,
1654+
) -> RustPSQLDriverPyResult<Py<PyAny>> {
1655+
if type_.name() == "vector" {
1656+
let vector = _composite_field_postgres_to_py::<Option<PgVector>>(type_, buf, is_simple)?;
1657+
match vector {
1658+
Some(real_vector) => {
1659+
return Ok(real_vector.to_vec().to_object(py));
1660+
}
1661+
None => return Ok(py.None()),
1662+
}
16241663
}
1664+
1665+
Err(RustPSQLDriverError::RustToPyValueConversionError(
1666+
format!("Cannot convert {type_} into Python type, please look at the custom_decoders functionality.")
1667+
))
16251668
}
16261669

16271670
/// Convert composite type from `PostgreSQL` to Python type.

0 commit comments

Comments
 (0)