diff --git a/vortex-array/src/dtype/struct_.rs b/vortex-array/src/dtype/struct_.rs index 95f73295d52..4b457d84d16 100644 --- a/vortex-array/src/dtype/struct_.rs +++ b/vortex-array/src/dtype/struct_.rs @@ -30,9 +30,8 @@ pub struct FieldDType { impl From for FieldDType { fn from(value: ViewedDType) -> Self { - Self { - inner: FieldDTypeInner::View(value), - } + let view = FieldDTypeInner::View(value, Arc::new(OnceLock::new())); + Self { inner: view } } } @@ -57,26 +56,38 @@ enum FieldDTypeInner { /// Owned DType instance // TODO(ngates): we should consider making this an Arc. Owned(DType), - /// A view over a flatbuffer, parsed only when accessed. - View(ViewedDType), + /// A view over a flatbuffer, parsed on first access and cached after. + /// We cache it because it's requested for every field for every file. + /// If there are many files with wide schemas (think clickbench), + /// re-parsing each field is costly + /// This form is quite ugly because in some cases we want to panic on + /// parsing but return an error on others + View(ViewedDType, Arc>), } impl PartialEq for FieldDTypeInner { fn eq(&self, other: &Self) -> bool { match (self, other) { (Self::Owned(lhs), Self::Owned(rhs)) => lhs == rhs, - (Self::View(lhs), Self::View(rhs)) => { - let lhs = DType::try_from(lhs.clone()) - .vortex_expect("Failed to parse FieldDType into DType"); - let rhs = DType::try_from(rhs.clone()) - .vortex_expect("Failed to parse FieldDType into DType"); + (Self::View(left_view, left_lock), Self::View(right_view, right_lock)) => { + let lhs = left_lock.get_or_init(|| { + DType::try_from(left_view.clone()) + .vortex_expect("Failed to parse FieldDType into DType") + }); + let rhs = right_lock.get_or_init(|| { + DType::try_from(right_view.clone()) + .vortex_expect("Failed to parse FieldDType into DType") + }); lhs == rhs } - (Self::View(view), Self::Owned(owned)) | (Self::Owned(owned), Self::View(view)) => { - let view = DType::try_from(view.clone()) - .vortex_expect("Failed to parse FieldDType into DType"); - owned == &view + (Self::View(view, lock), Self::Owned(owned)) + | (Self::Owned(owned), Self::View(view, lock)) => { + let view = lock.get_or_init(|| { + DType::try_from(view.clone()) + .vortex_expect("Failed to parse FieldDType into DType") + }); + owned == view } } } @@ -89,9 +100,11 @@ impl Hash for FieldDTypeInner { FieldDTypeInner::Owned(owned) => { owned.hash(state); } - FieldDTypeInner::View(view) => { - let owned = DType::try_from(view.clone()) - .vortex_expect("Failed to parse FieldDType into DType"); + FieldDTypeInner::View(view, lock) => { + let owned = lock.get_or_init(|| { + DType::try_from(view.clone()) + .vortex_expect("Failed to parse FieldDType into DType") + }); owned.hash(state); } } @@ -110,7 +123,13 @@ impl FieldDTypeInner { fn value(&self) -> VortexResult { match &self { FieldDTypeInner::Owned(owned) => Ok(owned.clone()), - FieldDTypeInner::View(view) => DType::try_from(view.clone()), + FieldDTypeInner::View(view, lock) => { + if let Some(dtype) = lock.get() { + return Ok(dtype.clone()); + } + let parsed = DType::try_from(view.clone())?; + Ok(lock.get_or_init(|| parsed).clone()) + } } } } @@ -491,12 +510,17 @@ mod test { use insta::assert_snapshot; use itertools::Itertools; + use vortex_error::VortexResult; + use vortex_flatbuffers::FlatBuffer; + use vortex_flatbuffers::WriteFlatBufferExt; + use super::FieldDTypeInner; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::dtype::test::SESSION; #[test] fn nullability() { @@ -575,6 +599,36 @@ mod test { assert_eq!(without_a.nfields(), 1); } + #[test] + fn field_dtype_cache() -> VortexResult<()> { + let inner = DType::Struct( + StructFields::from_iter([("x", DType::Bool(Nullability::NonNullable))]), + Nullability::NonNullable, + ); + let dtype = DType::Struct( + StructFields::from_iter([("a", inner)]), + Nullability::NonNullable, + ); + let buffer = FlatBuffer::from(dtype.write_flatbuffer_bytes()?); + let parsed = DType::from_flatbuffer(buffer, &SESSION)?; + let fields = parsed.as_struct_fields_opt().unwrap(); + + let field = &fields.0.dtypes[0]; + let FieldDTypeInner::View(_, cell) = &field.inner else { + panic!("flatbuffer-backed field must be a View"); + }; + assert!(cell.get().is_none()); + + let first = field.value()?; + assert!(cell.get().is_some()); + let second = field.value()?; + let (DType::Struct(first, _), DType::Struct(second, _)) = (&first, &second) else { + panic!("field must parse to a struct"); + }; + assert!(Arc::ptr_eq(&first.0, &second.0)); + Ok(()) + } + #[test] fn test_merge() { let child_a = DType::Primitive(PType::I32, Nullability::NonNullable);