Merge pull request #3 from blitss/feat/hash-indexes

feat: add support for hash indexes (requires aes now)

Author: blitss

.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"]

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
+harness = false

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::<i64>(
+            "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<Option<PgOid>, pgrx::spi::Error> {
+        use crate::pg_sys::Oid;
+
+        let oid = Spi::get_one_with_args::<u32>(
+            "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.

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<H: Hasher>(&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() {