diff --git a/crates/re_sdk/src/recording_stream.rs b/crates/re_sdk/src/recording_stream.rs index e386b51eace1..240a688b95f9 100644 --- a/crates/re_sdk/src/recording_stream.rs +++ b/crates/re_sdk/src/recording_stream.rs @@ -1144,21 +1144,7 @@ mod tests { _ => panic!("expected SetStoreInfo"), } - // The underlying batcher is never flushing: there's nothing else. - assert!(msgs.pop().is_none()); - } - - // The underlying batcher is never flushing: there's nothing else. - assert!(storage.take().is_empty()); - - rec_stream.flush_blocking(); // flush the entire hierarchy - - { - let mut msgs = { - let mut msgs = storage.take(); - msgs.reverse(); - msgs - }; + // MemorySinkStorage transparently handles flushing during `take()`! // The batched table itself, which was sent as a result of the explicit flush above. match msgs.pop().unwrap() { diff --git a/crates/re_types/src/datatypes/fuzzy.rs b/crates/re_types/src/datatypes/fuzzy.rs index dfbbe5b6c8de..85a660284199 100644 --- a/crates/re_types/src/datatypes/fuzzy.rs +++ b/crates/re_types/src/datatypes/fuzzy.rs @@ -1152,30 +1152,30 @@ impl crate::Datatype for AffixFuzzer2 { } #[derive(Debug, Clone, PartialEq)] -pub enum AffixFuzzer3 { - Degrees(f32), - Radians(Option), - Craziness(Vec), +pub enum AffixFuzzer4 { + SingleRequired(crate::datatypes::AffixFuzzer3), + ManyRequired(Vec), + ManyOptional(Option>), } -impl<'a> From for ::std::borrow::Cow<'a, AffixFuzzer3> { +impl<'a> From for ::std::borrow::Cow<'a, AffixFuzzer4> { #[inline] - fn from(value: AffixFuzzer3) -> Self { + fn from(value: AffixFuzzer4) -> Self { std::borrow::Cow::Owned(value) } } -impl<'a> From<&'a AffixFuzzer3> for ::std::borrow::Cow<'a, AffixFuzzer3> { +impl<'a> From<&'a AffixFuzzer4> for ::std::borrow::Cow<'a, AffixFuzzer4> { #[inline] - fn from(value: &'a AffixFuzzer3) -> Self { + fn from(value: &'a AffixFuzzer4) -> Self { std::borrow::Cow::Borrowed(value) } } -impl crate::Datatype for AffixFuzzer3 { +impl crate::Datatype for AffixFuzzer4 { #[inline] fn name() -> crate::DatatypeName { - crate::DatatypeName::Borrowed("rerun.testing.datatypes.AffixFuzzer3") + crate::DatatypeName::Borrowed("rerun.testing.datatypes.AffixFuzzer4") } #[allow(unused_imports, clippy::wildcard_imports)] @@ -1185,242 +1185,129 @@ impl crate::Datatype for AffixFuzzer3 { DataType::Union( vec![ Field { - name: "degrees".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "radians".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "craziness".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Struct(vec![ + name: "single_required".to_owned(), + data_type: DataType::Union( + vec![ Field { - name: "single_float_optional".to_owned(), + name: "degrees".to_owned(), data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "single_string_required".to_owned(), - data_type: DataType::Utf8, is_nullable: false, metadata: [].into(), }, Field { - name: "single_string_optional".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_floats_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - })), - is_nullable: true, + name: "radians".to_owned(), + data_type: DataType::Float32, + is_nullable: false, metadata: [].into(), }, Field { - name: "many_strings_required".to_owned(), + name: "craziness".to_owned(), data_type: DataType::List(Box::new(Field { name: "item".to_owned(), - data_type: DataType::Utf8, + data_type: DataType::Struct(vec![ + Field { + name: "single_float_optional".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "single_string_required".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "single_string_optional".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_floats_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_strings_required".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_strings_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "flattened_scalar".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "almost_flattened_scalar".to_owned(), + data_type: DataType::Struct(vec![Field { + name: "value".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }]), + is_nullable: false, + metadata: [].into(), + }, + ]), is_nullable: false, metadata: [].into(), })), is_nullable: false, metadata: [].into(), }, - Field { - name: "many_strings_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, + ], + None, + UnionMode::Dense, + ), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_required".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Union( + vec![ + Field { + name: "degrees".to_owned(), + data_type: DataType::Float32, + is_nullable: false, metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "flattened_scalar".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "almost_flattened_scalar".to_owned(), - data_type: DataType::Struct(vec![Field { - name: "value".to_owned(), + }, + Field { + name: "radians".to_owned(), data_type: DataType::Float32, is_nullable: false, metadata: [].into(), - }]), - is_nullable: false, - metadata: [].into(), - }, - ]), - is_nullable: false, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - ], - None, - UnionMode::Dense, - ) - } - - #[allow(unused_imports, clippy::wildcard_imports)] - fn try_to_arrow_opt<'a>( - data: impl IntoIterator>>>, - extension_wrapper: Option<&str>, - ) -> crate::SerializationResult> - where - Self: Clone + 'a, - { - use crate::{Component as _, Datatype as _}; - use ::arrow2::{array::*, datatypes::*}; - Ok({ - let data: Vec<_> = data - .into_iter() - .map(|datum| { - let datum: Option<::std::borrow::Cow<'a, Self>> = datum.map(Into::into); - datum - }) - .collect(); - UnionArray::new( - (if let Some(ext) = extension_wrapper { - DataType::Extension( - ext.to_owned(), - Box::new(::to_arrow_datatype()), - None, - ) - } else { - ::to_arrow_datatype() - }) - .to_logical_type() - .clone(), - { - data.iter() - .flatten() - .map(|v| match **v { - AffixFuzzer3::Degrees(_) => 0i8, - AffixFuzzer3::Radians(_) => 1i8, - AffixFuzzer3::Craziness(_) => 2i8, - }) - .collect() - }, - vec![ - { - let (somes, degrees): (Vec<_>, Vec<_>) = data - .iter() - .flatten() - .filter(|datum| matches!(***datum, AffixFuzzer3::Degrees(_))) - .map(|datum| { - let datum = match &**datum { - AffixFuzzer3::Degrees(v) => Some(v.clone()), - _ => None, - }; - (datum.is_some(), datum) - }) - .unzip(); - let degrees_bitmap: Option<::arrow2::bitmap::Bitmap> = { - let any_nones = somes.iter().any(|some| !*some); - any_nones.then(|| somes.into()) - }; - PrimitiveArray::new( - { - _ = extension_wrapper; - DataType::Float32.to_logical_type().clone() - }, - degrees.into_iter().map(|v| v.unwrap_or_default()).collect(), - degrees_bitmap, - ) - .boxed() - }, - { - let (somes, radians): (Vec<_>, Vec<_>) = data - .iter() - .flatten() - .filter(|datum| matches!(***datum, AffixFuzzer3::Radians(_))) - .map(|datum| { - let datum = match &**datum { - AffixFuzzer3::Radians(v) => Some(v.clone()), - _ => None, - } - .flatten(); - (datum.is_some(), datum) - }) - .unzip(); - let radians_bitmap: Option<::arrow2::bitmap::Bitmap> = { - let any_nones = somes.iter().any(|some| !*some); - any_nones.then(|| somes.into()) - }; - PrimitiveArray::new( - { - _ = extension_wrapper; - DataType::Float32.to_logical_type().clone() - }, - radians.into_iter().map(|v| v.unwrap_or_default()).collect(), - radians_bitmap, - ) - .boxed() - }, - { - let (somes, craziness): (Vec<_>, Vec<_>) = data - .iter() - .flatten() - .filter(|datum| matches!(***datum, AffixFuzzer3::Craziness(_))) - .map(|datum| { - let datum = match &**datum { - AffixFuzzer3::Craziness(v) => Some(v.clone()), - _ => None, - }; - (datum.is_some(), datum) - }) - .unzip(); - let craziness_bitmap: Option<::arrow2::bitmap::Bitmap> = { - let any_nones = somes.iter().any(|some| !*some); - any_nones.then(|| somes.into()) - }; - { - use arrow2::{buffer::Buffer, offset::OffsetsBuffer}; - let craziness_inner_data: Vec<_> = craziness - .iter() - .flatten() - .flatten() - .map(ToOwned::to_owned) - .map(Some) - .collect(); - let craziness_inner_bitmap: Option<::arrow2::bitmap::Bitmap> = { - let any_nones = craziness_inner_data.iter().any(|v| v.is_none()); - any_nones.then(|| { - craziness_inner_data.iter().map(|v| v.is_some()).collect() - }) - }; - let offsets = ::arrow2::offset::Offsets::::try_from_lengths( - craziness.iter().map(|opt| { - opt.as_ref().map(|datum| datum.len()).unwrap_or_default() - }), - ) - .unwrap() - .into(); - ListArray::new( - { - _ = extension_wrapper; - DataType::List(Box::new(Field { + }, + Field { + name: "craziness".to_owned(), + data_type: DataType::List(Box::new(Field { name: "item".to_owned(), data_type: DataType::Struct(vec![ Field { @@ -1494,744 +1381,444 @@ impl crate::Datatype for AffixFuzzer3 { ]), is_nullable: false, metadata: [].into(), - })) - .to_logical_type() - .clone() + })), + is_nullable: false, + metadata: [].into(), }, - offsets, - { - _ = craziness_inner_bitmap; - _ = extension_wrapper; - crate::datatypes::AffixFuzzer1::try_to_arrow_opt( - craziness_inner_data, - None::<&str>, - )? + ], + None, + UnionMode::Dense, + ), + is_nullable: false, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Union( + vec![ + Field { + name: "degrees".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), }, - craziness_bitmap, - ) - .boxed() - } - }, - ], - Some({ - let mut degrees_offset = 0; - let mut radians_offset = 0; - let mut craziness_offset = 0; - data.iter() - .flatten() - .map(|v| match **v { - AffixFuzzer3::Degrees(_) => { - let offset = degrees_offset; - degrees_offset += 1; - offset - } - - AffixFuzzer3::Radians(_) => { - let offset = radians_offset; - radians_offset += 1; - offset - } - - AffixFuzzer3::Craziness(_) => { - let offset = craziness_offset; - craziness_offset += 1; - offset - } - }) - .collect() - }), - ) - .boxed() - }) + Field { + name: "radians".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "craziness".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Struct(vec![ + Field { + name: "single_float_optional".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "single_string_required".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "single_string_optional".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_floats_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_strings_required".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_strings_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "flattened_scalar".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "almost_flattened_scalar".to_owned(), + data_type: DataType::Struct(vec![Field { + name: "value".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }]), + is_nullable: false, + metadata: [].into(), + }, + ]), + is_nullable: false, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + ], + None, + UnionMode::Dense, + ), + is_nullable: true, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + ], + None, + UnionMode::Dense, + ) } #[allow(unused_imports, clippy::wildcard_imports)] - fn try_from_arrow_opt( - data: &dyn ::arrow2::array::Array, - ) -> crate::DeserializationResult>> - where - Self: Sized, + fn try_to_arrow_opt<'a>( + data: impl IntoIterator>>>, + extension_wrapper: Option<&str>, + ) -> crate::SerializationResult> + where + Self: Clone + 'a, { use crate::{Component as _, Datatype as _}; use ::arrow2::{array::*, datatypes::*}; Ok({ - let data = data - .as_any() - .downcast_ref::<::arrow2::array::UnionArray>() - .ok_or_else(|| crate::DeserializationError::SchemaMismatch { - expected: data.data_type().clone(), - got: data.data_type().clone(), - })?; - let (data_types, data_arrays, data_offsets) = - (data.types(), data.fields(), data.offsets().unwrap()); - let degrees = { - let data = &*data_arrays[0usize]; - - data.as_any() - .downcast_ref::() - .unwrap() - .into_iter() - .map(|v| v.copied()) - .collect::>() - }; - let radians = { - let data = &*data_arrays[1usize]; - - data.as_any() - .downcast_ref::() - .unwrap() - .into_iter() - .map(|v| v.copied()) - .collect::>() - }; - let craziness = { - let data = &*data_arrays[2usize]; - + let data: Vec<_> = data + .into_iter() + .map(|datum| { + let datum: Option<::std::borrow::Cow<'a, Self>> = datum.map(Into::into); + datum + }) + .collect(); + UnionArray::new( + (if let Some(ext) = extension_wrapper { + DataType::Extension( + ext.to_owned(), + Box::new(::to_arrow_datatype()), + None, + ) + } else { + ::to_arrow_datatype() + }) + .to_logical_type() + .clone(), { - let datatype = data.data_type(); - let data = data - .as_any() - .downcast_ref::<::arrow2::array::ListArray>() - .unwrap(); - let bitmap = data.validity().cloned(); - let offsets = { - let offsets = data.offsets(); - offsets.iter().copied().zip(offsets.iter().copied().skip(1)) - }; - let data = &**data.values(); - let data = crate::datatypes::AffixFuzzer1::try_from_arrow_opt(data)? - .into_iter() - .map(|v| { - v.ok_or_else(|| crate::DeserializationError::MissingData { - datatype: DataType::Struct(vec![ - Field { - name: "single_float_optional".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "single_string_required".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "single_string_optional".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_floats_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_strings_required".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_strings_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "flattened_scalar".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "almost_flattened_scalar".to_owned(), - data_type: DataType::Struct(vec![Field { - name: "value".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }]), - is_nullable: false, - metadata: [].into(), - }, - ]), - }) + data.iter() + .flatten() + .map(|v| match **v { + AffixFuzzer4::SingleRequired(_) => 0i8, + AffixFuzzer4::ManyRequired(_) => 1i8, + AffixFuzzer4::ManyOptional(_) => 2i8, }) - .collect::>>()?; - offsets - .enumerate() - .map(move |(i, (start, end))| { - bitmap - .as_ref() - .map_or(true, |bitmap| bitmap.get_bit(i)) - .then(|| { - Ok(data - .get(start as usize..end as usize) - .ok_or_else(|| { - crate::DeserializationError::OffsetsMismatch { - bounds: (start as usize, end as usize), - len: data.len(), - datatype: datatype.clone(), - } - })? - .to_vec()) + .collect() + }, + vec![ + { + let (somes, single_required): (Vec<_>, Vec<_>) = data + .iter() + .flatten() + .filter(|datum| matches!(***datum, AffixFuzzer4::SingleRequired(_))) + .map(|datum| { + let datum = match &**datum { + AffixFuzzer4::SingleRequired(v) => Some(v.clone()), + _ => None, + }; + (datum.is_some(), datum) + }) + .unzip(); + let single_required_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + { + _ = single_required_bitmap; + _ = extension_wrapper; + crate::datatypes::AffixFuzzer3::try_to_arrow_opt( + single_required, + None::<&str>, + )? + } + }, + { + let (somes, many_required): (Vec<_>, Vec<_>) = data + .iter() + .flatten() + .filter(|datum| matches!(***datum, AffixFuzzer4::ManyRequired(_))) + .map(|datum| { + let datum = match &**datum { + AffixFuzzer4::ManyRequired(v) => Some(v.clone()), + _ => None, + }; + (datum.is_some(), datum) + }) + .unzip(); + let many_required_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + { + use arrow2::{buffer::Buffer, offset::OffsetsBuffer}; + let many_required_inner_data: Vec<_> = many_required + .iter() + .flatten() + .flatten() + .map(ToOwned::to_owned) + .map(Some) + .collect(); + let many_required_inner_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = + many_required_inner_data.iter().any(|v| v.is_none()); + any_nones.then(|| { + many_required_inner_data + .iter() + .map(|v| v.is_some()) + .collect() }) - .transpose() - }) - .collect::>>>()? - .into_iter() - } - .collect::>() - }; - data_types - .iter() - .enumerate() - .map(|(i, typ)| { - let offset = data_offsets[i]; - - Ok(Some(match typ { - 0i8 => AffixFuzzer3::Degrees( - degrees - .get(offset as usize) - .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { - bounds: (offset as usize, offset as usize), - len: degrees.len(), - datatype: data.data_type().clone(), - })? - .clone() - .unwrap(), - ), - 1i8 => AffixFuzzer3::Radians( - radians - .get(offset as usize) - .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { - bounds: (offset as usize, offset as usize), - len: radians.len(), - datatype: data.data_type().clone(), - })? - .clone(), - ), - 2i8 => AffixFuzzer3::Craziness( - craziness - .get(offset as usize) - .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { - bounds: (offset as usize, offset as usize), - len: craziness.len(), - datatype: data.data_type().clone(), - })? - .clone() - .unwrap(), - ), - _ => unreachable!(), - })) - }) - .collect::>>()? - }) - } -} - -#[derive(Debug, Clone, PartialEq)] -pub enum AffixFuzzer4 { - SingleRequired(crate::datatypes::AffixFuzzer3), - ManyRequired(Vec), - ManyOptional(Option>), -} - -impl<'a> From for ::std::borrow::Cow<'a, AffixFuzzer4> { - #[inline] - fn from(value: AffixFuzzer4) -> Self { - std::borrow::Cow::Owned(value) - } -} - -impl<'a> From<&'a AffixFuzzer4> for ::std::borrow::Cow<'a, AffixFuzzer4> { - #[inline] - fn from(value: &'a AffixFuzzer4) -> Self { - std::borrow::Cow::Borrowed(value) - } -} - -impl crate::Datatype for AffixFuzzer4 { - #[inline] - fn name() -> crate::DatatypeName { - crate::DatatypeName::Borrowed("rerun.testing.datatypes.AffixFuzzer4") - } - - #[allow(unused_imports, clippy::wildcard_imports)] - #[inline] - fn to_arrow_datatype() -> arrow2::datatypes::DataType { - use ::arrow2::datatypes::*; - DataType::Union( - vec![ - Field { - name: "single_required".to_owned(), - data_type: DataType::Union( - vec![ - Field { - name: "degrees".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "radians".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "craziness".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Struct(vec![ - Field { - name: "single_float_optional".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "single_string_required".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "single_string_optional".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_floats_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_strings_required".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_strings_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "flattened_scalar".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "almost_flattened_scalar".to_owned(), - data_type: DataType::Struct(vec![Field { - name: "value".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }]), - is_nullable: false, - metadata: [].into(), - }, - ]), - is_nullable: false, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - ], - None, - UnionMode::Dense, - ), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_required".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Union( - vec![ - Field { - name: "degrees".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "radians".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "craziness".to_owned(), - data_type: DataType::List(Box::new(Field { + }; + let offsets = ::arrow2::offset::Offsets::::try_from_lengths( + many_required.iter().map(|opt| { + opt.as_ref().map(|datum| datum.len()).unwrap_or_default() + }), + ) + .unwrap() + .into(); + ListArray::new( + { + _ = extension_wrapper; + DataType::List(Box::new(Field { name: "item".to_owned(), - data_type: DataType::Struct(vec![ - Field { - name: "single_float_optional".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "single_string_required".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "single_string_optional".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_floats_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), + data_type: DataType::Union( + vec![ + Field { + name: "degrees".to_owned(), data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_strings_required".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, is_nullable: false, metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_strings_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, + }, + Field { + name: "radians".to_owned(), + data_type: DataType::Float32, + is_nullable: false, metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "flattened_scalar".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "almost_flattened_scalar".to_owned(), - data_type: DataType::Struct(vec![Field { - name: "value".to_owned(), - data_type: DataType::Float32, + }, + Field { + name: "craziness".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Struct(vec![ + Field { + name: "single_float_optional" + .to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "single_string_required" + .to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "single_string_optional" + .to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_floats_optional" + .to_owned(), + data_type: DataType::List( + Box::new(Field { + name: "item".to_owned(), + data_type: + DataType::Float32, + is_nullable: true, + metadata: [].into(), + }), + ), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_strings_required" + .to_owned(), + data_type: DataType::List( + Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }), + ), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_strings_optional" + .to_owned(), + data_type: DataType::List( + Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }), + ), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "flattened_scalar".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "almost_flattened_scalar" + .to_owned(), + data_type: DataType::Struct(vec![ + Field { + name: "value".to_owned(), + data_type: + DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + ]), + is_nullable: false, + metadata: [].into(), + }, + ]), + is_nullable: false, + metadata: [].into(), + })), is_nullable: false, metadata: [].into(), - }]), - is_nullable: false, - metadata: [].into(), - }, - ]), + }, + ], + None, + UnionMode::Dense, + ), is_nullable: false, metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - ], - None, - UnionMode::Dense, - ), - is_nullable: false, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Union( - vec![ - Field { - name: "degrees".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), + })) + .to_logical_type() + .clone() }, - Field { - name: "radians".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), + offsets, + { + _ = many_required_inner_bitmap; + _ = extension_wrapper; + crate::datatypes::AffixFuzzer3::try_to_arrow_opt( + many_required_inner_data, + None::<&str>, + )? }, - Field { - name: "craziness".to_owned(), - data_type: DataType::List(Box::new(Field { + many_required_bitmap, + ) + .boxed() + } + }, + { + let (somes, many_optional): (Vec<_>, Vec<_>) = data + .iter() + .flatten() + .filter(|datum| matches!(***datum, AffixFuzzer4::ManyOptional(_))) + .map(|datum| { + let datum = match &**datum { + AffixFuzzer4::ManyOptional(v) => Some(v.clone()), + _ => None, + } + .flatten(); + (datum.is_some(), datum) + }) + .unzip(); + let many_optional_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + { + use arrow2::{buffer::Buffer, offset::OffsetsBuffer}; + let many_optional_inner_data: Vec<_> = many_optional + .iter() + .flatten() + .flatten() + .map(ToOwned::to_owned) + .map(Some) + .collect(); + let many_optional_inner_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = + many_optional_inner_data.iter().any(|v| v.is_none()); + any_nones.then(|| { + many_optional_inner_data + .iter() + .map(|v| v.is_some()) + .collect() + }) + }; + let offsets = ::arrow2::offset::Offsets::::try_from_lengths( + many_optional.iter().map(|opt| { + opt.as_ref().map(|datum| datum.len()).unwrap_or_default() + }), + ) + .unwrap() + .into(); + ListArray::new( + { + _ = extension_wrapper; + DataType::List(Box::new(Field { name: "item".to_owned(), - data_type: DataType::Struct(vec![ - Field { - name: "single_float_optional".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "single_string_required".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "single_string_optional".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_floats_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), + data_type: DataType::Union( + vec![ + Field { + name: "degrees".to_owned(), data_type: DataType::Float32, - is_nullable: true, + is_nullable: false, metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_strings_required".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_strings_optional".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - })), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "flattened_scalar".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "almost_flattened_scalar".to_owned(), - data_type: DataType::Struct(vec![Field { - name: "value".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }]), - is_nullable: false, - metadata: [].into(), - }, - ]), - is_nullable: false, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - ], - None, - UnionMode::Dense, - ), - is_nullable: true, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - ], - None, - UnionMode::Dense, - ) - } - - #[allow(unused_imports, clippy::wildcard_imports)] - fn try_to_arrow_opt<'a>( - data: impl IntoIterator>>>, - extension_wrapper: Option<&str>, - ) -> crate::SerializationResult> - where - Self: Clone + 'a, - { - use crate::{Component as _, Datatype as _}; - use ::arrow2::{array::*, datatypes::*}; - Ok({ - let data: Vec<_> = data - .into_iter() - .map(|datum| { - let datum: Option<::std::borrow::Cow<'a, Self>> = datum.map(Into::into); - datum - }) - .collect(); - UnionArray::new( - (if let Some(ext) = extension_wrapper { - DataType::Extension( - ext.to_owned(), - Box::new(::to_arrow_datatype()), - None, - ) - } else { - ::to_arrow_datatype() - }) - .to_logical_type() - .clone(), - { - data.iter() - .flatten() - .map(|v| match **v { - AffixFuzzer4::SingleRequired(_) => 0i8, - AffixFuzzer4::ManyRequired(_) => 1i8, - AffixFuzzer4::ManyOptional(_) => 2i8, - }) - .collect() - }, - vec![ - { - let (somes, single_required): (Vec<_>, Vec<_>) = data - .iter() - .flatten() - .filter(|datum| matches!(***datum, AffixFuzzer4::SingleRequired(_))) - .map(|datum| { - let datum = match &**datum { - AffixFuzzer4::SingleRequired(v) => Some(v.clone()), - _ => None, - }; - (datum.is_some(), datum) - }) - .unzip(); - let single_required_bitmap: Option<::arrow2::bitmap::Bitmap> = { - let any_nones = somes.iter().any(|some| !*some); - any_nones.then(|| somes.into()) - }; - { - _ = single_required_bitmap; - _ = extension_wrapper; - crate::datatypes::AffixFuzzer3::try_to_arrow_opt( - single_required, - None::<&str>, - )? - } - }, - { - let (somes, many_required): (Vec<_>, Vec<_>) = data - .iter() - .flatten() - .filter(|datum| matches!(***datum, AffixFuzzer4::ManyRequired(_))) - .map(|datum| { - let datum = match &**datum { - AffixFuzzer4::ManyRequired(v) => Some(v.clone()), - _ => None, - }; - (datum.is_some(), datum) - }) - .unzip(); - let many_required_bitmap: Option<::arrow2::bitmap::Bitmap> = { - let any_nones = somes.iter().any(|some| !*some); - any_nones.then(|| somes.into()) - }; - { - use arrow2::{buffer::Buffer, offset::OffsetsBuffer}; - let many_required_inner_data: Vec<_> = many_required - .iter() - .flatten() - .flatten() - .map(ToOwned::to_owned) - .map(Some) - .collect(); - let many_required_inner_bitmap: Option<::arrow2::bitmap::Bitmap> = { - let any_nones = - many_required_inner_data.iter().any(|v| v.is_none()); - any_nones.then(|| { - many_required_inner_data - .iter() - .map(|v| v.is_some()) - .collect() - }) - }; - let offsets = ::arrow2::offset::Offsets::::try_from_lengths( - many_required.iter().map(|opt| { - opt.as_ref().map(|datum| datum.len()).unwrap_or_default() - }), - ) - .unwrap() - .into(); - ListArray::new( - { - _ = extension_wrapper; - DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Union( - vec![ - Field { - name: "degrees".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "radians".to_owned(), - data_type: DataType::Float32, + }, + Field { + name: "radians".to_owned(), + data_type: DataType::Float32, is_nullable: false, metadata: [].into(), }, @@ -2336,7 +1923,7 @@ impl crate::Datatype for AffixFuzzer4 { None, UnionMode::Dense, ), - is_nullable: false, + is_nullable: true, metadata: [].into(), })) .to_logical_type() @@ -2344,57 +1931,687 @@ impl crate::Datatype for AffixFuzzer4 { }, offsets, { - _ = many_required_inner_bitmap; + _ = many_optional_inner_bitmap; _ = extension_wrapper; crate::datatypes::AffixFuzzer3::try_to_arrow_opt( - many_required_inner_data, + many_optional_inner_data, None::<&str>, )? }, - many_required_bitmap, + many_optional_bitmap, ) .boxed() } }, + ], + Some({ + let mut single_required_offset = 0; + let mut many_required_offset = 0; + let mut many_optional_offset = 0; + data.iter() + .flatten() + .map(|v| match **v { + AffixFuzzer4::SingleRequired(_) => { + let offset = single_required_offset; + single_required_offset += 1; + offset + } + + AffixFuzzer4::ManyRequired(_) => { + let offset = many_required_offset; + many_required_offset += 1; + offset + } + + AffixFuzzer4::ManyOptional(_) => { + let offset = many_optional_offset; + many_optional_offset += 1; + offset + } + }) + .collect() + }), + ) + .boxed() + }) + } + + #[allow(unused_imports, clippy::wildcard_imports)] + fn try_from_arrow_opt( + data: &dyn ::arrow2::array::Array, + ) -> crate::DeserializationResult>> + where + Self: Sized, + { + use crate::{Component as _, Datatype as _}; + use ::arrow2::{array::*, datatypes::*}; + Ok({ + let data = data + .as_any() + .downcast_ref::<::arrow2::array::UnionArray>() + .ok_or_else(|| crate::DeserializationError::SchemaMismatch { + expected: data.data_type().clone(), + got: data.data_type().clone(), + })?; + let (data_types, data_arrays, data_offsets) = + (data.types(), data.fields(), data.offsets().unwrap()); + let single_required = { + let data = &*data_arrays[0usize]; + + crate::datatypes::AffixFuzzer3::try_from_arrow_opt(data)? + .into_iter() + .collect::>() + }; + let many_required = { + let data = &*data_arrays[1usize]; + + { + let datatype = data.data_type(); + let data = data + .as_any() + .downcast_ref::<::arrow2::array::ListArray>() + .unwrap(); + let bitmap = data.validity().cloned(); + let offsets = { + let offsets = data.offsets(); + offsets.iter().copied().zip(offsets.iter().copied().skip(1)) + }; + let data = &**data.values(); + let data = crate::datatypes::AffixFuzzer3::try_from_arrow_opt(data)? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + datatype: DataType::Union( + vec![ + Field { + name: "degrees".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "radians".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "craziness".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Struct(vec![ + Field { + name: "single_float_optional".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "single_string_required".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "single_string_optional".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_floats_optional".to_owned(), + data_type: DataType::List(Box::new( + Field { + name: "item".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + )), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_strings_required".to_owned(), + data_type: DataType::List(Box::new( + Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + )), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_strings_optional".to_owned(), + data_type: DataType::List(Box::new( + Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + )), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "flattened_scalar".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "almost_flattened_scalar".to_owned(), + data_type: DataType::Struct(vec![Field { + name: "value".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }]), + is_nullable: false, + metadata: [].into(), + }, + ]), + is_nullable: false, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + ], + None, + UnionMode::Dense, + ), + }) + }) + .collect::>>()?; + offsets + .enumerate() + .map(move |(i, (start, end))| { + bitmap + .as_ref() + .map_or(true, |bitmap| bitmap.get_bit(i)) + .then(|| { + Ok(data + .get(start as usize..end as usize) + .ok_or_else(|| { + crate::DeserializationError::OffsetsMismatch { + bounds: (start as usize, end as usize), + len: data.len(), + datatype: datatype.clone(), + } + })? + .to_vec()) + }) + .transpose() + }) + .collect::>>>()? + .into_iter() + } + .collect::>() + }; + let many_optional = { + let data = &*data_arrays[2usize]; + + { + let datatype = data.data_type(); + let data = data + .as_any() + .downcast_ref::<::arrow2::array::ListArray>() + .unwrap(); + let bitmap = data.validity().cloned(); + let offsets = { + let offsets = data.offsets(); + offsets.iter().copied().zip(offsets.iter().copied().skip(1)) + }; + let data = &**data.values(); + let data = crate::datatypes::AffixFuzzer3::try_from_arrow_opt(data)? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + datatype: DataType::Union( + vec![ + Field { + name: "degrees".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "radians".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "craziness".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Struct(vec![ + Field { + name: "single_float_optional".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "single_string_required".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "single_string_optional".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_floats_optional".to_owned(), + data_type: DataType::List(Box::new( + Field { + name: "item".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + )), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_strings_required".to_owned(), + data_type: DataType::List(Box::new( + Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + )), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_strings_optional".to_owned(), + data_type: DataType::List(Box::new( + Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + )), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "flattened_scalar".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "almost_flattened_scalar".to_owned(), + data_type: DataType::Struct(vec![Field { + name: "value".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }]), + is_nullable: false, + metadata: [].into(), + }, + ]), + is_nullable: false, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + ], + None, + UnionMode::Dense, + ), + }) + }) + .collect::>>()?; + offsets + .enumerate() + .map(move |(i, (start, end))| { + bitmap + .as_ref() + .map_or(true, |bitmap| bitmap.get_bit(i)) + .then(|| { + Ok(data + .get(start as usize..end as usize) + .ok_or_else(|| { + crate::DeserializationError::OffsetsMismatch { + bounds: (start as usize, end as usize), + len: data.len(), + datatype: datatype.clone(), + } + })? + .to_vec()) + }) + .transpose() + }) + .collect::>>>()? + .into_iter() + } + .collect::>() + }; + data_types + .iter() + .enumerate() + .map(|(i, typ)| { + let offset = data_offsets[i]; + + Ok(Some(match typ { + 0i8 => AffixFuzzer4::SingleRequired( + single_required + .get(offset as usize) + .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { + bounds: (offset as usize, offset as usize), + len: single_required.len(), + datatype: data.data_type().clone(), + })? + .clone() + .unwrap(), + ), + 1i8 => AffixFuzzer4::ManyRequired( + many_required + .get(offset as usize) + .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { + bounds: (offset as usize, offset as usize), + len: many_required.len(), + datatype: data.data_type().clone(), + })? + .clone() + .unwrap(), + ), + 2i8 => AffixFuzzer4::ManyOptional( + many_optional + .get(offset as usize) + .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { + bounds: (offset as usize, offset as usize), + len: many_optional.len(), + datatype: data.data_type().clone(), + })? + .clone(), + ), + _ => unreachable!(), + })) + }) + .collect::>>()? + }) + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum AffixFuzzer3 { + Degrees(f32), + Radians(Option), + Craziness(Vec), +} + +impl<'a> From for ::std::borrow::Cow<'a, AffixFuzzer3> { + #[inline] + fn from(value: AffixFuzzer3) -> Self { + std::borrow::Cow::Owned(value) + } +} + +impl<'a> From<&'a AffixFuzzer3> for ::std::borrow::Cow<'a, AffixFuzzer3> { + #[inline] + fn from(value: &'a AffixFuzzer3) -> Self { + std::borrow::Cow::Borrowed(value) + } +} + +impl crate::Datatype for AffixFuzzer3 { + #[inline] + fn name() -> crate::DatatypeName { + crate::DatatypeName::Borrowed("rerun.testing.datatypes.AffixFuzzer3") + } + + #[allow(unused_imports, clippy::wildcard_imports)] + #[inline] + fn to_arrow_datatype() -> arrow2::datatypes::DataType { + use ::arrow2::datatypes::*; + DataType::Union( + vec![ + Field { + name: "degrees".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "radians".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "craziness".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Struct(vec![ + Field { + name: "single_float_optional".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "single_string_required".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "single_string_optional".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_floats_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_strings_required".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_strings_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "flattened_scalar".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "almost_flattened_scalar".to_owned(), + data_type: DataType::Struct(vec![Field { + name: "value".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }]), + is_nullable: false, + metadata: [].into(), + }, + ]), + is_nullable: false, + metadata: [].into(), + })), + is_nullable: false, + metadata: [].into(), + }, + ], + None, + UnionMode::Dense, + ) + } + + #[allow(unused_imports, clippy::wildcard_imports)] + fn try_to_arrow_opt<'a>( + data: impl IntoIterator>>>, + extension_wrapper: Option<&str>, + ) -> crate::SerializationResult> + where + Self: Clone + 'a, + { + use crate::{Component as _, Datatype as _}; + use ::arrow2::{array::*, datatypes::*}; + Ok({ + let data: Vec<_> = data + .into_iter() + .map(|datum| { + let datum: Option<::std::borrow::Cow<'a, Self>> = datum.map(Into::into); + datum + }) + .collect(); + UnionArray::new( + (if let Some(ext) = extension_wrapper { + DataType::Extension( + ext.to_owned(), + Box::new(::to_arrow_datatype()), + None, + ) + } else { + ::to_arrow_datatype() + }) + .to_logical_type() + .clone(), + { + data.iter() + .flatten() + .map(|v| match **v { + AffixFuzzer3::Degrees(_) => 0i8, + AffixFuzzer3::Radians(_) => 1i8, + AffixFuzzer3::Craziness(_) => 2i8, + }) + .collect() + }, + vec![ + { + let (somes, degrees): (Vec<_>, Vec<_>) = data + .iter() + .flatten() + .filter(|datum| matches!(***datum, AffixFuzzer3::Degrees(_))) + .map(|datum| { + let datum = match &**datum { + AffixFuzzer3::Degrees(v) => Some(v.clone()), + _ => None, + }; + (datum.is_some(), datum) + }) + .unzip(); + let degrees_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + PrimitiveArray::new( + { + _ = extension_wrapper; + DataType::Float32.to_logical_type().clone() + }, + degrees.into_iter().map(|v| v.unwrap_or_default()).collect(), + degrees_bitmap, + ) + .boxed() + }, { - let (somes, many_optional): (Vec<_>, Vec<_>) = data + let (somes, radians): (Vec<_>, Vec<_>) = data .iter() .flatten() - .filter(|datum| matches!(***datum, AffixFuzzer4::ManyOptional(_))) + .filter(|datum| matches!(***datum, AffixFuzzer3::Radians(_))) .map(|datum| { let datum = match &**datum { - AffixFuzzer4::ManyOptional(v) => Some(v.clone()), + AffixFuzzer3::Radians(v) => Some(v.clone()), _ => None, } .flatten(); (datum.is_some(), datum) }) .unzip(); - let many_optional_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let radians_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + PrimitiveArray::new( + { + _ = extension_wrapper; + DataType::Float32.to_logical_type().clone() + }, + radians.into_iter().map(|v| v.unwrap_or_default()).collect(), + radians_bitmap, + ) + .boxed() + }, + { + let (somes, craziness): (Vec<_>, Vec<_>) = data + .iter() + .flatten() + .filter(|datum| matches!(***datum, AffixFuzzer3::Craziness(_))) + .map(|datum| { + let datum = match &**datum { + AffixFuzzer3::Craziness(v) => Some(v.clone()), + _ => None, + }; + (datum.is_some(), datum) + }) + .unzip(); + let craziness_bitmap: Option<::arrow2::bitmap::Bitmap> = { let any_nones = somes.iter().any(|some| !*some); any_nones.then(|| somes.into()) }; { use arrow2::{buffer::Buffer, offset::OffsetsBuffer}; - let many_optional_inner_data: Vec<_> = many_optional + let craziness_inner_data: Vec<_> = craziness .iter() .flatten() .flatten() .map(ToOwned::to_owned) .map(Some) .collect(); - let many_optional_inner_bitmap: Option<::arrow2::bitmap::Bitmap> = { - let any_nones = - many_optional_inner_data.iter().any(|v| v.is_none()); + let craziness_inner_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = craziness_inner_data.iter().any(|v| v.is_none()); any_nones.then(|| { - many_optional_inner_data - .iter() - .map(|v| v.is_some()) - .collect() + craziness_inner_data.iter().map(|v| v.is_some()).collect() }) }; let offsets = ::arrow2::offset::Offsets::::try_from_lengths( - many_optional.iter().map(|opt| { + craziness.iter().map(|opt| { opt.as_ref().map(|datum| datum.len()).unwrap_or_default() }), ) @@ -2405,122 +2622,77 @@ impl crate::Datatype for AffixFuzzer4 { _ = extension_wrapper; DataType::List(Box::new(Field { name: "item".to_owned(), - data_type: DataType::Union( - vec![ - Field { - name: "degrees".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "radians".to_owned(), + data_type: DataType::Struct(vec![ + Field { + name: "single_float_optional".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "single_string_required".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "single_string_optional".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_floats_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), data_type: DataType::Float32, - is_nullable: false, + is_nullable: true, metadata: [].into(), - }, - Field { - name: "craziness".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Struct(vec![ - Field { - name: "single_float_optional" - .to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "single_string_required" - .to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "single_string_optional" - .to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_floats_optional" - .to_owned(), - data_type: DataType::List( - Box::new(Field { - name: "item".to_owned(), - data_type: - DataType::Float32, - is_nullable: true, - metadata: [].into(), - }), - ), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_strings_required" - .to_owned(), - data_type: DataType::List( - Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }), - ), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_strings_optional" - .to_owned(), - data_type: DataType::List( - Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }), - ), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "flattened_scalar".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "almost_flattened_scalar" - .to_owned(), - data_type: DataType::Struct(vec![ - Field { - name: "value".to_owned(), - data_type: - DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - ]), - is_nullable: false, - metadata: [].into(), - }, - ]), - is_nullable: false, - metadata: [].into(), - })), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_strings_required".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, is_nullable: false, metadata: [].into(), - }, - ], - None, - UnionMode::Dense, - ), - is_nullable: true, + })), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_strings_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "flattened_scalar".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "almost_flattened_scalar".to_owned(), + data_type: DataType::Struct(vec![Field { + name: "value".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }]), + is_nullable: false, + metadata: [].into(), + }, + ]), + is_nullable: false, metadata: [].into(), })) .to_logical_type() @@ -2528,229 +2700,91 @@ impl crate::Datatype for AffixFuzzer4 { }, offsets, { - _ = many_optional_inner_bitmap; + _ = craziness_inner_bitmap; _ = extension_wrapper; - crate::datatypes::AffixFuzzer3::try_to_arrow_opt( - many_optional_inner_data, + crate::datatypes::AffixFuzzer1::try_to_arrow_opt( + craziness_inner_data, None::<&str>, )? }, - many_optional_bitmap, + craziness_bitmap, ) .boxed() } }, ], Some({ - let mut single_required_offset = 0; - let mut many_required_offset = 0; - let mut many_optional_offset = 0; + let mut degrees_offset = 0; + let mut radians_offset = 0; + let mut craziness_offset = 0; data.iter() .flatten() .map(|v| match **v { - AffixFuzzer4::SingleRequired(_) => { - let offset = single_required_offset; - single_required_offset += 1; - offset - } - - AffixFuzzer4::ManyRequired(_) => { - let offset = many_required_offset; - many_required_offset += 1; - offset - } - - AffixFuzzer4::ManyOptional(_) => { - let offset = many_optional_offset; - many_optional_offset += 1; - offset - } - }) - .collect() - }), - ) - .boxed() - }) - } - - #[allow(unused_imports, clippy::wildcard_imports)] - fn try_from_arrow_opt( - data: &dyn ::arrow2::array::Array, - ) -> crate::DeserializationResult>> - where - Self: Sized, - { - use crate::{Component as _, Datatype as _}; - use ::arrow2::{array::*, datatypes::*}; - Ok({ - let data = data - .as_any() - .downcast_ref::<::arrow2::array::UnionArray>() - .ok_or_else(|| crate::DeserializationError::SchemaMismatch { - expected: data.data_type().clone(), - got: data.data_type().clone(), - })?; - let (data_types, data_arrays, data_offsets) = - (data.types(), data.fields(), data.offsets().unwrap()); - let single_required = { - let data = &*data_arrays[0usize]; - - crate::datatypes::AffixFuzzer3::try_from_arrow_opt(data)? - .into_iter() - .collect::>() - }; - let many_required = { - let data = &*data_arrays[1usize]; - - { - let datatype = data.data_type(); - let data = data - .as_any() - .downcast_ref::<::arrow2::array::ListArray>() - .unwrap(); - let bitmap = data.validity().cloned(); - let offsets = { - let offsets = data.offsets(); - offsets.iter().copied().zip(offsets.iter().copied().skip(1)) - }; - let data = &**data.values(); - let data = crate::datatypes::AffixFuzzer3::try_from_arrow_opt(data)? - .into_iter() - .map(|v| { - v.ok_or_else(|| crate::DeserializationError::MissingData { - datatype: DataType::Union( - vec![ - Field { - name: "degrees".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "radians".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "craziness".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Struct(vec![ - Field { - name: "single_float_optional".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "single_string_required".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "single_string_optional".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_floats_optional".to_owned(), - data_type: DataType::List(Box::new( - Field { - name: "item".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - )), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_strings_required".to_owned(), - data_type: DataType::List(Box::new( - Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - )), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_strings_optional".to_owned(), - data_type: DataType::List(Box::new( - Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - )), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "flattened_scalar".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "almost_flattened_scalar".to_owned(), - data_type: DataType::Struct(vec![Field { - name: "value".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }]), - is_nullable: false, - metadata: [].into(), - }, - ]), - is_nullable: false, - metadata: [].into(), - })), - is_nullable: false, - metadata: [].into(), - }, - ], - None, - UnionMode::Dense, - ), - }) - }) - .collect::>>()?; - offsets - .enumerate() - .map(move |(i, (start, end))| { - bitmap - .as_ref() - .map_or(true, |bitmap| bitmap.get_bit(i)) - .then(|| { - Ok(data - .get(start as usize..end as usize) - .ok_or_else(|| { - crate::DeserializationError::OffsetsMismatch { - bounds: (start as usize, end as usize), - len: data.len(), - datatype: datatype.clone(), - } - })? - .to_vec()) - }) - .transpose() + AffixFuzzer3::Degrees(_) => { + let offset = degrees_offset; + degrees_offset += 1; + offset + } + + AffixFuzzer3::Radians(_) => { + let offset = radians_offset; + radians_offset += 1; + offset + } + + AffixFuzzer3::Craziness(_) => { + let offset = craziness_offset; + craziness_offset += 1; + offset + } }) - .collect::>>>()? - .into_iter() - } - .collect::>() + .collect() + }), + ) + .boxed() + }) + } + + #[allow(unused_imports, clippy::wildcard_imports)] + fn try_from_arrow_opt( + data: &dyn ::arrow2::array::Array, + ) -> crate::DeserializationResult>> + where + Self: Sized, + { + use crate::{Component as _, Datatype as _}; + use ::arrow2::{array::*, datatypes::*}; + Ok({ + let data = data + .as_any() + .downcast_ref::<::arrow2::array::UnionArray>() + .ok_or_else(|| crate::DeserializationError::SchemaMismatch { + expected: data.data_type().clone(), + got: data.data_type().clone(), + })?; + let (data_types, data_arrays, data_offsets) = + (data.types(), data.fields(), data.offsets().unwrap()); + let degrees = { + let data = &*data_arrays[0usize]; + + data.as_any() + .downcast_ref::() + .unwrap() + .into_iter() + .map(|v| v.copied()) + .collect::>() }; - let many_optional = { + let radians = { + let data = &*data_arrays[1usize]; + + data.as_any() + .downcast_ref::() + .unwrap() + .into_iter() + .map(|v| v.copied()) + .collect::>() + }; + let craziness = { let data = &*data_arrays[2usize]; { @@ -2765,114 +2799,80 @@ impl crate::Datatype for AffixFuzzer4 { offsets.iter().copied().zip(offsets.iter().copied().skip(1)) }; let data = &**data.values(); - let data = crate::datatypes::AffixFuzzer3::try_from_arrow_opt(data)? + let data = crate::datatypes::AffixFuzzer1::try_from_arrow_opt(data)? .into_iter() .map(|v| { v.ok_or_else(|| crate::DeserializationError::MissingData { - datatype: DataType::Union( - vec![ - Field { - name: "degrees".to_owned(), + datatype: DataType::Struct(vec![ + Field { + name: "single_float_optional".to_owned(), + data_type: DataType::Float32, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "single_string_required".to_owned(), + data_type: DataType::Utf8, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "single_string_optional".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_floats_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), data_type: DataType::Float32, - is_nullable: false, + is_nullable: true, metadata: [].into(), - }, - Field { - name: "radians".to_owned(), - data_type: DataType::Float32, + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "many_strings_required".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, is_nullable: false, metadata: [].into(), - }, - Field { - name: "craziness".to_owned(), - data_type: DataType::List(Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Struct(vec![ - Field { - name: "single_float_optional".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "single_string_required".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "single_string_optional".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_floats_optional".to_owned(), - data_type: DataType::List(Box::new( - Field { - name: "item".to_owned(), - data_type: DataType::Float32, - is_nullable: true, - metadata: [].into(), - }, - )), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "many_strings_required".to_owned(), - data_type: DataType::List(Box::new( - Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: false, - metadata: [].into(), - }, - )), - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "many_strings_optional".to_owned(), - data_type: DataType::List(Box::new( - Field { - name: "item".to_owned(), - data_type: DataType::Utf8, - is_nullable: true, - metadata: [].into(), - }, - )), - is_nullable: true, - metadata: [].into(), - }, - Field { - name: "flattened_scalar".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }, - Field { - name: "almost_flattened_scalar".to_owned(), - data_type: DataType::Struct(vec![Field { - name: "value".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }]), - is_nullable: false, - metadata: [].into(), - }, - ]), - is_nullable: false, - metadata: [].into(), - })), + })), + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "many_strings_optional".to_owned(), + data_type: DataType::List(Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Utf8, + is_nullable: true, + metadata: [].into(), + })), + is_nullable: true, + metadata: [].into(), + }, + Field { + name: "flattened_scalar".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "almost_flattened_scalar".to_owned(), + data_type: DataType::Struct(vec![Field { + name: "value".to_owned(), + data_type: DataType::Float32, is_nullable: false, metadata: [].into(), - }, - ], - None, - UnionMode::Dense, - ), + }]), + is_nullable: false, + metadata: [].into(), + }, + ]), }) }) .collect::>>()?; @@ -2908,37 +2908,37 @@ impl crate::Datatype for AffixFuzzer4 { let offset = data_offsets[i]; Ok(Some(match typ { - 0i8 => AffixFuzzer4::SingleRequired( - single_required + 0i8 => AffixFuzzer3::Degrees( + degrees .get(offset as usize) .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { bounds: (offset as usize, offset as usize), - len: single_required.len(), + len: degrees.len(), datatype: data.data_type().clone(), })? .clone() .unwrap(), ), - 1i8 => AffixFuzzer4::ManyRequired( - many_required + 1i8 => AffixFuzzer3::Radians( + radians .get(offset as usize) .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { bounds: (offset as usize, offset as usize), - len: many_required.len(), + len: radians.len(), datatype: data.data_type().clone(), })? - .clone() - .unwrap(), + .clone(), ), - 2i8 => AffixFuzzer4::ManyOptional( - many_optional + 2i8 => AffixFuzzer3::Craziness( + craziness .get(offset as usize) .ok_or_else(|| crate::DeserializationError::OffsetsMismatch { bounds: (offset as usize, offset as usize), - len: many_optional.len(), + len: craziness.len(), datatype: data.data_type().clone(), })? - .clone(), + .clone() + .unwrap(), ), _ => unreachable!(), })) diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/fuzzy.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/fuzzy.py index 9428ae131a66..2ef1cdb01217 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/fuzzy.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/fuzzy.py @@ -187,6 +187,96 @@ def _native_to_pa_array(data: AffixFuzzer2ArrayLike, data_type: pa.DataType) -> # pa.register_extension_type(AffixFuzzer2Type()) +@define +class AffixFuzzer3: + degrees: float | None = None + radians: float | None = None + craziness: list[datatypes.AffixFuzzer1] | None = None + + +AffixFuzzer3Like = AffixFuzzer3 +AffixFuzzer3ArrayLike = Union[ + AffixFuzzer3, + Sequence[AffixFuzzer3Like], +] + + +# --- Arrow support --- + + +class AffixFuzzer3Type(BaseExtensionType): + def __init__(self) -> None: + pa.ExtensionType.__init__( + self, + pa.dense_union( + [ + pa.field("degrees", pa.float32(), False, {}), + pa.field("radians", pa.float32(), False, {}), + pa.field( + "craziness", + pa.list_( + pa.field( + "item", + pa.struct( + [ + pa.field("single_float_optional", pa.float32(), True, {}), + pa.field("single_string_required", pa.utf8(), False, {}), + pa.field("single_string_optional", pa.utf8(), True, {}), + pa.field( + "many_floats_optional", + pa.list_(pa.field("item", pa.float32(), True, {})), + True, + {}, + ), + pa.field( + "many_strings_required", + pa.list_(pa.field("item", pa.utf8(), False, {})), + False, + {}, + ), + pa.field( + "many_strings_optional", + pa.list_(pa.field("item", pa.utf8(), True, {})), + True, + {}, + ), + pa.field("flattened_scalar", pa.float32(), False, {}), + pa.field( + "almost_flattened_scalar", + pa.struct([pa.field("value", pa.float32(), False, {})]), + False, + {}, + ), + ] + ), + False, + {}, + ) + ), + False, + {}, + ), + ] + ), + "rerun.testing.datatypes.AffixFuzzer3", + ) + + +class AffixFuzzer3Array(BaseExtensionArray[AffixFuzzer3ArrayLike]): + _EXTENSION_NAME = "rerun.testing.datatypes.AffixFuzzer3" + _EXTENSION_TYPE = AffixFuzzer3Type + + @staticmethod + def _native_to_pa_array(data: AffixFuzzer3ArrayLike, data_type: pa.DataType) -> pa.Array: + raise NotImplementedError + + +AffixFuzzer3Type._ARRAY_TYPE = AffixFuzzer3Array + +# TODO(cmc): bring back registration to pyarrow once legacy types are gone +# pa.register_extension_type(AffixFuzzer3Type()) + + @define class AffixFuzzer4: single_required: datatypes.AffixFuzzer3 | None = None @@ -408,93 +498,3 @@ def _native_to_pa_array(data: AffixFuzzer4ArrayLike, data_type: pa.DataType) -> # TODO(cmc): bring back registration to pyarrow once legacy types are gone # pa.register_extension_type(AffixFuzzer4Type()) - - -@define -class AffixFuzzer3: - degrees: float | None = None - radians: float | None = None - craziness: list[datatypes.AffixFuzzer1] | None = None - - -AffixFuzzer3Like = AffixFuzzer3 -AffixFuzzer3ArrayLike = Union[ - AffixFuzzer3, - Sequence[AffixFuzzer3Like], -] - - -# --- Arrow support --- - - -class AffixFuzzer3Type(BaseExtensionType): - def __init__(self) -> None: - pa.ExtensionType.__init__( - self, - pa.dense_union( - [ - pa.field("degrees", pa.float32(), False, {}), - pa.field("radians", pa.float32(), False, {}), - pa.field( - "craziness", - pa.list_( - pa.field( - "item", - pa.struct( - [ - pa.field("single_float_optional", pa.float32(), True, {}), - pa.field("single_string_required", pa.utf8(), False, {}), - pa.field("single_string_optional", pa.utf8(), True, {}), - pa.field( - "many_floats_optional", - pa.list_(pa.field("item", pa.float32(), True, {})), - True, - {}, - ), - pa.field( - "many_strings_required", - pa.list_(pa.field("item", pa.utf8(), False, {})), - False, - {}, - ), - pa.field( - "many_strings_optional", - pa.list_(pa.field("item", pa.utf8(), True, {})), - True, - {}, - ), - pa.field("flattened_scalar", pa.float32(), False, {}), - pa.field( - "almost_flattened_scalar", - pa.struct([pa.field("value", pa.float32(), False, {})]), - False, - {}, - ), - ] - ), - False, - {}, - ) - ), - False, - {}, - ), - ] - ), - "rerun.testing.datatypes.AffixFuzzer3", - ) - - -class AffixFuzzer3Array(BaseExtensionArray[AffixFuzzer3ArrayLike]): - _EXTENSION_NAME = "rerun.testing.datatypes.AffixFuzzer3" - _EXTENSION_TYPE = AffixFuzzer3Type - - @staticmethod - def _native_to_pa_array(data: AffixFuzzer3ArrayLike, data_type: pa.DataType) -> pa.Array: - raise NotImplementedError - - -AffixFuzzer3Type._ARRAY_TYPE = AffixFuzzer3Array - -# TODO(cmc): bring back registration to pyarrow once legacy types are gone -# pa.register_extension_type(AffixFuzzer3Type())