risingwave_batch_executors/executor/
merge_sort.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::mem;
16use std::sync::Arc;
17
18use futures_async_stream::try_stream;
19use futures_util::StreamExt;
20use itertools::Itertools;
21use prometheus::core::Atomic;
22use risingwave_common::array::DataChunk;
23use risingwave_common::catalog::Schema;
24use risingwave_common::memory::{MemMonitoredHeap, MemoryContext, MonitoredGlobalAlloc};
25use risingwave_common::metrics::TrAdderAtomic;
26use risingwave_common::types::ToOwnedDatum;
27use risingwave_common::util::sort_util::{ColumnOrder, HeapElem};
28use risingwave_common_estimate_size::EstimateSize;
29
30use super::{BoxedDataChunkStream, BoxedExecutor, Executor};
31use crate::error::{BatchError, Result};
32
33pub struct MergeSortExecutor {
34    inputs: Vec<BoxedExecutor>,
35    column_orders: Arc<Vec<ColumnOrder>>,
36    identity: String,
37    schema: Schema,
38    chunk_size: usize,
39    mem_context: MemoryContext,
40    min_heap: MemMonitoredHeap<HeapElem>,
41    current_chunks: Vec<Option<DataChunk>, MonitoredGlobalAlloc>,
42}
43
44impl Executor for MergeSortExecutor {
45    fn schema(&self) -> &Schema {
46        &self.schema
47    }
48
49    fn identity(&self) -> &str {
50        &self.identity
51    }
52
53    fn execute(self: Box<Self>) -> BoxedDataChunkStream {
54        self.do_execute()
55    }
56}
57
58impl MergeSortExecutor {
59    #[try_stream(boxed, ok = DataChunk, error = BatchError)]
60    async fn do_execute(mut self: Box<Self>) {
61        let mut inputs = vec![];
62        mem::swap(&mut inputs, &mut self.inputs);
63        let mut input_streams = inputs
64            .into_iter()
65            .map(|input| input.execute())
66            .collect_vec();
67        debug_assert!(
68            self.current_chunks.is_empty(),
69            "merge-sort input slots must be empty before execution"
70        );
71        for input_idx in 0..input_streams.len() {
72            self.current_chunks.push(None);
73            // Initial chunks need the same charge as replacements: both are released on retirement.
74            self.get_input_chunk(&mut input_streams, input_idx).await?;
75            if let Some(chunk) = &self.current_chunks[input_idx] {
76                let next_row_idx = chunk.next_visible_row_idx(0);
77                self.push_row_into_heap(input_idx, next_row_idx.unwrap());
78            }
79        }
80
81        while !self.min_heap.is_empty() {
82            // It is possible that we cannot produce this much as
83            // we may run out of input data chunks from sources.
84            let mut want_to_produce = self.chunk_size;
85
86            let mut builders: Vec<_> = self
87                .schema
88                .fields
89                .iter()
90                .map(|field| field.data_type.create_array_builder(self.chunk_size))
91                .collect();
92            let mut array_len = 0;
93            while want_to_produce > 0 && !self.min_heap.is_empty() {
94                let top_elem = self.min_heap.pop().unwrap();
95                let child_idx = top_elem.chunk_idx();
96                let cur_chunk = top_elem.chunk();
97                let row_idx = top_elem.elem_idx();
98                for (idx, builder) in builders.iter_mut().enumerate() {
99                    let chunk_arr = cur_chunk.column_at(idx);
100                    let chunk_arr = chunk_arr.as_ref();
101                    let datum = chunk_arr.value_at(row_idx).to_owned_datum();
102                    builder.append(&datum);
103                }
104                want_to_produce -= 1;
105                array_len += 1;
106                // check whether we have another row from the same chunk being popped
107                let possible_next_row_idx = cur_chunk.next_visible_row_idx(row_idx + 1);
108                match possible_next_row_idx {
109                    Some(next_row_idx) => {
110                        self.push_row_into_heap(child_idx, next_row_idx);
111                    }
112                    None => {
113                        self.get_input_chunk(&mut input_streams, child_idx).await?;
114                        if let Some(chunk) = &self.current_chunks[child_idx] {
115                            let next_row_idx = chunk.next_visible_row_idx(0);
116                            self.push_row_into_heap(child_idx, next_row_idx.unwrap());
117                        }
118                    }
119                }
120            }
121
122            let columns = builders
123                .into_iter()
124                .map(|builder| builder.finish().into())
125                .collect::<Vec<_>>();
126            let chunk = DataChunk::new(columns, array_len);
127            yield chunk
128        }
129    }
130
131    async fn get_input_chunk(
132        &mut self,
133        input_streams: &mut Vec<BoxedDataChunkStream>,
134        input_idx: usize,
135    ) -> Result<()> {
136        assert!(input_idx < input_streams.len());
137        let res = input_streams[input_idx].next().await;
138        let old = match res {
139            Some(chunk) => {
140                let chunk = chunk?;
141                assert_ne!(chunk.cardinality(), 0);
142                let new_chunk_size = chunk.estimated_heap_size() as i64;
143                let old = self.current_chunks[input_idx].replace(chunk);
144                self.mem_context.add_unchecked(new_chunk_size);
145                old
146            }
147            None => std::mem::take(&mut self.current_chunks[input_idx]),
148        };
149
150        if let Some(chunk) = old {
151            // Reduce the heap size of retired chunk
152            self.mem_context
153                .add_unchecked(-(chunk.estimated_heap_size() as i64));
154        }
155
156        Ok(())
157    }
158
159    fn push_row_into_heap(&mut self, input_idx: usize, row_idx: usize) {
160        assert!(input_idx < self.current_chunks.len());
161        let chunk_ref = self.current_chunks[input_idx].as_ref().unwrap();
162        self.min_heap.push(HeapElem::new(
163            self.column_orders.clone(),
164            chunk_ref.clone(),
165            input_idx,
166            row_idx,
167            None,
168        ));
169    }
170}
171
172impl MergeSortExecutor {
173    pub fn new(
174        inputs: Vec<BoxedExecutor>,
175        column_orders: Arc<Vec<ColumnOrder>>,
176        schema: Schema,
177        identity: String,
178        chunk_size: usize,
179        mem_context: MemoryContext,
180    ) -> Self {
181        // Create the private context before allocating either container. Once this executor and
182        // its allocators are dropped, even unfinished chunk and heap charges leave the parent.
183        let mem_context = MemoryContext::new(Some(mem_context), TrAdderAtomic::new(0));
184        let inputs_num = inputs.len();
185        Self {
186            inputs,
187            column_orders,
188            identity,
189            schema,
190            chunk_size,
191            min_heap: MemMonitoredHeap::with_capacity(inputs_num, mem_context.clone()),
192            current_chunks: Vec::with_capacity_in(inputs_num, mem_context.global_allocator()),
193            mem_context,
194        }
195    }
196}
risingwave_batch_executors/executor/merge_sort.rs

risingwave_batch_executors/executor/
merge_sort.rs
