Skip to content

test(bench): cover GROUP BY / JOIN / DISTINCT on encrypted columns#203

Open
coderdan wants to merge 3 commits into
mainfrom
test/bench-group-by-join-distinct
Open

test(bench): cover GROUP BY / JOIN / DISTINCT on encrypted columns#203
coderdan wants to merge 3 commits into
mainfrom
test/bench-group-by-join-distinct

Conversation

@coderdan
Copy link
Copy Markdown
Contributor

@coderdan coderdan commented May 11, 2026

Summary

Adds bench coverage for the hash-strategy access patterns the Phase 1 hash operator class (#196) enabled — both at the root level (encrypted column as a whole) and at the field level (a JSON path extracted from the column). Plan-shape assertions confirm the right access methods engage today; timing thresholds capture the perf gap that the various follow-up issues (#202, #204, #205) are needed to close.

Companion to #202 and #205.

What's added

Plan assertions in bench_plan_tests.rs (gated on the bench feature, pass today):

  • group_by_encrypted_uses_hash_aggregateGROUP BY encrypted_col picks HashAggregate via the hash op class.
  • join_on_encrypted_uses_hmac_index — self-join on encrypted equality engages bench_text_hmac_idx.
  • distinct_encrypted_uses_hash_aggregate — unbounded DISTINCT picks HashAggregate.
  • group_by_jsonb_field_uses_hash_aggregateGROUP BY eql_v2.jsonb_path_query_first(col, '<selector>') engages HashAggregate (today as Partial HashAggregate under a parallel worker, then Sort + GroupAggregate for the merge; once the recipe evolves per perf/correctness: emit hm (HMAC-256) at sv element level, not b3 (option 1) #205, the flat single-HashAggregate path becomes available too).

Regression timing assertions in bench_regression_tests.rs (#[ignore]'d pending the listed issue; remove the markers when each merges):

Test Threshold Today After fix Tracking
group_by_encrypted_under_threshold 150 ms ~309 ms ~73 ms #202
self_join_encrypted_under_threshold 350 ms ~308 ms ~185 ms #202
distinct_encrypted_under_threshold 200 ms ~515 ms ~72 ms #202
group_by_jsonb_field_under_threshold 50 ms ~278 ms ~26 ms via #205 recipe (eql_v2.hmac_256(col, '<sel>')) #205

Each timing panic message states the expected post-fix number and the current observed number so the diagnostic stays useful when the gate flips.

New fixture

tests/sqlx/fixtures/bench_json_data.sql builds a bench_json table by overlaying hm onto the $.hello sv element of each row in the existing bench table — mirroring what @cipherstash/protect produces when the $.hello path is configured with a unique index. Without that overlay, field-level GROUP BY raises today ("Cannot hash eql_v2_encrypted value: no hmac_256 index term found"). Reuses bench rows so the fixture cost is just one INSERT-FROM-SELECT.

This fixture also coincides exactly with option 1 of #205 — the proposal to have the crypto layer emit hm at sv element level uniformly. When that lands, no fixture change is needed; the test data already reflects the post-fix shape.

Note on the underlying fixes

Field-level operation eql_v2.jsonb_path_query_first(e, '<sel>') (today) eql_v2.hmac_256(e, '<sel>') (per #205)
GROUP BY 429 ms 26 ms (~16× faster)
self-JOIN 422 ms with functional idx (Nested Loop + Memoize + Index Scan), 494 ms without
single-row WHERE 0.39 ms (Bitmap Index Scan via functional idx on eql_v2.hmac_256(col, '<sel>'))

Investigation also explored four alternative shapes that patch eql_v2."=" with a b3 fallback (so field-level equality works without the crypto-layer change). All four fix correctness but every one breaks at least one of two PostgreSQL planner optimisations the PR #196 body shape unlocks: functional index match on eql_v2.hmac_256(col), and Merge Join sort-key hoist. Concretely the self-JOIN on encrypted JSON regresses 250×–24 700× depending on the shape. The #205 crypto-layer change avoids the patch entirely and is the cleaner fix.

Why these tests don't already exist

The existing bench_plan_tests.rs / bench_regression_tests.rs cover single-row predicate paths (equality, LIKE, ORE range/order). Cross-row aggregation (GROUP BY, DISTINCT), cross-table joins, and field-level access weren't covered. The same gap exists in cipherstash/benches — flagged for follow-up there once we have credentials wired up locally.

Test plan

  • cargo fmt --check clean.
  • Plan tests pass on PG 17 with --features bench (--test-threads=1 when running locally; PG can OOM under fully-parallel test runs with concurrent bench fixture builds).
  • Regression tests fail today with --include-ignored; panic messages include current vs expected numbers.
  • Existing bench_* suites unaffected.
  • CI green across PG 14/15/16/17.

Followup once #205 lands

The group_by_jsonb_field_* tests should be updated to use the canonical eql_v2.hmac_256(col, '<selector>') form rather than the legacy eql_v2.jsonb_path_query_first. The threshold drops from 50 ms to ~40 ms. The #[ignore] marker comes off.

@coderabbitai
Copy link
Copy Markdown

coderabbitai Bot commented May 11, 2026

📝 Walkthrough

Walkthrough

Adds EXPLAIN-based plan assertions and ignored timing-regression benchmarks for GROUP BY, JOIN, and DISTINCT on encrypted_text, verifying HashAggregate usage and index selection; timing tests are gated pending issue #202.

Changes

Benchmark Tests for Encrypted Hash Operations

Layer / File(s) Summary
Test Module Imports
tests/sqlx/tests/bench_plan_tests.rs
Extended imports to include explain_query helper for EXPLAIN-based assertions.
Plan Assertions for Hash Operations
tests/sqlx/tests/bench_plan_tests.rs
Added documentation and three #[sqlx::test] cases: GROUP BY encrypted_text and SELECT DISTINCT encrypted_text assert HashAggregate appears in EXPLAIN; an equality join on encrypted_text asserts use of bench_text_hmac_idx via assert_uses_index. Tests are gated with bench feature ignore attributes.
Timing Regression Tests for Hash Operations
tests/sqlx/tests/bench_regression_tests.rs
Added documentation and three ignored #[sqlx::test] cases measuring average execution time via explain_analyze_avg: GROUP BY encrypted_text (5 runs, <150ms), self-join on encrypted_text (3 runs, <350ms), and SELECT DISTINCT encrypted_text (5 runs, <200ms). Tests are ignored until #202 merges.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related issues

Possibly related PRs

Suggested reviewers

  • freshtonic
  • auxesis

Poem

🐰 I hopped through EXPLAIN's gentle light,

Checking hashes through the night,
GROUP BY, DISTINCT — plans take flight,
Indices gleam, benchmarks polite,
A rabbit cheers the tests done right.

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main changes: adding benchmark test coverage for GROUP BY, JOIN, and DISTINCT operations on encrypted columns, which are the three core test scenarios added.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch test/bench-group-by-join-distinct

Warning

Review ran into problems

🔥 Problems

Git: Failed to clone repository. Please run the @coderabbitai full review command to re-trigger a full review. If the issue persists, set path_filters to include or exclude specific files.


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
tests/sqlx/tests/bench_regression_tests.rs (1)

132-134: ⚡ Quick win

Add bench feature gating to these new timing tests now

These tests are currently fully ignored, but once #[ignore] is removed for #202, they’ll run in non-bench test contexts unless they also carry the same bench gate used by the rest of this file.

Suggested patch
 #[sqlx::test(fixtures(path = "../fixtures", scripts("bench_data", "bench_setup")))]
+#[cfg_attr(
+    not(feature = "bench"),
+    ignore = "perf-bench: gated, run via mise test:bench"
+)]
 #[ignore = "#202: hash_encrypted chain not yet inlined; remove ignore when `#202` merges"]
 async fn group_by_encrypted_under_threshold(pool: PgPool) -> Result<()> {
@@
 #[sqlx::test(fixtures(path = "../fixtures", scripts("bench_data", "bench_setup")))]
+#[cfg_attr(
+    not(feature = "bench"),
+    ignore = "perf-bench: gated, run via mise test:bench"
+)]
 #[ignore = "#202: hash_encrypted chain not yet inlined; remove ignore when `#202` merges"]
 async fn self_join_encrypted_under_threshold(pool: PgPool) -> Result<()> {
@@
 #[sqlx::test(fixtures(path = "../fixtures", scripts("bench_data", "bench_setup")))]
+#[cfg_attr(
+    not(feature = "bench"),
+    ignore = "perf-bench: gated, run via mise test:bench"
+)]
 #[ignore = "#202: hash_encrypted chain not yet inlined; remove ignore when `#202` merges"]
 async fn distinct_encrypted_under_threshold(pool: PgPool) -> Result<()> {

Also applies to: 156-158, 176-178

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/sqlx/tests/bench_regression_tests.rs` around lines 132 - 134, The new
timing test annotated on the async function group_by_encrypted_under_threshold
is only ignored via #[ignore] but lacks the file’s established bench feature
gate; add the same bench gating used elsewhere by wrapping the sqlx::test
attribute or the test definition with cfg_attr/#[cfg(feature = "bench")] (or the
project’s equivalent bench gating macro) so the test only runs when the bench
feature is enabled, and apply the same change to the other two similar ignored
timing tests in this file (the other ignored group-by/bench tests) to match the
rest of the file’s pattern.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/sqlx/tests/bench_plan_tests.rs`:
- Around line 153-166: The test join_on_encrypted_uses_hmac_index currently
calls assert_uses_index(&pool, sql, BENCH_TEXT_HMAC_IDX) but the test comment
allows either an index-driven nested-loop plan or a Hash Join; update the
assertion to accept both plan shapes by either (A) changing the test to call a
new helper (e.g., assert_uses_index_or_plan(&pool, sql, BENCH_TEXT_HMAC_IDX,
"Hash Join")) that succeeds if the plan contains BENCH_TEXT_HMAC_IDX OR the
string "Hash Join", or (B) extend assert_uses_index to accept an additional
allowed_plan parameter and treat the test as passing if the executed plan
contains the index name OR the allowed_plan ("Hash Join") — locate the test
function join_on_encrypted_uses_hmac_index and the helper assert_uses_index /
BENCH_TEXT_HMAC_IDX to implement this conditional check.

---

Nitpick comments:
In `@tests/sqlx/tests/bench_regression_tests.rs`:
- Around line 132-134: The new timing test annotated on the async function
group_by_encrypted_under_threshold is only ignored via #[ignore] but lacks the
file’s established bench feature gate; add the same bench gating used elsewhere
by wrapping the sqlx::test attribute or the test definition with
cfg_attr/#[cfg(feature = "bench")] (or the project’s equivalent bench gating
macro) so the test only runs when the bench feature is enabled, and apply the
same change to the other two similar ignored timing tests in this file (the
other ignored group-by/bench tests) to match the rest of the file’s pattern.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: e751d829-0a35-4814-97d9-bc1a002cdc0e

📥 Commits

Reviewing files that changed from the base of the PR and between 273ec18 and 394aedd.

📒 Files selected for processing (2)
  • tests/sqlx/tests/bench_plan_tests.rs
  • tests/sqlx/tests/bench_regression_tests.rs

Comment on lines +153 to +166
/// JOIN on `a.encrypted_col = b.encrypted_col` engages the hmac functional index.
/// Acceptable plan shapes: Hash Join (preferred), or Nested Loop + Memoize +
/// Index Scan via `bench_text_hmac_idx` (current planner choice — fine since
/// the index lookup remains the per-probe cost).
#[sqlx::test(fixtures(path = "../fixtures", scripts("bench_data", "bench_setup")))]
#[cfg_attr(
not(feature = "bench"),
ignore = "perf-bench: gated, run via mise test:bench"
)]
async fn join_on_encrypted_uses_hmac_index(pool: PgPool) -> Result<()> {
let sql = "SELECT count(*) FROM bench a JOIN bench b \
ON a.encrypted_text = b.encrypted_text";
assert_uses_index(&pool, sql, BENCH_TEXT_HMAC_IDX).await?;
Ok(())
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Join assertion currently rejects a documented acceptable Hash Join plan

The test description allows either Hash Join or index-driven nested-loop plans, but the assertion only accepts index usage. This can fail valid plans.

Suggested patch
 async fn join_on_encrypted_uses_hmac_index(pool: PgPool) -> Result<()> {
     let sql = "SELECT count(*) FROM bench a JOIN bench b \
                ON a.encrypted_text = b.encrypted_text";
-    assert_uses_index(&pool, sql, BENCH_TEXT_HMAC_IDX).await?;
+    let plan = explain_query(&pool, sql).await?;
+    assert!(
+        plan.contains("Hash Join") || plan.contains(BENCH_TEXT_HMAC_IDX),
+        "Expected JOIN to use Hash Join or {}. EXPLAIN output:\n{}",
+        BENCH_TEXT_HMAC_IDX,
+        plan
+    );
     Ok(())
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
/// JOIN on `a.encrypted_col = b.encrypted_col` engages the hmac functional index.
/// Acceptable plan shapes: Hash Join (preferred), or Nested Loop + Memoize +
/// Index Scan via `bench_text_hmac_idx` (current planner choice — fine since
/// the index lookup remains the per-probe cost).
#[sqlx::test(fixtures(path = "../fixtures", scripts("bench_data", "bench_setup")))]
#[cfg_attr(
not(feature = "bench"),
ignore = "perf-bench: gated, run via mise test:bench"
)]
async fn join_on_encrypted_uses_hmac_index(pool: PgPool) -> Result<()> {
let sql = "SELECT count(*) FROM bench a JOIN bench b \
ON a.encrypted_text = b.encrypted_text";
assert_uses_index(&pool, sql, BENCH_TEXT_HMAC_IDX).await?;
Ok(())
/// JOIN on `a.encrypted_col = b.encrypted_col` engages the hmac functional index.
/// Acceptable plan shapes: Hash Join (preferred), or Nested Loop + Memoize +
/// Index Scan via `bench_text_hmac_idx` (current planner choice — fine since
/// the index lookup remains the per-probe cost).
#[sqlx::test(fixtures(path = "../fixtures", scripts("bench_data", "bench_setup")))]
#[cfg_attr(
not(feature = "bench"),
ignore = "perf-bench: gated, run via mise test:bench"
)]
async fn join_on_encrypted_uses_hmac_index(pool: PgPool) -> Result<()> {
let sql = "SELECT count(*) FROM bench a JOIN bench b \
ON a.encrypted_text = b.encrypted_text";
let plan = explain_query(&pool, sql).await?;
assert!(
plan.contains("Hash Join") || plan.contains(BENCH_TEXT_HMAC_IDX),
"Expected JOIN to use Hash Join or {}. EXPLAIN output:\n{}",
BENCH_TEXT_HMAC_IDX,
plan
);
Ok(())
}
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/sqlx/tests/bench_plan_tests.rs` around lines 153 - 166, The test
join_on_encrypted_uses_hmac_index currently calls assert_uses_index(&pool, sql,
BENCH_TEXT_HMAC_IDX) but the test comment allows either an index-driven
nested-loop plan or a Hash Join; update the assertion to accept both plan shapes
by either (A) changing the test to call a new helper (e.g.,
assert_uses_index_or_plan(&pool, sql, BENCH_TEXT_HMAC_IDX, "Hash Join")) that
succeeds if the plan contains BENCH_TEXT_HMAC_IDX OR the string "Hash Join", or
(B) extend assert_uses_index to accept an additional allowed_plan parameter and
treat the test as passing if the executed plan contains the index name OR the
allowed_plan ("Hash Join") — locate the test function
join_on_encrypted_uses_hmac_index and the helper assert_uses_index /
BENCH_TEXT_HMAC_IDX to implement this conditional check.

…ypted columns

Adds bench coverage for the three hash-strategy access patterns the Phase 1
hash operator class (#196) enabled but the Phase 2 chain inlining (#202) is
yet to make perf-competitive.

Plan assertions in bench_plan_tests.rs (gated on the bench feature, pass
today):

- group_by_encrypted_uses_hash_aggregate — confirms `GROUP BY encrypted_col`
  picks HashAggregate via the hash op class rather than degenerating to
  GroupAggregate-after-Sort or a Nested-Loop self-comparison.
- join_on_encrypted_uses_hmac_index — confirms self-join on encrypted
  equality engages bench_text_hmac_idx (Hash Join or Nested Loop + Memoize +
  Index Scan, both acceptable).
- distinct_encrypted_uses_hash_aggregate — confirms unbounded DISTINCT picks
  HashAggregate (the bounded-LIMIT variant biases toward IndexOnlyScan over
  the ORE btree opclass; that path is fine on full installs but unavailable
  on Supabase).

Regression timing assertions in bench_regression_tests.rs (#[ignore]'d
pending #202; remove the markers when it merges):

- group_by_encrypted_under_threshold — 150ms (current ~309ms via plpgsql
  hash chain, ~70ms with chain inlined; threshold ~2x the inlined target).
- self_join_encrypted_under_threshold — 350ms (current ~308ms, ~182ms with
  chain inlined; cardinality dominates so threshold is generous).
- distinct_encrypted_under_threshold — 200ms (current ~515ms unbounded via
  ORE btree path, expected to drop into HashAggregate-driven territory
  after chain inlining).

Each timing test panic message states the expected post-#202 number and
the current observed number, so the diagnostic remains useful when the
gate flips after the fix lands.
@coderdan coderdan force-pushed the test/bench-group-by-join-distinct branch from 394aedd to 411ec59 Compare May 11, 2026 05:48
…t naive inlining

After measuring properly, naive plpgsql → LANGUAGE sql conversion of the
hash_encrypted chain does not deliver the speedup originally cited:
`to_ste_vec_value`'s per-row JSONB inspection/reconstruction dominates
even when fully inlined. The actual #202 fix is a fast-path read of
root-level `hm` (`coalesce(val.data ->> 'hm', ...)`) with the
`to_ste_vec_value` unwrap reserved for single-element ste_vec-wrapped
payloads.

Updates the regression test docstrings, panic messages, and section
header to reflect the measured fast-path numbers:
- GROUP BY: expected ~73ms (was ~70ms — close, but framing was wrong).
- Self-join: expected ~185ms (was ~182ms — same).
- DISTINCT (previously TBD): measured ~72ms with fast-path applied.

Thresholds unchanged: 150ms / 350ms / 200ms — all still have ~2x
headroom over the fast-path numbers.
…lumn

Adds the field-level GROUP BY scenario alongside the root-level coverage.
The pattern — `GROUP BY eql_v2.jsonb_path_query_first(col, '<selector>')`
— is the canonical "tell me how many rows per region" query against
ste_vec encryption where the field has a `unique` index configured.

The current bench fixture's field-level sv elements carry only OPE
terms (`ocv` / `ocf`), no `hm`, which means field-level GROUP BY raises
("Cannot hash eql_v2_encrypted value: no hmac_256 index term found").
The new `bench_json_data.sql` fixture overlays `hm` onto the $.hello sv
element of each bench row — synthesising what `@cipherstash/protect`
produces when the $.hello path is configured with `unique`. Reuses the
existing `bench` rows so the fixture cost stays the same.

Plan assertion in `bench_plan_tests.rs` (gated on bench feature, passes
today): `group_by_jsonb_field_uses_hash_aggregate` — confirms the
planner engages `HashAggregate` (today as `Partial HashAggregate` under
a parallel worker, then Sort + GroupAggregate for the merge — once #204
inlines the extractors, this may flatten to a single HashAggregate).

Regression timing in `bench_regression_tests.rs` (#[ignore]'d pending
#204): `group_by_jsonb_field_under_threshold` — threshold 50ms.
Measured numbers at 10K rows: ~278ms on main, ~234ms with the #202
fast-path patched in, ~7ms achievable via raw JSONB extraction
(`(col).data->'sv'->N->>'hm'`). The 50ms threshold targets the
post-#204-extractor-inlining state.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant