[VQueues] Introduce sys_vqueues datafusion table

Expose vqueue entries as a single DataFusion table with stage-aware scanning.
When the query filters `stage`, only matching stage key kinds are scanned;
without a stage filter, all inbox stages are scanned and merged.

Also project the latest entry metadata for observability (status plus
EntryStatistics timestamps and counters), and add targeted tests for stage
predicate extraction and sys_vqueues stage filtering behavior.

Author: AhmedSoliman

crates/partition-store/src/tests/vqueue_table_test/mod.rs

@@ -25,12 +25,14 @@
 
 use restate_clock::time::MillisSinceEpoch;
 use restate_storage_api::Transaction;
+use restate_storage_api::vqueue_table::ScanVQueueInboxStages;
 use restate_storage_api::vqueue_table::{
     EntryKey, EntryMetadata, EntryValue, Stage, Status, VQueueCursor, VQueueStore,
     WriteVQueueTable, stats::EntryStatistics,
 };
 use restate_types::clock::UniqueTimestamp;
 use restate_types::identifiers::PartitionKey;
+use restate_types::sharding::KeyRange;
 use restate_types::vqueues::{EntryId, EntryKind, VQueueId};
 
 use crate::PartitionStore;
@@ -837,6 +839,67 @@ async fn concurrent_enqueue_and_delete(rocksdb: &mut PartitionStore) {
     );
 }
 
+/// Test: Stage scan reads only the requested stage key kind.
+///
+/// This validates the datafusion-oriented scan API and ensures stage-specific
+/// scans do not leak rows from adjacent stage key kinds or partition keys.
+async fn stage_scan_is_filtered_by_stage(rocksdb: &mut PartitionStore) {
+    let target_partition_key = PartitionKey::from(9_300u64);
+    let other_partition_key = PartitionKey::from(9_301u64);
+    let target_qid = VQueueId::custom(target_partition_key, "scan-target");
+    let other_qid = VQueueId::custom(other_partition_key, "scan-target");
+
+    let stages = [
+        Stage::Inbox,
+        Stage::Running,
+        Stage::Suspended,
+        Stage::Paused,
+        Stage::Finished,
+    ];
+
+    {
+        let mut txn = rocksdb.transaction();
+        for (index, stage) in stages.into_iter().enumerate() {
+            let entry_id = 100 + index as u8;
+            let target_entry = default_entry(entry_id);
+            let other_entry = default_entry(entry_id + 10);
+
+            txn.put_vqueue_inbox(&target_qid, stage, &target_entry.0, &target_entry.1);
+            txn.put_vqueue_inbox(&other_qid, stage, &other_entry.0, &other_entry.1);
+        }
+        txn.commit().await.expect("commit should succeed");
+    }
+
+    let range = KeyRange::from(target_partition_key..=target_partition_key);
+
+    for (index, stage) in stages.into_iter().enumerate() {
+        let expected_key = default_entry(100 + index as u8).0;
+        let rows = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
+        let rows_for_scan = rows.clone();
+
+        rocksdb
+            .for_each_vqueue_inbox_entry(range, stage, move |(qid, got_stage, key, _)| {
+                rows_for_scan
+                    .lock()
+                    .expect("stage scan lock should not be poisoned")
+                    .push((qid.clone(), got_stage, *key));
+                std::ops::ControlFlow::Continue(())
+            })
+            .expect("stage scan setup should succeed")
+            .await
+            .expect("stage scan should succeed");
+
+        let rows = rows
+            .lock()
+            .expect("stage scan lock should not be poisoned")
+            .clone();
+        assert_eq!(rows.len(), 1, "stage {stage} should return one row");
+        assert_eq!(rows[0].0, target_qid, "stage {stage} returned wrong qid");
+        assert_eq!(rows[0].1, stage, "stage {stage} returned wrong stage");
+        assert_eq!(rows[0].2, expected_key, "stage {stage} returned wrong key");
+    }
+}
+
 pub(crate) async fn run_tests(mut rocksdb: PartitionStore) {
     let mut txn = rocksdb.transaction();
 
@@ -863,6 +926,8 @@ pub(crate) async fn run_tests(mut rocksdb: PartitionStore) {
     verify_waiting_cursor_boundary_is_respected(db);
     verify_waiting_cursor_partition_prefix_boundary_is_respected(db);
 
+    stage_scan_is_filtered_by_stage(&mut rocksdb).await;
+
     // Tailing iterator tests - these need mutable access to rocksdb for writes
     tailing_iterator_sees_new_items_on_reseek(&mut rocksdb).await;
     reseek_shows_new_higher_order_items(&mut rocksdb).await;

crates/partition-store/src/vqueue_table/mod.rs

@@ -18,6 +18,7 @@ mod running_reader;
 mod waiting_reader;
 
 use std::io::Cursor;
+use std::pin::Pin;
 
 pub use entry::{EntryStatusKey, EntryStatusKeyRef, StatusHeaderRaw};
 pub use inbox::InboxKey;
@@ -32,12 +33,12 @@ use tracing::error;
 
 use restate_rocksdb::Priority;
 use restate_storage_api::StorageError;
-use restate_storage_api::vqueue_table::ScanVQueueMetaTable;
 use restate_storage_api::vqueue_table::metadata::{VQueueMeta, VQueueMetaRef, VQueueMetaUpdates};
 use restate_storage_api::vqueue_table::{
     EntryKey, EntryMetadata, EntryStatusHeader, EntryValue, LazyEntryStatus, ReadVQueueTable,
     ScanVQueueTable, Stage, Status, WriteVQueueTable, stats::EntryStatistics,
 };
+use restate_storage_api::vqueue_table::{ScanVQueueInboxStages, ScanVQueueMetaTable};
 use restate_types::sharding::{KeyRange, PartitionKey};
 use restate_types::vqueues::{EntryId, Seq, VQueueId};
 
@@ -474,6 +475,118 @@ impl ScanVQueueMetaTable for PartitionStore {
     }
 }
 
+fn stage_from_key_kind(key_kind: KeyKind) -> Option<Stage> {
+    match key_kind {
+        KeyKind::VQueueInboxStage => Some(Stage::Inbox),
+        KeyKind::VQueueRunningStage => Some(Stage::Running),
+        KeyKind::VQueueSuspendedStage => Some(Stage::Suspended),
+        KeyKind::VQueuePausedStage => Some(Stage::Paused),
+        KeyKind::VQueueFinishedStage => Some(Stage::Finished),
+        _ => None,
+    }
+}
+
+fn scan_vqueue_inbox_stage<'store, K, F>(
+    partition_store: &'store PartitionStore,
+    scanner_name: &'static str,
+    range: KeyRange,
+    stage: Stage,
+    mut f: F,
+) -> Result<Pin<Box<dyn Future<Output = Result<()>> + Send + 'store>>>
+where
+    K: EncodeTableKeyPrefix + 'store,
+    F: for<'row> FnMut(
+            (&'row VQueueId, Stage, &'row EntryKey, &'row EntryValue),
+        ) -> std::ops::ControlFlow<()>
+        + Send
+        + Sync
+        + 'static,
+{
+    let future = partition_store
+        .iterator_for_each(
+            scanner_name,
+            Priority::Low,
+            TableScan::FullScanPartitionKeyRange::<K>(range),
+            move |(mut key, mut value)| {
+                let key_kind = break_on_err(KeyKind::deserialize(&mut key))?;
+                let key_stage = break_on_err(
+                    stage_from_key_kind(key_kind).ok_or(StorageError::DataIntegrityError),
+                )?;
+                if key_stage != stage {
+                    return std::ops::ControlFlow::Break(Err(StorageError::DataIntegrityError));
+                }
+
+                let qid: VQueueId = break_on_err(crate::keys::deserialize(&mut key))?;
+                let entry_key: EntryKey = break_on_err(crate::keys::deserialize(&mut key))?;
+                let entry = break_on_err(
+                    EntryValue::decode(&mut value).map_err(StorageError::BilrostDecode),
+                )?;
+
+                f((&qid, stage, &entry_key, &entry)).map_break(Ok)
+            },
+        )
+        .map_err(|_| StorageError::OperationalError)?;
+
+    Ok(Box::pin(future))
+}
+
+impl ScanVQueueInboxStages for PartitionStore {
+    fn for_each_vqueue_inbox_entry<
+        F: for<'a> FnMut(
+                (&'a VQueueId, Stage, &'a EntryKey, &'a EntryValue),
+            ) -> std::ops::ControlFlow<()>
+            + Send
+            + Sync
+            + 'static,
+    >(
+        &self,
+        range: KeyRange,
+        stage: Stage,
+        f: F,
+    ) -> Result<impl Future<Output = Result<()>> + Send> {
+        match stage {
+            Stage::Unknown => Err(StorageError::Generic(anyhow::anyhow!(
+                "Unknown stage can't be scanned"
+            ))),
+            Stage::Inbox => scan_vqueue_inbox_stage::<inbox::InboxKey, _>(
+                self,
+                "df-vqueue-inbox",
+                range,
+                stage,
+                f,
+            ),
+            Stage::Running => scan_vqueue_inbox_stage::<inbox::RunningKey, _>(
+                self,
+                "df-vqueue-running",
+                range,
+                stage,
+                f,
+            ),
+            Stage::Suspended => scan_vqueue_inbox_stage::<inbox::SuspendedKey, _>(
+                self,
+                "df-vqueue-suspended",
+                range,
+                stage,
+                f,
+            ),
+            Stage::Paused => scan_vqueue_inbox_stage::<inbox::PausedKey, _>(
+                self,
+                "df-vqueue-paused",
+                range,
+                stage,
+                f,
+            ),
+            Stage::Finished => scan_vqueue_inbox_stage::<inbox::FinishedKey, _>(
+                self,
+                "df-vqueue-finished",
+                range,
+                stage,
+                f,
+            ),
+        }
+    }
+}
+
 // ## Safety
 // The iterator is guaranteed to be dropped before the database is dropped, we hold to the
 // PartitionDb in this struct for as long as the iterator is alive.

crates/storage-api/src/vqueue_table/entry.rs

@@ -205,7 +205,7 @@ impl<'a> From<&'a EntryMetadata> for EntryMetadataRef<'a> {
 pub struct EntryMetadata {
     // todo: This is temporary placeholder, type and name _will_ change.
     #[bilrost(tag(1))]
-    deployment: Option<String>,
+    pub deployment: Option<String>,
 }
 
 #[cfg(test)]

crates/storage-api/src/vqueue_table/tables.rs

@@ -221,3 +221,20 @@ pub trait ScanVQueueMetaTable {
         f: F,
     ) -> Result<impl Future<Output = Result<()>> + Send>;
 }
+
+pub trait ScanVQueueInboxStages {
+    /// Used for data-fusion queries
+    fn for_each_vqueue_inbox_entry<
+        F: for<'a> FnMut(
+                (&'a VQueueId, Stage, &'a EntryKey, &'a EntryValue),
+            ) -> std::ops::ControlFlow<()>
+            + Send
+            + Sync
+            + 'static,
+    >(
+        &self,
+        range: KeyRange,
+        stage: Stage,
+        f: F,
+    ) -> Result<impl Future<Output = Result<()>> + Send>;
+}

crates/storage-query-datafusion/src/context.rs

@@ -288,6 +288,12 @@ where
             self.partition_store_manager.clone(),
             &self.remote_scanner_manager,
         )?;
+        crate::vqueues::register_self(
+            ctx,
+            self.partition_selector.clone(),
+            self.partition_store_manager.clone(),
+            &self.remote_scanner_manager,
+        )?;
 
         ctx.datafusion_context.sql(SYS_INVOCATION_VIEW).await?;