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