pgvector compatibility: vector(n) type, distance operators and INSERT support - #419
Conversation
|
hello @xuchen-plus , thanks for the patch! That's great addition to the project! I haven't got time to walk through the content, but before that, could you please update the CI, to follow our current pattern, add pgvector feature to the feature matrix just like postgis, we will test it separately. |
Hi @sunng87 , I've added a separate pgvector test in For now all features test and postgis feature test failed because of mismatched arrow-rs version in GeoArrow. Thanks for helping review. |
|
Note that the Integration tests for default features and pgvector are expected to pass. please check the failure. |
|
Thanks @sunng87 Addressed both the CI failures and the review comments: CI: added pgvector to the feature matrix and an integration-pgvector job; fixed the default/pgvector integration failures (parameter decoding now prefers the client's type hint, so psycopg bound params work). Review Comments: operators (<->/<#>/<=>) are now a DataFusion ExprPlanner in datafusion-postgres (not pg_catalog rewrite rules), INSERT literals use a QueryHook, pg_catalog pgvector code moved to pg_catalog/pgvector, and OID params are decoded as i32. Verified locally: rustfmt, required clippy gate, default/full/pgvector tests, MSRV 1.94, and the python integration suite (CSV, transactions, parquet, SSL, pgvector) all pass. |
sunng87
left a comment
sunng87
left a comment
There was a problem hiding this comment.
Code-standards pass (non-blocking suggestions): inline notes on duplication, a latent sliced-array bug, and ungated changes riding along in the feature. Overall this is a solid, well-tested PR — none of these are merge blockers in my view.
Add opt-in pgvector support behind a `pgvector` cargo feature on
datafusion-postgres / datafusion-pg-catalog / arrow-pg:
- `vector(n)` / `vector` SQL type -> FixedSizeList/List(Float32) with
`pg.vector` field metadata, planned by the existing PgOidTypePlanner.
- Rewrite the pgvector distance operators `<->` / `<#>` / `<=>` onto
DataFusion's built-in array_distance / inner_product / cosine_distance,
and vector literals `'[1,2,3]'` (bare or `::vector`) to ARRAY literals.
- Schema-aware INSERT rewrite: `INSERT ... VALUES ('[1,2,3]')` into a
`vector(n)` column is rewritten to an ARRAY literal before planning.
- Wire type + encoding: vector columns report type OID 16385 and encode as
pgvector text `[1,2,3]`; fix FixedSizeList->ListArray downcast bug.
- Driver introspection: inject a `vector` row into pg_type (OID 16385),
tag pg_type.typtype as the internal `"char"` type, report parameters bound
to oid-alias / vector columns with their real wire types (recomputed from
the plan, since pgwire does not persist ParameterDescription types), and
decode OID / "char" / vector binary parameters.
- End-to-end tokio-postgres tests over a real TCP server: literal and
prepared-parameter vector INSERT, nearest-neighbour queries, psql-style
text reads, `prepare` introspection, and the canonical
`CREATE TABLE items (id int PRIMARY KEY, embedding vector(3))` DDL.
Server-side vector wire encoding: - Encode vector SELECT results in the real pgvector binary layout (big-endian u16 dimension + unused u16 + big-endian float32s) via a PgVectorValue ToSql/ToSqlText pair, so binary-format clients work; honor the requested result format instead of forcing text. - Match the official layout in parameter decoding (VectorParam) and key vector encoding off the resolved FieldInfo type, not just the arrow field metadata (optimizer rewrites can drop it). - Fix feature-gating nits (unused imports / mut) for default builds. E2e tests (datafusion-postgres/tests/pgvector.rs): - Switch the network tests from tokio-postgres to the synchronous rust-postgres client, driving a live server on a dedicated thread. - Use the official pgvector crate (pgvector::Vector) for parameter binds and result decoding; drop the hand-written PgVector/DecodedVector postgres-types codecs -- official-client round-trips now pin binary and text compatibility. - Cover DDL (CREATE TABLE ... vector(3)), literal and bound-parameter INSERT, prepare() type introspection (OID 16385), text simple-query reads, and nearest-neighbour queries. Dev-dependencies: replace tokio-postgres with postgres 0.19.
Add a pgvector entry to the test-features matrix (--features datafusion-postgres/pgvector, non-optional) so the pgvector unit/integration suite runs in CI independently of default-feature tests.
Add tests-integration/test_pgvector.py modeled on test_postgis.py, exercising pgvector against the CLI server over psycopg: DDL with a vector(3) column, pgtype registration, literal INSERT, nearest-neighbour ORDER BY embedding <-> and the <-> / <#> / <=> distance operators. Wire it into tests-integration/test.sh as a new pgvector section (port 5438) and extend the cleanup/summary output.
Mirror the existing --skip-postgis handling so the pgvector python integration test can be skipped independently, and only list pgvector in the summary output when it was run.
- tests-integration/test.sh: when the pgvector test is enabled (i.e. --skip-pgvector is not passed), include datafusion-postgres/pgvector in the cargo build features alongside the optional postgis feature. - ci.yml: add a dedicated non-advisory integration-pgvector job that builds the CLI and runs the pgvector python integration test; scope the base integration job (--skip-pgvector) and the advisory postgis job (--skip-pgvector) so each feature is tested in its own job.
The job now just runs test.sh --skip-postgis, which already builds with the pgvector feature and runs the python suite including the pgvector tests, reusing the script's wait/cleanup orchestration instead of duplicating it inline.
Decoding parameters gave the server-decided wire type priority over the client-provided type from Parse. For clients that send explicit parameter type OIDs (e.g. psycopg sends int2 for a Python int), this made the server decode the parameter as the plan's inferred type (float8) and read more bytes than the client sent, failing with 'failed to fill whole buffer'. Restore the original priority -- client hint first -- and only fall back to the server-decided type (which carries oid-alias / pgvector overrides for clients that send no types, e.g. tokio-postgres), then the inferred type.
Address review feedback by moving pgvector integration out of the generic
pg_catalog SQL-rewrite layer into datafusion-postgres, gated behind the
pgvector feature:
- Distance operators (<-> / <#> / <=>) are now implemented as a DataFusion
`ExprPlanner` (PgVectorExprPlanner) that rewrites them onto the built-in
array_distance / inner_product / cosine_distance functions while planning
SQL. This covers every expression position (projection, WHERE, ORDER BY,
subqueries) without AST rewrite rules. The session parser dialect is set to
Postgres so the operators tokenize; the planner is appended to the existing
expression planners.
- `INSERT ... VALUES ('[1,2,3]')` into a `vector(n)` column is now handled by
a `QueryHook` (PgVectorInsertHook), which has the target schema available:
the simple-protocol path rewrites and executes, the extended path returns
the rewritten logical plan.
- pg_catalog pgvector bits (vector row in pg_type, `"char"` typtype) move to
a dedicated `pg_catalog/pgvector` module.
- Decode OID parameters as i32 (Postgres has no unsigned integers), dropping
the u32/UInt32 path.
- Remove the pgvector SQL rewrite rule and vector_insert module from
datafusion-pg-catalog (and the cfg_attr hack in the parser).
- encoder: index FixedSizeList rows via value_offset(idx) so sliced arrays (offset > 0, e.g. LIMIT-sliced batches) read the right elements; add a regression test. - arrow-pg: expose parse_vector_text as the single pgvector text parser (rejects inf/NaN, which cannot round-trip through SQL number literals) and reuse it from VectorParam and from datafusion-postgres. - pg-catalog: derive the pgvector feature from arrow-pg/pgvector and reuse arrow_pg's PG_VECTOR_TYPE_OID / PG_VECTOR_KEY and a shared sql::is_vector_type predicate instead of duplicated constants. - datafusion-postgres: install the pgvector planner in serve_with_hooks (and fail loudly on error instead of silently skipping) rather than mutating the session inside DfSessionService::new; add a shared client::execute_statement helper so the INSERT hook no longer re-implements timeout/count/tag handling; drop the duplicated $N ordering helper.
4db41d4 to
904227f
Compare
sunng87
left a comment
sunng87
left a comment
There was a problem hiding this comment.
Thank you @xuchen-plus , we are almost there. Please keep the public new api HandlerFactory if it's possible.
| connection_manager: connection_manager.clone(), | ||
| }), | ||
| } | ||
| } |
There was a problem hiding this comment.
I see. There will be a deadcode warning if we keep this because HandlerFactory may not be exported.
It's ok to remove.
|
@xuchen-plus I will merge this. Thank you for contributing this. I would be nice if you can share with me about your use-case with this library. Note that I may do some refactoring later. Also due to datafusion-python's pending release, I can't make a release immediately for this. |
@sunng87 Thanks very much for the review and merge. We are building a PG compatibility layer for our lakehouse project LakeSoul, which already has its core built upon datafusion. And we are also working on vector search features (lakesoul-io/LakeSoul#798) so it would be nice to also provide pgvector compatibility. I'll continue to integrate |
3 participants
Close #418
Summary
Adds opt-in pgvector compatibility (a new
pgvectorCargo feature ondatafusion-postgres/datafusion-pg-catalog/arrow-pg) so PostgreSQLclients can use pgvector syntax against a DataFusion-backed server.
What's included
vector(n)/vectorSQL type ->FixedSizeList(Float32, n)/List(Float32)withpg.vectormetadata; reported on the wire as thepgvector
vectortype and encoded in pgvector text / binary format.<->->array_distance,<#>->-inner_product,<=>->cosine_distance;vector literals (
'[1,2,3]', bare or::vector) ->ARRAY[...].INSERT ... VALUES ('[1,2,3]')literals and bound vectorparameters over the extended protocol.
vectorrowin
pg_type, oid /"char"wire typing).CREATE TABLE ... (embedding vector(3))DDL works end to end.Testing
datafusion-postgres/tests/pgvector.rs: end-to-end tests with thesynchronous rust-postgres client and the official
pgvectorcrate(
pgvector::Vector) against a live server: DDL, literal / prepared-parameterINSERT,
prepare()introspection, psql-style text reads, binary resultdecoding and nearest-neighbour queries.
Notes
pgvector(default off), mirroring the existingpostgisfeature.vectortype only;halfvec/sparsevec/bit, indexes,and full constraint enforcement are follow-up work.