redis/go-redis
1
0
Fork 0
mirror of https://github.com/redis/go-redis.git synced 2026-10-22 19:49:00 +03:00
Code Activity

feat: sq compression hnsw (#4077)

Browse Source
This commit is contained in:
Vladislav Kotsev
2026-10-08 14:56:14 +03:00
committed by GitHub
parent f79ab09081
commit 760b98dc0b
3 changed files with 211 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
search_commands.go
View File

@@ -115,6 +115,16 @@ type FTHNSWOptions struct {
// false can be distinguished from unset (omitted).
Rerank bool
HasRerank bool
// Compression selects the vector quantization scheme for HNSW indexes
// (Redis 8.12+), e.g. "SQ8". Emitted as COMPRESSION when non-empty.
Compression string
// TrainingThreshold is the number of vectors the index collects before
// training the compression model. A positive value emits
// TRAINING_THRESHOLD on its own; to emit an explicit TRAINING_THRESHOLD 0,
// set HasTrainingThreshold=true with TrainingThreshold=0, so that an
// explicit zero can be distinguished from unset (omitted).
TrainingThreshold int
HasTrainingThreshold bool
}
type FTVamanaOptions struct {
@@ -622,12 +632,14 @@ type FTAttribute struct {
WithSuffixtrie bool
// Vector specific attributes
Algorithm string
DataType string
Dim int
DistanceMetric string
M int
EFConstruction int
Algorithm string
DataType string
Dim int
DistanceMetric string
M int
EFConstruction int
Compression string
TrainingThreshold int
}
type CursorStats struct {
@@ -1586,6 +1598,13 @@ func (c cmdable) FTCreate(ctx context.Context, index string, options *FTCreateOp
}
hnswArgs = append(hnswArgs, "RERANK", rerank)
}
if schema.VectorArgs.HNSWOptions.Compression != "" {
hnswArgs = append(hnswArgs, "COMPRESSION", schema.VectorArgs.HNSWOptions.Compression)
}
if schema.VectorArgs.HNSWOptions.TrainingThreshold > 0 || schema.VectorArgs.HNSWOptions.HasTrainingThreshold {
hnswArgs = append(hnswArgs, "TRAINING_THRESHOLD", schema.VectorArgs.HNSWOptions.TrainingThreshold)
}
args = append(args, len(hnswArgs))
args = append(args, hnswArgs...)
}
@@ -1818,6 +1837,10 @@ func parseFTAttributeFromMap(attrMap map[interface{}]interface{}) FTAttribute {
att.M = internal.ToInteger(v)
case "ef_construction":
att.EFConstruction = internal.ToInteger(v)
case "compression":
att.Compression = internal.ToString(v)
case "training_threshold":
att.TrainingThreshold = internal.ToInteger(v)
case "flags":
// flags is an array of strings like ["SORTABLE", "NOSTEM"]
if flags, ok := v.([]interface{}); ok {
@@ -2060,6 +2083,16 @@ func parseFTInfo(data map[string]interface{}) (FTInfoResult, error) {
i++
continue
}
if internal.ToLower(internal.ToString(attrSlice[i])) == "compression" && i+1 < attrLen {
att.Compression = internal.ToString(attrSlice[i+1])
i++
continue
}
if internal.ToLower(internal.ToString(attrSlice[i])) == "training_threshold" && i+1 < attrLen {
att.TrainingThreshold = internal.ToInteger(attrSlice[i+1])
i++
continue
}
}
ftInfo.Attributes = append(ftInfo.Attributes, att)
} else if attrMap, ok := attr.(map[interface{}]interface{}); ok {

121
search_commands_unit_test.go
View File

@@ -524,6 +524,127 @@ func TestFTCreateHNSWRerank_Args(t *testing.T) {
}
}
// TestFTCreateHNSWCompression_Args verifies the COMPRESSION and
// TRAINING_THRESHOLD attributes (Redis 8.12+, scalar-quantized HNSW) are
// emitted as key-value pairs counted in the attribute-count token.
// TRAINING_THRESHOLD is emitted for a positive value, or for an explicit zero
// when HasTrainingThreshold is set, and omitted otherwise.
func TestFTCreateHNSWCompression_Args(t *testing.T) {
tests := []struct {
name string
opts *FTHNSWOptions
want []any
}{
{
name: "unset",
opts: &FTHNSWOptions{Type: "FLOAT32", Dim: 2, DistanceMetric: "L2"},
want: []any{
"FT.CREATE", "idx", "SCHEMA", "v", "VECTOR", "HNSW", 6,
"TYPE", "FLOAT32", "DIM", 2, "DISTANCE_METRIC", "L2",
},
},
{
name: "compression_only",
opts: &FTHNSWOptions{Type: "FLOAT32", Dim: 2, DistanceMetric: "L2", Compression: "SQ8"},
want: []any{
"FT.CREATE", "idx", "SCHEMA", "v", "VECTOR", "HNSW", 8,
"TYPE", "FLOAT32", "DIM", 2, "DISTANCE_METRIC", "L2", "COMPRESSION", "SQ8",
},
},
{
name: "compression_and_training_threshold",
opts: &FTHNSWOptions{Type: "FLOAT32", Dim: 2, DistanceMetric: "L2", Compression: "SQ8", TrainingThreshold: 1024},
want: []any{
"FT.CREATE", "idx", "SCHEMA", "v", "VECTOR", "HNSW", 10,
"TYPE", "FLOAT32", "DIM", 2, "DISTANCE_METRIC", "L2", "COMPRESSION", "SQ8", "TRAINING_THRESHOLD", 1024,
},
},
{
name: "training_threshold_explicit_zero",
opts: &FTHNSWOptions{Type: "FLOAT32", Dim: 2, DistanceMetric: "L2", Compression: "SQ8", HasTrainingThreshold: true},
want: []any{
"FT.CREATE", "idx", "SCHEMA", "v", "VECTOR", "HNSW", 10,
"TYPE", "FLOAT32", "DIM", 2, "DISTANCE_METRIC", "L2", "COMPRESSION", "SQ8", "TRAINING_THRESHOLD", 0,
},
},
{
name: "training_threshold_zero_without_has_flag_omitted",
opts: &FTHNSWOptions{Type: "FLOAT32", Dim: 2, DistanceMetric: "L2", Compression: "SQ8", TrainingThreshold: 0},
want: []any{
"FT.CREATE", "idx", "SCHEMA", "v", "VECTOR", "HNSW", 8,
"TYPE", "FLOAT32", "DIM", 2, "DISTANCE_METRIC", "L2", "COMPRESSION", "SQ8",
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
m := &mockCmdable{}
c := m.asCmdable()
cmd := c.FTCreate(context.Background(), "idx", nil, &FieldSchema{
FieldName: "v",
FieldType: SearchFieldTypeVector,
VectorArgs: &FTVectorArgs{HNSWOptions: tt.opts},
})
if cmd.Err() != nil {
t.Fatalf("unexpected error: %v", cmd.Err())
}
if !reflect.DeepEqual(cmd.Args(), tt.want) {
t.Errorf("args mismatch\n got: %#v\nwant: %#v", cmd.Args(), tt.want)
}
})
}
}
// TestParseFTInfoHNSWCompression checks that FT.INFO populates
// FTAttribute.Compression and FTAttribute.TrainingThreshold from both the
// RESP2 flat-slice attribute shape and the RESP3 map attribute shape.
func TestParseFTInfoHNSWCompression(t *testing.T) {
t.Run("resp2 slice", func(t *testing.T) {
data := map[string]any{
"attributes": []any{
[]any{
"identifier", "v", "attribute", "v", "type", "VECTOR",
"algorithm", "HNSW", "data_type", "FLOAT32", "dim", int64(2),
"distance_metric", "L2", "M", int64(16), "ef_construction", int64(200),
"compression", "SQ8", "training_threshold", int64(1024),
},
},
}
info, err := parseFTInfo(data)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(info.Attributes) != 1 {
t.Fatalf("expected 1 attribute, got %d", len(info.Attributes))
}
att := info.Attributes[0]
if att.Compression != "SQ8" {
t.Errorf("Compression = %q, want %q", att.Compression, "SQ8")
}
if att.TrainingThreshold != 1024 {
t.Errorf("TrainingThreshold = %d, want %d", att.TrainingThreshold, 1024)
}
})
t.Run("resp3 map", func(t *testing.T) {
att := parseFTAttributeFromMap(map[any]any{
"identifier": "v",
"attribute": "v",
"type": "VECTOR",
"algorithm": "HNSW",
"compression": "SQ8",
"training_threshold": int64(1024),
})
if att.Compression != "SQ8" {
t.Errorf("Compression = %q, want %q", att.Compression, "SQ8")
}
if att.TrainingThreshold != 1024 {
t.Errorf("TrainingThreshold = %d, want %d", att.TrainingThreshold, 1024)
}
})
}
// TestParseFTHybridWarnings checks that FT.HYBRID warnings populate from either
// the "warning" or "warnings" key.
func TestParseFTHybridWarnings(t *testing.T) {

51
search_test.go
View File

@@ -2185,6 +2185,30 @@ var _ = Describe("RediSearch commands Resp 2", Label("search"), func() {
WaitForIndexing(rawClient, "idx1")
})
It("should FTCreate VECTOR with HNSW compression and read it back from FTInfo", Label("search", "ftcreate", "ftinfo"), func() {
SkipBeforeRedisVersion("8.12", "HNSW COMPRESSION requires Redis 8.12+")
hnswOptions := &redis.FTHNSWOptions{
Type: "FLOAT32",
Dim: 4,
DistanceMetric: "L2",
Compression: "SQ8",
TrainingThreshold: 1024,
}
val, err := client.FTCreate(ctx, "idx1",
&redis.FTCreateOptions{},
&redis.FieldSchema{FieldName: "v", FieldType: redis.SearchFieldTypeVector, VectorArgs: &redis.FTVectorArgs{HNSWOptions: hnswOptions}}).Result()
Expect(err).NotTo(HaveOccurred())
Expect(val).To(BeEquivalentTo("OK"))
WaitForIndexing(rawClient, "idx1")
resInfo, err := client.FTInfo(ctx, "idx1").Result()
Expect(err).NotTo(HaveOccurred())
Expect(resInfo.Attributes).To(HaveLen(1))
Expect(resInfo.Attributes[0].Algorithm).To(BeEquivalentTo("HNSW"))
Expect(resInfo.Attributes[0].Compression).To(BeEquivalentTo("SQ8"))
Expect(resInfo.Attributes[0].TrainingThreshold).To(BeEquivalentTo(1024))
})
It("should FTCreate VECTOR with VAMANA algorithm - advanced parameters", Label("search", "ftcreate"), func() {
SkipBeforeRedisVersion("8.2", "VAMANA requires Redis 8.2+")
vamanaOptions := &redis.FTVamanaOptions{
@@ -4871,4 +4895,31 @@ var _ = Describe("RediSearch commands Resp 3", Label("search"), func() {
Expect(res2).ToNot(BeEmpty())
}).ShouldNot(Panic())
})
It("should read HNSW compression back from FTInfo with RESP3", Label("search", "ftcreate", "ftinfo"), func() {
SkipBeforeRedisVersion("8.12", "HNSW COMPRESSION requires Redis 8.12+")
hnswOptions := &redis.FTHNSWOptions{
Type: "FLOAT32",
Dim: 4,
DistanceMetric: "L2",
Compression: "SQ8",
TrainingThreshold: 1024,
}
val, err := client.FTCreate(ctx, "idx1",
&redis.FTCreateOptions{},
&redis.FieldSchema{FieldName: "v", FieldType: redis.SearchFieldTypeVector, VectorArgs: &redis.FTVectorArgs{HNSWOptions: hnswOptions}}).Result()
Expect(err).NotTo(HaveOccurred())
Expect(val).To(BeEquivalentTo("OK"))
WaitForIndexing(rawClient, "idx1")
// RESP3 attributes arrive as maps and go through parseFTAttributeFromMap.
for _, c := range []redis.Cmdable{client, client2} {
resInfo, err := c.FTInfo(ctx, "idx1").Result()
Expect(err).NotTo(HaveOccurred())
Expect(resInfo.Attributes).To(HaveLen(1))
Expect(resInfo.Attributes[0].Algorithm).To(BeEquivalentTo("HNSW"))
Expect(resInfo.Attributes[0].Compression).To(BeEquivalentTo("SQ8"))
Expect(resInfo.Attributes[0].TrainingThreshold).To(BeEquivalentTo(1024))
}
})
})