use std::marker::PhantomData;
use std::sync::Arc;
use crypto::digest::Digest;
use crypto::md5::Md5;
use hex;
use crate::ingest::alias_method_fork::*;
use rand::Rng;
use rand::SeedableRng;
use rand::distributions::{Alphanumeric, Standard};
use rand;
use crate::mem_store::column::*;
use crate::mem_store::column_builder::{ColumnBuilder, IntColBuilder, StringColBuilder};
use crate::scheduler::inner_locustdb::InnerLocustDB;
pub trait ColumnGenerator: Sync + Send {
fn generate(&self, length: usize, name: &str, seed: u64) -> Arc<Column>;
}
pub fn int_markov_chain(
elements: Vec<i64>,
transition_probabilities: Vec<Vec<f64>>) -> Box<dyn ColumnGenerator> {
Box::new(MarkovChain {
elem: elements.into_iter().map(Some).collect(),
p_transition: transition_probabilities,
s: PhantomData::<IntColBuilder>,
})
}
pub fn int_uniform(low: i64, high: i64) -> Box<dyn ColumnGenerator> {
Box::new(UniformInteger { low, high })
}
pub fn splayed(offset: i64, coefficient: i64) -> Box<dyn ColumnGenerator> {
Box::new(Splayed { offset, coefficient })
}
pub fn int_weighted(values: Vec<i64>, weights: Vec<f64>) -> Box<dyn ColumnGenerator> {
Box::new(Weighted {
elem: values.into_iter().map(Some).collect(),
weights,
s: PhantomData::<IntColBuilder>,
})
}
pub fn incrementing_int() -> Box<dyn ColumnGenerator> {
Box::new(IncrementingInteger)
}
pub fn string_markov_chain(
elements: Vec<String>,
transition_probabilities: Vec<Vec<f64>>) -> Box<dyn ColumnGenerator> {
Box::new(MarkovChain {
elem: elements,
p_transition: transition_probabilities,
s: PhantomData::<StringColBuilder>,
})
}
pub fn string_weighted(values: Vec<String>, weights: Vec<f64>) -> Box<dyn ColumnGenerator> {
Box::new(Weighted {
elem: values,
weights,
s: PhantomData::<StringColBuilder>,
})
}
pub fn random_hex_string(length: usize) -> Box<dyn ColumnGenerator> {
Box::new(HexString { length })
}
pub fn random_string(min_length: usize, max_length: usize) -> Box<dyn ColumnGenerator> {
Box::new(RandomString {
min_length,
max_length,
})
}
pub fn partition_sparse(
null_probability: f64,
generator: Box<dyn ColumnGenerator>) -> Box<dyn ColumnGenerator> {
Box::new(PartitionSparse {
null_probability,
generator,
})
}
#[derive(Clone)]
struct MarkovChain<T, S> {
elem: Vec<T>,
p_transition: Vec<Vec<f64>>,
s: PhantomData<S>,
}
unsafe impl<T: Send, S: ColumnBuilder<T>> Send for MarkovChain<T, S> {}
unsafe impl<T: Sync, S: ColumnBuilder<T>> Sync for MarkovChain<T, S> {}
impl<T: Sync + Send, S: ColumnBuilder<T>> ColumnGenerator for MarkovChain<T, S> {
fn generate(&self, length: usize, name: &str, seed: u64) -> Arc<Column> {
let mut rng = seeded_rng(seed);
let mut builder = S::default();
let mut state = rng.gen_range(0, self.elem.len());
let p = self.p_transition.iter()
.map(|p| new_alias_table(p).unwrap())
.collect::<Vec<_>>();
let mut alias_method = AliasMethod::new(rng);
for _ in 0..length {
state = alias_method.random(&p[state]);
builder.push(&self.elem[state]);
}
builder.finalize(name, None)
}
}
#[derive(Clone)]
struct Weighted<T, S> {
elem: Vec<T>,
weights: Vec<f64>,
s: PhantomData<S>,
}
unsafe impl<T: Send, S: ColumnBuilder<T>> Send for Weighted<T, S> {}
unsafe impl<T: Sync, S: ColumnBuilder<T>> Sync for Weighted<T, S> {}
impl<T: Sync + Send, S: ColumnBuilder<T>> ColumnGenerator for Weighted<T, S> {
fn generate(&self, length: usize, name: &str, seed: u64) -> Arc<Column> {
let rng = seeded_rng(seed);
let mut builder = S::default();
let p = new_alias_table(&self.weights).unwrap();
let mut alias_method = AliasMethod::new(rng);
for _ in 0..length {
let i = alias_method.random(&p);
builder.push(&self.elem[i]);
}
builder.finalize(name, None)
}
}
struct UniformInteger {
low: i64,
high: i64,
}
impl ColumnGenerator for UniformInteger {
fn generate(&self, length: usize, name: &str, seed: u64) -> Arc<Column> {
let mut rng = seeded_rng(seed);
let mut builder = IntColBuilder::default();
for _ in 0..length {
builder.push(&Some(rng.gen_range::<i64>(self.low, self.high)));
}
ColumnBuilder::<Option<i64>>::finalize(builder, name, None)
}
}
struct Splayed {
offset: i64,
coefficient: i64,
}
impl ColumnGenerator for Splayed {
fn generate(&self, length: usize, name: &str, partition: u64) -> Arc<Column> {
let mut rng = seeded_rng(partition);
let mut builder = IntColBuilder::default();
for _ in 0..length {
builder.push(&Some(rng.gen_range::<i64>(
self.offset + self.coefficient * length as i64 * partition as i64,
self.offset + self.coefficient * length as i64 * (partition as i64 + 1),
)));
}
ColumnBuilder::<Option<i64>>::finalize(builder, name, None)
}
}
struct PartitionSparse {
null_probability: f64,
generator: Box<dyn ColumnGenerator>,
}
impl ColumnGenerator for PartitionSparse {
fn generate(&self, length: usize, name: &str, seed: u64) -> Arc<Column> {
let mut rng = seeded_rng(seed);
if rng.gen::<f64>() < self.null_probability {
Arc::new(Column::null(name, length))
} else {
self.generator.generate(length, name, seed)
}
}
}
struct HexString {
length: usize,
}
impl ColumnGenerator for HexString {
fn generate(&self, length: usize, name: &str, seed: u64) -> Arc<Column> {
let mut rng = seeded_rng(seed);
let mut builder = StringColBuilder::default();
for _ in 0..length {
let bytes: Vec<u8> = rng.sample_iter(&Standard).take(self.length).collect();
builder.push(&hex::encode(&bytes));
}
ColumnBuilder::<&str>::finalize(builder, name, None)
}
}
struct RandomString {
min_length: usize,
max_length: usize,
}
impl ColumnGenerator for RandomString {
fn generate(&self, length: usize, name: &str, seed: u64) -> Arc<Column> {
let mut rng = seeded_rng(seed);
let mut builder = StringColBuilder::default();
for _ in 0..length {
let len = rng.gen_range(self.min_length, self.max_length + 1);
let string: String = rng.sample_iter::<char, _>(&Alphanumeric).take(len).collect();
builder.push(&string);
}
ColumnBuilder::<&str>::finalize(builder, name, None)
}
}
struct IncrementingInteger;
impl ColumnGenerator for IncrementingInteger {
fn generate(&self, length: usize, name: &str, seed: u64) -> Arc<Column> {
let mut builder = IntColBuilder::default();
for i in seed as i64 * length as i64..length as i64 * (seed as i64 + 1) {
builder.push(&Some(i));
}
builder.finalize(name, None)
}
}
pub struct GenTable {
pub name: String,
pub partitions: usize,
pub partition_size: usize,
pub columns: Vec<(String, Box<dyn ColumnGenerator>)>,
}
impl GenTable {
pub fn gen(&self, db: &InnerLocustDB, partition_number: u64) {
let partition = self.columns
.iter()
.map(|(name, c)| c.generate(self.partition_size, &name, partition_number))
.collect();
db.store_partition(&self.name, partition);
}
}
fn seeded_rng(seed: u64) -> rand::XorShiftRng {
let mut seed_bytes = [0u8; 16];
let mut hasher = Md5::new();
hasher.input(&seed.to_ne_bytes());
hasher.result(&mut seed_bytes);
rand::XorShiftRng::from_seed(seed_bytes)
}