use std::{
cell::RefCell,
collections::VecDeque,
hash::{BuildHasherDefault, Hasher},
panic::Location,
sync::{Arc, RwLock},
};
use dbsp::{
Circuit, DynZWeight, OrdZSet, RootCircuit, Runtime, Stream, ZWeight,
circuit::{LocalStoreMarker, WorkerLocation, WorkerLocations},
dynamic::{DowncastTrait, DynData},
operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators},
storage::file::to_bytes,
trace::{
BatchReader, BatchReaderFactories, Cursor, OrdIndexedWSet as DynOrdIndexedWSet,
OrdIndexedWSetFactories, SpineSnapshot, aligned_deserialize,
},
utils::Tup1,
};
use quick_cache::{
OptionsBuilder, Weighter,
sync::{Cache, DefaultLifecycle, GuardResult},
};
use typedmap::TypedMapKey;
use crate::{SqlString, Uuid};
const CACHE_CAPACITY: usize = 1 << 26;
const DEFAULT_CACHE_CAPACITY_BYTES: u64 = 1 << 30;
pub type InternedStringId = Uuid;
type InternedString = (SqlString, bool);
#[derive(Default)]
struct IdentityHasher {
hash: u64,
}
impl Hasher for IdentityHasher {
#[inline]
fn write(&mut self, bytes: &[u8]) {
debug_assert_eq!(bytes.len(), 16, "Expected 16 bytes");
self.hash = u64::from_le_bytes(bytes[0..8].try_into().unwrap());
}
#[inline]
fn write_u64(&mut self, _i: u64) {}
#[inline]
fn write_usize(&mut self, _i: usize) {}
#[inline]
fn finish(&self) -> u64 {
self.hash
}
}
type BuildIdentityHasher = BuildHasherDefault<IdentityHasher>;
#[derive(Clone)]
struct StringWeighter;
impl Weighter<InternedStringId, InternedString> for StringWeighter {
fn weight(&self, _key: &InternedStringId, val: &InternedString) -> u64 {
if val.1 {
0
} else {
(val.0.len() + size_of::<InternedStringId>()) as u64
+ size_of::<InternedString>() as u64
}
}
}
thread_local! {
static CURRENT_STEP: RefCell<u64> = const { RefCell::new(0) };
static PINNED_STRINGS: RefCell<VecDeque<(InternedStringId, (SqlString, u64))>> = const { RefCell::new(VecDeque::new()) };
static INTERNED_STRING_CACHE: RefCell<Arc<InternedStringCache>> = RefCell::new(Arc::new(init_interned_string_cache(None)));
static INTERNED_STRING_BY_ID: RefCell<Arc<RwLock<InternedStringSpineSnapshot>>> = RefCell::new(Arc::new(RwLock::new(empty_by_id())));
}
pub type InternedStringSpineSnapshot =
SpineSnapshot<DynOrdIndexedWSet<DynData, DynData, DynZWeight>>;
fn empty_by_id() -> InternedStringSpineSnapshot {
let factories: OrdIndexedWSetFactories<DynData, DynData, DynZWeight> =
BatchReaderFactories::new::<InternedStringId, Tup1<SqlString>, ZWeight>();
SpineSnapshot::<DynOrdIndexedWSet<DynData, DynData, DynZWeight>>::new(factories)
}
type InternedStringCache = Cache<
InternedStringId,
InternedString,
StringWeighter,
BuildIdentityHasher,
DefaultLifecycle<InternedStringId, InternedString>,
>;
fn init_interned_string_cache(cache_capacity_bytes: Option<u64>) -> InternedStringCache {
Cache::with_options(
OptionsBuilder::new()
.estimated_items_capacity(CACHE_CAPACITY)
.weight_capacity(cache_capacity_bytes.unwrap_or(DEFAULT_CACHE_CAPACITY_BYTES))
.build()
.unwrap(),
StringWeighter,
BuildIdentityHasher::default(),
DefaultLifecycle::default(),
)
}
fn hash_string(s: &SqlString) -> Uuid {
let hash = blake3::hash(s.str().as_bytes());
let bytes = hash.as_bytes();
Uuid::from_bytes(bytes[0..16].try_into().unwrap())
}
pub fn intern_string(s: &SqlString) -> InternedStringId {
let id = hash_string(s);
INTERNED_STRING_CACHE.with_borrow(|cache| {
if let GuardResult::Guard(g) = cache.get_value_or_guard(&id, None) {
let current_step = CURRENT_STEP.with_borrow(|step| *step);
let val = (s.clone(), true);
g.insert(val)
.expect("Failed to insert into interned string cache");
PINNED_STRINGS.with_borrow_mut(|pinned| {
pinned.push_back((id.clone(), (s.clone(), current_step)))
});
}
});
id
}
pub fn unintern_string(id: &InternedStringId) -> Option<SqlString> {
INTERNED_STRING_CACHE.with_borrow(|cache| {
cache.get(id).map(|(string, _step)| string).or_else(|| {
INTERNED_STRING_BY_ID.with_borrow(|spine| {
let mut cursor = spine.read().unwrap().cursor();
if cursor.seek_key_exact(id, None) {
let val = unsafe { cursor.val().downcast::<Tup1<SqlString>>().0.clone() };
cache.insert(id.clone(), (val.clone(), false));
Some(val)
} else {
None
}
})
})
})
}
#[derive(Eq, PartialEq, Hash)]
struct InterneStringCacheKey;
impl TypedMapKey<LocalStoreMarker> for InterneStringCacheKey {
type Value = Arc<InternedStringCache>;
}
#[derive(Eq, PartialEq, Hash)]
struct InternedStringSpineKey;
impl TypedMapKey<LocalStoreMarker> for InternedStringSpineKey {
type Value = Arc<RwLock<InternedStringSpineSnapshot>>;
}
pub fn build_string_interner(
interned_strings: Stream<RootCircuit, OrdZSet<Tup1<SqlString>>>,
cache_capacity_bytes: Option<u64>,
) {
INTERNED_STRING_CACHE.with_borrow_mut(|cache| {
*cache = Runtime::runtime()
.unwrap()
.local_store()
.entry(InterneStringCacheKey)
.or_insert_with(|| Arc::new(init_interned_string_cache(cache_capacity_bytes)))
.clone()
});
INTERNED_STRING_BY_ID.with_borrow_mut(|by_id| {
*by_id = Runtime::runtime()
.unwrap()
.local_store()
.entry(InternedStringSpineKey)
.or_insert_with(|| Arc::new(RwLock::new(empty_by_id())))
.clone()
});
let by_id = interned_strings
.map_index(|s| (intern_string(&s.0), s.clone()))
.shard()
.set_persistent_id(Some("feldera_interned_string_by_id"))
.integrate_trace()
.inner()
.delay_trace();
let exchange = new_exchange_operators(
Some(Location::caller()),
empty_by_id,
move |spine: SpineSnapshot<_>, outputs| {
let mut locations = WorkerLocations::new();
match locations.next().unwrap() {
WorkerLocation::Local => outputs.push(Mailbox::Plain(spine.clone())),
WorkerLocation::Remote => outputs.push(Mailbox::Tx(to_bytes(&spine).unwrap())),
};
for location in locations {
match location {
WorkerLocation::Local => outputs.push(Mailbox::Plain(empty_by_id())),
WorkerLocation::Remote => {
outputs.push(Mailbox::Tx(to_bytes(&empty_by_id()).unwrap()))
}
}
}
},
|data| aligned_deserialize(&data[..]),
|snapshot, remote_snapshot| {
if Runtime::worker_index() == 0 {
snapshot.extend(remote_snapshot);
}
},
ExchangeActivity::AllSteps,
);
let by_id = match exchange {
Some((sender, receiver)) => interned_strings
.circuit()
.add_exchange(sender, receiver, &by_id),
None => by_id,
};
let interner_stream = by_id.apply(|spine| {
let current_step = CURRENT_STEP.with_borrow_mut(|step| {
*step += 1;
*step
});
if Runtime::worker_index() == 0 {
INTERNED_STRING_BY_ID.with_borrow(|by_id| *by_id.write().unwrap() = spine.clone());
}
PINNED_STRINGS.with_borrow_mut(|pinned| {
let first_pinned =
pinned.partition_point(|(_, (_, step))| *step <= current_step.saturating_sub(2));
for (id, val) in pinned.drain(..first_pinned) {
let _ = INTERNED_STRING_CACHE
.with_borrow(|cache| cache.replace(id, (val.0, false), true));
}
pinned.shrink_to(pinned.len() * 2);
});
});
interner_stream
.circuit()
.add_preprocessor(interner_stream.local_node_id());
}
#[cfg(test)]
mod interned_string_test {
use crate::string_interner::InterneStringCacheKey;
use crate::{SqlString, build_string_interner};
use crate::{intern_string, unintern_string};
use dbsp::circuit::{CircuitConfig, CircuitStorageConfig, StorageConfig, StorageOptions};
use dbsp::trace::{BatchReader, Cursor};
use dbsp::typed_batch::IndexedZSetReader;
use dbsp::utils::{Tup1, Tup2};
use dbsp::{
DBSPHandle, OrdZSet, OutputHandle, Runtime, ZSetHandle, typed_batch::SpineSnapshot,
};
use std::path::Path;
use uuid::Uuid;
#[allow(clippy::type_complexity)]
pub fn interner_test_circuit(
path: &Path,
checkpoint: Option<Uuid>,
) -> (
DBSPHandle,
(
ZSetHandle<SqlString>,
ZSetHandle<SqlString>,
OutputHandle<SpineSnapshot<OrdZSet<SqlString>>>,
),
) {
let (circuit, handles) = Runtime::init_circuit(
CircuitConfig::with_workers(8).with_storage(Some(
CircuitStorageConfig::for_config(
StorageConfig {
path: path.display().to_string(),
cache: Default::default(),
},
StorageOptions::default(),
)
.unwrap()
.with_init_checkpoint(checkpoint),
)),
move |circuit| {
let (input_strings, hinput_strings) = circuit.add_input_zset::<SqlString>();
input_strings.set_persistent_mir_id("input_strings");
let (queries, hqueries) = circuit.add_input_zset::<SqlString>();
queries.set_persistent_mir_id("queries");
build_string_interner(input_strings.map(|s| Tup1(s.clone())), Some(10_000));
let output_strings = input_strings
.map_index(|s| (s.clone(), intern_string(s)))
.join(
&queries.map_index(|q| (q.clone(), ())),
|_, intern_string_id, _| unintern_string(intern_string_id).unwrap(),
);
Ok((hinput_strings, hqueries, output_strings.accumulate_output()))
},
)
.unwrap();
(circuit, handles)
}
fn query<'a, I>(
circuit: &mut DBSPHandle,
hqueries: &ZSetHandle<SqlString>,
houtput_strings: &OutputHandle<SpineSnapshot<OrdZSet<SqlString>>>,
queries: I,
) where
I: IntoIterator<Item = &'a str>,
{
let mut queries = queries.into_iter().map(SqlString::from).collect::<Vec<_>>();
let mut tuples = queries
.iter()
.map(|s| Tup2(s.clone(), 1))
.collect::<Vec<_>>();
hqueries.append(&mut tuples);
circuit.transaction().unwrap();
let output = houtput_strings.concat().consolidate();
let mut output = output.iter().map(|(s, _, _)| s.clone()).collect::<Vec<_>>();
output.sort();
queries.sort();
assert_eq!(output, queries);
let mut tuples = queries
.iter()
.map(|s| Tup2(s.clone(), -1))
.collect::<Vec<_>>();
hqueries.append(&mut tuples);
circuit.transaction().unwrap();
}
#[test]
fn test_interner_basic() {
let path = tempfile::tempdir().unwrap().keep();
let (mut circuit, (hinput_strings, hqueries, houtput_strings)) =
interner_test_circuit(path.as_path(), None);
hinput_strings.push(SqlString::from("1"), 1);
query(&mut circuit, &hqueries, &houtput_strings, ["1"]);
hinput_strings.push(SqlString::from("2"), 1);
query(&mut circuit, &hqueries, &houtput_strings, ["2"]);
hinput_strings.push(SqlString::from("3"), 1);
query(&mut circuit, &hqueries, &houtput_strings, ["3"]);
hinput_strings.push(SqlString::from("4"), 1);
query(&mut circuit, &hqueries, &houtput_strings, ["4"]);
hinput_strings.push(SqlString::from("5"), 1);
query(&mut circuit, &hqueries, &houtput_strings, ["5"]);
query(
&mut circuit,
&hqueries,
&houtput_strings,
["1", "2", "3", "4", "5"],
);
let checkpoint = circuit.checkpoint().run().unwrap();
circuit.kill().unwrap();
let (mut circuit, (hinput_strings, hqueries, houtput_strings)) =
interner_test_circuit(path.as_path(), Some(checkpoint.uuid));
query(
&mut circuit,
&hqueries,
&houtput_strings,
["1", "2", "3", "4", "5"],
);
hinput_strings.push(SqlString::from("6"), 1);
hinput_strings.push(SqlString::from("7"), 1);
query(
&mut circuit,
&hqueries,
&houtput_strings,
["1", "2", "3", "4", "5", "6", "7"],
);
circuit.kill().unwrap();
}
#[test]
fn test_interner_small() {
let path = tempfile::tempdir().unwrap().keep();
let (mut circuit, (hinput_strings, hqueries, houtput_strings)) =
interner_test_circuit(path.as_path(), None);
for _batch in 0..1_000 {
let values = (0..1_000).map(|i| i.to_string()).collect::<Vec<_>>();
let mut chunk = values
.iter()
.map(|i| Tup2(SqlString::from(i.as_str()), 1))
.collect::<Vec<_>>();
hinput_strings.append(&mut chunk);
query(
&mut circuit,
&hqueries,
&houtput_strings,
values.iter().map(String::as_str),
);
}
let checkpoint = circuit.checkpoint().run().unwrap();
circuit.kill().unwrap();
let (mut circuit, (_hinput_strings, hqueries, houtput_strings)) =
interner_test_circuit(path.as_path(), Some(checkpoint.uuid));
query(
&mut circuit,
&hqueries,
&houtput_strings,
(0..1_000)
.map(|i| i.to_string())
.collect::<Vec<_>>()
.iter()
.map(|s| s.as_str()),
);
assert!(
circuit
.runtime()
.local_store()
.get(&InterneStringCacheKey)
.unwrap()
.len()
< 1000
);
assert!(
circuit
.runtime()
.local_store()
.get(&InterneStringCacheKey)
.unwrap()
.misses()
<= 10000
);
circuit.kill().unwrap();
}
#[test]
fn test_interner_deletions() {
let path = tempfile::tempdir().unwrap().keep();
let (mut circuit, (hinput_strings, hqueries, houtput_strings)) =
interner_test_circuit(path.as_path(), None);
for batch in 0..1_000 {
let values = (batch * 100..(batch + 1) * 100)
.map(|i| i.to_string())
.collect::<Vec<_>>();
let mut chunk = values
.iter()
.map(|i| Tup2(SqlString::from(i.as_str()), 1))
.collect::<Vec<_>>();
hinput_strings.append(&mut chunk);
query(
&mut circuit,
&hqueries,
&houtput_strings,
values.iter().map(String::as_str),
);
}
let checkpoint = circuit.checkpoint().run().unwrap();
circuit.kill().unwrap();
let (mut circuit, (hinput_strings, _hqueries, _houtput_strings)) =
interner_test_circuit(path.as_path(), Some(checkpoint.uuid));
for batch in 0..1_000 {
let values = (batch * 100..(batch + 1) * 100)
.map(|i| i.to_string())
.collect::<Vec<_>>();
let mut chunk = values
.iter()
.map(|i| Tup2(SqlString::from(i.as_str()), -1))
.collect::<Vec<_>>();
hinput_strings.append(&mut chunk);
circuit.transaction().unwrap();
}
circuit.transaction().unwrap();
circuit.transaction().unwrap();
super::INTERNED_STRING_BY_ID.with_borrow(|by_id| {
let mut cursor = by_id.read().unwrap().cursor();
while cursor.key_valid() {
while cursor.val_valid() {
assert_eq!(**cursor.weight(), 0);
cursor.step_val();
}
cursor.step_key();
}
circuit.kill().unwrap();
})
}
#[test]
fn test_interner_bulk() {
let path = tempfile::tempdir().unwrap().keep();
let (mut circuit, (hinput_strings, hqueries, houtput_strings)) =
interner_test_circuit(path.as_path(), None);
for batch in 0..100 {
let values = (batch * 1_000..(batch + 1) * 1_000)
.map(|i| i.to_string())
.collect::<Vec<_>>();
let mut chunk = values
.iter()
.map(|i| Tup2(SqlString::from(i.as_str()), 1))
.collect::<Vec<_>>();
hinput_strings.append(&mut chunk);
query(
&mut circuit,
&hqueries,
&houtput_strings,
values.iter().map(String::as_str),
);
}
let checkpoint = circuit.checkpoint().run().unwrap();
circuit.kill().unwrap();
let (mut circuit, (_hinput_strings, hqueries, houtput_strings)) =
interner_test_circuit(path.as_path(), Some(checkpoint.uuid));
query(
&mut circuit,
&hqueries,
&houtput_strings,
(0..100 * 1_000)
.map(|i| i.to_string())
.collect::<Vec<_>>()
.iter()
.map(|s| s.as_str()),
);
circuit.kill().unwrap();
}
}