executor: make tablesample work under different partition prune modes

pkg/executor/builder.go

@@ -5144,7 +5144,7 @@ func (b *executorBuilder) buildTableSample(v *plannercore.PhysicalTableSample) *
 		e.sampler = &emptySampler{}
 	} else if v.TableSampleInfo.AstNode.SampleMethod == ast.SampleMethodTypeTiDBRegion {
 		e.sampler = newTableRegionSampler(
-			b.ctx, v.TableInfo, startTS, v.TableSampleInfo.Partitions, v.Schema(),
+			b.ctx, v.TableInfo, startTS, v.PhysicalTableID, v.Schema(),
 			v.TableSampleInfo.FullSchema, e.RetFieldTypes(), v.Desc)
 	}
 

pkg/executor/sample.go

@@ -75,10 +75,11 @@ type rowSampler interface {
 }
 
 type tableRegionSampler struct {
-	ctx        sessionctx.Context
-	table      table.Table
-	startTS    uint64
-	partTables []table.PartitionedTable
+	ctx             sessionctx.Context
+	table           table.Table
+	startTS         uint64
+	physicalTableID int64
+
 	schema     *expression.Schema
 	fullSchema *expression.Schema
 	isDesc     bool
@@ -89,18 +90,18 @@ type tableRegionSampler struct {
 	isFinished   bool
 }
 
-func newTableRegionSampler(ctx sessionctx.Context, t table.Table, startTs uint64, partTables []table.PartitionedTable,
+func newTableRegionSampler(ctx sessionctx.Context, t table.Table, startTs uint64, pyhsicalTableID int64,
 	schema *expression.Schema, fullSchema *expression.Schema, retTypes []*types.FieldType, desc bool) *tableRegionSampler {
 	return &tableRegionSampler{
-		ctx:        ctx,
-		table:      t,
-		startTS:    startTs,
-		partTables: partTables,
-		schema:     schema,
-		fullSchema: fullSchema,
-		isDesc:     desc,
-		retTypes:   retTypes,
-		rowMap:     make(map[int64]types.Datum),
+		ctx:             ctx,
+		table:           t,
+		startTS:         startTs,
+		physicalTableID: pyhsicalTableID,
+		schema:          schema,
+		fullSchema:      fullSchema,
+		isDesc:          desc,
+		retTypes:        retTypes,
+		rowMap:          make(map[int64]types.Datum),
 	}
 }
 
@@ -176,23 +177,29 @@ func (s *tableRegionSampler) writeChunkFromRanges(ranges []kv.KeyRange, req *chu
 }
 
 func (s *tableRegionSampler) splitTableRanges() ([]kv.KeyRange, error) {
-	if len(s.partTables) != 0 {
-		var ranges []kv.KeyRange
-		for _, t := range s.partTables {
-			for _, pid := range t.GetAllPartitionIDs() {
-				start := tablecodec.GenTableRecordPrefix(pid)
-				end := start.PrefixNext()
-				rs, err := splitIntoMultiRanges(s.ctx.GetStore(), start, end)
-				if err != nil {
-					return nil, err
-				}
-				ranges = append(ranges, rs...)
-			}
+	partitionTable := s.table.GetPartitionedTable()
+	if partitionTable == nil {
+		startKey, endKey := s.table.RecordPrefix(), s.table.RecordPrefix().PrefixNext()
+		return splitIntoMultiRanges(s.ctx.GetStore(), startKey, endKey)
+	}
+
+	var partIDs []int64
+	if partitionTable.Meta().ID == s.physicalTableID {
+		partIDs = partitionTable.GetAllPartitionIDs()
+	} else {
+		partIDs = []int64{s.physicalTableID}
+	}
+	ranges := make([]kv.KeyRange, 0, len(partIDs))
+	for _, pid := range partIDs {
+		start := tablecodec.GenTableRecordPrefix(pid)
+		end := start.PrefixNext()
+		rs, err := splitIntoMultiRanges(s.ctx.GetStore(), start, end)
+		if err != nil {
+			return nil, err
 		}
-		return ranges, nil
+		ranges = append(ranges, rs...)
 	}
-	startKey, endKey := s.table.RecordPrefix(), s.table.RecordPrefix().PrefixNext()
-	return splitIntoMultiRanges(s.ctx.GetStore(), startKey, endKey)
+	return ranges, nil
 }
 
 func splitIntoMultiRanges(store kv.Storage, startKey, endKey kv.Key) ([]kv.KeyRange, error) {

pkg/planner/core/find_best_task.go

@@ -2541,6 +2541,7 @@ func (ds *DataSource) convertToSampleTable(prop *property.PhysicalProperty,
 	p := PhysicalTableSample{
 		TableSampleInfo: ds.SampleInfo,
 		TableInfo:       ds.table,
+		PhysicalTableID: ds.physicalTableID,
 		Desc:            candidate.isMatchProp && prop.SortItems[0].Desc,
 	}.Init(ds.SCtx(), ds.QueryBlockOffset())
 	p.schema = ds.schema

pkg/planner/core/logical_plan_builder.go

@@ -5017,7 +5017,7 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as
 	ds.SetSchema(schema)
 	ds.names = names
 	ds.setPreferredStoreType(b.TableHints())
-	ds.SampleInfo = NewTableSampleInfo(tn.TableSample, schema.Clone(), b.partitionedTable)
+	ds.SampleInfo = NewTableSampleInfo(tn.TableSample, schema)
 	b.isSampling = ds.SampleInfo != nil
 
 	for i, colExpr := range ds.Schema().Columns {

pkg/planner/core/physical_plans.go

@@ -2520,14 +2520,14 @@ type PhysicalTableSample struct {
 	physicalSchemaProducer
 	TableSampleInfo *TableSampleInfo
 	TableInfo       table.Table
+	PhysicalTableID int64
 	Desc            bool
 }
 
 // TableSampleInfo contains the information for PhysicalTableSample.
 type TableSampleInfo struct {
 	AstNode    *ast.TableSample
 	FullSchema *expression.Schema
-	Partitions []table.PartitionedTable
 }
 
 // MemoryUsage return the memory usage of TableSampleInfo
@@ -2536,7 +2536,7 @@ func (t *TableSampleInfo) MemoryUsage() (sum int64) {
 		return
 	}
 
-	sum = size.SizeOfPointer*2 + size.SizeOfSlice + int64(cap(t.Partitions))*size.SizeOfInterface
+	sum = size.SizeOfPointer*2 + size.SizeOfSlice
 	if t.AstNode != nil {
 		sum += int64(unsafe.Sizeof(ast.TableSample{}))
 	}
@@ -2547,14 +2547,13 @@ func (t *TableSampleInfo) MemoryUsage() (sum int64) {
 }
 
 // NewTableSampleInfo creates a new TableSampleInfo.
-func NewTableSampleInfo(node *ast.TableSample, fullSchema *expression.Schema, pt []table.PartitionedTable) *TableSampleInfo {
+func NewTableSampleInfo(node *ast.TableSample, fullSchema *expression.Schema) *TableSampleInfo {
 	if node == nil {
 		return nil
 	}
 	return &TableSampleInfo{
 		AstNode:    node,
-		FullSchema: fullSchema,
-		Partitions: pt,
+		FullSchema: fullSchema.Clone(),
 	}
 }
 

tests/integrationtest/r/executor/sample.result

@@ -83,7 +83,7 @@ delete from t;
 insert into t values (1, '1');
 select count(*) from t partition (p0) tablesample regions();
 count(*)
-0
+1
 select count(*) from t partition (p1) tablesample regions();
 count(*)
 1
@@ -207,3 +207,17 @@ pk	v
 500	a
 9223372036854775809	b
 set @@global.tidb_scatter_region=default;
+drop table if exists t;
+create table t (a int, b varchar(255), primary key (a)) partition by hash(a) partitions 2;
+insert into t values (1, '1'), (2, '2'), (3, '3');
+set @@tidb_partition_prune_mode='static';
+select * from t tablesample regions() order by a;
+a	b
+1	1
+2	2
+set @@tidb_partition_prune_mode='dynamic';
+select * from t tablesample regions() order by a;
+a	b
+1	1
+2	2
+set @@tidb_partition_prune_mode=default;

tests/integrationtest/t/executor/sample.test

@@ -131,3 +131,13 @@ SPLIT TABLE a BY (500);
 SELECT * FROM a TABLESAMPLE REGIONS() ORDER BY pk;
 
 set @@global.tidb_scatter_region=default;
+
+# TestTableSamplePartitionPruneMode
+drop table if exists t;
+create table t (a int, b varchar(255), primary key (a)) partition by hash(a) partitions 2;
+insert into t values (1, '1'), (2, '2'), (3, '3');
+set @@tidb_partition_prune_mode='static';
+select * from t tablesample regions() order by a;
+set @@tidb_partition_prune_mode='dynamic';
+select * from t tablesample regions() order by a;
+set @@tidb_partition_prune_mode=default;