risingwave_stream/executor/join/
asof_join.rs

1// Copyright 2026 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
15//! Hash map for `AsOf` join with optional LRU cache support.
16//!
17//! When cache is disabled, the hash map directly queries the state table.
18//! When cache is enabled, it maintains a `ManagedLruCache` that caches all rows for
19//! each join key, indexed by (`inequality_key`, `pk_suffix`) for efficient range queries.
20
21use std::collections::BTreeMap;
22use std::ops::Bound;
23use std::sync::Arc;
24
25use anyhow::Context;
26use futures::StreamExt;
27use risingwave_common::bitmap::Bitmap;
28use risingwave_common::catalog::TableId;
29use risingwave_common::row::{CompactedRow, OwnedRow, Row, RowExt};
30use risingwave_common::types::{DataType, ScalarImpl};
31use risingwave_common::util::epoch::EpochPair;
32use risingwave_common::util::row_serde::OrderedRowSerde;
33use risingwave_common::util::sort_util::OrderType;
34use risingwave_common_estimate_size::{EstimateSize, KvSize};
35use risingwave_storage::StateStore;
36use risingwave_storage::store::PrefetchOptions;
37
38use super::hash_join::{JoinEntryError, JoinHashMapMetrics, TableInner};
39use super::join_row_set::JoinRowSet;
40use crate::cache::ManagedLruCache;
41use crate::common::metrics::MetricsInfo;
42use crate::common::table::state_table::{RowStream, StateTable, StateTablePostCommit};
43use crate::consistency::{consistency_error, enable_strict_consistency};
44use crate::executor::error::StreamExecutorResult;
45use crate::executor::monitor::StreamingMetrics;
46use crate::task::{ActorId, AtomicU64Ref, FragmentId};
47
48/// Memcomparable encoding for inequality key and pk suffix.
49type InequalKeyBytes = Vec<u8>;
50type PkSuffixBytes = Vec<u8>;
51
52/// Trait for encoding/decoding cached `AsOf` join rows.
53/// Unlike `JoinEncoding`, this doesn't carry degree since `AsOf` join has no degree tracking.
54pub trait AsOfRowEncoding: 'static + Send + Sync {
55    type Encoded: EstimateSize + Clone + Send + Sync;
56    fn encode(row: &OwnedRow) -> Self::Encoded;
57    /// Decode an encoded value back to `OwnedRow`.
58    fn decode(encoded: &Self::Encoded, data_types: &[DataType]) -> OwnedRow;
59}
60
61/// CPU-optimized encoding: stores full `OwnedRow` in cache (no encode/decode cost).
62pub struct AsOfCpuEncoding;
63
64impl AsOfRowEncoding for AsOfCpuEncoding {
65    type Encoded = OwnedRow;
66
67    fn encode(row: &OwnedRow) -> OwnedRow {
68        // TODO: Avoid this clone
69        row.clone()
70    }
71
72    fn decode(encoded: &OwnedRow, _data_types: &[DataType]) -> OwnedRow {
73        // TODO: Avoid this clone
74        encoded.clone()
75    }
76}
77
78/// Memory-optimized encoding: stores `CompactedRow` (serialized bytes) in cache.
79pub struct AsOfMemoryEncoding;
80
81impl AsOfRowEncoding for AsOfMemoryEncoding {
82    type Encoded = CompactedRow;
83
84    fn encode(row: &OwnedRow) -> CompactedRow {
85        row.into()
86    }
87
88    fn decode(encoded: &CompactedRow, data_types: &[DataType]) -> OwnedRow {
89        encoded
90            .deserialize(data_types)
91            .expect("failed to decode compacted row in AsOf join cache")
92    }
93}
94
95/// A cache entry for a single join key, storing all encoded rows indexed by (`inequality_key`, `pk_suffix`).
96/// The outer `BTreeMap` gives efficient range queries on inequality keys;
97/// the inner `JoinRowSet` adaptively uses `Vec` for small pk sets and `BTreeMap` for larger ones.
98///
99/// `V` is the encoded row type: `OwnedRow` for CPU encoding, `CompactedRow` for memory encoding.
100pub struct AsOfJoinCacheEntry<V: EstimateSize + Clone> {
101    /// `inequality_key_bytes` -> { `pk_suffix_bytes` -> `encoded_row` }
102    inner: BTreeMap<InequalKeyBytes, JoinRowSet<PkSuffixBytes, V>>,
103    /// Incrementally tracked heap size of all keys and values.
104    kv_heap_size: KvSize,
105}
106
107impl<V: EstimateSize + Clone> Default for AsOfJoinCacheEntry<V> {
108    fn default() -> Self {
109        Self {
110            inner: BTreeMap::new(),
111            kv_heap_size: KvSize::new(),
112        }
113    }
114}
115
116impl<V: EstimateSize + Clone> EstimateSize for AsOfJoinCacheEntry<V> {
117    fn estimated_heap_size(&self) -> usize {
118        // Only counts the cached rows, not the overhead of the B-tree itself.
119        self.kv_heap_size.size()
120    }
121}
122
123impl<V: EstimateSize + Clone> AsOfJoinCacheEntry<V> {
124    /// Find the encoded row with the greatest inequality key <= bound (upper bound query).
125    /// Returns the first entry by pk order in the pk sub-map for that inequality key.
126    fn upper_bound(&self, bound: Bound<&[u8]>) -> Option<&V> {
127        self.inner
128            .range::<[u8], _>((Bound::Unbounded, bound))
129            .next_back()
130            .and_then(|(_, pk_map)| pk_map.first_key_sorted().and_then(|k| pk_map.get(k)))
131    }
132
133    /// Find the encoded row with the smallest inequality key >= bound (lower bound query).
134    /// Returns the first entry by pk order in the pk sub-map for that inequality key.
135    fn lower_bound(&self, bound: Bound<&[u8]>) -> Option<&V> {
136        self.inner
137            .range::<[u8], _>((bound, Bound::Unbounded))
138            .next()
139            .and_then(|(_, pk_map)| pk_map.first_key_sorted().and_then(|k| pk_map.get(k)))
140    }
141
142    /// Iterate all encoded rows within an inequality key range.
143    fn range(
144        &self,
145        range: (Bound<&[u8]>, Bound<&[u8]>),
146    ) -> impl Iterator<Item = &V> + '_ + use<'_, V> {
147        self.inner
148            .range::<[u8], _>(range)
149            .flat_map(|(_, pk_map)| pk_map.values())
150    }
151
152    /// Find the first encoded row (by pk order) with exact inequality key match.
153    fn first_by_inequality(&self, ineq_key_bytes: &[u8]) -> Option<&V> {
154        self.inner
155            .get(ineq_key_bytes)
156            .and_then(|pk_map| pk_map.first_key_sorted().and_then(|k| pk_map.get(k)))
157    }
158
159    /// Find the first two encoded rows (by PK order) with exact inequality key match.
160    fn first_two_by_inequality(&self, ineq_key_bytes: &[u8]) -> (Option<&V>, Option<&V>) {
161        if let Some(pk_map) = self.inner.get(ineq_key_bytes) {
162            let (first_key, second_key) = pk_map.first_two_key_sorted();
163            let first = first_key.and_then(|k| pk_map.get(k));
164            let second = second_key.and_then(|k| pk_map.get(k));
165            (first, second)
166        } else {
167            (None, None)
168        }
169    }
170
171    /// Insert an encoded row into the cache entry.
172    fn insert(
173        &mut self,
174        ineq_key_bytes: InequalKeyBytes,
175        pk_suffix_bytes: PkSuffixBytes,
176        val: V,
177    ) -> Result<(), JoinEntryError> {
178        use std::collections::btree_map::Entry;
179        let pk_map = match self.inner.entry(ineq_key_bytes) {
180            Entry::Vacant(e) => {
181                self.kv_heap_size.add_val(e.key());
182                e.insert(JoinRowSet::default())
183            }
184            Entry::Occupied(e) => e.into_mut(),
185        };
186        if !enable_strict_consistency()
187            && let Some(old_val) = pk_map.remove(&pk_suffix_bytes)
188        {
189            self.kv_heap_size.sub(&pk_suffix_bytes, &old_val);
190        }
191        self.kv_heap_size.add(&pk_suffix_bytes, &val);
192        let ret = pk_map.try_insert(pk_suffix_bytes.clone(), val);
193        if !enable_strict_consistency() {
194            assert!(ret.is_ok(), "we have removed existing entry, if any");
195        }
196        ret.map(|_| ()).map_err(|_| JoinEntryError::Occupied)
197    }
198
199    /// Delete an encoded row from the cache entry.
200    fn delete(
201        &mut self,
202        ineq_key_bytes: &InequalKeyBytes,
203        pk_suffix_bytes: &PkSuffixBytes,
204    ) -> Result<(), JoinEntryError> {
205        if let Some(pk_map) = self.inner.get_mut(ineq_key_bytes) {
206            if let Some(old_val) = pk_map.remove(pk_suffix_bytes) {
207                self.kv_heap_size.sub(pk_suffix_bytes, &old_val);
208                if pk_map.is_empty() {
209                    self.kv_heap_size.sub_val(ineq_key_bytes);
210                    self.inner.remove(ineq_key_bytes);
211                }
212                Ok(())
213            } else if enable_strict_consistency() {
214                Err(JoinEntryError::Remove)
215            } else {
216                consistency_error!(
217                    ?pk_suffix_bytes,
218                    "removing an asof join cache entry but it's not in the cache"
219                );
220                Ok(())
221            }
222        } else if enable_strict_consistency() {
223            Err(JoinEntryError::Remove)
224        } else {
225            consistency_error!(
226                ?pk_suffix_bytes,
227                "removing an asof join cache entry but it's not in the cache"
228            );
229            Ok(())
230        }
231    }
232}
233
234/// A hash map for `AsOf` join with optional LRU cache support.
235///
236/// When `cache` is `None`, it directly queries the state table (no-cache mode).
237/// When `cache` is `Some`, it maintains a `ManagedLruCache` keyed by join key,
238/// caching all rows for each join key indexed by (`inequality_key`, `pk_suffix`).
239///
240/// `E` controls row encoding in the cache: `AsOfCpuEncoding` (full `OwnedRow`) or
241/// `AsOfMemoryEncoding` (compact `CompactedRow`).
242pub struct AsOfJoinHashMap<S: StateStore, E: AsOfRowEncoding> {
243    /// State table. Contains the data from upstream.
244    state: TableInner<S>,
245    /// The index of the inequality key column.
246    inequality_key_idx: usize,
247    /// Serializer for inequality key (single column value -> ordered bytes).
248    inequality_key_serializer: OrderedRowSerde,
249    /// Serializer for pk suffix (deduped pk columns -> ordered bytes).
250    pk_suffix_serializer: OrderedRowSerde,
251    /// The indices of the deduped pk columns in the input row (for projecting pk suffix).
252    pk_suffix_indices: Vec<usize>,
253    /// LRU cache: `join_key` (`OwnedRow`) -> encoded cache entry. None if cache is disabled.
254    cache: Option<ManagedLruCache<OwnedRow, AsOfJoinCacheEntry<E::Encoded>>>,
255    /// The indices of join key columns in the input row (for projecting join key).
256    join_key_indices: Vec<usize>,
257    /// All data types of the state table rows, needed for decoding `CompactedRow`.
258    state_all_data_types: Vec<DataType>,
259    /// Metrics of the hash map
260    metrics: JoinHashMapMetrics,
261}
262
263impl<S: StateStore, E: AsOfRowEncoding> AsOfJoinHashMap<S, E> {
264    /// Create a [`AsOfJoinHashMap`] for `AsOf` join with optional cache.
265    ///
266    /// If `watermark_epoch` is `Some`, an LRU cache is enabled.
267    /// If `watermark_epoch` is `None`, the hash map directly queries the state table.
268    #[expect(clippy::too_many_arguments)]
269    pub fn new(
270        state_join_key_indices: Vec<usize>,
271        state_table: StateTable<S>,
272        state_pk_indices: Vec<usize>,
273        inequality_key_idx: usize,
274        state_all_data_types: Vec<DataType>,
275        watermark_epoch: Option<AtomicU64Ref>,
276        metrics: Arc<StreamingMetrics>,
277        actor_id: ActorId,
278        fragment_id: FragmentId,
279        side: &'static str,
280    ) -> Self {
281        let join_table_id = state_table.table_id();
282
283        // Build serializer for inequality key (single column).
284        let inequality_key_serializer = OrderedRowSerde::new(
285            vec![state_all_data_types[inequality_key_idx].clone()],
286            vec![OrderType::ascending()],
287        );
288
289        // Build serializer for pk suffix (deduped pk columns).
290        let pk_suffix_serializer = OrderedRowSerde::new(
291            state_pk_indices
292                .iter()
293                .map(|idx| state_all_data_types[*idx].clone())
294                .collect(),
295            vec![OrderType::ascending(); state_pk_indices.len()],
296        );
297
298        let pk_suffix_indices = state_pk_indices.clone();
299        let join_key_indices = state_join_key_indices.clone();
300
301        let cache = watermark_epoch.map(|epoch_ref| {
302            let metrics_info = MetricsInfo::new(
303                metrics.clone(),
304                join_table_id,
305                actor_id,
306                format!("asof join {}", side),
307            );
308            ManagedLruCache::unbounded(epoch_ref, metrics_info)
309        });
310
311        let state = TableInner::new(state_pk_indices, state_join_key_indices, state_table, None);
312
313        Self {
314            state,
315            inequality_key_idx,
316            inequality_key_serializer,
317            pk_suffix_serializer,
318            pk_suffix_indices,
319            cache,
320            join_key_indices,
321            state_all_data_types,
322            metrics: JoinHashMapMetrics::new(&metrics, actor_id, fragment_id, side, join_table_id),
323        }
324    }
325
326    pub async fn init(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
327        self.state.table.init_epoch(epoch).await?;
328        Ok(())
329    }
330
331    pub fn update_watermark(&mut self, watermark: ScalarImpl) {
332        self.state.table.update_watermark(watermark);
333    }
334
335    pub fn table_id(&self) -> TableId {
336        self.state.table.table_id()
337    }
338
339    /// Cached entry count (number of distinct join keys in the LRU cache).
340    /// Returns 0 if cache is disabled.
341    pub fn entry_count(&self) -> usize {
342        self.cache.as_ref().map_or(0, |c| c.len())
343    }
344
345    /// Serialize a row as an inequality key for cache lookup.
346    /// The row should already be the inequality column(s) (e.g. a projected single-column row).
347    fn serialize_inequal_key(&self, row: &impl Row) -> Vec<u8> {
348        row.memcmp_serialize(&self.inequality_key_serializer)
349    }
350
351    /// Serialize the inequality key from a full state table row for cache key.
352    fn serialize_inequal_key_from_state_row(&self, row: &impl Row) -> Vec<u8> {
353        row.project(&[self.inequality_key_idx])
354            .memcmp_serialize(&self.inequality_key_serializer)
355    }
356
357    /// Serialize the pk suffix from a row for cache key.
358    fn serialize_pk_suffix(&self, row: &impl Row) -> Vec<u8> {
359        row.project(&self.pk_suffix_indices)
360            .memcmp_serialize(&self.pk_suffix_serializer)
361    }
362
363    /// Ensure the cache entry for the given join key is populated.
364    /// Does nothing if cache is disabled or the entry is already present.
365    async fn ensure_cache_populated(&mut self, join_key: &OwnedRow) -> StreamExecutorResult<()> {
366        let needs_populate = match &self.cache {
367            Some(cache) => !cache.contains(join_key),
368            None => return Ok(()),
369        };
370        self.metrics.inc_lookup();
371        if !needs_populate {
372            return Ok(());
373        }
374        self.metrics.inc_lookup_miss();
375
376        // Cache miss - load all rows for this join key from state table.
377        let sub_range: (Bound<OwnedRow>, Bound<OwnedRow>) = (Bound::Unbounded, Bound::Unbounded);
378        let table_iter = self
379            .state
380            .table
381            .iter_with_prefix(join_key, &sub_range, PrefetchOptions::default())
382            .await?;
383        let mut pinned = std::pin::pin!(table_iter);
384        let mut entry = AsOfJoinCacheEntry::<E::Encoded>::default();
385        while let Some(row) = pinned.next().await.transpose()? {
386            let ineq_key = self.serialize_inequal_key_from_state_row(&row);
387            let pk_suffix = self.serialize_pk_suffix(&row);
388            entry
389                .insert(ineq_key, pk_suffix, E::encode(&row))
390                .with_context(|| format!("row: {}", row.display()))?;
391        }
392        self.cache.as_mut().unwrap().put(join_key.clone(), entry);
393        Ok(())
394    }
395
396    /// Insert a row into the state table and optionally update the cache.
397    pub fn insert(&mut self, value: impl Row) -> StreamExecutorResult<()> {
398        if self.cache.is_some() {
399            // Convert to OwnedRow first to avoid move issues with `impl Row`.
400            let owned_row = value.to_owned_row();
401            let ineq_key = self.serialize_inequal_key_from_state_row(&owned_row);
402            let pk_suffix = self.serialize_pk_suffix(&owned_row);
403            // Use reference to avoid moving owned_row.
404            let join_key = (&owned_row).project(&self.join_key_indices).to_owned_row();
405            // Update cache if the join key is cached.
406            if let Some(cache) = &mut self.cache
407                && let Some(mut entry) = cache.get_mut(&join_key)
408            {
409                entry
410                    .insert(ineq_key, pk_suffix, E::encode(&owned_row))
411                    .with_context(|| format!("row: {}", owned_row.display()))?;
412            } else {
413                self.metrics.inc_insert_cache_miss();
414            }
415            self.state.table.insert(owned_row);
416        } else {
417            self.state.table.insert(value);
418        }
419        Ok(())
420    }
421
422    /// Delete a row from the state table and optionally update the cache.
423    pub fn delete(&mut self, value: impl Row) -> StreamExecutorResult<()> {
424        if self.cache.is_some() {
425            // Convert to OwnedRow first to avoid move issues with `impl Row`.
426            let owned_row = value.to_owned_row();
427            let ineq_key = self.serialize_inequal_key_from_state_row(&owned_row);
428            let pk_suffix = self.serialize_pk_suffix(&owned_row);
429            // Use reference to avoid moving owned_row.
430            let join_key = (&owned_row).project(&self.join_key_indices).to_owned_row();
431            // Update cache if the join key is cached.
432            if let Some(cache) = &mut self.cache
433                && let Some(mut entry) = cache.get_mut(&join_key)
434            {
435                entry
436                    .delete(&ineq_key, &pk_suffix)
437                    .with_context(|| format!("row: {}", owned_row.display()))?;
438            }
439            self.state.table.delete(owned_row);
440        } else {
441            self.state.table.delete(value);
442        }
443        Ok(())
444    }
445
446    /// Evict the LRU cache. No-op if cache is not enabled.
447    pub fn evict_cache(&mut self) {
448        if let Some(cache) = &mut self.cache {
449            cache.evict();
450        }
451    }
452
453    /// Flush the state table.
454    pub async fn flush(
455        &mut self,
456        epoch: EpochPair,
457    ) -> StreamExecutorResult<AsOfJoinHashMapPostCommit<'_, S, E>> {
458        self.metrics.flush();
459        if let Some(cache) = &mut self.cache {
460            cache.evict();
461        }
462        let state_post_commit = self.state.table.commit(epoch).await?;
463        Ok(AsOfJoinHashMapPostCommit {
464            state: state_post_commit,
465            cache: self.cache.as_mut(),
466        })
467    }
468
469    pub async fn try_flush(&mut self) -> StreamExecutorResult<()> {
470        self.state.table.try_flush().await?;
471        Ok(())
472    }
473
474    pub async fn upper_bound_by_inequality_with_jk_prefix(
475        &mut self,
476        join_key: impl Row,
477        bound: Bound<&impl Row>,
478    ) -> StreamExecutorResult<Option<OwnedRow>> {
479        if self.cache.is_some() {
480            let join_key_owned = join_key.to_owned_row();
481            // Serialize bound bytes before borrowing cache to avoid borrow conflicts.
482            let bound_bytes = match &bound {
483                Bound::Included(r) => Bound::Included(self.serialize_inequal_key(r)),
484                Bound::Excluded(r) => Bound::Excluded(self.serialize_inequal_key(r)),
485                Bound::Unbounded => Bound::Unbounded,
486            };
487            self.ensure_cache_populated(&join_key_owned).await?;
488            let cache = self.cache.as_mut().unwrap();
489            if let Some(entry) = cache.get(&join_key_owned) {
490                let bound_ref = match &bound_bytes {
491                    Bound::Included(b) => Bound::Included(b.as_slice()),
492                    Bound::Excluded(b) => Bound::Excluded(b.as_slice()),
493                    Bound::Unbounded => Bound::Unbounded,
494                };
495                return Ok(entry
496                    .upper_bound(bound_ref)
497                    .map(|v| E::decode(v, &self.state_all_data_types)));
498            }
499            // Not in cache (shouldn't happen after ensure_cache_populated, but fallback)
500        }
501        // Non-cached path: use rev_iter to find the largest inequality key within the
502        // bound, then forward-scan for the smallest PK at that inequality key. This
503        // matches the cached path's behavior (smallest PK wins on tie), which is
504        // required for consistency with eq_join_right's first_two_by_inequality logic.
505        let sub_range = (Bound::<OwnedRow>::Unbounded, bound);
506        let table_iter = self
507            .state
508            .table
509            .rev_iter_with_prefix(&join_key, &sub_range, PrefetchOptions::default())
510            .await?;
511        let mut pinned_table_iter = std::pin::pin!(table_iter);
512        let Some(rev_row) = pinned_table_iter.next().await.transpose()? else {
513            return Ok(None);
514        };
515        // Extract the inequality key from the found row and forward-scan for smallest PK.
516        let found_ineq_key = (&rev_row)
517            .project(&[self.inequality_key_idx])
518            .to_owned_row();
519        let exact_range = (
520            Bound::Included(&found_ineq_key),
521            Bound::Included(&found_ineq_key),
522        );
523        let fwd_iter = self
524            .state
525            .table
526            .iter_with_prefix(&join_key, &exact_range, PrefetchOptions::default())
527            .await?;
528        let mut pinned_fwd = std::pin::pin!(fwd_iter);
529        // Forward scan returns smallest PK first; fall back to rev_row if empty (shouldn't happen).
530        let row = pinned_fwd.next().await.transpose()?;
531        Ok(row.or(Some(rev_row)))
532    }
533
534    pub async fn lower_bound_by_inequality_with_jk_prefix(
535        &mut self,
536        join_key: &impl Row,
537        bound: Bound<&impl Row>,
538    ) -> StreamExecutorResult<Option<OwnedRow>> {
539        if self.cache.is_some() {
540            let join_key_owned = join_key.to_owned_row();
541            // Serialize bound bytes before borrowing cache to avoid borrow conflicts.
542            let bound_bytes = match &bound {
543                Bound::Included(r) => Bound::Included(self.serialize_inequal_key(r)),
544                Bound::Excluded(r) => Bound::Excluded(self.serialize_inequal_key(r)),
545                Bound::Unbounded => Bound::Unbounded,
546            };
547            self.ensure_cache_populated(&join_key_owned).await?;
548            let cache = self.cache.as_mut().unwrap();
549            if let Some(entry) = cache.get(&join_key_owned) {
550                let bound_ref = match &bound_bytes {
551                    Bound::Included(b) => Bound::Included(b.as_slice()),
552                    Bound::Excluded(b) => Bound::Excluded(b.as_slice()),
553                    Bound::Unbounded => Bound::Unbounded,
554                };
555                return Ok(entry
556                    .lower_bound(bound_ref)
557                    .map(|v| E::decode(v, &self.state_all_data_types)));
558            }
559        }
560        // Non-cached path
561        let sub_range = (bound, Bound::<OwnedRow>::Unbounded);
562        let table_iter = self
563            .state
564            .table
565            .iter_with_prefix(join_key, &sub_range, PrefetchOptions::default())
566            .await?;
567        let mut pinned_table_iter = std::pin::pin!(table_iter);
568        let row = pinned_table_iter.next().await.transpose()?;
569        Ok(row)
570    }
571
572    async fn table_iter_by_inequality_with_jk_prefix<'a>(
573        &'a self,
574        join_key: &'a impl Row,
575        range: &'a (Bound<impl Row>, Bound<impl Row>),
576    ) -> StreamExecutorResult<impl RowStream<'a>> {
577        self.state
578            .table
579            .iter_with_prefix(join_key, range, PrefetchOptions::default())
580            .await
581    }
582
583    pub async fn first_by_inequality_with_jk_prefix(
584        &mut self,
585        join_key: &impl Row,
586        inequality_key: &impl Row,
587    ) -> StreamExecutorResult<Option<OwnedRow>> {
588        if self.cache.is_some() {
589            let join_key_owned = join_key.to_owned_row();
590            // Serialize before borrowing cache to avoid borrow conflicts.
591            let ineq_key_bytes = self.serialize_inequal_key(inequality_key);
592            self.ensure_cache_populated(&join_key_owned).await?;
593            let cache = self.cache.as_mut().unwrap();
594            if let Some(entry) = cache.get(&join_key_owned) {
595                return Ok(entry
596                    .first_by_inequality(&ineq_key_bytes)
597                    .map(|v| E::decode(v, &self.state_all_data_types)));
598            }
599        }
600        // Non-cached path
601        let range = (
602            Bound::Included(inequality_key),
603            Bound::Included(inequality_key),
604        );
605        let table_iter = self
606            .table_iter_by_inequality_with_jk_prefix(join_key, &range)
607            .await?;
608        let mut pinned_table_iter = std::pin::pin!(table_iter);
609        let row = pinned_table_iter.next().await.transpose()?;
610        Ok(row)
611    }
612
613    /// Return the first two rows (by PK order) with exact inequality key match.
614    /// Used by `eq_join_right` to determine replacement: Insert needs first row
615    /// for comparison, Delete needs first + second for fallback.
616    pub async fn first_two_by_inequality_with_jk_prefix(
617        &mut self,
618        join_key: &impl Row,
619        inequality_key: &impl Row,
620    ) -> StreamExecutorResult<(Option<OwnedRow>, Option<OwnedRow>)> {
621        if self.cache.is_some() {
622            let join_key_owned = join_key.to_owned_row();
623            let ineq_key_bytes = self.serialize_inequal_key(inequality_key);
624            self.ensure_cache_populated(&join_key_owned).await?;
625            let cache = self.cache.as_mut().unwrap();
626            if let Some(entry) = cache.get(&join_key_owned) {
627                let (first, second) = entry.first_two_by_inequality(&ineq_key_bytes);
628                return Ok((
629                    first.map(|v| E::decode(v, &self.state_all_data_types)),
630                    second.map(|v| E::decode(v, &self.state_all_data_types)),
631                ));
632            }
633        }
634        // Non-cached path
635        let range = (
636            Bound::Included(inequality_key),
637            Bound::Included(inequality_key),
638        );
639        let table_iter = self
640            .table_iter_by_inequality_with_jk_prefix(join_key, &range)
641            .await?;
642        let mut pinned_table_iter = std::pin::pin!(table_iter);
643        let first = pinned_table_iter.next().await.transpose()?;
644        let second = if first.is_some() {
645            pinned_table_iter.next().await.transpose()?
646        } else {
647            None
648        };
649        Ok((first, second))
650    }
651
652    pub async fn range_by_inequality_with_jk_prefix<'a>(
653        &'a mut self,
654        join_key: &'a impl Row,
655        inequality_key_range: &'a (Bound<impl Row>, Bound<impl Row>),
656    ) -> StreamExecutorResult<impl RowStream<'a>> {
657        if self.cache.is_some() {
658            let join_key_owned = join_key.to_owned_row();
659            // Serialize bound bytes before borrowing cache.
660            let bound_lo = match &inequality_key_range.0 {
661                Bound::Included(r) => Bound::Included(self.serialize_inequal_key(r)),
662                Bound::Excluded(r) => Bound::Excluded(self.serialize_inequal_key(r)),
663                Bound::Unbounded => Bound::Unbounded,
664            };
665            let bound_hi = match &inequality_key_range.1 {
666                Bound::Included(r) => Bound::Included(self.serialize_inequal_key(r)),
667                Bound::Excluded(r) => Bound::Excluded(self.serialize_inequal_key(r)),
668                Bound::Unbounded => Bound::Unbounded,
669            };
670            self.ensure_cache_populated(&join_key_owned).await?;
671            let cache = self.cache.as_mut().unwrap();
672            if let Some(entry) = cache.get(&join_key_owned) {
673                let range = (
674                    bound_lo.as_ref().map(|b| b.as_slice()),
675                    bound_hi.as_ref().map(|b| b.as_slice()),
676                );
677                // Cached: lazily decode cloned rows from cache and wrap as stream.
678                let data_types = &self.state_all_data_types;
679                return Ok(futures::future::Either::Left(futures::stream::iter(
680                    entry.range(range).map(|v| Ok(E::decode(v, data_types))),
681                )));
682            }
683        }
684        // Non-cached path: stream directly from state table without collecting.
685        let table_iter = self
686            .state
687            .table
688            .iter_with_prefix(join_key, inequality_key_range, PrefetchOptions::default())
689            .await?;
690        Ok(futures::future::Either::Right(table_iter))
691    }
692
693    /// Return true if the inequality key is null.
694    pub fn check_inequal_key_null(&self, row: &impl Row) -> bool {
695        row.datum_at(self.inequality_key_idx).is_none()
696    }
697
698    pub fn get_pk_from_row<'a>(&'a self, row: impl Row + 'a) -> impl Row + 'a {
699        row.project(&self.state.pk_indices)
700    }
701}
702
703#[must_use]
704pub struct AsOfJoinHashMapPostCommit<'a, S: StateStore, E: AsOfRowEncoding> {
705    state: StateTablePostCommit<'a, S>,
706    cache: Option<&'a mut ManagedLruCache<OwnedRow, AsOfJoinCacheEntry<E::Encoded>>>,
707}
708
709impl<S: StateStore, E: AsOfRowEncoding> AsOfJoinHashMapPostCommit<'_, S, E> {
710    pub async fn post_yield_barrier(
711        self,
712        vnode_bitmap: Option<Arc<Bitmap>>,
713    ) -> StreamExecutorResult<Option<bool>> {
714        let cache_may_stale = self.state.post_yield_barrier(vnode_bitmap.clone()).await?;
715        let keyed_cache_may_stale = cache_may_stale.map(|(_, changed)| changed);
716        // Clear cache if the keyed cache may be stale after the barrier update.
717        if keyed_cache_may_stale.unwrap_or(false)
718            && let Some(cache) = self.cache
719        {
720            cache.clear();
721        }
722        Ok(keyed_cache_may_stale)
723    }
724}
risingwave_stream/executor/join/asof_join.rs

risingwave_stream/executor/join/
asof_join.rs
