summaryrefslogtreecommitdiffstats
path: root/contrib/tsm_system_rows
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--contrib/tsm_system_rows/.gitignore4
-rw-r--r--contrib/tsm_system_rows/Makefile23
-rw-r--r--contrib/tsm_system_rows/expected/tsm_system_rows.out83
-rw-r--r--contrib/tsm_system_rows/sql/tsm_system_rows.sql39
-rw-r--r--contrib/tsm_system_rows/tsm_system_rows--1.0.sql9
-rw-r--r--contrib/tsm_system_rows/tsm_system_rows.c340
-rw-r--r--contrib/tsm_system_rows/tsm_system_rows.control6
7 files changed, 504 insertions, 0 deletions
diff --git a/contrib/tsm_system_rows/.gitignore b/contrib/tsm_system_rows/.gitignore
new file mode 100644
index 0000000..5dcb3ff
--- /dev/null
+++ b/contrib/tsm_system_rows/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/contrib/tsm_system_rows/Makefile b/contrib/tsm_system_rows/Makefile
new file mode 100644
index 0000000..3192ed5
--- /dev/null
+++ b/contrib/tsm_system_rows/Makefile
@@ -0,0 +1,23 @@
+# contrib/tsm_system_rows/Makefile
+
+MODULE_big = tsm_system_rows
+OBJS = \
+ $(WIN32RES) \
+ tsm_system_rows.o
+PGFILEDESC = "tsm_system_rows - TABLESAMPLE method which accepts number of rows as a limit"
+
+EXTENSION = tsm_system_rows
+DATA = tsm_system_rows--1.0.sql
+
+REGRESS = tsm_system_rows
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = contrib/tsm_system_rows
+top_builddir = ../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/contrib/tsm_system_rows/expected/tsm_system_rows.out b/contrib/tsm_system_rows/expected/tsm_system_rows.out
new file mode 100644
index 0000000..87b4a8f
--- /dev/null
+++ b/contrib/tsm_system_rows/expected/tsm_system_rows.out
@@ -0,0 +1,83 @@
+CREATE EXTENSION tsm_system_rows;
+CREATE TABLE test_tablesample (id int, name text);
+INSERT INTO test_tablesample SELECT i, repeat(i::text, 1000)
+ FROM generate_series(0, 30) s(i);
+ANALYZE test_tablesample;
+SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (0);
+ count
+-------
+ 0
+(1 row)
+
+SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (1);
+ count
+-------
+ 1
+(1 row)
+
+SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (10);
+ count
+-------
+ 10
+(1 row)
+
+SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (100);
+ count
+-------
+ 31
+(1 row)
+
+-- bad parameters should get through planning, but not execution:
+EXPLAIN (COSTS OFF)
+SELECT id FROM test_tablesample TABLESAMPLE system_rows (-1);
+ QUERY PLAN
+----------------------------------------
+ Sample Scan on test_tablesample
+ Sampling: system_rows ('-1'::bigint)
+(2 rows)
+
+SELECT id FROM test_tablesample TABLESAMPLE system_rows (-1);
+ERROR: sample size must not be negative
+-- fail, this method is not repeatable:
+SELECT * FROM test_tablesample TABLESAMPLE system_rows (10) REPEATABLE (0);
+ERROR: tablesample method system_rows does not support REPEATABLE
+LINE 1: SELECT * FROM test_tablesample TABLESAMPLE system_rows (10) ...
+ ^
+-- but a join should be allowed:
+EXPLAIN (COSTS OFF)
+SELECT * FROM
+ (VALUES (0),(10),(100)) v(nrows),
+ LATERAL (SELECT count(*) FROM test_tablesample
+ TABLESAMPLE system_rows (nrows)) ss;
+ QUERY PLAN
+----------------------------------------------------------
+ Nested Loop
+ -> Values Scan on "*VALUES*"
+ -> Aggregate
+ -> Sample Scan on test_tablesample
+ Sampling: system_rows ("*VALUES*".column1)
+(5 rows)
+
+SELECT * FROM
+ (VALUES (0),(10),(100)) v(nrows),
+ LATERAL (SELECT count(*) FROM test_tablesample
+ TABLESAMPLE system_rows (nrows)) ss;
+ nrows | count
+-------+-------
+ 0 | 0
+ 10 | 10
+ 100 | 31
+(3 rows)
+
+CREATE VIEW vv AS
+ SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (20);
+SELECT * FROM vv;
+ count
+-------
+ 20
+(1 row)
+
+DROP EXTENSION tsm_system_rows; -- fail, view depends on extension
+ERROR: cannot drop extension tsm_system_rows because other objects depend on it
+DETAIL: view vv depends on function system_rows(internal)
+HINT: Use DROP ... CASCADE to drop the dependent objects too.
diff --git a/contrib/tsm_system_rows/sql/tsm_system_rows.sql b/contrib/tsm_system_rows/sql/tsm_system_rows.sql
new file mode 100644
index 0000000..e3ab420
--- /dev/null
+++ b/contrib/tsm_system_rows/sql/tsm_system_rows.sql
@@ -0,0 +1,39 @@
+CREATE EXTENSION tsm_system_rows;
+
+CREATE TABLE test_tablesample (id int, name text);
+INSERT INTO test_tablesample SELECT i, repeat(i::text, 1000)
+ FROM generate_series(0, 30) s(i);
+ANALYZE test_tablesample;
+
+SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (0);
+SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (1);
+SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (10);
+SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (100);
+
+-- bad parameters should get through planning, but not execution:
+EXPLAIN (COSTS OFF)
+SELECT id FROM test_tablesample TABLESAMPLE system_rows (-1);
+
+SELECT id FROM test_tablesample TABLESAMPLE system_rows (-1);
+
+-- fail, this method is not repeatable:
+SELECT * FROM test_tablesample TABLESAMPLE system_rows (10) REPEATABLE (0);
+
+-- but a join should be allowed:
+EXPLAIN (COSTS OFF)
+SELECT * FROM
+ (VALUES (0),(10),(100)) v(nrows),
+ LATERAL (SELECT count(*) FROM test_tablesample
+ TABLESAMPLE system_rows (nrows)) ss;
+
+SELECT * FROM
+ (VALUES (0),(10),(100)) v(nrows),
+ LATERAL (SELECT count(*) FROM test_tablesample
+ TABLESAMPLE system_rows (nrows)) ss;
+
+CREATE VIEW vv AS
+ SELECT count(*) FROM test_tablesample TABLESAMPLE system_rows (20);
+
+SELECT * FROM vv;
+
+DROP EXTENSION tsm_system_rows; -- fail, view depends on extension
diff --git a/contrib/tsm_system_rows/tsm_system_rows--1.0.sql b/contrib/tsm_system_rows/tsm_system_rows--1.0.sql
new file mode 100644
index 0000000..de508ed
--- /dev/null
+++ b/contrib/tsm_system_rows/tsm_system_rows--1.0.sql
@@ -0,0 +1,9 @@
+/* contrib/tsm_system_rows/tsm_system_rows--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION tsm_system_rows" to load this file. \quit
+
+CREATE FUNCTION system_rows(internal)
+RETURNS tsm_handler
+AS 'MODULE_PATHNAME', 'tsm_system_rows_handler'
+LANGUAGE C STRICT;
diff --git a/contrib/tsm_system_rows/tsm_system_rows.c b/contrib/tsm_system_rows/tsm_system_rows.c
new file mode 100644
index 0000000..217a4f9
--- /dev/null
+++ b/contrib/tsm_system_rows/tsm_system_rows.c
@@ -0,0 +1,340 @@
+/*-------------------------------------------------------------------------
+ *
+ * tsm_system_rows.c
+ * support routines for SYSTEM_ROWS tablesample method
+ *
+ * The desire here is to produce a random sample with a given number of rows
+ * (or the whole relation, if that is fewer rows). We use a block-sampling
+ * approach. To ensure that the whole relation will be visited if necessary,
+ * we start at a randomly chosen block and then advance with a stride that
+ * is randomly chosen but is relatively prime to the relation's nblocks.
+ *
+ * Because of the dependence on nblocks, this method cannot be repeatable
+ * across queries. (Even if the user hasn't explicitly changed the relation,
+ * maintenance activities such as autovacuum might change nblocks.) However,
+ * we can at least make it repeatable across scans, by determining the
+ * sampling pattern only once on the first scan. This means that rescans
+ * won't visit blocks added after the first scan, but that is fine since
+ * such blocks shouldn't contain any visible tuples anyway.
+ *
+ * Portions Copyright (c) 1996-2022, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ * contrib/tsm_system_rows/tsm_system_rows.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "access/relscan.h"
+#include "access/tsmapi.h"
+#include "catalog/pg_type.h"
+#include "miscadmin.h"
+#include "optimizer/optimizer.h"
+#include "utils/sampling.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(tsm_system_rows_handler);
+
+
+/* Private state */
+typedef struct
+{
+ uint32 seed; /* random seed */
+ int64 ntuples; /* number of tuples to return */
+ OffsetNumber lt; /* last tuple returned from current block */
+ BlockNumber doneblocks; /* number of already-scanned blocks */
+ BlockNumber lb; /* last block visited */
+ /* these three values are not changed during a rescan: */
+ BlockNumber nblocks; /* number of blocks in relation */
+ BlockNumber firstblock; /* first block to sample from */
+ BlockNumber step; /* step size, or 0 if not set yet */
+} SystemRowsSamplerData;
+
+static void system_rows_samplescangetsamplesize(PlannerInfo *root,
+ RelOptInfo *baserel,
+ List *paramexprs,
+ BlockNumber *pages,
+ double *tuples);
+static void system_rows_initsamplescan(SampleScanState *node,
+ int eflags);
+static void system_rows_beginsamplescan(SampleScanState *node,
+ Datum *params,
+ int nparams,
+ uint32 seed);
+static BlockNumber system_rows_nextsampleblock(SampleScanState *node, BlockNumber nblocks);
+static OffsetNumber system_rows_nextsampletuple(SampleScanState *node,
+ BlockNumber blockno,
+ OffsetNumber maxoffset);
+static uint32 random_relative_prime(uint32 n, pg_prng_state *randstate);
+
+
+/*
+ * Create a TsmRoutine descriptor for the SYSTEM_ROWS method.
+ */
+Datum
+tsm_system_rows_handler(PG_FUNCTION_ARGS)
+{
+ TsmRoutine *tsm = makeNode(TsmRoutine);
+
+ tsm->parameterTypes = list_make1_oid(INT8OID);
+
+ /* See notes at head of file */
+ tsm->repeatable_across_queries = false;
+ tsm->repeatable_across_scans = true;
+
+ tsm->SampleScanGetSampleSize = system_rows_samplescangetsamplesize;
+ tsm->InitSampleScan = system_rows_initsamplescan;
+ tsm->BeginSampleScan = system_rows_beginsamplescan;
+ tsm->NextSampleBlock = system_rows_nextsampleblock;
+ tsm->NextSampleTuple = system_rows_nextsampletuple;
+ tsm->EndSampleScan = NULL;
+
+ PG_RETURN_POINTER(tsm);
+}
+
+/*
+ * Sample size estimation.
+ */
+static void
+system_rows_samplescangetsamplesize(PlannerInfo *root,
+ RelOptInfo *baserel,
+ List *paramexprs,
+ BlockNumber *pages,
+ double *tuples)
+{
+ Node *limitnode;
+ int64 ntuples;
+ double npages;
+
+ /* Try to extract an estimate for the limit rowcount */
+ limitnode = (Node *) linitial(paramexprs);
+ limitnode = estimate_expression_value(root, limitnode);
+
+ if (IsA(limitnode, Const) &&
+ !((Const *) limitnode)->constisnull)
+ {
+ ntuples = DatumGetInt64(((Const *) limitnode)->constvalue);
+ if (ntuples < 0)
+ {
+ /* Default ntuples if the value is bogus */
+ ntuples = 1000;
+ }
+ }
+ else
+ {
+ /* Default ntuples if we didn't obtain a non-null Const */
+ ntuples = 1000;
+ }
+
+ /* Clamp to the estimated relation size */
+ if (ntuples > baserel->tuples)
+ ntuples = (int64) baserel->tuples;
+ ntuples = clamp_row_est(ntuples);
+
+ if (baserel->tuples > 0 && baserel->pages > 0)
+ {
+ /* Estimate number of pages visited based on tuple density */
+ double density = baserel->tuples / (double) baserel->pages;
+
+ npages = ntuples / density;
+ }
+ else
+ {
+ /* For lack of data, assume one tuple per page */
+ npages = ntuples;
+ }
+
+ /* Clamp to sane value */
+ npages = clamp_row_est(Min((double) baserel->pages, npages));
+
+ *pages = npages;
+ *tuples = ntuples;
+}
+
+/*
+ * Initialize during executor setup.
+ */
+static void
+system_rows_initsamplescan(SampleScanState *node, int eflags)
+{
+ node->tsm_state = palloc0(sizeof(SystemRowsSamplerData));
+ /* Note the above leaves tsm_state->step equal to zero */
+}
+
+/*
+ * Examine parameters and prepare for a sample scan.
+ */
+static void
+system_rows_beginsamplescan(SampleScanState *node,
+ Datum *params,
+ int nparams,
+ uint32 seed)
+{
+ SystemRowsSamplerData *sampler = (SystemRowsSamplerData *) node->tsm_state;
+ int64 ntuples = DatumGetInt64(params[0]);
+
+ if (ntuples < 0)
+ ereport(ERROR,
+ (errcode(ERRCODE_INVALID_TABLESAMPLE_ARGUMENT),
+ errmsg("sample size must not be negative")));
+
+ sampler->seed = seed;
+ sampler->ntuples = ntuples;
+ sampler->lt = InvalidOffsetNumber;
+ sampler->doneblocks = 0;
+ /* lb will be initialized during first NextSampleBlock call */
+ /* we intentionally do not change nblocks/firstblock/step here */
+
+ /*
+ * We *must* use pagemode visibility checking in this module, so force
+ * that even though it's currently default.
+ */
+ node->use_pagemode = true;
+}
+
+/*
+ * Select next block to sample.
+ *
+ * Uses linear probing algorithm for picking next block.
+ */
+static BlockNumber
+system_rows_nextsampleblock(SampleScanState *node, BlockNumber nblocks)
+{
+ SystemRowsSamplerData *sampler = (SystemRowsSamplerData *) node->tsm_state;
+
+ /* First call within scan? */
+ if (sampler->doneblocks == 0)
+ {
+ /* First scan within query? */
+ if (sampler->step == 0)
+ {
+ /* Initialize now that we have scan descriptor */
+ pg_prng_state randstate;
+
+ /* If relation is empty, there's nothing to scan */
+ if (nblocks == 0)
+ return InvalidBlockNumber;
+
+ /* We only need an RNG during this setup step */
+ sampler_random_init_state(sampler->seed, &randstate);
+
+ /* Compute nblocks/firstblock/step only once per query */
+ sampler->nblocks = nblocks;
+
+ /* Choose random starting block within the relation */
+ /* (Actually this is the predecessor of the first block visited) */
+ sampler->firstblock = sampler_random_fract(&randstate) *
+ sampler->nblocks;
+
+ /* Find relative prime as step size for linear probing */
+ sampler->step = random_relative_prime(sampler->nblocks, &randstate);
+ }
+
+ /* Reinitialize lb */
+ sampler->lb = sampler->firstblock;
+ }
+
+ /* If we've read all blocks or returned all needed tuples, we're done */
+ if (++sampler->doneblocks > sampler->nblocks ||
+ node->donetuples >= sampler->ntuples)
+ return InvalidBlockNumber;
+
+ /*
+ * It's probably impossible for scan->rs_nblocks to decrease between scans
+ * within a query; but just in case, loop until we select a block number
+ * less than scan->rs_nblocks. We don't care if scan->rs_nblocks has
+ * increased since the first scan.
+ */
+ do
+ {
+ /* Advance lb, using uint64 arithmetic to forestall overflow */
+ sampler->lb = ((uint64) sampler->lb + sampler->step) % sampler->nblocks;
+ } while (sampler->lb >= nblocks);
+
+ return sampler->lb;
+}
+
+/*
+ * Select next sampled tuple in current block.
+ *
+ * In block sampling, we just want to sample all the tuples in each selected
+ * block.
+ *
+ * When we reach end of the block, return InvalidOffsetNumber which tells
+ * SampleScan to go to next block.
+ */
+static OffsetNumber
+system_rows_nextsampletuple(SampleScanState *node,
+ BlockNumber blockno,
+ OffsetNumber maxoffset)
+{
+ SystemRowsSamplerData *sampler = (SystemRowsSamplerData *) node->tsm_state;
+ OffsetNumber tupoffset = sampler->lt;
+
+ /* Quit if we've returned all needed tuples */
+ if (node->donetuples >= sampler->ntuples)
+ return InvalidOffsetNumber;
+
+ /* 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;
+}
+
+/*
+ * Compute greatest common divisor of two uint32's.
+ */
+static uint32
+gcd(uint32 a, uint32 b)
+{
+ uint32 c;
+
+ while (a != 0)
+ {
+ c = a;
+ a = b % a;
+ b = c;
+ }
+
+ return b;
+}
+
+/*
+ * Pick a random value less than and relatively prime to n, if possible
+ * (else return 1).
+ */
+static uint32
+random_relative_prime(uint32 n, pg_prng_state *randstate)
+{
+ uint32 r;
+
+ /* Safety check to avoid infinite loop or zero result for small n. */
+ if (n <= 1)
+ return 1;
+
+ /*
+ * This should only take 2 or 3 iterations as the probability of 2 numbers
+ * being relatively prime is ~61%; but just in case, we'll include a
+ * CHECK_FOR_INTERRUPTS in the loop.
+ */
+ do
+ {
+ CHECK_FOR_INTERRUPTS();
+ r = (uint32) (sampler_random_fract(randstate) * n);
+ } while (r == 0 || gcd(r, n) > 1);
+
+ return r;
+}
diff --git a/contrib/tsm_system_rows/tsm_system_rows.control b/contrib/tsm_system_rows/tsm_system_rows.control
new file mode 100644
index 0000000..b495fb1
--- /dev/null
+++ b/contrib/tsm_system_rows/tsm_system_rows.control
@@ -0,0 +1,6 @@
+# tsm_system_rows extension
+comment = 'TABLESAMPLE method which accepts number of rows as a limit'
+default_version = '1.0'
+module_pathname = '$libdir/tsm_system_rows'
+relocatable = true
+trusted = true