use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use bytes::Bytes;
use crate::value_borrowed::{BorrowedRowData, BorrowedValue};
use crate::Value;
#[derive(Debug, Default)]
pub struct ZeroCopyStats {
zero_copy_hits: AtomicU64,
fallback_copies: AtomicU64,
peak_rss_bytes: AtomicU64,
allocation_count: AtomicU64,
decimal_bytes_misses: AtomicU64,
json_bytes_misses: AtomicU64,
bytes_ref_misses: AtomicU64,
datetime_int_misses: AtomicU64,
array_ref_hits: AtomicU64,
object_ref_hits: AtomicU64,
array_heap_fallbacks: AtomicU64,
object_heap_fallbacks: AtomicU64,
}
impl ZeroCopyStats {
pub fn new() -> Self {
Self::default()
}
pub fn record_zero_copy_hit(&self) {
self.zero_copy_hits.fetch_add(1, Ordering::Relaxed);
}
pub fn record_fallback_copy(&self) {
self.fallback_copies.fetch_add(1, Ordering::Relaxed);
}
pub fn zero_copy_hits(&self) -> u64 {
self.zero_copy_hits.load(Ordering::Relaxed)
}
pub fn fallback_copies(&self) -> u64 {
self.fallback_copies.load(Ordering::Relaxed)
}
pub fn total(&self) -> u64 {
self.zero_copy_hits() + self.fallback_copies()
}
pub fn hit_rate(&self) -> f64 {
let total = self.total();
if total == 0 {
0.0
} else {
self.zero_copy_hits() as f64 / total as f64
}
}
pub fn record_rss(&self, rss_bytes: u64) {
let mut current = self.peak_rss_bytes.load(Ordering::Relaxed);
while rss_bytes > current {
match self.peak_rss_bytes.compare_exchange_weak(
current,
rss_bytes,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(actual) => current = actual,
}
}
}
pub fn record_allocation(&self) {
self.allocation_count.fetch_add(1, Ordering::Relaxed);
}
pub fn peak_rss_bytes(&self) -> u64 {
self.peak_rss_bytes.load(Ordering::Relaxed)
}
pub fn allocation_count(&self) -> u64 {
self.allocation_count.load(Ordering::Relaxed)
}
pub fn rss_reduction_pct(&self, baseline_rss: u64) -> f64 {
if baseline_rss == 0 {
return 0.0;
}
let current = self.peak_rss_bytes();
if current >= baseline_rss {
0.0
} else {
(baseline_rss - current) as f64 / baseline_rss as f64 * 100.0
}
}
pub fn record_decimal_bytes_miss(&self) {
self.decimal_bytes_misses.fetch_add(1, Ordering::Relaxed);
}
pub fn record_json_bytes_miss(&self) {
self.json_bytes_misses.fetch_add(1, Ordering::Relaxed);
}
pub fn record_bytes_ref_miss(&self) {
self.bytes_ref_misses.fetch_add(1, Ordering::Relaxed);
}
pub fn record_datetime_int_miss(&self) {
self.datetime_int_misses.fetch_add(1, Ordering::Relaxed);
}
pub fn decimal_bytes_misses(&self) -> u64 {
self.decimal_bytes_misses.load(Ordering::Relaxed)
}
pub fn json_bytes_misses(&self) -> u64 {
self.json_bytes_misses.load(Ordering::Relaxed)
}
pub fn bytes_ref_misses(&self) -> u64 {
self.bytes_ref_misses.load(Ordering::Relaxed)
}
pub fn datetime_int_misses(&self) -> u64 {
self.datetime_int_misses.load(Ordering::Relaxed)
}
pub fn total_type_misses(&self) -> u64 {
self.decimal_bytes_misses()
+ self.json_bytes_misses()
+ self.bytes_ref_misses()
+ self.datetime_int_misses()
}
pub fn record_array_ref_hit(&self) {
self.array_ref_hits.fetch_add(1, Ordering::Relaxed);
}
pub fn record_object_ref_hit(&self) {
self.object_ref_hits.fetch_add(1, Ordering::Relaxed);
}
pub fn record_array_heap_fallback(&self) {
self.array_heap_fallbacks.fetch_add(1, Ordering::Relaxed);
}
pub fn record_object_heap_fallback(&self) {
self.object_heap_fallbacks.fetch_add(1, Ordering::Relaxed);
}
pub fn array_ref_hits(&self) -> u64 {
self.array_ref_hits.load(Ordering::Relaxed)
}
pub fn object_ref_hits(&self) -> u64 {
self.object_ref_hits.load(Ordering::Relaxed)
}
pub fn array_heap_fallbacks(&self) -> u64 {
self.array_heap_fallbacks.load(Ordering::Relaxed)
}
pub fn object_heap_fallbacks(&self) -> u64 {
self.object_heap_fallbacks.load(Ordering::Relaxed)
}
pub fn heap_reduction_rate(&self) -> f64 {
let ref_hits = self.array_ref_hits() + self.object_ref_hits();
let heap_fallbacks = self.array_heap_fallbacks() + self.object_heap_fallbacks();
let total = ref_hits + heap_fallbacks;
if total == 0 {
0.0
} else {
ref_hits as f64 / total as f64
}
}
}
impl Clone for ZeroCopyStats {
fn clone(&self) -> Self {
Self {
zero_copy_hits: AtomicU64::new(self.zero_copy_hits()),
fallback_copies: AtomicU64::new(self.fallback_copies()),
peak_rss_bytes: AtomicU64::new(self.peak_rss_bytes()),
allocation_count: AtomicU64::new(self.allocation_count()),
decimal_bytes_misses: AtomicU64::new(self.decimal_bytes_misses()),
json_bytes_misses: AtomicU64::new(self.json_bytes_misses()),
bytes_ref_misses: AtomicU64::new(self.bytes_ref_misses()),
datetime_int_misses: AtomicU64::new(self.datetime_int_misses()),
array_ref_hits: AtomicU64::new(self.array_ref_hits()),
object_ref_hits: AtomicU64::new(self.object_ref_hits()),
array_heap_fallbacks: AtomicU64::new(self.array_heap_fallbacks()),
object_heap_fallbacks: AtomicU64::new(self.object_heap_fallbacks()),
}
}
}
#[derive(Debug, Clone)]
pub struct ZeroCopyRow {
pub columns: Arc<Vec<String>>,
pub data: Bytes,
pub offsets: Vec<(usize, usize)>,
}
impl ZeroCopyRow {
pub fn get(&self, col: &str) -> Option<&[u8]> {
let idx = self.columns.iter().position(|c| c == col)?;
let (start, len) = self.offsets[idx];
Some(&self.data[start..start + len])
}
pub fn column_count(&self) -> usize {
self.columns.len()
}
}
pub struct ZeroCopyRowStream {
rows: Vec<ZeroCopyRow>,
index: usize,
stats: Arc<ZeroCopyStats>,
}
impl ZeroCopyRowStream {
pub fn next_row(&mut self) -> Option<&ZeroCopyRow> {
if self.index < self.rows.len() {
let row = &self.rows[self.index];
self.index += 1;
Some(row)
} else {
None
}
}
pub fn remaining(&self) -> usize {
self.rows.len() - self.index
}
pub fn total_rows(&self) -> usize {
self.rows.len()
}
pub fn stats(&self) -> &ZeroCopyStats {
&self.stats
}
}
impl Iterator for ZeroCopyRowStream {
type Item = ZeroCopyRow;
fn next(&mut self) -> Option<Self::Item> {
if self.index < self.rows.len() {
let row = self.rows[self.index].clone();
self.index += 1;
Some(row)
} else {
None
}
}
}
pub struct ZeroCopyPipeline {
stats: Arc<ZeroCopyStats>,
}
impl ZeroCopyPipeline {
pub fn new() -> Self {
Self {
stats: Arc::new(ZeroCopyStats::new()),
}
}
pub fn stream_rows(
&self,
rows: Vec<std::collections::HashMap<String, Value>>,
columns: &[String],
) -> ZeroCopyRowStream {
let col_count = columns.len();
let col_arc: Arc<Vec<String>> = Arc::new(columns.to_vec());
let mut zero_copy_rows = Vec::with_capacity(rows.len());
for row in rows {
let mut buffer = Vec::new();
let mut offsets = Vec::with_capacity(col_count);
for col in columns.iter() {
let value = row.get(col);
let start = buffer.len();
match value {
Some(Value::String(s)) => {
buffer.extend_from_slice(s.as_bytes());
self.stats.record_zero_copy_hit();
}
Some(Value::Bytes(b)) => {
buffer.extend_from_slice(b);
self.stats.record_zero_copy_hit();
}
Some(v) => {
let formatted = format!("{}", v);
buffer.extend_from_slice(formatted.as_bytes());
self.stats.record_fallback_copy();
}
None => {
self.stats.record_fallback_copy();
}
}
let len = buffer.len() - start;
offsets.push((start, len));
}
zero_copy_rows.push(ZeroCopyRow {
columns: col_arc.clone(),
data: Bytes::from(buffer),
offsets,
});
}
ZeroCopyRowStream {
rows: zero_copy_rows,
index: 0,
stats: self.stats.clone(),
}
}
pub fn stream_borrowed(
&self,
rows: Vec<BorrowedRowData<'_>>,
columns: &[String],
) -> ZeroCopyRowStream {
let col_arc: Arc<Vec<String>> = Arc::new(columns.to_vec());
let mut zero_copy_rows = Vec::with_capacity(rows.len());
for row in rows {
let mut buffer = Vec::new();
let mut offsets = Vec::with_capacity(columns.len());
for col in columns.iter() {
let start = buffer.len();
if let Some(val) = row.get(col) {
match val {
BorrowedValue::String(s) => {
buffer.extend_from_slice(s.as_bytes());
}
BorrowedValue::Bytes(b) => {
buffer.extend_from_slice(b.as_ref());
}
_ => {
let owned = val.to_owned_value();
let formatted = format!("{}", owned);
buffer.extend_from_slice(formatted.as_bytes());
self.stats.record_fallback_copy();
}
}
self.stats.record_zero_copy_hit();
} else {
self.stats.record_fallback_copy();
}
let len = buffer.len() - start;
offsets.push((start, len));
}
zero_copy_rows.push(ZeroCopyRow {
columns: col_arc.clone(),
data: Bytes::from(buffer),
offsets,
});
}
ZeroCopyRowStream {
rows: zero_copy_rows,
index: 0,
stats: self.stats.clone(),
}
}
pub fn stats(&self) -> &ZeroCopyStats {
&self.stats
}
}
impl Default for ZeroCopyPipeline {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ZeroCopyTypeId {
I32,
I64,
F32,
F64,
Bool,
String,
Bytes,
}
#[derive(Debug, Clone)]
pub struct ZeroCopyTypeRegistry {
supported: std::collections::HashSet<ZeroCopyTypeId>,
}
impl Default for ZeroCopyTypeRegistry {
fn default() -> Self {
Self::with_builtins()
}
}
impl ZeroCopyTypeRegistry {
pub fn with_builtins() -> Self {
let mut supported = std::collections::HashSet::new();
supported.insert(ZeroCopyTypeId::I32);
supported.insert(ZeroCopyTypeId::I64);
supported.insert(ZeroCopyTypeId::F32);
supported.insert(ZeroCopyTypeId::F64);
supported.insert(ZeroCopyTypeId::Bool);
supported.insert(ZeroCopyTypeId::String);
supported.insert(ZeroCopyTypeId::Bytes);
Self { supported }
}
pub fn empty() -> Self {
Self {
supported: std::collections::HashSet::new(),
}
}
pub fn register(&mut self, type_id: ZeroCopyTypeId) {
self.supported.insert(type_id);
}
pub fn is_supported(&self, type_id: ZeroCopyTypeId) -> bool {
self.supported.contains(&type_id)
}
pub fn len(&self) -> usize {
self.supported.len()
}
pub fn is_empty(&self) -> bool {
self.supported.is_empty()
}
}
impl ZeroCopyPipeline {
pub fn try_parse_with_registry(
&self,
row: &std::collections::HashMap<String, Value>,
columns: &[String],
registry: &ZeroCopyTypeRegistry,
) -> Option<ZeroCopyRow> {
let col_arc: Arc<Vec<String>> = Arc::new(columns.to_vec());
let mut buffer = Vec::new();
let mut offsets = Vec::with_capacity(columns.len());
let mut all_supported = true;
for col in columns.iter() {
let start = buffer.len();
if let Some(value) = row.get(col) {
let type_id = value_to_type_id(value);
if registry.is_supported(type_id) {
match value {
Value::String(s) => buffer.extend_from_slice(s.as_bytes()),
Value::Bytes(b) => buffer.extend_from_slice(b),
_ => {
let formatted = format!("{}", value);
buffer.extend_from_slice(formatted.as_bytes());
}
}
self.stats.record_zero_copy_hit();
} else {
all_supported = false;
let formatted = format!("{}", value);
buffer.extend_from_slice(formatted.as_bytes());
self.stats.record_fallback_copy();
self.stats.record_allocation();
}
} else {
self.stats.record_fallback_copy();
self.stats.record_allocation();
}
let len = buffer.len() - start;
offsets.push((start, len));
}
if all_supported {
Some(ZeroCopyRow {
columns: col_arc,
data: Bytes::from(buffer),
offsets,
})
} else {
None
}
}
}
fn value_to_type_id(value: &Value) -> ZeroCopyTypeId {
match value {
Value::I32(_) => ZeroCopyTypeId::I32,
Value::I64(_) => ZeroCopyTypeId::I64,
Value::F32(_) => ZeroCopyTypeId::F32,
Value::F64(_) => ZeroCopyTypeId::F64,
Value::Bool(_) => ZeroCopyTypeId::Bool,
Value::String(_) => ZeroCopyTypeId::String,
Value::Bytes(_) => ZeroCopyTypeId::Bytes,
_ => ZeroCopyTypeId::String, }
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
fn make_rows(n: usize, cols: &[&str]) -> Vec<HashMap<String, Value>> {
(0..n)
.map(|i| {
let mut row = HashMap::new();
for col in cols {
row.insert(col.to_string(), Value::String(format!("val_{}_{}", i, col)));
}
row
})
.collect()
}
#[test]
fn stream_rows_preserves_count() {
let pipeline = ZeroCopyPipeline::new();
let columns = vec!["id".to_string(), "name".to_string()];
let rows = make_rows(100, &["id", "name"]);
let stream = pipeline.stream_rows(rows, &columns);
assert_eq!(stream.total_rows(), 100);
}
#[test]
fn stream_rows_data_accessible() {
let pipeline = ZeroCopyPipeline::new();
let columns = vec!["name".to_string()];
let mut rows = Vec::new();
let mut row = HashMap::new();
row.insert("name".to_string(), Value::String("Alice".into()));
rows.push(row);
let mut stream = pipeline.stream_rows(rows, &columns);
let first = stream.next().unwrap();
assert_eq!(first.get("name").unwrap(), b"Alice");
}
#[test]
fn stats_track_zero_copy_hits() {
let pipeline = ZeroCopyPipeline::new();
let columns = vec!["name".to_string()];
let rows = make_rows(10, &["name"]);
let _stream = pipeline.stream_rows(rows, &columns);
assert!(pipeline.stats().zero_copy_hits() > 0);
}
#[test]
fn stats_track_fallback_copies() {
let pipeline = ZeroCopyPipeline::new();
let columns = vec!["num".to_string()];
let mut rows = Vec::new();
let mut row = HashMap::new();
row.insert("num".to_string(), Value::I64(42));
rows.push(row);
let _stream = pipeline.stream_rows(rows, &columns);
assert!(pipeline.stats().fallback_copies() > 0);
}
#[test]
fn empty_rows_stream() {
let pipeline = ZeroCopyPipeline::new();
let columns = vec!["id".to_string()];
let stream = pipeline.stream_rows(Vec::new(), &columns);
assert_eq!(stream.total_rows(), 0);
assert_eq!(stream.remaining(), 0);
}
#[test]
fn iterator_interface() {
let pipeline = ZeroCopyPipeline::new();
let columns = vec!["id".to_string()];
let rows = make_rows(5, &["id"]);
let stream = pipeline.stream_rows(rows, &columns);
let collected: Vec<_> = stream.collect();
assert_eq!(collected.len(), 5);
}
#[test]
fn hit_rate_calculation() {
let stats = ZeroCopyStats::new();
for _ in 0..7 {
stats.record_zero_copy_hit();
}
for _ in 0..3 {
stats.record_fallback_copy();
}
assert!((stats.hit_rate() - 0.7).abs() < 0.001);
}
#[test]
fn large_resultset_streaming() {
let pipeline = ZeroCopyPipeline::new();
let columns: Vec<String> = (0..20).map(|i| format!("col_{}", i)).collect();
let rows: Vec<HashMap<String, Value>> = (0..10000)
.map(|i| {
let mut row = HashMap::new();
for j in 0..20 {
row.insert(
format!("col_{}", j),
Value::String(format!("val_{}_{}", i, j)),
);
}
row
})
.collect();
let stream = pipeline.stream_rows(rows, &columns);
assert_eq!(stream.total_rows(), 10000);
assert!(pipeline.stats().zero_copy_hits() > 0);
}
#[test]
fn test_array_ref_hits_tracking() {
let stats = ZeroCopyStats::new();
stats.record_array_ref_hit();
stats.record_array_ref_hit();
stats.record_array_ref_hit();
assert_eq!(stats.array_ref_hits(), 3);
}
#[test]
fn test_object_ref_hits_tracking() {
let stats = ZeroCopyStats::new();
stats.record_object_ref_hit();
stats.record_object_ref_hit();
assert_eq!(stats.object_ref_hits(), 2);
}
#[test]
fn test_array_heap_fallbacks_tracking() {
let stats = ZeroCopyStats::new();
stats.record_array_heap_fallback();
assert_eq!(stats.array_heap_fallbacks(), 1);
}
#[test]
fn test_object_heap_fallbacks_tracking() {
let stats = ZeroCopyStats::new();
stats.record_object_heap_fallback();
stats.record_object_heap_fallback();
assert_eq!(stats.object_heap_fallbacks(), 2);
}
#[test]
fn test_heap_reduction_rate_empty() {
let stats = ZeroCopyStats::new();
assert_eq!(stats.heap_reduction_rate(), 0.0);
}
#[test]
fn test_heap_reduction_rate_all_hits() {
let stats = ZeroCopyStats::new();
for _ in 0..8 {
stats.record_array_ref_hit();
}
for _ in 0..2 {
stats.record_object_ref_hit();
}
assert!((stats.heap_reduction_rate() - 1.0).abs() < 1e-9);
}
#[test]
fn test_heap_reduction_rate_mixed() {
let stats = ZeroCopyStats::new();
for _ in 0..8 {
stats.record_array_ref_hit();
}
for _ in 0..2 {
stats.record_array_heap_fallback();
}
let rate = stats.heap_reduction_rate();
assert!((rate - 0.8).abs() < 1e-9, "rate={}", rate);
}
#[test]
fn test_heap_reduction_rate_above_80_pct() {
let stats = ZeroCopyStats::new();
for _ in 0..85 {
stats.record_array_ref_hit();
}
for _ in 0..15 {
stats.record_array_heap_fallback();
}
assert!(stats.heap_reduction_rate() >= 0.80);
}
#[test]
fn test_zero_copy_stats_clone_preserves_v760_fields() {
let stats = ZeroCopyStats::new();
stats.record_array_ref_hit();
stats.record_object_ref_hit();
stats.record_array_heap_fallback();
stats.record_object_heap_fallback();
let cloned = stats.clone();
assert_eq!(cloned.array_ref_hits(), 1);
assert_eq!(cloned.object_ref_hits(), 1);
assert_eq!(cloned.array_heap_fallbacks(), 1);
assert_eq!(cloned.object_heap_fallbacks(), 1);
}
}