use std::collections::HashSet;
use crate::error::Result;
use crate::unified::{ColumnEncoder, EncodedBlock};
const DEFAULT_SAMPLE_SIZE: usize = 1024;
const NUM_STRATA: usize = 4;
#[derive(Debug, Clone)]
pub struct AdaptiveSelector {
pub sample_size: usize,
}
impl Default for AdaptiveSelector {
fn default() -> Self {
Self {
sample_size: DEFAULT_SAMPLE_SIZE,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FloatPattern {
Constant,
Periodic,
SlowlyVarying,
Random,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StringPattern {
LowCardinality,
HighCardinality,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IntegerPattern {
Constant,
RegularInterval,
NarrowRange,
Irregular,
}
impl AdaptiveSelector {
#[must_use]
pub fn new() -> Self {
Self::default()
}
fn stratified_sample<T: Copy>(values: &[T], budget: usize) -> Vec<T> {
let n = values.len();
if n <= budget {
return values.to_vec();
}
let per_stratum = budget / NUM_STRATA;
let stride = n / NUM_STRATA;
let mut out = Vec::with_capacity(budget);
for s in 0..NUM_STRATA {
let start = s * stride;
let end = (start + per_stratum).min(n);
out.extend_from_slice(&values[start..end]);
}
out
}
#[must_use]
pub fn analyze_floats(&self, values: &[f64]) -> FloatPattern {
if values.is_empty() {
return FloatPattern::Random;
}
let sample = Self::stratified_sample(values, self.sample_size);
let first = sample[0].to_bits();
if sample.iter().all(|v| v.to_bits() == first) {
return FloatPattern::Constant;
}
let mut total_leading: u64 = 0;
let mut prev = first;
for &v in &sample[1..] {
let bits = v.to_bits();
let xor = prev ^ bits;
total_leading += if xor == 0 {
64
} else {
xor.leading_zeros() as u64
};
prev = bits;
}
#[allow(clippy::cast_precision_loss)]
let avg_leading = total_leading as f64 / (sample.len() - 1) as f64;
if sample.len() > 128 {
let mut hits: usize = 0;
for i in 128..sample.len() {
let bits = sample[i].to_bits();
let window = &sample[i.saturating_sub(128)..i];
if window.iter().any(|v| v.to_bits() == bits) {
hits += 1;
}
}
#[allow(clippy::cast_precision_loss)]
let hit_rate = hits as f64 / (sample.len() - 128) as f64;
if hit_rate > 0.3 {
return FloatPattern::Periodic;
}
}
if avg_leading >= 16.0 {
FloatPattern::SlowlyVarying
} else {
FloatPattern::Random
}
}
#[must_use]
pub fn analyze_strings(&self, values: &[&str]) -> StringPattern {
if values.is_empty() {
return StringPattern::HighCardinality;
}
let sample = Self::stratified_sample(values, self.sample_size);
let distinct: HashSet<&str> = sample.iter().copied().collect();
if distinct.len() <= (sample.len() / 10).max(1) {
StringPattern::LowCardinality
} else {
StringPattern::HighCardinality
}
}
#[must_use]
pub fn analyze_integers(&self, values: &[i64]) -> IntegerPattern {
if values.is_empty() {
return IntegerPattern::Irregular;
}
let sample = Self::stratified_sample(values, self.sample_size);
let first = sample[0];
if sample.iter().all(|&v| v == first) {
return IntegerPattern::Constant;
}
if sample.len() >= 3 {
let interval = sample[1].wrapping_sub(sample[0]);
let is_regular = sample
.windows(2)
.all(|w| w[1].wrapping_sub(w[0]) == interval);
if is_regular {
return IntegerPattern::RegularInterval;
}
}
let min = sample.iter().copied().min().unwrap_or(0);
let max = sample.iter().copied().max().unwrap_or(0);
let range = (max as u128).wrapping_sub(min as u128);
if range <= u128::from(u16::MAX) {
return IntegerPattern::NarrowRange;
}
IntegerPattern::Irregular
}
pub fn encode_f64_adaptive(&self, values: &[f64]) -> Result<EncodedBlock> {
use crate::alp::AlpEncoder;
use crate::chimp::{Chimp128Encoder, ChimpEncoder};
use crate::gorilla::GorillaEncoder;
use crate::patas::PatasEncoder;
use crate::plain::PlainEncoder;
use crate::unified::{should_use_specialized, EncodingType};
let pattern = self.analyze_floats(values);
if matches!(pattern, FloatPattern::Constant) {
return self.encode_f64_direct(values, pattern);
}
if values.len() <= self.sample_size {
return self.encode_f64_direct(values, pattern);
}
let sample = Self::stratified_sample(values, self.sample_size);
let sample_raw_size = sample.len() * 8;
let candidates: &[EncodingType] = match pattern {
FloatPattern::Constant => &[
EncodingType::Rle,
EncodingType::Chimp,
EncodingType::Gorilla,
],
FloatPattern::Periodic => &[
EncodingType::Pco,
EncodingType::Alp,
EncodingType::Chimp128,
EncodingType::Patas,
EncodingType::Chimp,
EncodingType::Gorilla,
],
FloatPattern::SlowlyVarying => &[
EncodingType::Pco,
EncodingType::Alp,
EncodingType::Patas,
EncodingType::Chimp128,
EncodingType::Chimp,
EncodingType::Gorilla,
],
FloatPattern::Random => &[
EncodingType::Pco,
EncodingType::Alp,
EncodingType::Gorilla,
EncodingType::Patas,
],
};
let mut best: Option<(EncodingType, usize)> = None;
let mut best_any: Option<(EncodingType, usize)> = None;
for &enc in candidates {
let trial = match enc {
EncodingType::Rle => crate::rle::RleEncoder::encode_f64(&sample),
EncodingType::Pco => crate::pco::PcoEncoder::encode_f64(&sample),
EncodingType::Alp => AlpEncoder::encode(&sample),
EncodingType::Chimp128 => Chimp128Encoder::encode(&sample),
EncodingType::Patas => PatasEncoder::encode(&sample),
EncodingType::Chimp => ChimpEncoder::encode(&sample),
EncodingType::Gorilla => GorillaEncoder::encode(&sample),
_ => continue,
};
if let Ok(payload) = trial {
let len = payload.len();
if best_any.as_ref().is_none_or(|(_, bl)| len < *bl) {
best_any = Some((enc, len));
}
if should_use_specialized(&payload, sample_raw_size)
&& best.as_ref().is_none_or(|(_, best_len)| len < *best_len)
{
best = Some((enc, len));
}
}
}
let winner = best
.or(best_any)
.map_or(EncodingType::PlainF64, |(enc, _)| enc);
let payload = match winner {
EncodingType::Rle => crate::rle::RleEncoder::encode_f64(values)?,
EncodingType::Pco => crate::pco::PcoEncoder::encode_f64(values)?,
EncodingType::Alp => AlpEncoder::encode(values)?,
EncodingType::Chimp128 => Chimp128Encoder::encode(values)?,
EncodingType::Patas => PatasEncoder::encode(values)?,
EncodingType::Chimp => ChimpEncoder::encode(values)?,
EncodingType::Gorilla => GorillaEncoder::encode(values)?,
_ => PlainEncoder::encode_f64(values)?,
};
let full_raw_size = values.len() * 8;
if winner != EncodingType::PlainF64 && !should_use_specialized(&payload, full_raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::PlainF64,
payload: PlainEncoder::encode_f64(values)?,
});
}
Ok(EncodedBlock {
encoding: winner,
payload,
})
}
fn encode_f64_direct(&self, values: &[f64], pattern: FloatPattern) -> Result<EncodedBlock> {
use crate::alp::AlpEncoder;
use crate::chimp::{Chimp128Encoder, ChimpEncoder};
use crate::gorilla::GorillaEncoder;
use crate::patas::PatasEncoder;
use crate::plain::PlainEncoder;
use crate::unified::{should_use_specialized, EncodedBlock, EncodingType};
let raw_size = values.len() * 8;
match pattern {
FloatPattern::Constant => {
let rle_payload = crate::rle::RleEncoder::encode_f64(values)?;
if should_use_specialized(&rle_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Rle,
payload: rle_payload,
});
}
}
FloatPattern::Periodic | FloatPattern::SlowlyVarying => {
let pco_payload = crate::pco::PcoEncoder::encode_f64(values)?;
if should_use_specialized(&pco_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Pco,
payload: pco_payload,
});
}
let chimp128_payload = Chimp128Encoder::encode(values)?;
if should_use_specialized(&chimp128_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Chimp128,
payload: chimp128_payload,
});
}
let patas_payload = PatasEncoder::encode(values)?;
if should_use_specialized(&patas_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Patas,
payload: patas_payload,
});
}
}
FloatPattern::Random => {}
}
let pco_payload = crate::pco::PcoEncoder::encode_f64(values)?;
if should_use_specialized(&pco_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Pco,
payload: pco_payload,
});
}
let alp_payload = AlpEncoder::encode(values)?;
if should_use_specialized(&alp_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Alp,
payload: alp_payload,
});
}
let chimp_payload = ChimpEncoder::encode(values)?;
if should_use_specialized(&chimp_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Chimp,
payload: chimp_payload,
});
}
let patas_payload = PatasEncoder::encode(values)?;
if should_use_specialized(&patas_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Patas,
payload: patas_payload,
});
}
let gorilla_payload = GorillaEncoder::encode(values)?;
if should_use_specialized(&gorilla_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Gorilla,
payload: gorilla_payload,
});
}
Ok(EncodedBlock {
encoding: EncodingType::PlainF64,
payload: PlainEncoder::encode_f64(values)?,
})
}
pub fn encode_string_adaptive(&self, values: &[&str]) -> Result<EncodedBlock> {
let pattern = self.analyze_strings(values);
match pattern {
StringPattern::LowCardinality => ColumnEncoder::encode_string(values),
StringPattern::HighCardinality => {
ColumnEncoder::encode_string(values)
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::unified::EncodingType;
#[test]
fn detects_constant_floats() {
let selector = AdaptiveSelector::new();
let values = vec![42.0_f64; 500];
assert_eq!(selector.analyze_floats(&values), FloatPattern::Constant);
}
#[test]
fn detects_slowly_varying_floats() {
let selector = AdaptiveSelector::new();
let values: Vec<f64> = (0..500).map(|i| 50.0 + (i as f64) * 0.01).collect();
assert_eq!(
selector.analyze_floats(&values),
FloatPattern::SlowlyVarying
);
}
#[test]
fn detects_random_floats() {
let selector = AdaptiveSelector::new();
let values: Vec<f64> = (0..500)
.map(|i| (i as f64 * 123.456).sin() * 1e15)
.collect();
assert_eq!(selector.analyze_floats(&values), FloatPattern::Random);
}
#[test]
fn encode_f64_adaptive_selects_pco_for_slowly_varying_decimals() {
let selector = AdaptiveSelector::new();
let values: Vec<f64> = (0..500).map(|i| 50.0 + (i as f64) * 0.01).collect();
let block = selector.encode_f64_adaptive(&values).unwrap();
assert_eq!(
block.encoding,
EncodingType::Pco,
"decimal data should select pco over ALP and the XOR codecs, got {}",
block.encoding,
);
let decoded = crate::unified::ColumnDecoder::decode(&block).unwrap();
if let crate::unified::DecodedColumn::F64(vals) = decoded {
assert_eq!(vals.len(), values.len());
for (a, b) in values.iter().zip(vals.iter()) {
assert_eq!(a.to_bits(), b.to_bits());
}
} else {
panic!("Expected F64 decoded column");
}
}
#[test]
fn detects_low_cardinality_strings() {
let selector = AdaptiveSelector::new();
let values = vec!["us-east-1"; 100];
assert_eq!(
selector.analyze_strings(&values),
StringPattern::LowCardinality
);
}
#[test]
fn detects_high_cardinality_strings() {
let selector = AdaptiveSelector::new();
let owned: Vec<String> = (0..100).map(|i| format!("uuid-{i}")).collect();
let refs: Vec<&str> = owned.iter().map(String::as_str).collect();
assert_eq!(
selector.analyze_strings(&refs),
StringPattern::HighCardinality
);
}
#[test]
fn detects_constant_integers() {
let selector = AdaptiveSelector::new();
let values = vec![42_i64; 500];
assert_eq!(selector.analyze_integers(&values), IntegerPattern::Constant);
}
#[test]
fn detects_regular_interval_integers() {
let selector = AdaptiveSelector::new();
let values: Vec<i64> = (0..500).map(|i| 1000 + i * 10).collect();
assert_eq!(
selector.analyze_integers(&values),
IntegerPattern::RegularInterval
);
}
#[test]
fn detects_irregular_integers() {
let selector = AdaptiveSelector::new();
let values = vec![1, 100, 3, 999, 42, 7, 88, 100_000];
assert_eq!(
selector.analyze_integers(&values),
IntegerPattern::Irregular
);
}
#[test]
fn detects_narrow_range_integers() {
let selector = AdaptiveSelector::new();
let values = vec![200, 201, 204, 301, 302, 400, 404, 500, 502, 503];
assert_eq!(
selector.analyze_integers(&values),
IntegerPattern::NarrowRange
);
}
#[test]
fn periodic_data_dispatches_pco() {
let selector = AdaptiveSelector::new();
let base: Vec<f64> = (0..64).map(|i| 20.0 + (i as f64) * 0.5).collect();
let values: Vec<f64> = base.iter().cycle().take(5000).copied().collect();
assert_eq!(selector.analyze_floats(&values), FloatPattern::Periodic);
let block = selector.encode_f64_adaptive(&values).unwrap();
assert_eq!(
block.encoding,
EncodingType::Pco,
"Periodic data should use pco, got {}",
block.encoding,
);
let decoded = crate::unified::ColumnDecoder::decode(&block).unwrap();
if let crate::unified::DecodedColumn::F64(vals) = decoded {
assert_eq!(vals.len(), values.len());
for (a, b) in values.iter().zip(vals.iter()) {
assert_eq!(a.to_bits(), b.to_bits());
}
} else {
panic!("Expected F64 decoded column");
}
}
#[test]
fn slowly_varying_picks_a_specialised_float_codec() {
let selector = AdaptiveSelector::new();
let values: Vec<f64> = (0..5000).map(|i| 50.0 + (i as f64) * 0.001).collect();
let block = selector.encode_f64_adaptive(&values).unwrap();
assert!(
block.encoding == EncodingType::Pco
|| block.encoding == EncodingType::Alp
|| block.encoding == EncodingType::Chimp128
|| block.encoding == EncodingType::Chimp
|| block.encoding == EncodingType::Gorilla,
"Unexpected encoding: {}",
block.encoding,
);
let decoded = crate::unified::ColumnDecoder::decode(&block).unwrap();
if let crate::unified::DecodedColumn::F64(vals) = decoded {
assert_eq!(vals.len(), values.len());
} else {
panic!("Expected F64 decoded column");
}
}
}
#[cfg(test)]
mod proptests {
use super::*;
use crate::unified::{ColumnDecoder, DecodedColumn};
use proptest::prelude::*;
proptest! {
#[test]
fn roundtrip_f64_adaptive(values in proptest::collection::vec(any::<f64>(), 50..300)) {
let selector = AdaptiveSelector::new();
let block = selector.encode_f64_adaptive(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
match decoded {
DecodedColumn::F64(vals) => {
prop_assert_eq!(values.len(), vals.len());
for (a, b) in values.iter().zip(vals.iter()) {
prop_assert_eq!(a.to_bits(), b.to_bits());
}
}
other => prop_assert!(false, "Expected F64, got {other:?}"),
}
}
#[test]
fn roundtrip_string_adaptive(values in proptest::collection::vec("[a-z]{0,8}", 1..100)) {
let selector = AdaptiveSelector::new();
let refs: Vec<&str> = values.iter().map(String::as_str).collect();
let block = selector.encode_string_adaptive(&refs).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
match decoded {
DecodedColumn::String(vals) => {
prop_assert_eq!(&values, &vals);
}
other => prop_assert!(false, "Expected String, got {other:?}"),
}
}
}
}