reifydb_core/value/column/transform/
take.rs1use reifydb_type::{Result, util::cowvec::CowVec};
5
6use crate::value::column::columns::Columns;
7
8impl Columns {
9 pub fn take(&mut self, n: usize) -> Result<()> {
10 if !self.row_numbers.is_empty() {
11 let actual_n = n.min(self.row_numbers.len());
12 let new_row_numbers: Vec<_> = self.row_numbers.iter().take(actual_n).copied().collect();
13 self.row_numbers = CowVec::new(new_row_numbers);
14 }
15
16 if !self.created_at.is_empty() {
17 let actual_n = n.min(self.created_at.len());
18 let new_created_at: Vec<_> = self.created_at.iter().take(actual_n).copied().collect();
19 self.created_at = CowVec::new(new_created_at);
20 }
21
22 if !self.updated_at.is_empty() {
23 let actual_n = n.min(self.updated_at.len());
24 let new_updated_at: Vec<_> = self.updated_at.iter().take(actual_n).copied().collect();
25 self.updated_at = CowVec::new(new_updated_at);
26 }
27
28 let mut new_buffers = Vec::with_capacity(self.len());
29
30 for data in self.columns.iter() {
31 new_buffers.push(data.take(n));
32 }
33
34 self.columns = CowVec::new(new_buffers);
35
36 Ok(())
37 }
38}
39
40#[cfg(test)]
41pub mod tests {
42 use reifydb_type::value::{Value, r#type::Type};
43
44 use super::*;
45 use crate::value::column::{ColumnBuffer, ColumnWithName};
46
47 #[test]
48 fn test_bool_column() {
49 let mut test_instance = Columns::new(vec![ColumnWithName::bool_with_bitvec(
50 "flag",
51 [true, true, false],
52 [false, true, true],
53 )]);
54
55 test_instance.take(1).unwrap();
56
57 assert_eq!(test_instance[0], ColumnBuffer::bool_with_bitvec([true], [false]));
58 }
59
60 #[test]
61 fn test_float4_column() {
62 let mut test_instance = Columns::new(vec![ColumnWithName::float4_with_bitvec(
63 "a",
64 [1.0, 2.0, 3.0],
65 [true, false, true],
66 )]);
67
68 test_instance.take(2).unwrap();
69
70 assert_eq!(test_instance[0], ColumnBuffer::float4_with_bitvec([1.0, 2.0], [true, false]));
71 }
72
73 #[test]
74 fn test_float8_column() {
75 let mut test_instance = Columns::new(vec![ColumnWithName::float8_with_bitvec(
76 "a",
77 [1f64, 2.0, 3.0, 4.0],
78 [true, true, false, true],
79 )]);
80
81 test_instance.take(2).unwrap();
82
83 assert_eq!(test_instance[0], ColumnBuffer::float8_with_bitvec([1.0, 2.0], [true, true]));
84 }
85
86 #[test]
87 fn test_int1_column() {
88 let mut test_instance =
89 Columns::new(vec![ColumnWithName::int1_with_bitvec("a", [1, 2, 3], [true, false, true])]);
90
91 test_instance.take(2).unwrap();
92
93 assert_eq!(test_instance[0], ColumnBuffer::int1_with_bitvec([1, 2], [true, false]));
94 }
95
96 #[test]
97 fn test_int2_column() {
98 let mut test_instance = Columns::new(vec![ColumnWithName::int2_with_bitvec(
99 "a",
100 [1, 2, 3, 4],
101 [true, true, false, true],
102 )]);
103
104 test_instance.take(2).unwrap();
105
106 assert_eq!(test_instance[0], ColumnBuffer::int2_with_bitvec([1, 2], [true, true]));
107 }
108
109 #[test]
110 fn test_int4_column() {
111 let mut test_instance =
112 Columns::new(vec![ColumnWithName::int4_with_bitvec("a", [1, 2], [true, false])]);
113
114 test_instance.take(1).unwrap();
115
116 assert_eq!(test_instance[0], ColumnBuffer::int4_with_bitvec([1], [true]));
117 }
118
119 #[test]
120 fn test_int8_column() {
121 let mut test_instance =
122 Columns::new(vec![ColumnWithName::int8_with_bitvec("a", [1, 2, 3], [false, true, true])]);
123
124 test_instance.take(2).unwrap();
125
126 assert_eq!(test_instance[0], ColumnBuffer::int8_with_bitvec([1, 2], [false, true]));
127 }
128
129 #[test]
130 fn test_int16_column() {
131 let mut test_instance =
132 Columns::new(vec![ColumnWithName::int16_with_bitvec("a", [1, 2], [true, true])]);
133
134 test_instance.take(1).unwrap();
135
136 assert_eq!(test_instance[0], ColumnBuffer::int16_with_bitvec([1], [true]));
137 }
138
139 #[test]
140 fn test_uint1_column() {
141 let mut test_instance =
142 Columns::new(vec![ColumnWithName::uint1_with_bitvec("a", [1, 2, 3], [false, false, true])]);
143
144 test_instance.take(2).unwrap();
145
146 assert_eq!(test_instance[0], ColumnBuffer::uint1_with_bitvec([1, 2], [false, false]));
147 }
148
149 #[test]
150 fn test_uint2_column() {
151 let mut test_instance =
152 Columns::new(vec![ColumnWithName::uint2_with_bitvec("a", [1, 2], [true, false])]);
153
154 test_instance.take(1).unwrap();
155
156 assert_eq!(test_instance[0], ColumnBuffer::uint2_with_bitvec([1], [true]));
157 }
158
159 #[test]
160 fn test_uint4_column() {
161 let mut test_instance =
162 Columns::new(vec![ColumnWithName::uint4_with_bitvec("a", [10, 20], [false, true])]);
163
164 test_instance.take(1).unwrap();
165
166 assert_eq!(test_instance[0], ColumnBuffer::uint4_with_bitvec([10], [false]));
167 }
168
169 #[test]
170 fn test_uint8_column() {
171 let mut test_instance =
172 Columns::new(vec![ColumnWithName::uint8_with_bitvec("a", [10, 20, 30], [true, true, false])]);
173
174 test_instance.take(2).unwrap();
175
176 assert_eq!(test_instance[0], ColumnBuffer::uint8_with_bitvec([10, 20], [true, true]));
177 }
178
179 #[test]
180 fn test_uint16_column() {
181 let mut test_instance = Columns::new(vec![ColumnWithName::uint16_with_bitvec(
182 "a",
183 [100, 200, 300],
184 [true, false, true],
185 )]);
186
187 test_instance.take(1).unwrap();
188
189 assert_eq!(test_instance[0], ColumnBuffer::uint16_with_bitvec([100], [true]));
190 }
191
192 #[test]
193 fn test_text_column() {
194 let mut test_instance = Columns::new(vec![ColumnWithName::utf8_with_bitvec(
195 "t",
196 vec!["a".to_string(), "b".to_string(), "c".to_string()],
197 vec![true, false, true],
198 )]);
199
200 test_instance.take(2).unwrap();
201
202 assert_eq!(
203 test_instance[0],
204 ColumnBuffer::utf8_with_bitvec(["a".to_string(), "b".to_string()], [true, false])
205 );
206 }
207
208 #[test]
209 fn test_none_column() {
210 let mut test_instance = Columns::new(vec![ColumnWithName::undefined_typed("u", Type::Boolean, 3)]);
211
212 test_instance.take(2).unwrap();
213
214 assert_eq!(test_instance[0].len(), 2);
215 assert_eq!(test_instance[0].get_value(0), Value::none());
216 assert_eq!(test_instance[0].get_value(1), Value::none());
217 }
218
219 #[test]
220 fn test_handles_none() {
221 let mut test_instance = Columns::new(vec![ColumnWithName::undefined_typed("u", Type::Boolean, 5)]);
222
223 test_instance.take(3).unwrap();
224
225 assert_eq!(test_instance[0].len(), 3);
226 assert_eq!(test_instance[0].get_value(0), Value::none());
227 }
228
229 #[test]
230 fn test_n_larger_than_len_is_safe() {
231 let mut test_instance =
232 Columns::new(vec![ColumnWithName::int2_with_bitvec("a", [10, 20], [true, false])]);
233
234 test_instance.take(10).unwrap();
235
236 assert_eq!(test_instance[0], ColumnBuffer::int2_with_bitvec([10, 20], [true, false]));
237 }
238}