1use common_time::Timestamp;
18use futures::future::BoxFuture;
19use smallvec::{SmallVec, smallvec};
20use store_api::region_engine::PartitionRange;
21use store_api::storage::TimeSeriesDistribution;
22
23use crate::cache::CacheStrategy;
24use crate::memtable::{MemtableRange, MemtableStats};
25use crate::read::scan_region::ScanInput;
26use crate::sst::file::{FileHandle, FileTimeRange, overlaps};
27use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE;
28use crate::sst::parquet::file_range::{FileRange, FileRangeContextRef};
29use crate::sst::parquet::format::parquet_row_group_time_range;
30use crate::sst::parquet::row_selection::RowGroupSelection;
31
32const ALL_ROW_GROUPS: i64 = -1;
33
34#[derive(Debug, Clone, Copy, PartialEq)]
36pub(crate) struct SourceIndex {
37 pub(crate) index: usize,
39 pub(crate) num_row_groups: u64,
42}
43
44#[derive(Debug, Clone, Copy, PartialEq)]
46pub struct RowGroupIndex {
47 pub(crate) index: usize,
49 pub row_group_index: i64,
52}
53
54#[derive(Debug, PartialEq, Clone)]
58pub(crate) struct RangeMeta {
59 pub(crate) time_range: FileTimeRange,
61 pub(crate) indices: SmallVec<[SourceIndex; 2]>,
63 pub(crate) row_group_indices: SmallVec<[RowGroupIndex; 2]>,
65 pub(crate) num_rows: usize,
67}
68
69impl RangeMeta {
70 pub(crate) fn new_partition_range(&self, identifier: usize) -> PartitionRange {
73 PartitionRange {
74 start: self.time_range.0,
75 end: Timestamp::new(
76 self.time_range
79 .1
80 .value()
81 .checked_add(1)
82 .unwrap_or(self.time_range.1.value()),
83 self.time_range.1.unit(),
84 ),
85 num_rows: self.num_rows,
86 identifier,
87 }
88 }
89
90 pub(crate) fn seq_scan_ranges(input: &ScanInput) -> Vec<RangeMeta> {
93 let mut ranges = Vec::with_capacity(input.memtables.len() + input.files.len());
94 Self::push_seq_mem_ranges(&input.memtables, &mut ranges);
95 Self::push_seq_file_ranges(input.memtables.len(), &input.files, &mut ranges);
96
97 #[cfg(feature = "enterprise")]
98 Self::push_extension_ranges(input, &mut ranges);
99
100 let ranges = group_ranges_for_seq_scan(ranges);
101 if input.compaction || input.distribution == Some(TimeSeriesDistribution::PerSeries) {
102 return ranges;
104 }
105 maybe_split_ranges_for_seq_scan(ranges)
106 }
107
108 pub(crate) fn unordered_scan_ranges(input: &ScanInput) -> BoxFuture<'_, Vec<RangeMeta>> {
110 Box::pin(async move {
112 let mut ranges = Vec::with_capacity(input.memtables.len() + input.files.len());
113 Self::push_unordered_mem_ranges(&input.memtables, &mut ranges);
114 Self::push_unordered_file_ranges(
115 input.memtables.len(),
116 &input.files,
117 &input.cache_strategy,
118 &mut ranges,
119 )
120 .await;
121
122 #[cfg(feature = "enterprise")]
123 Self::push_extension_ranges(input, &mut ranges);
124
125 ranges
126 })
127 }
128
129 fn overlaps(&self, meta: &RangeMeta) -> bool {
131 overlaps(&self.time_range, &meta.time_range)
132 }
133
134 fn merge(&mut self, mut other: RangeMeta) {
137 debug_assert!(self.overlaps(&other));
138 debug_assert!(self.indices.iter().all(|idx| !other.indices.contains(idx)));
139 debug_assert!(
140 self.row_group_indices
141 .iter()
142 .all(|idx| !other.row_group_indices.contains(idx))
143 );
144
145 self.time_range = (
146 self.time_range.0.min(other.time_range.0),
147 self.time_range.1.max(other.time_range.1),
148 );
149 self.indices.append(&mut other.indices);
150 self.row_group_indices.append(&mut other.row_group_indices);
151 self.num_rows += other.num_rows;
152 }
153
154 fn can_split_preserve_order(&self) -> bool {
157 self.indices.len() == 1 && self.indices[0].num_row_groups > 1
158 }
159
160 fn maybe_split(self, output: &mut Vec<RangeMeta>) {
162 if self.can_split_preserve_order() {
163 let num_row_groups = self.indices[0].num_row_groups;
164 debug_assert_eq!(1, self.row_group_indices.len());
165 debug_assert_eq!(ALL_ROW_GROUPS, self.row_group_indices[0].row_group_index);
166
167 output.reserve(self.row_group_indices.len());
168 let num_rows = self.num_rows / num_row_groups as usize;
169 for row_group_index in 0..num_row_groups {
171 output.push(RangeMeta {
172 time_range: self.time_range,
173 indices: self.indices.clone(),
174 row_group_indices: smallvec![RowGroupIndex {
175 index: self.indices[0].index,
176 row_group_index: row_group_index as i64,
177 }],
178 num_rows,
179 });
180 }
181 } else {
182 output.push(self);
183 }
184 }
185
186 fn push_unordered_mem_ranges(memtables: &[MemRangeBuilder], ranges: &mut Vec<RangeMeta>) {
187 for (memtable_index, memtable) in memtables.iter().enumerate() {
189 let stats = memtable.stats();
190 let Some(time_range) = stats.time_range() else {
191 continue;
192 };
193 for row_group_index in 0..stats.num_ranges() {
194 let num_rows = stats.num_rows() / stats.num_ranges();
195 ranges.push(RangeMeta {
196 time_range,
197 indices: smallvec![SourceIndex {
198 index: memtable_index,
199 num_row_groups: stats.num_ranges() as u64,
200 }],
201 row_group_indices: smallvec![RowGroupIndex {
202 index: memtable_index,
203 row_group_index: row_group_index as i64,
204 }],
205 num_rows,
206 });
207 }
208 }
209 }
210
211 async fn push_unordered_file_ranges(
212 num_memtables: usize,
213 files: &[FileHandle],
214 cache: &CacheStrategy,
215 ranges: &mut Vec<RangeMeta>,
216 ) {
217 for (i, file) in files.iter().enumerate() {
219 let file_index = num_memtables + i;
220 let parquet_meta = cache
222 .get_parquet_meta_data_from_mem_cache(file.file_id())
223 .await;
224 if let Some(parquet_meta) = parquet_meta {
225 for row_group_index in 0..file.meta_ref().num_row_groups {
227 let time_range = parquet_row_group_time_range(
228 file.meta_ref(),
229 &parquet_meta,
230 row_group_index as usize,
231 );
232 let num_rows = parquet_meta.row_group(row_group_index as usize).num_rows();
233 ranges.push(RangeMeta {
234 time_range: time_range.unwrap_or_else(|| file.time_range()),
235 indices: smallvec![SourceIndex {
236 index: file_index,
237 num_row_groups: file.meta_ref().num_row_groups,
238 }],
239 row_group_indices: smallvec![RowGroupIndex {
240 index: file_index,
241 row_group_index: row_group_index as i64,
242 }],
243 num_rows: num_rows as usize,
244 });
245 }
246 } else if file.meta_ref().num_row_groups > 0 {
247 for row_group_index in 0..file.meta_ref().num_row_groups {
249 ranges.push(RangeMeta {
250 time_range: file.time_range(),
251 indices: smallvec![SourceIndex {
252 index: file_index,
253 num_row_groups: file.meta_ref().num_row_groups,
254 }],
255 row_group_indices: smallvec![RowGroupIndex {
256 index: file_index,
257 row_group_index: row_group_index as i64,
258 }],
259 num_rows: DEFAULT_ROW_GROUP_SIZE,
260 });
261 }
262 } else {
263 ranges.push(RangeMeta {
265 time_range: file.time_range(),
266 indices: smallvec![SourceIndex {
267 index: file_index,
268 num_row_groups: 0,
269 }],
270 row_group_indices: smallvec![RowGroupIndex {
271 index: file_index,
272 row_group_index: ALL_ROW_GROUPS,
273 }],
274 num_rows: file.meta_ref().num_rows as usize,
276 });
277 }
278 }
279 }
280
281 fn push_seq_mem_ranges(memtables: &[MemRangeBuilder], ranges: &mut Vec<RangeMeta>) {
282 for (i, memtable) in memtables.iter().enumerate() {
284 let stats = memtable.stats();
285 let Some(time_range) = stats.time_range() else {
286 continue;
287 };
288 ranges.push(RangeMeta {
289 time_range,
290 indices: smallvec![SourceIndex {
291 index: i,
292 num_row_groups: stats.num_ranges() as u64,
293 }],
294 row_group_indices: smallvec![RowGroupIndex {
295 index: i,
296 row_group_index: ALL_ROW_GROUPS,
297 }],
298 num_rows: stats.num_rows(),
299 });
300 }
301 }
302
303 fn push_seq_file_ranges(
304 num_memtables: usize,
305 files: &[FileHandle],
306 ranges: &mut Vec<RangeMeta>,
307 ) {
308 for (i, file) in files.iter().enumerate() {
310 let file_index = num_memtables + i;
311 ranges.push(RangeMeta {
312 time_range: file.time_range(),
313 indices: smallvec![SourceIndex {
314 index: file_index,
315 num_row_groups: file.meta_ref().num_row_groups,
316 }],
317 row_group_indices: smallvec![RowGroupIndex {
318 index: file_index,
319 row_group_index: ALL_ROW_GROUPS,
320 }],
321 num_rows: file.meta_ref().num_rows as usize,
322 });
323 }
324 }
325
326 #[cfg(feature = "enterprise")]
327 fn push_extension_ranges(input: &ScanInput, metas: &mut Vec<RangeMeta>) {
328 for (i, range) in input.extension_ranges().iter().enumerate() {
329 let index = input.num_memtables() + input.num_files() + i;
330 metas.push(RangeMeta {
331 time_range: range.time_range(),
332 indices: smallvec![SourceIndex {
333 index,
334 num_row_groups: range.num_row_groups(),
335 }],
336 row_group_indices: smallvec![RowGroupIndex {
337 index,
338 row_group_index: ALL_ROW_GROUPS,
339 }],
340 num_rows: range.num_rows() as usize,
341 });
342 }
343 }
344}
345
346fn group_ranges_for_seq_scan(mut ranges: Vec<RangeMeta>) -> Vec<RangeMeta> {
349 if ranges.is_empty() {
350 return ranges;
351 }
352
353 ranges.sort_unstable_by(|a, b| {
355 let l = a.time_range;
356 let r = b.time_range;
357 l.0.cmp(&r.0).then_with(|| r.1.cmp(&l.1))
358 });
359 let mut range_in_progress = None;
360 let mut exclusive_ranges = Vec::with_capacity(ranges.len());
362 for range in ranges {
363 let Some(mut prev_range) = range_in_progress.take() else {
364 range_in_progress = Some(range);
366 continue;
367 };
368
369 if prev_range.overlaps(&range) {
370 prev_range.merge(range);
371 range_in_progress = Some(prev_range);
372 } else {
373 exclusive_ranges.push(prev_range);
374 range_in_progress = Some(range);
375 }
376 }
377 if let Some(range) = range_in_progress {
378 exclusive_ranges.push(range);
379 }
380
381 exclusive_ranges
382}
383
384fn maybe_split_ranges_for_seq_scan(ranges: Vec<RangeMeta>) -> Vec<RangeMeta> {
387 let mut new_ranges = Vec::with_capacity(ranges.len());
388 for range in ranges {
389 range.maybe_split(&mut new_ranges);
390 }
391
392 new_ranges
393}
394
395#[derive(Default)]
397pub struct FileRangeBuilder {
398 context: Option<FileRangeContextRef>,
401 selection: RowGroupSelection,
403}
404
405impl FileRangeBuilder {
406 pub fn new(context: FileRangeContextRef, selection: RowGroupSelection) -> Self {
408 Self {
409 context: Some(context),
410 selection,
411 }
412 }
413
414 pub fn build_ranges(&self, row_group_index: i64, ranges: &mut SmallVec<[FileRange; 2]>) {
417 let Some(context) = self.context.clone() else {
418 return;
419 };
420 if row_group_index >= 0 {
421 let row_group_index = row_group_index as usize;
422 let Some(row_selection) = self.selection.get(row_group_index) else {
424 return;
425 };
426 ranges.push(FileRange::new(
427 context,
428 row_group_index,
429 Some(row_selection.clone()),
430 ));
431 } else {
432 ranges.extend(
434 self.selection
435 .iter()
436 .map(|(row_group_index, row_selection)| {
437 FileRange::new(
438 context.clone(),
439 *row_group_index,
440 Some(row_selection.clone()),
441 )
442 }),
443 );
444 }
445 }
446
447 pub(crate) fn memory_size(&self) -> usize {
449 let context_size = self
450 .context
451 .as_ref()
452 .map(|ctx| ctx.memory_size())
453 .unwrap_or(0);
454 let selection_size = self.selection.mem_usage();
455 context_size + selection_size
456 }
457}
458
459#[derive(Clone)]
461pub(crate) struct MemRangeBuilder {
462 range: MemtableRange,
464 stats: MemtableStats,
466}
467
468impl MemRangeBuilder {
469 pub(crate) fn new(range: MemtableRange, stats: MemtableStats) -> Self {
471 Self { range, stats }
472 }
473
474 pub(crate) fn build_ranges(
477 &self,
478 _row_group_index: i64,
479 ranges: &mut SmallVec<[MemtableRange; 2]>,
480 ) {
481 ranges.push(self.range.clone())
482 }
483
484 pub(crate) fn stats(&self) -> &MemtableStats {
486 &self.stats
487 }
488}
489
490#[cfg(test)]
491mod tests {
492 use common_time::Timestamp;
493 use common_time::timestamp::TimeUnit;
494
495 use super::*;
496
497 type Output = (Vec<usize>, i64, i64);
498
499 fn run_group_ranges_test(input: &[(usize, i64, i64)], expect: &[Output]) {
500 let ranges = input
501 .iter()
502 .map(|(idx, start, end)| {
503 let time_range = (
504 Timestamp::new(*start, TimeUnit::Second),
505 Timestamp::new(*end, TimeUnit::Second),
506 );
507 RangeMeta {
508 time_range,
509 indices: smallvec![SourceIndex {
510 index: *idx,
511 num_row_groups: 0,
512 }],
513 row_group_indices: smallvec![RowGroupIndex {
514 index: *idx,
515 row_group_index: 0
516 }],
517 num_rows: 1,
518 }
519 })
520 .collect();
521 let output = group_ranges_for_seq_scan(ranges);
522 let actual: Vec<_> = output
523 .iter()
524 .map(|range| {
525 let indices = range.indices.iter().map(|index| index.index).collect();
526 let group_indices: Vec<_> = range
527 .row_group_indices
528 .iter()
529 .map(|idx| idx.index)
530 .collect();
531 assert_eq!(indices, group_indices);
532 let range = range.time_range;
533 (indices, range.0.value(), range.1.value())
534 })
535 .collect();
536 assert_eq!(expect, actual);
537 }
538
539 #[test]
540 fn test_group_ranges() {
541 run_group_ranges_test(&[(1, 0, 2000)], &[(vec![1], 0, 2000)]);
543
544 run_group_ranges_test(
546 &[
547 (1, 1000, 2000),
548 (2, 6000, 7000),
549 (3, 0, 1500),
550 (4, 1500, 3000),
551 ],
552 &[(vec![3, 1, 4], 0, 3000), (vec![2], 6000, 7000)],
553 );
554
555 run_group_ranges_test(
557 &[(1, 3000, 4000), (2, 4001, 6000), (3, 0, 1000)],
558 &[
559 (vec![3], 0, 1000),
560 (vec![1], 3000, 4000),
561 (vec![2], 4001, 6000),
562 ],
563 );
564
565 run_group_ranges_test(
567 &[(1, 3000, 4000), (2, 4000, 6000), (3, 0, 1000)],
568 &[(vec![3], 0, 1000), (vec![1, 2], 3000, 6000)],
569 );
570 }
571
572 #[test]
573 fn test_merge_range() {
574 let mut left = RangeMeta {
575 time_range: (Timestamp::new_second(1000), Timestamp::new_second(2000)),
576 indices: smallvec![SourceIndex {
577 index: 1,
578 num_row_groups: 2,
579 }],
580 row_group_indices: smallvec![
581 RowGroupIndex {
582 index: 1,
583 row_group_index: 1
584 },
585 RowGroupIndex {
586 index: 1,
587 row_group_index: 2
588 }
589 ],
590 num_rows: 5,
591 };
592 let right = RangeMeta {
593 time_range: (Timestamp::new_second(800), Timestamp::new_second(1200)),
594 indices: smallvec![SourceIndex {
595 index: 2,
596 num_row_groups: 2,
597 }],
598 row_group_indices: smallvec![
599 RowGroupIndex {
600 index: 2,
601 row_group_index: 1
602 },
603 RowGroupIndex {
604 index: 2,
605 row_group_index: 2
606 }
607 ],
608 num_rows: 4,
609 };
610 left.merge(right);
611
612 assert_eq!(
613 left,
614 RangeMeta {
615 time_range: (Timestamp::new_second(800), Timestamp::new_second(2000)),
616 indices: smallvec![
617 SourceIndex {
618 index: 1,
619 num_row_groups: 2
620 },
621 SourceIndex {
622 index: 2,
623 num_row_groups: 2
624 }
625 ],
626 row_group_indices: smallvec![
627 RowGroupIndex {
628 index: 1,
629 row_group_index: 1
630 },
631 RowGroupIndex {
632 index: 1,
633 row_group_index: 2
634 },
635 RowGroupIndex {
636 index: 2,
637 row_group_index: 1
638 },
639 RowGroupIndex {
640 index: 2,
641 row_group_index: 2
642 },
643 ],
644 num_rows: 9,
645 }
646 );
647 }
648
649 #[test]
650 fn test_split_range() {
651 let range = RangeMeta {
652 time_range: (Timestamp::new_second(1000), Timestamp::new_second(2000)),
653 indices: smallvec![SourceIndex {
654 index: 1,
655 num_row_groups: 2,
656 }],
657 row_group_indices: smallvec![RowGroupIndex {
658 index: 1,
659 row_group_index: ALL_ROW_GROUPS,
660 }],
661 num_rows: 5,
662 };
663
664 assert!(range.can_split_preserve_order());
665 let mut output = Vec::new();
666 range.maybe_split(&mut output);
667
668 assert_eq!(
669 output,
670 &[
671 RangeMeta {
672 time_range: (Timestamp::new_second(1000), Timestamp::new_second(2000)),
673 indices: smallvec![SourceIndex {
674 index: 1,
675 num_row_groups: 2,
676 }],
677 row_group_indices: smallvec![RowGroupIndex {
678 index: 1,
679 row_group_index: 0
680 },],
681 num_rows: 2,
682 },
683 RangeMeta {
684 time_range: (Timestamp::new_second(1000), Timestamp::new_second(2000)),
685 indices: smallvec![SourceIndex {
686 index: 1,
687 num_row_groups: 2,
688 }],
689 row_group_indices: smallvec![RowGroupIndex {
690 index: 1,
691 row_group_index: 1
692 }],
693 num_rows: 2,
694 }
695 ]
696 );
697 }
698
699 #[test]
700 fn test_not_split_range() {
701 let range = RangeMeta {
702 time_range: (Timestamp::new_second(1000), Timestamp::new_second(2000)),
703 indices: smallvec![
704 SourceIndex {
705 index: 1,
706 num_row_groups: 1,
707 },
708 SourceIndex {
709 index: 2,
710 num_row_groups: 1,
711 }
712 ],
713 row_group_indices: smallvec![
714 RowGroupIndex {
715 index: 1,
716 row_group_index: 1
717 },
718 RowGroupIndex {
719 index: 2,
720 row_group_index: 1
721 }
722 ],
723 num_rows: 5,
724 };
725
726 assert!(!range.can_split_preserve_order());
727 let mut output = Vec::new();
728 range.maybe_split(&mut output);
729 assert_eq!(1, output.len());
730 }
731
732 #[test]
733 fn test_maybe_split_ranges() {
734 let ranges = vec![
735 RangeMeta {
736 time_range: (Timestamp::new_second(0), Timestamp::new_second(500)),
737 indices: smallvec![SourceIndex {
738 index: 0,
739 num_row_groups: 1,
740 }],
741 row_group_indices: smallvec![RowGroupIndex {
742 index: 0,
743 row_group_index: 0,
744 },],
745 num_rows: 4,
746 },
747 RangeMeta {
748 time_range: (Timestamp::new_second(1000), Timestamp::new_second(2000)),
749 indices: smallvec![SourceIndex {
750 index: 1,
751 num_row_groups: 2,
752 }],
753 row_group_indices: smallvec![RowGroupIndex {
754 index: 1,
755 row_group_index: ALL_ROW_GROUPS,
756 },],
757 num_rows: 4,
758 },
759 RangeMeta {
760 time_range: (Timestamp::new_second(3000), Timestamp::new_second(4000)),
761 indices: smallvec![
762 SourceIndex {
763 index: 2,
764 num_row_groups: 2,
765 },
766 SourceIndex {
767 index: 3,
768 num_row_groups: 0,
769 }
770 ],
771 row_group_indices: smallvec![
772 RowGroupIndex {
773 index: 2,
774 row_group_index: ALL_ROW_GROUPS,
775 },
776 RowGroupIndex {
777 index: 3,
778 row_group_index: ALL_ROW_GROUPS,
779 }
780 ],
781 num_rows: 5,
782 },
783 ];
784 let output = maybe_split_ranges_for_seq_scan(ranges);
785 assert_eq!(
786 output,
787 vec![
788 RangeMeta {
789 time_range: (Timestamp::new_second(0), Timestamp::new_second(500)),
790 indices: smallvec![SourceIndex {
791 index: 0,
792 num_row_groups: 1,
793 }],
794 row_group_indices: smallvec![RowGroupIndex {
795 index: 0,
796 row_group_index: 0
797 },],
798 num_rows: 4,
799 },
800 RangeMeta {
801 time_range: (Timestamp::new_second(1000), Timestamp::new_second(2000)),
802 indices: smallvec![SourceIndex {
803 index: 1,
804 num_row_groups: 2,
805 }],
806 row_group_indices: smallvec![RowGroupIndex {
807 index: 1,
808 row_group_index: 0
809 },],
810 num_rows: 2,
811 },
812 RangeMeta {
813 time_range: (Timestamp::new_second(1000), Timestamp::new_second(2000)),
814 indices: smallvec![SourceIndex {
815 index: 1,
816 num_row_groups: 2,
817 }],
818 row_group_indices: smallvec![RowGroupIndex {
819 index: 1,
820 row_group_index: 1
821 }],
822 num_rows: 2,
823 },
824 RangeMeta {
825 time_range: (Timestamp::new_second(3000), Timestamp::new_second(4000)),
826 indices: smallvec![
827 SourceIndex {
828 index: 2,
829 num_row_groups: 2
830 },
831 SourceIndex {
832 index: 3,
833 num_row_groups: 0,
834 }
835 ],
836 row_group_indices: smallvec![
837 RowGroupIndex {
838 index: 2,
839 row_group_index: ALL_ROW_GROUPS,
840 },
841 RowGroupIndex {
842 index: 3,
843 row_group_index: ALL_ROW_GROUPS,
844 }
845 ],
846 num_rows: 5,
847 },
848 ]
849 )
850 }
851}