test(bench): cover GROUP BY / JOIN / DISTINCT on encrypted columns#203
test(bench): cover GROUP BY / JOIN / DISTINCT on encrypted columns#203coderdan wants to merge 3 commits into
Conversation
📝 WalkthroughWalkthroughAdds EXPLAIN-based plan assertions and ignored timing-regression benchmarks for GROUP BY, JOIN, and DISTINCT on ChangesBenchmark Tests for Encrypted Hash Operations
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the 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. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/sqlx/tests/bench_regression_tests.rs (1)
132-134: ⚡ Quick winAdd
benchfeature gating to these new timing tests nowThese 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 samebenchgate 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
📒 Files selected for processing (2)
tests/sqlx/tests/bench_plan_tests.rstests/sqlx/tests/bench_regression_tests.rs
| /// 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(()) |
There was a problem hiding this comment.
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.
| /// 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.
394aedd to
411ec59
Compare
…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.
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 thebenchfeature, pass today):group_by_encrypted_uses_hash_aggregate—GROUP BY encrypted_colpicksHashAggregatevia the hash op class.join_on_encrypted_uses_hmac_index— self-join on encrypted equality engagesbench_text_hmac_idx.distinct_encrypted_uses_hash_aggregate— unboundedDISTINCTpicksHashAggregate.group_by_jsonb_field_uses_hash_aggregate—GROUP BY eql_v2.jsonb_path_query_first(col, '<selector>')engagesHashAggregate(today asPartial HashAggregateunder 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):group_by_encrypted_under_thresholdself_join_encrypted_under_thresholddistinct_encrypted_under_thresholdgroup_by_jsonb_field_under_thresholdeql_v2.hmac_256(col, '<sel>'))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.sqlbuilds abench_jsontable by overlayinghmonto the$.hellosv element of each row in the existingbenchtable — mirroring what@cipherstash/protectproduces when the$.hellopath is configured with auniqueindex. Without that overlay, field-level GROUP BY raises today ("Cannot hash eql_v2_encrypted value: no hmac_256 index term found"). Reusesbenchrows 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
hmat 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
hash_encryptedfast-path on root-levelhm): closes most of the gap on root-level GROUP BY / JOIN / DISTINCT. Naive plpgsql →LANGUAGE sqlconversion alone doesn't move the needle —to_ste_vec_value's per-row JSONB inspection dominates. The fast-pathcoalesce(val.data ->> 'hm', hmac_256(to_ste_vec_value(val))::text)reads root-levelhmdirectly.hm-at-sv +eql_v2.hmac_256(val, selector)API overload): the canonical fix for field-level GROUP BY / DISTINCT / JOIN on encrypted JSON paths. Measured on the samebench_hm_in_svshape:eql_v2.jsonb_path_query_first(e, '<sel>')(today)eql_v2.hmac_256(e, '<sel>')(per #205)eql_v2.hmac_256(col, '<sel>'))Investigation also explored four alternative shapes that patch
eql_v2."="with ab3fallback (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 oneql_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.rscover 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 incipherstash/benches— flagged for follow-up there once we have credentials wired up locally.Test plan
cargo fmt --checkclean.--features bench(--test-threads=1when running locally; PG can OOM under fully-parallel test runs with concurrent bench fixture builds).--include-ignored; panic messages include current vs expected numbers.bench_*suites unaffected.Followup once #205 lands
The
group_by_jsonb_field_*tests should be updated to use the canonicaleql_v2.hmac_256(col, '<selector>')form rather than the legacyeql_v2.jsonb_path_query_first. The threshold drops from 50 ms to ~40 ms. The#[ignore]marker comes off.