executor: estimate analyze NDV from singleton sketches

Author: 0xPoe

pkg/executor/analyze_col_sampling.go

@@ -35,7 +35,6 @@ import (
 	"github.com/pingcap/tidb/pkg/tablecodec"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util"
-	"github.com/pingcap/tidb/pkg/util/channel"
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/codec"
 	"github.com/pingcap/tidb/pkg/util/collate"
@@ -222,6 +221,17 @@ func (e *AnalyzeColumnsExec) buildSamplingStats(
 	for range l {
 		rootRowCollector.Base().FMSketches = append(rootRowCollector.Base().FMSketches, statistics.NewFMSketch(statistics.MaxSketchSize))
 	}
+	needNDVSampling := statistics.NeedNDVSampling(e.analyzePB.ColReq.GetNdvRate())
+	var nodeNDVSketches [][]*statistics.FMSketch
+	var nodeSingletonSketches [][]*statistics.FMSketch
+	var nodeSampleSizes []int
+	var nodeSketchSampleCounts []int
+	if needNDVSampling {
+		nodeNDVSketches = make([][]*statistics.FMSketch, 0, samplingStatsConcurrency)
+		nodeSingletonSketches = make([][]*statistics.FMSketch, 0, samplingStatsConcurrency)
+		nodeSampleSizes = make([]int, 0, samplingStatsConcurrency)
+		nodeSketchSampleCounts = make([]int, 0, samplingStatsConcurrency)
+	}
 
 	sc := e.ctx.GetSessionVars().StmtCtx
 
@@ -281,6 +291,12 @@ func (e *AnalyzeColumnsExec) buildSamplingStats(
 			printAnalyzeMergeCollectorLog(oldRootCollectorCount, newRootCollectorCount,
 				mergeResult.collector.Base().Count, e.tableID.TableID, e.tableID.PartitionID, e.tableID.IsPartitionTable(),
 				"merge subMergeWorker in AnalyzeColumnsExec", -1)
+			if needNDVSampling {
+				nodeNDVSketches = append(nodeNDVSketches, copySketches(mergeResult.collector.Base().FMSketches))
+				nodeSingletonSketches = append(nodeSingletonSketches, copySketches(mergeResult.collector.Base().SingletonSketches()))
+				nodeSampleSizes = append(nodeSampleSizes, mergeResult.collector.Base().Samples.Len())
+				nodeSketchSampleCounts = append(nodeSketchSampleCounts, int(mergeResult.collector.Base().SketchSampleCount))
+			}
 			e.memTracker.Consume(rootRowCollector.Base().MemSize - oldRootCollectorSize - mergeResult.collector.Base().MemSize)
 			mergeResult.collector.DestroyAndPutToPool()
 		}
@@ -304,6 +320,27 @@ func (e *AnalyzeColumnsExec) buildSamplingStats(
 		return 0, nil, nil, nil, err
 	}
 
+	colLen := len(e.colsInfo)
+	indexPushedDownResult := <-idxNDVPushDownCh
+	if indexPushedDownResult.err != nil {
+		return 0, nil, nil, nil, indexPushedDownResult.err
+	}
+	for _, offset := range indexesWithVirtualColOffsets {
+		ret := indexPushedDownResult.results[e.indexes[offset].ID]
+		rootRowCollector.Base().NullCount[colLen+offset] = ret.Count
+		rootRowCollector.Base().FMSketches[colLen+offset] = ret.Ars[0].Fms[0]
+	}
+	isSpecialIndex := make([]bool, len(e.indexes))
+	for _, offset := range indexesWithVirtualColOffsets {
+		if offset >= 0 && offset < len(isSpecialIndex) {
+			isSpecialIndex[offset] = true
+		}
+	}
+	estimateNDVs := make([]int64, l)
+	if needNDVSampling {
+		estimateNDVs = estimateNDVsBySketch(rootRowCollector, nodeNDVSketches, nodeSingletonSketches, nodeSampleSizes, nodeSketchSampleCounts, colLen, isSpecialIndex)
+	}
+
 	// Decode the data from sample collectors.
 	virtualColIdx := buildVirtualColumnIndex(e.schemaForVirtualColEval, e.colsInfo)
 	// Filling virtual columns is necessary here because these samples are used to build statistics for indexes that constructed by virtual columns.
@@ -335,7 +372,6 @@ func (e *AnalyzeColumnsExec) buildSamplingStats(
 			return 0, nil, nil, nil, err
 		}
 	}
-	colLen := len(e.colsInfo)
 	// The order of the samples are broken when merging samples from sub-collectors.
 	// So now we need to sort the samples according to the handle in order to calculate correlation.
 	slices.SortFunc(rootRowCollector.Base().Samples, func(i, j *statistics.ReservoirRowSampleItem) int {
@@ -358,7 +394,7 @@ func (e *AnalyzeColumnsExec) buildSamplingStats(
 	// Start workers to build stats.
 	for range samplingStatsConcurrency {
 		e.samplingBuilderWg.Run(func() {
-			e.subBuildWorker(ctx, buildResultChan, buildTaskChan, hists, topns, exitCh)
+			e.subBuildWorker(ctx, buildResultChan, buildTaskChan, hists, topns, estimateNDVs, exitCh)
 		})
 	}
 	// Generate tasks for building stats.
@@ -373,18 +409,6 @@ func (e *AnalyzeColumnsExec) buildSamplingStats(
 		fmSketches = append(fmSketches, rootRowCollector.Base().FMSketches[i])
 	}
 
-	indexPushedDownResult := <-idxNDVPushDownCh
-	if indexPushedDownResult.err != nil {
-		close(exitCh)
-		channel.Clear(buildResultChan)
-		return 0, nil, nil, nil, indexPushedDownResult.err
-	}
-	for _, offset := range indexesWithVirtualColOffsets {
-		ret := indexPushedDownResult.results[e.indexes[offset].ID]
-		rootRowCollector.Base().NullCount[colLen+offset] = ret.Count
-		rootRowCollector.Base().FMSketches[colLen+offset] = ret.Ars[0].Fms[0]
-	}
-
 	// Generate tasks for building stats for indexes.
 	for i, idx := range e.indexes {
 		buildTaskChan <- &samplingBuildTask{
@@ -713,7 +737,7 @@ func drainPendingSamplingMergeTasks(taskCh <-chan []byte, memTracker *memory.Tra
 	}
 }
 
-func (e *AnalyzeColumnsExec) subBuildWorker(ctx context.Context, resultCh chan error, taskCh chan *samplingBuildTask, hists []*statistics.Histogram, topns []*statistics.TopN, exitCh chan struct{}) {
+func (e *AnalyzeColumnsExec) subBuildWorker(ctx context.Context, resultCh chan error, taskCh chan *samplingBuildTask, hists []*statistics.Histogram, topns []*statistics.TopN, estimateNDVs []int64, exitCh chan struct{}) {
 	defer func() {
 		if r := recover(); r != nil {
 			logutil.BgLogger().Warn("analyze subBuildWorker panicked", zap.Any("recover", r), zap.Stack("stack"))
@@ -887,6 +911,9 @@ workLoop:
 				releaseCollectorMemory()
 				continue
 			}
+			if hist != nil && task.slicePos < len(estimateNDVs) && estimateNDVs[task.slicePos] > hist.NDV {
+				hist.NDV = estimateNDVs[task.slicePos]
+			}
 			finalMemSize := hist.MemoryUsage() + topn.MemoryUsage()
 			e.memTracker.Consume(finalMemSize)
 			hists[task.slicePos] = hist
@@ -920,6 +947,105 @@ type samplingBuildTask struct {
 	slicePos         int
 }
 
+func copySketches(src []*statistics.FMSketch) []*statistics.FMSketch {
+	if len(src) == 0 {
+		return nil
+	}
+	dst := make([]*statistics.FMSketch, len(src))
+	for i, sketch := range src {
+		if sketch != nil {
+			dst[i] = sketch.Copy()
+		}
+	}
+	return dst
+}
+
+func estimateNDVsBySketch(
+	rootCollector statistics.RowSampleCollector,
+	nodeNDVSketches [][]*statistics.FMSketch,
+	nodeSingletonSketches [][]*statistics.FMSketch,
+	nodeSampleSizes []int,
+	nodeSketchSampleCounts []int,
+	colLen int,
+	isSpecialIndex []bool,
+) []int64 {
+	totalLen := len(rootCollector.Base().FMSketches)
+	estimateNDVs := make([]int64, totalLen)
+	if totalLen == 0 || len(nodeNDVSketches) == 0 || len(nodeNDVSketches) != len(nodeSingletonSketches) {
+		return estimateNDVs
+	}
+	var sampleSize uint64
+	for _, size := range nodeSketchSampleCounts {
+		sampleSize += uint64(size)
+	}
+	if sampleSize == 0 {
+		for _, size := range nodeSampleSizes {
+			sampleSize += uint64(size)
+		}
+	}
+	if sampleSize == 0 {
+		return estimateNDVs
+	}
+	for i := range totalLen {
+		if i >= colLen {
+			idxOffset := i - colLen
+			if idxOffset >= 0 && idxOffset < len(isSpecialIndex) && isSpecialIndex[idxOffset] {
+				continue
+			}
+		}
+		sketch := rootCollector.Base().FMSketches[i]
+		if sketch == nil {
+			continue
+		}
+		sampleNDV := uint64(sketch.NDV())
+		if sampleNDV == 0 {
+			continue
+		}
+		rowCount := uint64(0)
+		if i < len(rootCollector.Base().NullCount) {
+			count := rootCollector.Base().Count - rootCollector.Base().NullCount[i]
+			if count > 0 {
+				rowCount = uint64(count)
+			}
+		} else if rootCollector.Base().Count > 0 {
+			rowCount = uint64(rootCollector.Base().Count)
+		}
+		if rowCount == 0 {
+			continue
+		}
+		if sampleNDV > rowCount {
+			sampleNDV = rowCount
+		}
+		ndvSketches := make([]*statistics.FMSketch, 0, len(nodeNDVSketches))
+		singletonSketches := make([]*statistics.FMSketch, 0, len(nodeSingletonSketches))
+		for nodeIdx := range nodeNDVSketches {
+			if i >= len(nodeNDVSketches[nodeIdx]) || i >= len(nodeSingletonSketches[nodeIdx]) {
+				continue
+			}
+			ndvSketch := nodeNDVSketches[nodeIdx][i]
+			singletonSketch := nodeSingletonSketches[nodeIdx][i]
+			if ndvSketch == nil || singletonSketch == nil {
+				continue
+			}
+			ndvSketches = append(ndvSketches, ndvSketch)
+			singletonSketches = append(singletonSketches, singletonSketch)
+		}
+		if len(ndvSketches) == 0 {
+			continue
+		}
+		singletonItems := statistics.EstimateGlobalSingletonBySketches(ndvSketches, singletonSketches)
+		logutil.BgLogger().Debug("analyze sampling ndv estimation",
+			zap.Int("item_index", i),
+			zap.Uint64("sample_ndv", sampleNDV),
+			zap.Uint64("singleton_items", singletonItems),
+			zap.Uint64("sample_size", sampleSize),
+			zap.Uint64("row_count", rowCount),
+		)
+		estimateNDVs[i] = int64(statistics.EstimateNDVByGEE(sampleNDV, singletonItems, sampleSize, rowCount))
+	}
+	return estimateNDVs
+}
+
 func readDataAndSendTask(ctx context.Context, sctx sessionctx.Context, handler *tableResultHandler, mergeTaskCh chan []byte, memTracker *memory.Tracker) error {
 	// After all tasks are sent, close the mergeTaskCh to notify the mergeWorker that all tasks have been sent.
 	defer close(mergeTaskCh)

pkg/executor/analyze_utils_test.go

@@ -20,7 +20,10 @@ import (
 	"testing"
 
 	"github.com/pingcap/tidb/pkg/planner/core"
+	"github.com/pingcap/tidb/pkg/statistics"
+	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/dbterror/exeerrors"
+	"github.com/pingcap/tidb/pkg/util/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -67,3 +70,35 @@ func BuildExecutorForTest(ctx context.Context, stmt *ExecStmt) error {
 	_, err := stmt.buildExecutor(ctx)
 	return err
 }
+
+func TestEstimateNDVsBySketch(t *testing.T) {
+	rootCollector := statistics.NewReservoirRowSampleCollector(1, 2)
+	rootCollector.Count = 100
+	rootCollector.FMSketches = append(rootCollector.FMSketches,
+		mustBuildFMSketch(t, 1, 2, 3, 4),
+		mustBuildFMSketch(t, 10, 11),
+	)
+
+	nodeNDVSketches := [][]*statistics.FMSketch{
+		{mustBuildFMSketch(t, 1, 2, 3), mustBuildFMSketch(t, 10)},
+		{mustBuildFMSketch(t, 2, 3, 4), mustBuildFMSketch(t, 11)},
+	}
+	nodeSingletonSketches := [][]*statistics.FMSketch{
+		{mustBuildFMSketch(t, 1, 2, 3), mustBuildFMSketch(t, 10)},
+		{mustBuildFMSketch(t, 2, 4), mustBuildFMSketch(t, 11)},
+	}
+
+	estimateNDVs := estimateNDVsBySketch(rootCollector, nodeNDVSketches, nodeSingletonSketches, nil, []int{3, 3}, 1, []bool{true})
+	require.Equal(t, int64(statistics.EstimateNDVByGEE(4, 2, 6, 100)), estimateNDVs[0])
+	require.Zero(t, estimateNDVs[1], "special index NDV should be left to the pushed-down index result")
+}
+
+func mustBuildFMSketch(t *testing.T, values ...int64) *statistics.FMSketch {
+	t.Helper()
+	sketch := statistics.NewFMSketch(1000)
+	sc := mock.NewContext().GetSessionVars().StmtCtx
+	for _, value := range values {
+		require.NoError(t, sketch.InsertValue(sc, types.NewIntDatum(value)))
+	}
+	return sketch
+}

pkg/executor/builder.go

@@ -3169,6 +3169,7 @@ func (b *executorBuilder) buildAnalyzeSamplingPushdown(
 		modifyCount = int64(val.(int))
 	})
 	sampleRate := new(float64)
+	ndvRate := float64(statistics.DefaultNDVSampleRate)
 	var sampleRateReason string
 	if opts[ast.AnalyzeOptNumSamples] == 0 {
 		*sampleRate = math.Float64frombits(opts[ast.AnalyzeOptSampleRate])
@@ -3229,6 +3230,7 @@ func (b *executorBuilder) buildAnalyzeSamplingPushdown(
 		BucketSize:   int64(opts[ast.AnalyzeOptNumBuckets]),
 		SampleSize:   int64(opts[ast.AnalyzeOptNumSamples]),
 		SampleRate:   sampleRate,
+		NdvRate:      &ndvRate,
 		SketchSize:   statistics.MaxSketchSize,
 		ColumnsInfo:  util.ColumnsToProto(task.ColsInfo, task.TblInfo.PKIsHandle, false, false),
 		ColumnGroups: colGroups,

pkg/store/mockstore/unistore/cophandler/analyze.go

@@ -457,6 +457,7 @@ func handleAnalyzeFullSamplingReq(
 		MaxSampleSize:   int(colReq.SampleSize),
 		MaxFMSketchSize: int(colReq.SketchSize),
 		SampleRate:      colReq.GetSampleRate(),
+		NDVSampleRate:   colReq.GetNdvRate(),
 		Rng:             rand.New(rand.NewSource(time.Now().UnixNano())),
 	}
 	collector, err := builder.Collect()