/Users/andrewlamb/Software/arrow-rs/arrow-array/src/array/fixed_size_list_array.rs
Line | Count | Source |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | |
18 | | use crate::array::print_long_array; |
19 | | use crate::builder::{FixedSizeListBuilder, PrimitiveBuilder}; |
20 | | use crate::iterator::FixedSizeListIter; |
21 | | use crate::{make_array, Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType}; |
22 | | use arrow_buffer::buffer::NullBuffer; |
23 | | use arrow_buffer::ArrowNativeType; |
24 | | use arrow_data::{ArrayData, ArrayDataBuilder}; |
25 | | use arrow_schema::{ArrowError, DataType, FieldRef}; |
26 | | use std::any::Any; |
27 | | use std::sync::Arc; |
28 | | |
29 | | /// An array of [fixed length lists], similar to JSON arrays |
30 | | /// (e.g. `["A", "B"]`). |
31 | | /// |
32 | | /// Lists are represented using a `values` child |
33 | | /// array where each list has a fixed size of `value_length`. |
34 | | /// |
35 | | /// Use [`FixedSizeListBuilder`] to construct a [`FixedSizeListArray`]. |
36 | | /// |
37 | | /// # Representation |
38 | | /// |
39 | | /// A [`FixedSizeListArray`] can represent a list of values of any other |
40 | | /// supported Arrow type. Each element of the `FixedSizeListArray` itself is |
41 | | /// a list which may contain NULL and non-null values, |
42 | | /// or may itself be NULL. |
43 | | /// |
44 | | /// For example, this `FixedSizeListArray` stores lists of strings: |
45 | | /// |
46 | | /// ```text |
47 | | /// ┌─────────────┐ |
48 | | /// │ [A,B] │ |
49 | | /// ├─────────────┤ |
50 | | /// │ NULL │ |
51 | | /// ├─────────────┤ |
52 | | /// │ [C,NULL] │ |
53 | | /// └─────────────┘ |
54 | | /// ``` |
55 | | /// |
56 | | /// The `values` of this `FixedSizeListArray`s are stored in a child |
57 | | /// [`StringArray`] where logical null values take up `values_length` slots in the array |
58 | | /// as shown in the following diagram. The logical values |
59 | | /// are shown on the left, and the actual `FixedSizeListArray` encoding on the right |
60 | | /// |
61 | | /// ```text |
62 | | /// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ |
63 | | /// ┌ ─ ─ ─ ─ ─ ─ ─ ─┐ |
64 | | /// ┌─────────────┐ │ ┌───┐ ┌───┐ ┌──────┐ │ |
65 | | /// │ [A,B] │ │ 1 │ │ │ 1 │ │ A │ │ 0 |
66 | | /// ├─────────────┤ │ ├───┤ ├───┤ ├──────┤ │ |
67 | | /// │ NULL │ │ 0 │ │ │ 1 │ │ B │ │ 1 |
68 | | /// ├─────────────┤ │ ├───┤ ├───┤ ├──────┤ │ |
69 | | /// │ [C,NULL] │ │ 1 │ │ │ 0 │ │ ???? │ │ 2 |
70 | | /// └─────────────┘ │ └───┘ ├───┤ ├──────┤ │ |
71 | | /// | │ 0 │ │ ???? │ │ 3 |
72 | | /// Logical Values │ Validity ├───┤ ├──────┤ │ |
73 | | /// (nulls) │ │ 1 │ │ C │ │ 4 |
74 | | /// │ ├───┤ ├──────┤ │ |
75 | | /// │ │ 0 │ │ ???? │ │ 5 |
76 | | /// │ └───┘ └──────┘ │ |
77 | | /// │ Values │ |
78 | | /// │ FixedSizeListArray (Array) │ |
79 | | /// └ ─ ─ ─ ─ ─ ─ ─ ─┘ |
80 | | /// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ |
81 | | /// ``` |
82 | | /// |
83 | | /// # Example |
84 | | /// |
85 | | /// ``` |
86 | | /// # use std::sync::Arc; |
87 | | /// # use arrow_array::{Array, FixedSizeListArray, Int32Array}; |
88 | | /// # use arrow_data::ArrayData; |
89 | | /// # use arrow_schema::{DataType, Field}; |
90 | | /// # use arrow_buffer::Buffer; |
91 | | /// // Construct a value array |
92 | | /// let value_data = ArrayData::builder(DataType::Int32) |
93 | | /// .len(9) |
94 | | /// .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8])) |
95 | | /// .build() |
96 | | /// .unwrap(); |
97 | | /// let list_data_type = DataType::FixedSizeList( |
98 | | /// Arc::new(Field::new_list_field(DataType::Int32, false)), |
99 | | /// 3, |
100 | | /// ); |
101 | | /// let list_data = ArrayData::builder(list_data_type.clone()) |
102 | | /// .len(3) |
103 | | /// .add_child_data(value_data.clone()) |
104 | | /// .build() |
105 | | /// .unwrap(); |
106 | | /// let list_array = FixedSizeListArray::from(list_data); |
107 | | /// let list0 = list_array.value(0); |
108 | | /// let list1 = list_array.value(1); |
109 | | /// let list2 = list_array.value(2); |
110 | | /// |
111 | | /// assert_eq!( &[0, 1, 2], list0.as_any().downcast_ref::<Int32Array>().unwrap().values()); |
112 | | /// assert_eq!( &[3, 4, 5], list1.as_any().downcast_ref::<Int32Array>().unwrap().values()); |
113 | | /// assert_eq!( &[6, 7, 8], list2.as_any().downcast_ref::<Int32Array>().unwrap().values()); |
114 | | /// ``` |
115 | | /// |
116 | | /// [`StringArray`]: crate::array::StringArray |
117 | | /// [fixed size arrays](https://arrow.apache.org/docs/format/Columnar.html#fixed-size-list-layout) |
118 | | #[derive(Clone)] |
119 | | pub struct FixedSizeListArray { |
120 | | data_type: DataType, // Must be DataType::FixedSizeList(value_length) |
121 | | values: ArrayRef, |
122 | | nulls: Option<NullBuffer>, |
123 | | value_length: i32, |
124 | | len: usize, |
125 | | } |
126 | | |
127 | | impl FixedSizeListArray { |
128 | | /// Create a new [`FixedSizeListArray`] with `size` element size, panicking on failure |
129 | | /// |
130 | | /// # Panics |
131 | | /// |
132 | | /// Panics if [`Self::try_new`] returns an error |
133 | 0 | pub fn new(field: FieldRef, size: i32, values: ArrayRef, nulls: Option<NullBuffer>) -> Self { |
134 | 0 | Self::try_new(field, size, values, nulls).unwrap() |
135 | 0 | } |
136 | | |
137 | | /// Create a new [`FixedSizeListArray`] from the provided parts, returning an error on failure |
138 | | /// |
139 | | /// # Errors |
140 | | /// |
141 | | /// * `size < 0` |
142 | | /// * `values.len() / size != nulls.len()` |
143 | | /// * `values.data_type() != field.data_type()` |
144 | | /// * `!field.is_nullable() && !nulls.expand(size).contains(values.logical_nulls())` |
145 | 0 | pub fn try_new( |
146 | 0 | field: FieldRef, |
147 | 0 | size: i32, |
148 | 0 | values: ArrayRef, |
149 | 0 | nulls: Option<NullBuffer>, |
150 | 0 | ) -> Result<Self, ArrowError> { |
151 | 0 | let s = size.to_usize().ok_or_else(|| { |
152 | 0 | ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {size}")) |
153 | 0 | })?; |
154 | | |
155 | 0 | let len = match s { |
156 | 0 | 0 => nulls.as_ref().map(|x| x.len()).unwrap_or_default(), |
157 | | _ => { |
158 | 0 | let len = values.len() / s.max(1); |
159 | 0 | if let Some(n) = nulls.as_ref() { |
160 | 0 | if n.len() != len { |
161 | 0 | return Err(ArrowError::InvalidArgumentError(format!( |
162 | 0 | "Incorrect length of null buffer for FixedSizeListArray, expected {} got {}", |
163 | 0 | len, |
164 | 0 | n.len(), |
165 | 0 | ))); |
166 | 0 | } |
167 | 0 | } |
168 | 0 | len |
169 | | } |
170 | | }; |
171 | | |
172 | 0 | if field.data_type() != values.data_type() { |
173 | 0 | return Err(ArrowError::InvalidArgumentError(format!( |
174 | 0 | "FixedSizeListArray expected data type {} got {} for {:?}", |
175 | 0 | field.data_type(), |
176 | 0 | values.data_type(), |
177 | 0 | field.name() |
178 | 0 | ))); |
179 | 0 | } |
180 | | |
181 | 0 | if let Some(a) = values.logical_nulls() { |
182 | 0 | let nulls_valid = field.is_nullable() |
183 | 0 | || nulls |
184 | 0 | .as_ref() |
185 | 0 | .map(|n| n.expand(size as _).contains(&a)) |
186 | 0 | .unwrap_or_default() |
187 | 0 | || (nulls.is_none() && a.null_count() == 0); |
188 | | |
189 | 0 | if !nulls_valid { |
190 | 0 | return Err(ArrowError::InvalidArgumentError(format!( |
191 | 0 | "Found unmasked nulls for non-nullable FixedSizeListArray field {:?}", |
192 | 0 | field.name() |
193 | 0 | ))); |
194 | 0 | } |
195 | 0 | } |
196 | | |
197 | 0 | let data_type = DataType::FixedSizeList(field, size); |
198 | 0 | Ok(Self { |
199 | 0 | data_type, |
200 | 0 | values, |
201 | 0 | value_length: size, |
202 | 0 | nulls, |
203 | 0 | len, |
204 | 0 | }) |
205 | 0 | } |
206 | | |
207 | | /// Create a new [`FixedSizeListArray`] of length `len` where all values are null |
208 | | /// |
209 | | /// # Panics |
210 | | /// |
211 | | /// Panics if |
212 | | /// |
213 | | /// * `size < 0` |
214 | | /// * `size * len` would overflow `usize` |
215 | 0 | pub fn new_null(field: FieldRef, size: i32, len: usize) -> Self { |
216 | 0 | let capacity = size.to_usize().unwrap().checked_mul(len).unwrap(); |
217 | 0 | Self { |
218 | 0 | values: make_array(ArrayData::new_null(field.data_type(), capacity)), |
219 | 0 | data_type: DataType::FixedSizeList(field, size), |
220 | 0 | nulls: Some(NullBuffer::new_null(len)), |
221 | 0 | value_length: size, |
222 | 0 | len, |
223 | 0 | } |
224 | 0 | } |
225 | | |
226 | | /// Deconstruct this array into its constituent parts |
227 | 0 | pub fn into_parts(self) -> (FieldRef, i32, ArrayRef, Option<NullBuffer>) { |
228 | 0 | let f = match self.data_type { |
229 | 0 | DataType::FixedSizeList(f, _) => f, |
230 | 0 | _ => unreachable!(), |
231 | | }; |
232 | 0 | (f, self.value_length, self.values, self.nulls) |
233 | 0 | } |
234 | | |
235 | | /// Returns a reference to the values of this list. |
236 | 0 | pub fn values(&self) -> &ArrayRef { |
237 | 0 | &self.values |
238 | 0 | } |
239 | | |
240 | | /// Returns a clone of the value type of this list. |
241 | 0 | pub fn value_type(&self) -> DataType { |
242 | 0 | self.values.data_type().clone() |
243 | 0 | } |
244 | | |
245 | | /// Returns ith value of this list array. |
246 | | /// |
247 | | /// Note: This method does not check for nulls and the value is arbitrary |
248 | | /// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index. |
249 | | /// |
250 | | /// # Panics |
251 | | /// Panics if index `i` is out of bounds |
252 | 0 | pub fn value(&self, i: usize) -> ArrayRef { |
253 | 0 | self.values |
254 | 0 | .slice(self.value_offset_at(i), self.value_length() as usize) |
255 | 0 | } |
256 | | |
257 | | /// Returns the offset for value at index `i`. |
258 | | /// |
259 | | /// Note this doesn't do any bound checking, for performance reason. |
260 | | #[inline] |
261 | 0 | pub fn value_offset(&self, i: usize) -> i32 { |
262 | 0 | self.value_offset_at(i) as i32 |
263 | 0 | } |
264 | | |
265 | | /// Returns the length for an element. |
266 | | /// |
267 | | /// All elements have the same length as the array is a fixed size. |
268 | | #[inline] |
269 | 0 | pub const fn value_length(&self) -> i32 { |
270 | 0 | self.value_length |
271 | 0 | } |
272 | | |
273 | | #[inline] |
274 | 0 | const fn value_offset_at(&self, i: usize) -> usize { |
275 | 0 | i * self.value_length as usize |
276 | 0 | } |
277 | | |
278 | | /// Returns a zero-copy slice of this array with the indicated offset and length. |
279 | 0 | pub fn slice(&self, offset: usize, len: usize) -> Self { |
280 | 0 | assert!( |
281 | 0 | offset.saturating_add(len) <= self.len, |
282 | 0 | "the length + offset of the sliced FixedSizeListArray cannot exceed the existing length" |
283 | | ); |
284 | 0 | let size = self.value_length as usize; |
285 | | |
286 | | Self { |
287 | 0 | data_type: self.data_type.clone(), |
288 | 0 | values: self.values.slice(offset * size, len * size), |
289 | 0 | nulls: self.nulls.as_ref().map(|n| n.slice(offset, len)), |
290 | 0 | value_length: self.value_length, |
291 | 0 | len, |
292 | | } |
293 | 0 | } |
294 | | |
295 | | /// Creates a [`FixedSizeListArray`] from an iterator of primitive values |
296 | | /// # Example |
297 | | /// ``` |
298 | | /// # use arrow_array::FixedSizeListArray; |
299 | | /// # use arrow_array::types::Int32Type; |
300 | | /// |
301 | | /// let data = vec![ |
302 | | /// Some(vec![Some(0), Some(1), Some(2)]), |
303 | | /// None, |
304 | | /// Some(vec![Some(3), None, Some(5)]), |
305 | | /// Some(vec![Some(6), Some(7), Some(45)]), |
306 | | /// ]; |
307 | | /// let list_array = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(data, 3); |
308 | | /// println!("{:?}", list_array); |
309 | | /// ``` |
310 | | pub fn from_iter_primitive<T, P, I>(iter: I, length: i32) -> Self |
311 | | where |
312 | | T: ArrowPrimitiveType, |
313 | | P: IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>, |
314 | | I: IntoIterator<Item = Option<P>>, |
315 | | { |
316 | | let l = length as usize; |
317 | | let iter = iter.into_iter(); |
318 | | let size_hint = iter.size_hint().0; |
319 | | let mut builder = FixedSizeListBuilder::with_capacity( |
320 | | PrimitiveBuilder::<T>::with_capacity(size_hint * l), |
321 | | length, |
322 | | size_hint, |
323 | | ); |
324 | | |
325 | | for i in iter { |
326 | | match i { |
327 | | Some(p) => { |
328 | | for t in p { |
329 | | builder.values().append_option(t); |
330 | | } |
331 | | builder.append(true); |
332 | | } |
333 | | None => { |
334 | | builder.values().append_nulls(l); |
335 | | builder.append(false) |
336 | | } |
337 | | } |
338 | | } |
339 | | builder.finish() |
340 | | } |
341 | | |
342 | | /// constructs a new iterator |
343 | 0 | pub fn iter(&self) -> FixedSizeListIter<'_> { |
344 | 0 | FixedSizeListIter::new(self) |
345 | 0 | } |
346 | | } |
347 | | |
348 | | impl From<ArrayData> for FixedSizeListArray { |
349 | 0 | fn from(data: ArrayData) -> Self { |
350 | 0 | let value_length = match data.data_type() { |
351 | 0 | DataType::FixedSizeList(_, len) => *len, |
352 | 0 | data_type => { |
353 | 0 | panic!("FixedSizeListArray data should contain a FixedSizeList data type, got {data_type:?}") |
354 | | } |
355 | | }; |
356 | | |
357 | 0 | let size = value_length as usize; |
358 | 0 | let values = |
359 | 0 | make_array(data.child_data()[0].slice(data.offset() * size, data.len() * size)); |
360 | 0 | Self { |
361 | 0 | data_type: data.data_type().clone(), |
362 | 0 | values, |
363 | 0 | nulls: data.nulls().cloned(), |
364 | 0 | value_length, |
365 | 0 | len: data.len(), |
366 | 0 | } |
367 | 0 | } |
368 | | } |
369 | | |
370 | | impl From<FixedSizeListArray> for ArrayData { |
371 | 0 | fn from(array: FixedSizeListArray) -> Self { |
372 | 0 | let builder = ArrayDataBuilder::new(array.data_type) |
373 | 0 | .len(array.len) |
374 | 0 | .nulls(array.nulls) |
375 | 0 | .child_data(vec![array.values.to_data()]); |
376 | | |
377 | 0 | unsafe { builder.build_unchecked() } |
378 | 0 | } |
379 | | } |
380 | | |
381 | | impl Array for FixedSizeListArray { |
382 | 0 | fn as_any(&self) -> &dyn Any { |
383 | 0 | self |
384 | 0 | } |
385 | | |
386 | 0 | fn to_data(&self) -> ArrayData { |
387 | 0 | self.clone().into() |
388 | 0 | } |
389 | | |
390 | 0 | fn into_data(self) -> ArrayData { |
391 | 0 | self.into() |
392 | 0 | } |
393 | | |
394 | 0 | fn data_type(&self) -> &DataType { |
395 | 0 | &self.data_type |
396 | 0 | } |
397 | | |
398 | 0 | fn slice(&self, offset: usize, length: usize) -> ArrayRef { |
399 | 0 | Arc::new(self.slice(offset, length)) |
400 | 0 | } |
401 | | |
402 | 0 | fn len(&self) -> usize { |
403 | 0 | self.len |
404 | 0 | } |
405 | | |
406 | 0 | fn is_empty(&self) -> bool { |
407 | 0 | self.len == 0 |
408 | 0 | } |
409 | | |
410 | 0 | fn shrink_to_fit(&mut self) { |
411 | 0 | self.values.shrink_to_fit(); |
412 | 0 | if let Some(nulls) = &mut self.nulls { |
413 | 0 | nulls.shrink_to_fit(); |
414 | 0 | } |
415 | 0 | } |
416 | | |
417 | 0 | fn offset(&self) -> usize { |
418 | 0 | 0 |
419 | 0 | } |
420 | | |
421 | 0 | fn nulls(&self) -> Option<&NullBuffer> { |
422 | 0 | self.nulls.as_ref() |
423 | 0 | } |
424 | | |
425 | 0 | fn logical_null_count(&self) -> usize { |
426 | | // More efficient that the default implementation |
427 | 0 | self.null_count() |
428 | 0 | } |
429 | | |
430 | 0 | fn get_buffer_memory_size(&self) -> usize { |
431 | 0 | let mut size = self.values.get_buffer_memory_size(); |
432 | 0 | if let Some(n) = self.nulls.as_ref() { |
433 | 0 | size += n.buffer().capacity(); |
434 | 0 | } |
435 | 0 | size |
436 | 0 | } |
437 | | |
438 | 0 | fn get_array_memory_size(&self) -> usize { |
439 | 0 | let mut size = std::mem::size_of::<Self>() + self.values.get_array_memory_size(); |
440 | 0 | if let Some(n) = self.nulls.as_ref() { |
441 | 0 | size += n.buffer().capacity(); |
442 | 0 | } |
443 | 0 | size |
444 | 0 | } |
445 | | } |
446 | | |
447 | | impl ArrayAccessor for FixedSizeListArray { |
448 | | type Item = ArrayRef; |
449 | | |
450 | 0 | fn value(&self, index: usize) -> Self::Item { |
451 | 0 | FixedSizeListArray::value(self, index) |
452 | 0 | } |
453 | | |
454 | 0 | unsafe fn value_unchecked(&self, index: usize) -> Self::Item { |
455 | 0 | FixedSizeListArray::value(self, index) |
456 | 0 | } |
457 | | } |
458 | | |
459 | | impl std::fmt::Debug for FixedSizeListArray { |
460 | 0 | fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { |
461 | 0 | write!(f, "FixedSizeListArray<{}>\n[\n", self.value_length())?; |
462 | 0 | print_long_array(self, f, |array, index, f| { |
463 | 0 | std::fmt::Debug::fmt(&array.value(index), f) |
464 | 0 | })?; |
465 | 0 | write!(f, "]") |
466 | 0 | } |
467 | | } |
468 | | |
469 | | impl ArrayAccessor for &FixedSizeListArray { |
470 | | type Item = ArrayRef; |
471 | | |
472 | 0 | fn value(&self, index: usize) -> Self::Item { |
473 | 0 | FixedSizeListArray::value(self, index) |
474 | 0 | } |
475 | | |
476 | 0 | unsafe fn value_unchecked(&self, index: usize) -> Self::Item { |
477 | 0 | FixedSizeListArray::value(self, index) |
478 | 0 | } |
479 | | } |
480 | | |
481 | | #[cfg(test)] |
482 | | mod tests { |
483 | | use arrow_buffer::{bit_util, BooleanBuffer, Buffer}; |
484 | | use arrow_schema::Field; |
485 | | |
486 | | use crate::cast::AsArray; |
487 | | use crate::types::Int32Type; |
488 | | use crate::{new_empty_array, Int32Array}; |
489 | | |
490 | | use super::*; |
491 | | |
492 | | #[test] |
493 | | fn test_fixed_size_list_array() { |
494 | | // Construct a value array |
495 | | let value_data = ArrayData::builder(DataType::Int32) |
496 | | .len(9) |
497 | | .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8])) |
498 | | .build() |
499 | | .unwrap(); |
500 | | |
501 | | // Construct a list array from the above two |
502 | | let list_data_type = |
503 | | DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3); |
504 | | let list_data = ArrayData::builder(list_data_type.clone()) |
505 | | .len(3) |
506 | | .add_child_data(value_data.clone()) |
507 | | .build() |
508 | | .unwrap(); |
509 | | let list_array = FixedSizeListArray::from(list_data); |
510 | | |
511 | | assert_eq!(value_data, list_array.values().to_data()); |
512 | | assert_eq!(DataType::Int32, list_array.value_type()); |
513 | | assert_eq!(3, list_array.len()); |
514 | | assert_eq!(0, list_array.null_count()); |
515 | | assert_eq!(6, list_array.value_offset(2)); |
516 | | assert_eq!(3, list_array.value_length()); |
517 | | assert_eq!(0, list_array.value(0).as_primitive::<Int32Type>().value(0)); |
518 | | for i in 0..3 { |
519 | | assert!(list_array.is_valid(i)); |
520 | | assert!(!list_array.is_null(i)); |
521 | | } |
522 | | |
523 | | // Now test with a non-zero offset |
524 | | let list_data = ArrayData::builder(list_data_type) |
525 | | .len(2) |
526 | | .offset(1) |
527 | | .add_child_data(value_data.clone()) |
528 | | .build() |
529 | | .unwrap(); |
530 | | let list_array = FixedSizeListArray::from(list_data); |
531 | | |
532 | | assert_eq!(value_data.slice(3, 6), list_array.values().to_data()); |
533 | | assert_eq!(DataType::Int32, list_array.value_type()); |
534 | | assert_eq!(2, list_array.len()); |
535 | | assert_eq!(0, list_array.null_count()); |
536 | | assert_eq!(3, list_array.value(0).as_primitive::<Int32Type>().value(0)); |
537 | | assert_eq!(3, list_array.value_offset(1)); |
538 | | assert_eq!(3, list_array.value_length()); |
539 | | } |
540 | | |
541 | | #[test] |
542 | | #[should_panic(expected = "assertion failed: (offset + length) <= self.len()")] |
543 | | // Different error messages, so skip for now |
544 | | // https://github.com/apache/arrow-rs/issues/1545 |
545 | | #[cfg(not(feature = "force_validate"))] |
546 | | fn test_fixed_size_list_array_unequal_children() { |
547 | | // Construct a value array |
548 | | let value_data = ArrayData::builder(DataType::Int32) |
549 | | .len(8) |
550 | | .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7])) |
551 | | .build() |
552 | | .unwrap(); |
553 | | |
554 | | // Construct a list array from the above two |
555 | | let list_data_type = |
556 | | DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3); |
557 | | let list_data = unsafe { |
558 | | ArrayData::builder(list_data_type) |
559 | | .len(3) |
560 | | .add_child_data(value_data) |
561 | | .build_unchecked() |
562 | | }; |
563 | | drop(FixedSizeListArray::from(list_data)); |
564 | | } |
565 | | |
566 | | #[test] |
567 | | fn test_fixed_size_list_array_slice() { |
568 | | // Construct a value array |
569 | | let value_data = ArrayData::builder(DataType::Int32) |
570 | | .len(10) |
571 | | .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])) |
572 | | .build() |
573 | | .unwrap(); |
574 | | |
575 | | // Set null buts for the nested array: |
576 | | // [[0, 1], null, null, [6, 7], [8, 9]] |
577 | | // 01011001 00000001 |
578 | | let mut null_bits: [u8; 1] = [0; 1]; |
579 | | bit_util::set_bit(&mut null_bits, 0); |
580 | | bit_util::set_bit(&mut null_bits, 3); |
581 | | bit_util::set_bit(&mut null_bits, 4); |
582 | | |
583 | | // Construct a fixed size list array from the above two |
584 | | let list_data_type = |
585 | | DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 2); |
586 | | let list_data = ArrayData::builder(list_data_type) |
587 | | .len(5) |
588 | | .add_child_data(value_data.clone()) |
589 | | .null_bit_buffer(Some(Buffer::from(null_bits))) |
590 | | .build() |
591 | | .unwrap(); |
592 | | let list_array = FixedSizeListArray::from(list_data); |
593 | | |
594 | | assert_eq!(value_data, list_array.values().to_data()); |
595 | | assert_eq!(DataType::Int32, list_array.value_type()); |
596 | | assert_eq!(5, list_array.len()); |
597 | | assert_eq!(2, list_array.null_count()); |
598 | | assert_eq!(6, list_array.value_offset(3)); |
599 | | assert_eq!(2, list_array.value_length()); |
600 | | |
601 | | let sliced_array = list_array.slice(1, 4); |
602 | | assert_eq!(4, sliced_array.len()); |
603 | | assert_eq!(2, sliced_array.null_count()); |
604 | | |
605 | | for i in 0..sliced_array.len() { |
606 | | if bit_util::get_bit(&null_bits, 1 + i) { |
607 | | assert!(sliced_array.is_valid(i)); |
608 | | } else { |
609 | | assert!(sliced_array.is_null(i)); |
610 | | } |
611 | | } |
612 | | |
613 | | // Check offset and length for each non-null value. |
614 | | let sliced_list_array = sliced_array |
615 | | .as_any() |
616 | | .downcast_ref::<FixedSizeListArray>() |
617 | | .unwrap(); |
618 | | assert_eq!(2, sliced_list_array.value_length()); |
619 | | assert_eq!(4, sliced_list_array.value_offset(2)); |
620 | | assert_eq!(6, sliced_list_array.value_offset(3)); |
621 | | } |
622 | | |
623 | | #[test] |
624 | | #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")] |
625 | | fn test_fixed_size_list_array_index_out_of_bound() { |
626 | | // Construct a value array |
627 | | let value_data = ArrayData::builder(DataType::Int32) |
628 | | .len(10) |
629 | | .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])) |
630 | | .build() |
631 | | .unwrap(); |
632 | | |
633 | | // Set null buts for the nested array: |
634 | | // [[0, 1], null, null, [6, 7], [8, 9]] |
635 | | // 01011001 00000001 |
636 | | let mut null_bits: [u8; 1] = [0; 1]; |
637 | | bit_util::set_bit(&mut null_bits, 0); |
638 | | bit_util::set_bit(&mut null_bits, 3); |
639 | | bit_util::set_bit(&mut null_bits, 4); |
640 | | |
641 | | // Construct a fixed size list array from the above two |
642 | | let list_data_type = |
643 | | DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 2); |
644 | | let list_data = ArrayData::builder(list_data_type) |
645 | | .len(5) |
646 | | .add_child_data(value_data) |
647 | | .null_bit_buffer(Some(Buffer::from(null_bits))) |
648 | | .build() |
649 | | .unwrap(); |
650 | | let list_array = FixedSizeListArray::from(list_data); |
651 | | |
652 | | list_array.value(10); |
653 | | } |
654 | | |
655 | | #[test] |
656 | | fn test_fixed_size_list_constructors() { |
657 | | let values = Arc::new(Int32Array::from_iter([ |
658 | | Some(1), |
659 | | Some(2), |
660 | | None, |
661 | | None, |
662 | | Some(3), |
663 | | Some(4), |
664 | | ])); |
665 | | |
666 | | let field = Arc::new(Field::new_list_field(DataType::Int32, true)); |
667 | | let list = FixedSizeListArray::new(field.clone(), 2, values.clone(), None); |
668 | | assert_eq!(list.len(), 3); |
669 | | |
670 | | let nulls = NullBuffer::new_null(3); |
671 | | let list = FixedSizeListArray::new(field.clone(), 2, values.clone(), Some(nulls)); |
672 | | assert_eq!(list.len(), 3); |
673 | | |
674 | | let list = FixedSizeListArray::new(field.clone(), 4, values.clone(), None); |
675 | | assert_eq!(list.len(), 1); |
676 | | |
677 | | let err = FixedSizeListArray::try_new(field.clone(), -1, values.clone(), None).unwrap_err(); |
678 | | assert_eq!( |
679 | | err.to_string(), |
680 | | "Invalid argument error: Size cannot be negative, got -1" |
681 | | ); |
682 | | |
683 | | let list = FixedSizeListArray::new(field.clone(), 0, values.clone(), None); |
684 | | assert_eq!(list.len(), 0); |
685 | | |
686 | | let nulls = NullBuffer::new_null(2); |
687 | | let err = FixedSizeListArray::try_new(field, 2, values.clone(), Some(nulls)).unwrap_err(); |
688 | | assert_eq!(err.to_string(), "Invalid argument error: Incorrect length of null buffer for FixedSizeListArray, expected 3 got 2"); |
689 | | |
690 | | let field = Arc::new(Field::new_list_field(DataType::Int32, false)); |
691 | | let err = FixedSizeListArray::try_new(field.clone(), 2, values.clone(), None).unwrap_err(); |
692 | | assert_eq!(err.to_string(), "Invalid argument error: Found unmasked nulls for non-nullable FixedSizeListArray field \"item\""); |
693 | | |
694 | | // Valid as nulls in child masked by parent |
695 | | let nulls = NullBuffer::new(BooleanBuffer::new(Buffer::from([0b0000101]), 0, 3)); |
696 | | FixedSizeListArray::new(field, 2, values.clone(), Some(nulls)); |
697 | | |
698 | | let field = Arc::new(Field::new_list_field(DataType::Int64, true)); |
699 | | let err = FixedSizeListArray::try_new(field, 2, values, None).unwrap_err(); |
700 | | assert_eq!(err.to_string(), "Invalid argument error: FixedSizeListArray expected data type Int64 got Int32 for \"item\""); |
701 | | } |
702 | | |
703 | | #[test] |
704 | | fn empty_fixed_size_list() { |
705 | | let field = Arc::new(Field::new_list_field(DataType::Int32, true)); |
706 | | let nulls = NullBuffer::new_null(2); |
707 | | let values = new_empty_array(&DataType::Int32); |
708 | | let list = FixedSizeListArray::new(field.clone(), 0, values, Some(nulls)); |
709 | | assert_eq!(list.len(), 2); |
710 | | } |
711 | | } |