akar_storage/
ice_format.rs

1use crate::parquet_reader::{ParquetReaderError, ParquetStreamReader, read_parquet, stream_parquet};
2use akar_catalog::CatalogColumn;
3use akar_common::file_system::VirtualFileSystemRegistry;
4use akar_common::types::Value;
5use std::path::{Path, PathBuf};
6
7/// Layout options for IceDiskRelTable
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum IceDiskRelTableLayout {
10    Flat,
11    Csr,
12}
13
14/// The ICE (IceDisk) native disk format for relationship tables.
15/// Based on Ladybug's IceDiskRelTable implementation, which stores
16/// relational data directly in Parquet format files (`indices.parquet` and optionally `indptr.parquet`).
17pub struct IceDiskRelTable {
18    pub name: String,
19    pub layout: IceDiskRelTableLayout,
20    pub indices_file_path: PathBuf,
21    pub indptr_file_path: Option<PathBuf>,
22}
23
24/// Scan state for IceDiskRelTable that streams rows on demand.
25///
26/// Instead of loading the entire Parquet file into a `Vec<Vec<Value>>`,
27/// this state holds a streaming reader and buffers one batch at a time.
28pub struct IceDiskRelTableScanState {
29    pub stream: ParquetStreamReader,
30    pub current_batch: Vec<Vec<Value>>,
31    pub current_row: usize,
32}
33
34impl IceDiskRelTable {
35    /// Initialize a new IceDiskRelTable pointing to its Parquet files.
36    pub fn new(name: String, base_path: &Path, layout: IceDiskRelTableLayout) -> Self {
37        let indices_path = match layout {
38            IceDiskRelTableLayout::Flat => base_path.join(format!("{}.flat.parquet", name)),
39            IceDiskRelTableLayout::Csr => base_path.join(format!("{}.indices.parquet", name)),
40        };
41
42        let indptr_path = match layout {
43            IceDiskRelTableLayout::Flat => None,
44            IceDiskRelTableLayout::Csr => Some(base_path.join(format!("{}.indptr.parquet", name))),
45        };
46
47        Self {
48            name,
49            layout,
50            indices_file_path: indices_path,
51            indptr_file_path: indptr_path,
52        }
53    }
54
55    /// Scan the indices parquet file and return a streaming scan state.
56    pub fn scan_indices(
57        &self,
58        vfs: &VirtualFileSystemRegistry,
59        columns: &[CatalogColumn],
60    ) -> Result<IceDiskRelTableScanState, ParquetReaderError> {
61        let stream = stream_parquet(self.indices_file_path.to_str().unwrap(), vfs, columns)?;
62        Ok(IceDiskRelTableScanState {
63            stream,
64            current_batch: Vec::new(),
65            current_row: 0,
66        })
67    }
68
69    /// Scan the indptr parquet file (if using CSR layout).
70    pub fn scan_indptr(
71        &self,
72        vfs: &VirtualFileSystemRegistry,
73        columns: &[CatalogColumn],
74    ) -> Result<Vec<Vec<Value>>, ParquetReaderError> {
75        if let Some(path) = &self.indptr_file_path {
76            read_parquet(path.to_str().unwrap(), vfs, columns)
77        } else {
78            Ok(vec![])
79        }
80    }
81}
82
83impl IceDiskRelTableScanState {
84    /// Read the next row from the streaming parquet reader.
85    ///
86    /// Rows are pulled from the underlying `ParquetStreamReader` one batch
87    /// at a time, avoiding materialization of the entire dataset in memory.
88    pub fn next_row(&mut self) -> Option<&Vec<Value>> {
89        // Advance to next batch if current one is exhausted
90        if self.current_row >= self.current_batch.len() {
91            match self.stream.next() {
92                Some(Ok(batch)) => {
93                    self.current_batch = batch;
94                    self.current_row = 0;
95                }
96                Some(Err(_)) | None => {
97                    self.current_batch = Vec::new();
98                    self.current_row = 0;
99                    return None;
100                }
101            }
102        }
103
104        if self.current_row < self.current_batch.len() {
105            let row = &self.current_batch[self.current_row];
106            self.current_row += 1;
107            Some(row)
108        } else {
109            None
110        }
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_ice_disk_paths() {
120        let base = Path::new("/tmp/akar");
121
122        let flat_table = IceDiskRelTable::new("knows".into(), base, IceDiskRelTableLayout::Flat);
123        assert_eq!(flat_table.indices_file_path, base.join("knows.flat.parquet"));
124        assert!(flat_table.indptr_file_path.is_none());
125
126        let csr_table = IceDiskRelTable::new("study_at".into(), base, IceDiskRelTableLayout::Csr);
127        assert_eq!(csr_table.indices_file_path, base.join("study_at.indices.parquet"));
128        assert_eq!(csr_table.indptr_file_path, Some(base.join("study_at.indptr.parquet")));
129    }
130}
akar_storage/ice_format.rs

akar_storage/
ice_format.rs
