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 {
21        /// Original Arrow data type.
22        data_type: DataType,
23        /// Byte length of the persisted backing data.
24        disk_bytes: usize,
25    },
26    /// Cached batch on disk as Arrow array.
27    DiskArrow {
28        /// Original Arrow data type.
29        data_type: DataType,
30        /// Byte length of the persisted backing data.
31        disk_bytes: usize,
32    },
33}
34
35impl CacheEntry {
36    /// Construct a cached batch stored as an in-memory Arrow array.
37    pub fn memory_arrow(array: ArrayRef) -> Self {
38        Self::MemoryArrow(array)
39    }
40
41    /// Construct a cached batch stored as an in-memory Liquid array.
42    pub fn memory_liquid(array: LiquidArrayRef) -> Self {
43        Self::MemoryLiquid(array)
44    }
45
46    /// Construct a cached batch stored as an in-memory squeezed Liquid array.
47    pub fn memory_squeezed_liquid(array: LiquidSqueezedArrayRef) -> Self {
48        Self::MemorySqueezedLiquid(array)
49    }
50
51    /// Construct a cached batch stored on disk as Liquid bytes.
52    pub fn disk_liquid(data_type: DataType, disk_bytes: usize) -> Self {
53        Self::DiskLiquid {
54            data_type,
55            disk_bytes,
56        }
57    }
58
59    /// Construct a cached batch stored on disk as Arrow bytes.
60    pub fn disk_arrow(data_type: DataType, disk_bytes: usize) -> Self {
61        Self::DiskArrow {
62            data_type,
63            disk_bytes,
64        }
65    }
66
67    /// Memory usage reported by the underlying representation.
68    pub fn memory_usage_bytes(&self) -> usize {
69        match self {
70            Self::MemoryArrow(array) => array.get_array_memory_size(),
71            Self::MemoryLiquid(array) => array.get_array_memory_size(),
72            Self::MemorySqueezedLiquid(array) => array.get_array_memory_size(),
73            Self::DiskLiquid { .. } | Self::DiskArrow { .. } => 0,
74        }
75    }
76
77    /// Reference count (if any) of the backing storage.
78    pub fn reference_count(&self) -> usize {
79        match self {
80            Self::MemoryArrow(array) => Arc::strong_count(array),
81            Self::MemoryLiquid(array) => Arc::strong_count(array),
82            Self::MemorySqueezedLiquid(array) => Arc::strong_count(array),
83            Self::DiskLiquid { .. } | Self::DiskArrow { .. } => 0,
84        }
85    }
86}
87
88impl Display for CacheEntry {
89    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90        match self {
91            Self::MemoryArrow(_) => write!(f, "MemoryArrow"),
92            Self::MemoryLiquid(_) => write!(f, "MemoryLiquid"),
93            Self::MemorySqueezedLiquid(_) => write!(f, "MemorySqueezedLiquid"),
94            Self::DiskLiquid { .. } => write!(f, "DiskLiquid"),
95            Self::DiskArrow { .. } => write!(f, "DiskArrow"),
96        }
97    }
98}
99
100/// The type of the cached batch.
101#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
102pub enum CachedBatchType {
103    /// Cached batch in memory as Arrow array.
104    MemoryArrow,
105    /// Cached batch in memory as liquid array.
106    MemoryLiquid,
107    /// Cached batch in memory as squeezed liquid array.
108    MemorySqueezedLiquid,
109    /// Cached batch on disk as liquid array.
110    DiskLiquid,
111    /// Cached batch on disk as Arrow array.
112    DiskArrow,
113}
114
115impl From<&CacheEntry> for CachedBatchType {
116    fn from(batch: &CacheEntry) -> Self {
117        match batch {
118            CacheEntry::MemoryArrow(_) => Self::MemoryArrow,
119            CacheEntry::MemoryLiquid(_) => Self::MemoryLiquid,
120            CacheEntry::MemorySqueezedLiquid(_) => Self::MemorySqueezedLiquid,
121            CacheEntry::DiskLiquid { .. } => Self::DiskLiquid,
122            CacheEntry::DiskArrow { .. } => Self::DiskArrow,
123        }
124    }
125}