use crate::error::TauqError;
use serde::Serialize;
#[derive(Debug)]
pub struct BatchEncoder<T: Serialize> {
items: Vec<T>,
parallel_threshold: usize,
}
impl<T: Serialize> BatchEncoder<T> {
pub fn new() -> Self {
Self {
items: Vec::new(),
parallel_threshold: 100,
}
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
items: Vec::with_capacity(capacity),
parallel_threshold: 100,
}
}
pub fn with_parallel_threshold(mut self, threshold: usize) -> Self {
self.parallel_threshold = threshold;
self
}
pub fn add_record(&mut self, item: T) {
self.items.push(item);
}
pub fn add_records(&mut self, items: impl IntoIterator<Item = T>) {
self.items.extend(items);
}
pub fn encode(&self) -> Result<Vec<u8>, TauqError> {
if self.items.is_empty() {
return super::to_bytes(&Vec::<T>::new());
}
super::to_bytes(&self.items)
}
pub fn len(&self) -> usize {
self.items.len()
}
pub fn is_empty(&self) -> bool {
self.items.is_empty()
}
pub fn clear(&mut self) {
self.items.clear();
}
pub fn items(&self) -> &[T] {
&self.items
}
}
impl<T: Serialize> Default for BatchEncoder<T> {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct BatchEncodingStats {
pub record_count: usize,
pub bytes: usize,
pub parallelized: bool,
pub bytes_per_record: f64,
}
impl BatchEncodingStats {
pub fn new(record_count: usize, bytes: usize, parallelized: bool) -> Self {
let bytes_per_record = if record_count > 0 {
bytes as f64 / record_count as f64
} else {
0.0
};
Self {
record_count,
bytes,
parallelized,
bytes_per_record,
}
}
pub fn compression_ratio_vs_json(&self, json_size: usize) -> f64 {
if json_size > 0 {
(self.bytes as f64 / json_size as f64) * 100.0
} else {
0.0
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_batch_encoder_empty() {
let encoder: BatchEncoder<i32> = BatchEncoder::new();
assert_eq!(encoder.len(), 0);
assert!(encoder.is_empty());
}
#[test]
fn test_batch_encoder_add_records() {
let mut encoder = BatchEncoder::new();
encoder.add_record(1);
encoder.add_record(2);
encoder.add_record(3);
assert_eq!(encoder.len(), 3);
assert!(!encoder.is_empty());
}
#[test]
fn test_batch_encoder_add_multiple() {
let mut encoder = BatchEncoder::new();
encoder.add_records(vec![1, 2, 3, 4, 5]);
assert_eq!(encoder.len(), 5);
}
#[test]
fn test_batch_encoder_clear() {
let mut encoder = BatchEncoder::new();
encoder.add_records(vec![1, 2, 3]);
encoder.clear();
assert_eq!(encoder.len(), 0);
assert!(encoder.is_empty());
}
#[test]
fn test_batch_encoder_encode_small() {
let mut encoder = BatchEncoder::new();
encoder.add_records(vec![1u32, 2u32, 3u32]);
let bytes = encoder.encode().unwrap();
assert!(!bytes.is_empty());
}
#[test]
fn test_batch_encoder_encode_large() {
let mut encoder = BatchEncoder::new();
for i in 0..1000 {
encoder.add_record(i);
}
let bytes = encoder.encode().unwrap();
assert!(!bytes.is_empty());
}
#[test]
fn test_batch_encoding_stats() {
let stats = BatchEncodingStats::new(1000, 5000, true);
assert_eq!(stats.record_count, 1000);
assert_eq!(stats.bytes, 5000);
assert!(stats.parallelized);
assert!(stats.bytes_per_record > 4.9 && stats.bytes_per_record < 5.1);
}
#[test]
fn test_compression_ratio_calculation() {
let stats = BatchEncodingStats::new(100, 1000, false);
let ratio = stats.compression_ratio_vs_json(10000);
assert_eq!(ratio, 10.0); }
}