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