Skip to content

test: cover the memory accounting paths of the grouped aggregate streams that no test reaches - #26562

Open
jayzhan211 wants to merge 1 commit into
apache:mainfrom
jayzhan211:memacct-call-sequence-tests
Open

jayzhan211 wants to merge 1 commit into
apache:mainfrom
jayzhan211:memacct-call-sequence-tests

Conversation

@jayzhan211

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

  • None. This adds tests only; it prepares a follow-up refactor of the aggregate streams' memory accounting.

Rationale for this change

The grouped aggregate streams decide how much memory to reserve in several places: when they spill, when they replay their spill files, when they hand out a materialized batch in slices, and when the pool refuses a request. To find out which of these decisions the existing tests protect, I broke them one at a time on main and ran the physical-plan tests, the memory_limit and aggregate fuzz integration tests, and the sqllogictests. Of 26 such changes, 14 broke no test:

  • Final, Single and ClusteredFinal: dropping the 4 bytes per group reserved for the sort indices of a spill.
  • Final: no longer shrinking the reservation after a spill, or no longer releasing it before the replay.
  • ClusteredSingle: no longer releasing the reservation before the replay.
  • Single with temporary files disabled: ignoring a refused reservation instead of returning the error.
  • Final: registering as spillable even when temporary files are disabled.
  • Final returning, and Single ignoring, a refusal while materializing the final output (each reacts the other way on main).
  • ClusteredPartial: returning an error instead of emitting the completed groups whole when they do not fit next to the table.
  • ClusteredPartial with full clustering: emitting early instead of returning the refusal.
  • ClusteredFinal and ClusteredSingle: no longer reserving the completed groups while they are handed out.

A pool refusal that is wrongly ignored or wrongly returned is a query that runs past its memory limit or fails when it should succeed, so these paths deserve tests before anyone changes them.

What changes are included in this PR?

One test file, aggregates/memory_call_sequence_tests.rs, with 11 tests. Each runs one AggregateExec under a pool that records every call made on behalf of the aggregate's consumers (register with its can_spill, grow, try_grow and whether it was refused, shrink, unregister), and snapshots that sequence together with the query result. Each test also asserts through metrics or output batch sizes that it really took its path (it spilled, or emitted whole batches, and so on).

Refused output materialization needs an aggregate whose output is larger than the memory it accounts for, which no built-in aggregate has, so the file defines a small test UDAF for it.

The early emit of the unclustered Partial and PartialReduce streams is left out: their own unit tests already catch every change I made there.

What is the testing strategy for this PR?

The PR is the tests. Each of the 14 changes above makes at least one of them fail, and all of them pass on main.

Are there any user-facing changes?

No.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant