mito2/sst/parquet/
row_selection.rs

1// Copyright 2023 Greptime Team
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::collections::{BTreeMap, BTreeSet};
16use std::ops::Range;
17
18use index::inverted_index::search::index_apply::ApplyOutput;
19use itertools::Itertools;
20use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
21
22/// A selection of row groups.
23#[derive(Debug, Clone, Default)]
24pub struct RowGroupSelection {
25    /// Row group id to row selection.
26    selection_in_rg: BTreeMap<usize, RowSelectionWithCount>,
27    /// Total number of rows in the selection.
28    row_count: usize,
29    /// Total length of the selectors.
30    selector_len: usize,
31}
32
33/// A row selection with its count.
34#[derive(Debug, Clone, Default)]
35struct RowSelectionWithCount {
36    /// Row selection.
37    selection: RowSelection,
38    /// Number of rows in the selection.
39    row_count: usize,
40    /// Length of the selectors.
41    selector_len: usize,
42}
43
44impl RowGroupSelection {
45    /// Creates a new `RowGroupSelection` with all row groups selected.
46    ///
47    /// # Arguments
48    /// * `row_group_size` - The number of rows in each row group (except possibly the last one)
49    /// * `total_row_count` - Total number of rows
50    pub fn new(row_group_size: usize, total_row_count: usize) -> Self {
51        let mut selection_in_rg = BTreeMap::new();
52
53        let row_group_count = total_row_count.div_ceil(row_group_size);
54        for rg_id in 0..row_group_count {
55            // The last row group may have fewer rows than `row_group_size`
56            let row_group_size = if rg_id == row_group_count - 1 {
57                total_row_count - (row_group_count - 1) * row_group_size
58            } else {
59                row_group_size
60            };
61
62            let selection = RowSelection::from(vec![RowSelector::select(row_group_size)]);
63            selection_in_rg.insert(
64                rg_id,
65                RowSelectionWithCount {
66                    selection,
67                    row_count: row_group_size,
68                    selector_len: 1,
69                },
70            );
71        }
72
73        Self {
74            selection_in_rg,
75            row_count: total_row_count,
76            selector_len: row_group_count,
77        }
78    }
79
80    /// Creates a new `RowGroupSelection` that selects all rows in the specified row groups.
81    ///
82    /// This is useful for fast construction after coarse pruning (e.g. min-max pruning),
83    /// avoiding building and then removing a full selection of all row groups.
84    pub fn from_full_row_group_ids<I>(
85        row_group_ids: I,
86        row_group_size: usize,
87        total_row_count: usize,
88    ) -> Self
89    where
90        I: IntoIterator<Item = usize>,
91    {
92        if row_group_size == 0 || total_row_count == 0 {
93            return Self::default();
94        }
95
96        let row_group_count = total_row_count.div_ceil(row_group_size);
97        if row_group_count == 0 {
98            return Self::default();
99        }
100
101        let last_row_group_size = total_row_count - (row_group_count - 1) * row_group_size;
102
103        let mut selection_in_rg = BTreeMap::new();
104        let mut row_count = 0usize;
105        let mut selector_len = 0usize;
106
107        for rg_id in row_group_ids {
108            if rg_id >= row_group_count {
109                continue;
110            }
111
112            let rg_row_count = if rg_id == row_group_count - 1 {
113                last_row_group_size
114            } else {
115                row_group_size
116            };
117
118            let selection = RowSelection::from(vec![RowSelector::select(rg_row_count)]);
119            if selection_in_rg
120                .insert(
121                    rg_id,
122                    RowSelectionWithCount {
123                        selection,
124                        row_count: rg_row_count,
125                        selector_len: 1,
126                    },
127                )
128                .is_none()
129            {
130                row_count += rg_row_count;
131                selector_len += 1;
132            }
133        }
134
135        Self {
136            selection_in_rg,
137            row_count,
138            selector_len,
139        }
140    }
141
142    /// Returns the row selection for a given row group.
143    ///
144    /// `None` indicates not selected.
145    pub fn get(&self, rg_id: usize) -> Option<&RowSelection> {
146        self.selection_in_rg.get(&rg_id).map(|x| &x.selection)
147    }
148
149    /// Creates a new `RowGroupSelection` from the output of inverted index application.
150    ///
151    /// # Arguments
152    /// * `row_group_size` - The number of rows in each row group (except possibly the last one)
153    /// * `num_row_groups` - Total number of row groups in the Parquet file
154    /// * `total_row_count` - Total number of rows in the Parquet file
155    /// * `apply_output` - The output from applying the inverted index
156    ///
157    /// # Assumptions
158    /// * All row groups (except possibly the last one) have the same number of rows
159    /// * The last row group may have fewer rows than `row_group_size`
160    ///
161    /// Returns `None` if the index and Parquet file have different row counts.
162    pub fn from_inverted_index_apply_output(
163        row_group_size: usize,
164        num_row_groups: usize,
165        total_row_count: usize,
166        apply_output: ApplyOutput,
167    ) -> Option<Self> {
168        if apply_output.total_row_count != total_row_count {
169            return None;
170        }
171
172        // Step 1: Convert segment IDs to row ranges within row groups.
173        // A segment may span more than one row group when `row_group_size` is not a multiple
174        // of `segment_row_count`, so we split each matched segment's global row range across
175        // every row group it overlaps instead of dropping the overflow.
176        let segment_row_count = apply_output.segment_row_count;
177        let row_group_ranges = apply_output
178            .matched_segment_ids
179            .iter_ones()
180            .flat_map(|seg_id| {
181                // Global row range covered by this segment, clamped to the file's row count.
182                let begin_row_id = seg_id * segment_row_count;
183                let end_row_id = (begin_row_id + segment_row_count).min(total_row_count);
184
185                // Walk the range, emitting one (row_group_id, in-group range) per row group.
186                let mut cursor = begin_row_id;
187                std::iter::from_fn(move || {
188                    if cursor >= end_row_id {
189                        return None;
190                    }
191                    let row_group_id = cursor / row_group_size;
192                    let rg_begin_row_id = cursor % row_group_size;
193                    let rg_boundary = (row_group_id + 1) * row_group_size;
194                    let chunk_end = end_row_id.min(rg_boundary);
195                    let rg_end_row_id = rg_begin_row_id + (chunk_end - cursor);
196                    cursor = chunk_end;
197                    Some((row_group_id, rg_begin_row_id..rg_end_row_id))
198                })
199            });
200
201        // Step 2: Group ranges by row group ID and create row selections
202        let mut selected_row_count = 0;
203        let mut total_selector_len = 0;
204        let mut selection_in_rg = row_group_ranges
205            .chunk_by(|(row_group_id, _)| *row_group_id)
206            .into_iter()
207            .map(|(row_group_id, group)| {
208                // Extract just the ranges from the group
209                let ranges = group.map(|(_, ranges)| ranges);
210                // Create row selection from the ranges
211                // Note: We use `row_group_size` here, which is safe because:
212                // 1. For non-last row groups, it's the actual size
213                // 2. For the last row group, any ranges beyond the actual size will be clipped
214                //    by the min() operation above
215                let selection = row_selection_from_row_ranges(ranges, row_group_size);
216                let row_count = selection.row_count();
217                let selector_len = selector_len(&selection);
218                selected_row_count += row_count;
219                total_selector_len += selector_len;
220                (
221                    row_group_id,
222                    RowSelectionWithCount {
223                        selection,
224                        row_count,
225                        selector_len,
226                    },
227                )
228            })
229            .collect::<BTreeMap<_, _>>();
230
231        Self::fill_missing_row_groups(&mut selection_in_rg, num_row_groups);
232
233        Some(Self {
234            selection_in_rg,
235            row_count: selected_row_count,
236            selector_len: total_selector_len,
237        })
238    }
239
240    /// Creates a new `RowGroupSelection` from a set of row IDs.
241    ///
242    /// # Arguments
243    /// * `row_ids` - Set of row IDs to select
244    /// * `row_group_size` - The number of rows in each row group (except possibly the last one)
245    /// * `num_row_groups` - Total number of row groups
246    ///
247    /// # Assumptions
248    /// * All row groups (except possibly the last one) have the same number of rows
249    /// * The last row group may have fewer rows than `row_group_size`
250    /// * All row IDs must within the range of [0, num_row_groups * row_group_size)
251    pub fn from_row_ids(
252        row_ids: BTreeSet<u32>,
253        row_group_size: usize,
254        num_row_groups: usize,
255    ) -> Self {
256        // Step 1: Group row IDs by their row group
257        let row_group_to_row_ids =
258            Self::group_row_ids_by_row_group(row_ids, row_group_size, num_row_groups);
259
260        // Step 2: Create row selections for each row group
261        let mut total_row_count = 0;
262        let mut total_selector_len = 0;
263        let mut selection_in_rg = row_group_to_row_ids
264            .into_iter()
265            .map(|(row_group_id, row_ids)| {
266                let selection =
267                    row_selection_from_sorted_row_ids(row_ids.into_iter(), row_group_size);
268                let row_count = selection.row_count();
269                let selector_len = selector_len(&selection);
270                total_row_count += row_count;
271                total_selector_len += selector_len;
272                (
273                    row_group_id,
274                    RowSelectionWithCount {
275                        selection,
276                        row_count,
277                        selector_len,
278                    },
279                )
280            })
281            .collect::<BTreeMap<_, _>>();
282
283        Self::fill_missing_row_groups(&mut selection_in_rg, num_row_groups);
284
285        Self {
286            selection_in_rg,
287            row_count: total_row_count,
288            selector_len: total_selector_len,
289        }
290    }
291
292    /// Creates a new `RowGroupSelection` from a set of row ranges.
293    ///
294    /// # Arguments
295    /// * `row_ranges` - A vector of (row_group_id, row_ranges) pairs
296    /// * `row_group_size` - The number of rows in each row group (except possibly the last one)
297    ///
298    /// # Assumptions
299    /// * All row groups (except possibly the last one) have the same number of rows
300    /// * The last row group may have fewer rows than `row_group_size`
301    /// * All ranges in `row_ranges` must be within the bounds of their respective row groups
302    ///   (i.e., for row group i, all ranges must be within [0, row_group_size) or [0, remaining_rows) for the last row group)
303    /// * Ranges within the same row group must not overlap. Overlapping ranges will result in undefined behavior.
304    pub fn from_row_ranges(
305        row_ranges: Vec<(usize, Vec<Range<usize>>)>,
306        row_group_size: usize,
307    ) -> Self {
308        let mut total_row_count = 0;
309        let mut total_selector_len = 0;
310        let selection_in_rg = row_ranges
311            .into_iter()
312            .map(|(row_group_id, ranges)| {
313                let selection = row_selection_from_row_ranges(ranges.into_iter(), row_group_size);
314                let row_count = selection.row_count();
315                let selector_len = selector_len(&selection);
316                total_row_count += row_count;
317                total_selector_len += selector_len;
318                (
319                    row_group_id,
320                    RowSelectionWithCount {
321                        selection,
322                        row_count,
323                        selector_len,
324                    },
325                )
326            })
327            .collect();
328
329        Self {
330            selection_in_rg,
331            row_count: total_row_count,
332            selector_len: total_selector_len,
333        }
334    }
335
336    /// Groups row IDs by their row group.
337    ///
338    /// # Arguments
339    /// * `row_ids` - Set of row IDs to group
340    /// * `row_group_size` - Size of each row group
341    /// * `num_row_groups` - Total number of row groups
342    ///
343    /// # Returns
344    /// A vector of (row_group_id, row_ids) pairs, where row_ids are the IDs within that row group.
345    fn group_row_ids_by_row_group(
346        row_ids: BTreeSet<u32>,
347        row_group_size: usize,
348        num_row_groups: usize,
349    ) -> Vec<(usize, Vec<usize>)> {
350        let est_rows_per_group = row_ids.len() / num_row_groups;
351        let mut row_group_to_row_ids: Vec<(usize, Vec<usize>)> = Vec::with_capacity(num_row_groups);
352
353        for row_id in row_ids {
354            let row_group_id = row_id as usize / row_group_size;
355            let row_id_in_group = row_id as usize % row_group_size;
356
357            if let Some((rg_id, row_ids)) = row_group_to_row_ids.last_mut()
358                && *rg_id == row_group_id
359            {
360                row_ids.push(row_id_in_group);
361            } else {
362                let mut row_ids = Vec::with_capacity(est_rows_per_group);
363                row_ids.push(row_id_in_group);
364                row_group_to_row_ids.push((row_group_id, row_ids));
365            }
366        }
367
368        row_group_to_row_ids
369    }
370
371    /// Intersects two `RowGroupSelection`s.
372    pub fn intersect(&self, other: &Self) -> Self {
373        let mut res = BTreeMap::new();
374        let mut total_row_count = 0;
375        let mut total_selector_len = 0;
376
377        for (rg_id, x) in other.selection_in_rg.iter() {
378            let Some(y) = self.selection_in_rg.get(rg_id) else {
379                continue;
380            };
381            let selection = intersect_row_selections(&x.selection, &y.selection);
382            let row_count = selection.row_count();
383            let selector_len = selector_len(&selection);
384            if row_count > 0 {
385                total_row_count += row_count;
386                total_selector_len += selector_len;
387                res.insert(
388                    *rg_id,
389                    RowSelectionWithCount {
390                        selection,
391                        row_count,
392                        selector_len,
393                    },
394                );
395            }
396        }
397
398        Self {
399            selection_in_rg: res,
400            row_count: total_row_count,
401            selector_len: total_selector_len,
402        }
403    }
404
405    /// Returns the number of row groups in the selection.
406    pub fn row_group_count(&self) -> usize {
407        self.selection_in_rg.len()
408    }
409
410    /// Returns the number of rows in the selection.
411    pub fn row_count(&self) -> usize {
412        self.row_count
413    }
414
415    /// Returns the first row group in the selection.
416    ///
417    /// Skip the row group if the row count is 0.
418    pub fn pop_first(&mut self) -> Option<(usize, RowSelection)> {
419        while let Some((
420            row_group_id,
421            RowSelectionWithCount {
422                selection,
423                row_count,
424                selector_len,
425            },
426        )) = self.selection_in_rg.pop_first()
427        {
428            if row_count > 0 {
429                self.row_count -= row_count;
430                self.selector_len -= selector_len;
431                return Some((row_group_id, selection));
432            }
433        }
434
435        None
436    }
437
438    /// Removes a row group from the selection.
439    pub fn remove_row_group(&mut self, row_group_id: usize) {
440        let Some(RowSelectionWithCount {
441            row_count,
442            selector_len,
443            ..
444        }) = self.selection_in_rg.remove(&row_group_id)
445        else {
446            return;
447        };
448        self.row_count -= row_count;
449        self.selector_len -= selector_len;
450    }
451
452    /// Returns true if the selection is empty.
453    pub fn is_empty(&self) -> bool {
454        self.selection_in_rg.is_empty()
455    }
456
457    /// Returns true if the selection contains a row group with the given ID.
458    pub fn contains_row_group(&self, row_group_id: usize) -> bool {
459        self.selection_in_rg.contains_key(&row_group_id)
460    }
461
462    /// Returns true if the selection contains a row group with the given ID and the row selection is not empty.
463    pub fn contains_non_empty_row_group(&self, row_group_id: usize) -> bool {
464        self.selection_in_rg
465            .get(&row_group_id)
466            .map(|r| r.row_count > 0)
467            .unwrap_or(false)
468    }
469
470    /// Returns an iterator over the row groups in the selection.
471    pub fn iter(&self) -> impl Iterator<Item = (&usize, &RowSelection)> {
472        self.selection_in_rg
473            .iter()
474            .map(|(row_group_id, x)| (row_group_id, &x.selection))
475    }
476
477    /// Returns the memory usage of the selection.
478    pub fn mem_usage(&self) -> usize {
479        self.selector_len * size_of::<RowSelector>()
480            + self.selection_in_rg.len() * size_of::<RowSelectionWithCount>()
481    }
482
483    /// Concatenates `other` into `self`. `other` must not contain row groups that `self` contains.
484    ///
485    /// Panics if `self` contains row groups that `other` contains.
486    pub fn concat(&mut self, other: &Self) {
487        for (rg_id, other_rs) in other.selection_in_rg.iter() {
488            if self.selection_in_rg.contains_key(rg_id) {
489                panic!("row group {} is already in `self`", rg_id);
490            }
491
492            self.selection_in_rg.insert(*rg_id, other_rs.clone());
493            self.row_count += other_rs.row_count;
494            self.selector_len += other_rs.selector_len;
495        }
496    }
497
498    /// Fills the missing row groups with empty selections.
499    /// This is to indicate that the row groups are searched even if no rows are found.
500    fn fill_missing_row_groups(
501        selection_in_rg: &mut BTreeMap<usize, RowSelectionWithCount>,
502        num_row_groups: usize,
503    ) {
504        for rg_id in 0..num_row_groups {
505            selection_in_rg.entry(rg_id).or_default();
506        }
507    }
508}
509
510/// Ported from `parquet` but trailing rows are removed.
511///
512/// Combine two lists of `RowSelection` return the intersection of them
513/// For example:
514/// self:      NNYYYYNNYYNYN
515/// other:     NYNNNNNNY
516///
517/// returned:  NNNNNNNNY     (modified)
518///            NNNNNNNNYYNYN (original)
519fn intersect_row_selections(left: &RowSelection, right: &RowSelection) -> RowSelection {
520    let mut l_iter = left.iter().copied().peekable();
521    let mut r_iter = right.iter().copied().peekable();
522
523    let iter = std::iter::from_fn(move || {
524        loop {
525            let l = l_iter.peek_mut();
526            let r = r_iter.peek_mut();
527
528            match (l, r) {
529                (Some(a), _) if a.row_count == 0 => {
530                    l_iter.next().unwrap();
531                }
532                (_, Some(b)) if b.row_count == 0 => {
533                    r_iter.next().unwrap();
534                }
535                (Some(l), Some(r)) => {
536                    return match (l.skip, r.skip) {
537                        // Keep both ranges
538                        (false, false) => {
539                            if l.row_count < r.row_count {
540                                r.row_count -= l.row_count;
541                                l_iter.next()
542                            } else {
543                                l.row_count -= r.row_count;
544                                r_iter.next()
545                            }
546                        }
547                        // skip at least one
548                        _ => {
549                            if l.row_count < r.row_count {
550                                let skip = l.row_count;
551                                r.row_count -= l.row_count;
552                                l_iter.next();
553                                Some(RowSelector::skip(skip))
554                            } else {
555                                let skip = r.row_count;
556                                l.row_count -= skip;
557                                r_iter.next();
558                                Some(RowSelector::skip(skip))
559                            }
560                        }
561                    };
562                }
563                (None, _) => return None,
564                (_, None) => return None,
565            }
566        }
567    });
568
569    iter.collect()
570}
571
572/// Converts an iterator of row ranges into a `RowSelection` by creating a sequence of `RowSelector`s.
573///
574/// This function processes each range in the input and either creates a new selector or merges
575/// with the existing one, depending on whether the current range is contiguous with the preceding one
576/// or if there's a gap that requires skipping rows. It handles both "select" and "skip" actions,
577/// optimizing the list of selectors by merging contiguous actions of the same type.
578///
579/// The returned selection intentionally stops at the end of the last matched range and may omit a
580/// trailing `skip` that would extend it to `total_row_count`. That is fine when the selection is
581/// used directly by the parquet reader, which simply stops once the selectors are exhausted.
582///
583/// Note: overlapping ranges are not supported and will result in an incorrect selection.
584pub(crate) fn row_selection_from_row_ranges(
585    row_ranges: impl Iterator<Item = Range<usize>>,
586    total_row_count: usize,
587) -> RowSelection {
588    let (selectors, _) = build_selectors_from_row_ranges(row_ranges, total_row_count);
589    RowSelection::from(selectors)
590}
591
592fn build_selectors_from_row_ranges(
593    row_ranges: impl Iterator<Item = Range<usize>>,
594    total_row_count: usize,
595) -> (Vec<RowSelector>, usize) {
596    let mut selectors: Vec<RowSelector> = Vec::new();
597    let mut last_processed_end = 0;
598
599    for Range { start, end } in row_ranges {
600        let end = end.min(total_row_count);
601        if start > last_processed_end {
602            add_or_merge_selector(&mut selectors, start - last_processed_end, true);
603        }
604
605        add_or_merge_selector(&mut selectors, end - start, false);
606        last_processed_end = end;
607    }
608
609    (selectors, last_processed_end)
610}
611
612/// Converts an iterator of sorted row IDs into a `RowSelection`.
613///
614/// Note: the input iterator must be sorted in ascending order and
615///       contain unique row IDs in the range [0, total_row_count).
616pub(crate) fn row_selection_from_sorted_row_ids(
617    row_ids: impl Iterator<Item = usize>,
618    total_row_count: usize,
619) -> RowSelection {
620    let mut selectors: Vec<RowSelector> = Vec::new();
621    let mut last_processed_end = 0;
622
623    for row_id in row_ids {
624        let start = row_id;
625        let end = start + 1;
626
627        if start > last_processed_end {
628            add_or_merge_selector(&mut selectors, start - last_processed_end, true);
629        }
630
631        add_or_merge_selector(&mut selectors, end - start, false);
632        last_processed_end = end;
633    }
634
635    if last_processed_end < total_row_count {
636        add_or_merge_selector(&mut selectors, total_row_count - last_processed_end, true);
637    }
638
639    RowSelection::from(selectors)
640}
641
642/// Helper function to either add a new `RowSelector` to `selectors` or merge it with the last one
643/// if they are of the same type (both skip or both select).
644fn add_or_merge_selector(selectors: &mut Vec<RowSelector>, count: usize, is_skip: bool) {
645    if let Some(last) = selectors.last_mut() {
646        // Merge with last if both actions are same
647        if last.skip == is_skip {
648            last.row_count += count;
649            return;
650        }
651    }
652    // Add new selector otherwise
653    let new_selector = if is_skip {
654        RowSelector::skip(count)
655    } else {
656        RowSelector::select(count)
657    };
658    selectors.push(new_selector);
659}
660
661/// Returns the length of the selectors in the selection.
662fn selector_len(selection: &RowSelection) -> usize {
663    selection.iter().size_hint().0
664}
665
666#[cfg(test)]
667#[allow(clippy::single_range_in_vec_init)]
668mod tests {
669    use index::bitmap::Bitmap;
670
671    use super::*;
672
673    #[test]
674    fn test_selector_len() {
675        let selection = RowSelection::from(vec![RowSelector::skip(5), RowSelector::select(5)]);
676        assert_eq!(selector_len(&selection), 2);
677
678        let selection = RowSelection::from(vec![
679            RowSelector::select(5),
680            RowSelector::skip(5),
681            RowSelector::select(5),
682        ]);
683        assert_eq!(selector_len(&selection), 3);
684
685        let selection = RowSelection::from(vec![]);
686        assert_eq!(selector_len(&selection), 0);
687    }
688
689    #[test]
690    fn test_single_contiguous_range() {
691        let selection = row_selection_from_row_ranges(Some(5..10).into_iter(), 10);
692        let expected = RowSelection::from(vec![RowSelector::skip(5), RowSelector::select(5)]);
693        assert_eq!(selection, expected);
694    }
695
696    #[test]
697    fn test_non_contiguous_ranges() {
698        let ranges = [1..3, 5..8];
699        let selection = row_selection_from_row_ranges(ranges.iter().cloned(), 10);
700        let expected = RowSelection::from(vec![
701            RowSelector::skip(1),
702            RowSelector::select(2),
703            RowSelector::skip(2),
704            RowSelector::select(3),
705        ]);
706        assert_eq!(selection, expected);
707    }
708
709    #[test]
710    fn test_empty_range() {
711        let ranges = [];
712        let selection = row_selection_from_row_ranges(ranges.iter().cloned(), 10);
713        let expected = RowSelection::from(vec![]);
714        assert_eq!(selection, expected);
715    }
716
717    #[test]
718    fn test_adjacent_ranges() {
719        let ranges = [1..2, 2..3];
720        let selection = row_selection_from_row_ranges(ranges.iter().cloned(), 10);
721        let expected = RowSelection::from(vec![RowSelector::skip(1), RowSelector::select(2)]);
722        assert_eq!(selection, expected);
723    }
724
725    #[test]
726    fn test_large_gap_between_ranges() {
727        let ranges = [1..2, 100..101];
728        let selection = row_selection_from_row_ranges(ranges.iter().cloned(), 10240);
729        let expected = RowSelection::from(vec![
730            RowSelector::skip(1),
731            RowSelector::select(1),
732            RowSelector::skip(98),
733            RowSelector::select(1),
734        ]);
735        assert_eq!(selection, expected);
736    }
737
738    fn apply_output(
739        matched: &[usize],
740        segment_row_count: usize,
741        total_row_count: usize,
742    ) -> ApplyOutput {
743        let mut matched_segment_ids = Bitmap::new_bitvec();
744        for &seg_id in matched {
745            matched_segment_ids.insert_range(seg_id..=seg_id);
746        }
747        ApplyOutput {
748            matched_segment_ids,
749            total_row_count,
750            segment_row_count,
751        }
752    }
753
754    #[test]
755    fn test_from_inverted_index_segment_within_row_group() {
756        // segment_row_count (4) divides row_group_size (8): each segment maps to one row group.
757        // segment 1 -> rows [4, 8) in row group 0; segment 2 -> rows [8, 12) in row group 1.
758        let selection = RowGroupSelection::from_inverted_index_apply_output(
759            8,
760            2,
761            16,
762            apply_output(&[1, 2], 4, 16),
763        )
764        .unwrap();
765        assert_eq!(selection.row_count(), 8);
766        assert_eq!(
767            selection.get(0),
768            Some(&RowSelection::from(vec![
769                RowSelector::skip(4),
770                RowSelector::select(4),
771            ]))
772        );
773        assert_eq!(
774            selection.get(1),
775            Some(&RowSelection::from(vec![RowSelector::select(4)]))
776        );
777    }
778
779    #[test]
780    fn test_from_inverted_index_segment_crosses_row_group_boundary() {
781        // row_group_size (10) is not a multiple of segment_row_count (4): segment 2 covers global
782        // rows [8, 12), straddling row group 0 (rows [8, 10)) and row group 1 (rows [0, 2)).
783        // Both halves must be selected; the old code dropped the row group 1 half.
784        let selection = RowGroupSelection::from_inverted_index_apply_output(
785            10,
786            2,
787            20,
788            apply_output(&[2], 4, 20),
789        )
790        .unwrap();
791        assert_eq!(selection.row_count(), 4);
792        assert_eq!(
793            selection.get(0),
794            Some(&RowSelection::from(vec![
795                RowSelector::skip(8),
796                RowSelector::select(2),
797            ]))
798        );
799        assert_eq!(
800            selection.get(1),
801            Some(&RowSelection::from(vec![RowSelector::select(2)]))
802        );
803    }
804
805    #[test]
806    fn test_from_inverted_index_last_segment_clamped_to_end() {
807        // Last segment 4 covers global rows [16, 20), but the last row group is partial and only
808        // 18 rows exist. The overflow rows [18, 20) must be clamped.
809        let selection = RowGroupSelection::from_inverted_index_apply_output(
810            10,
811            2,
812            18,
813            apply_output(&[4], 4, 18),
814        )
815        .unwrap();
816        assert_eq!(selection.row_count(), 2);
817        assert_eq!(selection.get(2), None);
818        // Rows [16, 18) land in row group 1 at offset [6, 8).
819        assert_eq!(
820            selection.get(1),
821            Some(&RowSelection::from(vec![
822                RowSelector::skip(6),
823                RowSelector::select(2),
824            ]))
825        );
826    }
827
828    #[test]
829    fn test_from_inverted_index_rejects_mismatched_row_count() {
830        let selection = RowGroupSelection::from_inverted_index_apply_output(
831            10,
832            2,
833            20,
834            apply_output(&[0], 4, 16),
835        );
836        assert!(selection.is_none());
837    }
838
839    #[test]
840    fn test_range_end_over_total_row_count() {
841        let ranges = Some(1..10);
842        let selection = row_selection_from_row_ranges(ranges.into_iter(), 5);
843        let expected = RowSelection::from(vec![RowSelector::skip(1), RowSelector::select(4)]);
844        assert_eq!(selection, expected);
845    }
846
847    #[test]
848    fn test_row_ids_to_selection() {
849        let row_ids = [1, 3, 5, 7, 9].into_iter();
850        let selection = row_selection_from_sorted_row_ids(row_ids, 10);
851        let expected = RowSelection::from(vec![
852            RowSelector::skip(1),
853            RowSelector::select(1),
854            RowSelector::skip(1),
855            RowSelector::select(1),
856            RowSelector::skip(1),
857            RowSelector::select(1),
858            RowSelector::skip(1),
859            RowSelector::select(1),
860            RowSelector::skip(1),
861            RowSelector::select(1),
862        ]);
863        assert_eq!(selection, expected);
864    }
865
866    #[test]
867    fn test_row_ids_to_selection_full() {
868        let row_ids = 0..10;
869        let selection = row_selection_from_sorted_row_ids(row_ids, 10);
870        let expected = RowSelection::from(vec![RowSelector::select(10)]);
871        assert_eq!(selection, expected);
872    }
873
874    #[test]
875    fn test_row_ids_to_selection_empty() {
876        let selection = row_selection_from_sorted_row_ids(None.into_iter(), 10);
877        let expected = RowSelection::from(vec![RowSelector::skip(10)]);
878        assert_eq!(selection, expected);
879    }
880
881    #[test]
882    fn test_group_row_ids() {
883        let row_ids = [0, 1, 2, 5, 6, 7, 8, 12].into_iter().collect();
884        let row_group_size = 5;
885        let num_row_groups = 3;
886
887        let row_group_to_row_ids =
888            RowGroupSelection::group_row_ids_by_row_group(row_ids, row_group_size, num_row_groups);
889
890        assert_eq!(
891            row_group_to_row_ids,
892            vec![(0, vec![0, 1, 2]), (1, vec![0, 1, 2, 3]), (2, vec![2])]
893        );
894    }
895
896    #[test]
897    fn test_row_group_selection_new() {
898        // Test with regular case
899        let selection = RowGroupSelection::new(100, 250);
900        assert_eq!(selection.row_count(), 250);
901        assert_eq!(selection.row_group_count(), 3);
902
903        // Check content of each row group
904        let row_selection = selection.get(0).unwrap();
905        assert_eq!(row_selection.row_count(), 100);
906
907        let row_selection = selection.get(1).unwrap();
908        assert_eq!(row_selection.row_count(), 100);
909
910        let row_selection = selection.get(2).unwrap();
911        assert_eq!(row_selection.row_count(), 50);
912
913        // Test with empty selection
914        let selection = RowGroupSelection::new(100, 0);
915        assert_eq!(selection.row_count(), 0);
916        assert_eq!(selection.row_group_count(), 0);
917        assert!(selection.get(0).is_none());
918
919        // Test with single row group
920        let selection = RowGroupSelection::new(100, 50);
921        assert_eq!(selection.row_count(), 50);
922        assert_eq!(selection.row_group_count(), 1);
923
924        let row_selection = selection.get(0).unwrap();
925        assert_eq!(row_selection.row_count(), 50);
926
927        // Test with row count that doesn't divide evenly
928        let selection = RowGroupSelection::new(100, 150);
929        assert_eq!(selection.row_count(), 150);
930        assert_eq!(selection.row_group_count(), 2);
931
932        let row_selection = selection.get(0).unwrap();
933        assert_eq!(row_selection.row_count(), 100);
934
935        let row_selection = selection.get(1).unwrap();
936        assert_eq!(row_selection.row_count(), 50);
937
938        // Test with row count that's just over a multiple of row_group_size
939        let selection = RowGroupSelection::new(100, 101);
940        assert_eq!(selection.row_count(), 101);
941        assert_eq!(selection.row_group_count(), 2);
942
943        let row_selection = selection.get(0).unwrap();
944        assert_eq!(row_selection.row_count(), 100);
945
946        let row_selection = selection.get(1).unwrap();
947        assert_eq!(row_selection.row_count(), 1);
948    }
949
950    #[test]
951    fn test_from_full_row_group_ids_dedup_duplicates() {
952        let selection = RowGroupSelection::from_full_row_group_ids([0, 0, 2, 2], 10, 25);
953        assert_eq!(selection.row_group_count(), 2);
954        assert_eq!(selection.row_count(), 15);
955
956        assert_eq!(selection.get(0).unwrap().row_count(), 10);
957        assert_eq!(selection.get(2).unwrap().row_count(), 5);
958    }
959
960    #[test]
961    fn test_from_row_ids() {
962        let row_group_size = 100;
963        let num_row_groups = 3;
964
965        // Test with regular case
966        let row_ids: BTreeSet<u32> = vec![5, 15, 25, 35, 105, 115, 205, 215]
967            .into_iter()
968            .collect();
969        let selection = RowGroupSelection::from_row_ids(row_ids, row_group_size, num_row_groups);
970        assert_eq!(selection.row_count(), 8);
971        assert_eq!(selection.row_group_count(), 3);
972
973        // Check content of each row group
974        let row_selection = selection.get(0).unwrap();
975        assert_eq!(row_selection.row_count(), 4); // 5, 15, 25, 35
976
977        let row_selection = selection.get(1).unwrap();
978        assert_eq!(row_selection.row_count(), 2); // 105, 115
979
980        let row_selection = selection.get(2).unwrap();
981        assert_eq!(row_selection.row_count(), 2); // 205, 215
982
983        // Test with empty row IDs
984        let empty_row_ids: BTreeSet<u32> = BTreeSet::new();
985        let selection =
986            RowGroupSelection::from_row_ids(empty_row_ids, row_group_size, num_row_groups);
987        assert_eq!(selection.row_count(), 0);
988        assert_eq!(selection.row_group_count(), 3);
989
990        // Test with consecutive row IDs
991        let consecutive_row_ids: BTreeSet<u32> = vec![5, 6, 7, 8, 9].into_iter().collect();
992        let selection =
993            RowGroupSelection::from_row_ids(consecutive_row_ids, row_group_size, num_row_groups);
994        assert_eq!(selection.row_count(), 5);
995        assert_eq!(selection.row_group_count(), 3);
996
997        let row_selection = selection.get(0).unwrap();
998        assert_eq!(row_selection.row_count(), 5); // 5, 6, 7, 8, 9
999
1000        // Test with row IDs at row group boundaries
1001        let boundary_row_ids: BTreeSet<u32> = vec![0, 99, 100, 199, 200, 249].into_iter().collect();
1002        let selection =
1003            RowGroupSelection::from_row_ids(boundary_row_ids, row_group_size, num_row_groups);
1004        assert_eq!(selection.row_count(), 6);
1005        assert_eq!(selection.row_group_count(), 3);
1006
1007        let row_selection = selection.get(0).unwrap();
1008        assert_eq!(row_selection.row_count(), 2); // 0, 99
1009
1010        let row_selection = selection.get(1).unwrap();
1011        assert_eq!(row_selection.row_count(), 2); // 100, 199
1012
1013        let row_selection = selection.get(2).unwrap();
1014        assert_eq!(row_selection.row_count(), 2); // 200, 249
1015
1016        // Test with single row group
1017        let single_group_row_ids: BTreeSet<u32> = vec![5, 10, 15].into_iter().collect();
1018        let selection = RowGroupSelection::from_row_ids(single_group_row_ids, row_group_size, 1);
1019        assert_eq!(selection.row_count(), 3);
1020        assert_eq!(selection.row_group_count(), 1);
1021
1022        let row_selection = selection.get(0).unwrap();
1023        assert_eq!(row_selection.row_count(), 3); // 5, 10, 15
1024    }
1025
1026    #[test]
1027    fn test_from_row_ranges() {
1028        let row_group_size = 100;
1029
1030        // Test with regular case
1031        let ranges = vec![
1032            (0, vec![5..15, 25..35]), // Within [0, 100)
1033            (1, vec![5..15]),         // Within [0, 100)
1034            (2, vec![0..5, 10..15]),  // Within [0, 50) for last row group
1035        ];
1036        let selection = RowGroupSelection::from_row_ranges(ranges, row_group_size);
1037        assert_eq!(selection.row_count(), 40);
1038        assert_eq!(selection.row_group_count(), 3);
1039
1040        // Check content of each row group
1041        let row_selection = selection.get(0).unwrap();
1042        assert_eq!(row_selection.row_count(), 20); // 5..15 (10) + 25..35 (10)
1043
1044        let row_selection = selection.get(1).unwrap();
1045        assert_eq!(row_selection.row_count(), 10); // 5..15 (10)
1046
1047        let row_selection = selection.get(2).unwrap();
1048        assert_eq!(row_selection.row_count(), 10); // 0..5 (5) + 10..15 (5)
1049
1050        // Test with empty ranges
1051        let empty_ranges: Vec<(usize, Vec<Range<usize>>)> = vec![];
1052        let selection = RowGroupSelection::from_row_ranges(empty_ranges, row_group_size);
1053        assert_eq!(selection.row_count(), 0);
1054        assert_eq!(selection.row_group_count(), 0);
1055        assert!(selection.get(0).is_none());
1056
1057        // Test with adjacent ranges within same row group
1058        let adjacent_ranges = vec![
1059            (0, vec![5..15, 15..25]), // Adjacent ranges within [0, 100)
1060        ];
1061        let selection = RowGroupSelection::from_row_ranges(adjacent_ranges, row_group_size);
1062        assert_eq!(selection.row_count(), 20);
1063        assert_eq!(selection.row_group_count(), 1);
1064
1065        let row_selection = selection.get(0).unwrap();
1066        assert_eq!(row_selection.row_count(), 20); // 5..15 (10) + 15..25 (10)
1067
1068        // Test with ranges at row group boundaries
1069        let boundary_ranges = vec![
1070            (0, vec![0..10, 90..100]), // Ranges at start and end of first row group
1071            (1, vec![0..100]),         // Full range of second row group
1072            (2, vec![0..50]),          // Full range of last row group
1073        ];
1074        let selection = RowGroupSelection::from_row_ranges(boundary_ranges, row_group_size);
1075        assert_eq!(selection.row_count(), 170);
1076        assert_eq!(selection.row_group_count(), 3);
1077
1078        let row_selection = selection.get(0).unwrap();
1079        assert_eq!(row_selection.row_count(), 20); // 0..10 (10) + 90..100 (10)
1080
1081        let row_selection = selection.get(1).unwrap();
1082        assert_eq!(row_selection.row_count(), 100); // 0..100 (100)
1083
1084        let row_selection = selection.get(2).unwrap();
1085        assert_eq!(row_selection.row_count(), 50); // 0..50 (50)
1086
1087        // Test with single row group
1088        let single_group_ranges = vec![
1089            (0, vec![0..50]), // Half of first row group
1090        ];
1091        let selection = RowGroupSelection::from_row_ranges(single_group_ranges, row_group_size);
1092        assert_eq!(selection.row_count(), 50);
1093        assert_eq!(selection.row_group_count(), 1);
1094
1095        let row_selection = selection.get(0).unwrap();
1096        assert_eq!(row_selection.row_count(), 50); // 0..50 (50)
1097    }
1098
1099    #[test]
1100    fn test_intersect() {
1101        let row_group_size = 100;
1102
1103        // Test case 1: Regular intersection with partial overlap
1104        let ranges1 = vec![
1105            (0, vec![5..15, 25..35]), // Within [0, 100)
1106            (1, vec![5..15]),         // Within [0, 100)
1107        ];
1108        let selection1 = RowGroupSelection::from_row_ranges(ranges1, row_group_size);
1109
1110        let ranges2 = vec![
1111            (0, vec![10..20]), // Within [0, 100)
1112            (1, vec![10..20]), // Within [0, 100)
1113            (2, vec![0..10]),  // Within [0, 50) for last row group
1114        ];
1115        let selection2 = RowGroupSelection::from_row_ranges(ranges2, row_group_size);
1116
1117        let intersection = selection1.intersect(&selection2);
1118        assert_eq!(intersection.row_count(), 10);
1119        assert_eq!(intersection.row_group_count(), 2);
1120
1121        let row_selection = intersection.get(0).unwrap();
1122        assert_eq!(row_selection.row_count(), 5); // 10..15 (5)
1123
1124        let row_selection = intersection.get(1).unwrap();
1125        assert_eq!(row_selection.row_count(), 5); // 10..15 (5)
1126
1127        // Test case 2: Empty intersection with empty selection
1128        let empty_ranges: Vec<(usize, Vec<Range<usize>>)> = vec![];
1129        let empty_selection = RowGroupSelection::from_row_ranges(empty_ranges, row_group_size);
1130        let intersection = selection1.intersect(&empty_selection);
1131        assert_eq!(intersection.row_count(), 0);
1132        assert_eq!(intersection.row_group_count(), 0);
1133        assert!(intersection.get(0).is_none());
1134
1135        // Test case 3: No overlapping row groups
1136        let non_overlapping_ranges = vec![
1137            (3, vec![0..10]), // Within [0, 50) for last row group
1138        ];
1139        let non_overlapping_selection =
1140            RowGroupSelection::from_row_ranges(non_overlapping_ranges, row_group_size);
1141        let intersection = selection1.intersect(&non_overlapping_selection);
1142        assert_eq!(intersection.row_count(), 0);
1143        assert_eq!(intersection.row_group_count(), 0);
1144        assert!(intersection.get(0).is_none());
1145
1146        // Test case 4: Full overlap within same row group
1147        let full_overlap_ranges1 = vec![
1148            (0, vec![0..50]), // Within [0, 100)
1149        ];
1150        let full_overlap_ranges2 = vec![
1151            (0, vec![0..50]), // Within [0, 100)
1152        ];
1153        let selection1 = RowGroupSelection::from_row_ranges(full_overlap_ranges1, row_group_size);
1154        let selection2 = RowGroupSelection::from_row_ranges(full_overlap_ranges2, row_group_size);
1155        let intersection = selection1.intersect(&selection2);
1156        assert_eq!(intersection.row_count(), 50);
1157        assert_eq!(intersection.row_group_count(), 1);
1158
1159        let row_selection = intersection.get(0).unwrap();
1160        assert_eq!(row_selection.row_count(), 50); // 0..50 (50)
1161
1162        // Test case 5: Partial overlap at row group boundaries
1163        let boundary_ranges1 = vec![
1164            (0, vec![0..10, 90..100]), // Within [0, 100)
1165            (1, vec![0..100]),         // Within [0, 100)
1166        ];
1167        let boundary_ranges2 = vec![
1168            (0, vec![5..15, 95..100]), // Within [0, 100)
1169            (1, vec![50..100]),        // Within [0, 100)
1170        ];
1171        let selection1 = RowGroupSelection::from_row_ranges(boundary_ranges1, row_group_size);
1172        let selection2 = RowGroupSelection::from_row_ranges(boundary_ranges2, row_group_size);
1173        let intersection = selection1.intersect(&selection2);
1174        assert_eq!(intersection.row_count(), 60);
1175        assert_eq!(intersection.row_group_count(), 2);
1176
1177        let row_selection = intersection.get(0).unwrap();
1178        assert_eq!(row_selection.row_count(), 10); // 5..10 (5) + 95..100 (5)
1179
1180        let row_selection = intersection.get(1).unwrap();
1181        assert_eq!(row_selection.row_count(), 50); // 50..100 (50)
1182
1183        // Test case 6: Multiple ranges with complex overlap
1184        let complex_ranges1 = vec![
1185            (0, vec![5..15, 25..35, 45..55]), // Within [0, 100)
1186            (1, vec![10..20, 30..40]),        // Within [0, 100)
1187        ];
1188        let complex_ranges2 = vec![
1189            (0, vec![10..20, 30..40, 50..60]), // Within [0, 100)
1190            (1, vec![15..25, 35..45]),         // Within [0, 100)
1191        ];
1192        let selection1 = RowGroupSelection::from_row_ranges(complex_ranges1, row_group_size);
1193        let selection2 = RowGroupSelection::from_row_ranges(complex_ranges2, row_group_size);
1194        let intersection = selection1.intersect(&selection2);
1195        assert_eq!(intersection.row_count(), 25);
1196        assert_eq!(intersection.row_group_count(), 2);
1197
1198        let row_selection = intersection.get(0).unwrap();
1199        assert_eq!(row_selection.row_count(), 15); // 10..15 (5) + 30..35 (5) + 50..55 (5)
1200
1201        let row_selection = intersection.get(1).unwrap();
1202        assert_eq!(row_selection.row_count(), 10); // 15..20 (5) + 35..40 (5)
1203
1204        // Test case 7: Intersection with last row group (smaller size)
1205        let last_rg_ranges1 = vec![
1206            (2, vec![0..25, 30..40]), // Within [0, 50) for last row group
1207        ];
1208        let last_rg_ranges2 = vec![
1209            (2, vec![20..30, 35..45]), // Within [0, 50) for last row group
1210        ];
1211        let selection1 = RowGroupSelection::from_row_ranges(last_rg_ranges1, row_group_size);
1212        let selection2 = RowGroupSelection::from_row_ranges(last_rg_ranges2, row_group_size);
1213        let intersection = selection1.intersect(&selection2);
1214        assert_eq!(intersection.row_count(), 10);
1215        assert_eq!(intersection.row_group_count(), 1);
1216
1217        let row_selection = intersection.get(2).unwrap();
1218        assert_eq!(row_selection.row_count(), 10); // 20..25 (5) + 35..40 (5)
1219
1220        // Test case 8: Intersection with empty ranges in one selection
1221        let empty_ranges = vec![
1222            (0, vec![]),      // Empty ranges
1223            (1, vec![5..15]), // Within [0, 100)
1224        ];
1225        let selection1 = RowGroupSelection::from_row_ranges(empty_ranges, row_group_size);
1226        let ranges2 = vec![
1227            (0, vec![5..15, 25..35]), // Within [0, 100)
1228            (1, vec![5..15]),         // Within [0, 100)
1229        ];
1230        let selection2 = RowGroupSelection::from_row_ranges(ranges2, row_group_size);
1231        let intersection = selection1.intersect(&selection2);
1232        assert_eq!(intersection.row_count(), 10);
1233        assert_eq!(intersection.row_group_count(), 1);
1234
1235        let row_selection = intersection.get(1).unwrap();
1236        assert_eq!(row_selection.row_count(), 10); // 5..15 (10)
1237    }
1238
1239    #[test]
1240    fn test_pop_first() {
1241        let row_group_size = 100;
1242        let ranges = vec![
1243            (0, vec![5..15]), // Within [0, 100)
1244            (1, vec![5..15]), // Within [0, 100)
1245            (2, vec![0..5]),  // Within [0, 50) for last row group
1246        ];
1247        let mut selection = RowGroupSelection::from_row_ranges(ranges, row_group_size);
1248
1249        // Test popping first row group
1250        let (rg_id, row_selection) = selection.pop_first().unwrap();
1251        assert_eq!(rg_id, 0);
1252        assert_eq!(row_selection.row_count(), 10); // 5..15 (10)
1253        assert_eq!(selection.row_count(), 15);
1254        assert_eq!(selection.row_group_count(), 2);
1255
1256        // Verify remaining row groups' content
1257        let row_selection = selection.get(1).unwrap();
1258        assert_eq!(row_selection.row_count(), 10); // 5..15 (10)
1259
1260        let row_selection = selection.get(2).unwrap();
1261        assert_eq!(row_selection.row_count(), 5); // 0..5 (5)
1262
1263        // Test popping second row group
1264        let (rg_id, row_selection) = selection.pop_first().unwrap();
1265        assert_eq!(rg_id, 1);
1266        assert_eq!(row_selection.row_count(), 10); // 5..15 (10)
1267        assert_eq!(selection.row_count(), 5);
1268        assert_eq!(selection.row_group_count(), 1);
1269
1270        // Verify remaining row group's content
1271        let row_selection = selection.get(2).unwrap();
1272        assert_eq!(row_selection.row_count(), 5); // 0..5 (5)
1273
1274        // Test popping last row group
1275        let (rg_id, row_selection) = selection.pop_first().unwrap();
1276        assert_eq!(rg_id, 2);
1277        assert_eq!(row_selection.row_count(), 5); // 0..5 (5)
1278        assert_eq!(selection.row_count(), 0);
1279        assert_eq!(selection.row_group_count(), 0);
1280        assert!(selection.is_empty());
1281
1282        // Test popping from empty selection
1283        let mut empty_selection = RowGroupSelection::from_row_ranges(vec![], row_group_size);
1284        assert!(empty_selection.pop_first().is_none());
1285        assert_eq!(empty_selection.row_count(), 0);
1286        assert_eq!(empty_selection.row_group_count(), 0);
1287        assert!(empty_selection.is_empty());
1288    }
1289
1290    #[test]
1291    fn test_remove_row_group() {
1292        let row_group_size = 100;
1293        let ranges = vec![
1294            (0, vec![5..15]), // Within [0, 100)
1295            (1, vec![5..15]), // Within [0, 100)
1296            (2, vec![0..5]),  // Within [0, 50) for last row group
1297        ];
1298        let mut selection = RowGroupSelection::from_row_ranges(ranges, row_group_size);
1299
1300        // Test removing existing row group
1301        selection.remove_row_group(1);
1302        assert_eq!(selection.row_count(), 15);
1303        assert_eq!(selection.row_group_count(), 2);
1304        assert!(!selection.contains_row_group(1));
1305
1306        // Verify remaining row groups' content
1307        let row_selection = selection.get(0).unwrap();
1308        assert_eq!(row_selection.row_count(), 10); // 5..15 (10)
1309
1310        let row_selection = selection.get(2).unwrap();
1311        assert_eq!(row_selection.row_count(), 5); // 0..5 (5)
1312
1313        // Test removing non-existent row group
1314        selection.remove_row_group(5);
1315        assert_eq!(selection.row_count(), 15);
1316        assert_eq!(selection.row_group_count(), 2);
1317
1318        // Test removing all row groups
1319        selection.remove_row_group(0);
1320        assert_eq!(selection.row_count(), 5);
1321        assert_eq!(selection.row_group_count(), 1);
1322
1323        let row_selection = selection.get(2).unwrap();
1324        assert_eq!(row_selection.row_count(), 5); // 0..5 (5)
1325
1326        selection.remove_row_group(2);
1327        assert_eq!(selection.row_count(), 0);
1328        assert_eq!(selection.row_group_count(), 0);
1329        assert!(selection.is_empty());
1330
1331        // Test removing from empty selection
1332        let mut empty_selection = RowGroupSelection::from_row_ranges(vec![], row_group_size);
1333        empty_selection.remove_row_group(0);
1334        assert_eq!(empty_selection.row_count(), 0);
1335        assert_eq!(empty_selection.row_group_count(), 0);
1336        assert!(empty_selection.is_empty());
1337    }
1338
1339    #[test]
1340    fn test_contains_row_group() {
1341        let row_group_size = 100;
1342        let ranges = vec![
1343            (0, vec![5..15]), // Within [0, 100)
1344            (1, vec![5..15]), // Within [0, 100)
1345        ];
1346        let selection = RowGroupSelection::from_row_ranges(ranges, row_group_size);
1347
1348        assert!(selection.contains_row_group(0));
1349        assert!(selection.contains_row_group(1));
1350        assert!(!selection.contains_row_group(2));
1351
1352        // Test empty selection
1353        let empty_selection = RowGroupSelection::from_row_ranges(vec![], row_group_size);
1354        assert!(!empty_selection.contains_row_group(0));
1355    }
1356
1357    #[test]
1358    fn test_concat() {
1359        let row_group_size = 100;
1360        let ranges1 = vec![
1361            (0, vec![5..15]), // Within [0, 100)
1362            (1, vec![5..15]), // Within [0, 100)
1363        ];
1364
1365        let ranges2 = vec![
1366            (2, vec![5..15]), // Within [0, 100)
1367            (3, vec![5..15]), // Within [0, 100)
1368        ];
1369
1370        let mut selection1 = RowGroupSelection::from_row_ranges(ranges1, row_group_size);
1371        let selection2 = RowGroupSelection::from_row_ranges(ranges2, row_group_size);
1372
1373        selection1.concat(&selection2);
1374        assert_eq!(selection1.row_count(), 40);
1375        assert_eq!(selection1.row_group_count(), 4);
1376    }
1377}
mito2/sst/parquet/row_selection.rs

mito2/sst/parquet/
row_selection.rs
