Skip to content

Commit f10a1ce

Browse files
yanghuayanghua
committed
refactor: rename RowIdSelection to RowAddrSelection
1 parent b7c8393 commit f10a1ce

File tree

1 file changed

+62
-61
lines changed

1 file changed

+62
-61
lines changed

‎rust/lance-core/src/utils/mask.rs‎

Lines changed: 62 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -374,16 +374,16 @@ pub struct RowIdTreeMap {
374374
/// The contents of the set. If there is a pair (k, Full) then the entire
375375
/// fragment k is selected. If there is a pair (k, Partial(v)) then the
376376
/// fragment k has the selected rows in v.
377-
inner: BTreeMap<u32, RowIdSelection>,
377+
inner: BTreeMap<u32, RowAddrSelection>,
378378
}
379379

380380
#[derive(Clone, Debug, PartialEq)]
381-
enum RowIdSelection {
381+
enum RowAddrSelection {
382382
Full,
383383
Partial(RoaringBitmap),
384384
}
385385

386-
impl DeepSizeOf for RowIdSelection {
386+
impl DeepSizeOf for RowAddrSelection {
387387
fn deep_size_of_children(&self, _context: &mut deepsize::Context) -> usize {
388388
match self {
389389
Self::Full => 0,
@@ -392,7 +392,7 @@ impl DeepSizeOf for RowIdSelection {
392392
}
393393
}
394394

395-
impl RowIdSelection {
395+
impl RowAddrSelection {
396396
fn union_all(selections: &[&Self]) -> Self {
397397
let mut is_full = false;
398398

@@ -434,8 +434,8 @@ impl RowIdTreeMap {
434434
self.inner
435435
.values()
436436
.map(|row_id_selection| match row_id_selection {
437-
RowIdSelection::Full => None,
438-
RowIdSelection::Partial(indices) => Some(indices.len()),
437+
RowAddrSelection::Full => None,
438+
RowAddrSelection::Partial(indices) => Some(indices.len()),
439439
})
440440
.try_fold(0_u64, |acc, next| next.map(|next| next + acc))
441441
}
@@ -449,8 +449,8 @@ impl RowIdTreeMap {
449449
.inner
450450
.iter()
451451
.filter_map(|(frag_id, row_id_selection)| match row_id_selection {
452-
RowIdSelection::Full => None,
453-
RowIdSelection::Partial(bitmap) => Some(
452+
RowAddrSelection::Full => None,
453+
RowAddrSelection::Partial(bitmap) => Some(
454454
bitmap
455455
.iter()
456456
.map(|row_offset| RowAddress::new_from_parts(*frag_id, row_offset)),
@@ -483,11 +483,11 @@ impl RowIdTreeMap {
483483
None => {
484484
let mut set = RoaringBitmap::new();
485485
set.insert(row_addr);
486-
self.inner.insert(fragment, RowIdSelection::Partial(set));
486+
self.inner.insert(fragment, RowAddrSelection::Partial(set));
487487
true
488488
}
489-
Some(RowIdSelection::Full) => false,
490-
Some(RowIdSelection::Partial(set)) => set.insert(row_addr),
489+
Some(RowAddrSelection::Full) => false,
490+
Some(RowAddrSelection::Partial(set)) => set.insert(row_addr),
491491
}
492492
}
493493

@@ -526,10 +526,10 @@ impl RowIdTreeMap {
526526
None => {
527527
let mut set = RoaringBitmap::new();
528528
count += set.insert_range(start..=end);
529-
self.inner.insert(fragment, RowIdSelection::Partial(set));
529+
self.inner.insert(fragment, RowAddrSelection::Partial(set));
530530
}
531-
Some(RowIdSelection::Full) => {}
532-
Some(RowIdSelection::Partial(set)) => {
531+
Some(RowAddrSelection::Full) => {}
532+
Some(RowAddrSelection::Partial(set)) => {
533533
count += set.insert_range(start..=end);
534534
}
535535
}
@@ -542,19 +542,20 @@ impl RowIdTreeMap {
542542

543543
/// Add a bitmap for a single fragment
544544
pub fn insert_bitmap(&mut self, fragment: u32, bitmap: RoaringBitmap) {
545-
self.inner.insert(fragment, RowIdSelection::Partial(bitmap));
545+
self.inner
546+
.insert(fragment, RowAddrSelection::Partial(bitmap));
546547
}
547548

548549
/// Add a whole fragment to the set
549550
pub fn insert_fragment(&mut self, fragment_id: u32) {
550-
self.inner.insert(fragment_id, RowIdSelection::Full);
551+
self.inner.insert(fragment_id, RowAddrSelection::Full);
551552
}
552553

553554
pub fn get_fragment_bitmap(&self, fragment_id: u32) -> Option<&RoaringBitmap> {
554555
match self.inner.get(&fragment_id) {
555556
None => None,
556-
Some(RowIdSelection::Full) => None,
557-
Some(RowIdSelection::Partial(set)) => Some(set),
557+
Some(RowAddrSelection::Full) => None,
558+
Some(RowAddrSelection::Partial(set)) => Some(set),
558559
}
559560
}
560561

@@ -564,8 +565,8 @@ impl RowIdTreeMap {
564565
let lower = value as u32;
565566
match self.inner.get(&upper) {
566567
None => false,
567-
Some(RowIdSelection::Full) => true,
568-
Some(RowIdSelection::Partial(fragment_set)) => fragment_set.contains(lower),
568+
Some(RowAddrSelection::Full) => true,
569+
Some(RowAddrSelection::Partial(fragment_set)) => fragment_set.contains(lower),
569570
}
570571
}
571572

@@ -574,13 +575,13 @@ impl RowIdTreeMap {
574575
let lower = value as u32;
575576
match self.inner.get_mut(&upper) {
576577
None => false,
577-
Some(RowIdSelection::Full) => {
578+
Some(RowAddrSelection::Full) => {
578579
let mut set = RoaringBitmap::full();
579580
set.remove(lower);
580-
self.inner.insert(upper, RowIdSelection::Partial(set));
581+
self.inner.insert(upper, RowAddrSelection::Partial(set));
581582
true
582583
}
583-
Some(RowIdSelection::Partial(lower_set)) => {
584+
Some(RowAddrSelection::Partial(lower_set)) => {
584585
let removed = lower_set.remove(lower);
585586
if lower_set.is_empty() {
586587
self.inner.remove(&upper);
@@ -603,7 +604,7 @@ impl RowIdTreeMap {
603604
for set in self.inner.values() {
604605
// Each entry is 8 bytes for the fragment id and the bitmap size
605606
size += 8;
606-
if let RowIdSelection::Partial(set) = set {
607+
if let RowAddrSelection::Partial(set) = set {
607608
size += set.serialized_size();
608609
}
609610
}
@@ -627,7 +628,7 @@ impl RowIdTreeMap {
627628
writer.write_u32::<byteorder::LittleEndian>(self.inner.len() as u32)?;
628629
for (fragment, set) in &self.inner {
629630
writer.write_u32::<byteorder::LittleEndian>(*fragment)?;
630-
if let RowIdSelection::Partial(set) = set {
631+
if let RowAddrSelection::Partial(set) = set {
631632
writer.write_u32::<byteorder::LittleEndian>(set.serialized_size() as u32)?;
632633
set.serialize_into(&mut writer)?;
633634
} else {
@@ -645,12 +646,12 @@ impl RowIdTreeMap {
645646
let fragment = reader.read_u32::<byteorder::LittleEndian>()?;
646647
let bitmap_size = reader.read_u32::<byteorder::LittleEndian>()?;
647648
if bitmap_size == 0 {
648-
inner.insert(fragment, RowIdSelection::Full);
649+
inner.insert(fragment, RowAddrSelection::Full);
649650
} else {
650651
let mut buffer = vec![0; bitmap_size as usize];
651652
reader.read_exact(&mut buffer)?;
652653
let set = RoaringBitmap::deserialize_from(&buffer[..])?;
653-
inner.insert(fragment, RowIdSelection::Partial(set));
654+
inner.insert(fragment, RowAddrSelection::Partial(set));
654655
}
655656
}
656657
Ok(Self { inner })
@@ -671,7 +672,7 @@ impl RowIdTreeMap {
671672

672673
let new_map = new_map
673674
.into_iter()
674-
.map(|(&fragment, selections)| (fragment, RowIdSelection::union_all(&selections)))
675+
.map(|(&fragment, selections)| (fragment, RowAddrSelection::union_all(&selections)))
675676
.collect();
676677

677678
Self { inner: new_map }
@@ -694,15 +695,15 @@ impl RowIdTreeMap {
694695
///
695696
/// # Safety
696697
///
697-
/// This is unsafe because if any of the inner RowIdSelection elements
698+
/// This is unsafe because if any of the inner RowAddrSelection elements
698699
/// is not a Partial then the iterator will panic because we don't know
699700
/// the size of the bitmap.
700701
pub unsafe fn into_addr_iter(self) -> impl Iterator<Item = u64> {
701702
self.inner
702703
.into_iter()
703704
.flat_map(|(fragment, selection)| match selection {
704-
RowIdSelection::Full => panic!("Size of full fragment is unknown"),
705-
RowIdSelection::Partial(bitmap) => bitmap.into_iter().map(move |val| {
705+
RowAddrSelection::Full => panic!("Size of full fragment is unknown"),
706+
RowAddrSelection::Partial(bitmap) => bitmap.into_iter().map(move |val| {
706707
let fragment = fragment as u64;
707708
let row_offset = val as u64;
708709
(fragment << 32) | row_offset
@@ -726,14 +727,14 @@ impl std::ops::BitOrAssign<Self> for RowIdTreeMap {
726727
let lhs_set = self.inner.get_mut(fragment);
727728
if let Some(lhs_set) = lhs_set {
728729
match lhs_set {
729-
RowIdSelection::Full => {
730+
RowAddrSelection::Full => {
730731
// If the fragment is already selected then there is nothing to do
731732
}
732-
RowIdSelection::Partial(lhs_bitmap) => match rhs_set {
733-
RowIdSelection::Full => {
734-
*lhs_set = RowIdSelection::Full;
733+
RowAddrSelection::Partial(lhs_bitmap) => match rhs_set {
734+
RowAddrSelection::Full => {
735+
*lhs_set = RowAddrSelection::Full;
735736
}
736-
RowIdSelection::Partial(rhs_set) => {
737+
RowAddrSelection::Partial(rhs_set) => {
737738
*lhs_bitmap |= rhs_set;
738739
}
739740
},
@@ -764,21 +765,21 @@ impl std::ops::BitAndAssign<&Self> for RowIdTreeMap {
764765
for (fragment, mut lhs_set) in &mut self.inner {
765766
match (&mut lhs_set, rhs.inner.get(fragment)) {
766767
(_, None) => {} // Already handled by retain
767-
(_, Some(RowIdSelection::Full)) => {
768+
(_, Some(RowAddrSelection::Full)) => {
768769
// Everything selected on RHS, so can leave LHS untouched.
769770
}
770-
(RowIdSelection::Partial(lhs_set), Some(RowIdSelection::Partial(rhs_set))) => {
771+
(RowAddrSelection::Partial(lhs_set), Some(RowAddrSelection::Partial(rhs_set))) => {
771772
*lhs_set &= rhs_set;
772773
}
773-
(RowIdSelection::Full, Some(RowIdSelection::Partial(rhs_set))) => {
774-
*lhs_set = RowIdSelection::Partial(rhs_set.clone());
774+
(RowAddrSelection::Full, Some(RowAddrSelection::Partial(rhs_set))) => {
775+
*lhs_set = RowAddrSelection::Partial(rhs_set.clone());
775776
}
776777
}
777778
}
778779
// Some bitmaps might now be empty. If they are, we should remove them.
779780
self.inner.retain(|_, set| match set {
780-
RowIdSelection::Partial(set) => !set.is_empty(),
781-
RowIdSelection::Full => true,
781+
RowAddrSelection::Partial(set) => !set.is_empty(),
782+
RowAddrSelection::Full => true,
782783
});
783784
}
784785
}
@@ -797,25 +798,25 @@ impl std::ops::SubAssign<&Self> for RowIdTreeMap {
797798
for (fragment, rhs_set) in &rhs.inner {
798799
match self.inner.get_mut(fragment) {
799800
None => {}
800-
Some(RowIdSelection::Full) => {
801+
Some(RowAddrSelection::Full) => {
801802
// If the fragment is already selected then there is nothing to do
802803
match rhs_set {
803-
RowIdSelection::Full => {
804+
RowAddrSelection::Full => {
804805
self.inner.remove(fragment);
805806
}
806-
RowIdSelection::Partial(rhs_set) => {
807+
RowAddrSelection::Partial(rhs_set) => {
807808
// This generally won't be hit.
808809
let mut set = RoaringBitmap::full();
809810
set -= rhs_set;
810-
self.inner.insert(*fragment, RowIdSelection::Partial(set));
811+
self.inner.insert(*fragment, RowAddrSelection::Partial(set));
811812
}
812813
}
813814
}
814-
Some(RowIdSelection::Partial(lhs_set)) => match rhs_set {
815-
RowIdSelection::Full => {
815+
Some(RowAddrSelection::Partial(lhs_set)) => match rhs_set {
816+
RowAddrSelection::Full => {
816817
self.inner.remove(fragment);
817818
}
818-
RowIdSelection::Partial(rhs_set) => {
819+
RowAddrSelection::Partial(rhs_set) => {
819820
*lhs_set -= rhs_set;
820821
if lhs_set.is_empty() {
821822
self.inner.remove(fragment);
@@ -837,12 +838,12 @@ impl FromIterator<u64> for RowIdTreeMap {
837838
None => {
838839
let mut set = RoaringBitmap::new();
839840
set.insert(lower);
840-
inner.insert(upper, RowIdSelection::Partial(set));
841+
inner.insert(upper, RowAddrSelection::Partial(set));
841842
}
842-
Some(RowIdSelection::Full) => {
843+
Some(RowAddrSelection::Full) => {
843844
// If the fragment is already selected then there is nothing to do
844845
}
845-
Some(RowIdSelection::Partial(set)) => {
846+
Some(RowAddrSelection::Partial(set)) => {
846847
set.insert(lower);
847848
}
848849
}
@@ -869,7 +870,7 @@ impl From<RoaringTreemap> for RowIdTreeMap {
869870
fn from(roaring: RoaringTreemap) -> Self {
870871
let mut inner = BTreeMap::new();
871872
for (fragment, set) in roaring.bitmaps() {
872-
inner.insert(fragment, RowIdSelection::Partial(set.clone()));
873+
inner.insert(fragment, RowAddrSelection::Partial(set.clone()));
873874
}
874875
Self { inner }
875876
}
@@ -884,12 +885,12 @@ impl Extend<u64> for RowIdTreeMap {
884885
None => {
885886
let mut set = RoaringBitmap::new();
886887
set.insert(lower);
887-
self.inner.insert(upper, RowIdSelection::Partial(set));
888+
self.inner.insert(upper, RowAddrSelection::Partial(set));
888889
}
889-
Some(RowIdSelection::Full) => {
890+
Some(RowAddrSelection::Full) => {
890891
// If the fragment is already selected then there is nothing to do
891892
}
892-
Some(RowIdSelection::Partial(set)) => {
893+
Some(RowAddrSelection::Partial(set)) => {
893894
set.insert(lower);
894895
}
895896
}
@@ -912,14 +913,14 @@ impl Extend<Self> for RowIdTreeMap {
912913
None => {
913914
self.inner.insert(fragment, set);
914915
}
915-
Some(RowIdSelection::Full) => {
916+
Some(RowAddrSelection::Full) => {
916917
// If the fragment is already selected then there is nothing to do
917918
}
918-
Some(RowIdSelection::Partial(lhs_set)) => match set {
919-
RowIdSelection::Full => {
920-
self.inner.insert(fragment, RowIdSelection::Full);
919+
Some(RowAddrSelection::Partial(lhs_set)) => match set {
920+
RowAddrSelection::Full => {
921+
self.inner.insert(fragment, RowAddrSelection::Full);
921922
}
922-
RowIdSelection::Partial(rhs_set) => {
923+
RowAddrSelection::Partial(rhs_set) => {
923924
*lhs_set |= rhs_set;
924925
}
925926
},

0 commit comments

Comments
 (0)