1- /*
2- * # Licensed to the LF AI & Data foundation under one
3- * # or more contributor license agreements. See the NOTICE file
4- * # distributed with this work for additional information
5- * # regarding copyright ownership. The ASF licenses this file
6- * # to you under the Apache License, Version 2.0 (the
7- * # "License"); you may not use this file except in compliance
8- * # with the License. You may obtain a copy of the License at
9- * #
10- * # http://www.apache.org/licenses/LICENSE-2.0
11- * #
12- * # Unless required by applicable law or agreed to in writing, software
13- * # distributed under the License is distributed on an "AS IS" BASIS,
14- * # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15- * # See the License for the specific language governing permissions and
16- * # limitations under the License.
17- */
18-
191package rerank
202
213import (
224 "context"
235 "fmt"
6+ "math"
247 "strings"
258
269 "github.com/expr-lang/expr"
@@ -44,6 +27,70 @@ const (
4427 ExprNormalizeKey = "normalize"
4528)
4629
30+ // toFloat64 converts various numeric types to float64
31+ func toFloat64 (v interface {}) float64 {
32+ switch val := v .(type ) {
33+ case float64 :
34+ return val
35+ case float32 :
36+ return float64 (val )
37+ case int :
38+ return float64 (val )
39+ case int32 :
40+ return float64 (val )
41+ case int64 :
42+ return float64 (val )
43+ default :
44+ return 0.0 // fallback for unsupported types
45+ }
46+ }
47+
48+ // mathFunctions contains all mathematical functions for expr-lang
49+ // Created once at package level to avoid recreating functions on every call
50+ var mathFunctions = map [string ]interface {}{
51+ // Basic math functions with type conversion
52+ "abs" : func (x interface {}) float64 { return math .Abs (toFloat64 (x )) },
53+ "ceil" : func (x interface {}) float64 { return math .Ceil (toFloat64 (x )) },
54+ "floor" : func (x interface {}) float64 { return math .Floor (toFloat64 (x )) },
55+ "round" : func (x interface {}) float64 { return math .Round (toFloat64 (x )) },
56+ "sqrt" : func (x interface {}) float64 { return math .Sqrt (toFloat64 (x )) },
57+ "pow" : func (x , y interface {}) float64 { return math .Pow (toFloat64 (x ), toFloat64 (y )) },
58+
59+ // Exponential and logarithmic functions with type conversion
60+ "exp" : func (x interface {}) float64 { return math .Exp (toFloat64 (x )) },
61+ "exp2" : func (x interface {}) float64 { return math .Exp2 (toFloat64 (x )) },
62+ "log" : func (x interface {}) float64 { return math .Log (toFloat64 (x )) },
63+ "log10" : func (x interface {}) float64 { return math .Log10 (toFloat64 (x )) },
64+ "log2" : func (x interface {}) float64 { return math .Log2 (toFloat64 (x )) },
65+
66+ // Min/Max functions with type conversion
67+ "min" : func (x , y interface {}) float64 { return math .Min (toFloat64 (x ), toFloat64 (y )) },
68+ "max" : func (x , y interface {}) float64 { return math .Max (toFloat64 (x ), toFloat64 (y )) },
69+
70+ // Utility functions with type conversion
71+ "mod" : func (x , y interface {}) float64 { return math .Mod (toFloat64 (x ), toFloat64 (y )) },
72+ "remainder" : func (x , y interface {}) float64 { return math .Remainder (toFloat64 (x ), toFloat64 (y )) },
73+ "trunc" : func (x interface {}) float64 { return math .Trunc (toFloat64 (x )) },
74+
75+ // Clamping function with type conversion
76+ "clamp" : func (x , min_val , max_val interface {}) float64 {
77+ xf := toFloat64 (x )
78+ minf := toFloat64 (min_val )
79+ maxf := toFloat64 (max_val )
80+ if xf < minf {
81+ return minf
82+ }
83+ if xf > maxf {
84+ return maxf
85+ }
86+ return xf
87+ },
88+
89+ // Mathematical constants
90+ "PI" : math .Pi ,
91+ "E" : math .E ,
92+ }
93+
4794func newExprRerank (collSchema * schemapb.CollectionSchema , funcSchema * schemapb.FunctionSchema ) (Reranker , error ) {
4895 base , err := newRerankBase (collSchema , funcSchema , ExprRerankName , false )
4996 if err != nil {
@@ -70,12 +117,21 @@ func newExprRerank(collSchema *schemapb.CollectionSchema, funcSchema *schemapb.F
70117 return nil , fmt .Errorf ("expr rerank requires %s parameter" , ExprCodeKey )
71118 }
72119
120+ // Create environment with mathematical functions
121+ // Pre-size map to avoid growth during insertion: 3 core variables + len(mathFunctions)
122+ env := make (map [string ]interface {}, 3 + len (mathFunctions ))
123+
124+ env ["score" ] = float32 (0 )
125+ env ["rank" ] = int (0 )
126+ env ["fields" ] = map [string ]interface {}{}
127+
128+ // Add mathematical functions to the environment (just copying references)
129+ for name , fn := range mathFunctions {
130+ env [name ] = fn
131+ }
132+
73133 // Compile the expression
74- program , err := expr .Compile (exprCode , expr .Env (map [string ]interface {}{
75- "score" : float32 (0 ),
76- "rank" : int (0 ),
77- "fields" : map [string ]interface {}{},
78- }))
134+ program , err := expr .Compile (exprCode , expr .Env (env ))
79135 if err != nil {
80136 return nil , fmt .Errorf ("failed to compile expression: %w" , err )
81137 }
@@ -95,7 +151,6 @@ func newExprRerank(collSchema *schemapb.CollectionSchema, funcSchema *schemapb.F
95151 needNormalize : needNormalize ,
96152 }, nil
97153}
98-
99154func (e * ExprRerank [T ]) processOneSearchData (ctx context.Context , searchParams * SearchParams , cols []* columns , idGroup map [any ]any ) (* IDScores [T ], error ) {
100155 newScores := map [T ]float32 {}
101156 idLocations := make (map [T ]IDLoc )
@@ -113,14 +168,21 @@ func (e *ExprRerank[T]) processOneSearchData(ctx context.Context, searchParams *
113168 continue // Already processed (use first occurrence)
114169 }
115170
116- // Prepare environment for expression evaluation
117- env := map [string ]interface {}{
118- "score" : scores [idx ],
119- "rank" : idx ,
171+ // Prepare environment with dynamic values and math functions
172+ // Pre-size map to avoid growth during insertion: 3 core variables + len(mathFunctions)
173+ env := make (map [string ]interface {}, 3 + len (mathFunctions ))
174+
175+ env ["score" ] = scores [idx ]
176+ env ["rank" ] = idx
177+
178+ // Add mathematical functions to the runtime environment (just copying references)
179+ for name , fn := range mathFunctions {
180+ env [name ] = fn
120181 }
121182
122183 // Add field values to environment
123- fields := make (map [string ]interface {})
184+ // Pre-size fields map for typical field count
185+ fields := make (map [string ]interface {}, len (e .inputFieldNames ))
124186 for fieldIdx , fieldName := range e .inputFieldNames {
125187 if fieldIdx < len (col .data ) {
126188 switch data := col .data [fieldIdx ].(type ) {
@@ -159,13 +221,13 @@ func (e *ExprRerank[T]) processOneSearchData(ctx context.Context, searchParams *
159221 return nil , fmt .Errorf ("failed to execute expression for id %v: %w" , id , err )
160222 }
161223
162- // Convert output to float32
224+ // Convert output to float32 (fast path for common types)
163225 var newScore float32
164226 switch v := output .(type ) {
165- case float32 :
166- newScore = v
167227 case float64 :
168228 newScore = float32 (v )
229+ case float32 :
230+ newScore = v
169231 case int :
170232 newScore = float32 (v )
171233 case int64 :
0 commit comments