feat(storage): change the prefix_hint to dist_key_hint for bloom_filter

src/batch/src/executor/insert.rs

@@ -307,7 +307,7 @@ mod tests {
                 epoch,
                 None,
                 ReadOptions {
-                    prefix_hint: None,
+                    dist_key_hint: None,
                     check_bloom_filter: false,
                     ignore_range_tombstone: false,
                     table_id: Default::default(),

src/common/src/catalog/internal_table.rs

@@ -14,6 +14,7 @@
 
 use std::sync::LazyLock;
 
+use itertools::Itertools;
 use regex::Regex;
 
 pub const RW_INTERNAL_TABLE_FUNCTION_NAME: &str = "rw_table";
@@ -38,3 +39,35 @@ pub fn valid_table_name(table_name: &str) -> bool {
         LazyLock::new(|| Regex::new(r"__internal_.*_\d+").unwrap());
     !INTERNAL_TABLE_NAME.is_match(table_name)
 }
+
+pub fn get_dist_key_in_pk_indices(dist_key_indices: &[usize], pk_indices: &[usize]) -> Vec<usize> {
+    let dist_key_in_pk_indices = dist_key_indices
+        .iter()
+        .map(|&di| {
+            pk_indices
+                .iter()
+                .position(|&pi| di == pi)
+                .unwrap_or_else(|| {
+                    panic!(
+                        "distribution key {:?} must be a subset of primary key {:?}",
+                        dist_key_indices, pk_indices
+                    )
+                })
+        })
+        .collect_vec();
+    dist_key_in_pk_indices
+}
+
+/// Get distribution key start index in pk, and return None if `dist_key_in_pk_indices` is not empty
+/// or continuous.
+/// Note that `dist_key_in_pk_indices` may be shuffled, the start index should be the
+/// minimum value.
+pub fn get_dist_key_start_index_in_pk(dist_key_in_pk_indices: &[usize]) -> Option<usize> {
+    match !dist_key_in_pk_indices.is_empty()
+        && *dist_key_in_pk_indices.iter().min().unwrap() + dist_key_in_pk_indices.len() - 1
+            == *dist_key_in_pk_indices.iter().max().unwrap()
+    {
+        true => Some(*dist_key_in_pk_indices.iter().min().unwrap()),
+        false => None,
+    }
+}

src/common/src/util/ordered/serde.rs

@@ -104,6 +104,34 @@ impl OrderedRowSerde {
 
         Ok(len)
     }
+
+    /// return the distribution key start position in serialized key and the distribution key
+    /// length.
+    pub fn deserialize_dist_key_position_with_column_indices(
+        &self,
+        key: &[u8],
+        dist_key_indices_pair: (usize, usize),
+    ) -> memcomparable::Result<(usize, usize)> {
+        let (dist_key_start_index, dist_key_end_index) = dist_key_indices_pair;
+        use crate::types::ScalarImpl;
+        let mut dist_key_start_position: usize = 0;
+        let mut len: usize = 0;
+        for index in 0..dist_key_end_index {
+            let data_type = &self.schema[index];
+            let order_type = &self.order_types[index];
+            let data = &key[len..];
+            let mut deserializer = memcomparable::Deserializer::new(data);
+            deserializer.set_reverse(*order_type == OrderType::Descending);
+
+            let field_length = ScalarImpl::encoding_data_size(data_type, &mut deserializer)?;
+            len += field_length;
+            if index < dist_key_start_index {
+                dist_key_start_position += field_length;
+            }
+        }
+
+        Ok((dist_key_start_position, (len - dist_key_start_position)))
+    }
 }
 
 #[cfg(test)]
@@ -237,6 +265,70 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_deserialize_dist_key_position_with_column_indices() {
+        let order_types = vec![
+            OrderType::Descending,
+            OrderType::Ascending,
+            OrderType::Descending,
+            OrderType::Ascending,
+        ];
+
+        let schema = vec![
+            DataType::Varchar,
+            DataType::Int16,
+            DataType::Varchar,
+            DataType::Varchar,
+        ];
+        let serde = OrderedRowSerde::new(schema, order_types);
+        let row1 = Row::new(vec![
+            Some(Utf8("aaa".to_string().into())),
+            Some(Int16(5)),
+            Some(Utf8("bbb".to_string().into())),
+            Some(Utf8("ccc".to_string().into())),
+        ]);
+        let rows = vec![row1];
+        let mut array = vec![];
+        for row in &rows {
+            let mut row_bytes = vec![];
+            serde.serialize(row, &mut row_bytes);
+            array.push(row_bytes);
+        }
+
+        {
+            let dist_key_indices = [1, 2];
+            let dist_key_start_index = 1;
+            let (dist_key_start_position, dist_key_len) = serde
+                .deserialize_dist_key_position_with_column_indices(
+                    &array[0],
+                    (
+                        dist_key_start_index,
+                        dist_key_start_index + dist_key_indices.len(),
+                    ),
+                )
+                .unwrap();
+
+            let schema = vec![DataType::Varchar];
+            let order_types = vec![OrderType::Descending];
+            let deserde = OrderedRowSerde::new(schema, order_types);
+            let prefix_slice = &array[0][0..dist_key_start_position];
+            assert_eq!(
+                deserde.deserialize(prefix_slice).unwrap(),
+                Row::new(vec![Some(Utf8("aaa".to_string().into()))])
+            );
+
+            let schema = vec![DataType::INT16, DataType::VARCHAR];
+            let order_types = vec![OrderType::Ascending, OrderType::Descending];
+            let deserde = OrderedRowSerde::new(schema, order_types);
+            let dist_key_slice =
+                &array[0][dist_key_start_position..dist_key_start_position + dist_key_len];
+            assert_eq!(
+                deserde.deserialize(dist_key_slice).unwrap(),
+                Row::new(vec![Some(Int16(5)), Some(Utf8("bbb".to_string().into()))])
+            );
+        }
+    }
+
     #[test]
     fn test_encoding_data_size() {
         use std::mem::size_of;

src/common/src/util/scan_range.rs

@@ -14,12 +14,12 @@
 
 use std::ops::{Bound, RangeBounds};
 
-use itertools::Itertools;
 use paste::paste;
 use risingwave_pb::batch_plan::scan_range::Bound as BoundProst;
 use risingwave_pb::batch_plan::ScanRange as ScanRangeProst;
 
 use super::value_encoding::serialize_datum_to_bytes;
+use crate::catalog::get_dist_key_in_pk_indices;
 use crate::hash::VirtualNode;
 use crate::row::{Row2, RowExt};
 use crate::types::{Datum, ScalarImpl};
@@ -84,20 +84,7 @@ impl ScanRange {
             return None;
         }
 
-        let dist_key_in_pk_indices = dist_key_indices
-            .iter()
-            .map(|&di| {
-                pk_indices
-                    .iter()
-                    .position(|&pi| di == pi)
-                    .unwrap_or_else(|| {
-                        panic!(
-                            "distribution keys {:?} must be a subset of primary keys {:?}",
-                            dist_key_indices, pk_indices
-                        )
-                    })
-            })
-            .collect_vec();
+        let dist_key_in_pk_indices = get_dist_key_in_pk_indices(dist_key_indices, pk_indices);
         let pk_prefix_len = self.eq_conds.len();
         if dist_key_in_pk_indices.iter().any(|&i| i >= pk_prefix_len) {
             return None;

src/ctl/src/cmd_impl/hummock/list_kv.rs

@@ -34,7 +34,7 @@ pub async fn list_kv(epoch: u64, table_id: u32) -> anyhow::Result<()> {
                 None,
                 ReadOptions {
                     ignore_range_tombstone: false,
-                    prefix_hint: None,
+                    dist_key_hint: None,
                     table_id: TableId { table_id },
                     retention_seconds: None,
                     check_bloom_filter: false,

src/frontend/src/catalog/table_catalog.rs

@@ -411,7 +411,7 @@ mod tests {
                 vnode_col_idx: None,
                 value_indices: vec![0],
                 definition: "".into(),
-                handle_pk_conflict: false
+                handle_pk_conflict: false,
             }
         );
         assert_eq!(table, TableCatalog::from(table.to_prost(0, 0)));

src/frontend/src/optimizer/plan_node/stream_materialize.rs

@@ -151,14 +151,16 @@ impl StreamMaterialize {
 
         let ctx = input.ctx();
         let properties = ctx.inner().with_options.internal_table_subset();
+        let distribution_key = base.dist.dist_column_indices().to_vec();
+
         let table = TableCatalog {
             id: TableId::placeholder(),
             associated_source_id: None,
             name: mv_name,
             columns,
             pk: pk_list,
             stream_key: pk_indices.clone(),
-            distribution_key: base.dist.dist_column_indices().to_vec(),
+            distribution_key,
             is_index,
             appendonly: input.append_only(),
             owner: risingwave_common::catalog::DEFAULT_SUPER_USER_ID,

src/storage/Cargo.toml

@@ -80,6 +80,7 @@ tokio-stream = "0.1"
 tonic = { version = "0.2", package = "madsim-tonic" }
 tracing = "0.1"
 twox-hash = "1"
+xxhash-rust = { version = "0.8.5", features = ["xxh32"] }
 zstd = "0.11.2"
 
 [target.'cfg(not(madsim))'.dependencies]