diff --git a/crates/partition-store/src/vqueue_table/mod.rs b/crates/partition-store/src/vqueue_table/mod.rs index 5f75d9c533..18b5db6924 100644 --- a/crates/partition-store/src/vqueue_table/mod.rs +++ b/crates/partition-store/src/vqueue_table/mod.rs @@ -35,7 +35,7 @@ use tracing::error; use restate_rocksdb::{Priority, StorageTaskKind}; use restate_storage_api::StorageError; -use restate_storage_api::vqueue_table::filters::ScanEntryIdFilter; +use restate_storage_api::vqueue_table::filters::{ScanEntryIdFilter, ScanMetaFilter}; use restate_storage_api::vqueue_table::metadata::{VQueueMeta, VQueueMetaRef}; use restate_storage_api::vqueue_table::{ EntryKey, EntryMetadata, EntryStatusHeader, EntryValue, ReadVQueueTable, ScanVQueueTable, @@ -398,31 +398,155 @@ impl ReadVQueueTable for PartitionStoreTransaction<'_> { } impl ScanVQueueMetaTable for PartitionStore { - fn for_each_vqueue_meta< + fn for_each_vqueue_meta( + &self, + filter: ScanMetaFilter, + mut f: F, + ) -> Result> + Send> + where F: for<'a> FnMut((&'a VQueueId, &'a VQueueMetaRef<'a>)) -> std::ops::ControlFlow<()> + Send + Sync + 'static, - >( - &self, - range: KeyRange, - mut f: F, - ) -> Result> + Send> { - self.iterator_for_each( - "df-vqueue-meta", - Priority::Low, - TableScan::ScanPartitionKeyRange::(range), - move |(mut key, value)| { - let meta_key = break_on_err(MetaKey::deserialize_from(&mut key))?; - let meta = break_on_err( - VQueueMetaRef::decode_borrowed(value).map_err(StorageError::BilrostDecode), - )?; + { + // Fast path: each vqueue id maps to exactly one fixed-length key, so an + // exact set is served via batched multi-get calls instead of scanning + // every metadata row in the partition-key range. + if let ScanMetaFilter::MetaIdSet(ids) = filter { + return Ok(multi_get_vqueue_meta(self, ids, f).boxed()); + } - let (vqueue_id,) = meta_key.split(); - f((&vqueue_id, &meta)).map_break(Ok) - }, - ) - .map_err(|_| StorageError::OperationalError) + let scan = match filter { + ScanMetaFilter::PartitionKey(range) => { + TableScan::ScanPartitionKeyRange::(range) + } + ScanMetaFilter::MetaIdRange(range) => { + let start = MetaKey::from(&range.start); + let end = MetaKey::from(&range.last); + TableScan::RangeInclusive(start, end) + } + ScanMetaFilter::MetaIdSet(_) => unreachable!("handled above"), + }; + + let scan_fut = self + .iterator_for_each( + "df-vqueue-meta", + Priority::Low, + scan, + move |(mut key, value)| { + let meta_key = break_on_err(MetaKey::deserialize_from(&mut key))?; + let meta = break_on_err( + VQueueMetaRef::decode_borrowed(value).map_err(StorageError::BilrostDecode), + )?; + + let (vqueue_id,) = meta_key.split(); + f((&vqueue_id, &meta)).map_break(Ok) + }, + ) + .map_err(|_| StorageError::OperationalError)?; + + Ok(scan_fut.boxed()) + } +} + +/// Maximum number of vqueue-metadata keys passed to one RocksDB multi-get call. +const VQUEUE_META_MULTI_GET_BATCH_SIZE: usize = 500; + +/// Serves a vqueue-metadata lookup for a known set of ids via batched +/// `batched_multi_get` calls, dispatched on the storage background thread-pool. +/// +/// `ids` is already sorted in on-disk key order (`VQueueId`'s `Ord` matches its +/// key byte encoding, which is prefixed by the partition key), which is what +/// `batched_multi_get_cf_opt`'s `sorted_input=true` requires. +/// +/// `MetaKey`s are fixed-length, so each batch is packed back-to-back into a +/// single buffer and handed to the multi-get as `chunks_exact` slices. +fn multi_get_vqueue_meta( + store: &PartitionStore, + ids: BTreeSet, + mut f: F, +) -> impl Future> + Send +where + F: for<'a> FnMut((&'a VQueueId, &'a VQueueMetaRef<'a>)) -> std::ops::ControlFlow<()> + + Send + + Sync + + 'static, +{ + const KEY_LEN: usize = MetaKey::serialized_length_fixed(); + + let rocksdb = store.partition_db().rocksdb().clone(); + let cf_name: restate_rocksdb::CfName = store.partition_db().partition().cf_name().into(); + + async move { + rocksdb + .run_background_read_op( + "df-vqueue-meta", + StorageTaskKind::MultiGet, + Priority::Low, + move |raw_db| -> Result<()> { + let Some(cf) = raw_db.cf_handle(cf_name.as_str()) else { + return Err(StorageError::Generic(anyhow::anyhow!( + "column family {cf_name} not found for vqueue meta multi-get" + ))); + }; + + let batch_capacity = ids.len().min(VQUEUE_META_MULTI_GET_BATCH_SIZE); + let mut key_buf = BytesMut::with_capacity(batch_capacity * KEY_LEN); + let mut batch_ids = Vec::with_capacity(batch_capacity); + + let mut readopts = ReadOptions::default(); + // future proofing to make use of parallel L0 reads and async-io + // if/when we build rocksdb with COROUTINES=1 and IO-URING support. + // by default, this will not do anything. + readopts.set_async_io(true); + readopts.set_optimize_multiget_for_io(true); + + let mut ids = ids.into_iter(); + loop { + key_buf.clear(); + batch_ids.clear(); + + for id in ids.by_ref().take(VQUEUE_META_MULTI_GET_BATCH_SIZE) { + EncodeTableKey::serialize_to(&MetaKey::from(&id), &mut key_buf); + batch_ids.push(id); + } + + if batch_ids.is_empty() { + break; + } + + let results = raw_db.batched_multi_get_cf_opt( + &cf, + key_buf.chunks_exact(KEY_LEN), + true, + &readopts, + ); + + for (id, result) in batch_ids.iter().zip(results) { + let Some(value) = + result.map_err(|e| StorageError::Generic(e.into()))? + else { + continue; + }; + + let meta = VQueueMetaRef::decode_borrowed(value.as_ref()) + .map_err(StorageError::BilrostDecode)?; + + if f((id, &meta)).is_break() { + return Ok(()); + } + } + + if batch_ids.len() < VQUEUE_META_MULTI_GET_BATCH_SIZE { + break; + } + } + + Ok(()) + }, + ) + .await + .map_err(|_| StorageError::OperationalError)? } } diff --git a/crates/storage-api/src/vqueue_table/filters.rs b/crates/storage-api/src/vqueue_table/filters.rs index d6dc25c790..0120e65112 100644 --- a/crates/storage-api/src/vqueue_table/filters.rs +++ b/crates/storage-api/src/vqueue_table/filters.rs @@ -12,7 +12,7 @@ use std::collections::BTreeSet; use std::range::RangeInclusive; use restate_sharding::KeyRange; -use restate_types::vqueues::VQueueEntryId; +use restate_types::vqueues::{VQueueEntryId, VQueueId}; /// Filter vqueue entries by partition keys, entry ID range, or an exact set of entry IDs. #[derive(Debug, Clone)] @@ -23,3 +23,18 @@ pub enum ScanEntryIdFilter { /// range scan. The set is sorted in on-disk key order. EntryIdSet(BTreeSet), } + +/// Filter vqueue metadata rows by partition keys, vqueue ID range, or an exact +/// set of vqueue IDs. +/// +/// Each vqueue id maps to exactly one metadata row, so [`ScanMetaFilter::MetaIdSet`] +/// is served via a batched multi-get instead of a partition-key-range scan. +#[derive(Debug, Clone)] +pub enum ScanMetaFilter { + PartitionKey(KeyRange), + MetaIdRange(RangeInclusive), + /// A known set of vqueue IDs served via batched multi-get calls. The set is + /// sorted in on-disk key order (`VQueueId`'s `Ord` matches its key byte + /// encoding, which is prefixed by the partition key). + MetaIdSet(BTreeSet), +} diff --git a/crates/storage-api/src/vqueue_table/tables.rs b/crates/storage-api/src/vqueue_table/tables.rs index a8802153f2..ecbd32e7a0 100644 --- a/crates/storage-api/src/vqueue_table/tables.rs +++ b/crates/storage-api/src/vqueue_table/tables.rs @@ -11,7 +11,7 @@ use restate_sharding::{KeyRange, PartitionKey}; use restate_types::vqueues::{Seq, VQueueId}; -use super::filters::ScanEntryIdFilter; +use super::filters::{ScanEntryIdFilter, ScanMetaFilter}; use super::metadata::{VQueueMeta, VQueueMetaRef}; use super::{ EntryId, EntryKey, EntryMetadata, EntryStatusHeader, EntryValue, stats::EntryStatistics, @@ -255,7 +255,7 @@ pub trait ScanVQueueMetaTable { + 'static, >( &self, - range: KeyRange, + filter: ScanMetaFilter, f: F, ) -> Result> + Send>; } diff --git a/crates/storage-query-datafusion/src/filter.rs b/crates/storage-query-datafusion/src/filter.rs index e52bc2a4dd..815b739588 100644 --- a/crates/storage-query-datafusion/src/filter.rs +++ b/crates/storage-query-datafusion/src/filter.rs @@ -28,7 +28,7 @@ use restate_types::PartitionedResourceId; use restate_types::identifiers::partitioner::HashPartitioner; use restate_types::identifiers::{InvocationId, PartitionKey, ResourceId, WithPartitionKey}; use restate_types::sharding::KeyRange; -use restate_types::vqueues::VQueueEntryId; +use restate_types::vqueues::{VQueueEntryId, VQueueId}; use crate::partition_store_scanner::ScanLocalPartitionFilter; @@ -101,7 +101,31 @@ impl FirstMatchingPartitionKeyExtractor { T: PartitionedResourceId + ResourceId + FromStr, ::Err: std::error::Error + Send + Sync + 'static, { - let e = MatchingColumnExtractor::new(column_name, |value: &ScalarValue| { + self.append(Self::create_partitioned_resource_id_extractor::( + column_name, + )) + } + + /// Adds a partitioned-resource-id extractor whose matches are grouped by Restate partition. + pub fn with_grouped_partitioned_resource_id(self, column_name: impl Into) -> Self + where + T: PartitionedResourceId + ResourceId + FromStr, + ::Err: std::error::Error + Send + Sync + 'static, + { + self.append_with_fanout( + Self::create_partitioned_resource_id_extractor::(column_name), + PointReadFanout::PerPartition, + ) + } + + fn create_partitioned_resource_id_extractor( + column_name: impl Into, + ) -> impl PartitionKeyExtractor + where + T: PartitionedResourceId + ResourceId + FromStr, + ::Err: std::error::Error + Send + Sync + 'static, + { + MatchingColumnExtractor::new(column_name, |value: &ScalarValue| { let value = value .try_as_str() .with_context(|| format!("expected string {:?}", T::RESOURCE_TYPE))? @@ -109,8 +133,7 @@ impl FirstMatchingPartitionKeyExtractor { let resource = T::from_str(value).with_context(|| format!("non valid {:?}", T::RESOURCE_TYPE))?; Ok(resource.partition_key()) - }); - self.append(e) + }) } pub fn with_service_key(self, column_name: impl Into) -> Self { @@ -639,6 +662,40 @@ impl ScanLocalPartitionFilter for VQueueEntryIdFilter { } } +/// Each vqueue id maps to exactly one metadata row, so an `id = / IN (...)` +/// predicate is served via a batched multi-get (the `Set`) instead of a full +/// partition-key-range scan. `VQueueId` is not `Copy`, but `IdSelection` only +/// requires `Ord + Clone`. +#[derive(Debug, Clone)] +pub struct VQueueMetaFilter { + pub partition_keys: KeyRange, + pub ids: Option>, +} + +impl ScanLocalPartitionFilter for VQueueMetaFilter { + fn new(range: KeyRange, predicate: Option>) -> Self { + if let Some(predicate) = predicate + && let Ok(predicate) = snapshot_physical_expr(predicate) + { + for conjunct in split_conjunction(&predicate) { + if let Some(ids) = + parse_id_selection("id", range, conjunct, |id: &VQueueId| id.partition_key()) + { + return Self { + partition_keys: range, + ids: Some(ids), + }; + } + } + } + + Self { + partition_keys: range, + ids: None, + } + } +} + #[cfg(test)] mod tests { use std::collections::BTreeSet; @@ -655,11 +712,11 @@ mod tests { use restate_types::identifiers::{InvocationId, ServiceId, StateMutationId, WithPartitionKey}; use restate_types::invocation::{InvocationTarget, VirtualObjectHandlerType}; use restate_types::sharding::KeyRange; - use restate_types::vqueues::VQueueEntryId; + use restate_types::vqueues::{VQueueEntryId, VQueueId}; use crate::filter::{ FirstMatchingPartitionKeyExtractor, InvocationIdFilter, PartitionKeyExtractor, - VQueueEntryIdFilter, VQueueFilter, + VQueueEntryIdFilter, VQueueFilter, VQueueMetaFilter, }; use crate::partition_store_scanner::ScanLocalPartitionFilter; @@ -1330,4 +1387,66 @@ mod tests { assert!(filter.stages.is_none()); assert_eq!(filter.partition_keys, FULL_RANGE); } + + #[test] + fn vqueue_meta_filter_set_and_rejections() { + let id1 = VQueueId::custom(1, "q1"); + let id2 = VQueueId::custom(2, "q2"); + + // `id = / IN (...)` yields an exact set served via multi-get. + let filter = VQueueMetaFilter::new( + FULL_RANGE, + Some(in_list( + "id", + vec![utf8_lit(id1.to_string()), utf8_lit(id2.to_string())], + )), + ); + let selection = filter.ids.expect("should extract vqueue-id set"); + assert_eq!(selection.ids, BTreeSet::from([id1.clone(), id2])); + + // No predicate and a negated list both fall back to a range scan. + assert!(VQueueMetaFilter::new(FULL_RANGE, None).ids.is_none()); + assert!( + VQueueMetaFilter::new( + FULL_RANGE, + Some(not_in_list("id", vec![utf8_lit(id1.to_string())])), + ) + .ids + .is_none() + ); + } + + #[test] + fn vqueue_meta_filter_excludes_out_of_range() { + let id = VQueueId::custom(1, "q1"); + let pk = id.partition_key(); + let narrow_range = if pk > 0 { + KeyRange::new(0, pk - 1) + } else { + KeyRange::new(1, 1) + }; + + let filter = + VQueueMetaFilter::new(narrow_range, Some(eq(col("id"), utf8_lit(id.to_string())))); + assert!(filter.ids.is_none()); + } + + #[test] + fn vqueue_meta_filter_keeps_large_in_list_as_set() { + let ids = (0..501) + .map(|id| VQueueId::custom(id, format!("q{id}"))) + .collect::>(); + let predicate = in_list( + "id", + ids.iter().map(|id| utf8_lit(id.to_string())).collect(), + ); + + let filter = VQueueMetaFilter::new(FULL_RANGE, Some(predicate)); + + let selection = filter.ids.expect("should extract vqueue-id set"); + assert_eq!(selection.ids.len(), ids.len()); + for id in ids { + assert!(selection.ids.contains(&id)); + } + } } diff --git a/crates/storage-query-datafusion/src/vqueue_meta/mod.rs b/crates/storage-query-datafusion/src/vqueue_meta/mod.rs index 057a963489..9e16f0e6d3 100644 --- a/crates/storage-query-datafusion/src/vqueue_meta/mod.rs +++ b/crates/storage-query-datafusion/src/vqueue_meta/mod.rs @@ -13,3 +13,6 @@ pub(crate) mod schema; mod table; pub(crate) use table::register_self; + +#[cfg(test)] +mod tests; diff --git a/crates/storage-query-datafusion/src/vqueue_meta/table.rs b/crates/storage-query-datafusion/src/vqueue_meta/table.rs index 3c0d259621..05a93121e9 100644 --- a/crates/storage-query-datafusion/src/vqueue_meta/table.rs +++ b/crates/storage-query-datafusion/src/vqueue_meta/table.rs @@ -13,16 +13,17 @@ use std::ops::ControlFlow; use std::sync::Arc; use restate_partition_store::{PartitionStore, PartitionStoreManager}; -use restate_sharding::KeyRange; use restate_storage_api::StorageError; use restate_storage_api::vqueue_table::ScanVQueueMetaTable; +use restate_storage_api::vqueue_table::filters::ScanMetaFilter; use restate_storage_api::vqueue_table::metadata::VQueueMetaRef; use restate_types::vqueues::VQueueId; use crate::context::{QueryContext, SelectPartitions}; -use crate::filter::FirstMatchingPartitionKeyExtractor; +use crate::filter::{FirstMatchingPartitionKeyExtractor, VQueueMetaFilter}; use crate::partition_store_scanner::{LocalPartitionsScanner, ScanLocalPartition}; use crate::remote_query_scanner_manager::RemoteScannerManager; +use crate::statistics::{RowEstimate, TableStatisticsBuilder}; use crate::table_providers::{PartitionedTableProvider, ScanPartition}; use crate::vqueue_meta::row::append_vqueues_meta_row; use crate::vqueue_meta::schema::{SysVqueueMetaBuilder, sys_vqueue_meta_sort_order}; @@ -40,15 +41,24 @@ pub(crate) fn register_self( VQueuesMetaScanner, )) as Arc; + let schema = SysVqueueMetaBuilder::schema(); + + // There are far fewer vqueues than vqueue entries, so this table is small. + let statistics = TableStatisticsBuilder::new(schema.clone()) + .with_num_rows_estimate(RowEstimate::Small) + .with_partition_key() + .with_primary_key("id"); + let vqueue_meta_table = PartitionedTableProvider::new( partition_selector, - SysVqueueMetaBuilder::schema(), + schema, sys_vqueue_meta_sort_order(), remote_scanner_manager.create_distributed_scanner(NAME, local_scanner), FirstMatchingPartitionKeyExtractor::default() .with_scope("scope") - .with_partitioned_resource_id::("id"), - ); + .with_grouped_partitioned_resource_id::("id"), + ) + .with_statistics(statistics.build()); ctx.register_partitioned_table(NAME, Arc::new(vqueue_meta_table)) } @@ -60,7 +70,7 @@ impl ScanLocalPartition for VQueuesMetaScanner { type Builder = SysVqueueMetaBuilder; type Item<'a> = (&'a VQueueId, &'a VQueueMetaRef<'a>); type ConversionError = std::convert::Infallible; - type Filter = KeyRange; + type Filter = VQueueMetaFilter; fn for_each_row< F: for<'a> FnMut(Self::Item<'a>) -> ControlFlow> @@ -69,10 +79,11 @@ impl ScanLocalPartition for VQueuesMetaScanner { + 'static, >( partition_store: &PartitionStore, - range: KeyRange, + filter: VQueueMetaFilter, mut f: F, ) -> Result> + Send, StorageError> { - partition_store.for_each_vqueue_meta(range, move |item| f(item).map_break(Result::unwrap)) + partition_store + .for_each_vqueue_meta(filter.into(), move |item| f(item).map_break(Result::unwrap)) } fn append_row<'a>( @@ -83,3 +94,12 @@ impl ScanLocalPartition for VQueuesMetaScanner { Ok(()) } } + +impl From for ScanMetaFilter { + fn from(value: VQueueMetaFilter) -> Self { + match value.ids { + Some(selection) => ScanMetaFilter::MetaIdSet(selection.ids), + None => ScanMetaFilter::PartitionKey(value.partition_keys), + } + } +} diff --git a/crates/storage-query-datafusion/src/vqueue_meta/tests.rs b/crates/storage-query-datafusion/src/vqueue_meta/tests.rs new file mode 100644 index 0000000000..dc1a9754c8 --- /dev/null +++ b/crates/storage-query-datafusion/src/vqueue_meta/tests.rs @@ -0,0 +1,121 @@ +// Copyright (c) 2023 - 2026 Restate Software, Inc., Restate GmbH. +// All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +use datafusion::arrow::array::LargeStringArray; +use datafusion::arrow::record_batch::RecordBatch; +use futures::StreamExt; + +use restate_limiter::LimitKey; +use restate_storage_api::Transaction; +use restate_storage_api::vqueue_table::WriteVQueueTable; +use restate_storage_api::vqueue_table::metadata::{VQueueLink, VQueueMeta}; +use restate_types::clock::UniqueTimestamp; +use restate_types::time::MillisSinceEpoch; +use restate_types::vqueues::VQueueId; + +use crate::mocks::*; + +fn meta() -> VQueueMeta { + let created_at = + UniqueTimestamp::try_from_unix_millis(MillisSinceEpoch::new(1_744_010_000_000)).unwrap(); + VQueueMeta::new(created_at, None, LimitKey::None, VQueueLink::None) +} + +async fn select_ids(engine: &mut MockQueryEngine, query: &str) -> Vec { + let records = engine + .execute(query.to_owned()) + .await + .unwrap() + .stream + .collect::>>() + .await + .remove(0) + .unwrap(); + + let mut ids: Vec = records + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .map(str::to_string) + .collect(); + ids.sort(); + ids +} + +/// `id IN (...)` is served via the batched multi-get fast path. Only the listed +/// vqueues must come back, even though several other metadata rows share the +/// same partition. +#[restate_core::test(flavor = "multi_thread", worker_threads = 2)] +async fn vqueue_meta_point_query_returns_only_matching_ids() { + let mut engine = MockQueryEngine::create().await; + + let qids: Vec = (0..4) + .map(|i| VQueueId::custom(3337, format!("q{i}"))) + .collect(); + + let mut tx = engine.partition_store().transaction(); + for qid in &qids { + tx.create_vqueue(qid, &meta()); + } + tx.commit().await.unwrap(); + drop(tx); + + let got = select_ids( + &mut engine, + &format!( + "SELECT id FROM sys_vqueue_meta WHERE id IN ('{}', '{}')", + qids[0], qids[2] + ), + ) + .await; + + let mut expected = vec![qids[0].to_string(), qids[2].to_string()]; + expected.sort(); + assert_eq!(got, expected); +} + +/// `id NOT IN (> 3 values)` must fall back to a full partition-key-range scan +/// (the negated list can't become a lookup set) and return the non-excluded +/// rows. Four values is the smallest list that survives as a negated +/// `InListExpr`. +#[restate_core::test(flavor = "multi_thread", worker_threads = 2)] +async fn vqueue_meta_not_in_returns_non_excluded_rows() { + let mut engine = MockQueryEngine::create().await; + + let qids: Vec = (0..6) + .map(|i| VQueueId::custom(3337, format!("q{i}"))) + .collect(); + + let mut tx = engine.partition_store().transaction(); + for qid in &qids { + tx.create_vqueue(qid, &meta()); + } + tx.commit().await.unwrap(); + drop(tx); + + let excluded = qids[0..4] + .iter() + .map(|qid| format!("'{qid}'")) + .collect::>() + .join(", "); + let got = select_ids( + &mut engine, + &format!("SELECT id FROM sys_vqueue_meta WHERE id NOT IN ({excluded})"), + ) + .await; + + let mut expected = vec![qids[4].to_string(), qids[5].to_string()]; + expected.sort(); + assert_eq!(got, expected); +}