Writing a Table Sampling Method
QHB's implementation of the TABLESAMPLE clause supports custom table sampling methods, in addition to the BERNOULLI and SYSTEM methods that are required by the SQL standard. The sampling method determines which rows of the table will be selected when the TABLESAMPLE clause is used.
At the SQL level, a table sampling method is represented by a single SQL function, typically implemented in C/RUST, having the signature
method_name(internal) RETURNS tsm_handler
The name of the function is the same method name appearing in the TABLESAMPLE clause. The internal argument is a dummy (always having value zero) that simply serves to prevent this function from being called directly from an SQL command. The result of the function must be a palloc'd struct of type TsmRoutine, which contains pointers to support functions for the sampling method. These support functions are plain C functions and are not visible or callable at the SQL level. The support functions are described in Section Sampling Method Support Functions.
In addition to function pointers, the TsmRoutine struct must provide these additional fields:
-
List *parameterTypes
This is an OID list containing the data type OIDs of the parameter(s) that will be accepted by the TABLESAMPLE clause when this sampling method is used. For example, for the built-in methods, this list contains a single item with value FLOAT4OID, which represents the sampling percentage. Custom sampling methods can have more or different parameters. -
bool repeatable_across_queries
If true, the sampling method can deliver identical samples across successive queries, if the same parameters and REPEATABLE seed value are supplied each time and the table contents have not changed. When this is false, the REPEATABLE clause is not accepted for use with the sampling method. -
bool repeatable_across_scans
If true, the sampling method can deliver identical samples across successive scans in the same query (assuming unchanging parameters, seed value, and snapshot). When this is false, the planner will not select plans that would require scanning the sampled table more than once, since that might result in inconsistent query output.
The TsmRoutine struct is declared as follows:
typedef struct TsmRoutine
{
NodeTag type;
/* List of datatype OIDs for the arguments of the TABLESAMPLE clause */
List *parameterTypes;
/* Can method produce repeatable samples across, or even within, queries? */
bool repeatable_across_queries;
bool repeatable_across_scans;
/* Functions for planning a SampleScan on a physical table */
SampleScanGetSampleSize_function SampleScanGetSampleSize;
/* Functions for executing a SampleScan on a physical table */
InitSampleScan_function InitSampleScan; /* can be NULL */
BeginSampleScan_function BeginSampleScan;
NextSampleBlock_function NextSampleBlock; /* can be NULL */
NextSampleTuple_function NextSampleTuple;
EndSampleScan_function EndSampleScan; /* can be NULL */
} TsmRoutine;
Note
More function pointers are likely to be added in the future. Therefore it's recommended that the handler initialize the struct withmakeNode(TsmRoutine)so that all fields are set to NULL. This will ensure that no fields are accidentally left undefined.
The table sampling methods included in the standard distribution are good references when trying to write your own.
Sampling Method Support Functions
The TSM (Table Sampling Method) handler function returns a palloc'd TsmRoutine struct containing pointers to the support functions described below. Most of the functions are required, but some are optional, and those pointers can be NULL.
void
SampleScanGetSampleSize (PlannerInfo *root,
RelOptInfo *baserel,
List *paramexprs,
BlockNumber *pages,
double *tuples);
This function is called during planning. It must estimate the number of relation pages that will be read during a sample scan, and the number of tuples that will be selected by the scan. (For example, these might be determined by estimating the sampling fraction, and then multiplying the baserel->pages and baserel->tuples numbers by that, being sure to round the results to integral values.) The paramexprs list holds the expression(s) that are parameters to the TABLESAMPLE clause. It is recommended to use estimate_expression_value() to try to reduce these expressions to constants, if their values are needed for estimation purposes; but the function must provide size estimates even if they cannot be reduced, and it should not fail even if the values appear invalid (remember that they're only estimates of what the run-time values will be). The pages and tuples parameters are outputs.
void
InitSampleScan (SampleScanState *node,
int eflags);
Initialize for execution of a SampleScan plan node. This is called during executor startup. It should perform any initialization needed before processing can start. The SampleScanState node has already been created, but its tsm_state field is NULL. The InitSampleScan function can palloc whatever internal state data is needed by the sampling method, and store a pointer to it in node->tsm_state. Information about the table to scan is accessible through other fields of the SampleScanState node (but note that the node->ss.ss_currentScanDesc scan descriptor is not set up yet). eflags contains flag bits describing the executor's operating mode for this plan node.
When (eflags & EXEC_FLAG_EXPLAIN_ONLY) is true, the scan will not actually
be performed, so this function should only do the minimum required to make the
node state valid for EXPLAIN and EndSampleScan.
This function can be omitted (set the pointer to NULL), in which case BeginSampleScan must perform all initialization needed by the sampling method.
void
BeginSampleScan (SampleScanState *node,
Datum *params,
int nparams,
uint32 seed);
Begin execution of a sampling scan. This is called just before the first attempt to fetch a tuple, and may be called again if the scan needs to be restarted. Information about the table to scan is accessible through fields of the SampleScanState node (but note that the node->ss.ss_currentScanDesc scan descriptor is not set up yet). The params array, of length nparams, contains the values of the parameters supplied in the TABLESAMPLE clause. These will have the number and types specified in the sampling method's parameterTypes list, and have been checked to not be null. seed contains a seed to use for any random numbers generated within the sampling method; it is either a hash derived from the REPEATABLE value if one was given, or the result of random() if not.
This function may adjust the fields node->use_bulkread and node->use_pagemode. If node->use_bulkread is true, which it is by default, the scan will use a buffer access strategy that encourages recycling buffers after use. It might be reasonable to set this to false if the scan will visit only a small fraction of the table's pages. If node->use_pagemode is true, which it is by default, the scan will perform visibility checking in a single pass for all tuples on each visited page. It might be reasonable to set this to false if the scan will select only a small fraction of the tuples on each visited page. That will result in fewer tuple visibility checks being performed, though each one will be more expensive because it will require more locking.
If the sampling method is marked repeatable_across_scans, it must be able to select the same set of tuples during a rescan as it did originally, that is a fresh call of BeginSampleScan must lead to selecting the same tuples as before (if the TABLESAMPLE parameters and seed don't change).
BlockNumber
NextSampleBlock (SampleScanState *node, BlockNumber nblocks);
Returns the block number of the next page to be scanned, or InvalidBlockNumber if no pages remain to be scanned.
This function can be omitted (set the pointer to NULL), in which case the core code will perform a sequential scan of the entire relation. Such a scan can use synchronized scanning, so that the sampling method cannot assume that the relation pages are visited in the same order on each scan.
OffsetNumber
NextSampleTuple (SampleScanState *node,
BlockNumber blockno,
OffsetNumber maxoffset);
Returns the offset number of the next tuple to be sampled on the specified page, or InvalidOffsetNumber if no tuples remain to be sampled. maxoffset is the largest offset number in use on the page.
Note
NextSampleTuple is not explicitly told which of the offset numbers in the range 1 .. maxoffset actually contain valid tuples. This is not normally a problem since the core code ignores requests to sample missing or invisible tuples; that should not result in any bias in the sample. However, if necessary, the function can use node->donetuples to examine how many of the tuples it returned were valid and visible.
Note
NextSampleTuple must not assume that blockno is the same page number returned by the most recent NextSampleBlock call. It was returned by some previous NextSampleBlock call, but the core code is allowed to call NextSampleBlock in advance of actually scanning pages, so as to support prefetching. It is OK to assume that once sampling of a given page begins, successive NextSampleTuple calls all refer to the same page until InvalidOffsetNumber is returned.
void
EndSampleScan (SampleScanState *node);
End the scan and release resources. It is normally not important to release palloc'd memory, but any externally-visible resources should be cleaned up. This function can be omitted (set the pointer to NULL) in the common case where no such resources exist.
Built-in Sampling Methods
BERNOULLI tablesample method
To ensure repeatability of samples, it is necessary that selection of a given tuple be history-independent; otherwise syncscanning would break repeatability, to say nothing of logically-irrelevant maintenance such as physical extension or shortening of the relation.
To achieve that, we proceed by hashing each candidate TID together with the active seed, and then selecting it if the hash is less than the cutoff value computed from the selection probability by BeginSampleScan.
/* Private state */
typedef struct
{
uint64 cutoff; /* select tuples with hash less than this */
uint32 seed; /* random seed */
OffsetNumber lt; /* last tuple returned from current block */
} BernoulliSamplerData;
static void bernoulli_samplescangetsamplesize(PlannerInfo *root,
RelOptInfo *baserel,
List *paramexprs,
BlockNumber *pages,
double *tuples);
static void bernoulli_initsamplescan(SampleScanState *node,
int eflags);
static void bernoulli_beginsamplescan(SampleScanState *node,
Datum *params,
int nparams,
uint32 seed);
static OffsetNumber bernoulli_nextsampletuple(SampleScanState *node,
BlockNumber blockno,
OffsetNumber maxoffset);
/*
* Create a TsmRoutine descriptor for the BERNOULLI method.
*/
Datum
tsm_bernoulli_handler(PG_FUNCTION_ARGS)
{
TsmRoutine *tsm = makeNode(TsmRoutine);
tsm->parameterTypes = list_make1_oid(FLOAT4OID);
tsm->repeatable_across_queries = true;
tsm->repeatable_across_scans = true;
tsm->SampleScanGetSampleSize = bernoulli_samplescangetsamplesize;
tsm->InitSampleScan = bernoulli_initsamplescan;
tsm->BeginSampleScan = bernoulli_beginsamplescan;
tsm->NextSampleBlock = NULL;
tsm->NextSampleTuple = bernoulli_nextsampletuple;
tsm->EndSampleScan = NULL;
PG_RETURN_POINTER(tsm);
}
/*
* Sample size estimation.
*/
static void
bernoulli_samplescangetsamplesize(PlannerInfo *root,
RelOptInfo *baserel,
List *paramexprs,
BlockNumber *pages,
double *tuples)
{
Node *pctnode;
float4 samplefract;
/* Try to extract an estimate for the sample percentage */
pctnode = (Node *) linitial(paramexprs);
pctnode = estimate_expression_value(root, pctnode);
if (IsA(pctnode, Const) &&
!((Const *) pctnode)->constisnull)
{
samplefract = DatumGetFloat4(((Const *) pctnode)->constvalue);
if (samplefract >= 0 && samplefract <= 100 && !isnan(samplefract))
samplefract /= 100.0f;
else
{
/* Default samplefract if the value is bogus */
samplefract = 0.1f;
}
}
else
{
/* Default samplefract if we didn't obtain a non-null Const */
samplefract = 0.1f;
}
/* We'll visit all pages of the baserel */
*pages = baserel->pages;
*tuples = clamp_row_est(baserel->tuples * samplefract);
}
/*
* Initialize during executor setup.
*/
static void
bernoulli_initsamplescan(SampleScanState *node, int eflags)
{
node->tsm_state = palloc0(sizeof(BernoulliSamplerData));
}
/*
* Examine parameters and prepare for a sample scan.
*/
static void
bernoulli_beginsamplescan(SampleScanState *node,
Datum *params,
int nparams,
uint32 seed)
{
BernoulliSamplerData *sampler = (BernoulliSamplerData *) node->tsm_state;
double percent = DatumGetFloat4(params[0]);
double dcutoff;
if (percent < 0 || percent > 100 || isnan(percent))
ereport(ERROR,
(errcode(ERRCODE_INVALID_TABLESAMPLE_ARGUMENT),
errmsg("sample percentage must be between 0 and 100")));
/*
* The cutoff is sample probability times (PG_UINT32_MAX + 1); we have to
* store that as a uint64, of course. Note that this gives strictly
* correct behavior at the limits of zero or one probability.
*/
dcutoff = rint(((double) PG_UINT32_MAX + 1) * percent / 100);
sampler->cutoff = (uint64) dcutoff;
sampler->seed = seed;
sampler->lt = InvalidOffsetNumber;
/*
* Use bulkread, since we're scanning all pages. But pagemode visibility
* checking is a win only at larger sampling fractions. The 25% cutoff
* here is based on very limited experimentation.
*/
node->use_bulkread = true;
node->use_pagemode = (percent >= 25);
}
/*
* Select next sampled tuple in current block.
*
* It is OK here to return an offset without knowing if the tuple is visible
* (or even exists). The reason is that we do the coinflip for every tuple
* offset in the table. Since all tuples have the same probability of being
* returned, it doesn't matter if we do extra coinflips for invisible tuples.
*
* When we reach end of the block, return InvalidOffsetNumber which tells
* SampleScan to go to next block.
*/
static OffsetNumber
bernoulli_nextsampletuple(SampleScanState *node,
BlockNumber blockno,
OffsetNumber maxoffset)
{
BernoulliSamplerData *sampler = (BernoulliSamplerData *) node->tsm_state;
OffsetNumber tupoffset = sampler->lt;
uint32 hashinput[3];
/* Advance to first/next tuple in block */
if (tupoffset == InvalidOffsetNumber)
tupoffset = FirstOffsetNumber;
else
tupoffset++;
/*
* We compute the hash by applying hash_any to an array of 3 uint32's
* containing the block, offset, and seed. This is efficient to set up,
* and with the current implementation of hash_any, it gives
* machine-independent results, which is a nice property for regression
* testing.
*
* These words in the hash input are the same throughout the block:
*/
hashinput[0] = blockno;
hashinput[2] = sampler->seed;
/*
* Loop over tuple offsets until finding suitable TID or reaching end of
* block.
*/
for (; tupoffset <= maxoffset; tupoffset++)
{
uint32 hash;
hashinput[1] = tupoffset;
hash = DatumGetUInt32(hash_any((const unsigned char *) hashinput,
(int) sizeof(hashinput)));
if (hash < sampler->cutoff)
break;
}
if (tupoffset > maxoffset)
tupoffset = InvalidOffsetNumber;
sampler->lt = tupoffset;
return tupoffset;
}
SYSTEM tablesample method
To ensure repeatability of samples, it is necessary that selection of a given tuple be history-independent; otherwise syncscanning would break repeatability, to say nothing of logically-irrelevant maintenance such as physical extension or shortening of the relation.
To achieve that, we proceed by hashing each candidate block number together with the active seed, and then selecting it if the hash is less than the cutoff value computed from the selection probability by BeginSampleScan.
/* Private state */
typedef struct
{
uint64 cutoff; /* select blocks with hash less than this */
uint32 seed; /* random seed */
BlockNumber nextblock; /* next block to consider sampling */
OffsetNumber lt; /* last tuple returned from current block */
} SystemSamplerData;
static void system_samplescangetsamplesize(PlannerInfo *root,
RelOptInfo *baserel,
List *paramexprs,
BlockNumber *pages,
double *tuples);
static void system_initsamplescan(SampleScanState *node,
int eflags);
static void system_beginsamplescan(SampleScanState *node,
Datum *params,
int nparams,
uint32 seed);
static BlockNumber system_nextsampleblock(SampleScanState *node, BlockNumber nblocks);
static OffsetNumber system_nextsampletuple(SampleScanState *node,
BlockNumber blockno,
OffsetNumber maxoffset);
/*
* Create a TsmRoutine descriptor for the SYSTEM method.
*/
Datum
tsm_system_handler(PG_FUNCTION_ARGS)
{
TsmRoutine *tsm = makeNode(TsmRoutine);
tsm->parameterTypes = list_make1_oid(FLOAT4OID);
tsm->repeatable_across_queries = true;
tsm->repeatable_across_scans = true;
tsm->SampleScanGetSampleSize = system_samplescangetsamplesize;
tsm->InitSampleScan = system_initsamplescan;
tsm->BeginSampleScan = system_beginsamplescan;
tsm->NextSampleBlock = system_nextsampleblock;
tsm->NextSampleTuple = system_nextsampletuple;
tsm->EndSampleScan = NULL;
PG_RETURN_POINTER(tsm);
}
/*
* Sample size estimation.
*/
static void
system_samplescangetsamplesize(PlannerInfo *root,
RelOptInfo *baserel,
List *paramexprs,
BlockNumber *pages,
double *tuples)
{
Node *pctnode;
float4 samplefract;
/* Try to extract an estimate for the sample percentage */
pctnode = (Node *) linitial(paramexprs);
pctnode = estimate_expression_value(root, pctnode);
if (IsA(pctnode, Const) &&
!((Const *) pctnode)->constisnull)
{
samplefract = DatumGetFloat4(((Const *) pctnode)->constvalue);
if (samplefract >= 0 && samplefract <= 100 && !isnan(samplefract))
samplefract /= 100.0f;
else
{
/* Default samplefract if the value is bogus */
samplefract = 0.1f;
}
}
else
{
/* Default samplefract if we didn't obtain a non-null Const */
samplefract = 0.1f;
}
/* We'll visit a sample of the pages ... */
*pages = clamp_row_est(baserel->pages * samplefract);
/* ... and hopefully get a representative number of tuples from them */
*tuples = clamp_row_est(baserel->tuples * samplefract);
}
/*
* Initialize during executor setup.
*/
static void
system_initsamplescan(SampleScanState *node, int eflags)
{
node->tsm_state = palloc0(sizeof(SystemSamplerData));
}
/*
* Examine parameters and prepare for a sample scan.
*/
static void
system_beginsamplescan(SampleScanState *node,
Datum *params,
int nparams,
uint32 seed)
{
SystemSamplerData *sampler = (SystemSamplerData *) node->tsm_state;
double percent = DatumGetFloat4(params[0]);
double dcutoff;
if (percent < 0 || percent > 100 || isnan(percent))
ereport(ERROR,
(errcode(ERRCODE_INVALID_TABLESAMPLE_ARGUMENT),
errmsg("sample percentage must be between 0 and 100")));
/*
* The cutoff is sample probability times (PG_UINT32_MAX + 1); we have to
* store that as a uint64, of course. Note that this gives strictly
* correct behavior at the limits of zero or one probability.
*/
dcutoff = rint(((double) PG_UINT32_MAX + 1) * percent / 100);
sampler->cutoff = (uint64) dcutoff;
sampler->seed = seed;
sampler->nextblock = 0;
sampler->lt = InvalidOffsetNumber;
/*
* Bulkread buffer access strategy probably makes sense unless we're
* scanning a very small fraction of the table. The 1% cutoff here is a
* guess. We should use pagemode visibility checking, since we scan all
* tuples on each selected page.
*/
node->use_bulkread = (percent >= 1);
node->use_pagemode = true;
}
/*
* Select next block to sample.
*/
static BlockNumber
system_nextsampleblock(SampleScanState *node, BlockNumber nblocks)
{
SystemSamplerData *sampler = (SystemSamplerData *) node->tsm_state;
BlockNumber nextblock = sampler->nextblock;
uint32 hashinput[2];
/*
* We compute the hash by applying hash_any to an array of 2 uint32's
* containing the block number and seed. This is efficient to set up, and
* with the current implementation of hash_any, it gives
* machine-independent results, which is a nice property for regression
* testing.
*
* These words in the hash input are the same throughout the block:
*/
hashinput[1] = sampler->seed;
/*
* Loop over block numbers until finding suitable block or reaching end of
* relation.
*/
for (; nextblock < nblocks; nextblock++)
{
uint32 hash;
hashinput[0] = nextblock;
hash = DatumGetUInt32(hash_any((const unsigned char *) hashinput,
(int) sizeof(hashinput)));
if (hash < sampler->cutoff)
break;
}
if (nextblock < nblocks)
{
/* Found a suitable block; remember where we should start next time */
sampler->nextblock = nextblock + 1;
return nextblock;
}
/* Done, but let's reset nextblock to 0 for safety. */
sampler->nextblock = 0;
return InvalidBlockNumber;
}
/*
* Select next sampled tuple in current block.
*
* In block sampling, we just want to sample all the tuples in each selected
* block.
*
* It is OK here to return an offset without knowing if the tuple is visible
* (or even exists); nodeSamplescan.c will deal with that.
*
* When we reach end of the block, return InvalidOffsetNumber which tells
* SampleScan to go to next block.
*/
static OffsetNumber
system_nextsampletuple(SampleScanState *node,
BlockNumber blockno,
OffsetNumber maxoffset)
{
SystemSamplerData *sampler = (SystemSamplerData *) node->tsm_state;
OffsetNumber tupoffset = sampler->lt;
/* Advance to next possible offset on page */
if (tupoffset == InvalidOffsetNumber)
tupoffset = FirstOffsetNumber;
else
tupoffset++;
/* Done? */
if (tupoffset > maxoffset)
tupoffset = InvalidOffsetNumber;
sampler->lt = tupoffset;
return tupoffset;
}