[Feat] add scan ycsb & [Fix] primary key prefix in key partition (#47)

* [Feat] add scan ycsb

* [Fix] primary key prefix in key partition

* [Fix] cargo fmt

Author: IHEII

src/location/ob_part_desc.rs

@@ -643,9 +643,19 @@ impl ObKeyPartDesc {
             is_min_max &= v.is_max();
         }
 
+        // For partition key, the first part of the row key is the partition key
+        let part_ref_column_size = self
+            .ob_part_desc_obj
+            .ordered_part_ref_column_row_key_relations
+            .len();
         // Note: Java / ODP may not query all the partitions, and will return an error
         // instead
-        if is_min_max || !self.is_equal_keys(start, end) {
+        if is_min_max
+            || !self.is_equal_keys(
+                &start[0..part_ref_column_size],
+                &end[0..part_ref_column_size],
+            )
+        {
             let mut part_ids: Vec<i64> = Vec::with_capacity(self.part_num as usize);
             for i in 0..self.part_num as i64 {
                 part_ids.push(i);

ycsb-rs/src/db.rs

@@ -12,6 +12,13 @@ pub trait DB {
     fn insert(&self, table: &str, key: &str, values: &HashMap<&str, String>) -> Result<()>;
     fn read(&self, table: &str, key: &str, result: &mut HashMap<String, String>) -> Result<()>;
     fn update(&self, table: &str, key: &str, values: &HashMap<&str, String>) -> Result<()>;
+    fn scan(
+        &self,
+        table: &str,
+        startkey: &str,
+        endkey: &str,
+        values: &mut HashMap<String, String>,
+    ) -> Result<()>;
 }
 
 pub fn create_db(db: &str, config: Arc<OBKVClientInitStruct>) -> Result<Rc<dyn DB>> {

ycsb-rs/src/obkv_client.rs

@@ -20,7 +20,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
 use anyhow::Result;
 #[allow(unused)]
 use obkv::error::CommonErrCode;
-use obkv::{Builder, ClientConfig, ObTableClient, RunningMode, Table, Value};
+use obkv::{Builder, ClientConfig, ObTableClient, RunningMode, Table, TableQuery, Value};
 
 use crate::{db::DB, properties::Properties};
 
@@ -180,4 +180,28 @@ impl DB for OBKVClient {
 
         Ok(())
     }
+
+    #[allow(unused)]
+    fn scan(
+        &self,
+        table: &str,
+        startkey: &str,
+        endkey: &str,
+        result: &mut HashMap<String, String>,
+    ) -> Result<()> {
+        let result = self
+            .client
+            .query(table)
+            .select(COLUMN_NAMES.iter().map(|s| s.to_string()).collect())
+            .primary_index()
+            .add_scan_range(
+                vec![Value::from(startkey)],
+                true,
+                vec![Value::from(endkey)],
+                true,
+            )
+            .execute();
+        assert!(result.is_ok());
+        Ok(())
+    }
 }

ycsb-rs/src/sqlite.rs

@@ -76,4 +76,15 @@ impl DB for SQLite {
     fn update(&self, table: &str, key: &str, values: &HashMap<&str, String>) -> Result<()> {
         todo!()
     }
+
+    #[allow(unused_variables)]
+    fn scan(
+        &self,
+        table: &str,
+        startkey: &str,
+        endkey: &str,
+        values: &mut HashMap<String, String>,
+    ) -> Result<()> {
+        todo!()
+    }
 }

ycsb-rs/src/workload/core_workload.rs

@@ -177,6 +177,14 @@ impl CoreWorkload {
         db.update(&self.table, &dbkey, &values).unwrap();
     }
 
+    fn ob_transaction_scan(&self, db: Arc<OBKVClient>) {
+        let start = self.next_key_num();
+        let dbstart = format!("{}", fnvhash64(start));
+        let dbend = format!("{}", fnvhash64(start));
+        let mut result = HashMap::new();
+        db.scan(&self.table, &dbstart, &dbend, &mut result).unwrap();
+    }
+
     fn next_key_num(&self) -> u64 {
         // FIXME: Handle case where keychooser is an ExponentialGenerator.
         // FIXME: Handle case where keynum is > transactioninsertkeysequence's last
@@ -267,6 +275,9 @@ impl Workload for CoreWorkload {
             CoreOperation::Update => {
                 self.ob_transaction_update(db);
             }
+            CoreOperation::Scan => {
+                self.ob_transaction_scan(db);
+            }
             _ => todo!(),
         }
     }