diff --git a/.cargo/config.toml b/.cargo/config.toml index 1dfbe381a5563..e59bb43bdf223 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -31,7 +31,10 @@ rustflags = [ # Flags for all targets. [target.'cfg(all())'] -rustflags = ["--cfg", "tokio_unstable"] +rustflags = ["--cfg", "tokio_unstable", "-Zhigher-ranked-assumptions"] + +[build] +rustdocflags = ["-Zhigher-ranked-assumptions"] # We have large git dependencies. This can make cloning faster. # https://doc.rust-lang.org/nightly/cargo/reference/unstable.html#git diff --git a/ci/scripts/doc.sh b/ci/scripts/doc.sh index 0459d4915bf07..f7b4e34d49e43 100755 --- a/ci/scripts/doc.sh +++ b/ci/scripts/doc.sh @@ -9,14 +9,14 @@ echo "--- Set openssl static link env vars" configure_static_openssl echo "--- Build documentation" -RUSTDOCFLAGS="-Dwarnings" cargo doc --document-private-items --no-deps +RUSTDOCFLAGS="-Dwarnings -Zhigher-ranked-assumptions" cargo doc --document-private-items --no-deps echo "--- Show sccache stats" sccache --show-stats sccache --zero-stats echo "--- Run doctest" -RUSTDOCFLAGS="-Clink-arg=-fuse-ld=lld" cargo test --doc +cargo test --doc echo "--- Show sccache stats" sccache --show-stats diff --git a/src/storage/hummock_test/src/hummock_vector_tests.rs b/src/storage/hummock_test/src/hummock_vector_tests.rs index e2faebb6c367e..98752f2488f24 100644 --- a/src/storage/hummock_test/src/hummock_vector_tests.rs +++ b/src/storage/hummock_test/src/hummock_vector_tests.rs @@ -82,7 +82,7 @@ async fn test_flat_vector() { let epoch1_vectors = (0..100).map(|_| next_input()).collect_vec(); for (vec, info) in &epoch1_vectors { - vector_writer.insert(vec.clone(), info.clone()).unwrap(); + vector_writer.insert(vec.to_ref(), info.clone()).unwrap(); vector_writer.try_flush().await.unwrap(); } @@ -95,7 +95,7 @@ async fn test_flat_vector() { let epoch2_vectors = (0..100).map(|_| next_input()).collect_vec(); for (vec, info) in &epoch2_vectors { - vector_writer.insert(vec.clone(), info.clone()).unwrap(); + vector_writer.insert(vec.to_ref(), info.clone()).unwrap(); vector_writer.try_flush().await.unwrap(); } @@ -152,7 +152,7 @@ async fn test_flat_vector() { let output = read_snapshot_epoch .nearest( - query.clone(), + query.to_ref(), VectorNearestOptions { top_n, measure: DistanceMeasurement::InnerProduct, @@ -189,7 +189,7 @@ async fn test_flat_vector() { vector_writer.init_for_test(epoch3).await.unwrap(); let epoch3_vectors = (0..100).map(|_| next_input()).collect_vec(); for (vec, info) in &epoch3_vectors { - vector_writer.insert(vec.clone(), info.clone()).unwrap(); + vector_writer.insert(vec.to_ref(), info.clone()).unwrap(); vector_writer.try_flush().await.unwrap(); } @@ -264,7 +264,7 @@ async fn test_hnsw_vector() { let epoch1_vectors = (0..100).map(|_| next_input()).collect_vec(); for (vec, info) in &epoch1_vectors { - vector_writer.insert(vec.clone(), info.clone()).unwrap(); + vector_writer.insert(vec.to_ref(), info.clone()).unwrap(); vector_writer.try_flush().await.unwrap(); } @@ -277,7 +277,7 @@ async fn test_hnsw_vector() { let epoch2_vectors = (0..100).map(|_| next_input()).collect_vec(); for (vec, info) in &epoch2_vectors { - vector_writer.insert(vec.clone(), info.clone()).unwrap(); + vector_writer.insert(vec.to_ref(), info.clone()).unwrap(); vector_writer.try_flush().await.unwrap(); } @@ -334,7 +334,7 @@ async fn test_hnsw_vector() { let top_n = 10; let output = read_snapshot_epoch .nearest( - query.clone(), + query.to_ref(), VectorNearestOptions { top_n, measure: DistanceMeasurement::InnerProduct, @@ -374,7 +374,7 @@ async fn test_hnsw_vector() { vector_writer.init_for_test(epoch3).await.unwrap(); let epoch3_vectors = (0..100).map(|_| next_input()).collect_vec(); for (vec, info) in &epoch3_vectors { - vector_writer.insert(vec.clone(), info.clone()).unwrap(); + vector_writer.insert(vec.to_ref(), info.clone()).unwrap(); vector_writer.try_flush().await.unwrap(); } diff --git a/src/storage/src/hummock/store/hummock_storage.rs b/src/storage/src/hummock/store/hummock_storage.rs index b938254e964d3..5c20b9cdbd9ac 100644 --- a/src/storage/src/hummock/store/hummock_storage.rs +++ b/src/storage/src/hummock/store/hummock_storage.rs @@ -20,6 +20,7 @@ use std::sync::Arc; use arc_swap::ArcSwap; use bytes::Bytes; use itertools::Itertools; +use risingwave_common::array::VectorRef; use risingwave_common::catalog::TableId; use risingwave_common::dispatch_distance_measurement; use risingwave_common::util::epoch::is_max_epoch; @@ -260,11 +261,11 @@ impl HummockStorageReadSnapshot { /// If `Ok(Some())` is returned, the key is found. If `Ok(None)` is returned, /// the key is not found. If `Err()` is returned, the searching for the key /// failed due to other non-EOF errors. - async fn get_inner( - &self, + async fn get_inner<'a, O>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { let key_range = (Bound::Included(key.clone()), Bound::Included(key.clone())); @@ -639,12 +640,12 @@ pub struct HummockStorageReadSnapshot { } impl StateStoreGet for HummockStorageReadSnapshot { - fn on_key_value( - &self, + fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, - ) -> impl StorageFuture<'_, Option> { + on_key_value_fn: impl KeyValueFn<'a, O>, + ) -> impl StorageFuture<'a, Option> { self.get_inner(key, read_options, on_key_value_fn) } } @@ -687,11 +688,11 @@ impl StateStoreRead for HummockStorageReadSnapshot { } impl StateStoreReadVector for HummockStorageReadSnapshot { - async fn nearest( - &self, - vec: Vector, + async fn nearest<'a, O: Send + 'a>( + &'a self, + vec: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, ) -> StorageResult> { let version = match self.epoch { HummockReadEpoch::Committed(epoch) diff --git a/src/storage/src/hummock/store/local_hummock_storage.rs b/src/storage/src/hummock/store/local_hummock_storage.rs index 84e4c60692c4b..f06d5a825a9c9 100644 --- a/src/storage/src/hummock/store/local_hummock_storage.rs +++ b/src/storage/src/hummock/store/local_hummock_storage.rs @@ -107,12 +107,12 @@ pub struct LocalHummockStorage { } impl LocalHummockFlushedSnapshotReader { - async fn get_flushed( - hummock_version_reader: &HummockVersionReader, + async fn get_flushed<'a, O>( + hummock_version_reader: &'a HummockVersionReader, read_version: &HummockReadVersionRef, user_key: UserKey, read_options: ReadOptions, - on_key_value_fn: impl crate::store::KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { let table_key_range = ( Bound::Included(user_key.table_key.clone()), @@ -266,11 +266,11 @@ pub struct LocalHummockFlushedSnapshotReader { } impl StateStoreGet for LocalHummockFlushedSnapshotReader { - async fn on_key_value( - &self, + async fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { let key = UserKey::new(self.table_id, key); Self::get_flushed( @@ -308,11 +308,11 @@ impl StateStoreRead for LocalHummockFlushedSnapshotReader { } impl StateStoreGet for LocalHummockStorage { - async fn on_key_value( - &self, + async fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { let key = UserKey::new(self.table_id, key); match self.mem_table.buffer.get(&key.table_key) { diff --git a/src/storage/src/hummock/store/vector_writer.rs b/src/storage/src/hummock/store/vector_writer.rs index e9a659b890aa5..fa51dc92d6723 100644 --- a/src/storage/src/hummock/store/vector_writer.rs +++ b/src/storage/src/hummock/store/vector_writer.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use bytes::Bytes; +use risingwave_common::array::VectorRef; use risingwave_common::catalog::TableId; use risingwave_common::util::epoch::EpochPair; use risingwave_hummock_sdk::HummockEpoch; @@ -171,7 +172,7 @@ impl StateStoreWriteEpochControl for HummockVectorWriter { } impl StateStoreWriteVector for HummockVectorWriter { - fn insert(&mut self, vec: Vector, info: Bytes) -> StorageResult<()> { + fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()> { Ok(self .state .as_mut() diff --git a/src/storage/src/hummock/store/version.rs b/src/storage/src/hummock/store/version.rs index 38c07f5a1c274..f19647bf00bde 100644 --- a/src/storage/src/hummock/store/version.rs +++ b/src/storage/src/hummock/store/version.rs @@ -22,6 +22,7 @@ use bytes::Bytes; use futures::future::try_join_all; use itertools::Itertools; use parking_lot::RwLock; +use risingwave_common::array::VectorRef; use risingwave_common::bitmap::Bitmap; use risingwave_common::catalog::TableId; use risingwave_common::hash::VirtualNode; @@ -67,7 +68,7 @@ use crate::monitor::{ GetLocalMetricsGuard, HummockStateStoreMetrics, IterLocalMetricsGuard, StoreLocalStatistic, }; use crate::store::{ - OnNearestItemFn, ReadLogOptions, ReadOptions, Vector, VectorNearestOptions, gen_min_epoch, + OnNearestItemFn, ReadLogOptions, ReadOptions, VectorNearestOptions, gen_min_epoch, }; use crate::vector::hnsw::nearest; use crate::vector::{MeasureDistanceBuilder, NearestBuilder}; @@ -588,14 +589,14 @@ impl HummockVersionReader { const SLOW_ITER_FETCH_META_DURATION_SECOND: f64 = 5.0; impl HummockVersionReader { - pub async fn get( - &self, + pub async fn get<'a, O>( + &'a self, table_key: TableKey, epoch: u64, table_id: TableId, read_options: ReadOptions, read_version_tuple: ReadVersionTuple, - on_key_value_fn: impl crate::store::KeyValueFn, + on_key_value_fn: impl crate::store::KeyValueFn<'a, O>, ) -> StorageResult> { let (imms, uncommitted_ssts, committed_version) = read_version_tuple; @@ -1201,13 +1202,13 @@ impl HummockVersionReader { .await } - pub async fn nearest( - &self, + pub async fn nearest<'a, M: MeasureDistanceBuilder, O: Send>( + &'a self, version: PinnedVersion, table_id: TableId, - target: Vector, + target: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, ) -> HummockResult> { let Some(index) = version.vector_indexes.get(&table_id) else { return Ok(vec![]); @@ -1221,7 +1222,7 @@ impl HummockVersionReader { } match &index.inner { VectorIndexImpl::Flat(flat) => { - let mut builder = NearestBuilder::<'_, O, M>::new(target.to_ref(), options.top_n); + let mut builder = NearestBuilder::<'_, O, M>::new(target, options.top_n); for vector_file in &flat.vector_store_info.vector_files { let meta = self.sstable_store.get_vector_file_meta(vector_file).await?; for (i, block_meta) in meta.block_metas.iter().enumerate() { @@ -1246,7 +1247,7 @@ impl HummockVersionReader { let (items, _stats) = nearest::( &vector_store, &*graph, - target.to_ref(), + target, on_nearest_item_fn, options.hnsw_ef_search, options.top_n, diff --git a/src/storage/src/hummock/vector/writer/hnsw.rs b/src/storage/src/hummock/vector/writer/hnsw.rs index 916e7a518c4be..c543974451bec 100644 --- a/src/storage/src/hummock/vector/writer/hnsw.rs +++ b/src/storage/src/hummock/vector/writer/hnsw.rs @@ -18,6 +18,7 @@ use bytes::{Bytes, BytesMut}; use prost::Message; use rand::SeedableRng; use rand::rngs::StdRng; +use risingwave_common::array::VectorRef; use risingwave_common::dispatch_distance_measurement; use risingwave_common::vector::distance::DistanceMeasurement; use risingwave_hummock_sdk::HummockObjectId; @@ -30,7 +31,6 @@ use crate::hummock::vector::file::FileVectorStore; use crate::hummock::vector::writer::VectorObjectIdManagerRef; use crate::hummock::{HummockResult, SstableStoreRef}; use crate::opts::StorageOpts; -use crate::store::Vector; use crate::vector::hnsw::{ HnswBuilderOptions, HnswGraphBuilder, VectorAccessor, insert_graph, new_node, }; @@ -90,13 +90,13 @@ impl HnswFlatIndexWriter { }) } - pub(crate) fn insert(&mut self, vec: Vector, info: Bytes) -> HummockResult<()> { + pub(crate) fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> HummockResult<()> { self.vector_store .building_vectors .as_mut() .expect("for write") .file_builder - .add(vec.to_ref(), &info); + .add(vec, &info); Ok(()) } diff --git a/src/storage/src/hummock/vector/writer/mod.rs b/src/storage/src/hummock/vector/writer/mod.rs index 55a56fb671af3..52ffa16685982 100644 --- a/src/storage/src/hummock/vector/writer/mod.rs +++ b/src/storage/src/hummock/vector/writer/mod.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use bytes::Bytes; use futures::FutureExt; use hnsw::HnswFlatIndexWriter; +use risingwave_common::array::VectorRef; use risingwave_common::vector::distance::DistanceMeasurement; use risingwave_hummock_sdk::vector_index::{ FlatIndex, FlatIndexAdd, VectorFileInfo, VectorIndex, VectorIndexAdd, VectorIndexImpl, @@ -29,7 +30,6 @@ use risingwave_hummock_sdk::{HummockObjectId, HummockRawObjectId}; use crate::hummock::vector::file::VectorFileBuilder; use crate::hummock::{HummockResult, ObjectIdManager, SstableStoreRef}; use crate::opts::StorageOpts; -use crate::vector::Vector; #[async_trait::async_trait] pub trait VectorObjectIdManager: Send + Sync { @@ -105,7 +105,7 @@ impl VectorWriterImpl { }) } - pub(crate) fn insert(&mut self, vec: Vector, info: Bytes) -> HummockResult<()> { + pub(crate) fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> HummockResult<()> { match self { VectorWriterImpl::Flat(writer) => writer.insert(vec, info), VectorWriterImpl::HnswFlat(writer) => writer.insert(vec, info), @@ -163,8 +163,8 @@ impl FlatIndexWriter { } } - pub(crate) fn insert(&mut self, vec: Vector, info: Bytes) -> HummockResult<()> { - self.vector_file_builder.add(vec.to_ref(), info.as_ref()); + pub(crate) fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> HummockResult<()> { + self.vector_file_builder.add(vec, info.as_ref()); Ok(()) } diff --git a/src/storage/src/memory.rs b/src/storage/src/memory.rs index facc08050a6b2..b4aabcb856ee1 100644 --- a/src/storage/src/memory.rs +++ b/src/storage/src/memory.rs @@ -22,10 +22,12 @@ use std::sync::{Arc, LazyLock}; use bytes::Bytes; use itertools::Itertools; use parking_lot::RwLock; +use risingwave_common::array::VectorRef; use risingwave_common::bitmap::{Bitmap, BitmapBuilder}; use risingwave_common::catalog::{TableId, TableOption}; use risingwave_common::dispatch_distance_measurement; use risingwave_common::hash::{VirtualNode, VnodeBitmapExt}; +use risingwave_common::types::ScalarRef; use risingwave_common::util::epoch::{EpochPair, MAX_EPOCH}; use risingwave_hummock_sdk::key::{ FullKey, TableKey, TableKeyRange, UserKey, prefixed_range_with_vnode, @@ -687,11 +689,11 @@ pub struct RangeKvStateStoreReadSnapshot { } impl StateStoreGet for RangeKvStateStoreReadSnapshot { - async fn on_key_value( - &self, + async fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, _read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { self.inner .get_keyed_row_impl(key, self.epoch, self.table_id) @@ -728,21 +730,21 @@ impl StateStoreRead for RangeKvStateStoreReadSnapshot { } impl StateStoreReadVector for RangeKvStateStoreReadSnapshot { - async fn nearest( - &self, - vec: Vector, + async fn nearest<'a, O: Send + 'a>( + &'a self, + vec: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, ) -> StorageResult> { - fn nearest_impl( - store: &InMemVectorStore, + fn nearest_impl<'a, M: MeasureDistanceBuilder, O>( + store: &'a InMemVectorStore, epoch: u64, table_id: TableId, - vec: Vector, + vec: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, ) -> Vec { - let mut builder = NearestBuilder::<'_, O, M>::new(vec.to_ref(), options.top_n); + let mut builder = NearestBuilder::<'_, O, M>::new(vec, options.top_n); builder.add( store .read() @@ -1007,11 +1009,11 @@ impl RangeKvLocalStateStore { } impl StateStoreGet for RangeKvLocalStateStore { - async fn on_key_value( - &self, + async fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, _read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { if let Some((key, value)) = match self.mem_table.buffer.get(&key) { None => self @@ -1285,8 +1287,8 @@ impl StateStoreWriteEpochControl for RangeKvLocalStateStore { } impl StateStoreWriteVector for RangeKvLocalStateStore { - fn insert(&mut self, vec: Vector, info: Bytes) -> StorageResult<()> { - self.vectors.push((vec, info)); + fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()> { + self.vectors.push((vec.to_owned_scalar(), info)); Ok(()) } } diff --git a/src/storage/src/monitor/monitored_store.rs b/src/storage/src/monitor/monitored_store.rs index d745438bb42e5..6135bfc976ec3 100644 --- a/src/storage/src/monitor/monitored_store.rs +++ b/src/storage/src/monitor/monitored_store.rs @@ -23,6 +23,7 @@ use futures::{Future, FutureExt, TryFutureExt}; use risingwave_common::bitmap::Bitmap; use risingwave_common::catalog::TableId; use risingwave_common::hash::VirtualNode; +use risingwave_common::types::VectorRef; use risingwave_hummock_sdk::key::{TableKey, TableKeyRange}; use risingwave_hummock_sdk::{HummockEpoch, HummockReadEpoch, SyncResult}; use thiserror_ext::AsReport; @@ -139,12 +140,12 @@ impl MonitoredStateStore { } impl StateStoreGet for MonitoredTableStateStore { - fn on_key_value( - &self, + fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, - ) -> impl StorageFuture<'_, Option> { + on_key_value_fn: impl KeyValueFn<'a, O>, + ) -> impl StorageFuture<'a, Option> { let table_id = self.table_id(); let key_len = key.len(); self.monitored_on_key_value( @@ -207,12 +208,12 @@ impl StateStoreReadLog for MonitoredStateStore { } impl StateStoreReadVector for MonitoredTableStateStore { - fn nearest( - &self, - vec: Vector, + fn nearest<'a, O: Send + 'a>( + &'a self, + vec: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, - ) -> impl StorageFuture<'_, Vec> { + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, + ) -> impl StorageFuture<'a, Vec> { // TODO: monitor self.inner.nearest(vec, options, on_nearest_item_fn) } @@ -300,7 +301,7 @@ impl StateStoreWriteEpochControl for MonitoredTa } impl StateStoreWriteVector for MonitoredTableStateStore { - fn insert(&mut self, vec: Vector, info: Bytes) -> StorageResult<()> { + fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()> { // TODO: monitor self.inner.insert(vec, info) } diff --git a/src/storage/src/monitor/traced_store.rs b/src/storage/src/monitor/traced_store.rs index ba2e56ea7c023..b7b7a6b8f9856 100644 --- a/src/storage/src/monitor/traced_store.rs +++ b/src/storage/src/monitor/traced_store.rs @@ -18,6 +18,7 @@ use std::sync::Arc; use bytes::Bytes; use futures::future::BoxFuture; use futures::{Future, FutureExt}; +use risingwave_common::array::VectorRef; use risingwave_common::bitmap::Bitmap; use risingwave_common::catalog::TableId; use risingwave_common::hash::VirtualNode; @@ -145,11 +146,11 @@ impl TracedStateStore { } impl StateStoreGet for TracedStateStore { - async fn on_key_value( - &self, + async fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { if let Some((key, value)) = self .traced_get_keyed_row( @@ -336,12 +337,12 @@ impl StateStore for TracedStateStore { } impl StateStoreReadVector for TracedStateStore { - fn nearest( - &self, - vec: Vector, + fn nearest<'a, O: Send + 'a>( + &'a self, + vec: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, - ) -> impl StorageFuture<'_, Vec> { + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, + ) -> impl StorageFuture<'a, Vec> { self.inner.nearest(vec, options, on_nearest_item_fn) } } diff --git a/src/storage/src/panic_store.rs b/src/storage/src/panic_store.rs index de774e924cb51..9c0b75610b429 100644 --- a/src/storage/src/panic_store.rs +++ b/src/storage/src/panic_store.rs @@ -16,6 +16,7 @@ use std::marker::PhantomData; use std::sync::Arc; use bytes::Bytes; +use risingwave_common::array::VectorRef; use risingwave_common::bitmap::Bitmap; use risingwave_common::hash::VirtualNode; use risingwave_hummock_sdk::HummockReadEpoch; @@ -30,12 +31,12 @@ use crate::store::*; pub struct PanicStateStore; impl StateStoreGet for PanicStateStore { - fn on_key_value( - &self, + fn on_key_value<'a, O: Send + 'a>( + &'a self, _key: TableKey, _read_options: ReadOptions, - _on_key_value_fn: impl KeyValueFn, - ) -> impl StorageFuture<'_, Option> { + _on_key_value_fn: impl KeyValueFn<'a, O>, + ) -> impl StorageFuture<'a, Option> { async { panic!("should not read from PanicStateStore") } } } @@ -146,17 +147,17 @@ impl StateStoreWriteEpochControl for PanicStateStore { } impl StateStoreWriteVector for PanicStateStore { - fn insert(&mut self, _vec: Vector, _info: Bytes) -> StorageResult<()> { + fn insert(&mut self, _vec: VectorRef<'_>, _info: Bytes) -> StorageResult<()> { panic!() } } impl StateStoreReadVector for PanicStateStore { - async fn nearest( - &self, - _vec: Vector, + async fn nearest<'a, O: Send + 'a>( + &'a self, + _vec: VectorRef<'a>, _options: VectorNearestOptions, - _on_nearest_item_fn: impl OnNearestItemFn, + _on_nearest_item_fn: impl OnNearestItemFn<'a, O>, ) -> StorageResult> { panic!() } diff --git a/src/storage/src/store.rs b/src/storage/src/store.rs index 346295a6be84d..e8abd45add720 100644 --- a/src/storage/src/store.rs +++ b/src/storage/src/store.rs @@ -22,7 +22,7 @@ use bytes::Bytes; use futures::{Stream, TryStreamExt}; use futures_async_stream::try_stream; use prost::Message; -use risingwave_common::array::Op; +use risingwave_common::array::{Op, VectorRef}; use risingwave_common::bitmap::Bitmap; use risingwave_common::catalog::{TableId, TableOption}; use risingwave_common::hash::VirtualNode; @@ -257,16 +257,16 @@ pub trait StateStoreReadLog: StaticSendSync { ) -> impl StorageFuture<'_, Self::ChangeLogIter>; } -pub trait KeyValueFn = - for<'kv> FnOnce(FullKey<&'kv [u8]>, &'kv [u8]) -> StorageResult + Send + 'static; +pub trait KeyValueFn<'a, O> = + for<'kv> FnOnce(FullKey<&'kv [u8]>, &'kv [u8]) -> StorageResult + Send + 'a; pub trait StateStoreGet: StaticSendSync { - fn on_key_value( - &self, + fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, - ) -> impl StorageFuture<'_, Option>; + on_key_value_fn: impl KeyValueFn<'a, O>, + ) -> impl StorageFuture<'a, Option>; } pub trait StateStoreRead: StateStoreGet + StaticSendSync { @@ -428,7 +428,7 @@ pub trait StateStoreWriteEpochControl: StaticSendSync { } pub trait StateStoreWriteVector: StateStoreWriteEpochControl + StaticSendSync { - fn insert(&mut self, vec: Vector, info: Bytes) -> StorageResult<()>; + fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()>; } pub struct VectorNearestOptions { @@ -437,15 +437,15 @@ pub struct VectorNearestOptions { pub hnsw_ef_search: usize, } -pub trait OnNearestItemFn = OnNearestItem + Send + Sync + 'static; +pub trait OnNearestItemFn<'a, O> = OnNearestItem + Send + Sync + 'a; pub trait StateStoreReadVector: StaticSendSync { - fn nearest( - &self, - vec: Vector, + fn nearest<'a, O: Send + 'a>( + &'a self, + vec: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, - ) -> impl StorageFuture<'_, Vec>; + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, + ) -> impl StorageFuture<'a, Vec>; } /// If `prefetch` is true, prefetch will be enabled. Prefetching may increase the memory diff --git a/src/storage/src/store_impl.rs b/src/storage/src/store_impl.rs index 609509738e081..6be788235ea75 100644 --- a/src/storage/src/store_impl.rs +++ b/src/storage/src/store_impl.rs @@ -286,6 +286,7 @@ pub mod verify { use std::sync::Arc; use bytes::Bytes; + use risingwave_common::array::VectorRef; use risingwave_common::bitmap::Bitmap; use risingwave_common::hash::VirtualNode; use risingwave_hummock_sdk::HummockReadEpoch; @@ -331,11 +332,11 @@ pub mod verify { } impl StateStoreGet for VerifyStateStore { - async fn on_key_value( - &self, + async fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { let actual: Option<(FullKey, Bytes)> = self .actual @@ -368,12 +369,12 @@ pub mod verify { impl StateStoreReadVector for VerifyStateStore { - fn nearest( - &self, - vec: Vector, + fn nearest<'a, O: Send + 'a>( + &'a self, + vec: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, - ) -> impl StorageFuture<'_, Vec> { + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, + ) -> impl StorageFuture<'a, Vec> { self.actual.nearest(vec, options, on_nearest_item_fn) } } @@ -932,6 +933,7 @@ mod dyn_state_store { use std::sync::Arc; use bytes::Bytes; + use risingwave_common::array::VectorRef; use risingwave_common::bitmap::Bitmap; use risingwave_common::hash::VirtualNode; use risingwave_hummock_sdk::HummockReadEpoch; @@ -1244,12 +1246,12 @@ mod dyn_state_store { #[async_trait::async_trait] pub trait DynStateStoreWriteVector: DynStateStoreWriteEpochControl + StaticSendSync { - fn insert(&mut self, vec: Vector, info: Bytes) -> StorageResult<()>; + fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()>; } #[async_trait::async_trait] impl DynStateStoreWriteVector for S { - fn insert(&mut self, vec: Vector, info: Bytes) -> StorageResult<()> { + fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()> { self.insert(vec, info) } } @@ -1257,7 +1259,7 @@ mod dyn_state_store { pub type BoxDynStateStoreWriteVector = StateStorePointer>; impl StateStoreWriteVector for BoxDynStateStoreWriteVector { - fn insert(&mut self, vec: Vector, info: Bytes) -> StorageResult<()> { + fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()> { self.0.insert(vec, info) } } @@ -1268,7 +1270,7 @@ mod dyn_state_store { pub trait DynStateStoreReadVector: StaticSendSync { async fn nearest( &self, - vec: Vector, + vec: VectorRef<'_>, options: VectorNearestOptions, ) -> StorageResult>; } @@ -1277,7 +1279,7 @@ mod dyn_state_store { impl DynStateStoreReadVector for S { async fn nearest( &self, - vec: Vector, + vec: VectorRef<'_>, options: VectorNearestOptions, ) -> StorageResult> { use risingwave_common::types::ScalarRef; @@ -1296,11 +1298,11 @@ mod dyn_state_store { where StateStorePointer

: AsRef + StaticSendSync, { - async fn nearest( - &self, - vec: Vector, + async fn nearest<'a, O: Send + 'a>( + &'a self, + vec: VectorRef<'a>, options: VectorNearestOptions, - on_nearest_item_fn: impl OnNearestItemFn, + on_nearest_item_fn: impl OnNearestItemFn<'a, O>, ) -> StorageResult> { let output = self.as_ref().nearest(vec, options).await?; Ok(output @@ -1421,11 +1423,11 @@ mod dyn_state_store { where StateStorePointer

: AsRef + StaticSendSync, { - async fn on_key_value( - &self, + async fn on_key_value<'a, O: Send + 'a>( + &'a self, key: TableKey, read_options: ReadOptions, - on_key_value_fn: impl KeyValueFn, + on_key_value_fn: impl KeyValueFn<'a, O>, ) -> StorageResult> { let option = self.as_ref().get_keyed_row(key, read_options).await?; option diff --git a/src/storage/src/table/batch_table/mod.rs b/src/storage/src/table/batch_table/mod.rs index 5efadc219a6ac..23a42e3d82e65 100644 --- a/src/storage/src/table/batch_table/mod.rs +++ b/src/storage/src/table/batch_table/mod.rs @@ -409,11 +409,9 @@ impl BatchTableInner { }, ) .await?; - // TODO: may avoid the clone here when making the `on_key_value_fn` non-static - let row_serde = self.row_serde.clone(); match read_snapshot .on_key_value(serialized_pk, read_options, move |key, value| { - let row = row_serde.deserialize(value)?; + let row = self.row_serde.deserialize(value)?; Ok((key.epoch_with_gap.pure_epoch(), row)) }) .await? diff --git a/src/storage/src/table/batch_table/vector_index_reader.rs b/src/storage/src/table/batch_table/vector_index_reader.rs index 9158ed7fc8a41..9c9a62d3112bc 100644 --- a/src/storage/src/table/batch_table/vector_index_reader.rs +++ b/src/storage/src/table/batch_table/vector_index_reader.rs @@ -21,7 +21,7 @@ use risingwave_common::array::{ }; use risingwave_common::catalog::TableId; use risingwave_common::row::RowDeserializer; -use risingwave_common::types::{DataType, ScalarImpl, ScalarRef, StructType}; +use risingwave_common::types::{DataType, ScalarImpl, StructType}; use risingwave_common::util::value_encoding::BasicDeserializer; use risingwave_common::vector::distance::DistanceMeasurement; use risingwave_hummock_sdk::HummockReadEpoch; @@ -164,7 +164,7 @@ impl VectorIndexSnapshot<'_, S> { let row_results: Vec> = self .snapshot .nearest( - vector.to_owned_scalar(), + vector, VectorNearestOptions { top_n: self.reader.top_n, measure: self.reader.measure, diff --git a/src/stream/src/common/table/state_table.rs b/src/stream/src/common/table/state_table.rs index 480ec6da69768..a021c2677a0f3 100644 --- a/src/stream/src/common/table/state_table.rs +++ b/src/stream/src/common/table/state_table.rs @@ -967,12 +967,9 @@ impl StateTableRowStore { ..Default::default() }; - // TODO: avoid clone when `on_key_value_fn` can be non-static - let row_serde = self.row_serde.clone(); - self.state_store .on_key_value(key_bytes, read_options, move |_, value| { - let row = row_serde.deserialize(value)?; + let row = self.row_serde.deserialize(value)?; Ok(OwnedRow::new(row)) }) .await diff --git a/src/stream/src/executor/vector_index/mod.rs b/src/stream/src/executor/vector_index/mod.rs index ba2cdd098b55c..aecdf3bbf3acb 100644 --- a/src/stream/src/executor/vector_index/mod.rs +++ b/src/stream/src/executor/vector_index/mod.rs @@ -18,7 +18,6 @@ use itertools::Itertools; use risingwave_common::array::Op; use risingwave_common::catalog::TableId; use risingwave_common::row::{Row, RowExt}; -use risingwave_common::types::ScalarRef; use risingwave_common::util::value_encoding::{BasicSerializer, ValueRowSerializer}; use risingwave_storage::StateStore; use risingwave_storage::store::{ @@ -99,7 +98,6 @@ impl VectorIndexWriteExecutor { continue; }; let vector = vector_datum.into_vector(); - let vector = vector.to_owned_scalar(); let info = self .serializer .serialize(row.project(&info_column_indices))