use std::sync::Arc;
use fnv::FnvHashMap;
use arrow::array::{Array, ArrayRef, Float64Array, Int64Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use crate::aggregate::AggFn;
use crate::error::{QueryError, Result};
#[allow(clippy::cast_precision_loss)]
fn extract_f64(col: &dyn Array, index: usize) -> Option<f64> {
crate::extract_f64(col, index)
}
#[allow(clippy::cast_precision_loss)]
pub fn downsample(
batch: &RecordBatch,
value_column: &str,
bucket: &chronix_core::TimeBucket,
function: &AggFn,
memory_tracker: Option<&crate::memory::MemoryTracker>,
) -> Result<RecordBatch> {
let schema = batch.schema();
let time_idx = schema
.index_of(chronix_core::TIME_COLUMN)
.map_err(|_| QueryError::Validation("'timestamp' column not found".into()))?;
let val_idx = schema
.index_of(value_column)
.map_err(|_| QueryError::FieldNotFound {
measurement: String::new(),
field: value_column.to_string(),
})?;
let timestamps = batch
.column(time_idx)
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| QueryError::Validation("'timestamp' column is not Int64".into()))?;
let value_col = batch.column(val_idx).as_ref();
match value_col.data_type() {
DataType::Float64 | DataType::Int64 | DataType::UInt64 | DataType::Decimal128(_, _) => {}
dt => {
return Err(QueryError::Validation(format!(
"'{value_column}' column has unsupported type {dt} \
(expected Float64, Int64, UInt64, or Decimal128)"
)));
}
}
let mut buckets: FnvHashMap<i64, Vec<usize>> = FnvHashMap::default();
for i in 0..timestamps.len() {
if timestamps.is_null(i) {
continue;
}
let ts = timestamps.value(i);
let bucket_start = bucket.start_of(ts);
buckets.entry(bucket_start).or_default().push(i);
}
let mut sorted_keys: Vec<i64> = buckets.keys().copied().collect();
sorted_keys.sort_unstable();
let mut out_times = Vec::with_capacity(buckets.len());
let mut out_values = Vec::with_capacity(buckets.len());
for bucket_start in &sorted_keys {
let indices = &buckets[bucket_start];
let raw_pairs: Vec<(i64, f64)> = indices
.iter()
.filter_map(|&i| {
let val = extract_f64(value_col, i)?;
Some((timestamps.value(i), val))
})
.collect();
if raw_pairs.is_empty() {
continue;
}
let agg_value = match function {
AggFn::Count => raw_pairs.iter().filter(|(_, v)| !v.is_nan()).count() as f64,
_ => {
let clean: Vec<(i64, f64)> =
raw_pairs.into_iter().filter(|(_, v)| !v.is_nan()).collect();
if clean.is_empty() {
continue;
}
match function {
AggFn::Sum => clean.iter().map(|(_, v)| v).sum(),
AggFn::Min => clean.iter().map(|(_, v)| *v).fold(f64::INFINITY, f64::min),
AggFn::Max => clean
.iter()
.map(|(_, v)| *v)
.fold(f64::NEG_INFINITY, f64::max),
AggFn::Avg => {
let sum: f64 = clean.iter().map(|(_, v)| v).sum();
sum / clean.len() as f64
}
AggFn::First => clean
.iter()
.min_by_key(|(ts, _)| *ts)
.map(|(_, v)| *v)
.expect("clean is non-empty"),
AggFn::Last => clean
.iter()
.max_by_key(|(ts, _)| *ts)
.map(|(_, v)| *v)
.expect("clean is non-empty"),
AggFn::Count => {
return Err(QueryError::Validation(
"Count aggregation reached inner branch unexpectedly".into(),
));
}
}
}
};
out_times.push(*bucket_start);
out_values.push(agg_value);
}
let out_schema = Arc::new(Schema::new(vec![
Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
Field::new(value_column, DataType::Float64, false),
]));
let result = RecordBatch::try_new(
out_schema,
vec![
Arc::new(Int64Array::from(out_times)) as ArrayRef,
Arc::new(Float64Array::from(out_values)) as ArrayRef,
],
)?;
if let Some(tracker) = memory_tracker {
tracker.try_allocate(result.get_array_memory_size())?;
}
Ok(result)
}
pub struct StreamingDownsampler {
value_column: String,
bucket: chronix_core::TimeBucket,
function: AggFn,
open_bucket: Option<i64>,
acc: BucketAcc,
pending: Vec<(i64, f64)>,
chunk_rows: usize,
}
#[derive(Default)]
struct BucketAcc {
count: u64,
sum: f64,
min: f64,
max: f64,
first_ts: i64,
first_val: f64,
last_ts: i64,
last_val: f64,
}
impl BucketAcc {
fn reset(&mut self) {
self.count = 0;
self.sum = 0.0;
self.min = f64::INFINITY;
self.max = f64::NEG_INFINITY;
self.first_ts = i64::MAX;
self.first_val = 0.0;
self.last_ts = i64::MIN;
self.last_val = 0.0;
}
fn update(&mut self, ts: i64, v: f64) {
if v.is_nan() {
return;
}
self.count += 1;
self.sum += v;
if v < self.min {
self.min = v;
}
if v > self.max {
self.max = v;
}
if ts < self.first_ts {
self.first_ts = ts;
self.first_val = v;
}
if ts >= self.last_ts {
self.last_ts = ts;
self.last_val = v;
}
}
#[allow(clippy::cast_precision_loss)]
fn finalize(&self, function: AggFn) -> Option<f64> {
if self.count == 0 {
return None;
}
Some(match function {
AggFn::Count => self.count as f64,
AggFn::Sum => self.sum,
AggFn::Min => self.min,
AggFn::Max => self.max,
AggFn::Avg => self.sum / self.count as f64,
AggFn::First => self.first_val,
AggFn::Last => self.last_val,
})
}
}
impl StreamingDownsampler {
const DEFAULT_CHUNK_ROWS: usize = 65_536;
pub fn new(
value_column: &str,
bucket: chronix_core::TimeBucket,
function: AggFn,
) -> Result<Self> {
let mut acc = BucketAcc::default();
acc.reset();
Ok(Self {
value_column: value_column.to_string(),
bucket,
function,
open_bucket: None,
acc,
pending: Vec::new(),
chunk_rows: Self::DEFAULT_CHUNK_ROWS,
})
}
pub fn push(&mut self, batch: &RecordBatch) -> Result<Vec<RecordBatch>> {
if batch.num_rows() == 0 {
return Ok(Vec::new());
}
let schema = batch.schema();
let time_idx = schema
.index_of(chronix_core::TIME_COLUMN)
.map_err(|_| QueryError::Validation("'timestamp' column not found".into()))?;
let val_idx =
schema
.index_of(&self.value_column)
.map_err(|_| QueryError::FieldNotFound {
measurement: String::new(),
field: self.value_column.clone(),
})?;
let timestamps = batch
.column(time_idx)
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| QueryError::Validation("'timestamp' column is not Int64".into()))?;
let value_col = batch.column(val_idx).as_ref();
match value_col.data_type() {
DataType::Float64 | DataType::Int64 | DataType::UInt64 | DataType::Decimal128(_, _) => {
}
dt => {
return Err(QueryError::Validation(format!(
"'{}' column has unsupported type {dt} \
(expected Float64, Int64, UInt64, or Decimal128)",
self.value_column
)));
}
}
for i in 0..timestamps.len() {
if timestamps.is_null(i) {
continue;
}
let ts = timestamps.value(i);
let bucket = self.bucket.start_of(ts);
match self.open_bucket {
Some(open) if bucket == open => {}
Some(open) if bucket < open => {
return Err(QueryError::Validation(format!(
"downsample input is not ascending by timestamp: bucket {bucket} \
arrived after {open} was closed"
)));
}
Some(_) => self.close_open_bucket(),
None => {}
}
self.open_bucket = Some(bucket);
if let Some(v) = extract_f64(value_col, i) {
self.acc.update(ts, v);
}
}
self.drain(false)
}
pub fn finish(mut self) -> Result<Vec<RecordBatch>> {
if self.open_bucket.is_some() {
self.close_open_bucket();
}
self.drain(true)
}
fn close_open_bucket(&mut self) {
if let (Some(start), Some(v)) = (self.open_bucket, self.acc.finalize(self.function)) {
self.pending.push((start, v));
}
self.acc.reset();
self.open_bucket = None;
}
fn drain(&mut self, flush: bool) -> Result<Vec<RecordBatch>> {
let mut out = Vec::new();
while self.pending.len() >= self.chunk_rows || (flush && !self.pending.is_empty()) {
let take = self.chunk_rows.min(self.pending.len());
let rows: Vec<(i64, f64)> = self.pending.drain(..take).collect();
out.push(build_batch(&self.value_column, &rows)?);
}
Ok(out)
}
}
fn build_batch(value_column: &str, rows: &[(i64, f64)]) -> Result<RecordBatch> {
let schema = Arc::new(Schema::new(vec![
Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
Field::new(value_column, DataType::Float64, false),
]));
let times: Vec<i64> = rows.iter().map(|(t, _)| *t).collect();
let values: Vec<f64> = rows.iter().map(|(_, v)| *v).collect();
RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(times)) as ArrayRef,
Arc::new(Float64Array::from(values)) as ArrayRef,
],
)
.map_err(Into::into)
}
#[cfg(test)]
mod tests {
use super::*;
use arrow::array::UInt64Array;
use chronix_core::TimeBucket;
fn make_batch(times: Vec<i64>, values: Vec<f64>) -> RecordBatch {
let schema = Arc::new(Schema::new(vec![
Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
Field::new("value", DataType::Float64, false),
]));
RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(times)),
Arc::new(Float64Array::from(values)),
],
)
.unwrap()
}
#[test]
fn downsample_avg_10s_buckets() {
let batch = make_batch(
vec![0, 5, 10, 15, 20, 25],
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
);
let result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Avg,
None,
)
.unwrap();
assert_eq!(result.num_rows(), 3);
let times = result
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let values = result
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert_eq!(times.value(0), 0);
assert!((values.value(0) - 1.5).abs() < f64::EPSILON);
assert_eq!(times.value(1), 10);
assert!((values.value(1) - 3.5).abs() < f64::EPSILON);
assert_eq!(times.value(2), 20);
assert!((values.value(2) - 5.5).abs() < f64::EPSILON);
}
#[test]
fn downsample_sum() {
let batch = make_batch(vec![0, 5, 10, 15], vec![10.0, 20.0, 30.0, 40.0]);
let result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Sum,
None,
)
.unwrap();
assert_eq!(result.num_rows(), 2);
let values = result
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert!((values.value(0) - 30.0).abs() < f64::EPSILON);
assert!((values.value(1) - 70.0).abs() < f64::EPSILON);
}
#[test]
fn downsample_count() {
let batch = make_batch(vec![0, 1, 2, 10, 20, 21, 22, 23], vec![1.0; 8]);
let result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Count,
None,
)
.unwrap();
assert_eq!(result.num_rows(), 3);
let values = result
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert!((values.value(0) - 3.0).abs() < f64::EPSILON); assert!((values.value(1) - 1.0).abs() < f64::EPSILON); assert!((values.value(2) - 4.0).abs() < f64::EPSILON); }
#[test]
fn downsample_min_max() {
let batch = make_batch(vec![0, 5, 10, 15], vec![3.0, 1.0, 4.0, 2.0]);
let min_result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Min,
None,
)
.unwrap();
let max_result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Max,
None,
)
.unwrap();
let min_vals = min_result
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
let max_vals = max_result
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert!((min_vals.value(0) - 1.0).abs() < f64::EPSILON); assert!((min_vals.value(1) - 2.0).abs() < f64::EPSILON); assert!((max_vals.value(0) - 3.0).abs() < f64::EPSILON); assert!((max_vals.value(1) - 4.0).abs() < f64::EPSILON); }
#[test]
fn downsample_first_last() {
let batch = make_batch(vec![0, 5, 10, 15], vec![10.0, 20.0, 30.0, 40.0]);
let first = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::First,
None,
)
.unwrap();
let last = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Last,
None,
)
.unwrap();
let first_vals = first
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
let last_vals = last
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert!((first_vals.value(0) - 10.0).abs() < f64::EPSILON);
assert!((first_vals.value(1) - 30.0).abs() < f64::EPSILON);
assert!((last_vals.value(0) - 20.0).abs() < f64::EPSILON);
assert!((last_vals.value(1) - 40.0).abs() < f64::EPSILON);
}
#[test]
fn empty_batch() {
let batch = make_batch(vec![], vec![]);
let result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Avg,
None,
)
.unwrap();
assert_eq!(result.num_rows(), 0);
}
#[test]
fn downsample_i64_column() {
let schema = Arc::new(Schema::new(vec![
Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
Field::new("value", DataType::Int64, false),
]));
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(vec![0, 5, 10, 15])) as ArrayRef,
Arc::new(Int64Array::from(vec![10, 20, 30, 40])) as ArrayRef,
],
)
.unwrap();
let result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Sum,
None,
)
.unwrap();
assert_eq!(result.num_rows(), 2);
let values = result
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert!((values.value(0) - 30.0).abs() < f64::EPSILON);
assert!((values.value(1) - 70.0).abs() < f64::EPSILON);
}
#[test]
fn downsample_u64_column() {
let schema = Arc::new(Schema::new(vec![
Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
Field::new("value", DataType::UInt64, false),
]));
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(vec![0, 5, 10, 15])) as ArrayRef,
Arc::new(UInt64Array::from(vec![100u64, 200, 300, 400])) as ArrayRef,
],
)
.unwrap();
let result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Avg,
None,
)
.unwrap();
assert_eq!(result.num_rows(), 2);
let values = result
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert!((values.value(0) - 150.0).abs() < f64::EPSILON);
assert!((values.value(1) - 350.0).abs() < f64::EPSILON);
}
#[test]
fn downsample_unsupported_type_fails() {
let schema = Arc::new(Schema::new(vec![
Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
Field::new("value", DataType::Boolean, false),
]));
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(vec![0i64])) as ArrayRef,
Arc::new(arrow::array::BooleanArray::from(vec![true])) as ArrayRef,
],
)
.unwrap();
let err = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Avg,
None,
)
.unwrap_err();
let msg = format!("{err}");
assert!(
msg.contains("unsupported"),
"expected unsupported type error, got: {msg}"
);
}
#[test]
fn downsample_respects_memory_tracker() {
let batch = make_batch(vec![0, 5, 10, 15], vec![1.0, 2.0, 3.0, 4.0]);
let tracker = crate::memory::MemoryTracker::new(1_000_000);
let result = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Avg,
Some(&tracker),
)
.unwrap();
assert!(tracker.allocated() > 0);
assert_eq!(result.num_rows(), 2);
}
#[test]
fn downsample_exceeds_memory_budget() {
let batch = make_batch(vec![0, 5, 10, 15], vec![1.0, 2.0, 3.0, 4.0]);
let tracker = crate::memory::MemoryTracker::new(1);
let err = downsample(
&batch,
"value",
&TimeBucket::fixed_ns(10),
&AggFn::Avg,
Some(&tracker),
)
.unwrap_err();
assert!(matches!(
err,
crate::error::QueryError::QueryMemoryExceeded { .. }
));
}
}