added filter for vector lengths

Author: darkcoderrises

posting/index.go

@@ -1373,6 +1373,40 @@ func rebuildVectorIndex(ctx context.Context, factorySpecs []*tok.FactoryCreateSp
 		return err
 	}
 
+	numVectorsToCheck := 100
+	lenFreq := make(map[int]int, numVectorsToCheck)
+	maxFreq := 0
+	dimension := 0
+	MemLayerInstance.IterateDisk(ctx, IterateDiskArgs{
+		Prefix:      pk.DataPrefix(),
+		ReadTs:      rb.StartTs,
+		AllVersions: false,
+		Reverse:     false,
+		CheckInclusion: func(uid uint64) error {
+			return nil
+		},
+		Function: func(l *List, pk x.ParsedKey) error {
+			val, err := l.Value(rb.StartTs)
+			if err != nil {
+				return err
+			}
+			inVec := types.BytesAsFloatArray(val.Value.([]byte))
+			lenFreq[len(inVec)] += 1
+			if lenFreq[len(inVec)] > maxFreq {
+				maxFreq = lenFreq[len(inVec)]
+				dimension = len(inVec)
+			}
+			numVectorsToCheck -= 1
+			if numVectorsToCheck <= 0 {
+				return ErrStopIteration
+			}
+			return nil
+		},
+		StartKey: x.DataKey(rb.Attr, 0),
+	})
+
+	fmt.Println("Selecting vector dimension to be:", dimension)
+
 	if indexer.NumSeedVectors() > 0 {
 		count := 0
 		MemLayerInstance.IterateDisk(ctx, IterateDiskArgs{
@@ -1389,6 +1423,9 @@ func rebuildVectorIndex(ctx context.Context, factorySpecs []*tok.FactoryCreateSp
 					return err
 				}
 				inVec := types.BytesAsFloatArray(val.Value.([]byte))
+				if len(inVec) != dimension {
+					return nil
+				}
 				count += 1
 				indexer.AddSeedVector(inVec)
 				if count == indexer.NumSeedVectors() {
@@ -1423,6 +1460,9 @@ func rebuildVectorIndex(ctx context.Context, factorySpecs []*tok.FactoryCreateSp
 			}
 
 			inVec := types.BytesAsFloatArray(val.Value.([]byte))
+			if len(inVec) != dimension {
+				return []*pb.DirectedEdge{}, nil
+			}
 			indexer.BuildInsert(ctx, uid, inVec)
 			return edges, nil
 		}
@@ -1449,6 +1489,12 @@ func rebuildVectorIndex(ctx context.Context, factorySpecs []*tok.FactoryCreateSp
 			}
 
 			inVec := types.BytesAsFloatArray(val.Value.([]byte))
+			if len(inVec) != dimension {
+				if pass_idx == 0 {
+					glog.Warningf("Skipping vector with invalid dimension uid: %d, dimension: %d", uid, len(inVec))
+				}
+				return []*pb.DirectedEdge{}, nil
+			}
 			indexer.BuildInsert(ctx, uid, inVec)
 			return edges, nil
 		}

tok/kmeans/kmeans.go

@@ -7,6 +7,7 @@ import (
 
 	c "github.com/hypermodeinc/dgraph/v25/tok/constraints"
 	"github.com/hypermodeinc/dgraph/v25/tok/index"
+	"github.com/hypermodeinc/dgraph/v25/x"
 )
 
 type Kmeans[T c.Float] struct {
@@ -106,8 +107,13 @@ func (vc *vectorCentroids[T]) addVector(vec []T) error {
 }
 
 func (vc *vectorCentroids[T]) updateCentroids() {
+	x.AssertTrue(len(vc.centroids) == vc.numCenters)
+	x.AssertTrue(len(vc.counts) == vc.numCenters)
+	x.AssertTrue(len(vc.weights) == vc.numCenters)
 	for i := 0; i < vc.numCenters; i++ {
 		for j := 0; j < vc.dimension; j++ {
+			x.AssertTrue(len(vc.centroids[i]) == vc.dimension)
+			x.AssertTrue(len(vc.weights[i]) == vc.dimension)
 			vc.centroids[i][j] = vc.weights[i][j] / T(vc.counts[i])
 			vc.weights[i][j] = 0
 		}
@@ -119,6 +125,9 @@ func (vc *vectorCentroids[T]) updateCentroids() {
 
 func (vc *vectorCentroids[T]) randomInit() {
 	vc.dimension = len(vc.centroids[0])
+	for i := range vc.centroids {
+		x.AssertTrue(len(vc.centroids[i]) == vc.dimension)
+	}
 	vc.numCenters = len(vc.centroids)
 	vc.counts = make([]int64, vc.numCenters)
 	vc.weights = make([][]T, vc.numCenters)