use std::{fmt::Display, sync::Arc};
use arrow::array::ArrayRef;
use arrow_schema::DataType;
use crate::liquid_array::{LiquidArrayRef, LiquidSqueezedArrayRef};
#[derive(Debug, Clone)]
pub enum CacheEntry {
MemoryArrow(ArrayRef),
MemoryLiquid(LiquidArrayRef),
MemorySqueezedLiquid(LiquidSqueezedArrayRef),
DiskLiquid(DataType),
DiskArrow(DataType),
}
impl CacheEntry {
pub fn memory_arrow(array: ArrayRef) -> Self {
Self::MemoryArrow(array)
}
pub fn memory_liquid(array: LiquidArrayRef) -> Self {
Self::MemoryLiquid(array)
}
pub fn memory_squeezed_liquid(array: LiquidSqueezedArrayRef) -> Self {
Self::MemorySqueezedLiquid(array)
}
pub fn disk_liquid(data_type: DataType) -> Self {
Self::DiskLiquid(data_type)
}
pub fn disk_arrow(data_type: DataType) -> Self {
Self::DiskArrow(data_type)
}
pub fn memory_usage_bytes(&self) -> usize {
match self {
Self::MemoryArrow(array) => array.get_array_memory_size(),
Self::MemoryLiquid(array) => array.get_array_memory_size(),
Self::MemorySqueezedLiquid(array) => array.get_array_memory_size(),
Self::DiskLiquid(_) | Self::DiskArrow(_) => 0,
}
}
pub fn reference_count(&self) -> usize {
match self {
Self::MemoryArrow(array) => Arc::strong_count(array),
Self::MemoryLiquid(array) => Arc::strong_count(array),
Self::MemorySqueezedLiquid(array) => Arc::strong_count(array),
Self::DiskLiquid(_) | Self::DiskArrow(_) => 0,
}
}
}
impl Display for CacheEntry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MemoryArrow(_) => write!(f, "MemoryArrow"),
Self::MemoryLiquid(_) => write!(f, "MemoryLiquid"),
Self::MemorySqueezedLiquid(_) => write!(f, "MemorySqueezedLiquid"),
Self::DiskLiquid(_) => write!(f, "DiskLiquid"),
Self::DiskArrow(_) => write!(f, "DiskArrow"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
pub enum CachedBatchType {
MemoryArrow,
MemoryLiquid,
MemorySqueezedLiquid,
DiskLiquid,
DiskArrow,
}
impl From<&CacheEntry> for CachedBatchType {
fn from(batch: &CacheEntry) -> Self {
match batch {
CacheEntry::MemoryArrow(_) => Self::MemoryArrow,
CacheEntry::MemoryLiquid(_) => Self::MemoryLiquid,
CacheEntry::MemorySqueezedLiquid(_) => Self::MemorySqueezedLiquid,
CacheEntry::DiskLiquid(_) => Self::DiskLiquid,
CacheEntry::DiskArrow(_) => Self::DiskArrow,
}
}
}