postgresql/contrib/tsm_system_rows/tsm_system_rows.c
2015-05-15 15:31:14 -04:00

271 lines
6.6 KiB
C

/*-------------------------------------------------------------------------
*
* tsm_system_rows.c
* interface routines for system_rows tablesample method
*
*
* Portions Copyright (c) 1996-2014, PostgreSQL Global Development Group
*
* IDENTIFICATION
* contrib/tsm_system_rows_rowlimit/tsm_system_rows.c
*
*-------------------------------------------------------------------------
*/
#include "postgres.h"
#include "fmgr.h"
#include "access/tablesample.h"
#include "access/relscan.h"
#include "miscadmin.h"
#include "nodes/execnodes.h"
#include "nodes/relation.h"
#include "optimizer/clauses.h"
#include "storage/bufmgr.h"
#include "utils/sampling.h"
PG_MODULE_MAGIC;
/*
* State
*/
typedef struct
{
SamplerRandomState randstate;
uint32 seed; /* random seed */
BlockNumber nblocks; /* number of block in relation */
int32 ntuples; /* number of tuples to return */
int32 donetuples; /* tuples already returned */
OffsetNumber lt; /* last tuple returned from current block */
BlockNumber step; /* step size */
BlockNumber lb; /* last block visited */
BlockNumber doneblocks; /* number of already returned blocks */
} SystemSamplerData;
PG_FUNCTION_INFO_V1(tsm_system_rows_init);
PG_FUNCTION_INFO_V1(tsm_system_rows_nextblock);
PG_FUNCTION_INFO_V1(tsm_system_rows_nexttuple);
PG_FUNCTION_INFO_V1(tsm_system_rows_examinetuple);
PG_FUNCTION_INFO_V1(tsm_system_rows_end);
PG_FUNCTION_INFO_V1(tsm_system_rows_reset);
PG_FUNCTION_INFO_V1(tsm_system_rows_cost);
static uint32 random_relative_prime(uint32 n, SamplerRandomState randstate);
/*
* Initializes the state.
*/
Datum
tsm_system_rows_init(PG_FUNCTION_ARGS)
{
TableSampleDesc *tsdesc = (TableSampleDesc *) PG_GETARG_POINTER(0);
uint32 seed = PG_GETARG_UINT32(1);
int32 ntuples = PG_ARGISNULL(2) ? -1 : PG_GETARG_INT32(2);
HeapScanDesc scan = tsdesc->heapScan;
SystemSamplerData *sampler;
if (ntuples < 1)
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("invalid sample size"),
errhint("Sample size must be positive integer value.")));
sampler = palloc0(sizeof(SystemSamplerData));
/* Remember initial values for reinit */
sampler->seed = seed;
sampler->nblocks = scan->rs_nblocks;
sampler->ntuples = ntuples;
sampler->donetuples = 0;
sampler->lt = InvalidOffsetNumber;
sampler->doneblocks = 0;
sampler_random_init_state(sampler->seed, sampler->randstate);
/* Find relative prime as step size for linear probing. */
sampler->step = random_relative_prime(sampler->nblocks, sampler->randstate);
/*
* Randomize start position so that blocks close to step size don't have
* higher probability of being chosen on very short scan.
*/
sampler->lb = sampler_random_fract(sampler->randstate) *
(sampler->nblocks / sampler->step);
tsdesc->tsmdata = (void *) sampler;
PG_RETURN_VOID();
}
/*
* Get next block number or InvalidBlockNumber when we're done.
*
* Uses linear probing algorithm for picking next block.
*/
Datum
tsm_system_rows_nextblock(PG_FUNCTION_ARGS)
{
TableSampleDesc *tsdesc = (TableSampleDesc *) PG_GETARG_POINTER(0);
SystemSamplerData *sampler = (SystemSamplerData *) tsdesc->tsmdata;
sampler->lb = (sampler->lb + sampler->step) % sampler->nblocks;
sampler->doneblocks++;
/* All blocks have been read, we're done */
if (sampler->doneblocks > sampler->nblocks ||
sampler->donetuples >= sampler->ntuples)
PG_RETURN_UINT32(InvalidBlockNumber);
PG_RETURN_UINT32(sampler->lb);
}
/*
* Get next tuple offset in current block or InvalidOffsetNumber if we are done
* with this block.
*/
Datum
tsm_system_rows_nexttuple(PG_FUNCTION_ARGS)
{
TableSampleDesc *tsdesc = (TableSampleDesc *) PG_GETARG_POINTER(0);
OffsetNumber maxoffset = PG_GETARG_UINT16(2);
SystemSamplerData *sampler = (SystemSamplerData *) tsdesc->tsmdata;
OffsetNumber tupoffset = sampler->lt;
if (tupoffset == InvalidOffsetNumber)
tupoffset = FirstOffsetNumber;
else
tupoffset++;
if (tupoffset > maxoffset ||
sampler->donetuples >= sampler->ntuples)
tupoffset = InvalidOffsetNumber;
sampler->lt = tupoffset;
PG_RETURN_UINT16(tupoffset);
}
/*
* Examine tuple and decide if it should be returned.
*/
Datum
tsm_system_rows_examinetuple(PG_FUNCTION_ARGS)
{
TableSampleDesc *tsdesc = (TableSampleDesc *) PG_GETARG_POINTER(0);
bool visible = PG_GETARG_BOOL(3);
SystemSamplerData *sampler = (SystemSamplerData *) tsdesc->tsmdata;
if (!visible)
PG_RETURN_BOOL(false);
sampler->donetuples++;
PG_RETURN_BOOL(true);
}
/*
* Cleanup method.
*/
Datum
tsm_system_rows_end(PG_FUNCTION_ARGS)
{
TableSampleDesc *tsdesc = (TableSampleDesc *) PG_GETARG_POINTER(0);
pfree(tsdesc->tsmdata);
PG_RETURN_VOID();
}
/*
* Reset state (called by ReScan).
*/
Datum
tsm_system_rows_reset(PG_FUNCTION_ARGS)
{
TableSampleDesc *tsdesc = (TableSampleDesc *) PG_GETARG_POINTER(0);
SystemSamplerData *sampler = (SystemSamplerData *) tsdesc->tsmdata;
sampler->lt = InvalidOffsetNumber;
sampler->donetuples = 0;
sampler->doneblocks = 0;
sampler_random_init_state(sampler->seed, sampler->randstate);
sampler->step = random_relative_prime(sampler->nblocks, sampler->randstate);
sampler->lb = sampler_random_fract(sampler->randstate) * (sampler->nblocks / sampler->step);
PG_RETURN_VOID();
}
/*
* Costing function.
*/
Datum
tsm_system_rows_cost(PG_FUNCTION_ARGS)
{
PlannerInfo *root = (PlannerInfo *) PG_GETARG_POINTER(0);
Path *path = (Path *) PG_GETARG_POINTER(1);
RelOptInfo *baserel = (RelOptInfo *) PG_GETARG_POINTER(2);
List *args = (List *) PG_GETARG_POINTER(3);
BlockNumber *pages = (BlockNumber *) PG_GETARG_POINTER(4);
double *tuples = (double *) PG_GETARG_POINTER(5);
Node *limitnode;
int32 ntuples;
limitnode = linitial(args);
limitnode = estimate_expression_value(root, limitnode);
if (IsA(limitnode, RelabelType))
limitnode = (Node *) ((RelabelType *) limitnode)->arg;
if (IsA(limitnode, Const))
ntuples = DatumGetInt32(((Const *) limitnode)->constvalue);
else
{
/* Default ntuples if the estimation didn't return Const. */
ntuples = 1000;
}
*pages = Min(baserel->pages, ntuples);
*tuples = ntuples;
path->rows = *tuples;
PG_RETURN_VOID();
}
static uint32
gcd (uint32 a, uint32 b)
{
uint32 c;
while (a != 0)
{
c = a;
a = b % a;
b = c;
}
return b;
}
static uint32
random_relative_prime(uint32 n, SamplerRandomState randstate)
{
/* Pick random starting number, with some limits on what it can be. */
uint32 r = (uint32) sampler_random_fract(randstate) * n/2 + n/4,
t;
/*
* This should only take 2 or 3 iterations as the probability of 2 numbers
* being relatively prime is ~61%.
*/
while ((t = gcd(r, n)) > 1)
{
CHECK_FOR_INTERRUPTS();
r /= t;
}
return r;
}