feat(streaming): arrangement backfill

src/batch/src/executor/row_seq_scan.rs

@@ -34,7 +34,7 @@ use risingwave_pb::common::BatchQueryEpoch;
 use risingwave_pb::plan_common::StorageTableDesc;
 use risingwave_storage::store::PrefetchOptions;
 use risingwave_storage::table::batch_table::storage_table::StorageTable;
-use risingwave_storage::table::{Distribution, TableIter};
+use risingwave_storage::table::{collect_data_chunk, get_second, Distribution};
 use risingwave_storage::{dispatch_state_store, StateStore};
 
 use crate::executor::{
@@ -416,14 +416,14 @@ impl<S: StateStore> RowSeqScanExecutor<S> {
                 ordered,
                 PrefetchOptions::new_for_exhaust_iter(),
             )
-            .await?;
+            .await?
+            .map(get_second);
 
         pin_mut!(iter);
         loop {
             let timer = histogram.as_ref().map(|histogram| histogram.start_timer());
 
-            let chunk = iter
-                .collect_data_chunk(table.schema(), Some(chunk_size))
+            let chunk = collect_data_chunk(&mut iter, table.schema(), Some(chunk_size))
                 .await
                 .map_err(RwError::from)?;
 

src/storage/src/table/batch_table/mod.rs

@@ -12,5 +12,4 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-pub mod iter_utils;
 pub mod storage_table;

src/storage/src/table/batch_table/storage_table.rs

@@ -39,14 +39,14 @@ use risingwave_hummock_sdk::key::{end_bound_of_prefix, next_key, prefixed_range}
 use risingwave_hummock_sdk::HummockReadEpoch;
 use tracing::trace;
 
-use super::iter_utils;
 use crate::error::{StorageError, StorageResult};
 use crate::hummock::CachePolicy;
 use crate::row_serde::row_serde_util::{
     parse_raw_key_to_vnode_and_key, serialize_pk, serialize_pk_with_vnode,
 };
 use crate::row_serde::{find_columns_by_ids, ColumnMapping};
 use crate::store::{PrefetchOptions, ReadOptions};
+use crate::table::merge_sort::merge_sort;
 use crate::table::{compute_vnode, Distribution, TableIter, DEFAULT_VNODE};
 use crate::StateStore;
 
@@ -511,7 +511,7 @@ impl<S: StateStore, SD: ValueRowSerde> StorageTableInner<S, SD> {
             // Concat all iterators if not to preserve order.
             _ if !ordered => futures::stream::iter(iterators).flatten(),
             // Merge all iterators if to preserve order.
-            _ => iter_utils::merge_sort(iterators.into_iter().map(Box::pin).collect()),
+            _ => merge_sort(iterators.into_iter().map(Box::pin).collect()),
         };
 
         Ok(iter)

src/storage/src/table/batch_table/iter_utils.rs → src/storage/src/table/merge_sort.rs

@@ -14,64 +14,60 @@
 
 use std::collections::binary_heap::PeekMut;
 use std::collections::BinaryHeap;
+use std::error::Error;
 
-use futures::StreamExt;
+use futures::{Stream, StreamExt};
 use futures_async_stream::try_stream;
 use risingwave_common::row::OwnedRow;
 
-use super::storage_table::PkAndRowStream;
-use crate::error::StorageError;
+pub trait MergeSortKey = Eq + PartialEq + Ord + PartialOrd;
 
-/// We use a binary heap to merge the results of the different streams in order.
-/// This is the node type of the heap.
-struct Node<S: PkAndRowStream> {
+struct Node<K: MergeSortKey, S> {
     stream: S,
 
     /// The next item polled from `stream` previously. Since the `eq` and `cmp` must be synchronous
     /// functions, we need to implement peeking manually.
-    peeked: (Vec<u8>, OwnedRow),
+    peeked: (K, OwnedRow),
 }
 
-impl<S: PkAndRowStream> PartialEq for Node<S> {
+impl<K: MergeSortKey, S> PartialEq for Node<K, S> {
     fn eq(&self, other: &Self) -> bool {
         match self.peeked.0 == other.peeked.0 {
             true => unreachable!("primary key from different iters should be unique"),
             false => false,
         }
     }
 }
-impl<S: PkAndRowStream> Eq for Node<S> {}
+impl<K: MergeSortKey, S> Eq for Node<K, S> {}
 
-impl<S: PkAndRowStream> PartialOrd for Node<S> {
+impl<K: MergeSortKey, S> PartialOrd for Node<K, S> {
     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
         Some(self.cmp(other))
     }
 }
-impl<S: PkAndRowStream> Ord for Node<S> {
+
+impl<K: MergeSortKey, S> Ord for Node<K, S> {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
         // The heap is a max heap, so we need to reverse the order.
         self.peeked.0.cmp(&other.peeked.0).reverse()
     }
 }
 
-/// Merge multiple streams of primary key and rows into a single stream, sorted by primary key.
-/// We should ensure that the primary key from different streams are unique.
-#[try_stream(ok = (Vec<u8>, OwnedRow), error = StorageError)]
-pub(super) async fn merge_sort<S>(streams: Vec<S>)
+#[try_stream(ok=(K, OwnedRow), error=E)]
+pub async fn merge_sort<'a, K, E, R>(streams: Vec<R>)
 where
-    S: PkAndRowStream + Unpin,
+    K: MergeSortKey + 'a,
+    E: Error + 'a,
+    R: Stream<Item = Result<(K, OwnedRow), E>> + 'a + Unpin,
 {
-    let mut heap = BinaryHeap::with_capacity(streams.len());
+    let mut heap = BinaryHeap::new();
     for mut stream in streams {
         if let Some(peeked) = stream.next().await.transpose()? {
             heap.push(Node { stream, peeked });
         }
     }
-
     while let Some(mut node) = heap.peek_mut() {
         // Note: If the `next` returns `Err`, we'll fail to yield the previous item.
-        // This is acceptable since we're not going to handle errors from cell-based table
-        // iteration, so where to fail does not matter. Or we need an `Option` for this.
         yield match node.stream.next().await.transpose()? {
             // There still remains data in the stream, take and update the peeked value.
             Some(new_peeked) => std::mem::replace(&mut node.peeked, new_peeked),

src/storage/src/table/mod.rs

@@ -13,9 +13,11 @@
 // limitations under the License.
 
 pub mod batch_table;
+pub mod merge_sort;
 
 use std::sync::{Arc, LazyLock};
 
+use futures::{Stream, StreamExt};
 use itertools::Itertools;
 use risingwave_common::array::DataChunk;
 use risingwave_common::buffer::{Bitmap, BitmapBuilder};
@@ -81,43 +83,52 @@ impl Distribution {
 #[async_trait::async_trait]
 pub trait TableIter: Send {
     async fn next_row(&mut self) -> StorageResult<Option<OwnedRow>>;
+}
 
-    async fn collect_data_chunk(
-        &mut self,
-        schema: &Schema,
-        chunk_size: Option<usize>,
-    ) -> StorageResult<Option<DataChunk>> {
-        let mut builders = schema.create_array_builders(chunk_size.unwrap_or(0));
-
-        let mut row_count = 0;
-        for _ in 0..chunk_size.unwrap_or(usize::MAX) {
-            match self.next_row().await? {
-                Some(row) => {
-                    for (datum, builder) in row.iter().zip_eq_fast(builders.iter_mut()) {
-                        builder.append(datum);
-                    }
-                    row_count += 1;
+/// Collects data chunks from stream of rows.
+pub async fn collect_data_chunk<E, S>(
+    stream: &mut S,
+    schema: &Schema,
+    chunk_size: Option<usize>,
+) -> Result<Option<DataChunk>, E>
+where
+    S: Stream<Item = Result<OwnedRow, E>> + Unpin,
+{
+    let mut builders = schema.create_array_builders(chunk_size.unwrap_or(0));
+
+    let mut row_count = 0;
+    for _ in 0..chunk_size.unwrap_or(usize::MAX) {
+        match stream.next().await.transpose()? {
+            Some(row) => {
+                for (datum, builder) in row.iter().zip_eq_fast(builders.iter_mut()) {
+                    builder.append(datum);
                 }
-                None => break,
             }
+            None => break,
         }
 
-        let chunk = {
-            let columns: Vec<_> = builders
-                .into_iter()
-                .map(|builder| builder.finish().into())
-                .collect();
-            DataChunk::new(columns, row_count)
-        };
-
-        if chunk.cardinality() == 0 {
-            Ok(None)
-        } else {
-            Ok(Some(chunk))
-        }
+        row_count += 1;
+    }
+
+    let chunk = {
+        let columns: Vec<_> = builders
+            .into_iter()
+            .map(|builder| builder.finish().into())
+            .collect();
+        DataChunk::new(columns, row_count)
+    };
+
+    if chunk.cardinality() == 0 {
+        Ok(None)
+    } else {
+        Ok(Some(chunk))
     }
 }
 
+pub fn get_second<T, U, E>(arg: Result<(T, U), E>) -> Result<U, E> {
+    arg.map(|x| x.1)
+}
+
 /// Get vnode value with `indices` on the given `row`.
 pub fn compute_vnode(row: impl Row, indices: &[usize], vnodes: &Bitmap) -> VirtualNode {
     let vnode = if indices.is_empty() {

src/stream/src/common/table/state_table.rs

@@ -45,7 +45,7 @@ use risingwave_storage::row_serde::row_serde_util::{
 use risingwave_storage::store::{
     LocalStateStore, NewLocalOptions, PrefetchOptions, ReadOptions, StateStoreIterItemStream,
 };
-use risingwave_storage::table::{compute_chunk_vnode, compute_vnode, Distribution};
+use risingwave_storage::table::{compute_chunk_vnode, compute_vnode, get_second, Distribution};
 use risingwave_storage::StateStore;
 use tracing::{trace, Instrument};
 
@@ -62,6 +62,7 @@ const STATE_CLEANING_PERIOD_EPOCH: usize = 5;
 pub struct StateTableInner<
     S,
     SD = BasicSerde,
+    const IS_REPLICATED: bool = false,
     W = WatermarkBufferByEpoch<STATE_CLEANING_PERIOD_EPOCH>,
 > where
     S: StateStore,
@@ -120,9 +121,10 @@ pub struct StateTableInner<
 
 /// `StateTable` will use `BasicSerde` as default
 pub type StateTable<S> = StateTableInner<S, BasicSerde>;
+pub type ReplicatedStateTable<S> = StateTableInner<S, BasicSerde, true>;
 
 // initialize
-impl<S, SD, W> StateTableInner<S, SD, W>
+impl<S, SD, const IS_REPLICATED: bool, W> StateTableInner<S, SD, IS_REPLICATED, W>
 where
     S: StateStore,
     SD: ValueRowSerde,
@@ -188,13 +190,12 @@ where
         };
 
         let table_option = TableOption::build_table_option(table_catalog.get_properties());
-        let local_state_store = store
-            .new_local(NewLocalOptions::new(
-                table_id,
-                is_consistent_op,
-                table_option,
-            ))
-            .await;
+        let new_local_options = if IS_REPLICATED {
+            NewLocalOptions::new_replicated(table_id, is_consistent_op, table_option)
+        } else {
+            NewLocalOptions::new(table_id, is_consistent_op, table_option)
+        };
+        let local_state_store = store.new_local(new_local_options).await;
 
         let pk_data_types = pk_indices
             .iter()
@@ -511,7 +512,7 @@ where
 }
 
 // point get
-impl<S, SD> StateTableInner<S, SD>
+impl<S, SD, const IS_REPLICATED: bool> StateTableInner<S, SD, IS_REPLICATED>
 where
     S: StateStore,
     SD: ValueRowSerde,
@@ -609,7 +610,7 @@ where
 }
 
 // write
-impl<S, SD> StateTableInner<S, SD>
+impl<S, SD, const IS_REPLICATED: bool> StateTableInner<S, SD, IS_REPLICATED>
 where
     S: StateStore,
     SD: ValueRowSerde,
@@ -870,12 +871,8 @@ where
     }
 }
 
-fn get_second<T, U>(arg: StreamExecutorResult<(T, U)>) -> StreamExecutorResult<U> {
-    arg.map(|x| x.1)
-}
-
 // Iterator functions
-impl<S, SD, W> StateTableInner<S, SD, W>
+impl<S, SD, const IS_REPLICATED: bool, W> StateTableInner<S, SD, IS_REPLICATED, W>
 where
     S: StateStore,
     SD: ValueRowSerde,