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
712 lines
21 KiB
Go
712 lines
21 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 (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"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"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
)
|
|
|
|
// =============================================================================
|
|
// Mock ModelProvider
|
|
// =============================================================================
|
|
|
|
type mockModelProvider struct {
|
|
maxBatch int
|
|
rerankFunc func(ctx context.Context, query string, docs []string) ([]float32, error)
|
|
}
|
|
|
|
func (m *mockModelProvider) Rerank(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
if m.rerankFunc != nil {
|
|
return m.rerankFunc(ctx, query, docs)
|
|
}
|
|
// Default: return sequential scores 0.0, 1.0, 2.0, ...
|
|
scores := make([]float32, len(docs))
|
|
for i := range docs {
|
|
scores[i] = float32(i)
|
|
}
|
|
return scores, nil
|
|
}
|
|
|
|
func (m *mockModelProvider) MaxBatch() int {
|
|
return m.maxBatch
|
|
}
|
|
|
|
// =============================================================================
|
|
// Test Suite
|
|
// =============================================================================
|
|
|
|
type RerankModelExprTestSuite struct {
|
|
suite.Suite
|
|
pool *memory.CheckedAllocator
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) SetupTest() {
|
|
s.pool = memory.NewCheckedAllocator(memory.NewGoAllocator())
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TearDownTest() {
|
|
s.pool.AssertSize(s.T(), 0)
|
|
}
|
|
|
|
func TestRerankModelExprTestSuite(t *testing.T) {
|
|
suite.Run(t, new(RerankModelExprTestSuite))
|
|
}
|
|
|
|
// =============================================================================
|
|
// Helper Functions
|
|
// =============================================================================
|
|
|
|
func strPtr(str string) *string {
|
|
return &str
|
|
}
|
|
|
|
// createStringChunked creates a String chunked array from the given slices.
|
|
// Each slice becomes a separate chunk in the chunked array.
|
|
func (s *RerankModelExprTestSuite) createStringChunked(chunks ...[]string) *arrow.Chunked {
|
|
arrs := make([]arrow.Array, len(chunks))
|
|
for i, chunk := range chunks {
|
|
builder := array.NewStringBuilder(s.pool)
|
|
for _, v := range chunk {
|
|
builder.Append(v)
|
|
}
|
|
arrs[i] = builder.NewArray()
|
|
builder.Release()
|
|
}
|
|
chunked := arrow.NewChunked(arrow.BinaryTypes.String, arrs)
|
|
for _, a := range arrs {
|
|
a.Release()
|
|
}
|
|
return chunked
|
|
}
|
|
|
|
// createStringChunkedWithNulls creates a String chunked array with null support.
|
|
// nil pointers represent null values.
|
|
func (s *RerankModelExprTestSuite) createStringChunkedWithNulls(chunks ...[]*string) *arrow.Chunked {
|
|
arrs := make([]arrow.Array, len(chunks))
|
|
for i, chunk := range chunks {
|
|
builder := array.NewStringBuilder(s.pool)
|
|
for _, v := range chunk {
|
|
if v == nil {
|
|
builder.AppendNull()
|
|
} else {
|
|
builder.Append(*v)
|
|
}
|
|
}
|
|
arrs[i] = builder.NewArray()
|
|
builder.Release()
|
|
}
|
|
chunked := arrow.NewChunked(arrow.BinaryTypes.String, arrs)
|
|
for _, a := range arrs {
|
|
a.Release()
|
|
}
|
|
return chunked
|
|
}
|
|
|
|
// =============================================================================
|
|
// Constructor Tests
|
|
// =============================================================================
|
|
|
|
func (s *RerankModelExprTestSuite) TestNewRerankModelExpr_Valid() {
|
|
provider := &mockModelProvider{maxBatch: 10}
|
|
expr, err := NewRerankModelExpr(provider, []string{"query1", "query2"})
|
|
s.Require().NoError(err)
|
|
s.NotNil(expr)
|
|
s.Equal(RerankModelFuncName, expr.Name())
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestNewRerankModelExpr_NilProvider() {
|
|
_, err := NewRerankModelExpr(nil, []string{"query1"})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "provider is nil")
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestNewRerankModelExpr_EmptyQueries() {
|
|
provider := &mockModelProvider{maxBatch: 10}
|
|
_, err := NewRerankModelExpr(provider, []string{})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "queries must not be empty")
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestNewRerankModelExpr_NilQueries() {
|
|
provider := &mockModelProvider{maxBatch: 10}
|
|
_, err := NewRerankModelExpr(provider, nil)
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "queries must not be empty")
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestNewRerankModelExprFromParams() {
|
|
paramtable.Init()
|
|
paramtable.Get().FunctionCfg.RerankModelProviders.GetFunc = func() map[string]string {
|
|
return map[string]string{}
|
|
}
|
|
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"results": [{"index": 0, "relevance_score": 0.1}]}`))
|
|
}))
|
|
defer ts.Close()
|
|
|
|
fn, err := NewRerankModelExprFromParams(types.FunctionBuildContext{}, types.FunctionConfig{Params: map[string]*schemapb.FunctionParamValue{
|
|
"provider": stringParam("vllm"),
|
|
"endpoint": stringParam(ts.URL),
|
|
"queries": arrayParam(stringParam("query")),
|
|
}})
|
|
s.Require().NoError(err)
|
|
s.Equal(RerankModelFuncName, fn.Name())
|
|
|
|
_, err = NewRerankModelExprFromParams(types.FunctionBuildContext{}, types.FunctionConfig{Params: map[string]*schemapb.FunctionParamValue{
|
|
"provider": stringParam("vllm"),
|
|
"endpoint": stringParam(ts.URL),
|
|
}})
|
|
s.ErrorContains(err, "missing required parameter")
|
|
|
|
_, err = NewRerankModelExprFromParams(types.FunctionBuildContext{}, types.FunctionConfig{Params: map[string]*schemapb.FunctionParamValue{
|
|
"provider": stringParam("vllm"),
|
|
"endpoint": stringParam(ts.URL),
|
|
"queries": arrayParam(stringParam("query")),
|
|
"invalid": arrayParam(stringParam("value")),
|
|
}})
|
|
s.ErrorContains(err, "must be scalar")
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestModelProviderParamsSkipsQueries() {
|
|
params, err := modelProviderParams(map[string]*schemapb.FunctionParamValue{
|
|
"provider": stringParam("vllm"),
|
|
"enabled": boolParam(true),
|
|
"batch": intParam(8),
|
|
"ratio": doubleParam(0.5),
|
|
"queries": arrayParam(stringParam("query")),
|
|
})
|
|
s.Require().NoError(err)
|
|
|
|
values := make(map[string]string, len(params))
|
|
for _, param := range params {
|
|
values[param.Key] = param.Value
|
|
}
|
|
s.Equal("vllm", values["provider"])
|
|
s.Equal("true", values["enabled"])
|
|
s.Equal("8", values["batch"])
|
|
s.Equal("0.5", values["ratio"])
|
|
s.NotContains(values, "queries")
|
|
}
|
|
|
|
// =============================================================================
|
|
// Interface Method Tests
|
|
// =============================================================================
|
|
|
|
func (s *RerankModelExprTestSuite) TestOutputDataTypes() {
|
|
provider := &mockModelProvider{maxBatch: 10}
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
outputTypes := expr.OutputDataTypes()
|
|
s.Len(outputTypes, 1)
|
|
s.Equal(arrow.PrimitiveTypes.Float32, outputTypes[0])
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestIsRunnable() {
|
|
provider := &mockModelProvider{maxBatch: 10}
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
// Only L2Rerank stage
|
|
s.True(expr.IsRunnable(types.StageL2Rerank))
|
|
s.False(expr.IsRunnable(types.StageL1Rerank))
|
|
s.False(expr.IsRunnable(types.StageL0Rerank))
|
|
s.False(expr.IsRunnable(types.StageIngestion))
|
|
s.False(expr.IsRunnable(types.StagePreProcess))
|
|
s.False(expr.IsRunnable(types.StagePostProcess))
|
|
s.False(expr.IsRunnable("unknown_stage"))
|
|
}
|
|
|
|
// =============================================================================
|
|
// Execute Tests
|
|
// =============================================================================
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_SingleChunk() {
|
|
provider := &mockModelProvider{
|
|
maxBatch: 100,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
s.Equal("test query", query)
|
|
scores := make([]float32, len(docs))
|
|
for i := range docs {
|
|
scores[i] = float32(i) * 0.1
|
|
}
|
|
return scores, nil
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"test query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked([]string{"doc1", "doc2", "doc3"})
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
outputs, err := expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Require().NoError(err)
|
|
s.Len(outputs, 1)
|
|
defer outputs[0].Release()
|
|
|
|
s.Equal(1, len(outputs[0].Chunks()))
|
|
result := outputs[0].Chunk(0).(*array.Float32)
|
|
s.Equal(3, result.Len())
|
|
s.InDelta(0.0, float64(result.Value(0)), 0.001)
|
|
s.InDelta(0.1, float64(result.Value(1)), 0.001)
|
|
s.InDelta(0.2, float64(result.Value(2)), 0.001)
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_MultipleChunks() {
|
|
callCount := 0
|
|
provider := &mockModelProvider{
|
|
maxBatch: 100,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
callCount++
|
|
scores := make([]float32, len(docs))
|
|
switch query {
|
|
case "query1":
|
|
for i := range docs {
|
|
scores[i] = float32(i) + 1.0
|
|
}
|
|
case "query2":
|
|
for i := range docs {
|
|
scores[i] = float32(i) + 10.0
|
|
}
|
|
}
|
|
return scores, nil
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"query1", "query2"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked(
|
|
[]string{"docA", "docB"},
|
|
[]string{"docC", "docD", "docE"},
|
|
)
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
outputs, err := expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Require().NoError(err)
|
|
s.Len(outputs, 1)
|
|
defer outputs[0].Release()
|
|
|
|
s.Equal(2, callCount)
|
|
s.Equal(2, len(outputs[0].Chunks()))
|
|
|
|
// Check chunk 0 (query1)
|
|
chunk0 := outputs[0].Chunk(0).(*array.Float32)
|
|
s.Equal(2, chunk0.Len())
|
|
s.InDelta(1.0, float64(chunk0.Value(0)), 0.001)
|
|
s.InDelta(2.0, float64(chunk0.Value(1)), 0.001)
|
|
|
|
// Check chunk 1 (query2)
|
|
chunk1 := outputs[0].Chunk(1).(*array.Float32)
|
|
s.Equal(3, chunk1.Len())
|
|
s.InDelta(10.0, float64(chunk1.Value(0)), 0.001)
|
|
s.InDelta(11.0, float64(chunk1.Value(1)), 0.001)
|
|
s.InDelta(12.0, float64(chunk1.Value(2)), 0.001)
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_EmptyChunk() {
|
|
provider := &mockModelProvider{maxBatch: 100}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked([]string{})
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
outputs, err := expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Require().NoError(err)
|
|
s.Len(outputs, 1)
|
|
defer outputs[0].Release()
|
|
|
|
s.Equal(0, outputs[0].Len())
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_WithNulls() {
|
|
var capturedDocs []string
|
|
provider := &mockModelProvider{
|
|
maxBatch: 100,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
capturedDocs = docs
|
|
scores := make([]float32, len(docs))
|
|
for i := range docs {
|
|
scores[i] = float32(i) * 0.5
|
|
}
|
|
return scores, nil
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunkedWithNulls([]*string{
|
|
strPtr("hello"),
|
|
nil,
|
|
strPtr("world"),
|
|
nil,
|
|
})
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
outputs, err := expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Require().NoError(err)
|
|
s.Len(outputs, 1)
|
|
defer outputs[0].Release()
|
|
|
|
// Nulls should be treated as empty strings
|
|
s.Require().Len(capturedDocs, 4)
|
|
s.Equal("hello", capturedDocs[0])
|
|
s.Equal("", capturedDocs[1])
|
|
s.Equal("world", capturedDocs[2])
|
|
s.Equal("", capturedDocs[3])
|
|
|
|
result := outputs[0].Chunk(0).(*array.Float32)
|
|
s.Equal(4, result.Len())
|
|
}
|
|
|
|
// =============================================================================
|
|
// Error Tests
|
|
// =============================================================================
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_WrongInputCount() {
|
|
provider := &mockModelProvider{maxBatch: 100}
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
|
|
// 0 inputs
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "expected 1 input column")
|
|
|
|
// 2 inputs
|
|
textCol1 := s.createStringChunked([]string{"a"})
|
|
defer textCol1.Release()
|
|
textCol2 := s.createStringChunked([]string{"b"})
|
|
defer textCol2.Release()
|
|
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{textCol1, textCol2})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "expected 1 input column")
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_QueryCountMismatch() {
|
|
provider := &mockModelProvider{maxBatch: 100}
|
|
// 1 query but 2 chunks
|
|
expr, err := NewRerankModelExpr(provider, []string{"only_one_query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked(
|
|
[]string{"doc1"},
|
|
[]string{"doc2"},
|
|
)
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Error(err)
|
|
s.ErrorIs(err, merr.ErrParameterInvalid)
|
|
s.Contains(err.Error(), "queries count")
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_ProviderError() {
|
|
provider := &mockModelProvider{
|
|
maxBatch: 100,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
return nil, fmt.Errorf("rerank service unavailable")
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked([]string{"doc1", "doc2"})
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "rerank service unavailable")
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_ProviderErrorOnSecondChunk() {
|
|
// First chunk succeeds, second chunk fails.
|
|
// The arrays from the first chunk should be properly released (no leak).
|
|
callCount := 0
|
|
provider := &mockModelProvider{
|
|
maxBatch: 100,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
callCount++
|
|
if callCount == 1 {
|
|
scores := make([]float32, len(docs))
|
|
for i := range docs {
|
|
scores[i] = float32(i)
|
|
}
|
|
return scores, nil
|
|
}
|
|
return nil, fmt.Errorf("second chunk failed")
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"q1", "q2"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked(
|
|
[]string{"doc1", "doc2"},
|
|
[]string{"doc3"},
|
|
)
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "second chunk failed")
|
|
// CheckedAllocator in TearDownTest ensures no memory leaks
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_WrongScoreCount() {
|
|
provider := &mockModelProvider{
|
|
maxBatch: 100,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
// Return fewer scores than docs
|
|
return []float32{0.5}, nil
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked([]string{"doc1", "doc2", "doc3"})
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "returned 1 scores for 3 docs")
|
|
}
|
|
|
|
// =============================================================================
|
|
// Batching Tests
|
|
// =============================================================================
|
|
|
|
func (s *RerankModelExprTestSuite) TestBatching_MaxBatch2With5Docs() {
|
|
batchCalls := 0
|
|
var batchSizes []int
|
|
provider := &mockModelProvider{
|
|
maxBatch: 2,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
batchCalls++
|
|
batchSizes = append(batchSizes, len(docs))
|
|
scores := make([]float32, len(docs))
|
|
for i := range docs {
|
|
scores[i] = float32(batchCalls*10 + i)
|
|
}
|
|
return scores, nil
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked([]string{"d1", "d2", "d3", "d4", "d5"})
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
outputs, err := expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Require().NoError(err)
|
|
s.Len(outputs, 1)
|
|
defer outputs[0].Release()
|
|
|
|
// 5 docs, maxBatch=2 -> batches of 2, 2, 1 = 3 calls
|
|
s.Equal(3, batchCalls)
|
|
s.Equal([]int{2, 2, 1}, batchSizes)
|
|
|
|
result := outputs[0].Chunk(0).(*array.Float32)
|
|
s.Equal(5, result.Len())
|
|
// Batch 1: scores 10, 11; Batch 2: scores 20, 21; Batch 3: scores 30
|
|
s.InDelta(10.0, float64(result.Value(0)), 0.001)
|
|
s.InDelta(11.0, float64(result.Value(1)), 0.001)
|
|
s.InDelta(20.0, float64(result.Value(2)), 0.001)
|
|
s.InDelta(21.0, float64(result.Value(3)), 0.001)
|
|
s.InDelta(30.0, float64(result.Value(4)), 0.001)
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestBatching_ZeroMaxBatch() {
|
|
batchCalls := 0
|
|
provider := &mockModelProvider{
|
|
maxBatch: 0, // should process all at once
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
batchCalls++
|
|
scores := make([]float32, len(docs))
|
|
for i := range docs {
|
|
scores[i] = float32(i) * 0.1
|
|
}
|
|
return scores, nil
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked([]string{"d1", "d2", "d3", "d4", "d5"})
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
outputs, err := expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Require().NoError(err)
|
|
s.Len(outputs, 1)
|
|
defer outputs[0].Release()
|
|
|
|
// maxBatch=0 means all docs in one call
|
|
s.Equal(1, batchCalls)
|
|
|
|
result := outputs[0].Chunk(0).(*array.Float32)
|
|
s.Equal(5, result.Len())
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestBatching_ErrorMidway() {
|
|
batchCalls := 0
|
|
provider := &mockModelProvider{
|
|
maxBatch: 2,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
batchCalls++
|
|
if batchCalls == 2 {
|
|
return nil, fmt.Errorf("batch 2 failed")
|
|
}
|
|
scores := make([]float32, len(docs))
|
|
return scores, nil
|
|
},
|
|
}
|
|
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked([]string{"d1", "d2", "d3", "d4", "d5"})
|
|
defer textCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "batch 2 failed")
|
|
// CheckedAllocator ensures no leaks
|
|
}
|
|
|
|
// =============================================================================
|
|
// Memory Leak Tests
|
|
// =============================================================================
|
|
|
|
func (s *RerankModelExprTestSuite) TestMemoryLeak_RepeatedExecution() {
|
|
provider := &mockModelProvider{
|
|
maxBatch: 100,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
scores := make([]float32, len(docs))
|
|
for i := range docs {
|
|
scores[i] = float32(i) * 0.1
|
|
}
|
|
return scores, nil
|
|
},
|
|
}
|
|
|
|
for range 10 {
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked([]string{"doc1", "doc2", "doc3"})
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
outputs, err := expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Require().NoError(err)
|
|
|
|
outputs[0].Release()
|
|
textCol.Release()
|
|
}
|
|
// Memory leak check happens in TearDownTest via CheckedAllocator
|
|
}
|
|
|
|
func (s *RerankModelExprTestSuite) TestMemoryLeak_OnError() {
|
|
provider := &mockModelProvider{
|
|
maxBatch: 100,
|
|
rerankFunc: func(ctx context.Context, query string, docs []string) ([]float32, error) {
|
|
return nil, fmt.Errorf("intentional error")
|
|
},
|
|
}
|
|
|
|
for range 10 {
|
|
expr, err := NewRerankModelExpr(provider, []string{"q1", "q2"})
|
|
s.Require().NoError(err)
|
|
|
|
textCol := s.createStringChunked(
|
|
[]string{"doc1", "doc2"},
|
|
[]string{"doc3"},
|
|
)
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{textCol})
|
|
s.Error(err)
|
|
|
|
textCol.Release()
|
|
}
|
|
// Memory leak check happens in TearDownTest
|
|
}
|
|
|
|
// =============================================================================
|
|
// Non-String Input Test
|
|
// =============================================================================
|
|
|
|
func (s *RerankModelExprTestSuite) TestExecute_NonStringInput() {
|
|
provider := &mockModelProvider{maxBatch: 100}
|
|
expr, err := NewRerankModelExpr(provider, []string{"query"})
|
|
s.Require().NoError(err)
|
|
|
|
// Create an Int64 chunked array instead of String
|
|
builder := array.NewInt64Builder(s.pool)
|
|
builder.AppendValues([]int64{1, 2, 3}, nil)
|
|
arr := builder.NewArray()
|
|
builder.Release()
|
|
|
|
intCol := arrow.NewChunked(arrow.PrimitiveTypes.Int64, []arrow.Array{arr})
|
|
arr.Release()
|
|
defer intCol.Release()
|
|
|
|
ctx := types.NewFuncContext(s.pool)
|
|
_, err = expr.Execute(ctx, []*arrow.Chunked{intCol})
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "must be String/VarChar")
|
|
}
|