Files
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

628 lines
20 KiB
Go

package agg
import (
"encoding/binary"
"fmt"
"hash"
"hash/fnv"
"math"
"unsafe"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
func NewFieldAccessor(fieldType schemapb.DataType) (FieldAccessor, error) {
switch fieldType {
case schemapb.DataType_Bool:
return newBoolFieldAccessor(), nil
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
return newInt32FieldAccessor(), nil
case schemapb.DataType_Int64:
return newInt64FieldAccessor(), nil
case schemapb.DataType_Timestamptz:
return newTimestamptzFieldAccessor(), nil
case schemapb.DataType_VarChar, schemapb.DataType_String:
return newStringFieldAccessor(), nil
case schemapb.DataType_Float:
return newFloat32FieldAccessor(), nil
case schemapb.DataType_Double:
return newFloat64FieldAccessor(), nil
default:
return nil, merr.WrapErrParameterInvalidMsg("unsupported data type for hasher")
}
}
type FieldAccessor interface {
Hash(idx int) uint64
ValAt(idx int) interface{}
IsNullAt(idx int) bool
SetVals(fieldData *schemapb.FieldData)
RowCount() int
}
// Special hash value for null - using a prime number unlikely to collide
const nullHashValue uint64 = 0x9E3779B97F4A7C15
type Int32FieldAccessor struct {
vals []int32
validData []bool
hasher hash.Hash64
buffer []byte
}
func (i32Field *Int32FieldAccessor) Hash(idx int) uint64 {
if idx < 0 || idx >= len(i32Field.vals) {
panic(fmt.Sprintf("Int32FieldAccessor.Hash: index %d out of range [0,%d)", idx, len(i32Field.vals)))
}
if i32Field.IsNullAt(idx) {
return nullHashValue
}
i32Field.hasher.Reset()
val := i32Field.vals[idx]
binary.LittleEndian.PutUint32(i32Field.buffer, uint32(val))
i32Field.hasher.Write(i32Field.buffer)
ret := i32Field.hasher.Sum64()
return ret
}
func (i32Field *Int32FieldAccessor) SetVals(fieldData *schemapb.FieldData) {
i32Field.vals = fieldData.GetScalars().GetIntData().GetData()
i32Field.validData = fieldData.GetValidData()
}
func (i32Field *Int32FieldAccessor) RowCount() int {
return len(i32Field.vals)
}
func (i32Field *Int32FieldAccessor) ValAt(idx int) interface{} {
return i32Field.vals[idx]
}
func (i32Field *Int32FieldAccessor) IsNullAt(idx int) bool {
if len(i32Field.validData) == 0 {
return false // No validity data means all values are valid
}
return !i32Field.validData[idx]
}
func newInt32FieldAccessor() FieldAccessor {
return &Int32FieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 4)}
}
type Int64FieldAccessor struct {
vals []int64
validData []bool
hasher hash.Hash64
buffer []byte
}
func (i64Field *Int64FieldAccessor) Hash(idx int) uint64 {
if idx < 0 || idx >= len(i64Field.vals) {
panic(fmt.Sprintf("Int64FieldAccessor.Hash: index %d out of range [0,%d)", idx, len(i64Field.vals)))
}
if i64Field.IsNullAt(idx) {
return nullHashValue
}
i64Field.hasher.Reset()
val := i64Field.vals[idx]
binary.LittleEndian.PutUint64(i64Field.buffer, uint64(val))
i64Field.hasher.Write(i64Field.buffer)
return i64Field.hasher.Sum64()
}
func (i64Field *Int64FieldAccessor) SetVals(fieldData *schemapb.FieldData) {
i64Field.vals = fieldData.GetScalars().GetLongData().GetData()
i64Field.validData = fieldData.GetValidData()
}
func (i64Field *Int64FieldAccessor) RowCount() int {
return len(i64Field.vals)
}
func (i64Field *Int64FieldAccessor) ValAt(idx int) interface{} {
return i64Field.vals[idx]
}
func (i64Field *Int64FieldAccessor) IsNullAt(idx int) bool {
if len(i64Field.validData) == 0 {
return false
}
return !i64Field.validData[idx]
}
func newInt64FieldAccessor() FieldAccessor {
return &Int64FieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 8)}
}
type TimestamptzFieldAccessor struct {
vals []int64
validData []bool
hasher hash.Hash64
buffer []byte
}
func (tzField *TimestamptzFieldAccessor) Hash(idx int) uint64 {
if idx < 0 || idx >= len(tzField.vals) {
panic(fmt.Sprintf("TimestamptzFieldAccessor.Hash: index %d out of range [0,%d)", idx, len(tzField.vals)))
}
if tzField.IsNullAt(idx) {
return nullHashValue
}
tzField.hasher.Reset()
val := tzField.vals[idx]
binary.LittleEndian.PutUint64(tzField.buffer, uint64(val))
tzField.hasher.Write(tzField.buffer)
return tzField.hasher.Sum64()
}
func (tzField *TimestamptzFieldAccessor) SetVals(fieldData *schemapb.FieldData) {
tzField.vals = fieldData.GetScalars().GetTimestamptzData().GetData()
tzField.validData = fieldData.GetValidData()
}
func (tzField *TimestamptzFieldAccessor) RowCount() int {
return len(tzField.vals)
}
func (tzField *TimestamptzFieldAccessor) ValAt(idx int) interface{} {
return tzField.vals[idx]
}
func (tzField *TimestamptzFieldAccessor) IsNullAt(idx int) bool {
if len(tzField.validData) == 0 {
return false
}
return !tzField.validData[idx]
}
func newTimestamptzFieldAccessor() FieldAccessor {
return &TimestamptzFieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 8)}
}
// BoolFieldAccessor
type BoolFieldAccessor struct {
vals []bool
validData []bool
hasher hash.Hash64
buffer []byte
}
func (boolField *BoolFieldAccessor) Hash(idx int) uint64 {
if idx < 0 || idx >= len(boolField.vals) {
panic(fmt.Sprintf("BoolFieldAccessor.Hash: index %d out of range [0,%d)", idx, len(boolField.vals)))
}
if boolField.IsNullAt(idx) {
return nullHashValue
}
boolField.hasher.Reset()
val := boolField.vals[idx]
if val {
boolField.buffer[0] = 1
} else {
boolField.buffer[0] = 0
}
boolField.hasher.Write(boolField.buffer[:1])
return boolField.hasher.Sum64()
}
func (boolField *BoolFieldAccessor) SetVals(fieldData *schemapb.FieldData) {
boolField.vals = fieldData.GetScalars().GetBoolData().GetData()
boolField.validData = fieldData.GetValidData()
}
func (boolField *BoolFieldAccessor) RowCount() int {
return len(boolField.vals)
}
func (boolField *BoolFieldAccessor) ValAt(idx int) interface{} {
return boolField.vals[idx]
}
func (boolField *BoolFieldAccessor) IsNullAt(idx int) bool {
if len(boolField.validData) == 0 {
return false
}
return !boolField.validData[idx]
}
func newBoolFieldAccessor() FieldAccessor {
return &BoolFieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 1)}
}
// Float32FieldAccessor
type Float32FieldAccessor struct {
vals []float32
validData []bool
hasher hash.Hash64
buffer []byte
}
func (f32FieldAccessor *Float32FieldAccessor) Hash(idx int) uint64 {
if idx < 0 || idx >= len(f32FieldAccessor.vals) {
panic(fmt.Sprintf("Float32FieldAccessor.Hash: index %d out of range [0,%d)", idx, len(f32FieldAccessor.vals)))
}
if f32FieldAccessor.IsNullAt(idx) {
return nullHashValue
}
f32FieldAccessor.hasher.Reset()
val := f32FieldAccessor.vals[idx]
binary.LittleEndian.PutUint32(f32FieldAccessor.buffer, math.Float32bits(val))
f32FieldAccessor.hasher.Write(f32FieldAccessor.buffer[:4])
return f32FieldAccessor.hasher.Sum64()
}
func (f32FieldAccessor *Float32FieldAccessor) SetVals(fieldData *schemapb.FieldData) {
f32FieldAccessor.vals = fieldData.GetScalars().GetFloatData().GetData()
f32FieldAccessor.validData = fieldData.GetValidData()
}
func (f32FieldAccessor *Float32FieldAccessor) RowCount() int {
return len(f32FieldAccessor.vals)
}
func (f32FieldAccessor *Float32FieldAccessor) ValAt(idx int) interface{} {
return f32FieldAccessor.vals[idx]
}
func (f32FieldAccessor *Float32FieldAccessor) IsNullAt(idx int) bool {
if len(f32FieldAccessor.validData) == 0 {
return false
}
return !f32FieldAccessor.validData[idx]
}
func newFloat32FieldAccessor() FieldAccessor {
return &Float32FieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 4)}
}
// Float64FieldAccessor
type Float64FieldAccessor struct {
vals []float64
validData []bool
hasher hash.Hash64
buffer []byte
}
func (f64Field *Float64FieldAccessor) Hash(idx int) uint64 {
if idx < 0 || idx >= len(f64Field.vals) {
panic(fmt.Sprintf("Float64FieldAccessor.Hash: index %d out of range [0,%d)", idx, len(f64Field.vals)))
}
if f64Field.IsNullAt(idx) {
return nullHashValue
}
f64Field.hasher.Reset()
val := f64Field.vals[idx]
binary.LittleEndian.PutUint64(f64Field.buffer, math.Float64bits(val))
f64Field.hasher.Write(f64Field.buffer)
return f64Field.hasher.Sum64()
}
func (f64Field *Float64FieldAccessor) SetVals(fieldData *schemapb.FieldData) {
f64Field.vals = fieldData.GetScalars().GetDoubleData().GetData()
f64Field.validData = fieldData.GetValidData()
}
func (f64Field *Float64FieldAccessor) RowCount() int {
return len(f64Field.vals)
}
func (f64Field *Float64FieldAccessor) ValAt(idx int) interface{} {
return f64Field.vals[idx]
}
func (f64Field *Float64FieldAccessor) IsNullAt(idx int) bool {
if len(f64Field.validData) == 0 {
return false
}
return !f64Field.validData[idx]
}
func newFloat64FieldAccessor() FieldAccessor {
return &Float64FieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 8)}
}
// StringFieldAccessor
type StringFieldAccessor struct {
vals []string
validData []bool
hasher hash.Hash64
}
func (stringField *StringFieldAccessor) Hash(idx int) uint64 {
if idx < 0 || idx >= len(stringField.vals) {
panic(fmt.Sprintf("StringFieldAccessor.Hash: index %d out of range [0,%d)", idx, len(stringField.vals)))
}
if stringField.IsNullAt(idx) {
return nullHashValue
}
stringField.hasher.Reset()
val := stringField.vals[idx]
b := unsafe.Slice(unsafe.StringData(val), len(val))
stringField.hasher.Write(b)
return stringField.hasher.Sum64()
}
func (stringField *StringFieldAccessor) SetVals(fieldData *schemapb.FieldData) {
stringField.vals = fieldData.GetScalars().GetStringData().GetData()
stringField.validData = fieldData.GetValidData()
}
func (stringField *StringFieldAccessor) RowCount() int {
return len(stringField.vals)
}
func (stringField *StringFieldAccessor) ValAt(idx int) interface{} {
return stringField.vals[idx]
}
func (stringField *StringFieldAccessor) IsNullAt(idx int) bool {
if len(stringField.validData) == 0 {
return false
}
return !stringField.validData[idx]
}
func newStringFieldAccessor() FieldAccessor {
return &StringFieldAccessor{hasher: fnv.New64a()}
}
func AssembleBucket(bucket *Bucket, fieldDatas []*schemapb.FieldData) error {
colCount := len(fieldDatas)
for r := 0; r < bucket.RowCount(); r++ {
row := bucket.RowAt(r)
if err := AssembleSingleRow(colCount, row, fieldDatas); err != nil {
return err
}
}
return nil
}
func AssembleSingleRow(colCount int, row *Row, fieldDatas []*schemapb.FieldData) error {
for c := 0; c < colCount; c++ {
err := AssembleSingleValue(row.FieldValueAt(c), fieldDatas[c])
if err != nil {
return err
}
}
return nil
}
func AssembleSingleValue(fv *FieldValue, fieldData *schemapb.FieldData) error {
isNull := fv.IsNull()
// Append validity data (true = valid, false = null)
fieldData.ValidData = append(fieldData.ValidData, !isNull)
// For null values, append zero/default values to maintain array alignment
if isNull {
switch fieldData.GetType() {
case schemapb.DataType_Bool:
fieldData.GetScalars().GetBoolData().Data = append(fieldData.GetScalars().GetBoolData().GetData(), false)
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
fieldData.GetScalars().GetIntData().Data = append(fieldData.GetScalars().GetIntData().GetData(), 0)
case schemapb.DataType_Int64:
fieldData.GetScalars().GetLongData().Data = append(fieldData.GetScalars().GetLongData().GetData(), 0)
case schemapb.DataType_Timestamptz:
fieldData.GetScalars().GetTimestamptzData().Data = append(fieldData.GetScalars().GetTimestamptzData().GetData(), 0)
case schemapb.DataType_Float:
fieldData.GetScalars().GetFloatData().Data = append(fieldData.GetScalars().GetFloatData().GetData(), 0)
case schemapb.DataType_Double:
fieldData.GetScalars().GetDoubleData().Data = append(fieldData.GetScalars().GetDoubleData().GetData(), 0)
case schemapb.DataType_VarChar, schemapb.DataType_String:
fieldData.GetScalars().GetStringData().Data = append(fieldData.GetScalars().GetStringData().GetData(), "")
default:
return merr.WrapErrParameterInvalidMsg("unsupported DataType:%d", fieldData.GetType())
}
return nil
}
// For non-null values, append the actual value
val := fv.val
switch fieldData.GetType() {
case schemapb.DataType_Bool:
boolVal, ok := val.(bool)
if !ok {
return merr.WrapErrServiceInternalMsg("type assertion failed: expected bool, got %T", val)
}
fieldData.GetScalars().GetBoolData().Data = append(fieldData.GetScalars().GetBoolData().GetData(), boolVal)
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
intVal, ok := val.(int32)
if !ok {
return merr.WrapErrServiceInternalMsg("type assertion failed: expected int32, got %T", val)
}
fieldData.GetScalars().GetIntData().Data = append(fieldData.GetScalars().GetIntData().GetData(), intVal)
case schemapb.DataType_Int64:
int64Val, ok := val.(int64)
if !ok {
return merr.WrapErrServiceInternalMsg("type assertion failed: expected int64, got %T", val)
}
fieldData.GetScalars().GetLongData().Data = append(fieldData.GetScalars().GetLongData().GetData(), int64Val)
case schemapb.DataType_Timestamptz:
timestampVal, ok := val.(int64)
if !ok {
return merr.WrapErrServiceInternalMsg("type assertion failed: expected int64 for Timestamptz, got %T", val)
}
fieldData.GetScalars().GetTimestamptzData().Data = append(fieldData.GetScalars().GetTimestamptzData().GetData(), timestampVal)
case schemapb.DataType_Float:
floatVal, ok := val.(float32)
if !ok {
return merr.WrapErrServiceInternalMsg("type assertion failed: expected float32, got %T", val)
}
fieldData.GetScalars().GetFloatData().Data = append(fieldData.GetScalars().GetFloatData().GetData(), floatVal)
case schemapb.DataType_Double:
doubleVal, ok := val.(float64)
if !ok {
return merr.WrapErrServiceInternalMsg("type assertion failed: expected float64, got %T", val)
}
fieldData.GetScalars().GetDoubleData().Data = append(fieldData.GetScalars().GetDoubleData().GetData(), doubleVal)
case schemapb.DataType_VarChar, schemapb.DataType_String:
stringVal, ok := val.(string)
if !ok {
return merr.WrapErrServiceInternalMsg("type assertion failed: expected string, got %T", val)
}
fieldData.GetScalars().GetStringData().Data = append(fieldData.GetScalars().GetStringData().GetData(), stringVal)
default:
return merr.WrapErrParameterInvalidMsg("unsupported DataType:%d", fieldData.GetType())
}
return nil
}
type AggregationFieldMap struct {
userOriginalOutputFields []string
userOriginalOutputFieldIdxes [][]int // Each user output field can map to multiple field indices (e.g., avg maps to sum and count)
}
func (aggMap *AggregationFieldMap) Count() int {
return len(aggMap.userOriginalOutputFields)
}
// IndexAt returns the first index for the given user output field index.
// For avg aggregation, this returns the sum index.
// For backward compatibility, this method is kept.
func (aggMap *AggregationFieldMap) IndexAt(idx int) int {
if len(aggMap.userOriginalOutputFieldIdxes[idx]) > 0 {
return aggMap.userOriginalOutputFieldIdxes[idx][0]
}
return -1
}
// IndexesAt returns all indices for the given user output field index.
// For avg aggregation, this returns both sum and count indices.
// For other aggregations, this returns a slice with a single index.
func (aggMap *AggregationFieldMap) IndexesAt(idx int) []int {
return aggMap.userOriginalOutputFieldIdxes[idx]
}
func (aggMap *AggregationFieldMap) NameAt(idx int) string {
return aggMap.userOriginalOutputFields[idx]
}
func NewAggregationFieldMap(originalUserOutputFields []string, groupByFields []string, aggs []AggregateBase) (*AggregationFieldMap, error) {
numGroupingKeys := len(groupByFields)
groupByFieldMap := make(map[string]int, len(groupByFields))
for i, field := range groupByFields {
groupByFieldMap[field] = i
}
// Build a map from originalName to all indices (for avg, this will include both sum and count indices)
aggFieldMap := make(map[string][]int, len(aggs))
for i, agg := range aggs {
originalName := agg.OriginalName()
idx := i + numGroupingKeys
// Check if this aggregate is part of an avg aggregation
var isAvg bool
switch a := agg.(type) {
case *SumAggregate:
isAvg = a.isAvg
case *CountAggregate:
isAvg = a.isAvg
}
if isAvg {
// For avg aggregates, both sum and count share the same originalName
// Add this index to the list for this originalName
aggFieldMap[originalName] = append(aggFieldMap[originalName], idx)
} else {
// For non-avg aggregates, each originalName maps to a single index
aggFieldMap[originalName] = []int{idx}
}
}
userOriginalOutputFieldIdxes := make([][]int, len(originalUserOutputFields))
for i, outputField := range originalUserOutputFields {
if idx, exist := groupByFieldMap[outputField]; exist {
// Group by field maps to a single index
userOriginalOutputFieldIdxes[i] = []int{idx}
} else if indices, exist := aggFieldMap[outputField]; exist {
// Aggregate field may map to multiple indices (for avg: sum and count)
userOriginalOutputFieldIdxes[i] = indices
} else {
// Field is neither a group_by field nor an aggregation — reject early.
// This covers two cases:
// 1. GROUP BY query: output_fields can only contain group_by columns or aggregation expressions
// 2. Global aggregation (no GROUP BY): output_fields can only contain aggregation expressions
// (e.g., "SELECT count(*), int64 FROM t" is invalid SQL — cannot mix aggregates with raw columns)
if numGroupingKeys > 0 {
return nil, merr.WrapErrParameterInvalidMsg(
"output field '%s' is not allowed: when using GROUP BY, output_fields can only contain "+
"group_by fields (%v) or aggregation expressions",
outputField, groupByFields,
)
}
return nil, merr.WrapErrParameterInvalidMsg(
"output field '%s' is not allowed: when using aggregation functions (e.g., count(*)), "+
"output_fields can only contain aggregation expressions, not regular columns",
outputField,
)
}
}
return &AggregationFieldMap{originalUserOutputFields, userOriginalOutputFieldIdxes}, nil
}
// ComputeAvgFromSumAndCount computes average from sum and count field data.
// It takes sumFieldData and countFieldData, computes avg = sum / count for each row,
// and returns a new Double FieldData containing the average values.
func ComputeAvgFromSumAndCount(sumFieldData *schemapb.FieldData, countFieldData *schemapb.FieldData) (*schemapb.FieldData, error) {
if sumFieldData == nil || countFieldData == nil {
return nil, merr.WrapErrServiceInternalMsg("sumFieldData and countFieldData cannot be nil")
}
sumType := sumFieldData.GetType()
countType := countFieldData.GetType()
if countType != schemapb.DataType_Int64 {
return nil, merr.WrapErrParameterInvalidMsg("count field must be Int64 type, got %s", countType.String())
}
countData := countFieldData.GetScalars().GetLongData().GetData()
rowCount := len(countData)
// Create result FieldData with Double type
result := &schemapb.FieldData{
Type: schemapb.DataType_Double,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{
DoubleData: &schemapb.DoubleArray{Data: make([]float64, 0, rowCount)},
},
},
},
}
resultData := make([]float64, 0, rowCount)
// Compute avg = sum / count for each row
switch sumType {
case schemapb.DataType_Int64:
sumData := sumFieldData.GetScalars().GetLongData().GetData()
if len(sumData) != rowCount {
return nil, merr.WrapErrParameterInvalidMsg("sum and count field data must have the same length, got sum:%d, count:%d", len(sumData), rowCount)
}
for i := 0; i < rowCount; i++ {
if countData[i] == 0 {
return nil, merr.WrapErrParameterInvalidMsg("division by zero: count is 0 at row %d", i)
}
resultData = append(resultData, float64(sumData[i])/float64(countData[i]))
}
case schemapb.DataType_Double:
sumData := sumFieldData.GetScalars().GetDoubleData().GetData()
if len(sumData) != rowCount {
return nil, merr.WrapErrParameterInvalidMsg("sum and count field data must have the same length, got sum:%d, count:%d", len(sumData), rowCount)
}
for i := 0; i < rowCount; i++ {
if countData[i] == 0 {
return nil, merr.WrapErrParameterInvalidMsg("division by zero: count is 0 at row %d", i)
}
resultData = append(resultData, sumData[i]/float64(countData[i]))
}
default:
return nil, merr.WrapErrParameterInvalidMsg("unsupported sum field type for avg computation: %s", sumType.String())
}
result.GetScalars().GetDoubleData().Data = resultData
return result, nil
}