Defining the Vector Structure

Our goal is to enable SQL statements like:

CREATE TABLE embeddings (embedding vector(3));

We start by defining a Rust struct that wraps a vector of floating-point numbers:

#[derive(Debug, serde::Deserialize, serde::Serialize)]
#[serde(transparent)]
pub struct VectorData {
    components: Vec<f64>,
}
</f64>

Creating a base type in PostgreSQL requires implementing four functions: input, output, type modifier input, and type modifier output. Since our type accepts a dimension parameter (e.g., vector(3)), we need the modifier functions to handle this constraint.

Shell Type Declaration

Before implementing the required functions, we must declare a shell type. This placeholder allows functions to reference the type before its full definition:

extension_sql!(
    r#"CREATE TYPE vector;"#,
    name = "vector_shell",
    bootstrap
);

Input Function Implementation

The input function deserializes text representation into the internal format. We except three parameters: the input string, the type OID, and the type modifier:

#[pg_extern(immutable, strict, parallel_safe, requires = ["vector_shell"])]
fn vector_input(
    text: &CStr,
    _oid: pg_sys::Oid,
    typmod: i32,
) -> VectorData {
    let parsed = match serde_json::from_str::<Vec<f64>>(
        text.to_str().expect("Input must be valid UTF-8"),
    ) {
        Ok(data) => data,
        Err(err) => pgrx::error!("Deserialization failed: {}", err),
    };

    let dim = match u16::try_from(parsed.len()) {
        Ok(d) => d,
        Err(_) => pgrx::error!("Dimension {} exceeds maximum allowed", parsed.len()),
    };

    if typmod != -1 {
        let target_dim = match u16::try_from(typmod) {
            Ok(d) => d,
            Err(_) => panic!("Invalid type modifier: {}", typmod),
        };

        if dim != target_dim {
            pgrx::error!("Dimension mismatch: expected {}, got {}", target_dim, dim);
        }
    }

    VectorData { components: parsed }
}

The pg_extern macro generates the corresponding SQL function definition. The requires parameter ensures proper ordering during SQL generation.

Output Function Implementation

The output function serializes the internal representation back to text:

#[pg_extern(immutable, strict, parallel_safe, requires = ["vector_shell"])]
fn vector_output(data: VectorData) -> CString {
    let serialized = serde_json::to_string(&data).unwrap();
    CString::new(serialized).expect("Serialization should not contain NUL bytes")
}

Type Modifier Functions

Type modifiers allow parameterized types like vector(3). The input function validates and encodes the modifier:

#[pg_extern(immutable, strict, parallel_safe, requires = ["vector_shell"])]
fn vector_typmod_in(args: pgrx::datum::Array<&CStr>) -> i32 {
    if args.len() != 1 {
        pgrx::error!("Expected exactly one type modifier");
    }

    let modifier = args
        .get(0)
        .expect("Element exists")
        .expect("Modifier cannot be null");

    let dim: u16 = modifier
        .to_str()
        .expect("Modifier must be UTF-8")
        .parse()
        .unwrap_or_else(|_| {
            pgrx::error!("Invalid dimension: must be between 1 and 65535")
        });

    dim as i32
}

The output function formats the modifier for display:

#[pg_extern(immutable, strict, parallel_safe, requires = ["vector_shell"])]
fn vector_typmod_out(typmod: i32) -> CString {
    CString::new(format!("({})", typmod)).expect("No NUL bytes in output")
}

Creating the Concrete Type

With all functions defined, we create the actual type:

extension_sql!(
    r#"
CREATE TYPE vector (
    INPUT = vector_input,
    OUTPUT = vector_output,
    TYPMOD_IN = vector_typmod_in,
    TYPMOD_OUT = vector_typmod_out,
    STORAGE = external
);
"#,
    name = "vector_concrete",
    creates = [Type(VectorData)],
    requires = [
        "vector_shell",
        vector_input,
        vector_output,
        vector_typmod_in,
        vector_typmod_out
    ]
);

The STORAGE = external option allows large vectors to be moved to TOAST tables when they cannot fit efficiently in main storage pages.

Implementing Required Traits

pgrx requires several traits for custom types. SqlTranslatable tells pgrx how to reference the type in generated SQL:

unsafe impl SqlTranslatable for VectorData {
    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
        Ok(SqlMapping::As("vector".into()))
    }

    fn return_sql() -> Result<Returns, ReturnsError> {
        Ok(Returns::One(SqlMapping::As("vector".into())))
    }
}

The FromDatum and IntoDatum traits enable conversion between PostgreSQL's enternal Datum format and our Rust type:

impl FromDatum for VectorData {
    unsafe fn from_polymorphic_datum(
        datum: pg_sys::Datum,
        is_null: bool,
        typoid: pg_sys::Oid,
    ) -> Option<Self> {
        let data = <Vec<f64> as FromDatum>::from_polymorphic_datum(datum, is_null, typoid)?;
        Some(Self { components: data })
    }
}

impl IntoDatum for VectorData {
    fn into_datum(self) -> Option<pg_sys::Datum> {
        self.components.into_datum()
    }

    fn type_oid() -> pg_sys::Oid {
        rust_regtypein::<Self>()
    }
}

The ArgAbi and BoxRet traits handle argument passing and return values in function calls:

unsafe impl<'fcx> ArgAbi<'fcx> for VectorData
where
    Self: 'fcx,
{
    unsafe fn unbox_arg_unchecked(arg: ::pgrx::callconv::Arg<'_, 'fcx>) -> Self {
        arg.unbox_arg_using_from_datum().expect("Expected non-null value")
    }
}

unsafe impl BoxRet for VectorData {
    unsafe fn box_into<'fcx>(
        self,
        fcinfo: &mut pgrx::callconv::FcInfo<'fcx>,
    ) -> pgrx::datum::Datum<'fcx> {
        match self.into_datum() {
            Some(datum) => fcinfo.return_raw_datum(datum),
            None => fcinfo.return_null(),
        }
    }
}

Handling the Type Modifier Quirk

PostgreSQL may pass -1 as the type modifier to the input function even when a modifier is defined. To ensure dimension validation occurs correctly, we implement a cast functon:

#[pg_extern(immutable, strict, parallel_safe, requires = ["vector_concrete"])]
fn cast_vector_to_vector(vec: VectorData, typmod: i32, _explicit: bool) -> VectorData {
    let target_dim = u16::try_from(typmod).expect("Invalid type modifier") as usize;
    if vec.components.len() != target_dim {
        pgrx::error!(
            "Dimension mismatch: expected {}, found {}",
            typmod,
            vec.components.len()
        );
    }
    vec
}

extension_sql!(
    r#"
CREATE CAST (vector AS vector)
WITH FUNCTION cast_vector_to_vector(vector, integer, boolean);
"#,
    name = "vector_cast",
    requires = ["vector_concrete", cast_vector_to_vector]
);

With all components in place, the extension can be compiled and installed:

$ cargo pgrx run
psql (17.2)

pg_extension=# CREATE EXTENSION vector_extension;
CREATE EXTENSION

pg_extension=# CREATE TABLE items (embedding vector(3));
CREATE TABLE

pg_extension=# INSERT INTO items VALUES ('[1.0, 2.0, 3.0]');
INSERT 0 1

pg_extension=# SELECT * FROM items;
    embedding
----------------
 [1.0,2.0,3.0]
(1 row)