reifydb_core/value/column/buffer/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use reifydb_value::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 = bitvec.take(num);
90
91				if new_bitvec.count_ones() == new_bitvec.len() && !new_bitvec.is_empty() {
92					inner.take(num)
93				} else {
94					ColumnBuffer::Option {
95						inner: Box::new(inner.take(num)),
96						bitvec: new_bitvec,
97					}
98				}
99			}
100			_ => map_container!(self, |c| c.take(num)),
101		}
102	}
103
104	pub fn slice(&self, start: usize, end: usize) -> ColumnBuffer {
105		match self {
106			ColumnBuffer::Option {
107				inner,
108				bitvec,
109			} => {
110				let len = end - start;
111				let mut new_bits = Vec::with_capacity(len);
112				for row in start..end {
113					new_bits.push(bitvec.get(row));
114				}
115				let new_bitvec = BitVec::from(new_bits);
116				ColumnBuffer::Option {
117					inner: Box::new(inner.slice(start, end)),
118					bitvec: new_bitvec,
119				}
120			}
121			_ => map_container!(self, |c| c.slice(start, end)),
122		}
123	}
124
125	pub fn gather(&self, indices: &[usize]) -> ColumnBuffer {
126		match self {
127			ColumnBuffer::Option {
128				inner,
129				bitvec,
130			} => {
131				let mut new_bits = Vec::with_capacity(indices.len());
132				for &i in indices {
133					new_bits.push(bitvec.get(i));
134				}
135				let new_bitvec = BitVec::from(new_bits);
136				ColumnBuffer::Option {
137					inner: Box::new(inner.gather(indices)),
138					bitvec: new_bitvec,
139				}
140			}
141			_ => {
142				let mut cloned = self.clone();
143				cloned.reorder(indices);
144				cloned
145			}
146		}
147	}
148}
reifydb_core/value/column/buffer/take.rs

reifydb_core/value/column/buffer/
take.rs
