fix: HashTable dynamic rehash and group count limit for GROUP BY aggregation (#47569)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: MrPresent-Han <chun.han@gmail.com>

Author: MrPresent-Han

configs/milvus.yaml

@@ -1055,6 +1055,8 @@ common:
     mode:
       queryNode: sync # File resource mode for query node, options: [sync, close]. Default is sync.
       dataNode: sync # File resource mode for data node, options: [sync, ref, close]. Default is sync.
+  groupBy:
+    maxGroups: 100000 # Maximum number of groups allowed in GROUP BY aggregation, enforced both per segment and during cross-segment merge. Exceeding this limit fails the query.
 
 # QuotaConfig, configurations of Milvus quota and limits.
 # By default, we enable:

internal/agg/aggregate_reducer.go

@@ -9,6 +9,8 @@ import (
 	"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/merr"
+	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
 )
 
@@ -345,11 +347,12 @@ func (reducer *GroupAggReducer) Reduce(ctx context.Context, results []*Aggregati
 	}
 
 	// 2. compute hash values for all rows in the result retrieved
-	var totalRowCount int64 = 0
+	var totalGroupCount int64 = 0
+	maxGroupByGroups := paramtable.Get().CommonCfg.GroupByMaxGroups.GetAsInt64()
 processResults:
 	for _, result := range results {
 		// Check limit before processing each shard to avoid unnecessary work
-		if reducer.groupLimit != -1 && totalRowCount >= reducer.groupLimit {
+		if reducer.groupLimit != -1 && totalGroupCount >= reducer.groupLimit {
 			break processResults
 		}
 
@@ -387,7 +390,7 @@ processResults:
 
 		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 {
+			if reducer.groupLimit != -1 && totalGroupCount >= reducer.groupLimit {
 				break processResults
 			}
 			rowFieldValues := make([]*FieldValue, outputColumnCount)
@@ -416,24 +419,29 @@ processResults:
 			if bucket := reducer.hashValsMap[hashVal]; bucket == nil {
 				newBucket := NewBucket()
 				newBucket.AddRow(newRow)
-				totalRowCount++
+				totalGroupCount++
 				reducer.hashValsMap[hashVal] = newBucket
 			} else {
 				if rowIdx := bucket.Find(newRow, numGroupingKeys); rowIdx == NONE {
 					bucket.AddRow(newRow)
-					totalRowCount++
+					totalGroupCount++
 				} else {
 					if err := bucket.Accumulate(newRow, rowIdx, numGroupingKeys, aggs); err != nil {
 						return nil, err
 					}
 				}
 			}
+			if totalGroupCount > maxGroupByGroups {
+				return nil, merr.WrapErrServiceInternal(fmt.Sprintf("GROUP BY produced too many groups (%d). "+
+					"Add filters or increase common.groupBy.maxGroups (current: %d)",
+					totalGroupCount, maxGroupByGroups))
+			}
 			// 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)
+	reducedResult.fieldDatas = typeutil.PrepareResultFieldData(firstFieldData, totalGroupCount)
 	for _, bucket := range reducer.hashValsMap {
 		err := AssembleBucket(bucket, reducedResult.GetFieldDatas())
 		if err != nil {

internal/agg/aggregate_reducer_test.go

@@ -2,15 +2,22 @@ package agg
 
 import (
 	"context"
+	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
 	"github.com/milvus-io/milvus/pkg/v2/proto/planpb"
+	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
 )
 
+func init() {
+	paramtable.Init()
+}
+
 func makeTestSchema() *schemapb.CollectionSchema {
 	return &schemapb.CollectionSchema{
 		Fields: []*schemapb.FieldSchema{
@@ -177,3 +184,116 @@ func TestReduceSingleResult(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, singleResult, out)
 }
+
+// buildTestSchema creates a simple schema with an INT64 groupBy field and an INT64 agg field.
+func buildTestSchema() *schemapb.CollectionSchema {
+	return &schemapb.CollectionSchema{
+		Fields: []*schemapb.FieldSchema{
+			{FieldID: 100, Name: "group_field", DataType: schemapb.DataType_Int64},
+			{FieldID: 101, Name: "agg_field", DataType: schemapb.DataType_Int64},
+		},
+	}
+}
+
+// buildAggResult creates an AggregationResult with N distinct groups.
+// Each group has group key = startKey+i and count = 1.
+func buildAggResult(startKey int64, numGroups int) *AggregationResult {
+	groupKeys := make([]int64, numGroups)
+	counts := make([]int64, numGroups)
+	for i := 0; i < numGroups; i++ {
+		groupKeys[i] = startKey + int64(i)
+		counts[i] = 1
+	}
+	return NewAggregationResult([]*schemapb.FieldData{
+		{
+			Type:      schemapb.DataType_Int64,
+			FieldName: "group_field",
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_LongData{
+						LongData: &schemapb.LongArray{Data: groupKeys},
+					},
+				},
+			},
+		},
+		{
+			Type:      schemapb.DataType_Int64,
+			FieldName: "agg_field",
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_LongData{
+						LongData: &schemapb.LongArray{Data: counts},
+					},
+				},
+			},
+		},
+	}, int64(numGroups))
+}
+
+func TestGroupAggReducer_MaxGroupByGroupsExceeded(t *testing.T) {
+	maxGroups := int64(10)
+	paramtable.Get().Save(paramtable.Get().CommonCfg.GroupByMaxGroups.Key, fmt.Sprintf("%d", maxGroups))
+	defer paramtable.Get().Reset(paramtable.Get().CommonCfg.GroupByMaxGroups.Key)
+
+	schema := buildTestSchema()
+	aggregates := []*planpb.Aggregate{
+		{Op: planpb.AggregateOp_count, FieldId: 101},
+	}
+	reducer := NewGroupAggReducer([]int64{100}, aggregates, -1, schema)
+
+	// Two results each with 10 distinct groups (20 total > 10 limit)
+	results := []*AggregationResult{
+		buildAggResult(0, 10),
+		buildAggResult(10, 10),
+	}
+
+	_, err := reducer.Reduce(context.Background(), results)
+	require.Error(t, err)
+	assert.True(t, strings.Contains(err.Error(), "too many groups"))
+}
+
+func TestGroupAggReducer_MaxGroupByGroupsExactlyAtLimit(t *testing.T) {
+	maxGroups := int64(10)
+	paramtable.Get().Save(paramtable.Get().CommonCfg.GroupByMaxGroups.Key, fmt.Sprintf("%d", maxGroups))
+	defer paramtable.Get().Reset(paramtable.Get().CommonCfg.GroupByMaxGroups.Key)
+
+	schema := buildTestSchema()
+	aggregates := []*planpb.Aggregate{
+		{Op: planpb.AggregateOp_count, FieldId: 101},
+	}
+	reducer := NewGroupAggReducer([]int64{100}, aggregates, -1, schema)
+
+	// Exactly 10 groups = limit, should succeed
+	// Use 2 results to force cross-segment merge path (single result fast-returns)
+	results := []*AggregationResult{
+		buildAggResult(0, 5),
+		buildAggResult(5, 5),
+	}
+
+	result, err := reducer.Reduce(context.Background(), results)
+	require.NoError(t, err)
+	assert.NotNil(t, result)
+}
+
+func TestGroupAggReducer_MaxGroupByGroupsJustOverLimit(t *testing.T) {
+	maxGroups := int64(10)
+	paramtable.Get().Save(paramtable.Get().CommonCfg.GroupByMaxGroups.Key, fmt.Sprintf("%d", maxGroups))
+	defer paramtable.Get().Reset(paramtable.Get().CommonCfg.GroupByMaxGroups.Key)
+
+	schema := buildTestSchema()
+	aggregates := []*planpb.Aggregate{
+		{Op: planpb.AggregateOp_count, FieldId: 101},
+	}
+	reducer := NewGroupAggReducer([]int64{100}, aggregates, -1, schema)
+
+	// 6 + 5 = 11 distinct groups > 10 limit, should fail
+	// Need 2 results to trigger cross-segment merge path (single result fast-returns)
+	results := []*AggregationResult{
+		buildAggResult(0, 6),
+		buildAggResult(6, 5),
+	}
+
+	_, err := reducer.Reduce(context.Background(), results)
+	require.Error(t, err)
+	assert.True(t, strings.Contains(err.Error(), "too many groups"))
+}

internal/core/src/exec/HashTable.cpp

@@ -22,6 +22,7 @@
 
 #include "common/SimdUtil.h"
 #include "exec/VectorHasher.h"
+#include "fmt/format.h"
 
 namespace milvus {
 namespace exec {
@@ -222,10 +223,20 @@ char*
 HashTable::insertEntry(milvus::exec::HashLookup& lookup,
                        uint64_t index,
                        milvus::vector_size_t row) {
+    if (numDistinct_ >= maxNumGroups_) {
+        ThrowInfo(
+            UnexpectedError,
+            fmt::format("GROUP BY produced too many groups ({}). "
+                        "Add filters or increase common.groupBy.maxGroups "
+                        "(current: {})",
+                        numDistinct_ + 1,
+                        maxNumGroups_));
+    }
     char* group = rows_->newRow();
     lookup.hits_[row] = group;
     storeKeys(lookup, row);
     storeRowPointer(index, lookup.hashes_[row], group);
+    rowHashes_.push_back(lookup.hashes_[row]);
     numDistinct_++;
     lookup.newGroups_.push_back(row);
     return group;
@@ -250,6 +261,9 @@ HashTable::groupProbe(milvus::exec::HashLookup& lookup) {
     checkSizeAndAllocateTable(0);
     ProbeState state;
     for (int32_t idx = 0; idx < lookup.hashes_.size(); idx++) {
+        if (numDistinct_ >= rehashSize()) {
+            rehash();
+        }
         state.preProbe(*this, lookup.hashes_[idx], idx);
         state.firstProbe<ProbeState::Operation::kInsert>(*this);
         fullProbe(lookup, state);
@@ -272,6 +286,40 @@ HashTable::clear(bool freeTable) {
     numBuckets_ = 0;
     sizeMask_ = 0;
     bucketOffsetMask_ = 0;
+    rowHashes_.clear();
+}
+
+void
+HashTable::insertForRehash(char* row, uint64_t hash) {
+    const auto tag = hashTag(hash);
+    const auto kEmptyGroup = TagVector::broadcast(0);
+    int64_t bktOffset = bucketOffset(hash);
+    for (int64_t i = 0; i < numBuckets_; i++) {
+        auto tags = loadTags(bktOffset);
+        uint16_t empty = toBitMask(tags == kEmptyGroup) & 0xffff;
+        if (empty > 0) {
+            auto pos = bits::getAndClearLastSetBit(empty);
+            auto* bucket = bucketAt(bktOffset);
+            bucket->setTag(pos, tag);
+            bucket->setPointer(pos, row);
+            return;
+        }
+        bktOffset = nextBucketOffset(bktOffset);
+    }
+    AssertInfo(false, "Failed to insert during rehash");
+}
+
+void
+HashTable::rehash() {
+    // allRows is safe to reference across allocateTables() because
+    // allocateTables() only rebuilds the hash bucket array (table_),
+    // it does not touch the RowContainer (rows_) that owns the row data.
+    const auto& allRows = rows_->allRows();
+    allocateTables(capacity_ * 2);
+    for (size_t i = 0; i < allRows.size(); i++) {
+        insertForRehash(allRows[i], rowHashes_[i]);
+    }
+    numRehashes_++;
 }
 
 }  // namespace exec

internal/core/src/exec/HashTable.h

@@ -34,6 +34,7 @@
 #include "common/Vector.h"
 #include "exec/operator/query-agg/RowContainer.h"
 #include "folly/CPortability.h"
+#include "segcore/SegcoreConfig.h"
 #include "xsimd/xsimd.hpp"
 
 namespace milvus {
@@ -174,9 +175,11 @@ class ProbeState;
 
 class HashTable : public BaseHashTable {
  public:
-    HashTable(std::vector<std::unique_ptr<VectorHasher>>&& hashers,
-              const std::vector<Accumulator>& accumulators)
-        : BaseHashTable(std::move(hashers)) {
+    HashTable(
+        std::vector<std::unique_ptr<VectorHasher>>&& hashers,
+        const std::vector<Accumulator>& accumulators,
+        int64_t maxNumGroups = segcore::SegcoreConfig::kDefaultMaxGroupByGroups)
+        : BaseHashTable(std::move(hashers)), maxNumGroups_(maxNumGroups) {
         std::vector<DataType> keyTypes;
         for (auto& hasher : hashers_) {
             keyTypes.push_back(hasher->ChannelDataType());
@@ -305,6 +308,12 @@ class HashTable : public BaseHashTable {
     void
     checkSizeAndAllocateTable(int32_t numNew);
 
+    void
+    rehash();
+
+    void
+    insertForRehash(char* row, uint64_t hash);
+
     // Returns the number of entries after which the table gets rehashed.
     static uint64_t
     rehashSize(int64_t size) {
@@ -344,6 +353,8 @@ class HashTable : public BaseHashTable {
 
     [[maybe_unused]] int64_t numRehashes_{0};
     char* table_ = nullptr;
+    std::vector<uint64_t> rowHashes_;
+    int64_t maxNumGroups_;
 
     HashMode
     hashMode() const override {

internal/core/src/exec/operator/query-agg/GroupingSet.cpp

@@ -27,6 +27,7 @@
 #include "exec/operator/query-agg/AggregateInfo.h"
 #include "exec/operator/query-agg/RowContainer.h"
 #include "folly/Range.h"
+#include "segcore/SegcoreConfig.h"
 
 namespace milvus {
 namespace exec {
@@ -248,8 +249,10 @@ initializeAggregates(const std::vector<AggregateInfo>& aggregates,
 
 void
 GroupingSet::createHashTable() {
-    hash_table_ =
-        std::make_unique<HashTable>(std::move(hashers_), accumulators());
+    auto maxGroups =
+        segcore::SegcoreConfig::default_config().get_max_group_by_groups();
+    hash_table_ = std::make_unique<HashTable>(
+        std::move(hashers_), accumulators(), maxGroups);
     auto& rows = *(hash_table_->rows());
     initializeAggregates(aggregates_, rows);
     lookup_ = std::make_unique<HashLookup>(hash_table_->hashers());

internal/core/src/segcore/SegcoreConfig.h

@@ -156,6 +156,18 @@ class SegcoreConfig {
         return enable_geometry_cache_;
     }
 
+    static constexpr int64_t kDefaultMaxGroupByGroups = 100000;
+
+    int64_t
+    get_max_group_by_groups() const {
+        return max_group_by_groups_;
+    }
+
+    void
+    set_max_group_by_groups(int64_t v) {
+        max_group_by_groups_ = v;
+    }
+
     void
     set_interim_index_mem_expansion_rate(float rate) {
         interim_index_mem_expansion_rate_ = rate;
@@ -186,6 +198,7 @@ class SegcoreConfig {
     inline static bool refine_with_quant_flag_ = false;
     inline static bool enable_geometry_cache_ = false;
     inline static float interim_index_mem_expansion_rate_ = 1.15f;
+    inline static int64_t max_group_by_groups_ = kDefaultMaxGroupByGroups;
 };
 
 }  // namespace milvus::segcore