Added distance functions for Ent - #23

Author: ankane

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.2.4 (unreleased)
+
+- Added distance functions for Ent
+
 ## 0.2.3 (2025-01-15)
 
 - Added support for Postgres arrays for pgx

ent_test.go

@@ -9,6 +9,7 @@ import (
 	_ "github.com/lib/pq"
 	"github.com/pgvector/pgvector-go"
 	"github.com/pgvector/pgvector-go/ent"
+	"github.com/pgvector/pgvector-go/entvec"
 )
 
 func TestEnt(t *testing.T) {
@@ -67,7 +68,7 @@ func TestEnt(t *testing.T) {
 	items, err := client.Item.
 		Query().
 		Order(func(s *sql.Selector) {
-			s.OrderExpr(sql.ExprP("embedding <-> $1", embedding))
+			s.OrderExpr(entvec.L2Distance("embedding", embedding))
 		}).
 		Limit(5).
 		All(ctx)

entvec/distance.go

@@ -0,0 +1,41 @@
+package entvec
+
+import (
+	"entgo.io/ent/dialect/sql"
+)
+
+func L2Distance(column string, value any) sql.Querier {
+	return sql.ExprFunc(func(b *sql.Builder) {
+		b.Ident(column).WriteString(" <-> ").Arg(value)
+	})
+}
+
+func MaxInnerProduct(column string, value any) sql.Querier {
+	return sql.ExprFunc(func(b *sql.Builder) {
+		b.Ident(column).WriteString(" <#> ").Arg(value)
+	})
+}
+
+func CosineDistance(column string, value any) sql.Querier {
+	return sql.ExprFunc(func(b *sql.Builder) {
+		b.Ident(column).WriteString(" <=> ").Arg(value)
+	})
+}
+
+func L1Distance(column string, value any) sql.Querier {
+	return sql.ExprFunc(func(b *sql.Builder) {
+		b.Ident(column).WriteString(" <+> ").Arg(value)
+	})
+}
+
+func HammingDistance(column string, value any) sql.Querier {
+	return sql.ExprFunc(func(b *sql.Builder) {
+		b.Ident(column).WriteString(" <~> ").Arg(value)
+	})
+}
+
+func JaccardDistance(column string, value any) sql.Querier {
+	return sql.ExprFunc(func(b *sql.Builder) {
+		b.Ident(column).WriteString(" <%> ").Arg(value)
+	})
+}