cache DTypes parsed from flatbuffers

vortex-array/src/dtype/struct_.rs

@@ -30,9 +30,8 @@ pub struct FieldDType {
 
 impl From<ViewedDType> 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<DType>.
     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<OnceLock<DType>>),
 }
 
 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<DType> {
         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);