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

reifydb_core/value/column/buffer/
take.rs
