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
203 lines
6.7 KiB
Go
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)
|
|
}
|