Skip to main content

risingwave_storage/hummock/compactor/
fast_compactor_runner.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::cmp::Ordering;
16use std::collections::HashSet;
17use std::marker::PhantomData;
18use std::sync::atomic::AtomicU64;
19use std::sync::{Arc, atomic};
20use std::time::Instant;
21
22use await_tree::{InstrumentAwait, SpanExt};
23use bytes::Bytes;
24use itertools::Itertools;
25use risingwave_common::catalog::TableId;
26use risingwave_hummock_sdk::compact_task::CompactTask;
27use risingwave_hummock_sdk::key::FullKey;
28use risingwave_hummock_sdk::key_range::KeyRange;
29use risingwave_hummock_sdk::sstable_info::SstableInfo;
30use risingwave_hummock_sdk::table_stats::TableStats;
31use risingwave_hummock_sdk::{EpochWithGap, LocalSstableInfo, can_concat, compact_task_to_string};
32use risingwave_pb::hummock::PbSstableFilterLayout;
33
34use crate::compaction_catalog_manager::CompactionCatalogAgentRef;
35use crate::hummock::compactor::block_stream::SstableBlockStream;
36use crate::hummock::compactor::compaction_utils::{
37    blocked_xor_filter_key_count_threshold, estimate_output_key_count_for_task,
38};
39use crate::hummock::compactor::task_progress::TaskProgress;
40use crate::hummock::compactor::{
41    CompactionFilter, CompactionStatistics, Compactor, CompactorContext, MultiCompactionFilter,
42    RemoteBuilderFactory, TaskConfig,
43};
44use crate::hummock::iterator::{
45    NonPkPrefixSkipWatermarkState, PkPrefixSkipWatermarkState, SkipWatermarkState,
46    ValueSkipWatermarkState,
47};
48use crate::hummock::multi_builder::{CapacitySplitTableBuilder, TableBuilderFactory};
49use crate::hummock::sstable_store::SstableStoreRef;
50use crate::hummock::value::HummockValue;
51use crate::hummock::{
52    Block, BlockHolder, BlockIterator, BlockMeta, BlockedXor16FilterBuilder, CachePolicy,
53    FilterBuilder, GetObjectId, HummockResult, SstableBuilderOptions,
54    StreamingSstableWriterFactory, TableHolder,
55};
56use crate::monitor::{CompactorMetrics, StoreLocalStatistic};
57
58/// Streams physical SST blocks for raw copy or decoded compaction.
59pub(crate) struct BlockStreamIterator {
60    block_stream: SstableBlockStream,
61    /// When present, this is the decoded block immediately before the stream cursor.
62    /// Otherwise the next block is still eligible for raw copy, or the SST is exhausted.
63    iter: Option<BlockIterator>,
64    task_progress: Arc<TaskProgress>,
65    stats_ptr: Arc<AtomicU64>,
66}
67
68impl BlockStreamIterator {
69    pub(crate) fn new(
70        sstable: TableHolder,
71        task_progress: Arc<TaskProgress>,
72        sstable_store: SstableStoreRef,
73        sstable_info: SstableInfo,
74        max_io_retry_times: usize,
75        stats_ptr: Arc<AtomicU64>,
76    ) -> Self {
77        // Fast compaction streams the physical SST. The executor decides whether each block
78        // can be copied or needs decoding, including table-id pruning.
79        let block_count = sstable.meta.block_metas.len();
80        task_progress.inc_num_pending_read_io();
81        Self {
82            block_stream: SstableBlockStream::new(
83                sstable,
84                0..block_count,
85                sstable_info,
86                sstable_store,
87                max_io_retry_times,
88            ),
89            iter: None,
90            task_progress,
91            stats_ptr,
92        }
93    }
94
95    pub(crate) async fn download_next_block(&mut self) -> HummockResult<Option<(Bytes, usize)>> {
96        let now = Instant::now();
97        let _time_stat = scopeguard::guard(self.stats_ptr.clone(), |stats_ptr: Arc<AtomicU64>| {
98            let add = (now.elapsed().as_secs_f64() * 1000.0).ceil();
99            stats_ptr.fetch_add(add as u64, atomic::Ordering::Relaxed);
100        });
101        let block = self.block_stream.next_block().await?;
102        if block.is_none() {
103            self.iter = None;
104        }
105        Ok(block)
106    }
107
108    /// Materialize writer metadata only when copying the most recently downloaded block.
109    /// Decoded compaction uses the uncompressed size returned by `download_next_block` instead.
110    fn current_block_raw_metadata(&self) -> (Vec<u8>, BlockMeta) {
111        let block_index = self.block_stream.next_block_index() - 1;
112        let sstable = &self.block_stream.sstable;
113        (
114            sstable.filter_reader.get_block_raw_filter(block_index),
115            sstable.meta.block_metas[block_index].clone(),
116        )
117    }
118
119    fn has_decoded_block(&self) -> bool {
120        self.iter.is_some()
121    }
122
123    async fn ensure_block_iter(&mut self) -> HummockResult<()> {
124        if self.iter.is_none() {
125            let (buf, uncompressed_size) = self.download_next_block().await?.unwrap();
126            self.init_block_iter(buf, uncompressed_size)?;
127        }
128        Ok(())
129    }
130
131    pub(crate) fn init_block_iter(
132        &mut self,
133        buf: Bytes,
134        uncompressed_capacity: usize,
135    ) -> HummockResult<()> {
136        let block = Block::decode(buf, uncompressed_capacity)?;
137        let mut iter = BlockIterator::new(BlockHolder::from_owned_block(Box::new(block)));
138        iter.seek_to_first();
139        self.iter = Some(iter);
140        Ok(())
141    }
142
143    fn next_block_smallest(&self) -> &[u8] {
144        self.block_stream.sstable.meta.block_metas[self.block_stream.next_block_index()]
145            .smallest_key
146            .as_ref()
147    }
148
149    /// Upper bound used to decide whether the unread block precedes the other input.
150    /// The next block's lower bound is exclusive; the SST's largest key is inclusive.
151    /// Callers must compare user keys strictly to avoid copying across versions of one key.
152    fn next_block_upper_bound(&self) -> &[u8] {
153        let sstable = &self.block_stream.sstable;
154        let next_block_index = self.block_stream.next_block_index();
155        if next_block_index + 1 < sstable.meta.block_metas.len() {
156            sstable.meta.block_metas[next_block_index + 1]
157                .smallest_key
158                .as_ref()
159        } else {
160            sstable.meta.largest_key.as_ref()
161        }
162    }
163
164    /// Builder boundary for the just-downloaded block, not necessarily its last stored key.
165    fn current_block_largest(&self) -> Vec<u8> {
166        let sstable = &self.block_stream.sstable;
167        let next_block_index = self.block_stream.next_block_index();
168        if self.block_stream.has_next_block() {
169            let mut largest_key = FullKey::decode(
170                sstable.meta.block_metas[next_block_index]
171                    .smallest_key
172                    .as_ref(),
173            );
174            // do not include this key because it is the smallest key of next block.
175            largest_key.epoch_with_gap = EpochWithGap::new_max_epoch();
176            largest_key.encode()
177        } else {
178            sstable.meta.largest_key.clone()
179        }
180    }
181
182    /// Current decoded key, or a possibly shortened lower bound of the unread block.
183    fn key(&self) -> FullKey<&[u8]> {
184        match self.iter.as_ref() {
185            Some(iter) => iter.key(),
186            None => FullKey::decode(self.next_block_smallest()),
187        }
188    }
189
190    pub(crate) fn is_valid(&self) -> bool {
191        self.iter.is_some() || self.block_stream.has_next_block()
192    }
193
194    #[cfg(test)]
195    #[cfg(feature = "failpoints")]
196    pub(crate) fn iter_mut(&mut self) -> &mut BlockIterator {
197        self.iter.as_mut().unwrap()
198    }
199}
200
201impl Drop for BlockStreamIterator {
202    fn drop(&mut self) {
203        self.task_progress.dec_num_pending_read_io();
204    }
205}
206
207/// Iterates over the KV-pairs of a given list of SSTs. The key-ranges of these SSTs are assumed to
208/// be consecutive and non-overlapping.
209struct ConcatSstableIterator {
210    /// The iterator of the current table.
211    sstable_iter: Option<BlockStreamIterator>,
212
213    /// Current table index.
214    cur_idx: usize,
215
216    /// All non-overlapping tables.
217    sstables: Vec<SstableInfo>,
218
219    sstable_store: SstableStoreRef,
220
221    stats: StoreLocalStatistic,
222    task_progress: Arc<TaskProgress>,
223
224    max_io_retry_times: usize,
225}
226
227impl ConcatSstableIterator {
228    /// The SSTs must have non-overlapping key ranges in ascending order.
229    fn new(
230        sst_infos: Vec<SstableInfo>,
231        sstable_store: SstableStoreRef,
232        task_progress: Arc<TaskProgress>,
233        max_io_retry_times: usize,
234    ) -> Self {
235        Self {
236            sstable_iter: None,
237            cur_idx: 0,
238            sstables: sst_infos,
239            sstable_store,
240            task_progress,
241            stats: StoreLocalStatistic::default(),
242            max_io_retry_times,
243        }
244    }
245
246    async fn rewind(&mut self) -> HummockResult<()> {
247        self.seek_idx(0).await
248    }
249
250    async fn next_sstable(&mut self) -> HummockResult<()> {
251        self.seek_idx(self.cur_idx + 1).await
252    }
253
254    fn current_sstable(&mut self) -> &mut BlockStreamIterator {
255        self.sstable_iter.as_mut().unwrap()
256    }
257
258    fn is_valid(&self) -> bool {
259        self.cur_idx < self.sstables.len()
260    }
261
262    /// Replaces the current iterator with one at the start of the specified SST.
263    async fn seek_idx(&mut self, idx: usize) -> HummockResult<()> {
264        self.sstable_iter.take();
265        self.cur_idx = idx;
266        if self.cur_idx < self.sstables.len() {
267            let sstable_info = &self.sstables[self.cur_idx];
268            let sstable = self
269                .sstable_store
270                .sstable(sstable_info, &mut self.stats)
271                .instrument_await("stream_iter_sstable".verbose())
272                .await?;
273            let sstable_iter = BlockStreamIterator::new(
274                sstable,
275                self.task_progress.clone(),
276                self.sstable_store.clone(),
277                sstable_info.clone(),
278                self.max_io_retry_times,
279                self.stats.remote_io_time.clone(),
280            );
281            self.sstable_iter = Some(sstable_iter);
282        }
283        Ok(())
284    }
285}
286
287pub struct CompactorRunner<
288    B: FilterBuilder = BlockedXor16FilterBuilder,
289    C: CompactionFilter = MultiCompactionFilter,
290> {
291    left: Box<ConcatSstableIterator>,
292    right: Box<ConcatSstableIterator>,
293    task_id: u64,
294    executor: CompactTaskExecutor<RemoteBuilderFactory<StreamingSstableWriterFactory, B>, C>,
295    metrics: Arc<CompactorMetrics>,
296}
297
298impl<B: FilterBuilder, C: CompactionFilter> CompactorRunner<B, C> {
299    pub fn new(
300        context: CompactorContext,
301        task: CompactTask,
302        compaction_catalog_agent_ref: CompactionCatalogAgentRef,
303        object_id_getter: Arc<dyn GetObjectId>,
304        task_progress: Arc<TaskProgress>,
305        compaction_filter: C,
306    ) -> Self {
307        let mut options: SstableBuilderOptions = context.storage_opts.as_ref().into();
308        options.compression_algorithm = task.compression_algorithm.into();
309        options.capacity = task.target_file_size as usize;
310        let estimated_output_key_count =
311            estimate_output_key_count_for_task(&task, options.capacity);
312        options.estimated_output_key_count = Some(estimated_output_key_count);
313        options.filter_hash_prealloc_key_count_cap =
314            blocked_xor_filter_key_count_threshold(task.blocked_xor_filter_kv_count_threshold);
315        // Disable vnode key-range hints for fast compaction path by default.
316        options.max_vnode_key_range_bytes = None;
317        let get_id_time = Arc::new(AtomicU64::new(0));
318
319        let key_range = KeyRange::inf();
320        let read_table_ids = HashSet::from_iter(task.get_table_ids_from_input_ssts());
321
322        let task_config = TaskConfig {
323            key_range,
324            cache_policy: CachePolicy::NotFill,
325            gc_delete_keys: task.gc_delete_keys,
326            retain_multiple_version: false,
327            table_vnode_partition: task.table_vnode_partition.clone(),
328            sstable_filter_layout: PbSstableFilterLayout::Blocked,
329            sstable_filter_type: task.sstable_filter_type,
330            table_schemas: Default::default(),
331            disable_drop_column_optimization: false,
332        };
333        let factory = StreamingSstableWriterFactory::new(context.sstable_store.clone());
334
335        let builder_factory = RemoteBuilderFactory::<_, B> {
336            object_id_getter,
337            limiter: context.memory_limiter.clone(),
338            options,
339            policy: task_config.cache_policy,
340            remote_rpc_cost: get_id_time,
341            compaction_catalog_agent_ref: compaction_catalog_agent_ref.clone(),
342            sstable_writer_factory: factory,
343            _phantom: PhantomData,
344        };
345        let sst_builder = CapacitySplitTableBuilder::new(
346            builder_factory,
347            context.compactor_metrics.clone(),
348            Some(task_progress.clone()),
349            task_config.table_vnode_partition.clone(),
350            context
351                .storage_opts
352                .compactor_concurrent_uploading_sst_count,
353            compaction_catalog_agent_ref.clone(),
354        );
355        assert_eq!(
356            task.input_ssts.len(),
357            2,
358            "TaskId {} target_level {:?} task {:?}",
359            task.task_id,
360            task.target_level,
361            compact_task_to_string(&task)
362        );
363        let left_ssts = task.input_ssts[0]
364            .read_sstable_infos()
365            .cloned()
366            .collect_vec();
367        let right_ssts = task.input_ssts[1]
368            .read_sstable_infos()
369            .cloned()
370            .collect_vec();
371        assert!(
372            left_ssts
373                .iter()
374                .chain(right_ssts.iter())
375                .all(|sst| sst.filter_layout == PbSstableFilterLayout::Blocked),
376            "fast compaction requires blocked-filter SSTs: {}",
377            compact_task_to_string(&task)
378        );
379        let left = Box::new(ConcatSstableIterator::new(
380            left_ssts,
381            context.sstable_store.clone(),
382            task_progress.clone(),
383            context.storage_opts.compactor_iter_max_io_retry_times,
384        ));
385        let right = Box::new(ConcatSstableIterator::new(
386            right_ssts,
387            context.sstable_store,
388            task_progress.clone(),
389            context.storage_opts.compactor_iter_max_io_retry_times,
390        ));
391
392        // Can not consume the watermarks because the watermarks may be used by `check_compact_result`.
393        let pk_prefix_state = PkPrefixSkipWatermarkState::from_safe_epoch_watermarks(
394            task.pk_prefix_table_watermarks.clone(),
395        );
396        let non_pk_prefix_state = NonPkPrefixSkipWatermarkState::from_safe_epoch_watermarks(
397            task.non_pk_prefix_table_watermarks.clone(),
398            compaction_catalog_agent_ref.clone(),
399        );
400        let value_skip_watermark_state = ValueSkipWatermarkState::from_safe_epoch_watermarks(
401            task.value_table_watermarks.clone(),
402            compaction_catalog_agent_ref,
403        );
404
405        Self {
406            executor: CompactTaskExecutor::new(
407                sst_builder,
408                task_config,
409                task_progress,
410                pk_prefix_state,
411                non_pk_prefix_state,
412                value_skip_watermark_state,
413                compaction_filter,
414                read_table_ids,
415            ),
416            left,
417            right,
418            task_id: task.task_id,
419            metrics: context.compactor_metrics,
420        }
421    }
422
423    pub async fn run(mut self) -> HummockResult<(Vec<LocalSstableInfo>, CompactionStatistics)> {
424        self.left.rewind().await?;
425        self.right.rewind().await?;
426        self.merge_inputs().await?;
427        self.drain_remaining().await?;
428        let mut total_read_bytes = 0;
429        for sst in &self.left.sstables {
430            total_read_bytes += sst.sst_size;
431        }
432        for sst in &self.right.sstables {
433            total_read_bytes += sst.sst_size;
434        }
435        self.metrics
436            .compact_fast_runner_bytes
437            .inc_by(self.executor.skip_raw_block_size);
438        tracing::info!(
439            "OPTIMIZATION: skip {} blocks for task-{}, optimize {}% data compression",
440            self.executor.skip_raw_block_count,
441            self.task_id,
442            self.executor.skip_raw_block_size * 100 / total_read_bytes,
443        );
444
445        let statistic = self.executor.take_statistics();
446        let output_ssts = self.executor.builder.finish().await?;
447        Compactor::report_progress(
448            self.metrics.clone(),
449            Some(self.executor.task_progress.clone()),
450            &output_ssts,
451            false,
452        );
453        let sst_infos = output_ssts
454            .iter()
455            .map(|sst| sst.sst_info.clone())
456            .collect_vec();
457        assert!(can_concat(&sst_infos));
458        Ok((output_ssts, statistic))
459    }
460
461    /// Merge while both inputs have data, retaining partially consumed decoded blocks.
462    async fn merge_inputs(&mut self) -> HummockResult<()> {
463        while self.left.is_valid() && self.right.is_valid() {
464            let mut ret = self
465                .left
466                .current_sstable()
467                .key()
468                .cmp(&self.right.current_sstable().key());
469            if ret == Ordering::Equal {
470                // Equal block lower bounds need not represent duplicate keys. Compare the
471                // actual keys before choosing an input or asserting that the ranges overlap.
472                self.left.current_sstable().ensure_block_iter().await?;
473                self.right.current_sstable().ensure_block_iter().await?;
474                ret = self
475                    .left
476                    .current_sstable()
477                    .key()
478                    .cmp(&self.right.current_sstable().key());
479            }
480            let (first, second) = if ret == Ordering::Less {
481                (&mut self.left, &mut self.right)
482            } else {
483                (&mut self.right, &mut self.left)
484            };
485            assert!(
486                ret != Ordering::Equal,
487                "sst range overlap equal_key {:?}",
488                self.left.current_sstable().key()
489            );
490            if !first.current_sstable().has_decoded_block() {
491                let right_key = second.current_sstable().key();
492                while first.current_sstable().is_valid() && !self.executor.builder.need_flush() {
493                    let full_key =
494                        FullKey::decode(first.current_sstable().next_block_upper_bound());
495                    // Equality may hide more versions of the other input's user key.
496                    // Only a strictly smaller upper bound is safe for raw copy.
497                    if full_key.user_key.ge(&right_key.user_key) {
498                        break;
499                    }
500                    let smallest_key =
501                        FullKey::decode(first.current_sstable().next_block_smallest());
502                    if !self.executor.shall_copy_raw_block(&smallest_key) {
503                        break;
504                    }
505                    self.executor
506                        .append_raw_block(first.current_sstable())
507                        .await?;
508                }
509                if !first.current_sstable().is_valid() {
510                    first.next_sstable().await?;
511                    continue;
512                }
513            }
514
515            let target_key = second.current_sstable().key();
516            self.executor
517                .compact_block(first.current_sstable(), Some(target_key))
518                .await?;
519            if !first.current_sstable().is_valid() {
520                first.next_sstable().await?;
521            }
522        }
523        Ok(())
524    }
525
526    /// Finish the remaining input after the two-way merge, including any decoded suffix.
527    async fn drain_remaining(&mut self) -> HummockResult<()> {
528        let rest_data = if !self.left.is_valid() {
529            &mut self.right
530        } else {
531            &mut self.left
532        };
533        // The stream cursor may already be at EOF while the current decoded block still
534        // has rows. Consume those rows before asking for another block or SST.
535        if rest_data.is_valid() && rest_data.current_sstable().has_decoded_block() {
536            self.executor
537                .compact_block(rest_data.current_sstable(), None)
538                .await?;
539        }
540
541        while rest_data.is_valid() {
542            let sstable_iter = rest_data.current_sstable();
543            while sstable_iter.is_valid() {
544                let smallest_key = FullKey::decode(sstable_iter.next_block_smallest());
545                if self.executor.builder.need_flush()
546                    || !self.executor.shall_copy_raw_block(&smallest_key)
547                {
548                    self.executor.compact_block(sstable_iter, None).await?;
549                } else {
550                    self.executor.append_raw_block(sstable_iter).await?;
551                }
552            }
553            rest_data.next_sstable().await?;
554        }
555        Ok(())
556    }
557}
558
559struct CompactTaskExecutor<F: TableBuilderFactory, C: CompactionFilter> {
560    last_key: FullKey<Vec<u8>>,
561    compaction_statistics: CompactionStatistics,
562    last_table_id: Option<TableId>,
563    last_table_stats: TableStats,
564    builder: CapacitySplitTableBuilder<F>,
565    task_config: TaskConfig,
566    task_progress: Arc<TaskProgress>,
567    pk_prefix_skip_watermark_state: PkPrefixSkipWatermarkState,
568    last_key_is_delete: bool,
569    progress_key_num: u32,
570    skip_raw_block_count: u64,
571    skip_raw_block_size: u64,
572    non_pk_prefix_skip_watermark_state: NonPkPrefixSkipWatermarkState,
573    value_skip_watermark_state: ValueSkipWatermarkState,
574    compaction_filter: C,
575    read_table_ids: HashSet<TableId>,
576}
577
578impl<F: TableBuilderFactory, C: CompactionFilter> CompactTaskExecutor<F, C> {
579    fn new(
580        builder: CapacitySplitTableBuilder<F>,
581        task_config: TaskConfig,
582        task_progress: Arc<TaskProgress>,
583        pk_prefix_skip_watermark_state: PkPrefixSkipWatermarkState,
584        non_pk_prefix_skip_watermark_state: NonPkPrefixSkipWatermarkState,
585        value_skip_watermark_state: ValueSkipWatermarkState,
586        compaction_filter: C,
587        read_table_ids: HashSet<TableId>,
588    ) -> Self {
589        Self {
590            builder,
591            task_config,
592            last_key: FullKey::default(),
593            last_key_is_delete: false,
594            compaction_statistics: CompactionStatistics::default(),
595            last_table_id: None,
596            last_table_stats: TableStats::default(),
597            task_progress,
598            pk_prefix_skip_watermark_state,
599            progress_key_num: 0,
600            skip_raw_block_count: 0,
601            skip_raw_block_size: 0,
602            non_pk_prefix_skip_watermark_state,
603            value_skip_watermark_state,
604            compaction_filter,
605            read_table_ids,
606        }
607    }
608
609    fn take_statistics(&mut self) -> CompactionStatistics {
610        if let Some(last_table_id) = self.last_table_id.take() {
611            self.compaction_statistics
612                .delta_drop_stat
613                .insert(last_table_id, std::mem::take(&mut self.last_table_stats));
614        }
615        std::mem::take(&mut self.compaction_statistics)
616    }
617
618    /// Read and append the next block after the caller has checked raw-copy eligibility.
619    /// The builder may coalesce it; only an actual raw copy counts as skipped work.
620    async fn append_raw_block(
621        &mut self,
622        sstable_iter: &mut BlockStreamIterator,
623    ) -> HummockResult<()> {
624        // Preserve the lower bound before advancing; metadata then refers to the read block.
625        let smallest_key = FullKey::decode(sstable_iter.next_block_smallest()).to_vec();
626        let (block, _) = sstable_iter.download_next_block().await?.unwrap();
627        let (filter_data, meta) = sstable_iter.current_block_raw_metadata();
628        let largest_key = sstable_iter.current_block_largest();
629        let key_count = meta.total_key_count;
630        if let Some(block_len) = self
631            .builder
632            .add_raw_block(block, filter_data, smallest_key, largest_key, meta)
633            .await?
634        {
635            self.skip_raw_block_count += 1;
636            self.skip_raw_block_size += block_len as u64;
637        }
638        self.may_report_process_key(key_count);
639        // Decoded-key state precedes this block and must not affect subsequent rows.
640        // Raw-copy eligibility has already ruled out a deleted key crossing into it.
641        if !self.last_key.is_empty() {
642            self.last_key = FullKey::default();
643        }
644        self.last_key_is_delete = false;
645        Ok(())
646    }
647
648    fn reset_watermark(&mut self) {
649        self.pk_prefix_skip_watermark_state.reset_watermark();
650        self.non_pk_prefix_skip_watermark_state.reset_watermark();
651        self.value_skip_watermark_state.reset_watermark();
652    }
653
654    #[inline(always)]
655    fn should_skip_block(&self, table_id: TableId) -> bool {
656        !self.read_table_ids.contains(&table_id)
657    }
658
659    #[inline(always)]
660    fn may_report_process_key(&mut self, key_count: u32) {
661        const PROGRESS_KEY_INTERVAL: u32 = 100;
662        self.progress_key_num += key_count;
663        if self.progress_key_num > PROGRESS_KEY_INTERVAL {
664            self.task_progress
665                .inc_progress_key(self.progress_key_num as u64);
666            self.progress_key_num = 0;
667        }
668    }
669
670    /// Decode on demand and compact the current block up to `target_key`. With no target,
671    /// consume the whole block. Keep a partially consumed block for the next merge step.
672    async fn compact_block(
673        &mut self,
674        sstable_iter: &mut BlockStreamIterator,
675        target_key: Option<FullKey<&[u8]>>,
676    ) -> HummockResult<()> {
677        sstable_iter.ensure_block_iter().await?;
678        let consume_whole_block = target_key.is_none();
679        let target_key = target_key.unwrap_or_else(|| {
680            FullKey::decode(&sstable_iter.block_stream.sstable.meta.largest_key)
681        });
682        let iter = sstable_iter.iter.as_mut().unwrap();
683        self.reset_watermark();
684        self.run(iter, target_key).await?;
685        if consume_whole_block {
686            assert!(
687                !iter.is_valid(),
688                "iter should not be valid key {:?}",
689                iter.key()
690            );
691        }
692        if !iter.is_valid() {
693            sstable_iter.iter = None;
694        }
695        Ok(())
696    }
697
698    async fn run(
699        &mut self,
700        iter: &mut BlockIterator,
701        target_key: FullKey<&[u8]>,
702    ) -> HummockResult<()> {
703        if self.should_skip_block(iter.table_id()) {
704            iter.finish_block();
705            return Ok(());
706        }
707
708        while iter.is_valid() && iter.key().le(&target_key) {
709            let is_new_user_key =
710                !self.last_key.is_empty() && iter.key().user_key != self.last_key.user_key.as_ref();
711            self.compaction_statistics.iter_total_key_counts += 1;
712            self.may_report_process_key(1);
713
714            let mut drop = false;
715            let value = HummockValue::from_slice(iter.value()).unwrap();
716            let is_first_or_new_user_key = is_new_user_key || self.last_key.is_empty();
717            if is_first_or_new_user_key {
718                self.last_key.set(iter.key());
719                self.last_key_is_delete = false;
720            }
721
722            // See note in `compactor_runner.rs`.
723            if !self.task_config.retain_multiple_version
724                && self.task_config.gc_delete_keys
725                && value.is_delete()
726            {
727                drop = true;
728                self.last_key_is_delete = true;
729            } else if !self.task_config.retain_multiple_version && !is_first_or_new_user_key {
730                drop = true;
731            }
732
733            if !drop && self.compaction_filter.should_delete(iter.key()) {
734                drop = true;
735            }
736
737            if !drop && self.watermark_should_delete(&iter.key(), value) {
738                drop = true;
739                self.last_key_is_delete = true;
740            }
741
742            if self.last_table_id != Some(self.last_key.user_key.table_id) {
743                if let Some(last_table_id) = self.last_table_id.take() {
744                    self.compaction_statistics
745                        .delta_drop_stat
746                        .insert(last_table_id, std::mem::take(&mut self.last_table_stats));
747                }
748                self.last_table_id = Some(self.last_key.user_key.table_id);
749            }
750
751            if drop {
752                self.compaction_statistics.iter_drop_key_counts += 1;
753
754                self.last_table_stats.total_key_count -= 1;
755                self.last_table_stats.total_key_size -= self.last_key.encoded_len() as i64;
756                self.last_table_stats.total_value_size -= value.encoded_len() as i64;
757                iter.next();
758                continue;
759            }
760            self.builder
761                .add_full_key(iter.key(), value, is_new_user_key)
762                .instrument_await("fast_add_full_key".verbose())
763                .await?;
764            iter.next();
765        }
766        Ok(())
767    }
768
769    fn shall_copy_raw_block(&mut self, smallest_key: &FullKey<&[u8]>) -> bool {
770        if self.should_skip_block(smallest_key.user_key.table_id) {
771            // If the table id of smallest key is not in read_table_ids, we can not copy the raw block.
772            return false;
773        }
774
775        if self.last_key_is_delete && self.last_key.user_key.as_ref().eq(&smallest_key.user_key) {
776            // If the last key is delete tombstone, we can not append the origin block
777            // because it would cause a deleted key could be see by user again.
778            return false;
779        }
780
781        if self.watermark_may_delete(smallest_key) {
782            return false;
783        }
784
785        // Check compaction filter
786        if self.compaction_filter.should_delete(*smallest_key) {
787            return false;
788        }
789
790        true
791    }
792
793    fn watermark_may_delete(&mut self, key: &FullKey<&[u8]>) -> bool {
794        // Correctness requires the assumption that these PkPrefixSkipWatermarkState and NonPkPrefixSkipWatermarkState never use the `unused_put`.
795        let pk_prefix_has_watermark = self.pk_prefix_skip_watermark_state.has_watermark();
796        let non_pk_prefix_has_watermark = self.non_pk_prefix_skip_watermark_state.has_watermark();
797        if pk_prefix_has_watermark || non_pk_prefix_has_watermark {
798            let unused = vec![];
799            let unused_put = HummockValue::Put(unused.as_slice());
800            if (pk_prefix_has_watermark
801                && self
802                    .pk_prefix_skip_watermark_state
803                    .should_delete(key, unused_put))
804                || (non_pk_prefix_has_watermark
805                    && self
806                        .non_pk_prefix_skip_watermark_state
807                        .should_delete(key, unused_put))
808            {
809                return true;
810            }
811        }
812        self.value_skip_watermark_state.has_watermark()
813            && self.value_skip_watermark_state.may_delete(key)
814    }
815
816    fn watermark_should_delete(
817        &mut self,
818        key: &FullKey<&[u8]>,
819        value: HummockValue<&[u8]>,
820    ) -> bool {
821        (self.pk_prefix_skip_watermark_state.has_watermark()
822            && self
823                .pk_prefix_skip_watermark_state
824                .should_delete(key, value))
825            || (self.non_pk_prefix_skip_watermark_state.has_watermark()
826                && self
827                    .non_pk_prefix_skip_watermark_state
828                    .should_delete(key, value))
829            || (self.value_skip_watermark_state.has_watermark()
830                && self.value_skip_watermark_state.should_delete(key, value))
831    }
832}
833
834#[cfg(test)]
835mod tests {
836    use std::collections::{HashMap, VecDeque};
837    use std::sync::Arc;
838
839    use risingwave_common::catalog::TableId;
840    use risingwave_common::hash::VirtualNode;
841    use risingwave_common::util::epoch::test_epoch;
842    use risingwave_hummock_sdk::compact_task::CompactTask;
843    use risingwave_hummock_sdk::key::FullKey;
844    use risingwave_hummock_sdk::level::InputLevel;
845    use risingwave_pb::hummock::compact_task::TaskType;
846    use risingwave_pb::hummock::{LevelType, PbSstableFilterLayout, PbSstableFilterType};
847
848    use super::CompactorRunner;
849    use crate::compaction_catalog_manager::CompactionCatalogAgent;
850    use crate::hummock::compactor::compaction_utils::optimize_by_copy_block;
851    use crate::hummock::compactor::task_progress::TaskProgress;
852    use crate::hummock::compactor::{CompactorContext, MultiCompactionFilter};
853    use crate::hummock::iterator::test_utils::mock_sstable_store;
854    use crate::hummock::test_utils::{
855        default_builder_opt_for_test, default_opts_for_test, gen_test_sstable_impl, test_value_of,
856    };
857    use crate::hummock::value::HummockValue;
858    use crate::hummock::{
859        BlockedXor16FilterBuilder, CachePolicy, SharedComapctorObjectIdManager, Xor16FilterBuilder,
860    };
861    use crate::monitor::CompactorMetrics;
862
863    fn test_key(table_id: u32, idx: usize) -> FullKey<Vec<u8>> {
864        let mut table_key = VirtualNode::ZERO.to_be_bytes().to_vec();
865        table_key.extend_from_slice(format!("key_test_{idx:05}").as_bytes());
866        FullKey::for_test(TableId::new(table_id), table_key, test_epoch(1))
867    }
868
869    #[tokio::test]
870    async fn test_fast_compact_skips_empty_table_id_sst() {
871        let sstable_store = mock_sstable_store().await;
872        let table_id_to_vnode = HashMap::from([
873            (1, VirtualNode::COUNT_FOR_TEST),
874            (2, VirtualNode::COUNT_FOR_TEST),
875        ]);
876        let table_id_to_watermark_serde = HashMap::from([(1, None), (2, None)]);
877
878        let mut dropped_only_sst = gen_test_sstable_impl::<_, Xor16FilterBuilder>(
879            default_builder_opt_for_test(),
880            1,
881            (0..2).map(|idx| (test_key(1, idx), HummockValue::put(test_value_of(idx)))),
882            sstable_store.clone(),
883            CachePolicy::NotFill,
884            table_id_to_vnode.clone(),
885            table_id_to_watermark_serde.clone(),
886        )
887        .await;
888        assert_eq!(dropped_only_sst.filter_layout, PbSstableFilterLayout::Plain);
889        let mut inner = dropped_only_sst.get_inner();
890        inner.table_ids.clear();
891        dropped_only_sst.set_inner(inner);
892
893        let live_left_sst = gen_test_sstable_impl::<_, BlockedXor16FilterBuilder>(
894            default_builder_opt_for_test(),
895            2,
896            (0..2).map(|idx| (test_key(2, idx), HummockValue::put(test_value_of(idx)))),
897            sstable_store.clone(),
898            CachePolicy::NotFill,
899            table_id_to_vnode.clone(),
900            table_id_to_watermark_serde.clone(),
901        )
902        .await;
903        let live_right_sst = gen_test_sstable_impl::<_, BlockedXor16FilterBuilder>(
904            default_builder_opt_for_test(),
905            3,
906            (2..4).map(|idx| (test_key(2, idx), HummockValue::put(test_value_of(idx)))),
907            sstable_store.clone(),
908            CachePolicy::NotFill,
909            table_id_to_vnode,
910            table_id_to_watermark_serde,
911        )
912        .await;
913
914        let mut storage_opts = default_opts_for_test();
915        storage_opts.enable_fast_compaction = true;
916        storage_opts.compactor_fast_max_compact_task_size = u64::MAX;
917        storage_opts.compactor_fast_max_compact_delete_ratio = 100;
918        let context = CompactorContext::new_local_compact_context(
919            Arc::new(storage_opts),
920            sstable_store,
921            Arc::new(CompactorMetrics::unused()),
922            None,
923        );
924
925        let task = CompactTask {
926            input_ssts: vec![
927                InputLevel {
928                    level_idx: 1,
929                    level_type: LevelType::Nonoverlapping,
930                    table_infos: vec![dropped_only_sst, live_left_sst],
931                },
932                InputLevel {
933                    level_idx: 2,
934                    level_type: LevelType::Nonoverlapping,
935                    table_infos: vec![live_right_sst],
936                },
937            ],
938            task_id: 42,
939            target_level: 2,
940            existing_table_ids: vec![TableId::new(2)],
941            target_file_size: 1 << 20,
942            task_type: TaskType::Dynamic,
943            blocked_xor_filter_kv_count_threshold: Some(0),
944            sstable_filter_type: PbSstableFilterType::SstableFilterXor16,
945            ..Default::default()
946        };
947
948        assert_eq!(task.input_ssts[0].read_sstable_infos().count(), 1);
949        assert!(optimize_by_copy_block(&task, &context));
950
951        let runner = CompactorRunner::<BlockedXor16FilterBuilder, _>::new(
952            context,
953            task,
954            CompactionCatalogAgent::for_test(vec![1, 2]),
955            SharedComapctorObjectIdManager::for_test(VecDeque::from([100])),
956            Arc::new(TaskProgress::default()),
957            MultiCompactionFilter::default(),
958        );
959        runner.run().await.unwrap();
960    }
961}