risingwave_batch_executors/executor/
top_n.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::cmp::Ordering;
16use std::sync::Arc;
17
18use futures_async_stream::try_stream;
19use prometheus::core::Atomic;
20use risingwave_common::array::DataChunk;
21use risingwave_common::catalog::Schema;
22use risingwave_common::memory::{MemMonitoredHeap, MemoryContext};
23use risingwave_common::metrics::TrAdderAtomic;
24use risingwave_common::row::{OwnedRow, Row};
25use risingwave_common::util::chunk_coalesce::DataChunkBuilder;
26use risingwave_common::util::memcmp_encoding::{MemcmpEncoded, encode_chunk};
27use risingwave_common::util::sort_util::ColumnOrder;
28use risingwave_common_estimate_size::EstimateSize;
29use risingwave_pb::batch_plan::plan_node::NodeBody;
30
31use crate::error::{BatchError, Result};
32use crate::executor::{
33    BoxedDataChunkStream, BoxedExecutor, BoxedExecutorBuilder, Executor, ExecutorBuilder,
34};
35
36/// Top-N Executor
37///
38/// Use a N-heap to store the smallest N rows.
39pub struct TopNExecutor {
40    child: BoxedExecutor,
41    column_orders: Vec<ColumnOrder>,
42    offset: usize,
43    limit: usize,
44    with_ties: bool,
45    schema: Schema,
46    identity: String,
47    chunk_size: usize,
48    mem_ctx: MemoryContext,
49}
50
51impl BoxedExecutorBuilder for TopNExecutor {
52    async fn new_boxed_executor(
53        source: &ExecutorBuilder<'_>,
54        inputs: Vec<BoxedExecutor>,
55    ) -> Result<BoxedExecutor> {
56        let [child]: [_; 1] = inputs.try_into().unwrap();
57
58        let top_n_node =
59            try_match_expand!(source.plan_node().get_node_body().unwrap(), NodeBody::TopN)?;
60
61        let column_orders = top_n_node
62            .column_orders
63            .iter()
64            .map(ColumnOrder::from_protobuf)
65            .collect();
66
67        let identity = source.plan_node().get_identity();
68
69        Ok(Box::new(Self::new(
70            child,
71            column_orders,
72            top_n_node.get_offset() as usize,
73            top_n_node.get_limit() as usize,
74            top_n_node.get_with_ties(),
75            identity.clone(),
76            source.context().get_config().developer.chunk_size,
77            source.context().create_executor_mem_context(identity),
78        )))
79    }
80}
81
82impl TopNExecutor {
83    pub fn new(
84        child: BoxedExecutor,
85        column_orders: Vec<ColumnOrder>,
86        offset: usize,
87        limit: usize,
88        with_ties: bool,
89        identity: String,
90        chunk_size: usize,
91        mem_ctx: MemoryContext,
92    ) -> Self {
93        let schema = child.schema().clone();
94        Self {
95            child,
96            column_orders,
97            offset,
98            limit,
99            with_ties,
100            schema,
101            identity,
102            chunk_size,
103            mem_ctx,
104        }
105    }
106}
107
108impl Executor for TopNExecutor {
109    fn schema(&self) -> &Schema {
110        &self.schema
111    }
112
113    fn identity(&self) -> &str {
114        &self.identity
115    }
116
117    fn execute(self: Box<Self>) -> BoxedDataChunkStream {
118        self.do_execute()
119    }
120}
121
122pub const MAX_TOPN_INIT_HEAP_CAPACITY: usize = 1024;
123
124/// A max-heap used to find the smallest `limit+offset` items.
125pub struct TopNHeap {
126    heap: MemMonitoredHeap<HeapElem>,
127    limit: usize,
128    offset: usize,
129    with_ties: bool,
130}
131
132impl TopNHeap {
133    pub fn new(limit: usize, offset: usize, with_ties: bool, mem_ctx: MemoryContext) -> Self {
134        assert!(limit > 0);
135        Self {
136            heap: MemMonitoredHeap::with_capacity(
137                (limit + offset).min(MAX_TOPN_INIT_HEAP_CAPACITY),
138                mem_ctx,
139            ),
140            limit,
141            offset,
142            with_ties,
143        }
144    }
145
146    // Only used for swapping out the heap in hashmap, due to a bug in hashmap which forbids us from
147    // using `into_iter`. We should remove this after Hashmap upgraded and fixed the bug.
148    pub fn empty() -> Self {
149        Self {
150            heap: MemMonitoredHeap::with_capacity(0, MemoryContext::none()),
151            limit: 0,
152            offset: 0,
153            with_ties: false,
154        }
155    }
156
157    pub fn push(&mut self, elem: HeapElem) {
158        if self.heap.len() < self.limit + self.offset {
159            self.heap.push(elem);
160        } else {
161            // heap is full
162            if !self.with_ties {
163                let peek = self.heap.pop().unwrap();
164                if elem < peek {
165                    self.heap.push(elem);
166                } else {
167                    self.heap.push(peek);
168                }
169                // let inner = self.heap.inner();
170                // let mut peek = inner.peek_mut().unwrap();
171                // if elem < *peek {
172                //     *peek = elem;
173                // }
174            } else {
175                let peek = self.heap.peek().unwrap().clone();
176                match elem.cmp(&peek) {
177                    Ordering::Less => {
178                        let mut ties_with_peek = vec![];
179                        // pop all the ties with peek
180                        ties_with_peek.push(self.heap.pop().unwrap());
181                        while let Some(e) = self.heap.peek()
182                            && e.encoded_row == peek.encoded_row
183                        {
184                            ties_with_peek.push(self.heap.pop().unwrap());
185                        }
186                        self.heap.push(elem);
187                        // If the size is smaller than limit, we can push all the elements back.
188                        if self.heap.len() < self.limit {
189                            self.heap.extend(ties_with_peek);
190                        }
191                    }
192                    Ordering::Equal => {
193                        // It's a tie.
194                        self.heap.push(elem);
195                    }
196                    Ordering::Greater => {}
197                }
198            }
199        }
200    }
201
202    /// Returns the elements in the range `[offset, offset+limit)`.
203    ///
204    /// # Warning
205    ///
206    /// `deallocate` only subtracts the backing-storage size. The elements may already have been
207    /// dropped, so it cannot determine the size of their separately allocated memory. Those charges
208    /// need separate cleanup, either explicitly or through a private memory context. A future
209    /// redesign should make this handling automatic.
210    ///
211    /// In this executor example, the heap is charged directly to a new private `mem_ctx`. Skipped
212    /// and consumed elements stay charged until the context is dropped on normal return, errors,
213    /// or cancellation, even though the elements themselves are dropped earlier.
214    ///
215    /// ```rust,ignore
216    /// let mem_ctx = MemoryContext::new(Some(parent), TrAdderAtomic::new(0));
217    /// let mut heap = TopNHeap::new(3, 1, false, mem_ctx.clone());
218    /// for elem in input_elements {
219    ///     heap.push(elem);
220    /// }
221    /// for elem in heap.dump() {
222    ///     let output = chunk_builder.append_one_row(elem.row());
223    ///     drop(elem);
224    ///     if let Some(output) = output {
225    ///         yield output;
226    ///     }
227    /// }
228    /// ```
229    pub fn dump(self) -> impl Iterator<Item = HeapElem> {
230        self.heap
231            .into_sorted_vec()
232            .into_iter()
233            .rev()
234            .skip(self.offset)
235    }
236}
237
238#[derive(Clone, EstimateSize)]
239pub struct HeapElem {
240    encoded_row: MemcmpEncoded,
241    row: OwnedRow,
242}
243
244impl PartialEq for HeapElem {
245    fn eq(&self, other: &Self) -> bool {
246        self.encoded_row.eq(&other.encoded_row)
247    }
248}
249
250impl Eq for HeapElem {}
251
252impl PartialOrd for HeapElem {
253    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
254        Some(self.cmp(other))
255    }
256}
257
258impl Ord for HeapElem {
259    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
260        self.encoded_row.cmp(&other.encoded_row)
261    }
262}
263
264impl HeapElem {
265    pub fn new(encoded_row: MemcmpEncoded, row: impl Row) -> Self {
266        Self {
267            encoded_row,
268            row: row.into_owned_row(),
269        }
270    }
271
272    pub fn row(&self) -> impl Row + '_ {
273        &self.row
274    }
275}
276
277impl TopNExecutor {
278    #[try_stream(boxed, ok = DataChunk, error = BatchError)]
279    async fn do_execute(self: Box<Self>) {
280        if self.limit == 0 {
281            return Ok(());
282        }
283        // Keep payload charges until this execution exits, including errors and cancellation.
284        let mem_ctx = MemoryContext::new(Some(self.mem_ctx.clone()), TrAdderAtomic::new(0));
285        let mut heap = TopNHeap::new(self.limit, self.offset, self.with_ties, mem_ctx.clone());
286
287        #[for_await]
288        for chunk in self.child.execute() {
289            let chunk = Arc::new(chunk?.compact_vis());
290            for (row_id, encoded_row) in encode_chunk(&chunk, &self.column_orders)?
291                .into_iter()
292                .enumerate()
293            {
294                heap.push(HeapElem {
295                    encoded_row,
296                    row: chunk.row_at(row_id).0.to_owned_row(),
297                });
298            }
299        }
300
301        let mut chunk_builder = DataChunkBuilder::new(self.schema.data_types(), self.chunk_size);
302        for elem in heap.dump() {
303            let output = chunk_builder.append_one_row(elem.row());
304            drop(elem);
305            if let Some(output) = output {
306                yield output
307            }
308        }
309        if let Some(spilled) = chunk_builder.consume_all() {
310            yield spilled
311        }
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use futures::stream::StreamExt;
318    use itertools::Itertools;
319    use risingwave_common::array::Array;
320    use risingwave_common::catalog::Field;
321    use risingwave_common::test_prelude::DataChunkTestExt;
322    use risingwave_common::types::DataType;
323    use risingwave_common::util::sort_util::OrderType;
324
325    use super::*;
326    use crate::executor::test_utils::MockExecutor;
327
328    const CHUNK_SIZE: usize = 1024;
329
330    #[tokio::test]
331    async fn test_simple_top_n_executor() {
332        let schema = Schema {
333            fields: vec![
334                Field::unnamed(DataType::Int32),
335                Field::unnamed(DataType::Int32),
336            ],
337        };
338        let mut mock_executor = MockExecutor::new(schema);
339        mock_executor.add(DataChunk::from_pretty(
340            "i i
341             1 5
342             2 4
343             3 3
344             4 2
345             5 1",
346        ));
347        let column_orders = vec![
348            ColumnOrder {
349                column_index: 1,
350                order_type: OrderType::ascending(),
351            },
352            ColumnOrder {
353                column_index: 0,
354                order_type: OrderType::ascending(),
355            },
356        ];
357        let top_n_executor = Box::new(TopNExecutor::new(
358            Box::new(mock_executor),
359            column_orders,
360            1,
361            3,
362            false,
363            "TopNExecutor".to_owned(),
364            CHUNK_SIZE,
365            MemoryContext::none(),
366        ));
367        let fields = &top_n_executor.schema().fields;
368        assert_eq!(fields[0].data_type, DataType::Int32);
369        assert_eq!(fields[1].data_type, DataType::Int32);
370
371        let mut stream = top_n_executor.execute();
372        let res = stream.next().await;
373
374        assert!(res.is_some());
375        if let Some(res) = res {
376            let res = res.unwrap();
377            assert_eq!(res.cardinality(), 3);
378            assert_eq!(
379                res.column_at(0).as_int32().iter().collect_vec(),
380                vec![Some(4), Some(3), Some(2)]
381            );
382        }
383
384        let res = stream.next().await;
385        assert!(res.is_none());
386    }
387
388    #[tokio::test]
389    async fn test_limit_0() {
390        let schema = Schema {
391            fields: vec![
392                Field::unnamed(DataType::Int32),
393                Field::unnamed(DataType::Int32),
394            ],
395        };
396        let mut mock_executor = MockExecutor::new(schema);
397        mock_executor.add(DataChunk::from_pretty(
398            "i i
399             1 5
400             2 4
401             3 3
402             4 2
403             5 1",
404        ));
405        let column_orders = vec![
406            ColumnOrder {
407                column_index: 1,
408                order_type: OrderType::ascending(),
409            },
410            ColumnOrder {
411                column_index: 0,
412                order_type: OrderType::ascending(),
413            },
414        ];
415        let top_n_executor = Box::new(TopNExecutor::new(
416            Box::new(mock_executor),
417            column_orders,
418            1,
419            0,
420            false,
421            "TopNExecutor".to_owned(),
422            CHUNK_SIZE,
423            MemoryContext::none(),
424        ));
425        let fields = &top_n_executor.schema().fields;
426        assert_eq!(fields[0].data_type, DataType::Int32);
427        assert_eq!(fields[1].data_type, DataType::Int32);
428
429        let mut stream = top_n_executor.execute();
430        let res = stream.next().await;
431
432        assert!(res.is_none());
433    }
434}
risingwave_batch_executors/executor/top_n.rs

risingwave_batch_executors/executor/
top_n.rs
