add group agg spill

pkg/sql/colexec/group/exec.go

@@ -63,6 +63,14 @@ func (group *Group) Prepare(proc *process.Process) (err error) {
 	if err = group.prepareGroup(proc); err != nil {
 		return err
 	}
+
+	if group.SpillManager == nil {
+		group.SpillManager = NewMemorySpillManager()
+	}
+	if group.SpillThreshold <= 0 {
+		group.SpillThreshold = 256 * 1024 //TODO configurable
+	}
+
 	return group.PrepareProjection(proc)
 }
 
@@ -188,6 +196,12 @@ func (group *Group) callToGetFinalResult(proc *process.Process) (*batch.Batch, e
 
 	for {
 		if group.ctr.state == vm.Eval {
+			if len(group.ctr.spilledStates) > 0 {
+				if err := group.mergeSpilledResults(proc); err != nil {
+					return nil, err
+				}
+			}
+
 			if group.ctr.result1.IsEmpty() {
 				group.ctr.state = vm.End
 				return nil, nil
@@ -203,7 +217,7 @@ func (group *Group) callToGetFinalResult(proc *process.Process) (*batch.Batch, e
 			group.ctr.state = vm.Eval
 
 			if group.ctr.isDataSourceEmpty() && len(group.Exprs) == 0 {
-				if err = group.generateInitialResult1WithoutGroupBy(proc); err != nil {
+				if err := group.generateInitialResult1WithoutGroupBy(proc); err != nil {
 					return nil, err
 				}
 				group.ctr.result1.ToPopped[0].SetRowCount(1)
@@ -215,6 +229,16 @@ func (group *Group) callToGetFinalResult(proc *process.Process) (*batch.Batch, e
 		}
 
 		group.ctr.dataSourceIsEmpty = false
+
+		// Check if we need to spill before processing this batch
+		group.updateMemoryUsage(proc)
+		if group.shouldSpill() {
+			if err := group.spillPartialResults(proc); err != nil {
+				return nil, err
+			}
+			continue
+		}
+
 		if err = group.consumeBatchToGetFinalResult(proc, res); err != nil {
 			return nil, err
 		}
@@ -245,7 +269,6 @@ func (group *Group) consumeBatchToGetFinalResult(
 
 	switch group.ctr.mtyp {
 	case H0:
-		// without group by.
 		if group.ctr.result1.IsEmpty() {
 			if err := group.generateInitialResult1WithoutGroupBy(proc); err != nil {
 				return err
@@ -260,7 +283,6 @@ func (group *Group) consumeBatchToGetFinalResult(
 		}
 
 	default:
-		// with group by.
 		if group.ctr.result1.IsEmpty() {
 			err := group.ctr.hr.BuildHashTable(false, group.ctr.mtyp == HStr, group.ctr.keyNullable, group.PreAllocSize)
 			if err != nil {
@@ -315,6 +337,16 @@ func (group *Group) consumeBatchToGetFinalResult(
 		}
 	}
 
+	// Update memory usage after processing the batch
+	group.updateMemoryUsage(proc)
+
+	// Check if we need to spill after processing this batch
+	if group.shouldSpill() {
+		if err := group.spillPartialResults(proc); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 

pkg/sql/colexec/group/exec_test.go

@@ -46,6 +46,10 @@ func (h *hackAggExecToTest) GetOptResult() aggexec.SplitResult {
 	return nil
 }
 
+func (h *hackAggExecToTest) Size() int64 {
+	return 0
+}
+
 func (h *hackAggExecToTest) GroupGrow(more int) error {
 	h.groupNumber += more
 	return nil

pkg/sql/colexec/group/execctx.go

@@ -123,7 +123,7 @@ type GroupResultBuffer struct {
 }
 
 func (buf *GroupResultBuffer) IsEmpty() bool {
-	return cap(buf.ToPopped) == 0
+	return len(buf.ToPopped) == 0
 }
 
 func (buf *GroupResultBuffer) InitOnlyAgg(chunkSize int, aggList []aggexec.AggFuncExec) {