1use crate::array::print_long_array;
19use crate::builder::{FixedSizeListBuilder, PrimitiveBuilder};
20use crate::iterator::FixedSizeListIter;
21use crate::{make_array, Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType};
22use arrow_buffer::buffer::NullBuffer;
23use arrow_buffer::ArrowNativeType;
24use arrow_data::{ArrayData, ArrayDataBuilder};
25use arrow_schema::{ArrowError, DataType, FieldRef};
26use std::any::Any;
27use std::sync::Arc;
28
29#[derive(Clone)]
119pub struct FixedSizeListArray {
120 data_type: DataType, values: ArrayRef,
122 nulls: Option<NullBuffer>,
123 value_length: i32,
124 len: usize,
125}
126
127impl FixedSizeListArray {
128 pub fn new(field: FieldRef, size: i32, values: ArrayRef, nulls: Option<NullBuffer>) -> Self {
134 Self::try_new(field, size, values, nulls).unwrap()
135 }
136
137 pub fn try_new(
146 field: FieldRef,
147 size: i32,
148 values: ArrayRef,
149 nulls: Option<NullBuffer>,
150 ) -> Result<Self, ArrowError> {
151 let s = size.to_usize().ok_or_else(|| {
152 ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {}", size))
153 })?;
154
155 let len = match s {
156 0 => nulls.as_ref().map(|x| x.len()).unwrap_or_default(),
157 _ => {
158 let len = values.len() / s.max(1);
159 if let Some(n) = nulls.as_ref() {
160 if n.len() != len {
161 return Err(ArrowError::InvalidArgumentError(format!(
162 "Incorrect length of null buffer for FixedSizeListArray, expected {} got {}",
163 len,
164 n.len(),
165 )));
166 }
167 }
168 len
169 }
170 };
171
172 if field.data_type() != values.data_type() {
173 return Err(ArrowError::InvalidArgumentError(format!(
174 "FixedSizeListArray expected data type {} got {} for {:?}",
175 field.data_type(),
176 values.data_type(),
177 field.name()
178 )));
179 }
180
181 if let Some(a) = values.logical_nulls() {
182 let nulls_valid = field.is_nullable()
183 || nulls
184 .as_ref()
185 .map(|n| n.expand(size as _).contains(&a))
186 .unwrap_or_default()
187 || (nulls.is_none() && a.null_count() == 0);
188
189 if !nulls_valid {
190 return Err(ArrowError::InvalidArgumentError(format!(
191 "Found unmasked nulls for non-nullable FixedSizeListArray field {:?}",
192 field.name()
193 )));
194 }
195 }
196
197 let data_type = DataType::FixedSizeList(field, size);
198 Ok(Self {
199 data_type,
200 values,
201 value_length: size,
202 nulls,
203 len,
204 })
205 }
206
207 pub fn new_null(field: FieldRef, size: i32, len: usize) -> Self {
216 let capacity = size.to_usize().unwrap().checked_mul(len).unwrap();
217 Self {
218 values: make_array(ArrayData::new_null(field.data_type(), capacity)),
219 data_type: DataType::FixedSizeList(field, size),
220 nulls: Some(NullBuffer::new_null(len)),
221 value_length: size,
222 len,
223 }
224 }
225
226 pub fn into_parts(self) -> (FieldRef, i32, ArrayRef, Option<NullBuffer>) {
228 let f = match self.data_type {
229 DataType::FixedSizeList(f, _) => f,
230 _ => unreachable!(),
231 };
232 (f, self.value_length, self.values, self.nulls)
233 }
234
235 pub fn values(&self) -> &ArrayRef {
237 &self.values
238 }
239
240 pub fn value_type(&self) -> DataType {
242 self.values.data_type().clone()
243 }
244
245 pub fn value(&self, i: usize) -> ArrayRef {
247 self.values
248 .slice(self.value_offset_at(i), self.value_length() as usize)
249 }
250
251 #[inline]
255 pub fn value_offset(&self, i: usize) -> i32 {
256 self.value_offset_at(i) as i32
257 }
258
259 #[inline]
263 pub const fn value_length(&self) -> i32 {
264 self.value_length
265 }
266
267 #[inline]
268 const fn value_offset_at(&self, i: usize) -> usize {
269 i * self.value_length as usize
270 }
271
272 pub fn slice(&self, offset: usize, len: usize) -> Self {
274 assert!(
275 offset.saturating_add(len) <= self.len,
276 "the length + offset of the sliced FixedSizeListArray cannot exceed the existing length"
277 );
278 let size = self.value_length as usize;
279
280 Self {
281 data_type: self.data_type.clone(),
282 values: self.values.slice(offset * size, len * size),
283 nulls: self.nulls.as_ref().map(|n| n.slice(offset, len)),
284 value_length: self.value_length,
285 len,
286 }
287 }
288
289 pub fn from_iter_primitive<T, P, I>(iter: I, length: i32) -> Self
305 where
306 T: ArrowPrimitiveType,
307 P: IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>,
308 I: IntoIterator<Item = Option<P>>,
309 {
310 let l = length as usize;
311 let iter = iter.into_iter();
312 let size_hint = iter.size_hint().0;
313 let mut builder = FixedSizeListBuilder::with_capacity(
314 PrimitiveBuilder::<T>::with_capacity(size_hint * l),
315 length,
316 size_hint,
317 );
318
319 for i in iter {
320 match i {
321 Some(p) => {
322 for t in p {
323 builder.values().append_option(t);
324 }
325 builder.append(true);
326 }
327 None => {
328 builder.values().append_nulls(l);
329 builder.append(false)
330 }
331 }
332 }
333 builder.finish()
334 }
335
336 pub fn iter(&self) -> FixedSizeListIter<'_> {
338 FixedSizeListIter::new(self)
339 }
340}
341
342impl From<ArrayData> for FixedSizeListArray {
343 fn from(data: ArrayData) -> Self {
344 let value_length = match data.data_type() {
345 DataType::FixedSizeList(_, len) => *len,
346 _ => {
347 panic!("FixedSizeListArray data should contain a FixedSizeList data type")
348 }
349 };
350
351 let size = value_length as usize;
352 let values =
353 make_array(data.child_data()[0].slice(data.offset() * size, data.len() * size));
354 Self {
355 data_type: data.data_type().clone(),
356 values,
357 nulls: data.nulls().cloned(),
358 value_length,
359 len: data.len(),
360 }
361 }
362}
363
364impl From<FixedSizeListArray> for ArrayData {
365 fn from(array: FixedSizeListArray) -> Self {
366 let builder = ArrayDataBuilder::new(array.data_type)
367 .len(array.len)
368 .nulls(array.nulls)
369 .child_data(vec![array.values.to_data()]);
370
371 unsafe { builder.build_unchecked() }
372 }
373}
374
375impl Array for FixedSizeListArray {
376 fn as_any(&self) -> &dyn Any {
377 self
378 }
379
380 fn to_data(&self) -> ArrayData {
381 self.clone().into()
382 }
383
384 fn into_data(self) -> ArrayData {
385 self.into()
386 }
387
388 fn data_type(&self) -> &DataType {
389 &self.data_type
390 }
391
392 fn slice(&self, offset: usize, length: usize) -> ArrayRef {
393 Arc::new(self.slice(offset, length))
394 }
395
396 fn len(&self) -> usize {
397 self.len
398 }
399
400 fn is_empty(&self) -> bool {
401 self.len == 0
402 }
403
404 fn offset(&self) -> usize {
405 0
406 }
407
408 fn nulls(&self) -> Option<&NullBuffer> {
409 self.nulls.as_ref()
410 }
411
412 fn logical_null_count(&self) -> usize {
413 self.null_count()
415 }
416
417 fn get_buffer_memory_size(&self) -> usize {
418 let mut size = self.values.get_buffer_memory_size();
419 if let Some(n) = self.nulls.as_ref() {
420 size += n.buffer().capacity();
421 }
422 size
423 }
424
425 fn get_array_memory_size(&self) -> usize {
426 let mut size = std::mem::size_of::<Self>() + self.values.get_array_memory_size();
427 if let Some(n) = self.nulls.as_ref() {
428 size += n.buffer().capacity();
429 }
430 size
431 }
432}
433
434impl ArrayAccessor for FixedSizeListArray {
435 type Item = ArrayRef;
436
437 fn value(&self, index: usize) -> Self::Item {
438 FixedSizeListArray::value(self, index)
439 }
440
441 unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
442 FixedSizeListArray::value(self, index)
443 }
444}
445
446impl std::fmt::Debug for FixedSizeListArray {
447 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
448 write!(f, "FixedSizeListArray<{}>\n[\n", self.value_length())?;
449 print_long_array(self, f, |array, index, f| {
450 std::fmt::Debug::fmt(&array.value(index), f)
451 })?;
452 write!(f, "]")
453 }
454}
455
456impl ArrayAccessor for &FixedSizeListArray {
457 type Item = ArrayRef;
458
459 fn value(&self, index: usize) -> Self::Item {
460 FixedSizeListArray::value(self, index)
461 }
462
463 unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
464 FixedSizeListArray::value(self, index)
465 }
466}
467
468#[cfg(test)]
469mod tests {
470 use arrow_buffer::{bit_util, BooleanBuffer, Buffer};
471 use arrow_schema::Field;
472
473 use crate::cast::AsArray;
474 use crate::types::Int32Type;
475 use crate::{new_empty_array, Int32Array};
476
477 use super::*;
478
479 #[test]
480 fn test_fixed_size_list_array() {
481 let value_data = ArrayData::builder(DataType::Int32)
483 .len(9)
484 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8]))
485 .build()
486 .unwrap();
487
488 let list_data_type =
490 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, false)), 3);
491 let list_data = ArrayData::builder(list_data_type.clone())
492 .len(3)
493 .add_child_data(value_data.clone())
494 .build()
495 .unwrap();
496 let list_array = FixedSizeListArray::from(list_data);
497
498 assert_eq!(value_data, list_array.values().to_data());
499 assert_eq!(DataType::Int32, list_array.value_type());
500 assert_eq!(3, list_array.len());
501 assert_eq!(0, list_array.null_count());
502 assert_eq!(6, list_array.value_offset(2));
503 assert_eq!(3, list_array.value_length());
504 assert_eq!(0, list_array.value(0).as_primitive::<Int32Type>().value(0));
505 for i in 0..3 {
506 assert!(list_array.is_valid(i));
507 assert!(!list_array.is_null(i));
508 }
509
510 let list_data = ArrayData::builder(list_data_type)
512 .len(2)
513 .offset(1)
514 .add_child_data(value_data.clone())
515 .build()
516 .unwrap();
517 let list_array = FixedSizeListArray::from(list_data);
518
519 assert_eq!(value_data.slice(3, 6), list_array.values().to_data());
520 assert_eq!(DataType::Int32, list_array.value_type());
521 assert_eq!(2, list_array.len());
522 assert_eq!(0, list_array.null_count());
523 assert_eq!(3, list_array.value(0).as_primitive::<Int32Type>().value(0));
524 assert_eq!(3, list_array.value_offset(1));
525 assert_eq!(3, list_array.value_length());
526 }
527
528 #[test]
529 #[should_panic(expected = "assertion failed: (offset + length) <= self.len()")]
530 #[cfg(not(feature = "force_validate"))]
533 fn test_fixed_size_list_array_unequal_children() {
534 let value_data = ArrayData::builder(DataType::Int32)
536 .len(8)
537 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
538 .build()
539 .unwrap();
540
541 let list_data_type =
543 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, false)), 3);
544 let list_data = unsafe {
545 ArrayData::builder(list_data_type)
546 .len(3)
547 .add_child_data(value_data)
548 .build_unchecked()
549 };
550 drop(FixedSizeListArray::from(list_data));
551 }
552
553 #[test]
554 fn test_fixed_size_list_array_slice() {
555 let value_data = ArrayData::builder(DataType::Int32)
557 .len(10)
558 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
559 .build()
560 .unwrap();
561
562 let mut null_bits: [u8; 1] = [0; 1];
566 bit_util::set_bit(&mut null_bits, 0);
567 bit_util::set_bit(&mut null_bits, 3);
568 bit_util::set_bit(&mut null_bits, 4);
569
570 let list_data_type =
572 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, false)), 2);
573 let list_data = ArrayData::builder(list_data_type)
574 .len(5)
575 .add_child_data(value_data.clone())
576 .null_bit_buffer(Some(Buffer::from(null_bits)))
577 .build()
578 .unwrap();
579 let list_array = FixedSizeListArray::from(list_data);
580
581 assert_eq!(value_data, list_array.values().to_data());
582 assert_eq!(DataType::Int32, list_array.value_type());
583 assert_eq!(5, list_array.len());
584 assert_eq!(2, list_array.null_count());
585 assert_eq!(6, list_array.value_offset(3));
586 assert_eq!(2, list_array.value_length());
587
588 let sliced_array = list_array.slice(1, 4);
589 assert_eq!(4, sliced_array.len());
590 assert_eq!(2, sliced_array.null_count());
591
592 for i in 0..sliced_array.len() {
593 if bit_util::get_bit(&null_bits, 1 + i) {
594 assert!(sliced_array.is_valid(i));
595 } else {
596 assert!(sliced_array.is_null(i));
597 }
598 }
599
600 let sliced_list_array = sliced_array
602 .as_any()
603 .downcast_ref::<FixedSizeListArray>()
604 .unwrap();
605 assert_eq!(2, sliced_list_array.value_length());
606 assert_eq!(4, sliced_list_array.value_offset(2));
607 assert_eq!(6, sliced_list_array.value_offset(3));
608 }
609
610 #[test]
611 #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
612 fn test_fixed_size_list_array_index_out_of_bound() {
613 let value_data = ArrayData::builder(DataType::Int32)
615 .len(10)
616 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
617 .build()
618 .unwrap();
619
620 let mut null_bits: [u8; 1] = [0; 1];
624 bit_util::set_bit(&mut null_bits, 0);
625 bit_util::set_bit(&mut null_bits, 3);
626 bit_util::set_bit(&mut null_bits, 4);
627
628 let list_data_type =
630 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, false)), 2);
631 let list_data = ArrayData::builder(list_data_type)
632 .len(5)
633 .add_child_data(value_data)
634 .null_bit_buffer(Some(Buffer::from(null_bits)))
635 .build()
636 .unwrap();
637 let list_array = FixedSizeListArray::from(list_data);
638
639 list_array.value(10);
640 }
641
642 #[test]
643 fn test_fixed_size_list_constructors() {
644 let values = Arc::new(Int32Array::from_iter([
645 Some(1),
646 Some(2),
647 None,
648 None,
649 Some(3),
650 Some(4),
651 ]));
652
653 let field = Arc::new(Field::new("item", DataType::Int32, true));
654 let list = FixedSizeListArray::new(field.clone(), 2, values.clone(), None);
655 assert_eq!(list.len(), 3);
656
657 let nulls = NullBuffer::new_null(3);
658 let list = FixedSizeListArray::new(field.clone(), 2, values.clone(), Some(nulls));
659 assert_eq!(list.len(), 3);
660
661 let list = FixedSizeListArray::new(field.clone(), 4, values.clone(), None);
662 assert_eq!(list.len(), 1);
663
664 let err = FixedSizeListArray::try_new(field.clone(), -1, values.clone(), None).unwrap_err();
665 assert_eq!(
666 err.to_string(),
667 "Invalid argument error: Size cannot be negative, got -1"
668 );
669
670 let list = FixedSizeListArray::new(field.clone(), 0, values.clone(), None);
671 assert_eq!(list.len(), 0);
672
673 let nulls = NullBuffer::new_null(2);
674 let err = FixedSizeListArray::try_new(field, 2, values.clone(), Some(nulls)).unwrap_err();
675 assert_eq!(err.to_string(), "Invalid argument error: Incorrect length of null buffer for FixedSizeListArray, expected 3 got 2");
676
677 let field = Arc::new(Field::new("item", DataType::Int32, false));
678 let err = FixedSizeListArray::try_new(field.clone(), 2, values.clone(), None).unwrap_err();
679 assert_eq!(err.to_string(), "Invalid argument error: Found unmasked nulls for non-nullable FixedSizeListArray field \"item\"");
680
681 let nulls = NullBuffer::new(BooleanBuffer::new(Buffer::from([0b0000101]), 0, 3));
683 FixedSizeListArray::new(field, 2, values.clone(), Some(nulls));
684
685 let field = Arc::new(Field::new("item", DataType::Int64, true));
686 let err = FixedSizeListArray::try_new(field, 2, values, None).unwrap_err();
687 assert_eq!(err.to_string(), "Invalid argument error: FixedSizeListArray expected data type Int64 got Int32 for \"item\"");
688 }
689
690 #[test]
691 fn empty_fixed_size_list() {
692 let field = Arc::new(Field::new("item", DataType::Int32, true));
693 let nulls = NullBuffer::new_null(2);
694 let values = new_empty_array(&DataType::Int32);
695 let list = FixedSizeListArray::new(field.clone(), 0, values, Some(nulls));
696 assert_eq!(list.len(), 2);
697 }
698}