diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b68cd7b..5c04673 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -60,6 +60,15 @@ jobs: - name: Run tests run: cargo pgrx test ${{ matrix.pg }} + - name: Run multi-session tests + run: | + PG_CONFIG="$(cargo pgrx info pg-config ${{ matrix.pg }})" + MAJOR="${{ matrix.pg }}"; MAJOR="${MAJOR#pg}" + cargo pgrx install --no-default-features --features ${{ matrix.pg }} --pg-config "$PG_CONFIG" + cargo pgrx start ${{ matrix.pg }} + PATH="$(dirname "$PG_CONFIG"):$PATH" PGHOST=localhost PGPORT=$((28800 + MAJOR)) \ + tests/multisession.sh + postgis: name: postgis (pg16) runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..73c5843 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,58 @@ +# Changelog + +## 0.1.2 + +A correctness release. 0.1.1 could return wrong results (rows missing from query +results) and could crash the server. + +Summaries built by 0.1.1 are not trusted: 0.1.2 never prunes with them, so each +table_range index keeps all partitions (results correct, just unpruned) until it is +rebuilt. To upgrade: + +1. Install 0.1.2 and **restart PostgreSQL**, so no backend still runs 0.1.1 code (a + 0.1.1 backend cannot read 0.1.2 summaries and would skip widening them). +2. `ALTER EXTENSION pg_table_range UPDATE;` (warns how many indexes need rebuilding). +3. `REINDEX INDEX` each table_range index (on a partitioned table, the parent index). + +### Crashes fixed + +- Any error raised while planning a query (e.g. `SELECT 1/0`) in a backend that had + loaded the extension aborted the backend, restarting the whole cluster. +- `REINDEX` of a table_range index crashed the server the same way. + +### Wrong results fixed + +- Prepared statements and PL/pgSQL functions kept using plans that pruned a partition + after an insert had widened its summary. +- A second widening inside one transaction was invisible to that transaction's later + queries; after a rolled-back savepoint, a later widening was never seen by other + backends. +- Concurrent inserts could overwrite each other's widening. +- A widening that failed (e.g. a value too large for the metapage) was silently skipped. +- Partial table_range indexes ignored inserts outside their predicate. They are now + rejected. +- `CREATE INDEX` inside a `REPEATABLE READ` transaction missed rows committed after the + transaction's snapshot, and builds omitted rows still visible to older snapshots. +- Summary bounds depended on `DateStyle` and `extra_float_digits` at write and read time. +- Cross-type comparisons (e.g. `float8_col > 1.5::float4`) re-parsed the constant as the + column's type. +- Comparisons under a different collation than the summary's (e.g. `COLLATE "C"` on an + `en_US` column) were pruned with the wrong ordering. + +### Other fixes + +- `CREATE INDEX` / `REINDEX` set the table's `reltuples` to 0, so the planner estimated + one row per partition until the next `ANALYZE`. +- Indexes including a `uuid` or `boolean` column silently built no summary at all (those + types have no `min()`/`max()`). Builds no longer use SQL aggregates. +- On PostgreSQL 17+, geometry summaries could not be built (index builds run with a + restricted `search_path`). All SQL the extension runs is now schema-qualified, which + also stops a caller's `search_path` from substituting functions used to maintain + summaries. +- Updates that change only table_range-indexed columns can be HOT updates again + (`amsummarizing`). +- The README no longer claims `VACUUM` re-tightens summaries; `REINDEX` does. + +### Format + +Summaries are stored in on-page format version 2. Version-1 summaries are ignored. diff --git a/Cargo.toml b/Cargo.toml index 5cee9d2..6a37779 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pg_table_range" -version = "0.1.1" +version = "0.1.2" edition = "2021" license = "MIT" description = "PostgreSQL extension for planning-time partition pruning by per-partition data ranges, including non-key columns." diff --git a/META.json b/META.json index 02c0006..fb6ce6a 100644 --- a/META.json +++ b/META.json @@ -2,7 +2,7 @@ "name": "pg_table_range", "abstract": "Planning-time partition pruning by per-partition data ranges, including non-key columns", "description": "pg_table_range prunes partitions at planning time from a compact per-partition summary of each column's actual data: its min/max range for btree-comparable scalar types, or a covering extent for range types and PostGIS geometry. Unlike native PostgreSQL partition pruning, it works on columns that are not the partition key. Summaries are stored in the index's own metapage and maintained incrementally on insert (like BRIN), so no REINDEX is needed, and pruning is conservative so results are always identical to running without it.", - "version": "0.1.1", + "version": "0.1.2", "maintainer": [ "David Bitner " ], @@ -12,7 +12,7 @@ "abstract": "Planning-time partition pruning by per-partition data ranges, including non-key columns", "file": "pg_table_range.control", "docfile": "README.md", - "version": "0.1.1" + "version": "0.1.2" } }, "prereqs": { diff --git a/README.md b/README.md index 5fe2ef5..bf48ea6 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,8 @@ cargo pgrx init --pg18 /usr/lib/postgresql/18/bin/pg_config cargo pgrx install --release ``` +When upgrading, follow the steps for each version in [CHANGELOG.md](CHANGELOG.md). + The distribution is also on [PGXN](https://pgxn.org/dist/pg_table_range/): `pgxn download pg_table_range`, unzip, then `make install` (or the cargo-pgrx commands above). `pgxn install` runs the bundled `Makefile`, so it works too — but still needs cargo-pgrx present, @@ -93,8 +95,8 @@ CREATE INDEX events_tr ON events USING table_range (val, created_at); -- Verify with EXPLAIN: non-matching partitions disappear from the plan. EXPLAIN (COSTS OFF) SELECT * FROM events WHERE val >= 250; --- Inserts maintain the summary automatically; REINDEX only re-tightens after many --- deletes. DROP INDEX removes the summary with the index. +-- Inserts maintain the summary automatically; REINDEX re-tightens it after deletes. +-- DROP INDEX removes the summary with the index. REINDEX INDEX events_tr; DROP INDEX events_tr; ``` @@ -122,31 +124,42 @@ EXPLAIN (COSTS OFF) SELECT * FROM places WHERE geom && ST_MakeEnvelope(0,0,10,10 that partition's index — one record per indexed column on the index's **metapage**, not in any side table. It holds the `has_nulls` / `all_nulls` flags plus either the column's btree `min`/`max` (scalar columns) or a single covering **extent** — a covering range - for range types (`range_merge(range_agg(col))`) or the bounding box for PostGIS geometry - (`ST_Extent(col)`). + for range types (via `range_merge`) or the bounding box for PostGIS geometry (via + `ST_Extent`). `CREATE INDEX` / `REINDEX` compute it with the table AM's index-build + scan (as BRIN does), which also covers rows that older snapshots can still see. - **Planning.** For each partition the planner builds, a `set_rel_pathlist_hook` reads the summary from that partition's index and evaluates the partition's restriction clauses against it, calling `mark_dummy_rel` on any partition that provably cannot match — - eliminating it before child paths are generated. Deserialized summaries are cached for + eliminating it from the plan. Deserialized summaries are cached for the life of the backend (kept coherent by a relcache-invalidation callback), so warm plans skip the per-partition page read; see [Performance](#performance). -- **Typed comparisons.** Min/max vs. constant comparisons use each column type's own - btree compare function, so **any btree-comparable type works**: `bigint` / `int` / - `smallint`, `numeric`, `real` / `double precision`, `text` / `varchar`, `date`, - `time`, `timestamp`, `timestamptz`, `uuid`, `boolean`, `oid`, etc. Any conversion - problem degrades safely to "keep". +- **Typed comparisons.** Min/max are compared with the query's constant as given, using + the btree operator family's compare function for the operator's (column, constant) + type pair — so cross-type predicates such as `bigint_col = 5` or + `float8_col > 1.5::float4` compare exactly. **Any btree-comparable type works**: + `bigint` / `int` / `smallint`, `numeric`, `real` / `double precision`, `text` / + `varchar`, `date`, `time`, `timestamp`, `timestamptz`, `uuid`, `boolean`, enums, `oid`, + etc. Text-like columns prune only when the query compares under the collation the + summary was built with (the index column's collation). Any conversion problem degrades + safely to "keep". - **Overlap (`&&`).** For range types and PostGIS geometry, an `&&` (overlaps) predicate - is pruned by testing the constant against the partition's stored extent with - PostgreSQL's own `&&` operator — so a partition is eliminated when its extent cannot - overlap the query. + is pruned by calling the clause's own `&&` operator on the partition's stored extent + and the constant — so a partition is eliminated when its extent cannot overlap the + query. Only an `&&` implemented in `pg_catalog` or in the column type's own schema is + trusted. - **Incremental maintenance (no REINDEX).** `aminsert` widens the summary in place as rows are inserted — the same way BRIN maintains its ranges. Because the summary only ever needs to be over-inclusive, these updates need no MVCC: an insert within the - existing range writes nothing; one that extends it grows the min/max/extent. Pruning - therefore stays correct **and** active across inserts without any rebuild. Deletes only - shrink a partition's true range, leaving the summary conservatively wide (still safe); - `VACUUM`/`REINDEX` can re-tighten it for selectivity. `DROP INDEX` removes the summary - with the index's storage — there is no side table to clean up. + existing range writes nothing; one that extends it grows the min/max/extent, as a + compare-and-swap on the metapage so concurrent inserts never overwrite each other's + widening. A widening invalidates cached copies of the summary and cached plans on the + partition in every backend. If a value cannot be represented (for example it is too + large for the metapage), that column becomes *unbounded* — it no longer prunes by + value, but is never too narrow. Pruning therefore stays correct **and** active across + inserts without any rebuild. Deletes only shrink a partition's true range, leaving the + summary conservatively wide (still safe); `REINDEX` re-tightens it once the deleted rows + are no longer visible to any transaction (`VACUUM` does not). `DROP INDEX` removes the + summary with the index's storage — there is no side table to clean up. ## Performance @@ -155,20 +168,21 @@ selective predicate on a non-key column scans only the matching partition instea partition. The planner pays a little to evaluate each partition's summary, but that cost is small and — warm — close to free (see the cache note below). -The numbers below are reproducible with `bench/benchmark.sql` (`cargo pgrx run pg18`, then -`\i bench/benchmark.sql`); they report `EXPLAIN (ANALYZE)` planning and execution time -separately, warm, on PostgreSQL 18. +The numbers below are reproducible with `bench/benchmark.sql` (`cargo pgrx run pg18 +--release`, then `\i bench/benchmark.sql`); they are medians of the warm `EXPLAIN +(ANALYZE)` planning and execution times over three runs of the script, on PostgreSQL 18.4 +with a release build of the extension. Absolute numbers depend on hardware. **Faster execution.** 300 partitions × 8,000 rows (2.4M rows), `WHERE nk = `: | | Planning | Execution | Total | |---|---|---|---| -| pruning **off** (scans all 300 partitions) | ~4 ms | ~110 ms | ~114 ms | -| pruning **on** (scans 1 partition) | ~4 ms | ~0.4 ms | **~4 ms** | +| pruning **off** (scans all 300 partitions) | ~1.5 ms | ~55 ms | ~57 ms | +| pruning **on** (scans 1 partition) | ~1.0 ms | ~0.2 ms | **~1.2 ms** | -Execution is ~250× faster, total time drops ~25×, and warm the planning overhead is in the -noise. The win grows with how much data the eliminated partitions hold; measure your +Execution is ~275× faster, total time drops ~45×, and planning is no slower (pruned +partitions drop out before the parent's paths are built). The win grows with how much data the eliminated partitions hold; measure your workload with `table_range.enable_pruning`. **Honest comparison to native pruning.** When a predicate is on the *partition key*, @@ -179,15 +193,16 @@ natively) and `nk` (the same values, not the key, pruned by table_range): | Same `=` predicate, 2,000 partitions | Planning | Execution | |---|---|---| -| native pruning — column **is** the partition key | **~0.15 ms** | ~0.05 ms | -| table_range — column is **not** the partition key | ~34 ms | ~0.06 ms | -| no pruning — scans all 2,000 partitions | ~28 ms | ~27 ms | +| native pruning — column **is** the partition key | **~0.05 ms** | ~0.02 ms | +| table_range — column is **not** the partition key | ~15 ms | ~0.02 ms | +| no pruning — scans all 2,000 partitions | ~15 ms | ~17 ms | Native pruning is *hundreds of times* cheaper to plan and is effectively constant in the partition count. table_range cannot match that (see [Scaling](#scaling-and-partition-count)): its job is the case native pruning **can't** do -— eliminating partitions by a non-key column. Note that table_range's overhead over the -no-pruning baseline (~28 ms to expand 2,000 partitions) is now small (~6 ms, ~3 µs/part). +— eliminating partitions by a non-key column. table_range's own cost is within measurement +noise of the no-pruning baseline: nearly all of the ~15 ms is PostgreSQL expanding 2,000 +partitions. **Comparison to `CHECK` constraint exclusion.** The built-in way to prune on a non-key column is to put a data-range `CHECK (col BETWEEN lo AND hi)` on each partition and let the @@ -196,12 +211,12 @@ baseline. Same table, 2,000 partitions, same `nk = ` predicate: | Same `=` predicate, 2,000 partitions | Planning | Execution | Scans | |---|---|---|---| -| `CHECK` constraint exclusion (`constraint_exclusion=on`) | ~37 ms | ~0.08 ms | 1 partition | -| table_range pruning | ~34 ms | ~0.08 ms | 1 partition | -| no pruning | ~26 ms | ~25 ms | all 2,000 | +| `CHECK` constraint exclusion (`constraint_exclusion=on`) | ~17.6 ms | ~0.04 ms | 1 partition | +| table_range pruning | ~14.7 ms | ~0.04 ms | 1 partition | +| no pruning | ~14.2 ms | ~15.6 ms | all 2,000 | -Both are O(partitions) and give the **identical execution win**, and **table_range now -plans on par with — and warm, slightly faster than — constraint exclusion.** (Constraint +Both are O(partitions) and give the **identical execution win**, and **table_range plans +faster than constraint exclusion.** (Constraint exclusion re-parses each partition's `CHECK` expression on every plan; table_range serves warm plans from a cached summary, see below.) On top of matching the speed, table_range avoids everything `CHECK` constraints make you give up: @@ -213,18 +228,16 @@ avoids everything `CHECK` constraints make you give up: - **Incremental maintenance** — changing a `CHECK` means `DROP`/`ADD CONSTRAINT` with a full-partition revalidation scan; table_range widens in place in `aminsert`, no rescan. -**How the per-partition cost got small.** Two optimizations took the per-partition planning -cost from ~31 µs to ~3–4 µs: +**Why the per-partition cost is small.** -1. *Per-plan compilation.* The compare function, type-input function, and operator strategy - are identical across a column's partitions, so they are resolved once per plan - (cached `FmgrInfo`s) instead of re-looked-up for each partition. +1. *Per-plan compilation.* An operator's strategy and compare function, and a type's input + function, are the same for every partition, so they are resolved once per plan (cached + `FmgrInfo`s). 2. *Backend summary cache.* Each index's deserialized summary is cached for the life of the - backend, so warm/repeated plans skip the per-partition index open and metapage - read+deserialize entirely. The cache is kept coherent by a relcache invalidation - callback: `aminsert` only ever *widens* a summary, and when it does it invalidates the - cached copy everywhere — so a cached summary is never narrower than reality (a wider one - prunes correctly). A cold first plan still reads each page; every plan after is cached. + backend, so warm plans skip the per-partition index open and metapage read. `aminsert` + only ever *widens* a summary, and when it does it invalidates the cached copy (and + cached plans on that partition) everywhere — so a cached summary is never narrower than + reality. A cold first plan still reads each page. ## Scaling and partition count @@ -269,9 +282,8 @@ Two practical consequences and how to handle them: concurrency. With default `max_connections`, a few thousand (e.g. `4096`) covers tens of thousands of partitions; each lock slot costs only a few hundred bytes of shared memory. - **Planning time grows with partition count.** Even below the lock wall, planning scales - linearly — though the per-partition constant is now small (~3–4 µs warm, on par with - `CHECK` constraint exclusion) thanks to the per-plan compilation and backend summary - cache described above. **Mitigations:** prefer **fewer, larger partitions** (table_range's + linearly, dominated by PostgreSQL's own partition expansion (~15 ms for 2,000 partitions + in the benchmark above); table_range's evaluation adds little on top. **Mitigations:** prefer **fewer, larger partitions** (table_range's sweet spot — the execution win is biggest there anyway); use **prepared statements** so a plan is reused across executions; and where you can, **align the hot filter column with the partition key** so native pruning handles it. @@ -286,7 +298,9 @@ support in PostgreSQL core. Everything not listed is conservatively **kept** (never mispruned): - Comparisons `col < c`, `<=`, `=`, `>=`, `>` (either operand order), and `BETWEEN` - (the planner expands it into two comparisons). + (the planner expands it into two comparisons), including cross-type operators of the + column's btree operator family. Comparisons under a collation other than the index + column's (e.g. `col < 'x' COLLATE "C"` on an `en_US` column) are kept. - `col IS NULL` / `col IS NOT NULL`. - `col IN (c1, c2, …)` / `col = ANY()` — pruned when no listed value falls in the partition's range. @@ -306,6 +320,10 @@ Everything not listed is conservatively **kept** (never mispruned): There is no catalog table — each partition's summary lives on its `table_range` index's metapage (block 0), written by `ambuild` and updated in place by `aminsert`, like BRIN. +Bounds are stored as each type's text form, written and read with the type I/O settings +that affect it (`DateStyle`, `IntervalStyle`, `extra_float_digits`) pinned to +round-trip-exact values. The on-page format is versioned: a summary in any other format +version is never used for pruning until `REINDEX` rebuilds it. ## Project layout @@ -316,14 +334,21 @@ metapage (block 0), written by `ambuild` and updated in place by `aminsert`, lik | `src/summary_build.rs` | build a leaf's summary by scanning its data (used by `ambuild`) | | `src/prune_hook.rs` | planner + pathlist hooks, per-plan compilation cache, typed in-memory evaluation | | `src/summary_cache.rs` | backend-lifetime per-index summary cache + relcache-invalidation coherence | +| `src/stable_io.rs` | pins the GUCs that change a type's text form while summaries are written or read | | `src/index_am.rs` | `table_range` index AM: build, incremental `aminsert` widening, opclass provisioning | | `src/e2e_tests.rs`, `src/index_am_tests.rs` | end-to-end tests | +| `tests/multisession.sh` | tests that need several concurrent backends (run against a live server) | +| `sql/` | extension upgrade scripts | ## Building and testing ```sh cargo pgrx test pg18 # run the end-to-end test suite (PostgreSQL 18) cargo pgrx run pg18 # open psql with the extension installed + +# multi-backend tests, against a running server with the extension installed +cargo pgrx install && cargo pgrx start pg18 +PGHOST=localhost PGPORT=28818 tests/multisession.sh ``` Supported targets: PostgreSQL 16, 17, 18. The test suite is entirely end-to-end — @@ -379,9 +404,19 @@ the published tarball via `.gitattributes`. `table_range.enable_pruning`. - `NOT IN` / `<> ALL`, `NOT (...)`, expression predicates, and parameterized prepared-statement plans are kept rather than pruned. +- Partial indexes (`CREATE INDEX … USING table_range … WHERE …`) are rejected, and + expression columns in a table_range index are not summarized. - Inserts keep summaries current incrementally, but deletes only relax them (the summary - can stay wider than the live data until a `VACUUM`/`REINDEX` re-tightens it) — always - correct, just potentially less selective. + can stay wider than the live data until a `REINDEX` re-tightens it) — always correct, + just potentially less selective. +- Every insert that widens a summary writes a full metapage image to WAL and invalidates + cached plans on that partition. Append-only workloads whose values keep growing (e.g. + timestamps) widen on nearly every insert. +- Pruning is decided when a query is planned. Under `READ COMMITTED`, an insert that + widens a summary and commits after a query is planned but before it starts executing + can be missed by that one execution (cached plans are replanned once the widening + commits). A query planned and run in one `REPEATABLE READ` / `SERIALIZABLE` snapshot is + not affected. - Pruning engages only while the index is **valid** (`indisvalid`); the planner ignores invalid indexes, so anything that invalidates a table_range index silently disables its pruning until rebuilt. diff --git a/bench/benchmark.sql b/bench/benchmark.sql index 4d309e3..c7f886b 100644 --- a/bench/benchmark.sql +++ b/bench/benchmark.sql @@ -1,6 +1,6 @@ -- Reproducible benchmarks for pg_table_range. -- --- cargo pgrx run pg18 +-- cargo pgrx run pg18 --release -- \i bench/benchmark.sql -- -- Reports planning and execution time separately (warm) via EXPLAIN (ANALYZE), for: @@ -8,19 +8,12 @@ -- 2. Honest comparison to native partition pruning: two identical columns, one the -- partition key (pruned natively) and one not (pruned by table_range). -- --- Note: cargo pgrx run reuses its database. If you have previously loaded older versions --- of this extension, stale event triggers can interfere; this script drops the extension --- and recreates it for a clean baseline. +-- The script drops and recreates the extension for a clean baseline. \set ECHO none \pset pager off DROP EXTENSION IF EXISTS pg_table_range CASCADE; --- Clear any event triggers left behind by older builds (no-ops on a clean DB). -DROP EVENT TRIGGER IF EXISTS table_range_hide_indexes_trg; -DROP EVENT TRIGGER IF EXISTS table_range_drop_trg; -DROP EVENT TRIGGER IF EXISTS table_range_opclass_sync_trg; -DROP FUNCTION IF EXISTS table_range_hide_indexes() CASCADE; CREATE EXTENSION pg_table_range; -- ==================================================================================== diff --git a/sql/pg_table_range--0.1.1--0.1.2.sql b/sql/pg_table_range--0.1.1--0.1.2.sql new file mode 100644 index 0000000..74d7f41 --- /dev/null +++ b/sql/pg_table_range--0.1.1--0.1.2.sql @@ -0,0 +1,22 @@ +-- pg_table_range 0.1.1 -> 0.1.2 +-- +-- No SQL objects change. The on-page summary format does (version 2), and summaries +-- written by 0.1.1 can be narrower than the data they describe, so 0.1.2 never prunes +-- with them. Pruning resumes for each table_range index once it is rebuilt. +DO $$ +DECLARE + n bigint; +BEGIN + SELECT count(*) INTO n + FROM pg_catalog.pg_index i + JOIN pg_catalog.pg_class c ON c.oid = i.indexrelid + JOIN pg_catalog.pg_am am ON am.oid = c.relam + WHERE am.amname = 'table_range' AND c.relkind = 'i'; + IF n > 0 THEN + -- Extension scripts run with client_min_messages >= WARNING, so a NOTICE would + -- never reach the user. + RAISE WARNING 'pg_table_range: % table_range index(es) must be rebuilt with REINDEX before they prune again', n + USING HINT = 'Summaries built by 0.1.1 are ignored by 0.1.2. REINDEX each table_range index (on a partitioned table, REINDEX the parent index).'; + END IF; +END +$$; diff --git a/src/e2e_tests.rs b/src/e2e_tests.rs index 999b183..9b7fe81 100644 --- a/src/e2e_tests.rs +++ b/src/e2e_tests.rs @@ -71,6 +71,34 @@ fn e2e_count_where(pred: &str) -> i64 { .expect("count value") } +/// Count rows of `table` matching `pred`. +fn e2e_count_in(table: &str, pred: &str) -> i64 { + Spi::get_one::(&format!("SELECT count(*)::bigint FROM {table} WHERE {pred}")) + .expect("count") + .expect("count value") +} + +/// Count matching rows with pruning on and off, assert they agree, and return the count. +fn e2e_same_on_off(table: &str, pred: &str) -> i64 { + e2e_set_pruning(false); + let off = e2e_count_in(table, pred); + e2e_set_pruning(true); + let on = e2e_count_in(table, pred); + assert_eq!(on, off, "pruning changed results for `{pred}`: on={on} off={off}"); + on +} + +/// A table `name` LIST-partitioned on `k` into `{name}_1` (k=1) and `{name}_2` (k=2). +fn e2e_two_partitions(name: &str, columns: &str) { + Spi::run(&format!( + "DROP TABLE IF EXISTS {name} CASCADE; + CREATE TABLE {name} (k int, {columns}) PARTITION BY LIST (k); + CREATE TABLE {name}_1 PARTITION OF {name} FOR VALUES IN (1); + CREATE TABLE {name}_2 PARTITION OF {name} FOR VALUES IN (2);" + )) + .expect("create partitioned table"); +} + #[pg_test] fn e2e_prunes_lower_partitions_for_high_predicate() { e2e_setup_events(); @@ -168,22 +196,22 @@ fn e2e_works_on_plain_unpartitioned_table() { } #[pg_test] -fn e2e_insert_keeps_results_correct_via_staleness() { +fn e2e_insert_widens_summary() { e2e_setup_events(); e2e_set_pruning(true); // Sanity: before the insert, val=500 prunes everything (no partition covers it). assert_eq!(e2e_count_where("val = 500"), 0); - // Insert a value far outside r1's summarized range. aminsert marks r1 stale, so the - // new row is still found — no false negative despite a now-stale summary. + // Insert a value far outside r1's summarized range. aminsert widens r1's summary to + // cover it, so the new row is found and r1 is kept for it. Spi::run("INSERT INTO events VALUES (1, 500)").expect("insert"); assert_eq!( e2e_count_where("val = 500"), 1, - "stale summary must not prune away newly inserted matching rows" + "widened summary must cover the newly inserted row" ); let plan = e2e_explain("SELECT * FROM events WHERE val = 500"); - assert!(plan.contains("events_r1"), "r1 kept while stale:\n{plan}"); + assert!(plan.contains("events_r1"), "r1 kept after widening:\n{plan}"); } #[pg_test] @@ -563,3 +591,144 @@ fn e2e_is_null_pruning() { assert_eq!(on, off, "null mismatch for `{pred}`"); } } + +#[pg_test(error = "division by zero")] +fn e2e_planning_error_is_reported_as_error() { + // A planning-time ERROR passes through the planner hooks and must not abort the backend. + Spi::run("SELECT 1/0").unwrap(); +} + +#[pg_test] +fn e2e_cached_plan_replanned_after_widening() { + // PL/pgSQL static SQL uses a cached generic plan, like PREPARE. + e2e_two_partitions("cp", "val bigint"); + Spi::run( + "INSERT INTO cp VALUES (1, 1), (2, 100); + CREATE INDEX cp_tr ON cp USING table_range (val); + CREATE FUNCTION cp_count() RETURNS bigint LANGUAGE plpgsql AS + $$ BEGIN RETURN (SELECT count(*) FROM cp WHERE val = 5000); END $$;", + ) + .unwrap(); + e2e_set_pruning(true); + for _ in 0..6 { + assert_eq!(Spi::get_one::("SELECT cp_count()").unwrap(), Some(0)); + } + Spi::run("INSERT INTO cp VALUES (2, 5000)").unwrap(); + assert_eq!( + Spi::get_one::("SELECT cp_count()").unwrap(), + Some(1), + "cached plan must be replanned after the summary widens" + ); +} + +#[pg_test] +fn e2e_each_widening_visible_later_in_transaction() { + + e2e_two_partitions("sx", "val bigint"); + Spi::run( + "INSERT INTO sx VALUES (1, 1), (2, 100); + CREATE INDEX sx_tr ON sx USING table_range (val);", + ) + .unwrap(); + e2e_set_pruning(true); + Spi::run("INSERT INTO sx VALUES (1, 7000)").unwrap(); + assert_eq!(e2e_count_in("sx", "val = 7000"), 1); + Spi::run("INSERT INTO sx VALUES (1, 8000)").unwrap(); + assert_eq!(e2e_count_in("sx", "val = 8000"), 1); +} + +#[pg_test] +fn e2e_cross_type_comparisons() { + // `f > 100.1::float4` means f > 100.09999847..., not f > 100.1. + e2e_two_partitions("xt", "f float8, b bigint"); + Spi::run( + "INSERT INTO xt VALUES (1, 1.0, 1), (1, 100.09999999, 2), (2, 500, 100); + CREATE INDEX xt_tr ON xt USING table_range (f, b);", + ) + .unwrap(); + assert_eq!(e2e_same_on_off("xt", "f > 100.1::float4"), 2); + assert_eq!(e2e_same_on_off("xt", "f >= 100.1::float4 AND f < 101"), 1); + let plan = e2e_explain_on("xt", "b = 100::int4"); + assert!(!plan.contains("xt_1"), "int8 = int4 must prune:\n{plan}"); + let plan = e2e_explain_on("xt", "f > 400::float4"); + assert!(!plan.contains("xt_1"), "float8 > float4 must prune:\n{plan}"); +} + +#[pg_test] +fn e2e_other_collation_is_not_pruned() { + + let Some(coll) = Spi::get_one::( + "SELECT collname::text FROM pg_collation + WHERE collname IN ('und-x-icu', 'en_US.utf8', 'en_US.UTF-8', 'en_US') + AND collencoding IN (-1, pg_char_to_encoding(current_setting('server_encoding'))) + ORDER BY collname = 'und-x-icu' DESC LIMIT 1", + ) + .unwrap() else { + return; // no linguistic collation in this build + }; + e2e_two_partitions("co", &format!("s text COLLATE \"{coll}\"")); + Spi::run( + "INSERT INTO co VALUES (1, 'apple'), (1, 'Zebra'), (2, 'mango'); + CREATE INDEX co_tr ON co USING table_range (s);", + ) + .unwrap(); + // Linguistically 'apple' < 'Zebra'; bytewise 'Zebra' < 'apple'. + assert_eq!(e2e_same_on_off("co", "s > 'Zebra' COLLATE \"C\""), 2); + let plan = e2e_explain_on("co", "s > 'Zebra' COLLATE \"C\""); + assert!(plan.contains("co_1"), "another collation must not prune:\n{plan}"); + // Under the summary's own collation, 'mango' > 'b'. + let plan = e2e_explain_on("co", "s < 'b'"); + assert!(!plan.contains("co_2"), "the summary's collation must prune:\n{plan}"); +} + +#[pg_test] +fn e2e_bounds_independent_of_session_io_settings() { + // Bounds are stored as text, which these session settings would otherwise change. + e2e_two_partitions("ss", "d date, f float8"); + Spi::run( + "INSERT INTO ss VALUES (1, '2024-03-05', 1), (2, '2025-01-01', 500); + CREATE INDEX ss_tr ON ss USING table_range (d, f); + SET DateStyle = 'SQL, DMY'; + INSERT INTO ss VALUES (1, '01/12/2024', NULL); + SET DateStyle = 'ISO, MDY'; + SET extra_float_digits = -10; + INSERT INTO ss VALUES (1, NULL, 123444.4); + RESET extra_float_digits;", + ) + .unwrap(); + assert_eq!(e2e_same_on_off("ss", "d >= '2024-11-01'"), 2); + assert_eq!(e2e_same_on_off("ss", "f > 123441 AND f < 200000"), 1); + Spi::run("RESET DateStyle").unwrap(); +} + +#[pg_test] +fn e2e_types_without_min_max_aggregates_prune() { + // uuid/bool lack min()/max(); varchar/enum compare via text/anyenum procs. + Spi::run( + "DROP TYPE IF EXISTS e2e_mood CASCADE; + CREATE TYPE e2e_mood AS ENUM ('sad', 'ok', 'happy');", + ) + .unwrap(); + e2e_two_partitions("ty", "u uuid, fl bool, m e2e_mood, v varchar"); + Spi::run( + "INSERT INTO ty VALUES + (1, '00000000-0000-0000-0000-000000000001', false, 'sad', 'aa'), + (2, 'ffffffff-0000-0000-0000-000000000001', true, 'happy', 'zz'); + CREATE INDEX ty_tr ON ty USING table_range (u, fl, m, v);", + ) + .unwrap(); + e2e_set_pruning(true); + for pred in [ + "u = 'ffffffff-0000-0000-0000-000000000001'", + "m = 'happy'", + "v = 'zz'", + ] { + let plan = e2e_explain_on("ty", pred); + assert!(!plan.contains("ty_1"), "`{pred}` must prune ty_1:\n{plan}"); + assert_eq!(e2e_same_on_off("ty", pred), 1); + } + // Widening on insert works for these types too. + Spi::run("INSERT INTO ty VALUES (1, NULL, NULL, 'happy', 'zzz')").unwrap(); + assert_eq!(e2e_same_on_off("ty", "v = 'zzz'"), 1); + assert_eq!(e2e_same_on_off("ty", "m = 'happy'"), 2); +} diff --git a/src/index_am.rs b/src/index_am.rs index d8a1852..02d0ddb 100644 --- a/src/index_am.rs +++ b/src/index_am.rs @@ -1,5 +1,7 @@ use crate::index_storage::{self, ColSummary}; use crate::prune_hook::{btree_cmp_proc, datum_cmp, datum_to_text, text_to_datum}; +use crate::stable_io::with_stable_io; +use crate::summary_build::{geometry_schema, spi_value}; use pgrx::pg_sys; use pgrx::prelude::*; @@ -37,6 +39,8 @@ pub unsafe extern "C-unwind" fn table_range_amhandler( amroutine.amcanmulticol = true; // CREATE INDEX may list several columns amroutine.amoptionalkey = true; // queries need not constrain the first column amroutine.amstorage = true; // opclasses declare a STORAGE type + // Summarizing (like BRIN): updates of indexed columns can stay HOT. + amroutine.amsummarizing = true; amroutine.ambuild = Some(am_build); amroutine.ambuildempty = Some(am_buildempty); @@ -51,41 +55,29 @@ pub unsafe extern "C-unwind" fn table_range_amhandler( } } -/// Build the summary for one (leaf) relation from its actual data. +/// Errors propagate: a failed build must fail `CREATE INDEX`, not silently disable pruning. #[pg_guard] unsafe extern "C-unwind" fn am_build( heap: pg_sys::Relation, index: pg_sys::Relation, index_info: *mut pg_sys::IndexInfo, ) -> *mut pg_sys::IndexBuildResult { - let heap_relid = (*heap).rd_id; - - let nattrs = (*index_info).ii_NumIndexAttrs.max(0) as usize; - let attnums: Vec = (0..nattrs) - .map(|i| (*index_info).ii_IndexAttrNumbers[i]) - .collect(); - - // Compute the summary via SPI under a freshly pushed snapshot so the snapshot the - // CREATE INDEX portal relies on is restored afterwards. Any failure degrades to - // "no summary" (KEEP at planning time). - pg_sys::PushActiveSnapshot(pg_sys::GetTransactionSnapshot()); - let summary = pgrx::PgTryBuilder::new(|| { - let names = crate::summary_build::column_names_for_attnums(heap_relid, &attnums).ok()?; - crate::summary_build::build_one_leaf(heap_relid, &names).ok() - }) - .catch_others(|_| None) - .execute(); - pg_sys::PopActiveSnapshot(); - - // Persist the summary into the index's own metapage. Page writes use no snapshot, - // so this happens after the SPI section. - if let Some(summary) = summary { - let _ = crate::index_storage::write_summary(index, &summary); + // A partial index never sees inserts outside its predicate, so could not cover them. + if !(*index_info).ii_Predicate.is_null() { + ereport!( + ERROR, + PgSqlErrorCode::ERRCODE_FEATURE_NOT_SUPPORTED, + "table_range indexes cannot be partial" + ); } + let mut built = crate::summary_build::build_leaf_summary(heap, index, index_info); + index_storage::write_summary(index, &mut built.summary).unwrap_or_else(|e| error!("{e}")); + let result = pg_sys::palloc0(std::mem::size_of::()) as *mut pg_sys::IndexBuildResult; - (*result).heap_tuples = 0.0; + // index_build stores this as the heap's reltuples. + (*result).heap_tuples = built.heap_tuples; (*result).index_tuples = 0.0; result } @@ -97,6 +89,8 @@ unsafe extern "C-unwind" fn am_buildempty(_index: pg_sys::Relation) {} /// the BRIN-style maintenance path. Because the summary only needs to be over-inclusive, /// this update is in place and needs no MVCC; an insert that stays within the existing /// range writes nothing. This keeps pruning correct and active without any REINDEX. +/// +/// Errors are not caught: an insert the summary does not cover must not succeed. #[pg_guard] #[allow(clippy::too_many_arguments)] // signature is fixed by PostgreSQL's aminsert_function unsafe extern "C-unwind" fn am_insert( @@ -104,29 +98,57 @@ unsafe extern "C-unwind" fn am_insert( values: *mut pg_sys::Datum, isnull: *mut bool, _heap_tid: pg_sys::ItemPointer, - _heap: pg_sys::Relation, + heap: pg_sys::Relation, _check_unique: pg_sys::IndexUniqueCheck::Type, _index_unchanged: bool, index_info: *mut pg_sys::IndexInfo, ) -> bool { - pgrx::PgTryBuilder::new(|| widen_on_insert(index, index_info, values, isnull)) - .catch_others(|_| ()) - .execute(); + widen_on_insert(index, heap, index_info, values, isnull); false } -/// Read the metapage summary, widen each indexed column to cover the new value, and -/// write it back only if something changed. +/// Compare-and-swap against the bytes read, so concurrent widenings never overwrite each other. unsafe fn widen_on_insert( index: pg_sys::Relation, + heap: pg_sys::Relation, index_info: *mut pg_sys::IndexInfo, values: *mut pg_sys::Datum, isnull: *mut bool, ) { - let mut summary = match index_storage::read_summary(index) { - Some(s) => s, - None => return, - }; + loop { + pgrx::check_for_interrupts!(); + // No summary, or an unreadable one: the planner never prunes with it. + let Some(before) = index_storage::read_blob(index) else { + return; + }; + let Some(mut summary) = index_storage::deserialize(&before) else { + return; + }; + let changed = + with_stable_io(|| widen_summary(&mut summary, index, index_info, values, isnull)) + .unwrap_or_else(|| { + error!("table_range: cannot maintain a summary in parallel mode") + }); + if !changed { + return; + } + let after = index_storage::encode_fitting(&mut summary); + if index_storage::replace_blob_if(index, &before, &after).unwrap_or_else(|e| error!("{e}")) + { + crate::summary_cache::note_widened((*index).rd_id, (*heap).rd_id); + return; + } + // Changed by another backend since we read it: retry. + } +} + +unsafe fn widen_summary( + summary: &mut index_storage::IndexSummary, + index: pg_sys::Relation, + index_info: *mut pg_sys::IndexInfo, + values: *mut pg_sys::Datum, + isnull: *mut bool, +) -> bool { let nattrs = (*index_info).ii_NumIndexAttrs.max(0) as usize; let mut changed = false; for i in 0..nattrs { @@ -143,29 +165,24 @@ unsafe fn widen_on_insert( }; changed |= widen_column(col, typoid, collation, *values.add(i), *isnull.add(i)); } - if changed { - let _ = index_storage::write_summary(index, &summary); - // The on-page summary just widened; drop any cached (now-too-narrow) copy in every - // backend so planning never prunes away the newly covered values. - crate::summary_cache::note_widened((*index).rd_id); - } + changed } - /// The type OID of a tuple descriptor's `i`-th attribute, portable across PG versions. /// PG18 made `TupleDescAttr` an inline function (bound by pgrx) and moved attributes to /// `compact_attrs` (which has no `atttypid`); PG13–17 expose `attrs` directly and only a /// `TupleDescAttr` macro (which bindgen does not surface as `pg_sys::TupleDescAttr`). #[cfg(feature = "pg18")] -unsafe fn att_typid(tupdesc: pg_sys::TupleDesc, i: usize) -> pg_sys::Oid { +pub(crate) unsafe fn att_typid(tupdesc: pg_sys::TupleDesc, i: usize) -> pg_sys::Oid { (*pg_sys::TupleDescAttr(tupdesc, i as i32)).atttypid } #[cfg(not(feature = "pg18"))] -unsafe fn att_typid(tupdesc: pg_sys::TupleDesc, i: usize) -> pg_sys::Oid { +pub(crate) unsafe fn att_typid(tupdesc: pg_sys::TupleDesc, i: usize) -> pg_sys::Oid { let natts = (*tupdesc).natts as usize; (*tupdesc).attrs.as_slice(natts)[i].atttypid } -/// Widen one column's summary for a single inserted value. Returns whether it changed. +/// Returns whether the summary changed. A value that cannot be folded in makes the column +/// unbounded rather than leaving it too narrow. unsafe fn widen_column( col: &mut ColSummary, typoid: pg_sys::Oid, @@ -186,79 +203,116 @@ unsafe fn widen_column( col.all_nulls = false; changed = true; } - - if col.overlap { - return changed | widen_overlap(col, typoid, value); + if col.unbounded { + return changed; } - // btree min/max: widen using the column type's compare support function. - let new_text = match datum_to_text(typoid, value) { - Some(t) => t, - None => return changed, + let widened = if col.overlap { + widen_overlap(col, typoid, value) + } else { + widen_minmax(col, typoid, collation, value) }; + match widened { + Some(c) => changed | c, + None => { + col.make_unbounded(); + true + } + } +} + +/// `None`: the value could not be compared against the stored bounds. +unsafe fn widen_minmax( + col: &mut ColSummary, + typoid: pg_sys::Oid, + collation: pg_sys::Oid, + value: pg_sys::Datum, +) -> Option { + // Domains compare and parse as their base type. + let base = pg_sys::getBaseType(typoid); + let new_text = datum_to_text(typoid, value)?; match (&col.min, &col.max) { (Some(min), Some(max)) => { - let cmp = match btree_cmp_proc(typoid) { - Some(c) => c, - None => return changed, - }; - if let (Some(min_d), Some(max_d)) = - (text_to_datum(typoid, min), text_to_datum(typoid, max)) - { - if datum_cmp(cmp, collation, value, min_d) < 0 { - col.min = Some(new_text.clone()); - changed = true; - } - if datum_cmp(cmp, collation, value, max_d) > 0 { - col.max = Some(new_text); - changed = true; - } + let cmp = btree_cmp_proc(base)?; + let min_d = text_to_datum(base, min)?; + let max_d = text_to_datum(base, max)?; + let mut changed = false; + if datum_cmp(cmp, collation, value, min_d) < 0 { + col.min = Some(new_text.clone()); + changed = true; } + if datum_cmp(cmp, collation, value, max_d) > 0 { + col.max = Some(new_text); + changed = true; + } + Some(changed) } // No range yet (was all-null/empty): seed it with this value. - _ => { + (None, None) => { col.min = Some(new_text.clone()); col.max = Some(new_text); - changed = true; + Some(true) } + _ => None, } - changed } -/// Widen an overlap (range/geometry) extent to cover a new value, via the type's own -/// union operator (delegated to SQL). Writes nothing if the value is already covered. -unsafe fn widen_overlap(col: &mut ColSummary, typoid: pg_sys::Oid, value: pg_sys::Datum) -> bool { - let new_text = match datum_to_text(typoid, value) { - Some(t) => t, - None => return false, +/// `None`: the new extent could not be computed. +unsafe fn widen_overlap( + col: &mut ColSummary, + typoid: pg_sys::Oid, + value: pg_sys::Datum, +) -> Option { + let base = pg_sys::getBaseType(typoid); + let is_geometry = col.type_name.ends_with(".geometry"); + let widened = if is_geometry { + geometry_extent_with(col, typoid, value)? + } else { + let merged = match &col.min { + None => value, + Some(ext) => { + let ext = text_to_datum(base, ext)?; + pg_sys::OidFunctionCall2Coll( + pg_sys::Oid::from(pg_sys::F_RANGE_MERGE_ANYRANGE_ANYRANGE), + pg_sys::Oid::INVALID, + ext, + value, + ) + } + }; + Some(datum_to_text(typoid, merged)?) }; - let tn = &col.type_name; - let is_geometry = tn == "geometry" || tn.ends_with(".geometry"); - let lit = |s: &str| format!("'{}'", s.replace('\'', "''")); - let new_lit = lit(&new_text); - let sql = match &col.min { - None if is_geometry => format!( - "SELECT ST_Extent(g)::geometry::text FROM (VALUES (CAST({new_lit} AS {tn}))) v(g)" - ), - None => format!("SELECT CAST({new_lit} AS {tn})::text"), - Some(ext) if is_geometry => format!( - "SELECT ST_Extent(g)::geometry::text FROM \ - (VALUES (CAST({} AS {tn})), (CAST({new_lit} AS {tn}))) v(g)", - lit(ext) - ), - Some(ext) => format!( - "SELECT range_merge(CAST({} AS {tn}), CAST({new_lit} AS {tn}))::text", - lit(ext) - ), - }; + match widened { + Some(w) if col.min.as_ref() != Some(&w) => { + col.min = Some(w); + Some(true) + } + Some(_) => Some(false), + // An empty geometry has no extent and can never satisfy `&&`: nothing to cover. + None => Some(false), + } +} - let widened = Spi::get_one::(&sql).ok().flatten(); - if widened.is_some() && widened != col.min { - col.min = widened; - return true; +/// `Some(None)` if neither the stored extent nor `value` has one (empty geometries). +unsafe fn geometry_extent_with( + col: &ColSummary, + typoid: pg_sys::Oid, + value: pg_sys::Datum, +) -> Option> { + let new_text = datum_to_text(typoid, value)?; + let tn = &col.type_name; + let ns = geometry_schema(tn); + let lit = |s: &str| format!("'{}'", s.replace('\'', "''")); + let mut rows = format!("(CAST({} AS {tn}))", lit(&new_text)); + if let Some(ext) = &col.min { + rows.push_str(&format!(", (CAST({} AS {tn}))", lit(ext))); } - false + let sql = format!("SELECT {ns}.st_extent(g)::{tn}::pg_catalog.text FROM (VALUES {rows}) v(g)"); + pg_sys::SPI_connect(); + let extent = spi_value(&sql); + pg_sys::SPI_finish(); + Some(extent) } #[pg_guard] diff --git a/src/index_am_tests.rs b/src/index_am_tests.rs index f42d4ec..188d5df 100644 --- a/src/index_am_tests.rs +++ b/src/index_am_tests.rs @@ -87,35 +87,34 @@ fn am_multicolumn_index_prunes() { fn am_insert_then_correct_without_rebuild() { am_setup(); Spi::run("CREATE INDEX amt_tr ON amt USING table_range (val)").unwrap(); - // Insert a value outside r1's built range; staleness must keep results correct. + // Insert a value outside r1's built range; aminsert widens the summary to cover it. Spi::run("INSERT INTO amt VALUES (1, 500, date '2025-01-01')").unwrap(); assert_eq!( Spi::get_one::("SELECT count(*)::bigint FROM amt WHERE val = 500") .unwrap() .unwrap(), 1, - "stale summary must not prune away the newly inserted row" + "widened summary must cover the newly inserted row" ); } #[pg_test] -fn am_bulk_insert_stays_correct_under_stale_memo() { +fn am_bulk_insert_widens_summary() { am_setup(); Spi::run("CREATE INDEX amt_tr ON amt USING table_range (val)").unwrap(); Spi::run("SET table_range.enable_pruning = on").unwrap(); - // Bulk insert many out-of-range rows into r1 in one statement. This exercises the - // per-transaction stale memo (r1 is marked stale once, not once per row); all rows - // must still be found despite the now-stale summary. + // Bulk insert many out-of-range rows into r1 in one statement. Each row widens the + // summary as needed, so all of them are found. Spi::run("INSERT INTO amt SELECT 1, g FROM generate_series(1000, 1099) g").unwrap(); assert_eq!( Spi::get_one::("SELECT count(*)::bigint FROM amt WHERE val BETWEEN 1000 AND 1099") .unwrap() .unwrap(), 100, - "stale summary must not prune away bulk-inserted rows" + "widened summary must cover bulk-inserted rows" ); - // A predicate matching original r3 data still works (no over-pruning). + // A predicate matching r3's data still works (no over-pruning). assert_eq!( Spi::get_one::("SELECT count(*)::bigint FROM amt WHERE val = 250") .unwrap() @@ -190,3 +189,107 @@ fn ambuild_writes_summary_to_page() { "ambuild must persist the summary to the index page: {out}" ); } + +#[pg_test] +fn am_reindex_rebuilds_summary() { + // REINDEX rebuilds the summary from the table's data. + Spi::run( + "DROP TABLE IF EXISTS rx CASCADE; + CREATE TABLE rx (val bigint); + INSERT INTO rx SELECT g FROM generate_series(1, 100) g; + CREATE INDEX rx_tr ON rx USING table_range (val); + INSERT INTO rx VALUES (500); + REINDEX INDEX rx_tr;", + ) + .unwrap(); + let summary = + Spi::get_one::("SELECT table_range_test_read_summary('rx_tr'::regclass::oid)") + .unwrap() + .unwrap(); + assert!( + summary.contains("min=Some(\"1\")") && summary.contains("max=Some(\"500\")"), + "REINDEX must rebuild the summary from the data: {summary}" + ); + assert_eq!(e2e_same_on_off("rx", "val = 500"), 1); + assert_eq!(e2e_same_on_off("rx", "val = 300"), 0); +} + +#[pg_test(error = "table_range indexes cannot be partial")] +fn am_partial_index_rejected() { + Spi::run( + "DROP TABLE IF EXISTS pi CASCADE; + CREATE TABLE pi (val bigint, flag bool); + CREATE INDEX pi_tr ON pi USING table_range (val) WHERE flag;", + ) + .unwrap(); +} + +#[pg_test] +fn am_build_reports_heap_row_count() { + // index_build stores the reported heap_tuples as the table's reltuples. + Spi::run( + "DROP TABLE IF EXISTS rt CASCADE; + CREATE TABLE rt (val bigint); + INSERT INTO rt SELECT g FROM generate_series(1, 500) g; + ANALYZE rt; + CREATE INDEX rt_tr ON rt USING table_range (val);", + ) + .unwrap(); + let reltuples = Spi::get_one::("SELECT reltuples FROM pg_class WHERE relname = 'rt'") + .unwrap() + .unwrap(); + assert_eq!(reltuples, 500.0); +} + +#[pg_test] +fn am_oversized_value_makes_column_unbounded() { + e2e_two_partitions("ov", "s text"); + Spi::run( + "INSERT INTO ov VALUES (1, 'm'), (2, 'a'); + CREATE INDEX ov_tr ON ov USING table_range (s); + INSERT INTO ov VALUES (1, 'z' || repeat('x', 9000));", + ) + .unwrap(); + assert_eq!(e2e_same_on_off("ov", "s > 'y'"), 1); + let summary = Spi::get_one::( + "SELECT table_range_test_read_summary(i.indexrelid) + FROM pg_index i WHERE i.indrelid = 'ov_1'::regclass", + ) + .unwrap() + .unwrap(); + assert!(summary.contains("unbounded=true"), "{summary}"); +} + +#[pg_test] +fn am_range_extent_widens_on_insert() { + e2e_two_partitions("rg", "r int4range"); + Spi::run( + "INSERT INTO rg VALUES (1, '[0,10)'), (2, '[100,110)'); + CREATE INDEX rg_tr ON rg USING table_range (r); + INSERT INTO rg VALUES (1, '[500,501)'), (1, 'empty');", + ) + .unwrap(); + assert_eq!(e2e_same_on_off("rg", "r && '[500,500]'::int4range"), 1); + // After the insert rg_1's extent is [0,501): past it, both partitions prune. + let plan = e2e_explain_on("rg", "r && '[600,601)'::int4range"); + assert!(!plan.contains("rg_1") && !plan.contains("rg_2"), "{plan}"); +} + +#[pg_test] +fn am_unknown_summary_format_is_ignored() { + // A well-formed summary in another format version claiming val in [0, 0]. + Spi::run( + "DROP TABLE IF EXISTS fmt CASCADE; + CREATE TABLE fmt (val bigint); + INSERT INTO fmt VALUES (5); + CREATE INDEX fmt_tr ON fmt USING table_range (val);", + ) + .unwrap(); + // version 1 | ncols=1 | attnum=1 | overlap=0 | flags=0 | "int8" | min "0" | max "0" + Spi::run( + "SELECT table_range_test_write_blob('fmt_tr'::regclass::oid, + '\\x010100010000000400696e74380101003001010030')", + ) + .unwrap(); + assert_eq!(e2e_same_on_off("fmt", "val = 5"), 1); +} diff --git a/src/index_storage.rs b/src/index_storage.rs index 5a852f7..06f8ae1 100644 --- a/src/index_storage.rs +++ b/src/index_storage.rs @@ -5,35 +5,34 @@ use pgrx::prelude::*; // Low-level storage for a per-index summary, kept in the index's own metapage (block 0) // and updated in place with the Generic WAL API — the same model BRIN uses. Because a // table_range summary only ever needs to be *over-inclusive*, these in-place page -// updates need no MVCC/transactionality: a rolled-back or concurrent widening that -// over-covers is always safe. -// -// This module currently provides the raw page round-trip (length-prefixed byte blob in -// the metapage). Typed summary (de)serialization and the ambuild/aminsert/planner wiring -// build on top of it. +// updates need no MVCC/transactionality: a rolled-back widening that over-covers is +// always safe. // Usable bytes in the metapage content area (after the page header + length prefix). A // generous margin below BLCKSZ keeps us clear of the page header regardless of alignment. -fn max_blob_len() -> usize { +pub(crate) fn max_blob_len() -> usize { pg_sys::BLCKSZ as usize - 64 } -/// Write a byte blob into the index's metapage (block 0), creating the block if needed. -/// WAL-logged via Generic WAL. Caller must hold a lock on the index relation appropriate -/// for the calling context (ambuild owns it; aminsert holds the row's locks). -pub unsafe fn write_blob(index: pg_sys::Relation, data: &[u8]) -> Result<(), &'static str> { - if data.len() > max_blob_len() { - return Err("table_range: summary blob exceeds one page"); +unsafe fn page_blob(page: pg_sys::Page) -> Option> { + if pg_sys::PageIsNew(page) { + return None; } - let nblocks = pg_sys::RelationGetNumberOfBlocksInFork(index, pg_sys::ForkNumber::MAIN_FORKNUM); - let is_new = nblocks == 0; - let buffer = if is_new { - pg_sys::ReadBuffer(index, pg_sys::InvalidBlockNumber) // P_NEW -> extend - } else { - pg_sys::ReadBuffer(index, 0) - }; - pg_sys::LockBuffer(buffer, pg_sys::BUFFER_LOCK_EXCLUSIVE as i32); + let contents = pg_sys::PageGetContents(page) as *const u8; + let mut len_bytes = [0u8; 4]; + std::ptr::copy_nonoverlapping(contents, len_bytes.as_mut_ptr(), 4); + let len = u32::from_ne_bytes(len_bytes) as usize; + if len == 0 || len > max_blob_len() { + return None; + } + let mut data = vec![0u8; len]; + std::ptr::copy_nonoverlapping(contents.add(4), data.as_mut_ptr(), len); + Some(data) +} +// shortcut: logs a full page image per write (~8 KB WAL per widening); log a delta once +// pd_lower marks the end of the data, if insert-heavy workloads need less WAL. +unsafe fn put_blob(index: pg_sys::Relation, buffer: pg_sys::Buffer, is_new: bool, data: &[u8]) { let state = pg_sys::GenericXLogStart(index); let page = pg_sys::GenericXLogRegisterBuffer(state, buffer, pg_sys::GENERIC_XLOG_FULL_IMAGE as i32); @@ -53,10 +52,50 @@ pub unsafe fn write_blob(index: pg_sys::Relation, data: &[u8]) -> Result<(), &'s (*header).pd_lower = (*header).pd_upper; pg_sys::GenericXLogFinish(state); +} + +/// For `ambuild`, which holds the index exclusively; `aminsert` uses `replace_blob_if`. +pub unsafe fn write_blob(index: pg_sys::Relation, data: &[u8]) -> Result<(), &'static str> { + if data.len() > max_blob_len() { + return Err("table_range: summary blob exceeds one page"); + } + let nblocks = pg_sys::RelationGetNumberOfBlocksInFork(index, pg_sys::ForkNumber::MAIN_FORKNUM); + let is_new = nblocks == 0; + let buffer = if is_new { + pg_sys::ReadBuffer(index, pg_sys::InvalidBlockNumber) // P_NEW -> extend + } else { + pg_sys::ReadBuffer(index, 0) + }; + pg_sys::LockBuffer(buffer, pg_sys::BUFFER_LOCK_EXCLUSIVE as i32); + put_blob(index, buffer, is_new, data); pg_sys::UnlockReleaseBuffer(buffer); Ok(()) } +/// Writes only if the blob still equals `expected`; `Ok(false)` means re-read and retry. +pub unsafe fn replace_blob_if( + index: pg_sys::Relation, + expected: &[u8], + data: &[u8], +) -> Result { + if data.len() > max_blob_len() { + return Err("table_range: summary blob exceeds one page"); + } + let nblocks = pg_sys::RelationGetNumberOfBlocksInFork(index, pg_sys::ForkNumber::MAIN_FORKNUM); + if nblocks == 0 { + return Ok(false); + } + let buffer = pg_sys::ReadBuffer(index, 0); + pg_sys::LockBuffer(buffer, pg_sys::BUFFER_LOCK_EXCLUSIVE as i32); + let current = page_blob(pg_sys::BufferGetPage(buffer)); + let matches = current.as_deref() == Some(expected); + if matches { + put_blob(index, buffer, false, data); + } + pg_sys::UnlockReleaseBuffer(buffer); + Ok(matches) +} + /// Read the byte blob from the index's metapage, or `None` if the index has no metapage. pub unsafe fn read_blob(index: pg_sys::Relation) -> Option> { let nblocks = pg_sys::RelationGetNumberOfBlocksInFork(index, pg_sys::ForkNumber::MAIN_FORKNUM); @@ -65,22 +104,7 @@ pub unsafe fn read_blob(index: pg_sys::Relation) -> Option> { } let buffer = pg_sys::ReadBuffer(index, 0); pg_sys::LockBuffer(buffer, pg_sys::BUFFER_LOCK_SHARE as i32); - let page = pg_sys::BufferGetPage(buffer); - let result = if pg_sys::PageIsNew(page) { - None - } else { - let contents = pg_sys::PageGetContents(page) as *const u8; - let mut len_bytes = [0u8; 4]; - std::ptr::copy_nonoverlapping(contents, len_bytes.as_mut_ptr(), 4); - let len = u32::from_ne_bytes(len_bytes) as usize; - if len == 0 || len > max_blob_len() { - None - } else { - let mut data = vec![0u8; len]; - std::ptr::copy_nonoverlapping(contents.add(4), data.as_mut_ptr(), len); - Some(data) - } - }; + let result = page_blob(pg_sys::BufferGetPage(buffer)); pg_sys::UnlockReleaseBuffer(buffer); result } @@ -95,12 +119,24 @@ pub struct ColSummary { /// `true` -> `min` holds a covering extent for `&&` pruning (range/geometry); /// `false` -> `min`/`max` hold the column's btree min/max. pub overlap: bool, - /// SQL type name (for casting in overlap evaluation). + /// Schema-qualified, so the SQL that maintains extents ignores the caller's search_path. pub type_name: String, + /// Min/max order; only comparisons under this collation are pruned. + pub collation: pg_sys::Oid, pub min: Option, pub max: Option, pub has_nulls: bool, pub all_nulls: bool, + /// Bounds unknown (too large or not representable): only null flags prune. + pub unbounded: bool, +} + +impl ColSummary { + pub fn make_unbounded(&mut self) { + self.unbounded = true; + self.min = None; + self.max = None; + } } /// The whole index's summary: one entry per indexed column. @@ -109,7 +145,12 @@ pub struct IndexSummary { pub cols: Vec, } -const SUMMARY_VERSION: u8 = 1; +// Other format versions are unreadable on purpose: they may not cover their data. +const SUMMARY_VERSION: u8 = 2; + +const FLAG_HAS_NULLS: u8 = 1; +const FLAG_ALL_NULLS: u8 = 2; +const FLAG_UNBOUNDED: u8 = 4; fn put_str(out: &mut Vec, s: &str) { out.extend_from_slice(&(s.len() as u16).to_le_bytes()); @@ -142,6 +183,11 @@ impl Reader<'_> { self.pos += 2; Some(u16::from_le_bytes([bytes[0], bytes[1]])) } + fn u32(&mut self) -> Option { + let bytes = self.buf.get(self.pos..self.pos + 4)?; + self.pos += 4; + Some(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])) + } fn i16(&mut self) -> Option { self.u16().map(|v| v as i16) } @@ -166,7 +212,18 @@ pub fn serialize(summary: &IndexSummary) -> Vec { for c in &summary.cols { out.extend_from_slice(&c.attnum.to_le_bytes()); out.push(c.overlap as u8); - out.push((c.has_nulls as u8) | ((c.all_nulls as u8) << 1)); + let mut flags = 0; + if c.has_nulls { + flags |= FLAG_HAS_NULLS; + } + if c.all_nulls { + flags |= FLAG_ALL_NULLS; + } + if c.unbounded { + flags |= FLAG_UNBOUNDED; + } + out.push(flags); + out.extend_from_slice(&u32::from(c.collation).to_le_bytes()); put_str(&mut out, &c.type_name); put_opt_str(&mut out, &c.min); put_opt_str(&mut out, &c.max); @@ -185,6 +242,7 @@ pub fn deserialize(buf: &[u8]) -> Option { let attnum = r.i16()?; let overlap = r.u8()? != 0; let flags = r.u8()?; + let collation = pg_sys::Oid::from(r.u32()?); let type_name = r.str()?; let min = r.opt_str()?; let max = r.opt_str()?; @@ -192,21 +250,46 @@ pub fn deserialize(buf: &[u8]) -> Option { attnum, overlap, type_name, + collation, min, max, - has_nulls: flags & 1 != 0, - all_nulls: flags & 2 != 0, + has_nulls: flags & FLAG_HAS_NULLS != 0, + all_nulls: flags & FLAG_ALL_NULLS != 0, + unbounded: flags & FLAG_UNBOUNDED != 0, }); } Some(IndexSummary { cols }) } +/// Drops the largest columns' bounds until it fits: less selective, never too narrow. +pub fn encode_fitting(summary: &mut IndexSummary) -> Vec { + loop { + let bytes = serialize(summary); + if bytes.len() <= max_blob_len() { + return bytes; + } + let bound_len = |c: &ColSummary| { + c.min.as_ref().map_or(0, String::len) + c.max.as_ref().map_or(0, String::len) + }; + match summary + .cols + .iter_mut() + .filter(|c| bound_len(c) > 0) + .max_by_key(|c| bound_len(c)) + { + Some(c) => c.make_unbounded(), + // Nothing left to drop; the caller's write will report the size error. + None => return bytes, + } + } +} + /// Persist the typed summary into the index metapage. pub unsafe fn write_summary( index: pg_sys::Relation, - summary: &IndexSummary, + summary: &mut IndexSummary, ) -> Result<(), &'static str> { - write_blob(index, &serialize(summary)) + write_blob(index, &encode_fitting(summary)) } /// Read the typed summary from the index metapage, if present. @@ -218,6 +301,20 @@ pub unsafe fn read_summary(index: pg_sys::Relation) -> Option { mod serde_tests { use super::*; + fn col(attnum: i16, min: &str, max: &str) -> ColSummary { + ColSummary { + attnum, + overlap: false, + type_name: "pg_catalog.text".into(), + collation: pg_sys::Oid::from(100), + min: Some(min.into()), + max: Some(max.into()), + has_nulls: false, + all_nulls: false, + unbounded: false, + } + } + #[test] fn summary_roundtrips() { let s = IndexSummary { @@ -225,20 +322,31 @@ mod serde_tests { ColSummary { attnum: 2, overlap: false, - type_name: "bigint".into(), + type_name: "pg_catalog.int8".into(), + collation: pg_sys::Oid::INVALID, min: Some("0".into()), max: Some("99".into()), has_nulls: true, all_nulls: false, + unbounded: false, }, ColSummary { attnum: 3, overlap: true, - type_name: "int8range".into(), + type_name: "pg_catalog.int8range".into(), + collation: pg_sys::Oid::INVALID, min: Some("[0,100)".into()), max: None, has_nulls: false, all_nulls: false, + unbounded: false, + }, + ColSummary { + unbounded: true, + min: None, + max: None, + has_nulls: true, + ..col(4, "", "") }, ], }; @@ -249,6 +357,25 @@ mod serde_tests { fn rejects_garbage_and_wrong_version() { assert_eq!(deserialize(&[]), None); assert_eq!(deserialize(&[99, 0, 0]), None); + // Another format version is never trusted, even if well-formed. + let mut v1 = serialize(&IndexSummary { + cols: vec![col(1, "a", "b")], + }); + v1[0] = 1; + assert_eq!(deserialize(&v1), None); + } + + #[test] + fn oversized_summary_drops_largest_bounds_first() { + let big = "x".repeat(max_blob_len()); + let mut s = IndexSummary { + cols: vec![col(1, "a", "b"), col(2, "a", &big)], + }; + let bytes = encode_fitting(&mut s); + assert!(bytes.len() <= max_blob_len()); + assert!(!s.cols[0].unbounded, "small column keeps its bounds"); + assert!(s.cols[1].unbounded && s.cols[1].max.is_none()); + assert_eq!(deserialize(&bytes), Some(s)); } } @@ -269,8 +396,9 @@ fn table_range_test_read_summary(index: pg_sys::Oid) -> String { .iter() .map(|c| { format!( - "attnum={} overlap={} min={:?} max={:?} has_nulls={} all_nulls={}", - c.attnum, c.overlap, c.min, c.max, c.has_nulls, c.all_nulls + "attnum={} overlap={} min={:?} max={:?} has_nulls={} all_nulls={} \ + unbounded={}", + c.attnum, c.overlap, c.min, c.max, c.has_nulls, c.all_nulls, c.unbounded ) }) .collect::>() @@ -293,3 +421,13 @@ fn table_range_test_page_roundtrip(index: pg_sys::Oid, payload: String) -> Strin } } } + +#[cfg(any(test, feature = "pg_test"))] +#[pg_extern] +fn table_range_test_write_blob(index: pg_sys::Oid, payload: &[u8]) { + unsafe { + let rel = pg_sys::index_open(index, pg_sys::AccessExclusiveLock as i32); + write_blob(rel, payload).unwrap_or_else(|e| error!("write_blob: {e}")); + pg_sys::index_close(rel, pg_sys::AccessExclusiveLock as i32); + } +} diff --git a/src/lib.rs b/src/lib.rs index eda56a1..5181d64 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,6 +8,7 @@ use pgrx::{GucContext, GucFlags, GucRegistry}; mod index_am; mod index_storage; mod prune_hook; +mod stable_io; mod summary_build; mod summary_cache; diff --git a/src/prune_hook.rs b/src/prune_hook.rs index e5ec95a..c924d3c 100644 --- a/src/prune_hook.rs +++ b/src/prune_hook.rs @@ -1,5 +1,5 @@ +use pgrx::pg_guard; use pgrx::pg_sys; -use pgrx::prelude::*; use std::cell::{Cell, RefCell}; use std::collections::HashMap; use std::ffi::{CStr, CString}; @@ -13,11 +13,9 @@ use crate::{TABLE_RANGE_ENABLE_PRUNING, TABLE_RANGE_LOG_PRUNING_DEBUG}; // // For every base/partition-member relation the planner builds, we look up that // relation's cached summaries and evaluate its restriction clauses (`Var op Const`, -// `Var IS [NOT] NULL`, `Var IN (...)`). Typed comparisons are performed in-memory with -// the column type's btree compare support function, so any error (e.g. an -// incompatible cast) is caught and degrades to KEEP. If any single clause proves the -// partition cannot contain a matching row, we call `mark_dummy_rel` so the planner -// eliminates it before generating child paths. +// `Var IS [NOT] NULL`, `Var IN (...)`, `Var && Const`) in memory against the query's own +// constant Datums. Anything that cannot be decided exactly degrades to KEEP. If any single +// clause proves the partition cannot contain a matching row, we call `mark_dummy_rel`. // // Each partition's summary is read from its table_range index's metapage and cached for // the duration of one top-level planner invocation. @@ -44,14 +42,8 @@ thread_local! { static AM_OID: Cell> = const { Cell::new(None) }; /// Nesting depth of planner invocations (SPI during planning re-enters). static PLAN_DEPTH: Cell = const { Cell::new(0) }; - /// Guards against re-entering pruning logic from the SPI overlap evaluation issues. + /// Guards against re-entering pruning logic from catalog access during evaluation. static IN_HOOK: Cell = const { Cell::new(false) }; - /// Per-plan memo of each column type's btree compare proc OID (the lookup is three - /// syscache hits, identical for every partition of a column, so we do it once). - static CMP_PROC_MEMO: RefCell>> = RefCell::new(HashMap::new()); - /// Per-plan memo of parsed query constants, keyed by (type oid, text). The same - /// constant is otherwise re-rendered and re-parsed once per partition. - static CONST_MEMO: RefCell> = RefCell::new(HashMap::new()); /// Per-plan cache of resolved `FmgrInfo` structs (one palloc per function), so each /// comparison / input-function call across partitions skips the `fmgr_info` syscache /// lookup. Pointers are into the planner memory context and cleared per top-level plan. @@ -59,8 +51,8 @@ thread_local! { /// Per-plan memo of a type's input function + ioparam (from `getTypeInputInfo`). static INPUT_INFO_MEMO: RefCell> = RefCell::new(HashMap::new()); - /// Per-plan memo of the btree strategy for an (operator, left type) pair. - static STRATEGY_MEMO: RefCell>> = RefCell::new(HashMap::new()); + /// Per-plan memo keyed by (operator, column is right operand): same for every partition. + static OP_MEMO: RefCell>> = RefCell::new(HashMap::new()); } /// A planner-cached `FmgrInfo` for `proc_oid`, valid for the current plan. Avoids the @@ -117,45 +109,65 @@ unsafe fn parse_cached(typ: pg_sys::Oid, text: &str) -> Option { )) } -/// btree strategy for an (operator, left type) pair, memoized for the current plan. -unsafe fn strategy_cached(opno: pg_sys::Oid, lefttype: pg_sys::Oid) -> Option { - let key = (opno.into(), lefttype.into()); - if let Some(v) = STRATEGY_MEMO.with(|m| m.borrow().get(&key).copied()) { +#[derive(Clone, Copy)] +struct CompareOp { + /// Strategy (1..=5) with the column on the left. + strategy: i16, + /// May be polymorphic (e.g. anyenum). + var_type: pg_sys::Oid, + /// Compare proc for (column-side type, constant-side type). + cmpproc: pg_sys::Oid, +} + +unsafe fn compare_op_cached(opno: pg_sys::Oid, commuted: bool) -> Option { + let key = (opno.into(), commuted); + if let Some(v) = OP_MEMO.with(|m| m.borrow().get(&key).copied()) { return v; } - let v = btree_strategy(opno, lefttype); - STRATEGY_MEMO.with(|m| { + let v = compare_op(opno, commuted); + OP_MEMO.with(|m| { m.borrow_mut().insert(key, v); }); v } -/// Compare proc for `typ`, memoized for the current plan. See [`btree_cmp_proc`]. -unsafe fn cmp_proc_cached(typ: pg_sys::Oid) -> Option { - let key: u32 = typ.into(); - if let Some(v) = CMP_PROC_MEMO.with(|m| m.borrow().get(&key).copied()) { - return v; +unsafe fn compare_op(opno: pg_sys::Oid, commuted: bool) -> Option { + let mut left = pg_sys::Oid::INVALID; + let mut right = pg_sys::Oid::INVALID; + pg_sys::op_input_types(opno, &mut left, &mut right); + let (var_type, const_type) = if commuted { + (right, left) + } else { + (left, right) + }; + let opfamily = default_btree_family(var_type)?; + let strategy = match pg_sys::get_op_opfamily_strategy(opno, opfamily) { + s @ 1..=5 => s as i16, + _ => return None, + }; + // Cross-type operators (float8 > float4) compare via the family's cross-type proc. + let cmpproc = pg_sys::get_opfamily_proc(opfamily, var_type, const_type, BTORDER_PROC as i16); + if cmpproc == pg_sys::Oid::INVALID { + return None; } - let v = btree_cmp_proc(typ); - CMP_PROC_MEMO.with(|m| { - m.borrow_mut().insert(key, v); - }); - v + Some(CompareOp { + strategy: if commuted { + commute_strategy(strategy) + } else { + strategy + }, + var_type, + cmpproc, + }) } -/// Parse `text` to a Datum of `typ`, memoized for the current plan. The cached Datum is -/// allocated in the planner's memory context (which outlives the plan) and the memo is -/// cleared per top-level plan, so the pointer stays valid for its lifetime. -unsafe fn const_datum_cached(typ: pg_sys::Oid, text: &str) -> Option { - let key = (typ.into(), text.to_string()); - if let Some(d) = CONST_MEMO.with(|m| m.borrow().get(&key).copied()) { - return Some(d); +/// A polymorphic operator input type (anyenum, anyarray) cannot parse text; use the column's. +unsafe fn summary_parse_type(op_var_type: pg_sys::Oid, vartype: pg_sys::Oid) -> pg_sys::Oid { + if pg_sys::get_typtype(op_var_type) as u8 == b'p' { + pg_sys::getBaseType(vartype) + } else { + op_var_type } - let d = parse_cached(typ, text)?; - CONST_MEMO.with(|m| { - m.borrow_mut().insert(key, d); - }); - Some(d) } /// Install our planner and pathlist hooks, preserving any previously-registered hooks. @@ -168,6 +180,8 @@ pub fn install() { } } +// `#[pg_guard]`: planner ERRORs reach Rust as panics, which must not unwind into C. +#[pg_guard] unsafe extern "C-unwind" fn table_range_planner_hook( parse: *mut pg_sys::Query, query_string: *const std::ffi::c_char, @@ -184,22 +198,25 @@ unsafe extern "C-unwind" fn table_range_planner_hook( clear_cache(); } - let result = match PREV_PLANNER_HOOK { + // Per-plan state points into planner memory: reset it even when planning errors. + pgrx::PgTryBuilder::new(|| match PREV_PLANNER_HOOK { Some(prev) => prev(parse, query_string, cursor_options, bound_params), None => pg_sys::standard_planner(parse, query_string, cursor_options, bound_params), - }; - - let last = PLAN_DEPTH.with(|d| { - let v = d.get() - 1; - d.set(v); - v == 0 - }); - if last { - clear_cache(); - } - result + }) + .finally(|| { + let last = PLAN_DEPTH.with(|d| { + let v = d.get().saturating_sub(1); + d.set(v); + v == 0 + }); + if last { + clear_cache(); + } + }) + .execute() } +#[pg_guard] unsafe extern "C-unwind" fn table_range_pathlist_hook( root: *mut pg_sys::PlannerInfo, rel: *mut pg_sys::RelOptInfo, @@ -231,11 +248,16 @@ unsafe extern "C-unwind" fn table_range_pathlist_hook( } IN_HOOK.with(|f| f.set(true)); + // shortcut: an ERROR here is caught without a subtransaction; switch input parsing to + // soft errors (InputFunctionCallSafe) and drop the catch once nothing else can raise. let prune = pgrx::PgTryBuilder::new(|| evaluate_relation(rel, relid)) .catch_others(|_| false) .execute(); IN_HOOK.with(|f| f.set(false)); + // shortcut: decided at plan time, so under READ COMMITTED a widening that commits + // between planning and execution is missed by that execution; recheck at executor + // start if that window matters. if prune { if TABLE_RANGE_LOG_PRUNING_DEBUG.get() { let oid: u32 = relid.into(); @@ -248,13 +270,11 @@ unsafe extern "C-unwind" fn table_range_pathlist_hook( fn clear_cache() { CACHE.with(|c| c.borrow_mut().clear()); AM_OID.with(|c| c.set(None)); - CMP_PROC_MEMO.with(|c| c.borrow_mut().clear()); - CONST_MEMO.with(|c| c.borrow_mut().clear()); // FmgrInfo structs are palloc'd in the planner context and freed when it resets; we // just drop the (now-dangling-after-reset) pointers at the start/end of each plan. FMGR_MEMO.with(|c| c.borrow_mut().clear()); INPUT_INFO_MEMO.with(|c| c.borrow_mut().clear()); - STRATEGY_MEMO.with(|c| c.borrow_mut().clear()); + OP_MEMO.with(|c| c.borrow_mut().clear()); } /// The table_range access-method OID, resolved once per planner invocation. @@ -325,6 +345,10 @@ unsafe fn read_index_summary(rel: *mut pg_sys::RelOptInfo) -> Rc /// Returns true iff some restriction clause proves the partition cannot match. unsafe fn evaluate_relation(rel: *mut pg_sys::RelOptInfo, relid: pg_sys::Oid) -> bool { + crate::stable_io::with_stable_io(|| evaluate_relation_inner(rel, relid)).unwrap_or(false) +} + +unsafe fn evaluate_relation_inner(rel: *mut pg_sys::RelOptInfo, relid: pg_sys::Oid) -> bool { let relid_u32: u32 = relid.into(); load_summary(rel, relid_u32); @@ -395,29 +419,35 @@ unsafe fn clause_proves_prune(node: *mut pg_sys::Node, rows: &[ColSummary], dept // --------------------------------------------------------------------------------- enum QualSpec { - /// `Var op Const`: prune decision from the column-typed btree comparison. + /// `Var op Const` with a btree operator. Compare { attnum: i16, strategy: i16, - typ: pg_sys::Oid, + parse_type: pg_sys::Oid, + cmpproc: pg_sys::Oid, collation: pg_sys::Oid, - const_text: String, + konst: pg_sys::Datum, }, /// `Var IS [NOT] NULL`: decided from the summary's null flags. Null { attnum: i16, is_null: bool }, /// `Var IN (const array)` / `= ANY`: prune iff no listed value falls in [min, max]. InList { attnum: i16, - typ: pg_sys::Oid, + parse_type: pg_sys::Oid, + cmpproc: pg_sys::Oid, collation: pg_sys::Oid, - elems: Vec>, + elems: Vec, }, /// `Var && Const` (overlap): prune iff the partition's extent does not overlap the - /// constant. Used for range types and PostGIS geometry. + /// constant, evaluated with the clause's own operator. Used for range types and + /// PostGIS geometry. Overlap { attnum: i16, - const_text: String, - const_type_name: String, + parse_type: pg_sys::Oid, + opfunc: pg_sys::Oid, + collation: pg_sys::Oid, + konst: pg_sys::Datum, + commuted: bool, }, } @@ -442,26 +472,42 @@ impl QualSpec { row.all_nulls } } - // Scalar comparisons only apply to btree min/max summaries. + // min/max order says nothing about another collation's order. QualSpec::Compare { strategy, - typ, + parse_type, + cmpproc, collation, - const_text, + konst, .. - } => !row.overlap && eval_compare(row, *strategy, *typ, *collation, const_text), + } => { + !row.overlap + && row.collation == *collation + && eval_compare(row, *strategy, *parse_type, *cmpproc, *collation, *konst) + } QualSpec::InList { - typ, + parse_type, + cmpproc, collation, elems, .. - } => !row.overlap && eval_in_list(row, *typ, *collation, elems), + } => { + !row.overlap + && row.collation == *collation + && eval_in_list(row, *parse_type, *cmpproc, *collation, elems) + } // Overlap only applies to extent summaries. QualSpec::Overlap { - const_text, - const_type_name, + parse_type, + opfunc, + collation, + konst, + commuted, .. - } => row.overlap && eval_overlap(row, const_text, const_type_name), + } => { + row.overlap + && eval_overlap(row, *parse_type, *opfunc, *collation, *konst, *commuted) + } } } } @@ -498,41 +544,43 @@ unsafe fn extract_opexpr(opexpr: *mut pg_sys::OpExpr) -> Option { if attnum <= 0 || (*con).constisnull { return None; } + let opno = (*opexpr).opno; // Scalar btree comparison (`<`, `<=`, `=`, `>=`, `>`). - if let Some(strategy) = strategy_cached((*opexpr).opno, (*var).vartype) { - let strategy = if commuted { - commute_strategy(strategy) - } else { - strategy - }; - if (1..=5).contains(&strategy) { - let const_text = datum_to_text((*con).consttype, (*con).constvalue)?; - return Some(QualSpec::Compare { - attnum, - strategy, - typ: (*var).vartype, - collation: (*opexpr).inputcollid, - const_text, - }); - } - return None; + if let Some(op) = compare_op_cached(opno, commuted) { + return Some(QualSpec::Compare { + attnum, + strategy: op.strategy, + parse_type: summary_parse_type(op.var_type, (*var).vartype), + cmpproc: op.cmpproc, + collation: (*opexpr).inputcollid, + konst: (*con).constvalue, + }); } - // Overlap operator (`&&`) for range types / PostGIS geometry. `&&` is commutative, - // so operand order does not matter. - if operator_name((*opexpr).opno).as_deref() == Some("&&") { - let const_text = datum_to_text((*con).consttype, (*con).constvalue)?; - let const_type_name = type_name((*con).consttype)?; + if is_trusted_overlap_op(opno, (*var).vartype) { return Some(QualSpec::Overlap { attnum, - const_text, - const_type_name, + parse_type: pg_sys::getBaseType((*var).vartype), + opfunc: pg_sys::get_opcode(opno), + collation: (*opexpr).inputcollid, + konst: (*con).constvalue, + commuted, }); } None } +/// Only an `&&` from pg_catalog or the column type's schema is trusted to mean "overlaps". +unsafe fn is_trusted_overlap_op(opno: pg_sys::Oid, vartype: pg_sys::Oid) -> bool { + if operator_name(opno).as_deref() != Some("&&") { + return false; + } + let ns = pg_sys::get_func_namespace(pg_sys::get_opcode(opno)); + ns == pg_sys::Oid::from(pg_sys::PG_CATALOG_NAMESPACE) + || Some(ns) == type_namespace(pg_sys::getBaseType(vartype)) +} + unsafe fn extract_nulltest(nulltest: *mut pg_sys::NullTest) -> Option { let arg = strip_relabel((*nulltest).arg as *mut pg_sys::Node); if !is_tag(arg, pg_sys::NodeTag::T_Var) { @@ -571,15 +619,16 @@ unsafe fn extract_saop(saop: *mut pg_sys::ScalarArrayOpExpr) -> Option return None; } // Only the equality operator gives the "any element in range" semantics. - if strategy_cached((*saop).opno, (*var).vartype)? != 3 { + let op = compare_op_cached((*saop).opno, false)?; + if op.strategy != 3 { return None; } - let elems = array_const_texts(con)?; Some(QualSpec::InList { attnum, - typ: (*var).vartype, + parse_type: summary_parse_type(op.var_type, (*var).vartype), + cmpproc: op.cmpproc, collation: (*saop).inputcollid, - elems, + elems: array_const_datums(con)?, }) } @@ -587,112 +636,79 @@ unsafe fn extract_saop(saop: *mut pg_sys::ScalarArrayOpExpr) -> Option // Typed in-memory evaluation // --------------------------------------------------------------------------------- +unsafe fn parsed_bounds( + row: &ColSummary, + parse_type: pg_sys::Oid, +) -> Option<(pg_sys::Datum, pg_sys::Datum)> { + match (&row.min, &row.max) { + (Some(min), Some(max)) if !row.unbounded => Some(( + parse_cached(parse_type, min)?, + parse_cached(parse_type, max)?, + )), + _ => None, // no usable range -> conservative KEEP + } +} + unsafe fn eval_compare( row: &ColSummary, strategy: i16, - typ: pg_sys::Oid, + parse_type: pg_sys::Oid, + cmpproc: pg_sys::Oid, collation: pg_sys::Oid, - const_text: &str, + k: pg_sys::Datum, ) -> bool { - let (min, max) = match (&row.min, &row.max) { - (Some(min), Some(max)) => (min, max), - _ => return false, // no usable range -> conservative KEEP - }; - let cmpproc = match cmp_proc_cached(typ) { - Some(p) => p, - None => return false, - }; - let k = match const_datum_cached(typ, const_text) { - Some(d) => d, - None => return false, - }; - let min_d = match parse_cached(typ, min) { - Some(d) => d, - None => return false, - }; - let max_d = match parse_cached(typ, max) { - Some(d) => d, - None => return false, + let Some((min_d, max_d)) = parsed_bounds(row, parse_type) else { + return false; }; - let cmp = |a, b| cmp_cached(cmpproc, collation, a, b); + let cmp = |bound| cmp_cached(cmpproc, collation, bound, k); match strategy { - 1 => cmp(min_d, k) >= 0, // col < K : prune iff min >= K - 2 => cmp(min_d, k) > 0, // col <= K : prune iff min > K - 3 => cmp(k, min_d) < 0 || cmp(k, max_d) > 0, // col = K : prune iff K outside [min,max] - 4 => cmp(max_d, k) < 0, // col >= K : prune iff max < K - 5 => cmp(max_d, k) <= 0, // col > K : prune iff max <= K + 1 => cmp(min_d) >= 0, // col < K : prune iff min >= K + 2 => cmp(min_d) > 0, // col <= K : prune iff min > K + 3 => cmp(min_d) > 0 || cmp(max_d) < 0, // col = K : prune iff K outside [min,max] + 4 => cmp(max_d) < 0, // col >= K : prune iff max < K + 5 => cmp(max_d) <= 0, // col > K : prune iff max <= K _ => false, } } unsafe fn eval_in_list( row: &ColSummary, - typ: pg_sys::Oid, + parse_type: pg_sys::Oid, + cmpproc: pg_sys::Oid, collation: pg_sys::Oid, - elems: &[Option], + elems: &[pg_sys::Datum], ) -> bool { - let (min, max) = match (&row.min, &row.max) { - (Some(min), Some(max)) => (min, max), - _ => return false, - }; - let cmpproc = match cmp_proc_cached(typ) { - Some(p) => p, - None => return false, - }; - let min_d = match parse_cached(typ, min) { - Some(d) => d, - None => return false, - }; - let max_d = match parse_cached(typ, max) { - Some(d) => d, - None => return false, + let Some((min_d, max_d)) = parsed_bounds(row, parse_type) else { + return false; }; - // Prune iff no listed value falls within [min, max]. A conversion failure on any - // element is conservative: treat it as "could match" -> KEEP. - for elem in elems { - let text = match elem { - Some(t) => t, - None => continue, // NULL element never matches a value - }; - let d = match const_datum_cached(typ, text) { - Some(d) => d, - None => return false, - }; - if cmp_cached(cmpproc, collation, d, min_d) >= 0 - && cmp_cached(cmpproc, collation, d, max_d) <= 0 - { - return false; // this value is in range -> KEEP - } - } - true + // Prune iff no listed value falls within [min, max]. + !elems.iter().any(|&d| { + cmp_cached(cmpproc, collation, min_d, d) <= 0 + && cmp_cached(cmpproc, collation, max_d, d) >= 0 + }) } /// `Var && Const`: prune iff the partition's stored extent does not overlap the -/// constant. The overlap test is delegated to PostgreSQL's own `&&` operator on the -/// column's type (range types, PostGIS geometry), so it works wherever that operator -/// is defined. A missing extent or any error is conservative (KEEP). -unsafe fn eval_overlap(row: &ColSummary, const_text: &str, const_type_name: &str) -> bool { +/// constant. Every row's value lies within the extent, so if the extent does not overlap +/// the constant, no row can. +unsafe fn eval_overlap( + row: &ColSummary, + parse_type: pg_sys::Oid, + opfunc: pg_sys::Oid, + collation: pg_sys::Oid, + k: pg_sys::Datum, + commuted: bool, +) -> bool { let extent = match &row.min { - Some(e) => e, - None => return false, + Some(e) if !row.unbounded => e, + _ => return false, }; - let type_name = &row.type_name; - if type_name.is_empty() { + let Some(ext) = parse_cached(parse_type, extent) else { return false; - } - let sql = format!( - "SELECT NOT (CAST({ext} AS {ext_t}) && CAST({k} AS {k_t}))", - ext = sql_literal(extent), - ext_t = type_name, - k = sql_literal(const_text), - k_t = const_type_name, - ); - Spi::get_one::(&sql).ok().flatten().unwrap_or(false) -} - -/// Minimal single-quote escaping for SQL string literals built in the planner hook. -fn sql_literal(s: &str) -> String { - format!("'{}'", s.replace('\'', "''")) + }; + let (a, b) = if commuted { (k, ext) } else { (ext, k) }; + let overlaps = pg_sys::FunctionCall2Coll(fmgr_cached(opfunc), collation, a, b); + overlaps.value() & 0xff == 0 } // --------------------------------------------------------------------------------- @@ -719,22 +735,16 @@ pub(crate) unsafe fn datum_cmp( pg_sys::OidFunctionCall2Coll(cmpproc, collation, a, b).value() as i32 } -/// Default btree "compare" support proc for a type, or `None` if unavailable. +/// Looked up under the opclass's input type, which may differ from `typ` (varchar -> text). pub(crate) unsafe fn btree_cmp_proc(typ: pg_sys::Oid) -> Option { let opclass = pg_sys::GetDefaultOpClass(typ, pg_sys::BTREE_AM_OID); if opclass == pg_sys::Oid::INVALID { return None; } let opfamily = pg_sys::get_opclass_family(opclass); - if opfamily == pg_sys::Oid::INVALID { - return None; - } - let proc = pg_sys::get_opfamily_proc(opfamily, typ, typ, BTORDER_PROC as i16); - if proc == pg_sys::Oid::INVALID { - None - } else { - Some(proc) - } + let intype = pg_sys::get_opclass_input_type(opclass); + let proc = pg_sys::get_opfamily_proc(opfamily, intype, intype, BTORDER_PROC as i16); + (proc != pg_sys::Oid::INVALID).then_some(proc) } /// Convert text to a Datum of `typ` via the type's input function. @@ -767,19 +777,21 @@ pub(crate) unsafe fn datum_to_text(typ: pg_sys::Oid, datum: pg_sys::Datum) -> Op text } -/// SQL type name for a type OID (e.g. "geometry", "int4range"). -unsafe fn type_name(typ: pg_sys::Oid) -> Option { - let ptr = pg_sys::format_type_be(typ); - if ptr.is_null() { +unsafe fn type_namespace(typ: pg_sys::Oid) -> Option { + let tup = pg_sys::SearchSysCache1( + pg_sys::SysCacheIdentifier::TYPEOID as i32, + pg_sys::Datum::from(typ), + ); + if tup.is_null() { return None; } - let name = CStr::from_ptr(ptr).to_str().ok().map(|s| s.to_owned()); - pg_sys::pfree(ptr as *mut _); - name + let ns = (*pg_sys::heap_tuple_get_struct::(tup)).typnamespace; + pg_sys::ReleaseSysCache(tup); + Some(ns) } -/// Decompose a non-null array Const into per-element text values. -unsafe fn array_const_texts(con: *mut pg_sys::Const) -> Option>> { +/// NULL elements are dropped: they never match. +unsafe fn array_const_datums(con: *mut pg_sys::Const) -> Option> { let elem_type = pg_sys::get_element_type((*con).consttype); if elem_type == pg_sys::Oid::INVALID { return None; @@ -807,15 +819,12 @@ unsafe fn array_const_texts(con: *mut pg_sys::Const) -> Option *mut pg_sys::Node { @@ -829,22 +838,13 @@ unsafe fn is_tag(node: *mut pg_sys::Node, tag: pg_sys::NodeTag) -> bool { !node.is_null() && (*node).type_ == tag } -/// Map an operator to a btree strategy number (1..5) over `lefttype`'s default family. -unsafe fn btree_strategy(opno: pg_sys::Oid, lefttype: pg_sys::Oid) -> Option { - let opclass = pg_sys::GetDefaultOpClass(lefttype, pg_sys::BTREE_AM_OID); +unsafe fn default_btree_family(typ: pg_sys::Oid) -> Option { + let opclass = pg_sys::GetDefaultOpClass(typ, pg_sys::BTREE_AM_OID); if opclass == pg_sys::Oid::INVALID { return None; } let opfamily = pg_sys::get_opclass_family(opclass); - if opfamily == pg_sys::Oid::INVALID { - return None; - } - let s = pg_sys::get_op_opfamily_strategy(opno, opfamily); - if s == 0 { - None - } else { - Some(s as i16) - } + (opfamily != pg_sys::Oid::INVALID).then_some(opfamily) } fn commute_strategy(strategy: i16) -> i16 { diff --git a/src/stable_io.rs b/src/stable_io.rs new file mode 100644 index 0000000..74241af --- /dev/null +++ b/src/stable_io.rs @@ -0,0 +1,49 @@ +use pgrx::pg_sys; + +// shortcut: summary bounds are stored as text, so these GUCs must be pinned around every +// write and read; drop this module when bounds are stored as binary datums. + +/// Run `f` with DateStyle, IntervalStyle and extra_float_digits pinned to values whose +/// text output round-trips exactly. `None` (f not run) in parallel mode: treat as unknown. +pub(crate) unsafe fn with_stable_io(f: impl FnOnce() -> T) -> Option { + // ISO output parses the same under any DateStyle field order. + let need_date = pg_sys::DateStyle != pg_sys::USE_ISO_DATES as i32; + let need_interval = pg_sys::IntervalStyle != pg_sys::INTSTYLE_POSTGRES as i32; + // Values < 1 drop float precision; any value >= 1 is shortest-exact. + let need_float = pg_sys::extra_float_digits < 1; + if !(need_date || need_interval || need_float) { + return Some(f()); + } + if pg_sys::IsInParallelMode() { + return None; + } + + let level = pg_sys::NewGUCNestLevel(); + if need_date { + set_guc(c"DateStyle", c"ISO"); + } + if need_interval { + set_guc(c"IntervalStyle", c"postgres"); + } + if need_float { + set_guc(c"extra_float_digits", c"1"); + } + Some( + pgrx::PgTryBuilder::new(std::panic::AssertUnwindSafe(f)) + .finally(|| pg_sys::AtEOXact_GUC(true, level)) + .execute(), + ) +} + +unsafe fn set_guc(name: &std::ffi::CStr, value: &std::ffi::CStr) { + pg_sys::set_config_option( + name.as_ptr(), + value.as_ptr(), + pg_sys::GucContext::PGC_USERSET, + pg_sys::GucSource::PGC_S_SESSION, + pg_sys::GucAction::GUC_ACTION_SAVE, + true, + 0, + false, + ); +} diff --git a/src/summary_build.rs b/src/summary_build.rs index b485d7b..c978bcc 100644 --- a/src/summary_build.rs +++ b/src/summary_build.rs @@ -1,183 +1,430 @@ +use crate::index_am::att_typid; use crate::index_storage::{ColSummary, IndexSummary}; +use crate::prune_hook::{btree_cmp_proc, datum_to_text, text_to_datum}; use pgrx::prelude::*; -use pgrx::spi::SpiError; - -// Computes a leaf partition's summary by scanning its data. Called by the index access -// method's `ambuild` (see `index_am.rs`); the result is persisted into the index's own -// metapage (see `index_storage.rs`) and maintained incrementally by `aminsert`. There is -// no side table — the summary lives in the index, like BRIN. - -fn oid_u32(oid: pg_sys::Oid) -> u32 { - oid.into() -} - -/// Build the summary for a single leaf relation's named columns by scanning its data. -/// Returns the typed summary that `ambuild` persists into the index's own metapage. -pub(crate) fn build_one_leaf( - leaf: pg_sys::Oid, - columns: &[String], -) -> Result { - let leaf_name = match relation_name(leaf)? { - Some(n) => n, - None => return Ok(IndexSummary::default()), +use pgrx::PgMemoryContexts; +use std::ffi::{CStr, CString}; + +// Index-build scan, not SQL: it sees recently-dead tuples and works during REINDEX. + +/// `heap_tuples` becomes the heap's `reltuples` via `index_build`. +pub(crate) struct LeafBuild { + pub summary: IndexSummary, + pub heap_tuples: f64, +} + +// shortcut: geometry extents go through SQL `ST_Extent` per batch; call PostGIS box +// functions directly if geometry builds or inserts become a bottleneck. +const GEOMETRY_BATCH: usize = 1000; + +enum Accumulator { + MinMax { + cmp: *mut pg_sys::FmgrInfo, + min: Option, + max: Option, + }, + Range { + merge: *mut pg_sys::FmgrInfo, + extent: Option, + }, + Geometry { + batch: Vec, + extent: Option, + }, + /// Stored as unbounded. + Unsupported, +} + +struct ColumnBuild { + /// Index into the build callback's `values`. + index_pos: usize, + attnum: i16, + typ: pg_sys::Oid, + type_name: String, + collation: pg_sys::Oid, + typlen: i16, + typbyval: bool, + nonnull: i64, + nulls: i64, + acc: Accumulator, +} + +struct BuildState { + cols: Vec, + cxt: pg_sys::MemoryContext, +} + +pub(crate) unsafe fn build_leaf_summary( + heap: pg_sys::Relation, + index: pg_sys::Relation, + index_info: *mut pg_sys::IndexInfo, +) -> LeafBuild { + let build_cxt = PgMemoryContexts::new("table_range build"); + let cxt = build_cxt.value(); + let mut state = BuildState { + cols: PgMemoryContexts::For(cxt).switch_to(|_| plan_columns(index, index_info)), + cxt, }; - let mut summary = IndexSummary::default(); - for col in columns { - // Resolve this leaf's attnum for the column by name. - let attnum = match leaf_attnum(leaf, col)? { - Some(a) => a, - None => continue, // column absent on this leaf; KEEP behavior + let scan = (*(*heap).rd_tableam) + .index_build_range_scan + .unwrap_or_else(|| error!("table_range: table access method cannot build indexes")); + let heap_tuples = scan( + heap, + index, + index_info, + true, // allow_sync + false, // anyvisible + true, // progress + 0, + pg_sys::InvalidBlockNumber, + Some(build_callback), + &mut state as *mut BuildState as *mut std::ffi::c_void, + std::ptr::null_mut(), + ); + + let summary = crate::stable_io::with_stable_io(|| finish(&mut state)) + .unwrap_or_else(|| error!("table_range: cannot build a summary in parallel mode")); + drop(build_cxt); + LeafBuild { + summary, + heap_tuples, + } +} + +unsafe fn plan_columns( + index: pg_sys::Relation, + index_info: *mut pg_sys::IndexInfo, +) -> Vec { + let nattrs = (*index_info).ii_NumIndexAttrs.max(0) as usize; + let mut cols = Vec::new(); + for index_pos in 0..nattrs { + let attnum = (*index_info).ii_IndexAttrNumbers[index_pos]; + if attnum <= 0 { + continue; // expression column: not summarized, so never used for pruning + } + let typ = att_typid((*index).rd_att, index_pos); + let base = pg_sys::getBaseType(typ); + let collation = if (*index).rd_indcollation.is_null() { + pg_sys::Oid::INVALID + } else { + *(*index).rd_indcollation.add(index_pos) }; + let mut typlen: i16 = 0; + let mut typbyval = false; + let mut typalign: std::ffi::c_char = 0; + pg_sys::get_typlenbyvalalign(typ, &mut typlen, &mut typbyval, &mut typalign); - let (kind, type_name) = column_kind(leaf, col)?; - let qcol = quote_ident(col); - let (e1, e2) = kind.exprs(&qcol); - - let stats_sql = - format!("SELECT count(*)::bigint, count({qcol})::bigint, {e1}, {e2} FROM {leaf_name}"); - let (total, nonnull, s1, s2) = Spi::connect(|client| { - let table = client.select(&stats_sql, Some(1), &[])?; - let res = match table.into_iter().next() { - Some(r) => ( - r.get::(1).ok().flatten().unwrap_or(0), - r.get::(2).ok().flatten().unwrap_or(0), - r.get::(3).ok().flatten(), - r.get::(4).ok().flatten(), - ), - None => (0, 0, None, None), - }; - Ok::<(i64, i64, Option, Option), SpiError>(res) - })?; - - let has_nulls = total > nonnull; - let all_nulls = total > 0 && nonnull == 0; - let overlap = kind.overlap(); - - summary.cols.push(ColSummary { + let acc = if typlen == -2 { + Accumulator::Unsupported // cstring-like values are not copied + } else if pg_sys::get_typtype(base) as u8 == b'r' { + Accumulator::Range { + merge: fmgr_in_current(pg_sys::Oid::from(pg_sys::F_RANGE_MERGE_ANYRANGE_ANYRANGE)), + extent: None, + } + } else if type_name_parts(base).is_some_and(|(_, name)| name == "geometry") { + Accumulator::Geometry { + batch: Vec::new(), + extent: None, + } + } else { + match btree_cmp_proc(base) { + Some(proc) => Accumulator::MinMax { + cmp: fmgr_in_current(proc), + min: None, + max: None, + }, + None => Accumulator::Unsupported, + } + }; + cols.push(ColumnBuild { + index_pos, attnum, - overlap, - type_name, - min: s1, - max: s2, - has_nulls, - all_nulls, + typ, + type_name: qualified_type_name(typ), + collation, + typlen, + typbyval, + nonnull: 0, + nulls: 0, + acc, }); } - Ok(summary) + cols } -/// How a column is summarized. -enum ColumnKind { - /// Default: btree min/max, used for scalar comparison pruning. - MinMax, - /// Range type: a covering range is stored for `&&` overlap pruning. - Range, - /// PostGIS geometry: the bounding extent is stored for `&&` overlap pruning. - Geometry, +#[pg_guard] +unsafe extern "C-unwind" fn build_callback( + _index: pg_sys::Relation, + _tid: pg_sys::ItemPointer, + values: *mut pg_sys::Datum, + isnull: *mut bool, + _tuple_is_alive: bool, + state: *mut std::ffi::c_void, +) { + // Recently-dead tuples are folded in too: an older snapshot may still see them. + let state = &mut *(state as *mut BuildState); + let cxt = state.cxt; + for col in state.cols.iter_mut() { + if *isnull.add(col.index_pos) { + col.nulls += 1; + continue; + } + col.nonnull += 1; + let value = *values.add(col.index_pos); + PgMemoryContexts::For(cxt).switch_to(|_| col.add(value)); + } } -impl ColumnKind { - fn overlap(&self) -> bool { - !matches!(self, ColumnKind::MinMax) +impl ColumnBuild { + /// `value` lives in the scan's per-tuple context, so anything kept is copied. + unsafe fn add(&mut self, value: pg_sys::Datum) { + let (typlen, typbyval, collation) = (self.typlen, self.typbyval, self.collation); + match &mut self.acc { + Accumulator::MinMax { cmp, min, max } => match (*min, *max) { + (Some(lo), Some(hi)) => { + if call_cmp(*cmp, collation, value, lo) < 0 { + free_datum(lo, typbyval); + *min = Some(copy_datum(value, typlen, typbyval)); + } + if call_cmp(*cmp, collation, value, hi) > 0 { + free_datum(hi, typbyval); + *max = Some(copy_datum(value, typlen, typbyval)); + } + } + _ => { + *min = Some(copy_datum(value, typlen, typbyval)); + *max = Some(copy_datum(value, typlen, typbyval)); + } + }, + Accumulator::Range { merge, extent } => { + *extent = Some(match *extent { + None => copy_datum(value, typlen, typbyval), + Some(old) => { + let merged = + pg_sys::FunctionCall2Coll(*merge, pg_sys::Oid::INVALID, old, value); + // range_merge may hand back one of its inputs unchanged. + if merged == value { + free_datum(old, typbyval); + copy_datum(value, typlen, typbyval) + } else { + if merged != old { + free_datum(old, typbyval); + } + merged + } + } + }); + } + Accumulator::Geometry { batch, .. } => { + batch.push(copy_datum(value, typlen, typbyval)); + if batch.len() >= GEOMETRY_BATCH { + self.flush_geometry(); + } + } + Accumulator::Unsupported => {} + } + } + + /// `ST_Extent` over an array parameter, so no table is planned during the build. + unsafe fn flush_geometry(&mut self) { + let typ = pg_sys::getBaseType(self.typ); + let (typlen, typbyval) = (self.typlen, self.typbyval); + let sql = format!( + "SELECT {ns}.st_extent(g)::{ty}::pg_catalog.text FROM pg_catalog.unnest($1) AS g", + ns = geometry_schema(&self.type_name), + ty = self.type_name + ); + let Accumulator::Geometry { batch, extent } = &mut self.acc else { + return; + }; + if let Some(text) = extent.as_deref() { + match text_to_datum(typ, text) { + Some(d) => batch.push(d), + None => error!("table_range: cannot re-read geometry extent"), + } + } + if batch.is_empty() { + return; + } + let mut typalign: std::ffi::c_char = 0; + let (mut l, mut b) = (0i16, false); + pg_sys::get_typlenbyvalalign(typ, &mut l, &mut b, &mut typalign); + let array = pg_sys::construct_array( + batch.as_mut_ptr(), + batch.len() as i32, + typ, + typlen as i32, + typbyval, + typalign, + ); + let mut argtype = pg_sys::get_array_type(typ); + let mut arg = pg_sys::Datum::from(array); + let cmd = CString::new(sql).expect("SQL contains no NUL bytes"); + + pg_sys::PushActiveSnapshot(pg_sys::GetTransactionSnapshot()); + pg_sys::SPI_connect(); + let rc = pg_sys::SPI_execute_with_args( + cmd.as_ptr(), + 1, + &mut argtype, + &mut arg, + std::ptr::null(), + true, + 1, + ); + if rc != pg_sys::SPI_OK_SELECT as i32 { + error!("table_range: ST_Extent batch failed ({rc})"); + } + let result = first_value(); + pg_sys::SPI_finish(); + pg_sys::PopActiveSnapshot(); + + for d in batch.drain(..) { + free_datum(d, typbyval); + } + pg_sys::pfree(array.cast()); + // NULL: every value was empty, so there is nothing to cover yet. + *extent = result; } - /// SQL expressions for (min/extent, max) text given the quoted column name. - fn exprs(&self, qcol: &str) -> (String, String) { - match self { - ColumnKind::MinMax => (format!("min({qcol})::text"), format!("max({qcol})::text")), - ColumnKind::Range => ( - format!("range_merge(range_agg({qcol}))::text"), - "NULL::text".to_string(), +} + +unsafe fn finish(state: &mut BuildState) -> IndexSummary { + let mut summary = IndexSummary::default(); + let cxt = state.cxt; + for col in state.cols.iter_mut() { + if matches!(col.acc, Accumulator::Geometry { .. }) { + PgMemoryContexts::For(cxt).switch_to(|_| col.flush_geometry()); + } + let (overlap, min, max, unbounded) = match &col.acc { + Accumulator::MinMax { min, max, .. } => ( + false, + min.and_then(|d| datum_to_text(col.typ, d)), + max.and_then(|d| datum_to_text(col.typ, d)), + false, ), - ColumnKind::Geometry => ( - format!("ST_Extent({qcol})::geometry::text"), - "NULL::text".to_string(), + Accumulator::Range { extent, .. } => ( + true, + extent.and_then(|d| datum_to_text(col.typ, d)), + None, + false, ), + Accumulator::Geometry { extent, .. } => (true, extent.clone(), None, false), + Accumulator::Unsupported => (false, None, None, true), + }; + let mut out = ColSummary { + attnum: col.attnum, + overlap, + type_name: col.type_name.clone(), + collation: col.collation, + min, + max, + has_nulls: col.nulls > 0, + all_nulls: col.nonnull == 0 && col.nulls > 0, + unbounded, + }; + // A bound that could not be rendered as UTF-8 is unknown, not absent. + let lost_bound = match &col.acc { + Accumulator::MinMax { .. } => out.min.is_none() || out.max.is_none(), + Accumulator::Range { .. } => out.min.is_none(), + _ => false, + }; + if out.unbounded || (col.nonnull > 0 && lost_bound) { + out.make_unbounded(); } + summary.cols.push(out); } + summary } -/// Classify a leaf column and resolve its SQL type name: range types and PostGIS -/// geometry get extent/overlap summaries; everything else gets btree min/max. -fn column_kind(leaf: pg_sys::Oid, col: &str) -> Result<(ColumnKind, String), SpiError> { - let row = Spi::connect(|client| { - let table = client.select( - &format!( - "SELECT a.atttypid::regtype::text AS typename, t.typtype::text \ - FROM pg_attribute a JOIN pg_type t ON t.oid = a.atttypid \ - WHERE a.attrelid = {}::oid AND a.attname = {} AND NOT a.attisdropped", - oid_u32(leaf), - quote_literal(col) - ), - Some(1), - &[], - )?; - let res = table.into_iter().next().map(|r| { - ( - r.get::(1).ok().flatten(), - r.get::(2).ok().flatten(), - ) - }); - Ok::, Option)>, SpiError>(res) - })?; +unsafe fn fmgr_in_current(proc: pg_sys::Oid) -> *mut pg_sys::FmgrInfo { + let finfo = pg_sys::palloc0(std::mem::size_of::()) as *mut pg_sys::FmgrInfo; + pg_sys::fmgr_info(proc, finfo); + finfo +} - let (type_name, typtype) = match row { - Some((Some(tn), tt)) => (tn, tt.unwrap_or_default()), - _ => return Ok((ColumnKind::MinMax, String::new())), - }; +unsafe fn call_cmp( + cmp: *mut pg_sys::FmgrInfo, + collation: pg_sys::Oid, + a: pg_sys::Datum, + b: pg_sys::Datum, +) -> i32 { + pg_sys::FunctionCall2Coll(cmp, collation, a, b).value() as i32 +} - let kind = if typtype == "r" { - ColumnKind::Range - } else if type_name == "geometry" || type_name.ends_with(".geometry") { - ColumnKind::Geometry +unsafe fn copy_datum(value: pg_sys::Datum, typlen: i16, typbyval: bool) -> pg_sys::Datum { + if typbyval { + value + } else if typlen == -1 { + pg_sys::Datum::from(pg_sys::pg_detoast_datum_copy(value.cast_mut_ptr())) } else { - ColumnKind::MinMax - }; - Ok((kind, type_name)) + let len = typlen as usize; + let copy = pg_sys::palloc(len); + std::ptr::copy_nonoverlapping(value.cast_mut_ptr::(), copy as *mut u8, len); + pg_sys::Datum::from(copy) + } } -/// Resolve a relation's column names for the given heap attnums (skips dropped/missing). -pub(crate) fn column_names_for_attnums( - relid: pg_sys::Oid, - attnums: &[i16], -) -> Result, SpiError> { - let mut names = Vec::new(); - for &attnum in attnums { - if attnum <= 0 { - continue; - } - let name = Spi::get_one::(&format!( - "SELECT attname::text FROM pg_attribute \ - WHERE attrelid = {}::oid AND attnum = {} AND NOT attisdropped", - oid_u32(relid), - attnum - ))?; - if let Some(name) = name { - names.push(name); - } +unsafe fn free_datum(value: pg_sys::Datum, typbyval: bool) { + if !typbyval { + pg_sys::pfree(value.cast_mut_ptr()); } - Ok(names) } -fn relation_name(relid: pg_sys::Oid) -> Result, SpiError> { - Spi::get_one::(&format!("SELECT {}::oid::regclass::text", oid_u32(relid))) +unsafe fn type_name_parts(typ: pg_sys::Oid) -> Option<(String, String)> { + let tup = pg_sys::SearchSysCache1( + pg_sys::SysCacheIdentifier::TYPEOID as i32, + pg_sys::Datum::from(typ), + ); + if tup.is_null() { + return None; + } + let form = pg_sys::heap_tuple_get_struct::(tup); + let name = CStr::from_ptr((*form).typname.data.as_ptr()) + .to_string_lossy() + .into_owned(); + let nsp = pg_sys::get_namespace_name((*form).typnamespace); + pg_sys::ReleaseSysCache(tup); + if nsp.is_null() { + return None; + } + Some((CStr::from_ptr(nsp).to_string_lossy().into_owned(), name)) } -fn leaf_attnum(leaf: pg_sys::Oid, col: &str) -> Result, SpiError> { - Spi::get_one::(&format!( - "SELECT attnum FROM pg_attribute \ - WHERE attrelid = {}::oid AND attname = {} AND attnum > 0 AND NOT attisdropped", - oid_u32(leaf), - quote_literal(col) - )) +unsafe fn qualified_type_name(typ: pg_sys::Oid) -> String { + let (nsp, name) = type_name_parts(typ) + .unwrap_or_else(|| error!("table_range: type {} not found", u32::from(typ))); + let nsp = CString::new(nsp).expect("no NUL in identifier"); + let name = CString::new(name).expect("no NUL in identifier"); + let q = pg_sys::quote_qualified_identifier(nsp.as_ptr(), name.as_ptr()); + CStr::from_ptr(q).to_string_lossy().into_owned() } -/// Minimal single-quote escaping for SQL string literals. -fn quote_literal(s: &str) -> String { - format!("'{}'", s.replace('\'', "''")) +/// PostGIS functions live in the same schema as its `geometry` type. +pub(crate) fn geometry_schema(type_name: &str) -> &str { + type_name.strip_suffix(".geometry").unwrap_or("public") } -/// Minimal identifier quoting (always double-quote, escaping embedded quotes). -fn quote_ident(s: &str) -> String { - format!("\"{}\"", s.replace('"', "\"\"")) +/// `None` for NULL, no row, or a value that is not valid UTF-8. +unsafe fn first_value() -> Option { + if pg_sys::SPI_processed == 0 || pg_sys::SPI_tuptable.is_null() { + return None; + } + let table = pg_sys::SPI_tuptable; + let p = pg_sys::SPI_getvalue(*(*table).vals, (*table).tupdesc, 1); + if p.is_null() { + None + } else { + CStr::from_ptr(p).to_str().ok().map(str::to_owned) + } +} + +/// Read-only, so SPI uses the active snapshot rather than the transaction snapshot. +pub(crate) unsafe fn spi_value(sql: &str) -> Option { + let cmd = CString::new(sql).expect("SQL contains no NUL bytes"); + let rc = pg_sys::SPI_execute(cmd.as_ptr(), true, 1); + if rc != pg_sys::SPI_OK_SELECT as i32 { + error!("table_range: SPI_execute failed ({rc}): {sql}"); + } + first_value() } diff --git a/src/summary_cache.rs b/src/summary_cache.rs index 1e27e70..3ac12a9 100644 --- a/src/summary_cache.rs +++ b/src/summary_cache.rs @@ -4,26 +4,18 @@ // // Correctness rests on the over-inclusive invariant: a cached summary is safe as long as // it is never *narrower* than the real data. Only `aminsert` widens a summary, and it -// calls `note_widened`, which sends a relcache invalidation for the index. That clears the -// cache in every backend (via `relcache_callback`) — locally at the next command boundary, -// and in other backends when the widening transaction commits, matching row visibility. -// Operations that only *narrow* a summary (deletes, vacuum re-tighten) need no -// invalidation, because an over-wide cached summary still prunes correctly. +// calls `note_widened`. Narrowing (deletes) needs no invalidation: an over-wide cached +// summary still prunes correctly. use crate::index_storage::ColSummary; use pgrx::pg_sys; -use std::cell::{Cell, RefCell}; -use std::collections::{HashMap, HashSet}; +use std::cell::RefCell; +use std::collections::HashMap; use std::rc::Rc; thread_local! { /// indexoid -> its deserialized summary (empty vec = "no summary / not a leaf"). static CACHE: RefCell>>> = RefCell::new(HashMap::new()); - /// Indexes already invalidated in the current transaction, to coalesce a bulk insert's - /// repeated widenings into one invalidation per index per transaction. - static INVALIDATED: RefCell> = RefCell::new(HashSet::new()); - /// TransactionId the `INVALIDATED` set belongs to; on change we clear the set. - static INVALIDATED_XID: Cell = const { Cell::new(0) }; } /// Look up a cached summary for `indexoid`, if present. @@ -40,28 +32,16 @@ pub fn put(indexoid: pg_sys::Oid, summary: Rc>) { }); } -/// Called by `aminsert` when it widens an index's on-page summary. Invalidates the cached -/// copy everywhere (once per index per transaction). -pub unsafe fn note_widened(indexoid: pg_sys::Oid) { - let key: u32 = indexoid.into(); - let xid: u32 = pg_sys::GetTopTransactionIdIfAny().into(); - if xid == 0 { - // No assigned xid to scope dedup to; invalidate unconditionally (safe). - pg_sys::CacheInvalidateRelcacheByRelid(indexoid); - return; - } - if INVALIDATED_XID.with(|x| x.get()) != xid { - INVALIDATED_XID.with(|x| x.set(xid)); - INVALIDATED.with(|s| s.borrow_mut().clear()); - } - let first = INVALIDATED.with(|s| s.borrow_mut().insert(key)); - if first { - pg_sys::CacheInvalidateRelcacheByRelid(indexoid); - } +/// The heap invalidation is what replans cached plans: they depend on tables, not indexes. +/// Not deduplicated: PostgreSQL coalesces per command, and coarser dedup loses invalidations. +pub unsafe fn note_widened(indexoid: pg_sys::Oid, heapoid: pg_sys::Oid) { + pg_sys::CacheInvalidateRelcacheByRelid(indexoid); + pg_sys::CacheInvalidateRelcacheByRelid(heapoid); } /// Relcache invalidation callback: drop the cached summary for `relid` (or all of them /// when `relid` is invalid, PG's "flush everything" signal). +#[pgrx::pg_guard] unsafe extern "C-unwind" fn relcache_callback(_arg: pg_sys::Datum, relid: pg_sys::Oid) { if relid == pg_sys::Oid::INVALID { CACHE.with(|c| c.borrow_mut().clear()); diff --git a/tests/multisession.sh b/tests/multisession.sh new file mode 100755 index 0000000..25c63e2 --- /dev/null +++ b/tests/multisession.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +# Tests that need more than one backend, which the single-transaction #[pg_test] +# harness cannot express: cross-backend invalidation, snapshots taken before an index +# build, and concurrent summary updates. +# +# Runs against a server that already has the extension installed. Connection settings +# come from the usual libpq environment (PGHOST, PGPORT, PGUSER); psql is taken from +# PATH. A scratch database is created and dropped. +# +# PGHOST=localhost PGPORT=28818 tests/multisession.sh + +set -euo pipefail + +DB=table_range_multisession +PSQL=(psql -X -q -v ON_ERROR_STOP=1 -At) +failures=0 + +psql -X -q -d postgres -c "SET client_min_messages = warning" \ + -c "DROP DATABASE IF EXISTS $DB" -c "CREATE DATABASE $DB" +trap 'psql -X -q -d postgres -c "DROP DATABASE IF EXISTS $DB" >/dev/null 2>&1 || true' EXIT + +q() { "${PSQL[@]}" -d "$DB" "$@"; } + +# Compare a count with pruning on and off; record a failure on mismatch. +check() { + local label=$1 from_where=$2 on off + on=$(q -c "SET table_range.enable_pruning = on" -c "SELECT count(*) FROM $from_where") + off=$(q -c "SET table_range.enable_pruning = off" -c "SELECT count(*) FROM $from_where") + if [[ "$on" == "$off" ]]; then + echo "ok $label (rows=$on)" + else + echo "FAIL $label: pruning on=$on off=$off" + failures=$((failures + 1)) + fi +} + +q -c "CREATE EXTENSION pg_table_range" + +# 1. A widening in a subtransaction that is rolled back, followed by another widening +# in the parent, must still reach other backends that cached the older summary. +q <<'SQL' +CREATE TABLE sp (k int, n bigint) PARTITION BY LIST (k); +CREATE TABLE sp_1 PARTITION OF sp FOR VALUES IN (1); +INSERT INTO sp VALUES (1, 1); +CREATE INDEX sp_tr ON sp USING table_range (n); +SQL +# A long-lived backend that caches the summary before, and queries after, the widening. +coproc WATCHER { "${PSQL[@]}" -d "$DB"; } +echo "SELECT count(*) FROM sp WHERE n = 600;" >&"${WATCHER[1]}" +read -r -u "${WATCHER[0]}" before +q <<'SQL' +BEGIN; +SAVEPOINT s; +INSERT INTO sp VALUES (1, 500); +ROLLBACK TO s; +INSERT INTO sp VALUES (1, 600); +COMMIT; +SQL +echo "SELECT count(*) FROM sp WHERE n = 600;" >&"${WATCHER[1]}" +read -r -u "${WATCHER[0]}" after +exec {WATCHER[1]}>&- +wait "$WATCHER_PID" || true +if [[ "$before" == 0 && "$after" == 1 ]]; then + echo "ok other backend sees widening after savepoint rollback" +else + echo "FAIL other backend after savepoint rollback: before=$before after=$after (want 0, 1)" + failures=$((failures + 1)) +fi + +# 2. CREATE INDEX inside a REPEATABLE READ transaction must cover rows committed after +# that transaction's snapshot was taken. +q <<'SQL' +CREATE TABLE rr (k int, n bigint) PARTITION BY LIST (k); +CREATE TABLE rr_1 PARTITION OF rr FOR VALUES IN (1); +INSERT INTO rr VALUES (1, 1); +SQL +coproc BUILDER { "${PSQL[@]}" -d "$DB"; } +printf '%s\n' "BEGIN ISOLATION LEVEL REPEATABLE READ;" "SELECT 1;" >&"${BUILDER[1]}" +read -r -u "${BUILDER[0]}" _ +q -c "INSERT INTO rr VALUES (1, 777)" +printf '%s\n' "CREATE INDEX rr_tr ON rr USING table_range (n);" "COMMIT;" "SELECT 'built';" >&"${BUILDER[1]}" +read -r -u "${BUILDER[0]}" _ +exec {BUILDER[1]}>&- +wait "$BUILDER_PID" || true +check "index built under REPEATABLE READ covers later commit" "rr WHERE n = 777" + +# 3. Concurrent widenings must not overwrite each other. A range type over a domain +# with a slow CHECK holds each insert between reading and writing the summary, so +# the writers overlap deterministically. +q <<'SQL' +CREATE FUNCTION slow(int) RETURNS bool LANGUAGE sql AS 'SELECT pg_sleep(0.3) IS NOT NULL'; +CREATE DOMAIN slowint AS int CHECK (slow(VALUE)); +CREATE TYPE slowrange AS RANGE (subtype = slowint); +DO $$ BEGIN PERFORM table_range_sync_opclasses(); END $$; +CREATE TABLE cc (k int, r slowrange) PARTITION BY LIST (k); +CREATE TABLE cc_1 PARTITION OF cc FOR VALUES IN (1); +INSERT INTO cc VALUES (1, '[0,1)'); +CREATE INDEX cc_tr ON cc USING table_range (r); +SQL +pids=() +for i in 1 2 3; do + q -c "INSERT INTO cc VALUES (1, '[$((i * 100)),$((i * 100 + 1)))')" & + pids+=($!) + q -c "INSERT INTO cc VALUES (1, '[-$((i * 100)),-$((i * 100 - 1)))')" & + pids+=($!) +done +wait "${pids[@]}" +for v in 100 -100 200 -200 300 -300; do + check "concurrent widening to $v kept" "cc WHERE r && '[$v,$((v + 1)))'::slowrange" +done + +if ((failures > 0)); then + echo "$failures multi-session check(s) failed" + exit 1 +fi +echo "all multi-session checks passed"