Line data Source code
1 : /*-------------------------------------------------------------------------
2 : *
3 : * tsm_system_rows.c
4 : * support routines for SYSTEM_ROWS tablesample method
5 : *
6 : * The desire here is to produce a random sample with a given number of rows
7 : * (or the whole relation, if that is fewer rows). We use a block-sampling
8 : * approach. To ensure that the whole relation will be visited if necessary,
9 : * we start at a randomly chosen block and then advance with a stride that
10 : * is randomly chosen but is relatively prime to the relation's nblocks.
11 : *
12 : * Because of the dependence on nblocks, this method cannot be repeatable
13 : * across queries. (Even if the user hasn't explicitly changed the relation,
14 : * maintenance activities such as autovacuum might change nblocks.) However,
15 : * we can at least make it repeatable across scans, by determining the
16 : * sampling pattern only once on the first scan. This means that rescans
17 : * won't visit blocks added after the first scan, but that is fine since
18 : * such blocks shouldn't contain any visible tuples anyway.
19 : *
20 : * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
21 : * Portions Copyright (c) 1994, Regents of the University of California
22 : *
23 : * IDENTIFICATION
24 : * contrib/tsm_system_rows/tsm_system_rows.c
25 : *
26 : *-------------------------------------------------------------------------
27 : */
28 :
29 : #include "postgres.h"
30 :
31 : #include "access/relscan.h"
32 : #include "access/tsmapi.h"
33 : #include "catalog/pg_type.h"
34 : #include "miscadmin.h"
35 : #include "optimizer/optimizer.h"
36 : #include "utils/sampling.h"
37 :
38 2 : PG_MODULE_MAGIC;
39 :
40 4 : PG_FUNCTION_INFO_V1(tsm_system_rows_handler);
41 :
42 :
43 : /* Private state */
44 : typedef struct
45 : {
46 : uint32 seed; /* random seed */
47 : int64 ntuples; /* number of tuples to return */
48 : OffsetNumber lt; /* last tuple returned from current block */
49 : BlockNumber doneblocks; /* number of already-scanned blocks */
50 : BlockNumber lb; /* last block visited */
51 : /* these three values are not changed during a rescan: */
52 : BlockNumber nblocks; /* number of blocks in relation */
53 : BlockNumber firstblock; /* first block to sample from */
54 : BlockNumber step; /* step size, or 0 if not set yet */
55 : } SystemRowsSamplerData;
56 :
57 : static void system_rows_samplescangetsamplesize(PlannerInfo *root,
58 : RelOptInfo *baserel,
59 : List *paramexprs,
60 : BlockNumber *pages,
61 : double *tuples);
62 : static void system_rows_initsamplescan(SampleScanState *node,
63 : int eflags);
64 : static void system_rows_beginsamplescan(SampleScanState *node,
65 : Datum *params,
66 : int nparams,
67 : uint32 seed);
68 : static BlockNumber system_rows_nextsampleblock(SampleScanState *node, BlockNumber nblocks);
69 : static OffsetNumber system_rows_nextsampletuple(SampleScanState *node,
70 : BlockNumber blockno,
71 : OffsetNumber maxoffset);
72 : static uint32 random_relative_prime(uint32 n, pg_prng_state *randstate);
73 :
74 :
75 : /*
76 : * Create a TsmRoutine descriptor for the SYSTEM_ROWS method.
77 : */
78 : Datum
79 80 : tsm_system_rows_handler(PG_FUNCTION_ARGS)
80 : {
81 80 : TsmRoutine *tsm = makeNode(TsmRoutine);
82 :
83 80 : tsm->parameterTypes = list_make1_oid(INT8OID);
84 :
85 : /* See notes at head of file */
86 80 : tsm->repeatable_across_queries = false;
87 80 : tsm->repeatable_across_scans = true;
88 :
89 80 : tsm->SampleScanGetSampleSize = system_rows_samplescangetsamplesize;
90 80 : tsm->InitSampleScan = system_rows_initsamplescan;
91 80 : tsm->BeginSampleScan = system_rows_beginsamplescan;
92 80 : tsm->NextSampleBlock = system_rows_nextsampleblock;
93 80 : tsm->NextSampleTuple = system_rows_nextsampletuple;
94 80 : tsm->EndSampleScan = NULL;
95 :
96 80 : PG_RETURN_POINTER(tsm);
97 : }
98 :
99 : /*
100 : * Sample size estimation.
101 : */
102 : static void
103 18 : system_rows_samplescangetsamplesize(PlannerInfo *root,
104 : RelOptInfo *baserel,
105 : List *paramexprs,
106 : BlockNumber *pages,
107 : double *tuples)
108 : {
109 : Node *limitnode;
110 : int64 ntuples;
111 : double npages;
112 :
113 : /* Try to extract an estimate for the limit rowcount */
114 18 : limitnode = (Node *) linitial(paramexprs);
115 18 : limitnode = estimate_expression_value(root, limitnode);
116 :
117 18 : if (IsA(limitnode, Const) &&
118 14 : !((Const *) limitnode)->constisnull)
119 : {
120 14 : ntuples = DatumGetInt64(((Const *) limitnode)->constvalue);
121 14 : if (ntuples < 0)
122 : {
123 : /* Default ntuples if the value is bogus */
124 4 : ntuples = 1000;
125 : }
126 : }
127 : else
128 : {
129 : /* Default ntuples if we didn't obtain a non-null Const */
130 4 : ntuples = 1000;
131 : }
132 :
133 : /* Clamp to the estimated relation size */
134 18 : if (ntuples > baserel->tuples)
135 10 : ntuples = (int64) baserel->tuples;
136 18 : ntuples = clamp_row_est(ntuples);
137 :
138 18 : if (baserel->tuples > 0 && baserel->pages > 0)
139 18 : {
140 : /* Estimate number of pages visited based on tuple density */
141 18 : double density = baserel->tuples / (double) baserel->pages;
142 :
143 18 : npages = ntuples / density;
144 : }
145 : else
146 : {
147 : /* For lack of data, assume one tuple per page */
148 0 : npages = ntuples;
149 : }
150 :
151 : /* Clamp to sane value */
152 18 : npages = clamp_row_est(Min((double) baserel->pages, npages));
153 :
154 18 : *pages = npages;
155 18 : *tuples = ntuples;
156 18 : }
157 :
158 : /*
159 : * Initialize during executor setup.
160 : */
161 : static void
162 18 : system_rows_initsamplescan(SampleScanState *node, int eflags)
163 : {
164 18 : node->tsm_state = palloc0(sizeof(SystemRowsSamplerData));
165 : /* Note the above leaves tsm_state->step equal to zero */
166 18 : }
167 :
168 : /*
169 : * Examine parameters and prepare for a sample scan.
170 : */
171 : static void
172 18 : system_rows_beginsamplescan(SampleScanState *node,
173 : Datum *params,
174 : int nparams,
175 : uint32 seed)
176 : {
177 18 : SystemRowsSamplerData *sampler = (SystemRowsSamplerData *) node->tsm_state;
178 18 : int64 ntuples = DatumGetInt64(params[0]);
179 :
180 18 : if (ntuples < 0)
181 2 : ereport(ERROR,
182 : (errcode(ERRCODE_INVALID_TABLESAMPLE_ARGUMENT),
183 : errmsg("sample size must not be negative")));
184 :
185 16 : sampler->seed = seed;
186 16 : sampler->ntuples = ntuples;
187 16 : sampler->lt = InvalidOffsetNumber;
188 16 : sampler->doneblocks = 0;
189 : /* lb will be initialized during first NextSampleBlock call */
190 : /* we intentionally do not change nblocks/firstblock/step here */
191 :
192 : /*
193 : * We *must* use pagemode visibility checking in this module, so force
194 : * that even though it's currently default.
195 : */
196 16 : node->use_pagemode = true;
197 16 : }
198 :
199 : /*
200 : * Select next block to sample.
201 : *
202 : * Uses linear probing algorithm for picking next block.
203 : */
204 : static BlockNumber
205 66 : system_rows_nextsampleblock(SampleScanState *node, BlockNumber nblocks)
206 : {
207 66 : SystemRowsSamplerData *sampler = (SystemRowsSamplerData *) node->tsm_state;
208 :
209 : /* First call within scan? */
210 66 : if (sampler->doneblocks == 0)
211 : {
212 : /* First scan within query? */
213 16 : if (sampler->step == 0)
214 : {
215 : /* Initialize now that we have scan descriptor */
216 : pg_prng_state randstate;
217 :
218 : /* If relation is empty, there's nothing to scan */
219 12 : if (nblocks == 0)
220 0 : return InvalidBlockNumber;
221 :
222 : /* We only need an RNG during this setup step */
223 12 : sampler_random_init_state(sampler->seed, &randstate);
224 :
225 : /* Compute nblocks/firstblock/step only once per query */
226 12 : sampler->nblocks = nblocks;
227 :
228 : /* Choose random starting block within the relation */
229 : /* (Actually this is the predecessor of the first block visited) */
230 12 : sampler->firstblock = sampler_random_fract(&randstate) *
231 12 : sampler->nblocks;
232 :
233 : /* Find relative prime as step size for linear probing */
234 12 : sampler->step = random_relative_prime(sampler->nblocks, &randstate);
235 : }
236 :
237 : /* Reinitialize lb */
238 16 : sampler->lb = sampler->firstblock;
239 : }
240 :
241 : /* If we've read all blocks or returned all needed tuples, we're done */
242 66 : if (++sampler->doneblocks > sampler->nblocks ||
243 62 : node->donetuples >= sampler->ntuples)
244 16 : return InvalidBlockNumber;
245 :
246 : /*
247 : * It's probably impossible for scan->rs_nblocks to decrease between scans
248 : * within a query; but just in case, loop until we select a block number
249 : * less than scan->rs_nblocks. We don't care if scan->rs_nblocks has
250 : * increased since the first scan.
251 : */
252 : do
253 : {
254 : /* Advance lb, using uint64 arithmetic to forestall overflow */
255 50 : sampler->lb = ((uint64) sampler->lb + sampler->step) % sampler->nblocks;
256 50 : } while (sampler->lb >= nblocks);
257 :
258 50 : return sampler->lb;
259 : }
260 :
261 : /*
262 : * Select next sampled tuple in current block.
263 : *
264 : * In block sampling, we just want to sample all the tuples in each selected
265 : * block.
266 : *
267 : * When we reach end of the block, return InvalidOffsetNumber which tells
268 : * SampleScan to go to next block.
269 : */
270 : static OffsetNumber
271 256 : system_rows_nextsampletuple(SampleScanState *node,
272 : BlockNumber blockno,
273 : OffsetNumber maxoffset)
274 : {
275 256 : SystemRowsSamplerData *sampler = (SystemRowsSamplerData *) node->tsm_state;
276 256 : OffsetNumber tupoffset = sampler->lt;
277 :
278 : /* Quit if we've returned all needed tuples */
279 256 : if (node->donetuples >= sampler->ntuples)
280 8 : return InvalidOffsetNumber;
281 :
282 : /* Advance to next possible offset on page */
283 248 : if (tupoffset == InvalidOffsetNumber)
284 50 : tupoffset = FirstOffsetNumber;
285 : else
286 198 : tupoffset++;
287 :
288 : /* Done? */
289 248 : if (tupoffset > maxoffset)
290 42 : tupoffset = InvalidOffsetNumber;
291 :
292 248 : sampler->lt = tupoffset;
293 :
294 248 : return tupoffset;
295 : }
296 :
297 : /*
298 : * Compute greatest common divisor of two uint32's.
299 : */
300 : static uint32
301 12 : gcd(uint32 a, uint32 b)
302 : {
303 : uint32 c;
304 :
305 38 : while (a != 0)
306 : {
307 26 : c = a;
308 26 : a = b % a;
309 26 : b = c;
310 : }
311 :
312 12 : return b;
313 : }
314 :
315 : /*
316 : * Pick a random value less than and relatively prime to n, if possible
317 : * (else return 1).
318 : */
319 : static uint32
320 12 : random_relative_prime(uint32 n, pg_prng_state *randstate)
321 : {
322 : uint32 r;
323 :
324 : /* Safety check to avoid infinite loop or zero result for small n. */
325 12 : if (n <= 1)
326 0 : return 1;
327 :
328 : /*
329 : * This should only take 2 or 3 iterations as the probability of 2 numbers
330 : * being relatively prime is ~61%; but just in case, we'll include a
331 : * CHECK_FOR_INTERRUPTS in the loop.
332 : */
333 : do
334 : {
335 14 : CHECK_FOR_INTERRUPTS();
336 14 : r = (uint32) (sampler_random_fract(randstate) * n);
337 14 : } while (r == 0 || gcd(r, n) > 1);
338 :
339 12 : return r;
340 : }
|
