risingwave_hummock_sdk/
table_watermark.rs

1// Copyright 2023 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::collections::hash_map::Entry;
16use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
17use std::fmt::Display;
18use std::mem::size_of;
19use std::ops::Bound::{Excluded, Included, Unbounded};
20use std::sync::Arc;
21
22use bytes::Bytes;
23use itertools::Itertools;
24use risingwave_common::bitmap::{Bitmap, BitmapBuilder};
25use risingwave_common::catalog::TableId;
26use risingwave_common::hash::{VirtualNode, VnodeBitmapExt};
27use risingwave_common::types::ToDatumRef;
28use risingwave_common::util::sort_util::{OrderType, cmp_datum};
29use risingwave_common_estimate_size::EstimateSize;
30use risingwave_pb::hummock::table_watermarks::PbEpochNewWatermarks;
31use risingwave_pb::hummock::{
32    PbVnodeWatermark, PbWatermarkSerdeType, TableWatermarks as PbTableWatermarks,
33};
34use tracing::{debug, warn};
35
36use crate::HummockEpoch;
37use crate::key::{TableKey, TableKeyRange, prefix_slice_with_vnode, vnode};
38
39/// Narrow `key_range`, which lies in a single vnode, to the keys not cleaned by the pk-prefix
40/// `watermark` in `direction`.
41pub fn rewrite_range_with_watermark(
42    direction: WatermarkDirection,
43    watermark: &[u8],
44    key_range: &mut TableKeyRange,
45) {
46    let vnode = vnode(key_range);
47    match direction {
48        WatermarkDirection::Ascending => {
49            let overwrite_start_key = match &key_range.0 {
50                Included(start_key) | Excluded(start_key) => start_key.key_part() < watermark,
51                Unbounded => true,
52            };
53            if overwrite_start_key {
54                let watermark_key = TableKey(prefix_slice_with_vnode(vnode, watermark));
55                let fully_filtered = match &key_range.1 {
56                    Included(end_key) => end_key < &watermark_key,
57                    Excluded(end_key) => end_key <= &watermark_key,
58                    Unbounded => false,
59                };
60                if fully_filtered {
61                    key_range.1 = Excluded(watermark_key.clone());
62                }
63                key_range.0 = Included(watermark_key);
64            }
65        }
66        WatermarkDirection::Descending => {
67            let overwrite_end_key = match &key_range.1 {
68                Included(end_key) | Excluded(end_key) => end_key.key_part() > watermark,
69                Unbounded => true,
70            };
71            if overwrite_end_key {
72                let watermark_key = TableKey(prefix_slice_with_vnode(vnode, watermark));
73                let fully_filtered = match &key_range.0 {
74                    Included(start_key) => start_key > &watermark_key,
75                    Excluded(start_key) => start_key >= &watermark_key,
76                    Unbounded => false,
77                };
78                if fully_filtered {
79                    *key_range = (Included(watermark_key.clone()), Excluded(watermark_key));
80                } else {
81                    key_range.1 = Included(watermark_key);
82                }
83            }
84        }
85    }
86}
87
88#[derive(Clone)]
89pub struct ReadTableWatermark {
90    pub direction: WatermarkDirection,
91    pub vnode_watermarks: BTreeMap<VirtualNode, Bytes>,
92}
93
94#[derive(Clone)]
95pub struct TableWatermarksIndex {
96    pub watermark_direction: WatermarkDirection,
97    // later epoch at the back
98    pub staging_watermarks: VecDeque<(HummockEpoch, Arc<[VnodeWatermark]>)>,
99    pub committed_watermarks: Option<Arc<TableWatermarks>>,
100    latest_epoch: HummockEpoch,
101    committed_epoch: Option<HummockEpoch>,
102    watermark_type: WatermarkSerdeType,
103}
104
105impl TableWatermarksIndex {
106    pub fn new(
107        watermark_direction: WatermarkDirection,
108        first_epoch: HummockEpoch,
109        first_vnode_watermark: Vec<VnodeWatermark>,
110        committed_epoch: Option<HummockEpoch>,
111        watermark_type: WatermarkSerdeType,
112    ) -> Self {
113        if let Some(committed_epoch) = committed_epoch {
114            assert!(first_epoch > committed_epoch);
115        }
116        Self {
117            watermark_direction,
118            staging_watermarks: VecDeque::from_iter([(
119                first_epoch,
120                Arc::from(first_vnode_watermark),
121            )]),
122            committed_watermarks: None,
123            latest_epoch: first_epoch,
124            committed_epoch,
125            watermark_type,
126        }
127    }
128
129    pub fn new_committed(
130        committed_watermarks: Arc<TableWatermarks>,
131        committed_epoch: HummockEpoch,
132        watermark_type: WatermarkSerdeType,
133    ) -> Self {
134        Self {
135            watermark_direction: committed_watermarks.direction,
136            staging_watermarks: VecDeque::new(),
137            committed_epoch: Some(committed_epoch),
138            latest_epoch: committed_epoch,
139            committed_watermarks: Some(committed_watermarks),
140            watermark_type,
141        }
142    }
143
144    pub fn read_watermark(&self, vnode: VirtualNode, epoch: HummockEpoch) -> Option<Bytes> {
145        // iterate from new epoch to old epoch
146        for (watermark_epoch, vnode_watermark_list) in self.staging_watermarks.iter().rev().chain(
147            self.committed_watermarks
148                .iter()
149                .flat_map(|watermarks| watermarks.watermarks.iter().rev()),
150        ) {
151            if *watermark_epoch > epoch {
152                continue;
153            }
154            for vnode_watermark in vnode_watermark_list.as_ref() {
155                if vnode_watermark.vnode_bitmap.is_set(vnode.to_index()) {
156                    return Some(vnode_watermark.watermark.clone());
157                }
158            }
159        }
160        None
161    }
162
163    pub fn latest_watermark(&self, vnode: VirtualNode) -> Option<Bytes> {
164        self.read_watermark(vnode, HummockEpoch::MAX)
165    }
166
167    pub fn rewrite_range_with_table_watermark(
168        &self,
169        epoch: HummockEpoch,
170        key_range: &mut TableKeyRange,
171    ) {
172        if !matches!(self.watermark_type, WatermarkSerdeType::PkPrefix) {
173            // TODO: Storage read respecting watermark will be implemented in a later PR in StateTable.
174            return;
175        }
176        if let Some(watermark) = self.read_watermark(vnode(key_range), epoch) {
177            rewrite_range_with_watermark(self.watermark_direction, &watermark, key_range);
178        }
179    }
180
181    pub fn filter_regress_watermarks(&self, watermarks: &mut Vec<VnodeWatermark>) {
182        if !matches!(self.watermark_type, WatermarkSerdeType::PkPrefix) {
183            // The current kv log store can emit regressed watermarks, but only for WatermarkSerdeType::PkPrefix.
184            return;
185        }
186        let mut ret = Vec::with_capacity(watermarks.len());
187        for watermark in watermarks.drain(..) {
188            let vnode_count = watermark.vnode_count();
189
190            let mut regress_vnodes = None;
191            for vnode in watermark.vnode_bitmap.iter_vnodes() {
192                if let Some(prev_watermark) = self.latest_watermark(vnode) {
193                    let is_regress = match self.direction() {
194                        WatermarkDirection::Ascending => prev_watermark > watermark.watermark,
195                        WatermarkDirection::Descending => prev_watermark < watermark.watermark,
196                    };
197                    if is_regress {
198                        warn!(
199                            "table watermark regress: {:?} {} {:?} {:?}",
200                            self.direction(),
201                            vnode,
202                            watermark.watermark,
203                            prev_watermark
204                        );
205                        regress_vnodes
206                            .get_or_insert_with(|| BitmapBuilder::zeroed(vnode_count))
207                            .set(vnode.to_index(), true);
208                    }
209                }
210            }
211            if let Some(regress_vnodes) = regress_vnodes {
212                let mut bitmap_builder = None;
213                for vnode in watermark.vnode_bitmap.iter_vnodes() {
214                    let vnode_index = vnode.to_index();
215                    if !regress_vnodes.is_set(vnode_index) {
216                        bitmap_builder
217                            .get_or_insert_with(|| BitmapBuilder::zeroed(vnode_count))
218                            .set(vnode_index, true);
219                    }
220                }
221                if let Some(bitmap_builder) = bitmap_builder {
222                    ret.push(VnodeWatermark::new(
223                        Arc::new(bitmap_builder.finish()),
224                        watermark.watermark,
225                    ));
226                }
227            } else {
228                // no vnode has regress watermark
229                ret.push(watermark);
230            }
231        }
232        *watermarks = ret;
233    }
234
235    pub fn direction(&self) -> WatermarkDirection {
236        self.watermark_direction
237    }
238
239    pub fn add_epoch_watermark(
240        &mut self,
241        epoch: HummockEpoch,
242        vnode_watermark_list: Arc<[VnodeWatermark]>,
243        direction: WatermarkDirection,
244    ) {
245        assert!(epoch > self.latest_epoch);
246        assert_eq!(self.watermark_direction, direction);
247        self.latest_epoch = epoch;
248        #[cfg(debug_assertions)]
249        if !vnode_watermark_list.is_empty() {
250            let vnode_count = vnode_watermark_list[0].vnode_count();
251            let mut vnode_is_set = BitmapBuilder::zeroed(vnode_count);
252            for vnode_watermark in vnode_watermark_list.as_ref() {
253                for vnode in vnode_watermark.vnode_bitmap.iter_ones() {
254                    assert!(!vnode_is_set.is_set(vnode));
255                    vnode_is_set.set(vnode, true);
256                    let vnode = VirtualNode::from_index(vnode);
257                    if let Some(prev_watermark) = self.latest_watermark(vnode) {
258                        match self.watermark_direction {
259                            WatermarkDirection::Ascending => {
260                                assert!(vnode_watermark.watermark >= prev_watermark);
261                            }
262                            WatermarkDirection::Descending => {
263                                assert!(vnode_watermark.watermark <= prev_watermark);
264                            }
265                        };
266                    }
267                }
268            }
269        }
270        self.staging_watermarks
271            .push_back((epoch, vnode_watermark_list));
272    }
273
274    pub fn apply_committed_watermarks(
275        &mut self,
276        committed_watermark: Arc<TableWatermarks>,
277        committed_epoch: HummockEpoch,
278    ) {
279        assert_eq!(self.watermark_direction, committed_watermark.direction);
280        if let Some(prev_committed_epoch) = self.committed_epoch {
281            assert!(prev_committed_epoch <= committed_epoch);
282            if prev_committed_epoch == committed_epoch {
283                return;
284            }
285        }
286        if self.latest_epoch < committed_epoch {
287            debug!(
288                latest_epoch = self.latest_epoch,
289                committed_epoch, "committed_epoch exceed table watermark latest_epoch"
290            );
291            self.latest_epoch = committed_epoch;
292        }
293        self.committed_epoch = Some(committed_epoch);
294        self.committed_watermarks = Some(committed_watermark);
295        // keep only watermark higher than committed epoch
296        while let Some((old_epoch, _)) = self.staging_watermarks.front()
297            && *old_epoch <= committed_epoch
298        {
299            let _ = self.staging_watermarks.pop_front();
300        }
301    }
302}
303
304#[derive(Debug, Clone, Copy, Eq, PartialEq)]
305pub enum WatermarkDirection {
306    Ascending,
307    Descending,
308}
309
310impl Display for WatermarkDirection {
311    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
312        match self {
313            WatermarkDirection::Ascending => write!(f, "Ascending"),
314            WatermarkDirection::Descending => write!(f, "Descending"),
315        }
316    }
317}
318
319impl WatermarkDirection {
320    pub fn key_filter_by_watermark(
321        &self,
322        key: impl AsRef<[u8]>,
323        watermark: impl AsRef<[u8]>,
324    ) -> bool {
325        let key = key.as_ref();
326        let watermark = watermark.as_ref();
327        match self {
328            WatermarkDirection::Ascending => key < watermark,
329            WatermarkDirection::Descending => key > watermark,
330        }
331    }
332
333    pub fn datum_filter_by_watermark(
334        &self,
335        watermark_col: impl ToDatumRef,
336        watermark: impl ToDatumRef,
337        order_type: OrderType,
338    ) -> bool {
339        let watermark_col = watermark_col.to_datum_ref();
340        let watermark = watermark.to_datum_ref();
341        match self {
342            WatermarkDirection::Ascending => {
343                // watermark_col < watermark
344                cmp_datum(watermark_col, watermark, order_type).is_lt()
345            }
346            WatermarkDirection::Descending => {
347                //  watermark_col > watermark
348                cmp_datum(watermark_col, watermark, order_type).is_gt()
349            }
350        }
351    }
352
353    pub fn is_ascending(&self) -> bool {
354        match self {
355            WatermarkDirection::Ascending => true,
356            WatermarkDirection::Descending => false,
357        }
358    }
359}
360
361#[derive(Clone, Debug, PartialEq, EstimateSize)]
362pub struct VnodeWatermark {
363    vnode_bitmap: Arc<Bitmap>,
364    watermark: Bytes,
365}
366
367impl VnodeWatermark {
368    pub fn new(vnode_bitmap: Arc<Bitmap>, watermark: Bytes) -> Self {
369        Self {
370            vnode_bitmap,
371            watermark,
372        }
373    }
374
375    pub fn vnode_bitmap(&self) -> &Bitmap {
376        &self.vnode_bitmap
377    }
378
379    /// Vnode count derived from the bitmap.
380    pub fn vnode_count(&self) -> usize {
381        self.vnode_bitmap.len()
382    }
383
384    pub fn watermark(&self) -> &Bytes {
385        &self.watermark
386    }
387
388    pub fn to_protobuf(&self) -> PbVnodeWatermark {
389        self.into()
390    }
391}
392
393impl From<PbVnodeWatermark> for VnodeWatermark {
394    fn from(pb: PbVnodeWatermark) -> Self {
395        Self {
396            vnode_bitmap: Arc::new(Bitmap::from(pb.vnode_bitmap.as_ref().unwrap())),
397            watermark: Bytes::from(pb.watermark),
398        }
399    }
400}
401
402impl From<&PbVnodeWatermark> for VnodeWatermark {
403    fn from(pb: &PbVnodeWatermark) -> Self {
404        Self {
405            vnode_bitmap: Arc::new(Bitmap::from(pb.vnode_bitmap.as_ref().unwrap())),
406            watermark: Bytes::from(pb.watermark.clone()),
407        }
408    }
409}
410
411impl From<VnodeWatermark> for PbVnodeWatermark {
412    fn from(watermark: VnodeWatermark) -> Self {
413        Self {
414            watermark: watermark.watermark.into(),
415            vnode_bitmap: Some(watermark.vnode_bitmap.to_protobuf()),
416        }
417    }
418}
419
420impl From<&VnodeWatermark> for PbVnodeWatermark {
421    fn from(watermark: &VnodeWatermark) -> Self {
422        Self {
423            watermark: watermark.watermark.to_vec(),
424            vnode_bitmap: Some(watermark.vnode_bitmap.to_protobuf()),
425        }
426    }
427}
428
429#[derive(Clone, Debug, PartialEq)]
430pub struct TableWatermarks {
431    // later epoch at the back
432    pub watermarks: Vec<(HummockEpoch, Arc<[VnodeWatermark]>)>,
433    pub direction: WatermarkDirection,
434    pub watermark_type: WatermarkSerdeType,
435}
436
437impl TableWatermarks {
438    pub fn single_epoch(
439        epoch: HummockEpoch,
440        watermarks: Vec<VnodeWatermark>,
441        direction: WatermarkDirection,
442        watermark_type: WatermarkSerdeType,
443    ) -> Self {
444        let mut this = Self {
445            direction,
446            watermarks: Vec::new(),
447            watermark_type,
448        };
449        this.add_new_epoch_watermarks(epoch, watermarks.into(), direction, watermark_type);
450        this
451    }
452
453    pub fn add_new_epoch_watermarks(
454        &mut self,
455        epoch: HummockEpoch,
456        watermarks: Arc<[VnodeWatermark]>,
457        direction: WatermarkDirection,
458        watermark_type: WatermarkSerdeType,
459    ) {
460        assert_eq!(self.direction, direction);
461        assert_eq!(self.watermark_type, watermark_type);
462
463        if let Some((prev_epoch, _)) = self.watermarks.last() {
464            assert!(*prev_epoch < epoch);
465        }
466        if !watermarks.is_empty() {
467            let vnode_count = watermarks[0].vnode_count();
468            for watermark in &*watermarks {
469                assert_eq!(watermark.vnode_count(), vnode_count);
470            }
471            if let Some(existing_vnode_count) = self.vnode_count() {
472                assert_eq!(existing_vnode_count, vnode_count);
473            }
474        }
475        self.watermarks.push((epoch, watermarks));
476    }
477
478    /// Vnode count derived from existing watermarks. Returns `None` if there is no watermark.
479    fn vnode_count(&self) -> Option<usize> {
480        self.watermarks
481            .iter()
482            .flat_map(|(_, watermarks)| watermarks.as_ref())
483            .next()
484            .map(|w| w.vnode_count())
485    }
486}
487
488pub fn merge_multiple_new_table_watermarks(
489    table_watermarks_list: impl IntoIterator<Item = HashMap<TableId, TableWatermarks>>,
490) -> HashMap<TableId, TableWatermarks> {
491    #[expect(clippy::type_complexity)]
492    let mut ret: HashMap<
493        TableId,
494        (
495            WatermarkDirection,
496            BTreeMap<u64, Vec<VnodeWatermark>>,
497            WatermarkSerdeType,
498        ),
499    > = HashMap::new();
500    for table_watermarks in table_watermarks_list {
501        for (table_id, new_table_watermarks) in table_watermarks {
502            let epoch_watermarks = match ret.entry(table_id) {
503                Entry::Occupied(entry) => {
504                    let (direction, epoch_watermarks, watermark_type) = entry.into_mut();
505                    assert_eq!(&new_table_watermarks.direction, direction);
506                    assert_eq!(&new_table_watermarks.watermark_type, watermark_type);
507                    epoch_watermarks
508                }
509                Entry::Vacant(entry) => {
510                    let (_, epoch_watermarks, _) = entry.insert((
511                        new_table_watermarks.direction,
512                        BTreeMap::new(),
513                        new_table_watermarks.watermark_type,
514                    ));
515                    epoch_watermarks
516                }
517            };
518            for (new_epoch, new_epoch_watermarks) in new_table_watermarks.watermarks {
519                epoch_watermarks
520                    .entry(new_epoch)
521                    .or_insert_with(Vec::new)
522                    .extend(new_epoch_watermarks.iter().cloned());
523            }
524        }
525    }
526    ret.into_iter()
527        .map(
528            |(table_id, (direction, epoch_watermarks, watermark_type))| {
529                (
530                    table_id,
531                    TableWatermarks {
532                        direction,
533                        // ordered from earlier epoch to later epoch
534                        watermarks: epoch_watermarks
535                            .into_iter()
536                            .map(|(epoch, watermarks)| (epoch, Arc::from(watermarks)))
537                            .collect(),
538                        watermark_type,
539                    },
540                )
541            },
542        )
543        .collect()
544}
545
546impl TableWatermarks {
547    pub fn apply_new_table_watermarks(&mut self, newly_added_watermarks: &TableWatermarks) {
548        assert_eq!(self.direction, newly_added_watermarks.direction);
549        assert!(self.watermarks.iter().map(|(epoch, _)| epoch).is_sorted());
550        assert!(
551            newly_added_watermarks
552                .watermarks
553                .iter()
554                .map(|(epoch, _)| epoch)
555                .is_sorted()
556        );
557        // ensure that the newly added watermarks have a later epoch than the previous latest epoch.
558        if let Some((prev_last_epoch, _)) = self.watermarks.last()
559            && let Some((new_first_epoch, _)) = newly_added_watermarks.watermarks.first()
560        {
561            assert!(prev_last_epoch < new_first_epoch);
562        }
563        self.watermarks.extend(
564            newly_added_watermarks
565                .watermarks
566                .iter()
567                .map(|(epoch, new_watermarks)| (*epoch, new_watermarks.clone())),
568        );
569    }
570
571    pub fn clear_stale_epoch_watermark(&mut self, safe_epoch: u64) {
572        match self.watermarks.first() {
573            None => {
574                // return on empty watermark
575                return;
576            }
577            Some((earliest_epoch, _)) => {
578                if *earliest_epoch >= safe_epoch {
579                    // No stale epoch watermark needs to be cleared.
580                    return;
581                }
582            }
583        }
584        debug!("clear stale table watermark below epoch {}", safe_epoch);
585        let mut result_epoch_watermark = Vec::with_capacity(self.watermarks.len());
586        let mut set_vnode: HashSet<VirtualNode> = HashSet::new();
587        let mut vnode_count: Option<usize> = None; // lazy initialized on first occurrence of vnode watermark
588        while let Some((epoch, _)) = self.watermarks.last() {
589            if *epoch >= safe_epoch {
590                let (epoch, watermarks) = self.watermarks.pop().expect("have check Some");
591                for watermark in watermarks.as_ref() {
592                    vnode_count.get_or_insert_with(|| watermark.vnode_count());
593                    for vnode in watermark.vnode_bitmap.iter_vnodes() {
594                        set_vnode.insert(vnode);
595                    }
596                }
597                result_epoch_watermark.push((epoch, watermarks));
598            } else {
599                break;
600            }
601        }
602        while vnode_count != Some(set_vnode.len())
603            && let Some((_, watermarks)) = self.watermarks.pop()
604        {
605            let mut new_vnode_watermarks = Vec::new();
606            for vnode_watermark in watermarks.as_ref() {
607                let mut new_set_vnode = Vec::new();
608                vnode_count.get_or_insert_with(|| vnode_watermark.vnode_count());
609                for vnode in vnode_watermark.vnode_bitmap.iter_vnodes() {
610                    if set_vnode.insert(vnode) {
611                        new_set_vnode.push(vnode);
612                    }
613                }
614                if !new_set_vnode.is_empty() {
615                    let mut builder = BitmapBuilder::zeroed(vnode_watermark.vnode_count());
616                    for vnode in new_set_vnode {
617                        builder.set(vnode.to_index(), true);
618                    }
619                    let bitmap = Arc::new(builder.finish());
620                    new_vnode_watermarks.push(VnodeWatermark {
621                        vnode_bitmap: bitmap,
622                        watermark: vnode_watermark.watermark.clone(),
623                    })
624                }
625            }
626            if !new_vnode_watermarks.is_empty() {
627                if let Some((last_epoch, last_watermarks)) = result_epoch_watermark.last_mut()
628                    && *last_epoch == safe_epoch
629                {
630                    *last_watermarks = Arc::from(
631                        last_watermarks
632                            .iter()
633                            .cloned()
634                            .chain(new_vnode_watermarks)
635                            .collect_vec(),
636                    );
637                } else {
638                    result_epoch_watermark.push((safe_epoch, Arc::from(new_vnode_watermarks)));
639                }
640            }
641        }
642        // epoch watermark are added from later epoch to earlier epoch.
643        // reverse to ensure that earlier epochs are at the front
644        result_epoch_watermark.reverse();
645        assert!(
646            result_epoch_watermark
647                .is_sorted_by(|(first_epoch, _), (second_epoch, _)| { first_epoch < second_epoch })
648        );
649        *self = TableWatermarks {
650            watermarks: result_epoch_watermark,
651            direction: self.direction,
652            watermark_type: self.watermark_type,
653        }
654    }
655}
656
657impl TableWatermarks {
658    pub fn estimated_encode_len(&self) -> usize {
659        self.watermarks.len() * size_of::<HummockEpoch>()
660            + self
661                .watermarks
662                .iter()
663                .map(|(_, watermarks)| {
664                    watermarks
665                        .iter()
666                        .map(|watermark| watermark.estimated_size())
667                        .sum::<usize>()
668                })
669                .sum::<usize>()
670            + size_of::<bool>() // for direction
671    }
672
673    pub fn to_protobuf(&self) -> PbTableWatermarks {
674        self.into()
675    }
676}
677
678impl From<&PbTableWatermarks> for WatermarkSerdeType {
679    fn from(pb: &PbTableWatermarks) -> Self {
680        match pb.raw_watermark_serde_type() {
681            PbWatermarkSerdeType::TypeUnspecified => {
682                // For backward compatibility.
683                #[expect(deprecated)]
684                if pb.is_non_pk_prefix {
685                    WatermarkSerdeType::NonPkPrefix
686                } else {
687                    WatermarkSerdeType::PkPrefix
688                }
689            }
690            PbWatermarkSerdeType::PkPrefix => WatermarkSerdeType::PkPrefix,
691            PbWatermarkSerdeType::NonPkPrefix => WatermarkSerdeType::NonPkPrefix,
692            PbWatermarkSerdeType::Value => WatermarkSerdeType::Value,
693        }
694    }
695}
696impl From<&PbTableWatermarks> for TableWatermarks {
697    fn from(pb: &PbTableWatermarks) -> Self {
698        let watermark_type = WatermarkSerdeType::from(pb);
699        Self {
700            watermarks: pb
701                .epoch_watermarks
702                .iter()
703                .map(|epoch_watermark| {
704                    let epoch = epoch_watermark.epoch;
705                    let watermarks = epoch_watermark
706                        .watermarks
707                        .iter()
708                        .map(VnodeWatermark::from)
709                        .collect();
710                    (epoch, watermarks)
711                })
712                .collect(),
713            direction: if pb.is_ascending {
714                WatermarkDirection::Ascending
715            } else {
716                WatermarkDirection::Descending
717            },
718            watermark_type,
719        }
720    }
721}
722
723impl From<WatermarkSerdeType> for PbWatermarkSerdeType {
724    fn from(s: WatermarkSerdeType) -> Self {
725        match s {
726            WatermarkSerdeType::PkPrefix => PbWatermarkSerdeType::PkPrefix,
727            WatermarkSerdeType::NonPkPrefix => PbWatermarkSerdeType::NonPkPrefix,
728            WatermarkSerdeType::Value => PbWatermarkSerdeType::Value,
729        }
730    }
731}
732
733impl From<&TableWatermarks> for PbTableWatermarks {
734    fn from(table_watermarks: &TableWatermarks) -> Self {
735        let is_non_pk_prefix = match table_watermarks.watermark_type {
736            WatermarkSerdeType::PkPrefix => false,
737            WatermarkSerdeType::NonPkPrefix => true,
738            WatermarkSerdeType::Value => false,
739        };
740        Self {
741            epoch_watermarks: table_watermarks
742                .watermarks
743                .iter()
744                .map(|(epoch, watermarks)| PbEpochNewWatermarks {
745                    watermarks: watermarks.iter().map(|wm| wm.into()).collect(),
746                    epoch: *epoch,
747                })
748                .collect(),
749            is_ascending: match table_watermarks.direction {
750                WatermarkDirection::Ascending => true,
751                WatermarkDirection::Descending => false,
752            },
753            #[expect(deprecated)]
754            is_non_pk_prefix,
755            raw_watermark_serde_type: PbWatermarkSerdeType::from(table_watermarks.watermark_type)
756                as i32,
757        }
758    }
759}
760
761impl From<PbTableWatermarks> for TableWatermarks {
762    fn from(pb: PbTableWatermarks) -> Self {
763        let watermark_type = WatermarkSerdeType::from(&pb);
764        Self {
765            watermarks: pb
766                .epoch_watermarks
767                .into_iter()
768                .map(|epoch_watermark| {
769                    let epoch = epoch_watermark.epoch;
770                    let watermarks = epoch_watermark
771                        .watermarks
772                        .into_iter()
773                        .map(VnodeWatermark::from)
774                        .collect();
775                    (epoch, watermarks)
776                })
777                .collect(),
778            direction: if pb.is_ascending {
779                WatermarkDirection::Ascending
780            } else {
781                WatermarkDirection::Descending
782            },
783            watermark_type,
784        }
785    }
786}
787
788impl From<TableWatermarks> for PbTableWatermarks {
789    fn from(table_watermarks: TableWatermarks) -> Self {
790        Self {
791            epoch_watermarks: table_watermarks
792                .watermarks
793                .into_iter()
794                .map(|(epoch, watermarks)| PbEpochNewWatermarks {
795                    watermarks: watermarks.iter().map(PbVnodeWatermark::from).collect(),
796                    epoch,
797                })
798                .collect(),
799            is_ascending: match table_watermarks.direction {
800                WatermarkDirection::Ascending => true,
801                WatermarkDirection::Descending => false,
802            },
803            #[expect(deprecated)]
804            is_non_pk_prefix: match table_watermarks.watermark_type {
805                WatermarkSerdeType::NonPkPrefix => true,
806                WatermarkSerdeType::PkPrefix => false,
807                WatermarkSerdeType::Value => false,
808            },
809            raw_watermark_serde_type: PbWatermarkSerdeType::from(table_watermarks.watermark_type)
810                as i32,
811        }
812    }
813}
814
815#[derive(Debug, Clone, Copy, PartialEq, Eq)]
816pub enum WatermarkSerdeType {
817    PkPrefix,
818    NonPkPrefix,
819    Value,
820}
821
822#[cfg(test)]
823mod tests {
824    use std::collections::Bound::Included;
825    use std::collections::{Bound, HashMap};
826    use std::ops::Bound::Excluded;
827    use std::sync::Arc;
828    use std::vec;
829
830    use bytes::Bytes;
831    use itertools::Itertools;
832    use risingwave_common::bitmap::{Bitmap, BitmapBuilder};
833    use risingwave_common::catalog::TableId;
834    use risingwave_common::hash::VirtualNode;
835    use risingwave_common::util::epoch::{EpochExt, test_epoch};
836    use risingwave_pb::hummock::{PbHummockVersion, StateTableInfo};
837
838    use crate::compaction_group::StaticCompactionGroupId;
839    use crate::key::{TableKeyRange, is_empty_key_range, prefixed_range_with_vnode};
840    use crate::table_watermark::{
841        TableWatermarks, TableWatermarksIndex, VnodeWatermark, WatermarkDirection,
842        WatermarkSerdeType, merge_multiple_new_table_watermarks,
843    };
844    use crate::version::HummockVersion;
845
846    fn build_bitmap(vnodes: impl IntoIterator<Item = usize>) -> Arc<Bitmap> {
847        let mut builder = BitmapBuilder::zeroed(VirtualNode::COUNT_FOR_TEST);
848        for vnode in vnodes {
849            builder.set(vnode, true);
850        }
851        Arc::new(builder.finish())
852    }
853
854    #[test]
855    fn test_apply_new_table_watermark() {
856        let epoch1 = test_epoch(1);
857        let direction = WatermarkDirection::Ascending;
858        let watermark1 = Bytes::from("watermark1");
859        let watermark2 = Bytes::from("watermark2");
860        let watermark3 = Bytes::from("watermark3");
861        let watermark4 = Bytes::from("watermark4");
862        let watermark_type = WatermarkSerdeType::PkPrefix;
863        let mut table_watermarks = TableWatermarks::single_epoch(
864            epoch1,
865            vec![VnodeWatermark::new(build_bitmap(vec![0, 1, 2]), watermark1)],
866            direction,
867            watermark_type,
868        );
869        let epoch2 = epoch1.next_epoch();
870        table_watermarks.add_new_epoch_watermarks(
871            epoch2,
872            vec![VnodeWatermark::new(
873                build_bitmap(vec![0, 1, 2, 3]),
874                watermark2,
875            )]
876            .into(),
877            direction,
878            watermark_type,
879        );
880
881        let mut table_watermark_checkpoint = table_watermarks.clone();
882
883        let epoch3 = epoch2.next_epoch();
884        let mut second_table_watermark = TableWatermarks::single_epoch(
885            epoch3,
886            vec![VnodeWatermark::new(
887                build_bitmap(0..VirtualNode::COUNT_FOR_TEST),
888                watermark3.clone(),
889            )],
890            direction,
891            watermark_type,
892        );
893        table_watermarks.add_new_epoch_watermarks(
894            epoch3,
895            vec![VnodeWatermark::new(
896                build_bitmap(0..VirtualNode::COUNT_FOR_TEST),
897                watermark3,
898            )]
899            .into(),
900            direction,
901            watermark_type,
902        );
903        let epoch4 = epoch3.next_epoch();
904        let epoch5 = epoch4.next_epoch();
905        table_watermarks.add_new_epoch_watermarks(
906            epoch5,
907            vec![VnodeWatermark::new(
908                build_bitmap(vec![0, 3, 4]),
909                watermark4.clone(),
910            )]
911            .into(),
912            direction,
913            watermark_type,
914        );
915        second_table_watermark.add_new_epoch_watermarks(
916            epoch5,
917            vec![VnodeWatermark::new(build_bitmap(vec![0, 3, 4]), watermark4)].into(),
918            direction,
919            watermark_type,
920        );
921
922        table_watermark_checkpoint.apply_new_table_watermarks(&second_table_watermark);
923        assert_eq!(table_watermarks, table_watermark_checkpoint);
924    }
925
926    #[test]
927    fn test_clear_stale_epoch_watmermark() {
928        let epoch1 = test_epoch(1);
929        let direction = WatermarkDirection::Ascending;
930        let watermark1 = Bytes::from("watermark1");
931        let watermark2 = Bytes::from("watermark2");
932        let watermark3 = Bytes::from("watermark3");
933        let watermark4 = Bytes::from("watermark4");
934        let watermark_type = WatermarkSerdeType::PkPrefix;
935        let mut table_watermarks = TableWatermarks::single_epoch(
936            epoch1,
937            vec![VnodeWatermark::new(build_bitmap(vec![0, 1, 2]), watermark1)],
938            direction,
939            watermark_type,
940        );
941        let epoch2 = epoch1.next_epoch();
942        table_watermarks.add_new_epoch_watermarks(
943            epoch2,
944            vec![VnodeWatermark::new(
945                build_bitmap(vec![0, 1, 2, 3]),
946                watermark2.clone(),
947            )]
948            .into(),
949            direction,
950            watermark_type,
951        );
952        let epoch3 = epoch2.next_epoch();
953        table_watermarks.add_new_epoch_watermarks(
954            epoch3,
955            vec![VnodeWatermark::new(
956                build_bitmap(0..VirtualNode::COUNT_FOR_TEST),
957                watermark3.clone(),
958            )]
959            .into(),
960            direction,
961            watermark_type,
962        );
963        let epoch4 = epoch3.next_epoch();
964        let epoch5 = epoch4.next_epoch();
965        table_watermarks.add_new_epoch_watermarks(
966            epoch5,
967            vec![VnodeWatermark::new(
968                build_bitmap(vec![0, 3, 4]),
969                watermark4.clone(),
970            )]
971            .into(),
972            direction,
973            watermark_type,
974        );
975
976        let mut table_watermarks_checkpoint = table_watermarks.clone();
977        table_watermarks_checkpoint.clear_stale_epoch_watermark(epoch1);
978        assert_eq!(table_watermarks_checkpoint, table_watermarks);
979
980        table_watermarks_checkpoint.clear_stale_epoch_watermark(epoch2);
981        assert_eq!(
982            table_watermarks_checkpoint,
983            TableWatermarks {
984                watermarks: vec![
985                    (
986                        epoch2,
987                        vec![VnodeWatermark::new(
988                            build_bitmap(vec![0, 1, 2, 3]),
989                            watermark2,
990                        )]
991                        .into()
992                    ),
993                    (
994                        epoch3,
995                        vec![VnodeWatermark::new(
996                            build_bitmap(0..VirtualNode::COUNT_FOR_TEST),
997                            watermark3.clone(),
998                        )]
999                        .into()
1000                    ),
1001                    (
1002                        epoch5,
1003                        vec![VnodeWatermark::new(
1004                            build_bitmap(vec![0, 3, 4]),
1005                            watermark4.clone(),
1006                        )]
1007                        .into()
1008                    )
1009                ],
1010                direction,
1011                watermark_type,
1012            }
1013        );
1014
1015        table_watermarks_checkpoint.clear_stale_epoch_watermark(epoch3);
1016        assert_eq!(
1017            table_watermarks_checkpoint,
1018            TableWatermarks {
1019                watermarks: vec![
1020                    (
1021                        epoch3,
1022                        vec![VnodeWatermark::new(
1023                            build_bitmap(0..VirtualNode::COUNT_FOR_TEST),
1024                            watermark3.clone(),
1025                        )]
1026                        .into()
1027                    ),
1028                    (
1029                        epoch5,
1030                        vec![VnodeWatermark::new(
1031                            build_bitmap(vec![0, 3, 4]),
1032                            watermark4.clone(),
1033                        )]
1034                        .into()
1035                    )
1036                ],
1037                direction,
1038                watermark_type,
1039            }
1040        );
1041
1042        table_watermarks_checkpoint.clear_stale_epoch_watermark(epoch4);
1043        assert_eq!(
1044            table_watermarks_checkpoint,
1045            TableWatermarks {
1046                watermarks: vec![
1047                    (
1048                        epoch4,
1049                        vec![VnodeWatermark::new(
1050                            build_bitmap((1..3).chain(5..VirtualNode::COUNT_FOR_TEST)),
1051                            watermark3.clone()
1052                        )]
1053                        .into()
1054                    ),
1055                    (
1056                        epoch5,
1057                        vec![VnodeWatermark::new(
1058                            build_bitmap(vec![0, 3, 4]),
1059                            watermark4.clone(),
1060                        )]
1061                        .into()
1062                    )
1063                ],
1064                direction,
1065                watermark_type,
1066            }
1067        );
1068
1069        table_watermarks_checkpoint.clear_stale_epoch_watermark(epoch5);
1070        assert_eq!(
1071            table_watermarks_checkpoint,
1072            TableWatermarks {
1073                watermarks: vec![(
1074                    epoch5,
1075                    vec![
1076                        VnodeWatermark::new(build_bitmap(vec![0, 3, 4]), watermark4),
1077                        VnodeWatermark::new(
1078                            build_bitmap((1..3).chain(5..VirtualNode::COUNT_FOR_TEST)),
1079                            watermark3
1080                        )
1081                    ]
1082                    .into()
1083                )],
1084                direction,
1085                watermark_type,
1086            }
1087        );
1088    }
1089
1090    #[test]
1091    fn test_merge_multiple_new_table_watermarks() {
1092        fn epoch_new_watermark(epoch: u64, bitmaps: Vec<&Bitmap>) -> (u64, Arc<[VnodeWatermark]>) {
1093            (
1094                epoch,
1095                bitmaps
1096                    .into_iter()
1097                    .map(|bitmap| VnodeWatermark {
1098                        watermark: Bytes::from(vec![1, 2, epoch as _]),
1099                        vnode_bitmap: Arc::new(bitmap.clone()),
1100                    })
1101                    .collect_vec()
1102                    .into(),
1103            )
1104        }
1105        fn build_table_watermark(
1106            vnodes: impl IntoIterator<Item = usize>,
1107            epochs: impl IntoIterator<Item = u64>,
1108        ) -> TableWatermarks {
1109            let bitmap = build_bitmap(vnodes);
1110            TableWatermarks {
1111                watermarks: epochs
1112                    .into_iter()
1113                    .map(|epoch: u64| epoch_new_watermark(epoch, vec![&bitmap]))
1114                    .collect(),
1115                direction: WatermarkDirection::Ascending,
1116                watermark_type: WatermarkSerdeType::PkPrefix,
1117            }
1118        }
1119        let table1_watermark1 = build_table_watermark(0..3, vec![1, 2, 4]);
1120        let table1_watermark2 = build_table_watermark(4..6, vec![1, 2, 5]);
1121        let table2_watermark = build_table_watermark(0..4, 1..3);
1122        let table3_watermark = build_table_watermark(0..4, 3..5);
1123        let mut first = HashMap::new();
1124        first.insert(TableId::new(1), table1_watermark1);
1125        first.insert(TableId::new(2), table2_watermark.clone());
1126        let mut second = HashMap::new();
1127        second.insert(TableId::new(1), table1_watermark2);
1128        second.insert(TableId::new(3), table3_watermark.clone());
1129        let result = merge_multiple_new_table_watermarks(vec![first, second]);
1130        let mut expected = HashMap::new();
1131        expected.insert(
1132            TableId::new(1),
1133            TableWatermarks {
1134                watermarks: vec![
1135                    epoch_new_watermark(1, vec![&build_bitmap(0..3), &build_bitmap(4..6)]),
1136                    epoch_new_watermark(2, vec![&build_bitmap(0..3), &build_bitmap(4..6)]),
1137                    epoch_new_watermark(4, vec![&build_bitmap(0..3)]),
1138                    epoch_new_watermark(5, vec![&build_bitmap(4..6)]),
1139                ],
1140                direction: WatermarkDirection::Ascending,
1141                watermark_type: WatermarkSerdeType::PkPrefix,
1142            },
1143        );
1144        expected.insert(TableId::new(2), table2_watermark);
1145        expected.insert(TableId::new(3), table3_watermark);
1146        assert_eq!(result, expected);
1147    }
1148
1149    const COMMITTED_EPOCH: u64 = test_epoch(1);
1150    const EPOCH1: u64 = test_epoch(2);
1151    const EPOCH2: u64 = test_epoch(3);
1152    const TEST_SINGLE_VNODE: VirtualNode = VirtualNode::from_index(1);
1153
1154    fn build_watermark_range(
1155        direction: WatermarkDirection,
1156        (low, high): (Bound<Bytes>, Bound<Bytes>),
1157    ) -> TableKeyRange {
1158        let range = match direction {
1159            WatermarkDirection::Ascending => (low, high),
1160            WatermarkDirection::Descending => (high, low),
1161        };
1162        prefixed_range_with_vnode(range, TEST_SINGLE_VNODE)
1163    }
1164
1165    /// Build and return a watermark index with the following watermarks
1166    /// EPOCH1 bitmap(0, 1, 2, 3) watermark1
1167    /// EPOCH2 bitmap(1, 2, 3, 4) watermark2
1168    fn build_and_test_watermark_index(
1169        direction: WatermarkDirection,
1170        watermark1: Bytes,
1171        watermark2: Bytes,
1172        watermark3: Bytes,
1173    ) -> TableWatermarksIndex {
1174        let mut index = TableWatermarksIndex::new(
1175            direction,
1176            EPOCH1,
1177            vec![VnodeWatermark::new(build_bitmap(0..4), watermark1.clone())],
1178            Some(COMMITTED_EPOCH),
1179            WatermarkSerdeType::PkPrefix,
1180        );
1181        index.add_epoch_watermark(
1182            EPOCH2,
1183            vec![VnodeWatermark::new(build_bitmap(1..5), watermark2.clone())].into(),
1184            direction,
1185        );
1186
1187        assert_eq!(
1188            index.read_watermark(VirtualNode::from_index(0), EPOCH1),
1189            Some(watermark1.clone())
1190        );
1191        assert_eq!(
1192            index.read_watermark(VirtualNode::from_index(1), EPOCH1),
1193            Some(watermark1.clone())
1194        );
1195        assert_eq!(
1196            index.read_watermark(VirtualNode::from_index(4), EPOCH1),
1197            None
1198        );
1199        assert_eq!(
1200            index.read_watermark(VirtualNode::from_index(0), EPOCH2),
1201            Some(watermark1.clone())
1202        );
1203        assert_eq!(
1204            index.read_watermark(VirtualNode::from_index(1), EPOCH2),
1205            Some(watermark2.clone())
1206        );
1207        assert_eq!(
1208            index.read_watermark(VirtualNode::from_index(4), EPOCH2),
1209            Some(watermark2.clone())
1210        );
1211        assert_eq!(
1212            index.latest_watermark(VirtualNode::from_index(0)),
1213            Some(watermark1.clone())
1214        );
1215        assert_eq!(
1216            index.latest_watermark(VirtualNode::from_index(1)),
1217            Some(watermark2.clone())
1218        );
1219        assert_eq!(
1220            index.latest_watermark(VirtualNode::from_index(4)),
1221            Some(watermark2.clone())
1222        );
1223
1224        // watermark is watermark2
1225        let check_watermark_range =
1226            |query_range: (Bound<Bytes>, Bound<Bytes>),
1227             output_range: Option<(Bound<Bytes>, Bound<Bytes>)>| {
1228                let mut range = build_watermark_range(direction, query_range);
1229                index.rewrite_range_with_table_watermark(EPOCH2, &mut range);
1230                if let Some(output_range) = output_range {
1231                    assert_eq!(range, build_watermark_range(direction, output_range));
1232                } else {
1233                    assert!(is_empty_key_range(&range));
1234                }
1235            };
1236
1237        // test read from single vnode and truncate begin key range
1238        check_watermark_range(
1239            (Included(watermark1.clone()), Excluded(watermark3.clone())),
1240            Some((Included(watermark2.clone()), Excluded(watermark3.clone()))),
1241        );
1242
1243        // test read from single vnode and begin key right at watermark
1244        check_watermark_range(
1245            (Included(watermark2.clone()), Excluded(watermark3.clone())),
1246            Some((Included(watermark2.clone()), Excluded(watermark3.clone()))),
1247        );
1248        check_watermark_range(
1249            (Excluded(watermark2.clone()), Excluded(watermark3.clone())),
1250            Some((Excluded(watermark2.clone()), Excluded(watermark3))),
1251        );
1252
1253        // test read from single vnode and end key right at watermark
1254        check_watermark_range(
1255            (Excluded(watermark1.clone()), Excluded(watermark2.clone())),
1256            None,
1257        );
1258        check_watermark_range(
1259            (Excluded(watermark1), Included(watermark2.clone())),
1260            Some((Included(watermark2.clone()), Included(watermark2))),
1261        );
1262
1263        index
1264    }
1265
1266    #[test]
1267    fn test_watermark_index_ascending() {
1268        let watermark1 = Bytes::from_static(b"watermark1");
1269        let watermark2 = Bytes::from_static(b"watermark2");
1270        let watermark3 = Bytes::from_static(b"watermark3");
1271        build_and_test_watermark_index(
1272            WatermarkDirection::Ascending,
1273            watermark1,
1274            watermark2,
1275            watermark3,
1276        );
1277    }
1278
1279    #[test]
1280    fn test_watermark_index_descending() {
1281        let watermark1 = Bytes::from_static(b"watermark254");
1282        let watermark2 = Bytes::from_static(b"watermark253");
1283        let watermark3 = Bytes::from_static(b"watermark252");
1284        build_and_test_watermark_index(
1285            WatermarkDirection::Descending,
1286            watermark1,
1287            watermark2,
1288            watermark3,
1289        );
1290    }
1291
1292    #[test]
1293    fn test_apply_committed_index() {
1294        let watermark1 = Bytes::from_static(b"watermark1");
1295        let watermark2 = Bytes::from_static(b"watermark2");
1296        let watermark3 = Bytes::from_static(b"watermark3");
1297        let mut index = build_and_test_watermark_index(
1298            WatermarkDirection::Ascending,
1299            watermark1.clone(),
1300            watermark2.clone(),
1301            watermark3,
1302        );
1303
1304        let test_table_id = TableId::from(233);
1305
1306        let mut version = HummockVersion::from_rpc_protobuf(&PbHummockVersion {
1307            state_table_info: HashMap::from_iter([(
1308                test_table_id,
1309                StateTableInfo {
1310                    committed_epoch: EPOCH1,
1311                    compaction_group_id: StaticCompactionGroupId::StateDefault,
1312                },
1313            )]),
1314            ..Default::default()
1315        });
1316        version.table_watermarks.insert(
1317            test_table_id,
1318            TableWatermarks {
1319                watermarks: vec![(
1320                    EPOCH1,
1321                    vec![VnodeWatermark {
1322                        watermark: watermark1.clone(),
1323                        vnode_bitmap: build_bitmap(0..VirtualNode::COUNT_FOR_TEST),
1324                    }]
1325                    .into(),
1326                )],
1327                direction: WatermarkDirection::Ascending,
1328                watermark_type: WatermarkSerdeType::PkPrefix,
1329            }
1330            .into(),
1331        );
1332        index.apply_committed_watermarks(
1333            version
1334                .table_watermarks
1335                .get(&test_table_id)
1336                .unwrap()
1337                .clone(),
1338            EPOCH1,
1339        );
1340        assert_eq!(EPOCH1, index.committed_epoch.unwrap());
1341        assert_eq!(EPOCH2, index.latest_epoch);
1342        for vnode in 0..VirtualNode::COUNT_FOR_TEST {
1343            let vnode = VirtualNode::from_index(vnode);
1344            if (1..5).contains(&vnode.to_index()) {
1345                assert_eq!(watermark1, index.read_watermark(vnode, EPOCH1).unwrap());
1346                assert_eq!(watermark2, index.read_watermark(vnode, EPOCH2).unwrap());
1347            } else {
1348                assert_eq!(watermark1, index.read_watermark(vnode, EPOCH1).unwrap());
1349            }
1350        }
1351    }
1352
1353    #[test]
1354    fn test_filter_regress_watermark() {
1355        let watermark1 = Bytes::from_static(b"watermark1");
1356        let watermark2 = Bytes::from_static(b"watermark2");
1357        let watermark3 = Bytes::from_static(b"watermark3");
1358        let index = build_and_test_watermark_index(
1359            WatermarkDirection::Ascending,
1360            watermark1.clone(),
1361            watermark2,
1362            watermark3.clone(),
1363        );
1364
1365        let mut new_watermarks = vec![
1366            // Partial regress
1367            VnodeWatermark {
1368                vnode_bitmap: build_bitmap(0..2),
1369                watermark: watermark1.clone(),
1370            },
1371            // All not regress
1372            VnodeWatermark {
1373                vnode_bitmap: build_bitmap(2..4),
1374                watermark: watermark3.clone(),
1375            },
1376            // All regress
1377            VnodeWatermark {
1378                vnode_bitmap: build_bitmap(4..5),
1379                watermark: watermark1.clone(),
1380            },
1381            // All newly set vnode
1382            VnodeWatermark {
1383                vnode_bitmap: build_bitmap(5..6),
1384                watermark: watermark3.clone(),
1385            },
1386        ];
1387
1388        index.filter_regress_watermarks(&mut new_watermarks);
1389
1390        assert_eq!(
1391            new_watermarks,
1392            vec![
1393                VnodeWatermark {
1394                    vnode_bitmap: build_bitmap(0..1),
1395                    watermark: watermark1,
1396                },
1397                VnodeWatermark {
1398                    vnode_bitmap: build_bitmap(2..4),
1399                    watermark: watermark3.clone(),
1400                },
1401                VnodeWatermark {
1402                    vnode_bitmap: build_bitmap(5..6),
1403                    watermark: watermark3,
1404                },
1405            ]
1406        );
1407    }
1408}
risingwave_hummock_sdk/table_watermark.rs

risingwave_hummock_sdk/
table_watermark.rs
