risingwave_batch_executors/executor/
hash_agg.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::hash::BuildHasher;
16use std::marker::PhantomData;
17use std::sync::Arc;
18
19use bytes::Bytes;
20use futures_async_stream::try_stream;
21use hashbrown::hash_map::Entry;
22use itertools::Itertools;
23use prometheus::core::Atomic;
24use risingwave_common::array::{DataChunk, StreamChunk};
25use risingwave_common::bitmap::{Bitmap, FilterByBitmap};
26use risingwave_common::catalog::{Field, Schema};
27use risingwave_common::hash::{HashKey, HashKeyDispatcher, PrecomputedBuildHasher};
28use risingwave_common::memory::MemoryContext;
29use risingwave_common::metrics::TrAdderAtomic;
30use risingwave_common::row::{OwnedRow, Row, RowExt};
31use risingwave_common::types::{DataType, ToOwnedDatum};
32use risingwave_common::util::chunk_coalesce::DataChunkBuilder;
33use risingwave_common::util::iter_util::ZipEqFast;
34use risingwave_common_estimate_size::EstimateSize;
35use risingwave_expr::aggregate::{AggCall, AggregateState, BoxedAggregateFunction};
36use risingwave_pb::Message;
37use risingwave_pb::batch_plan::HashAggNode;
38use risingwave_pb::batch_plan::plan_node::NodeBody;
39use risingwave_pb::data::DataChunk as PbDataChunk;
40
41use crate::error::{BatchError, Result};
42use crate::executor::aggregation::build as build_agg;
43use crate::executor::{
44    BoxedDataChunkStream, BoxedExecutor, BoxedExecutorBuilder, Executor, ExecutorBuilder,
45    WrapStreamExecutor,
46};
47use crate::monitor::BatchSpillMetrics;
48use crate::spill::spill_op::SpillBackend::Disk;
49use crate::spill::spill_op::{
50    DEFAULT_SPILL_PARTITION_NUM, SPILL_AT_LEAST_MEMORY, SpillBackend, SpillBuildHasher, SpillOp,
51};
52use crate::task::{ShutdownToken, TaskId};
53
54type AggHashMap<K, A> = hashbrown::HashMap<K, Vec<AggregateState>, PrecomputedBuildHasher, A>;
55
56/// A dispatcher to help create specialized hash agg executor.
57impl HashKeyDispatcher for HashAggExecutorBuilder {
58    type Output = BoxedExecutor;
59
60    fn dispatch_impl<K: HashKey>(self) -> Self::Output {
61        Box::new(HashAggExecutor::<K>::new(
62            Arc::new(self.aggs),
63            self.group_key_columns,
64            self.group_key_types,
65            self.schema,
66            self.child,
67            self.identity,
68            self.chunk_size,
69            self.mem_context,
70            self.spill_backend,
71            self.spill_metrics,
72            self.shutdown_rx,
73        ))
74    }
75
76    fn data_types(&self) -> &[DataType] {
77        &self.group_key_types
78    }
79}
80
81pub struct HashAggExecutorBuilder {
82    aggs: Vec<BoxedAggregateFunction>,
83    group_key_columns: Vec<usize>,
84    group_key_types: Vec<DataType>,
85    child: BoxedExecutor,
86    schema: Schema,
87    #[expect(dead_code)]
88    task_id: TaskId,
89    identity: String,
90    chunk_size: usize,
91    mem_context: MemoryContext,
92    spill_backend: Option<SpillBackend>,
93    spill_metrics: Arc<BatchSpillMetrics>,
94    shutdown_rx: ShutdownToken,
95}
96
97impl HashAggExecutorBuilder {
98    fn deserialize(
99        hash_agg_node: &HashAggNode,
100        child: BoxedExecutor,
101        task_id: TaskId,
102        identity: String,
103        chunk_size: usize,
104        mem_context: MemoryContext,
105        spill_backend: Option<SpillBackend>,
106        spill_metrics: Arc<BatchSpillMetrics>,
107        shutdown_rx: ShutdownToken,
108    ) -> Result<BoxedExecutor> {
109        let aggs: Vec<_> = hash_agg_node
110            .get_agg_calls()
111            .iter()
112            .map(|agg| AggCall::from_protobuf(agg).and_then(|agg| build_agg(&agg)))
113            .try_collect()?;
114
115        let group_key_columns = hash_agg_node
116            .get_group_key()
117            .iter()
118            .map(|x| *x as usize)
119            .collect_vec();
120
121        let child_schema = child.schema();
122
123        let group_key_types = group_key_columns
124            .iter()
125            .map(|i| child_schema.fields[*i].data_type.clone())
126            .collect_vec();
127
128        let fields = group_key_types
129            .iter()
130            .cloned()
131            .chain(aggs.iter().map(|e| e.return_type()))
132            .map(Field::unnamed)
133            .collect::<Vec<Field>>();
134
135        let builder = HashAggExecutorBuilder {
136            aggs,
137            group_key_columns,
138            group_key_types,
139            child,
140            schema: Schema { fields },
141            task_id,
142            identity,
143            chunk_size,
144            mem_context,
145            spill_backend,
146            spill_metrics,
147            shutdown_rx,
148        };
149
150        Ok(builder.dispatch())
151    }
152}
153
154impl BoxedExecutorBuilder for HashAggExecutorBuilder {
155    async fn new_boxed_executor(
156        source: &ExecutorBuilder<'_>,
157        inputs: Vec<BoxedExecutor>,
158    ) -> Result<BoxedExecutor> {
159        let [child]: [_; 1] = inputs.try_into().unwrap();
160
161        let hash_agg_node = try_match_expand!(
162            source.plan_node().get_node_body().unwrap(),
163            NodeBody::HashAgg
164        )?;
165
166        let identity = source.plan_node().get_identity();
167
168        let spill_metrics = source.context().spill_metrics();
169
170        Self::deserialize(
171            hash_agg_node,
172            child,
173            source.task_id.clone(),
174            identity.clone(),
175            source.context().get_config().developer.chunk_size,
176            source.context().create_executor_mem_context(identity),
177            if source.context().get_config().enable_spill {
178                Some(Disk)
179            } else {
180                None
181            },
182            spill_metrics,
183            source.shutdown_rx().clone(),
184        )
185    }
186}
187
188/// `HashAggExecutor` implements the hash aggregate algorithm.
189pub struct HashAggExecutor<K> {
190    /// Aggregate functions.
191    aggs: Arc<Vec<BoxedAggregateFunction>>,
192    /// Column indexes that specify a group
193    group_key_columns: Vec<usize>,
194    /// Data types of group key columns
195    group_key_types: Vec<DataType>,
196    /// Output schema
197    schema: Schema,
198    child: BoxedExecutor,
199    /// Used to initialize the state of the aggregation from the spilled files.
200    init_agg_state_executor: Option<BoxedExecutor>,
201    identity: String,
202    chunk_size: usize,
203    mem_context: MemoryContext,
204    spill_backend: Option<SpillBackend>,
205    spill_metrics: Arc<BatchSpillMetrics>,
206    /// The upper bound of memory usage for this executor.
207    memory_upper_bound: Option<u64>,
208    shutdown_rx: ShutdownToken,
209    _phantom: PhantomData<K>,
210}
211
212impl<K> HashAggExecutor<K> {
213    #[expect(clippy::too_many_arguments)]
214    pub fn new(
215        aggs: Arc<Vec<BoxedAggregateFunction>>,
216        group_key_columns: Vec<usize>,
217        group_key_types: Vec<DataType>,
218        schema: Schema,
219        child: BoxedExecutor,
220        identity: String,
221        chunk_size: usize,
222        mem_context: MemoryContext,
223        spill_backend: Option<SpillBackend>,
224        spill_metrics: Arc<BatchSpillMetrics>,
225        shutdown_rx: ShutdownToken,
226    ) -> Self {
227        Self::new_inner(
228            aggs,
229            group_key_columns,
230            group_key_types,
231            schema,
232            child,
233            None,
234            identity,
235            chunk_size,
236            mem_context,
237            spill_backend,
238            spill_metrics,
239            None,
240            shutdown_rx,
241        )
242    }
243
244    #[expect(clippy::too_many_arguments)]
245    fn new_inner(
246        aggs: Arc<Vec<BoxedAggregateFunction>>,
247        group_key_columns: Vec<usize>,
248        group_key_types: Vec<DataType>,
249        schema: Schema,
250        child: BoxedExecutor,
251        init_agg_state_executor: Option<BoxedExecutor>,
252        identity: String,
253        chunk_size: usize,
254        mem_context: MemoryContext,
255        spill_backend: Option<SpillBackend>,
256        spill_metrics: Arc<BatchSpillMetrics>,
257        memory_upper_bound: Option<u64>,
258        shutdown_rx: ShutdownToken,
259    ) -> Self {
260        HashAggExecutor {
261            aggs,
262            group_key_columns,
263            group_key_types,
264            schema,
265            child,
266            init_agg_state_executor,
267            identity,
268            chunk_size,
269            mem_context,
270            spill_backend,
271            spill_metrics,
272            memory_upper_bound,
273            shutdown_rx,
274            _phantom: PhantomData,
275        }
276    }
277}
278
279impl<K: HashKey + Send + Sync> Executor for HashAggExecutor<K> {
280    fn schema(&self) -> &Schema {
281        &self.schema
282    }
283
284    fn identity(&self) -> &str {
285        &self.identity
286    }
287
288    fn execute(self: Box<Self>) -> BoxedDataChunkStream {
289        self.do_execute()
290    }
291}
292
293/// `AggSpillManager` is used to manage how to write spill data file and read them back.
294/// The spill data first need to be partitioned. Each partition contains 2 files: `agg_state_file` and `input_chunks_file`.
295/// The spill file consume a data chunk and serialize the chunk into a protobuf bytes.
296/// Finally, spill file content will look like the below.
297/// The file write pattern is append-only and the read pattern is sequential scan.
298/// This can maximize the disk IO performance.
299///
300/// ```text
301/// [proto_len]
302/// [proto_bytes]
303/// ...
304/// [proto_len]
305/// [proto_bytes]
306/// ```
307pub struct AggSpillManager {
308    op: SpillOp,
309    partition_num: usize,
310    agg_state_writers: Vec<opendal::Writer>,
311    agg_state_chunk_builder: Vec<DataChunkBuilder>,
312    input_writers: Vec<opendal::Writer>,
313    input_chunk_builders: Vec<DataChunkBuilder>,
314    spill_build_hasher: SpillBuildHasher,
315    group_key_types: Vec<DataType>,
316    child_data_types: Vec<DataType>,
317    agg_data_types: Vec<DataType>,
318    spill_chunk_size: usize,
319    spill_metrics: Arc<BatchSpillMetrics>,
320}
321
322impl AggSpillManager {
323    fn new(
324        spill_backend: SpillBackend,
325        agg_identity: &String,
326        partition_num: usize,
327        group_key_types: Vec<DataType>,
328        agg_data_types: Vec<DataType>,
329        child_data_types: Vec<DataType>,
330        spill_chunk_size: usize,
331        spill_metrics: Arc<BatchSpillMetrics>,
332    ) -> Result<Self> {
333        let suffix_uuid = uuid::Uuid::new_v4();
334        let dir = format!("{}-{}/", agg_identity, suffix_uuid);
335        let op = SpillOp::create(dir, spill_backend)?;
336        let agg_state_writers = Vec::with_capacity(partition_num);
337        let agg_state_chunk_builder = Vec::with_capacity(partition_num);
338        let input_writers = Vec::with_capacity(partition_num);
339        let input_chunk_builders = Vec::with_capacity(partition_num);
340        // Use uuid to generate an unique hasher so that when recursive spilling happens they would use a different hasher to avoid data skew.
341        let spill_build_hasher = SpillBuildHasher(suffix_uuid.as_u64_pair().1);
342        Ok(Self {
343            op,
344            partition_num,
345            agg_state_writers,
346            agg_state_chunk_builder,
347            input_writers,
348            input_chunk_builders,
349            spill_build_hasher,
350            group_key_types,
351            child_data_types,
352            agg_data_types,
353            spill_chunk_size,
354            spill_metrics,
355        })
356    }
357
358    async fn init_writers(&mut self) -> Result<()> {
359        for i in 0..self.partition_num {
360            let agg_state_partition_file_name = format!("agg-state-p{}", i);
361            let w = self.op.writer_with(&agg_state_partition_file_name).await?;
362            self.agg_state_writers.push(w);
363
364            let partition_file_name = format!("input-chunks-p{}", i);
365            let w = self.op.writer_with(&partition_file_name).await?;
366            self.input_writers.push(w);
367            self.input_chunk_builders.push(DataChunkBuilder::new(
368                self.child_data_types.clone(),
369                self.spill_chunk_size,
370            ));
371            self.agg_state_chunk_builder.push(DataChunkBuilder::new(
372                self.group_key_types
373                    .iter()
374                    .cloned()
375                    .chain(self.agg_data_types.iter().cloned())
376                    .collect(),
377                self.spill_chunk_size,
378            ));
379        }
380        Ok(())
381    }
382
383    async fn write_agg_state_row(&mut self, row: impl Row, hash_code: u64) -> Result<()> {
384        let partition = hash_code as usize % self.partition_num;
385        if let Some(output_chunk) = self.agg_state_chunk_builder[partition].append_one_row(row) {
386            let chunk_pb: PbDataChunk = output_chunk.to_protobuf();
387            let buf = Message::encode_to_vec(&chunk_pb);
388            let len_bytes = Bytes::copy_from_slice(&(buf.len() as u32).to_le_bytes());
389            self.spill_metrics
390                .batch_spill_write_bytes
391                .inc_by((buf.len() + len_bytes.len()) as u64);
392            self.agg_state_writers[partition].write(len_bytes).await?;
393            self.agg_state_writers[partition].write(buf).await?;
394        }
395        Ok(())
396    }
397
398    async fn write_input_chunk(&mut self, chunk: DataChunk, hash_codes: Vec<u64>) -> Result<()> {
399        let (columns, vis) = chunk.into_parts_v2();
400        for partition in 0..self.partition_num {
401            let new_vis = vis.clone()
402                & Bitmap::from_iter(
403                    hash_codes
404                        .iter()
405                        .map(|hash_code| (*hash_code as usize % self.partition_num) == partition),
406                );
407            let new_chunk = DataChunk::from_parts(columns.clone(), new_vis);
408            for output_chunk in self.input_chunk_builders[partition].append_chunk(new_chunk) {
409                let chunk_pb: PbDataChunk = output_chunk.to_protobuf();
410                let buf = Message::encode_to_vec(&chunk_pb);
411                let len_bytes = Bytes::copy_from_slice(&(buf.len() as u32).to_le_bytes());
412                self.spill_metrics
413                    .batch_spill_write_bytes
414                    .inc_by((buf.len() + len_bytes.len()) as u64);
415                self.input_writers[partition].write(len_bytes).await?;
416                self.input_writers[partition].write(buf).await?;
417            }
418        }
419        Ok(())
420    }
421
422    async fn close_writers(&mut self) -> Result<()> {
423        for partition in 0..self.partition_num {
424            if let Some(output_chunk) = self.agg_state_chunk_builder[partition].consume_all() {
425                let chunk_pb: PbDataChunk = output_chunk.to_protobuf();
426                let buf = Message::encode_to_vec(&chunk_pb);
427                let len_bytes = Bytes::copy_from_slice(&(buf.len() as u32).to_le_bytes());
428                self.spill_metrics
429                    .batch_spill_write_bytes
430                    .inc_by((buf.len() + len_bytes.len()) as u64);
431                self.agg_state_writers[partition].write(len_bytes).await?;
432                self.agg_state_writers[partition].write(buf).await?;
433            }
434
435            if let Some(output_chunk) = self.input_chunk_builders[partition].consume_all() {
436                let chunk_pb: PbDataChunk = output_chunk.to_protobuf();
437                let buf = Message::encode_to_vec(&chunk_pb);
438                let len_bytes = Bytes::copy_from_slice(&(buf.len() as u32).to_le_bytes());
439                self.spill_metrics
440                    .batch_spill_write_bytes
441                    .inc_by((buf.len() + len_bytes.len()) as u64);
442                self.input_writers[partition].write(len_bytes).await?;
443                self.input_writers[partition].write(buf).await?;
444            }
445        }
446
447        for mut w in self.agg_state_writers.drain(..) {
448            w.close().await?;
449        }
450        for mut w in self.input_writers.drain(..) {
451            w.close().await?;
452        }
453        Ok(())
454    }
455
456    async fn read_agg_state_partition(&mut self, partition: usize) -> Result<BoxedDataChunkStream> {
457        let agg_state_partition_file_name = format!("agg-state-p{}", partition);
458        let r = self.op.reader_with(&agg_state_partition_file_name).await?;
459        Ok(SpillOp::read_stream(r, self.spill_metrics.clone()))
460    }
461
462    async fn read_input_partition(&mut self, partition: usize) -> Result<BoxedDataChunkStream> {
463        let input_partition_file_name = format!("input-chunks-p{}", partition);
464        let r = self.op.reader_with(&input_partition_file_name).await?;
465        Ok(SpillOp::read_stream(r, self.spill_metrics.clone()))
466    }
467
468    async fn estimate_partition_size(&self, partition: usize) -> Result<u64> {
469        let agg_state_partition_file_name = format!("agg-state-p{}", partition);
470        let agg_state_size = self
471            .op
472            .stat(&agg_state_partition_file_name)
473            .await?
474            .content_length();
475        let input_partition_file_name = format!("input-chunks-p{}", partition);
476        let input_size = self
477            .op
478            .stat(&input_partition_file_name)
479            .await?
480            .content_length();
481        Ok(agg_state_size + input_size)
482    }
483
484    async fn clear_partition(&mut self, partition: usize) -> Result<()> {
485        let agg_state_partition_file_name = format!("agg-state-p{}", partition);
486        self.op.delete(&agg_state_partition_file_name).await?;
487        let input_partition_file_name = format!("input-chunks-p{}", partition);
488        self.op.delete(&input_partition_file_name).await?;
489        Ok(())
490    }
491}
492
493impl<K: HashKey + Send + Sync> HashAggExecutor<K> {
494    #[try_stream(boxed, ok = DataChunk, error = BatchError)]
495    async fn do_execute(self: Box<Self>) {
496        // Keep this partition's charges separate from the shared parent. Its private context
497        // releases remaining state charges on completion, errors, or cancellation.
498        let mem_context = MemoryContext::new(Some(self.mem_context.clone()), TrAdderAtomic::new(0));
499        let child_schema = self.child.schema().clone();
500        let mut need_to_spill = false;
501        // If the memory upper bound is less than 1MB, we don't need to check memory usage.
502        let check_memory = match self.memory_upper_bound {
503            Some(upper_bound) => upper_bound > SPILL_AT_LEAST_MEMORY,
504            None => true,
505        };
506
507        // Track only this executor's manual charges; spill partitions share the parent counter.
508        let mut states_heap_size = 0;
509        // hash map for each agg groups
510        let mut groups = AggHashMap::<K, _>::with_hasher_in(
511            PrecomputedBuildHasher,
512            mem_context.global_allocator(),
513        );
514
515        if let Some(init_agg_state_executor) = self.init_agg_state_executor {
516            // `init_agg_state_executor` exists which means this is a sub `HashAggExecutor` used to consume spilling data.
517            // The spilled agg states by its parent executor need to be recovered first.
518            let mut init_agg_state_stream = init_agg_state_executor.execute();
519            #[for_await]
520            for chunk in &mut init_agg_state_stream {
521                let chunk = chunk?;
522                let group_key_indices = (0..self.group_key_columns.len()).collect_vec();
523                let keys = K::build_many(&group_key_indices, &chunk);
524                let mut memory_usage_diff = 0;
525                for (row_id, key) in keys.into_iter().enumerate() {
526                    let mut agg_states = vec![];
527                    for i in 0..self.aggs.len() {
528                        let agg = &self.aggs[i];
529                        let datum = chunk
530                            .row_at(row_id)
531                            .0
532                            .datum_at(self.group_key_columns.len() + i)
533                            .to_owned_datum();
534                        let agg_state = agg.decode_state(datum)?;
535                        memory_usage_diff += agg_state.estimated_size() as i64;
536                        agg_states.push(agg_state);
537                    }
538                    groups.try_insert(key, agg_states).unwrap();
539                }
540
541                // Restored states are retained even when loading the partition exceeds the limit.
542                states_heap_size += memory_usage_diff;
543                mem_context.add_unchecked(memory_usage_diff);
544                if check_memory && !mem_context.check_memory_usage() {
545                    warn!(
546                        "not enough memory to load one partition agg state after spill which is not a normal case, so keep going"
547                    );
548                }
549            }
550        }
551
552        let mut input_stream = self.child.execute();
553        // consume all chunks to compute the agg result
554        #[for_await]
555        for chunk in &mut input_stream {
556            let chunk = StreamChunk::from(chunk?);
557            let keys = K::build_many(self.group_key_columns.as_slice(), &chunk);
558            let mut memory_usage_diff = 0;
559            for (row_id, key) in keys
560                .into_iter()
561                .enumerate()
562                .filter_by_bitmap(chunk.visibility())
563            {
564                let mut new_group = false;
565                let states = match groups.entry(key) {
566                    Entry::Occupied(entry) => entry.into_mut(),
567                    Entry::Vacant(entry) => {
568                        new_group = true;
569                        let states = self
570                            .aggs
571                            .iter()
572                            .map(|agg| agg.create_state())
573                            .try_collect()?;
574                        entry.insert(states)
575                    }
576                };
577
578                // TODO: currently not a vectorized implementation
579                for (agg, state) in self.aggs.iter().zip_eq_fast(states) {
580                    if !new_group {
581                        memory_usage_diff -= state.estimated_size() as i64;
582                    }
583                    agg.update_range(state, &chunk, row_id..row_id + 1).await?;
584                    memory_usage_diff += state.estimated_size() as i64;
585                }
586            }
587            // States have already been updated and must remain charged until they are released.
588            states_heap_size += memory_usage_diff;
589            mem_context.add_unchecked(memory_usage_diff);
590            if check_memory && !mem_context.check_memory_usage() {
591                if self.spill_backend.is_some() {
592                    need_to_spill = true;
593                    break;
594                } else {
595                    Err(BatchError::OutOfMemory(self.mem_context.mem_limit()))?;
596                }
597            }
598        }
599
600        if need_to_spill {
601            // A spilling version of aggregation based on the RFC: Spill Hash Aggregation https://github.com/risingwavelabs/rfcs/pull/89
602            // When HashAggExecutor told memory is insufficient, AggSpillManager will start to partition the hash table and spill to disk.
603            // After spilling the hash table, AggSpillManager will consume all chunks from the input executor,
604            // partition and spill to disk with the same hash function as the hash table spilling.
605            // Finally, we would get e.g. 20 partitions. Each partition should contain a portion of the original hash table and input data.
606            // A sub HashAggExecutor would be used to consume each partition one by one.
607            // If memory is still not enough in the sub HashAggExecutor, it will spill its hash table and input recursively.
608            info!(
609                "batch hash agg executor {} starts to spill out",
610                &self.identity
611            );
612            let mut agg_spill_manager = AggSpillManager::new(
613                self.spill_backend.clone().unwrap(),
614                &self.identity,
615                DEFAULT_SPILL_PARTITION_NUM,
616                self.group_key_types.clone(),
617                self.aggs.iter().map(|agg| agg.return_type()).collect(),
618                child_schema.data_types(),
619                self.chunk_size,
620                self.spill_metrics.clone(),
621            )?;
622            agg_spill_manager.init_writers().await?;
623
624            let mut memory_usage_diff = 0;
625            // Spill agg states.
626            for (key, states) in groups {
627                let key_row = key.deserialize(&self.group_key_types)?;
628                let mut agg_datums = vec![];
629                for (agg, state) in self.aggs.iter().zip_eq_fast(states) {
630                    let encode_state = agg.encode_state(&state)?;
631                    memory_usage_diff -= state.estimated_size() as i64;
632                    agg_datums.push(encode_state);
633                }
634                let agg_state_row = OwnedRow::from_iter(agg_datums.into_iter());
635                let hash_code = agg_spill_manager.spill_build_hasher.hash_one(key);
636                agg_spill_manager
637                    .write_agg_state_row(key_row.chain(agg_state_row), hash_code)
638                    .await?;
639            }
640
641            // Release memory occupied by agg hash map.
642            debug_assert_eq!(memory_usage_diff, -states_heap_size);
643            mem_context.add_unchecked(memory_usage_diff);
644
645            // Spill input chunks.
646            #[for_await]
647            for chunk in input_stream {
648                let chunk: DataChunk = chunk?;
649                let hash_codes = chunk.get_hash_values(
650                    self.group_key_columns.as_slice(),
651                    agg_spill_manager.spill_build_hasher,
652                );
653                agg_spill_manager
654                    .write_input_chunk(
655                        chunk,
656                        hash_codes
657                            .into_iter()
658                            .map(|hash_code| hash_code.value())
659                            .collect(),
660                    )
661                    .await?;
662            }
663
664            agg_spill_manager.close_writers().await?;
665
666            // Process each partition one by one.
667            for i in 0..agg_spill_manager.partition_num {
668                let partition_size = agg_spill_manager.estimate_partition_size(i).await?;
669
670                let agg_state_stream = agg_spill_manager.read_agg_state_partition(i).await?;
671                let input_stream = agg_spill_manager.read_input_partition(i).await?;
672
673                let sub_hash_agg_executor: HashAggExecutor<K> = HashAggExecutor::new_inner(
674                    self.aggs.clone(),
675                    self.group_key_columns.clone(),
676                    self.group_key_types.clone(),
677                    self.schema.clone(),
678                    Box::new(WrapStreamExecutor::new(child_schema.clone(), input_stream)),
679                    Some(Box::new(WrapStreamExecutor::new(
680                        self.schema.clone(),
681                        agg_state_stream,
682                    ))),
683                    format!("{}-sub{}", self.identity.clone(), i),
684                    self.chunk_size,
685                    self.mem_context.clone(),
686                    self.spill_backend.clone(),
687                    self.spill_metrics.clone(),
688                    Some(partition_size),
689                    self.shutdown_rx.clone(),
690                );
691
692                debug!(
693                    "create sub_hash_agg {} for hash_agg {} to spill",
694                    sub_hash_agg_executor.identity, self.identity
695                );
696
697                let sub_hash_agg_stream = Box::new(sub_hash_agg_executor).execute();
698
699                #[for_await]
700                for chunk in sub_hash_agg_stream {
701                    let chunk = chunk?;
702                    yield chunk;
703                }
704
705                // Clear files of the current partition.
706                agg_spill_manager.clear_partition(i).await?;
707            }
708        } else {
709            // Don't use `into_iter` here, it may cause memory leak.
710            let mut result = groups.iter_mut();
711            let cardinality = self.chunk_size;
712            loop {
713                let mut group_builders: Vec<_> = self
714                    .group_key_types
715                    .iter()
716                    .map(|datatype| datatype.create_array_builder(cardinality))
717                    .collect();
718
719                let mut agg_builders: Vec<_> = self
720                    .aggs
721                    .iter()
722                    .map(|agg| agg.return_type().create_array_builder(cardinality))
723                    .collect();
724
725                let mut has_next = false;
726                let mut array_len = 0;
727                for (key, states) in result.by_ref().take(cardinality) {
728                    self.shutdown_rx.check()?;
729                    has_next = true;
730                    array_len += 1;
731                    key.deserialize_to_builders(&mut group_builders[..], &self.group_key_types)?;
732                    for ((agg, state), builder) in (self.aggs.iter())
733                        .zip_eq_fast(states)
734                        .zip_eq_fast(&mut agg_builders)
735                    {
736                        let result = agg.get_result(state).await?;
737                        builder.append(result);
738                    }
739                }
740                if !has_next {
741                    break; // exit loop
742                }
743
744                let columns = group_builders
745                    .into_iter()
746                    .chain(agg_builders)
747                    .map(|b| b.finish().into())
748                    .collect::<Vec<_>>();
749
750                let output = DataChunk::new(columns, array_len);
751                yield output;
752            }
753        }
754    }
755}
756
757#[cfg(test)]
758mod tests {
759    use std::alloc::{AllocError, Allocator, Global, Layout};
760    use std::ptr::NonNull;
761    use std::sync::atomic::{AtomicBool, Ordering};
762
763    use allocator_api2::alloc::{AllocError as AllocErrorApi2, Allocator as AllocatorApi2};
764    use futures_async_stream::for_await;
765    use risingwave_common::metrics::LabelGuardedIntGauge;
766    use risingwave_common::test_prelude::DataChunkTestExt;
767    use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
768    use risingwave_pb::data::PbDataType;
769    use risingwave_pb::data::data_type::TypeName;
770    use risingwave_pb::expr::agg_call::PbKind as PbAggKind;
771    use risingwave_pb::expr::{AggCall, InputRef};
772
773    use super::*;
774    use crate::executor::SortExecutor;
775    use crate::executor::test_utils::{MockExecutor, diff_executor_output};
776
777    const CHUNK_SIZE: usize = 1024;
778
779    #[tokio::test]
780    async fn execute_int32_grouped() {
781        let parent_mem = MemoryContext::root(LabelGuardedIntGauge::test_int_gauge::<4>(), u64::MAX);
782        {
783            let src_exec = Box::new(MockExecutor::with_chunk(
784                DataChunk::from_pretty(
785                    "i i i
786                 0 1 1
787                 1 1 1
788                 0 0 1
789                 1 1 2
790                 1 0 1
791                 0 0 2
792                 1 1 3
793                 0 1 2",
794                ),
795                Schema::new(vec![
796                    Field::unnamed(DataType::Int32),
797                    Field::unnamed(DataType::Int32),
798                    Field::unnamed(DataType::Int64),
799                ]),
800            ));
801
802            let agg_call = AggCall {
803                kind: PbAggKind::Sum as i32,
804                args: vec![InputRef {
805                    index: 2,
806                    r#type: Some(PbDataType {
807                        type_name: TypeName::Int32 as i32,
808                        ..Default::default()
809                    }),
810                }],
811                return_type: Some(PbDataType {
812                    type_name: TypeName::Int64 as i32,
813                    ..Default::default()
814                }),
815                distinct: false,
816                order_by: vec![],
817                filter: None,
818                direct_args: vec![],
819                udf: None,
820                scalar: None,
821            };
822
823            let agg_prost = HashAggNode {
824                group_key: vec![0, 1],
825                agg_calls: vec![agg_call],
826            };
827
828            let mem_context = MemoryContext::new(
829                Some(parent_mem.clone()),
830                LabelGuardedIntGauge::test_int_gauge::<4>(),
831            );
832            let actual_exec = HashAggExecutorBuilder::deserialize(
833                &agg_prost,
834                src_exec,
835                TaskId::default(),
836                "HashAggExecutor".to_owned(),
837                CHUNK_SIZE,
838                mem_context.clone(),
839                None,
840                BatchSpillMetrics::for_test(),
841                ShutdownToken::empty(),
842            )
843            .unwrap();
844
845            // TODO: currently the order is fixed unless the hasher is changed
846            let expect_exec = Box::new(MockExecutor::with_chunk(
847                DataChunk::from_pretty(
848                    "i i I
849                 1 0 1
850                 0 0 3
851                 0 1 3
852                 1 1 6",
853                ),
854                Schema::new(vec![
855                    Field::unnamed(DataType::Int32),
856                    Field::unnamed(DataType::Int32),
857                    Field::unnamed(DataType::Int64),
858                ]),
859            ));
860            diff_executor_output(actual_exec, expect_exec).await;
861
862            // Finishing a partition promptly clears its memory usage accounting.
863            assert_eq!(mem_context.get_bytes_used(), 0);
864        }
865
866        // Ensure that agg memory counter has been dropped.
867        assert_eq!(0, parent_mem.get_bytes_used());
868    }
869
870    #[tokio::test]
871    async fn execute_count_star() {
872        let src_exec = MockExecutor::with_chunk(
873            DataChunk::from_pretty(
874                "i
875                 0
876                 1
877                 0
878                 1
879                 1
880                 0
881                 1
882                 0",
883            ),
884            Schema::new(vec![Field::unnamed(DataType::Int32)]),
885        );
886
887        let agg_call = AggCall {
888            kind: PbAggKind::Count as i32,
889            args: vec![],
890            return_type: Some(PbDataType {
891                type_name: TypeName::Int64 as i32,
892                ..Default::default()
893            }),
894            distinct: false,
895            order_by: vec![],
896            filter: None,
897            direct_args: vec![],
898            udf: None,
899            scalar: None,
900        };
901
902        let agg_prost = HashAggNode {
903            group_key: vec![],
904            agg_calls: vec![agg_call],
905        };
906
907        let actual_exec = HashAggExecutorBuilder::deserialize(
908            &agg_prost,
909            Box::new(src_exec),
910            TaskId::default(),
911            "HashAggExecutor".to_owned(),
912            CHUNK_SIZE,
913            MemoryContext::none(),
914            None,
915            BatchSpillMetrics::for_test(),
916            ShutdownToken::empty(),
917        )
918        .unwrap();
919
920        let expect_exec = MockExecutor::with_chunk(
921            DataChunk::from_pretty(
922                "I
923                 8",
924            ),
925            Schema::new(vec![Field::unnamed(DataType::Int64)]),
926        );
927        diff_executor_output(actual_exec, Box::new(expect_exec)).await;
928    }
929
930    /// A test to verify that `HashMap` may leak memory counter when using `into_iter`.
931    #[test]
932    #[should_panic] // TODO(MrCroxx): This bug is fixed and the test should panic. Remove the test and fix the related code later.
933    fn test_hashmap_into_iter_bug() {
934        let dropped: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
935
936        {
937            struct MyAllocInner {
938                drop_flag: Arc<AtomicBool>,
939            }
940
941            #[derive(Clone)]
942            struct MyAlloc {
943                #[expect(dead_code)]
944                inner: Arc<MyAllocInner>,
945            }
946
947            impl Drop for MyAllocInner {
948                fn drop(&mut self) {
949                    println!("MyAlloc freed.");
950                    self.drop_flag.store(true, Ordering::SeqCst);
951                }
952            }
953
954            unsafe impl Allocator for MyAlloc {
955                fn allocate(
956                    &self,
957                    layout: Layout,
958                ) -> std::result::Result<NonNull<[u8]>, AllocError> {
959                    let g = Global;
960                    g.allocate(layout)
961                }
962
963                unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
964                    unsafe {
965                        let g = Global;
966                        g.deallocate(ptr, layout)
967                    }
968                }
969            }
970
971            unsafe impl AllocatorApi2 for MyAlloc {
972                fn allocate(
973                    &self,
974                    layout: Layout,
975                ) -> std::result::Result<NonNull<[u8]>, AllocErrorApi2> {
976                    let g = Global;
977                    g.allocate(layout).map_err(|_| AllocErrorApi2)
978                }
979
980                unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
981                    unsafe {
982                        let g = Global;
983                        g.deallocate(ptr, layout)
984                    }
985                }
986            }
987
988            let mut map = hashbrown::HashMap::with_capacity_in(
989                10,
990                MyAlloc {
991                    inner: Arc::new(MyAllocInner {
992                        drop_flag: dropped.clone(),
993                    }),
994                },
995            );
996            for i in 0..10 {
997                map.entry(i).or_insert_with(|| "i".to_owned());
998            }
999
1000            for (k, v) in map {
1001                println!("{}, {}", k, v);
1002            }
1003        }
1004
1005        assert!(!dropped.load(Ordering::SeqCst));
1006    }
1007
1008    #[tokio::test]
1009    async fn test_shutdown() {
1010        let src_exec = MockExecutor::with_chunk(
1011            DataChunk::from_pretty(
1012                "i i i
1013                 0 1 1",
1014            ),
1015            Schema::new(vec![Field::unnamed(DataType::Int32); 3]),
1016        );
1017
1018        let agg_call = AggCall {
1019            kind: PbAggKind::Sum as i32,
1020            args: vec![InputRef {
1021                index: 2,
1022                r#type: Some(PbDataType {
1023                    type_name: TypeName::Int32 as i32,
1024                    ..Default::default()
1025                }),
1026            }],
1027            return_type: Some(PbDataType {
1028                type_name: TypeName::Int64 as i32,
1029                ..Default::default()
1030            }),
1031            distinct: false,
1032            order_by: vec![],
1033            filter: None,
1034            direct_args: vec![],
1035            udf: None,
1036            scalar: None,
1037        };
1038
1039        let agg_prost = HashAggNode {
1040            group_key: vec![0, 1],
1041            agg_calls: vec![agg_call],
1042        };
1043
1044        let (shutdown_tx, shutdown_rx) = ShutdownToken::new();
1045        let actual_exec = HashAggExecutorBuilder::deserialize(
1046            &agg_prost,
1047            Box::new(src_exec),
1048            TaskId::default(),
1049            "HashAggExecutor".to_owned(),
1050            CHUNK_SIZE,
1051            MemoryContext::none(),
1052            None,
1053            BatchSpillMetrics::for_test(),
1054            shutdown_rx,
1055        )
1056        .unwrap();
1057
1058        shutdown_tx.cancel();
1059
1060        #[for_await]
1061        for data in actual_exec.execute() {
1062            assert!(data.is_err());
1063            break;
1064        }
1065    }
1066
1067    fn create_order_by_executor(child: BoxedExecutor) -> BoxedExecutor {
1068        let column_orders = child
1069            .schema()
1070            .fields
1071            .iter()
1072            .enumerate()
1073            .map(|(i, _)| ColumnOrder {
1074                column_index: i,
1075                order_type: OrderType::ascending(),
1076            })
1077            .collect_vec();
1078
1079        Box::new(SortExecutor::new(
1080            child,
1081            Arc::new(column_orders),
1082            "SortExecutor".into(),
1083            CHUNK_SIZE,
1084            MemoryContext::none(),
1085            None,
1086            BatchSpillMetrics::for_test(),
1087        ))
1088    }
1089
1090    #[tokio::test]
1091    async fn test_spill_hash_agg() {
1092        let src_exec = Box::new(MockExecutor::with_chunk(
1093            DataChunk::from_pretty(
1094                "i i i
1095                 0 1 1
1096                 1 1 1
1097                 0 0 1
1098                 1 1 2
1099                 1 0 1
1100                 0 0 2
1101                 1 1 3
1102                 0 1 2",
1103            ),
1104            Schema::new(vec![
1105                Field::unnamed(DataType::Int32),
1106                Field::unnamed(DataType::Int32),
1107                Field::unnamed(DataType::Int64),
1108            ]),
1109        ));
1110
1111        let agg_call = AggCall {
1112            kind: PbAggKind::Sum as i32,
1113            args: vec![InputRef {
1114                index: 2,
1115                r#type: Some(PbDataType {
1116                    type_name: TypeName::Int32 as i32,
1117                    ..Default::default()
1118                }),
1119            }],
1120            return_type: Some(PbDataType {
1121                type_name: TypeName::Int64 as i32,
1122                ..Default::default()
1123            }),
1124            distinct: false,
1125            order_by: vec![],
1126            filter: None,
1127            direct_args: vec![],
1128            udf: None,
1129            scalar: None,
1130        };
1131
1132        let agg_prost = HashAggNode {
1133            group_key: vec![0, 1],
1134            agg_calls: vec![agg_call],
1135        };
1136
1137        let mem_context =
1138            MemoryContext::new_with_mem_limit(None, LabelGuardedIntGauge::test_int_gauge::<4>(), 0);
1139        let actual_exec = HashAggExecutorBuilder::deserialize(
1140            &agg_prost,
1141            src_exec,
1142            TaskId::default(),
1143            "HashAggExecutor".to_owned(),
1144            CHUNK_SIZE,
1145            mem_context.clone(),
1146            Some(SpillBackend::Memory),
1147            BatchSpillMetrics::for_test(),
1148            ShutdownToken::empty(),
1149        )
1150        .unwrap();
1151
1152        let actual_exec = create_order_by_executor(actual_exec);
1153
1154        let expect_exec = Box::new(MockExecutor::with_chunk(
1155            DataChunk::from_pretty(
1156                "i i I
1157                 1 0 1
1158                 0 0 3
1159                 0 1 3
1160                 1 1 6",
1161            ),
1162            Schema::new(vec![
1163                Field::unnamed(DataType::Int32),
1164                Field::unnamed(DataType::Int32),
1165                Field::unnamed(DataType::Int64),
1166            ]),
1167        ));
1168
1169        let expect_exec = create_order_by_executor(expect_exec);
1170        diff_executor_output(actual_exec, expect_exec).await;
1171    }
1172}
risingwave_batch_executors/executor/hash_agg.rs

risingwave_batch_executors/executor/
hash_agg.rs
