diff --git a/internal/agg/aggregate.go b/internal/agg/aggregate.go new file mode 100644 index 0000000000..34829d6a4d --- /dev/null +++ b/internal/agg/aggregate.go @@ -0,0 +1,240 @@ +package agg + +import ( + "fmt" + "reflect" + "regexp" + "strings" + + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" + "github.com/milvus-io/milvus/pkg/v2/proto/planpb" +) + +const ( + kSum = "sum" + kCount = "count" + kAvg = "avg" + kMin = "min" + kMax = "max" +) + +var ( + // Define the regular expression pattern once to avoid repeated concatenation. + aggregationTypes = kSum + `|` + kCount + `|` + kAvg + `|` + kMin + `|` + kMax + aggregationPattern = regexp.MustCompile(`(?i)^(` + aggregationTypes + `)\s*\(\s*([\w\*]*)\s*\)$`) +) + +// MatchAggregationExpression return isAgg, operator name, operator parameter +func MatchAggregationExpression(expression string) (bool, string, string) { + // FindStringSubmatch returns the full match and submatches. + matches := aggregationPattern.FindStringSubmatch(expression) + if len(matches) > 0 { + // Return true, the operator, and the captured parameter. + return true, strings.ToLower(matches[1]), strings.TrimSpace(matches[2]) + } + return false, "", "" +} + +type AggregateBase interface { + Name() string + Update(target *FieldValue, new *FieldValue) error + ToPB() *planpb.Aggregate + FieldID() int64 + OriginalName() string +} + +func isSupportedAggregateName(aggregateName string) bool { + switch aggregateName { + case kCount, kSum, kAvg, kMin, kMax: + return true + default: + return false + } +} + +func NewAggregate(aggregateName string, aggFieldID int64, originalName string, fieldType schemapb.DataType) ([]AggregateBase, error) { + if !isSupportedAggregateName(aggregateName) { + return nil, fmt.Errorf("invalid Aggregation operator %s", aggregateName) + } + + if err := ValidateAggFieldType(aggregateName, fieldType); err != nil { + return nil, err + } + + switch aggregateName { + case kCount: + return []AggregateBase{&CountAggregate{fieldID: aggFieldID, originalName: originalName, isAvg: false}}, nil + case kSum: + return []AggregateBase{&SumAggregate{fieldID: aggFieldID, originalName: originalName, isAvg: false}}, nil + case kAvg: + // avg is implemented as sum and count, which will be computed as sum/count later + return []AggregateBase{ + &SumAggregate{fieldID: aggFieldID, originalName: originalName, isAvg: true}, + &CountAggregate{fieldID: aggFieldID, originalName: originalName, isAvg: true}, + }, nil + case kMin: + return []AggregateBase{&MinAggregate{fieldID: aggFieldID, originalName: originalName}}, nil + case kMax: + return []AggregateBase{&MaxAggregate{fieldID: aggFieldID, originalName: originalName}}, nil + default: + // should never happen due to isSupportedAggregateName check + return nil, fmt.Errorf("invalid Aggregation operator %s", aggregateName) + } +} + +func FromPB(pb *planpb.Aggregate) (AggregateBase, error) { + switch pb.Op { + case planpb.AggregateOp_count: + return &CountAggregate{fieldID: pb.GetFieldId(), isAvg: false}, nil + case planpb.AggregateOp_sum: + return &SumAggregate{fieldID: pb.GetFieldId(), isAvg: false}, nil + case planpb.AggregateOp_min: + return &MinAggregate{fieldID: pb.GetFieldId()}, nil + case planpb.AggregateOp_max: + return &MaxAggregate{fieldID: pb.GetFieldId()}, nil + default: + return nil, fmt.Errorf("invalid Aggregation operator %d", pb.Op) + } +} + +func AccumulateFieldValue(target *FieldValue, new *FieldValue) error { + if target == nil || new == nil { + return fmt.Errorf("target or new field value is nil") + } + + // Handle nil `val` for initialization + if target.val == nil { + target.val = new.val + return nil + } + // Verify types match + if reflect.TypeOf(target.val) != reflect.TypeOf(new.val) { + return fmt.Errorf("type mismatch: target is %T, new is %T", target.val, new.val) + } + + switch target.val.(type) { + case int: + target.val = target.val.(int) + new.val.(int) + case int32: + target.val = target.val.(int32) + new.val.(int32) + case int64: + target.val = target.val.(int64) + new.val.(int64) + case float32: + target.val = target.val.(float32) + new.val.(float32) + case float64: + target.val = target.val.(float64) + new.val.(float64) + default: + return fmt.Errorf("unsupported type: %T", target.val) + } + return nil +} + +func AggregatesToPB(aggregates []AggregateBase) []*planpb.Aggregate { + ret := make([]*planpb.Aggregate, len(aggregates)) + for idx, agg := range aggregates { + ret[idx] = agg.ToPB() + } + return ret +} + +type FieldValue struct { + val interface{} +} + +func NewFieldValue(v interface{}) *FieldValue { + return &FieldValue{val: v} +} + +type Row struct { + fieldValues []*FieldValue +} + +func (r *Row) ColumnCount() int { + return len(r.fieldValues) +} + +func (r *Row) ValAt(col int) interface{} { + return r.fieldValues[col].val +} + +func (r *Row) Equal(other *Row, keyCount int) bool { + // Ensure both rows have enough fields for key comparison + if len(r.fieldValues) < keyCount || len(other.fieldValues) < keyCount || len(r.fieldValues) != len(other.fieldValues) { + return false + } + // Compare each field value for equality + for i := 0; i < keyCount; i++ { + if r.fieldValues[i].val != other.fieldValues[i].val { + return false + } + } + return true +} + +func (r *Row) UpdateFieldValue(newRow *Row, col int, agg AggregateBase) { + agg.Update(r.fieldValues[col], newRow.fieldValues[col]) +} + +func (r *Row) ToString() string { + var builder strings.Builder + builder.WriteString("agg-row:") + for _, fv := range r.fieldValues { + builder.WriteString(fmt.Sprintf("%v,", fv.val)) + } + return builder.String() +} + +func NewRow(fieldValues []*FieldValue) *Row { + return &Row{fieldValues: fieldValues} +} + +type Bucket struct { + rows []*Row +} + +func (bucket *Bucket) AddRow(row *Row) { + bucket.rows = append(bucket.rows, row) +} + +func (bucket *Bucket) RowAt(idx int) *Row { + return bucket.rows[idx] +} + +func (bucket *Bucket) RowCount() int { + return len(bucket.rows) +} + +func (bucket *Bucket) Accumulate(row *Row, idx int, keyCount int, aggs []AggregateBase) error { + if idx >= len(bucket.rows) || idx < 0 { + return fmt.Errorf("wrong idx:%d for bucket", idx) + } + targetRow := bucket.rows[idx] + if targetRow == nil { + return fmt.Errorf("nil row at the target idx:%d, cannot accumulate the row", idx) + } + if row.ColumnCount() != targetRow.ColumnCount() { + return fmt.Errorf("column count:%d in the row must be equal to the target row:%d", row.ColumnCount(), bucket.rows[idx].ColumnCount()) + } + if row.ColumnCount() != keyCount+len(aggs) { + return fmt.Errorf("column count:%d in the row must be sum of keyCount:%d and the number of aggs:%d", row.ColumnCount(), keyCount, len(aggs)) + } + for col := keyCount; col < row.ColumnCount(); col++ { + targetRow.UpdateFieldValue(row, col, aggs[col-keyCount]) + } + return nil +} + +const NONE int = -1 + +func (bucket *Bucket) Find(row *Row, keyCount int) int { + for idx, existingRow := range bucket.rows { + if existingRow.Equal(row, keyCount) { + return idx + } + } + return NONE +} + +func NewBucket() *Bucket { + return &Bucket{rows: make([]*Row, 0, 1)} +} diff --git a/internal/agg/aggregate_reducer.go b/internal/agg/aggregate_reducer.go new file mode 100644 index 0000000000..0c0719a48d --- /dev/null +++ b/internal/agg/aggregate_reducer.go @@ -0,0 +1,459 @@ +package agg + +import ( + "context" + "fmt" + + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" + typeutil2 "github.com/milvus-io/milvus/internal/util/typeutil" + "github.com/milvus-io/milvus/pkg/v2/proto/internalpb" + "github.com/milvus-io/milvus/pkg/v2/proto/planpb" + "github.com/milvus-io/milvus/pkg/v2/proto/segcorepb" + "github.com/milvus-io/milvus/pkg/v2/util/typeutil" +) + +type GroupAggReducer struct { + groupByFieldIds []int64 + aggregates []*planpb.Aggregate + hashValsMap map[uint64]*Bucket + groupLimit int64 + schema *schemapb.CollectionSchema +} + +func NewGroupAggReducer(groupByFieldIds []int64, aggregates []*planpb.Aggregate, groupLimit int64, schema *schemapb.CollectionSchema) *GroupAggReducer { + return &GroupAggReducer{ + groupByFieldIds: groupByFieldIds, + aggregates: aggregates, + hashValsMap: make(map[uint64]*Bucket), // Initialize hashValsMap + groupLimit: groupLimit, + schema: schema, + } +} + +type AggregationResult struct { + fieldDatas []*schemapb.FieldData + allRetrieveCount int64 +} + +func NewAggregationResult(fieldDatas []*schemapb.FieldData, allRetrieveCount int64) *AggregationResult { + if fieldDatas == nil { + fieldDatas = make([]*schemapb.FieldData, 0) + } + return &AggregationResult{ + fieldDatas: fieldDatas, + allRetrieveCount: allRetrieveCount, + } +} + +// GetFieldDatas returns the fieldDatas slice +func (ar *AggregationResult) GetFieldDatas() []*schemapb.FieldData { + return ar.fieldDatas +} + +func (ar *AggregationResult) GetAllRetrieveCount() int64 { + return ar.allRetrieveCount +} + +func (reducer *GroupAggReducer) EmptyAggResult() (*AggregationResult, error) { + helper, err := typeutil.CreateSchemaHelper(reducer.schema) + if err != nil { + return nil, err + } + ret := NewAggregationResult(nil, 0) + appendEmptyField := func(fieldId int64) error { + field, err := helper.GetFieldFromID(fieldId) + if err != nil { + return err + } + emptyFieldData, err := typeutil.GenEmptyFieldData(field) + if err != nil { + return err + } + ret.fieldDatas = append(ret.fieldDatas, emptyFieldData) + return nil + } + + for _, grpFid := range reducer.groupByFieldIds { + err := appendEmptyField(grpFid) + if err != nil { + return nil, err + } + } + for _, agg := range reducer.aggregates { + if agg.GetOp() == planpb.AggregateOp_count { + countField := genEmptyLongFieldData(schemapb.DataType_Int64, []int64{0}) + ret.fieldDatas = append(ret.fieldDatas, countField) + } else { + field, err := helper.GetFieldFromID(agg.GetFieldId()) + if err != nil { + return nil, fmt.Errorf("failed to get field schema for aggregate fieldID %d: %w", agg.GetFieldId(), err) + } + resultType, err := getAggregateResultType(agg.GetOp(), field.GetDataType()) + if err != nil { + return nil, fmt.Errorf("failed to get result type for aggregate fieldID %d: %w", agg.GetFieldId(), err) + } + emptyFieldData, err := genEmptyFieldDataByType(resultType) + if err != nil { + return nil, fmt.Errorf("failed to generate empty field data for result type %s: %w", resultType.String(), err) + } + ret.fieldDatas = append(ret.fieldDatas, emptyFieldData) + } + } + return ret, nil +} + +func genEmptyLongFieldData(dataType schemapb.DataType, data []int64) *schemapb.FieldData { + return &schemapb.FieldData{ + Type: dataType, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: data}}, + }, + }, + } +} + +// genEmptyFieldDataByType generates empty field data based on the data type +func genEmptyFieldDataByType(dataType schemapb.DataType) (*schemapb.FieldData, error) { + switch dataType { + case schemapb.DataType_Int64: + return genEmptyLongFieldData(dataType, []int64{0}), nil + case schemapb.DataType_Double: + return &schemapb.FieldData{ + Type: dataType, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{0}}}, + }, + }, + }, nil + case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32: + return &schemapb.FieldData{ + Type: dataType, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{0}}}, + }, + }, + }, nil + case schemapb.DataType_Float: + return &schemapb.FieldData{ + Type: dataType, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{0}}}, + }, + }, + }, nil + case schemapb.DataType_VarChar, schemapb.DataType_String, schemapb.DataType_Text: + return &schemapb.FieldData{ + Type: dataType, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{}}}, + }, + }, + }, nil + case schemapb.DataType_Timestamptz: + return genEmptyLongFieldData(dataType, []int64{0}), nil + default: + // For other types, try to use the original field's GenEmptyFieldData + return nil, fmt.Errorf("unsupported data type for aggregate result: %s", dataType.String()) + } +} + +// getAggregateResultType returns the expected result type for an aggregate operation +// based on the aggregate operator type and the input field type. +func getAggregateResultType(op planpb.AggregateOp, inputType schemapb.DataType) (schemapb.DataType, error) { + switch op { + case planpb.AggregateOp_count: + // count aggregation always returns Int64 + return schemapb.DataType_Int64, nil + case planpb.AggregateOp_avg: + // avg aggregation always returns Double + return schemapb.DataType_Double, nil + case planpb.AggregateOp_min, planpb.AggregateOp_max: + // min/max keep the original field type + return inputType, nil + case planpb.AggregateOp_sum: + // sum returns Int64 for integer types, Double for float types + switch inputType { + case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32, schemapb.DataType_Int64: + return schemapb.DataType_Int64, nil + case schemapb.DataType_Timestamptz: + return schemapb.DataType_Timestamptz, nil + case schemapb.DataType_Float, schemapb.DataType_Double: + return schemapb.DataType_Double, nil + default: + return schemapb.DataType_None, fmt.Errorf("unsupported input type %s for sum aggregation", inputType.String()) + } + default: + return schemapb.DataType_None, fmt.Errorf("unknown aggregate operator: %d", op) + } +} + +// validateAggregationResults validates the input AggregationResult slice +// It checks: +// 1. Each result's fieldDatas length equals numGroupingKeys + numAggs +// 2. No nil fieldData in any result +// 3. Each fieldData's Type matches the expected type from schema +func (reducer *GroupAggReducer) validateAggregationResults(results []*AggregationResult) error { + if reducer.schema == nil { + return fmt.Errorf("schema is nil, cannot validate field types") + } + + helper, err := typeutil.CreateSchemaHelper(reducer.schema) + if err != nil { + return fmt.Errorf("failed to create schema helper: %w", err) + } + + numGroupingKeys := len(reducer.groupByFieldIds) + numAggs := len(reducer.aggregates) + expectedColumnCount := numGroupingKeys + numAggs + + // Build expected types for each column + expectedTypes := make([]schemapb.DataType, 0, expectedColumnCount) + + // Add types for grouping keys + for _, fieldID := range reducer.groupByFieldIds { + field, err := helper.GetFieldFromID(fieldID) + if err != nil { + return fmt.Errorf("failed to get field schema for groupBy fieldID %d: %w", fieldID, err) + } + expectedTypes = append(expectedTypes, field.GetDataType()) + } + + // Add types for aggregates + for _, agg := range reducer.aggregates { + var expectedType schemapb.DataType + if agg.GetOp() == planpb.AggregateOp_count { + // count aggregation always returns Int64 + expectedType = schemapb.DataType_Int64 + } else { + field, err := helper.GetFieldFromID(agg.GetFieldId()) + if err != nil { + return fmt.Errorf("failed to get field schema for aggregate fieldID %d: %w", agg.GetFieldId(), err) + } + expectedType, err = getAggregateResultType(agg.GetOp(), field.GetDataType()) + if err != nil { + return fmt.Errorf("failed to get aggregate result type for aggregate fieldID %d: %w", agg.GetFieldId(), err) + } + } + expectedTypes = append(expectedTypes, expectedType) + } + + // Validate each result + for resultIdx, result := range results { + if result == nil { + return fmt.Errorf("result at index %d is nil", resultIdx) + } + + fieldDatas := result.GetFieldDatas() + + // Check 1: fieldDatas length + if len(fieldDatas) != expectedColumnCount { + return fmt.Errorf("result at index %d has fieldDatas length %d, expected %d (numGroupingKeys=%d, numAggs=%d)", + resultIdx, len(fieldDatas), expectedColumnCount, numGroupingKeys, numAggs) + } + + // Check 2: no nil fieldData and Check 3: type matching + for colIdx, fieldData := range fieldDatas { + if fieldData == nil { + return fmt.Errorf("result at index %d has nil fieldData at column %d", resultIdx, colIdx) + } + + expectedType := expectedTypes[colIdx] + actualType := fieldData.GetType() + if actualType != expectedType { + return fmt.Errorf("result at index %d, column %d has type %s, expected %s", + resultIdx, colIdx, schemapb.DataType_name[int32(actualType)], schemapb.DataType_name[int32(expectedType)]) + } + } + } + + return nil +} + +func (reducer *GroupAggReducer) Reduce(ctx context.Context, results []*AggregationResult) (*AggregationResult, error) { + if len(results) == 0 { + return reducer.EmptyAggResult() + } + + // Validate input results before processing + if err := reducer.validateAggregationResults(results); err != nil { + return nil, err + } + + if len(results) == 1 { + return results[0], nil + } + + // 0. set up aggregates + aggs := make([]AggregateBase, len(reducer.aggregates)) + for idx, aggPb := range reducer.aggregates { + agg, err := FromPB(aggPb) + if err != nil { + return nil, err + } + aggs[idx] = agg + } + + // 1. set up hashers and accumulators + numGroupingKeys := len(reducer.groupByFieldIds) + numAggs := len(reducer.aggregates) + hashers := make([]FieldAccessor, numGroupingKeys) + accumulators := make([]FieldAccessor, numAggs) + firstFieldData := results[0].GetFieldDatas() + outputColumnCount := len(firstFieldData) + for idx, fieldData := range firstFieldData { + accessor, err := NewFieldAccessor(fieldData.GetType()) + if err != nil { + return nil, err + } + if idx < numGroupingKeys { + hashers[idx] = accessor + } else { + accumulators[idx-numGroupingKeys] = accessor + } + } + reducedResult := NewAggregationResult(nil, 0) + isGlobal := numGroupingKeys == 0 + if isGlobal { + reducedResult.fieldDatas = typeutil.PrepareResultFieldData(firstFieldData, 1) + rows := make([]*Row, len(results)) + for idx, result := range results { + reducedResult.allRetrieveCount += result.GetAllRetrieveCount() + fieldValues := make([]*FieldValue, outputColumnCount) + for col := 0; col < outputColumnCount; col++ { + fieldData := result.GetFieldDatas()[col] + accumulators[col].SetVals(fieldData) + fieldValues[col] = NewFieldValue(accumulators[col].ValAt(0)) + } + rows[idx] = NewRow(fieldValues) + } + for r := 1; r < len(rows); r++ { + for c := 0; c < outputColumnCount; c++ { + rows[0].UpdateFieldValue(rows[r], c, aggs[c]) + } + } + AssembleSingleRow(outputColumnCount, rows[0], reducedResult.fieldDatas) + return reducedResult, nil + } + + // 2. compute hash values for all rows in the result retrieved + var totalRowCount int64 = 0 +processResults: + for _, result := range results { + // Check limit before processing each shard to avoid unnecessary work + if reducer.groupLimit != -1 && totalRowCount >= reducer.groupLimit { + break processResults + } + + reducedResult.allRetrieveCount += result.GetAllRetrieveCount() + if result == nil { + return nil, fmt.Errorf("input result from any sources cannot be nil") + } + fieldDatas := result.GetFieldDatas() + if outputColumnCount != len(fieldDatas) { + return nil, fmt.Errorf("retrieved results from different segments have different size of columns") + } + if outputColumnCount == 0 { + return nil, fmt.Errorf("retrieved results have no column data") + } + rowCount := -1 + for i := 0; i < outputColumnCount; i++ { + fieldData := fieldDatas[i] + if i < numGroupingKeys { + hashers[i].SetVals(fieldData) + } else { + accumulators[i-numGroupingKeys].SetVals(fieldData) + } + if rowCount == -1 { + rowCount = hashers[i].RowCount() + } else if i < numGroupingKeys { + if rowCount != hashers[i].RowCount() { + return nil, fmt.Errorf("field data:%d for different columns have different row count, %d vs %d, wrong state", + i, rowCount, hashers[i].RowCount()) + } + } else if rowCount != accumulators[i-numGroupingKeys].RowCount() { + return nil, fmt.Errorf("field data:%d for different columns have different row count, %d vs %d, wrong state", + i, rowCount, accumulators[i-numGroupingKeys].RowCount()) + } + } + + for row := 0; row < rowCount; row++ { + // Check limit before processing each row to avoid unnecessary hashing and copying + if reducer.groupLimit != -1 && totalRowCount >= reducer.groupLimit { + break processResults + } + rowFieldValues := make([]*FieldValue, outputColumnCount) + var hashVal uint64 + for col := 0; col < outputColumnCount; col++ { + if col < numGroupingKeys { + if col > 0 { + hashVal = typeutil2.HashMix(hashVal, hashers[col].Hash(row)) + } else { + hashVal = hashers[col].Hash(row) + } + rowFieldValues[col] = NewFieldValue(hashers[col].ValAt(row)) + } else { + rowFieldValues[col] = NewFieldValue(accumulators[col-numGroupingKeys].ValAt(row)) + } + } + newRow := NewRow(rowFieldValues) + if bucket := reducer.hashValsMap[hashVal]; bucket == nil { + newBucket := NewBucket() + newBucket.AddRow(newRow) + totalRowCount++ + reducer.hashValsMap[hashVal] = newBucket + } else { + if rowIdx := bucket.Find(newRow, numGroupingKeys); rowIdx == NONE { + bucket.AddRow(newRow) + totalRowCount++ + } else { + bucket.Accumulate(newRow, rowIdx, numGroupingKeys, aggs) + } + } + // Don't guarantee specific groups to be returned before milvus support order by + } + } + + // 3. assemble reduced buckets into retrievedResult + reducedResult.fieldDatas = typeutil.PrepareResultFieldData(firstFieldData, totalRowCount) + for _, bucket := range reducer.hashValsMap { + err := AssembleBucket(bucket, reducedResult.GetFieldDatas()) + if err != nil { + return nil, err + } + } + return reducedResult, nil +} + +func InternalResult2AggResult(results []*internalpb.RetrieveResults) []*AggregationResult { + aggResults := make([]*AggregationResult, len(results)) + for i := 0; i < len(results); i++ { + aggResults[i] = NewAggregationResult(results[i].GetFieldsData(), results[i].GetAllRetrieveCount()) + } + return aggResults +} + +func AggResult2internalResult(aggRes *AggregationResult) *internalpb.RetrieveResults { + return &internalpb.RetrieveResults{FieldsData: aggRes.GetFieldDatas(), AllRetrieveCount: aggRes.GetAllRetrieveCount()} +} + +func SegcoreResults2AggResult(results []*segcorepb.RetrieveResults) ([]*AggregationResult, error) { + aggResults := make([]*AggregationResult, len(results)) + for i := 0; i < len(results); i++ { + if results[i] == nil { + return nil, fmt.Errorf("input segcore query results from any sources cannot be nil") + } + fieldsData := results[i].GetFieldsData() + allRetrieveCount := results[i].GetAllRetrieveCount() + aggResults[i] = NewAggregationResult(fieldsData, allRetrieveCount) + } + return aggResults, nil +} + +func AggResult2segcoreResult(aggRes *AggregationResult) *segcorepb.RetrieveResults { + return &segcorepb.RetrieveResults{FieldsData: aggRes.GetFieldDatas(), AllRetrieveCount: aggRes.GetAllRetrieveCount()} +} diff --git a/internal/agg/aggregate_util.go b/internal/agg/aggregate_util.go new file mode 100644 index 0000000000..05a2319350 --- /dev/null +++ b/internal/agg/aggregate_util.go @@ -0,0 +1,494 @@ +package agg + +import ( + "encoding/binary" + "fmt" + "hash" + "hash/fnv" + "math" + "unsafe" + + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" +) + +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, fmt.Errorf("unsupported data type for hasher") + } +} + +type FieldAccessor interface { + Hash(idx int) uint64 + ValAt(idx int) interface{} + SetVals(fieldData *schemapb.FieldData) + RowCount() int +} + +type Int32FieldAccessor struct { + vals []int32 + 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))) + } + 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() +} + +func (i32Field *Int32FieldAccessor) RowCount() int { + return len(i32Field.vals) +} + +func (i32Field *Int32FieldAccessor) ValAt(idx int) interface{} { + return i32Field.vals[idx] +} + +func newInt32FieldAccessor() FieldAccessor { + return &Int32FieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 4)} +} + +type Int64FieldAccessor struct { + vals []int64 + 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))) + } + 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() +} + +func (i64Field *Int64FieldAccessor) RowCount() int { + return len(i64Field.vals) +} + +func (i64Field *Int64FieldAccessor) ValAt(idx int) interface{} { + return i64Field.vals[idx] +} + +func newInt64FieldAccessor() FieldAccessor { + return &Int64FieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 8)} +} + +type TimestamptzFieldAccessor struct { + vals []int64 + 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))) + } + 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() +} + +func (tzField *TimestamptzFieldAccessor) RowCount() int { + return len(tzField.vals) +} + +func (tzField *TimestamptzFieldAccessor) ValAt(idx int) interface{} { + return tzField.vals[idx] +} + +func newTimestamptzFieldAccessor() FieldAccessor { + return &TimestamptzFieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 8)} +} + +// BoolFieldAccessor +type BoolFieldAccessor struct { + vals []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))) + } + 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() +} + +func (boolField *BoolFieldAccessor) RowCount() int { + return len(boolField.vals) +} + +func (boolField *BoolFieldAccessor) ValAt(idx int) interface{} { + return boolField.vals[idx] +} + +func newBoolFieldAccessor() FieldAccessor { + return &BoolFieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 1)} +} + +// Float32FieldAccessor +type Float32FieldAccessor struct { + vals []float32 + 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))) + } + 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() +} + +func (f32FieldAccessor *Float32FieldAccessor) RowCount() int { + return len(f32FieldAccessor.vals) +} + +func (f32FieldAccessor *Float32FieldAccessor) ValAt(idx int) interface{} { + return f32FieldAccessor.vals[idx] +} + +func newFloat32FieldAccessor() FieldAccessor { + return &Float32FieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 4)} +} + +// Float64FieldAccessor +type Float64FieldAccessor struct { + vals []float64 + 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))) + } + 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() +} + +func (f64Field *Float64FieldAccessor) RowCount() int { + return len(f64Field.vals) +} + +func (f64Field *Float64FieldAccessor) ValAt(idx int) interface{} { + return f64Field.vals[idx] +} + +func newFloat64FieldAccessor() FieldAccessor { + return &Float64FieldAccessor{hasher: fnv.New64a(), buffer: make([]byte, 8)} +} + +// StringFieldAccessor +type StringFieldAccessor struct { + vals []string + 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))) + } + 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() +} + +func (stringField *StringFieldAccessor) RowCount() int { + return len(stringField.vals) +} + +func (stringField *StringFieldAccessor) ValAt(idx int) interface{} { + return stringField.vals[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.ValAt(c), fieldDatas[c]) + if err != nil { + return err + } + } + return nil +} + +func AssembleSingleValue(val interface{}, fieldData *schemapb.FieldData) error { + switch fieldData.GetType() { + case schemapb.DataType_Bool: + boolVal, ok := val.(bool) + if !ok { + return fmt.Errorf("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 fmt.Errorf("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 fmt.Errorf("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 fmt.Errorf("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 fmt.Errorf("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 fmt.Errorf("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 fmt.Errorf("type assertion failed: expected string, got %T", val) + } + fieldData.GetScalars().GetStringData().Data = append(fieldData.GetScalars().GetStringData().GetData(), stringVal) + default: + return fmt.Errorf("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 { + 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 not found, set empty slice + userOriginalOutputFieldIdxes[i] = []int{} + } + } + + return &AggregationFieldMap{originalUserOutputFields, userOriginalOutputFieldIdxes} +} + +// 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, fmt.Errorf("sumFieldData and countFieldData cannot be nil") + } + + sumType := sumFieldData.GetType() + countType := countFieldData.GetType() + + if countType != schemapb.DataType_Int64 { + return nil, fmt.Errorf("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, fmt.Errorf("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, fmt.Errorf("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, fmt.Errorf("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, fmt.Errorf("division by zero: count is 0 at row %d", i) + } + resultData = append(resultData, sumData[i]/float64(countData[i])) + } + default: + return nil, fmt.Errorf("unsupported sum field type for avg computation: %s", sumType.String()) + } + + result.GetScalars().GetDoubleData().Data = resultData + return result, nil +} diff --git a/internal/agg/operators.go b/internal/agg/operators.go new file mode 100644 index 0000000000..c99ae68151 --- /dev/null +++ b/internal/agg/operators.go @@ -0,0 +1,235 @@ +package agg + +import ( + "fmt" + + "github.com/milvus-io/milvus/pkg/v2/proto/planpb" +) + +type SumAggregate struct { + fieldID int64 + originalName string + isAvg bool +} + +func (sum *SumAggregate) Name() string { + return kSum +} + +func (sum *SumAggregate) Update(target *FieldValue, new *FieldValue) error { + return AccumulateFieldValue(target, new) +} + +func (sum *SumAggregate) ToPB() *planpb.Aggregate { + return &planpb.Aggregate{Op: planpb.AggregateOp_sum, FieldId: sum.FieldID()} +} + +func (sum *SumAggregate) FieldID() int64 { + return sum.fieldID +} + +func (sum *SumAggregate) OriginalName() string { + return sum.originalName +} + +type CountAggregate struct { + fieldID int64 + originalName string + isAvg bool +} + +func (count *CountAggregate) Name() string { + return kCount +} + +func (count *CountAggregate) Update(target *FieldValue, new *FieldValue) error { + return AccumulateFieldValue(target, new) +} + +func (count *CountAggregate) ToPB() *planpb.Aggregate { + return &planpb.Aggregate{Op: planpb.AggregateOp_count, FieldId: count.FieldID()} +} + +func (count *CountAggregate) FieldID() int64 { + return count.fieldID +} + +func (count *CountAggregate) OriginalName() string { + return count.originalName +} + +type MinAggregate struct { + fieldID int64 + originalName string +} + +func (min *MinAggregate) Name() string { + return kMin +} + +func (min *MinAggregate) Update(target *FieldValue, new *FieldValue) error { + if target == nil || new == nil { + return fmt.Errorf("target or new field value is nil") + } + + // Handle nil `val` for initialization + if target.val == nil { + target.val = new.val + return nil + } + + // Compare and keep the minimum value + switch targetVal := target.val.(type) { + case int: + newVal, ok := new.val.(int) + if !ok { + return fmt.Errorf("type mismatch: target is int, new is %T", new.val) + } + if newVal < targetVal { + target.val = newVal + } + case int32: + newVal, ok := new.val.(int32) + if !ok { + return fmt.Errorf("type mismatch: target is int32, new is %T", new.val) + } + if newVal < targetVal { + target.val = newVal + } + case int64: + newVal, ok := new.val.(int64) + if !ok { + return fmt.Errorf("type mismatch: target is int64, new is %T", new.val) + } + if newVal < targetVal { + target.val = newVal + } + case float32: + newVal, ok := new.val.(float32) + if !ok { + return fmt.Errorf("type mismatch: target is float32, new is %T", new.val) + } + if newVal < targetVal { + target.val = newVal + } + case float64: + newVal, ok := new.val.(float64) + if !ok { + return fmt.Errorf("type mismatch: target is float64, new is %T", new.val) + } + if newVal < targetVal { + target.val = newVal + } + case string: + newVal, ok := new.val.(string) + if !ok { + return fmt.Errorf("type mismatch: target is string, new is %T", new.val) + } + if newVal < targetVal { + target.val = newVal + } + default: + return fmt.Errorf("unsupported type for min aggregation: %T", target.val) + } + return nil +} + +func (min *MinAggregate) ToPB() *planpb.Aggregate { + return &planpb.Aggregate{Op: planpb.AggregateOp_min, FieldId: min.FieldID()} +} + +func (min *MinAggregate) FieldID() int64 { + return min.fieldID +} + +func (min *MinAggregate) OriginalName() string { + return min.originalName +} + +type MaxAggregate struct { + fieldID int64 + originalName string +} + +func (max *MaxAggregate) Name() string { + return kMax +} + +func (max *MaxAggregate) Update(target *FieldValue, new *FieldValue) error { + if target == nil || new == nil { + return fmt.Errorf("target or new field value is nil") + } + + // Handle nil `val` for initialization + if target.val == nil { + target.val = new.val + return nil + } + + // Compare and keep the maximum value + switch targetVal := target.val.(type) { + case int: + newVal, ok := new.val.(int) + if !ok { + return fmt.Errorf("type mismatch: target is int, new is %T", new.val) + } + if newVal > targetVal { + target.val = newVal + } + case int32: + newVal, ok := new.val.(int32) + if !ok { + return fmt.Errorf("type mismatch: target is int32, new is %T", new.val) + } + if newVal > targetVal { + target.val = newVal + } + case int64: + newVal, ok := new.val.(int64) + if !ok { + return fmt.Errorf("type mismatch: target is int64, new is %T", new.val) + } + if newVal > targetVal { + target.val = newVal + } + case float32: + newVal, ok := new.val.(float32) + if !ok { + return fmt.Errorf("type mismatch: target is float32, new is %T", new.val) + } + if newVal > targetVal { + target.val = newVal + } + case float64: + newVal, ok := new.val.(float64) + if !ok { + return fmt.Errorf("type mismatch: target is float64, new is %T", new.val) + } + if newVal > targetVal { + target.val = newVal + } + case string: + newVal, ok := new.val.(string) + if !ok { + return fmt.Errorf("type mismatch: target is string, new is %T", new.val) + } + if newVal > targetVal { + target.val = newVal + } + default: + return fmt.Errorf("unsupported type for max aggregation: %T", target.val) + } + return nil +} + +func (max *MaxAggregate) ToPB() *planpb.Aggregate { + return &planpb.Aggregate{Op: planpb.AggregateOp_max, FieldId: max.FieldID()} +} + +func (max *MaxAggregate) FieldID() int64 { + return max.fieldID +} + +func (max *MaxAggregate) OriginalName() string { + return max.originalName +} diff --git a/internal/agg/type_check.go b/internal/agg/type_check.go new file mode 100644 index 0000000000..4932f0b4f9 --- /dev/null +++ b/internal/agg/type_check.go @@ -0,0 +1,57 @@ +package agg + +import ( + "fmt" + + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" +) + +func isSupportedAggFieldType(aggregateName string, dt schemapb.DataType) bool { + switch aggregateName { + case kCount: + // count(*) uses DataType_None; count(field) can work for any type since it + // doesn't depend on the field's scalar representation in reducer. + return true + case kSum, kAvg: + switch dt { + case schemapb.DataType_Int8, + schemapb.DataType_Int16, + schemapb.DataType_Int32, + schemapb.DataType_Int64, + schemapb.DataType_Float, + schemapb.DataType_Double: + return true + default: + return false + } + case kMin, kMax: + switch dt { + case schemapb.DataType_Int8, + schemapb.DataType_Int16, + schemapb.DataType_Int32, + schemapb.DataType_Int64, + schemapb.DataType_Float, + schemapb.DataType_Double, + schemapb.DataType_VarChar, + schemapb.DataType_String, + schemapb.DataType_Timestamptz: + return true + default: + return false + } + default: + // operator validity is handled by NewAggregate; keep this conservative. + return false + } +} + +// ValidateAggFieldType checks whether the input field type can be used by the +// given aggregation operator. +// +// Note: operator name is expected to be normalized (lowercase) by caller. +func ValidateAggFieldType(aggregateName string, dt schemapb.DataType) error { + if !isSupportedAggFieldType(aggregateName, dt) { + return fmt.Errorf("aggregation operator %s does not support data type %s", aggregateName, dt.String()) + } + return nil +} diff --git a/internal/core/src/CMakeLists.txt b/internal/core/src/CMakeLists.txt index 17d66f77a4..02beca53e4 100644 --- a/internal/core/src/CMakeLists.txt +++ b/internal/core/src/CMakeLists.txt @@ -52,6 +52,7 @@ add_subdirectory( exec ) add_subdirectory( bitset ) add_subdirectory( futures ) add_subdirectory( rescores ) +add_subdirectory( plan ) milvus_add_pkg_config("milvus_core") @@ -70,6 +71,7 @@ add_library(milvus_core SHARED $ $ $ + $ ) set(LINK_TARGETS diff --git a/internal/core/src/common/Array.h b/internal/core/src/common/Array.h index 26b3160426..66663faa3e 100644 --- a/internal/core/src/common/Array.h +++ b/internal/core/src/common/Array.h @@ -637,6 +637,17 @@ class ArrayView { } } + void + output_data(Array& array) const { + // Create a new Array object from ArrayView's data and assign it to the + // output array + array = Array(data_, + length_, + static_cast(size_), + element_type_, + offsets_ptr_); + } + ScalarFieldProto output_data() const { ScalarFieldProto data_array; diff --git a/internal/core/src/common/BitUtil.h b/internal/core/src/common/BitUtil.h new file mode 100644 index 0000000000..e84b69e5a4 --- /dev/null +++ b/internal/core/src/common/BitUtil.h @@ -0,0 +1,472 @@ +// 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. + +#pragma once +#include +#include +#include +#include + +namespace milvus { +namespace bits { + +/** + * @brief Safely rounds up a value to the nearest multiple of factor. + * @tparam T Integer type of the value. + * @tparam U Integer type of the factor. + * @param value The number to be rounded. + * @param factor The multiple to round up to. + * @return Rounded value. For positive values, rounds up; for negative values, + * rounds towards zero (e.g., roundUp(-5, 4) = -4). + * @note This function uses assert() to detect overflow conditions, which + * should not occur in normal usage. If overflow is detected, the program + * will terminate in debug builds. + */ +template < + typename T, + typename U, + typename std::enable_if_t && std::is_integral_v, + int> = 0> +constexpr T +roundUp(T value, U factor) noexcept { + // 1. Handle base cases for safety + if (factor <= 0) { + return value; + } + if (value == 0) { + return 0; + } + + // 2. Handle signed integer edge cases (T_MIN) before any type conversion + // This must be done before converting to CommonT to avoid sign issues + if constexpr (std::is_signed_v) { + if (value == std::numeric_limits::min()) { + T factor_t = static_cast(factor); + // For T_MIN, if it's already a multiple of factor, return it + // Otherwise, we cannot round up without overflow + if (value % factor_t == 0) { + return value; + } + // Overflow: T_MIN that's not a multiple cannot be rounded up + assert(false && "roundUp: T_MIN value cannot be rounded up"); + __builtin_trap(); + } + } + + // 3. Choose appropriate common type that preserves sign of T + // If T is signed, use a signed common type to avoid sign conversion issues + using CommonT = + std::conditional_t, + std::make_signed_t>, + std::common_type_t>; + + const CommonT v = static_cast(value); + const CommonT f = static_cast(factor); + + // 4. Optimization for Power of Two (Common in memory alignment) + // Check if factor is a power of two: (f & (f - 1)) == 0 + if (f > 0 && (f & (f - 1)) == 0) { + // Check for potential overflow before bitwise operations + if constexpr (std::is_signed_v) { + if (v > 0 && v > (std::numeric_limits::max() - f + 1)) { + assert(false && + "roundUp: overflow detected in power-of-two path"); + __builtin_trap(); + } + } else { + if (v > (std::numeric_limits::max() - f + 1)) { + assert(false && + "roundUp: overflow detected in power-of-two path"); + __builtin_trap(); + } + } + return static_cast((v + f - 1) & ~(f - 1)); + } + + // 5. General safe implementation for positive and negative values + if constexpr (std::is_signed_v) { + if (v > 0) { + // Positive values: round up using ((v-1)/f + 1) * f + CommonT num_factors = (v - 1) / f + 1; + if (num_factors > std::numeric_limits::max() / f) { + assert(false && + "roundUp: overflow detected in positive value path"); + __builtin_trap(); + } + return static_cast(num_factors * f); + } else { + // Negative values: round towards zero (e.g., -5 -> -4 for factor 4) + // This is the standard behavior for alignment operations + return static_cast((v / f) * f); + } + } else { + // Unsigned types: only positive values possible + CommonT num_factors = (v - 1) / f + 1; + if (num_factors > std::numeric_limits::max() / f) { + assert(false && "roundUp: overflow detected in unsigned path"); + __builtin_trap(); + } + return static_cast(num_factors * f); + } +} + +constexpr uint64_t +nBytes(int32_t value) { + return roundUp(value, 8) / 8; +} + +constexpr inline uint64_t +lowMask(int32_t bits) { + if (bits >= 64) + return ~0ULL; + if (bits <= 0) + return 0ULL; + return (1UL << bits) - 1; +} + +/** + * @brief Gets and clears the last (least significant) set bit. + * @param bits Reference to the bitmask to modify. + * @return The position (0-based) of the cleared bit, counting from the least significant bit. + * @note The input bits must have at least one bit set. If bits is 0, the behavior is undefined + * and will trigger an assertion failure. + */ +inline int32_t +getAndClearLastSetBit(uint16_t& bits) { + // __builtin_ctz(0) is undefined behavior, so we must check for zero + // we will not call this function if bits is 0 + assert(bits != 0 && "getAndClearLastSetBit: input bits must not be zero"); + if (bits == 0) { + __builtin_trap(); + } + int32_t trailingZeros = __builtin_ctz(bits); + bits &= bits - 1; + return trailingZeros; +} + +constexpr inline uint64_t +highMask(int32_t bits) { + if (bits >= 64) + return ~0ULL; + if (bits <= 0) + return 0ULL; + return lowMask(bits) << (64 - bits); +} + +/** + * Invokes a function for each batch of bits (partial or full words) + * in a given range. + * + * @param begin first bit to check (inclusive) + * @param end last bit to check (exclusive) + * @param partialWordFunc function to invoke for a partial word; + * takes index of the word and mask + * @param fullWordFunc function to invoke for a full word; + * takes index of the word + */ +template +inline void +forEachWord(int32_t begin, + int32_t end, + PartialWordFunc partialWordFunc, + FullWordFunc fullWordFunc) { + if (begin >= end) { + return; + } + int32_t firstWord = roundUp(begin, 64); + int32_t lastWord = end & ~63L; + if (lastWord < firstWord) { + partialWordFunc(lastWord / 64, + lowMask(end - lastWord) & highMask(firstWord - begin)); + return; + } + if (begin != firstWord) { + partialWordFunc(begin / 64, highMask(firstWord - begin)); + } + for (int32_t i = firstWord; i + 64 <= lastWord; i += 64) { + fullWordFunc(i / 64); + } + if (end != lastWord) { + partialWordFunc(lastWord / 64, lowMask(end - lastWord)); + } +} + +inline int32_t +countBitsScalar(const uint64_t* bits, int32_t begin, int32_t end) { + int32_t count = 0; + forEachWord( + begin, + end, + [&count, bits](int32_t idx, uint64_t mask) { + count += __builtin_popcountll(bits[idx] & mask); + }, + [&count, bits](int32_t idx) { + count += __builtin_popcountll(bits[idx]); + }); + return count; +} + +#ifdef __AVX2__ +#include +inline int32_t +countBitsAVX2(const uint64_t* bits, int32_t begin, int32_t end) { + int32_t count = 0; + + // Convert bit indices to word-aligned bit indices (same logic as forEachWord) + int32_t firstWordBit = roundUp(begin, 64); // first word-aligned bit index + int32_t lastWordBit = end & ~63L; // last word-aligned bit index + + // Handle case where range spans less than one word + if (lastWordBit < firstWordBit) { + // begin and end are in the same word + int32_t wordIdx = begin >> 6; + uint64_t offset = begin & 63; + uint64_t mask = highMask(64 - offset) & lowMask(end - (wordIdx << 6)); + return __builtin_popcountll(bits[wordIdx] & mask); + } + + // Handle partial word at the beginning (if begin is not word-aligned) + if (begin != firstWordBit) { + int32_t wordIdx = begin >> 6; + uint64_t offset = begin & 63; + uint64_t mask = highMask(64 - offset); + count += __builtin_popcountll(bits[wordIdx] & mask); + } + + // Process full words using unrolled scalar loop + // Note: Modern CPUs have fast scalar POPCNT (1 cycle latency, 1 per cycle throughput). + // The previous AVX2 implementation only used SIMD for load/store, then did scalar + // popcount anyway, adding overhead without benefit. This simple unrolled loop is + // faster and more maintainable. + int32_t wordIdx = firstWordBit >> 6; + int32_t lastWordIdx = lastWordBit >> 6; + + for (; wordIdx + 4 <= lastWordIdx; wordIdx += 4) { + count += __builtin_popcountll(bits[wordIdx]); + count += __builtin_popcountll(bits[wordIdx + 1]); + count += __builtin_popcountll(bits[wordIdx + 2]); + count += __builtin_popcountll(bits[wordIdx + 3]); + } + + // Handle remaining full words + for (; wordIdx < lastWordIdx; ++wordIdx) { + count += __builtin_popcountll(bits[wordIdx]); + } + + // Handle partial word at the end (if end is not word-aligned) + if (end != lastWordBit) { + int32_t wordIdx = lastWordBit >> 6; + uint64_t mask = lowMask(end & 63); + count += __builtin_popcountll(bits[wordIdx] & mask); + } + + return count; +} +#endif + +#ifdef __AVX512F__ +#include +inline int32_t +countBitsAVX512(const uint64_t* bits, int32_t begin, int32_t end) { + if (begin >= end) { + return 0; + } + + int32_t count = 0; + + // Convert bit indices to word-aligned bit indices (same logic as forEachWord) + int32_t firstWordBit = roundUp(begin, 64); // first word-aligned bit index + int32_t lastWordBit = end & ~63L; // last word-aligned bit index + + // Handle case where range spans less than one word + if (lastWordBit < firstWordBit) { + // begin and end are in the same word + int32_t wordIdx = begin >> 6; + uint64_t offset = begin & 63; + uint64_t mask = highMask(64 - offset) & lowMask(end - (wordIdx << 6)); + return __builtin_popcountll(bits[wordIdx] & mask); + } + + // Handle partial word at the beginning (if begin is not word-aligned) + if (begin != firstWordBit) { + int32_t wordIdx = begin >> 6; + uint64_t offset = begin & 63; + uint64_t mask = highMask(64 - offset); + count += __builtin_popcountll(bits[wordIdx] & mask); + } + + // Process full words using AVX-512 (8 words at a time) + int32_t wordIdx = firstWordBit >> 6; + int32_t lastWordIdx = lastWordBit >> 6; + + for (; wordIdx + 8 <= lastWordIdx; wordIdx += 8) { + __m512i v = _mm512_loadu_si512(bits + wordIdx); + __m512i c = _mm512_popcnt_epi64(v); + + count += static_cast(_mm512_reduce_add_epi64(c)); + } + + // Handle remaining full words + for (; wordIdx < lastWordIdx; ++wordIdx) { + count += __builtin_popcountll(bits[wordIdx]); + } + + // Handle partial word at the end (if end is not word-aligned) + if (end != lastWordBit) { + int32_t wordIdx = lastWordBit >> 6; + uint64_t mask = lowMask(end & 63); + count += __builtin_popcountll(bits[wordIdx] & mask); + } + + return count; +} +#endif + +inline int32_t +countBits(const uint64_t* bits, int32_t begin, int32_t end) { + const int32_t range = end - begin; + + if (range < 256) { + return countBitsScalar(bits, begin, end); + } + +#ifdef __AVX512F__ + if (range >= 2048) { + return countBitsAVX512(bits, begin, end); + } +#endif + +#ifdef __AVX2__ + return countBitsAVX2(bits, begin, end); +#else + return countBitsScalar(bits, begin, end); +#endif +} + +inline bool +isPowerOfTwo(uint64_t size) { + return (size > 0) && ((size & (size - 1)) == 0); +} + +template +inline int32_t +countLeadingZeros(T word) { + static_assert(std::is_same_v || + std::is_same_v); + /// Built-in Function: int __builtin_clz (unsigned int x) returns the number + /// of leading 0-bits in x, starting at the most significant bit position. If + /// x is 0, the result is undefined. + if (word == 0) { + return sizeof(T) * 8; + } + if constexpr (std::is_same_v) { + return __builtin_clzll(word); + } else { + uint64_t hi = word >> 64; + uint64_t lo = static_cast(word); + return (hi == 0) ? 64 + __builtin_clzll(lo) : __builtin_clzll(hi); + } +} + +inline uint64_t +nextPowerOfTwo(uint64_t size) { + if (size == 0) { + return 0; + } + uint32_t bits = 63 - countLeadingZeros(size); + uint64_t lower = 1ULL << bits; + // Size is a power of 2. + if (lower == size) { + return size; + } + return 2 * lower; +} + +// This is the Hash128to64 function from Google's cityhash (available +// under the MIT License). We use it to reduce multiple 64 bit hashes +// into a single hash. +#if defined(FOLLY_DISABLE_UNDEFINED_BEHAVIOR_SANITIZER) +FOLLY_DISABLE_UNDEFINED_BEHAVIOR_SANITIZER("unsigned-integer-overflow") +#endif +inline uint64_t +hashMix(const uint64_t upper, const uint64_t lower) noexcept { + // Murmur-inspired hashing. + const uint64_t kMul = 0x9ddfea08eb382d69ULL; + uint64_t a = (lower ^ upper) * kMul; + a ^= (a >> 47); + uint64_t b = (upper ^ a) * kMul; + b ^= (b >> 47); + b *= kMul; + return b; +} + +/// Extract bits from integer 'a' at the corresponding bit locations specified +/// by 'mask' to contiguous low bits in return value; the remaining upper bits +/// in return value are set to zero. +template +inline T +extractBits(T a, T mask); + +#ifdef __BMI2__ +template <> +inline uint32_t +extractBits(uint32_t a, uint32_t mask) { + return _pext_u32(a, mask); +} +template <> +inline uint64_t +extractBits(uint64_t a, uint64_t mask) { + return _pext_u64(a, mask); +} +#else +template +inline T +extractBits(T a, T mask) { + static_assert(std::is_unsigned_v, "extractBits requires unsigned type"); + + // 1. first try to use BMI2 intrinsic +#ifdef __BMI2__ + if constexpr (sizeof(T) == 8) { + return _pext_u64(static_cast(a), static_cast(mask)); + } else { + // For uint32_t, uint16_t, uint8_t, use 32-bit PEXT + return static_cast( + _pext_u32(static_cast(a), static_cast(mask))); + } +#else + // 2. Scalar fallback implementation + T dst = 0; + for (int k = 0; mask != 0; ++k) { + // Use ctz implementation for different widths + int shift; + if constexpr (sizeof(T) <= sizeof(unsigned int)) { + shift = __builtin_ctz(static_cast(mask)); + } else { + shift = __builtin_ctzll(static_cast(mask)); + } + + // Core extraction logic + dst |= ((a >> shift) & 1) << k; + + // Clear the lowest set bit in mask (efficient clearing method) + mask &= (mask - 1); + } + return dst; +#endif +} +#endif +} // namespace bits +} // namespace milvus \ No newline at end of file diff --git a/internal/core/src/common/ComplexVector.cpp b/internal/core/src/common/ComplexVector.cpp new file mode 100644 index 0000000000..eaa33f5edc --- /dev/null +++ b/internal/core/src/common/ComplexVector.cpp @@ -0,0 +1,60 @@ +// 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. + +#include "Vector.h" + +namespace milvus { + +void +BaseVector::prepareForReuse(milvus::VectorPtr& vector, + milvus::vector_size_t size) { + // Guard against null shared_ptr dereference + AssertInfo(vector != nullptr, + "BaseVector::prepareForReuse: vector cannot be null"); + + // Use use_count() instead of unique() (deprecated in C++17, removed in C++20) + if (vector.use_count() != 1) { + // When vector is non-unique (shared by multiple owners), create a new + // instance of the same dynamic type using the virtual factory method to + // preserve the subclass type (e.g., ColumnVector, RowVector) instead of + // creating a BaseVector. + vector = vector->cloneEmpty(size); + } else { + // When vector is unique (only one owner), reuse it by resetting state + // and resizing. This preserves subclass state (e.g., ColumnVector's + // values_, RowVector's children_values_). + vector->prepareForReuse(); + vector->resize(size); + } +} + +void +BaseVector::prepareForReuse() { + null_count_ = std::nullopt; +} + +void +RowVector::resize(milvus::vector_size_t new_size, bool setNotNull) { + BaseVector::resize(new_size, setNotNull); + // Always propagate resize to all children to maintain the invariant that + // child->size() == row->size(). This ensures consistency when shrinking + // as well as growing, preventing children from being longer than the row. + for (auto& child : childrens()) { + child->resize(new_size, setNotNull); + } +} + +} // namespace milvus diff --git a/internal/core/src/common/FieldData.cpp b/internal/core/src/common/FieldData.cpp index 0251cc8add..db36afff8b 100644 --- a/internal/core/src/common/FieldData.cpp +++ b/internal/core/src/common/FieldData.cpp @@ -731,4 +731,75 @@ InitScalarFieldData(const DataType& type, bool nullable, int64_t cap_rows) { } } +void +ResizeScalarFieldData(const DataType& type, + int64_t new_num_rows, + FieldDataPtr& field_data) { + switch (type) { + case DataType::BOOL: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + inner_field_data->resize_field_data(new_num_rows); + return; + } + case DataType::INT8: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + inner_field_data->resize_field_data(new_num_rows); + return; + } + case DataType::INT16: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + inner_field_data->resize_field_data(new_num_rows); + return; + } + case DataType::INT32: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + inner_field_data->resize_field_data(new_num_rows); + return; + } + case DataType::TIMESTAMPTZ: + case DataType::INT64: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + inner_field_data->resize_field_data(new_num_rows); + return; + } + case DataType::FLOAT: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + inner_field_data->resize_field_data(new_num_rows); + return; + } + case DataType::DOUBLE: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + inner_field_data->resize_field_data(new_num_rows); + return; + } + case DataType::STRING: + case DataType::VARCHAR: + case DataType::TEXT: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + AssertInfo(inner_field_data != nullptr, + "Failed to cast field_data to FieldData"); + inner_field_data->resize_field_data(new_num_rows); + return; + } + case DataType::JSON: { + auto inner_field_data = + std::dynamic_pointer_cast>(field_data); + inner_field_data->resize_field_data(new_num_rows); + return; + } + default: + ThrowInfo(DataTypeInvalid, + "ResizeScalarFieldData not support data type " + + GetDataTypeName(type)); + } +} + } // namespace milvus diff --git a/internal/core/src/common/FieldData.h b/internal/core/src/common/FieldData.h index fe9c87e742..48696dcf00 100644 --- a/internal/core/src/common/FieldData.h +++ b/internal/core/src/common/FieldData.h @@ -446,4 +446,9 @@ using FieldDataChannelPtr = std::shared_ptr; FieldDataPtr InitScalarFieldData(const DataType& type, bool nullable, int64_t cap_rows); -} // namespace milvus +void +ResizeScalarFieldData(const DataType& type, + int64_t new_size, + FieldDataPtr& field_data); + +} // namespace milvus \ No newline at end of file diff --git a/internal/core/src/common/FieldDataInterface.h b/internal/core/src/common/FieldDataInterface.h index 39deab12f9..b72ca47668 100644 --- a/internal/core/src/common/FieldDataInterface.h +++ b/internal/core/src/common/FieldDataInterface.h @@ -362,6 +362,7 @@ class FieldDataImpl : public FieldDataBase { data_ = std::move(data); Assert(data_.size() % dim == 0); num_rows_ = data_.size() / dim; + length_ = num_rows_; } explicit FieldDataImpl(size_t dim, @@ -376,6 +377,7 @@ class FieldDataImpl : public FieldDataBase { valid_data_ = std::move(valid_data); Assert(data_.size() % dim == 0); num_rows_ = data_.size() / dim; + length_ = num_rows_; } void diff --git a/internal/core/src/common/Schema.h b/internal/core/src/common/Schema.h index 8317f87b56..d82d8cc0dd 100644 --- a/internal/core/src/common/Schema.h +++ b/internal/core/src/common/Schema.h @@ -291,7 +291,9 @@ class Schema { FieldId get_field_id(const FieldName& field_name) const { - AssertInfo(name_ids_.count(field_name), "Cannot find field_name"); + AssertInfo(name_ids_.count(field_name), + "Cannot find field_name:{}", + field_name.get()); return name_ids_.at(field_name); } @@ -424,6 +426,24 @@ class Schema { const FieldMeta& GetFirstArrayFieldInStruct(const std::string& struct_name) const; + DataType + GetFieldType(const FieldId& field_id) const { + AssertInfo(fields_.count(field_id), + "field_id:{} does not exist in the schema", + field_id.get()); + auto& meta = fields_.at(field_id); + return meta.get_data_type(); + } + + const std::string& + GetFieldName(const FieldId& field_id) const { + AssertInfo(fields_.count(field_id), + "field_id:{} does not exist in the schema", + field_id.get()); + auto& meta = fields_.at(field_id); + return meta.get_name().get(); + } + private: int64_t debug_id = START_USER_FIELDID; std::vector field_ids_; diff --git a/internal/core/src/common/SimdUtil.h b/internal/core/src/common/SimdUtil.h new file mode 100644 index 0000000000..a5fd8eb491 --- /dev/null +++ b/internal/core/src/common/SimdUtil.h @@ -0,0 +1,173 @@ +// 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. +#pragma once + +#include +#include "common/BitUtil.h" + +namespace milvus { + +// Generic fallback implementation for toBitMask when no architecture-specific +// implementation is available +template +int +genericToBitMask(xsimd::batch_bool mask) { + constexpr size_t size = xsimd::batch_bool::size; + int result = 0; + // Convert batch_bool to integer batch and extract each element + // Using select to convert boolean mask to integer values + auto ones = xsimd::batch(T(1)); + auto zeros = xsimd::batch(T(0)); + auto int_batch = xsimd::select(mask, ones, zeros); + // Extract each element and set corresponding bit in result + alignas(A::alignment()) T values[size]; + int_batch.store_aligned(values); + for (size_t i = 0; i < size; ++i) { + if (values[i] != T(0)) { + result |= (1 << i); + } + } + return result; +} + +template +struct BitMask; + +template +struct BitMask { + static constexpr int kAllSet = + milvus::bits::lowMask(xsimd::batch_bool::size); + +#if XSIMD_WITH_AVX2 + static int + toBitMask(xsimd::batch_bool mask, const xsimd::avx2&) { + return _mm256_movemask_epi8(mask); + } +#endif + +#if XSIMD_WITH_SSE2 + static int + toBitMask(xsimd::batch_bool mask, const xsimd::sse2&) { + return _mm_movemask_epi8(mask); + } +#endif + +#if XSIMD_WITH_NEON + static int + toBitMask(xsimd::batch_bool mask, const xsimd::neon&) { + alignas(A::alignment()) static const int8_t kShift[] = { + -7, -6, -5, -4, -3, -2, -1, 0, -7, -6, -5, -4, -3, -2, -1, 0}; + int8x16_t vshift = vld1q_s8(kShift); + uint8x16_t vmask = vshlq_u8(vandq_u8(mask, vdupq_n_u8(0x80)), vshift); + return (vaddv_u8(vget_high_u8(vmask)) << 8) | + vaddv_u8(vget_low_u8(vmask)); + } +#endif + + static int + toBitMask(xsimd::batch_bool mask, const xsimd::generic&) { + return genericToBitMask(mask); + } +}; + +template +struct BitMask { + static constexpr int kAllSet = + milvus::bits::lowMask(xsimd::batch_bool::size); + +#if XSIMD_WITH_AVX2 + static int + toBitMask(xsimd::batch_bool mask, const xsimd::avx2&) { + // There is no intrinsic for extracting high bits of a 16x16 + // vector. Hence take every second bit of the high bits of a 32x1 + // vector. + // + // NOTE: TVL might have a more efficient implementation for this. + return bits::extractBits(_mm256_movemask_epi8(mask), + 0xAAAAAAAA); + } +#endif + +#if XSIMD_WITH_SSE2 + static int + toBitMask(xsimd::batch_bool mask, const xsimd::sse2&) { + return milvus::bits::extractBits(_mm_movemask_epi8(mask), + 0xAAAA); + } +#endif + + static int + toBitMask(xsimd::batch_bool mask, const xsimd::generic&) { + return genericToBitMask(mask); + } +}; + +template +struct BitMask { + static constexpr int kAllSet = + milvus::bits::lowMask(xsimd::batch_bool::size); + +#if XSIMD_WITH_AVX + static int + toBitMask(xsimd::batch_bool mask, const xsimd::avx&) { + return _mm256_movemask_ps(reinterpret_cast<__m256>(mask.data)); + } +#endif + +#if XSIMD_WITH_SSE2 + static int + toBitMask(xsimd::batch_bool mask, const xsimd::sse2&) { + return _mm_movemask_ps(reinterpret_cast<__m128>(mask.data)); + } +#endif + + static int + toBitMask(xsimd::batch_bool mask, const xsimd::generic&) { + return genericToBitMask(mask); + } +}; + +template +struct BitMask { + static constexpr int kAllSet = + milvus::bits::lowMask(xsimd::batch_bool::size); + +#if XSIMD_WITH_AVX + static int + toBitMask(xsimd::batch_bool mask, const xsimd::avx&) { + return _mm256_movemask_pd(reinterpret_cast<__m256d>(mask.data)); + } +#endif + +#if XSIMD_WITH_SSE2 + static int + toBitMask(xsimd::batch_bool mask, const xsimd::sse2&) { + return _mm_movemask_pd(reinterpret_cast<__m128d>(mask.data)); + } +#endif + + static int + toBitMask(xsimd::batch_bool mask, const xsimd::generic&) { + return genericToBitMask(mask); + } +}; + +template +auto +toBitMask(xsimd::batch_bool mask, const A& arch = {}) { + return BitMask::toBitMask(mask, arch); +} +} // namespace milvus diff --git a/internal/core/src/common/Types.cpp b/internal/core/src/common/Types.cpp new file mode 100644 index 0000000000..d6be880215 --- /dev/null +++ b/internal/core/src/common/Types.cpp @@ -0,0 +1,77 @@ +// 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. + +#include "Types.h" +#include "common/Exception.h" +#include "common/EasyAssert.h" + +const RowTypePtr RowType::None = std::make_shared( + std::vector{}, std::vector{}); +namespace milvus { +bool +IsFixedSizeType(DataType type) { + switch (type) { + case DataType::NONE: + return false; + case DataType::BOOL: + return TypeTraits::IsFixedWidth; + case DataType::INT8: + return TypeTraits::IsFixedWidth; + case DataType::INT16: + return TypeTraits::IsFixedWidth; + case DataType::INT32: + return TypeTraits::IsFixedWidth; + case DataType::INT64: + return TypeTraits::IsFixedWidth; + case DataType::TIMESTAMPTZ: + return TypeTraits::IsFixedWidth; + case DataType::FLOAT: + return TypeTraits::IsFixedWidth; + case DataType::DOUBLE: + return TypeTraits::IsFixedWidth; + case DataType::STRING: + return TypeTraits::IsFixedWidth; + case DataType::VARCHAR: + return TypeTraits::IsFixedWidth; + case DataType::ARRAY: + return TypeTraits::IsFixedWidth; + case DataType::JSON: + return TypeTraits::IsFixedWidth; + case DataType::ROW: + return TypeTraits::IsFixedWidth; + case DataType::VECTOR_BINARY: + return TypeTraits::IsFixedWidth; + case DataType::VECTOR_FLOAT: + return TypeTraits::IsFixedWidth; + case DataType::GEOMETRY: + return TypeTraits::IsFixedWidth; + case DataType::TEXT: + return TypeTraits::IsFixedWidth; + case DataType::VECTOR_FLOAT16: + return TypeTraits::IsFixedWidth; + case DataType::VECTOR_BFLOAT16: + return TypeTraits::IsFixedWidth; + case DataType::VECTOR_SPARSE_U32_F32: + return TypeTraits::IsFixedWidth; + case DataType::VECTOR_INT8: + return TypeTraits::IsFixedWidth; + case DataType::VECTOR_ARRAY: + return TypeTraits::IsFixedWidth; + default: + ThrowInfo(DataTypeInvalid, "unknown data type: {}", type); + } +} +} // namespace milvus diff --git a/internal/core/src/common/Types.h b/internal/core/src/common/Types.h index c92a52dcaf..50efb6326f 100644 --- a/internal/core/src/common/Types.h +++ b/internal/core/src/common/Types.h @@ -66,6 +66,8 @@ using int8 = knowhere::int8; using sparse_u32_f32 = knowhere::sparse_u32_f32; // See also: https://github.com/milvus-io/milvus-proto/blob/master/proto/schema.proto +using vector_size_t = int32_t; + enum class DataType { NONE = 0, BOOL = 1, @@ -806,6 +808,42 @@ struct TypeTraits { static constexpr const char* Name = "VECTOR_FLOAT"; }; +template <> +struct TypeTraits { + using NativeType = float16; + static constexpr DataType TypeKind = DataType::VECTOR_FLOAT16; + static constexpr bool IsPrimitiveType = false; + static constexpr bool IsFixedWidth = true; + static constexpr const char* Name = "VECTOR_FLOAT16"; +}; + +template <> +struct TypeTraits { + using NativeType = bfloat16; + static constexpr DataType TypeKind = DataType::VECTOR_BFLOAT16; + static constexpr bool IsPrimitiveType = false; + static constexpr bool IsFixedWidth = true; + static constexpr const char* Name = "VECTOR_BFLOAT16"; +}; + +template <> +struct TypeTraits { + using NativeType = void; + static constexpr DataType TypeKind = DataType::VECTOR_SPARSE_U32_F32; + static constexpr bool IsPrimitiveType = false; + static constexpr bool IsFixedWidth = false; + static constexpr const char* Name = "VECTOR_SPARSE_U32_F32"; +}; + +template <> +struct TypeTraits { + using NativeType = int8_t; + static constexpr DataType TypeKind = DataType::VECTOR_INT8; + static constexpr bool IsPrimitiveType = false; + static constexpr bool IsFixedWidth = true; + static constexpr const char* Name = "VECTOR_INT8"; +}; + template <> struct TypeTraits { using NativeType = void; @@ -856,6 +894,9 @@ vector_bytes_per_element(const DataType data_type, int64_t dim) { } } +bool +IsFixedSizeType(DataType type); + } // namespace milvus template <> struct fmt::formatter : formatter { @@ -1382,3 +1423,91 @@ struct fmt::formatter : fmt::formatter { return fmt::formatter::format(name, ctx); } }; + +using column_index_t = uint32_t; +class RowType final { + public: + RowType(std::vector&& names, + std::vector&& types) + : names_(std::move(names)), columns_types_(std::move(types)) { + AssertInfo(names_.size() == columns_types_.size(), + "Name count:{} and column count:{} must be the same", + names_.size(), + columns_types_.size()); + }; + + static const std::shared_ptr None; + + column_index_t + GetChildIndex(const std::string& name) const { + std::optional idx; + for (auto i = 0; i < names_.size(); i++) { + if (names_[i] == name) { + idx = i; + break; + } + } + AssertInfo(idx.has_value(), + "Cannot find target column in the rowType list"); + return idx.value(); + } + + milvus::DataType + column_type(uint32_t idx) const { + AssertInfo(idx >= 0 && idx < names_.size(), + "try to get column type from row type with wrong index:{}", + idx); + return columns_types_.at(idx); + } + + size_t + column_count() const { + return names_.size(); + } + + private: + const std::vector names_; + const std::vector columns_types_; +}; + +using RowTypePtr = std::shared_ptr; + +#define MILVUS_DYNAMIC_TYPE_DISPATCH(TEMPLATE_FUNC, DATETYPE, ...) \ + MILVUS_DYNAMIC_TYPE_DISPATCH_IMPL(TEMPLATE_FUNC, , DATETYPE, __VA_ARGS__) + +#define MILVUS_DYNAMIC_TYPE_DISPATCH_IMPL(PREFIX, SUFFIX, DATATYPE, ...) \ + [&]() { \ + switch (DATATYPE) { \ + case milvus::DataType::BOOL: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::INT8: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::INT16: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::INT32: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::INT64: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::TIMESTAMPTZ: \ + return PREFIX SUFFIX( \ + __VA_ARGS__); \ + case milvus::DataType::FLOAT: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::DOUBLE: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::VARCHAR: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::STRING: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::JSON: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::ARRAY: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + case milvus::DataType::ROW: \ + return PREFIX SUFFIX(__VA_ARGS__); \ + default: \ + ThrowInfo(milvus::DataTypeInvalid, \ + "UnsupportedDataType for " \ + "MILVUS_DYNAMIC_TYPE_DISPATCH_IMPL"); \ + } \ + }() diff --git a/internal/core/src/common/Utils.h b/internal/core/src/common/Utils.h index a891740c3e..0e956e45d7 100644 --- a/internal/core/src/common/Utils.h +++ b/internal/core/src/common/Utils.h @@ -20,9 +20,11 @@ #include #include #include +#include #include #include #include +#include #include #include "common/Consts.h" @@ -352,4 +354,222 @@ class Defer { #define DeferLambda(fn) Defer Defer_##__COUNTER__(fn); +template +FOLLY_ALWAYS_INLINE int +comparePrimitiveAsc(const T& left, const T& right) { + if constexpr (std::is_floating_point::value) { + bool leftNan = std::isnan(left); + bool rightNan = std::isnan(right); + if (leftNan) { + return rightNan ? 0 : 1; + } + if (rightNan) { + return -1; + } + } + return left < right ? -1 : left == right ? 0 : 1; +} + +inline std::string +lowerString(const std::string& str) { + std::string ret; + ret.resize(str.size()); + std::transform(str.begin(), str.end(), ret.begin(), [](unsigned char c) { + return std::tolower(c); + }); + return ret; +} + +template +T +checkPlus(const T& a, const T& b, const char* typeName = "integer") { + static_assert(std::is_integral_v, "checkPlus requires integral type"); +#if defined(__GNUC__) || defined(__clang__) + // Use compiler builtin for GCC/Clang + T result; + bool overflow = __builtin_add_overflow(a, b, &result); + if (UNLIKELY(overflow)) { + ThrowInfo(DataTypeInvalid, "{} overflow: {} + {}", typeName, a, b); + } + return result; +#else + // Portable fallback for MSVC and other compilers + constexpr T max_val = std::numeric_limits::max(); + constexpr T min_val = std::numeric_limits::min(); + + bool overflow = false; + if (b > 0) { + // Positive addition: check if a > max - b + if (a > max_val - b) { + overflow = true; + } + } else if (b < 0) { + // Negative addition: check if a < min - b + if (a < min_val - b) { + overflow = true; + } + } + // If b == 0, no overflow possible + + if (UNLIKELY(overflow)) { + ThrowInfo(DataTypeInvalid, "{} overflow: {} + {}", typeName, a, b); + } + return a + b; +#endif +} + +template +T +checkedMultiply(const T& a, const T& b, const char* typeName = "integer") { + static_assert(std::is_integral_v, + "checkedMultiply requires integral type"); +#if defined(__GNUC__) || defined(__clang__) + // Use compiler builtin for GCC/Clang + T result; + bool overflow = __builtin_mul_overflow(a, b, &result); + if (UNLIKELY(overflow)) { + ThrowInfo(DataTypeInvalid, "{} overflow: {} * {}", typeName, a, b); + } + return result; +#else + // Portable fallback for MSVC and other compilers + constexpr T max_val = std::numeric_limits::max(); + constexpr T min_val = std::numeric_limits::min(); + + // Handle zero case: no overflow if either operand is zero + if (a == 0 || b == 0) { + return 0; + } + + bool overflow = false; + + if constexpr (std::is_signed_v) { + // Signed type: handle MIN edge case and use absolute values + if (a == min_val) { + // Special case: if a is MIN, only safe multiplications are with + // 0 (already handled), 1, or -1 + if (b != 1 && b != -1) { + overflow = true; + } + } else if (b == min_val) { + // Special case: if b is MIN, only safe multiplications are with + // 0 (already handled), 1, or -1 + if (a != 1 && a != -1) { + overflow = true; + } + } else { + // Use division-based check: |a| > max_val / |b| implies overflow + // Compute absolute values carefully to avoid overflow + T abs_a = (a < 0) ? -a : a; + T abs_b = (b < 0) ? -b : b; + + // Check: abs_a > max_val / abs_b + // This works because both operands are not MIN (handled above) + // and not zero (handled above) + if (abs_a > max_val / abs_b) { + overflow = true; + } + } + } else { + // Unsigned type: simpler case + // Check: a > max_val / b + if (a > max_val / b) { + overflow = true; + } + } + + if (UNLIKELY(overflow)) { + ThrowInfo(DataTypeInvalid, "{} overflow: {} * {}", typeName, a, b); + } + return a * b; +#endif +} + +inline const char* const KSum = "sum"; +inline const char* const KMin = "min"; +inline const char* const KMax = "max"; +inline const char* const KCount = "count"; +inline const char* const KAvg = "avg"; + +inline DataType +GetAggResultType(std::string func_name, DataType input_type) { + if (func_name == KSum) { + switch (input_type) { + case DataType::INT8: + case DataType::INT16: + case DataType::INT32: + case DataType::INT64: { + return DataType::INT64; + } + case DataType::TIMESTAMPTZ: { + return DataType::TIMESTAMPTZ; + } + case DataType::FLOAT: { + return DataType::DOUBLE; + } + case DataType::DOUBLE: { + return DataType::DOUBLE; + } + default: { + ThrowInfo(DataTypeInvalid, + "Unsupported data type for type:{}", + input_type); + } + } + } + if (func_name == KAvg) { + switch (input_type) { + case DataType::INT8: + case DataType::INT16: + case DataType::INT32: + case DataType::INT64: + case DataType::FLOAT: + case DataType::DOUBLE: { + return DataType::DOUBLE; + } + default: { + ThrowInfo(DataTypeInvalid, + "Unsupported data type for {} aggregation: {}", + func_name, + input_type); + } + } + } + if (func_name == KMin || func_name == KMax) { + // min/max keep the original scalar type. + switch (input_type) { + case DataType::INT8: + case DataType::INT16: + case DataType::INT32: + case DataType::INT64: + case DataType::FLOAT: + case DataType::DOUBLE: + case DataType::VARCHAR: + case DataType::STRING: + case DataType::TEXT: + case DataType::TIMESTAMPTZ: { + return input_type; + } + default: { + ThrowInfo(DataTypeInvalid, + "Unsupported data type for {} aggregation: {}", + func_name, + input_type); + } + } + } + if (func_name == KCount) { + return DataType::INT64; + } + ThrowInfo(OpTypeInvalid, "Unsupported func type:{}", func_name); +} + +inline int32_t +Align(int32_t number, int32_t alignment) { + AssertInfo(alignment > 0 && (alignment & (alignment - 1)) == 0, + "Alignment must be a power of 2, got {}", + alignment); + return (number + alignment - 1) & ~(alignment - 1); +} + } // namespace milvus diff --git a/internal/core/src/common/Vector.h b/internal/core/src/common/Vector.h index 052fe5cd68..bd577237a7 100644 --- a/internal/core/src/common/Vector.h +++ b/internal/core/src/common/Vector.h @@ -26,6 +26,8 @@ #include "common/Types.h" namespace milvus { +class BaseVector; +using VectorPtr = std::shared_ptr; /** * @brief base class for different type vector @@ -41,15 +43,45 @@ class BaseVector { virtual ~BaseVector() = default; int64_t - size() { + size() const { return length_; } DataType - type() { + type() const { return type_kind_; } + int32_t + elementSize() const { + return GetDataTypeSize(type_kind_); + }; + + size_t + nullCount() const { + return null_count_.has_value() ? null_count_.value() : 0; + } + + virtual void + resize(vector_size_t newSize, bool setNotNull = true) { + length_ = newSize; + } + + static void + prepareForReuse(VectorPtr& vector, vector_size_t size); + + /// Resets non-reusable buffers and updates child vectors by calling + /// BaseVector::prepareForReuse. + /// Base implementation checks and resets nulls buffer if needed. Keeps the + /// nulls buffer if singly-referenced, mutable and has at least one null bit set. + virtual void + prepareForReuse(); + + /// Creates a new empty instance of the same dynamic type with the given size. + /// Used when cloning non-unique vectors to preserve the dynamic type. + virtual VectorPtr + cloneEmpty(vector_size_t size) const = 0; + protected: DataType type_kind_; size_t length_; @@ -57,8 +89,6 @@ class BaseVector { std::optional null_count_; }; -using VectorPtr = std::shared_ptr; - /** * SimpleVector abstracts over various Columnar Storage Formats, * it is used in custom functions. @@ -110,6 +140,23 @@ class ColumnVector final : public SimpleVector { std::move(bitmap)); } + ColumnVector(FieldDataPtr&& value, TargetBitmap&& valid_bitmap) + : SimpleVector(value ? value->get_data_type() : DataType::NONE, + value ? value->Length() : 0), + is_bitmap_(false), + valid_values_(std::move(valid_bitmap)) { + AssertInfo(value, "ColumnVector value cannot be null"); + values_ = std::move(value); + // Compute null_count_ from valid_bitmap: count false bits (nulls) + size_t nulls = 0; + for (size_t i = 0; i < valid_values_.size(); ++i) { + if (!valid_values_[i]) { + ++nulls; + } + } + null_count_ = (nulls > 0) ? std::optional(nulls) : std::nullopt; + } + virtual ~ColumnVector() override { values_.reset(); valid_values_.reset(); @@ -120,6 +167,43 @@ class ColumnVector final : public SimpleVector { return reinterpret_cast(GetRawData()) + index * size_of_element; } + template + T + ValueAt(size_t index) const { + AssertInfo(index < length_, "ValueAt index out of range"); + return *(reinterpret_cast(GetRawData()) + index); + } + + template + void + SetValueAt(size_t index, const T& value) { + AssertInfo(index < length_, "SetValueAt index out of range"); + *(reinterpret_cast(values_->Data()) + index) = value; + } + + void + nullAt(size_t index) { + // Update null_count_ when setting a value to null + if (valid_values_[index]) { + // Was valid, now null: increment count + null_count_ = null_count_.value_or(0) + 1; + } + valid_values_.set(index, false); + } + + void + clearNullAt(size_t index) { + // Update null_count_ when clearing a null value + if (!valid_values_[index]) { + // Was null, now valid: decrement count + if (null_count_.has_value()) { + null_count_ = + (null_count_.value() > 0) ? null_count_.value() - 1 : 0; + } + } + valid_values_.set(index, true); + } + bool ValidAt(size_t index) override { return valid_values_[index]; @@ -146,6 +230,92 @@ class ColumnVector final : public SimpleVector { return is_bitmap_; } + void + resize(vector_size_t new_size, bool setNotNull = true) override { + AssertInfo(!is_bitmap_, "Cannot resize bitmap column vector"); + auto old_size = length_; + BaseVector::resize(new_size, setNotNull); + ResizeScalarFieldData(type(), new_size, values_); + + if (new_size > old_size) { + // Growing: new bits are added + valid_values_.resize(new_size); + if (setNotNull) { + // All new bits are set to valid (true), no nulls added + for (auto i = old_size; i < new_size; ++i) { + valid_values_.set(i, true); + } + // null_count_ unchanged (no new nulls) + } else { + // New bits default to null (false) + for (auto i = old_size; i < new_size; ++i) { + valid_values_.set(i, false); + } + // Update null_count_: add (new_size - old_size) nulls + null_count_ = null_count_.value_or(0) + (new_size - old_size); + } + } else if (new_size < old_size) { + // Shrinking: need to count nulls in removed range and subtract + size_t removed_nulls = 0; + for (auto i = new_size; i < old_size; ++i) { + if (!valid_values_[i]) { + ++removed_nulls; + } + } + valid_values_.resize(new_size); + // Update null_count_: subtract removed nulls + if (removed_nulls > 0 && null_count_.has_value()) { + null_count_ = (null_count_.value() >= removed_nulls) + ? null_count_.value() - removed_nulls + : 0; + } + } + // If new_size == old_size, no change needed + } + + void + append(const ColumnVector& other) { + // Validate that both vectors have the same type + AssertInfo(type() == other.type(), + "Cannot append ColumnVector with different type: {} != {}", + static_cast(type()), + static_cast(other.type())); + + auto old_size = length_; + values_->FillFieldData(other.GetRawData(), other.size()); + length_ += other.size(); + valid_values_.resize(length_); + + // Copy validity from other + for (size_t i = 0; i < other.size(); ++i) { + valid_values_.set(old_size + i, other.valid_values_[i]); + } + + // Update null_count_ by accumulating the original null_count_ and other's null count + size_t other_nulls = 0; + if (other.null_count_.has_value()) { + // Use cached null_count_ if available + other_nulls = other.null_count_.value(); + } else { + // Compute null count from other.valid_values_ if not cached + for (size_t i = 0; i < other.size(); ++i) { + if (!other.valid_values_[i]) { + ++other_nulls; + } + } + } + + // Accumulate null counts + size_t total_nulls = null_count_.value_or(0) + other_nulls; + null_count_ = (total_nulls > 0) ? std::optional{total_nulls} + : std::nullopt; + } + + VectorPtr + cloneEmpty(vector_size_t size) const override { + return std::make_shared(type(), size, std::nullopt); + } + private: bool is_bitmap_; // TODO: remove the field after implementing BitmapVector FieldDataPtr values_; @@ -186,6 +356,12 @@ class ConstantVector : public SimpleVector { return is_null_; } + VectorPtr + cloneEmpty(vector_size_t size) const override { + return std::make_shared>( + type(), size, val_, std::nullopt); + } + private: T val_; bool is_null_; @@ -226,6 +402,18 @@ class RowVector : public BaseVector { } } + RowVector(const RowTypePtr& rowType, + size_t size, + std::optional nullCount = std::nullopt) + : BaseVector(DataType::ROW, size, nullCount) { + auto column_count = rowType->column_count(); + for (auto i = 0; i < column_count; i++) { + auto column_type = rowType->column_type(i); + children_values_.emplace_back( + std::make_shared(column_type, size)); + } + } + const std::vector& childrens() const { return children_values_; @@ -237,6 +425,19 @@ class RowVector : public BaseVector { return children_values_[index]; } + void + resize(vector_size_t new_size, bool setNotNull = true) override; + + VectorPtr + cloneEmpty(vector_size_t size) const override { + // Extract data types from existing children to preserve structure + std::vector data_types; + for (const auto& child : children_values_) { + data_types.push_back(child->type()); + } + return std::make_shared(data_types, size, std::nullopt); + } + private: std::vector children_values_; }; diff --git a/internal/core/src/common/float_util_c.h b/internal/core/src/common/float_util_c.h new file mode 100644 index 0000000000..093e1b0f8b --- /dev/null +++ b/internal/core/src/common/float_util_c.h @@ -0,0 +1,40 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +namespace milvus { + +template ::value, bool> = true> +struct NaNAwareHash { + std::size_t + operator()(const FLOAT& val) const noexcept { + static const std::size_t kNanHash = + folly::hasher{}(std::numeric_limits::quiet_NaN()); + if (std::isnan(val)) { + return kNanHash; + } + return folly::hasher{}(val); + } +}; + +} // namespace milvus \ No newline at end of file diff --git a/internal/core/src/exec/Driver.cpp b/internal/core/src/exec/Driver.cpp index b267932eb6..679faab661 100644 --- a/internal/core/src/exec/Driver.cpp +++ b/internal/core/src/exec/Driver.cpp @@ -20,10 +20,10 @@ #include #include +#include #include "common/EasyAssert.h" #include "exec/operator/CallbackSink.h" -#include "exec/operator/CountNode.h" #include "exec/operator/FilterBitsNode.h" #include "exec/operator/IterativeFilterNode.h" #include "exec/operator/MvccNode.h" @@ -31,9 +31,11 @@ #include "exec/operator/RescoresNode.h" #include "exec/operator/VectorSearchNode.h" #include "exec/operator/RandomSampleNode.h" -#include "exec/operator/GroupByNode.h" #include "exec/operator/ElementFilterNode.h" #include "exec/operator/ElementFilterBitsNode.h" +#include "exec/operator/SearchGroupByNode.h" +#include "exec/operator/AggregationNode.h" +#include "exec/operator/ProjectNode.h" #include "exec/Task.h" #include "plan/PlanNode.h" @@ -76,24 +78,30 @@ DriverFactory::CreateDriver(std::unique_ptr ctx, tracer::AddEvent("create_operator: MvccNode"); operators.push_back( std::make_unique(id, ctx.get(), mvccnode)); - } else if (auto countnode = - std::dynamic_pointer_cast( - plannode)) { - tracer::AddEvent("create_operator: CountNode"); - operators.push_back( - std::make_unique(id, ctx.get(), countnode)); } else if (auto vectorsearchnode = std::dynamic_pointer_cast( plannode)) { tracer::AddEvent("create_operator: VectorSearchNode"); operators.push_back(std::make_unique( id, ctx.get(), vectorsearchnode)); - } else if (auto groupbynode = - std::dynamic_pointer_cast( + } else if (auto searchGroupByNode = + std::dynamic_pointer_cast( plannode)) { - tracer::AddEvent("create_operator: GroupByNode"); + tracer::AddEvent("create_operator: SearchGroupByNode"); + operators.push_back(std::make_unique( + id, ctx.get(), searchGroupByNode)); + } else if (auto queryGroupByNode = + std::dynamic_pointer_cast( + plannode)) { + tracer::AddEvent("create_operator: AggregationNode"); + operators.push_back(std::make_unique( + id, ctx.get(), queryGroupByNode)); + } else if (auto projectNode = + std::dynamic_pointer_cast( + plannode)) { + tracer::AddEvent("create_operator: ProjectNode"); operators.push_back( - std::make_unique(id, ctx.get(), groupbynode)); + std::make_unique(id, ctx.get(), projectNode)); } else if (auto samplenode = std::dynamic_pointer_cast( plannode)) { @@ -171,6 +179,31 @@ Driver::Run(std::shared_ptr self) { } } +void +Driver::initializeOperators() { + // Atomically check and set: only the first thread to call this will + // get false (the previous value) and proceed with initialization. + // Other threads will get true and skip initialization. + // Use memory barriers to ensure initialization writes are visible + // before the flag is seen by other threads. + if (!operatorsInitialized_.exchange(true, std::memory_order_acq_rel)) { + // This thread won the initialization race. Perform initialization. + for (auto& op : operators_) { + op->initialize(); + } + // Use release semantics to ensure all initialization writes are + // visible to other threads before they see the flag as true. + // The flag was already set to true by exchange above, but this + // barrier ensures all writes from initialization are visible. + std::atomic_thread_fence(std::memory_order_release); + } else { + // Another thread is initializing or has completed initialization. + // Use acquire semantics to ensure we see all initialization writes + // after the flag becomes true. + operatorsInitialized_.load(std::memory_order_acquire); + } +} + void Driver::Init(std::unique_ptr ctx, std::vector> operators) { @@ -230,13 +263,13 @@ Driver::RunInternal(std::shared_ptr& self, std::shared_ptr& blocking_state, RowVectorPtr& result) { try { + initializeOperators(); int num_operators = operators_.size(); ContinueFuture future; for (;;) { for (int32_t i = num_operators - 1; i >= 0; --i) { auto op = operators_[i].get(); - current_operator_index_ = i; CALL_OPERATOR( blocking_reason_ = op->IsBlocked(&future), op, "IsBlocked"); diff --git a/internal/core/src/exec/Driver.h b/internal/core/src/exec/Driver.h index ef513b88de..c92e4efb76 100644 --- a/internal/core/src/exec/Driver.h +++ b/internal/core/src/exec/Driver.h @@ -16,6 +16,7 @@ #pragma once +#include #include #include #include @@ -219,6 +220,11 @@ class Driver : public std::enable_shared_from_this { EnqueueInternal() { } + /// Invoked to initialize the operators from this driver once on its first + /// execution. + void + initializeOperators(); + static void Run(std::shared_ptr self); @@ -238,6 +244,8 @@ class Driver : public std::enable_shared_from_this { size_t current_operator_index_{0}; + std::atomic_bool operatorsInitialized_{false}; + BlockingReason blocking_reason_{BlockingReason::kNotBlocked}; friend struct DriverFactory; diff --git a/internal/core/src/exec/HashTable.cpp b/internal/core/src/exec/HashTable.cpp new file mode 100644 index 0000000000..57808e6b27 --- /dev/null +++ b/internal/core/src/exec/HashTable.cpp @@ -0,0 +1,275 @@ +// 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. + +#include "HashTable.h" +#include +#include +#include +#include "common/SimdUtil.h" + +namespace milvus { +namespace exec { +void +BaseHashTable::prepareForGroupProbe(HashLookup& lookup, + const RowVectorPtr& input) { + auto& hashers = lookup.hashers_; + int numKeys = hashers.size(); + // set up column vector to each column + for (auto i = 0; i < numKeys; i++) { + auto& hasher = hashers[i]; + auto column_idx = hasher->ChannelIndex(); + ColumnVectorPtr column_ptr = + std::dynamic_pointer_cast(input->child(column_idx)); + AssertInfo(column_ptr != nullptr, + "Failed to get column vector from row vector input"); + hashers[i]->setColumnData(column_ptr); + } + lookup.reset(input->size()); + + const auto mode = hashMode(); + for (auto i = 0; i < hashers.size(); i++) { + if (mode == BaseHashTable::HashMode::kHash) { + hashers[i]->hash(i > 0, lookup.hashes_); + } else { + ThrowInfo( + milvus::OpTypeInvalid, + "Not support target hashMode, only support kHash for now"); + } + } +} + +class ProbeState { + public: + enum class Operation { kProbe, kInsert, kErase }; + static constexpr int32_t kFullMask = 0xffff; + + int32_t + row() const { + return row_; + } + + template + inline void + preProbe(const Table& table, uint64_t hash, int32_t row) { + row_ = row; + bucketOffset_ = table.bucketOffset(hash); + const auto tag = BaseHashTable::hashTag(hash); + wantedTags_ = BaseHashTable::TagVector::broadcast(tag); + group_ = nullptr; + __builtin_prefetch(reinterpret_cast(table.table_) + + bucketOffset_); + } + + template + inline void + firstProbe(const Table& table) { + tagsInTable_ = BaseHashTable::loadTags( + reinterpret_cast(table.table_), bucketOffset_); + hits_ = milvus::toBitMask(tagsInTable_ == wantedTags_); + if (hits_) { + loadNextHit(table); + } + } + + template + inline char* + fullProbe(Table& table, Compare compare, Insert insert) { + AssertInfo(op == Operation::kInsert, + "Only support insert operation for group cases"); + if (group_ && compare(group_, row_)) { + return group_; + } + const auto kEmptyGroup = BaseHashTable::TagVector::broadcast(0); + for (int64_t numProbedBuckets = 0; + numProbedBuckets < table.numBuckets(); + ++numProbedBuckets) { + while (hits_ > 0) { + loadNextHit(table); + if (compare(group_, row_)) { + return group_; + } + } + + uint16_t empty = + milvus::toBitMask(tagsInTable_ == kEmptyGroup) & kFullMask; + // if there are still empty slot available, try to insert into existing empty slot or tombstone slot + if (empty > 0) { + auto pos = milvus::bits::getAndClearLastSetBit(empty); + return insert(row_, bucketOffset_ + pos); + } + bucketOffset_ = table.nextBucketOffset(bucketOffset_); + tagsInTable_ = table.loadTags(bucketOffset_); + hits_ = milvus::toBitMask(tagsInTable_ == wantedTags_); + } + ThrowInfo(UnexpectedError, + "Slots in hash table is not enough for hash operation, fail " + "the request"); + } + + private: + static constexpr uint8_t kNotSet = 0xff; + template + inline void + loadNextHit(Table& table) { + const int32_t hit = milvus::bits::getAndClearLastSetBit(hits_); + group_ = table.row(bucketOffset_, hit); + __builtin_prefetch(group_); + } + + char* group_; + BaseHashTable::TagVector wantedTags_; + BaseHashTable::TagVector tagsInTable_; + int32_t row_; + int64_t bucketOffset_; + BaseHashTable::MaskType hits_; + //uint8_t indexInTags_ = kNotSet; +}; + +void +HashTable::allocateTables(uint64_t size) { + AssertInfo(milvus::bits::isPowerOfTwo(size), + "Size:{} for allocating tables must be a power of two", + size); + AssertInfo(size > 0, + "Size:{} for allocating tables must be larger than zero", + size); + // Free existing table if present + if (table_ != nullptr) { + ::operator delete(table_, std::align_val_t(64)); + table_ = nullptr; + } + capacity_ = size; + const uint64_t byteSize = capacity_ * tableSlotSize(); + AssertInfo(byteSize % kBucketSize == 0, + "byteSize:{} for hashTable must be a multiple of kBucketSize:{}", + byteSize, + kBucketSize); + numBuckets_ = byteSize / kBucketSize; + sizeMask_ = byteSize - 1; + sizeBits_ = __builtin_popcountll(sizeMask_); + bucketOffsetMask_ = sizeMask_ & ~(kBucketSize - 1); + // The total size is 8 bytes per slot, in groups of 16 slots with 16 bytes of + // tags and 16 * 6 bytes of pointers and a padding of 16 bytes to round up the + // cache line. + // Allocate aligned memory (64-byte cache line alignment) for the table buffer. + // TODO support memory pool here to avoid OOM + table_ = static_cast(::operator new(byteSize, std::align_val_t(64))); + std::memset(table_, 0, byteSize); +} + +void +HashTable::checkSizeAndAllocateTable(int32_t numNew) { + AssertInfo(capacity_ == 0 || capacity_ > numDistinct_, + "capacity_ {}, numDistinct {}", + capacity_, + numDistinct_); + if (table_ == nullptr || capacity_ == 0) { + const auto newSize = newHashTableEntriesNumber(numDistinct_, numNew); + allocateTables(newSize); + } +} + +bool +HashTable::compareKeys(const char* group, + milvus::exec::HashLookup& lookup, + milvus::vector_size_t row) { + int32_t numKeys = lookup.hashers_.size(); + for (int32_t i = 0; i < numKeys; i++) { + auto& hasher = lookup.hashers_[i]; + if (!rows_->equals( + group, rows()->columnAt(i), hasher->columnData(), row)) { + return false; + } + } + return true; +} + +void +HashTable::storeKeys(milvus::exec::HashLookup& lookup, + milvus::vector_size_t row) { + for (int32_t i = 0; i < lookup.hashers_.size(); i++) { + auto& hasher = lookup.hashers_[i]; + rows_->store(hasher->columnData(), row, lookup.hits_[row], i); + } +} + +void +HashTable::storeRowPointer(uint64_t index, uint64_t hash, char* row) { + const int64_t bktOffset = bucketOffset(index); + auto* bucket = bucketAt(bktOffset); + const auto slotIndex = index & (sizeof(TagVector) - 1); + bucket->setTag(slotIndex, hashTag(hash)); + bucket->setPointer(slotIndex, row); +} + +char* +HashTable::insertEntry(milvus::exec::HashLookup& lookup, + uint64_t index, + milvus::vector_size_t row) { + char* group = rows_->newRow(); + lookup.hits_[row] = group; + storeKeys(lookup, row); + storeRowPointer(index, lookup.hashes_[row], group); + numDistinct_++; + lookup.newGroups_.push_back(row); + return group; +} + +FOLLY_ALWAYS_INLINE void +HashTable::fullProbe(HashLookup& lookup, ProbeState& state) { + constexpr ProbeState::Operation op = ProbeState::Operation::kInsert; + lookup.hits_[state.row()] = state.fullProbe( + *this, + [&](char* group, int32_t row) { + return compareKeys(group, lookup, row); + }, + [&](int32_t row, uint64_t index) { + return insertEntry(lookup, index, row); + }); +} + +void +HashTable::groupProbe(milvus::exec::HashLookup& lookup) { + AssertInfo(hashMode_ == HashMode::kHash, "Only support kHash mode for now"); + checkSizeAndAllocateTable(0); + ProbeState state; + for (int32_t idx = 0; idx < lookup.hashes_.size(); idx++) { + state.preProbe(*this, lookup.hashes_[idx], idx); + state.firstProbe(*this); + fullProbe(lookup, state); + } +} + +void +HashTable::setHashMode(HashMode mode, int32_t numNew) { + // TODO set hash mode kArray/kHash/kNormalizedKey +} + +void +HashTable::clear(bool freeTable) { + if (table_) { + ::operator delete(table_, std::align_val_t(64)); + table_ = nullptr; + } + numDistinct_ = 0; + capacity_ = 0; + numBuckets_ = 0; + sizeMask_ = 0; + bucketOffsetMask_ = 0; +} + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/HashTable.h b/internal/core/src/exec/HashTable.h new file mode 100644 index 0000000000..6a471a0eac --- /dev/null +++ b/internal/core/src/exec/HashTable.h @@ -0,0 +1,344 @@ +// 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. +#pragma once +#include +#include +#include +#include + +#include "VectorHasher.h" +#include "exec/operator/query-agg/RowContainer.h" +#include "xsimd/xsimd.hpp" +#include "common/BitUtil.h" + +namespace milvus { +namespace exec { + +struct HashLookup { + explicit HashLookup( + const std::vector>& hashers) + : hashers_(hashers) { + } + + void + reset(vector_size_t size) { + rows_.resize(size); + hashes_.resize(size); + hits_.resize(size); + newGroups_.clear(); + } + + /// One entry per group-by + const std::vector>& hashers_; + + /// Set of row numbers of row to probe. + std::vector rows_; + + /// Hashes or value IDs for rows in 'rows'. Not aligned with 'rows'. Index is + /// the row number. + std::vector hashes_; + + /// Contains one entry for each row in 'rows'. Index is the row number. + /// For groupProbe, a pointer to an existing or new row with matching grouping + /// keys. + std::vector hits_; + + /// For groupProbe, row numbers for which a new entry was inserted (didn't + /// exist before the groupProbe). Empty for joinProbe. + std::vector newGroups_; +}; + +class BaseHashTable { + public: +#if XSIMD_WITH_SSE2 + using TagVector = xsimd::batch; +#elif XSIMD_WITH_NEON + using TagVector = xsimd::batch; +#else + // Fallback for non-SIMD builds + using TagVector = std::array; +#endif + using MaskType = uint16_t; + + enum class HashMode { kHash, kArray, kNormalizedKey }; + + explicit BaseHashTable(std::vector>&& hashers) + : hashers_(std::move(hashers)) { + } + + virtual ~BaseHashTable() = default; + + RowContainer* + rows() const { + return rows_.get(); + } + + /// Extracts a 7 bit tag from a hash number. The high bit is always set. + static uint8_t + hashTag(uint64_t hash) { + // This is likely all 0 for small key types (<= 32 bits). Not an issue + // because small types have a range that makes them normalized key cases. + // If there are multiple small type keys, they are mixed which makes them a + // 64 bit hash. Normalized keys are mixed before being used as hash + // numbers. + return static_cast(hash >> 38) | 0x80; + } + + static FOLLY_ALWAYS_INLINE size_t + tableSlotSize() { + // Each slot is 8 bytes. + return sizeof(void*); + } + + static TagVector + loadTags(uint8_t* tags, int64_t tagIndex) { + auto src = tags + tagIndex; +#if XSIMD_WITH_SSE2 + return TagVector( + _mm_loadu_si128(reinterpret_cast<__m128i const*>(src))); +#elif XSIMD_WITH_NEON + return TagVector(vld1q_u8(src)); +#else + // Generic fallback: load bytes into std::array + TagVector result; + std::memcpy(result.data(), src, result.size()); + return result; +#endif + } + + const std::vector>& + hashers() const { + return hashers_; + } + + /// Returns the hash mode. This is needed for the caller to calculate + /// the hash numbers using the appropriate method of the + /// VectorHashers of 'this'. + virtual HashMode + hashMode() const = 0; + + virtual void + setHashMode(HashMode mode, int32_t numNew) = 0; + + /// Disables use of array or normalized key hash modes. + void + forceGenericHashMode() { + setHashMode(HashMode::kHash, 0); + } + + /// Populates 'hashes' and 'rows' fields in 'lookup' in preparation for + /// 'groupProbe' call. Rehashes the table if necessary. Uses lookup.hashes to + /// decode grouping keys from 'input'. + void + prepareForGroupProbe(HashLookup& lookup, const RowVectorPtr& input); + + /// Finds or creates a group for each key in 'lookup'. The keys are + /// returned in 'lookup.hits'. + virtual void + groupProbe(HashLookup& lookup) = 0; + + virtual void + clear(bool freeTable = false) = 0; + + protected: + std::vector> hashers_; + std::unique_ptr rows_; +}; + +class ProbeState; + +class HashTable : public BaseHashTable { + public: + HashTable(std::vector>&& hashers, + const std::vector& accumulators) + : BaseHashTable(std::move(hashers)) { + std::vector keyTypes; + for (auto& hasher : hashers_) { + keyTypes.push_back(hasher->ChannelDataType()); + } + hashMode_ = HashMode::kHash; + rows_ = std::make_unique(keyTypes, accumulators); + }; + + ~HashTable() override { + clear(); + } + + void + setHashMode(HashMode mode, int32_t numNew) override; + + void + groupProbe(HashLookup& lookup) override; + + // The table in non-kArray mode has a power of two number of buckets each with + // 16 slots. Each slot has a 1 byte tag (a field of hash number) and a 48 bit + // pointer. All the tags are in a 16 byte SIMD word followed by the 6 byte + // pointers. There are 16 bytes of padding at the end to make the bucket + // occupy exactly two (64 bytes) cache lines. + class Bucket { + public: + uint8_t + tagAt(int32_t slotIndex) { + return reinterpret_cast(&tags_)[slotIndex]; + } + + char* + pointerAt(int32_t slotIndex) { + return reinterpret_cast( + *reinterpret_cast( + &pointers_[kPointerSize * slotIndex]) & + kPointerMask); + } + + void + setTag(int32_t slotIndex, uint8_t tag) { + reinterpret_cast(&tags_)[slotIndex] = tag; + } + + void + setPointer(int32_t slotIndex, void* pointer) { + auto* const slot = reinterpret_cast( + &pointers_[slotIndex * kPointerSize]); + *slot = + (*slot & ~kPointerMask) | reinterpret_cast(pointer); + } + + private: + static constexpr uint8_t kPointerSignificantBits = 48; + static constexpr uint64_t kPointerMask = + milvus::bits::lowMask(kPointerSignificantBits); + static constexpr int32_t kPointerSize = kPointerSignificantBits / 8; + + TagVector tags_; + char pointers_[sizeof(TagVector) * kPointerSize]; + char padding_[16]; + }; + static_assert(sizeof(Bucket) == 128); + static constexpr uint64_t kBucketSize = sizeof(Bucket); + + Bucket* + bucketAt(int64_t offset) const { + //AssertInfo(offset&(kBucketSize-1)==0, "Invalid offset:{} and kBucketSize:{}", offset, kBucketSize); + return reinterpret_cast(table_ + offset); + } + + int64_t + bucketOffset(uint64_t hash) const { + return hash & bucketOffsetMask_; + } + + int64_t + nextBucketOffset(int64_t bucketOffset) const { + //AssertInfo(bucketOffset&(kBucketSize - 1) == 0, "Invalid bucketOffset:{} for nextBucketOffset", bucketOffset); + AssertInfo(bucketOffset < sizeMask_, + "BucketOffset:{} must be less than sizeMask_:{} for " + "nextBucketOffset", + bucketOffset, + sizeMask_); + return sizeMask_ & (bucketOffset + kBucketSize); + } + + bool + compareKeys(const char* group, HashLookup& lookup, vector_size_t row); + + char* + row(int64_t bucketOffset, int32_t slotIndex) const { + return bucketAt(bucketOffset)->pointerAt(slotIndex); + } + + int64_t + numBuckets() const { + return numBuckets_; + } + + TagVector + loadTags(int64_t bucketOffset) const { + return BaseHashTable::loadTags(reinterpret_cast(table_), + bucketOffset); + } + + char* + insertEntry(HashLookup& lookup, uint64_t index, vector_size_t row); + + void + storeKeys(HashLookup& lookup, vector_size_t row); + + void + storeRowPointer(uint64_t index, uint64_t hash, char* row); + + // Allocates new tables for tags and payload pointers. The size must + // a power of 2. + void + allocateTables(uint64_t size); + + void + fullProbe(HashLookup& lookup, ProbeState& state); + + void + clear(bool freeTable = false) override; + + void + checkSizeAndAllocateTable(int32_t numNew); + + // Returns the number of entries after which the table gets rehashed. + static uint64_t + rehashSize(int64_t size) { + // This implements the F14 load factor: Resize if less than 1/8 unoccupied. + return size - (size / 8); + } + + uint64_t + rehashSize() const { + return rehashSize(capacity_); + } + + static uint64_t + newHashTableEntriesNumber(uint64_t numDistinct, uint64_t numNew) { + auto numNewEntries = + std::max((uint64_t)2048, + milvus::bits::nextPowerOfTwo(numNew * 2 + numDistinct)); + const auto newNumDistinct = numDistinct + numNew; + if (newNumDistinct > rehashSize(numNewEntries)) { + numNewEntries *= 2; + } + return numNewEntries; + } + + private: + HashMode hashMode_ = HashMode::kHash; + int64_t bucketOffsetMask_{0}; + int64_t numBuckets_{0}; + int64_t numDistinct_{0}; + + // Number of slots across all buckets. + int64_t capacity_{0}; + // Mask for extracting low bits of hash number for use as byte offsets into + // the table. This is set to 'capacity_ * sizeof(void*) - 1'. + int64_t sizeMask_{0}; + int8_t sizeBits_; + + int64_t numRehashes_{0}; + char* table_ = nullptr; + + HashMode + hashMode() const override { + return hashMode_; + } + friend class ProbeState; +}; + +} // namespace exec +} // namespace milvus \ No newline at end of file diff --git a/internal/core/src/exec/VectorHasher.cpp b/internal/core/src/exec/VectorHasher.cpp new file mode 100644 index 0000000000..54db20568e --- /dev/null +++ b/internal/core/src/exec/VectorHasher.cpp @@ -0,0 +1,77 @@ +// 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. + +#include "VectorHasher.h" +#include "common/float_util_c.h" +#include +#include "common/BitUtil.h" + +namespace milvus { +namespace exec { +std::vector> +createVectorHashers(const RowTypePtr& rowType, + const std::vector& exprs) { + std::vector> hashers; + hashers.reserve(exprs.size()); + for (const auto& expr : exprs) { + auto column_idx = rowType->GetChildIndex(expr->name()); + hashers.emplace_back(VectorHasher::create(expr->type(), column_idx)); + } + return hashers; +} + +template +void +VectorHasher::hashValues(const ColumnVectorPtr& column_data, + bool mix, + uint64_t* result) { + if constexpr (Type == DataType::ROW || Type == DataType::ARRAY || + Type == DataType::JSON) { + ThrowInfo(milvus::DataTypeInvalid, + "hash not supported for complex types ROW/ARRAY/JSON: {}", + Type); + } else { + using T = typename TypeTraits::NativeType; + for (size_t row_idx = 0; row_idx < column_data->size(); ++row_idx) { + if (!column_data->ValidAt(row_idx)) { + result[row_idx] = + mix ? milvus::bits::hashMix(result[row_idx], kNullHash) + : kNullHash; + } else { + T raw_value = column_data->ValueAt(row_idx); + uint64_t hash_value = kNullHash; + if constexpr (std::is_floating_point_v) { + hash_value = milvus::NaNAwareHash()(raw_value); + } else { + hash_value = folly::hasher()(raw_value); + } + result[row_idx] = + mix ? milvus::bits::hashMix(result[row_idx], hash_value) + : hash_value; + } + } + } +} + +void +VectorHasher::hash(bool mix, std::vector& result) { + auto element_data_type = ChannelDataType(); + MILVUS_DYNAMIC_TYPE_DISPATCH( + hashValues, element_data_type, columnData(), mix, result.data()); +} + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/VectorHasher.h b/internal/core/src/exec/VectorHasher.h new file mode 100644 index 0000000000..c360f0e5d7 --- /dev/null +++ b/internal/core/src/exec/VectorHasher.h @@ -0,0 +1,92 @@ +// 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. + +#pragma once + +#include "common/Vector.h" +#include "common/Types.h" +#include "expr/ITypeExpr.h" + +namespace milvus { +namespace exec { +class VectorHasher { + public: + VectorHasher(DataType data_type, column_index_t column_idx) + : channel_type_(data_type), channel_idx_(column_idx) { + } + + static std::unique_ptr + create(DataType data_type, column_index_t col_idx) { + return std::make_unique(data_type, col_idx); + } + + column_index_t + ChannelIndex() const { + return channel_idx_; + } + + DataType + ChannelDataType() const { + return channel_type_; + } + + void + hash(bool mix, std::vector& result); + + static constexpr uint64_t kNullHash = 1; + + static bool + typeSupportValueIds(DataType type) { + switch (type) { + case DataType::BOOL: + case DataType::INT8: + case DataType::INT16: + case DataType::INT32: + case DataType::INT64: + case DataType::VARCHAR: + case DataType::STRING: + return true; + default: + return false; + } + } + + template + void + hashValues(const ColumnVectorPtr& column_data, bool mix, uint64_t* result); + + void + setColumnData(const ColumnVectorPtr& column_data) { + column_data_ = column_data; + } + + const ColumnVectorPtr& + columnData() const { + return column_data_; + } + + private: + const column_index_t channel_idx_; + const DataType channel_type_; + ColumnVectorPtr column_data_; +}; + +std::vector> +createVectorHashers(const RowTypePtr& rowType, + const std::vector& exprs); + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/expression/Utils.h b/internal/core/src/exec/expression/Utils.h index c5be68c617..99e484f60c 100644 --- a/internal/core/src/exec/expression/Utils.h +++ b/internal/core/src/exec/expression/Utils.h @@ -16,6 +16,8 @@ #pragma once +#include + #include #include "common/EasyAssert.h" @@ -289,5 +291,25 @@ GetValueWithCastNumber(const milvus::proto::plan::GenericValue& value_proto) { } } +// Locale-independent ASCII lowercase conversion +// Converts only ASCII uppercase letters (A-Z) to lowercase (a-z) +// Non-ASCII characters and non-uppercase characters remain unchanged +inline unsigned char +asciiToLower(unsigned char c) { + if (c >= 'A' && c <= 'Z') { + return static_cast('a' + (c - 'A')); + } + return c; +} + +inline std::string +sanitizeName(const std::string& name) { + std::string sanitizedName; + sanitizedName.resize(name.size()); + std::transform( + name.begin(), name.end(), sanitizedName.begin(), asciiToLower); + return sanitizedName; +} + } // namespace exec } // namespace milvus \ No newline at end of file diff --git a/internal/core/src/exec/expression/function/FunctionFactory.cpp b/internal/core/src/exec/expression/function/FunctionFactory.cpp index 4f621f6506..fd38435825 100644 --- a/internal/core/src/exec/expression/function/FunctionFactory.cpp +++ b/internal/core/src/exec/expression/function/FunctionFactory.cpp @@ -15,9 +15,12 @@ // limitations under the License. #include "exec/expression/function/FunctionFactory.h" -#include #include "exec/expression/function/impl/StringFunctions.h" #include "log/Log.h" +#include "exec/operator/query-agg/CountAggregateBase.h" +#include "exec/operator/query-agg/MinAggregateBase.h" +#include "exec/operator/query-agg/MaxAggregateBase.h" +#include "exec/operator/query-agg/SumAggregateBase.h" namespace milvus { namespace exec { @@ -56,7 +59,8 @@ FunctionFactory::RegisterAllFunctions() { RegisterFilterFunction("starts_with", {DataType::VARCHAR, DataType::VARCHAR}, function::StartsWithVarchar); - LOG_INFO("{} functions registered", GetFilterFunctionNum()); + LOG_INFO("{} filter functions registered", GetFilterFunctionNum()); + RegisterAggregateFunction(); } void @@ -68,6 +72,14 @@ FunctionFactory::RegisterFilterFunction( func_name, func_param_type_list}] = func; } +void +FunctionFactory::RegisterAggregateFunction() { + milvus::exec::registerCountAggregate(); + milvus::exec::registerMinAggregate(); + milvus::exec::registerMaxAggregate(); + milvus::exec::registerSumAggregate(); +} + const FilterFunctionPtr FunctionFactory::GetFilterFunction( const FilterFunctionRegisterKey& func_sig) const { diff --git a/internal/core/src/exec/expression/function/FunctionFactory.h b/internal/core/src/exec/expression/function/FunctionFactory.h index 0563408c9d..e4056bcb39 100644 --- a/internal/core/src/exec/expression/function/FunctionFactory.h +++ b/internal/core/src/exec/expression/function/FunctionFactory.h @@ -78,6 +78,9 @@ class FunctionFactory { std::vector func_param_type_list, FilterFunctionPtr func); + void + RegisterAggregateFunction(); + const FilterFunctionPtr GetFilterFunction(const FilterFunctionRegisterKey& func_sig) const; diff --git a/internal/core/src/exec/operator/AggregationNode.cpp b/internal/core/src/exec/operator/AggregationNode.cpp new file mode 100644 index 0000000000..9dd5f4d1a5 --- /dev/null +++ b/internal/core/src/exec/operator/AggregationNode.cpp @@ -0,0 +1,71 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed 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 + +// +// Created by hanchun on 24-10-18. +// + +#include "AggregationNode.h" +#include "common/Utils.h" + +namespace milvus { +namespace exec { + +PhyAggregationNode::PhyAggregationNode( + int32_t operator_id, + milvus::exec::DriverContext* ctx, + const std::shared_ptr& node) + : Operator( + ctx, node->output_type(), operator_id, node->id(), "AggregationNode"), + aggregationNode_(node), + isGlobal_(node->GroupingKeys().empty()) { +} + +void +PhyAggregationNode::initialize() { + Operator::initialize(); + // aggregation operator will always have single one source + const auto& input_type = aggregationNode_->sources()[0]->output_type(); + auto hashers = + createVectorHashers(input_type, aggregationNode_->GroupingKeys()); + auto numHashers = hashers.size(); + std::vector aggregateInfos = + toAggregateInfo(*aggregationNode_, *operator_context_, numHashers); + grouping_set_ = std::make_unique( + input_type, std::move(hashers), std::move(aggregateInfos)); + aggregationNode_.reset(); +} + +void +PhyAggregationNode::AddInput(RowVectorPtr& input) { + grouping_set_->addInput(input); + numInputRows_ += input->size(); +} + +RowVectorPtr +PhyAggregationNode::GetOutput() { + if (finished_ || !no_more_input_) { + input_ = nullptr; + return nullptr; + } + DeferLambda([&]() { finished_ = true; }); + const auto outputRowCount = isGlobal_ ? 1 : grouping_set_->outputRowCount(); + output_ = std::make_shared(output_type_, outputRowCount); + const bool hasData = grouping_set_->getOutput(output_); + if (!hasData) { + return nullptr; + } + numOutputRows_ += output_->size(); + return output_; +} + +}; // namespace exec +}; // namespace milvus diff --git a/internal/core/src/exec/operator/AggregationNode.h b/internal/core/src/exec/operator/AggregationNode.h new file mode 100644 index 0000000000..68e26739af --- /dev/null +++ b/internal/core/src/exec/operator/AggregationNode.h @@ -0,0 +1,86 @@ +// 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. + +#pragma once +#include "exec/operator/Operator.h" +#include "exec/operator/query-agg/GroupingSet.h" +#include "common/Types.h" + +namespace milvus { +namespace exec { +class PhyAggregationNode : public Operator { + public: + PhyAggregationNode( + int32_t operator_id, + DriverContext* ctx, + const std::shared_ptr& node); + + bool + NeedInput() const override { + return true; + } + + void + AddInput(RowVectorPtr& input) override; + + RowVectorPtr + GetOutput() override; + + bool + IsFinished() override { + return finished_; + } + + bool + IsFilter() const override { + return false; + } + + BlockingReason + IsBlocked(ContinueFuture* future) override { + return BlockingReason::kNotBlocked; + } + + void + Close() override { + input_ = nullptr; + results_.clear(); + } + + void + initialize() override; + + std::string + ToString() const override { + return "PhyAggregationNode"; + } + + private: + RowVectorPtr output_; + std::unique_ptr grouping_set_; + std::shared_ptr aggregationNode_; + const bool isGlobal_; + + // Count the number of input rows. It is reset on partial aggregation output + // flush. + int64_t numInputRows_ = 0; + // Count the number of output rows. It is reset on partial aggregation output + // flush. + int64_t numOutputRows_ = 0; + bool finished_ = false; +}; +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/operator/CallbackSink.h b/internal/core/src/exec/operator/CallbackSink.h index d0f5e2d37a..a069b68a1f 100644 --- a/internal/core/src/exec/operator/CallbackSink.h +++ b/internal/core/src/exec/operator/CallbackSink.h @@ -26,7 +26,7 @@ class CallbackSink : public Operator { int32_t operator_id, DriverContext* ctx, std::function callback) - : Operator(ctx, DataType::NONE, operator_id, "N/A", "CallbackSink"), + : Operator(ctx, RowType::None, operator_id, "N/A", "CallbackSink"), callback_(callback) { } @@ -52,7 +52,7 @@ class CallbackSink : public Operator { } bool - IsFilter() override { + IsFilter() const override { return false; } diff --git a/internal/core/src/exec/operator/CountNode.cpp b/internal/core/src/exec/operator/CountNode.cpp deleted file mode 100644 index 00d7aa0739..0000000000 --- a/internal/core/src/exec/operator/CountNode.cpp +++ /dev/null @@ -1,80 +0,0 @@ -// 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. - -#include "CountNode.h" -#include "common/Tracer.h" -#include "fmt/format.h" -namespace milvus { -namespace exec { - -static std::unique_ptr -wrap_num_entities(int64_t cnt, int64_t size) { - auto retrieve_result = std::make_unique(); - DataArray arr; - arr.set_type(milvus::proto::schema::Int64); - auto scalar = arr.mutable_scalars(); - scalar->mutable_long_data()->mutable_data()->Add(cnt); - retrieve_result->field_data_ = {arr}; - retrieve_result->total_data_cnt_ = size; - return retrieve_result; -} - -PhyCountNode::PhyCountNode(int32_t operator_id, - DriverContext* driverctx, - const std::shared_ptr& node) - : Operator(driverctx, - node->output_type(), - operator_id, - node->id(), - "PhyCountNode") { - ExecContext* exec_context = operator_context_->get_exec_context(); - query_context_ = exec_context->get_query_context(); - segment_ = query_context_->get_segment(); - query_timestamp_ = query_context_->get_query_timestamp(); - active_count_ = query_context_->get_active_count(); -} - -void -PhyCountNode::AddInput(RowVectorPtr& input) { - input_ = std::move(input); -} - -RowVectorPtr -PhyCountNode::GetOutput() { - milvus::exec::checkCancellation(query_context_); - - if (is_finished_ || !no_more_input_) { - return nullptr; - } - tracer::AutoSpan span("PhyCountNode::Execute", tracer::GetRootSpan(), true); - auto col_input = GetColumnVector(input_); - TargetBitmapView view(col_input->GetRawData(), col_input->size()); - auto cnt = view.size() - view.count(); - query_context_->set_retrieve_result( - std::move(*(wrap_num_entities(cnt, view.size())))); - is_finished_ = true; - - tracer::AddEvent(fmt::format("count_result: {}", cnt)); - return input_; -} - -bool -PhyCountNode::IsFinished() { - return is_finished_; -} - -} // namespace exec -} // namespace milvus \ No newline at end of file diff --git a/internal/core/src/exec/operator/ElementFilterBitsNode.cpp b/internal/core/src/exec/operator/ElementFilterBitsNode.cpp index d31b296c56..c1034c4203 100644 --- a/internal/core/src/exec/operator/ElementFilterBitsNode.cpp +++ b/internal/core/src/exec/operator/ElementFilterBitsNode.cpp @@ -29,9 +29,9 @@ PhyElementFilterBitsNode::PhyElementFilterBitsNode( const std::shared_ptr& element_filter_bits_node) : Operator(driverctx, - DataType::NONE, + element_filter_bits_node->output_type(), operator_id, - "element_filter_bits_plan_node", + element_filter_bits_node->id(), "PhyElementFilterBitsNode"), struct_name_(element_filter_bits_node->struct_name()) { ExecContext* exec_context = operator_context_->get_exec_context(); diff --git a/internal/core/src/exec/operator/ElementFilterBitsNode.h b/internal/core/src/exec/operator/ElementFilterBitsNode.h index f68b4a3967..3b1fcfb55e 100644 --- a/internal/core/src/exec/operator/ElementFilterBitsNode.h +++ b/internal/core/src/exec/operator/ElementFilterBitsNode.h @@ -38,7 +38,7 @@ class PhyElementFilterBitsNode : public Operator { element_filter_bits_node); bool - IsFilter() override { + IsFilter() const override { return true; } diff --git a/internal/core/src/exec/operator/ElementFilterNode.h b/internal/core/src/exec/operator/ElementFilterNode.h index 76853a69a1..5f56531bb1 100644 --- a/internal/core/src/exec/operator/ElementFilterNode.h +++ b/internal/core/src/exec/operator/ElementFilterNode.h @@ -38,7 +38,7 @@ class PhyElementFilterNode : public Operator { element_filter_node); bool - IsFilter() override { + IsFilter() const override { return true; } diff --git a/internal/core/src/exec/operator/FilterBitsNode.cpp b/internal/core/src/exec/operator/FilterBitsNode.cpp index ca189a3b13..ce29f1932b 100644 --- a/internal/core/src/exec/operator/FilterBitsNode.cpp +++ b/internal/core/src/exec/operator/FilterBitsNode.cpp @@ -17,7 +17,6 @@ #include "FilterBitsNode.h" #include "common/Tracer.h" #include "fmt/format.h" - #include "monitor/Monitor.h" namespace milvus { diff --git a/internal/core/src/exec/operator/FilterBitsNode.h b/internal/core/src/exec/operator/FilterBitsNode.h index de6d472a50..9eea755287 100644 --- a/internal/core/src/exec/operator/FilterBitsNode.h +++ b/internal/core/src/exec/operator/FilterBitsNode.h @@ -34,7 +34,7 @@ class PhyFilterBitsNode : public Operator { const std::shared_ptr& filter); bool - IsFilter() override { + IsFilter() const override { return true; } diff --git a/internal/core/src/exec/operator/IterativeFilterNode.h b/internal/core/src/exec/operator/IterativeFilterNode.h index 07404d974b..23eb6052ae 100644 --- a/internal/core/src/exec/operator/IterativeFilterNode.h +++ b/internal/core/src/exec/operator/IterativeFilterNode.h @@ -37,7 +37,7 @@ class PhyIterativeFilterNode : public Operator { const std::shared_ptr& filter); bool - IsFilter() override { + IsFilter() const override { return true; } diff --git a/internal/core/src/exec/operator/MvccNode.cpp b/internal/core/src/exec/operator/MvccNode.cpp index d65df9be0d..b1c3a492e0 100644 --- a/internal/core/src/exec/operator/MvccNode.cpp +++ b/internal/core/src/exec/operator/MvccNode.cpp @@ -17,6 +17,7 @@ #include "MvccNode.h" #include "common/Tracer.h" #include "fmt/format.h" +#include namespace milvus { namespace exec { @@ -54,12 +55,12 @@ PhyMvccNode::GetOutput() { tracer::AutoSpan span("PhyMvccNode::Execute", tracer::GetRootSpan(), true); - if (!is_source_node_ && input_ == nullptr) { + if (active_count_ == 0) { + is_finished_ = true; return nullptr; } - if (active_count_ == 0) { - is_finished_ = true; + if (!is_source_node_ && input_ == nullptr) { return nullptr; } diff --git a/internal/core/src/exec/operator/MvccNode.h b/internal/core/src/exec/operator/MvccNode.h index 6890455786..b286e00bcd 100644 --- a/internal/core/src/exec/operator/MvccNode.h +++ b/internal/core/src/exec/operator/MvccNode.h @@ -34,7 +34,7 @@ class PhyMvccNode : public Operator { const std::shared_ptr& mvcc_node); bool - IsFilter() override { + IsFilter() const override { return false; } diff --git a/internal/core/src/exec/operator/Operator.cpp b/internal/core/src/exec/operator/Operator.cpp index 972482c797..6af1b46721 100644 --- a/internal/core/src/exec/operator/Operator.cpp +++ b/internal/core/src/exec/operator/Operator.cpp @@ -17,5 +17,10 @@ #include "Operator.h" namespace milvus { -namespace exec {} +namespace exec { +void +Operator::initialize() { + // TODO check memory and set up memory pool in the future +} +} // namespace exec } // namespace milvus diff --git a/internal/core/src/exec/operator/Operator.h b/internal/core/src/exec/operator/Operator.h index 0e3e1503be..5b813619f3 100644 --- a/internal/core/src/exec/operator/Operator.h +++ b/internal/core/src/exec/operator/Operator.h @@ -94,16 +94,26 @@ class OperatorContext { class Operator { public: Operator(DriverContext* ctx, - DataType output_type, + RowTypePtr output_type, int32_t operator_id, const std::string& plannode_id, const std::string& operator_type) : operator_context_(std::make_unique( - ctx, plannode_id, operator_id, operator_type)) { + ctx, plannode_id, operator_id, operator_type)), + output_type_(output_type) { } virtual ~Operator() = default; + /// Does initialization work for this operator which requires memory + /// allocation from memory pool that can't be done under operator constructor. + /// + /// NOTE: the default implementation set 'initialized_' to true to ensure we + /// never call this more than once. The overload initialize() implementation + /// must call this base implementation first. + virtual void + initialize(); + virtual bool NeedInput() const = 0; @@ -122,7 +132,7 @@ class Operator { IsFinished() = 0; virtual bool - IsFilter() = 0; + IsFilter() const = 0; virtual BlockingReason IsBlocked(ContinueFuture* future) = 0; @@ -158,10 +168,15 @@ class Operator { return "Base Operator"; } + virtual const RowTypePtr& + OutputType() const { + return output_type_; + } + protected: std::unique_ptr operator_context_; - DataType output_type_; + RowTypePtr output_type_; RowVectorPtr input_; @@ -173,7 +188,7 @@ class Operator { class SourceOperator : public Operator { public: SourceOperator(DriverContext* driver_ctx, - DataType out_type, + RowTypePtr out_type, int32_t operator_id, const std::string& plannode_id, const std::string& operator_type) diff --git a/internal/core/src/exec/operator/ProjectNode.cpp b/internal/core/src/exec/operator/ProjectNode.cpp new file mode 100644 index 0000000000..533391ba6b --- /dev/null +++ b/internal/core/src/exec/operator/ProjectNode.cpp @@ -0,0 +1,82 @@ +// 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. + +#include "ProjectNode.h" +#include "exec/expression/Utils.h" +#include "segcore/Utils.h" + +namespace milvus { +namespace exec { +PhyProjectNode::PhyProjectNode( + int32_t operator_id, + milvus::exec::DriverContext* ctx, + const std::shared_ptr& projectNode) + : Operator(ctx, + projectNode->output_type(), + operator_id, + projectNode->id(), + "Project"), + fields_to_project_(projectNode->FieldsToProject()) { + auto exec_context = operator_context_->get_exec_context(); + segment_ = exec_context->get_query_context()->get_segment(); + op_context_ = exec_context->get_query_context()->get_op_context(); + AssertInfo(op_context_, "op_context_ cannot be nullptr for ProjectNode"); + AssertInfo(segment_, "segment_ cannot be nullptr for ProjectNode"); +} + +void +PhyProjectNode::AddInput(milvus::RowVectorPtr& input) { + input_ = std::move(input); +} + +RowVectorPtr +PhyProjectNode::GetOutput() { + if (is_finished_ || input_ == nullptr) { + return nullptr; + } + auto col_input = GetColumnVector(input_); + // raw data view + TargetBitmapView raw_data_view(col_input->GetRawData(), col_input->size()); + auto result_pair = segment_->find_first(-1, raw_data_view); + auto& selected_offsets = result_pair.first; + auto selected_count = selected_offsets.size(); + auto row_type = OutputType(); + std::vector column_vectors; + column_vectors.reserve(fields_to_project_.size()); + for (int i = 0; i < fields_to_project_.size(); i++) { + auto column_type = row_type->column_type(i); + auto field_id = fields_to_project_.at(i); + + TargetBitmap valid_map(selected_count); + auto field_data = bulk_script_field_data(op_context_, + field_id, + column_type, + selected_offsets.data(), + selected_count, + segment_, + valid_map, + true); + auto column_vector = std::make_shared( + std::move(field_data), std::move(valid_map)); + column_vectors.emplace_back(column_vector); + } + is_finished_ = true; + auto row_vector = std::make_shared(std::move(column_vectors)); + return row_vector; +} + +}; // namespace exec +}; // namespace milvus \ No newline at end of file diff --git a/internal/core/src/exec/operator/CountNode.h b/internal/core/src/exec/operator/ProjectNode.h similarity index 65% rename from internal/core/src/exec/operator/CountNode.h rename to internal/core/src/exec/operator/ProjectNode.h index cfb9512a55..517748105b 100644 --- a/internal/core/src/exec/operator/CountNode.h +++ b/internal/core/src/exec/operator/ProjectNode.h @@ -15,45 +15,36 @@ // limitations under the License. #pragma once - -#include -#include - -#include "exec/Driver.h" -#include "exec/expression/Expr.h" -#include "exec/operator/Operator.h" -#include "exec/QueryContext.h" +#include "Operator.h" +#include "plan/PlanNode.h" namespace milvus { namespace exec { - -class PhyCountNode : public Operator { +class PhyProjectNode : public Operator { public: - PhyCountNode(int32_t operator_id, - DriverContext* ctx, - const std::shared_ptr& node); + PhyProjectNode(int32_t operator_id, + DriverContext* ctx, + const std::shared_ptr& projectNode); bool - IsFilter() override { + IsFilter() const override { return false; } bool NeedInput() const override { - return !is_finished_; + return true; } void - AddInput(RowVectorPtr& input); + AddInput(RowVectorPtr& input) override; RowVectorPtr GetOutput() override; bool - IsFinished() override; - - void - Close() override { + IsFinished() override { + return is_finished_; } BlockingReason @@ -61,18 +52,16 @@ class PhyCountNode : public Operator { return BlockingReason::kNotBlocked; } - virtual std::string + std::string ToString() const override { - return "PhyCountNode"; + return "Project Operator"; } private: const segcore::SegmentInternalInterface* segment_; - milvus::Timestamp query_timestamp_; - int64_t active_count_; - QueryContext* query_context_; bool is_finished_{false}; + const std::vector fields_to_project_; + OpContext* op_context_; }; - } // namespace exec -} // namespace milvus \ No newline at end of file +} // namespace milvus diff --git a/internal/core/src/exec/operator/RandomSampleNode.h b/internal/core/src/exec/operator/RandomSampleNode.h index 2c3dc243c7..398a7e0c8c 100644 --- a/internal/core/src/exec/operator/RandomSampleNode.h +++ b/internal/core/src/exec/operator/RandomSampleNode.h @@ -31,7 +31,7 @@ class PhyRandomSampleNode : public Operator { random_sample_node); bool - IsFilter() override { + IsFilter() const override { return false; } diff --git a/internal/core/src/exec/operator/RescoresNode.h b/internal/core/src/exec/operator/RescoresNode.h index 27cb576453..3e1fd72cea 100644 --- a/internal/core/src/exec/operator/RescoresNode.h +++ b/internal/core/src/exec/operator/RescoresNode.h @@ -37,7 +37,7 @@ class PhyRescoresNode : public Operator { const std::shared_ptr& scorer); bool - IsFilter() override { + IsFilter() const override { return true; } diff --git a/internal/core/src/exec/operator/GroupByNode.cpp b/internal/core/src/exec/operator/SearchGroupByNode.cpp similarity index 91% rename from internal/core/src/exec/operator/GroupByNode.cpp rename to internal/core/src/exec/operator/SearchGroupByNode.cpp index bf296c1790..3f962b2ecf 100644 --- a/internal/core/src/exec/operator/GroupByNode.cpp +++ b/internal/core/src/exec/operator/SearchGroupByNode.cpp @@ -14,24 +14,23 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "GroupByNode.h" #include "common/Tracer.h" #include "fmt/format.h" - -#include "exec/operator/groupby/SearchGroupByOperator.h" +#include "SearchGroupByNode.h" +#include "exec/operator/search-groupby/SearchGroupByOperator.h" #include "monitor/Monitor.h" namespace milvus { namespace exec { -PhyGroupByNode::PhyGroupByNode( +PhySearchGroupByNode::PhySearchGroupByNode( int32_t operator_id, DriverContext* driverctx, - const std::shared_ptr& node) + const std::shared_ptr& node) : Operator(driverctx, node->output_type(), operator_id, node->id(), - "PhyGroupByNode") { + "PhySearchGroupByNode") { ExecContext* exec_context = operator_context_->get_exec_context(); query_context_ = exec_context->get_query_context(); segment_ = query_context_->get_segment(); @@ -39,14 +38,13 @@ PhyGroupByNode::PhyGroupByNode( } void -PhyGroupByNode::AddInput(RowVectorPtr& input) { +PhySearchGroupByNode::AddInput(RowVectorPtr& input) { input_ = std::move(input); } RowVectorPtr -PhyGroupByNode::GetOutput() { +PhySearchGroupByNode::GetOutput() { milvus::exec::checkCancellation(query_context_); - if (is_finished_ || !no_more_input_) { return nullptr; } @@ -104,7 +102,7 @@ PhyGroupByNode::GetOutput() { } bool -PhyGroupByNode::IsFinished() { +PhySearchGroupByNode::IsFinished() { return is_finished_; } diff --git a/internal/core/src/exec/operator/GroupByNode.h b/internal/core/src/exec/operator/SearchGroupByNode.h similarity index 86% rename from internal/core/src/exec/operator/GroupByNode.h rename to internal/core/src/exec/operator/SearchGroupByNode.h index a5d05898f9..ee0b64c122 100644 --- a/internal/core/src/exec/operator/GroupByNode.h +++ b/internal/core/src/exec/operator/SearchGroupByNode.h @@ -27,14 +27,15 @@ namespace milvus { namespace exec { -class PhyGroupByNode : public Operator { +class PhySearchGroupByNode : public Operator { public: - PhyGroupByNode(int32_t operator_id, - DriverContext* ctx, - const std::shared_ptr& node); + PhySearchGroupByNode( + int32_t operator_id, + DriverContext* ctx, + const std::shared_ptr& node); bool - IsFilter() override { + IsFilter() const override { return false; } @@ -63,7 +64,7 @@ class PhyGroupByNode : public Operator { virtual std::string ToString() const override { - return "PhyGroupByNode"; + return "PhySearchGroupByNode"; } private: diff --git a/internal/core/src/exec/operator/VectorSearchNode.h b/internal/core/src/exec/operator/VectorSearchNode.h index e6ec630eed..7dd7b32dfb 100644 --- a/internal/core/src/exec/operator/VectorSearchNode.h +++ b/internal/core/src/exec/operator/VectorSearchNode.h @@ -35,7 +35,7 @@ class PhyVectorSearchNode : public Operator { const std::shared_ptr& search_node); bool - IsFilter() override { + IsFilter() const override { return false; } diff --git a/internal/core/src/exec/operator/query-agg/Aggregate.cpp b/internal/core/src/exec/operator/query-agg/Aggregate.cpp new file mode 100644 index 0000000000..c79d7c451c --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/Aggregate.cpp @@ -0,0 +1,70 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed 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 +#include "Aggregate.h" +#include "common/Utils.h" +#include "exec/expression/Utils.h" + +namespace milvus { +namespace exec { + +void +Aggregate::setOffsetsInternal(int32_t offset, + int32_t nullByte, + uint8_t nullMask, + int32_t rowSizeOffset) { + offset_ = offset; + nullByte_ = nullByte; + nullMask_ = nullMask; + rowSizeOffset_ = rowSizeOffset; +} + +const AggregateFunctionFactory* +getAggregateFunctionEntry(const std::string& name) { + // Use the same normalization as registration (lowerString) to ensure + // consistent lookup for non-ASCII and locale-sensitive characters. + auto normalizedName = lowerString(name); + + return aggregateFunctions().withRLock( + [&](const auto& functionsMap) -> const AggregateFunctionFactory* { + auto it = functionsMap.find(normalizedName); + if (it != functionsMap.end()) { + return &it->second; + } + return nullptr; + }); +} + +std::unique_ptr +Aggregate::create(const std::string& name, + const std::vector& argTypes, + const QueryConfig& query_config) { + if (auto func = getAggregateFunctionEntry(name)) { + return (*func)(argTypes, query_config); + } + ThrowInfo(UnexpectedError, "Aggregate function not registered: {}", name); +} + +void +registerAggregateFunction(const std::string& name, + const AggregateFunctionFactory& factory) { + auto realName = lowerString(name); + aggregateFunctions().withWLock( + [&](auto& aggFunctionMap) { aggFunctionMap[realName] = factory; }); +} + +AggregateFunctionMap& +aggregateFunctions() { + static AggregateFunctionMap aggFunctionMap; + return aggFunctionMap; +} + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/operator/query-agg/Aggregate.h b/internal/core/src/exec/operator/query-agg/Aggregate.h new file mode 100644 index 0000000000..b303d7d055 --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/Aggregate.h @@ -0,0 +1,180 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed 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 +#pragma once + +#include "common/Types.h" +#include "plan/PlanNode.h" +#include "exec/QueryContext.h" +#include + +namespace milvus { +namespace exec { +class Aggregate { + protected: + explicit Aggregate(DataType result_type) : result_type_(result_type) { + } + + private: + const DataType result_type_; + + // Byte position of null flag in group row. + int32_t nullByte_; + uint8_t nullMask_; + // Offset of fixed length accumulator state in group row. + int32_t offset_; + // Offset of uint32_t row byte size of row. 0 if there are no + // variable width fields or accumulators on the row. The size is + // capped at 4G and will stay at 4G and not wrap around if growing + // past this. This serves to track the batch size when extracting + // rows. A size in excess of 4G would finish the batch in any case, + // so larger values need not be represented. + int32_t rowSizeOffset_ = 0; + + public: + virtual ~Aggregate() = default; + + DataType + resultType() const { + return result_type_; + } + + static std::unique_ptr + create(const std::string& name, + const std::vector& argTypes, + const QueryConfig& query_config); + + void + setOffsets(int32_t offset, + int32_t nullByte, + uint8_t nullMask, + int32_t rowSizeOffset) { + setOffsetsInternal(offset, nullByte, nullMask, rowSizeOffset); + } + + virtual void + initializeNewGroups(char** groups, + folly::Range indices) { + initializeNewGroupsInternal(groups, indices); + } + + virtual void + addSingleGroupRawInput(char* group, + const std::vector& input) = 0; + + virtual void + addRawInput(char** groups, + int numGroups, + const std::vector& input) = 0; + + virtual void + extractValues(char** groups, int32_t numGroups, VectorPtr* result) = 0; + + template + T* + value(char* group) const { + AssertInfo(reinterpret_cast(group + offset_) % + accumulatorAlignmentSize() == + 0, + "aggregation value in the groups is not aligned"); + return reinterpret_cast(group + offset_); + } + + bool + isNull(char* group) const { + return numNulls_ && (group[nullByte_] & nullMask_); + } + + // Returns true if the accumulator never takes more than + // accumulatorFixedWidthSize() bytes. If this is false, the + // accumulator needs to track its changing variable length footprint + // using RowSizeTracker (Aggregate::trackRowSize), see ArrayAggAggregate for + // sample usage. A group row with at least one variable length key or + // aggregate will have a 32-bit slot at offset RowContainer::rowSize_ for + // keeping track of per-row size. The size is relevant for keeping caps on + // result set and spilling batch sizes with skewed data. + virtual bool + isFixedSize() const { + return true; + } + + // Returns the fixed number of bytes the accumulator takes on a group + // row. Variable width accumulators will reference the variable + // width part of the state from the fixed part. + virtual int32_t + accumulatorFixedWidthSize() const = 0; + + /// Returns the alignment size of the accumulator. Some types such as + /// int128_t require aligned access. This value must be a power of 2. + virtual int32_t + accumulatorAlignmentSize() const { + return 1; + } + + protected: + virtual void + setOffsetsInternal(int32_t offset, + int32_t nullByte, + uint8_t nullMask, + int32_t rowSizeOffset); + + virtual void + initializeNewGroupsInternal(char** groups, + folly::Range indices) = 0; + // Number of null accumulators in the current state of the aggregation + // operator for this aggregate. If 0, clearing the null as part of update + // is not needed. + uint64_t numNulls_ = 0; + + inline bool + clearNull(char* group) { + if (numNulls_) { + uint8_t mask = group[nullByte_]; + if (mask & nullMask_) { + group[nullByte_] = mask & ~nullMask_; + numNulls_--; + return true; + } + } + return false; + } + + void + setAllNulls(char** groups, folly::Range indices) { + for (auto i : indices) { + groups[i][nullByte_] |= nullMask_; + } + numNulls_ += indices.size(); + } +}; + +using AggregateFunctionFactory = std::function( + const std::vector& argTypes, const QueryConfig& config)>; + +const AggregateFunctionFactory* +getAggregateFunctionEntry(const std::string& name); + +using AggregateFunctionMap = folly::Synchronized< + std::unordered_map>; + +AggregateFunctionMap& +aggregateFunctions(); + +/// Register an aggregate function with the specified name and signatures. If +/// registerCompanionFunctions is true, also register companion aggregate and +/// scalar functions with it. When functions with `name` already exist, if +/// overwrite is true, existing registration will be replaced. Otherwise, return +/// false without overwriting the registry. +void +registerAggregateFunction(const std::string& name, + const AggregateFunctionFactory& factory); + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/operator/query-agg/AggregateInfo.cpp b/internal/core/src/exec/operator/query-agg/AggregateInfo.cpp new file mode 100644 index 0000000000..483c770a47 --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/AggregateInfo.cpp @@ -0,0 +1,61 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed 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 + +// +// Created by hanchun on 24-10-22. +// +#include "AggregateInfo.h" +#include "common/Types.h" + +namespace milvus { +namespace exec { + +std::vector +toAggregateInfo(const plan::AggregationNode& aggregationNode, + const milvus::exec::OperatorContext& operatorCtx, + uint32_t numKeys) { + const auto numAggregates = aggregationNode.aggregates().size(); + std::vector aggregates; + aggregates.reserve(numAggregates); + const auto& inputType = aggregationNode.sources()[0]->output_type(); + + for (auto i = 0; i < numAggregates; i++) { + const auto& aggregate = aggregationNode.aggregates()[i]; + AggregateInfo info; + auto& inputColumnIdxes = info.input_column_idxes_; + for (const auto& inputExpr : aggregate.call_->inputs()) { + if (auto fieldExpr = dynamic_cast( + inputExpr.get())) { + const auto& fieldName = fieldExpr->name(); + if (fieldName.empty()) { + ThrowInfo(ExprInvalid, + "Field name cannot be empty in aggregate input"); + } + inputColumnIdxes.emplace_back( + inputType->GetChildIndex(fieldName)); + } else if (inputExpr != nullptr) { + ThrowInfo(ExprInvalid, + "Only support aggregation towards column for now"); + } + } + auto index = numKeys + i; + info.function_ = Aggregate::create( + aggregate.call_->fun_name(), + aggregate.rawInputTypes_, + *(operatorCtx.get_exec_context()->get_query_config())); + info.output_ = index; + aggregates.emplace_back(std::move(info)); + } + return aggregates; +} + +} // namespace exec +} // namespace milvus \ No newline at end of file diff --git a/internal/core/src/exec/operator/query-agg/AggregateInfo.h b/internal/core/src/exec/operator/query-agg/AggregateInfo.h new file mode 100644 index 0000000000..22b21c75d3 --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/AggregateInfo.h @@ -0,0 +1,42 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed 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 +#pragma once + +#include +#include + +#include "common/Types.h" +#include "Aggregate.h" +#include "plan/PlanNode.h" +#include "exec/operator/Operator.h" + +namespace milvus { +namespace exec { + +/// Information needed to evaluate an aggregate function. +struct AggregateInfo { + /// Instance of the Aggregate class. + std::unique_ptr function_; + + /// Indices of the input columns in the input RowVector. + std::vector input_column_idxes_; + + /// Index of the result column in the output RowVector. + column_index_t output_; +}; + +std::vector +toAggregateInfo(const plan::AggregationNode& aggregationNode, + const milvus::exec::OperatorContext& operatorCtx, + uint32_t numKeys); + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/operator/query-agg/AggregateUtil.h b/internal/core/src/exec/operator/query-agg/AggregateUtil.h new file mode 100644 index 0000000000..a70d38e29e --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/AggregateUtil.h @@ -0,0 +1,41 @@ +// 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. +#pragma once + +#include "common/Types.h" + +namespace milvus { +namespace exec { +// The result of aggregation function registration. +struct AggregateRegistrationResult { + bool mainFunction{false}; + bool partialFunction{false}; + bool mergeFunction{false}; + bool extractFunction{false}; + bool mergeExtractFunction{false}; + + bool + operator==(const AggregateRegistrationResult& other) const { + return mainFunction == other.mainFunction && + partialFunction == other.partialFunction && + mergeFunction == other.mergeFunction && + extractFunction == other.extractFunction && + mergeExtractFunction == other.mergeExtractFunction; + } +}; + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/operator/query-agg/CountAggregateBase.cpp b/internal/core/src/exec/operator/query-agg/CountAggregateBase.cpp new file mode 100644 index 0000000000..da404d7152 --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/CountAggregateBase.cpp @@ -0,0 +1,34 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed 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 +#include "CountAggregateBase.h" +#include "common/Utils.h" +#include "log/Log.h" + +namespace milvus { +namespace exec { + +void +registerCount(const std::string& name) { + exec::registerAggregateFunction( + name, + [name](const std::vector& /*argumentTypes*/, + const QueryConfig& /*config*/) -> std::unique_ptr { + return std::make_unique(); + }); +} + +void +registerCountAggregate() { + registerCount(milvus::KCount); + LOG_INFO("Registered Count Aggregate Function"); +} +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/operator/query-agg/CountAggregateBase.h b/internal/core/src/exec/operator/query-agg/CountAggregateBase.h new file mode 100644 index 0000000000..2b234046e7 --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/CountAggregateBase.h @@ -0,0 +1,111 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed 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 +#pragma once + +#include "SimpleNumericAggregate.h" + +namespace milvus { +namespace exec { +class CountAggregate : public SimpleNumericAggregate { + using BaseAggregate = SimpleNumericAggregate; + + public: + explicit CountAggregate() : BaseAggregate(DataType::INT64) { + } + + int32_t + accumulatorFixedWidthSize() const override { + return sizeof(int64_t); + } + + void + extractValues(char** groups, + int32_t numGroups, + VectorPtr* result) override { + BaseAggregate::doExtractValues( + groups, numGroups, result, [&](char* group) { + return *value(group); + }); + } + + void + addRawInput(char** groups, + int numGroups, + const std::vector& input) override { + if (!input.empty()) { + ColumnVectorPtr input_column = nullptr; + AssertInfo(input.size() == 1, + fmt::format("input column count for count aggregation " + "must be one , but got:{}", + input.size())); + input_column = std::dynamic_pointer_cast(input[0]); + AssertInfo(input_column != nullptr, + "input[0] must be ColumnVector for count aggregation"); + for (auto i = 0; i < input_column->size(); i++) { + if (input_column->ValidAt(i)) { + addToGroup(groups[i], 1); + } + } + return; + } + for (auto i = 0; i < numGroups; i++) { + addToGroup(groups[i], 1); + } + } + + void + addSingleGroupRawInput(char* group, + const std::vector& input) override { + AssertInfo(input.size() == 1, + fmt::format("input column count for count aggregation " + "must be exactly one for now, but got:{}", + input.size())); + const auto& column = std::dynamic_pointer_cast(input[0]); + AssertInfo(column != nullptr, + "input[0] must be ColumnVector for count aggregation"); + if (column->IsBitmap()) { + // Use validity bitmap to count non-null values + BitsetTypeView view(column->GetRawData(), column->size()); + addToGroup(group, view.size() - view.count()); + } else { + for (auto i = 0; i < column->size(); i++) { + if (column->ValidAt(i)) { + addToGroup(group, 1); + } + } + } + } + + void + initializeNewGroupsInternal( + char** groups, folly::Range indices) override { + // no need to set nulls for count aggregate as count result is always a number + for (auto i : indices) { + // initialized result of count is always zero + *value(groups[i]) = static_cast(0); + } + } + + private: + inline void + addToGroup(char* group, int64_t count) { + *value(group) += count; + } +}; + +void +registerCount(const std::string& name); + +void +registerCountAggregate(); + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/operator/query-agg/GroupingSet.cpp b/internal/core/src/exec/operator/query-agg/GroupingSet.cpp new file mode 100644 index 0000000000..a4e1ab664a --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/GroupingSet.cpp @@ -0,0 +1,247 @@ +// 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. + +#include "GroupingSet.h" +#include "common/Utils.h" +#include "SumAggregateBase.h" + +namespace milvus { +namespace exec { +GroupingSet::~GroupingSet() { + if (isGlobal_ && lookup_) { + AssertInfo(lookup_->hits_.size() == 1, + "GlobalAggregation should have exactly one output line"); + // globalAggregationBuffer_ is automatically cleaned up by unique_ptr + // No need to manually delete[] or set to nullptr + lookup_->hits_[0] = nullptr; + } +} + +void +GroupingSet::addInput(const RowVectorPtr& input) { + auto numRows = input->size(); + numInputRows_ += numRows; + if (isGlobal_) { + addGlobalAggregationInput(input); + return; + } + addInputForActiveRows(input); +} + +void +GroupingSet::initializeGlobalAggregation() { + if (globalAggregationInitialized_) { + return; + } + lookup_ = std::make_unique(hashers_); + lookup_->reset(1); + + // Row layout is: + // - alternating null flag - one bit per flag, one pair per + // aggregation, + // - uint32_t row size, + // - fixed-width accumulators - one per aggregate + // + // Here we always make space for a row size since we only have one row and no + // RowContainer. The whole row is allocated to guarantee that alignment + // requirements of all aggregate functions are satisfied. + + // Allocate space for the null and initialized flags. + size_t numAggregates = aggregates_.size(); + int32_t rowSizeOffset = milvus::bits::nBytes(numAggregates); + int32_t offset = rowSizeOffset + sizeof(int32_t); + int32_t accumulatorFlagsOffset = 0; + int32_t alignment = 1; + + for (auto& aggregate : aggregates_) { + auto& function = aggregate.function_; + Accumulator accumulator(function.get()); + // Accumulator offset must be aligned by their alignment size. + offset = milvus::bits::roundUp(offset, accumulator.alignment()); + function->setOffsets(offset, + RowContainer::nullByte(accumulatorFlagsOffset), + RowContainer::nullMask(accumulatorFlagsOffset), + rowSizeOffset); + offset += accumulator.fixedWidthSize(); + accumulatorFlagsOffset += 1; + alignment = + RowContainer::combineAlignments(accumulator.alignment(), alignment); + } + AssertInfo(__builtin_popcount(alignment) == 1, + "alignment of aggregations must be power of two"); + offset = milvus::Align(offset, alignment); + // Use unique_ptr for exception-safe memory management + // If initializeNewGroups throws, the buffer will be automatically cleaned up + globalAggregationBuffer_ = std::make_unique(offset); + lookup_->hits_[0] = globalAggregationBuffer_.get(); + const auto singleGroup = std::vector{0}; + for (auto& aggregate : aggregates_) { + aggregate.function_->initializeNewGroups(lookup_->hits_.data(), + singleGroup); + } + globalAggregationInitialized_ = true; +} + +void +GroupingSet::addGlobalAggregationInput(const milvus::RowVectorPtr& input) { + initializeGlobalAggregation(); + auto* group = lookup_->hits_[0]; + for (auto i = 0; i < aggregates_.size(); i++) { + auto& function = aggregates_[i].function_; + populateTempVectors(i, input); + function->addSingleGroupRawInput(group, tempVectors_); + } + tempVectors_.clear(); +} + +bool +GroupingSet::getGlobalAggregationOutput(milvus::RowVectorPtr& result) { + initializeGlobalAggregation(); // when input from upstream operator is empty, we need to initialize the accumulators for global aggregation + AssertInfo(lookup_->hits_.size() == 1, + "GlobalAggregation should have exactly one output line"); + auto groups = lookup_->hits_.data(); + for (auto i = 0; i < aggregates_.size(); i++) { + auto& function = aggregates_[i].function_; + auto resultVector = result->child(aggregates_[i].output_); + function->extractValues(groups, 1, &resultVector); + } + return true; +} + +bool +GroupingSet::getOutput(milvus::RowVectorPtr& result) { + if (isGlobal_) { + return getGlobalAggregationOutput(result); + } + AssertInfo(hash_table_ != nullptr, + "hash_table_ should not be nullptr for non-global aggregation"); + const auto& all_rows = hash_table_->rows()->allRows(); + if (!all_rows.empty()) { + extractGroups(result); + return true; + } + return false; +} + +std::vector +GroupingSet::accumulators() { + std::vector accumulators; + accumulators.reserve(aggregates_.size()); + for (auto& aggregate : aggregates_) { + accumulators.emplace_back(Accumulator{aggregate.function_.get()}); + } + return accumulators; +} + +void +GroupingSet::ensureInputFits(const RowVectorPtr& input) { + //TODO memory check +} + +void +GroupingSet::extractGroups(const milvus::RowVectorPtr& result) { + RowContainer* rows = hash_table_->rows(); + const auto& groups = rows->allRows(); + result->resize(groups.size()); + auto totalKeys = rows->KeyTypes().size(); + auto groups_range = + folly::Range(const_cast(groups.data()), groups.size()); + for (auto i = 0; i < totalKeys; i++) { + auto keyVector = result->child(i); + rows->extractColumn( + groups_range.data(), groups_range.size(), i, keyVector); + } + for (auto i = 0; i < aggregates_.size(); i++) { + auto& function = aggregates_[i].function_; + auto aggregateVector = result->child(totalKeys + i); + function->extractValues( + groups_range.data(), groups_range.size(), &aggregateVector); + } +} + +void +GroupingSet::addInputForActiveRows(const RowVectorPtr& input) { + AssertInfo( + !isGlobal_, + "Global aggregations should not reach add input for active rows"); + if (!hash_table_) { + createHashTable(); + } + ensureInputFits(input); + hash_table_->prepareForGroupProbe(*lookup_, input); + hash_table_->groupProbe(*lookup_); + auto& hits = lookup_->hits_; + auto* groups = hits.data(); + auto numGroups = hits.size(); + const auto& newGroups = lookup_->newGroups_; + for (auto i = 0; i < aggregates_.size(); i++) { + auto& function = aggregates_[i].function_; + if (!newGroups.empty()) { + function->initializeNewGroups(groups, newGroups); + } + populateTempVectors(i, input); + function->addRawInput(groups, numGroups, tempVectors_); + } + tempVectors_.clear(); +} + +void +GroupingSet::populateTempVectors(int32_t aggregateIndex, + const milvus::RowVectorPtr& input) { + const auto& channel_idxes = aggregates_[aggregateIndex].input_column_idxes_; + if (channel_idxes.empty() && input->childrens().size() == 1) { + tempVectors_.resize(1); + tempVectors_[0] = input->child(0); + } else { + tempVectors_.resize(channel_idxes.size()); + for (auto i = 0; i < channel_idxes.size(); i++) { + tempVectors_[i] = input->child(channel_idxes[i]); + } + } +} + +int32_t +GroupingSet::outputRowCount() const { + return lookup_->newGroups_.size(); +} + +void +initializeAggregates(const std::vector& aggregates, + RowContainer& rows) { + const auto numKeys = rows.KeyTypes().size(); + int i = 0; + for (auto& aggregate : aggregates) { + auto& function = aggregate.function_; + const auto& rowColumn = rows.columnAt(numKeys + i); + function->setOffsets(rowColumn.offset(), + rowColumn.nullByte(), + rowColumn.nullMask(), + rows.rowSizeOffset()); + i++; + } +} + +void +GroupingSet::createHashTable() { + hash_table_ = + std::make_unique(std::move(hashers_), accumulators()); + auto& rows = *(hash_table_->rows()); + initializeAggregates(aggregates_, rows); + lookup_ = std::make_unique(hash_table_->hashers()); +} + +} // namespace exec +} // namespace milvus \ No newline at end of file diff --git a/internal/core/src/exec/operator/query-agg/GroupingSet.h b/internal/core/src/exec/operator/query-agg/GroupingSet.h new file mode 100644 index 0000000000..e104465045 --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/GroupingSet.h @@ -0,0 +1,99 @@ +// 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. + +#pragma once +#include "common/Types.h" +#include "exec/VectorHasher.h" +#include "AggregateInfo.h" +#include "exec/HashTable.h" +#include "plan/PlanNode.h" +#include "RowContainer.h" + +namespace milvus { +namespace exec { + +class GroupingSet { + public: + GroupingSet(const RowTypePtr& input_type, + std::vector>&& hashers, + std::vector&& aggregates) + : hashers_(std::move(hashers)), aggregates_(std::move(aggregates)) { + isGlobal_ = hashers_.empty(); + } + + ~GroupingSet(); + + void + addInput(const RowVectorPtr& input); + + void + initializeGlobalAggregation(); + + void + addGlobalAggregationInput(const RowVectorPtr& input); + + void + addInputForActiveRows(const RowVectorPtr& input); + + void + createHashTable(); + + std::vector + accumulators(); + + // Checks if input will fit in the existing memory and increases reservation + // if not. If reservation cannot be increased, spills enough to make 'input' + // fit. + void + ensureInputFits(const RowVectorPtr& input); + + bool + getOutput(RowVectorPtr& result); + + void + extractGroups(const RowVectorPtr& result); + + void + populateTempVectors(int32_t aggregateIndex, const RowVectorPtr& input); + + bool + getGlobalAggregationOutput(RowVectorPtr& result); + + int32_t + outputRowCount() const; + + private: + bool isGlobal_; + + std::vector> hashers_; + std::vector aggregates_; + + // Place for the arguments of the aggregate being updated. + std::vector tempVectors_; + std::unique_ptr hash_table_; + std::unique_ptr lookup_; + uint64_t numInputRows_ = 0; + + // RAII-managed buffer for global aggregation to ensure exception safety + std::unique_ptr globalAggregationBuffer_; + + // Boolean indicating whether accumulators for a global aggregation (i.e. hashers_.empty()) are initialized + // This is used to avoid segv when getting output directly without input for empty output of upstream operator + bool globalAggregationInitialized_{false}; +}; + +} // namespace exec +} // namespace milvus diff --git a/internal/core/src/exec/operator/query-agg/MaxAggregate.cpp b/internal/core/src/exec/operator/query-agg/MaxAggregate.cpp new file mode 100644 index 0000000000..0c5bfc8b08 --- /dev/null +++ b/internal/core/src/exec/operator/query-agg/MaxAggregate.cpp @@ -0,0 +1,76 @@ +// 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. + +#include "MaxAggregateBase.h" +#include "log/Log.h" + +namespace milvus { +namespace exec { + +template