risingwave_batch_executors/executor/
order_by.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::sync::Arc;
16
17use bytes::Bytes;
18use futures_async_stream::try_stream;
19use itertools::Itertools;
20use prometheus::core::Atomic;
21use risingwave_common::array::DataChunk;
22use risingwave_common::catalog::Schema;
23use risingwave_common::memory::MemoryContext;
24use risingwave_common::metrics::TrAdderAtomic;
25use risingwave_common::types::DataType;
26use risingwave_common::util::chunk_coalesce::DataChunkBuilder;
27use risingwave_common::util::memcmp_encoding::encode_chunk;
28use risingwave_common::util::sort_util::ColumnOrder;
29use risingwave_common_estimate_size::EstimateSize;
30use risingwave_pb::Message;
31use risingwave_pb::batch_plan::plan_node::NodeBody;
32use risingwave_pb::data::DataChunk as PbDataChunk;
33
34use super::{
35    BoxedDataChunkStream, BoxedExecutor, BoxedExecutorBuilder, Executor, ExecutorBuilder,
36    WrapStreamExecutor,
37};
38use crate::error::{BatchError, Result};
39use crate::executor::merge_sort::MergeSortExecutor;
40use crate::monitor::BatchSpillMetrics;
41use crate::spill::spill_op::SpillBackend::Disk;
42use crate::spill::spill_op::{
43    DEFAULT_SPILL_PARTITION_NUM, SPILL_AT_LEAST_MEMORY, SpillBackend, SpillOp,
44};
45
46/// Sort Executor
47///
48/// High-level idea:
49/// 1. Load data chunks from child executor
50/// 2. Serialize each row into memcomparable format
51/// 3. Sort the serialized rows by quicksort
52/// 4. Build and yield data chunks according to the row order
53pub struct SortExecutor {
54    child: BoxedExecutor,
55    column_orders: Arc<Vec<ColumnOrder>>,
56    identity: String,
57    schema: Schema,
58    chunk_size: usize,
59    mem_context: MemoryContext,
60    spill_backend: Option<SpillBackend>,
61    spill_metrics: Arc<BatchSpillMetrics>,
62    /// The upper bound of memory usage for this executor.
63    memory_upper_bound: Option<u64>,
64}
65
66impl Executor for SortExecutor {
67    fn schema(&self) -> &Schema {
68        &self.schema
69    }
70
71    fn identity(&self) -> &str {
72        &self.identity
73    }
74
75    fn execute(self: Box<Self>) -> BoxedDataChunkStream {
76        self.do_execute()
77    }
78}
79
80impl BoxedExecutorBuilder for SortExecutor {
81    async fn new_boxed_executor(
82        source: &ExecutorBuilder<'_>,
83        inputs: Vec<BoxedExecutor>,
84    ) -> Result<BoxedExecutor> {
85        let [child]: [_; 1] = inputs.try_into().unwrap();
86
87        let order_by_node =
88            try_match_expand!(source.plan_node().get_node_body().unwrap(), NodeBody::Sort)?;
89
90        let column_orders = order_by_node
91            .column_orders
92            .iter()
93            .map(ColumnOrder::from_protobuf)
94            .collect_vec();
95
96        let identity = source.plan_node().get_identity();
97        Ok(Box::new(SortExecutor::new(
98            child,
99            Arc::new(column_orders),
100            identity.clone(),
101            source.context().get_config().developer.chunk_size,
102            source.context().create_executor_mem_context(identity),
103            if source.context().get_config().enable_spill {
104                Some(Disk)
105            } else {
106                None
107            },
108            source.context().spill_metrics(),
109        )))
110    }
111}
112
113impl SortExecutor {
114    #[try_stream(boxed, ok = DataChunk, error = BatchError)]
115    async fn do_execute(self: Box<Self>) {
116        // Each sort invocation owns a private context. Completion, errors, and cancellation release
117        // its remaining charges without waiting for the shared parent or sibling sorts to finish.
118        let mem_context = MemoryContext::new(Some(self.mem_context.clone()), TrAdderAtomic::new(0));
119        let child_schema = self.child.schema().clone();
120        let mut need_to_spill = false;
121        // If the memory upper bound is less than 1MB, we don't need to check memory usage.
122        let check_memory = match self.memory_upper_bound {
123            Some(upper_bound) => upper_bound > SPILL_AT_LEAST_MEMORY,
124            None => true,
125        };
126
127        let mut chunk_builder = DataChunkBuilder::new(self.schema.data_types(), self.chunk_size);
128        let mut chunks = Vec::new_in(mem_context.global_allocator());
129
130        let mut input_stream = self.child.execute();
131        #[for_await]
132        for chunk in &mut input_stream {
133            let chunk = chunk?.compact_vis();
134            let chunk_estimated_heap_size = chunk.estimated_heap_size();
135            chunks.push(chunk);
136            // The chunk is retained for sorting or spilling regardless of the budget check.
137            mem_context.add_unchecked(chunk_estimated_heap_size as i64);
138            if check_memory && !mem_context.check_memory_usage() {
139                if self.spill_backend.is_some() {
140                    need_to_spill = true;
141                    break;
142                } else {
143                    Err(BatchError::OutOfMemory(self.mem_context.mem_limit()))?;
144                }
145            }
146        }
147
148        let mut encoded_rows_heap_size = 0;
149        let mut encoded_rows = Vec::with_capacity_in(chunks.len(), mem_context.global_allocator());
150
151        for chunk in &chunks {
152            let encoded_chunk = encode_chunk(chunk, &self.column_orders)?;
153            let chunk_estimated_heap_size = encoded_chunk
154                .iter()
155                .map(|x| x.estimated_heap_size())
156                .sum::<usize>();
157            encoded_rows.extend(
158                encoded_chunk
159                    .into_iter()
160                    .enumerate()
161                    .map(|(row_id, row)| (chunk.row_at_unchecked_vis(row_id), row)),
162            );
163            encoded_rows_heap_size += chunk_estimated_heap_size as i64;
164            mem_context.add_unchecked(chunk_estimated_heap_size as i64);
165            if check_memory && !mem_context.check_memory_usage() {
166                if self.spill_backend.is_some() {
167                    need_to_spill = true;
168                    break;
169                } else {
170                    Err(BatchError::OutOfMemory(self.mem_context.mem_limit()))?;
171                }
172            }
173        }
174
175        if need_to_spill {
176            // A spilling version of sort, a.k.a. external sort.
177            // When SortExecutor told memory is insufficient, SortSpillManager will start to partition the sort buffer and spill to disk.
178            // After spilling the sort buffer, SortSpillManager will consume all chunks from its input executor.
179            // Finally, we would get e.g. 20 partitions. Each partition should contain a portion of the original input data.
180            // A sub SortExecutor would be used to sort each partition respectively and then a MergeSortExecutor would be used to merge all sorted partitions.
181            // If memory is still not enough in the sub SortExecutor, it will spill its inputs recursively.
182            info!("batch sort executor {} starts to spill out", &self.identity);
183            let mut sort_spill_manager = SortSpillManager::new(
184                self.spill_backend.clone().unwrap(),
185                &self.identity,
186                DEFAULT_SPILL_PARTITION_NUM,
187                child_schema.data_types(),
188                self.chunk_size,
189                self.spill_metrics.clone(),
190            )?;
191            sort_spill_manager.init_writers().await?;
192
193            // Release memory
194            drop(encoded_rows);
195            mem_context.add_unchecked(-encoded_rows_heap_size);
196
197            // Release manual buffer charges before recursive sorts charge the shared context.
198            for chunk in chunks {
199                let chunk_heap_size = chunk.estimated_heap_size() as i64;
200                sort_spill_manager.write_input_chunk(chunk).await?;
201                mem_context.add_unchecked(-chunk_heap_size);
202            }
203
204            // Spill input chunks.
205            #[for_await]
206            for chunk in input_stream {
207                let chunk: DataChunk = chunk?;
208                sort_spill_manager.write_input_chunk(chunk).await?;
209            }
210
211            sort_spill_manager.close_writers().await?;
212
213            let partition_num = sort_spill_manager.partition_num;
214            // Merge sorted-partitions
215            let mut sorted_inputs: Vec<BoxedExecutor> = Vec::with_capacity(partition_num);
216            for i in 0..partition_num {
217                let partition_size = sort_spill_manager.estimate_partition_size(i).await?;
218
219                let input_stream = sort_spill_manager.read_input_partition(i).await?;
220
221                let sub_sort_executor: SortExecutor = SortExecutor::new_inner(
222                    Box::new(WrapStreamExecutor::new(child_schema.clone(), input_stream)),
223                    self.column_orders.clone(),
224                    format!("{}-sub{}", self.identity.clone(), i),
225                    self.chunk_size,
226                    self.mem_context.clone(),
227                    self.spill_backend.clone(),
228                    self.spill_metrics.clone(),
229                    Some(partition_size),
230                );
231
232                debug!(
233                    "create sub_sort {} for sort {} to spill",
234                    sub_sort_executor.identity, self.identity
235                );
236
237                sorted_inputs.push(Box::new(sub_sort_executor));
238            }
239
240            let merge_sort = MergeSortExecutor::new(
241                sorted_inputs,
242                self.column_orders.clone(),
243                self.schema.clone(),
244                format!("{}-merge-sort", self.identity.clone()),
245                self.chunk_size,
246                self.mem_context.clone(),
247            );
248
249            #[for_await]
250            for chunk in Box::new(merge_sort).execute() {
251                yield chunk?;
252            }
253        } else {
254            encoded_rows.sort_unstable_by(|(_, a), (_, b)| a.cmp(b));
255
256            for (row, _) in encoded_rows {
257                if let Some(spilled) = chunk_builder.append_one_row(row) {
258                    yield spilled
259                }
260            }
261
262            if let Some(spilled) = chunk_builder.consume_all() {
263                yield spilled
264            }
265        }
266    }
267}
268
269impl SortExecutor {
270    pub fn new(
271        child: BoxedExecutor,
272        column_orders: Arc<Vec<ColumnOrder>>,
273        identity: String,
274        chunk_size: usize,
275        mem_context: MemoryContext,
276        spill_backend: Option<SpillBackend>,
277        spill_metrics: Arc<BatchSpillMetrics>,
278    ) -> Self {
279        Self::new_inner(
280            child,
281            column_orders,
282            identity,
283            chunk_size,
284            mem_context,
285            spill_backend,
286            spill_metrics,
287            None,
288        )
289    }
290
291    fn new_inner(
292        child: BoxedExecutor,
293        column_orders: Arc<Vec<ColumnOrder>>,
294        identity: String,
295        chunk_size: usize,
296        mem_context: MemoryContext,
297        spill_backend: Option<SpillBackend>,
298        spill_metrics: Arc<BatchSpillMetrics>,
299        memory_upper_bound: Option<u64>,
300    ) -> Self {
301        let schema = child.schema().clone();
302        Self {
303            child,
304            column_orders,
305            identity,
306            schema,
307            chunk_size,
308            mem_context,
309            spill_backend,
310            spill_metrics,
311            memory_upper_bound,
312        }
313    }
314}
315
316/// `SortSpillManager` is used to manage how to write spill data file and read them back.
317/// The spill data first need to be partitioned in a round-robin way. Each partition contains 1 file: `input_chunks_file`
318/// The spill file consume a data chunk and serialize the chunk into a protobuf bytes.
319/// Finally, spill file content will look like the below.
320/// The file write pattern is append-only and the read pattern is sequential scan.
321/// This can maximize the disk IO performance.
322///
323/// ```text
324/// [proto_len]
325/// [proto_bytes]
326/// ...
327/// [proto_len]
328/// [proto_bytes]
329/// ```
330struct SortSpillManager {
331    op: SpillOp,
332    partition_num: usize,
333    round_robin_idx: usize,
334    input_writers: Vec<opendal::Writer>,
335    input_chunk_builders: Vec<DataChunkBuilder>,
336    child_data_types: Vec<DataType>,
337    spill_chunk_size: usize,
338    spill_metrics: Arc<BatchSpillMetrics>,
339}
340
341impl SortSpillManager {
342    fn new(
343        spill_backend: SpillBackend,
344        agg_identity: &String,
345        partition_num: usize,
346        child_data_types: Vec<DataType>,
347        spill_chunk_size: usize,
348        spill_metrics: Arc<BatchSpillMetrics>,
349    ) -> Result<Self> {
350        let suffix_uuid = uuid::Uuid::new_v4();
351        let dir = format!("{}-{}/", agg_identity, suffix_uuid);
352        let op = SpillOp::create(dir, spill_backend)?;
353        let input_writers = Vec::with_capacity(partition_num);
354        let input_chunk_builders = Vec::with_capacity(partition_num);
355        Ok(Self {
356            op,
357            partition_num,
358            input_writers,
359            input_chunk_builders,
360            round_robin_idx: 0,
361            child_data_types,
362            spill_chunk_size,
363            spill_metrics,
364        })
365    }
366
367    async fn init_writers(&mut self) -> Result<()> {
368        for i in 0..self.partition_num {
369            let partition_file_name = format!("input-chunks-p{}", i);
370            let w = self.op.writer_with(&partition_file_name).await?;
371            self.input_writers.push(w);
372            self.input_chunk_builders.push(DataChunkBuilder::new(
373                self.child_data_types.clone(),
374                self.spill_chunk_size,
375            ));
376        }
377        Ok(())
378    }
379
380    async fn write_input_chunk(&mut self, chunk: DataChunk) -> Result<()> {
381        for row in chunk.rows() {
382            let partition = self.round_robin_idx;
383            if let Some(chunk) = self.input_chunk_builders[partition].append_one_row(row) {
384                let chunk_pb: PbDataChunk = chunk.to_protobuf();
385                let buf = Message::encode_to_vec(&chunk_pb);
386                let len_bytes = Bytes::copy_from_slice(&(buf.len() as u32).to_le_bytes());
387                self.spill_metrics
388                    .batch_spill_write_bytes
389                    .inc_by((buf.len() + len_bytes.len()) as u64);
390                self.input_writers[partition].write(len_bytes).await?;
391                self.input_writers[partition].write(buf).await?;
392            }
393            self.round_robin_idx = (self.round_robin_idx + 1) % self.partition_num;
394        }
395        Ok(())
396    }
397
398    async fn close_writers(&mut self) -> Result<()> {
399        for partition in 0..self.partition_num {
400            if let Some(output_chunk) = self.input_chunk_builders[partition].consume_all() {
401                let chunk_pb: PbDataChunk = output_chunk.to_protobuf();
402                let buf = Message::encode_to_vec(&chunk_pb);
403                let len_bytes = Bytes::copy_from_slice(&(buf.len() as u32).to_le_bytes());
404                self.spill_metrics
405                    .batch_spill_write_bytes
406                    .inc_by((buf.len() + len_bytes.len()) as u64);
407                self.input_writers[partition].write(len_bytes).await?;
408                self.input_writers[partition].write(buf).await?;
409            }
410        }
411
412        for mut w in self.input_writers.drain(..) {
413            w.close().await?;
414        }
415        Ok(())
416    }
417
418    async fn read_input_partition(&mut self, partition: usize) -> Result<BoxedDataChunkStream> {
419        let input_partition_file_name = format!("input-chunks-p{}", partition);
420        let r = self.op.reader_with(&input_partition_file_name).await?;
421        Ok(SpillOp::read_stream(r, self.spill_metrics.clone()))
422    }
423
424    async fn estimate_partition_size(&self, partition: usize) -> Result<u64> {
425        let input_partition_file_name = format!("input-chunks-p{}", partition);
426        let input_size = self
427            .op
428            .stat(&input_partition_file_name)
429            .await?
430            .content_length();
431        Ok(input_size)
432    }
433}
434
435#[cfg(test)]
436mod tests {
437    use futures::StreamExt;
438    use risingwave_common::array::*;
439    use risingwave_common::catalog::Field;
440    use risingwave_common::types::{Date, F32, Interval, Scalar, StructType, Time, Timestamp};
441    use risingwave_common::util::sort_util::OrderType;
442
443    use super::*;
444    use crate::executor::test_utils::MockExecutor;
445
446    const CHUNK_SIZE: usize = 1024;
447
448    #[tokio::test]
449    async fn test_simple_order_by_executor() {
450        let schema = Schema {
451            fields: vec![
452                Field::unnamed(DataType::Int32),
453                Field::unnamed(DataType::Int32),
454            ],
455        };
456        let mut mock_executor = MockExecutor::new(schema);
457        mock_executor.add(DataChunk::from_pretty(
458            "i i
459             1 3
460             2 2
461             3 1",
462        ));
463        let column_orders = vec![
464            ColumnOrder {
465                column_index: 1,
466                order_type: OrderType::ascending(),
467            },
468            ColumnOrder {
469                column_index: 0,
470                order_type: OrderType::ascending(),
471            },
472        ];
473
474        let order_by_executor = Box::new(SortExecutor::new(
475            Box::new(mock_executor),
476            Arc::new(column_orders),
477            "SortExecutor2".to_owned(),
478            CHUNK_SIZE,
479            MemoryContext::none(),
480            None,
481            BatchSpillMetrics::for_test(),
482        ));
483        let fields = &order_by_executor.schema().fields;
484        assert_eq!(fields[0].data_type, DataType::Int32);
485        assert_eq!(fields[1].data_type, DataType::Int32);
486
487        let mut stream = order_by_executor.execute();
488        let res = stream.next().await;
489        assert!(res.is_some());
490        if let Some(res) = res {
491            let res = res.unwrap();
492            let col0 = res.column_at(0);
493            assert_eq!(col0.as_int32().value_at(0), Some(3));
494            assert_eq!(col0.as_int32().value_at(1), Some(2));
495            assert_eq!(col0.as_int32().value_at(2), Some(1));
496        }
497    }
498
499    #[tokio::test]
500    async fn test_encoding_for_float() {
501        let schema = Schema {
502            fields: vec![
503                Field::unnamed(DataType::Float32),
504                Field::unnamed(DataType::Float64),
505            ],
506        };
507        let mut mock_executor = MockExecutor::new(schema);
508        mock_executor.add(DataChunk::from_pretty(
509            " f    F
510             -2.2  3.3
511             -1.1  2.2
512              1.1  1.1
513              2.2 -1.1
514              3.3 -2.2",
515        ));
516        let column_orders = vec![
517            ColumnOrder {
518                column_index: 1,
519                order_type: OrderType::ascending(),
520            },
521            ColumnOrder {
522                column_index: 0,
523                order_type: OrderType::ascending(),
524            },
525        ];
526        let order_by_executor = Box::new(SortExecutor::new(
527            Box::new(mock_executor),
528            Arc::new(column_orders),
529            "SortExecutor2".to_owned(),
530            CHUNK_SIZE,
531            MemoryContext::none(),
532            None,
533            BatchSpillMetrics::for_test(),
534        ));
535        let fields = &order_by_executor.schema().fields;
536        assert_eq!(fields[0].data_type, DataType::Float32);
537        assert_eq!(fields[1].data_type, DataType::Float64);
538
539        let mut stream = order_by_executor.execute();
540        let res = stream.next().await;
541        assert!(res.is_some());
542        if let Some(res) = res {
543            let res = res.unwrap();
544            let col0 = res.column_at(0);
545            assert_eq!(col0.as_float32().value_at(0), Some(3.3.into()));
546            assert_eq!(col0.as_float32().value_at(1), Some(2.2.into()));
547            assert_eq!(col0.as_float32().value_at(2), Some(1.1.into()));
548            assert_eq!(col0.as_float32().value_at(3), Some((-1.1).into()));
549            assert_eq!(col0.as_float32().value_at(4), Some((-2.2).into()));
550        }
551    }
552
553    #[tokio::test]
554    async fn test_bsc_for_string() {
555        let schema = Schema {
556            fields: vec![
557                Field::unnamed(DataType::Varchar),
558                Field::unnamed(DataType::Varchar),
559            ],
560        };
561        let mut mock_executor = MockExecutor::new(schema);
562        mock_executor.add(DataChunk::from_pretty(
563            "T   T
564             1.1 3.3
565             2.2 2.2
566             3.3 1.1",
567        ));
568        let column_orders = vec![
569            ColumnOrder {
570                column_index: 1,
571                order_type: OrderType::ascending(),
572            },
573            ColumnOrder {
574                column_index: 0,
575                order_type: OrderType::ascending(),
576            },
577        ];
578        let order_by_executor = Box::new(SortExecutor::new(
579            Box::new(mock_executor),
580            Arc::new(column_orders),
581            "SortExecutor2".to_owned(),
582            CHUNK_SIZE,
583            MemoryContext::none(),
584            None,
585            BatchSpillMetrics::for_test(),
586        ));
587        let fields = &order_by_executor.schema().fields;
588        assert_eq!(fields[0].data_type, DataType::Varchar);
589        assert_eq!(fields[1].data_type, DataType::Varchar);
590
591        let mut stream = order_by_executor.execute();
592        let res = stream.next().await;
593        assert!(res.is_some());
594        if let Some(res) = res {
595            let res = res.unwrap();
596            let col0 = res.column_at(0);
597            assert_eq!(col0.as_utf8().value_at(0), Some("3.3"));
598            assert_eq!(col0.as_utf8().value_at(1), Some("2.2"));
599            assert_eq!(col0.as_utf8().value_at(2), Some("1.1"));
600        }
601    }
602
603    // TODO: write following tests in a more concise way
604    #[tokio::test]
605    async fn test_encoding_for_boolean_int32_float64() {
606        let schema = Schema {
607            fields: vec![
608                Field::unnamed(DataType::Boolean),
609                Field::unnamed(DataType::Int32),
610                Field::unnamed(DataType::Float64),
611            ],
612        };
613        // f   3    .
614        // t   3    .
615        // .   .    3.5
616        // .   .    -4.3
617        // .   .    .
618        let input_chunk = DataChunk::new(
619            vec![
620                BoolArray::from_iter([Some(false), Some(true), None, None, None]).into_ref(),
621                I32Array::from_iter([Some(3), Some(3), None, None, None]).into_ref(),
622                F64Array::from_iter([None, None, Some(3.5), Some(-4.3), None]).into_ref(),
623            ],
624            5,
625        );
626        // .   .   -4.3
627        // .   .   3.5
628        // .   .   .
629        // f   3   .
630        // t   3   .
631        let output_chunk = DataChunk::new(
632            vec![
633                BoolArray::from_iter([None, None, None, Some(false), Some(true)]).into_ref(),
634                I32Array::from_iter([None, None, None, Some(3), Some(3)]).into_ref(),
635                F64Array::from_iter([Some(-4.3), Some(3.5), None, None, None]).into_ref(),
636            ],
637            5,
638        );
639        let mut mock_executor = MockExecutor::new(schema);
640        mock_executor.add(input_chunk);
641        let column_orders = vec![
642            ColumnOrder {
643                column_index: 2,
644                order_type: OrderType::ascending(),
645            },
646            ColumnOrder {
647                column_index: 1,
648                order_type: OrderType::descending(),
649            },
650            ColumnOrder {
651                column_index: 0,
652                order_type: OrderType::ascending(),
653            },
654        ];
655        let order_by_executor = Box::new(SortExecutor::new(
656            Box::new(mock_executor),
657            Arc::new(column_orders),
658            "SortExecutor".to_owned(),
659            CHUNK_SIZE,
660            MemoryContext::none(),
661            None,
662            BatchSpillMetrics::for_test(),
663        ));
664
665        let mut stream = order_by_executor.execute();
666        let res = stream.next().await;
667        assert_eq!(res.unwrap().unwrap(), output_chunk)
668    }
669
670    #[tokio::test]
671    async fn test_encoding_for_decimal_date_varchar() {
672        let schema = Schema {
673            fields: vec![
674                Field::unnamed(DataType::Varchar),
675                Field::unnamed(DataType::Decimal),
676                Field::unnamed(DataType::Date),
677            ],
678        };
679        // abc       .     123
680        // b         -3    789
681        // abc       .     456
682        // abcdefgh  .     .
683        // b         7     345
684        let input_chunk = DataChunk::new(
685            vec![
686                Utf8Array::from_iter(["abc", "b", "abc", "abcdefgh", "b"]).into_ref(),
687                DecimalArray::from_iter([None, Some((-3).into()), None, None, Some(7.into())])
688                    .into_ref(),
689                DateArray::from_iter([
690                    Some(Date::with_days_since_ce(123).unwrap()),
691                    Some(Date::with_days_since_ce(789).unwrap()),
692                    Some(Date::with_days_since_ce(456).unwrap()),
693                    None,
694                    Some(Date::with_days_since_ce(345).unwrap()),
695                ])
696                .into_ref(),
697            ],
698            5,
699        );
700        // b         7     345
701        // b         -3    789
702        // abcdefgh  .     .
703        // abc       .     123
704        // abc       .     456
705        let output_chunk = DataChunk::new(
706            vec![
707                Utf8Array::from_iter(["b", "b", "abcdefgh", "abc", "abc"]).into_ref(),
708                DecimalArray::from_iter([Some(7.into()), Some((-3).into()), None, None, None])
709                    .into_ref(),
710                DateArray::from_iter([
711                    Some(Date::with_days_since_ce(345).unwrap()),
712                    Some(Date::with_days_since_ce(789).unwrap()),
713                    None,
714                    Some(Date::with_days_since_ce(123).unwrap()),
715                    Some(Date::with_days_since_ce(456).unwrap()),
716                ])
717                .into_ref(),
718            ],
719            5,
720        );
721        let mut mock_executor = MockExecutor::new(schema);
722        mock_executor.add(input_chunk);
723        let column_orders = vec![
724            ColumnOrder {
725                column_index: 0,
726                order_type: OrderType::descending(),
727            },
728            ColumnOrder {
729                column_index: 1,
730                order_type: OrderType::descending(),
731            },
732            ColumnOrder {
733                column_index: 2,
734                order_type: OrderType::ascending(),
735            },
736        ];
737        let order_by_executor = Box::new(SortExecutor::new(
738            Box::new(mock_executor),
739            Arc::new(column_orders),
740            "SortExecutor".to_owned(),
741            CHUNK_SIZE,
742            MemoryContext::none(),
743            None,
744            BatchSpillMetrics::for_test(),
745        ));
746
747        let mut stream = order_by_executor.execute();
748        let res = stream.next().await;
749        assert_eq!(res.unwrap().unwrap(), output_chunk)
750    }
751
752    #[tokio::test]
753    async fn test_encoding_for_time_timestamp_interval() {
754        let schema = Schema {
755            fields: vec![
756                Field::unnamed(DataType::Time),
757                Field::unnamed(DataType::Timestamp),
758                Field::unnamed(DataType::Interval),
759            ],
760        };
761        // .     1:23  .
762        // 4:56  4:56  1:2:3
763        // .     7:89  .
764        // 4:56  4:56  4:5:6
765        // 7:89  .     .
766        let input_chunk = DataChunk::new(
767            vec![
768                TimeArray::from_iter([
769                    None,
770                    Some(Time::with_secs_nano(4, 56).unwrap()),
771                    None,
772                    Some(Time::with_secs_nano(4, 56).unwrap()),
773                    Some(Time::with_secs_nano(7, 89).unwrap()),
774                ])
775                .into_ref(),
776                TimestampArray::from_iter([
777                    Some(Timestamp::with_secs_nsecs(1, 23).unwrap()),
778                    Some(Timestamp::with_secs_nsecs(4, 56).unwrap()),
779                    Some(Timestamp::with_secs_nsecs(7, 89).unwrap()),
780                    Some(Timestamp::with_secs_nsecs(4, 56).unwrap()),
781                    None,
782                ])
783                .into_ref(),
784                IntervalArray::from_iter([
785                    None,
786                    Some(Interval::from_month_day_usec(1, 2, 3)),
787                    None,
788                    Some(Interval::from_month_day_usec(4, 5, 6)),
789                    None,
790                ])
791                .into_ref(),
792            ],
793            5,
794        );
795        // 4:56  4:56  4:5:6
796        // 4:56  4:56  1:2:3
797        // 7:89  .     .
798        // .     1:23  .
799        // .     7:89  .
800        let output_chunk = DataChunk::new(
801            vec![
802                TimeArray::from_iter([
803                    Some(Time::with_secs_nano(4, 56).unwrap()),
804                    Some(Time::with_secs_nano(4, 56).unwrap()),
805                    Some(Time::with_secs_nano(7, 89).unwrap()),
806                    None,
807                    None,
808                ])
809                .into_ref(),
810                TimestampArray::from_iter([
811                    Some(Timestamp::with_secs_nsecs(4, 56).unwrap()),
812                    Some(Timestamp::with_secs_nsecs(4, 56).unwrap()),
813                    None,
814                    Some(Timestamp::with_secs_nsecs(1, 23).unwrap()),
815                    Some(Timestamp::with_secs_nsecs(7, 89).unwrap()),
816                ])
817                .into_ref(),
818                IntervalArray::from_iter([
819                    Some(Interval::from_month_day_usec(4, 5, 6)),
820                    Some(Interval::from_month_day_usec(1, 2, 3)),
821                    None,
822                    None,
823                    None,
824                ])
825                .into_ref(),
826            ],
827            5,
828        );
829        let mut mock_executor = MockExecutor::new(schema);
830        mock_executor.add(input_chunk);
831        let column_orders = vec![
832            ColumnOrder {
833                column_index: 0,
834                order_type: OrderType::ascending(),
835            },
836            ColumnOrder {
837                column_index: 1,
838                order_type: OrderType::ascending(),
839            },
840            ColumnOrder {
841                column_index: 2,
842                order_type: OrderType::descending(),
843            },
844        ];
845        let order_by_executor = Box::new(SortExecutor::new(
846            Box::new(mock_executor),
847            column_orders.into(),
848            "SortExecutor".to_owned(),
849            CHUNK_SIZE,
850            MemoryContext::none(),
851            None,
852            BatchSpillMetrics::for_test(),
853        ));
854
855        let mut stream = order_by_executor.execute();
856        let res = stream.next().await;
857        assert_eq!(res.unwrap().unwrap(), output_chunk)
858    }
859
860    #[tokio::test]
861    async fn test_encoding_for_struct_list() {
862        let schema = Schema {
863            fields: vec![
864                Field::unnamed(
865                    StructType::unnamed(vec![DataType::Varchar, DataType::Float32]).into(),
866                ),
867                Field::unnamed(DataType::Int64.list()),
868            ],
869        };
870        let mut struct_builder = StructArrayBuilder::with_type(
871            0,
872            DataType::Struct(StructType::unnamed(vec![
873                DataType::Varchar,
874                DataType::Float32,
875            ])),
876        );
877        let mut list_builder = ListArrayBuilder::with_type(0, DataType::Int64.list());
878        // {abcd, -1.2}   .
879        // {c, 0}         [1, ., 3]
880        // {c, .}         .
881        // {c, 0}         [2]
882        // {., 3.4}       .
883        let input_chunk = DataChunk::new(
884            vec![
885                {
886                    struct_builder.append(Some(StructRef::ValueRef {
887                        val: &StructValue::new(vec![
888                            Some("abcd".into()),
889                            Some(F32::from(-1.2).to_scalar_value()),
890                        ]),
891                    }));
892                    struct_builder.append(Some(StructRef::ValueRef {
893                        val: &StructValue::new(vec![
894                            Some("c".into()),
895                            Some(F32::from(0.0).to_scalar_value()),
896                        ]),
897                    }));
898                    struct_builder.append(Some(StructRef::ValueRef {
899                        val: &StructValue::new(vec![Some("c".into()), None]),
900                    }));
901                    struct_builder.append(Some(StructRef::ValueRef {
902                        val: &StructValue::new(vec![
903                            Some("c".into()),
904                            Some(F32::from(0.0).to_scalar_value()),
905                        ]),
906                    }));
907                    struct_builder.append(Some(StructRef::ValueRef {
908                        val: &StructValue::new(vec![None, Some(F32::from(3.4).to_scalar_value())]),
909                    }));
910                    struct_builder.finish().into_ref()
911                },
912                {
913                    list_builder.append(None);
914                    list_builder.append(Some(
915                        ListValue::from_iter([Some(1i64), None, Some(3i64)]).as_scalar_ref(),
916                    ));
917                    list_builder.append(None);
918                    list_builder.append(Some(ListValue::from_iter([2i64]).as_scalar_ref()));
919                    list_builder.append(None);
920                    list_builder.finish().into_ref()
921                },
922            ],
923            5,
924        );
925        let mut struct_builder = StructArrayBuilder::with_type(
926            0,
927            DataType::Struct(StructType::unnamed(vec![
928                DataType::Varchar,
929                DataType::Float32,
930            ])),
931        );
932        let mut list_builder = ListArrayBuilder::with_type(0, DataType::Int64.list());
933        // {abcd, -1.2}   .
934        // {c, 0}         [2]
935        // {c, 0}         [1, ., 3]
936        // {c, .}         .
937        // {., 3.4}       .
938        let output_chunk = DataChunk::new(
939            vec![
940                {
941                    struct_builder.append(Some(StructRef::ValueRef {
942                        val: &StructValue::new(vec![
943                            Some("abcd".into()),
944                            Some(F32::from(-1.2).to_scalar_value()),
945                        ]),
946                    }));
947                    struct_builder.append(Some(StructRef::ValueRef {
948                        val: &StructValue::new(vec![
949                            Some("c".into()),
950                            Some(F32::from(0.0).to_scalar_value()),
951                        ]),
952                    }));
953                    struct_builder.append(Some(StructRef::ValueRef {
954                        val: &StructValue::new(vec![
955                            Some("c".into()),
956                            Some(F32::from(0.0).to_scalar_value()),
957                        ]),
958                    }));
959                    struct_builder.append(Some(StructRef::ValueRef {
960                        val: &StructValue::new(vec![Some("c".into()), None]),
961                    }));
962                    struct_builder.append(Some(StructRef::ValueRef {
963                        val: &StructValue::new(vec![None, Some(F32::from(3.4).to_scalar_value())]),
964                    }));
965                    struct_builder.finish().into_ref()
966                },
967                {
968                    list_builder.append(None);
969                    list_builder.append(Some(ListValue::from_iter([2i64]).as_scalar_ref()));
970                    list_builder.append(Some(
971                        ListValue::from_iter([Some(1i64), None, Some(3i64)]).as_scalar_ref(),
972                    ));
973                    list_builder.append(None);
974                    list_builder.append(None);
975                    list_builder.finish().into_ref()
976                },
977            ],
978            5,
979        );
980        let mut mock_executor = MockExecutor::new(schema);
981        mock_executor.add(input_chunk);
982        let column_orders = vec![
983            ColumnOrder {
984                column_index: 0,
985                order_type: OrderType::ascending(),
986            },
987            ColumnOrder {
988                column_index: 1,
989                order_type: OrderType::descending(),
990            },
991        ];
992        let order_by_executor = Box::new(SortExecutor::new(
993            Box::new(mock_executor),
994            Arc::new(column_orders),
995            "SortExecutor".to_owned(),
996            CHUNK_SIZE,
997            MemoryContext::none(),
998            None,
999            BatchSpillMetrics::for_test(),
1000        ));
1001
1002        let mut stream = order_by_executor.execute();
1003        let res = stream.next().await;
1004        assert_eq!(res.unwrap().unwrap(), output_chunk)
1005    }
1006
1007    #[tokio::test]
1008    async fn test_spill_out() {
1009        let schema = Schema {
1010            fields: vec![
1011                Field::unnamed(DataType::Float32),
1012                Field::unnamed(DataType::Float64),
1013            ],
1014        };
1015        let mut mock_executor = MockExecutor::new(schema);
1016        mock_executor.add(DataChunk::from_pretty(
1017            " f    F
1018             -2.2  3.3
1019             -1.1  2.2
1020              1.1  1.1
1021              2.2 -1.1
1022              3.3 -2.2",
1023        ));
1024        let column_orders = vec![
1025            ColumnOrder {
1026                column_index: 1,
1027                order_type: OrderType::ascending(),
1028            },
1029            ColumnOrder {
1030                column_index: 0,
1031                order_type: OrderType::ascending(),
1032            },
1033        ];
1034        let order_by_executor = Box::new(SortExecutor::new(
1035            Box::new(mock_executor),
1036            Arc::new(column_orders),
1037            "SortExecutor2".to_owned(),
1038            CHUNK_SIZE,
1039            MemoryContext::for_spill_test(),
1040            Some(SpillBackend::Memory),
1041            BatchSpillMetrics::for_test(),
1042        ));
1043        let fields = &order_by_executor.schema().fields;
1044        assert_eq!(fields[0].data_type, DataType::Float32);
1045        assert_eq!(fields[1].data_type, DataType::Float64);
1046
1047        let mut stream = order_by_executor.execute();
1048        let res = stream.next().await;
1049        assert!(res.is_some());
1050        if let Some(res) = res {
1051            let res = res.unwrap();
1052            let col0 = res.column_at(0);
1053            assert_eq!(col0.as_float32().value_at(0), Some(3.3.into()));
1054            assert_eq!(col0.as_float32().value_at(1), Some(2.2.into()));
1055            assert_eq!(col0.as_float32().value_at(2), Some(1.1.into()));
1056            assert_eq!(col0.as_float32().value_at(3), Some((-1.1).into()));
1057            assert_eq!(col0.as_float32().value_at(4), Some((-2.2).into()));
1058        }
1059    }
1060}
risingwave_batch_executors/executor/order_by.rs

risingwave_batch_executors/executor/
order_by.rs
