[improvement](agg) bound streaming pre-aggregation memory by the query memory limit

[mrhhsg]

What problem does this PR solve?

Issue Number: None

Problem Summary:

Streaming pre-aggregation (StreamingAggOperatorX) only limited its hash table
memory when spilling was enabled, using the fixed session variable
spill_streaming_agg_mem_limit per task; without spilling there was no memory
gate at all, so a first-phase agg with high cardinality could grow its hash
tables until the query hit its memory limit. DistinctStreamingAggOperatorX
had no memory gate at all, only the low_memory_mode() switch.

This PR derives the per-task pre-aggregation budget from the live query memory
limit instead (streaming_agg_memory_limit() in
exec/operator/streaming_agg_memory_limit.h):

• budget = query mem_limit / parallel_tasks / 5, re-read on every block so a
  limit lowered or restored by the workload group manager takes effect at once;
• the budget never drops below twice the last cache tier of the min-reduction
  table (32MB) so a small query limit does not disable pre-aggregation
  altogether, but the floor is capped by the per-task share of the query limit;
• when spilling is enabled, spill_streaming_agg_mem_limit stays an explicit
  upper bound applied on top of that budget (0 = no explicit bound); it is not
  consulted otherwise.

Both operators use the same helper, so the distinct pre-agg now honours the
same rule. The low-memory switch of StreamingAggOperatorX is turned into an
std::atomic_bool (it was a plain size_t written by any pipeline task while
others read it), _spill_streaming_agg_mem_limit gets an initializer (the mock
operator in unit tests never calls init()), and the effective limit is
exposed in the profile as MemoryUseLimit.

Follow-ups from the first review round (all of them are pre-existing on the
pass-through paths, which this PR makes reachable under memory pressure):

• A LIMIT pushed down to the local distinct stage (visitPhysicalLimit) was
  applied to raw pass-through rows, so duplicates consumed the allowance and the
  global stage could receive too few distinct keys. Both operators now skip the
  local truncation / reached_limit() for pass-through rows, and the distinct
  operator never enters memory-triggered pass-through while such a limit is
  active (the limit already bounds its hash set).
• The TopN pre-agg state stays undecided (need_do_sort_limit == -1) while the
  hash table is smaller than the limit, so the heap can still be built after a
  transient pass-through.
• Pass-through serialization uses a per-block arena instead of the
  operator-lifetime _agg_arena_pool (collect/map/foreach/... allocate
  from it per row).
• Low-memory mode only ever tightens the budget (min(1MB, budget)).
• The helper never collapses a positive query limit to the "no cap" sentinel.
• spill_streaming_agg_mem_limit is forwarded to the planning master
  (needForward = true), like enable_spill.

Release note

Streaming pre-aggregation memory is now bounded by one fifth of the query
memory limit per operator (shared among its parallel tasks, floor 32MB per
task); with enable_spill, spill_streaming_agg_mem_limit additionally caps
that budget.

Check List (For Author)

• Test

  • Unit Test
  • Regression test
  • Manual test
  • No need to test or manual test

• Behavior changed:

  • Yes.
    • non-spill queries: the pre-agg hash table is now capped at
      mem_limit / parallel_tasks / 5 (min 32MB) per task instead of
      unbounded.
    • spill-enabled queries: the cap is min(spill_streaming_agg_mem_limit, dynamic budget) instead of the fixed session value.

• Does this need documentation?

  • Yes. The streaming pre-agg memory budget and the MemoryUseLimit
    profile counter.

Validation

./run-be-ut.sh -j 160 --run --filter='StreamingAggMemoryLimitTest.*:StreamingAggOperatorTest.*:DistinctStreamingAggOperatorTest.*'
[  PASSED  ] 14 tests.

• StreamingAggMemoryLimitTest.budget_floor_and_fixed_limit (new) covers the 1/5 rule, the
  32MB floor, the floor capped by the per-task share, the explicit bound, and unknown limits.
• StreamingAggOperatorTest.test1 / DistinctStreamingAggOperatorTest.refresh_memory_limit
  check the effective limit reported in MemoryUseLimit while the query limit changes, that
  spill_streaming_agg_mem_limit only applies with enable_spill, that low-memory mode only
  tightens, and that the distinct pre-agg passes rows through once the limit is exceeded.
• StreamingAggOperatorTest.memory_limit_pass_through_and_recover drives the production
  _should_not_do_pre_agg() decision: pass-through under a lowered limit (no hash-table
  growth, duplicates not counted against a pushed-down limit), aggregation resumed after the
  limit is restored.
• DistinctStreamingAggOperatorTest.pushed_limit_keeps_deduplicating /
  pass_through_does_not_consume_pushed_limit cover the pushed-down LIMIT interaction.
• Not covered by a unit test: the TopN pre-agg state after a transient pass-through (no TopN
  fixture exists for this operator); the change there is confined to not latching
  need_do_sort_limit = 0.
• No regression test: the behaviour depends on the query memory limit and the hash-table size
  and cannot be asserted deterministically from SQL output.

Check List (For Reviewer who merge this PR)

• Confirm the release note
• Confirm test cases
• Confirm document
• Add branch pick label

https://claude.ai/code/session_01DzWxKNMrrsmpogSDPEP6Hi

[hello-stephen]

Thank you for your contribution to Apache Doris.
Don't know what should be done next? See How to process your PR.

Please clearly describe your PR:

1. What problem was fixed (it's best to include specific error reporting information). How it was fixed.
2. Which behaviors were modified. What was the previous behavior, what is it now, why was it modified, and what possible impacts might there be.
3. What features were added. Why was this function added?
4. Which code was refactored and why was this part of the code refactored?
5. Which functions were optimized and what is the difference before and after the optimization?

[mrhhsg]

/review

[github-actions]

Round 3 capped/incomplete: request changes. The final permitted review round still produced two newly accepted findings, so this review cannot claim convergence.

I found nine blocking issues: two reachable wrong-result paths where cap-triggered pass-through consumes a pushed local distinct limit with duplicates; incomplete state accounting and retained pass-through-arena growth that defeat the core memory bound; low-memory precedence that can raise a tighter cap; a positive-limit-to-zero/no-cap arithmetic edge; non-restoring distinct and TopN lifecycle states; loss of the explicit cap during observer-to-master forwarding; and a unit test that bypasses the production predicate.

Critical checkpoints:

• Goal and scope: the shared helper and both operators implement the intended policy shape, but the accepted correctness and memory-bound failures mean the OOM-prevention goal is not safely achieved. The eight-file scope is otherwise focused.
• Concurrency and lifecycle: the standalone relaxed atomic flag and atomic live query limit are safe; no new race, deadlock, static-lifecycle, or cleanup-status defect survived. The blocking lifecycle problems are the permanent distinct latch, stale TopN state, invisible aggregate-owned heap, and lifetime pass-through arena.
• Configuration and compatibility: no new variable or wire/storage format was added, and the existing optional Thrift field remains mixed-version compatible. However, the documented explicit upper bound is not forwarded to the planning master. Direct-master FE-to-BE serialization is otherwise correct.
• Parallel and conditional paths: normal and specialized distinct operators, spill/non-spill, low-memory, live lower/restore, helper boundaries, pure distinct, TopN, fixed maps, and generic/specialized serialization were traced. Both pure-distinct operator selections have an early-EOS wrong-result path.
• Tests and observability: MemoryUseLimit has correct byte units and refresh points, but it cannot make incomplete accounting safe. The changed normal test discards the production decision, and adverse restore, boundary, owning-state, pass-through, pushed-limit, TopN, and forwarding paths are unproved. No build or test suite was run by this review because the authoritative bundle explicitly prohibited builds; the read-only BE header-hygiene gate passed.
• Performance and error handling: per-block policy overhead is constant-time and status propagation remains intact, but the distinct latch and stale TopN state can expand retained work and shuffle volume for the rest of a query.
• Transactions and persistence: no transaction, persisted metadata, data-write format, or storage compatibility surface is changed.
• Documentation and focus: the comments describe the broadened behavior, but the implementation does not honor it on all supported paths and the PR identifies documentation work without supplying a doc link. No additional user focus was supplied.

All nine comments were independently validated and deduplicated against the live PR, which had no existing reviews, inline comments, or review threads at submission time.

[github-actions]

[P1] Do not count raw pass-through duplicates toward a distinct limit. A plain local LIMIT N is pushed to this lower distinct aggregation. Once this cap condition switches it to pass-through, duplicates are emitted raw, but push() still truncates and reaches _limit by raw row count and then stops child input. With limit 2, after key A was emitted, a triggering [A, B] block lets duplicate A consume the final allowance and prevents B from reaching the global distinct, returning too few rows. Keep deduplicating while the pushed distinct limit is active or remove/disable that local limit for pass-through, and add this transition case.

[mrhhsg]

Fixed. Two changes: (1) the memory budget never switches this operator to pass-through while a LIMIT is pushed down to it (parent.limit() != -1) — the limit already bounds the hash set; (2) push() no longer truncates against _limit once _stop_emplace_flag is set (low reduction / low-memory mode), so raw rows cannot consume the allowance; the global stage applies the limit. Covered by DistinctStreamingAggOperatorTest.pushed_limit_keeps_deduplicating and pass_through_does_not_consume_pushed_limit. Note that the truncation itself pre-dates this PR (the low-reduction path passed rows through before as well).

[github-actions]

[P2] Re-evaluate the cap stop after a restored query limit. This assignment permanently latches _stop_emplace_flag, but workload-group query limits can fall and later rise, and _memory_limit() is intentionally refreshed every block. After one transient breach, every later block skips distinct emplacement even when retained usage is below the restored cap, so the remainder of a long query loses all local deduplication and can flood the exchange. Keep cap-triggered pass-through separate from the permanent low-reduction/low-memory stop, and add a lower-then-restore test.

[mrhhsg]

Fixed. The budget decision is now per block (pass_through local), and only the pre-existing switches (low reduction rate, low_memory_mode()) latch _stop_emplace_flag. When the query limit is restored, the next block deduplicates against the retained hash set. refresh_memory_limit now asserts the lower-then-restore sequence.

[mrhhsg]

Revisited: we decided to keep the permanent latch after all and reverted the per-block decision — the extra state made _distinct_pre_agg_with_serialized_key() noticeably harder to follow for a benefit that is narrow in practice. Rationale: this operator already latches forever on the two pre-existing switches (low reduction rate, low_memory_mode()), and the "lowered then restored" workload-group scenario overlaps with low_memory_mode(true) being raised for the whole group (workload_group_manager.cpp, low-watermark path), which would latch it anyway. The memory-triggered latch is therefore just a third trigger of existing behaviour, and it is skipped entirely while a LIMIT is pushed down (see the sibling thread). StreamingAggOperatorX keeps re-evaluating per block as before.

[github-actions]

[P2] Keep a positive share distinct from the no-cap sentinel. For query_memory_limit=1 and parallel_tasks=2, integer division makes per_task_query_limit zero, and the helper returns zero even if fixed_limit is positive. Both callers interpret zero as unlimited, so a known positive query limit disables the new gate. Clamp a known-positive per-task share to at least one byte or represent unknown/no-cap separately, and add query_limit < parallel_tasks cases.

[mrhhsg]

Fixed: the per-task share is clamped to at least one byte when the query limit is positive, so a known limit can no longer collapse to the "no cap" sentinel. budget_floor_and_fixed_limit covers (1, 2, 0) and (1, 2, 256MB).

[github-actions]

[P1] Protect the regular pure-distinct fallback from the same early-EOS bug. With enable_distinct_streaming_aggregation=false, the lower pure-distinct stage is still streaming and uses this regular operator. When this dynamic cap selects pass-through, it copies duplicate group keys verbatim, while generic reached_limit() counts those raw rows against the pushed local distinct limit and reports EOS. With limit 2, retained A, and input [A, A, B], the duplicates stop the child before B, so the global distinct returns too few keys. Fix the pushed-limit/pass-through interaction in this operator too and cover the disabled-specialized-operator plan.

[mrhhsg]

Fixed. pull() only applies reached_limit() to the aggregated output (_get_results_with_serialized_key); pass-through rows from _pre_aggregated_block are neither aggregated nor deduplicated and must not count toward a limit pushed to this local stage. StreamingAggOperatorTest.memory_limit_pass_through_and_recover pushes [2, 3, 3] through with _limit = 2 and asserts the pulled block keeps its 3 rows with eos == false.

[github-actions]

[P2] Restore TopN pruning after a transient cap pass-through. If a pushed TopN aggregate hits the cap before its hash table reaches _sort_limit, the pass-through branch changes need_do_sort_limit from -1 to 0. After workload-group limit restoration, aggregation resumes, but heap construction only runs while state is -1, so state 0 can never activate _do_limit_filter(); the local stage retains and shuffles all later groups. Keep the state re-eligible or rebuild it on recovery, and test lower/restore below the TopN threshold.

[mrhhsg]

Fixed. The pass-through branch only latches need_do_sort_limit = 1 (and builds the heap) once the hash table has reached limit; a smaller table leaves the state at -1, so the aggregating branch can still build the heap after aggregation resumes. There is no TopN fixture for this operator in the unit tests, so this one is covered by inspection only; the change is confined to not assigning 0.

[github-actions]

[P1] Include all retained state in this memory cap. This comparison is not an upper bound for supported aggregates. _memory_usage() counts hash buffers and arenas but not heap owned by states stored in fixed slots; for example, bitmap_union's per-group BitmapValue/Roaring buffers grow invisibly. After this condition selects pass-through, generic serializers such as sum_foreach also allocate nested temporary states from the lifetime _agg_arena_pool, so repeated blocks keep increasing retained memory until close. Account state-owned heap and use a resettable pass-through arena, with tests for both families.

[mrhhsg]

Partially addressed. The arena part is fixed: pass-through serialization now uses a per-block Arena instead of the operator-lifetime _agg_arena_pool, so collect/map/foreach/... no longer accumulate dead bytes for the rest of the query (that growth also kept _memory_usage() above the budget forever). Heap owned by aggregate states (e.g. BitmapValue inside fixed slots) is not counted — that is the existing accounting model shared with the blocking agg sink and the previous spill-only limit, and is out of scope here; happy to track it in a separate issue.

[github-actions]

[P1] Never let low-memory mode raise a tighter cap. This early return replaces rather than constrains the live query/task and explicit spill bounds. For enable_spill=true and spill_streaming_agg_mem_limit=512 KiB, entering low-memory mode increases the allowed retention to 1 MiB; the live per-task query share can also legitimately be below 1 MiB. Compute the normal effective limit first and apply the low-memory value with min, then cover both sub-MiB cases.

[mrhhsg]

Fixed: the normal budget is computed first and low-memory mode applies min(1MB, budget). test1 now sets spill_streaming_agg_mem_limit = 512KB before entering low-memory mode and expects 512KB.

[github-actions]

[P1] Exercise the production cap decision in this test. MockStreamingAggLocalState::_should_not_do_pre_agg() calls the base method but discards its return value and then returns should_not_do_pre_agg, which remains false in test1. Consequently this push() never follows the new cap-triggered pass-through path; these additions only prove the profile value and hash-table growth, and remain green if production enforcement is broken. Use a real local state or a fixture that returns the base decision, then lower the live/fixed cap below measured usage and assert pass-through/no new hash entries (plus restoration).

[mrhhsg]

Fixed. MockStreamingAggLocalState gained use_real_decision; the new memory_limit_pass_through_and_recover test drives push() with the production _should_not_do_pre_agg() using a real sum evaluator: a lowered query limit yields pass-through (no hash-table growth, 3 raw rows pulled, eos == false despite _limit = 2), and restoring the limit resumes aggregation (hash table grows to 4).

[github-actions]

[P1] Forward this explicit cap to the planning master. The annotation lacks needForward=true, so getForwardVariables() omits the setting. A follower session with enable_spill=true and a 1 MiB cap that forwards an insert/select sends enable_spill but reconstructs this value as the master's 256 MiB default before toThrift(); both newly covered streaming operators then use the wrong upper bound. Mark the variable for forwarding and add an observer-to-master non-default-value test.

[mrhhsg]

Fixed: needForward = true added, consistent with enable_spill / enable_force_spill.

[mrhhsg]

/review

[github-actions]

Round 2 converged: request changes.

Three new issues are attached inline. Two earlier blockers also remain unresolved:

• The specialized-distinct cap still latches permanently after a reversible workload-group limit reduction, even when that query never enters low-memory mode and its live tracker limit is later restored. This is already covered by thread 3860482275, so I have not duplicated it inline.
• The regular cap is still not an upper bound over retained state: aggregate-owned heap remains omitted, and the same comparison also misses TopN owning columns/history and admits a whole current block before observing its allocations. The broad same-line request is already in thread 3860482290, so I have kept those instances duplicate-fenced.

Critical checkpoints:

• Goal and scope: the eight-file change is focused and the shared policy shape is reasonable, but the specialized distinct path bypasses the cap for pushed LIMIT plans and omits retained hash-method scratch. Those failures mean the stated query-derived memory bound is not achieved on supported paths.
• Correctness and parallel/special paths: specialized distinct, regular pure-distinct fallback, ordinary aggregation, pushed LIMIT, TopN before/after heap construction, mem-reuse/fresh output, cache draining, child EOS, spill/non-spill, low-memory, and lower/restore transitions were traced. No additional row-result issue survived beyond the live threads; the new LIMIT finding is a memory-bound failure, not a duplicate of the repaired early-EOS thread.
• Memory, lifecycle, and concurrency: aggregate/state/key arenas, hash-method scratch, TopN owners, query-tracker charging, serializer ownership, close/error paths, and task publication were checked. The block-scoped serializer arena is safe because all reachable serializers copy or move into owning columns. The relaxed atomic flag is an independent monotonic signal and associated configuration is published before execution; no new race, deadlock, or dangling-lifetime issue survived.
• Configuration and compatibility: no new setting or wire field is introduced. Existing Thrift field 105, FE toThrift(), needForward=true, old/new decoding, zero/negative/tiny sentinels, spill-only fixed bounds, and per-pipeline task counts were rechecked without another runtime defect. The explicit non-default observer-to-master test requested in the existing forwarding thread is still absent.
• Tests, observability, and performance: the PR reports 14 filtered BE unit tests, but this review did not run builds or tests because the authoritative review prompt prohibited them. Current style, license, secrets, title, and large-file checks pass; the automated code-review check is still pending and no BE test check is present. The changed tests do not cover the three inline cases, dynamic distinct restoration, TopN recovery, aggregate-owned heap, or end-to-end SQL behavior. Per-block policy overhead is constant-time, but pass-through now under-reports exported ProcessRows, and a permanent distinct latch can amplify shuffle work for the rest of a query.
• Error handling, persistence, and data writes: new Status results are propagated and impossible variants follow adjacent invariant handling. No transaction, persisted metadata, storage format, or data-write compatibility surface changes.
• Documentation and user focus: the PR marks documentation as required but changes no documentation. The release-note phrase floor 32MB per task should state that the target floor is clipped by the per-task query share and any smaller explicit spill cap. No additional user-provided focus was supplied.

Review status: complete and converged after two rounds. All three Round 2 reviewers returned NO_NEW_VALUABLE_FINDINGS; every candidate is accepted, duplicate-fenced, or dismissed with code evidence.

[github-actions]

[P1] Keep the byte cap active when a LIMIT is pushed down. The LIMIT is enforced only after _emplace_into_hash_table_to_distinct() has inserted the entire block, so it does not bound the retained hash set or arena: with LIMIT 1 and one block of many unique large strings, lazy_emplace_batch_void() copies every key before push() truncates the output, and this non-spillable state can exhaust the query tracker despite MemoryUseLimit. Let the existing pass-through/no-local-truncation logic handle a cap breach (or stop emplacement at the remaining limit/budget), and cover variable-length keys.

[github-actions]

[P2] Preserve process-row accounting for the pass-through branch. reached_limit() also increments the query IOContext::process_rows; after moving it only here, every raw block returned above for a cap or low-reduction fallback disappears from TQueryStatistics.process_rows. Split that accounting from local LIMIT truncation, or explicitly update process_rows for pass-through rows, and extend the recovery test to assert it while keeping eos == false.

[github-actions]

[P1] Count the hash method's retained serialization scratch here. Composite-key probing fills MethodSerialized::Base::arena, stored_keys, and hash_values for the whole block before deduplication, and Arena::clear() retains its largest head chunk. A duplicate-heavy 4096-row block with roughly 2 KiB serialized keys can therefore leave about 8 MiB retained while the persistent distinct set contains one key; this helper still reports below a 4 MiB MemoryUseLimit, so later blocks keep aggregating. Include the method's retained buffers (or release oversized scratch) and cover duplicate composite keys.