Repository: pola-rs/polars
File: crates/polars-expr/src/state/execution_state.rs (Rust)

Code before:

use bitflags::bitflags;
use polars_core::config::verbose;
use polars_core::prelude::*;
use polars_ops::prelude::ChunkJoinOptIds;
use polars_utils::relaxed_cell::RelaxedCell;
use polars_utils::unique_id::UniqueId;

use super::NodeTimer;

pub type JoinTuplesCache = Arc<Mutex<PlHashMap<String, ChunkJoinOptIds>>>;

#[derive(Default)]
pub struct WindowCache {
    groups: RwLock<PlHashMap<String, GroupPositions>>,
    join_tuples: RwLock<PlHashMap<String, Arc<ChunkJoinOptIds>>>,
    map_idx: RwLock<PlHashMap<String, Arc<IdxCa>>>,
}

impl WindowCache {
    pub(crate) fn clear(&self) {
        let mut g = self.groups.write().unwrap();
        g.clear();
        let mut t = self.join_tuples.write().unwrap();
        t.clear();
        let mut i = self.map_idx.write().unwrap();
        i.clear();
    }

    pub fn get_groups(&self, key: &str) -> Option<GroupPositions> {
        let g = self.groups.read().unwrap();
        g.get(key).cloned()
    }

    pub fn insert_groups(&self, key: String, groups: GroupPositions) {
        let mut g = self.groups.write().unwrap();
        g.insert(key, groups);
    }

    pub fn get_join(&self, key: &str) -> Option<Arc<ChunkJoinOptIds>> {
        let g = self.join_tuples.read().unwrap();
        g.get(key).cloned()
    }

    pub fn insert_join(&self, key: String, join_tuples: Arc<ChunkJoinOptIds>) {
        let mut g = self.join_tuples.write().unwrap();
        g.insert(key, join_tuples);
    }

    pub fn get_map(&self, key: &str) -> Option<Arc<IdxCa>> {
        let g = self.map_idx.read().unwrap();

Review comment (suggestion):
Could we use destructuring syntax here? That gives us a safety net in case new fields are added in the future.


Code after:

use bitflags::bitflags;
use polars_core::config::verbose;
use polars_core::prelude::*;
use polars_ops::prelude::ChunkJoinOptIds;
use polars_utils::relaxed_cell::RelaxedCell;
use polars_utils::unique_id::UniqueId;

use super::NodeTimer;

pub type JoinTuplesCache = Arc<Mutex<PlHashMap<String, ChunkJoinOptIds>>>;

#[derive(Default)]
pub struct WindowCache {
    groups: RwLock<PlHashMap<String, GroupPositions>>,
    join_tuples: RwLock<PlHashMap<String, Arc<ChunkJoinOptIds>>>,
    map_idx: RwLock<PlHashMap<String, Arc<IdxCa>>>,
}

impl WindowCache {
    pub(crate) fn clear(&self) {
        let Self {
            groups,
            join_tuples,
            map_idx,
        } = self;
        groups.write().unwrap().clear();
        join_tuples.write().unwrap().clear();
        map_idx.write().unwrap().clear();
    }

    pub fn get_groups(&self, key: &str) -> Option<GroupPositions> {
        let g = self.groups.read().unwrap();
        g.get(key).cloned()
    }

    pub fn insert_groups(&self, key: String, groups: GroupPositions) {
        let mut g = self.groups.write().unwrap();
        g.insert(key, groups);
    }

    pub fn get_join(&self, key: &str) -> Option<Arc<ChunkJoinOptIds>> {
        let g = self.join_tuples.read().unwrap();
        g.get(key).cloned()
    }

    pub fn insert_join(&self, key: String, join_tuples: Arc<ChunkJoinOptIds>) {
        let mut g = self.join_tuples.write().unwrap();
        g.insert(key, join_tuples);
    }