use itertools::Itertools;
use vortex_buffer::ByteBuffer;
use vortex_dtype::DType;
use vortex_error::{VortexExpect, VortexResult, vortex_bail};
use crate::EmptyMetadata;
use crate::arrays::struct_::{StructArray, StructEncoding, StructVTable};
use crate::serde::ArrayChildren;
use crate::validity::Validity;
use crate::vtable::SerdeVTable;
impl SerdeVTable<StructVTable> for StructVTable {
type Metadata = EmptyMetadata;
fn metadata(_array: &StructArray) -> VortexResult<Option<Self::Metadata>> {
Ok(Some(EmptyMetadata))
}
fn build(
_encoding: &StructEncoding,
dtype: &DType,
len: usize,
_metadata: &Self::Metadata,
_buffers: &[ByteBuffer],
children: &dyn ArrayChildren,
) -> VortexResult<StructArray> {
let DType::Struct(struct_dtype, nullability) = dtype else {
vortex_bail!("Expected struct dtype, found {:?}", dtype)
};
let (validity, non_data_children) = if children.len() == struct_dtype.nfields() {
(Validity::from(*nullability), 0_usize)
} else if children.len() == struct_dtype.nfields() + 1 {
let validity = children.get(0, &Validity::DTYPE, len)?;
(Validity::Array(validity), 1_usize)
} else {
vortex_bail!(
"Expected {} or {} children, found {}",
struct_dtype.nfields(),
struct_dtype.nfields() + 1,
children.len()
);
};
let children: Vec<_> = (0..struct_dtype.nfields())
.map(|i| {
let child_dtype = struct_dtype
.field_by_index(i)
.vortex_expect("no out of bounds");
children.get(non_data_children + i, &child_dtype, len)
})
.try_collect()?;
StructArray::try_new_with_dtype(children, struct_dtype.clone(), len, validity)
}
}