Skip to content

Commit

Permalink
Use FixedSizeListArray::new in FixedSizeListBuilder
Browse files Browse the repository at this point in the history
  • Loading branch information
@tustvold
tustvold committed Apr 9, 2024
1 parent 1b3d1a9 commit 999c746
Showing 1 changed file with 22 additions and 92 deletions.
114 changes: 22 additions & 92 deletions arrow-array/src/builder/fixed_size_list_builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
use crate::builder::ArrayBuilder;
use crate::{ArrayRef, FixedSizeListArray};
use arrow_buffer::NullBufferBuilder;
use arrow_data::ArrayData;
use arrow_schema::{DataType, Field, FieldRef};
use std::any::Any;
use std::sync::Arc;
Expand Down Expand Up @@ -169,110 +168,45 @@ where
/// Builds the [`FixedSizeListBuilder`] and reset this builder.
pub fn finish(&mut self) -> FixedSizeListArray {
let len = self.len();
let values_arr = self.values_builder.finish();
let values_data = values_arr.to_data();
let values = self.values_builder.finish();
let nulls = self.null_buffer_builder.finish();

assert_eq!(
values_data.len(), len * self.list_len as usize,
values.len(), len * self.list_len as usize,
"Length of the child array ({}) must be the multiple of the value length ({}) and the array length ({}).",
values_data.len(),
values.len(),
self.list_len,
len,
);

let nulls = self.null_buffer_builder.finish();
let field = self
.field
.clone()
.unwrap_or_else(|| Arc::new(Field::new("item", values.data_type().clone(), true)));

let field = match &self.field {
Some(f) => {
let size = self.value_length();
assert_eq!(
f.data_type(),
values_data.data_type(),
"DataType of field ({}) should be the same as the values_builder DataType ({})",
f.data_type(),
values_data.data_type()
);

if let Some(a) = values_arr.logical_nulls() {
let nulls_valid = f.is_nullable()
|| nulls
.as_ref()
.map(|n| n.expand(size as _).contains(&a))
.unwrap_or_default();

assert!(
nulls_valid,
"Found unmasked nulls for non-nullable FixedSizeListBuilder field {:?}",
f.name()
);
}
f.clone()
}
None => Arc::new(Field::new("item", values_data.data_type().clone(), true)),
};

let array_data = ArrayData::builder(DataType::FixedSizeList(field, self.list_len))
.len(len)
.add_child_data(values_data)
.nulls(nulls);

let array_data = unsafe { array_data.build_unchecked() };

FixedSizeListArray::from(array_data)
FixedSizeListArray::new(field, self.list_len, values, nulls)
}

/// Builds the [`FixedSizeListBuilder`] without resetting the builder.
pub fn finish_cloned(&self) -> FixedSizeListArray {
let len = self.len();
let values_arr = self.values_builder.finish_cloned();
let values_data = values_arr.to_data();
let values = self.values_builder.finish_cloned();
let nulls = self.null_buffer_builder.finish_cloned();

assert_eq!(
values_data.len(), len * self.list_len as usize,
values.len(), len * self.list_len as usize,
"Length of the child array ({}) must be the multiple of the value length ({}) and the array length ({}).",
values_data.len(),
values.len(),
self.list_len,
len,
);

let nulls = self.null_buffer_builder.finish_cloned();
let field = self
.field
.clone()
.unwrap_or_else(|| Arc::new(Field::new("item", values.data_type().clone(), true)));

let field = match &self.field {
Some(f) => {
let size = self.value_length();
assert_eq!(
f.data_type(),
values_data.data_type(),
"DataType of field ({}) should be the same as the values_builder DataType ({})",
f.data_type(),
values_data.data_type()
);
if let Some(a) = values_arr.logical_nulls() {
let nulls_valid = f.is_nullable()
|| nulls
.as_ref()
.map(|n| n.expand(size as _).contains(&a))
.unwrap_or_default();

assert!(
nulls_valid,
"Found unmasked nulls for non-nullable FixedSizeListBuilder field {:?}",
f.name()
);
}
f.clone()
}
None => Arc::new(Field::new("item", values_data.data_type().clone(), true)),
};

let array_data = ArrayData::builder(DataType::FixedSizeList(field, self.list_len))
.len(len)
.add_child_data(values_data)
.nulls(nulls);

let array_data = unsafe { array_data.build_unchecked() };

FixedSizeListArray::from(array_data)
FixedSizeListArray::new(field, self.list_len, values, nulls)
}
}

Expand Down Expand Up @@ -368,7 +302,7 @@ mod tests {
}

#[test]
#[should_panic(expected = "Found unmasked nulls for non-nullable FixedSizeListBuilder field")]
#[should_panic(expected = "Found unmasked nulls for non-nullable FixedSizeListArray")]
fn test_fixed_size_list_array_builder_with_field_null_panic() {
let builder = make_list_builder(true, true);
let mut builder = builder.with_field(Field::new("list_item", DataType::Int32, false));
Expand All @@ -377,9 +311,7 @@ mod tests {
}

#[test]
#[should_panic(
expected = "DataType of field (Int64) should be the same as the values_builder DataType (Int32)"
)]
#[should_panic(expected = "FixedSizeListArray expected data type Int64 got Int32")]
fn test_fixed_size_list_array_builder_with_field_type_panic() {
let values_builder = Int32Builder::new();
let builder = FixedSizeListBuilder::new(values_builder, 3);
Expand Down Expand Up @@ -417,7 +349,7 @@ mod tests {
}

#[test]
#[should_panic(expected = "Found unmasked nulls for non-nullable FixedSizeListBuilder field")]
#[should_panic(expected = "Found unmasked nulls for non-nullable FixedSizeListArray")]
fn test_fixed_size_list_array_builder_cloned_with_field_null_panic() {
let builder = make_list_builder(true, true);
let builder = builder.with_field(Field::new("list_item", DataType::Int32, false));
Expand All @@ -439,9 +371,7 @@ mod tests {
}

#[test]
#[should_panic(
expected = "DataType of field (Int64) should be the same as the values_builder DataType (Int32)"
)]
#[should_panic(expected = "FixedSizeListArray expected data type Int64 got Int32")]
fn test_fixed_size_list_array_builder_cloned_with_field_type_panic() {
let builder = make_list_builder(false, false);
let builder = builder.with_field(Field::new("list_item", DataType::Int64, true));
Expand Down

0 comments on commit 999c746

Please sign in to comment.