mz_storage/upsert/
memory.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! An `UpsertStateBackend` that stores values in memory.
11
12// Allow usage of `std::collections::HashMap`.
13// We need to iterate through all the values in the map, so we can't use `mz_ore` wrapper.
14// Also, we don't need any ordering for the values fetched, so using std HashMap.
15#![allow(clippy::disallowed_types)]
16
17use std::collections::HashMap;
18
19use itertools::Itertools;
20use mz_ore::cast::CastFrom;
21
22use super::UpsertKey;
23use super::types::{
24    GetStats, MergeStats, MergeValue, PutStats, PutValue, StateValue, UpsertStateBackend,
25    UpsertValueAndSize, ValueMetadata,
26};
27
28/// A `HashMap` tracking its total size
29pub struct InMemoryHashMap<T, O> {
30    state: HashMap<UpsertKey, StateValue<T, O>>,
31    total_size: i64,
32}
33
34impl<T, O> Default for InMemoryHashMap<T, O> {
35    fn default() -> Self {
36        Self {
37            state: HashMap::new(),
38            total_size: 0,
39        }
40    }
41}
42
43#[async_trait::async_trait(?Send)]
44impl<T, O> UpsertStateBackend<T, O> for InMemoryHashMap<T, O>
45where
46    O: Clone + 'static,
47    T: Clone + 'static,
48{
49    fn supports_merge(&self) -> bool {
50        false
51    }
52
53    async fn multi_put<P>(&mut self, puts: P) -> Result<PutStats, anyhow::Error>
54    where
55        P: IntoIterator<Item = (UpsertKey, PutValue<StateValue<T, O>>)>,
56    {
57        let mut stats = PutStats::default();
58        for (key, p_value) in puts {
59            stats.processed_puts += 1;
60            match p_value.value {
61                Some(value) => {
62                    let size: i64 = value.memory_size().try_into().expect("less than i64 size");
63                    stats.adjust(Some(&value), Some(size), &p_value.previous_value_metadata);
64                    self.state.insert(key, value);
65                }
66                None => {
67                    stats.adjust::<T, O>(None, None, &p_value.previous_value_metadata);
68                    self.state.remove(&key);
69                }
70            }
71        }
72        self.total_size += stats.size_diff;
73        Ok(stats)
74    }
75
76    async fn multi_merge<M>(&mut self, _merges: M) -> Result<MergeStats, anyhow::Error>
77    where
78        M: IntoIterator<Item = (UpsertKey, MergeValue<StateValue<T, O>>)>,
79    {
80        anyhow::bail!("InMemoryHashMap does not support merging");
81    }
82
83    async fn multi_get<'r, G, R>(
84        &mut self,
85        gets: G,
86        results_out: R,
87    ) -> Result<GetStats, anyhow::Error>
88    where
89        G: IntoIterator<Item = UpsertKey>,
90        R: IntoIterator<Item = &'r mut UpsertValueAndSize<T, O>>,
91    {
92        let mut stats = GetStats::default();
93        for (key, result_out) in gets.into_iter().zip_eq(results_out) {
94            stats.processed_gets += 1;
95            let value = self.state.get(&key).cloned();
96            let metadata = value.as_ref().map(|v| ValueMetadata {
97                size: u64::cast_from(v.memory_size()),
98                is_tombstone: v.is_tombstone(),
99            });
100            stats.processed_gets_size += metadata.map_or(0, |m| m.size);
101            stats.returned_gets += metadata.map_or(0, |_| 1);
102            *result_out = UpsertValueAndSize { value, metadata };
103        }
104        Ok(stats)
105    }
106}
mz_storage/upsert/memory.rs

mz_storage/upsert/
memory.rs
