statistics: estimate NDV from FM sketches of sampled rows

DEPS.bzl

@@ -3805,8 +3805,8 @@ def go_deps():
         name = "com_github_pingcap_tipb",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/tipb",
-        sum = "h1:jQf0oStCMU6ZXGQmz5CmOIBsAwMWutorDHI49u8tURU=",
-        version = "v0.0.0-20260908093239-fed7bc47c39d",
+        sum = "h1:GalvMypJTiQvfc6p8fs7OS0z+X9O5ZAN+hTZ67QpTYM=",
+        version = "v0.0.0-20261006155659-59b678398b17",
     )
     go_repository(
         name = "com_github_pkg_browser",

go.mod

@@ -103,7 +103,7 @@ require (
 	github.com/pingcap/metering_sdk v0.0.0-20260814062708-9e3b68cd9adf
 	github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5
 	github.com/pingcap/tidb/pkg/parser v0.0.0-20211011031125-9b13dc409c5e
-	github.com/pingcap/tipb v0.0.0-20260908093239-fed7bc47c39d
+	github.com/pingcap/tipb v0.0.0-20261006155659-59b678398b17
 	github.com/prometheus/client_golang v1.23.0
 	github.com/prometheus/client_model v0.6.2
 	github.com/prometheus/common v0.65.0

go.sum

@@ -2186,8 +2186,8 @@ github.com/pingcap/metering_sdk v0.0.0-20260814062708-9e3b68cd9adf h1:8gzQg9J2EN
 github.com/pingcap/metering_sdk v0.0.0-20260814062708-9e3b68cd9adf/go.mod h1:IP3HAzDbY5sdoLTQSkF5HjZukuO2koKzbX2FglVSaIA=
 github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5 h1:T4pXRhBflzDeAhmOQHNPRRogMYxP13V7BkYw3ZsoSfE=
 github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5/go.mod h1:rlimy0GcTvjiJqvD5mXTRr8O2eNZPBrcUgiWVYp9530=
-github.com/pingcap/tipb v0.0.0-20260908093239-fed7bc47c39d h1:jQf0oStCMU6ZXGQmz5CmOIBsAwMWutorDHI49u8tURU=
-github.com/pingcap/tipb v0.0.0-20260908093239-fed7bc47c39d/go.mod h1:RM8iRcMalzOthG2XJxnNBniM4xFGb/lDwHUwqkaVzt4=
+github.com/pingcap/tipb v0.0.0-20261006155659-59b678398b17 h1:GalvMypJTiQvfc6p8fs7OS0z+X9O5ZAN+hTZ67QpTYM=
+github.com/pingcap/tipb v0.0.0-20261006155659-59b678398b17/go.mod h1:RM8iRcMalzOthG2XJxnNBniM4xFGb/lDwHUwqkaVzt4=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=

pkg/statistics/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
         "index.go",
         "row_sampler.go",
         "sample.go",
+        "sampled_ndv.go",
         "scalar.go",
         "table.go",
     ],
@@ -116,6 +117,7 @@ go_test(
         "//pkg/util/sqlkiller",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
+        "@com_github_pingcap_tipb//go-tipb",
         "@com_github_stretchr_testify//require",
         "@org_uber_go_goleak//:goleak",
     ],

pkg/statistics/fmsketch.go

@@ -54,21 +54,25 @@ const MaxSketchSize = 10000
 //  1. https://www.vldb.org/conf/2001/P541.pdf
 //  2. https://algo.inria.fr/flajolet/Publications/FlMa85.pdf
 type FMSketch struct {
-	// A set to store unique hashed values.
-	hashset map[uint64]struct{}
+	// A set to store unique hashed values. The value is true only in a sketch
+	// of sampled rows, for a hash seen more than once.
+	hashset map[uint64]bool
+	// ndvCounts is set only in a sketch of sampled rows. A sketch is sampled if
+	// and only if ndvCounts is not nil.
+	ndvCounts *ndvCounts
 	// A binary mask used to track the maximum number of trailing zeroes in the hashed values.
 	// Also used to track the level of the sketch.
-	// Every time the size of the hashset exceeds the maximum size, the mask will be moved to the next level.
+	// Every time the size of the hashset exceeds the maximum size, the mask will have one more bit set.
 	mask uint64
-	// The maximum size of the hashset. If the size exceeds this value, the mask will be moved to the next level.
+	// The maximum size of the hashset. If the size exceeds this value, the mask will have one more bit set.
 	// And the hashset will only keep the hashed values with trailing zeroes greater than or equal to the new mask.
 	maxSize int
 }
 
 // NewFMSketch returns a new FM sketch.
 func NewFMSketch(maxSize int) *FMSketch {
 	return &FMSketch{
-		hashset: make(map[uint64]struct{}, maxSize),
+		hashset: make(map[uint64]bool, maxSize),
 		maxSize: maxSize,
 	}
 }
@@ -78,18 +82,26 @@ func (s *FMSketch) Copy() *FMSketch {
 	if s == nil {
 		return nil
 	}
-	return &FMSketch{
+	copied := &FMSketch{
 		hashset: maps.Clone(s.hashset),
 		mask:    s.mask,
 		maxSize: s.maxSize,
 	}
+	if s.ndvCounts != nil {
+		value := *s.ndvCounts
+		copied.ndvCounts = &value
+	}
+	return copied
 }
 
 // NDV returns the estimated number of distinct values (NDV) in the sketch.
 func (s *FMSketch) NDV() int64 {
 	if s == nil {
 		return 0
 	}
+	if s.ndvCounts != nil {
+		return s.sampledNDV()
+	}
 	// The estimated count of distinct values is 2^r * count, where 'r' is the maximum number of trailing zeroes observed and 'count' is the number of unique hashed values.
 	// The fundamental idea is that the hash function maps the input domain onto a logarithmic scale.
 	// This is achieved by hashing the input value and counting the number of trailing zeroes in the binary representation of the hash value.
@@ -100,13 +112,29 @@ func (s *FMSketch) NDV() int64 {
 
 // insertHashValue inserts a hashed value into the sketch.
 func (s *FMSketch) insertHashValue(hashVal uint64) {
+	s.mergeHashValue(hashVal, false)
+}
+
+// mergeHashValue adds a hash from a source: a row, which has the hash one time, or a merged
+// sketch. repeated is true if the source saw the hash more than once.
+func (s *FMSketch) mergeHashValue(hashVal uint64, repeated bool) {
 	// If the hashed value is already covered by the mask, we can skip it.
 	// This is because the number of trailing zeroes in the hashed value is less than the mask.
 	if (hashVal & s.mask) != 0 {
 		return
 	}
-	// Put the hashed value into the hashset.
-	s.hashset[hashVal] = struct{}{}
+	// Do not skip this lookup in a sketch that is not sampled. Assigning to a hash that is
+	// already present costs two to three times the lookup, and it made merges about 35% slower.
+	alreadyRepeated, ok := s.hashset[hashVal]
+	if ok {
+		// A sampled sketch sees the hash again here, so the hash is now repeated.
+		if s.ndvCounts != nil && !alreadyRepeated {
+			s.hashset[hashVal] = true
+		}
+		return
+	}
+	// A new hash keeps the flag of its source.
+	s.hashset[hashVal] = repeated
 	// We track the unique hashed values level by level to ensure a minimum count of distinct values at each level.
 	// This way, the final estimation is less likely to be skewed by outliers.
 	if len(s.hashset) > s.maxSize {
@@ -118,7 +146,7 @@ func (s *FMSketch) insertHashValue(hashVal uint64) {
 }
 
 func (s *FMSketch) filterHashes() {
-	rejected := func(hash uint64, _ struct{}) bool { return hash&s.mask != 0 }
+	rejected := func(hash uint64, _ bool) bool { return hash&s.mask != 0 }
 	maps.DeleteFunc(s.hashset, rejected)
 }
 
@@ -180,17 +208,28 @@ func hashRow(sc *stmtctx.StatementContext, values []types.Datum) (uint64, error)
 	return hashFunc.Sum64(), nil
 }
 
-// MergeFMSketch merges two FM Sketch.
+// MergeFMSketch merges two FM Sketch. Sampled sketches come from the responses
+// of one ANALYZE request, so they are all sampled at the same rate. ANALYZE
+// allocates the sketch that collects them before the first response says
+// whether TiKV sampled, so an empty s becomes sampled with the first of them.
 func (s *FMSketch) MergeFMSketch(rs *FMSketch) {
 	if s == nil || rs == nil {
 		return
 	}
+	if rs.ndvCounts != nil {
+		if s.ndvCounts == nil {
+			s.ndvCounts = &ndvCounts{}
+		}
+		s.ndvCounts.rows += rs.ndvCounts.rows
+		s.ndvCounts.samples += rs.ndvCounts.samples
+		s.ndvCounts.nulls += rs.ndvCounts.nulls
+	}
 	if s.mask < rs.mask {
 		s.mask = rs.mask
 		s.filterHashes()
 	}
-	for key := range rs.hashset {
-		s.insertHashValue(key)
+	for key, repeated := range rs.hashset {
+		s.mergeHashValue(key, repeated)
 	}
 }
 
@@ -199,8 +238,11 @@ func FMSketchToProto(s *FMSketch) *tipb.FMSketch {
 	protoSketch := new(tipb.FMSketch)
 	if s != nil {
 		protoSketch.Mask = s.mask
-		for val := range s.hashset {
-			protoSketch.Hashset = append(protoSketch.Hashset, val)
+		for val, repeated := range s.hashset {
+			// The hashset field of a sampled sketch has only the hashes seen once.
+			if !repeated {
+				protoSketch.Hashset = append(protoSketch.Hashset, val)
+			}
 		}
 	}
 	return protoSketch
@@ -212,11 +254,11 @@ func FMSketchFromProto(protoSketch *tipb.FMSketch) *FMSketch {
 		return nil
 	}
 	sketch := &FMSketch{
-		hashset: make(map[uint64]struct{}, len(protoSketch.Hashset)),
+		hashset: make(map[uint64]bool, len(protoSketch.Hashset)),
 		mask:    protoSketch.Mask,
 	}
 	for _, val := range protoSketch.Hashset {
-		sketch.hashset[val] = struct{}{}
+		sketch.hashset[val] = false
 	}
 	return sketch
 }
@@ -248,8 +290,12 @@ func DecodeFMSketch(data []byte) (*FMSketch, error) {
 
 // MemoryUsage returns the total memory usage of a FMSketch.
 func (s *FMSketch) MemoryUsage() (sum int64) {
-	// As for the variables mask(uint64) and maxSize(int) each will consume 8 bytes. This is the origin of the constant 16.
-	// And for the variables hashset(map[uint64]struct{}), we estimate 8 bytes per entry (key size only, excluding Go map overhead).
-	sum = int64(16 + 8*len(s.hashset))
+	// 80 is the 32-byte struct plus the 48-byte Go map header. An entry takes about 30 bytes: its slot takes 18
+	// (8-byte key, value padded to 8, a control byte, allocator rounding), and tables are 7/16 to 7/8 full.
+	sum = int64(80 + 30*len(s.hashset))
+	// A sampled sketch also keeps three 8-byte row counts.
+	if s.ndvCounts != nil {
+		sum += 24
+	}
 	return
 }

pkg/statistics/fmsketch_test.go

@@ -19,8 +19,12 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/pkg/types"
+	"github.com/pingcap/tidb/pkg/util/memory"
+	"github.com/pingcap/tidb/pkg/util/mock"
+	"github.com/pingcap/tipb/go-tipb"
 	"github.com/stretchr/testify/require"
 )
 
@@ -128,3 +132,67 @@ func SubTestFMSketchCoding() func(*testing.T) {
 		require.Equal(t, fmsketch.NDV(), pkSketch.NDV())
 	}
 }
+
+func SubTestSampledNDV() func(*testing.T) {
+	return func(t *testing.T) {
+		left := newSampledFMSketch(&tipb.FMSketch{Hashset: []uint64{1, 3}, MultiHashset: []uint64{2}},
+			ndvCounts{rows: 50, samples: 5, nulls: 10})
+		right := newSampledFMSketch(&tipb.FMSketch{Hashset: []uint64{1, 4}},
+			ndvCounts{rows: 30, samples: 3, nulls: 10})
+		for _, inputs := range [][2]*FMSketch{{left, right}, {right, left}} {
+			merged := inputs[0].Copy()
+			merged.MergeFMSketch(inputs[1])
+			require.Equal(t, map[uint64]bool{1: true, 2: true, 3: false, 4: false}, merged.hashset)
+			require.Equal(t, int64(8), merged.NDV())
+			require.Equal(t, int64(20), merged.ndvCounts.nulls)
+			require.Equal(t, int64(80+30*4+24), merged.MemoryUsage())
+		}
+
+		// ANALYZE merges the responses into sketches it allocated beforehand.
+		// A worker that got no response returns such a sketch unchanged.
+		root := NewFMSketch(MaxSketchSize)
+		for _, sketch := range []*FMSketch{left, NewFMSketch(MaxSketchSize), right} {
+			root.MergeFMSketch(sketch)
+		}
+		require.Equal(t, int64(80), root.ndvCounts.rows)
+		require.Equal(t, int64(8), root.NDV())
+
+		// Only a response that carries a sample count was sampled, even if empty.
+		zero := int64(0)
+		for _, sampled := range []bool{false, true} {
+			pb := &tipb.RowSampleCollector{FmSketch: []*tipb.FMSketch{{}}, NullCounts: []int64{0}, TotalSize: []int64{0}}
+			if sampled {
+				pb.NdvSampleCount = &zero
+			}
+			collector := NewBernoulliRowSampleCollector(1, 1)
+			collector.Base().FromProto(pb, memory.NewTracker(-1, -1))
+			require.Equal(t, sampled, collector.Base().FMSketches[0].ndvCounts != nil)
+		}
+		empty := newSampledFMSketch(&tipb.FMSketch{}, ndvCounts{rows: 10})
+		require.Equal(t, int64(0), empty.NDV())
+		empty.MergeFMSketch(left)
+		require.Equal(t, int64(60), empty.ndvCounts.rows)
+		require.Equal(t, int64(5), empty.ndvCounts.samples)
+		require.Equal(t, int64(8), empty.NDV())
+
+		// Without sampled rows or non-NULL rows, NDV is zero.
+		for _, counts := range []ndvCounts{{rows: 200}, {rows: 10, samples: 2, nulls: 10}} {
+			require.Zero(t, newSampledFMSketch(&tipb.FMSketch{}, counts).NDV())
+		}
+
+		// Values the schema keeps unique need no estimate.
+		singles := make([]uint64, 0, 14)
+		for i := range 14 {
+			singles = append(singles, uint64(i+1))
+		}
+		distinct := make([]*SampleItem, 0, 300)
+		for i := range 300 {
+			distinct = append(distinct, &SampleItem{Value: types.NewIntDatum(int64(i)), Ordinal: i})
+		}
+		sketch := newSampledFMSketch(&tipb.FMSketch{Hashset: singles}, ndvCounts{rows: 3000, samples: 15})
+		collector := &SampleCollector{Samples: distinct, FMSketch: sketch, Count: 3000, Unique: true}
+		hist, _, err := BuildHistAndTopN(mock.NewContext(), 256, 0, 1, collector, types.NewFieldType(mysql.TypeLonglong), true, nil)
+		require.NoError(t, err)
+		require.Equal(t, int64(3000), hist.NDV)
+	}
+}

pkg/statistics/main_test.go

@@ -71,6 +71,7 @@ func TestStatistics(t *testing.T) {
 	t.Run("SubTestSketch", SubTestSketch())
 	t.Run("SubTestSketchProtoConversion", SubTestSketchProtoConversion())
 	t.Run("SubTestFMSketchCoding", SubTestFMSketchCoding())
+	t.Run("SubTestSampledNDV", SubTestSampledNDV())
 
 	// statistics_test.go
 	t.Run("SubTestColumnRange", SubTestColumnRange())

pkg/statistics/row_sampler.go

@@ -296,7 +296,14 @@ func (s *baseCollector) FromProto(pbCollector *tipb.RowSampleCollector, memTrack
 	s.Count = pbCollector.Count
 	s.NullCount = pbCollector.NullCounts
 	s.FMSketches = make([]*FMSketch, 0, len(pbCollector.FmSketch))
-	for _, pbSketch := range pbCollector.FmSketch {
+	for i, pbSketch := range pbCollector.FmSketch {
+		// TiKV sends the sample count only when it sampled rows for NDV.
+		if pbCollector.NdvSampleCount != nil {
+			s.FMSketches = append(s.FMSketches, newSampledFMSketch(pbSketch, ndvCounts{
+				rows: s.Count, samples: *pbCollector.NdvSampleCount, nulls: s.NullCount[i],
+			}))
+			continue
+		}
 		s.FMSketches = append(s.FMSketches, FMSketchFromProto(pbSketch))
 	}
 	s.TotalSizes = pbCollector.TotalSize