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 {
253 self.values
254 .slice(self.value_offset_at(i), self.value_length() as usize)
255 }
256
257 #[inline]
261 pub fn value_offset(&self, i: usize) -> i32 {
262 self.value_offset_at(i) as i32
263 }
264
265 #[inline]
269 pub const fn value_length(&self) -> i32 {
270 self.value_length
271 }
272
273 #[inline]
274 const fn value_offset_at(&self, i: usize) -> usize {
275 i * self.value_length as usize
276 }
277
278 pub fn slice(&self, offset: usize, len: usize) -> Self {
280 assert!(
281 offset.saturating_add(len) <= self.len,
282 "the length + offset of the sliced FixedSizeListArray cannot exceed the existing length"
283 );
284 let size = self.value_length as usize;
285
286 Self {
287 data_type: self.data_type.clone(),
288 values: self.values.slice(offset * size, len * size),
289 nulls: self.nulls.as_ref().map(|n| n.slice(offset, len)),
290 value_length: self.value_length,
291 len,
292 }
293 }
294
295 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 pub fn iter(&self) -> FixedSizeListIter<'_> {
344 FixedSizeListIter::new(self)
345 }
346}
347
348impl From<ArrayData> for FixedSizeListArray {
349 fn from(data: ArrayData) -> Self {
350 let value_length = match data.data_type() {
351 DataType::FixedSizeList(_, len) => *len,
352 data_type => {
353 panic!("FixedSizeListArray data should contain a FixedSizeList data type, got {data_type}")
354 }
355 };
356
357 let size = value_length as usize;
358 let values =
359 make_array(data.child_data()[0].slice(data.offset() * size, data.len() * size));
360 Self {
361 data_type: data.data_type().clone(),
362 values,
363 nulls: data.nulls().cloned(),
364 value_length,
365 len: data.len(),
366 }
367 }
368}
369
370impl From<FixedSizeListArray> for ArrayData {
371 fn from(array: FixedSizeListArray) -> Self {
372 let builder = ArrayDataBuilder::new(array.data_type)
373 .len(array.len)
374 .nulls(array.nulls)
375 .child_data(vec![array.values.to_data()]);
376
377 unsafe { builder.build_unchecked() }
378 }
379}
380
381impl Array for FixedSizeListArray {
382 fn as_any(&self) -> &dyn Any {
383 self
384 }
385
386 fn to_data(&self) -> ArrayData {
387 self.clone().into()
388 }
389
390 fn into_data(self) -> ArrayData {
391 self.into()
392 }
393
394 fn data_type(&self) -> &DataType {
395 &self.data_type
396 }
397
398 fn slice(&self, offset: usize, length: usize) -> ArrayRef {
399 Arc::new(self.slice(offset, length))
400 }
401
402 fn len(&self) -> usize {
403 self.len
404 }
405
406 fn is_empty(&self) -> bool {
407 self.len == 0
408 }
409
410 fn shrink_to_fit(&mut self) {
411 self.values.shrink_to_fit();
412 if let Some(nulls) = &mut self.nulls {
413 nulls.shrink_to_fit();
414 }
415 }
416
417 fn offset(&self) -> usize {
418 0
419 }
420
421 fn nulls(&self) -> Option<&NullBuffer> {
422 self.nulls.as_ref()
423 }
424
425 fn logical_null_count(&self) -> usize {
426 self.null_count()
428 }
429
430 fn get_buffer_memory_size(&self) -> usize {
431 let mut size = self.values.get_buffer_memory_size();
432 if let Some(n) = self.nulls.as_ref() {
433 size += n.buffer().capacity();
434 }
435 size
436 }
437
438 fn get_array_memory_size(&self) -> usize {
439 let mut size = std::mem::size_of::<Self>() + self.values.get_array_memory_size();
440 if let Some(n) = self.nulls.as_ref() {
441 size += n.buffer().capacity();
442 }
443 size
444 }
445}
446
447impl ArrayAccessor for FixedSizeListArray {
448 type Item = ArrayRef;
449
450 fn value(&self, index: usize) -> Self::Item {
451 FixedSizeListArray::value(self, index)
452 }
453
454 unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
455 FixedSizeListArray::value(self, index)
456 }
457}
458
459impl std::fmt::Debug for FixedSizeListArray {
460 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
461 write!(f, "FixedSizeListArray<{}>\n[\n", self.value_length())?;
462 print_long_array(self, f, |array, index, f| {
463 std::fmt::Debug::fmt(&array.value(index), f)
464 })?;
465 write!(f, "]")
466 }
467}
468
469impl ArrayAccessor for &FixedSizeListArray {
470 type Item = ArrayRef;
471
472 fn value(&self, index: usize) -> Self::Item {
473 FixedSizeListArray::value(self, index)
474 }
475
476 unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
477 FixedSizeListArray::value(self, index)
478 }
479}
480
481#[cfg(test)]
482mod 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 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 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 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 #[cfg(not(feature = "force_validate"))]
546 fn test_fixed_size_list_array_unequal_children() {
547 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 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 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 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 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 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 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 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 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 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}