From 285cbc4bec901ff5aa129af7fd518e38849fd8c9 Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Fri, 4 Sep 2026 14:27:19 +0800 Subject: [PATCH 01/10] feat: pgvector compatibility (vector(n) type, distance ops, INSERT) 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. --- .gitignore | 1 + Cargo.lock | 146 +++++- arrow-pg/Cargo.toml | 3 + arrow-pg/src/datatypes.rs | 156 +++++- arrow-pg/src/datatypes/df.rs | 204 ++++++++ arrow-pg/src/encoder.rs | 321 +++++++++++- datafusion-pg-catalog/Cargo.toml | 6 + datafusion-pg-catalog/src/pg_catalog.rs | 149 +++++- .../src/pg_catalog/oid_type_planner.rs | 114 +++++ datafusion-pg-catalog/src/sql.rs | 5 + datafusion-pg-catalog/src/sql/parser.rs | 52 +- datafusion-pg-catalog/src/sql/rules.rs | 484 ++++++++++++++++++ .../src/sql/vector_insert.rs | 398 ++++++++++++++ datafusion-postgres-cli/Cargo.toml | 2 +- datafusion-postgres/Cargo.toml | 5 + datafusion-postgres/src/handlers.rs | 80 ++- datafusion-postgres/src/planner.rs | 81 +++ datafusion-postgres/tests/pgvector.rs | 455 ++++++++++++++++ 18 files changed, 2584 insertions(+), 78 deletions(-) create mode 100644 datafusion-pg-catalog/src/sql/vector_insert.rs create mode 100644 datafusion-postgres/tests/pgvector.rs diff --git a/.gitignore b/.gitignore index c68b88b2..1b98d041 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ .direnv .envrc .vscode +.idea .aider* /tests-integration/test_env diff --git a/Cargo.lock b/Cargo.lock index 649a9a65..eb0a77d7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -950,7 +950,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" dependencies = [ "chrono", - "phf", + "phf 0.12.1", ] [[package]] @@ -2497,6 +2497,7 @@ dependencies = [ "rustls-pemfile", "rustls-pki-types", "tokio", + "tokio-postgres", "tokio-rustls", ] @@ -3071,7 +3072,7 @@ checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", ] [[package]] @@ -3528,11 +3529,12 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.83" +version = "0.3.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "464a3709c7f55f1f721e5389aa6ea4e3bc6aba669353300af094b29ffbdde1d8" +checksum = "0e0c1080212aad755ea003d18543e8768dd432c48819efd73a7bf1e39b7a5a3a" dependencies = [ - "once_cell", + "cfg-if", + "futures-util", "wasm-bindgen", ] @@ -3666,6 +3668,15 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "libredox" +version = "0.1.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d8f1ea3f21fd3405dcaf6c9b5c1630af9afc422d9073ea39c5f6d6c772e08ed" +dependencies = [ + "libc", +] + [[package]] name = "libtest-mimic" version = "0.8.2" @@ -3772,7 +3783,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" dependencies = [ "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -3846,6 +3857,24 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.10.0", +] + +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.32.2" @@ -4051,7 +4080,17 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" dependencies = [ - "phf_shared", + "phf_shared 0.12.1", +] + +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_shared 0.13.1", + "serde", ] [[package]] @@ -4063,6 +4102,15 @@ dependencies = [ "siphasher", ] +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.13" @@ -5072,6 +5120,32 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "tokio-postgres" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf 0.13.1", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand 0.10.2", + "socket2", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-rustls" version = "0.26.5" @@ -5294,6 +5368,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.1+wasi-0.2.4" @@ -5312,11 +5395,20 @@ dependencies = [ "wit-bindgen 0.51.0", ] +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" -version = "0.2.106" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d759f433fa64a2d763d1340820e46e111a7a5ab75f993d1852d70b03dbb80fd" +checksum = "1b70935747edd64d89de3efa29d73789b806c15798f8e7dca4d8ac356b50ce70" dependencies = [ "cfg-if", "once_cell", @@ -5328,22 +5420,19 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.56" +version = "0.4.77" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "836d9622d604feee9e5de25ac10e3ea5f2d65b41eac0d9ce72eb5deae707ce7c" +checksum = "6b7777d5cc23d0e91404e53ce2d5e8ec7acae3026b16233dba62cd3246457950" dependencies = [ - "cfg-if", "js-sys", - "once_cell", "wasm-bindgen", - "web-sys", ] [[package]] name = "wasm-bindgen-macro" -version = "0.2.106" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48cb0d2638f8baedbc542ed444afc0644a29166f1595371af4fecf8ce1e7eeb3" +checksum = "77775f8f3f7217702089053b94958f8f54061a3f663417df76e19cbdcca29bc1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -5351,9 +5440,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.106" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cefb59d5cd5f92d9dcf80e4683949f15ca4b511f4ac0a6e14d4e1ac60c6ecd40" +checksum = "e11d33f857dc2fb11b8bc75aee111aa9cbeb12cd9f25efd3d4c2a3dd4e235284" dependencies = [ "bumpalo", "proc-macro2", @@ -5364,9 +5453,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.106" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbc538057e648b67f72a982e708d485b2efa771e1ac05fec311f9f63e5800db4" +checksum = "7ef64dbcc55df09c7e5a46182d181c2cfa3e925f3da937ea764728b4bbb9dcbf" dependencies = [ "unicode-ident", ] @@ -5407,9 +5496,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.83" +version = "0.3.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b32828d774c412041098d182a8b38b16ea816958e07cf40eec2bc080ae137ac" +checksum = "c435338968042f4f59a557f690a253676d47ce13ceb55d70100e7facf6620a30" dependencies = [ "js-sys", "wasm-bindgen", @@ -5425,6 +5514,19 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "whoami" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite", + "web-sys", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/arrow-pg/Cargo.toml b/arrow-pg/Cargo.toml index 98a0ddfd..3d0b21ea 100644 --- a/arrow-pg/Cargo.toml +++ b/arrow-pg/Cargo.toml @@ -16,6 +16,9 @@ rust-version.workspace = true default = ["arrow"] arrow = ["dep:arrow"] datafusion = ["dep:datafusion"] +# pgvector compatibility: teach the type mapping/encoding layer about the +# pgvector `vector` wire type. See datafusion-pg-catalog `pgvector` feature. +pgvector = [] postgis = ["postgres-types/with-geo-types-0_7", "dep:geoarrow", "dep:geoarrow-schema", "dep:postgis", "dep:geo-postgis", "pgwire/pg-type-postgis", "dep:geo-traits"] [dependencies] diff --git a/arrow-pg/src/datatypes.rs b/arrow-pg/src/datatypes.rs index e80095a9..4de2d435 100644 --- a/arrow-pg/src/datatypes.rs +++ b/arrow-pg/src/datatypes.rs @@ -136,15 +136,47 @@ pub fn into_pg_type(arrow_type: &DataType) -> PgWireResult { /// so untagged columns are unaffected. pub const PG_OID_ALIAS_KEY: &str = "pg.oid_alias"; -/// Map a field's [`PG_OID_ALIAS_KEY`] metadata, when present and recognized, to -/// the matching Postgres alias [`Type`] (e.g. `regtype` -> OID 2206). +/// Field metadata key marking an Arrow list/fixed-size-list field as a +/// pgvector `vector` column. /// -/// Returns `None` when the metadata is absent or names an unrecognized alias, -/// so the caller can fall back to the physical-type mapping (`Int32` -> -/// `INT4`). -fn pg_alias_type(field: &Field) -> Option { - let kind = field.metadata().get(PG_OID_ALIAS_KEY)?; - Some(match kind.as_str() { +/// This is a **cross-crate contract**: the lower-level `arrow-pg` crate does +/// not depend on `datafusion-pg-catalog`, which writes the metadata when it +/// plans the `vector(n)` SQL type. Columns carrying this marker are reported to +/// the client as the pgvector `vector` type and encoded using the pgvector text +/// format `[1,2,3]` instead of the ordinary Postgres array format `{1,2,3}`. +/// +/// The value is currently always `"vector"`. +#[cfg(feature = "pgvector")] +pub const PG_VECTOR_KEY: &str = "pg.vector"; + +/// Fixed OID reported for pgvector `vector` columns. +/// +/// Real pgvector assigns OIDs dynamically when the extension is installed, so +/// there is no universally canonical value. We pick the OID a default install +/// commonly lands on (the first user-defined type, right above +/// `FirstNormalObjectId`), which is what most name-based clients / hardcoded +/// drivers expect. Adjust if you run against clients that hardcode another OID. +#[cfg(feature = "pgvector")] +pub const PG_VECTOR_TYPE_OID: u32 = 16385; + +/// Field metadata key marking a UTF-8 Arrow field that models a Postgres +/// internal `"char"` column (e.g. `pg_type.typtype`). +/// +/// Postgres' internal `"char"` is a one-byte type whose wire encoding differs +/// from both `text` and `int2`: the binary form is the raw byte of the +/// character while the text form is the single character itself. Clients +/// (e.g. tokio-postgres) introspecting `pg_catalog` decode these columns as +/// `CHAR`. Columns exported from a real database carry the character in a UTF-8 +/// string, so this marker lets `arrow-pg` report the wire type as `CHAR` and +/// pick the right encoding per result format. +pub const PG_CHAR_KEY: &str = "pg.char"; + +/// Map an oid-alias kind name ([`PG_OID_ALIAS_KEY`] value) to its Postgres +/// [`Type`], e.g. `"oid"` -> OID 26, `"regtype"` -> OID 2206. +/// +/// Returns `None` for unrecognized kinds. +pub fn pg_type_for_alias_kind(kind: &str) -> Option { + Some(match kind { "oid" => Type::OID, "regproc" => Type::REGPROC, "regprocedure" => Type::REGPROCEDURE, @@ -161,6 +193,39 @@ fn pg_alias_type(field: &Field) -> Option { }) } +/// Map a field's [`PG_OID_ALIAS_KEY`] metadata, when present and recognized, to +/// the matching Postgres alias [`Type`] (e.g. `regtype` -> OID 2206). +/// +/// Returns `None` when the metadata is absent or names an unrecognized alias, +/// so the caller can fall back to the physical-type mapping (`Int32` -> +/// `INT4`). +fn pg_alias_type(field: &Field) -> Option { + let kind = field.metadata().get(PG_OID_ALIAS_KEY)?; + pg_type_for_alias_kind(kind.as_str()) +} + +/// True when `field` carries the [`PG_CHAR_KEY`] marker (Postgres `"char"`). +pub fn is_pg_char_field(field: &Field) -> bool { + field.metadata().get(PG_CHAR_KEY).is_some() +} + +/// True when `field` carries the pgvector [`PG_VECTOR_KEY`] metadata marker. +#[cfg(feature = "pgvector")] +pub fn is_pg_vector_field(field: &Field) -> bool { + matches!(field.metadata().get(PG_VECTOR_KEY), Some(kind) if kind == "vector") +} + +/// The pgwire [`Type`] reported for a pgvector `vector` column. +#[cfg(feature = "pgvector")] +pub fn pg_vector_type() -> Type { + Type::new( + "vector".to_string(), + PG_VECTOR_TYPE_OID, + Kind::Simple, + "public".to_string(), + ) +} + pub fn field_into_pg_type(field: &Arc) -> PgWireResult { // A `pg.oid_alias`-tagged Int32 field is reported as its Postgres alias // type (regtype -> OID 2206, regclass -> OID 2205, ...) instead of the @@ -172,6 +237,20 @@ pub fn field_into_pg_type(field: &Arc) -> PgWireResult { return Ok(alias_type); } + // A pg.char-tagged UTF-8 field is reported as the Postgres internal + // `"char"` type (OID 18) so clients decode it as a one-byte char. + if is_pg_char_field(field) { + return Ok(Type::CHAR); + } + + // A pg.vector-tagged list column is reported as the pgvector `vector` type + // instead of the physical float4[] array, so RowDescription and result + // schemas match a real pgvector backend. + #[cfg(feature = "pgvector")] + if is_pg_vector_field(field) { + return Ok(pg_vector_type()); + } + let arrow_type = field.data_type(); match field.extension_type_name() { @@ -219,8 +298,20 @@ pub fn arrow_schema_to_pg_fields( .enumerate() .map(|(idx, f)| { let pg_type = field_into_pg_type(f)?; - let mut field_info = - FieldInfo::new(f.name().into(), None, None, pg_type, format.format_for(idx)); + + // pgvector `vector` has no binary wire encoding implemented yet, so + // always negotiate the text format (`[1,2,3]`) for vector columns, + // even when the client asked for binary. + #[cfg(feature = "pgvector")] + let col_format = if is_pg_vector_field(f) { + pgwire::api::results::FieldFormat::Text + } else { + format.format_for(idx) + }; + #[cfg(not(feature = "pgvector"))] + let col_format = format.format_for(idx); + + let mut field_info = FieldInfo::new(f.name().into(), None, None, pg_type, col_format); if let Some(data_format_options) = &data_format_options { field_info = field_info.with_format_options(data_format_options.clone()); } @@ -361,4 +452,49 @@ mod tests { let field = Arc::new(Field::new("c", DataType::Int32, false)); assert_eq!(field_into_pg_type(&field).unwrap(), Type::INT4); } + + #[cfg(feature = "pgvector")] + mod vector { + use super::*; + + fn vector_field() -> Arc { + use std::collections::HashMap; + Arc::new( + Field::new( + "embedding", + DataType::FixedSizeList( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + ), + true, + ) + .with_metadata(HashMap::from([( + PG_VECTOR_KEY.to_string(), + "vector".to_string(), + )])), + ) + } + + #[test] + fn pg_vector_field_maps_to_custom_vector_type() { + let ty = field_into_pg_type(&vector_field()).unwrap(); + assert_eq!(ty.name(), "vector"); + assert_eq!(ty.schema(), "public"); + assert_eq!(ty.oid(), PG_VECTOR_TYPE_OID); + } + + #[test] + fn plain_fixed_size_list_stays_float4_array() { + let field = Arc::new(Field::new( + "embedding", + DataType::FixedSizeList( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + ), + true, + )); + assert!(!is_pg_vector_field(&field)); + assert_eq!(field_into_pg_type(&field).unwrap(), Type::FLOAT4_ARRAY); + } + } } diff --git a/arrow-pg/src/datatypes/df.rs b/arrow-pg/src/datatypes/df.rs index ae7a34bb..f32558c0 100644 --- a/arrow-pg/src/datatypes/df.rs +++ b/arrow-pg/src/datatypes/df.rs @@ -272,6 +272,151 @@ fn coerce_interval_value(value: Option, target: &DataType) -> PgWireRe } } +/// Decode a pgvector `vector` parameter value sent by a client over the wire. +#[cfg(feature = "pgvector")] +#[derive(Debug)] +struct VectorParam(Vec); + +#[cfg(feature = "pgvector")] +impl VectorParam { + /// The pgvector binary layout is a big-endian `int16` dimension followed by + /// that many big-endian IEEE float32 elements. + fn from_binary(raw: &[u8]) -> Result> { + if raw.len() < 2 { + return Err("vector parameter binary payload too short".into()); + } + let dim = i16::from_be_bytes([raw[0], raw[1]]); + if dim <= 0 { + return Err("vector parameter dimension must be positive".into()); + } + let dim = dim as usize; + if raw.len() != 2 + dim * 4 { + return Err("vector parameter binary payload has wrong length".into()); + } + let mut values = Vec::with_capacity(dim); + for i in 0..dim { + let off = 2 + i * 4; + values.push(f32::from_be_bytes([ + raw[off], + raw[off + 1], + raw[off + 2], + raw[off + 3], + ])); + } + Ok(VectorParam(values)) + } + + /// The pgvector text form is `[1,2,3]`. + fn from_text(raw: &[u8]) -> Result> { + let text = std::str::from_utf8(raw)?; + let text = text.trim(); + if !(text.starts_with('[') && text.ends_with(']')) { + return Err("vector parameter text must look like [1,2,3]".into()); + } + let inner = text[1..text.len() - 1].trim(); + if inner.is_empty() { + return Err("vector parameter must not be empty".into()); + } + let mut values = Vec::new(); + for part in inner.split(',') { + values.push( + part.trim().parse::().map_err(|_| { + format!("invalid vector element '{}' in parameter", part.trim()) + })?, + ); + } + Ok(VectorParam(values)) + } +} + +#[cfg(feature = "pgvector")] +impl<'a> postgres_types::FromSql<'a> for VectorParam { + fn from_sql( + _ty: &Type, + raw: &'a [u8], + ) -> Result> { + VectorParam::from_binary(raw) + } + + fn accepts(ty: &Type) -> bool { + ty.oid() == crate::datatypes::PG_VECTOR_TYPE_OID + } +} + +#[cfg(feature = "pgvector")] +impl<'a> pgwire::types::FromSqlText<'a> for VectorParam { + fn from_sql_text( + _ty: &Type, + input: &'a [u8], + _format_options: &FormatOptions, + ) -> Result> { + VectorParam::from_text(input) + } +} + +/// Turn a decoded vector parameter into the DataFusion scalar expected by the +/// `vector(n)` column it is bound to. +#[cfg(feature = "pgvector")] +fn vector_param_to_scalar( + value: Option, + inferred: Option<&DataType>, +) -> PgWireResult { + use datafusion::arrow::array::FixedSizeListArray; + use datafusion::arrow::buffer::NullBuffer; + use datafusion::arrow::datatypes::Field as ArrowField; + + let (dim, inner) = match inferred { + Some(DataType::FixedSizeList(field, n)) => (*n as usize, Arc::new(field.as_ref().clone())), + Some(DataType::List(field)) => (0, Arc::new(field.as_ref().clone())), + _ => ( + 0, + Arc::new(ArrowField::new_list_field(DataType::Float32, true)), + ), + }; + + let Some(VectorParam(values)) = value else { + // SQL NULL vector: a single null FixedSizeList element. + let inner = ArrowField::new_list_field(DataType::Float32, true); + let values = datafusion::arrow::array::Float32Array::from(vec![0.0f32; dim]); + let array = FixedSizeListArray::try_new( + Arc::new(inner), + dim as i32, + Arc::new(values), + Some(NullBuffer::from(vec![false])), + ) + .map_err(|e| invalid_parameter_error(e.to_string()))?; + return Ok(ScalarValue::FixedSizeList(Arc::new(array))); + }; + + // Fixed-dimension vector column (`vector(n)`): honor the declared size. + if dim > 0 { + if values.len() != dim { + return Err(invalid_parameter_error(format!( + "expected {dim} dimensions for vector parameter, got {}", + values.len() + ))); + } + let array = FixedSizeListArray::try_new( + Arc::new(inner.as_ref().clone()), + dim as i32, + Arc::new(datafusion::arrow::array::Float32Array::from(values)), + None, + ) + .map_err(|e| invalid_parameter_error(e.to_string()))?; + return Ok(ScalarValue::FixedSizeList(Arc::new(array))); + } + + // Bare `vector` (no declared dimension): a plain list of float32. + let scalars: Vec = values + .into_iter() + .map(|v| ScalarValue::Float32(Some(v))) + .collect(); + Ok(ScalarValue::List(ScalarValue::new_list_nullable( + &scalars, + &DataType::Float32, + ))) +} + /// Deserialize client provided parameter data. /// /// First we try to use the type information from `pg_type_hint`, which is @@ -279,17 +424,53 @@ fn coerce_interval_value(value: Option, target: &DataType) -> PgWireRe /// If the type is empty or unknown, we fallback to datafusion inferenced type /// from `inferenced_types`. /// An error will be raised when neither sources can provide type information. +/// +/// This is a convenience wrapper that passes no server-decided parameter types +/// (see [`deserialize_parameters_with_server_types`]). pub fn deserialize_parameters( portal: &Portal, inferenced_types: &[Option<&DataType>], ) -> PgWireResult +where + S: Clone, +{ + deserialize_parameters_with_server_types(portal, inferenced_types, &[]) +} + +/// Deserialize client provided parameter data using the server-decided +/// parameter wire types. +/// +/// pgwire does not persist the parameter types it advertised in +/// `ParameterDescription` onto the portal, so the caller recomputes them from +/// the logical plan (see the planner's parameter overrides) and passes them in +/// `server_types` (positionally aligned with the parameters). A value of `None` +/// for a position means "no server-decided type", falling back to the client's +/// `pg_type_hint` and then the DataFusion-inferred type. +/// +/// This is what lets semantically-typed parameters (e.g. a pgvector `vector` +/// bound to a `vector(n)` column) be decoded with their real wire format rather +/// than the physical Arrow type mapping (`FixedSizeList` -> `float4[]`). +pub fn deserialize_parameters_with_server_types( + portal: &Portal, + inferenced_types: &[Option<&DataType>], + server_types: &[Option<&Type>], +) -> PgWireResult where S: Clone, { fn get_pg_type( + server_type: Option<&Type>, pg_type_hint: Option, inferenced_type: Option<&DataType>, ) -> PgWireResult { + // A concrete server-decided type wins -- it reflects the parameter's + // semantically-typed column (oid-alias, pgvector, ...). UNKNOWN is not + // authoritative so the historical hint/inferred fallbacks still apply. + if let Some(ty) = server_type + && *ty != Type::UNKNOWN + { + return Ok(ty.clone()); + } if let Some(ty) = pg_type_hint { Ok(ty.clone()) } else if let Some(infer_type) = inferenced_type { @@ -303,7 +484,9 @@ where let mut deserialized_params = Vec::with_capacity(param_len); for i in 0..param_len { let inferenced_type = inferenced_types.get(i).and_then(|v| v.to_owned()); + let server_type = server_types.get(i).and_then(|t| *t); let pg_type = get_pg_type( + server_type, portal .statement .parameter_types @@ -314,6 +497,15 @@ where // enumerate all supported parameter types and deserialize the // type to ScalarValue, with data coercion when server-inferred // types are available + + // pgvector `vector` parameters (binary or text form). + #[cfg(feature = "pgvector")] + if pg_type.oid() == crate::datatypes::PG_VECTOR_TYPE_OID { + let value: Option = portal.parameter(i, &pg_type)?; + deserialized_params.push(vector_param_to_scalar(value, inferenced_type)?); + continue; + } + match pg_type { Type::BOOL => { let value = portal.parameter::(i, &pg_type)?; @@ -846,6 +1038,18 @@ where // Store MAC addresses as strings for now deserialized_params.push(ScalarValue::Utf8(value)); } + // PostgreSQL `oid` (unsigned 32-bit). Drivers bind these with u32; + // DataFusion catalog oid columns are Int32, so coerce when the + // inferred type is Int32. + Type::OID => { + let value = portal.parameter::(i, &pg_type)?; + match inferenced_type { + Some(DataType::Int32) => { + deserialized_params.push(ScalarValue::Int32(value.map(|v| v as i32))); + } + _ => deserialized_params.push(ScalarValue::UInt32(value)), + } + } // TODO: add more advanced types (composite types, ranges, etc.) _ => { // the client didn't provide type information and we are also diff --git a/arrow-pg/src/encoder.rs b/arrow-pg/src/encoder.rs index 91f35f67..4f81233e 100644 --- a/arrow-pg/src/encoder.rs +++ b/arrow-pg/src/encoder.rs @@ -8,7 +8,7 @@ use chrono::{NaiveDate, NaiveDateTime}; #[cfg(feature = "datafusion")] use datafusion::arrow::{array::*, datatypes::*}; use pg_interval::Interval as PgInterval; -use pgwire::api::results::{CopyEncoder, DataRowEncoder, FieldInfo}; +use pgwire::api::results::{CopyEncoder, DataRowEncoder, FieldFormat, FieldInfo}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::copy::CopyData; use pgwire::messages::data::DataRow; @@ -159,6 +159,130 @@ fn get_large_binary_value(arr: &Arc, idx: usize) -> Option<&[u8]> { }) } +/// Encode a Postgres `oid` column (stored as Arrow `Int32` with `pg.oid_alias +/// = "oid"` metadata) whose result format is binary. +/// +/// PostgreSQL's `oid` is unsigned 32-bit, and drivers (tokio-postgres, ...) +/// encode/decode it with `u32`. Arrow stores our catalog oids as `Int32`, so +/// for the binary protocol the value must be re-encoded as `u32`, otherwise the +/// `i32` ToSql impl rejects the OID type. +fn encode_pg_oid_binary( + encoder: &mut T, + arr: &Arc, + idx: usize, + pg_field: &FieldInfo, +) -> PgWireResult<()> { + if arr.is_null(idx) { + return encoder.encode_field(&None::, pg_field); + } + let value = get_i32_value(arr, idx).unwrap_or(0) as u32; + encoder.encode_field(&Some(value), pg_field) +} + +/// Encode a Postgres internal `"char"` column (stored as a single-character +/// UTF-8 string tagged with the `pg.char` metadata). +/// +/// Binary format is the raw byte of the character (decoded by clients as an +/// `i8`/`char`); text format is the character itself. +fn encode_pg_char( + encoder: &mut T, + arr: &Arc, + idx: usize, + pg_field: &FieldInfo, +) -> PgWireResult<()> { + if arr.is_null(idx) { + return encoder.encode_field(&None::, pg_field); + } + let text = get_utf8_value(arr, idx) + .ok_or_else(|| PgWireError::ApiError(ToSqlError::from("pg.char column must be UTF-8")))?; + if pg_field.format() == FieldFormat::Binary { + let byte = text + .as_bytes() + .first() + .copied() + .ok_or_else(|| PgWireError::ApiError(ToSqlError::from("pg.char value is empty")))? + as i8; + encoder.encode_field(&Some(byte), pg_field) + } else { + encoder.encode_field(&Some(text.to_string()), pg_field) + } +} + +/// Render the pgvector text form of a vector: `[1,2,3]`. +/// +/// Element formatting uses Rust's shortest round-trip `Display` for `f32` +/// (`1.0` -> `1`), matching pgvector's `vector_out`. +#[cfg(feature = "pgvector")] +fn format_pg_vector(values: &[f32]) -> String { + let inner = values + .iter() + .map(|v| v.to_string()) + .collect::>() + .join(","); + format!("[{inner}]") +} + +/// Encode a pgvector `vector` column (an Arrow `List`/`FixedSizeList` of +/// `Float32` tagged with the `pg.vector` field metadata) for a single row. +/// +/// The value is emitted in the pgvector text format `[1,2,3]`. The result +/// `FieldInfo` for vector columns is forced to the text format by +/// `arrow_schema_to_pg_fields`, so pgwire serializes the string verbatim via +/// `ToSqlText` regardless of the client's requested result format. +#[cfg(feature = "pgvector")] +fn encode_pg_vector( + encoder: &mut T, + arr: &Arc, + idx: usize, + pg_field: &FieldInfo, +) -> PgWireResult<()> { + if arr.is_null(idx) { + return encoder.encode_field(&None::, pg_field); + } + + fn row_values(arr: &Arc, idx: usize) -> PgWireResult> { + let values = match arr.data_type() { + DataType::FixedSizeList(_, _) => { + let list = arr.as_any().downcast_ref::().unwrap(); + let values = list + .values() + .as_any() + .downcast_ref::() + .ok_or_else(|| { + PgWireError::ApiError(ToSqlError::from( + "vector FixedSizeList values must be Float32", + )) + })?; + let size = list.value_length() as usize; + let start = idx * size; + (0..size).map(|i| values.value(start + i)).collect() + } + DataType::List(_) => { + let list = arr.as_any().downcast_ref::().unwrap(); + let value = list.value(idx); + let values = value + .as_any() + .downcast_ref::() + .ok_or_else(|| { + PgWireError::ApiError(ToSqlError::from( + "vector List values must be Float32", + )) + })?; + values.values().to_vec() + } + other => { + return Err(PgWireError::ApiError(ToSqlError::from(format!( + "vector column has unsupported arrow type {other}" + )))); + } + }; + Ok(values) + } + + let text = format_pg_vector(&row_values(arr, idx)?); + encoder.encode_field(&Some(text), pg_field) +} + fn get_date32_value(arr: &Arc, idx: usize) -> Option { if arr.is_null(idx) { return None; @@ -283,6 +407,34 @@ pub fn encode_value( ); } + // pgvector `vector` columns are tagged with the `pg.vector` field + // metadata. Route them through the vector encoder (text `[1,2,3]`) before + // the generic list handling below, which would otherwise emit them as a + // Postgres float4[] (`{1,2,3}`). + #[cfg(feature = "pgvector")] + if crate::datatypes::is_pg_vector_field(arrow_field) { + return encode_pg_vector(encoder, arr, idx, pg_field); + } + + // Postgres internal `"char"` columns (pg.typtype, ...): wire type CHAR with + // char-specific binary/text encoding. + if crate::datatypes::is_pg_char_field(arrow_field) { + return encode_pg_char(encoder, arr, idx, pg_field); + } + + // Postgres `oid` columns are `u32` on the wire but stored as Arrow Int32. + // Over the binary protocol the value must be sent as `u32` (the `i32` ToSql + // impl does not accept the OID type); text keeps the existing Int32 path. + if pg_field.format() == FieldFormat::Binary + && matches!(arrow_type, DataType::Int32) + && arrow_field + .metadata() + .get(crate::datatypes::PG_OID_ALIAS_KEY) + .is_some_and(|kind| kind == "oid") + { + return encode_pg_oid_binary(encoder, arr, idx, pg_field); + } + match arrow_type { DataType::Null => encoder.encode_field(&None::, pg_field)?, DataType::Boolean => encoder.encode_field(&get_bool_value(arr, idx), pg_field)?, @@ -485,7 +637,22 @@ pub fn encode_value( if arr.is_null(idx) { return encoder.encode_field(&None::<&[i8]>, pg_field); } - let array = arr.as_any().downcast_ref::().unwrap().value(idx); + // Extract this row's element slice from the actual list flavour. + // (FixedSizeList / LargeList were previously downcast to ListArray, + // which panics -- these are not ListArrays.) + let array = match arrow_type { + DataType::FixedSizeList(_, _) => arr + .as_any() + .downcast_ref::() + .unwrap() + .value(idx), + DataType::LargeList(_) => arr + .as_any() + .downcast_ref::() + .unwrap() + .value(idx), + _ => arr.as_any().downcast_ref::().unwrap().value(idx), + }; encode_list(encoder, array, pg_field)? } DataType::Struct(arrow_fields) => encode_struct(encoder, arr, idx, arrow_fields, pg_field)?, @@ -735,4 +902,154 @@ mod tests { Some(NaiveTime::from_hms_nano_opt(1, 2, 3, 1001001)).unwrap() ); } + + #[cfg(feature = "pgvector")] + mod vector { + use super::*; + use arrow::buffer::NullBuffer; + use bytes::BytesMut; + use pgwire::{api::results::FieldFormat, types::format::FormatOptions}; + use postgres_types::Type; + use std::collections::HashMap; + + #[derive(Default)] + struct TextCapture { + encoded: String, + } + + impl Encoder for TextCapture { + type Item = String; + + fn encode_field(&mut self, value: &T, pg_field: &FieldInfo) -> PgWireResult<()> + where + T: ToSql + ToSqlText + Sized, + { + let mut bytes = BytesMut::new(); + value + .to_sql_text(pg_field.datatype(), &mut bytes, &FormatOptions::default()) + .unwrap(); + self.encoded = String::from_utf8(bytes.to_vec()).unwrap(); + Ok(()) + } + + fn take_row(&mut self) -> Self::Item { + std::mem::take(&mut self.encoded) + } + } + + fn vector_arrow_field(metadata: bool) -> Field { + let mut field = Field::new( + "embedding", + DataType::FixedSizeList( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + ), + true, + ); + if metadata { + field = field.with_metadata(HashMap::from([( + crate::datatypes::PG_VECTOR_KEY.to_string(), + "vector".to_string(), + )])); + } + field + } + + #[test] + fn encodes_vector_fixed_size_list_as_pgvector_text() { + // Two rows: [1,2,3] and [4.5,0,7] + let values = Float32Array::from(vec![1.0, 2.0, 3.0, 4.5, 0.0, 7.0]); + let array: Arc = Arc::new( + FixedSizeListArray::try_new( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + Arc::new(values), + None, + ) + .unwrap(), + ); + + let arrow_field = vector_arrow_field(true); + let pg_field = FieldInfo::new( + "embedding".to_string(), + None, + None, + crate::datatypes::pg_vector_type(), + FieldFormat::Text, + ); + + let mut encoder = TextCapture::default(); + encode_value(&mut encoder, &array, 0, &arrow_field, &pg_field).unwrap(); + assert_eq!(encoder.encoded, "[1,2,3]"); + + let mut encoder = TextCapture::default(); + encode_value(&mut encoder, &array, 1, &arrow_field, &pg_field).unwrap(); + assert_eq!(encoder.encoded, "[4.5,0,7]"); + } + + #[test] + fn encodes_null_vector_as_null() { + // `NullBuffer::from(Vec)` treats `true` as valid. + let nulls = NullBuffer::from(vec![false, true]); + let values = Float32Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let array: Arc = Arc::new( + FixedSizeListArray::try_new( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + Arc::new(values), + Some(nulls), + ) + .unwrap(), + ); + + let arrow_field = vector_arrow_field(true); + let pg_field = FieldInfo::new( + "embedding".to_string(), + None, + None, + crate::datatypes::pg_vector_type(), + FieldFormat::Text, + ); + + // Row 0 is NULL: no value bytes. + let mut encoder = TextCapture::default(); + encode_value(&mut encoder, &array, 0, &arrow_field, &pg_field).unwrap(); + assert!(encoder.encoded.is_empty(), "NULL must emit no bytes"); + + // Row 1 is a valid vector. + let mut encoder = TextCapture::default(); + encode_value(&mut encoder, &array, 1, &arrow_field, &pg_field).unwrap(); + assert_eq!(encoder.encoded, "[4,5,6]"); + } + + #[test] + fn plain_fixed_size_list_encodes_as_float4_array() { + // Regression: FixedSizeList columns used to be downcast to + // ListArray and panic. A non-pgvector fixed-size float list must + // still encode (as a Postgres array). + let values = Float32Array::from(vec![1.0, 2.0, 3.0]); + let array: Arc = Arc::new( + FixedSizeListArray::try_new( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + Arc::new(values), + None, + ) + .unwrap(), + ); + + let arrow_field = vector_arrow_field(false); + let pg_field = FieldInfo::new( + "embedding".to_string(), + None, + None, + Type::FLOAT4_ARRAY, + FieldFormat::Text, + ); + + let mut encoder = TextCapture::default(); + encode_value(&mut encoder, &array, 0, &arrow_field, &pg_field).unwrap(); + assert_eq!(encoder.encoded, "{1.0,2.0,3.0}"); + } + } } diff --git a/datafusion-pg-catalog/Cargo.toml b/datafusion-pg-catalog/Cargo.toml index 9a6f8bf8..964bc7e9 100644 --- a/datafusion-pg-catalog/Cargo.toml +++ b/datafusion-pg-catalog/Cargo.toml @@ -26,6 +26,12 @@ log = "0.4" postgres-types.workspace = true tokio = { version = "1.52", features = ["sync"] } +[features] +# pgvector compatibility: recognize the pgvector `vector(n)` SQL type and the +# `<->` / `<#>` / `<=>` distance operators. Reuses DataFusion's built-in +# array_distance / inner_product / cosine_distance functions. +pgvector = [] + [dev-dependencies] env_logger = "0.11" criterion = { version = "0.8", features = ["html_reports"] } diff --git a/datafusion-pg-catalog/src/pg_catalog.rs b/datafusion-pg-catalog/src/pg_catalog.rs index 46cf3ccb..0f04332f 100644 --- a/datafusion-pg-catalog/src/pg_catalog.rs +++ b/datafusion-pg-catalog/src/pg_catalog.rs @@ -7,7 +7,7 @@ use datafusion::arrow::array::{ ArrayRef, AsArray, BooleanBuilder, Int32Builder, RecordBatch, StringArray, StringBuilder, as_boolean_array, }; -use datafusion::arrow::datatypes::{DataType, Field, Int32Type, SchemaRef}; +use datafusion::arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; use datafusion::arrow::ipc::reader::FileReader; use datafusion::catalog::streaming::StreamingTable; use datafusion::catalog::{MemTable, SchemaProvider, TableFunctionImpl}; @@ -652,7 +652,7 @@ pub struct PgCatalogStaticTables { impl PgCatalogStaticTables { pub fn try_new() -> Result { - Ok(Self { + let tables = Self { pg_aggregate: Self::create_arrow_table( include_bytes!(concat!( env!("CARGO_MANIFEST_DIR"), @@ -1061,13 +1061,156 @@ impl PgCatalogStaticTables { )) .to_vec(), )?, - }) + }; + + // pgvector support: expose the `vector` type (OID 16385) in `pg_type` + // so clients that resolve unknown result types (tokio-postgres, JDBC, + // ...) find it, and tag the internal `"char"` column for correct wire + // encoding. + #[cfg(feature = "pgvector")] + let tables = tables.with_pg_vector_support()?; + + Ok(tables) } /// Create table from dumped arrow data fn create_arrow_table(data_bytes: Vec) -> Result> { ArrowTable::from_ipc_data(data_bytes).map(Arc::new) } + + /// Return `self` with the pg_catalog pieces needed for pgvector-aware + /// clients that resolve unknown result types via `pg_type` introspection: + /// + /// * a `vector` type row (OID 16385) in `pg_type`, and + /// * the internal `"char"` wire type on `pg_type.typtype`. + /// + /// The vector row intentionally lives in `pg_catalog` (namespace OID 11) so + /// the `pg_type`/`pg_namespace` join the introspection query performs + /// resolves without depending on user-schema oids. + #[cfg(feature = "pgvector")] + fn with_pg_vector_support(mut self) -> Result { + const VECTOR_OID: i32 = 16385; + + // 1. Tag pg_type.typtype as a Postgres internal `"char"` column so the + // wire layer encodes it correctly for clients decoding `typtype`. + let pg_type = Arc::new(arrow_table_with_field_metadata( + &self.pg_type, + "typtype", + "pg.char", + "char", + )?); + + // 2. Append the `vector` row to pg_type (all other columns get safe + // defaults; the introspection query only reads the ones we set). + let pg_type = Arc::new(arrow_table_append_row(&pg_type, |field, scalar| { + match field.name().as_str() { + "oid" => *scalar = datafusion::scalar::ScalarValue::Int32(Some(VECTOR_OID)), + "typname" => { + *scalar = datafusion::scalar::ScalarValue::Utf8(Some("vector".to_string())) + } + "typtype" => *scalar = datafusion::scalar::ScalarValue::Utf8(Some("b".to_string())), + // pg_catalog namespace OID 11 (see the doc comment above). + "typnamespace" => *scalar = datafusion::scalar::ScalarValue::Int32(Some(11)), + _ => {} + } + })?); + + self.pg_type = pg_type; + Ok(self) + } +} + +/// Rebuild `table` with `key = value` metadata added to the field named `name`. +/// +/// Arrow keeps field metadata in the schema; this rebuilds the schema and every +/// record batch so the wire layer (`arrow-pg`) sees the marker when encoding. +#[cfg(feature = "pgvector")] +fn arrow_table_with_field_metadata( + table: &ArrowTable, + name: &str, + key: &str, + value: &str, +) -> Result { + let schema = table.schema(); + let fields = schema + .fields() + .iter() + .map(|field| { + if field.name() == name { + let mut metadata = field.metadata().clone(); + metadata.insert(key.to_string(), value.to_string()); + (**field).clone().with_metadata(metadata) + } else { + (**field).clone() + } + }) + .collect::>(); + let new_schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); + let mut batches = Vec::with_capacity(table.data().len()); + for batch in table.data() { + batches.push( + RecordBatch::try_new(Arc::clone(&new_schema), batch.columns().to_vec()) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?, + ); + } + Ok(ArrowTable { + schema: new_schema, + data: batches, + }) +} + +/// Return a copy of `table` with one extra row appended. `fill` lets callers +/// override the default scalar produced for each column. +/// +/// `fill` is invoked for every field with a fresh default scalar (int 0 / +/// string "" / bool false / ...), matching the field's data type; non-default +/// catalog values are set by the caller. +#[cfg(feature = "pgvector")] +fn arrow_table_append_row( + table: &ArrowTable, + mut fill: impl FnMut(&Field, &mut datafusion::scalar::ScalarValue), +) -> Result { + let schema = table.schema(); + let columns = schema + .fields() + .iter() + .map(|field| { + let mut scalar = default_scalar(field.data_type())?; + fill(field, &mut scalar); + scalar.to_array_of_size(1) + }) + .collect::>>()?; + let row = RecordBatch::try_new(Arc::clone(&schema), columns) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + let mut data = table.data().to_vec(); + data.push(row); + Ok(ArrowTable { + schema: Arc::clone(&schema), + data, + }) +} + +/// A single default [`ScalarValue`] for a catalog column of `data_type`. +/// +/// Only the data types present in the exported `pg_type` schema are handled; +/// anything else yields a `NotImplemented` error rather than guessing. +#[cfg(feature = "pgvector")] +fn default_scalar(data_type: &DataType) -> Result { + use datafusion::scalar::ScalarValue; + Ok(match data_type { + DataType::Null => ScalarValue::Null, + DataType::Boolean => ScalarValue::Boolean(Some(false)), + DataType::Int16 => ScalarValue::Int16(Some(0)), + DataType::Int32 => ScalarValue::Int32(Some(0)), + DataType::Int64 => ScalarValue::Int64(Some(0)), + DataType::Utf8 => ScalarValue::Utf8(Some(String::new())), + DataType::LargeUtf8 => ScalarValue::LargeUtf8(Some(String::new())), + other => { + return Err(DataFusionError::NotImplemented(format!( + "no default catalog scalar for {other:?}" + ))); + } + }) } pub fn create_current_schemas_udf() -> ScalarUDF { diff --git a/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs b/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs index 08a48b4c..1e10ebb9 100644 --- a/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs +++ b/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs @@ -44,6 +44,11 @@ use datafusion::sql::sqlparser::ast::{DataType as SQLDataType, ObjectNamePart}; use crate::pg_catalog::oid_field::{self, OID_ALIAS_TYPE_NAMES}; +/// Field metadata key (arrow-pg contract) marking an Arrow list field as a +/// pgvector `vector`. Must match `arrow_pg::datatypes::PG_VECTOR_KEY`. +#[cfg(feature = "pgvector")] +const PG_VECTOR_KEY: &str = "pg.vector"; + /// Recognize Postgres type names DataFusion rejects and map them to Arrow /// types/metadata at planning time. #[derive(Debug, Default)] @@ -125,6 +130,60 @@ impl PgOidTypePlanner { } builtin_arrow_type(type_name) } + + /// True when `sql_type` names the pgvector `vector` type (optionally + /// schema-qualified, e.g. `public.vector`). + #[cfg(feature = "pgvector")] + fn is_vector_type(sql_type: &SQLDataType) -> bool { + let SQLDataType::Custom(name, _) = sql_type else { + return false; + }; + name.0 + .last() + .and_then(|part| part.as_ident()) + .is_some_and(|ident| ident.value.eq_ignore_ascii_case("vector")) + } + + /// Map the pgvector `vector` / `vector(n)` SQL type to an Arrow field. + /// + /// `vector(n)` is a fixed-dimension vector and maps to + /// `FixedSizeList(Float32, n)`; a bare `vector` (no declared dimension) + /// maps to `List(Float32)`. Both carry the `pg.vector` field metadata that + /// `arrow-pg` uses to report the pgwire `vector` type and encode values in + /// the pgvector text format. + /// + /// Returns `None` when the type is not a vector we recognize, so the caller + /// can fall back to the other planners. A `vector` with an unparseable / + /// non-positive dimension is also left untouched (DataFusion will then + /// reject the unknown type with its own error). + #[cfg(feature = "pgvector")] + fn vector_field(sql_type: &SQLDataType) -> Option> { + if !Self::is_vector_type(sql_type) { + return None; + } + let SQLDataType::Custom(_, modifiers) = sql_type else { + return None; + }; + + let element = Field::new_list_field(DataType::Float32, true); + let arrow_type = match modifiers.as_slice() { + [] => DataType::List(Arc::new(element)), + [dim] => { + let dim: i32 = dim.trim().parse().ok()?; + if dim <= 0 { + return None; + } + DataType::FixedSizeList(Arc::new(element), dim) + } + _ => return None, + }; + + let mut metadata = std::collections::HashMap::new(); + metadata.insert(PG_VECTOR_KEY.to_string(), "vector".to_string()); + Some(Arc::new( + Field::new("", arrow_type, true).with_metadata(metadata), + )) + } } /// Arrow [`DataType`] for a Postgres pg_catalog builtin type name, or `None` @@ -159,6 +218,12 @@ fn builtin_arrow_type(name: &str) -> Option { impl TypePlanner for PgOidTypePlanner { fn plan_type_field(&self, sql_type: &SQLDataType) -> Result>> { + // 0. pgvector `vector(n)` / `vector` -> FixedSizeList/List of Float32 + // tagged with `pg.vector` metadata (see the arrow-pg contract). + #[cfg(feature = "pgvector")] + if let Some(field) = Self::vector_field(sql_type) { + return Ok(Some(field)); + } // 1. Scalar oid-alias types (regclass, oid, ...) -> int4 with kind // metadata. Arrays of these (`regtype[]`) are handled by DataFusion // recursing into this planner for the element type, so no Array arm @@ -319,4 +384,53 @@ mod tests { let field = planner.plan_type_field(&dt).unwrap().unwrap(); assert_eq!(field.data_type(), &DataType::Int16); } + + #[cfg(feature = "pgvector")] + #[test] + fn vector_with_dimension_is_fixed_size_list_of_float32() { + let planner = PgOidTypePlanner; + let dt = cast_target_type("SELECT 'x'::vector(3) AS c"); + let field = planner.plan_type_field(&dt).unwrap().unwrap(); + + let expected = + DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Float32, true)), 3); + assert_eq!(field.data_type(), &expected); + assert_eq!( + field + .metadata() + .get(super::PG_VECTOR_KEY) + .map(String::as_str), + Some("vector") + ); + } + + #[cfg(feature = "pgvector")] + #[test] + fn bare_vector_is_a_list_of_float32() { + let planner = PgOidTypePlanner; + let dt = cast_target_type("SELECT 'x'::vector AS c"); + let field = planner.plan_type_field(&dt).unwrap().unwrap(); + + let expected = DataType::List(Arc::new(Field::new_list_field(DataType::Float32, true))); + assert_eq!(field.data_type(), &expected); + assert_eq!( + field + .metadata() + .get(super::PG_VECTOR_KEY) + .map(String::as_str), + Some("vector") + ); + } + + #[cfg(feature = "pgvector")] + #[test] + fn vector_with_invalid_dimension_falls_through() { + let planner = PgOidTypePlanner; + let dt = cast_target_type("SELECT 'x'::vector(0) AS c"); + assert_eq!(planner.plan_type_field(&dt).unwrap(), None); + + // A non-vector custom type is untouched. + let dt = cast_target_type("SELECT 'x'::public.my_type AS c"); + assert_eq!(planner.plan_type_field(&dt).unwrap(), None); + } } diff --git a/datafusion-pg-catalog/src/sql.rs b/datafusion-pg-catalog/src/sql.rs index 42c6a178..452b0807 100644 --- a/datafusion-pg-catalog/src/sql.rs +++ b/datafusion-pg-catalog/src/sql.rs @@ -1,3 +1,8 @@ mod parser; pub use parser::PostgresCompatibilityParser; pub mod rules; + +#[cfg(feature = "pgvector")] +mod vector_insert; +#[cfg(feature = "pgvector")] +pub use vector_insert::rewrite_vector_insert; diff --git a/datafusion-pg-catalog/src/sql/parser.rs b/datafusion-pg-catalog/src/sql/parser.rs index d8e9a519..03a3eda7 100644 --- a/datafusion-pg-catalog/src/sql/parser.rs +++ b/datafusion-pg-catalog/src/sql/parser.rs @@ -18,6 +18,8 @@ use super::rules::ResolveUnqualifiedIdentifier; use super::rules::RewriteArrayAnyAllOperation; use super::rules::RewritePgCatalogOperator; use super::rules::RewriteRegCastToSubquery; +#[cfg(feature = "pgvector")] +use super::rules::RewriteVectorOperators; use super::rules::SqlStatementRewriteRule; use super::rules::StripCallableQualifier; use super::rules::StripCollate; @@ -309,28 +311,38 @@ impl PostgresCompatibilityParser { )); } + #[cfg_attr(not(feature = "pgvector"), allow(unused_mut))] + let mut rewrite_rules: Vec> = vec![ + // The blacklist substitution in `parse()` runs before any of + // these rules, so by the time they see the statement any + // blacklisted fragment has already been replaced. + Arc::new(AliasDuplicatedProjectionRewrite), + Arc::new(ResolveUnqualifiedIdentifier), + Arc::new(RewriteArrayAnyAllOperation), + Arc::new(PrependUnqualifiedPgTableName), + Arc::new(StripCallableQualifier), + Arc::new(FixArrayLiteral), + Arc::new(CurrentUserVariableToSessionUserFunctionCall), + Arc::new(StripCollate), + Arc::new(RewritePgCatalogOperator), + // Resolve forward oid-alias casts (`'x'::regclass`, ...) to oid + // values BEFORE RemoveSubqueryFromProjection runs, so the + // emitted scalar subqueries it produces get its LIMIT 1 stamp. + Arc::new(RewriteRegCastToSubquery::new()), + Arc::new(RemoveSubqueryFromProjection), + Arc::new(FixVersionColumnName), + ]; + + // pgvector support: distance operators (`<->` / `<#>` / `<=>`) and + // vector literals. Runs last -- it needs to see oid/array rewrites in + // operands already applied, and rewrites operators that no other rule + // touches. + #[cfg(feature = "pgvector")] + rewrite_rules.push(Arc::new(RewriteVectorOperators)); + Self { blacklist: mapping, - rewrite_rules: vec![ - // The blacklist substitution in `parse()` runs before any of - // these rules, so by the time they see the statement any - // blacklisted fragment has already been replaced. - Arc::new(AliasDuplicatedProjectionRewrite), - Arc::new(ResolveUnqualifiedIdentifier), - Arc::new(RewriteArrayAnyAllOperation), - Arc::new(PrependUnqualifiedPgTableName), - Arc::new(StripCallableQualifier), - Arc::new(FixArrayLiteral), - Arc::new(CurrentUserVariableToSessionUserFunctionCall), - Arc::new(StripCollate), - Arc::new(RewritePgCatalogOperator), - // Resolve forward oid-alias casts (`'x'::regclass`, ...) to oid - // values BEFORE RemoveSubqueryFromProjection runs, so the - // emitted scalar subqueries it produces get its LIMIT 1 stamp. - Arc::new(RewriteRegCastToSubquery::new()), - Arc::new(RemoveSubqueryFromProjection), - Arc::new(FixVersionColumnName), - ], + rewrite_rules, } } diff --git a/datafusion-pg-catalog/src/sql/rules.rs b/datafusion-pg-catalog/src/sql/rules.rs index b019dac8..798e69b2 100644 --- a/datafusion-pg-catalog/src/sql/rules.rs +++ b/datafusion-pg-catalog/src/sql/rules.rs @@ -1232,6 +1232,277 @@ impl SqlStatementRewriteRule for FixVersionColumnName { } } +/// Rewrite pgvector distance operators and vector literals into plain SQL that +/// DataFusion can plan, backed by its built-in array distance functions. +/// +/// # Operators +/// +/// | pgvector | meaning | rewrite | +/// | -------- | --------------- | ------------------------------------ | +/// | `<->` | L2 distance | `array_distance(l, r)` | +/// | `<#>` | negative dot | `-inner_product(l, r)` | +/// | `<=>` | cosine distance | `cosine_distance(l, r)` | +/// +/// sqlparser parses `<->` / `<=>` / `<#>` as the `LtDashGt`, `Spaceship` and +/// `Custom("<#>")` [`BinaryOperator`]s (see the pgvector support notes). +/// +/// # Vector literals +/// +/// pgvector accepts a vector as a bracket string, either bare (`'[1,2,3]'`) or +/// cast (`'[1,2,3]'::vector`). Both appear in the canonical usage: +/// `ORDER BY embedding <-> '[1,2,3]'`. Because pgvector treats an unknown +/// literal as the `vector` type for these operators, each operand that is such +/// a string literal is rewritten to an `ARRAY[...]` of float literals. The +/// built-in distance functions then coerce the `FixedSizeList(Float32, n)` +/// column and the `List(Float64)` literal to a common `List(Float64)` and +/// enforce the equal-length (dimension) requirement at runtime. +/// +/// # Why a rule (not UDFs / casts) +/// +/// * No schema is needed: a literal is recognized syntactically, exactly like +/// [`RewriteArrayAnyAllOperation`] / [`FixArrayLiteral`]. +/// * Dimension of the literal is known at rewrite time, so no dynamically-typed +/// `parse_vector`-style UDF is required. +/// * The rewritten SQL is re-serialised and re-parsed by DataFusion (see +/// `handlers.rs`), so the emitted form must be plain parseable SQL. +/// +/// # Coverage +/// +/// The generic AST visitor reaches vector operators in projections, `WHERE`, +/// `HAVING` and joins; sqlparser's visitor does *not* descend into `ORDER BY`, +/// so [`RewriteVectorOperatorsVisitor::rewrite_order_by`] handles those +/// expressions explicitly (this is where pgvector queries place the operator). +#[cfg(feature = "pgvector")] +#[derive(Debug)] +pub struct RewriteVectorOperators; + +#[cfg(feature = "pgvector")] +#[derive(Clone, Copy)] +enum DistanceFunc { + /// `<->` L2 distance. + ArrayDistance, + /// `<#>` negative inner product. + InnerProduct, + /// `<=>` cosine distance. + CosineDistance, +} + +#[cfg(feature = "pgvector")] +impl DistanceFunc { + fn name(self) -> &'static str { + match self { + DistanceFunc::ArrayDistance => "array_distance", + DistanceFunc::InnerProduct => "inner_product", + DistanceFunc::CosineDistance => "cosine_distance", + } + } + + /// `<#>` is the *negative* inner product, so the call must be negated. + fn negate(self) -> bool { + matches!(self, DistanceFunc::InnerProduct) + } +} + +#[cfg(feature = "pgvector")] +impl RewriteVectorOperators { + /// The distance function backing a binary operator, if it is a pgvector + /// distance operator we handle. + fn operator_to_func(op: &BinaryOperator) -> Option { + match op { + BinaryOperator::LtDashGt => Some(DistanceFunc::ArrayDistance), + BinaryOperator::Spaceship => Some(DistanceFunc::CosineDistance), + BinaryOperator::Custom(name) if name == "<#>" => Some(DistanceFunc::InnerProduct), + _ => None, + } + } + + /// True when `data_type` is the pgvector `vector` type name. + pub(crate) fn is_vector_data_type(data_type: &DataType) -> bool { + let DataType::Custom(name, _) = data_type else { + return false; + }; + name.0 + .last() + .and_then(|part| part.as_ident()) + .is_some_and(|ident| ident.value.eq_ignore_ascii_case("vector")) + } + + /// Build a SQL `ARRAY[]` literal from the numeric text between the + /// brackets of a pgvector literal like `'[1,-2.5,3]'`. Returns `None` if the + /// text does not look like a (possibly empty-forbidden) float vector. + pub(crate) fn vector_literal_to_array(text: &str) -> Option { + let text = text.trim(); + if !(text.starts_with('[') && text.ends_with(']') && text.len() >= 2) { + return None; + } + let inner = &text[1..text.len() - 1]; + if inner.trim().is_empty() { + return None; + } + + let mut elems = Vec::new(); + for part in inner.split(',') { + elems.push(Self::float_literal(part.trim())?); + } + Some(Expr::Array(Array { + elem: elems, + named: true, + })) + } + + /// A float SQL literal for `s`. Negative values become a unary-minus + /// number; non-numeric or special (`nan`, `infinity`) values return `None`. + fn float_literal(s: &str) -> Option { + if s.is_empty() { + return None; + } + let lower = s.to_lowercase(); + if matches!( + lower.as_str(), + "nan" | "inf" | "+inf" | "-inf" | "infinity" | "+infinity" | "-infinity" + ) { + return None; + } + // Parse through f32 (pgvector stores float4) and re-render with the + // shortest decimal form (Rust Display never uses scientific notation, + // so the re-parsed SQL is a plain float literal). + let value: f32 = s.parse().ok()?; + let rendered = value.to_string(); + let (neg, digits) = match rendered.strip_prefix('-') { + Some(d) => (true, d), + None => (false, rendered.as_str()), + }; + let num = Expr::Value(Value::Number(digits.to_string(), false).with_empty_span()); + if neg { + Some(Expr::UnaryOp { + op: UnaryOperator::Minus, + expr: Box::new(num), + }) + } else { + Some(num) + } + } + + /// Normalize a vector-typed operand for a distance operator: unwrap a + /// `'[...]'::vector` cast or a bare `'[...]'` string into an `ARRAY[...]` + /// literal. Anything else (columns, placeholders, functions, ...) passes + /// through unchanged. + fn coerce_vector_operand(expr: &Expr) -> Expr { + // `'[1,2,3]'::vector` / `'[1,2,3]'::public.vector` + if let Expr::Cast { + expr: inner, + data_type, + .. + } = expr + && Self::is_vector_data_type(data_type) + && let Expr::Value(ValueWithSpan { + value: Value::SingleQuotedString(text), + .. + }) = inner.as_ref() + && let Some(array) = Self::vector_literal_to_array(text) + { + return array; + } + // Bare `'[1,2,3]'` + if let Expr::Value(ValueWithSpan { + value: Value::SingleQuotedString(text), + .. + }) = expr + && let Some(array) = Self::vector_literal_to_array(text) + { + return array; + } + expr.clone() + } + + /// Rewrite `expr` in place when it is a pgvector distance `BinaryOp`, + /// recursing into the (already-rewritten-safe) operands afterwards is not + /// needed because the generic visitor will visit them separately; this + /// method handles a single operator node. + fn rewrite_binary_op(expr: &mut Expr) -> bool { + let Expr::BinaryOp { left, op, right } = expr else { + return false; + }; + let Some(func) = Self::operator_to_func(op) else { + return false; + }; + + let left = Self::coerce_vector_operand(left); + let right = Self::coerce_vector_operand(right); + let call = Expr::Function(Function { + name: ObjectName::from(vec![Ident::new(func.name())]), + args: FunctionArguments::List(FunctionArgumentList { + args: vec![ + FunctionArg::Unnamed(FunctionArgExpr::Expr(left)), + FunctionArg::Unnamed(FunctionArgExpr::Expr(right)), + ], + duplicate_treatment: None, + clauses: vec![], + }), + uses_odbc_syntax: false, + parameters: FunctionArguments::None, + filter: None, + null_treatment: None, + over: None, + within_group: vec![], + }); + + *expr = if func.negate() { + Expr::UnaryOp { + op: UnaryOperator::Minus, + expr: Box::new(call), + } + } else { + call + }; + true + } +} + +#[cfg(feature = "pgvector")] +#[derive(Debug)] +struct RewriteVectorOperatorsVisitor; + +#[cfg(feature = "pgvector")] +impl RewriteVectorOperatorsVisitor { + /// Rewrite the `ORDER BY` expressions of `query`. sqlparser's visitor does + /// not descend into `order_by`, so this is done explicitly -- pgvector's + /// canonical query places the distance operator exactly there. + fn rewrite_order_by(query: &mut Query) { + if let Some(order_by) = query.order_by.as_mut() + && let OrderByKind::Expressions(exprs) = &mut order_by.kind + { + for order_by_expr in exprs { + RewriteVectorOperators::rewrite_binary_op(&mut order_by_expr.expr); + } + } + } +} + +#[cfg(feature = "pgvector")] +impl VisitorMut for RewriteVectorOperatorsVisitor { + type Break = (); + + fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { + RewriteVectorOperators::rewrite_binary_op(expr); + ControlFlow::Continue(()) + } + + fn pre_visit_query(&mut self, query: &mut Query) -> ControlFlow { + Self::rewrite_order_by(query); + ControlFlow::Continue(()) + } +} + +#[cfg(feature = "pgvector")] +impl SqlStatementRewriteRule for RewriteVectorOperators { + fn rewrite(&self, mut s: Statement) -> Statement { + let mut visitor = RewriteVectorOperatorsVisitor; + let _ = s.visit(&mut visitor); + s + } +} + #[cfg(test)] mod tests { use super::*; @@ -1728,4 +1999,217 @@ mod tests { // A numeric operand cast to a non-oid type is left alone. assert_rewrite!(&rules, "SELECT '1'::int4", "SELECT '1'::INT4"); } + + #[cfg(feature = "pgvector")] + mod vector_rewrite { + use super::*; + use crate::sql::PostgresCompatibilityParser; + use datafusion::arrow::array::{FixedSizeListArray, Float32Array, Int64Array}; + use datafusion::arrow::datatypes::{Field as ArrowField, Schema as ArrowSchema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + use std::sync::Arc; + + fn rewrite_vector(sql: &str) -> String { + let statement = parse(sql).expect("Failed to parse").remove(0); + let rules: Vec> = + vec![Arc::new(RewriteVectorOperators)]; + rewrite(statement, &rules).to_string() + } + + #[test] + fn l2_operator_becomes_array_distance() { + let out = + rewrite_vector("SELECT * FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 5"); + assert!( + out.contains("array_distance(embedding, ARRAY[1, 2, 3])"), + "unexpected rewrite: {out}" + ); + } + + #[test] + fn inner_product_operator_is_negated() { + let out = + rewrite_vector("SELECT * FROM items ORDER BY embedding <#> '[1,2,3]' LIMIT 5"); + assert!( + out.contains("- inner_product(embedding, ARRAY[1, 2, 3])") + || out.contains("-inner_product(embedding, ARRAY[1, 2, 3])"), + "unexpected rewrite: {out}" + ); + } + + #[test] + fn cosine_operator_becomes_cosine_distance() { + let out = + rewrite_vector("SELECT * FROM items ORDER BY embedding <=> '[1,2,3]' LIMIT 5"); + assert!( + out.contains("cosine_distance(embedding, ARRAY[1, 2, 3])"), + "unexpected rewrite: {out}" + ); + } + + #[test] + fn cast_vector_literal_is_rewritten() { + let out = rewrite_vector( + "SELECT * FROM items ORDER BY embedding <-> '[1,2,3]'::vector LIMIT 5", + ); + assert!( + out.contains("array_distance(embedding, ARRAY[1, 2, 3])"), + "unexpected rewrite: {out}" + ); + } + + #[test] + fn negative_and_float_elements_are_preserved() { + let out = rewrite_vector("SELECT embedding <-> '[-1.5, 2, 0]' FROM items"); + assert!( + out.contains("array_distance(embedding, ARRAY[-1.5, 2, 0])"), + "unexpected rewrite: {out}" + ); + } + + #[test] + fn operator_in_where_is_rewritten() { + let out = rewrite_vector("SELECT id FROM items WHERE embedding <-> '[1,2,3]' < 1.0"); + assert!( + out.contains("array_distance(embedding, ARRAY[1, 2, 3]) < 1.0"), + "unexpected rewrite: {out}" + ); + } + + /// End-to-end: parse the pgvector query through the compatibility + /// parser (as the server does), then execute the rewritten SQL against + /// a DataFusion context that holds a `FixedSizeList(Float32, 3)` column. + #[tokio::test] + async fn executes_pgvector_queries() { + let ctx = SessionContext::new(); + + // items(id, embedding vector(3)) + let element = + ArrowField::new_list_field(datafusion::arrow::datatypes::DataType::Float32, true); + let fsl_type = + datafusion::arrow::datatypes::DataType::FixedSizeList(Arc::new(element), 3); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", datafusion::arrow::datatypes::DataType::Int64, false), + ArrowField::new("embedding", fsl_type, false), + ])); + let id = Int64Array::from(vec![1, 2, 3]); + let values = Float32Array::from(vec![ + 1.0, 2.0, 3.0, // row 0 + 4.0, 5.0, 6.0, // row 1 + 7.0, 8.0, 9.0, // row 2 + ]); + let embedding = FixedSizeListArray::try_new( + Arc::new(ArrowField::new_list_field( + datafusion::arrow::datatypes::DataType::Float32, + true, + )), + 3, + Arc::new(values), + None, + ) + .unwrap(); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(id), Arc::new(embedding)]).unwrap(); + ctx.register_batch("items", batch).unwrap(); + + let parser = PostgresCompatibilityParser::new(); + + // SELECT id ORDER BY L2 distance to the query vector. + let sql = parser + .parse("SELECT id FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 2") + .unwrap() + .remove(0) + .to_string(); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let ids: Vec = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(ids, vec![1, 2], "nearest rows by L2 distance"); + + // Project the computed distances for each operator. + let checks = [ + ( + "SELECT embedding <-> '[1,2,3]' AS d FROM items ORDER BY id", + 0.0, + ), + ( + "SELECT embedding <#> '[1,2,3]' AS d FROM items ORDER BY id", + -(14.0), + ), + ( + "SELECT embedding <=> '[1,2,3]' AS d FROM items ORDER BY id", + 0.0, + ), + ]; + for (query, expected_first) in checks { + let sql = parser.parse(query).unwrap().remove(0).to_string(); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let d = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let got = d.value(0); + assert!( + (got - expected_first).abs() < 1e-6, + "query {query:?} first distance = {got}, expected ~{expected_first}" + ); + } + } + + /// The `pg.vector` field metadata must survive DataFusion planning and + /// execution -- arrow-pg keys the wire type + text encoding off it, so + /// if it were dropped the server would report the column as float4[]. + #[tokio::test] + async fn vector_field_metadata_survives_planning_and_execution() { + use std::collections::HashMap; + + let ctx = SessionContext::new(); + let element = + ArrowField::new_list_field(datafusion::arrow::datatypes::DataType::Float32, true); + let field = ArrowField::new( + "embedding", + datafusion::arrow::datatypes::DataType::FixedSizeList(Arc::new(element.clone()), 3), + false, + ) + .with_metadata(HashMap::from([( + "pg.vector".to_string(), + "vector".to_string(), + )])); + let schema = Arc::new(ArrowSchema::new(vec![field])); + let values = Float32Array::from(vec![1.0, 2.0, 3.0]); + let embedding = + FixedSizeListArray::try_new(Arc::new(element), 3, Arc::new(values), None).unwrap(); + let batch = RecordBatch::try_new(schema, vec![Arc::new(embedding)]).unwrap(); + ctx.register_batch("items_meta", batch).unwrap(); + + let df = ctx.sql("SELECT embedding FROM items_meta").await.unwrap(); + assert_eq!( + df.schema() + .field(0) + .metadata() + .get("pg.vector") + .map(String::as_str), + Some("vector"), + "metadata must be present on the planned output schema" + ); + + let batches = df.collect().await.unwrap(); + assert_eq!( + batches[0] + .schema() + .field(0) + .metadata() + .get("pg.vector") + .map(String::as_str), + Some("vector"), + "metadata must be present on the executed record batch schema" + ); + } + } } diff --git a/datafusion-pg-catalog/src/sql/vector_insert.rs b/datafusion-pg-catalog/src/sql/vector_insert.rs new file mode 100644 index 00000000..129c732a --- /dev/null +++ b/datafusion-pg-catalog/src/sql/vector_insert.rs @@ -0,0 +1,398 @@ +//! Schema-aware rewrite that lets `INSERT ... VALUES ('[1,2,3]')` write into a +//! pgvector `vector` column. +//! +//! DataFusion plans `INSERT` by coercing every provided value to the target +//! column's type with [`Expr::cast_to`], which only allows casts DataFusion +//! already knows. There is no `Utf8 -> FixedSizeList(Float32, n)` cast, so a +//! bare pgvector literal `'[1,2,3]'` fails to plan with "Cannot automatically +//! convert Utf8 to FixedSizeList(...)". +//! +//! Unlike the operator rules (which run before the query has a schema), this +//! rewrite is invoked from the server handlers with the live `SessionContext`, +//! so it can resolve the INSERT target table's schema. For every value that is +//! bound to a pgvector `vector` column (Arrow `List(Float32)` / +//! `FixedSizeList(Float32, n)`), the string literal is replaced with an +//! `ARRAY[...]` of floats. DataFusion then inserts via its supported +//! `List -> FixedSizeList(Float32, n)` coercion (mismatched dimensions surface +//! as a cast error at runtime, matching pgvector's dimension enforcement). +//! +//! The rewrite is conservative: +//! * only `INSERT ... VALUES` statements whose target table resolves are +//! touched; +//! * only positions whose target column is a float vector/list type; +//! * only string literals that actually parse as a numeric `[...]` vector +//! (optionally written as `'[...]'::vector`). +//! +//! Everything else is left untouched, so a plain text column receiving a +//! bracket-looking string still works. + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::TableReference; +use datafusion::prelude::SessionContext; +use datafusion::sql::sqlparser::ast::{ + Expr, ObjectName, SetExpr, Statement, TableObject, Value, ValueWithSpan, +}; + +use super::rules::RewriteVectorOperators; + +/// The Arrow [`DataType`] of a pgvector `vector` column, if `field_type` is one. +/// +/// `Some(Some(n))` for `vector(n)` (`FixedSizeList(Float32, n)`), +/// `Some(None)` for a dimension-less `vector` (`List(Float32)`), `None` +/// otherwise. +fn vector_dimension(field_type: &DataType) -> Option> { + match field_type { + DataType::FixedSizeList(field, n) if field.data_type() == &DataType::Float32 => { + Some(Some(*n)) + } + DataType::List(field) if field.data_type() == &DataType::Float32 => Some(None), + _ => None, + } +} + +/// Convert a table-name `ObjectName` into a [`TableReference`]. +fn object_name_to_table_reference(name: &ObjectName) -> TableReference { + let parts = name + .0 + .iter() + .filter_map(|part| part.as_ident().map(|ident| ident.value.clone())) + .collect::>(); + match parts.as_slice() { + [catalog, schema, table] => { + TableReference::full(catalog.as_str(), schema.as_str(), table.as_str()) + } + [schema, table] => TableReference::partial(schema.as_str(), table.as_str()), + [table] => TableReference::bare(table.as_str()), + _ => TableReference::bare(name.to_string().as_str()), + } +} + +/// Return the pgvector literal text carried by `expr` (a bare `'[1,2,3]'` or +/// `'[1,2,3]'::vector`), or `None` if the expression is not one of those. +fn vector_literal_text(expr: &Expr) -> Option { + match expr { + Expr::Value(ValueWithSpan { + value: Value::SingleQuotedString(text), + .. + }) => Some(text.clone()), + Expr::Cast { + expr: inner, + data_type, + .. + } if RewriteVectorOperators::is_vector_data_type(data_type) => { + if let Expr::Value(ValueWithSpan { + value: Value::SingleQuotedString(text), + .. + }) = inner.as_ref() + { + Some(text.clone()) + } else { + None + } + } + _ => None, + } +} + +/// Rewrite the pgvector string literals of an `INSERT ... VALUES` statement so +/// DataFusion can write them into `vector` columns. +/// +/// Returns `true` if any value was rewritten. When the target table cannot be +/// resolved, the columns are not vector columns, or no value parses as a vector +/// literal, the statement is left untouched and `false` is returned so the +/// caller can fall back to DataFusion's normal (erroring) handling. +pub async fn rewrite_vector_insert( + session_context: &SessionContext, + statement: &mut Statement, +) -> bool { + let Statement::Insert(insert) = statement else { + return false; + }; + let TableObject::TableName(table_name) = &insert.table else { + return false; + }; + let Some(source) = insert.source.as_mut() else { + return false; + }; + let SetExpr::Values(values) = source.body.as_mut() else { + return false; + }; + // Only single-part column names are supported by DataFusion INSERT anyway. + if insert.columns.iter().any(|col| col.0.len() != 1) { + return false; + } + + let Ok(provider) = session_context + .table_provider(object_name_to_table_reference(table_name)) + .await + else { + return false; + }; + let target_schema = provider.schema(); + + // Map each provided value position to the Arrow type of the target column. + let target_types: Vec> = if insert.columns.is_empty() { + target_schema + .fields() + .iter() + .map(|field| Some(field.data_type().clone())) + .collect() + } else { + insert + .columns + .iter() + .map(|col| { + let ident = col.0[0].as_ident()?; + target_schema + .fields() + .iter() + .find(|field| field.name().eq_ignore_ascii_case(&ident.value)) + .map(|field| field.data_type().clone()) + }) + .collect() + }; + + let mut changed = false; + for row in &mut values.rows { + for (pos, target_type) in target_types.iter().enumerate() { + let Some(target_type) = target_type else { + continue; + }; + // Only float vector/list columns accept the bracket-string form. + let Some(dim) = vector_dimension(target_type) else { + continue; + }; + let Some(expr) = row.content.get_mut(pos) else { + continue; + }; + let Some(text) = vector_literal_text(expr) else { + continue; + }; + let Some(array) = RewriteVectorOperators::vector_literal_to_array(&text) else { + continue; + }; + // When the column fixes a dimension, honor it: leave mismatched + // values for DataFusion to reject rather than inserting silently. + let count = match &array { + Expr::Array(array) => array.elem.len(), + _ => unreachable!("vector_literal_to_array returns an Array"), + }; + if dim.is_some_and(|expected| expected as usize != count) { + continue; + } + *expr = array; + changed = true; + } + } + changed +} + +#[cfg(all(test, feature = "pgvector"))] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::array::{FixedSizeListArray, Float32Array, Int64Array}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + use datafusion::sql::sqlparser::dialect::PostgreSqlDialect; + use datafusion::sql::sqlparser::parser::Parser; + + use super::*; + + fn vector_field() -> Field { + Field::new( + "embedding", + DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Float32, true)), 3), + false, + ) + .with_metadata(std::collections::HashMap::from([( + "pg.vector".to_string(), + "vector".to_string(), + )])) + } + + async fn register_items(ctx: &SessionContext) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + vector_field(), + ])); + let id = Int64Array::from(vec![0]); // placeholder row; only the schema matters + let values = Float32Array::from(vec![0.0, 0.0, 0.0]); + let embedding = FixedSizeListArray::try_new( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + Arc::new(values), + None, + ) + .unwrap(); + let batch = RecordBatch::try_new(schema, vec![Arc::new(id), Arc::new(embedding)]).unwrap(); + ctx.register_batch("items", batch).unwrap(); + } + + async fn parse_insert(ctx: &SessionContext, sql: &str) -> (bool, String) { + let mut stmts = Parser::parse_sql(&PostgreSqlDialect {}, sql).unwrap(); + let mut statement = stmts.remove(0); + let changed = rewrite_vector_insert(ctx, &mut statement).await; + (changed, statement.to_string()) + } + + #[tokio::test] + async fn rewrites_vector_string_to_array_literal() { + let ctx = SessionContext::new(); + register_items(&ctx).await; + + let (changed, sql) = parse_insert( + &ctx, + "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]')", + ) + .await; + assert!(changed, "vector literal must be rewritten"); + assert!( + sql.contains("VALUES (1, ARRAY[1, 2, 3])") || sql.contains("VALUES (1, ARRAY[1,2,3])"), + "unexpected rewrite: {sql}" + ); + } + + #[tokio::test] + async fn rewrites_vector_cast_literal() { + let ctx = SessionContext::new(); + register_items(&ctx).await; + + let (changed, sql) = parse_insert( + &ctx, + "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'::vector)", + ) + .await; + assert!(changed, "vector cast literal must be rewritten"); + assert!(sql.contains("ARRAY[1, 2, 3]"), "unexpected rewrite: {sql}"); + } + + #[tokio::test] + async fn leaves_text_and_wrong_dimension_inserts_alone() { + let ctx = SessionContext::new(); + register_items(&ctx).await; + + // A non-vector target (table missing) is untouched. + let (changed, _) = parse_insert( + &ctx, + "INSERT INTO nope (id, embedding) VALUES (1, '[1,2,3]')", + ) + .await; + assert!(!changed); + + // A bracket string into a plain text column is untouched. + ctx.register_batch( + "logs", + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("tag", DataType::Utf8, true), + ])), + vec![ + Arc::new(Int64Array::from(vec![0])), + Arc::new(datafusion::arrow::array::StringArray::from(vec![""])), + ], + ) + .unwrap(), + ) + .unwrap(); + let (changed, _) = + parse_insert(&ctx, "INSERT INTO logs (id, tag) VALUES (1, '[1,2,3]')").await; + assert!(!changed, "text column must not be rewritten"); + + // A vector column with a mismatched dimension is untouched (DataFusion + // will reject it as a cast error rather than silently truncating). + let (changed, _) = parse_insert( + &ctx, + "INSERT INTO items (id, embedding) VALUES (1, '[1,2]')", + ) + .await; + assert!(!changed); + } + + #[tokio::test] + async fn insert_vector_literal_stores_row() { + let ctx = SessionContext::new(); + register_items(&ctx).await; + + let sql = parse_insert( + &ctx, + "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]')", + ) + .await + .1; + let affected = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let count = affected[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, 1, "one row must be inserted"); + + let batches = ctx + .sql("SELECT id FROM items WHERE id = 1") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].num_rows(), 1); + + // And the nearest-neighbour query finds it at distance 0. + let batches = ctx + .sql("SELECT id FROM items ORDER BY array_distance(embedding, ARRAY[1.0, 2.0, 3.0]) LIMIT 1") + .await + .unwrap() + .collect() + .await + .unwrap(); + let id = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(id, 1); + } + + #[tokio::test] + async fn insert_multiple_rows() { + let ctx = SessionContext::new(); + register_items(&ctx).await; + + let sql = parse_insert( + &ctx, + "INSERT INTO items (id, embedding) VALUES (1, '[1,0,0]'), (2, '[0,1,0]')", + ) + .await + .1; + let affected = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let count = affected[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, 2); + + // Order by distance to [1,0,0]: the placeholder row 0 ([0,0,0]) sorts + // between row 1 (exact match) and row 2. + let batches = ctx + .sql("SELECT id FROM items ORDER BY array_distance(embedding, ARRAY[1.0, 0.0, 0.0])") + .await + .unwrap() + .collect() + .await + .unwrap(); + let ids: Vec = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(ids, vec![1, 0, 2]); + } +} diff --git a/datafusion-postgres-cli/Cargo.toml b/datafusion-postgres-cli/Cargo.toml index cc4eaf3b..d67cf597 100644 --- a/datafusion-postgres-cli/Cargo.toml +++ b/datafusion-postgres-cli/Cargo.toml @@ -15,7 +15,7 @@ rust-version.workspace = true [dependencies] datafusion = { workspace = true, features = ["avro"] } tokio = { workspace = true, features = ["full"] } -datafusion-postgres = { path = "../datafusion-postgres", version = "0.18.0" } +datafusion-postgres = { path = "../datafusion-postgres", version = "0.18.0", features = ["pgvector"] } structopt = { version = "0.3", default-features = false } log = "0.4" env_logger = "0.11" diff --git a/datafusion-postgres/Cargo.toml b/datafusion-postgres/Cargo.toml index f6de2f8b..c82555cf 100644 --- a/datafusion-postgres/Cargo.toml +++ b/datafusion-postgres/Cargo.toml @@ -34,7 +34,12 @@ rustls-pki-types = "1.14" [dev-dependencies] env_logger = "0.11" +tokio-postgres = "0.7" [features] default = [] postgis = ["geodatafusion", "arrow-pg/postgis"] +# pgvector compatibility: `vector(n)` column type plus the `<->` / `<#>` / +# `<=>` distance operators, backed by DataFusion's built-in array distance +# functions. See datafusion-pg-catalog / arrow-pg `pgvector` features. +pgvector = ["datafusion-pg-catalog/pgvector", "arrow-pg/pgvector"] diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs index c66236cb..21e6c227 100644 --- a/datafusion-postgres/src/handlers.rs +++ b/datafusion-postgres/src/handlers.rs @@ -178,7 +178,17 @@ impl SimpleQueryHandler for DfSessionService { } let mut results = vec![]; - 'stmt: for statement in statements { + 'stmt: for mut statement in statements { + // pgvector: `INSERT ... VALUES ('[1,2,3]')` into a `vector` column + // needs the string literal rewritten to an ARRAY literal against the + // target table's schema (see datafusion_pg_catalog::sql). + #[cfg(feature = "pgvector")] + datafusion_pg_catalog::sql::rewrite_vector_insert( + &self.session_context, + &mut statement, + ) + .await; + // Call query hooks with the parsed statement for hook in &self.query_hooks { if let Some(result) = hook @@ -268,9 +278,14 @@ impl ExtendedQueryHandler for DfSessionService { // TODO: in the case where query hooks all return None, we do the param handling again later. let param_types = planner::get_inferred_parameter_types(plan) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; + let wire_types = parameter_wire_types(plan)?; + let wire_type_refs: Vec> = wire_types.iter().map(Some).collect(); - let param_values: ParamValues = - df::deserialize_parameters(portal, &ordered_param_types(¶m_types))?; + let param_values: ParamValues = df::deserialize_parameters_with_server_types( + portal, + &ordered_param_types(¶m_types), + &wire_type_refs, + )?; for hook in &self.query_hooks { if let Some(result) = hook @@ -291,9 +306,14 @@ impl ExtendedQueryHandler for DfSessionService { if let (_, Some((statement, plan))) = &portal.statement.statement { let param_types = planner::get_inferred_parameter_types(plan) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; + let wire_types = parameter_wire_types(plan)?; + let wire_type_refs: Vec> = wire_types.iter().map(Some).collect(); - let param_values = - df::deserialize_parameters(portal, &ordered_param_types(¶m_types))?; + let param_values = df::deserialize_parameters_with_server_types( + portal, + &ordered_param_types(¶m_types), + &wire_type_refs, + )?; let plan = plan .clone() @@ -403,7 +423,14 @@ impl QueryParser for Parser { return Ok(None); } - let statement = statements.remove(0); + let mut statement = statements.remove(0); + + // pgvector: rewrite vector string literals of INSERT ... VALUES against + // the target table's schema before DataFusion plans the statement. + #[cfg(feature = "pgvector")] + datafusion_pg_catalog::sql::rewrite_vector_insert(&self.session_context, &mut statement) + .await; + let query = statement.to_string(); let context = &self.session_context; @@ -427,20 +454,7 @@ impl QueryParser for Parser { fn get_parameter_types(&self, stmt: &Self::Statement) -> PgWireResult> { if let (_, Some((_, plan))) = stmt { - let params = planner::get_inferred_parameter_types(plan) - .map_err(|e| PgWireError::ApiError(Box::new(e)))?; - - let mut param_types = Vec::with_capacity(params.len()); - for param_type in ordered_param_types(¶ms).iter() { - if let Some(datatype) = param_type { - let pgtype = into_pg_type(datatype)?; - param_types.push(pgtype); - } else { - param_types.push(Type::UNKNOWN); - } - } - - Ok(param_types) + parameter_wire_types(plan) } else { Ok(vec![]) } @@ -470,6 +484,32 @@ impl QueryParser for Parser { } } +/// The parameter wire types the server reports in `ParameterDescription` for a +/// prepared statement's plan: the physical Arrow type mapping by default, with +/// overrides for semantically-typed columns (`pg.oid_alias` catalog columns, +/// pgvector `vector`, ...). +/// +/// The same list is reused when decoding bound parameters at Execute time, +/// because pgwire does not persist the advertised types onto the portal. +fn parameter_wire_types(plan: &LogicalPlan) -> PgWireResult> { + let params = planner::get_inferred_parameter_types(plan) + .map_err(|e| PgWireError::ApiError(Box::new(e)))?; + let overrides = planner::parameter_override_types(plan); + + let mut types = Vec::with_capacity(params.len()); + for (id, datatype) in planner::ordered_parameter_entries(¶ms) { + if let Some(ty) = overrides.get(&id) { + types.push(ty.clone()); + } else { + match datatype { + Some(datatype) => types.push(into_pg_type(&datatype)?), + None => types.push(Type::UNKNOWN), + } + } + } + Ok(types) +} + fn ordered_param_types(types: &HashMap>) -> Vec> { // Datafusion stores the parameters as a map. In our case, the keys will be // `$1`, `$2` etc. The values will be the parameter types. diff --git a/datafusion-postgres/src/planner.rs b/datafusion-postgres/src/planner.rs index 002e872d..faa740b9 100644 --- a/datafusion-postgres/src/planner.rs +++ b/datafusion-postgres/src/planner.rs @@ -5,6 +5,7 @@ use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::error::Result; use datafusion::logical_expr::LogicalPlan; use datafusion::prelude::Expr; +use pgwire::api::Type; fn extract_placeholder_cast_types(plan: &LogicalPlan) -> Result>> { let mut placeholder_types = HashMap::new(); @@ -64,3 +65,83 @@ pub fn get_inferred_parameter_types( Ok(merged) } } + +/// For each prepared-statement parameter whose type was resolved from a +/// semantically-typed catalog column, return the Postgres wire type the client +/// must use to bind it. +/// +/// DataFusion records the compared/assigned column's field -- including its +/// metadata -- on the placeholder. Two cases need an override over the physical +/// Arrow type mapping: +/// +/// * `pg.oid_alias` columns (`pg_type.oid = $1`): report `OID`, not the `INT4` +/// the Int32 storage would imply, so drivers can bind their `u32` OID; +/// * `pg.vector` columns (`INSERT ... VALUES ($1, $2)` into a `vector(n) +/// column): report the pgvector `vector` type (OID 16385) instead of the +/// physical `float4[]`, so drivers can binary-encode the vector. +/// +/// See the `pg.oid_alias` / `pg.vector` cross-crate contracts in arrow-pg. +pub fn parameter_override_types(plan: &LogicalPlan) -> HashMap { + let mut overrides = HashMap::new(); + + let _ = plan.apply(|node| { + for expr in node.expressions() { + let _ = expr.apply(|e| { + if let Some((id, field)) = placeholder_field(e) { + // oid-alias kinds. Only `oid` (stored as Int32) gets an + // override -- the `reg*` aliases are stored/displayed as + // their name strings and keep the plain TEXT mapping so + // clients can bind them with ordinary string values. + if field.data_type() == &DataType::Int32 + && let Some(kind) = + field.metadata().get(arrow_pg::datatypes::PG_OID_ALIAS_KEY) + && let Some(ty) = arrow_pg::datatypes::pg_type_for_alias_kind(kind.as_str()) + { + overrides.insert(id.to_string(), ty); + } + // pgvector `vector` + #[cfg(feature = "pgvector")] + if arrow_pg::datatypes::is_pg_vector_field(field) { + overrides.insert(id.to_string(), arrow_pg::datatypes::pg_vector_type()); + } + } + Ok(TreeNodeRecursion::Continue) + }); + } + Ok(TreeNodeRecursion::Continue) + }); + + overrides +} + +/// Unwrap a `$N` placeholder (possibly wrapped in a single cast) and return its +/// id (e.g. `"$1"`) together with the inferred field, if DataFusion resolved +/// one from the compared/assigned column. +fn placeholder_field(expr: &Expr) -> Option<(&str, &datafusion::arrow::datatypes::Field)> { + let placeholder = match expr { + Expr::Placeholder(ph) => ph, + Expr::Cast(cast) => match cast.expr.as_ref() { + Expr::Placeholder(ph) => ph, + _ => return None, + }, + _ => return None, + }; + Some((placeholder.id.as_str(), placeholder.field.as_deref()?)) +} + +/// Sort a parameter-type map (`$1`, `$2`, ...) into positional order, returning +/// each entry's placeholder id alongside its inferred type. +pub fn ordered_parameter_entries( + params: &HashMap>, +) -> Vec<(String, Option)> { + let mut entries = params.iter().collect::>(); + entries.sort_by_key(|(key, _)| { + key.trim_start_matches('$') + .parse::() + .unwrap_or(u32::MAX) + }); + entries + .into_iter() + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} diff --git a/datafusion-postgres/tests/pgvector.rs b/datafusion-postgres/tests/pgvector.rs new file mode 100644 index 00000000..6a1ccafc --- /dev/null +++ b/datafusion-postgres/tests/pgvector.rs @@ -0,0 +1,455 @@ +//! End-to-end pgvector tests through the wire-protocol handler: INSERT of a +//! `'[...]'` string literal into a `vector(n)` column and vector distance +//! queries, driven exactly like a real PostgreSQL client would send them. +#![cfg(feature = "pgvector")] + +use std::sync::Arc; +use std::time::Duration; + +use bytes::{BufMut, BytesMut}; +use datafusion::arrow::array::{FixedSizeListArray, Float32Array, Int64Array}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::prelude::SessionContext; +use datafusion_pg_catalog::setup_pg_catalog; +use pgwire::api::query::SimpleQueryHandler; +use postgres_types::{IsNull, ToSql, Type}; +use tokio_postgres::NoTls; + +use datafusion_postgres::DfSessionService; +use datafusion_postgres::auth::AuthManager; +use datafusion_postgres::testing::MockClient; +use datafusion_postgres::{ServerOptions, serve}; + +/// pgvector `vector` type OID, matching `arrow_pg::datatypes::PG_VECTOR_TYPE_OID`. +const VECTOR_OID: u32 = 16385; + +/// A client-side pgvector `vector` value that binary-encodes like real pgvector +/// does: big-endian `int16` dimension followed by big-endian IEEE float32s. +#[derive(Debug)] +struct PgVector(Vec); + +impl ToSql for PgVector { + fn to_sql( + &self, + ty: &Type, + out: &mut BytesMut, + ) -> Result> { + if !PgVector::accepts(ty) { + return Err("vector value bound to a non-vector parameter".into()); + } + out.put_i16(self.0.len() as i16); + for v in &self.0 { + out.put_slice(&v.to_be_bytes()); + } + Ok(IsNull::No) + } + + fn accepts(ty: &Type) -> bool { + ty.oid() == VECTOR_OID + } + + fn to_sql_checked( + &self, + ty: &Type, + out: &mut BytesMut, + ) -> Result> { + self.to_sql(ty, out) + } +} + +/// Register `items(id bigint, embedding vector(3))` as an empty table whose +/// `embedding` field carries the `pg.vector` metadata. +fn register_items(ctx: &SessionContext) { + let embedding = Field::new( + "embedding", + DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Float32, true)), 3), + false, + ) + .with_metadata( + [("pg.vector".to_string(), "vector".to_string())] + .into_iter() + .collect(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + embedding, + ])); + + let values = Float32Array::from(Vec::::new()); + let embedding_array = FixedSizeListArray::try_new( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + Arc::new(values), + None, + ) + .unwrap(); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(Vec::::new())), + Arc::new(embedding_array), + ], + ) + .unwrap(); + ctx.register_batch("items", batch).unwrap(); +} + +async fn service() -> (SessionContext, DfSessionService) { + let session_context = SessionContext::new(); + setup_pg_catalog( + &session_context, + "datafusion", + Arc::new(AuthManager::default()), + ) + .expect("failed to setup pg_catalog"); + register_items(&session_context); + + let service = DfSessionService::new(Arc::new(session_context.clone())); + (session_context, service) +} + +#[tokio::test] +async fn insert_vector_literal_over_wire_protocol() { + let (ctx, service) = service().await; + let mut client = MockClient::new(); + + // A plain INSERT of a pgvector bracket literal, exactly as psql would send + // it. The handler must rewrite the value so DataFusion can store it. + let responses = SimpleQueryHandler::do_query( + &service, + &mut client, + "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'), (2, '[4,5,6]')", + ) + .await + .expect("INSERT of vector literals should succeed"); + + assert_eq!(responses.len(), 1); + assert!( + matches!(responses[0], pgwire::api::results::Response::Execution(_)), + "INSERT must return an execution response" + ); + + // Rows actually landed in the table. + let batches = ctx + .sql("SELECT count(*) FROM items") + .await + .unwrap() + .collect() + .await + .unwrap(); + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, 2, "two rows must have been inserted"); + + // Distance operators keep working against the stored vectors. + let responses = SimpleQueryHandler::do_query( + &service, + &mut client, + "SELECT id FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 1", + ) + .await + .expect("distance query should succeed"); + assert_eq!(responses.len(), 1); +} + +fn free_port() -> u16 { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() +} + +/// A real PostgreSQL client (`tokio-postgres`) driving the full pgwire +/// protocol -- startup/authentication over TCP, then INSERT of pgvector +/// literals and a vector distance query over the *extended* query protocol. +#[tokio::test] +async fn real_pgwire_client_inserts_and_queries_vectors() { + let session_context = SessionContext::new(); + setup_pg_catalog( + &session_context, + "datafusion", + Arc::new(AuthManager::default()), + ) + .expect("failed to setup pg_catalog"); + register_items(&session_context); + + let port = free_port(); + let server = tokio::spawn(async move { + let ctx = Arc::new(session_context); + let options = ServerOptions::new() + .with_host("127.0.0.1".to_string()) + .with_port(port); + let _ = serve(ctx, &options).await; + }); + + // Connect over a real TCP socket, retrying briefly while the listener + // comes up. + let (client, connection) = loop { + let mut config = tokio_postgres::Config::new(); + config.host("127.0.0.1"); + config.port(port); + config.user("postgres"); + config.dbname("datafusion"); + if let Ok(connected) = config.connect(NoTls).await { + break connected; + } + tokio::time::sleep(Duration::from_millis(20)).await; + }; + tokio::spawn(async move { + let _ = connection.await; + }); + + // INSERT of pgvector bracket literals (extended protocol). + let inserted = client + .execute( + "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'), (2, '[4,5,6]')", + &[], + ) + .await + .expect("client INSERT of vector literals should succeed"); + assert_eq!(inserted, 2, "both rows must be inserted"); + + // Nearest-neighbour query with the pgvector operator. + let rows = client + .query( + "SELECT id FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 1", + &[], + ) + .await + .expect("client distance query should succeed"); + assert_eq!(rows.len(), 1); + let nearest: i64 = rows[0].get(0); + assert_eq!(nearest, 1, "row 1 embeds [1,2,3] and must be the closest"); + + // A psql-style simple-protocol read of the stored vector column returns the + // pgvector text form. (psql does not introspect pg_type for a result's + // unknown-type columns, so this exercises the same path psql uses.) + let messages = client + .simple_query("SELECT embedding FROM items ORDER BY id LIMIT 1") + .await + .expect("simple query of the vector column should succeed"); + use tokio_postgres::SimpleQueryMessage; + let mut cells = Vec::new(); + for message in messages { + if let SimpleQueryMessage::Row(row) = message { + cells.push(row.get(0).map(str::to_owned)); + } + } + assert_eq!( + cells, + vec![Some("[1,2,3]".to_string())], + "stored vector must round-trip as pgvector text" + ); + + // Typed drivers introspect unknown-type result columns by querying + // pg_catalog.pg_type for the column's OID during `prepare`. With the + // vector row injected and oid/"char" columns wired correctly this must + // succeed and report the pgvector type. + let prepared = client + .prepare("SELECT embedding FROM items LIMIT 1") + .await + .expect("prepare of a vector result column should succeed"); + assert_eq!(prepared.columns().len(), 1); + let vector_type = prepared.columns()[0].type_(); + assert_eq!(vector_type.name(), "vector"); + assert_eq!(vector_type.oid(), 16385); // matches arrow-pg PG_VECTOR_TYPE_OID + + let rows = client + .query(&prepared, &[]) + .await + .expect("executing the prepared vector select should succeed"); + assert_eq!(rows.len(), 1); + + server.abort(); +} + +/// A prepared INSERT binding a vector parameter over the extended protocol: +/// the server must report the parameter's type as pgvector `vector` (OID +/// 16385), and a client binary-encoded vector must be accepted and stored. +#[tokio::test] +async fn prepared_insert_binds_vector_parameter() { + let session_context = SessionContext::new(); + setup_pg_catalog( + &session_context, + "datafusion", + Arc::new(AuthManager::default()), + ) + .expect("failed to setup pg_catalog"); + register_items(&session_context); + + let port = free_port(); + let server = tokio::spawn(async move { + let ctx = Arc::new(session_context); + let options = ServerOptions::new() + .with_host("127.0.0.1".to_string()) + .with_port(port); + let _ = serve(ctx, &options).await; + }); + + let (client, connection) = loop { + let mut config = tokio_postgres::Config::new(); + config.host("127.0.0.1"); + config.port(port); + config.user("postgres"); + config.dbname("datafusion"); + if let Ok(connected) = config.connect(NoTls).await { + break connected; + } + tokio::time::sleep(Duration::from_millis(20)).await; + }; + tokio::spawn(async move { + let _ = connection.await; + }); + + // Prepare an INSERT with a bound vector. The server's ParameterDescription + // must advertise the vector parameter as the pgvector type (OID 16385) so a + // typed client knows how to binary-encode it. + let statement = client + .prepare("INSERT INTO items (id, embedding) VALUES ($1, $2)") + .await + .expect("prepare INSERT with a vector parameter should succeed"); + + assert_eq!(statement.params().len(), 2, "two parameters expected"); + let vector_param = &statement.params()[1]; + assert_eq!( + vector_param.oid(), + VECTOR_OID, + "vector parameter must be reported with the pgvector type OID" + ); + assert_eq!(vector_param.name(), "vector"); + + // Bind an id and a binary-encoded vector and execute. + let id: i64 = 42; + let vector = PgVector(vec![1.0, 2.0, 3.0]); + let affected = client + .execute(&statement, &[&id, &vector]) + .await + .expect("executing the prepared INSERT should succeed"); + assert_eq!(affected, 1, "one row must be inserted"); + + // The stored vector round-trips: the closest row to [1,2,3] is the one we + // just inserted, and its text form is correct. + let rows = client + .query( + "SELECT id FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 1", + &[], + ) + .await + .unwrap(); + let nearest: i64 = rows[0].get(0); + assert_eq!(nearest, 42, "inserted row must be the nearest match"); + + let messages = client + .simple_query("SELECT embedding FROM items WHERE id = 42") + .await + .expect("simple query of the vector column should succeed"); + use tokio_postgres::SimpleQueryMessage; + let mut cells = Vec::new(); + for message in messages { + if let SimpleQueryMessage::Row(row) = message { + cells.push(row.get(0).map(str::to_owned)); + } + } + assert_eq!( + cells, + vec![Some("[1,2,3]".to_string())], + "binary-encoded vector parameter must be stored correctly" + ); + + server.abort(); +} + +/// Start a server over `pg_catalog` (so the pgvector type planner and the +/// injected `vector` pg_type row are active) on a free port, returning the port +/// and server task. +fn spawn_pgvector_server() -> (u16, tokio::task::JoinHandle<()>) { + let session_context = SessionContext::new(); + setup_pg_catalog( + &session_context, + "datafusion", + Arc::new(AuthManager::default()), + ) + .expect("failed to setup pg_catalog"); + + let port = free_port(); + let server = tokio::spawn(async move { + let ctx = Arc::new(session_context); + let options = ServerOptions::new() + .with_host("127.0.0.1".to_string()) + .with_port(port); + let _ = serve(ctx, &options).await; + }); + (port, server) +} + +/// Connect a tokio-postgres client to `port`, retrying while the listener starts. +async fn connect_pgwire_client(port: u16) -> tokio_postgres::Client { + let (client, connection) = loop { + let mut config = tokio_postgres::Config::new(); + config.host("127.0.0.1"); + config.port(port); + config.user("postgres"); + config.dbname("datafusion"); + if let Ok(connected) = config.connect(NoTls).await { + break connected; + } + tokio::time::sleep(Duration::from_millis(20)).await; + }; + tokio::spawn(async move { + let _ = connection.await; + }); + client +} + +/// The canonical pgvector DDL -- `CREATE TABLE items (id int PRIMARY KEY, +/// embedding vector(3))` -- must work over the wire, and the created table must +/// accept both literal and bound-parameter vector INSERTs. +#[tokio::test] +async fn create_table_ddl_with_vector_column() { + let (port, server) = spawn_pgvector_server(); + let client = connect_pgwire_client(port).await; + + client + .batch_execute("CREATE TABLE items (id int PRIMARY KEY, embedding vector(3))") + .await + .expect("CREATE TABLE with a vector(3) column should succeed"); + + // Insert a pgvector literal ... + client + .execute( + "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'), (2, '[4,5,6]')", + &[], + ) + .await + .expect("literal INSERT into the DDL-created table should succeed"); + + // ... and a bound vector parameter via a prepared statement. + let statement = client + .prepare("INSERT INTO items (id, embedding) VALUES ($1, $2)") + .await + .expect("prepare INSERT should succeed"); + assert_eq!(statement.params()[1].oid(), VECTOR_OID); + let id: i32 = 3; + let vector = PgVector(vec![7.0, 8.0, 9.0]); + client + .execute(&statement, &[&id, &vector]) + .await + .expect("prepared vector INSERT into the DDL-created table should succeed"); + + // Rows landed and nearest-neighbour search works over the DDL-created table. + let rows = client + .query( + "SELECT id FROM items ORDER BY embedding <-> '[7,8,9]' LIMIT 1", + &[], + ) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + let nearest: i32 = rows[0].get(0); + assert_eq!(nearest, 3, "nearest row must be the prepared vector insert"); + + server.abort(); +} From ae1b54d531620843b3a81f21ee497de8d9e16209 Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Mon, 7 Sep 2026 11:56:02 +0800 Subject: [PATCH 02/10] feat(pgvector): vector binary results + official client e2e tests 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. --- Cargo.lock | 27 ++- arrow-pg/src/datatypes.rs | 11 - arrow-pg/src/datatypes/df.rs | 19 +- arrow-pg/src/encoder.rs | 163 +++++++++++-- datafusion-postgres/Cargo.toml | 3 +- datafusion-postgres/src/handlers.rs | 4 +- datafusion-postgres/tests/pgvector.rs | 321 ++++++++++++-------------- 7 files changed, 332 insertions(+), 216 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index eb0a77d7..430d2821 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2491,13 +2491,14 @@ dependencies = [ "geodatafusion", "getset", "log", + "pgvector", "pgwire", + "postgres", "postgres-types", "rust_decimal", "rustls-pemfile", "rustls-pki-types", "tokio", - "tokio-postgres", "tokio-rustls", ] @@ -4040,6 +4041,16 @@ dependencies = [ "postgres-types", ] +[[package]] +name = "pgvector" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3673cba5b9a124916096a423b806a9f29620972c6c97b08db5f2053e9428b481" +dependencies = [ + "bytes", + "postgres-types", +] + [[package]] name = "pgwire" version = "0.41.0" @@ -4197,6 +4208,20 @@ dependencies = [ "postgres-types", ] +[[package]] +name = "postgres" +version = "0.19.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ad20e0aa0b24f5a394eab4f78c781d248982b22b25cecc7e3aa46a681605bd" +dependencies = [ + "bytes", + "fallible-iterator", + "futures-util", + "log", + "tokio", + "tokio-postgres", +] + [[package]] name = "postgres-protocol" version = "0.6.12" diff --git a/arrow-pg/src/datatypes.rs b/arrow-pg/src/datatypes.rs index 4de2d435..93f659b8 100644 --- a/arrow-pg/src/datatypes.rs +++ b/arrow-pg/src/datatypes.rs @@ -298,17 +298,6 @@ pub fn arrow_schema_to_pg_fields( .enumerate() .map(|(idx, f)| { let pg_type = field_into_pg_type(f)?; - - // pgvector `vector` has no binary wire encoding implemented yet, so - // always negotiate the text format (`[1,2,3]`) for vector columns, - // even when the client asked for binary. - #[cfg(feature = "pgvector")] - let col_format = if is_pg_vector_field(f) { - pgwire::api::results::FieldFormat::Text - } else { - format.format_for(idx) - }; - #[cfg(not(feature = "pgvector"))] let col_format = format.format_for(idx); let mut field_info = FieldInfo::new(f.name().into(), None, None, pg_type, col_format); diff --git a/arrow-pg/src/datatypes/df.rs b/arrow-pg/src/datatypes/df.rs index f32558c0..8c58b763 100644 --- a/arrow-pg/src/datatypes/df.rs +++ b/arrow-pg/src/datatypes/df.rs @@ -279,23 +279,26 @@ struct VectorParam(Vec); #[cfg(feature = "pgvector")] impl VectorParam { - /// The pgvector binary layout is a big-endian `int16` dimension followed by - /// that many big-endian IEEE float32 elements. + /// The pgvector binary layout is a big-endian `u16` dimension, an unused + /// `u16` that must be 0, then that many big-endian IEEE float32 elements. fn from_binary(raw: &[u8]) -> Result> { - if raw.len() < 2 { + if raw.len() < 4 { return Err("vector parameter binary payload too short".into()); } - let dim = i16::from_be_bytes([raw[0], raw[1]]); - if dim <= 0 { + let dim = u16::from_be_bytes([raw[0], raw[1]]) as usize; + let unused = u16::from_be_bytes([raw[2], raw[3]]); + if unused != 0 { + return Err("vector parameter binary payload has a non-zero unused word".into()); + } + if dim == 0 { return Err("vector parameter dimension must be positive".into()); } - let dim = dim as usize; - if raw.len() != 2 + dim * 4 { + if raw.len() != 4 + dim * 4 { return Err("vector parameter binary payload has wrong length".into()); } let mut values = Vec::with_capacity(dim); for i in 0..dim { - let off = 2 + i * 4; + let off = 4 + i * 4; values.push(f32::from_be_bytes([ raw[off], raw[off + 1], diff --git a/arrow-pg/src/encoder.rs b/arrow-pg/src/encoder.rs index 4f81233e..42857e63 100644 --- a/arrow-pg/src/encoder.rs +++ b/arrow-pg/src/encoder.rs @@ -1,6 +1,9 @@ use std::str::FromStr; use std::sync::Arc; +#[cfg(feature = "pgvector")] +use bytes::BytesMut; + #[cfg(not(feature = "datafusion"))] use arrow::{array::*, datatypes::*}; use chrono::NaiveTime; @@ -208,27 +211,75 @@ fn encode_pg_char( } } -/// Render the pgvector text form of a vector: `[1,2,3]`. -/// -/// Element formatting uses Rust's shortest round-trip `Display` for `f32` -/// (`1.0` -> `1`), matching pgvector's `vector_out`. +/// A pgvector `vector` value that pgwire serializes in the format the client +/// requested: binary (pgvector layout, big-endian `int16` dimension followed by +/// big-endian IEEE float32s) or text (`[1,2,3]`). +#[cfg(feature = "pgvector")] +#[derive(Debug)] +struct PgVectorValue(Vec); + +#[cfg(feature = "pgvector")] +impl postgres_types::ToSql for PgVectorValue { + fn to_sql( + &self, + ty: &postgres_types::Type, + out: &mut BytesMut, + ) -> Result> { + use bytes::BufMut as _; + if !PgVectorValue::accepts(ty) { + return Err("vector value bound to a non-vector result type".into()); + } + // pgvector binary layout: big-endian u16 dimension, an unused u16 that + // must be 0, then big-endian IEEE float32 elements. + out.put_u16(self.0.len() as u16); + out.put_u16(0); + for v in &self.0 { + out.put_f32(*v); + } + Ok(postgres_types::IsNull::No) + } + + fn accepts(ty: &postgres_types::Type) -> bool { + ty.oid() == crate::datatypes::PG_VECTOR_TYPE_OID + } + + fn to_sql_checked( + &self, + ty: &postgres_types::Type, + out: &mut BytesMut, + ) -> Result> { + self.to_sql(ty, out) + } +} + #[cfg(feature = "pgvector")] -fn format_pg_vector(values: &[f32]) -> String { - let inner = values - .iter() - .map(|v| v.to_string()) - .collect::>() - .join(","); - format!("[{inner}]") +impl pgwire::types::ToSqlText for PgVectorValue { + fn to_sql_text( + &self, + _ty: &postgres_types::Type, + out: &mut BytesMut, + _format_options: &pgwire::types::format::FormatOptions, + ) -> Result> { + use bytes::BufMut as _; + // Element formatting uses Rust's shortest round-trip `Display` for + // `f32` (`1.0` -> `1`), matching pgvector's `vector_out`. + let inner = self + .0 + .iter() + .map(|v| v.to_string()) + .collect::>() + .join(","); + out.put_slice(format!("[{inner}]").as_bytes()); + Ok(postgres_types::IsNull::No) + } } /// Encode a pgvector `vector` column (an Arrow `List`/`FixedSizeList` of /// `Float32` tagged with the `pg.vector` field metadata) for a single row. /// -/// The value is emitted in the pgvector text format `[1,2,3]`. The result -/// `FieldInfo` for vector columns is forced to the text format by -/// `arrow_schema_to_pg_fields`, so pgwire serializes the string verbatim via -/// `ToSqlText` regardless of the client's requested result format. +/// pgwire picks the encoding from the result `FieldInfo` format: text produces +/// `[1,2,3]` (as psql shows), binary produces the pgvector wire layout that +/// typed drivers decode. #[cfg(feature = "pgvector")] fn encode_pg_vector( encoder: &mut T, @@ -237,7 +288,7 @@ fn encode_pg_vector( pg_field: &FieldInfo, ) -> PgWireResult<()> { if arr.is_null(idx) { - return encoder.encode_field(&None::, pg_field); + return encoder.encode_field(&None::, pg_field); } fn row_values(arr: &Arc, idx: usize) -> PgWireResult> { @@ -279,8 +330,7 @@ fn encode_pg_vector( Ok(values) } - let text = format_pg_vector(&row_values(arr, idx)?); - encoder.encode_field(&Some(text), pg_field) + encoder.encode_field(&Some(PgVectorValue(row_values(arr, idx)?)), pg_field) } fn get_date32_value(arr: &Arc, idx: usize) -> Option { @@ -407,12 +457,16 @@ pub fn encode_value( ); } - // pgvector `vector` columns are tagged with the `pg.vector` field - // metadata. Route them through the vector encoder (text `[1,2,3]`) before - // the generic list handling below, which would otherwise emit them as a - // Postgres float4[] (`{1,2,3}`). + // pgvector `vector` columns are tagged with the `pg.vector` field metadata. + // Route them through the vector encoder (text `[1,2,3]` / pgvector binary) + // before the generic list handling below, which would otherwise emit them + // as a Postgres float4[] (`{1,2,3}`). The resolved FieldInfo wire type is + // the authoritative signal: some optimizer rewrites drop the arrow field + // metadata while keeping the logical vector type. #[cfg(feature = "pgvector")] - if crate::datatypes::is_pg_vector_field(arrow_field) { + if crate::datatypes::is_pg_vector_field(arrow_field) + || pg_field.datatype().oid() == crate::datatypes::PG_VECTOR_TYPE_OID + { return encode_pg_vector(encoder, arr, idx, pg_field); } @@ -987,6 +1041,69 @@ mod tests { assert_eq!(encoder.encoded, "[4.5,0,7]"); } + #[derive(Default)] + struct BinaryCapture { + encoded: Vec, + } + + impl Encoder for BinaryCapture { + type Item = Vec; + + fn encode_field(&mut self, value: &T, pg_field: &FieldInfo) -> PgWireResult<()> + where + T: ToSql + ToSqlText + Sized, + { + use postgres_types::ToSql as _; + let mut bytes = BytesMut::new(); + value.to_sql(pg_field.datatype(), &mut bytes).unwrap(); + self.encoded = bytes.to_vec(); + Ok(()) + } + + fn take_row(&mut self) -> Self::Item { + std::mem::take(&mut self.encoded) + } + } + + #[test] + fn encodes_vector_fixed_size_list_as_pgvector_binary() { + // pgvector binary layout: big-endian u16 dimension, an unused u16 + // (0), then big-endian IEEE float32s. [1,2,3] -> + // 00 03 | 00 00 | 3f800000 40000000 40400000. + let values = Float32Array::from(vec![1.0, 2.0, 3.0]); + let array: Arc = Arc::new( + FixedSizeListArray::try_new( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + Arc::new(values), + None, + ) + .unwrap(), + ); + + let arrow_field = vector_arrow_field(true); + let pg_field = FieldInfo::new( + "embedding".to_string(), + None, + None, + crate::datatypes::pg_vector_type(), + FieldFormat::Binary, + ); + + let mut encoder = BinaryCapture::default(); + encode_value(&mut encoder, &array, 0, &arrow_field, &pg_field).unwrap(); + assert_eq!( + encoder.encoded, + vec![ + 0x00, 0x03, // dim = 3 + 0x00, 0x00, // unused + 0x3f, 0x80, 0x00, 0x00, // 1.0 + 0x40, 0x00, 0x00, 0x00, // 2.0 + 0x40, 0x40, 0x00, 0x00, // 3.0 + ] + ); + } + #[test] fn encodes_null_vector_as_null() { // `NullBuffer::from(Vec)` treats `true` as valid. diff --git a/datafusion-postgres/Cargo.toml b/datafusion-postgres/Cargo.toml index c82555cf..7cb3fb6a 100644 --- a/datafusion-postgres/Cargo.toml +++ b/datafusion-postgres/Cargo.toml @@ -34,7 +34,8 @@ rustls-pki-types = "1.14" [dev-dependencies] env_logger = "0.11" -tokio-postgres = "0.7" +pgvector = { version = "0.4", features = ["postgres"] } +postgres = "0.19" [features] default = [] diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs index 21e6c227..5a3e3b15 100644 --- a/datafusion-postgres/src/handlers.rs +++ b/datafusion-postgres/src/handlers.rs @@ -178,11 +178,13 @@ impl SimpleQueryHandler for DfSessionService { } let mut results = vec![]; - 'stmt: for mut statement in statements { + 'stmt: for statement in statements { // pgvector: `INSERT ... VALUES ('[1,2,3]')` into a `vector` column // needs the string literal rewritten to an ARRAY literal against the // target table's schema (see datafusion_pg_catalog::sql). #[cfg(feature = "pgvector")] + let mut statement = statement; + #[cfg(feature = "pgvector")] datafusion_pg_catalog::sql::rewrite_vector_insert( &self.session_context, &mut statement, diff --git a/datafusion-postgres/tests/pgvector.rs b/datafusion-postgres/tests/pgvector.rs index 6a1ccafc..be53ca3b 100644 --- a/datafusion-postgres/tests/pgvector.rs +++ b/datafusion-postgres/tests/pgvector.rs @@ -1,20 +1,24 @@ -//! End-to-end pgvector tests through the wire-protocol handler: INSERT of a -//! `'[...]'` string literal into a `vector(n)` column and vector distance -//! queries, driven exactly like a real PostgreSQL client would send them. +//! End-to-end pgvector tests through the wire protocol. +//! +//! The in-process handler tests drive SQL through the server's query handlers +//! with a mock client; the network tests connect a real PostgreSQL client +//! ([`postgres`], the synchronous rust-postgres crate) to a live server and +//! exercise DDL, INSERT and queries using the official [`pgvector`] client +//! types -- proving binary/text compatibility against the real pgvector wire +//! format without re-implementing any pgvector encoding here. #![cfg(feature = "pgvector")] use std::sync::Arc; use std::time::Duration; -use bytes::{BufMut, BytesMut}; use datafusion::arrow::array::{FixedSizeListArray, Float32Array, Int64Array}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::prelude::SessionContext; use datafusion_pg_catalog::setup_pg_catalog; use pgwire::api::query::SimpleQueryHandler; -use postgres_types::{IsNull, ToSql, Type}; -use tokio_postgres::NoTls; +use postgres::NoTls; +use tokio::sync::oneshot; use datafusion_postgres::DfSessionService; use datafusion_postgres::auth::AuthManager; @@ -24,40 +28,6 @@ use datafusion_postgres::{ServerOptions, serve}; /// pgvector `vector` type OID, matching `arrow_pg::datatypes::PG_VECTOR_TYPE_OID`. const VECTOR_OID: u32 = 16385; -/// A client-side pgvector `vector` value that binary-encodes like real pgvector -/// does: big-endian `int16` dimension followed by big-endian IEEE float32s. -#[derive(Debug)] -struct PgVector(Vec); - -impl ToSql for PgVector { - fn to_sql( - &self, - ty: &Type, - out: &mut BytesMut, - ) -> Result> { - if !PgVector::accepts(ty) { - return Err("vector value bound to a non-vector parameter".into()); - } - out.put_i16(self.0.len() as i16); - for v in &self.0 { - out.put_slice(&v.to_be_bytes()); - } - Ok(IsNull::No) - } - - fn accepts(ty: &Type) -> bool { - ty.oid() == VECTOR_OID - } - - fn to_sql_checked( - &self, - ty: &Type, - out: &mut BytesMut, - ) -> Result> { - self.to_sql(ty, out) - } -} - /// Register `items(id bigint, embedding vector(3))` as an empty table whose /// `embedding` field carries the `pg.vector` metadata. fn register_items(ctx: &SessionContext) { @@ -162,11 +132,49 @@ fn free_port() -> u16 { listener.local_addr().unwrap().port() } -/// A real PostgreSQL client (`tokio-postgres`) driving the full pgwire -/// protocol -- startup/authentication over TCP, then INSERT of pgvector -/// literals and a vector distance query over the *extended* query protocol. -#[tokio::test] -async fn real_pgwire_client_inserts_and_queries_vectors() { +/// Run a server over `session_context` on its own thread. The server is shut +/// down by sending on the returned oneshot channel. +fn spawn_server(session_context: SessionContext) -> (u16, oneshot::Sender<()>) { + let port = free_port(); + let (stop_tx, stop_rx) = oneshot::channel(); + + std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("failed to build server runtime"); + let ctx = Arc::new(session_context); + let options = ServerOptions::new() + .with_host("127.0.0.1".to_string()) + .with_port(port); + runtime.block_on(async move { + tokio::select! { + _ = serve(ctx, &options) => {} + _ = stop_rx => {} + } + }); + }); + + (port, stop_tx) +} + +/// Connect a synchronous rust-postgres client, retrying while the server's +/// listener comes up. +fn connect(port: u16) -> postgres::Client { + loop { + let config = format!("host=127.0.0.1 port={port} user=postgres dbname=datafusion"); + match postgres::Client::connect(&config, NoTls) { + Ok(client) => return client, + Err(_) => std::thread::sleep(Duration::from_millis(20)), + } + } +} + +/// A real PostgreSQL client driving the full pgwire protocol -- startup over +/// TCP, INSERT of pgvector literals, prepared vector parameters and vector +/// distance queries. +#[test] +fn real_pgwire_client_inserts_and_queries_vectors() { let session_context = SessionContext::new(); setup_pg_catalog( &session_context, @@ -176,31 +184,8 @@ async fn real_pgwire_client_inserts_and_queries_vectors() { .expect("failed to setup pg_catalog"); register_items(&session_context); - let port = free_port(); - let server = tokio::spawn(async move { - let ctx = Arc::new(session_context); - let options = ServerOptions::new() - .with_host("127.0.0.1".to_string()) - .with_port(port); - let _ = serve(ctx, &options).await; - }); - - // Connect over a real TCP socket, retrying briefly while the listener - // comes up. - let (client, connection) = loop { - let mut config = tokio_postgres::Config::new(); - config.host("127.0.0.1"); - config.port(port); - config.user("postgres"); - config.dbname("datafusion"); - if let Ok(connected) = config.connect(NoTls).await { - break connected; - } - tokio::time::sleep(Duration::from_millis(20)).await; - }; - tokio::spawn(async move { - let _ = connection.await; - }); + let (port, stop_tx) = spawn_server(session_context); + let mut client = connect(port); // INSERT of pgvector bracket literals (extended protocol). let inserted = client @@ -208,7 +193,6 @@ async fn real_pgwire_client_inserts_and_queries_vectors() { "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'), (2, '[4,5,6]')", &[], ) - .await .expect("client INSERT of vector literals should succeed"); assert_eq!(inserted, 2, "both rows must be inserted"); @@ -218,23 +202,19 @@ async fn real_pgwire_client_inserts_and_queries_vectors() { "SELECT id FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 1", &[], ) - .await .expect("client distance query should succeed"); assert_eq!(rows.len(), 1); let nearest: i64 = rows[0].get(0); assert_eq!(nearest, 1, "row 1 embeds [1,2,3] and must be the closest"); // A psql-style simple-protocol read of the stored vector column returns the - // pgvector text form. (psql does not introspect pg_type for a result's - // unknown-type columns, so this exercises the same path psql uses.) + // pgvector text form. let messages = client .simple_query("SELECT embedding FROM items ORDER BY id LIMIT 1") - .await .expect("simple query of the vector column should succeed"); - use tokio_postgres::SimpleQueryMessage; let mut cells = Vec::new(); for message in messages { - if let SimpleQueryMessage::Row(row) = message { + if let postgres::SimpleQueryMessage::Row(row) = message { cells.push(row.get(0).map(str::to_owned)); } } @@ -245,32 +225,32 @@ async fn real_pgwire_client_inserts_and_queries_vectors() { ); // Typed drivers introspect unknown-type result columns by querying - // pg_catalog.pg_type for the column's OID during `prepare`. With the - // vector row injected and oid/"char" columns wired correctly this must - // succeed and report the pgvector type. - let prepared = client + // pg_catalog.pg_type for the column's OID during `prepare`. The official + // pgvector::Vector decoder then reads the binary result. + let statement = client .prepare("SELECT embedding FROM items LIMIT 1") - .await .expect("prepare of a vector result column should succeed"); - assert_eq!(prepared.columns().len(), 1); - let vector_type = prepared.columns()[0].type_(); + assert_eq!(statement.columns().len(), 1); + let vector_type = &statement.columns()[0].type_(); assert_eq!(vector_type.name(), "vector"); - assert_eq!(vector_type.oid(), 16385); // matches arrow-pg PG_VECTOR_TYPE_OID + assert_eq!(vector_type.oid(), VECTOR_OID); let rows = client - .query(&prepared, &[]) - .await + .query(&statement, &[]) .expect("executing the prepared vector select should succeed"); assert_eq!(rows.len(), 1); + let decoded: pgvector::Vector = rows[0].get(0); + assert_eq!(decoded, pgvector::Vector::from(vec![1.0, 2.0, 3.0])); - server.abort(); + let _ = stop_tx.send(()); } /// A prepared INSERT binding a vector parameter over the extended protocol: /// the server must report the parameter's type as pgvector `vector` (OID -/// 16385), and a client binary-encoded vector must be accepted and stored. -#[tokio::test] -async fn prepared_insert_binds_vector_parameter() { +/// 16385), and an official `pgvector::Vector` binary-encoded value must be +/// accepted and stored. +#[test] +fn prepared_insert_binds_vector_parameter() { let session_context = SessionContext::new(); setup_pg_catalog( &session_context, @@ -280,36 +260,14 @@ async fn prepared_insert_binds_vector_parameter() { .expect("failed to setup pg_catalog"); register_items(&session_context); - let port = free_port(); - let server = tokio::spawn(async move { - let ctx = Arc::new(session_context); - let options = ServerOptions::new() - .with_host("127.0.0.1".to_string()) - .with_port(port); - let _ = serve(ctx, &options).await; - }); - - let (client, connection) = loop { - let mut config = tokio_postgres::Config::new(); - config.host("127.0.0.1"); - config.port(port); - config.user("postgres"); - config.dbname("datafusion"); - if let Ok(connected) = config.connect(NoTls).await { - break connected; - } - tokio::time::sleep(Duration::from_millis(20)).await; - }; - tokio::spawn(async move { - let _ = connection.await; - }); + let (port, stop_tx) = spawn_server(session_context); + let mut client = connect(port); // Prepare an INSERT with a bound vector. The server's ParameterDescription // must advertise the vector parameter as the pgvector type (OID 16385) so a // typed client knows how to binary-encode it. let statement = client .prepare("INSERT INTO items (id, embedding) VALUES ($1, $2)") - .await .expect("prepare INSERT with a vector parameter should succeed"); assert_eq!(statement.params().len(), 2, "two parameters expected"); @@ -321,12 +279,11 @@ async fn prepared_insert_binds_vector_parameter() { ); assert_eq!(vector_param.name(), "vector"); - // Bind an id and a binary-encoded vector and execute. + // Bind an id and an official pgvector::Vector, then execute. let id: i64 = 42; - let vector = PgVector(vec![1.0, 2.0, 3.0]); + let vector = pgvector::Vector::from(vec![1.0, 2.0, 3.0]); let affected = client .execute(&statement, &[&id, &vector]) - .await .expect("executing the prepared INSERT should succeed"); assert_eq!(affected, 1, "one row must be inserted"); @@ -337,19 +294,16 @@ async fn prepared_insert_binds_vector_parameter() { "SELECT id FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 1", &[], ) - .await .unwrap(); let nearest: i64 = rows[0].get(0); assert_eq!(nearest, 42, "inserted row must be the nearest match"); let messages = client .simple_query("SELECT embedding FROM items WHERE id = 42") - .await .expect("simple query of the vector column should succeed"); - use tokio_postgres::SimpleQueryMessage; let mut cells = Vec::new(); for message in messages { - if let SimpleQueryMessage::Row(row) = message { + if let postgres::SimpleQueryMessage::Row(row) = message { cells.push(row.get(0).map(str::to_owned)); } } @@ -359,13 +313,14 @@ async fn prepared_insert_binds_vector_parameter() { "binary-encoded vector parameter must be stored correctly" ); - server.abort(); + let _ = stop_tx.send(()); } -/// Start a server over `pg_catalog` (so the pgvector type planner and the -/// injected `vector` pg_type row are active) on a free port, returning the port -/// and server task. -fn spawn_pgvector_server() -> (u16, tokio::task::JoinHandle<()>) { +/// The canonical pgvector DDL -- `CREATE TABLE items (id int PRIMARY KEY, +/// embedding vector(3))` -- must work over the wire, and the created table must +/// accept both literal and bound-parameter vector INSERTs. +#[test] +fn create_table_ddl_with_vector_column() { let session_context = SessionContext::new(); setup_pg_catalog( &session_context, @@ -374,47 +329,11 @@ fn spawn_pgvector_server() -> (u16, tokio::task::JoinHandle<()>) { ) .expect("failed to setup pg_catalog"); - let port = free_port(); - let server = tokio::spawn(async move { - let ctx = Arc::new(session_context); - let options = ServerOptions::new() - .with_host("127.0.0.1".to_string()) - .with_port(port); - let _ = serve(ctx, &options).await; - }); - (port, server) -} - -/// Connect a tokio-postgres client to `port`, retrying while the listener starts. -async fn connect_pgwire_client(port: u16) -> tokio_postgres::Client { - let (client, connection) = loop { - let mut config = tokio_postgres::Config::new(); - config.host("127.0.0.1"); - config.port(port); - config.user("postgres"); - config.dbname("datafusion"); - if let Ok(connected) = config.connect(NoTls).await { - break connected; - } - tokio::time::sleep(Duration::from_millis(20)).await; - }; - tokio::spawn(async move { - let _ = connection.await; - }); - client -} - -/// The canonical pgvector DDL -- `CREATE TABLE items (id int PRIMARY KEY, -/// embedding vector(3))` -- must work over the wire, and the created table must -/// accept both literal and bound-parameter vector INSERTs. -#[tokio::test] -async fn create_table_ddl_with_vector_column() { - let (port, server) = spawn_pgvector_server(); - let client = connect_pgwire_client(port).await; + let (port, stop_tx) = spawn_server(session_context); + let mut client = connect(port); client .batch_execute("CREATE TABLE items (id int PRIMARY KEY, embedding vector(3))") - .await .expect("CREATE TABLE with a vector(3) column should succeed"); // Insert a pgvector literal ... @@ -423,20 +342,17 @@ async fn create_table_ddl_with_vector_column() { "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'), (2, '[4,5,6]')", &[], ) - .await .expect("literal INSERT into the DDL-created table should succeed"); // ... and a bound vector parameter via a prepared statement. let statement = client .prepare("INSERT INTO items (id, embedding) VALUES ($1, $2)") - .await .expect("prepare INSERT should succeed"); assert_eq!(statement.params()[1].oid(), VECTOR_OID); let id: i32 = 3; - let vector = PgVector(vec![7.0, 8.0, 9.0]); + let vector = pgvector::Vector::from(vec![7.0, 8.0, 9.0]); client .execute(&statement, &[&id, &vector]) - .await .expect("prepared vector INSERT into the DDL-created table should succeed"); // Rows landed and nearest-neighbour search works over the DDL-created table. @@ -445,11 +361,74 @@ async fn create_table_ddl_with_vector_column() { "SELECT id FROM items ORDER BY embedding <-> '[7,8,9]' LIMIT 1", &[], ) - .await .unwrap(); assert_eq!(rows.len(), 1); let nearest: i32 = rows[0].get(0); assert_eq!(nearest, 3, "nearest row must be the prepared vector insert"); - server.abort(); + let _ = stop_tx.send(()); +} + +/// Full pgvector-rust (`pgvector::Vector`) client round-trip: DDL, INSERT of a +/// bound `Vector` parameter, and SELECT decoding the binary result back into +/// `Vector`. Because the official client both encodes and decodes the real +/// pgvector wire format, this pins binary/text compatibility. +#[test] +fn pgvector_rust_official_client_roundtrip() { + let session_context = SessionContext::new(); + setup_pg_catalog( + &session_context, + "datafusion", + Arc::new(AuthManager::default()), + ) + .expect("failed to setup pg_catalog"); + + let (port, stop_tx) = spawn_server(session_context); + let mut client = connect(port); + + // DDL (inline PRIMARY KEY + vector(3)). + client + .execute( + "CREATE TABLE items (id int PRIMARY KEY, embedding vector(3))", + &[], + ) + .expect("CREATE TABLE with a vector(3) column should succeed"); + + // Insert rows with pgvector string literals ... + let inserted = client + .execute( + "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'), (2, '[4,5,6]')", + &[], + ) + .expect("literal vector INSERT should succeed"); + assert_eq!(inserted, 2); + + // ... and with an official pgvector::Vector bound parameter (binary). + let official_vec = pgvector::Vector::from(vec![7.0, 8.0, 9.0]); + let inserted = client + .execute( + "INSERT INTO items (id, embedding) VALUES ($1, $2)", + &[&3i32, &official_vec], + ) + .expect("official pgvector Vector parameter INSERT should succeed"); + assert_eq!(inserted, 1); + + // Read the stored vector back with the official client's binary decoder. + let row = client + .query_one("SELECT embedding FROM items WHERE id = 3", &[]) + .expect("SELECT of the vector column should succeed"); + let decoded: pgvector::Vector = row.get(0); + assert_eq!(decoded, official_vec, "official Vector must round-trip"); + + // Nearest-neighbour search via the pgvector operator over the same data. + let row = client + .query_one( + "SELECT id FROM items ORDER BY embedding <-> '[7,8,9]' LIMIT 1", + &[], + ) + .expect("distance query should succeed"); + let nearest: i32 = row.get(0); + assert_eq!(nearest, 3, "nearest row must be the [7,8,9] vector"); + + let _ = stop_tx.send(()); } From 48ae88d215122efafc2fd72c51ab1530f6cbe688 Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Wed, 9 Sep 2026 12:04:28 +0800 Subject: [PATCH 03/10] ci: add pgvector feature test job 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. --- .github/workflows/ci.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c29878b3..0bb49ddf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -92,6 +92,9 @@ jobs: - name: functions-full flags: "--features datafusion-pg-functions/full" optional: false + - name: pgvector + flags: "--features datafusion-postgres/pgvector" + optional: false - name: postgis flags: "--features datafusion-postgres/postgis" optional: true From f19a7aed2ae17960250685bf7dc3951ed32da646 Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Wed, 9 Sep 2026 12:14:41 +0800 Subject: [PATCH 04/10] test(integration): add pgvector python integration test 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. --- tests-integration/test.sh | 29 ++++++++- tests-integration/test_pgvector.py | 99 ++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 tests-integration/test_pgvector.py diff --git a/tests-integration/test.sh b/tests-integration/test.sh index 2ef76b23..fe53b46a 100755 --- a/tests-integration/test.sh +++ b/tests-integration/test.sh @@ -15,7 +15,7 @@ done # Function to cleanup processes cleanup() { echo "🧹 Cleaning up processes..." - for pid in $CSV_PID $TRANSACTION_PID $PARQUET_PID $RBAC_PID $SSL_PID $POSTGIS_PID $FDW_PID; do + for pid in $CSV_PID $TRANSACTION_PID $PARQUET_PID $RBAC_PID $SSL_PID $POSTGIS_PID $FDW_PID $PGVECTOR_PID; do if [ ! -z "$pid" ]; then kill -9 $pid 2>/dev/null || true fi @@ -248,6 +248,32 @@ else echo "⏭️ Skipped (--skip-postgis)" fi +# Test 7: pgvector +echo "" +echo "🧪 Test 7: pgvector Support" +echo "---------------------------" +wait_for_port 5438 +../target/debug/datafusion-postgres-cli -p 5438 --csv delhi:delhiclimate.csv & +PGVECTOR_PID=$! +sleep 5 + +# Check if server is actually running +if ! ps -p $PGVECTOR_PID > /dev/null 2>&1; then + echo "❌ pgvector server failed to start" + exit 1 +fi + +if python test_pgvector.py; then + echo "✅ pgvector test passed" +else + echo "❌ pgvector test failed" + kill -9 $PGVECTOR_PID 2>/dev/null || true + exit 1 +fi + +kill -9 $PGVECTOR_PID 2>/dev/null || true +sleep 3 + echo "" echo "🎉 All enhanced integration tests passed!" echo "==========================================" @@ -261,6 +287,7 @@ echo " ✅ Array types and complex data type support" echo " ✅ Improved pg_catalog system tables" echo " ✅ PostgreSQL function compatibility" echo " ✅ SSL/TLS encryption support" +echo " ✅ pgvector support (vector columns, literals, distance operators)" if [ -z "$SKIP_POSTGIS" ]; then echo " ✅ PostGIS spatial functions support" fi diff --git a/tests-integration/test_pgvector.py b/tests-integration/test_pgvector.py new file mode 100644 index 00000000..4062bb1d --- /dev/null +++ b/tests-integration/test_pgvector.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +""" +Integration tests for pgvector functionality. +Tests typical pgvector workflows against the datafusion-postgres server: +DDL with a `vector(n)` column, INSERT of pgvector literals, the distance +operators (`<->` L2, `<#>` negative inner product, `<=>` cosine distance), +and nearest-neighbour queries. +""" + +import psycopg + + +def main(): + print("🧪 Testing pgvector Queries") + print("=" * 50) + + conn = psycopg.connect("host=127.0.0.1 port=5438 user=postgres dbname=public") + conn.autocommit = True + + with conn.cursor() as cur: + print("\n📋 Test 1: Create table with vector(3) column") + test_create_table(cur) + + print("\n📋 Test 2: INSERT vector literals") + test_insert_literals(cur) + + print("\n📋 Test 3: Nearest-neighbour (ORDER BY <->)") + test_nearest_neighbour(cur) + + print("\n📋 Test 4: Distance operators") + test_distance_operators(cur) + + conn.close() + print("\n✅ All pgvector tests passed!") + + +def test_create_table(cur): + """The canonical pgvector DDL must succeed and the vector type be known.""" + cur.execute("DROP TABLE IF EXISTS items") + cur.execute( + "CREATE TABLE items (id int PRIMARY KEY, embedding vector(3))" + ) + # The pgvector `vector` type must be registered in pg_catalog.pg_type so + # drivers can introspect it. + cur.execute( + "SELECT count(*) FROM pg_catalog.pg_type WHERE typname = 'vector'" + ) + count = cur.fetchone()[0] + assert count >= 1, "expected a vector row in pg_catalog.pg_type" + print(" ✓ CREATE TABLE with vector(3); vector type registered in pg_type") + + +def test_insert_literals(cur): + """INSERT ... VALUES with pgvector bracket literals.""" + cur.execute( + "INSERT INTO items (id, embedding) VALUES " + "(1, '[1,2,3]'), (2, '[4,5,6]')" + ) + cur.execute("SELECT count(*) FROM items") + count = cur.fetchone()[0] + assert count == 2, f"expected 2 rows, got {count}" + print(f" ✓ Inserted 2 vector rows (count = {count})") + + +def test_nearest_neighbour(cur): + """ORDER BY embedding <-> '[..]' LIMIT n returns the closest rows.""" + cur.execute( + "SELECT id FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 1" + ) + nearest = cur.fetchone()[0] + assert nearest == 1, f"expected id 1 to be nearest, got {nearest}" + print(f" ✓ Nearest neighbour to [1,2,3] is id {nearest}") + + +def test_distance_operators(cur): + """The L2, negative inner product and cosine operators return numbers.""" + # <-> : L2 distance; [1,2,3] vs itself is 0 + cur.execute("SELECT embedding <-> '[1,2,3]' FROM items WHERE id = 1") + l2 = cur.fetchone()[0] + assert abs(l2 - 0.0) < 1e-6, f"expected L2 distance 0, got {l2}" + print(f" ✓ <-> L2 distance: {l2}") + + # <#>: negative inner product; [4,5,6] . [4,5,6] = 77 -> -77 + cur.execute("SELECT embedding <#> '[4,5,6]' FROM items WHERE id = 2") + neg_ip = cur.fetchone()[0] + assert abs(neg_ip - (-77.0)) < 1e-4, ( + f"expected negative inner product -77, got {neg_ip}" + ) + print(f" ✓ <#> negative inner product: {neg_ip}") + + # <=>: cosine distance; identical vectors -> 0 + cur.execute("SELECT embedding <=> '[1,2,3]' FROM items WHERE id = 1") + cosine = cur.fetchone()[0] + assert abs(cosine - 0.0) < 1e-6, f"expected cosine distance 0, got {cosine}" + print(f" ✓ <=> cosine distance: {cosine}") + + +if __name__ == "__main__": + main() From d3edf2bf57107d1a7bf9d840cea3d2c6505224b5 Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Wed, 9 Sep 2026 12:28:02 +0800 Subject: [PATCH 05/10] test(integration): add --skip-pgvector flag to test.sh 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 | 49 +++++++++++++++++++++++---------------- 1 file changed, 29 insertions(+), 20 deletions(-) diff --git a/tests-integration/test.sh b/tests-integration/test.sh index fe53b46a..119ec7a5 100755 --- a/tests-integration/test.sh +++ b/tests-integration/test.sh @@ -2,13 +2,16 @@ set -e -# Optional flag: --skip-postgis -# Skips the PostGIS integration test and builds without the postgis -# feature, for use when the geodatafusion dependency is not ready. +# Optional flags: +# --skip-postgis Skips the PostGIS integration test and builds without the +# postgis feature (for when geodatafusion lags upstream). +# --skip-pgvector Skips the pgvector integration test. SKIP_POSTGIS="" +SKIP_PGVECTOR="" for arg in "$@"; do case $arg in --skip-postgis) SKIP_POSTGIS=1 ;; + --skip-pgvector) SKIP_PGVECTOR=1 ;; esac done @@ -252,26 +255,30 @@ fi echo "" echo "🧪 Test 7: pgvector Support" echo "---------------------------" -wait_for_port 5438 -../target/debug/datafusion-postgres-cli -p 5438 --csv delhi:delhiclimate.csv & -PGVECTOR_PID=$! -sleep 5 +if [ -z "$SKIP_PGVECTOR" ]; then + wait_for_port 5438 + ../target/debug/datafusion-postgres-cli -p 5438 --csv delhi:delhiclimate.csv & + PGVECTOR_PID=$! + sleep 5 -# Check if server is actually running -if ! ps -p $PGVECTOR_PID > /dev/null 2>&1; then - echo "❌ pgvector server failed to start" - exit 1 -fi + # Check if server is actually running + if ! ps -p $PGVECTOR_PID > /dev/null 2>&1; then + echo "❌ pgvector server failed to start" + exit 1 + fi + + if python test_pgvector.py; then + echo "✅ pgvector test passed" + else + echo "❌ pgvector test failed" + kill -9 $PGVECTOR_PID 2>/dev/null || true + exit 1 + fi -if python test_pgvector.py; then - echo "✅ pgvector test passed" -else - echo "❌ pgvector test failed" kill -9 $PGVECTOR_PID 2>/dev/null || true - exit 1 +else + echo "⏭️ Skipped (--skip-pgvector)" fi - -kill -9 $PGVECTOR_PID 2>/dev/null || true sleep 3 echo "" @@ -287,7 +294,9 @@ echo " ✅ Array types and complex data type support" echo " ✅ Improved pg_catalog system tables" echo " ✅ PostgreSQL function compatibility" echo " ✅ SSL/TLS encryption support" -echo " ✅ pgvector support (vector columns, literals, distance operators)" +if [ -z "$SKIP_PGVECTOR" ]; then + echo " ✅ pgvector support (vector columns, literals, distance operators)" +fi if [ -z "$SKIP_POSTGIS" ]; then echo " ✅ PostGIS spatial functions support" fi From cc54e8cb30f9c1cca531ec723d4570059e96b1a6 Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Wed, 9 Sep 2026 13:01:59 +0800 Subject: [PATCH 06/10] ci: build/test pgvector feature and add integration-pgvector job - 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. --- .github/workflows/ci.yml | 33 +++++++++++++++++++++++++++++++-- tests-integration/test.sh | 15 +++++++++++---- 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0bb49ddf..c330c084 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -139,7 +139,36 @@ jobs: run: nix develop --command podman info >/dev/null 2>&1 || true - run: | cd tests-integration - nix develop --command bash test.sh --skip-postgis + nix develop --command bash test.sh --skip-postgis --skip-pgvector + + # pgvector has its own dedicated integration job so it can run (and fail) + # independently of the base suite. + integration-pgvector: + name: Integration tests (pgvector) + runs-on: ubuntu-latest + timeout-minutes: 15 + needs: [test] + steps: + - uses: actions/checkout@v6 + - uses: insightsengineering/disk-space-reclaimer@v1 + - uses: cachix/install-nix-action@v31 + - name: Build CLI and run pgvector integration tests + run: | + nix develop --command bash -ec ' + set -e + cargo build -p datafusion-postgres-cli + cd tests-integration + ../target/debug/datafusion-postgres-cli -p 5438 --csv delhi:delhiclimate.csv & + pid=$! + trap "kill -9 $pid 2>/dev/null || true" EXIT + for i in $(seq 1 30); do + if python -c "import socket,sys; s=socket.socket(); sys.exit(0 if s.connect_ex((\"127.0.0.1\", 5438)) == 0 else 1)" 2>/dev/null; then + break + fi + sleep 1 + done + python test_pgvector.py + ' # Advisory: full integration suite including the PostGIS tests, built # with the postgis feature. Allowed to fail while geodatafusion lags. @@ -157,7 +186,7 @@ jobs: run: nix develop --command podman info >/dev/null 2>&1 || true - run: | cd tests-integration - nix develop --command bash test.sh + nix develop --command bash test.sh --skip-pgvector msrv: name: MSRV diff --git a/tests-integration/test.sh b/tests-integration/test.sh index 119ec7a5..c0712cfb 100755 --- a/tests-integration/test.sh +++ b/tests-integration/test.sh @@ -51,13 +51,20 @@ wait_for_port() { echo "🚀 Running DataFusion PostgreSQL Integration Tests" echo "==================================================" -# Build the project +# Build the project, enabling the optional features whose tests will run. echo "Building datafusion-postgres..." cd .. -if [ -n "$SKIP_POSTGIS" ]; then - cargo build +BUILD_FEATURES="" +if [ -z "$SKIP_PGVECTOR" ]; then + BUILD_FEATURES="${BUILD_FEATURES:+$BUILD_FEATURES }datafusion-postgres/pgvector" +fi +if [ -z "$SKIP_POSTGIS" ]; then + BUILD_FEATURES="${BUILD_FEATURES:+$BUILD_FEATURES }datafusion-postgres/postgis" +fi +if [ -n "$BUILD_FEATURES" ]; then + cargo build --features "$BUILD_FEATURES" else - cargo build --features datafusion-postgres/postgis + cargo build fi cd tests-integration From c2550c6ba9f6a4bfb00fee61b9d42c56460eb99f Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Wed, 9 Sep 2026 16:46:23 +0800 Subject: [PATCH 07/10] ci: simplify integration-pgvector job to reuse test.sh 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. --- .github/workflows/ci.yml | 25 +++++++------------------ 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c330c084..32151773 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -142,7 +142,8 @@ jobs: nix develop --command bash test.sh --skip-postgis --skip-pgvector # pgvector has its own dedicated integration job so it can run (and fail) - # independently of the base suite. + # independently of the base suite. `--skip-postgis` builds with the pgvector + # feature and runs the full python suite including the pgvector tests. integration-pgvector: name: Integration tests (pgvector) runs-on: ubuntu-latest @@ -152,23 +153,11 @@ jobs: - uses: actions/checkout@v6 - uses: insightsengineering/disk-space-reclaimer@v1 - uses: cachix/install-nix-action@v31 - - name: Build CLI and run pgvector integration tests - run: | - nix develop --command bash -ec ' - set -e - cargo build -p datafusion-postgres-cli - cd tests-integration - ../target/debug/datafusion-postgres-cli -p 5438 --csv delhi:delhiclimate.csv & - pid=$! - trap "kill -9 $pid 2>/dev/null || true" EXIT - for i in $(seq 1 30); do - if python -c "import socket,sys; s=socket.socket(); sys.exit(0 if s.connect_ex((\"127.0.0.1\", 5438)) == 0 else 1)" 2>/dev/null; then - break - fi - sleep 1 - done - python test_pgvector.py - ' + - name: Initialize podman + run: nix develop --command podman info >/dev/null 2>&1 || true + - run: | + cd tests-integration + nix develop --command bash test.sh --skip-postgis # Advisory: full integration suite including the PostGIS tests, built # with the postgis feature. Allowed to fail while geodatafusion lags. From 4b5fb58d4299ded10f4e95e6e06bdedd447017ce Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Thu, 10 Sep 2026 09:40:57 +0800 Subject: [PATCH 08/10] fix(param): prefer client parameter type hint over server override 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. --- arrow-pg/src/datatypes/df.rs | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/arrow-pg/src/datatypes/df.rs b/arrow-pg/src/datatypes/df.rs index 8c58b763..5d76cf0e 100644 --- a/arrow-pg/src/datatypes/df.rs +++ b/arrow-pg/src/datatypes/df.rs @@ -466,15 +466,17 @@ where pg_type_hint: Option, inferenced_type: Option<&DataType>, ) -> PgWireResult { - // A concrete server-decided type wins -- it reflects the parameter's - // semantically-typed column (oid-alias, pgvector, ...). UNKNOWN is not - // authoritative so the historical hint/inferred fallbacks still apply. - if let Some(ty) = server_type + // The client-provided hint (Parse parameter type OIDs) wins: it is the + // type the client actually encoded the parameter with. The + // server-decided type is only a fallback for clients that send no + // types (e.g. tokio-postgres); it carries the semantically-typed + // overrides (oid-alias, pgvector) that the physical Arrow mapping + // would otherwise lose. + if let Some(ty) = pg_type_hint { + Ok(ty.clone()) + } else if let Some(ty) = server_type && *ty != Type::UNKNOWN { - return Ok(ty.clone()); - } - if let Some(ty) = pg_type_hint { Ok(ty.clone()) } else if let Some(infer_type) = inferenced_type { into_pg_type(infer_type) From c17b46de0485aa1cab75de008a3f7993a5c1de75 Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Thu, 10 Sep 2026 12:28:46 +0800 Subject: [PATCH 09/10] refactor(pgvector): use DataFusion ExprPlanner and a QueryHook 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). --- arrow-pg/src/datatypes/df.rs | 55 +- datafusion-pg-catalog/src/pg_catalog.rs | 140 +---- .../src/pg_catalog/pgvector.rs | 147 +++++ datafusion-pg-catalog/src/sql.rs | 5 - datafusion-pg-catalog/src/sql/parser.rs | 12 +- datafusion-pg-catalog/src/sql/rules.rs | 484 ----------------- .../src/sql/vector_insert.rs | 398 -------------- datafusion-postgres/src/handlers.rs | 28 +- datafusion-postgres/src/lib.rs | 2 + datafusion-postgres/src/pgvector.rs | 512 ++++++++++++++++++ 10 files changed, 722 insertions(+), 1061 deletions(-) create mode 100644 datafusion-pg-catalog/src/pg_catalog/pgvector.rs delete mode 100644 datafusion-pg-catalog/src/sql/vector_insert.rs create mode 100644 datafusion-postgres/src/pgvector.rs diff --git a/arrow-pg/src/datatypes/df.rs b/arrow-pg/src/datatypes/df.rs index 5d76cf0e..d5d6dc8b 100644 --- a/arrow-pg/src/datatypes/df.rs +++ b/arrow-pg/src/datatypes/df.rs @@ -272,6 +272,44 @@ fn coerce_interval_value(value: Option, target: &DataType) -> PgWireRe } } +/// A Postgres `oid` parameter decoded as a signed `i32`. +/// +/// Postgres has no unsigned integer types, and DataFusion catalog `oid` columns +/// are stored as `Int32`, so decode the 4-byte OID value directly as `i32` +/// rather than going through `u32`. +#[derive(Debug)] +struct OidParam(i32); + +impl<'a> postgres_types::FromSql<'a> for OidParam { + fn from_sql( + _ty: &Type, + raw: &'a [u8], + ) -> Result> { + if raw.len() != 4 { + return Err("oid parameter must be exactly 4 bytes".into()); + } + Ok(OidParam(i32::from_be_bytes([ + raw[0], raw[1], raw[2], raw[3], + ]))) + } + + fn accepts(ty: &Type) -> bool { + ty == &Type::OID + } +} + +impl<'a> pgwire::types::FromSqlText<'a> for OidParam { + fn from_sql_text( + _ty: &Type, + input: &'a [u8], + _format_options: &FormatOptions, + ) -> Result> { + let text = std::str::from_utf8(input)?; + let value: i64 = text.trim().parse()?; + Ok(OidParam(value as i32)) + } +} + /// Decode a pgvector `vector` parameter value sent by a client over the wire. #[cfg(feature = "pgvector")] #[derive(Debug)] @@ -1043,16 +1081,19 @@ where // Store MAC addresses as strings for now deserialized_params.push(ScalarValue::Utf8(value)); } - // PostgreSQL `oid` (unsigned 32-bit). Drivers bind these with u32; - // DataFusion catalog oid columns are Int32, so coerce when the - // inferred type is Int32. + // PostgreSQL `oid` (unsigned 32-bit on the wire, but Postgres has no + // unsigned integer types). DataFusion catalog oid columns are Int32, + // so decode as i32 and coerce if the inferred type differs. Type::OID => { - let value = portal.parameter::(i, &pg_type)?; + let value = portal.parameter::(i, &pg_type)?; match inferenced_type { - Some(DataType::Int32) => { - deserialized_params.push(ScalarValue::Int32(value.map(|v| v as i32))); + Some(target) if !matches!(target, DataType::Int32) => { + deserialized_params + .push(coerce_int_value(value.map(|v| v.0 as i64), target)?); + } + _ => { + deserialized_params.push(ScalarValue::Int32(value.map(|v| v.0))); } - _ => deserialized_params.push(ScalarValue::UInt32(value)), } } // TODO: add more advanced types (composite types, ranges, etc.) diff --git a/datafusion-pg-catalog/src/pg_catalog.rs b/datafusion-pg-catalog/src/pg_catalog.rs index 0f04332f..c90db8e0 100644 --- a/datafusion-pg-catalog/src/pg_catalog.rs +++ b/datafusion-pg-catalog/src/pg_catalog.rs @@ -7,7 +7,7 @@ use datafusion::arrow::array::{ ArrayRef, AsArray, BooleanBuilder, Int32Builder, RecordBatch, StringArray, StringBuilder, as_boolean_array, }; -use datafusion::arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; +use datafusion::arrow::datatypes::{DataType, Field, Int32Type, SchemaRef}; use datafusion::arrow::ipc::reader::FileReader; use datafusion::catalog::streaming::StreamingTable; use datafusion::catalog::{MemTable, SchemaProvider, TableFunctionImpl}; @@ -46,6 +46,8 @@ pub mod pg_settings; pub mod pg_stat_gssapi; pub mod pg_tables; pub mod pg_views; +#[cfg(feature = "pgvector")] +pub mod pgvector; pub mod quote_ident_udf; const PG_CATALOG_TABLE_PG_AGGREGATE: &str = "pg_aggregate"; @@ -1068,7 +1070,7 @@ impl PgCatalogStaticTables { // ...) find it, and tag the internal `"char"` column for correct wire // encoding. #[cfg(feature = "pgvector")] - let tables = tables.with_pg_vector_support()?; + let tables = pgvector::with_pg_vector_support(tables)?; Ok(tables) } @@ -1077,140 +1079,6 @@ impl PgCatalogStaticTables { fn create_arrow_table(data_bytes: Vec) -> Result> { ArrowTable::from_ipc_data(data_bytes).map(Arc::new) } - - /// Return `self` with the pg_catalog pieces needed for pgvector-aware - /// clients that resolve unknown result types via `pg_type` introspection: - /// - /// * a `vector` type row (OID 16385) in `pg_type`, and - /// * the internal `"char"` wire type on `pg_type.typtype`. - /// - /// The vector row intentionally lives in `pg_catalog` (namespace OID 11) so - /// the `pg_type`/`pg_namespace` join the introspection query performs - /// resolves without depending on user-schema oids. - #[cfg(feature = "pgvector")] - fn with_pg_vector_support(mut self) -> Result { - const VECTOR_OID: i32 = 16385; - - // 1. Tag pg_type.typtype as a Postgres internal `"char"` column so the - // wire layer encodes it correctly for clients decoding `typtype`. - let pg_type = Arc::new(arrow_table_with_field_metadata( - &self.pg_type, - "typtype", - "pg.char", - "char", - )?); - - // 2. Append the `vector` row to pg_type (all other columns get safe - // defaults; the introspection query only reads the ones we set). - let pg_type = Arc::new(arrow_table_append_row(&pg_type, |field, scalar| { - match field.name().as_str() { - "oid" => *scalar = datafusion::scalar::ScalarValue::Int32(Some(VECTOR_OID)), - "typname" => { - *scalar = datafusion::scalar::ScalarValue::Utf8(Some("vector".to_string())) - } - "typtype" => *scalar = datafusion::scalar::ScalarValue::Utf8(Some("b".to_string())), - // pg_catalog namespace OID 11 (see the doc comment above). - "typnamespace" => *scalar = datafusion::scalar::ScalarValue::Int32(Some(11)), - _ => {} - } - })?); - - self.pg_type = pg_type; - Ok(self) - } -} - -/// Rebuild `table` with `key = value` metadata added to the field named `name`. -/// -/// Arrow keeps field metadata in the schema; this rebuilds the schema and every -/// record batch so the wire layer (`arrow-pg`) sees the marker when encoding. -#[cfg(feature = "pgvector")] -fn arrow_table_with_field_metadata( - table: &ArrowTable, - name: &str, - key: &str, - value: &str, -) -> Result { - let schema = table.schema(); - let fields = schema - .fields() - .iter() - .map(|field| { - if field.name() == name { - let mut metadata = field.metadata().clone(); - metadata.insert(key.to_string(), value.to_string()); - (**field).clone().with_metadata(metadata) - } else { - (**field).clone() - } - }) - .collect::>(); - let new_schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); - let mut batches = Vec::with_capacity(table.data().len()); - for batch in table.data() { - batches.push( - RecordBatch::try_new(Arc::clone(&new_schema), batch.columns().to_vec()) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?, - ); - } - Ok(ArrowTable { - schema: new_schema, - data: batches, - }) -} - -/// Return a copy of `table` with one extra row appended. `fill` lets callers -/// override the default scalar produced for each column. -/// -/// `fill` is invoked for every field with a fresh default scalar (int 0 / -/// string "" / bool false / ...), matching the field's data type; non-default -/// catalog values are set by the caller. -#[cfg(feature = "pgvector")] -fn arrow_table_append_row( - table: &ArrowTable, - mut fill: impl FnMut(&Field, &mut datafusion::scalar::ScalarValue), -) -> Result { - let schema = table.schema(); - let columns = schema - .fields() - .iter() - .map(|field| { - let mut scalar = default_scalar(field.data_type())?; - fill(field, &mut scalar); - scalar.to_array_of_size(1) - }) - .collect::>>()?; - let row = RecordBatch::try_new(Arc::clone(&schema), columns) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; - let mut data = table.data().to_vec(); - data.push(row); - Ok(ArrowTable { - schema: Arc::clone(&schema), - data, - }) -} - -/// A single default [`ScalarValue`] for a catalog column of `data_type`. -/// -/// Only the data types present in the exported `pg_type` schema are handled; -/// anything else yields a `NotImplemented` error rather than guessing. -#[cfg(feature = "pgvector")] -fn default_scalar(data_type: &DataType) -> Result { - use datafusion::scalar::ScalarValue; - Ok(match data_type { - DataType::Null => ScalarValue::Null, - DataType::Boolean => ScalarValue::Boolean(Some(false)), - DataType::Int16 => ScalarValue::Int16(Some(0)), - DataType::Int32 => ScalarValue::Int32(Some(0)), - DataType::Int64 => ScalarValue::Int64(Some(0)), - DataType::Utf8 => ScalarValue::Utf8(Some(String::new())), - DataType::LargeUtf8 => ScalarValue::LargeUtf8(Some(String::new())), - other => { - return Err(DataFusionError::NotImplemented(format!( - "no default catalog scalar for {other:?}" - ))); - } - }) } pub fn create_current_schemas_udf() -> ScalarUDF { diff --git a/datafusion-pg-catalog/src/pg_catalog/pgvector.rs b/datafusion-pg-catalog/src/pg_catalog/pgvector.rs new file mode 100644 index 00000000..3e67444a --- /dev/null +++ b/datafusion-pg-catalog/src/pg_catalog/pgvector.rs @@ -0,0 +1,147 @@ +//! pgvector support for the `pg_catalog` schema. +//! +//! DataFusion-backed pgvector clients resolve unknown result-type OIDs through +//! `pg_catalog.pg_type`. To make the pgvector `vector` type resolvable, this +//! module injects a `vector` row into the static `pg_type` table and tags +//! `pg_type.typtype` with the internal Postgres `"char"` wire type. +//! +//! The vector row intentionally lives in `pg_catalog` (namespace OID 11) so the +//! `pg_type`/`pg_namespace` join performed by driver introspection resolves +//! without depending on user-schema OIDs. + +use std::sync::Arc; + +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::error::{DataFusionError, Result}; +use datafusion::scalar::ScalarValue; + +use super::{ArrowTable, PgCatalogStaticTables}; + +/// Fixed OID reported for pgvector `vector` columns, matching +/// `arrow_pg::datatypes::PG_VECTOR_TYPE_OID`. +const VECTOR_OID: i32 = 16385; + +/// Canonical OID of the `pg_catalog` namespace (`PG_CATALOG_NAMESPACE`). +const PG_CATALOG_NAMESPACE_OID: i32 = 11; + +/// Return `tables` with the pgvector pieces needed for driver introspection: +/// +/// * a `vector` type row (OID 16385) in `pg_type`, and +/// * the internal `"char"` wire type on `pg_type.typtype`. +pub(crate) fn with_pg_vector_support( + mut tables: PgCatalogStaticTables, +) -> Result { + // 1. Tag pg_type.typtype as a Postgres internal `"char"` column so the wire + // layer encodes it correctly for clients decoding `typtype`. + let pg_type = Arc::new(with_field_metadata( + &tables.pg_type, + "typtype", + "pg.char", + "char", + )?); + + // 2. Append the `vector` row to pg_type (all other columns get safe + // defaults; the introspection query only reads the ones we set). + let pg_type = Arc::new(append_row(&pg_type, |field, scalar| { + match field.name().as_str() { + "oid" => *scalar = ScalarValue::Int32(Some(VECTOR_OID)), + "typname" => *scalar = ScalarValue::Utf8(Some("vector".to_string())), + "typtype" => *scalar = ScalarValue::Utf8(Some("b".to_string())), + "typnamespace" => *scalar = ScalarValue::Int32(Some(PG_CATALOG_NAMESPACE_OID)), + _ => {} + } + })?); + + tables.pg_type = pg_type; + Ok(tables) +} + +/// Rebuild `table` with `key = value` metadata added to the field named `name`. +/// +/// Arrow keeps field metadata in the schema; this rebuilds the schema and every +/// record batch so the wire layer (`arrow-pg`) sees the marker when encoding. +fn with_field_metadata( + table: &ArrowTable, + name: &str, + key: &str, + value: &str, +) -> Result { + let schema = table.schema(); + let fields = schema + .fields() + .iter() + .map(|field| { + if field.name() == name { + let mut metadata = field.metadata().clone(); + metadata.insert(key.to_string(), value.to_string()); + (**field).clone().with_metadata(metadata) + } else { + (**field).clone() + } + }) + .collect::>(); + let new_schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); + let mut batches = Vec::with_capacity(table.data().len()); + for batch in table.data() { + batches.push( + RecordBatch::try_new(Arc::clone(&new_schema), batch.columns().to_vec()) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?, + ); + } + Ok(ArrowTable { + schema: new_schema, + data: batches, + }) +} + +/// Return a copy of `table` with one extra row appended. `fill` lets callers +/// override the default scalar produced for each column. +/// +/// `fill` is invoked for every field with a fresh default scalar (int 0 / +/// string "" / bool false / ...), matching the field's data type; non-default +/// catalog values are set by the caller. +fn append_row( + table: &ArrowTable, + mut fill: impl FnMut(&Field, &mut ScalarValue), +) -> Result { + let schema = table.schema(); + let columns = schema + .fields() + .iter() + .map(|field| { + let mut scalar = default_scalar(field.data_type())?; + fill(field, &mut scalar); + scalar.to_array_of_size(1) + }) + .collect::>>()?; + let row = RecordBatch::try_new(Arc::clone(&schema), columns) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + let mut data = table.data().to_vec(); + data.push(row); + Ok(ArrowTable { + schema: Arc::clone(&schema), + data, + }) +} + +/// A single default [`ScalarValue`] for a catalog column of `data_type`. +/// +/// Only the data types present in the exported `pg_type` schema are handled; +/// anything else yields a `NotImplemented` error rather than guessing. +fn default_scalar(data_type: &DataType) -> Result { + Ok(match data_type { + DataType::Null => ScalarValue::Null, + DataType::Boolean => ScalarValue::Boolean(Some(false)), + DataType::Int16 => ScalarValue::Int16(Some(0)), + DataType::Int32 => ScalarValue::Int32(Some(0)), + DataType::Int64 => ScalarValue::Int64(Some(0)), + DataType::Utf8 => ScalarValue::Utf8(Some(String::new())), + DataType::LargeUtf8 => ScalarValue::LargeUtf8(Some(String::new())), + other => { + return Err(DataFusionError::NotImplemented(format!( + "no default catalog scalar for {other:?}" + ))); + } + }) +} diff --git a/datafusion-pg-catalog/src/sql.rs b/datafusion-pg-catalog/src/sql.rs index 452b0807..42c6a178 100644 --- a/datafusion-pg-catalog/src/sql.rs +++ b/datafusion-pg-catalog/src/sql.rs @@ -1,8 +1,3 @@ mod parser; pub use parser::PostgresCompatibilityParser; pub mod rules; - -#[cfg(feature = "pgvector")] -mod vector_insert; -#[cfg(feature = "pgvector")] -pub use vector_insert::rewrite_vector_insert; diff --git a/datafusion-pg-catalog/src/sql/parser.rs b/datafusion-pg-catalog/src/sql/parser.rs index 03a3eda7..59966100 100644 --- a/datafusion-pg-catalog/src/sql/parser.rs +++ b/datafusion-pg-catalog/src/sql/parser.rs @@ -18,8 +18,6 @@ use super::rules::ResolveUnqualifiedIdentifier; use super::rules::RewriteArrayAnyAllOperation; use super::rules::RewritePgCatalogOperator; use super::rules::RewriteRegCastToSubquery; -#[cfg(feature = "pgvector")] -use super::rules::RewriteVectorOperators; use super::rules::SqlStatementRewriteRule; use super::rules::StripCallableQualifier; use super::rules::StripCollate; @@ -311,8 +309,7 @@ impl PostgresCompatibilityParser { )); } - #[cfg_attr(not(feature = "pgvector"), allow(unused_mut))] - let mut rewrite_rules: Vec> = vec![ + let rewrite_rules: Vec> = vec![ // The blacklist substitution in `parse()` runs before any of // these rules, so by the time they see the statement any // blacklisted fragment has already been replaced. @@ -333,13 +330,6 @@ impl PostgresCompatibilityParser { Arc::new(FixVersionColumnName), ]; - // pgvector support: distance operators (`<->` / `<#>` / `<=>`) and - // vector literals. Runs last -- it needs to see oid/array rewrites in - // operands already applied, and rewrites operators that no other rule - // touches. - #[cfg(feature = "pgvector")] - rewrite_rules.push(Arc::new(RewriteVectorOperators)); - Self { blacklist: mapping, rewrite_rules, diff --git a/datafusion-pg-catalog/src/sql/rules.rs b/datafusion-pg-catalog/src/sql/rules.rs index 798e69b2..b019dac8 100644 --- a/datafusion-pg-catalog/src/sql/rules.rs +++ b/datafusion-pg-catalog/src/sql/rules.rs @@ -1232,277 +1232,6 @@ impl SqlStatementRewriteRule for FixVersionColumnName { } } -/// Rewrite pgvector distance operators and vector literals into plain SQL that -/// DataFusion can plan, backed by its built-in array distance functions. -/// -/// # Operators -/// -/// | pgvector | meaning | rewrite | -/// | -------- | --------------- | ------------------------------------ | -/// | `<->` | L2 distance | `array_distance(l, r)` | -/// | `<#>` | negative dot | `-inner_product(l, r)` | -/// | `<=>` | cosine distance | `cosine_distance(l, r)` | -/// -/// sqlparser parses `<->` / `<=>` / `<#>` as the `LtDashGt`, `Spaceship` and -/// `Custom("<#>")` [`BinaryOperator`]s (see the pgvector support notes). -/// -/// # Vector literals -/// -/// pgvector accepts a vector as a bracket string, either bare (`'[1,2,3]'`) or -/// cast (`'[1,2,3]'::vector`). Both appear in the canonical usage: -/// `ORDER BY embedding <-> '[1,2,3]'`. Because pgvector treats an unknown -/// literal as the `vector` type for these operators, each operand that is such -/// a string literal is rewritten to an `ARRAY[...]` of float literals. The -/// built-in distance functions then coerce the `FixedSizeList(Float32, n)` -/// column and the `List(Float64)` literal to a common `List(Float64)` and -/// enforce the equal-length (dimension) requirement at runtime. -/// -/// # Why a rule (not UDFs / casts) -/// -/// * No schema is needed: a literal is recognized syntactically, exactly like -/// [`RewriteArrayAnyAllOperation`] / [`FixArrayLiteral`]. -/// * Dimension of the literal is known at rewrite time, so no dynamically-typed -/// `parse_vector`-style UDF is required. -/// * The rewritten SQL is re-serialised and re-parsed by DataFusion (see -/// `handlers.rs`), so the emitted form must be plain parseable SQL. -/// -/// # Coverage -/// -/// The generic AST visitor reaches vector operators in projections, `WHERE`, -/// `HAVING` and joins; sqlparser's visitor does *not* descend into `ORDER BY`, -/// so [`RewriteVectorOperatorsVisitor::rewrite_order_by`] handles those -/// expressions explicitly (this is where pgvector queries place the operator). -#[cfg(feature = "pgvector")] -#[derive(Debug)] -pub struct RewriteVectorOperators; - -#[cfg(feature = "pgvector")] -#[derive(Clone, Copy)] -enum DistanceFunc { - /// `<->` L2 distance. - ArrayDistance, - /// `<#>` negative inner product. - InnerProduct, - /// `<=>` cosine distance. - CosineDistance, -} - -#[cfg(feature = "pgvector")] -impl DistanceFunc { - fn name(self) -> &'static str { - match self { - DistanceFunc::ArrayDistance => "array_distance", - DistanceFunc::InnerProduct => "inner_product", - DistanceFunc::CosineDistance => "cosine_distance", - } - } - - /// `<#>` is the *negative* inner product, so the call must be negated. - fn negate(self) -> bool { - matches!(self, DistanceFunc::InnerProduct) - } -} - -#[cfg(feature = "pgvector")] -impl RewriteVectorOperators { - /// The distance function backing a binary operator, if it is a pgvector - /// distance operator we handle. - fn operator_to_func(op: &BinaryOperator) -> Option { - match op { - BinaryOperator::LtDashGt => Some(DistanceFunc::ArrayDistance), - BinaryOperator::Spaceship => Some(DistanceFunc::CosineDistance), - BinaryOperator::Custom(name) if name == "<#>" => Some(DistanceFunc::InnerProduct), - _ => None, - } - } - - /// True when `data_type` is the pgvector `vector` type name. - pub(crate) fn is_vector_data_type(data_type: &DataType) -> bool { - let DataType::Custom(name, _) = data_type else { - return false; - }; - name.0 - .last() - .and_then(|part| part.as_ident()) - .is_some_and(|ident| ident.value.eq_ignore_ascii_case("vector")) - } - - /// Build a SQL `ARRAY[]` literal from the numeric text between the - /// brackets of a pgvector literal like `'[1,-2.5,3]'`. Returns `None` if the - /// text does not look like a (possibly empty-forbidden) float vector. - pub(crate) fn vector_literal_to_array(text: &str) -> Option { - let text = text.trim(); - if !(text.starts_with('[') && text.ends_with(']') && text.len() >= 2) { - return None; - } - let inner = &text[1..text.len() - 1]; - if inner.trim().is_empty() { - return None; - } - - let mut elems = Vec::new(); - for part in inner.split(',') { - elems.push(Self::float_literal(part.trim())?); - } - Some(Expr::Array(Array { - elem: elems, - named: true, - })) - } - - /// A float SQL literal for `s`. Negative values become a unary-minus - /// number; non-numeric or special (`nan`, `infinity`) values return `None`. - fn float_literal(s: &str) -> Option { - if s.is_empty() { - return None; - } - let lower = s.to_lowercase(); - if matches!( - lower.as_str(), - "nan" | "inf" | "+inf" | "-inf" | "infinity" | "+infinity" | "-infinity" - ) { - return None; - } - // Parse through f32 (pgvector stores float4) and re-render with the - // shortest decimal form (Rust Display never uses scientific notation, - // so the re-parsed SQL is a plain float literal). - let value: f32 = s.parse().ok()?; - let rendered = value.to_string(); - let (neg, digits) = match rendered.strip_prefix('-') { - Some(d) => (true, d), - None => (false, rendered.as_str()), - }; - let num = Expr::Value(Value::Number(digits.to_string(), false).with_empty_span()); - if neg { - Some(Expr::UnaryOp { - op: UnaryOperator::Minus, - expr: Box::new(num), - }) - } else { - Some(num) - } - } - - /// Normalize a vector-typed operand for a distance operator: unwrap a - /// `'[...]'::vector` cast or a bare `'[...]'` string into an `ARRAY[...]` - /// literal. Anything else (columns, placeholders, functions, ...) passes - /// through unchanged. - fn coerce_vector_operand(expr: &Expr) -> Expr { - // `'[1,2,3]'::vector` / `'[1,2,3]'::public.vector` - if let Expr::Cast { - expr: inner, - data_type, - .. - } = expr - && Self::is_vector_data_type(data_type) - && let Expr::Value(ValueWithSpan { - value: Value::SingleQuotedString(text), - .. - }) = inner.as_ref() - && let Some(array) = Self::vector_literal_to_array(text) - { - return array; - } - // Bare `'[1,2,3]'` - if let Expr::Value(ValueWithSpan { - value: Value::SingleQuotedString(text), - .. - }) = expr - && let Some(array) = Self::vector_literal_to_array(text) - { - return array; - } - expr.clone() - } - - /// Rewrite `expr` in place when it is a pgvector distance `BinaryOp`, - /// recursing into the (already-rewritten-safe) operands afterwards is not - /// needed because the generic visitor will visit them separately; this - /// method handles a single operator node. - fn rewrite_binary_op(expr: &mut Expr) -> bool { - let Expr::BinaryOp { left, op, right } = expr else { - return false; - }; - let Some(func) = Self::operator_to_func(op) else { - return false; - }; - - let left = Self::coerce_vector_operand(left); - let right = Self::coerce_vector_operand(right); - let call = Expr::Function(Function { - name: ObjectName::from(vec![Ident::new(func.name())]), - args: FunctionArguments::List(FunctionArgumentList { - args: vec![ - FunctionArg::Unnamed(FunctionArgExpr::Expr(left)), - FunctionArg::Unnamed(FunctionArgExpr::Expr(right)), - ], - duplicate_treatment: None, - clauses: vec![], - }), - uses_odbc_syntax: false, - parameters: FunctionArguments::None, - filter: None, - null_treatment: None, - over: None, - within_group: vec![], - }); - - *expr = if func.negate() { - Expr::UnaryOp { - op: UnaryOperator::Minus, - expr: Box::new(call), - } - } else { - call - }; - true - } -} - -#[cfg(feature = "pgvector")] -#[derive(Debug)] -struct RewriteVectorOperatorsVisitor; - -#[cfg(feature = "pgvector")] -impl RewriteVectorOperatorsVisitor { - /// Rewrite the `ORDER BY` expressions of `query`. sqlparser's visitor does - /// not descend into `order_by`, so this is done explicitly -- pgvector's - /// canonical query places the distance operator exactly there. - fn rewrite_order_by(query: &mut Query) { - if let Some(order_by) = query.order_by.as_mut() - && let OrderByKind::Expressions(exprs) = &mut order_by.kind - { - for order_by_expr in exprs { - RewriteVectorOperators::rewrite_binary_op(&mut order_by_expr.expr); - } - } - } -} - -#[cfg(feature = "pgvector")] -impl VisitorMut for RewriteVectorOperatorsVisitor { - type Break = (); - - fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { - RewriteVectorOperators::rewrite_binary_op(expr); - ControlFlow::Continue(()) - } - - fn pre_visit_query(&mut self, query: &mut Query) -> ControlFlow { - Self::rewrite_order_by(query); - ControlFlow::Continue(()) - } -} - -#[cfg(feature = "pgvector")] -impl SqlStatementRewriteRule for RewriteVectorOperators { - fn rewrite(&self, mut s: Statement) -> Statement { - let mut visitor = RewriteVectorOperatorsVisitor; - let _ = s.visit(&mut visitor); - s - } -} - #[cfg(test)] mod tests { use super::*; @@ -1999,217 +1728,4 @@ mod tests { // A numeric operand cast to a non-oid type is left alone. assert_rewrite!(&rules, "SELECT '1'::int4", "SELECT '1'::INT4"); } - - #[cfg(feature = "pgvector")] - mod vector_rewrite { - use super::*; - use crate::sql::PostgresCompatibilityParser; - use datafusion::arrow::array::{FixedSizeListArray, Float32Array, Int64Array}; - use datafusion::arrow::datatypes::{Field as ArrowField, Schema as ArrowSchema}; - use datafusion::arrow::record_batch::RecordBatch; - use datafusion::prelude::SessionContext; - use std::sync::Arc; - - fn rewrite_vector(sql: &str) -> String { - let statement = parse(sql).expect("Failed to parse").remove(0); - let rules: Vec> = - vec![Arc::new(RewriteVectorOperators)]; - rewrite(statement, &rules).to_string() - } - - #[test] - fn l2_operator_becomes_array_distance() { - let out = - rewrite_vector("SELECT * FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 5"); - assert!( - out.contains("array_distance(embedding, ARRAY[1, 2, 3])"), - "unexpected rewrite: {out}" - ); - } - - #[test] - fn inner_product_operator_is_negated() { - let out = - rewrite_vector("SELECT * FROM items ORDER BY embedding <#> '[1,2,3]' LIMIT 5"); - assert!( - out.contains("- inner_product(embedding, ARRAY[1, 2, 3])") - || out.contains("-inner_product(embedding, ARRAY[1, 2, 3])"), - "unexpected rewrite: {out}" - ); - } - - #[test] - fn cosine_operator_becomes_cosine_distance() { - let out = - rewrite_vector("SELECT * FROM items ORDER BY embedding <=> '[1,2,3]' LIMIT 5"); - assert!( - out.contains("cosine_distance(embedding, ARRAY[1, 2, 3])"), - "unexpected rewrite: {out}" - ); - } - - #[test] - fn cast_vector_literal_is_rewritten() { - let out = rewrite_vector( - "SELECT * FROM items ORDER BY embedding <-> '[1,2,3]'::vector LIMIT 5", - ); - assert!( - out.contains("array_distance(embedding, ARRAY[1, 2, 3])"), - "unexpected rewrite: {out}" - ); - } - - #[test] - fn negative_and_float_elements_are_preserved() { - let out = rewrite_vector("SELECT embedding <-> '[-1.5, 2, 0]' FROM items"); - assert!( - out.contains("array_distance(embedding, ARRAY[-1.5, 2, 0])"), - "unexpected rewrite: {out}" - ); - } - - #[test] - fn operator_in_where_is_rewritten() { - let out = rewrite_vector("SELECT id FROM items WHERE embedding <-> '[1,2,3]' < 1.0"); - assert!( - out.contains("array_distance(embedding, ARRAY[1, 2, 3]) < 1.0"), - "unexpected rewrite: {out}" - ); - } - - /// End-to-end: parse the pgvector query through the compatibility - /// parser (as the server does), then execute the rewritten SQL against - /// a DataFusion context that holds a `FixedSizeList(Float32, 3)` column. - #[tokio::test] - async fn executes_pgvector_queries() { - let ctx = SessionContext::new(); - - // items(id, embedding vector(3)) - let element = - ArrowField::new_list_field(datafusion::arrow::datatypes::DataType::Float32, true); - let fsl_type = - datafusion::arrow::datatypes::DataType::FixedSizeList(Arc::new(element), 3); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", datafusion::arrow::datatypes::DataType::Int64, false), - ArrowField::new("embedding", fsl_type, false), - ])); - let id = Int64Array::from(vec![1, 2, 3]); - let values = Float32Array::from(vec![ - 1.0, 2.0, 3.0, // row 0 - 4.0, 5.0, 6.0, // row 1 - 7.0, 8.0, 9.0, // row 2 - ]); - let embedding = FixedSizeListArray::try_new( - Arc::new(ArrowField::new_list_field( - datafusion::arrow::datatypes::DataType::Float32, - true, - )), - 3, - Arc::new(values), - None, - ) - .unwrap(); - let batch = - RecordBatch::try_new(schema, vec![Arc::new(id), Arc::new(embedding)]).unwrap(); - ctx.register_batch("items", batch).unwrap(); - - let parser = PostgresCompatibilityParser::new(); - - // SELECT id ORDER BY L2 distance to the query vector. - let sql = parser - .parse("SELECT id FROM items ORDER BY embedding <-> '[1,2,3]' LIMIT 2") - .unwrap() - .remove(0) - .to_string(); - let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); - let ids: Vec = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec(); - assert_eq!(ids, vec![1, 2], "nearest rows by L2 distance"); - - // Project the computed distances for each operator. - let checks = [ - ( - "SELECT embedding <-> '[1,2,3]' AS d FROM items ORDER BY id", - 0.0, - ), - ( - "SELECT embedding <#> '[1,2,3]' AS d FROM items ORDER BY id", - -(14.0), - ), - ( - "SELECT embedding <=> '[1,2,3]' AS d FROM items ORDER BY id", - 0.0, - ), - ]; - for (query, expected_first) in checks { - let sql = parser.parse(query).unwrap().remove(0).to_string(); - let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); - let d = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let got = d.value(0); - assert!( - (got - expected_first).abs() < 1e-6, - "query {query:?} first distance = {got}, expected ~{expected_first}" - ); - } - } - - /// The `pg.vector` field metadata must survive DataFusion planning and - /// execution -- arrow-pg keys the wire type + text encoding off it, so - /// if it were dropped the server would report the column as float4[]. - #[tokio::test] - async fn vector_field_metadata_survives_planning_and_execution() { - use std::collections::HashMap; - - let ctx = SessionContext::new(); - let element = - ArrowField::new_list_field(datafusion::arrow::datatypes::DataType::Float32, true); - let field = ArrowField::new( - "embedding", - datafusion::arrow::datatypes::DataType::FixedSizeList(Arc::new(element.clone()), 3), - false, - ) - .with_metadata(HashMap::from([( - "pg.vector".to_string(), - "vector".to_string(), - )])); - let schema = Arc::new(ArrowSchema::new(vec![field])); - let values = Float32Array::from(vec![1.0, 2.0, 3.0]); - let embedding = - FixedSizeListArray::try_new(Arc::new(element), 3, Arc::new(values), None).unwrap(); - let batch = RecordBatch::try_new(schema, vec![Arc::new(embedding)]).unwrap(); - ctx.register_batch("items_meta", batch).unwrap(); - - let df = ctx.sql("SELECT embedding FROM items_meta").await.unwrap(); - assert_eq!( - df.schema() - .field(0) - .metadata() - .get("pg.vector") - .map(String::as_str), - Some("vector"), - "metadata must be present on the planned output schema" - ); - - let batches = df.collect().await.unwrap(); - assert_eq!( - batches[0] - .schema() - .field(0) - .metadata() - .get("pg.vector") - .map(String::as_str), - Some("vector"), - "metadata must be present on the executed record batch schema" - ); - } - } } diff --git a/datafusion-pg-catalog/src/sql/vector_insert.rs b/datafusion-pg-catalog/src/sql/vector_insert.rs deleted file mode 100644 index 129c732a..00000000 --- a/datafusion-pg-catalog/src/sql/vector_insert.rs +++ /dev/null @@ -1,398 +0,0 @@ -//! Schema-aware rewrite that lets `INSERT ... VALUES ('[1,2,3]')` write into a -//! pgvector `vector` column. -//! -//! DataFusion plans `INSERT` by coercing every provided value to the target -//! column's type with [`Expr::cast_to`], which only allows casts DataFusion -//! already knows. There is no `Utf8 -> FixedSizeList(Float32, n)` cast, so a -//! bare pgvector literal `'[1,2,3]'` fails to plan with "Cannot automatically -//! convert Utf8 to FixedSizeList(...)". -//! -//! Unlike the operator rules (which run before the query has a schema), this -//! rewrite is invoked from the server handlers with the live `SessionContext`, -//! so it can resolve the INSERT target table's schema. For every value that is -//! bound to a pgvector `vector` column (Arrow `List(Float32)` / -//! `FixedSizeList(Float32, n)`), the string literal is replaced with an -//! `ARRAY[...]` of floats. DataFusion then inserts via its supported -//! `List -> FixedSizeList(Float32, n)` coercion (mismatched dimensions surface -//! as a cast error at runtime, matching pgvector's dimension enforcement). -//! -//! The rewrite is conservative: -//! * only `INSERT ... VALUES` statements whose target table resolves are -//! touched; -//! * only positions whose target column is a float vector/list type; -//! * only string literals that actually parse as a numeric `[...]` vector -//! (optionally written as `'[...]'::vector`). -//! -//! Everything else is left untouched, so a plain text column receiving a -//! bracket-looking string still works. - -use datafusion::arrow::datatypes::DataType; -use datafusion::common::TableReference; -use datafusion::prelude::SessionContext; -use datafusion::sql::sqlparser::ast::{ - Expr, ObjectName, SetExpr, Statement, TableObject, Value, ValueWithSpan, -}; - -use super::rules::RewriteVectorOperators; - -/// The Arrow [`DataType`] of a pgvector `vector` column, if `field_type` is one. -/// -/// `Some(Some(n))` for `vector(n)` (`FixedSizeList(Float32, n)`), -/// `Some(None)` for a dimension-less `vector` (`List(Float32)`), `None` -/// otherwise. -fn vector_dimension(field_type: &DataType) -> Option> { - match field_type { - DataType::FixedSizeList(field, n) if field.data_type() == &DataType::Float32 => { - Some(Some(*n)) - } - DataType::List(field) if field.data_type() == &DataType::Float32 => Some(None), - _ => None, - } -} - -/// Convert a table-name `ObjectName` into a [`TableReference`]. -fn object_name_to_table_reference(name: &ObjectName) -> TableReference { - let parts = name - .0 - .iter() - .filter_map(|part| part.as_ident().map(|ident| ident.value.clone())) - .collect::>(); - match parts.as_slice() { - [catalog, schema, table] => { - TableReference::full(catalog.as_str(), schema.as_str(), table.as_str()) - } - [schema, table] => TableReference::partial(schema.as_str(), table.as_str()), - [table] => TableReference::bare(table.as_str()), - _ => TableReference::bare(name.to_string().as_str()), - } -} - -/// Return the pgvector literal text carried by `expr` (a bare `'[1,2,3]'` or -/// `'[1,2,3]'::vector`), or `None` if the expression is not one of those. -fn vector_literal_text(expr: &Expr) -> Option { - match expr { - Expr::Value(ValueWithSpan { - value: Value::SingleQuotedString(text), - .. - }) => Some(text.clone()), - Expr::Cast { - expr: inner, - data_type, - .. - } if RewriteVectorOperators::is_vector_data_type(data_type) => { - if let Expr::Value(ValueWithSpan { - value: Value::SingleQuotedString(text), - .. - }) = inner.as_ref() - { - Some(text.clone()) - } else { - None - } - } - _ => None, - } -} - -/// Rewrite the pgvector string literals of an `INSERT ... VALUES` statement so -/// DataFusion can write them into `vector` columns. -/// -/// Returns `true` if any value was rewritten. When the target table cannot be -/// resolved, the columns are not vector columns, or no value parses as a vector -/// literal, the statement is left untouched and `false` is returned so the -/// caller can fall back to DataFusion's normal (erroring) handling. -pub async fn rewrite_vector_insert( - session_context: &SessionContext, - statement: &mut Statement, -) -> bool { - let Statement::Insert(insert) = statement else { - return false; - }; - let TableObject::TableName(table_name) = &insert.table else { - return false; - }; - let Some(source) = insert.source.as_mut() else { - return false; - }; - let SetExpr::Values(values) = source.body.as_mut() else { - return false; - }; - // Only single-part column names are supported by DataFusion INSERT anyway. - if insert.columns.iter().any(|col| col.0.len() != 1) { - return false; - } - - let Ok(provider) = session_context - .table_provider(object_name_to_table_reference(table_name)) - .await - else { - return false; - }; - let target_schema = provider.schema(); - - // Map each provided value position to the Arrow type of the target column. - let target_types: Vec> = if insert.columns.is_empty() { - target_schema - .fields() - .iter() - .map(|field| Some(field.data_type().clone())) - .collect() - } else { - insert - .columns - .iter() - .map(|col| { - let ident = col.0[0].as_ident()?; - target_schema - .fields() - .iter() - .find(|field| field.name().eq_ignore_ascii_case(&ident.value)) - .map(|field| field.data_type().clone()) - }) - .collect() - }; - - let mut changed = false; - for row in &mut values.rows { - for (pos, target_type) in target_types.iter().enumerate() { - let Some(target_type) = target_type else { - continue; - }; - // Only float vector/list columns accept the bracket-string form. - let Some(dim) = vector_dimension(target_type) else { - continue; - }; - let Some(expr) = row.content.get_mut(pos) else { - continue; - }; - let Some(text) = vector_literal_text(expr) else { - continue; - }; - let Some(array) = RewriteVectorOperators::vector_literal_to_array(&text) else { - continue; - }; - // When the column fixes a dimension, honor it: leave mismatched - // values for DataFusion to reject rather than inserting silently. - let count = match &array { - Expr::Array(array) => array.elem.len(), - _ => unreachable!("vector_literal_to_array returns an Array"), - }; - if dim.is_some_and(|expected| expected as usize != count) { - continue; - } - *expr = array; - changed = true; - } - } - changed -} - -#[cfg(all(test, feature = "pgvector"))] -mod tests { - use std::sync::Arc; - - use datafusion::arrow::array::{FixedSizeListArray, Float32Array, Int64Array}; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion::arrow::record_batch::RecordBatch; - use datafusion::prelude::SessionContext; - use datafusion::sql::sqlparser::dialect::PostgreSqlDialect; - use datafusion::sql::sqlparser::parser::Parser; - - use super::*; - - fn vector_field() -> Field { - Field::new( - "embedding", - DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Float32, true)), 3), - false, - ) - .with_metadata(std::collections::HashMap::from([( - "pg.vector".to_string(), - "vector".to_string(), - )])) - } - - async fn register_items(ctx: &SessionContext) { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int64, false), - vector_field(), - ])); - let id = Int64Array::from(vec![0]); // placeholder row; only the schema matters - let values = Float32Array::from(vec![0.0, 0.0, 0.0]); - let embedding = FixedSizeListArray::try_new( - Arc::new(Field::new_list_field(DataType::Float32, true)), - 3, - Arc::new(values), - None, - ) - .unwrap(); - let batch = RecordBatch::try_new(schema, vec![Arc::new(id), Arc::new(embedding)]).unwrap(); - ctx.register_batch("items", batch).unwrap(); - } - - async fn parse_insert(ctx: &SessionContext, sql: &str) -> (bool, String) { - let mut stmts = Parser::parse_sql(&PostgreSqlDialect {}, sql).unwrap(); - let mut statement = stmts.remove(0); - let changed = rewrite_vector_insert(ctx, &mut statement).await; - (changed, statement.to_string()) - } - - #[tokio::test] - async fn rewrites_vector_string_to_array_literal() { - let ctx = SessionContext::new(); - register_items(&ctx).await; - - let (changed, sql) = parse_insert( - &ctx, - "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]')", - ) - .await; - assert!(changed, "vector literal must be rewritten"); - assert!( - sql.contains("VALUES (1, ARRAY[1, 2, 3])") || sql.contains("VALUES (1, ARRAY[1,2,3])"), - "unexpected rewrite: {sql}" - ); - } - - #[tokio::test] - async fn rewrites_vector_cast_literal() { - let ctx = SessionContext::new(); - register_items(&ctx).await; - - let (changed, sql) = parse_insert( - &ctx, - "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'::vector)", - ) - .await; - assert!(changed, "vector cast literal must be rewritten"); - assert!(sql.contains("ARRAY[1, 2, 3]"), "unexpected rewrite: {sql}"); - } - - #[tokio::test] - async fn leaves_text_and_wrong_dimension_inserts_alone() { - let ctx = SessionContext::new(); - register_items(&ctx).await; - - // A non-vector target (table missing) is untouched. - let (changed, _) = parse_insert( - &ctx, - "INSERT INTO nope (id, embedding) VALUES (1, '[1,2,3]')", - ) - .await; - assert!(!changed); - - // A bracket string into a plain text column is untouched. - ctx.register_batch( - "logs", - RecordBatch::try_new( - Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int64, false), - Field::new("tag", DataType::Utf8, true), - ])), - vec![ - Arc::new(Int64Array::from(vec![0])), - Arc::new(datafusion::arrow::array::StringArray::from(vec![""])), - ], - ) - .unwrap(), - ) - .unwrap(); - let (changed, _) = - parse_insert(&ctx, "INSERT INTO logs (id, tag) VALUES (1, '[1,2,3]')").await; - assert!(!changed, "text column must not be rewritten"); - - // A vector column with a mismatched dimension is untouched (DataFusion - // will reject it as a cast error rather than silently truncating). - let (changed, _) = parse_insert( - &ctx, - "INSERT INTO items (id, embedding) VALUES (1, '[1,2]')", - ) - .await; - assert!(!changed); - } - - #[tokio::test] - async fn insert_vector_literal_stores_row() { - let ctx = SessionContext::new(); - register_items(&ctx).await; - - let sql = parse_insert( - &ctx, - "INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]')", - ) - .await - .1; - let affected = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); - let count = affected[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, 1, "one row must be inserted"); - - let batches = ctx - .sql("SELECT id FROM items WHERE id = 1") - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!(batches[0].num_rows(), 1); - - // And the nearest-neighbour query finds it at distance 0. - let batches = ctx - .sql("SELECT id FROM items ORDER BY array_distance(embedding, ARRAY[1.0, 2.0, 3.0]) LIMIT 1") - .await - .unwrap() - .collect() - .await - .unwrap(); - let id = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(id, 1); - } - - #[tokio::test] - async fn insert_multiple_rows() { - let ctx = SessionContext::new(); - register_items(&ctx).await; - - let sql = parse_insert( - &ctx, - "INSERT INTO items (id, embedding) VALUES (1, '[1,0,0]'), (2, '[0,1,0]')", - ) - .await - .1; - let affected = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); - let count = affected[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, 2); - - // Order by distance to [1,0,0]: the placeholder row 0 ([0,0,0]) sorts - // between row 1 (exact match) and row 2. - let batches = ctx - .sql("SELECT id FROM items ORDER BY array_distance(embedding, ARRAY[1.0, 0.0, 0.0])") - .await - .unwrap() - .collect() - .await - .unwrap(); - let ids: Vec = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec(); - assert_eq!(ids, vec![1, 0, 2]); - } -} diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs index 5a3e3b15..c068c685 100644 --- a/datafusion-postgres/src/handlers.rs +++ b/datafusion-postgres/src/handlers.rs @@ -129,6 +129,8 @@ impl DfSessionService { Arc::new(CursorStatementHook), Arc::new(SetShowHook), Arc::new(TransactionStatementHook), + #[cfg(feature = "pgvector")] + Arc::new(crate::pgvector::PgVectorInsertHook), ]; Self::new_with_hooks(session_context, hooks) } @@ -137,6 +139,11 @@ impl DfSessionService { session_context: Arc, query_hooks: Vec>, ) -> DfSessionService { + // Install the pgvector expression planner (distance operators) before + // any statement is planned. + #[cfg(feature = "pgvector")] + crate::pgvector::install(&session_context); + let parser = Arc::new(Parser { session_context: session_context.clone(), sql_parser: PostgresCompatibilityParser::new(), @@ -179,18 +186,6 @@ impl SimpleQueryHandler for DfSessionService { let mut results = vec![]; 'stmt: for statement in statements { - // pgvector: `INSERT ... VALUES ('[1,2,3]')` into a `vector` column - // needs the string literal rewritten to an ARRAY literal against the - // target table's schema (see datafusion_pg_catalog::sql). - #[cfg(feature = "pgvector")] - let mut statement = statement; - #[cfg(feature = "pgvector")] - datafusion_pg_catalog::sql::rewrite_vector_insert( - &self.session_context, - &mut statement, - ) - .await; - // Call query hooks with the parsed statement for hook in &self.query_hooks { if let Some(result) = hook @@ -425,14 +420,7 @@ impl QueryParser for Parser { return Ok(None); } - let mut statement = statements.remove(0); - - // pgvector: rewrite vector string literals of INSERT ... VALUES against - // the target table's schema before DataFusion plans the statement. - #[cfg(feature = "pgvector")] - datafusion_pg_catalog::sql::rewrite_vector_insert(&self.session_context, &mut statement) - .await; - + let statement = statements.remove(0); let query = statement.to_string(); let context = &self.session_context; diff --git a/datafusion-postgres/src/lib.rs b/datafusion-postgres/src/lib.rs index 1451da7d..705a82e7 100644 --- a/datafusion-postgres/src/lib.rs +++ b/datafusion-postgres/src/lib.rs @@ -2,6 +2,8 @@ pub mod auth; pub(crate) mod client; mod handlers; pub mod hooks; +#[cfg(feature = "pgvector")] +mod pgvector; mod planner; #[cfg(any(test, debug_assertions))] pub mod testing; diff --git a/datafusion-postgres/src/pgvector.rs b/datafusion-postgres/src/pgvector.rs new file mode 100644 index 00000000..210de81f --- /dev/null +++ b/datafusion-postgres/src/pgvector.rs @@ -0,0 +1,512 @@ +//! pgvector support for the `datafusion-postgres` frontend, gated behind the +//! `pgvector` Cargo feature. +//! +//! Two integration points live here, deliberately kept out of the generic +//! `pg_catalog` compatibility layer: +//! +//! * [`PgVectorExprPlanner`] -- a DataFusion [`ExprPlanner`] that rewrites the +//! pgvector distance operators (`<->`, `<#>`, `<=>`) into DataFusion's +//! built-in array distance functions (`array_distance`, `inner_product`, +//! `cosine_distance`) while planning SQL, so they work in every expression +//! position (projection, `WHERE`, `ORDER BY`, subqueries, ...). +//! * [`PgVectorInsertHook`] -- a [`QueryHook`] that rewrites +//! `INSERT ... VALUES ('[1,2,3]')` string literals into `ARRAY[...]` for +//! `vector(n)` columns, which need the target table's schema. + +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::array::UInt64Array; +use datafusion::arrow::datatypes::DataType; +use datafusion::execution::FunctionRegistry; +use datafusion::execution::session_state::SessionStateBuilder; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr}; +use datafusion::logical_expr::{Expr, LogicalPlan, ScalarUDF}; +use datafusion::prelude::SessionContext; +use datafusion::scalar::ScalarValue; +use datafusion::sql::sqlparser::ast::{ + Array, BinaryOperator, DataType as SQLDataType, Expr as SQLExpr, ObjectName, SetExpr, + Statement, TableObject, UnaryOperator, Value, ValueWithSpan, +}; +use pgwire::api::ClientInfo; +use pgwire::api::results::{Response, Tag}; +use pgwire::error::{PgWireError, PgWireResult}; + +use crate::hooks::{HookClient, QueryHook}; + +/// Install the pgvector expression planner into `session_context`. +/// +/// Looks up the distance UDFs the planner rewrites onto; if they are missing +/// (e.g. a context without the default nested functions) installation is +/// skipped so the server keeps working, just without pgvector operators. +pub fn install(session_context: &SessionContext) { + let state = session_context.state(); + let (Ok(array_distance), Ok(inner_product), Ok(cosine_distance)) = ( + state.udf("array_distance"), + state.udf("inner_product"), + state.udf("cosine_distance"), + ) else { + return; + }; + + let planner = Arc::new(PgVectorExprPlanner { + array_distance, + inner_product, + cosine_distance, + }); + + // Append to (do not replace) the existing expression planners: DataFusion + // registers its own (e.g. the nested-function array literal planner) that + // must keep working. + let mut planners = state.expr_planners().to_vec(); + planners.push(planner); + + // The distance operators are only tokenized by the Postgres SQL dialect + // (DataFusion's default Generic dialect rejects `<->`). Queries sent over + // the simple protocol are re-serialized and parsed by DataFusion, so the + // session parser must use the Postgres dialect for the planner to see them. + let mut config = session_context.copied_config(); + let _ = config + .options_mut() + .set("datafusion.sql_parser.dialect", "postgres"); + + let state_ref = session_context.state_ref(); + let existing = state_ref.read().clone(); + let new_state = SessionStateBuilder::new_from_existing(existing) + .with_config(config) + .with_expr_planners(planners) + .build(); + *state_ref.write() = new_state; +} + +/// A DataFusion [`ExprPlanner`] implementing the pgvector distance operators. +#[derive(Debug)] +pub struct PgVectorExprPlanner { + array_distance: Arc, + inner_product: Arc, + cosine_distance: Arc, +} + +impl ExprPlanner for PgVectorExprPlanner { + fn plan_binary_op( + &self, + expr: RawBinaryExpr, + _schema: &datafusion::common::DFSchema, + ) -> datafusion::error::Result> { + // (function, whether the result must be negated) + let (func, negate) = match &expr.op { + BinaryOperator::LtDashGt => (&self.array_distance, false), + BinaryOperator::Spaceship => (&self.cosine_distance, false), + // `<#>` is the *negative* inner product. + BinaryOperator::Custom(name) if name == "<#>" => (&self.inner_product, true), + _ => return Ok(PlannerResult::Original(expr)), + }; + + let left = coerce_vector_operand(expr.left); + let right = coerce_vector_operand(expr.right); + let call = + Expr::ScalarFunction(ScalarFunction::new_udf(Arc::clone(func), vec![left, right])); + + let planned = if negate { + Expr::Negative(Box::new(call)) + } else { + call + }; + Ok(PlannerResult::Planned(planned)) + } +} + +/// True when `data_type` is a float vector list (`List(Float32)` / +/// `FixedSizeList(Float32, n)`). +fn is_float_list(data_type: &DataType) -> bool { + match data_type { + DataType::FixedSizeList(field, _) | DataType::List(field) => { + field.data_type() == &DataType::Float32 + } + _ => false, + } +} + +/// The UTF-8 text of a string scalar literal, if it is one. +fn utf8_literal(scalar: &ScalarValue) -> Option<&str> { + match scalar { + ScalarValue::Utf8(Some(s)) + | ScalarValue::LargeUtf8(Some(s)) + | ScalarValue::Utf8View(Some(s)) => Some(s), + _ => None, + } +} + +/// Parse `[1,2,3]` into a `List(Float32)` scalar literal. +fn parse_vector_list(text: &str) -> Option { + let values = parse_vector_floats(text)?; + let scalars: Vec = values + .into_iter() + .map(|v| ScalarValue::Float32(Some(v))) + .collect(); + Some(ScalarValue::List(ScalarValue::new_list_nullable( + &scalars, + &DataType::Float32, + ))) +} + +/// Normalize an operator operand: turn a pgvector string literal (`'[1,2,3]'`) +/// or a cast of one (`'[1,2,3]'::vector`) into an `ARRAY[...]`-style list +/// literal. Anything else is returned unchanged. +fn coerce_vector_operand(expr: Expr) -> Expr { + // `'[1,2,3]'::vector` -- the cast target is a float list. + if let Expr::Cast(cast) = expr { + if is_float_list(cast.field.data_type()) + && let Expr::Literal(scalar, _) = cast.expr.as_ref() + && let Some(text) = utf8_literal(scalar) + && let Some(list) = parse_vector_list(text) + { + return Expr::Literal(list, None); + } + return Expr::Cast(cast); + } + + // Bare `'[1,2,3]'`. + if let Expr::Literal(scalar, metadata) = &expr + && let Some(text) = utf8_literal(scalar) + && let Some(list) = parse_vector_list(text) + { + return Expr::Literal(list, metadata.clone()); + } + + expr +} + +// --------------------------------------------------------------------------- +// INSERT ... VALUES ('[1,2,3]') support +// --------------------------------------------------------------------------- + +/// A [`QueryHook`] rewriting pgvector string literals in `INSERT ... VALUES` +/// against the target table's schema. +#[derive(Debug)] +pub struct PgVectorInsertHook; + +#[async_trait] +impl QueryHook for PgVectorInsertHook { + async fn handle_simple_query( + &self, + statement: &Statement, + session_context: &SessionContext, + client: &mut dyn HookClient, + ) -> Option> { + let mut statement = statement.clone(); + if !rewrite_insert(session_context, &mut statement).await { + return None; + } + + let query = statement.to_string(); + let timeout = crate::client::get_statement_timeout(client); + let result = async { + let df = match timeout { + Some(duration) => tokio::time::timeout(duration, session_context.sql(&query)) + .await + .map_err(|_| { + PgWireError::UserError(Box::new(pgwire::error::ErrorInfo::new( + "ERROR".to_string(), + "57014".to_string(), + "canceling statement due to statement timeout".to_string(), + ))) + })? + .map_err(|e| PgWireError::ApiError(Box::new(e)))?, + None => session_context + .sql(&query) + .await + .map_err(|e| PgWireError::ApiError(Box::new(e)))?, + }; + let batches = df + .collect() + .await + .map_err(|e| PgWireError::ApiError(Box::new(e)))?; + let rows_affected = batches + .first() + .and_then(|batch| batch.column_by_name("count")) + .and_then(|col| col.as_any().downcast_ref::()) + .map_or(0, |array| array.value(0) as usize); + Ok::<_, PgWireError>(Response::Execution( + Tag::new("INSERT").with_oid(0).with_rows(rows_affected), + )) + } + .await; + + Some(result) + } + + async fn handle_extended_parse_query( + &self, + sql: &Statement, + session_context: &SessionContext, + _client: &(dyn ClientInfo + Send + Sync), + ) -> Option> { + let mut statement = sql.clone(); + if !rewrite_insert(session_context, &mut statement).await { + return None; + } + + let state = session_context.state(); + let plan = state + .statement_to_plan(datafusion::sql::parser::Statement::Statement(Box::new( + statement, + ))) + .await + .map_err(|e| PgWireError::ApiError(Box::new(e))); + Some(plan) + } + + async fn handle_extended_query( + &self, + _statement: &Statement, + _logical_plan: &LogicalPlan, + _params: &datafusion::common::ParamValues, + _session_context: &SessionContext, + _client: &mut dyn HookClient, + ) -> Option> { + None + } +} + +/// Rewrite the pgvector string literals of an `INSERT ... VALUES` statement so +/// DataFusion can write them into `vector` columns. +/// +/// Returns `true` if any value was rewritten. When the target table cannot be +/// resolved, no column is a vector column, or no value parses as a vector +/// literal, the statement is left untouched and `false` is returned. +async fn rewrite_insert(session_context: &SessionContext, statement: &mut Statement) -> bool { + let Statement::Insert(insert) = statement else { + return false; + }; + let TableObject::TableName(table_name) = &insert.table else { + return false; + }; + let Some(source) = insert.source.as_mut() else { + return false; + }; + let SetExpr::Values(values) = source.body.as_mut() else { + return false; + }; + // DataFusion INSERT only supports single-part column names. + if insert.columns.iter().any(|col| col.0.len() != 1) { + return false; + } + + let Ok(provider) = session_context + .table_provider(object_name_to_table_reference(table_name)) + .await + else { + return false; + }; + let target_schema = provider.schema(); + + // Map each provided value position to the Arrow type of the target column. + let target_types: Vec> = if insert.columns.is_empty() { + target_schema + .fields() + .iter() + .map(|field| Some(field.data_type().clone())) + .collect() + } else { + insert + .columns + .iter() + .map(|col| { + let ident = col.0[0].as_ident()?; + target_schema + .fields() + .iter() + .find(|field| field.name().eq_ignore_ascii_case(&ident.value)) + .map(|field| field.data_type().clone()) + }) + .collect() + }; + + let mut changed = false; + for row in &mut values.rows { + for (pos, target_type) in target_types.iter().enumerate() { + let Some(target_type) = target_type else { + continue; + }; + // Only float vector/list columns accept the bracket-string form. + let Some(dim) = vector_dimension(target_type) else { + continue; + }; + let Some(expr) = row.content.get_mut(pos) else { + continue; + }; + let Some(text) = vector_literal_text(expr) else { + continue; + }; + let Some(array) = vector_literal_to_array(&text) else { + continue; + }; + // When the column fixes a dimension, honor it: leave mismatched + // values for DataFusion to reject rather than inserting silently. + let count = match &array { + SQLExpr::Array(array) => array.elem.len(), + _ => unreachable!("vector_literal_to_array returns an Array"), + }; + if dim.is_some_and(|expected| expected as usize != count) { + continue; + } + *expr = array; + changed = true; + } + } + changed +} + +/// The Arrow [`DataType`] of a pgvector `vector` column, if `field_type` is one. +/// +/// `Some(Some(n))` for `vector(n)` (`FixedSizeList(Float32, n)`), +/// `Some(None)` for a dimension-less `vector` (`List(Float32)`), `None` +/// otherwise. +fn vector_dimension(field_type: &DataType) -> Option> { + match field_type { + DataType::FixedSizeList(field, n) if field.data_type() == &DataType::Float32 => { + Some(Some(*n)) + } + DataType::List(field) if field.data_type() == &DataType::Float32 => Some(None), + _ => None, + } +} + +/// Convert a table-name `ObjectName` into a DataFusion table reference. +fn object_name_to_table_reference(name: &ObjectName) -> datafusion::common::TableReference { + let parts = name + .0 + .iter() + .filter_map(|part| part.as_ident().map(|ident| ident.value.clone())) + .collect::>(); + match parts.as_slice() { + [catalog, schema, table] => datafusion::common::TableReference::full( + catalog.as_str(), + schema.as_str(), + table.as_str(), + ), + [schema, table] => { + datafusion::common::TableReference::partial(schema.as_str(), table.as_str()) + } + [table] => datafusion::common::TableReference::bare(table.as_str()), + _ => datafusion::common::TableReference::bare(name.to_string().as_str()), + } +} + +/// Return the pgvector literal text carried by `expr` (a bare `'[1,2,3]'` or +/// `'[1,2,3]'::vector`), or `None` if the expression is not one of those. +fn vector_literal_text(expr: &SQLExpr) -> Option { + match expr { + SQLExpr::Value(ValueWithSpan { + value: Value::SingleQuotedString(text), + .. + }) => Some(text.clone()), + SQLExpr::Cast { + expr: inner, + data_type, + .. + } if is_vector_sql_type(data_type) => { + if let SQLExpr::Value(ValueWithSpan { + value: Value::SingleQuotedString(text), + .. + }) = inner.as_ref() + { + Some(text.clone()) + } else { + None + } + } + _ => None, + } +} + +/// True when `data_type` is the pgvector `vector` type name. +fn is_vector_sql_type(data_type: &SQLDataType) -> bool { + let SQLDataType::Custom(name, _) = data_type else { + return false; + }; + name.0 + .last() + .and_then(|part| part.as_ident()) + .is_some_and(|ident| ident.value.eq_ignore_ascii_case("vector")) +} + +/// Parse the floats of a pgvector literal `[1,-2.5,3]`. +fn parse_vector_floats(text: &str) -> Option> { + let text = text.trim(); + if !(text.starts_with('[') && text.ends_with(']') && text.len() >= 2) { + return None; + } + let inner = &text[1..text.len() - 1]; + if inner.trim().is_empty() { + return None; + } + let mut values = Vec::new(); + for part in inner.split(',') { + values.push(part.trim().parse::().ok()?); + } + Some(values) +} + +/// Build a SQL `ARRAY[]` literal from a pgvector literal string. +fn vector_literal_to_array(text: &str) -> Option { + let values = parse_vector_floats(text)?; + let mut elems = Vec::with_capacity(values.len()); + for value in values { + elems.push(float_literal(value)); + } + Some(SQLExpr::Array(Array { + elem: elems, + named: true, + })) +} + +/// A float SQL literal for `value` (negative values become a unary minus). +fn float_literal(value: f32) -> SQLExpr { + let rendered = value.to_string(); + let (negative, digits) = match rendered.strip_prefix('-') { + Some(digits) => (true, digits), + None => (false, rendered.as_str()), + }; + let number = SQLExpr::Value(Value::Number(digits.to_string(), false).with_empty_span()); + if negative { + SQLExpr::UnaryOp { + op: UnaryOperator::Minus, + expr: Box::new(number), + } + } else { + number + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_vector_literal_floats() { + assert_eq!( + parse_vector_floats("[1, -2.5, 3]"), + Some(vec![1.0, -2.5, 3.0]) + ); + } + + #[test] + fn rejects_non_vector_literals() { + assert_eq!(parse_vector_floats("[a,b]"), None); + assert_eq!(parse_vector_floats("1,2,3"), None); + assert_eq!(parse_vector_floats("[]"), None); + } + + #[test] + fn detects_float_list_columns() { + assert!(is_float_list(&DataType::List(std::sync::Arc::new( + datafusion::arrow::datatypes::Field::new_list_field(DataType::Float32, true) + )))); + assert!(!is_float_list(&DataType::List(std::sync::Arc::new( + datafusion::arrow::datatypes::Field::new_list_field(DataType::Int32, true) + )))); + } +} From 904227fee4a3e0227064449d4802019a19467a84 Mon Sep 17 00:00:00 2001 From: dmetasoul01 Date: Mon, 14 Sep 2026 10:42:33 +0800 Subject: [PATCH 10/10] refactor(pgvector): address second review round - 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. --- arrow-pg/src/datatypes/df.rs | 44 +++--- arrow-pg/src/encoder.rs | 35 ++++- datafusion-pg-catalog/Cargo.toml | 9 +- .../src/pg_catalog/oid_type_planner.rs | 29 +--- .../src/pg_catalog/pgvector.rs | 7 +- datafusion-pg-catalog/src/sql.rs | 14 ++ datafusion-postgres/src/client.rs | 31 +++++ datafusion-postgres/src/handlers.rs | 127 +++++------------- datafusion-postgres/src/lib.rs | 22 +-- datafusion-postgres/src/pgvector.rs | 106 ++++----------- datafusion-postgres/tests/pgvector.rs | 11 +- 11 files changed, 205 insertions(+), 230 deletions(-) diff --git a/arrow-pg/src/datatypes/df.rs b/arrow-pg/src/datatypes/df.rs index d5d6dc8b..38574258 100644 --- a/arrow-pg/src/datatypes/df.rs +++ b/arrow-pg/src/datatypes/df.rs @@ -350,24 +350,36 @@ impl VectorParam { /// The pgvector text form is `[1,2,3]`. fn from_text(raw: &[u8]) -> Result> { let text = std::str::from_utf8(raw)?; - let text = text.trim(); - if !(text.starts_with('[') && text.ends_with(']')) { - return Err("vector parameter text must look like [1,2,3]".into()); - } - let inner = text[1..text.len() - 1].trim(); - if inner.is_empty() { - return Err("vector parameter must not be empty".into()); - } - let mut values = Vec::new(); - for part in inner.split(',') { - values.push( - part.trim().parse::().map_err(|_| { - format!("invalid vector element '{}' in parameter", part.trim()) - })?, - ); + parse_vector_text(text) + .map(VectorParam) + .ok_or_else(|| "invalid pgvector text parameter (expected e.g. [1,2,3])".into()) + } +} + +/// Parse a pgvector text literal such as `[1,2,3]` into its `f32` elements. +/// +/// This is the single home for the pgvector text format: it is used by the +/// wire parameter decoder above and by the SQL-level rewrite in +/// `datafusion-postgres`. It rejects malformed input and non-finite elements +/// (`inf`/`NaN`), which cannot be re-rendered as valid SQL number literals. +pub fn parse_vector_text(text: &str) -> Option> { + let text = text.trim(); + if !(text.starts_with('[') && text.ends_with(']') && text.len() >= 2) { + return None; + } + let inner = &text[1..text.len() - 1]; + if inner.trim().is_empty() { + return None; + } + let mut values = Vec::new(); + for part in inner.split(',') { + let value: f32 = part.trim().parse().ok()?; + if !value.is_finite() { + return None; } - Ok(VectorParam(values)) + values.push(value); } + Some(values) } #[cfg(feature = "pgvector")] diff --git a/arrow-pg/src/encoder.rs b/arrow-pg/src/encoder.rs index 42857e63..a0cca636 100644 --- a/arrow-pg/src/encoder.rs +++ b/arrow-pg/src/encoder.rs @@ -305,7 +305,10 @@ fn encode_pg_vector( )) })?; let size = list.value_length() as usize; - let start = idx * size; + // Respect the array's own offset: a sliced FixedSizeListArray + // (offset > 0, e.g. a batch sliced by LIMIT) starts its rows + // further into the flattened values. + let start = list.value_offset(idx) as usize; (0..size).map(|i| values.value(start + i)).collect() } DataType::List(_) => { @@ -1041,6 +1044,36 @@ mod tests { assert_eq!(encoder.encoded, "[4.5,0,7]"); } + #[test] + fn encodes_sliced_fixed_size_list_vector() { + // A sliced FixedSizeListArray has a non-zero offset; row indexing + // must account for it instead of using `idx * size`. + let values = Float32Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]); + let array: Arc = Arc::new( + FixedSizeListArray::try_new( + Arc::new(Field::new_list_field(DataType::Float32, true)), + 3, + Arc::new(values), + None, + ) + .unwrap(), + ); + let sliced = array.slice(1, 2); // rows [4,5,6] and [7,8,9] + + let arrow_field = vector_arrow_field(true); + let pg_field = FieldInfo::new( + "embedding".to_string(), + None, + None, + crate::datatypes::pg_vector_type(), + FieldFormat::Text, + ); + + let mut encoder = TextCapture::default(); + encode_value(&mut encoder, &sliced, 0, &arrow_field, &pg_field).unwrap(); + assert_eq!(encoder.encoded, "[4,5,6]"); + } + #[derive(Default)] struct BinaryCapture { encoded: Vec, diff --git a/datafusion-pg-catalog/Cargo.toml b/datafusion-pg-catalog/Cargo.toml index 964bc7e9..35aa3646 100644 --- a/datafusion-pg-catalog/Cargo.toml +++ b/datafusion-pg-catalog/Cargo.toml @@ -27,10 +27,11 @@ postgres-types.workspace = true tokio = { version = "1.52", features = ["sync"] } [features] -# pgvector compatibility: recognize the pgvector `vector(n)` SQL type and the -# `<->` / `<#>` / `<=>` distance operators. Reuses DataFusion's built-in -# array_distance / inner_product / cosine_distance functions. -pgvector = [] +# pgvector compatibility: plan the pgvector `vector(n)` SQL type (as a +# FixedSizeList/List of Float32) and expose the `vector` row in `pg_catalog`. +# The distance operators (`<->` / `<#>` / `<=>`) and INSERT literal rewriting +# are implemented in datafusion-postgres, not here. +pgvector = ["arrow-pg/pgvector"] [dev-dependencies] env_logger = "0.11" diff --git a/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs b/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs index 1e10ebb9..ea22217c 100644 --- a/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs +++ b/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs @@ -44,11 +44,6 @@ use datafusion::sql::sqlparser::ast::{DataType as SQLDataType, ObjectNamePart}; use crate::pg_catalog::oid_field::{self, OID_ALIAS_TYPE_NAMES}; -/// Field metadata key (arrow-pg contract) marking an Arrow list field as a -/// pgvector `vector`. Must match `arrow_pg::datatypes::PG_VECTOR_KEY`. -#[cfg(feature = "pgvector")] -const PG_VECTOR_KEY: &str = "pg.vector"; - /// Recognize Postgres type names DataFusion rejects and map them to Arrow /// types/metadata at planning time. #[derive(Debug, Default)] @@ -131,19 +126,6 @@ impl PgOidTypePlanner { builtin_arrow_type(type_name) } - /// True when `sql_type` names the pgvector `vector` type (optionally - /// schema-qualified, e.g. `public.vector`). - #[cfg(feature = "pgvector")] - fn is_vector_type(sql_type: &SQLDataType) -> bool { - let SQLDataType::Custom(name, _) = sql_type else { - return false; - }; - name.0 - .last() - .and_then(|part| part.as_ident()) - .is_some_and(|ident| ident.value.eq_ignore_ascii_case("vector")) - } - /// Map the pgvector `vector` / `vector(n)` SQL type to an Arrow field. /// /// `vector(n)` is a fixed-dimension vector and maps to @@ -158,7 +140,7 @@ impl PgOidTypePlanner { /// reject the unknown type with its own error). #[cfg(feature = "pgvector")] fn vector_field(sql_type: &SQLDataType) -> Option> { - if !Self::is_vector_type(sql_type) { + if !crate::sql::is_vector_type(sql_type) { return None; } let SQLDataType::Custom(_, modifiers) = sql_type else { @@ -179,7 +161,10 @@ impl PgOidTypePlanner { }; let mut metadata = std::collections::HashMap::new(); - metadata.insert(PG_VECTOR_KEY.to_string(), "vector".to_string()); + metadata.insert( + arrow_pg::datatypes::PG_VECTOR_KEY.to_string(), + "vector".to_string(), + ); Some(Arc::new( Field::new("", arrow_type, true).with_metadata(metadata), )) @@ -398,7 +383,7 @@ mod tests { assert_eq!( field .metadata() - .get(super::PG_VECTOR_KEY) + .get(arrow_pg::datatypes::PG_VECTOR_KEY) .map(String::as_str), Some("vector") ); @@ -416,7 +401,7 @@ mod tests { assert_eq!( field .metadata() - .get(super::PG_VECTOR_KEY) + .get(arrow_pg::datatypes::PG_VECTOR_KEY) .map(String::as_str), Some("vector") ); diff --git a/datafusion-pg-catalog/src/pg_catalog/pgvector.rs b/datafusion-pg-catalog/src/pg_catalog/pgvector.rs index 3e67444a..83991113 100644 --- a/datafusion-pg-catalog/src/pg_catalog/pgvector.rs +++ b/datafusion-pg-catalog/src/pg_catalog/pgvector.rs @@ -11,6 +11,7 @@ use std::sync::Arc; +use arrow_pg::datatypes::PG_VECTOR_TYPE_OID; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::error::{DataFusionError, Result}; @@ -18,10 +19,6 @@ use datafusion::scalar::ScalarValue; use super::{ArrowTable, PgCatalogStaticTables}; -/// Fixed OID reported for pgvector `vector` columns, matching -/// `arrow_pg::datatypes::PG_VECTOR_TYPE_OID`. -const VECTOR_OID: i32 = 16385; - /// Canonical OID of the `pg_catalog` namespace (`PG_CATALOG_NAMESPACE`). const PG_CATALOG_NAMESPACE_OID: i32 = 11; @@ -45,7 +42,7 @@ pub(crate) fn with_pg_vector_support( // defaults; the introspection query only reads the ones we set). let pg_type = Arc::new(append_row(&pg_type, |field, scalar| { match field.name().as_str() { - "oid" => *scalar = ScalarValue::Int32(Some(VECTOR_OID)), + "oid" => *scalar = ScalarValue::Int32(Some(PG_VECTOR_TYPE_OID as i32)), "typname" => *scalar = ScalarValue::Utf8(Some("vector".to_string())), "typtype" => *scalar = ScalarValue::Utf8(Some("b".to_string())), "typnamespace" => *scalar = ScalarValue::Int32(Some(PG_CATALOG_NAMESPACE_OID)), diff --git a/datafusion-pg-catalog/src/sql.rs b/datafusion-pg-catalog/src/sql.rs index 42c6a178..2374a86b 100644 --- a/datafusion-pg-catalog/src/sql.rs +++ b/datafusion-pg-catalog/src/sql.rs @@ -1,3 +1,17 @@ mod parser; pub use parser::PostgresCompatibilityParser; pub mod rules; + +use datafusion::sql::sqlparser::ast::DataType as SQLDataType; + +/// True when `data_type` names the pgvector `vector` type (optionally +/// schema-qualified, e.g. `public.vector`). +pub fn is_vector_type(data_type: &SQLDataType) -> bool { + let SQLDataType::Custom(name, _) = data_type else { + return false; + }; + name.0 + .last() + .and_then(|part| part.as_ident()) + .is_some_and(|ident| ident.value.eq_ignore_ascii_case("vector")) +} diff --git a/datafusion-postgres/src/client.rs b/datafusion-postgres/src/client.rs index 7c1bab02..9784259a 100644 --- a/datafusion-postgres/src/client.rs +++ b/datafusion-postgres/src/client.rs @@ -1,9 +1,40 @@ +use datafusion::prelude::{DataFrame, SessionContext}; use pgwire::api::ClientInfo; +use pgwire::error::{PgWireError, PgWireResult}; // Metadata keys for session-level settings const METADATA_STATEMENT_TIMEOUT: &str = "statement_timeout_ms"; const METADATA_TIMEZONE: &str = "timezone"; +/// Run `query` against `session_context`, honoring the client's +/// `statement_timeout` (mapping expiry to the `57014` query-canceled error). +/// +/// Shared by the default simple-query path and by query hooks that execute a +/// rewritten statement themselves (e.g. the pgvector `INSERT` hook). +pub(crate) async fn execute_statement( + client: &C, + session_context: &SessionContext, + query: &str, +) -> PgWireResult +where + C: ClientInfo + ?Sized, +{ + let result = match get_statement_timeout(client) { + Some(duration) => tokio::time::timeout(duration, session_context.sql(query)) + .await + .map_err(|_| { + PgWireError::UserError(Box::new(pgwire::error::ErrorInfo::new( + "ERROR".to_string(), + "57014".to_string(), // query_canceled error code + "canceling statement due to statement timeout".to_string(), + ))) + })?, + None => session_context.sql(query).await, + }; + + result.map_err(|e| PgWireError::ApiError(Box::new(e))) +} + /// Get statement timeout from client metadata pub fn get_statement_timeout(client: &C) -> Option where diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs index c068c685..e4d3b8f7 100644 --- a/datafusion-postgres/src/handlers.rs +++ b/datafusion-postgres/src/handlers.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; @@ -52,18 +51,6 @@ pub struct HandlerFactory { } impl HandlerFactory { - pub fn new(session_context: Arc) -> Self { - let session_service = Arc::new(DfSessionService::new(session_context)); - let connection_manager = Arc::new(ConnectionManager::new()); - HandlerFactory { - session_service, - cancel_handler: Arc::new(DefaultCancelHandler::new(connection_manager.clone())), - startup_handler: Arc::new(SimpleStartupHandler { - connection_manager: connection_manager.clone(), - }), - } - } - pub fn new_with_hooks( session_context: Arc, query_hooks: Vec>, @@ -125,25 +112,13 @@ pub struct DfSessionService { impl DfSessionService { pub fn new(session_context: Arc) -> DfSessionService { - let hooks: Vec> = vec![ - Arc::new(CursorStatementHook), - Arc::new(SetShowHook), - Arc::new(TransactionStatementHook), - #[cfg(feature = "pgvector")] - Arc::new(crate::pgvector::PgVectorInsertHook), - ]; - Self::new_with_hooks(session_context, hooks) + Self::new_with_hooks(session_context, default_query_hooks()) } pub fn new_with_hooks( session_context: Arc, query_hooks: Vec>, ) -> DfSessionService { - // Install the pgvector expression planner (distance operators) before - // any statement is planned. - #[cfg(feature = "pgvector")] - crate::pgvector::install(&session_context); - let parser = Arc::new(Parser { session_context: session_context.clone(), sql_parser: PostgresCompatibilityParser::new(), @@ -157,6 +132,19 @@ impl DfSessionService { } } +/// The built-in query hooks (cursor, `SET`/`SHOW`, transactions). +/// +/// pgvector's `INSERT` hook and expression planner are installed by the +/// `serve*` entry points rather than here, so a `DfSessionService` built by +/// hand is not silently mutated. +pub(crate) fn default_query_hooks() -> Vec> { + vec![ + Arc::new(CursorStatementHook), + Arc::new(SetShowHook), + Arc::new(TransactionStatementHook), + ] +} + #[async_trait] impl SimpleQueryHandler for DfSessionService { async fn do_query(&self, client: &mut C, query: &str) -> PgWireResult> @@ -197,32 +185,8 @@ impl SimpleQueryHandler for DfSessionService { } } - let df_result = { - let query = statement.to_string(); - - let timeout = client::get_statement_timeout(client); - if let Some(timeout_duration) = timeout { - tokio::time::timeout(timeout_duration, self.session_context.sql(&query)) - .await - .map_err(|_| { - PgWireError::UserError(Box::new(pgwire::error::ErrorInfo::new( - "ERROR".to_string(), - "57014".to_string(), // query_canceled error code - "canceling statement due to statement timeout".to_string(), - ))) - })? - } else { - self.session_context.sql(&query).await - } - }; - - // Handle query execution errors and transaction state - let df = match df_result { - Ok(df) => df, - Err(e) => { - return Err(PgWireError::ApiError(Box::new(e))); - } - }; + let query = statement.to_string(); + let df = client::execute_statement(client, &self.session_context, &query).await?; if matches!(statement, sqlparser::ast::Statement::Insert(_)) { let resp = map_rows_affected_for_insert(&df).await?; @@ -277,10 +241,19 @@ impl ExtendedQueryHandler for DfSessionService { .map_err(|e| PgWireError::ApiError(Box::new(e)))?; let wire_types = parameter_wire_types(plan)?; let wire_type_refs: Vec> = wire_types.iter().map(Some).collect(); + let inferenced: Vec> = + planner::ordered_parameter_entries(¶m_types) + .into_iter() + .map(|(_, datatype)| datatype) + .collect(); + let inferenced_refs: Vec> = inferenced + .iter() + .map(|datatype| datatype.as_ref()) + .collect(); let param_values: ParamValues = df::deserialize_parameters_with_server_types( portal, - &ordered_param_types(¶m_types), + &inferenced_refs, &wire_type_refs, )?; @@ -305,10 +278,19 @@ impl ExtendedQueryHandler for DfSessionService { .map_err(|e| PgWireError::ApiError(Box::new(e)))?; let wire_types = parameter_wire_types(plan)?; let wire_type_refs: Vec> = wire_types.iter().map(Some).collect(); + let inferenced: Vec> = + planner::ordered_parameter_entries(¶m_types) + .into_iter() + .map(|(_, datatype)| datatype) + .collect(); + let inferenced_refs: Vec> = inferenced + .iter() + .map(|datatype| datatype.as_ref()) + .collect(); let param_values = df::deserialize_parameters_with_server_types( portal, - &ordered_param_types(¶m_types), + &inferenced_refs, &wire_type_refs, )?; @@ -368,7 +350,7 @@ impl ExtendedQueryHandler for DfSessionService { } } -async fn map_rows_affected_for_insert(df: &DataFrame) -> PgWireResult { +pub(crate) async fn map_rows_affected_for_insert(df: &DataFrame) -> PgWireResult { // For INSERT queries, we need to execute the query to get the row count // and return an Execution response with the proper tag let result = df @@ -500,18 +482,6 @@ fn parameter_wire_types(plan: &LogicalPlan) -> PgWireResult> { Ok(types) } -fn ordered_param_types(types: &HashMap>) -> Vec> { - // Datafusion stores the parameters as a map. In our case, the keys will be - // `$1`, `$2` etc. The values will be the parameter types. - let mut types = types.iter().collect::>(); - types.sort_by_key(|(key, _)| { - key.trim_start_matches('$') - .parse::() - .unwrap_or(u32::MAX) - }); - types.into_iter().map(|pt| pt.1.as_ref()).collect() -} - #[cfg(test)] mod tests { use datafusion::prelude::SessionContext; @@ -559,29 +529,6 @@ mod tests { } } - #[test] - fn test_ordered_param_types_sorts_placeholders_numerically() { - let params = HashMap::from([ - ("$1".to_string(), Some(DataType::Boolean)), - ("$2".to_string(), Some(DataType::Int64)), - ("$10".to_string(), Some(DataType::Utf8)), - ]); - - let ordered = ordered_param_types(¶ms) - .into_iter() - .map(|ty| ty.cloned()) - .collect::>(); - - assert_eq!( - ordered, - vec![ - Some(DataType::Boolean), - Some(DataType::Int64), - Some(DataType::Utf8) - ] - ); - } - #[tokio::test] async fn test_query_hooks() { let hook = TestHook; diff --git a/datafusion-postgres/src/lib.rs b/datafusion-postgres/src/lib.rs index 705a82e7..1c7e6e5a 100644 --- a/datafusion-postgres/src/lib.rs +++ b/datafusion-postgres/src/lib.rs @@ -3,7 +3,7 @@ pub(crate) mod client; mod handlers; pub mod hooks; #[cfg(feature = "pgvector")] -mod pgvector; +pub mod pgvector; mod planner; #[cfg(any(test, debug_assertions))] pub mod testing; @@ -90,17 +90,15 @@ pub async fn serve( session_context: Arc, opts: &ServerOptions, ) -> Result<(), std::io::Error> { - #[cfg(feature = "postgis")] - geodatafusion::register(&session_context); - - // Create the handler factory with authentication - let factory = Arc::new(HandlerFactory::new(session_context)); - - serve_with_handlers(factory, opts).await + serve_with_hooks(session_context, opts, handlers::default_query_hooks()).await } /// Serve the Datafusion `SessionContext` with Postgres protocol, using custom /// query processing hooks. +/// +/// The optional pgvector support (the distance-operator expression planner and +/// the `INSERT` literal hook) is installed here, where the session context can +/// be modified, and appended to `hooks`. pub async fn serve_with_hooks( session_context: Arc, opts: &ServerOptions, @@ -109,6 +107,14 @@ pub async fn serve_with_hooks( #[cfg(feature = "postgis")] geodatafusion::register(&session_context); + #[cfg(feature = "pgvector")] + let hooks = { + let mut hooks = hooks; + crate::pgvector::install(&session_context).map_err(std::io::Error::other)?; + hooks.push(Arc::new(crate::pgvector::PgVectorInsertHook)); + hooks + }; + // Create the handler factory with authentication let factory = Arc::new(HandlerFactory::new_with_hooks(session_context, hooks)); diff --git a/datafusion-postgres/src/pgvector.rs b/datafusion-postgres/src/pgvector.rs index 210de81f..99d1a919 100644 --- a/datafusion-postgres/src/pgvector.rs +++ b/datafusion-postgres/src/pgvector.rs @@ -16,7 +16,6 @@ use std::sync::Arc; use async_trait::async_trait; -use datafusion::arrow::array::UInt64Array; use datafusion::arrow::datatypes::DataType; use datafusion::execution::FunctionRegistry; use datafusion::execution::session_state::SessionStateBuilder; @@ -26,29 +25,27 @@ use datafusion::logical_expr::{Expr, LogicalPlan, ScalarUDF}; use datafusion::prelude::SessionContext; use datafusion::scalar::ScalarValue; use datafusion::sql::sqlparser::ast::{ - Array, BinaryOperator, DataType as SQLDataType, Expr as SQLExpr, ObjectName, SetExpr, - Statement, TableObject, UnaryOperator, Value, ValueWithSpan, + Array, BinaryOperator, Expr as SQLExpr, ObjectName, SetExpr, Statement, TableObject, + UnaryOperator, Value, ValueWithSpan, }; use pgwire::api::ClientInfo; -use pgwire::api::results::{Response, Tag}; +use pgwire::api::results::Response; use pgwire::error::{PgWireError, PgWireResult}; +use crate::arrow_pg::datatypes::df::parse_vector_text; use crate::hooks::{HookClient, QueryHook}; /// Install the pgvector expression planner into `session_context`. /// -/// Looks up the distance UDFs the planner rewrites onto; if they are missing -/// (e.g. a context without the default nested functions) installation is -/// skipped so the server keeps working, just without pgvector operators. -pub fn install(session_context: &SessionContext) { +/// Looks up the distance UDFs the planner rewrites onto (they come from +/// DataFusion's default nested functions) and appends the planner to the +/// session's expression planners, and switches the SQL parser dialect to +/// Postgres. Errors are returned rather than ignored. +pub fn install(session_context: &SessionContext) -> datafusion::error::Result<()> { let state = session_context.state(); - let (Ok(array_distance), Ok(inner_product), Ok(cosine_distance)) = ( - state.udf("array_distance"), - state.udf("inner_product"), - state.udf("cosine_distance"), - ) else { - return; - }; + let array_distance = state.udf("array_distance")?; + let inner_product = state.udf("inner_product")?; + let cosine_distance = state.udf("cosine_distance")?; let planner = Arc::new(PgVectorExprPlanner { array_distance, @@ -67,9 +64,9 @@ pub fn install(session_context: &SessionContext) { // the simple protocol are re-serialized and parsed by DataFusion, so the // session parser must use the Postgres dialect for the planner to see them. let mut config = session_context.copied_config(); - let _ = config + config .options_mut() - .set("datafusion.sql_parser.dialect", "postgres"); + .set("datafusion.sql_parser.dialect", "postgres")?; let state_ref = session_context.state_ref(); let existing = state_ref.read().clone(); @@ -78,6 +75,8 @@ pub fn install(session_context: &SessionContext) { .with_expr_planners(planners) .build(); *state_ref.write() = new_state; + + Ok(()) } /// A DataFusion [`ExprPlanner`] implementing the pgvector distance operators. @@ -140,7 +139,7 @@ fn utf8_literal(scalar: &ScalarValue) -> Option<&str> { /// Parse `[1,2,3]` into a `List(Float32)` scalar literal. fn parse_vector_list(text: &str) -> Option { - let values = parse_vector_floats(text)?; + let values = parse_vector_text(text)?; let scalars: Vec = values .into_iter() .map(|v| ScalarValue::Float32(Some(v))) @@ -201,36 +200,9 @@ impl QueryHook for PgVectorInsertHook { } let query = statement.to_string(); - let timeout = crate::client::get_statement_timeout(client); let result = async { - let df = match timeout { - Some(duration) => tokio::time::timeout(duration, session_context.sql(&query)) - .await - .map_err(|_| { - PgWireError::UserError(Box::new(pgwire::error::ErrorInfo::new( - "ERROR".to_string(), - "57014".to_string(), - "canceling statement due to statement timeout".to_string(), - ))) - })? - .map_err(|e| PgWireError::ApiError(Box::new(e)))?, - None => session_context - .sql(&query) - .await - .map_err(|e| PgWireError::ApiError(Box::new(e)))?, - }; - let batches = df - .collect() - .await - .map_err(|e| PgWireError::ApiError(Box::new(e)))?; - let rows_affected = batches - .first() - .and_then(|batch| batch.column_by_name("count")) - .and_then(|col| col.as_any().downcast_ref::()) - .map_or(0, |array| array.value(0) as usize); - Ok::<_, PgWireError>(Response::Execution( - Tag::new("INSERT").with_oid(0).with_rows(rows_affected), - )) + let df = crate::client::execute_statement(client, session_context, &query).await?; + crate::handlers::map_rows_affected_for_insert(&df).await } .await; @@ -407,7 +379,7 @@ fn vector_literal_text(expr: &SQLExpr) -> Option { expr: inner, data_type, .. - } if is_vector_sql_type(data_type) => { + } if datafusion_pg_catalog::sql::is_vector_type(data_type) => { if let SQLExpr::Value(ValueWithSpan { value: Value::SingleQuotedString(text), .. @@ -422,37 +394,9 @@ fn vector_literal_text(expr: &SQLExpr) -> Option { } } -/// True when `data_type` is the pgvector `vector` type name. -fn is_vector_sql_type(data_type: &SQLDataType) -> bool { - let SQLDataType::Custom(name, _) = data_type else { - return false; - }; - name.0 - .last() - .and_then(|part| part.as_ident()) - .is_some_and(|ident| ident.value.eq_ignore_ascii_case("vector")) -} - -/// Parse the floats of a pgvector literal `[1,-2.5,3]`. -fn parse_vector_floats(text: &str) -> Option> { - let text = text.trim(); - if !(text.starts_with('[') && text.ends_with(']') && text.len() >= 2) { - return None; - } - let inner = &text[1..text.len() - 1]; - if inner.trim().is_empty() { - return None; - } - let mut values = Vec::new(); - for part in inner.split(',') { - values.push(part.trim().parse::().ok()?); - } - Some(values) -} - /// Build a SQL `ARRAY[]` literal from a pgvector literal string. fn vector_literal_to_array(text: &str) -> Option { - let values = parse_vector_floats(text)?; + let values = parse_vector_text(text)?; let mut elems = Vec::with_capacity(values.len()); for value in values { elems.push(float_literal(value)); @@ -488,16 +432,16 @@ mod tests { #[test] fn parses_vector_literal_floats() { assert_eq!( - parse_vector_floats("[1, -2.5, 3]"), + parse_vector_text("[1, -2.5, 3]"), Some(vec![1.0, -2.5, 3.0]) ); } #[test] fn rejects_non_vector_literals() { - assert_eq!(parse_vector_floats("[a,b]"), None); - assert_eq!(parse_vector_floats("1,2,3"), None); - assert_eq!(parse_vector_floats("[]"), None); + assert_eq!(parse_vector_text("[a,b]"), None); + assert_eq!(parse_vector_text("1,2,3"), None); + assert_eq!(parse_vector_text("[]"), None); } #[test] diff --git a/datafusion-postgres/tests/pgvector.rs b/datafusion-postgres/tests/pgvector.rs index be53ca3b..e37e1d9c 100644 --- a/datafusion-postgres/tests/pgvector.rs +++ b/datafusion-postgres/tests/pgvector.rs @@ -25,8 +25,8 @@ use datafusion_postgres::auth::AuthManager; use datafusion_postgres::testing::MockClient; use datafusion_postgres::{ServerOptions, serve}; -/// pgvector `vector` type OID, matching `arrow_pg::datatypes::PG_VECTOR_TYPE_OID`. -const VECTOR_OID: u32 = 16385; +/// pgvector `vector` type OID (the canonical constant from `arrow-pg`). +use datafusion_postgres::arrow_pg::datatypes::PG_VECTOR_TYPE_OID as VECTOR_OID; /// Register `items(id bigint, embedding vector(3))` as an empty table whose /// `embedding` field carries the `pg.vector` metadata. @@ -75,7 +75,12 @@ async fn service() -> (SessionContext, DfSessionService) { .expect("failed to setup pg_catalog"); register_items(&session_context); - let service = DfSessionService::new(Arc::new(session_context.clone())); + // The pgvector planner/hook are installed by the `serve*` entry points; + // when building a `DfSessionService` directly they must be wired manually. + datafusion_postgres::pgvector::install(&session_context).expect("install pgvector planner"); + let hooks: Vec> = + vec![Arc::new(datafusion_postgres::pgvector::PgVectorInsertHook)]; + let service = DfSessionService::new_with_hooks(Arc::new(session_context.clone()), hooks); (session_context, service) }