linkrs_query/executor/streaming/operators/blocking/
aggregate.rs

1use std::collections::HashMap;
2
3use crate::executor::streaming::helpers::accumulator_states::{
4    decode_partial_with_args, AggregateAccumulator,
5};
6use crate::executor::streaming::spill::{HashPartitionSpiller, SpilledRun};
7use linkrs_core::columnar::{MaterializedBatch, RowKey};
8use linkrs_core::types::expr::Expression;
9use linkrs_core::types::operators::AggregateFunction;
10use linkrs_core::Value;
11
12/// Estimated in-memory overhead of one `AggregateAccumulator` instance
13/// (enum tag plus internal state) used for memory accounting. Charged once
14/// per aggregate function per group so workloads with many small keys are
15/// accounted for beyond the group key itself.
16pub const ACCUMULATOR_OVERHEAD_BYTES: usize = 64;
17
18#[derive(Debug)]
19pub struct AggregateState {
20    /// Accumulator state per group key (computed key values, already minimal;
21    /// `RowKey` documents the key-projection contract without changing layout).
22    pub group_map: HashMap<RowKey, Vec<AggregateAccumulator>>,
23    /// Per-group memory-budget overhead charged for accumulator instances.
24    pub accumulator_overhead: usize,
25    /// Output buffer replacing the drained row iterator.
26    pub result_batch: MaterializedBatch,
27    /// Output cursor into `result_batch`.
28    pub emitted_offset: usize,
29    pub partition_spiller: Option<HashPartitionSpiller>,
30    pub spilled_runs: Vec<Option<SpilledRun>>,
31    pub current_partition: usize,
32    pub has_spilled: bool,
33    /// True once the final aggregate result has been fully emitted.
34    pub output_complete: bool,
35    pub col_names: Vec<String>,
36}
37
38#[derive(Debug)]
39pub struct GroupByState {
40    /// Column-oriented input buffer (was `Vec<Vec<Value>>`).
41    pub batch: MaterializedBatch,
42    pub col_names: Vec<String>,
43    /// Output buffer replacing the drained row iterator.
44    pub result_batch: MaterializedBatch,
45    /// Output cursor into `result_batch`.
46    pub emitted_offset: usize,
47    pub partition_spiller: Option<HashPartitionSpiller>,
48    pub spilled_runs: Vec<Option<SpilledRun>>,
49    pub current_partition: usize,
50    pub has_spilled: bool,
51    /// True once all groups have been emitted.
52    pub output_complete: bool,
53}
54
55#[derive(Debug)]
56pub struct PartialAggregateState {
57    pub group_map: HashMap<RowKey, Vec<AggregateAccumulator>>,
58    pub col_names: Vec<String>,
59    /// Output buffer replacing the drained row iterator.
60    pub result_batch: MaterializedBatch,
61    /// Output cursor into `result_batch`.
62    pub emitted_offset: usize,
63}
64
65#[derive(Debug)]
66pub struct FinalAggregateState {
67    pub group_map: HashMap<RowKey, Vec<AggregateAccumulator>>,
68    pub col_names: Vec<String>,
69    /// Output buffer replacing the drained row iterator.
70    pub result_batch: MaterializedBatch,
71    /// Output cursor into `result_batch`.
72    pub emitted_offset: usize,
73}
74
75pub(crate) fn value_to_partial_accumulator(
76    func: &AggregateFunction,
77    args: &[Expression],
78    value: &Value,
79) -> Option<AggregateAccumulator> {
80    decode_partial_with_args(func, args, value)
81}
linkrs_query/executor/streaming/operators/blocking/aggregate.rs

linkrs_query/executor/streaming/operators/blocking/
aggregate.rs
