alopex_core/columnar/
memory.rs

1//! インメモリカラムナーストア。
2
3use std::collections::HashMap;
4use std::fs::{rename, OpenOptions};
5use std::io::Write;
6use std::path::Path;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::RwLock;
9
10use crate::columnar::error::{ColumnarError, Result};
11use crate::columnar::kvs_bridge::ColumnarKvsBridge;
12use crate::columnar::segment_v2::{
13    ColumnSegmentV2, InMemorySegmentSource, RecordBatch, Schema, SegmentReaderV2,
14};
15use crate::storage::format::{AlopexFileWriter, ColumnarSectionWriter};
16
17/// インメモリにカラムナーセグメントを保持する。
18pub struct InMemorySegmentStore {
19    segments: RwLock<HashMap<(u32, u64), ColumnSegmentV2>>,
20    memory_usage: AtomicU64,
21    memory_limit: Option<u64>,
22}
23
24impl InMemorySegmentStore {
25    /// `memory_limit` バイト上限付きで生成する。`None` は無制限。
26    pub fn new(memory_limit: Option<u64>) -> Self {
27        Self {
28            segments: RwLock::new(HashMap::new()),
29            memory_usage: AtomicU64::new(0),
30            memory_limit,
31        }
32    }
33
34    /// セグメントを書き込み、割り当てたセグメントIDを返す。
35    pub fn write_segment(&self, table_id: u32, segment: ColumnSegmentV2) -> Result<u64> {
36        let size = segment.data.len() as u64;
37        let current = self.memory_usage.load(Ordering::Relaxed);
38        let requested = current.saturating_add(size);
39        if let Some(limit) = self.memory_limit {
40            if requested > limit {
41                return Err(ColumnarError::MemoryLimitExceeded {
42                    limit: std::cmp::min(limit, usize::MAX as u64) as usize,
43                    requested: std::cmp::min(requested, usize::MAX as u64) as usize,
44                });
45            }
46        }
47
48        let mut guard = self.segments.write().unwrap();
49        let next_id = guard
50            .keys()
51            .filter(|(tid, _)| *tid == table_id)
52            .map(|(_, sid)| *sid)
53            .max()
54            .map(|id| id.saturating_add(1))
55            .unwrap_or(0);
56        guard.insert((table_id, next_id), segment);
57        drop(guard);
58        self.memory_usage.store(requested, Ordering::Relaxed);
59        Ok(next_id)
60    }
61
62    /// セグメントを読み取り、指定カラムを返す。
63    pub fn read_segment(
64        &self,
65        table_id: u32,
66        segment_id: u64,
67        columns: &[usize],
68    ) -> Result<Vec<RecordBatch>> {
69        let guard = self.segments.read().unwrap();
70        let segment = guard
71            .get(&(table_id, segment_id))
72            .ok_or(ColumnarError::NotFound)?
73            .clone();
74        drop(guard);
75        let reader =
76            SegmentReaderV2::open(Box::new(InMemorySegmentSource::new(segment.data.clone())))?;
77        reader.read_columns(columns)
78    }
79
80    /// 現在のメモリ使用量（バイト）を返す。
81    pub fn memory_usage(&self) -> u64 {
82        self.memory_usage.load(Ordering::Relaxed)
83    }
84
85    /// セグメントをファイルへフラッシュする（fsync + rename）。
86    pub fn flush_to_segment_file<P: AsRef<Path>>(
87        &self,
88        table_id: u32,
89        segment_id: u64,
90        path: P,
91    ) -> Result<()> {
92        let guard = self.segments.read().unwrap();
93        let segment = guard
94            .get(&(table_id, segment_id))
95            .ok_or(ColumnarError::NotFound)?
96            .clone();
97        drop(guard);
98
99        let tmp_path = path.as_ref().with_extension("tmp");
100        let mut file = OpenOptions::new()
101            .create(true)
102            .write(true)
103            .truncate(true)
104            .open(&tmp_path)?;
105        file.write_all(&segment.data)?;
106        file.sync_all()?;
107        rename(&tmp_path, path)?;
108        Ok(())
109    }
110
111    /// KVS へフラッシュする（ColumnarKvsBridge 経由）。
112    pub fn flush_to_kvs(
113        &self,
114        table_id: u32,
115        segment_id: u64,
116        bridge: &ColumnarKvsBridge,
117    ) -> Result<u64> {
118        let guard = self.segments.read().unwrap();
119        let segment = guard
120            .get(&(table_id, segment_id))
121            .ok_or(ColumnarError::NotFound)?
122            .clone();
123        drop(guard);
124        bridge.write_segment(table_id, &segment)
125    }
126
127    /// `.alopex` ファイルへフラッシュする。
128    pub fn flush_to_alopex(
129        &self,
130        table_id: u32,
131        segment_id: u64,
132        writer: &mut AlopexFileWriter,
133    ) -> Result<u32> {
134        let guard = self.segments.read().unwrap();
135        let segment = guard
136            .get(&(table_id, segment_id))
137            .ok_or(ColumnarError::NotFound)?
138            .clone();
139        drop(guard);
140
141        ColumnarSectionWriter::write_section(writer, &segment)
142            .map_err(|e| ColumnarError::InvalidFormat(e.to_string()))
143    }
144
145    /// カラム数を返す（メタデータから取得）。
146    pub fn column_count(&self, table_id: u32, segment_id: u64) -> Result<usize> {
147        let guard = self.segments.read().unwrap();
148        let segment = guard
149            .get(&(table_id, segment_id))
150            .ok_or(ColumnarError::NotFound)?;
151        Ok(segment.meta.schema.column_count())
152    }
153
154    /// スキーマを返す（メタデータから取得）。
155    pub fn schema(&self, table_id: u32, segment_id: u64) -> Result<Schema> {
156        let guard = self.segments.read().unwrap();
157        let segment = guard
158            .get(&(table_id, segment_id))
159            .ok_or(ColumnarError::NotFound)?;
160        Ok(segment.meta.schema.clone())
161    }
162
163    /// すべてのセグメント (table_id, segment_id) を返す。
164    pub fn list_segments(&self) -> Vec<(u32, u64)> {
165        let guard = self.segments.read().unwrap();
166        guard.keys().cloned().collect()
167    }
168}
169
170#[cfg(all(test, not(target_arch = "wasm32")))]
171mod tests {
172    use super::*;
173    use crate::columnar::encoding::{Column, LogicalType};
174    use crate::columnar::segment_v2::{ColumnSchema, Schema, SegmentWriterV2};
175    use crate::kv::memory::MemoryKV;
176    use crate::storage::format::{
177        AlopexFileReader, FileFlags, FileReader, FileSource, FileVersion, SectionType,
178    };
179    use tempfile::tempdir;
180
181    fn make_segment() -> ColumnSegmentV2 {
182        let schema = Schema {
183            columns: vec![
184                ColumnSchema {
185                    name: "id".into(),
186                    logical_type: LogicalType::Int64,
187                    nullable: false,
188                    fixed_len: None,
189                },
190                ColumnSchema {
191                    name: "val".into(),
192                    logical_type: LogicalType::Int64,
193                    nullable: false,
194                    fixed_len: None,
195                },
196            ],
197        };
198        let batch = RecordBatch::new(
199            schema,
200            vec![Column::Int64(vec![1, 2]), Column::Int64(vec![10, 20])],
201            vec![None, None],
202        );
203        let mut writer = SegmentWriterV2::new(Default::default());
204        writer.write_batch(batch).unwrap();
205        writer.finish().unwrap()
206    }
207
208    #[test]
209    fn test_memory_limit_enforcement() {
210        let store = InMemorySegmentStore::new(Some(1));
211        let segment = make_segment();
212        let err = store.write_segment(1, segment).unwrap_err();
213        assert!(matches!(err, ColumnarError::MemoryLimitExceeded { .. }));
214    }
215
216    #[test]
217    fn test_write_and_read_segment() {
218        let store = InMemorySegmentStore::new(None);
219        let id = store.write_segment(2, make_segment()).unwrap();
220        let batches = store.read_segment(2, id, &[0, 1]).unwrap();
221        assert_eq!(batches[0].num_rows(), 2);
222    }
223
224    #[test]
225    fn test_flush_to_segment_file() {
226        let dir = tempdir().unwrap();
227        let path = dir.path().join("seg.bin");
228        let store = InMemorySegmentStore::new(None);
229        let id = store.write_segment(1, make_segment()).unwrap();
230        store
231            .flush_to_segment_file(1, id, &path)
232            .expect("flush succeeds");
233        let bytes = std::fs::read(&path).unwrap();
234        assert!(!bytes.is_empty());
235    }
236
237    #[test]
238    fn test_flush_to_kvs() {
239        let store = InMemorySegmentStore::new(None);
240        let seg = make_segment();
241        let id = store.write_segment(4, seg.clone()).unwrap();
242        let kv = MemoryKV::new();
243        let bridge = ColumnarKvsBridge::new(std::sync::Arc::new(crate::kv::AnyKV::Memory(kv)));
244        let new_id = store.flush_to_kvs(4, id, &bridge).unwrap();
245        assert_eq!(new_id, 0);
246        let batches = bridge.read_segment(4, new_id, &[0, 1]).unwrap();
247        assert_eq!(batches[0].num_rows(), 2);
248    }
249
250    #[test]
251    fn test_flush_to_alopex() {
252        let dir = tempdir().unwrap();
253        let path = dir.path().join("test.alopex");
254        let mut writer =
255            AlopexFileWriter::new(path.clone(), FileVersion::CURRENT, FileFlags(0)).unwrap();
256        let store = InMemorySegmentStore::new(None);
257        let id = store.write_segment(1, make_segment()).unwrap();
258        let section_id = store.flush_to_alopex(1, id, &mut writer).unwrap();
259        assert_eq!(section_id, 0);
260        writer.finalize().unwrap();
261
262        let reader =
263            AlopexFileReader::open(FileSource::Path(path)).expect("alopex file should open");
264        let entry = reader
265            .section_index()
266            .find_by_id(section_id)
267            .expect("entry exists");
268        assert_eq!(entry.section_type, SectionType::ColumnarSegment);
269    }
270}
alopex_core/columnar/memory.rs

alopex_core/columnar/
memory.rs
