1use crate::column::Column;
14use crate::column_chunk::{ColumnChunk, NODE_GROUP_SIZE};
15use crate::spiller::{MultiWayStreamMerge, SpillFile, Spiller};
16use crate::version_info::VersionInfo;
17use akar_common::error::StorageError;
18use akar_common::types::Value;
19use std::sync::Arc;
20
21#[derive(Debug, Clone)]
37pub struct NodeGroup {
38 pub columns: Vec<ColumnChunk>,
40 pub start_offset: u64,
42 pub num_nodes: u64,
44 pub version_info: Option<VersionInfo>,
46 spiller: Option<Arc<Spiller>>,
48 spill_files: Vec<SpillFile>,
50}
51
52impl NodeGroup {
53 pub fn new(num_columns: usize, start_offset: u64) -> Self {
59 let columns = (0..num_columns).map(|_| ColumnChunk::new()).collect();
60 Self {
61 columns,
62 start_offset,
63 num_nodes: 0,
64 version_info: None,
65 spiller: None,
66 spill_files: Vec::new(),
67 }
68 }
69
70 pub fn with_capacity(num_columns: usize, start_offset: u64, capacity: usize) -> Self {
72 let columns = (0..num_columns).map(|_| ColumnChunk::with_capacity(capacity)).collect();
73 Self {
74 columns,
75 start_offset,
76 num_nodes: 0,
77 version_info: None,
78 spiller: None,
79 spill_files: Vec::new(),
80 }
81 }
82
83 pub fn with_spiller(mut self, spiller: Arc<Spiller>) -> Self {
90 self.spiller = Some(spiller);
91 self
92 }
93
94 pub fn set_spiller(&mut self, spiller: Arc<Spiller>) {
96 self.spiller = Some(spiller);
97 }
98
99 pub fn enable_version_info(&mut self) {
102 if self.version_info.is_none() {
103 self.version_info = Some(VersionInfo::new(NODE_GROUP_SIZE));
104 }
105 }
106
107 pub fn append_row(&mut self, row: Vec<Value>) -> Result<(), StorageError> {
119 self.append_row_with_txn(row, None)
120 }
121
122 pub fn append_row_with_txn(&mut self, row: Vec<Value>, txn_id: Option<u64>) -> Result<(), StorageError> {
129 if row.len() != self.columns.len() {
130 return Err(StorageError::Page(format!(
131 "column count mismatch: expected {} values, got {}",
132 self.columns.len(),
133 row.len()
134 )));
135 }
136 if self.is_full() {
137 return Err(StorageError::Page("node group is already full".to_string()));
138 }
139
140 if let Some(ref spiller) = self.spiller
142 && !self.columns.is_empty()
143 && spiller.should_spill(&self.columns[0])
144 {
145 self.spill_and_clear()?;
146 }
147
148 for (col_idx, value) in row.into_iter().enumerate() {
149 self.columns[col_idx].append(value);
150 }
151 if let Some(ref vi) = self.version_info
153 && let Some(txn) = txn_id
154 {
155 vi.insert(txn, self.num_nodes as u32);
156 }
157 self.num_nodes += 1;
158 Ok(())
159 }
160
161 pub fn spill_and_clear(&mut self) -> Result<(), StorageError> {
166 let spiller = self
167 .spiller
168 .as_ref()
169 .ok_or_else(|| StorageError::Spiller("No spiller attached to NodeGroup".to_string()))?;
170
171 if self.is_empty() {
172 return Ok(());
173 }
174
175 let spill = spiller.spill_columns(&mut self.columns)?;
176 if let Some(sf) = spill {
177 self.spill_files.push(sf);
178 }
179 self.num_nodes = 0;
180 Ok(())
181 }
182
183 pub fn restore_spilled(&mut self) -> Result<(), StorageError> {
192 if self.spill_files.is_empty() {
193 return Ok(());
194 }
195 let spiller = self
196 .spiller
197 .clone()
198 .ok_or_else(|| StorageError::Spiller("No spiller attached to NodeGroup".to_string()))?;
199 let num_cols = self.columns.len();
200
201 let mut rows: Vec<Vec<Value>> = Vec::new();
202 let files = std::mem::take(&mut self.spill_files);
203 for sf in &files {
204 let chunks = spiller.restore_columns(sf, num_cols)?;
205 let n = chunks.first().map(|c| c.num_values()).unwrap_or(0);
206 for r in 0..n {
207 let mut row = Vec::with_capacity(num_cols);
208 for c in &chunks {
209 row.push(c.get(r).cloned().unwrap_or(Value::Null));
210 }
211 rows.push(row);
212 }
213 }
214 for row in self.scan() {
215 rows.push(row);
216 }
217
218 let mut columns: Vec<ColumnChunk> = (0..num_cols).map(|_| ColumnChunk::new()).collect();
219 for row in &rows {
220 for (ci, value) in row.iter().enumerate() {
221 columns[ci].append(value.clone());
222 }
223 }
224 self.columns = columns;
225 self.num_nodes = rows.len() as u64;
226 for sf in &files {
227 let _ = spiller.cleanup(sf);
228 }
229 Ok(())
230 }
231
232 pub fn flush_with_spiller(
243 &mut self,
244 columns: &mut [Column],
245 sort_key_column: Option<usize>,
246 dedup: bool,
247 ) -> std::io::Result<usize> {
248 if self.spill_files.is_empty() {
249 return self.flush(columns);
251 }
252
253 assert_eq!(
254 columns.len(),
255 self.columns.len(),
256 "NodeGroup::flush_with_spiller: column count mismatch"
257 );
258
259 let in_memory_rows = self.scan();
261 self.clear();
262
263 let sort_col = sort_key_column.unwrap_or(0);
265 let mut merger = MultiWayStreamMerge::new(&self.spill_files, Some(in_memory_rows), sort_col, dedup)
266 .map_err(std::io::Error::other)?;
267
268 let mut total: usize = 0;
270 while let Some(row) = merger.next_tuple() {
271 for (col_idx, value) in row.into_iter().enumerate() {
272 if col_idx < columns.len() {
273 columns[col_idx].append_value(&value)?;
274 }
275 }
276 total += 1;
277 }
278
279 if let Some(ref spiller) = self.spiller {
281 let files = std::mem::take(&mut self.spill_files);
282 for sf in &files {
283 let _ = spiller.cleanup(sf);
284 }
285 }
286
287 Ok(total)
288 }
289
290 pub fn is_full(&self) -> bool {
292 self.num_nodes as usize >= NODE_GROUP_SIZE
293 }
294
295 pub fn is_empty(&self) -> bool {
297 self.num_nodes == 0
298 }
299
300 pub fn num_columns(&self) -> usize {
302 self.columns.len()
303 }
304
305 pub fn has_spill_files(&self) -> bool {
307 !self.spill_files.is_empty()
308 }
309
310 pub fn remaining(&self) -> usize {
312 NODE_GROUP_SIZE.saturating_sub(self.num_nodes as usize)
313 }
314
315 pub fn flush(&mut self, columns: &mut [Column]) -> std::io::Result<usize> {
327 assert_eq!(
328 columns.len(),
329 self.columns.len(),
330 "NodeGroup::flush: column count mismatch"
331 );
332 let mut total = 0;
333 for (chunk, col) in self.columns.iter_mut().zip(columns.iter_mut()) {
334 let n = chunk.flush_to_column(col)?;
335 if total == 0 {
337 total = n;
338 }
339 debug_assert!(n == 0 || n == total, "inconsistent flush count");
340 }
341 self.num_nodes = 0;
342 Ok(total)
343 }
344
345 pub fn flush_copy(&self, columns: &mut [Column]) -> std::io::Result<usize> {
347 assert_eq!(
348 columns.len(),
349 self.columns.len(),
350 "NodeGroup::flush_copy: column count mismatch"
351 );
352 let mut total = 0;
353 for (chunk, col) in self.columns.iter().zip(columns.iter_mut()) {
354 let n = chunk.flush_copy_to_column(col)?;
355 if total == 0 {
356 total = n;
357 }
358 }
359 Ok(total)
360 }
361
362 pub fn scan(&self) -> Vec<Vec<Value>> {
367 let n_rows = self.num_nodes as usize;
368 let n_cols = self.columns.len();
369 let mut result = Vec::with_capacity(n_rows);
370
371 for row in 0..n_rows {
372 let mut row_data = Vec::with_capacity(n_cols);
373 for chunk in &self.columns {
374 match chunk.get(row) {
375 Some(v) => row_data.push(v.clone()),
376 None => row_data.push(Value::Null),
377 }
378 }
379 result.push(row_data);
380 }
381 result
382 }
383
384 pub fn scan_range(&self, start: usize, count: usize) -> Vec<Vec<Value>> {
388 let end = (start + count).min(self.num_nodes as usize);
389 if start >= end {
390 return Vec::new();
391 }
392 let n_cols = self.columns.len();
393 let mut result = Vec::with_capacity(end - start);
394
395 for row in start..end {
396 let mut row_data = Vec::with_capacity(n_cols);
397 for chunk in &self.columns {
398 match chunk.get(row) {
399 Some(v) => row_data.push(v.clone()),
400 None => row_data.push(Value::Null),
401 }
402 }
403 result.push(row_data);
404 }
405 result
406 }
407
408 pub fn get_value(&self, local_row: usize, col_idx: usize) -> Option<&Value> {
410 self.columns.get(col_idx).and_then(|chunk| chunk.get(local_row))
411 }
412
413 pub fn get_value_with_snapshot(
419 &self,
420 local_row: usize,
421 col_idx: usize,
422 snapshot_ts: Option<u64>,
423 commit_history: &[(u64, u64)],
424 ) -> Option<&Value> {
425 if let Some(ts) = snapshot_ts
427 && !self.is_row_visible(local_row, ts, commit_history)
428 {
429 return None;
430 }
431 self.columns
433 .get(col_idx)
434 .and_then(|chunk| chunk.get_value_with_snapshot(local_row, snapshot_ts, commit_history))
435 }
436
437 pub fn get_value_owned_with_snapshot(
443 &self,
444 local_row: usize,
445 col_idx: usize,
446 snapshot_ts: Option<u64>,
447 commit_history: &[(u64, u64)],
448 ) -> Option<Value> {
449 if let Some(ts) = snapshot_ts
451 && !self.is_row_visible(local_row, ts, commit_history)
452 {
453 return None;
454 }
455 self.columns
457 .get(col_idx)
458 .and_then(|chunk| chunk.get_value_owned_with_snapshot(local_row, snapshot_ts, commit_history))
459 }
460
461 pub fn is_row_visible(&self, local_row: usize, snapshot_ts: u64, commit_history: &[(u64, u64)]) -> bool {
464 match &self.version_info {
465 Some(vi) => vi.is_visible(local_row as u32, snapshot_ts, commit_history),
466 None => true, }
468 }
469
470 pub fn clear(&mut self) {
472 for chunk in &mut self.columns {
473 chunk.clear();
474 }
475 self.num_nodes = 0;
476 }
477}
478
479#[cfg(test)]
484mod tests {
485 use super::*;
486 use crate::buffer_manager::BufferManagerConfig;
487 use crate::page::DEFAULT_PAGE_SIZE;
488 use akar_common::memory::MemoryManager;
489 use akar_common::types::LogicalTypeID;
490 use std::sync::{Arc, Mutex};
491
492 fn setup_columns(num_cols: usize, db_path: &std::path::Path) -> Vec<Column> {
493 let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
494 let config = BufferManagerConfig::default();
495 let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
496 db_path.to_path_buf(),
497 mm,
498 config,
499 )));
500 (0..num_cols)
501 .map(|i| {
502 Column::new(
503 LogicalTypeID::Int64,
504 0,
505 i as u32,
506 db_path,
507 bm.clone(),
508 DEFAULT_PAGE_SIZE,
509 )
510 })
511 .collect()
512 }
513
514 #[test]
515 fn test_empty_group() {
516 let group = NodeGroup::new(3, 0);
517 assert_eq!(group.num_columns(), 3);
518 assert_eq!(group.num_nodes, 0);
519 assert!(group.is_empty());
520 assert!(!group.is_full());
521 assert_eq!(group.start_offset, 0);
522 }
523
524 #[test]
525 fn test_append_row() {
526 let mut group = NodeGroup::new(2, 100);
527 group.append_row(vec![Value::Int64(1), Value::Int64(2)]).unwrap();
528 assert_eq!(group.num_nodes, 1);
529 assert!(!group.is_empty());
530
531 group.append_row(vec![Value::Int64(3), Value::Int64(4)]).unwrap();
532 assert_eq!(group.num_nodes, 2);
533 }
534
535 #[test]
536 fn test_append_wrong_column_count() {
537 let mut group = NodeGroup::new(2, 0);
538 let result = group.append_row(vec![Value::Int64(1)]);
539 assert!(result.is_err());
540 assert!(result.unwrap_err().to_string().contains("column count mismatch"));
541 }
542
543 #[test]
544 fn test_append_when_full() {
545 let mut group = NodeGroup::with_capacity(1, 0, 5);
548 for _ in 0..10 {
549 group.append_row(vec![Value::Int64(1)]).unwrap();
550 }
551 assert!(!group.is_full());
552 assert_eq!(group.num_nodes, 10);
553 }
554
555 #[test]
556 fn test_scan() {
557 let mut group = NodeGroup::new(3, 0);
558 group
559 .append_row(vec![Value::Int64(10), Value::Int64(20), Value::Int64(30)])
560 .unwrap();
561 group
562 .append_row(vec![Value::Int64(11), Value::Int64(21), Value::Int64(31)])
563 .unwrap();
564
565 let data = group.scan();
566 assert_eq!(data.len(), 2);
567 assert_eq!(data[0][0], Value::Int64(10));
568 assert_eq!(data[0][1], Value::Int64(20));
569 assert_eq!(data[1][2], Value::Int64(31));
570 }
571
572 #[test]
573 fn test_scan_range() {
574 let mut group = NodeGroup::new(2, 0);
575 for i in 0..10 {
576 group.append_row(vec![Value::Int64(i), Value::Int64(i * 10)]).unwrap();
577 }
578
579 let slice = group.scan_range(3, 4);
580 assert_eq!(slice.len(), 4);
581 assert_eq!(slice[0][0], Value::Int64(3));
582 assert_eq!(slice[3][0], Value::Int64(6));
583 }
584
585 #[test]
586 fn test_get_value() {
587 let mut group = NodeGroup::new(2, 50);
588 group.append_row(vec![Value::Int64(100), Value::Int64(200)]).unwrap();
589
590 assert_eq!(group.get_value(0, 0), Some(&Value::Int64(100)));
591 assert_eq!(group.get_value(0, 1), Some(&Value::Int64(200)));
592 assert_eq!(group.get_value(1, 0), None);
593 }
594
595 #[test]
596 fn test_flush_to_columns() {
597 let dir = tempfile::tempdir().unwrap();
598 let mut cols = setup_columns(2, dir.path());
599 let mut group = NodeGroup::new(2, 0);
600
601 for i in 0i64..50 {
602 group.append_row(vec![Value::Int64(i), Value::Int64(i * 10)]).unwrap();
603 }
604
605 let flushed = group.flush(&mut cols).unwrap();
606 assert_eq!(flushed, 50);
607 assert_eq!(group.num_nodes, 0);
608 assert!(group.is_empty());
609
610 for i in 0i64..50 {
612 assert_eq!(cols[0].get_value(i as u64).unwrap(), Value::Int64(i));
613 assert_eq!(cols[1].get_value(i as u64).unwrap(), Value::Int64(i * 10));
614 }
615 }
616
617 #[test]
618 fn test_flush_copy_preserves_buffer() {
619 let dir = tempfile::tempdir().unwrap();
620 let mut cols = setup_columns(2, dir.path());
621 let mut group = NodeGroup::new(2, 0);
622
623 group.append_row(vec![Value::Int64(1), Value::Int64(2)]).unwrap();
624
625 let flushed = group.flush_copy(&mut cols).unwrap();
626 assert_eq!(flushed, 1);
627 assert_eq!(group.num_nodes, 1);
629 assert_eq!(cols[0].get_value(0).unwrap(), Value::Int64(1));
630 }
631
632 #[test]
633 fn test_restore_spilled_reconstructs_full_group() {
634 let dir = tempfile::tempdir().unwrap();
635 let spiller = Arc::new(crate::spiller::Spiller::new(dir.path(), 64));
636 let mut group = NodeGroup::new(2, 0);
637 group.set_spiller(spiller.clone());
638
639 for i in 0i64..20 {
641 group.append_row(vec![Value::Int64(i), Value::Int64(i * 10)]).unwrap();
642 }
643 assert!(!group.spill_files.is_empty(), "low threshold must spill");
644 assert!(group.num_nodes < 20, "spill must have evicted rows from memory");
645
646 group.restore_spilled().unwrap();
647 assert_eq!(group.num_nodes, 20, "restore must recover the full row set");
648 assert!(group.spill_files.is_empty(), "spill files cleaned up after restore");
649 for i in 0i64..20 {
650 assert_eq!(group.get_value(i as usize, 0), Some(&Value::Int64(i)));
651 assert_eq!(group.get_value(i as usize, 1), Some(&Value::Int64(i * 10)));
652 }
653 }
654
655 #[test]
656 fn test_clear() {
657 let mut group = NodeGroup::new(2, 0);
658 group.append_row(vec![Value::Int64(1), Value::Int64(2)]).unwrap();
659 group.clear();
660 assert_eq!(group.num_nodes, 0);
661 assert!(group.is_empty());
662 }
663
664 #[test]
665 fn test_remaining() {
666 let mut group = NodeGroup::with_capacity(3, 0, 10);
668 assert_eq!(group.remaining(), NODE_GROUP_SIZE);
669 group
670 .append_row(vec![Value::Int64(1), Value::Int64(2), Value::Int64(3)])
671 .unwrap();
672 assert_eq!(group.remaining(), NODE_GROUP_SIZE - 1);
673 group
674 .append_row(vec![Value::Int64(4), Value::Int64(5), Value::Int64(6)])
675 .unwrap();
676 assert_eq!(group.remaining(), NODE_GROUP_SIZE - 2);
677 }
678
679 #[test]
680 fn test_start_offset() {
681 let group = NodeGroup::new(2, 12345);
682 assert_eq!(group.start_offset, 12345);
683 }
684
685 #[test]
686 fn test_multi_column_scan() {
687 let mut group = NodeGroup::new(4, 0);
688 group
689 .append_row(vec![
690 Value::String("Alice".into()),
691 Value::Int64(30),
692 Value::Double(1.65),
693 Value::Bool(true),
694 ])
695 .unwrap();
696
697 let data = group.scan();
698 assert_eq!(data.len(), 1);
699 assert_eq!(data[0][0], Value::String("Alice".into()));
700 assert_eq!(data[0][1], Value::Int64(30));
701 assert_eq!(data[0][3], Value::Bool(true));
702 }
703
704 #[test]
705 fn test_multiple_flush_cycles() {
706 let dir = tempfile::tempdir().unwrap();
707 let mut cols = setup_columns(2, dir.path());
708 let mut group = NodeGroup::with_capacity(2, 0, 20);
709
710 for i in 0i64..15 {
712 group.append_row(vec![Value::Int64(i), Value::Int64(-i)]).unwrap();
713 }
714 assert_eq!(group.flush(&mut cols).unwrap(), 15);
715
716 for i in 15i64..30 {
718 group.append_row(vec![Value::Int64(i), Value::Int64(-i)]).unwrap();
719 }
720 assert_eq!(group.flush(&mut cols).unwrap(), 15);
721
722 assert_eq!(cols[0].num_values, 30);
724 for i in 0i64..30 {
725 assert_eq!(cols[0].get_value(i as u64).unwrap(), Value::Int64(i));
726 assert_eq!(cols[1].get_value(i as u64).unwrap(), Value::Int64(-i));
727 }
728 }
729}