Skip to content

Commit f906e20

Browse files
maesenkabeikov
authored andcommitted
HHH-19054 - pgvector support for CockroachDB
1 parent cf0f2e4 commit f906e20

File tree

4 files changed

+53
-4
lines changed

4 files changed

+53
-4
lines changed

docker_db.sh

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -812,7 +812,44 @@ hana() {
812812
}
813813

814814
cockroachdb() {
815-
cockroachdb_24_1
815+
cockroachdb_24_3
816+
}
817+
818+
cockroachdb_24_3() {
819+
$CONTAINER_CLI rm -f cockroach || true
820+
LOG_CONFIG="
821+
sinks:
822+
stderr:
823+
channels: all
824+
filter: ERROR
825+
redact: false
826+
exit-on-error: true
827+
"
828+
$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 \
829+
--insecure --store=type=mem,size=0.25 --advertise-addr=localhost --log="$LOG_CONFIG"
830+
OUTPUT=
831+
while [[ $OUTPUT != *"CockroachDB node starting"* ]]; do
832+
echo "Waiting for CockroachDB to start..."
833+
sleep 10
834+
# Note we need to redirect stderr to stdout to capture the logs
835+
OUTPUT=$($CONTAINER_CLI logs cockroach 2>&1)
836+
done
837+
echo "Enabling experimental box2d operators and some optimized settings for running the tests"
838+
#settings documented in https://www.cockroachlabs.com/docs/v24.1/local-testing#use-a-local-single-node-cluster-with-in-memory-storage
839+
$CONTAINER_CLI exec cockroach bash -c "cat <<EOF | ./cockroach sql --insecure
840+
SET CLUSTER SETTING sql.spatial.experimental_box2d_comparison_operators.enabled = on;
841+
SET CLUSTER SETTING kv.range_merge.queue_interval = '50ms';
842+
SET CLUSTER SETTING jobs.registry.interval.gc = '30s';
843+
SET CLUSTER SETTING jobs.registry.interval.cancel = '180s';
844+
SET CLUSTER SETTING jobs.retention_time = '5s';
845+
SET CLUSTER SETTING sql.stats.automatic_collection.enabled = false;
846+
ALTER RANGE default CONFIGURE ZONE USING "gc.ttlseconds" = 300;
847+
ALTER DATABASE system CONFIGURE ZONE USING "gc.ttlseconds" = 300;
848+
849+
quit
850+
EOF
851+
"
852+
echo "Cockroachdb successfully started"
816853
}
817854

818855
cockroachdb_24_1() {

hibernate-vector/src/main/java/org/hibernate/vector/PGVectorFunctionContributor.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
import org.hibernate.boot.model.FunctionContributions;
88
import org.hibernate.boot.model.FunctionContributor;
9+
import org.hibernate.dialect.CockroachDialect;
910
import org.hibernate.dialect.Dialect;
1011
import org.hibernate.dialect.PostgreSQLDialect;
1112
import org.hibernate.query.sqm.function.SqmFunctionRegistry;
@@ -24,7 +25,8 @@ public void contributeFunctions(FunctionContributions functionContributions) {
2425
final TypeConfiguration typeConfiguration = functionContributions.getTypeConfiguration();
2526
final BasicTypeRegistry basicTypeRegistry = typeConfiguration.getBasicTypeRegistry();
2627
final Dialect dialect = functionContributions.getDialect();
27-
if ( dialect instanceof PostgreSQLDialect ) {
28+
if ( dialect instanceof PostgreSQLDialect ||
29+
dialect instanceof CockroachDialect ) {
2830
final BasicType<Double> doubleType = basicTypeRegistry.resolve( StandardBasicTypes.DOUBLE );
2931
final BasicType<Integer> integerType = basicTypeRegistry.resolve( StandardBasicTypes.INTEGER );
3032
functionRegistry.patternDescriptorBuilder( "cosine_distance", "?1<=>?2" )

hibernate-vector/src/main/java/org/hibernate/vector/PGVectorTypeContributor.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
import org.hibernate.boot.model.TypeContributions;
1010
import org.hibernate.boot.model.TypeContributor;
11+
import org.hibernate.dialect.CockroachDialect;
1112
import org.hibernate.dialect.Dialect;
1213
import org.hibernate.dialect.PostgreSQLDialect;
1314
import org.hibernate.engine.jdbc.Size;
@@ -34,7 +35,8 @@ public class PGVectorTypeContributor implements TypeContributor {
3435
@Override
3536
public void contribute(TypeContributions typeContributions, ServiceRegistry serviceRegistry) {
3637
final Dialect dialect = serviceRegistry.requireService( JdbcServices.class ).getDialect();
37-
if ( dialect instanceof PostgreSQLDialect ) {
38+
if ( dialect instanceof PostgreSQLDialect ||
39+
dialect instanceof CockroachDialect ) {
3840
final TypeConfiguration typeConfiguration = typeContributions.getTypeConfiguration();
3941
final JavaTypeRegistry javaTypeRegistry = typeConfiguration.getJavaTypeRegistry();
4042
final JdbcTypeRegistry jdbcTypeRegistry = typeConfiguration.getJdbcTypeRegistry();

hibernate-vector/src/test/java/org/hibernate/vector/PGVectorTest.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88

99
import org.hibernate.annotations.Array;
1010
import org.hibernate.annotations.JdbcTypeCode;
11+
import org.hibernate.dialect.CockroachDialect;
1112
import org.hibernate.dialect.PostgreSQLDialect;
13+
import org.hibernate.testing.orm.junit.RequiresDialects;
14+
import org.hibernate.testing.orm.junit.SkipForDialect;
1215
import org.hibernate.type.SqlTypes;
1316

1417
import org.hibernate.testing.orm.junit.DomainModel;
@@ -32,7 +35,10 @@
3235
*/
3336
@DomainModel(annotatedClasses = PGVectorTest.VectorEntity.class)
3437
@SessionFactory
35-
@RequiresDialect(value = PostgreSQLDialect.class, matchSubTypes = false)
38+
@RequiresDialects({
39+
@RequiresDialect(value = PostgreSQLDialect.class, matchSubTypes = false),
40+
@RequiresDialect(value = CockroachDialect.class, majorVersion = 24, minorVersion = 2)
41+
})
3642
public class PGVectorTest {
3743

3844
private static final float[] V1 = new float[]{ 1, 2, 3 };
@@ -166,6 +172,7 @@ public void testVectorNorm(SessionFactoryScope scope) {
166172
}
167173

168174
@Test
175+
@SkipForDialect(dialectClass = CockroachDialect.class, reason = "CockroachDB does not currently support the sum() function on vector type" )
169176
public void testVectorSum(SessionFactoryScope scope) {
170177
scope.inTransaction( em -> {
171178
//tag::vector-sum-example[]
@@ -178,6 +185,7 @@ public void testVectorSum(SessionFactoryScope scope) {
178185
}
179186

180187
@Test
188+
@SkipForDialect(dialectClass = CockroachDialect.class, reason = "CockroachDB does not currently support the avg() function on vector type" )
181189
public void testVectorAvg(SessionFactoryScope scope) {
182190
scope.inTransaction( em -> {
183191
//tag::vector-avg-example[]

0 commit comments

Comments
 (0)