Skip to main content

liquid_cache/cache/
cached_batch.rs

1//! Cached batch types.
2
3use std::{fmt::Display, sync::Arc};
4
5use arrow::array::ArrayRef;
6use arrow_schema::DataType;
7
8use crate::liquid_array::{LiquidArrayRef, LiquidSqueezedArrayRef};
9
10/// A cached entry storing data in various formats.
11#[derive(Debug, Clone)]
12pub enum CacheEntry {
13    /// Cached batch in memory as Arrow array.
14    MemoryArrow(ArrayRef),
15    /// Cached batch in memory as liquid array.
16    MemoryLiquid(LiquidArrayRef),
17    /// Cached batch in memory as squeezed liquid array.
18    MemorySqueezedLiquid(LiquidSqueezedArrayRef),
19    /// Cached batch on disk as liquid array.
20    DiskLiquid(DataType),
21    /// Cached batch on disk as Arrow array.
22    DiskArrow(DataType),
23}
24
25impl CacheEntry {
26    /// Construct a cached batch stored as an in-memory Arrow array.
27    pub fn memory_arrow(array: ArrayRef) -> Self {
28        Self::MemoryArrow(array)
29    }
30
31    /// Construct a cached batch stored as an in-memory Liquid array.
32    pub fn memory_liquid(array: LiquidArrayRef) -> Self {
33        Self::MemoryLiquid(array)
34    }
35
36    /// Construct a cached batch stored as an in-memory squeezed Liquid array.
37    pub fn memory_squeezed_liquid(array: LiquidSqueezedArrayRef) -> Self {
38        Self::MemorySqueezedLiquid(array)
39    }
40
41    /// Construct a cached batch stored on disk as Liquid bytes.
42    pub fn disk_liquid(data_type: DataType) -> Self {
43        Self::DiskLiquid(data_type)
44    }
45
46    /// Construct a cached batch stored on disk as Arrow bytes.
47    pub fn disk_arrow(data_type: DataType) -> Self {
48        Self::DiskArrow(data_type)
49    }
50
51    /// Memory usage reported by the underlying representation.
52    pub fn memory_usage_bytes(&self) -> usize {
53        match self {
54            Self::MemoryArrow(array) => array.get_array_memory_size(),
55            Self::MemoryLiquid(array) => array.get_array_memory_size(),
56            Self::MemorySqueezedLiquid(array) => array.get_array_memory_size(),
57            Self::DiskLiquid(_) | Self::DiskArrow(_) => 0,
58        }
59    }
60
61    /// Reference count (if any) of the backing storage.
62    pub fn reference_count(&self) -> usize {
63        match self {
64            Self::MemoryArrow(array) => Arc::strong_count(array),
65            Self::MemoryLiquid(array) => Arc::strong_count(array),
66            Self::MemorySqueezedLiquid(array) => Arc::strong_count(array),
67            Self::DiskLiquid(_) | Self::DiskArrow(_) => 0,
68        }
69    }
70}
71
72impl Display for CacheEntry {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        match self {
75            Self::MemoryArrow(_) => write!(f, "MemoryArrow"),
76            Self::MemoryLiquid(_) => write!(f, "MemoryLiquid"),
77            Self::MemorySqueezedLiquid(_) => write!(f, "MemorySqueezedLiquid"),
78            Self::DiskLiquid(_) => write!(f, "DiskLiquid"),
79            Self::DiskArrow(_) => write!(f, "DiskArrow"),
80        }
81    }
82}
83
84/// The type of the cached batch.
85#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
86pub enum CachedBatchType {
87    /// Cached batch in memory as Arrow array.
88    MemoryArrow,
89    /// Cached batch in memory as liquid array.
90    MemoryLiquid,
91    /// Cached batch in memory as squeezed liquid array.
92    MemorySqueezedLiquid,
93    /// Cached batch on disk as liquid array.
94    DiskLiquid,
95    /// Cached batch on disk as Arrow array.
96    DiskArrow,
97}
98
99impl From<&CacheEntry> for CachedBatchType {
100    fn from(batch: &CacheEntry) -> Self {
101        match batch {
102            CacheEntry::MemoryArrow(_) => Self::MemoryArrow,
103            CacheEntry::MemoryLiquid(_) => Self::MemoryLiquid,
104            CacheEntry::MemorySqueezedLiquid(_) => Self::MemorySqueezedLiquid,
105            CacheEntry::DiskLiquid(_) => Self::DiskLiquid,
106            CacheEntry::DiskArrow(_) => Self::DiskArrow,
107        }
108    }
109}