perf: use SlateDB multi_get for batch sample loading

Replace individual get() calls and narrow range scans with a new
multi_get API that does a single scan from min to max key with a
HashSet filter. This amortizes SST index block reads and data block
reads across all requested keys.

Changes:
- Add multi_get to StorageRead trait with default fallback to individual gets
- Override in SlateDbStorageReader to use DbReader::multi_get (scan-based)
- Update get_time_series_batch and get_forward_index_batch to use multi_get
- Update preload_all_samples to batch by bucket using samples_batch
- Switch SlateDB dependency to fork with multi_get + tracing instrumentation

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: apurvam

Cargo.lock

@@ -43,8 +43,8 @@ serde_json = "1.0"
 serde_with = { version = "3", features = ["base64"] }
 serde_yaml = "0.9"
 foyer = "0.18"
-slatedb = "0.11.1"
-slatedb-common = "0.11.1"
+slatedb = { git = "https://github.com/apurvam/slatedb", branch = "multi-get-prototype" }
+slatedb-common = { git = "https://github.com/apurvam/slatedb", branch = "multi-get-prototype" }
 thiserror = "2.0"
 tokio = { version = "1.0", features = ["full"] }
 tokio-util = "0.7"

Cargo.toml

@@ -5,6 +5,7 @@ pub mod loader;
 pub mod slate;
 pub mod util;
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -253,6 +254,24 @@ pub trait StorageRead: Send + Sync {
         tracing::Span::current().record("num_records", records.len());
         Ok(records)
     }
+
+    /// Retrieves multiple records by their keys in a single batch operation.
+    ///
+    /// Returns a map of key to record for all keys that were found.
+    /// Keys not present in storage are omitted from the result.
+    ///
+    /// The default implementation falls back to individual `get()` calls.
+    /// Implementations may override this for better performance (e.g., using
+    /// a scan from min to max key with a filter).
+    async fn multi_get(&self, keys: &[Bytes]) -> StorageResult<HashMap<Bytes, Record>> {
+        let mut results = HashMap::with_capacity(keys.len());
+        for key in keys {
+            if let Some(record) = self.get(key.clone()).await? {
+                results.insert(key.clone(), record);
+            }
+        }
+        Ok(results)
+    }
 }
 
 /// A point-in-time snapshot of the storage layer.

common/src/storage/mod.rs

@@ -450,6 +450,25 @@ impl StorageRead for SlateDbStorageReader {
             .map_err(StorageError::from_storage)?;
         Ok(Box::new(SlateDbIterator { iter }))
     }
+
+    #[tracing::instrument(level = "debug", skip_all, fields(num_keys = keys.len(), num_found, num_scanned))]
+    async fn multi_get(
+        &self,
+        keys: &[Bytes],
+    ) -> StorageResult<std::collections::HashMap<Bytes, Record>> {
+        let result = self
+            .reader
+            .multi_get(keys)
+            .await
+            .map_err(StorageError::from_storage)?;
+
+        tracing::Span::current().record("num_found", result.len());
+
+        Ok(result
+            .into_iter()
+            .map(|(k, v)| (k.clone(), Record::new(k, v)))
+            .collect())
+    }
 }
 
 #[cfg(test)]

common/src/storage/slate.rs

@@ -709,48 +709,54 @@ impl<'reader, R: QueryReader> CachedQueryReader<'reader, R> {
     /// Preload samples for all (bucket, series_ids) pairs concurrently across
     /// all buckets in a single parallel batch.
     ///
-    /// Uses a sequential scan per bucket (via `samples_batch`) when the number
-    /// of uncached series is large, which is dramatically faster than thousands
-    /// of individual point reads.
+    /// Uses high concurrency (128) for individual gets, flattened across all
+    /// buckets into a single stream for maximum parallelism.
     pub(crate) async fn preload_all_samples(
         &mut self,
         bucket_candidates: &[(TimeBucket, Vec<SeriesId>)],
     ) -> Result<()> {
-        // Group uncached series by bucket.
-        let mut uncached_by_bucket: Vec<(TimeBucket, Vec<SeriesId>)> = Vec::new();
-        for (bucket, series_ids) in bucket_candidates {
-            let uncached_ids: Vec<SeriesId> = series_ids
-                .iter()
-                .filter(|id| self.cache.get_samples(bucket, id).is_none())
-                .copied()
-                .collect();
-            if !uncached_ids.is_empty() {
-                uncached_by_bucket.push((*bucket, uncached_ids));
-            }
-        }
+        // Group uncached series by bucket for batch loading.
+        let uncached_by_bucket: Vec<(TimeBucket, Vec<SeriesId>)> = bucket_candidates
+            .iter()
+            .filter_map(|(bucket, series_ids)| {
+                let uncached: Vec<SeriesId> = series_ids
+                    .iter()
+                    .filter(|id| self.cache.get_samples(bucket, id).is_none())
+                    .copied()
+                    .collect();
+                if uncached.is_empty() {
+                    None
+                } else {
+                    Some((*bucket, uncached))
+                }
+            })
+            .collect();
 
         if uncached_by_bucket.is_empty() {
             return Ok(());
         }
 
-        // Use batch scan per bucket (scan vs individual gets decided by the reader)
+        let total_uncached: usize = uncached_by_bucket.iter().map(|(_, ids)| ids.len()).sum();
+
+        // Load all buckets concurrently using samples_batch (which uses multi_get internally).
         let reader = self.reader;
         let results: Vec<std::result::Result<_, crate::error::Error>> =
             futures::stream::iter(uncached_by_bucket)
                 .map(|(bucket, series_ids)| async move {
-                    let batch = reader
+                    let samples_map = reader
                         .samples_batch(&bucket, &series_ids, i64::MIN, i64::MAX)
                         .await?;
-                    Ok((bucket, batch))
+                    Ok((bucket, samples_map))
                 })
-                .buffer_unordered(8)
+                .buffer_unordered(4) // one per bucket, not per series
                 .collect()
                 .await;
 
+        self.stats.samples_cache_misses += total_uncached as u64;
+
         for result in results {
-            let (bucket, batch) = result?;
-            self.stats.samples_cache_misses += batch.len() as u64;
-            for (series_id, samples) in batch {
+            let (bucket, samples_map) = result?;
+            for (series_id, samples) in samples_map {
                 self.cache.cache_samples(bucket, series_id, samples);
             }
         }

timeseries/src/promql/evaluator.rs

@@ -4,7 +4,7 @@ use common::{Record, Storage, StorageRead};
 use futures::stream::{self, StreamExt};
 use roaring::RoaringBitmap;
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 
 use crate::index::{InvertedIndex, SeriesSpec};
 use crate::model::{Sample, SeriesFingerprint, SeriesId, TimeBucket};
@@ -255,114 +255,73 @@ pub(crate) trait OpenTsdbStorageReadExt: StorageRead {
         Ok(max_series_id)
     }
 
-    /// Load samples for a batch of series using a narrowed sequential scan.
+    /// Load samples for a batch of series using multi_get.
     ///
-    /// Scans from the minimum to maximum series_id key, filtering to only
-    /// the requested IDs via HashSet lookup. This is faster than N individual
-    /// `get()` calls because sequential scan has much lower per-record overhead
-    /// (~0.35ms vs ~6ms for point reads due to LSM tree traversal).
-    ///
-    /// Stops early once all requested series have been found.
-    #[tracing::instrument(level = "info", skip(self, bucket, series_ids), fields(bucket_start = bucket.start, wanted = series_ids.len(), scanned, range_size))]
+    /// Uses the storage layer's multi_get which internally performs a
+    /// sequential scan from min to max key. This is much faster than
+    /// individual get() calls because it amortizes SST index reads and
+    /// data block reads across all requested keys.
+    #[tracing::instrument(level = "info", skip(self, bucket, series_ids), fields(bucket_start = bucket.start, wanted = series_ids.len(), found))]
     async fn get_time_series_batch(
         &self,
         bucket: &TimeBucket,
         series_ids: &[SeriesId],
     ) -> Result<HashMap<SeriesId, Vec<Sample>>> {
-        let wanted: HashSet<SeriesId> = series_ids.iter().copied().collect();
-        let min_id = *series_ids.iter().min().unwrap();
-        let max_id = *series_ids.iter().max().unwrap();
-        let range_size = (max_id - min_id + 1) as u64;
-        tracing::Span::current().record("range_size", range_size);
-
-        // Narrow scan from min to max+1 series_id key
-        let start_key = TimeSeriesKey {
-            time_bucket: bucket.start,
-            bucket_size: bucket.size,
-            series_id: min_id,
-        }
-        .encode();
-        let end_key = TimeSeriesKey {
-            time_bucket: bucket.start,
-            bucket_size: bucket.size,
-            series_id: max_id.wrapping_add(1),
-        }
-        .encode();
-        let range = common::BytesRange::new(
-            std::ops::Bound::Included(start_key),
-            std::ops::Bound::Excluded(end_key),
-        );
-
-        let mut iter = self.scan_iter(range).await?;
-        let mut result = HashMap::with_capacity(series_ids.len());
-        let mut scanned = 0u64;
-        let mut found = 0usize;
-        while let Some(record) = iter.next().await? {
-            scanned += 1;
-            let key = TimeSeriesKey::decode(record.key.as_ref())?;
-            if wanted.contains(&key.series_id) {
-                let samples: Vec<Sample> = match TimeSeriesIterator::new(record.value.as_ref()) {
-                    Some(iter) => iter.filter_map(|r| r.ok()).collect(),
-                    None => Vec::new(),
-                };
-                result.insert(key.series_id, samples);
-                found += 1;
-                if found == wanted.len() {
-                    break; // All requested series found
+        let keys: Vec<bytes::Bytes> = series_ids
+            .iter()
+            .map(|&series_id| {
+                TimeSeriesKey {
+                    time_bucket: bucket.start,
+                    bucket_size: bucket.size,
+                    series_id,
                 }
-            }
+                .encode()
+            })
+            .collect();
+
+        let records = self.multi_get(&keys).await?;
+        let mut result = HashMap::with_capacity(records.len());
+
+        for (key_bytes, record) in records {
+            let key = TimeSeriesKey::decode(key_bytes.as_ref())?;
+            let samples: Vec<Sample> = match TimeSeriesIterator::new(record.value.as_ref()) {
+                Some(iter) => iter.filter_map(|r| r.ok()).collect(),
+                None => Vec::new(),
+            };
+            result.insert(key.series_id, samples);
         }
-        tracing::Span::current().record("scanned", scanned);
+        tracing::Span::current().record("found", result.len());
         Ok(result)
     }
 
-    /// Load forward index entries for a batch of series using a narrowed scan.
-    #[tracing::instrument(level = "info", skip(self, bucket, series_ids), fields(bucket_start = bucket.start, wanted = series_ids.len(), scanned, range_size))]
+    /// Load forward index entries for a batch of series using multi_get.
+    #[tracing::instrument(level = "info", skip(self, bucket, series_ids), fields(bucket_start = bucket.start, wanted = series_ids.len(), found))]
     async fn get_forward_index_batch(
         &self,
         bucket: &TimeBucket,
         series_ids: &[SeriesId],
     ) -> Result<ForwardIndex> {
-        let wanted: HashSet<SeriesId> = series_ids.iter().copied().collect();
-        let min_id = *series_ids.iter().min().unwrap();
-        let max_id = *series_ids.iter().max().unwrap();
-        let range_size = (max_id - min_id + 1) as u64;
-        tracing::Span::current().record("range_size", range_size);
-
-        let start_key = ForwardIndexKey {
-            time_bucket: bucket.start,
-            bucket_size: bucket.size,
-            series_id: min_id,
-        }
-        .encode();
-        let end_key = ForwardIndexKey {
-            time_bucket: bucket.start,
-            bucket_size: bucket.size,
-            series_id: max_id.wrapping_add(1),
-        }
-        .encode();
-        let range = common::BytesRange::new(
-            std::ops::Bound::Included(start_key),
-            std::ops::Bound::Excluded(end_key),
-        );
+        let keys: Vec<bytes::Bytes> = series_ids
+            .iter()
+            .map(|&series_id| {
+                ForwardIndexKey {
+                    time_bucket: bucket.start,
+                    bucket_size: bucket.size,
+                    series_id,
+                }
+                .encode()
+            })
+            .collect();
 
-        let mut iter = self.scan_iter(range).await?;
+        let records = self.multi_get(&keys).await?;
         let result = ForwardIndex::default();
-        let mut scanned = 0u64;
-        let mut found = 0usize;
-        while let Some(record) = iter.next().await? {
-            scanned += 1;
-            let key = ForwardIndexKey::decode(record.key.as_ref())?;
-            if wanted.contains(&key.series_id) {
-                let value = ForwardIndexValue::decode(record.value.as_ref())?;
-                result.series.insert(key.series_id, value.into());
-                found += 1;
-                if found == wanted.len() {
-                    break;
-                }
-            }
+
+        for (key_bytes, record) in records {
+            let key = ForwardIndexKey::decode(key_bytes.as_ref())?;
+            let value = ForwardIndexValue::decode(record.value.as_ref())?;
+            result.series.insert(key.series_id, value.into());
         }
-        tracing::Span::current().record("scanned", scanned);
+        tracing::Span::current().record("found", result.series.len());
         Ok(result)
     }