risingwave_stream/executor/top_n/
group_top_n_appendonly.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::HashMap;
16
17use risingwave_common::array::Op;
18use risingwave_common::hash::HashKey;
19use risingwave_common::row::{RowDeserializer, RowExt};
20use risingwave_common::util::epoch::EpochPair;
21use risingwave_common::util::iter_util::ZipEqDebug;
22use risingwave_common::util::sort_util::{ColumnOrder, topn_watermark_forwardable_order_key};
23
24use super::group_top_n::GroupTopNCache;
25use super::top_n_cache::AppendOnlyTopNCacheTrait;
26use super::utils::*;
27use super::{ManagedTopNState, TopNCache};
28use crate::common::metrics::MetricsInfo;
29use crate::common::table::state_table::StateTablePostCommit;
30use crate::executor::monitor::GroupTopNMetrics;
31use crate::executor::prelude::*;
32use crate::executor::top_n::top_n_cache::TopNStaging;
33
34/// If the input is append-only, `AppendOnlyGroupTopNExecutor` does not need
35/// to keep all the rows seen. As long as a record
36/// is no longer in the result set, it can be deleted.
37pub type AppendOnlyGroupTopNExecutor<K, S, const WITH_TIES: bool> =
38    TopNExecutorWrapper<InnerAppendOnlyGroupTopNExecutor<K, S, WITH_TIES>>;
39
40impl<K: HashKey, S: StateStore, const WITH_TIES: bool>
41    AppendOnlyGroupTopNExecutor<K, S, WITH_TIES>
42{
43    #[expect(clippy::too_many_arguments)]
44    pub fn new(
45        input: Executor,
46        ctx: ActorContextRef,
47        schema: Schema,
48        storage_key: Vec<ColumnOrder>,
49        offset_and_limit: (usize, usize),
50        order_by: Vec<ColumnOrder>,
51        group_by: Vec<usize>,
52        state_table: StateTable<S>,
53        watermark_epoch: AtomicU64Ref,
54    ) -> StreamResult<Self> {
55        let inner = InnerAppendOnlyGroupTopNExecutor::new(
56            schema,
57            storage_key,
58            offset_and_limit,
59            order_by,
60            group_by,
61            state_table,
62            watermark_epoch,
63            &ctx,
64        )?;
65        Ok(TopNExecutorWrapper { input, ctx, inner })
66    }
67}
68
69pub struct InnerAppendOnlyGroupTopNExecutor<K: HashKey, S: StateStore, const WITH_TIES: bool> {
70    schema: Schema,
71
72    /// `LIMIT XXX`. None means no limit.
73    limit: usize,
74
75    /// `OFFSET XXX`. `0` means no offset.
76    offset: usize,
77
78    /// The storage key indices of the `AppendOnlyGroupTopNExecutor`
79    storage_key_indices: Vec<usize>,
80
81    managed_state: ManagedTopNState<S>,
82
83    /// which column we used to group the data.
84    group_by: Vec<usize>,
85
86    /// The `ORDER BY` column whose watermarks can be forwarded, if any.
87    watermark_order_key: Option<usize>,
88
89    /// group key -> cache for this group
90    caches: GroupTopNCache<K, WITH_TIES>,
91
92    /// Used for serializing pk into `CacheKey`.
93    cache_key_serde: CacheKeySerde,
94
95    /// Minimum cache capacity per group from config
96    topn_cache_min_capacity: usize,
97
98    metrics: GroupTopNMetrics,
99}
100
101impl<K: HashKey, S: StateStore, const WITH_TIES: bool>
102    InnerAppendOnlyGroupTopNExecutor<K, S, WITH_TIES>
103{
104    #[expect(clippy::too_many_arguments)]
105    pub fn new(
106        schema: Schema,
107        storage_key: Vec<ColumnOrder>,
108        offset_and_limit: (usize, usize),
109        order_by: Vec<ColumnOrder>,
110        group_by: Vec<usize>,
111        state_table: StateTable<S>,
112        watermark_epoch: AtomicU64Ref,
113        ctx: &ActorContext,
114    ) -> StreamResult<Self> {
115        let metrics_info = MetricsInfo::new(
116            ctx.streaming_metrics.clone(),
117            state_table.table_id(),
118            ctx.id,
119            "AppendOnlyGroupTopN",
120        );
121        let metrics = ctx.streaming_metrics.new_append_only_group_top_n_metrics(
122            state_table.table_id(),
123            ctx.id,
124            ctx.fragment_id,
125        );
126
127        let cache_key_serde = create_cache_key_serde(&storage_key, &schema, &order_by, &group_by);
128        let managed_state = ManagedTopNState::<S>::new(state_table, cache_key_serde.clone());
129
130        Ok(Self {
131            schema,
132            offset: offset_and_limit.0,
133            limit: offset_and_limit.1,
134            managed_state,
135            storage_key_indices: storage_key.into_iter().map(|op| op.column_index).collect(),
136            watermark_order_key: topn_watermark_forwardable_order_key(&order_by),
137            group_by,
138            caches: GroupTopNCache::new(watermark_epoch, metrics_info),
139            cache_key_serde,
140            topn_cache_min_capacity: ctx.config.developer.topn_cache_min_capacity,
141            metrics,
142        })
143    }
144}
145
146impl<K: HashKey, S: StateStore, const WITH_TIES: bool> TopNExecutorBase
147    for InnerAppendOnlyGroupTopNExecutor<K, S, WITH_TIES>
148where
149    TopNCache<WITH_TIES>: AppendOnlyTopNCacheTrait,
150{
151    type State = S;
152
153    async fn apply_chunk(
154        &mut self,
155        chunk: StreamChunk,
156    ) -> StreamExecutorResult<Option<StreamChunk>> {
157        let keys = K::build_many(&self.group_by, chunk.data_chunk());
158        let mut stagings: HashMap<K, TopNStaging> = HashMap::new(); // K -> `TopNStaging`
159
160        let data_types = self.schema.data_types();
161        let deserializer = RowDeserializer::new(data_types.clone());
162        for (r, group_cache_key) in chunk.rows_with_holes().zip_eq_debug(keys.iter()) {
163            let Some((op, row_ref)) = r else {
164                continue;
165            };
166
167            // The pk without group by
168            let pk_row = row_ref.project(&self.storage_key_indices[self.group_by.len()..]);
169            let cache_key = serialize_pk_to_cache_key(pk_row, &self.cache_key_serde);
170
171            let group_key = row_ref.project(&self.group_by);
172            self.metrics.group_top_n_total_query_cache_count.inc();
173            // If 'self.caches' does not already have a cache for the current group, create a new
174            // cache for it and insert it into `self.caches`
175            if !self.caches.contains(group_cache_key) {
176                self.metrics.group_top_n_cache_miss_count.inc();
177                let mut topn_cache = TopNCache::with_min_capacity(
178                    self.offset,
179                    self.limit,
180                    data_types.clone(),
181                    self.topn_cache_min_capacity,
182                );
183                self.managed_state
184                    .init_append_only_topn_cache(Some(group_key), &mut topn_cache)
185                    .await?;
186                self.caches.put(group_cache_key.clone(), topn_cache);
187            }
188
189            let mut cache = self.caches.get_mut(group_cache_key).unwrap();
190            let staging = stagings.entry(group_cache_key.clone()).or_default();
191
192            debug_assert_eq!(op, Op::Insert);
193            cache.insert(
194                cache_key,
195                row_ref,
196                staging,
197                &mut self.managed_state,
198                &deserializer,
199            )?;
200        }
201
202        self.metrics
203            .group_top_n_cached_entry_count
204            .set(self.caches.len() as i64);
205
206        let mut chunk_builder = StreamChunkBuilder::unlimited(data_types, Some(chunk.capacity()));
207        for staging in stagings.into_values() {
208            for res in staging.into_deserialized_changes(&deserializer) {
209                let record = res?;
210                let _none = chunk_builder.append_record(record);
211            }
212        }
213
214        Ok(chunk_builder.take())
215    }
216
217    async fn flush_data(
218        &mut self,
219        epoch: EpochPair,
220    ) -> StreamExecutorResult<StateTablePostCommit<'_, S>> {
221        self.managed_state.flush(epoch).await
222    }
223
224    async fn try_flush_data(&mut self) -> StreamExecutorResult<()> {
225        self.managed_state.try_flush().await
226    }
227
228    fn clear_cache(&mut self) {
229        self.caches.clear();
230    }
231
232    fn evict(&mut self) {
233        self.caches.evict()
234    }
235
236    async fn init(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
237        self.managed_state.init_epoch(epoch).await
238    }
239
240    async fn handle_watermark(&mut self, watermark: Watermark) -> Option<Watermark> {
241        if watermark.col_idx == self.group_by[0] {
242            // The state table is ordered by the first group key column, so the watermark on it can
243            // also be used to clean up the states of the groups below it.
244            self.managed_state.update_watermark(watermark.val.clone());
245        }
246        // A row can only change the output of its own group, so watermarks on group key columns
247        // can always be forwarded. Watermarks on the first `ORDER BY` column can be forwarded if
248        // it's ordered `ASC NULLS LAST`. See `topn_watermark_forwardable_order_key` for the
249        // reasoning.
250        (self.group_by.contains(&watermark.col_idx)
251            || Some(watermark.col_idx) == self.watermark_order_key)
252            .then_some(watermark)
253    }
254}
risingwave_stream/executor/top_n/group_top_n_appendonly.rs

risingwave_stream/executor/top_n/
group_top_n_appendonly.rs
