feat(aggexec): add approx_percentile aggregate function for p95/p99

pkg/container/batch/batch_test.go

@@ -18,8 +18,6 @@ import (
 	"bytes"
 	"testing"
 
-	"github.com/matrixorigin/matrixone/pkg/sql/colexec/aggexec"
-
 	"github.com/matrixorigin/matrixone/pkg/common/mpool"
 	"github.com/matrixorigin/matrixone/pkg/container/types"
 	"github.com/matrixorigin/matrixone/pkg/container/vector"
@@ -155,7 +153,6 @@ func newBatch(ts []types.Type, rows int) *Batch {
 	}
 
 	bat.ExtraBuf = []byte("extra buf")
-	aggexec.RegisterGroupConcatAgg(0, ",")
 	bat.Attrs = []string{"1"}
 	return bat
 }

pkg/sql/colexec/aggexec/aggState.go

@@ -46,6 +46,10 @@ type MarshalerUnmarshaler interface {
 	UnmarshalFromReader(io.Reader) error
 }
 
+type freeableMarshalerUnmarshaler interface {
+	Free()
+}
+
 type aggInfo struct {
 	aggId                    int64
 	isDistinct               bool
@@ -648,6 +652,12 @@ func (ag *aggState) free(mp *mpool.MPool) {
 		vec.Free(mp)
 	}
 	ag.vecs = nil
+	for _, mob := range ag.mobs {
+		if freeable, ok := mob.(freeableMarshalerUnmarshaler); ok {
+			freeable.Free()
+		}
+	}
+	ag.mobs = nil
 }
 
 type aggExec struct {
@@ -822,9 +832,15 @@ func checkAggStateMagic(reader io.Reader) {
 	}
 }
 
-func (ae *aggExec) UnmarshalFromReader(reader io.Reader, mp *mpool.MPool) error {
+func (ae *aggExec) UnmarshalFromReader(reader io.Reader, mp *mpool.MPool) (retErr error) {
 	checkAggStateMagic(reader)
 	defer checkAggStateMagic(reader)
+	defer func() {
+		if retErr != nil {
+			ae.Free()
+			ae.state = nil
+		}
+	}()
 
 	// read number of chunks
 	cnt, err := types.ReadInt32(reader)

pkg/sql/colexec/aggexec/aggexec_test.go

@@ -30,11 +30,9 @@ func TestMakeAggSpecialAgg(t *testing.T) {
 		require.Equal(t, int64(0), mp.CurrNB())
 	}()
 
-	const testAggID = -9901
 	param := types.T_int64.ToType()
-	RegisterAvgTwCache(testAggID)
 
-	exec, err := MakeAgg(mp, testAggID, false, param)
+	exec, err := MakeAgg(mp, AggIdOfAvgTwCache, false, param)
 	require.NoError(t, err)
 
 	require.NoError(t, exec.GroupGrow(2))
@@ -73,6 +71,59 @@ func TestVectorsUnmarshalFromReader(t *testing.T) {
 	exec2.Free()
 }
 
+func TestAggExecUnmarshalReplacesExistingState(t *testing.T) {
+	mp := mpool.MustNewZero()
+	defer func() {
+		require.Equal(t, int64(0), mp.CurrNB())
+	}()
+	newExec := func() *countColumnExec {
+		return newCountColumnExec(
+			mp,
+			AggIdOfCountColumn,
+			false,
+			[]types.Type{types.T_int64.ToType()},
+		).(*countColumnExec)
+	}
+
+	t.Run("empty", func(t *testing.T) {
+		target := newExec()
+		defer target.Free()
+		require.NoError(t, target.GroupGrow(1))
+
+		empty := newExec()
+		var buf bytes.Buffer
+		require.NoError(t, empty.SaveIntermediateResult(0, nil, &buf))
+		empty.Free()
+		require.NoError(t, target.UnmarshalFromReader(bytes.NewReader(buf.Bytes()), mp))
+		require.Zero(t, target.GetNumGroups())
+	})
+
+	t.Run("multiple_chunks", func(t *testing.T) {
+		target := newExec()
+		defer target.Free()
+		require.NoError(t, target.GroupGrow(2))
+
+		source := newExec()
+		defer source.Free()
+		require.NoError(t, source.GroupGrow(AggBatchSize+1))
+		flags := [][]uint8{make([]uint8, AggBatchSize), {1}}
+		for i := range flags[0] {
+			flags[0][i] = 1
+		}
+		var buf bytes.Buffer
+		require.NoError(t, source.SaveIntermediateResult(AggBatchSize+1, flags, &buf))
+		require.NoError(t, target.UnmarshalFromReader(bytes.NewReader(buf.Bytes()), mp))
+		require.Equal(t, AggBatchSize+1, target.GetNumGroups())
+
+		results, err := target.Flush()
+		require.NoError(t, err)
+		require.Len(t, results, 2)
+		for _, result := range results {
+			result.Free(mp)
+		}
+	})
+}
+
 func TestIntermediateResultCompactsAndReusesSingleChunk(t *testing.T) {
 	mp := mpool.MustNewZero()
 	defer func() {