Pull out message cache from layout reader (#1773)

* More closely aligns to the poll_read pattern where the thing being
read is passed in.
* Working towards moving to MessageIds instead of layout paths, so this
is a temporary intermediate state.

Author: gatesn

vortex-datafusion/src/persistent/format.rs

@@ -1,5 +1,5 @@
 use std::any::Any;
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 
 use arrow_schema::{Schema, SchemaRef};
 use async_trait::async_trait;
@@ -24,8 +24,7 @@ use vortex_array::Context;
 use vortex_error::VortexResult;
 use vortex_file::metadata::fetch_metadata;
 use vortex_file::{
-    LayoutContext, LayoutDeserializer, LayoutMessageCache, RelativeLayoutCache, Scan,
-    VORTEX_FILE_EXTENSION,
+    LayoutContext, LayoutDeserializer, LayoutMessageCache, LayoutPath, Scan, VORTEX_FILE_EXTENSION,
 };
 use vortex_io::{IoDispatcher, ObjectStoreReadAt};
 
@@ -145,23 +144,23 @@ impl FileFormat for VortexFormat {
 
         let layout_deserializer =
             LayoutDeserializer::new(self.context.clone(), LayoutContext::default().into());
-        let layout_message_cache = Arc::new(RwLock::new(LayoutMessageCache::new()));
-        let relative_message_cache = RelativeLayoutCache::new(layout_message_cache.clone());
 
         let root_layout = layout_deserializer.read_layout(
+            LayoutPath::default(),
             initial_read.fb_layout(),
             Scan::empty(),
             initial_read.lazy_dtype().into(),
-            relative_message_cache,
         )?;
 
         let os_read_at = ObjectStoreReadAt::new(store.clone(), object.location.clone());
         let io = IoDispatcher::default();
         let mut stats = Statistics::new_unknown(&table_schema);
         stats.num_rows = Precision::Exact(row_count as usize);
 
+        let msgs = LayoutMessageCache::default();
+
         if let Some(metadata_table) =
-            fetch_metadata(os_read_at, io.into(), root_layout, layout_message_cache).await?
+            fetch_metadata(os_read_at, io.into(), root_layout, msgs).await?
         {
             let mut column_statistics = Vec::with_capacity(table_schema.fields().len());
             let mut total_size = 0_u64;

vortex-file/src/read/buffered.rs

@@ -2,17 +2,19 @@ use std::collections::VecDeque;
 use std::io;
 use std::io::ErrorKind;
 use std::pin::Pin;
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 use std::task::{Context, Poll, Waker};
 
 use futures::Stream;
 use futures_util::future::BoxFuture;
 use futures_util::{FutureExt, StreamExt};
 use vortex_array::ArrayData;
-use vortex_error::{vortex_err, vortex_panic, VortexExpect, VortexResult};
+use vortex_error::{vortex_err, VortexExpect, VortexResult};
 use vortex_io::{Dispatch, IoDispatcher, VortexReadAt, VortexReadRanges};
 
-use crate::{LayoutMessageCache, LayoutReader, Message, MessageLocator, PollRead, RowMask};
+use crate::{
+    LayoutMessageCache, LayoutReader, Message, MessageCache, MessageLocator, PollRead, RowMask,
+};
 
 const NUM_TO_COALESCE: usize = 8;
 
@@ -21,7 +23,11 @@ pub(crate) trait ReadMasked {
 
     /// Read a Layout into a `V`, applying the given bitmask. Only entries corresponding to positions
     /// where mask is `true` will be included in the output.
-    fn read_masked(&self, mask: &RowMask) -> VortexResult<Option<PollRead<Self::Value>>>;
+    fn read_masked(
+        &self,
+        mask: &RowMask,
+        msgs: &dyn MessageCache,
+    ) -> VortexResult<Option<PollRead<Self::Value>>>;
 }
 
 /// Read an array with a [`RowMask`].
@@ -39,8 +45,12 @@ impl ReadMasked for ReadArray {
     type Value = ArrayData;
 
     /// Read given mask out of the reader
-    fn read_masked(&self, mask: &RowMask) -> VortexResult<Option<PollRead<ArrayData>>> {
-        self.layout.poll_read(mask)
+    fn read_masked(
+        &self,
+        mask: &RowMask,
+        msgs: &dyn MessageCache,
+    ) -> VortexResult<Option<PollRead<ArrayData>>> {
+        self.layout.poll_read(mask, msgs)
     }
 }
 
@@ -58,7 +68,7 @@ pub struct BufferedLayoutReader<R, S, V, RM> {
     queued: VecDeque<RowMaskState<V>>,
     io_read: VortexReadRanges<R>,
     dispatcher: Arc<IoDispatcher>,
-    cache: Arc<RwLock<LayoutMessageCache>>,
+    msgs: LayoutMessageCache,
 }
 
 impl<R, S, V, RM> BufferedLayoutReader<R, S, V, RM>
@@ -72,7 +82,7 @@ where
         dispatcher: Arc<IoDispatcher>,
         read_masks: S,
         row_mask_reader: RM,
-        cache: Arc<RwLock<LayoutMessageCache>>,
+        msgs: LayoutMessageCache,
     ) -> Self {
         Self {
             read_masks,
@@ -81,18 +91,13 @@ where
             queued: VecDeque::new(),
             io_read: VortexReadRanges::new(read, dispatcher.clone(), 1 << 20),
             dispatcher,
-            cache,
+            msgs,
         }
     }
 
-    fn store_messages(&self, messages: Vec<Message>) {
-        let mut write_cache_guard = self
-            .cache
-            .write()
-            .unwrap_or_else(|poison| vortex_panic!("Failed to write to message cache: {poison}"));
-        for Message(message_id, buf) in messages {
-            write_cache_guard.set(message_id, buf);
-        }
+    fn store_messages(&mut self, messages: Vec<Message>) {
+        self.msgs
+            .set_many(messages.into_iter().map(|msg| (msg.0, msg.1)))
     }
 
     fn gather_read_messages(
@@ -106,7 +111,9 @@ where
         for queued_res in self.queued.iter_mut() {
             match queued_res {
                 RowMaskState::Pending(pending_mask) => {
-                    if let Some(pending_read) = self.row_mask_reader.read_masked(pending_mask)? {
+                    if let Some(pending_read) =
+                        self.row_mask_reader.read_masked(pending_mask, &self.msgs)?
+                    {
                         match pending_read {
                             PollRead::ReadMore(m) => {
                                 to_read.extend(m);
@@ -129,7 +136,9 @@ where
         while read_more_count < NUM_TO_COALESCE {
             match self.read_masks.poll_next_unpin(cx) {
                 Poll::Ready(Some(Ok(next_mask))) => {
-                    if let Some(read_result) = self.row_mask_reader.read_masked(&next_mask)? {
+                    if let Some(read_result) =
+                        self.row_mask_reader.read_masked(&next_mask, &self.msgs)?
+                    {
                         match read_result {
                             PollRead::ReadMore(m) => {
                                 self.queued.push_back(RowMaskState::Pending(next_mask));

vortex-file/src/read/builder/mod.rs

@@ -1,4 +1,4 @@
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 
 use initial_read::read_initial_bytes;
 use vortex_array::{ArrayDType, ArrayData};
@@ -8,11 +8,12 @@ use vortex_io::{IoDispatcher, VortexReadAt};
 
 use super::handle::VortexReadHandle;
 use super::InitialRead;
-use crate::read::cache::{LayoutMessageCache, RelativeLayoutCache};
+use crate::read::cache::LayoutMessageCache;
 use crate::read::context::LayoutDeserializer;
 use crate::read::filtering::RowFilter;
 use crate::read::projection::Projection;
 use crate::read::{RowMask, Scan};
+use crate::LayoutPath;
 
 pub(crate) mod initial_read;
 
@@ -147,25 +148,25 @@ impl<R: VortexReadAt + Unpin> VortexReadBuilder<R> {
             Projection::Flat(ref fields) => lazy_dtype.project(fields)?,
         };
 
-        let message_cache = Arc::new(RwLock::new(LayoutMessageCache::default()));
+        let message_cache = LayoutMessageCache::default();
         let layout_reader = self.layout_serde.read_layout(
+            LayoutPath::default(),
             initial_read.fb_layout(),
             match self.projection {
                 Projection::All => Scan::empty(),
                 Projection::Flat(p) => Scan::new(Arc::new(Select::include(p))),
             },
             lazy_dtype.clone(),
-            RelativeLayoutCache::new(message_cache.clone()),
         )?;
 
         let filter_reader = self
             .row_filter
             .map(|row_filter| {
                 self.layout_serde.read_layout(
+                    LayoutPath::default(),
                     initial_read.fb_layout(),
                     Scan::new(Arc::new(row_filter)),
                     lazy_dtype,
-                    RelativeLayoutCache::new(message_cache.clone()),
                 )
             })
             .transpose()?;

vortex-file/src/read/cache.rs

@@ -9,34 +9,59 @@ use vortex_buffer::ByteBuffer;
 use vortex_dtype::field::Field;
 use vortex_dtype::flatbuffers::{extract_field, project_and_deserialize, resolve_field};
 use vortex_dtype::{DType, FieldNames};
-use vortex_error::{vortex_bail, vortex_err, vortex_panic, VortexResult};
+use vortex_error::{vortex_bail, vortex_err, VortexExpect, VortexResult};
 use vortex_flatbuffers::dtype as fbd;
 
 use crate::read::projection::Projection;
 use crate::read::{LayoutPartId, MessageId};
 
-#[derive(Default, Debug)]
+/// A read-only cache of messages.
+pub trait MessageCache {
+    fn get(&self, path: &[LayoutPartId]) -> Option<Bytes>;
+}
+
+#[derive(Default, Debug, Clone)]
 pub struct LayoutMessageCache {
-    cache: HashMap<MessageId, Bytes>,
+    cache: Arc<RwLock<HashMap<MessageId, Bytes>>>,
 }
 
 impl LayoutMessageCache {
-    pub fn new() -> Self {
-        Self {
-            cache: HashMap::new(),
-        }
+    pub fn remove(&mut self, path: &[LayoutPartId]) -> Option<Bytes> {
+        self.cache
+            .write()
+            .map_err(|_| vortex_err!("Poisoned cache"))
+            .vortex_expect("poisoned")
+            .remove(path)
     }
 
-    pub fn get(&self, path: &[LayoutPartId]) -> Option<Bytes> {
-        self.cache.get(path).cloned()
+    pub fn set(&mut self, path: MessageId, value: Bytes) {
+        self.cache
+            .write()
+            .map_err(|_| vortex_err!("Poisoned cache"))
+            .vortex_expect("poisoned")
+            .insert(path, value);
     }
 
-    pub fn remove(&mut self, path: &[LayoutPartId]) -> Option<Bytes> {
-        self.cache.remove(path)
+    pub fn set_many<I: IntoIterator<Item = (MessageId, Bytes)>>(&mut self, iter: I) {
+        let mut guard = self
+            .cache
+            .write()
+            .map_err(|_| vortex_err!("Poisoned cache"))
+            .vortex_expect("poisoned");
+        for (id, bytes) in iter.into_iter() {
+            guard.insert(id, bytes);
+        }
     }
+}
 
-    pub fn set(&mut self, path: MessageId, value: Bytes) {
-        self.cache.insert(path, value);
+impl MessageCache for LayoutMessageCache {
+    fn get(&self, path: &[LayoutPartId]) -> Option<Bytes> {
+        self.cache
+            .read()
+            .map_err(|_| vortex_err!("Poisoned cache"))
+            .vortex_expect("poisoned")
+            .get(path)
+            .cloned()
     }
 }
 
@@ -258,61 +283,3 @@ fn fb_struct(bytes: &[u8]) -> VortexResult<fbd::Struct_> {
 fn fb_dtype(bytes: &[u8]) -> fbd::DType {
     unsafe { root_unchecked::<fbd::DType>(bytes) }
 }
-
-#[derive(Debug, Clone)]
-pub struct RelativeLayoutCache {
-    root: Arc<RwLock<LayoutMessageCache>>,
-    path: MessageId,
-}
-
-impl RelativeLayoutCache {
-    pub fn new(root: Arc<RwLock<LayoutMessageCache>>) -> Self {
-        Self {
-            root,
-            path: Vec::new(),
-        }
-    }
-
-    pub fn relative(&self, id: LayoutPartId) -> Self {
-        let mut new_path = Vec::with_capacity(self.path.len() + 1);
-        new_path.clone_from(&self.path);
-        new_path.push(id);
-        Self {
-            root: self.root.clone(),
-            path: new_path,
-        }
-    }
-
-    pub fn get(&self, path: &[LayoutPartId]) -> Option<Bytes> {
-        self.root
-            .read()
-            .unwrap_or_else(|poison| {
-                vortex_panic!(
-                    "Failed to read from layout cache at path {:?} with error {}",
-                    path,
-                    poison
-                );
-            })
-            .get(&self.absolute_id(path))
-    }
-
-    pub fn remove(&mut self, path: &[LayoutPartId]) -> Option<Bytes> {
-        self.root
-            .write()
-            .unwrap_or_else(|poison| {
-                vortex_panic!(
-                    "Failed to write to layout cache at path {:?} with error {}",
-                    path,
-                    poison
-                )
-            })
-            .remove(&self.absolute_id(path))
-    }
-
-    pub fn absolute_id(&self, path: &[LayoutPartId]) -> MessageId {
-        let mut lookup_key = Vec::with_capacity(self.path.len() + path.len());
-        lookup_key.clone_from(&self.path);
-        lookup_key.extend_from_slice(path);
-        lookup_key
-    }
-}

vortex-file/src/read/context.rs

@@ -7,7 +7,7 @@ use vortex_error::{vortex_err, VortexResult};
 use vortex_flatbuffers::footer as fb;
 
 use crate::layouts::{ChunkedLayout, ColumnarLayout, FlatLayout};
-use crate::{LayoutReader, LazyDType, RelativeLayoutCache, Scan};
+use crate::{LayoutPath, LayoutReader, LazyDType, Scan};
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
 pub struct LayoutId(pub u16);
@@ -23,11 +23,11 @@ pub trait Layout: Debug + Send + Sync {
 
     fn reader(
         &self,
+        path: LayoutPath,
         layout: fb::Layout,
         dtype: Arc<LazyDType>,
         scan: Scan,
         layout_serde: LayoutDeserializer,
-        message_cache: RelativeLayoutCache,
     ) -> VortexResult<Arc<dyn LayoutReader>>;
 }
 
@@ -72,16 +72,16 @@ impl LayoutDeserializer {
 
     pub fn read_layout(
         &self,
+        path: LayoutPath,
         layout: fb::Layout,
         scan: Scan,
         dtype: Arc<LazyDType>,
-        message_cache: RelativeLayoutCache,
     ) -> VortexResult<Arc<dyn LayoutReader>> {
         let layout_id = LayoutId(layout.encoding());
         self.layout_ctx
             .lookup_layout(&layout_id)
             .ok_or_else(|| vortex_err!("Unknown layout definition {layout_id}"))?
-            .reader(layout, dtype, scan, self.clone(), message_cache)
+            .reader(path, layout, dtype, scan, self.clone())
     }
 
     pub(crate) fn ctx(&self) -> Arc<Context> {