feat(cache): add tier probes and selective entry reads

rust/lance-core/src/cache/backend.rs

@@ -45,6 +45,20 @@ use super::{CacheCodec, InternalCacheKey};
 /// A type-erased cache entry.
 pub type CacheEntry = Arc<dyn Any + Send + Sync>;
 
+/// The tier that would serve a read of an entry, as
+/// [`CacheBackend::peek_tier`] predicts it.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum CacheTier {
+    /// In RAM: a read is a memory hit.
+    Resident,
+    /// Not in RAM but held by a local tier below it, such as a persistent
+    /// store or a buffer of recently evicted entries: a read costs local I/O
+    /// at most, never the caller's loader.
+    Local,
+    /// Held by no tier: a read runs the caller's loader.
+    Absent,
+}
+
 /// Low-level pluggable cache backend.
 ///
 /// Implementations store entries keyed by [`InternalCacheKey`] and return
@@ -53,6 +67,62 @@ pub type CacheEntry = Arc<dyn Any + Send + Sync>;
 /// backend authors only need to implement storage and eviction.
 #[async_trait]
 pub trait CacheBackend: Send + Sync + std::fmt::Debug {
+    /// Look up RAM only, without reading or promoting a persistent entry.
+    /// Backends without a resident lookup return a miss safely.
+    async fn get_resident(&self, _key: &InternalCacheKey) -> Option<CacheEntry> {
+        None
+    }
+
+    /// Report whether `key` is resident in RAM without counting as an access,
+    /// so a residency probe does not change what the backend evicts next.
+    /// The default uses [`get_resident`](Self::get_resident), which backends
+    /// with recency state may treat as an access.
+    async fn peek_resident(&self, key: &InternalCacheKey) -> bool {
+        self.get_resident(key).await.is_some()
+    }
+
+    /// Predict which tier would serve a read of `key`, without reading the
+    /// entry or counting as an access, so the check changes neither what the
+    /// backend evicts next nor its hit statistics. The default knows RAM only,
+    /// through [`peek_resident`](Self::peek_resident), and reports every other
+    /// entry as [`CacheTier::Absent`]; backends with a local tier below RAM
+    /// override it.
+    async fn peek_tier(&self, key: &InternalCacheKey) -> CacheTier {
+        if self.peek_resident(key).await {
+            CacheTier::Resident
+        } else {
+            CacheTier::Absent
+        }
+    }
+
+    /// Whether layered index planes should gate lower-plane RAM admission on
+    /// their sign plane being resident. Backends that admit plane entries
+    /// through their ordinary policy return `false`, so lower planes are
+    /// loaded and admitted like any other entry.
+    fn plane_admission_gated(&self) -> bool {
+        true
+    }
+
+    /// Read an entry without admitting a persistent hit into RAM.
+    /// Backends without this capability safely fall back to resident entries.
+    async fn get_without_promotion(
+        &self,
+        key: &InternalCacheKey,
+        _codec: Option<CacheCodec>,
+    ) -> Option<CacheEntry> {
+        self.get_resident(key).await
+    }
+
+    /// Gather selected rows from persistent storage, without RAM admission.
+    async fn get_rows(
+        &self,
+        _key: &InternalCacheKey,
+        _rows: &[u32],
+        _codec: Option<CacheCodec>,
+    ) -> Option<CacheEntry> {
+        None
+    }
+
     /// Look up an entry by its key.
     ///
     /// `codec` is provided so that persistent backends can deserialize the

rust/lance-core/src/cache/codec.rs

@@ -188,6 +188,17 @@ impl<T> CacheDecode<T> {
 // CacheCodecImpl — trait for serializable cache entry types
 // ---------------------------------------------------------------------------
 
+/// Bounded random reads from a persistent cache payload. Backends may coalesce
+/// and align physical I/O while codecs request logical byte ranges.
+pub trait CacheRangeReader {
+    fn read_range(&self, range: std::ops::Range<usize>) -> Result<Bytes>;
+}
+impl<F: Fn(std::ops::Range<usize>) -> Result<Bytes>> CacheRangeReader for F {
+    fn read_range(&self, range: std::ops::Range<usize>) -> Result<Bytes> {
+        self(range)
+    }
+}
+
 /// Serialization trait for cache entries.
 ///
 /// **Experimental**: the serialized format is not yet covered by a stability
@@ -207,6 +218,40 @@ impl<T> CacheDecode<T> {
 /// The read sequence mirroring the write sequence for each `type_version` is
 /// the invariant the implementor owns.
 pub trait CacheCodecImpl: Send + Sync {
+    /// Whether this body supports gathering sorted, unique row offsets.
+    const SUPPORTS_ROW_SELECTION: bool = false;
+    fn deserialize_rows(
+        _reader: &dyn CacheRangeReader,
+        _body_offset: usize,
+        _version: u32,
+        _rows: &[u32],
+    ) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        Err(Error::invalid_input(
+            "cache codec does not support row selection",
+        ))
+    }
+
+    /// Payload byte ranges that [`deserialize_rows`](Self::deserialize_rows)
+    /// requests for `rows`, in request order, excluding the envelope and any
+    /// body header it reads through `reader` while planning.
+    ///
+    /// Lets a backend issue those reads concurrently before decoding. The
+    /// default `None` means the body cannot plan, and backends read serially.
+    fn plan_row_ranges(
+        _reader: &dyn CacheRangeReader,
+        _body_offset: usize,
+        _version: u32,
+        _rows: &[u32],
+    ) -> Result<Option<Vec<std::ops::Range<usize>>>>
+    where
+        Self: Sized,
+    {
+        Ok(None)
+    }
+
     /// Stable identity for this entry type. **Must not change once shipped.**
     /// This is a deliberate author-assigned string, not `std::any::type_name`
     /// (which is not stable across compiler versions).
@@ -234,15 +279,28 @@ pub trait CacheCodecImpl: Send + Sync {
 
 pub(crate) type ArcAny = Arc<dyn std::any::Any + Send + Sync>;
 
+type RowDecoder = fn(&dyn CacheRangeReader, usize, u32, &[u32]) -> Result<ArcAny>;
+type RowPlanner =
+    fn(&dyn CacheRangeReader, usize, u32, &[u32]) -> Result<Option<Vec<std::ops::Range<usize>>>>;
+
+/// Envelope bytes before `type_id`: magic, envelope version and `type_id_len`.
+const ENVELOPE_ID_OFFSET: usize = 4 + 1 + 2;
+/// Envelope bytes that follow `type_id`: `type_version`.
+const ENVELOPE_VERSION_BYTES: usize = 4;
+
 /// Type-erased codec for serializing and deserializing cache entries.
 ///
-/// `CacheCodec` carries the entry's stable `type_id`/`version` plus two plain
+/// `CacheCodec` carries the entry's stable `type_id`/`version` plus plain
 /// function pointers — it is `Copy` and has no heap allocation. Construct one
 /// via [`CacheCodec::from_impl`] for types that implement [`CacheCodecImpl`],
 /// or [`CacheCodec::new`] for custom cases (e.g. when the orphan rule prevents
 /// a direct impl).
 #[derive(Copy, Clone)]
 pub struct CacheCodec {
+    row_decoder: Option<RowDecoder>,
+    row_planner: Option<RowPlanner>,
+    memory_priority: u8,
+    plane_tag: Option<u8>,
     type_id: &'static str,
     version: u32,
     serialize_body: fn(&ArcAny, &mut CacheEntryWriter<'_>) -> Result<()>,
@@ -289,6 +347,10 @@ impl CacheCodec {
         deserialize_body: fn(&mut CacheEntryReader<'_>) -> Result<ArcAny>,
     ) -> Self {
         Self {
+            row_decoder: None,
+            row_planner: None,
+            memory_priority: 0,
+            plane_tag: None,
             type_id,
             version,
             serialize_body,
@@ -299,13 +361,49 @@ impl CacheCodec {
     /// Create a `CacheCodec` from a [`CacheCodecImpl`] implementation.
     pub fn from_impl<T: CacheCodecImpl + 'static>() -> Self {
         Self {
+            row_decoder: if T::SUPPORTS_ROW_SELECTION {
+                Some(|reader, offset, version, rows| {
+                    Ok(Arc::new(T::deserialize_rows(reader, offset, version, rows)?) as ArcAny)
+                })
+            } else {
+                None
+            },
+            row_planner: if T::SUPPORTS_ROW_SELECTION {
+                Some(T::plan_row_ranges)
+            } else {
+                None
+            },
+            memory_priority: 0,
+            plane_tag: None,
             type_id: T::TYPE_ID,
             version: T::CURRENT_VERSION,
             serialize_body: serialize_via_impl::<T>,
             deserialize_body: deserialize_via_impl::<T>,
         }
     }
 
+    /// Hint for priority-aware memory tiers. Larger values survive pressure first.
+    pub fn with_memory_priority(mut self, priority: u8) -> Self {
+        self.memory_priority = priority;
+        self
+    }
+
+    pub fn memory_priority(&self) -> u8 {
+        self.memory_priority
+    }
+
+    /// Tag the entry as one plane of a layered index so backends can report
+    /// per-plane metrics without deriving the plane from an admission priority.
+    pub fn with_plane_tag(mut self, plane: u8) -> Self {
+        self.plane_tag = Some(plane);
+        self
+    }
+
+    /// The layered plane this entry belongs to, or `None` for other entries.
+    pub fn plane_tag(&self) -> Option<u8> {
+        self.plane_tag
+    }
+
     /// Return the stable entry type identity.
     ///
     /// Persistent and tiered backends can use this metadata to route encoded
@@ -321,6 +419,62 @@ impl CacheCodec {
         (self.serialize_body)(value, &mut entry_writer)
     }
 
+    pub fn supports_row_selection(&self) -> bool {
+        self.row_decoder.is_some()
+    }
+
+    /// Read and validate the envelope of a range-read entry, returning the
+    /// body offset and the entry's `type_version`.
+    fn read_row_envelope(&self, reader: &dyn CacheRangeReader) -> Result<(usize, u32)> {
+        let prefix = reader.read_range(0..ENVELOPE_ID_OFFSET)?;
+        let id_bytes = prefix
+            .get(ENVELOPE_ID_OFFSET - 2..ENVELOPE_ID_OFFSET)
+            .ok_or_else(|| Error::invalid_input("short cache envelope"))?;
+        let id_len = u16::from_le_bytes(id_bytes.try_into().unwrap()) as usize;
+        let header =
+            reader.read_range(0..(ENVELOPE_ID_OFFSET + id_len + ENVELOPE_VERSION_BYTES))?;
+        let envelope = parse_envelope(&header)
+            .ok_or_else(|| Error::invalid_input("invalid cache envelope"))?;
+        if envelope.type_id != self.type_id || envelope.type_version > self.version {
+            return Err(Error::invalid_input("incompatible cache entry"));
+        }
+        Ok((envelope.body_offset, envelope.type_version))
+    }
+
+    /// Gather rows without loading the complete persistent entry.
+    pub fn deserialize_rows(
+        &self,
+        reader: &dyn CacheRangeReader,
+        rows: &[u32],
+    ) -> CacheDecode<ArcAny> {
+        let decode = || -> Result<ArcAny> {
+            let decoder = self
+                .row_decoder
+                .ok_or_else(|| Error::invalid_input("row selection unsupported"))?;
+            let (body_offset, version) = self.read_row_envelope(reader)?;
+            decoder(reader, body_offset, version, rows)
+        };
+        match decode() {
+            Ok(entry) => CacheDecode::Hit(entry),
+            Err(_) => CacheDecode::Miss(CacheMissReason::BodyError),
+        }
+    }
+
+    /// Payload ranges that [`deserialize_rows`](Self::deserialize_rows) would
+    /// request for `rows`, beyond the envelope and body header it reads from
+    /// `reader` here. `None` when the codec cannot plan, the entry is not
+    /// compatible, or the rows are invalid; callers then decode serially,
+    /// which reports invalid rows as a miss.
+    pub fn plan_rows(
+        &self,
+        reader: &dyn CacheRangeReader,
+        rows: &[u32],
+    ) -> Option<Vec<std::ops::Range<usize>>> {
+        let planner = self.row_planner?;
+        let (body_offset, version) = self.read_row_envelope(reader).ok()?;
+        planner(reader, body_offset, version, rows).ok().flatten()
+    }
+
     /// Deserialize an entry from `data`.
     ///
     /// Never fails: any non-fatal failure to interpret the bytes becomes a
@@ -536,6 +690,91 @@ mod tests {
         );
     }
 
+    #[test]
+    fn plane_tag_is_metadata_only() {
+        let plain = CacheCodec::from_impl::<Widget>();
+        assert_eq!(plain.plane_tag(), None);
+
+        let tagged = plain.with_plane_tag(2).with_memory_priority(1);
+        assert_eq!(tagged.plane_tag(), Some(2));
+        assert_eq!(tagged.memory_priority(), 1);
+        let copied = tagged;
+        assert_eq!(copied.plane_tag(), Some(2));
+
+        // The tag never reaches the wire format.
+        let value: ArcAny = Arc::new(Widget { n: 11 });
+        let (mut plain_bytes, mut tagged_bytes) = (Vec::new(), Vec::new());
+        plain.serialize(&value, &mut plain_bytes).unwrap();
+        tagged.serialize(&value, &mut tagged_bytes).unwrap();
+        assert_eq!(plain_bytes, tagged_bytes);
+        let decoded = tagged
+            .deserialize(&Bytes::from(tagged_bytes))
+            .hit()
+            .unwrap();
+        assert_eq!(decoded.downcast_ref::<Widget>(), Some(&Widget { n: 11 }));
+    }
+
+    /// Row-selectable body that keeps the default (absent) range planner.
+    struct Rows;
+
+    impl CacheCodecImpl for Rows {
+        const TYPE_ID: &'static str = "test.Rows";
+        const CURRENT_VERSION: u32 = 1;
+        const SUPPORTS_ROW_SELECTION: bool = true;
+
+        fn serialize(&self, _writer: &mut CacheEntryWriter<'_>) -> Result<()> {
+            Ok(())
+        }
+
+        fn deserialize(_reader: &mut CacheEntryReader<'_>) -> Result<Self> {
+            Ok(Self)
+        }
+
+        fn deserialize_rows(
+            _reader: &dyn CacheRangeReader,
+            _body_offset: usize,
+            _version: u32,
+            _rows: &[u32],
+        ) -> Result<Self> {
+            Ok(Self)
+        }
+    }
+
+    #[test]
+    fn plan_rows_defaults_to_serial_decode() {
+        let value: ArcAny = Arc::new(Widget { n: 1 });
+        let widget = CacheCodec::from_impl::<Widget>();
+        let mut widget_bytes = Vec::new();
+        widget.serialize(&value, &mut widget_bytes).unwrap();
+        let widget_bytes = Bytes::from(widget_bytes);
+        let read = |range: std::ops::Range<usize>| -> Result<Bytes> {
+            widget_bytes
+                .get(range.clone())
+                .map(|_| widget_bytes.slice(range))
+                .ok_or_else(|| Error::invalid_input("out of range"))
+        };
+        assert!(!widget.supports_row_selection());
+        assert!(widget.plan_rows(&read, &[0]).is_none());
+
+        let rows = CacheCodec::from_impl::<Rows>();
+        let mut rows_bytes = Vec::new();
+        rows.serialize(&(Arc::new(Rows) as ArcAny), &mut rows_bytes)
+            .unwrap();
+        let rows_bytes = Bytes::from(rows_bytes);
+        let read = |range: std::ops::Range<usize>| -> Result<Bytes> {
+            rows_bytes
+                .get(range.clone())
+                .map(|_| rows_bytes.slice(range))
+                .ok_or_else(|| Error::invalid_input("out of range"))
+        };
+        assert!(rows.supports_row_selection());
+        assert!(rows.plan_rows(&read, &[0]).is_none());
+        assert!(matches!(
+            rows.deserialize_rows(&read, &[0]),
+            CacheDecode::Hit(_)
+        ));
+    }
+
     #[test]
     fn reader_exposes_envelope_version() {
         // type_version travels through the envelope to reader.version().