From 995da5b8a6df8a98f587b9482e65e6d839944546 Mon Sep 17 00:00:00 2001 From: Karel Maesen Date: Sun, 19 Jan 2025 17:23:52 +0100 Subject: [PATCH] HHH-19054 - pgvector support for CockroachDB --- docker_db.sh | 39 ++++++++++++++++++- .../vector/PGVectorFunctionContributor.java | 4 +- .../vector/PGVectorTypeContributor.java | 4 +- .../org/hibernate/vector/PGVectorTest.java | 10 ++++- 4 files changed, 53 insertions(+), 4 deletions(-) diff --git a/docker_db.sh b/docker_db.sh index e6a78a7b2b1c..ba8252e55ca6 100755 --- a/docker_db.sh +++ b/docker_db.sh @@ -812,7 +812,44 @@ hana() { } cockroachdb() { - cockroachdb_24_1 + cockroachdb_24_3 +} + +cockroachdb_24_3() { + $CONTAINER_CLI rm -f cockroach || true + LOG_CONFIG=" +sinks: + stderr: + channels: all + filter: ERROR + redact: false + exit-on-error: true +" + $CONTAINER_CLI run -d --name=cockroach -m 6g -p 26257:26257 -p 8080:8080 ${DB_IMAGE_COCKROACHDB_24_3:-cockroachdb/cockroach:v24.3.3} start-single-node \ + --insecure --store=type=mem,size=0.25 --advertise-addr=localhost --log="$LOG_CONFIG" + OUTPUT= + while [[ $OUTPUT != *"CockroachDB node starting"* ]]; do + echo "Waiting for CockroachDB to start..." + sleep 10 + # Note we need to redirect stderr to stdout to capture the logs + OUTPUT=$($CONTAINER_CLI logs cockroach 2>&1) + done + echo "Enabling experimental box2d operators and some optimized settings for running the tests" + #settings documented in https://www.cockroachlabs.com/docs/v24.1/local-testing#use-a-local-single-node-cluster-with-in-memory-storage + $CONTAINER_CLI exec cockroach bash -c "cat < doubleType = basicTypeRegistry.resolve( StandardBasicTypes.DOUBLE ); final BasicType integerType = basicTypeRegistry.resolve( StandardBasicTypes.INTEGER ); functionRegistry.patternDescriptorBuilder( "cosine_distance", "?1<=>?2" ) diff --git a/hibernate-vector/src/main/java/org/hibernate/vector/PGVectorTypeContributor.java b/hibernate-vector/src/main/java/org/hibernate/vector/PGVectorTypeContributor.java index d18287cf8de5..497d7a50dd97 100644 --- a/hibernate-vector/src/main/java/org/hibernate/vector/PGVectorTypeContributor.java +++ b/hibernate-vector/src/main/java/org/hibernate/vector/PGVectorTypeContributor.java @@ -8,6 +8,7 @@ import org.hibernate.boot.model.TypeContributions; import org.hibernate.boot.model.TypeContributor; +import org.hibernate.dialect.CockroachDialect; import org.hibernate.dialect.Dialect; import org.hibernate.dialect.PostgreSQLDialect; import org.hibernate.engine.jdbc.Size; @@ -34,7 +35,8 @@ public class PGVectorTypeContributor implements TypeContributor { @Override public void contribute(TypeContributions typeContributions, ServiceRegistry serviceRegistry) { final Dialect dialect = serviceRegistry.requireService( JdbcServices.class ).getDialect(); - if ( dialect instanceof PostgreSQLDialect ) { + if ( dialect instanceof PostgreSQLDialect || + dialect instanceof CockroachDialect ) { final TypeConfiguration typeConfiguration = typeContributions.getTypeConfiguration(); final JavaTypeRegistry javaTypeRegistry = typeConfiguration.getJavaTypeRegistry(); final JdbcTypeRegistry jdbcTypeRegistry = typeConfiguration.getJdbcTypeRegistry(); diff --git a/hibernate-vector/src/test/java/org/hibernate/vector/PGVectorTest.java b/hibernate-vector/src/test/java/org/hibernate/vector/PGVectorTest.java index 67de8ea6b099..fc26b59b7521 100644 --- a/hibernate-vector/src/test/java/org/hibernate/vector/PGVectorTest.java +++ b/hibernate-vector/src/test/java/org/hibernate/vector/PGVectorTest.java @@ -8,7 +8,10 @@ import org.hibernate.annotations.Array; import org.hibernate.annotations.JdbcTypeCode; +import org.hibernate.dialect.CockroachDialect; import org.hibernate.dialect.PostgreSQLDialect; +import org.hibernate.testing.orm.junit.RequiresDialects; +import org.hibernate.testing.orm.junit.SkipForDialect; import org.hibernate.type.SqlTypes; import org.hibernate.testing.orm.junit.DomainModel; @@ -32,7 +35,10 @@ */ @DomainModel(annotatedClasses = PGVectorTest.VectorEntity.class) @SessionFactory -@RequiresDialect(value = PostgreSQLDialect.class, matchSubTypes = false) +@RequiresDialects({ + @RequiresDialect(value = PostgreSQLDialect.class, matchSubTypes = false), + @RequiresDialect(value = CockroachDialect.class, majorVersion = 24, minorVersion = 2) +}) public class PGVectorTest { private static final float[] V1 = new float[]{ 1, 2, 3 }; @@ -166,6 +172,7 @@ public void testVectorNorm(SessionFactoryScope scope) { } @Test + @SkipForDialect(dialectClass = CockroachDialect.class, reason = "CockroachDB does not currently support the sum() function on vector type" ) public void testVectorSum(SessionFactoryScope scope) { scope.inTransaction( em -> { //tag::vector-sum-example[] @@ -178,6 +185,7 @@ public void testVectorSum(SessionFactoryScope scope) { } @Test + @SkipForDialect(dialectClass = CockroachDialect.class, reason = "CockroachDB does not currently support the avg() function on vector type" ) public void testVectorAvg(SessionFactoryScope scope) { scope.inTransaction( em -> { //tag::vector-avg-example[]