Files
milvus-io--milvus/internal/util/function/chain/expr/round_decimal_expr_test.go
T
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

203 lines
6.7 KiB
Go

/*
* # Licensed to the LF AI & Data foundation under one
* # or more contributor license agreements. See the NOTICE file
* # distributed with this work for additional information
* # regarding copyright ownership. The ASF licenses this file
* # to you 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 expr
import (
"testing"
"github.com/apache/arrow/go/v17/arrow"
"github.com/apache/arrow/go/v17/arrow/array"
"github.com/apache/arrow/go/v17/arrow/memory"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/util/function/chain/types"
)
type RoundDecimalExprTestSuite struct {
suite.Suite
pool *memory.CheckedAllocator
}
func (s *RoundDecimalExprTestSuite) SetupTest() {
s.pool = memory.NewCheckedAllocator(memory.NewGoAllocator())
}
func (s *RoundDecimalExprTestSuite) TearDownTest() {
s.pool.AssertSize(s.T(), 0)
}
func TestRoundDecimalExprTestSuite(t *testing.T) {
suite.Run(t, new(RoundDecimalExprTestSuite))
}
func (s *RoundDecimalExprTestSuite) buildScoreInput(scores []float32) *arrow.Chunked {
builder := array.NewFloat32Builder(s.pool)
defer builder.Release()
builder.AppendValues(scores, nil)
arr := builder.NewArray()
chunked := arrow.NewChunked(arrow.PrimitiveTypes.Float32, []arrow.Array{arr})
arr.Release()
return chunked
}
func (s *RoundDecimalExprTestSuite) TestNewRoundDecimalExpr_ValidDecimals() {
for _, d := range []int64{0, 1, 2, 3, 4, 5, 6} {
expr, err := NewRoundDecimalExpr(d)
s.Require().NoError(err)
s.Equal(RoundDecimalFuncName, expr.Name())
s.Equal([]arrow.DataType{arrow.PrimitiveTypes.Float32}, expr.OutputDataTypes())
s.True(expr.IsRunnable(types.StageL2Rerank))
}
}
func (s *RoundDecimalExprTestSuite) TestNewRoundDecimalExpr_InvalidDecimals() {
_, err := NewRoundDecimalExpr(-1)
s.Error(err)
_, err = NewRoundDecimalExpr(7)
s.Error(err)
}
func (s *RoundDecimalExprTestSuite) TestNewRoundDecimalExprFromParams() {
expr, err := NewRoundDecimalExprFromParams(types.FunctionBuildContext{}, types.FunctionConfig{Params: map[string]*schemapb.FunctionParamValue{"decimal": intParam(2)}})
s.Require().NoError(err)
s.Equal(RoundDecimalFuncName, expr.Name())
_, err = NewRoundDecimalExprFromParams(types.FunctionBuildContext{}, types.FunctionConfig{Params: map[string]*schemapb.FunctionParamValue{}})
s.ErrorContains(err, "missing required parameter")
_, err = NewRoundDecimalExprFromParams(types.FunctionBuildContext{}, types.FunctionConfig{Params: map[string]*schemapb.FunctionParamValue{"decimal": doubleParam(1.5)}})
s.ErrorContains(err, "must be an integer")
_, err = NewRoundDecimalExprFromParams(types.FunctionBuildContext{}, types.FunctionConfig{Params: map[string]*schemapb.FunctionParamValue{"decimal": intParam(7)}})
s.ErrorContains(err, "decimal must be in range")
}
func (s *RoundDecimalExprTestSuite) TestExecute_RoundTo2Decimals() {
expr, err := NewRoundDecimalExpr(2)
s.Require().NoError(err)
ctx := types.NewFuncContext(s.pool)
input := s.buildScoreInput([]float32{0.956, 0.834, 0.715, 0.123, 0.999})
defer input.Release()
results, err := expr.Execute(ctx, []*arrow.Chunked{input})
s.Require().NoError(err)
s.Len(results, 1)
defer results[0].Release()
chunk := results[0].Chunk(0).(*array.Float32)
s.Equal(5, chunk.Len())
s.InDelta(float32(0.96), chunk.Value(0), 0.001)
s.InDelta(float32(0.83), chunk.Value(1), 0.001)
// 0.715 as float32 is ~0.71499997, so floor(71.499997+0.5)/100 = 0.71
s.InDelta(float32(0.71), chunk.Value(2), 0.001)
s.InDelta(float32(0.12), chunk.Value(3), 0.001)
s.InDelta(float32(1.00), chunk.Value(4), 0.001)
}
func (s *RoundDecimalExprTestSuite) TestExecute_RoundTo0Decimals() {
expr, err := NewRoundDecimalExpr(0)
s.Require().NoError(err)
ctx := types.NewFuncContext(s.pool)
input := s.buildScoreInput([]float32{0.5, 1.4, 2.5, 3.6})
defer input.Release()
results, err := expr.Execute(ctx, []*arrow.Chunked{input})
s.Require().NoError(err)
defer results[0].Release()
chunk := results[0].Chunk(0).(*array.Float32)
s.InDelta(float32(1.0), chunk.Value(0), 0.001) // 0.5 -> 1
s.InDelta(float32(1.0), chunk.Value(1), 0.001) // 1.4 -> 1
s.InDelta(float32(3.0), chunk.Value(2), 0.001) // 2.5 -> 3
s.InDelta(float32(4.0), chunk.Value(3), 0.001) // 3.6 -> 4
}
func (s *RoundDecimalExprTestSuite) TestExecute_MultipleChunks() {
expr, err := NewRoundDecimalExpr(1)
s.Require().NoError(err)
ctx := types.NewFuncContext(s.pool)
b1 := array.NewFloat32Builder(s.pool)
b1.AppendValues([]float32{0.96, 0.84}, nil)
arr1 := b1.NewArray()
b1.Release()
b2 := array.NewFloat32Builder(s.pool)
b2.AppendValues([]float32{0.71}, nil)
arr2 := b2.NewArray()
b2.Release()
input := arrow.NewChunked(arrow.PrimitiveTypes.Float32, []arrow.Array{arr1, arr2})
arr1.Release()
arr2.Release()
defer input.Release()
results, err := expr.Execute(ctx, []*arrow.Chunked{input})
s.Require().NoError(err)
defer results[0].Release()
s.Equal(2, len(results[0].Chunks()))
chunk0 := results[0].Chunk(0).(*array.Float32)
s.InDelta(float32(1.0), chunk0.Value(0), 0.001)
s.InDelta(float32(0.8), chunk0.Value(1), 0.001)
chunk1 := results[0].Chunk(1).(*array.Float32)
s.InDelta(float32(0.7), chunk1.Value(0), 0.001)
}
func (s *RoundDecimalExprTestSuite) TestExecute_WrongInputCount() {
expr, err := NewRoundDecimalExpr(2)
s.Require().NoError(err)
ctx := types.NewFuncContext(s.pool)
input1 := s.buildScoreInput([]float32{0.5})
defer input1.Release()
input2 := s.buildScoreInput([]float32{0.5})
defer input2.Release()
_, err = expr.Execute(ctx, []*arrow.Chunked{input1, input2})
s.Error(err)
s.Contains(err.Error(), "expected 1 input")
}
func (s *RoundDecimalExprTestSuite) TestExecute_MatchesOldFormula() {
// Verify the formula matches the old rerank implementation:
// score = floor(score * 10^decimal + 0.5) / 10^decimal
expr, err := NewRoundDecimalExpr(3)
s.Require().NoError(err)
ctx := types.NewFuncContext(s.pool)
input := s.buildScoreInput([]float32{0.12345, 0.98765, 0.55555})
defer input.Release()
results, err := expr.Execute(ctx, []*arrow.Chunked{input})
s.Require().NoError(err)
defer results[0].Release()
chunk := results[0].Chunk(0).(*array.Float32)
s.InDelta(float32(0.123), chunk.Value(0), 0.0001)
s.InDelta(float32(0.988), chunk.Value(1), 0.0001)
s.InDelta(float32(0.556), chunk.Value(2), 0.0001)
}