reifydb_core/value/column/buffer/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_type::{storage::DataBitVec, util::bitvec::BitVec};
5
6use crate::value::column::ColumnBuffer;
7
8macro_rules! map_container {
9	($self:expr, |$c:ident| $body:expr) => {
10		match $self {
11			ColumnBuffer::Bool($c) => ColumnBuffer::Bool($body),
12			ColumnBuffer::Float4($c) => ColumnBuffer::Float4($body),
13			ColumnBuffer::Float8($c) => ColumnBuffer::Float8($body),
14			ColumnBuffer::Int1($c) => ColumnBuffer::Int1($body),
15			ColumnBuffer::Int2($c) => ColumnBuffer::Int2($body),
16			ColumnBuffer::Int4($c) => ColumnBuffer::Int4($body),
17			ColumnBuffer::Int8($c) => ColumnBuffer::Int8($body),
18			ColumnBuffer::Int16($c) => ColumnBuffer::Int16($body),
19			ColumnBuffer::Uint1($c) => ColumnBuffer::Uint1($body),
20			ColumnBuffer::Uint2($c) => ColumnBuffer::Uint2($body),
21			ColumnBuffer::Uint4($c) => ColumnBuffer::Uint4($body),
22			ColumnBuffer::Uint8($c) => ColumnBuffer::Uint8($body),
23			ColumnBuffer::Uint16($c) => ColumnBuffer::Uint16($body),
24			ColumnBuffer::Utf8 {
25				container: $c,
26				max_bytes,
27			} => ColumnBuffer::Utf8 {
28				container: $body,
29				max_bytes: *max_bytes,
30			},
31			ColumnBuffer::Date($c) => ColumnBuffer::Date($body),
32			ColumnBuffer::DateTime($c) => ColumnBuffer::DateTime($body),
33			ColumnBuffer::Time($c) => ColumnBuffer::Time($body),
34			ColumnBuffer::Duration($c) => ColumnBuffer::Duration($body),
35
36			ColumnBuffer::IdentityId($c) => ColumnBuffer::IdentityId($body),
37			ColumnBuffer::DictionaryId($c) => ColumnBuffer::DictionaryId($body),
38			ColumnBuffer::Uuid4($c) => ColumnBuffer::Uuid4($body),
39			ColumnBuffer::Uuid7($c) => ColumnBuffer::Uuid7($body),
40			ColumnBuffer::Blob {
41				container: $c,
42				max_bytes,
43			} => ColumnBuffer::Blob {
44				container: $body,
45				max_bytes: *max_bytes,
46			},
47			ColumnBuffer::Int {
48				container: $c,
49				max_bytes,
50			} => ColumnBuffer::Int {
51				container: $body,
52				max_bytes: *max_bytes,
53			},
54			ColumnBuffer::Uint {
55				container: $c,
56				max_bytes,
57			} => ColumnBuffer::Uint {
58				container: $body,
59				max_bytes: *max_bytes,
60			},
61			ColumnBuffer::Decimal {
62				container: $c,
63				precision,
64				scale,
65			} => ColumnBuffer::Decimal {
66				container: $body,
67				precision: *precision,
68				scale: *scale,
69			},
70			ColumnBuffer::Any($c) => ColumnBuffer::Any($body),
71			ColumnBuffer::Option {
72				..
73			} => {
74				unreachable!(
75					"map_container! must not be called on Option variant directly; handle it explicitly"
76				)
77			}
78		}
79	};
80}
81
82impl ColumnBuffer {
83	pub fn take(&self, num: usize) -> ColumnBuffer {
84		match self {
85			ColumnBuffer::Option {
86				inner,
87				bitvec,
88			} => {
89				let new_bitvec = DataBitVec::take(bitvec, num);
90				// If all bits in the taken bitvec are set (all defined),
91				// unwrap the Option and return the bare inner data.
92				if DataBitVec::count_ones(&new_bitvec) == DataBitVec::len(&new_bitvec)
93					&& DataBitVec::len(&new_bitvec) > 0
94				{
95					inner.take(num)
96				} else {
97					ColumnBuffer::Option {
98						inner: Box::new(inner.take(num)),
99						bitvec: new_bitvec,
100					}
101				}
102			}
103			_ => map_container!(self, |c| c.take(num)),
104		}
105	}
106
107	// Contiguous-range slice: returns rows [start, end).
108	pub fn slice(&self, start: usize, end: usize) -> ColumnBuffer {
109		match self {
110			ColumnBuffer::Option {
111				inner,
112				bitvec,
113			} => {
114				let len = end - start;
115				let mut new_bits = Vec::with_capacity(len);
116				for row in start..end {
117					new_bits.push(DataBitVec::get(bitvec, row));
118				}
119				let new_bitvec = BitVec::from(new_bits);
120				ColumnBuffer::Option {
121					inner: Box::new(inner.slice(start, end)),
122					bitvec: new_bitvec,
123				}
124			}
125			_ => map_container!(self, |c| c.slice(start, end)),
126		}
127	}
128
129	// Gather rows at the given indices.
130	pub fn gather(&self, indices: &[usize]) -> ColumnBuffer {
131		match self {
132			ColumnBuffer::Option {
133				inner,
134				bitvec,
135			} => {
136				let mut new_bits = Vec::with_capacity(indices.len());
137				for &i in indices {
138					new_bits.push(DataBitVec::get(bitvec, i));
139				}
140				let new_bitvec = BitVec::from(new_bits);
141				ColumnBuffer::Option {
142					inner: Box::new(inner.gather(indices)),
143					bitvec: new_bitvec,
144				}
145			}
146			_ => {
147				let mut cloned = self.clone();
148				cloned.reorder(indices);
149				cloned
150			}
151		}
152	}
153}
reifydb_core/value/column/buffer/take.rs

reifydb_core/value/column/buffer/
take.rs
