From 63c4a3891168cdedefdf8a39d4b931c4f7ec72e5 Mon Sep 17 00:00:00 2001 From: Andrew Vasilyev Date: Wed, 26 Jun 2024 19:48:59 +0000 Subject: [PATCH] feat: add support for hash indexes (requires aes now) --- .cargo/config.toml | 5 +- Cargo.toml | 6 ++- src/lib.rs | 114 +++++++++++++++++++++++++++++++++++++++++++++ src/typeid.rs | 12 ++++- 4 files changed, 133 insertions(+), 4 deletions(-) diff --git a/.cargo/config.toml b/.cargo/config.toml index 13c456b..b6bdb88 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,3 +1,6 @@ [target.'cfg(target_os="macos")'] # Postgres symbols won't be available until runtime -rustflags = ["-Clink-arg=-Wl,-undefined,dynamic_lookup"] +rustflags = ["-Clink-arg=-Wl,-undefined,dynamic_lookup", "-C target-feature=+aes,+sse2"] + +[target.'cfg(target_arch = "x86_64")'] +rustflags = ["-C", "target-feature=+aes"] \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index fc17c1b..5db76fc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,9 @@ name = "typeid" version = "0.0.0" edition = "2021" +[build] +rustflags = ["-C", "target-feature=+aes"] + [lib] crate-type = ["cdylib", "lib"] @@ -17,6 +20,7 @@ pg16 = ["pgrx/pg16", "pgrx-tests/pg16" ] pg_test = [] [dependencies] +gxhash = { version = "3.4.1" } pgrx = "=0.11.4" serde = "1.0.203" thiserror = "1.0.61" @@ -39,4 +43,4 @@ codegen-units = 1 [[test]] name = "spec" path = "tests/spec.rs" -harness = false \ No newline at end of file +harness = false diff --git a/src/lib.rs b/src/lib.rs index 1eea36b..6eb83e8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,6 +8,8 @@ use uuid::Uuid; use pgrx::prelude::*; +use std::hash::{Hash, Hasher}; + pgrx::pg_module_magic!(); #[pg_extern] @@ -66,6 +68,21 @@ fn typeid_ne(a: TypeID, b: TypeID) -> bool { typeid_cmp(a, b) != 0 } +#[pg_extern] +fn typeid_hash(typeid: TypeID) -> i32 { + let mut hasher = gxhash::GxHasher::default(); + typeid.hash(&mut hasher); + hasher.finish() as i32 +} + +#[pg_extern] +fn typeid_hash_extended(typeid: TypeID, seed: i64) -> i64 { + let mut hasher = gxhash::GxHasher::with_seed(seed); + + typeid.hash(&mut hasher); + hasher.finish() as i64 +} + extension_sql! { r#" CREATE OPERATOR < ( @@ -115,6 +132,13 @@ r#" OPERATOR 4 >= (typeid, typeid), OPERATOR 5 > (typeid, typeid), FUNCTION 1 typeid_cmp(typeid, typeid); + + CREATE OPERATOR FAMILY typeid_hash_ops USING hash; + + CREATE OPERATOR CLASS typeid_hash_ops DEFAULT FOR TYPE typeid USING hash AS + OPERATOR 1 = (typeid, typeid), + FUNCTION 1 typeid_hash(typeid), + FUNCTION 2 typeid_hash_extended(typeid, bigint); "#, name = "create_typeid_operator_class", finalize, @@ -129,6 +153,7 @@ fn uuid_generate_v7() -> pgrx::Uuid { #[cfg(any(test, feature = "pg_test"))] #[pg_schema] mod tests { + use crate::TypeID; use pgrx::prelude::*; use uuid::Uuid; @@ -147,6 +172,95 @@ mod tests { assert_eq!(converted.get_version_num(), 7); } + + #[pg_test] + fn test_hashing() { + use crate::typeid_hash; + use crate::TypeID; + + let id = TypeID::from_string("qual_01j1acv2aeehk8hcapaw7qyjvq").unwrap(); + let id2 = TypeID::from_string("qual_01j1acv2aeehk8hcapaw7qyjvq").unwrap(); + + let hash = typeid_hash(id); + let hash2 = typeid_hash(id2); + println!("UUID: {:?}", hash); + + assert_eq!( + hash, hash2, + "Hashes should be consistent for the same input" + ); + } + + #[pg_test] + fn test_custom_type_in_query() { + use crate::typeid_generate; + // Create tables + Spi::run("CREATE TABLE question (id typeid);").unwrap(); + Spi::run("CREATE TABLE answer (id typeid, question typeid);").unwrap(); + + println!("Creating tables"); + // Generate and insert test data + let typeid1 = typeid_generate("qual"); + let typeid2 = typeid_generate("answer"); + let typeid3 = typeid_generate("answer"); + + insert_into_table("question", &typeid1); + + insert_answer(&typeid2, &typeid1); + insert_answer(&typeid3, &typeid1); + + // Execute the query and check results + let result = Spi::get_one::( + "SELECT COUNT(*) FROM answer WHERE question IN (SELECT id FROM question)", + ) + .unwrap(); + assert_eq!(result, Some(2)); + } + + fn oid_for_type(type_name: &str) -> Result, pgrx::spi::Error> { + use crate::pg_sys::Oid; + + let oid = Spi::get_one_with_args::( + "SELECT oid FROM pg_type WHERE typname = $1", + vec![(PgBuiltInOids::TEXTOID.oid(), type_name.into_datum())], + )?; + Ok(oid.map(|oid| PgOid::from(Oid::from(oid)))) + } + + fn insert_answer(typeid: &TypeID, reference: &TypeID) { + let query = format!( + "INSERT INTO {} (id, question) VALUES ($1::typeid, $2::typeid)", + "answer" + ); + let oid = oid_for_type("typeid") + .unwrap() + .expect("expected to find oid"); + + println!("Inserting into table: {:?}", oid); + Spi::run_with_args( + &query, + Some(vec![ + (oid, typeid.clone().into_datum()), + (oid, reference.clone().into_datum()), + ]), + ) + .unwrap(); + } + + fn insert_into_table(table_name: &str, typeid: &TypeID) { + let query = format!("INSERT INTO {} (id) VALUES ($1::typeid)", table_name); + let oid = oid_for_type("typeid").unwrap(); + + println!("Inserting into table: {:?}", oid.unwrap()); + Spi::run_with_args( + &query, + Some(vec![( + oid.expect("expected to find oid"), + typeid.clone().into_datum(), + )]), + ) + .unwrap(); + } } /// This module is required by `cargo pgrx test` invocations. diff --git a/src/typeid.rs b/src/typeid.rs index 7540184..eb100ba 100644 --- a/src/typeid.rs +++ b/src/typeid.rs @@ -3,6 +3,7 @@ use std::borrow::Cow; use pgrx::prelude::*; use serde::{Deserialize, Serialize}; +use std::hash::{Hash, Hasher}; use uuid::Uuid; use crate::base32::{decode_base32_uuid, encode_base32_uuid}; @@ -24,7 +25,7 @@ pub enum Error { InvalidData, } -#[derive(Serialize, Deserialize)] +#[derive(Serialize, Deserialize, PartialEq, Eq, Clone)] pub struct TypeIDPrefix(String); impl TypeIDPrefix { @@ -72,7 +73,7 @@ impl TypeIDPrefix { } } -#[derive(Serialize, Deserialize, PostgresType)] +#[derive(Serialize, Deserialize, Clone, PostgresType, PartialEq, Eq)] #[inoutfuncs] pub struct TypeID(TypeIDPrefix, Uuid); @@ -107,6 +108,13 @@ impl TypeID { } } +impl Hash for TypeID { + fn hash(&self, state: &mut H) { + self.type_prefix().as_bytes().hash(state); + self.uuid().hash(state); + } +} + impl fmt::Display for TypeID { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { if self.type_prefix().is_empty() {