perf(over window): locality-based range cache for general over window executor (#11576)

Signed-off-by: Richard Chien <[email protected]>

Author: stdrc

Cargo.lock

@@ -591,12 +591,20 @@ message EowcOverWindowNode {
   catalog.Table state_table = 4;
 }
 
-// Note this is not exactly the same as EowcOverWindowNode in terms of future extension, so better to split as 2 different nodes.
+enum OverWindowCachePolicy {
+  OVER_WINDOW_CACHE_POLICY_UNSPECIFIED = 0;
+  OVER_WINDOW_CACHE_POLICY_FULL = 1;
+  OVER_WINDOW_CACHE_POLICY_RECENT = 2;
+  OVER_WINDOW_CACHE_POLICY_RECENT_FIRST_N = 3;
+  OVER_WINDOW_CACHE_POLICY_RECENT_LAST_N = 4;
+}
+
 message OverWindowNode {
   repeated expr.WindowFunction calls = 1;
   repeated uint32 partition_by = 2;
   repeated common.ColumnOrder order_by = 3;
   catalog.Table state_table = 4;
+  OverWindowCachePolicy cache_policy = 5;
 }
 
 message StreamNode {

proto/stream_plan.proto

@@ -124,7 +124,16 @@ pub trait Row: Sized + std::fmt::Debug + PartialEq + Eq {
     /// Determines whether the datums of this row are equal to those of another.
     #[inline]
     fn eq(this: &Self, other: impl Row) -> bool {
-        this.iter().eq(other.iter())
+        if this.len() != other.len() {
+            return false;
+        }
+        for i in (0..this.len()).rev() {
+            // compare from the end to the start, as it's more likely to have same prefix
+            if this.datum_at(i) != other.datum_at(i) {
+                return false;
+            }
+        }
+        true
     }
 }
 

src/common/src/row/mod.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod over_window;
 mod query_mode;
 mod search_path;
 mod transaction_isolation_level;
@@ -23,6 +24,7 @@ use std::ops::Deref;
 use chrono_tz::Tz;
 use educe::{self, Educe};
 use itertools::Itertools;
+pub use over_window::OverWindowCachePolicy;
 pub use query_mode::QueryMode;
 pub use search_path::{SearchPath, USER_NAME_WILD_CARD};
 use tracing::info;
@@ -34,7 +36,7 @@ use crate::util::epoch::Epoch;
 
 // This is a hack, &'static str is not allowed as a const generics argument.
 // TODO: refine this using the adt_const_params feature.
-const CONFIG_KEYS: [&str; 36] = [
+const CONFIG_KEYS: [&str; 37] = [
     "RW_IMPLICIT_FLUSH",
     "CREATE_COMPACTION_GROUP_FOR_MV",
     "QUERY_MODE",
@@ -71,6 +73,7 @@ const CONFIG_KEYS: [&str; 36] = [
     "STANDARD_CONFORMING_STRINGS",
     "RW_STREAMING_RATE_LIMIT",
     "CDC_BACKFILL",
+    "RW_STREAMING_OVER_WINDOW_CACHE_POLICY",
 ];
 
 // MUST HAVE 1v1 relationship to CONFIG_KEYS. e.g. CONFIG_KEYS[IMPLICIT_FLUSH] =
@@ -111,6 +114,7 @@ const ROW_SECURITY: usize = 32;
 const STANDARD_CONFORMING_STRINGS: usize = 33;
 const RW_STREAMING_RATE_LIMIT: usize = 34;
 const CDC_BACKFILL: usize = 35;
+const STREAMING_OVER_WINDOW_CACHE_POLICY: usize = 36;
 
 trait ConfigEntry: Default + for<'a> TryFrom<&'a [&'a str], Error = RwError> {
     fn entry_name() -> &'static str;
@@ -478,6 +482,10 @@ pub struct ConfigMap {
     streaming_rate_limit: StreamingRateLimit,
 
     cdc_backfill: CdcBackfill,
+
+    /// Cache policy for partition cache in streaming over window.
+    /// Can be "full", "recent", "recent_first_n" or "recent_last_n".
+    streaming_over_window_cache_policy: OverWindowCachePolicy,
 }
 
 impl ConfigMap {
@@ -593,6 +601,8 @@ impl ConfigMap {
             self.streaming_rate_limit = val.as_slice().try_into()?;
         } else if key.eq_ignore_ascii_case(CdcBackfill::entry_name()) {
             self.cdc_backfill = val.as_slice().try_into()?
+        } else if key.eq_ignore_ascii_case(OverWindowCachePolicy::entry_name()) {
+            self.streaming_over_window_cache_policy = val.as_slice().try_into()?;
         } else {
             return Err(ErrorCode::UnrecognizedConfigurationParameter(key.to_string()).into());
         }
@@ -678,6 +688,8 @@ impl ConfigMap {
             Ok(self.streaming_rate_limit.to_string())
         } else if key.eq_ignore_ascii_case(CdcBackfill::entry_name()) {
             Ok(self.cdc_backfill.to_string())
+        } else if key.eq_ignore_ascii_case(OverWindowCachePolicy::entry_name()) {
+            Ok(self.streaming_over_window_cache_policy.to_string())
         } else {
             Err(ErrorCode::UnrecognizedConfigurationParameter(key.to_string()).into())
         }
@@ -864,7 +876,12 @@ impl ConfigMap {
                 name: CdcBackfill::entry_name().to_lowercase(),
                 setting: self.cdc_backfill.to_string(),
                 description: String::from("Enable backfill for CDC table to allow lock-free and incremental snapshot"),
-            }
+            },
+            VariableInfo{
+                name: OverWindowCachePolicy::entry_name().to_lowercase(),
+                setting: self.streaming_over_window_cache_policy.to_string(),
+                description: String::from(r#"Cache policy for partition cache in streaming over window. Can be "full", "recent", "recent_first_n" or "recent_last_n"."#),
+            },
         ]
     }
 
@@ -999,4 +1016,8 @@ impl ConfigMap {
     pub fn get_cdc_backfill(&self) -> bool {
         self.cdc_backfill.0
     }
+
+    pub fn get_streaming_over_window_cache_policy(&self) -> OverWindowCachePolicy {
+        self.streaming_over_window_cache_policy
+    }
 }

src/common/src/session_config/mod.rs

@@ -0,0 +1,125 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::str::FromStr;
+
+use enum_as_inner::EnumAsInner;
+use parse_display::{Display, FromStr};
+use risingwave_pb::stream_plan::PbOverWindowCachePolicy;
+
+use super::{ConfigEntry, CONFIG_KEYS, STREAMING_OVER_WINDOW_CACHE_POLICY};
+use crate::error::ErrorCode::{self, InvalidConfigValue};
+use crate::error::RwError;
+
+#[derive(Copy, Default, Debug, Clone, PartialEq, Eq, FromStr, Display, EnumAsInner)]
+#[display(style = "snake_case")]
+pub enum OverWindowCachePolicy {
+    /// Cache all entries.
+    #[default]
+    Full,
+    /// Cache only recently accessed range of entries.
+    Recent,
+    /// Cache only the first N entries in recently accessed range.
+    RecentFirstN,
+    /// Cache only the last N entries in recently accessed range.
+    RecentLastN,
+}
+
+impl TryFrom<&[&str]> for OverWindowCachePolicy {
+    type Error = RwError;
+
+    fn try_from(value: &[&str]) -> Result<Self, Self::Error> {
+        if value.len() != 1 {
+            return Err(ErrorCode::InternalError(format!(
+                "SET {} takes only one argument",
+                Self::entry_name()
+            ))
+            .into());
+        }
+
+        let s = value[0].to_ascii_lowercase().replace('-', "_");
+        OverWindowCachePolicy::from_str(&s).map_err(|_| {
+            InvalidConfigValue {
+                config_entry: Self::entry_name().to_string(),
+                config_value: s.to_string(),
+            }
+            .into()
+        })
+    }
+}
+
+impl ConfigEntry for OverWindowCachePolicy {
+    fn entry_name() -> &'static str {
+        CONFIG_KEYS[STREAMING_OVER_WINDOW_CACHE_POLICY]
+    }
+}
+
+impl OverWindowCachePolicy {
+    pub fn to_protobuf(self) -> PbOverWindowCachePolicy {
+        match self {
+            Self::Full => PbOverWindowCachePolicy::Full,
+            Self::Recent => PbOverWindowCachePolicy::Recent,
+            Self::RecentFirstN => PbOverWindowCachePolicy::RecentFirstN,
+            Self::RecentLastN => PbOverWindowCachePolicy::RecentLastN,
+        }
+    }
+
+    pub fn from_protobuf(pb: PbOverWindowCachePolicy) -> Self {
+        match pb {
+            PbOverWindowCachePolicy::Unspecified => Self::default(),
+            PbOverWindowCachePolicy::Full => Self::Full,
+            PbOverWindowCachePolicy::Recent => Self::Recent,
+            PbOverWindowCachePolicy::RecentFirstN => Self::RecentFirstN,
+            PbOverWindowCachePolicy::RecentLastN => Self::RecentLastN,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parse_over_window_cache_policy() {
+        assert_eq!(
+            OverWindowCachePolicy::try_from(["full"].as_slice()).unwrap(),
+            OverWindowCachePolicy::Full
+        );
+        assert_eq!(
+            OverWindowCachePolicy::try_from(["recent"].as_slice()).unwrap(),
+            OverWindowCachePolicy::Recent
+        );
+        assert_eq!(
+            OverWindowCachePolicy::try_from(["RECENT"].as_slice()).unwrap(),
+            OverWindowCachePolicy::Recent
+        );
+        assert_eq!(
+            OverWindowCachePolicy::try_from(["recent_first_n"].as_slice()).unwrap(),
+            OverWindowCachePolicy::RecentFirstN
+        );
+        assert_eq!(
+            OverWindowCachePolicy::try_from(["recent_last_n"].as_slice()).unwrap(),
+            OverWindowCachePolicy::RecentLastN
+        );
+        assert_eq!(
+            OverWindowCachePolicy::try_from(["recent-last-n"].as_slice()).unwrap(),
+            OverWindowCachePolicy::RecentLastN
+        );
+        assert_eq!(
+            OverWindowCachePolicy::try_from(["recent_last_N"].as_slice()).unwrap(),
+            OverWindowCachePolicy::RecentLastN
+        );
+        assert!(OverWindowCachePolicy::try_from(["foo"].as_slice()).is_err());
+    }
+}

src/common/src/session_config/over_window.rs

@@ -78,6 +78,10 @@ impl Frame {
             exclusion,
         }
     }
+
+    pub fn is_unbounded(&self) -> bool {
+        self.bounds.is_unbounded()
+    }
 }
 
 impl Frame {
@@ -127,6 +131,10 @@ impl FrameBounds {
             Self::Rows(_, end) => matches!(end, FrameBound::UnboundedFollowing),
         }
     }
+
+    pub fn is_unbounded(&self) -> bool {
+        self.start_is_unbounded() || self.end_is_unbounded()
+    }
 }
 
 impl Display for FrameBounds {

src/expr/src/window_function/call.rs

@@ -58,22 +58,19 @@ impl StreamOverWindow {
 
         let mut order_cols = HashSet::new();
         for idx in self.logical.partition_key_indices() {
-            if !order_cols.contains(&idx) {
+            if order_cols.insert(idx) {
                 tbl_builder.add_order_column(idx, OrderType::ascending());
-                order_cols.insert(idx);
             }
         }
         let read_prefix_len_hint = tbl_builder.get_current_pk_len();
         for o in self.logical.order_key() {
-            if !order_cols.contains(&o.column_index) {
+            if order_cols.insert(o.column_index) {
                 tbl_builder.add_order_column(o.column_index, o.order_type);
-                order_cols.insert(o.column_index);
             }
         }
-        for idx in self.logical.input.logical_pk() {
-            if !order_cols.contains(idx) {
-                tbl_builder.add_order_column(*idx, OrderType::ascending());
-                order_cols.insert(*idx);
+        for &idx in self.logical.input.logical_pk() {
+            if order_cols.insert(idx) {
+                tbl_builder.add_order_column(idx, OrderType::ascending());
             }
         }
 
@@ -123,12 +120,19 @@ impl StreamNode for StreamOverWindow {
             .infer_state_table()
             .with_id(state.gen_table_id_wrapped())
             .to_internal_table_prost();
+        let cache_policy = self
+            .base
+            .ctx
+            .session_ctx()
+            .config()
+            .get_streaming_over_window_cache_policy();
 
         PbNodeBody::OverWindow(OverWindowNode {
             calls,
             partition_by,
             order_by,
             state_table: Some(state_table),
+            cache_policy: cache_policy.to_protobuf() as _,
         })
     }
 }

src/frontend/src/optimizer/plan_node/stream_over_window.rs

@@ -38,6 +38,7 @@ memcomparable = "0.2"
 multimap = "0.8"
 num-traits = "0.2"
 parking_lot = "0.12"
+parse-display = "0.8"
 pin-project = "1"
 prometheus = { version = "0.13", features = ["process"] }
 prost = "0.11"