Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion internal/core/diskann.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ func DefaultDiskANNBuildOptions(metric Metric) DiskANNBuildOptions {

// Validate checks invariants that do not depend on vector dimension.
func (o DiskANNBuildOptions) Validate() error {
if !o.Metric.valid() {
if !o.Metric.Valid() {
return fmt.Errorf("%w: invalid metric", ErrInvalidDiskANNOptions)
}
if o.MaxDegree <= 0 || o.MaxDegree > MaxVamanaDegree {
Expand Down
2 changes: 1 addition & 1 deletion internal/core/diskann_storage.go
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ type DiskANNLayout struct {

// NewDiskANNLayout calculates the pinned packed-or-multi-sector node layout.
func NewDiskANNLayout(metric Metric, count, dimension, maxDegree int) (DiskANNLayout, error) {
if !metric.valid() {
if !metric.Valid() {
return DiskANNLayout{}, fmt.Errorf("%w: invalid metric", ErrInvalidDiskANNLayout)
}
if count < 0 || uint64(count) > math.MaxUint32 {
Expand Down
8 changes: 2 additions & 6 deletions internal/core/flat.go
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ func NewDenseFlatIndex(dimension int, metric Metric) (*DenseFlatIndex, error) {
if dimension <= 0 {
return nil, fmt.Errorf("%w: got %d", ErrInvalidDimension, dimension)
}
if !metric.valid() {
if !metric.Valid() {
return nil, errors.New("core: invalid metric")
}
return &DenseFlatIndex{
Expand Down Expand Up @@ -244,10 +244,6 @@ func (b *DenseFlatIndexBuilder) Build(ctx context.Context) (DenseIndex, error) {
return b.index, nil
}

func (m Metric) valid() bool {
return m >= MetricL2 && m <= MetricMIPSL2
}

var (
_ DenseProvider = (*DenseFlatIndex)(nil)
_ DenseSearcher = (*DenseFlatIndex)(nil)
Expand Down Expand Up @@ -698,7 +694,7 @@ func QueryDenseGroups(
if err := ctx.Err(); err != nil {
return nil, err
}
if !metric.valid() {
if !metric.Valid() {
return nil, errors.New("core: invalid group-by metric")
}
if err := options.Validate(); err != nil {
Expand Down
8 changes: 4 additions & 4 deletions internal/core/hnsw.go
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ func DefaultHNSWBuildOptions(metric Metric) HNSWBuildOptions {

// Validate checks graph degree and construction-search invariants.
func (o HNSWBuildOptions) Validate() error {
if !o.Metric.valid() {
if !o.Metric.Valid() {
return fmt.Errorf("%w: invalid metric", ErrInvalidHNSWOptions)
}
if o.M <= 0 || o.M > MaxHNSWM {
Expand Down Expand Up @@ -173,7 +173,7 @@ func (b *HNSWBuilder) build(ctx context.Context, workers int) (*HNSWIndex, error
if err := ctx.Err(); err != nil {
return nil, err
}
distance, err := b.options.Metric.prevalidatedDistance()
distance, err := b.options.Metric.PrevalidatedDistance()
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -559,7 +559,7 @@ func (i *HNSWIndex) computeDistance(left, right []float32) (float32, error) {
// Keep package-local literal fixtures usable while production indexes
// always install the scorer at build or open time.
var err error
distance, err = i.options.Metric.prevalidatedDistance()
distance, err = i.options.Metric.PrevalidatedDistance()
if err != nil {
return 0, err
}
Expand Down Expand Up @@ -1203,7 +1203,7 @@ func decodeHNSWIndex(ctx context.Context, encoded []byte) (*HNSWIndex, error) {
if count > maxPlatformInt()/dimension {
return nil, fmt.Errorf("%w: vector storage exceeds platform capacity", ErrInvalidHNSWFile)
}
distance, err := options.Metric.prevalidatedDistance()
distance, err := options.Metric.PrevalidatedDistance()
if err != nil {
return nil, fmt.Errorf("%w: invalid metric", ErrInvalidHNSWFile)
}
Expand Down
2 changes: 1 addition & 1 deletion internal/core/ivf.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ func DefaultIVFBuildOptions(metric Metric) IVFBuildOptions {

// Validate checks IVF build invariants.
func (o IVFBuildOptions) Validate() error {
if !o.Metric.valid() {
if !o.Metric.Valid() {
return fmt.Errorf("%w: invalid metric", ErrInvalidIVFOptions)
}
if o.NList <= 0 {
Expand Down
4 changes: 2 additions & 2 deletions internal/core/kmeans.go
Original file line number Diff line number Diff line change
Expand Up @@ -299,7 +299,7 @@ func validateKMeansOptions(options KMeansOptions) error {
if options.Tolerance < 0 || math.IsNaN(options.Tolerance) || math.IsInf(options.Tolerance, 0) {
return fmt.Errorf("%w: Tolerance must be finite and non-negative", ErrInvalidKMeansOptions)
}
if !options.Metric.valid() {
if !options.Metric.Valid() {
return fmt.Errorf("%w: invalid metric", ErrInvalidKMeansOptions)
}
if options.Initializer != KMeansInitReservoir && options.Initializer != KMeansInitPlusPlus {
Expand Down Expand Up @@ -512,7 +512,7 @@ func cloneVectorsContext(ctx context.Context, vectors [][]float32) ([][]float32,
}

func (m *KMeansModel) validate() error {
if m == nil || !m.metric.valid() || m.dimension <= 0 || len(m.centroids) == 0 {
if m == nil || !m.metric.Valid() || m.dimension <= 0 || len(m.centroids) == 0 {
return errors.New("core: invalid k-means model")
}
return nil
Expand Down
78 changes: 78 additions & 0 deletions internal/core/metric/metric.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
// Copyright 2026-present the xvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package metric provides dense-vector score computation and ordering.
package metric

import (
"errors"

"github.com/gorse-io/xvec/internal/ailego/math"
)

// Metric selects score computation and ordering for vector search.
type Metric uint8

const (
L2 Metric = iota + 1
IP
Cosine
MIPSL2
)

// Compute calculates the score for left and right.
func (m Metric) Compute(left, right []float32) (float32, error) {
switch m {
case L2:
return mathutil.L2Squared(left, right)
case IP:
return mathutil.InnerProduct(left, right)
case Cosine:
return mathutil.CosineDistance(left, right)
case MIPSL2:
return mathutil.MIPSL2Squared(left, right)
default:
return 0, errors.New("core: invalid metric")
}
}

// PrevalidatedDistance selects the allocation-free kernel used by index hot
// paths after vectors have passed their storage or query boundary validation.
func (m Metric) PrevalidatedDistance() (mathutil.DenseDistance, error) {
switch m {
case L2:
return mathutil.L2SquaredPrevalidated, nil
case IP:
return mathutil.InnerProductPrevalidated, nil
case Cosine:
return mathutil.CosineDistancePrevalidated, nil
case MIPSL2:
return mathutil.MIPSL2SquaredPrevalidated, nil
default:
return nil, errors.New("core: invalid metric")
}
}

// Better reports whether left should rank before right.
func (m Metric) Better(left, right float32) bool {
if m == IP {
return left > right
}
return left < right
}

// Valid reports whether m identifies a supported metric.
func (m Metric) Valid() bool {
return m >= L2 && m <= MIPSL2
}
37 changes: 37 additions & 0 deletions internal/core/metric/metric_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
// Copyright 2026-present the xvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");

package metric

import (
"testing"

"github.com/stretchr/testify/require"
)

func TestMetricComputeOrderingAndValidation(t *testing.T) {
t.Parallel()

left := []float32{0.2, 0.9, -0.4, 0.7}
right := []float32{0.3, 0.5, 0.8, -0.1}
for _, value := range []Metric{L2, IP, Cosine, MIPSL2} {
require.True(t, value.Valid())
expected, err := value.Compute(left, right)
require.NoError(t, err)
distance, err := value.PrevalidatedDistance()
require.NoError(t, err)
actual, err := distance(left, right)
require.NoError(t, err)
require.Equal(t, expected, actual)
}

require.False(t, Metric(0).Valid())
_, err := Metric(0).Compute(left, right)
require.Error(t, err)
_, err = Metric(0).PrevalidatedDistance()
require.Error(t, err)

require.True(t, IP.Better(2, 1))
require.True(t, L2.Better(1, 2))
}
2 changes: 1 addition & 1 deletion internal/core/quantization.go
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ func (v QuantizedVector) Decode() ([]float32, error) {
// QuantizedDistance calculates a metric directly from scalar codes. Both
// vectors must use the same encoding and logical dimension.
func QuantizedDistance(metric Metric, left, right QuantizedVector) (float32, error) {
if !metric.valid() {
if !metric.Valid() {
return 0, errors.New("core: invalid metric")
}
if err := left.validate(); err != nil {
Expand Down
2 changes: 1 addition & 1 deletion internal/core/quantized_index.go
Original file line number Diff line number Diff line change
Expand Up @@ -197,7 +197,7 @@ func newScalarQuantizedVectors(
if dimension <= 0 || dimension > MaxRotationDimension {
return nil, fmt.Errorf("%w: got %d", ErrInvalidDimension, dimension)
}
if !metric.valid() {
if !metric.Valid() {
return nil, errors.New("core: invalid scalar-quantized index metric")
}
if !kind.valid() {
Expand Down
2 changes: 1 addition & 1 deletion internal/core/query.go
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ func QueryDense(
if err := ctx.Err(); err != nil {
return nil, err
}
if !metric.valid() {
if !metric.Valid() {
return nil, errors.New("core: invalid metric")
}
if err := options.Validate(); err != nil {
Expand Down
2 changes: 1 addition & 1 deletion internal/core/refiner.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ func NewOriginalVectorRefiner(provider DenseProvider, metric Metric) (*OriginalV
if provider.Dimension() <= 0 {
return nil, ErrInvalidDimension
}
if !metric.valid() {
if !metric.Valid() {
return nil, errors.New("core: invalid metric")
}
return &OriginalVectorRefiner{provider: provider, metric: metric}, nil
Expand Down
52 changes: 6 additions & 46 deletions internal/core/topk.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,16 +23,17 @@ import (
"github.com/gorse-io/xvec/internal/ailego/container"
"github.com/gorse-io/xvec/internal/ailego/math"
"github.com/gorse-io/xvec/internal/ailego/parallel"
"github.com/gorse-io/xvec/internal/core/metric"
)

// Metric selects score computation and ordering for exact search.
type Metric uint8
type Metric = metric.Metric

const (
MetricL2 Metric = iota + 1
MetricIP
MetricCosine
MetricMIPSL2
MetricL2 = metric.L2
MetricIP = metric.IP
MetricCosine = metric.Cosine
MetricMIPSL2 = metric.MIPSL2
)

// Candidate is one immutable dense vector considered by exact search.
Expand All @@ -48,47 +49,6 @@ type Result struct {
Score float32
}

// Compute calculates the score for left and right.
func (m Metric) Compute(left, right []float32) (float32, error) {
switch m {
case MetricL2:
return mathutil.L2Squared(left, right)
case MetricIP:
return mathutil.InnerProduct(left, right)
case MetricCosine:
return mathutil.CosineDistance(left, right)
case MetricMIPSL2:
return mathutil.MIPSL2Squared(left, right)
default:
return 0, errors.New("core: invalid metric")
}
}

// prevalidatedDistance selects the allocation-free kernel used by index hot
// paths after vectors have passed their storage or query boundary validation.
func (m Metric) prevalidatedDistance() (mathutil.DenseDistance, error) {
switch m {
case MetricL2:
return mathutil.L2SquaredPrevalidated, nil
case MetricIP:
return mathutil.InnerProductPrevalidated, nil
case MetricCosine:
return mathutil.CosineDistancePrevalidated, nil
case MetricMIPSL2:
return mathutil.MIPSL2SquaredPrevalidated, nil
default:
return nil, errors.New("core: invalid metric")
}
}

// Better reports whether left should rank before right.
func (m Metric) Better(left, right float32) bool {
if m == MetricIP {
return left > right
}
return left < right
}

// TopK computes exact scores and returns at most k results. It uses O(k)
// memory and checks ctx between candidates.
func TopK(
Expand Down
18 changes: 0 additions & 18 deletions internal/core/topk_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -85,24 +85,6 @@ func TestTopKMetricOrdering(t *testing.T) {
}
}

func TestMetricPrevalidatedDistanceMatchesCheckedCompute(t *testing.T) {
t.Parallel()

left := []float32{0.2, 0.9, -0.4, 0.7}
right := []float32{0.3, 0.5, 0.8, -0.1}
for _, metric := range []Metric{MetricL2, MetricIP, MetricCosine, MetricMIPSL2} {
distance, err := metric.prevalidatedDistance()
require.NoError(t, err)
expected, err := metric.Compute(left, right)
require.NoError(t, err)
actual, err := distance(left, right)
require.NoError(t, err)
require.Equal(t, expected, actual)
}
_, err := Metric(0).prevalidatedDistance()
require.Error(t, err)
}

func TestTopKStableAcrossCandidateOrder(t *testing.T) {
t.Parallel()

Expand Down
2 changes: 1 addition & 1 deletion internal/core/vamana.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ func DefaultVamanaBuildOptions(metric Metric) VamanaBuildOptions {

// Validate checks graph degree, construction width, and RobustPrune settings.
func (o VamanaBuildOptions) Validate() error {
if !o.Metric.valid() {
if !o.Metric.Valid() {
return fmt.Errorf("%w: invalid metric", ErrInvalidVamanaOptions)
}
if o.MaxDegree <= 0 || o.MaxDegree > MaxVamanaDegree {
Expand Down
Loading