diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c29878b3..32151773 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 @@ -132,6 +135,24 @@ jobs: - uses: actions/checkout@v6 - uses: insightsengineering/disk-space-reclaimer@v1 - uses: cachix/install-nix-action@v31 + - 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 --skip-pgvector + + # pgvector has its own dedicated integration job so it can run (and fail) + # 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 + timeout-minutes: 15 + needs: [test] + steps: + - uses: actions/checkout@v6 + - uses: insightsengineering/disk-space-reclaimer@v1 + - uses: cachix/install-nix-action@v31 - name: Initialize podman run: nix develop --command podman info >/dev/null 2>&1 || true - run: | @@ -154,7 +175,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/.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..430d2821 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]] @@ -2491,7 +2491,9 @@ dependencies = [ "geodatafusion", "getset", "log", + "pgvector", "pgwire", + "postgres", "postgres-types", "rust_decimal", "rustls-pemfile", @@ -3071,7 +3073,7 @@ checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", ] [[package]] @@ -3528,11 +3530,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 +3669,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 +3784,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 +3858,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" @@ -4011,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" @@ -4051,7 +4091,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 +4113,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" @@ -4149,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" @@ -5072,6 +5145,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 +5393,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 +5420,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 +5445,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 +5465,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 +5478,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 +5521,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 +5539,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..93f659b8 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,9 @@ 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)); + 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 +441,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..38574258 100644 --- a/arrow-pg/src/datatypes/df.rs +++ b/arrow-pg/src/datatypes/df.rs @@ -272,6 +272,204 @@ 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)] +struct VectorParam(Vec); + +#[cfg(feature = "pgvector")] +impl VectorParam { + /// 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() < 4 { + return Err("vector parameter binary payload too short".into()); + } + 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()); + } + 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 = 4 + 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)?; + 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; + } + values.push(value); + } + Some(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,19 +477,57 @@ 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 { + // 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 + { + Ok(ty.clone()) } else if let Some(infer_type) = inferenced_type { into_pg_type(infer_type) } else { @@ -303,7 +539,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 +552,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 +1093,21 @@ where // Store MAC addresses as strings for now deserialized_params.push(ScalarValue::Utf8(value)); } + // 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)?; + match inferenced_type { + 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))); + } + } + } // 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..a0cca636 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; @@ -8,7 +11,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 +162,180 @@ 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) + } +} + +/// 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")] +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. +/// +/// 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, + 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; + // 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(_) => { + 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) + } + + encoder.encode_field(&Some(PgVectorValue(row_values(arr, idx)?)), pg_field) +} + fn get_date32_value(arr: &Arc, idx: usize) -> Option { if arr.is_null(idx) { return None; @@ -283,6 +460,38 @@ 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]` / 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) + || pg_field.datatype().oid() == crate::datatypes::PG_VECTOR_TYPE_OID + { + 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 +694,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 +959,247 @@ 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_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, + } + + 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. + 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..35aa3646 100644 --- a/datafusion-pg-catalog/Cargo.toml +++ b/datafusion-pg-catalog/Cargo.toml @@ -26,6 +26,13 @@ log = "0.4" postgres-types.workspace = true tokio = { version = "1.52", features = ["sync"] } +[features] +# 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" 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..c90db8e0 100644 --- a/datafusion-pg-catalog/src/pg_catalog.rs +++ b/datafusion-pg-catalog/src/pg_catalog.rs @@ -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"; @@ -652,7 +654,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,7 +1063,16 @@ 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 = pgvector::with_pg_vector_support(tables)?; + + Ok(tables) } /// Create table from dumped arrow data 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..ea22217c 100644 --- a/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs +++ b/datafusion-pg-catalog/src/pg_catalog/oid_type_planner.rs @@ -125,6 +125,50 @@ impl PgOidTypePlanner { } builtin_arrow_type(type_name) } + + /// 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 !crate::sql::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( + arrow_pg::datatypes::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 +203,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 +369,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(arrow_pg::datatypes::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(arrow_pg::datatypes::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/pg_catalog/pgvector.rs b/datafusion-pg-catalog/src/pg_catalog/pgvector.rs new file mode 100644 index 00000000..83991113 --- /dev/null +++ b/datafusion-pg-catalog/src/pg_catalog/pgvector.rs @@ -0,0 +1,144 @@ +//! 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 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}; +use datafusion::scalar::ScalarValue; + +use super::{ArrowTable, PgCatalogStaticTables}; + +/// 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(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)), + _ => {} + } + })?); + + 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 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-pg-catalog/src/sql/parser.rs b/datafusion-pg-catalog/src/sql/parser.rs index d8e9a519..59966100 100644 --- a/datafusion-pg-catalog/src/sql/parser.rs +++ b/datafusion-pg-catalog/src/sql/parser.rs @@ -309,28 +309,30 @@ impl PostgresCompatibilityParser { )); } + 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. + 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), + ]; + 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-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..7cb3fb6a 100644 --- a/datafusion-postgres/Cargo.toml +++ b/datafusion-postgres/Cargo.toml @@ -34,7 +34,13 @@ rustls-pki-types = "1.14" [dev-dependencies] env_logger = "0.11" +pgvector = { version = "0.4", features = ["postgres"] } +postgres = "0.19" [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/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 c66236cb..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,12 +112,7 @@ 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), - ]; - Self::new_with_hooks(session_context, hooks) + Self::new_with_hooks(session_context, default_query_hooks()) } pub fn new_with_hooks( @@ -150,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> @@ -190,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?; @@ -268,9 +239,23 @@ 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 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(portal, &ordered_param_types(¶m_types))?; + let param_values: ParamValues = df::deserialize_parameters_with_server_types( + portal, + &inferenced_refs, + &wire_type_refs, + )?; for hook in &self.query_hooks { if let Some(result) = hook @@ -291,9 +276,23 @@ 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 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(portal, &ordered_param_types(¶m_types))?; + let param_values = df::deserialize_parameters_with_server_types( + portal, + &inferenced_refs, + &wire_type_refs, + )?; let plan = plan .clone() @@ -351,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 @@ -427,20 +426,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,16 +456,30 @@ impl QueryParser for Parser { } } -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() +/// 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) } #[cfg(test)] @@ -529,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 1451da7d..1c7e6e5a 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")] +pub mod pgvector; mod planner; #[cfg(any(test, debug_assertions))] pub mod testing; @@ -88,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, @@ -107,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 new file mode 100644 index 00000000..99d1a919 --- /dev/null +++ b/datafusion-postgres/src/pgvector.rs @@ -0,0 +1,456 @@ +//! 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::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, Expr as SQLExpr, ObjectName, SetExpr, Statement, TableObject, + UnaryOperator, Value, ValueWithSpan, +}; +use pgwire::api::ClientInfo; +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 (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 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, + 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(); + 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; + + Ok(()) +} + +/// 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_text(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 result = async { + let df = crate::client::execute_statement(client, session_context, &query).await?; + crate::handlers::map_rows_affected_for_insert(&df).await + } + .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 datafusion_pg_catalog::sql::is_vector_type(data_type) => { + if let SQLExpr::Value(ValueWithSpan { + value: Value::SingleQuotedString(text), + .. + }) = inner.as_ref() + { + Some(text.clone()) + } else { + None + } + } + _ => None, + } +} + +/// Build a SQL `ARRAY[]` literal from a pgvector literal string. +fn vector_literal_to_array(text: &str) -> Option { + let values = parse_vector_text(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_text("[1, -2.5, 3]"), + Some(vec![1.0, -2.5, 3.0]) + ); + } + + #[test] + fn rejects_non_vector_literals() { + assert_eq!(parse_vector_text("[a,b]"), None); + assert_eq!(parse_vector_text("1,2,3"), None); + assert_eq!(parse_vector_text("[]"), 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) + )))); + } +} 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..e37e1d9c --- /dev/null +++ b/datafusion-postgres/tests/pgvector.rs @@ -0,0 +1,439 @@ +//! 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 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::NoTls; +use tokio::sync::oneshot; + +use datafusion_postgres::DfSessionService; +use datafusion_postgres::auth::AuthManager; +use datafusion_postgres::testing::MockClient; +use datafusion_postgres::{ServerOptions, serve}; + +/// 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. +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); + + // 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) +} + +#[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() +} + +/// 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, + "datafusion", + Arc::new(AuthManager::default()), + ) + .expect("failed to setup pg_catalog"); + register_items(&session_context); + + let (port, stop_tx) = spawn_server(session_context); + let mut client = connect(port); + + // 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]')", + &[], + ) + .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", + &[], + ) + .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. + let messages = client + .simple_query("SELECT embedding FROM items ORDER BY id LIMIT 1") + .expect("simple query of the vector column should succeed"); + let mut cells = Vec::new(); + for message in messages { + if let postgres::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`. The official + // pgvector::Vector decoder then reads the binary result. + let statement = client + .prepare("SELECT embedding FROM items LIMIT 1") + .expect("prepare of a vector result column should succeed"); + assert_eq!(statement.columns().len(), 1); + let vector_type = &statement.columns()[0].type_(); + assert_eq!(vector_type.name(), "vector"); + assert_eq!(vector_type.oid(), VECTOR_OID); + + let rows = client + .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])); + + 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 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, + "datafusion", + Arc::new(AuthManager::default()), + ) + .expect("failed to setup pg_catalog"); + register_items(&session_context); + + 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)") + .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 an official pgvector::Vector, then execute. + let id: i64 = 42; + let vector = pgvector::Vector::from(vec![1.0, 2.0, 3.0]); + let affected = client + .execute(&statement, &[&id, &vector]) + .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", + &[], + ) + .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") + .expect("simple query of the vector column should succeed"); + let mut cells = Vec::new(); + for message in messages { + if let postgres::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" + ); + + let _ = stop_tx.send(()); +} + +/// 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, + "datafusion", + Arc::new(AuthManager::default()), + ) + .expect("failed to setup pg_catalog"); + + 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))") + .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]')", + &[], + ) + .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)") + .expect("prepare INSERT should succeed"); + assert_eq!(statement.params()[1].oid(), VECTOR_OID); + let id: i32 = 3; + let vector = pgvector::Vector::from(vec![7.0, 8.0, 9.0]); + client + .execute(&statement, &[&id, &vector]) + .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", + &[], + ) + .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"); + + 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(()); +} diff --git a/tests-integration/test.sh b/tests-integration/test.sh index 2ef76b23..c0712cfb 100755 --- a/tests-integration/test.sh +++ b/tests-integration/test.sh @@ -2,20 +2,23 @@ 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 # 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 @@ -48,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 @@ -248,6 +258,36 @@ else echo "โญ๏ธ Skipped (--skip-postgis)" fi +# Test 7: pgvector +echo "" +echo "๐Ÿงช Test 7: pgvector Support" +echo "---------------------------" +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 + + 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 +else + echo "โญ๏ธ Skipped (--skip-pgvector)" +fi +sleep 3 + echo "" echo "๐ŸŽ‰ All enhanced integration tests passed!" echo "==========================================" @@ -261,6 +301,9 @@ echo " โœ… Array types and complex data type support" echo " โœ… Improved pg_catalog system tables" echo " โœ… PostgreSQL function compatibility" echo " โœ… SSL/TLS encryption support" +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 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()