Merge pull request #758 from wprzytula/stateful-partitioners

Stateful partitioners
scylladb · Aug 22, 2023 · 83f4050 · 83f4050
2 parents ffd374b + 0e21804
commit 83f4050
Show file tree

Hide file tree

Showing 10 changed files with 660 additions and 426 deletions.
diff --git a/examples/compare-tokens.rs b/examples/compare-tokens.rs
@@ -1,7 +1,6 @@
 use anyhow::Result;
 use scylla::frame::value::ValueList;
 use scylla::routing::Token;
-use scylla::transport::partitioner::{Murmur3Partitioner, Partitioner};
 use scylla::transport::NodeAddr;
 use scylla::{Session, SessionBuilder};
 use std::env;
@@ -31,7 +30,7 @@ async fn main() -> Result<()> {
             .await?;
 
         let serialized_pk = (pk,).serialized()?.into_owned();
-        let t = Murmur3Partitioner::hash(&prepared.compute_partition_key(&serialized_pk)?).value;
+        let t = prepared.calculate_token(&serialized_pk)?.unwrap().value;
 
         println!(
             "Token endpoints for query: {:?}",

diff --git a/scylla-cql/src/frame/response/result.rs b/scylla-cql/src/frame/response/result.rs
@@ -34,7 +34,7 @@ pub struct SchemaChange {
     pub event: SchemaChangeEvent,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TableSpec {
     pub ks_name: String,
     pub table_name: String,
@@ -329,7 +329,7 @@ impl CqlValue {
     // TODO
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ColumnSpec {
     pub table_spec: TableSpec,
     pub name: String,

diff --git a/scylla/benches/benchmark.rs b/scylla/benches/benchmark.rs
@@ -1,7 +1,11 @@
 use criterion::{criterion_group, criterion_main, Criterion};
 
 use bytes::BytesMut;
-use scylla::frame::types;
+use scylla::{
+    frame::types,
+    frame::value::ValueList,
+    transport::partitioner::{calculate_token_for_partition_key, Murmur3Partitioner},
+};
 
 fn types_benchmark(c: &mut Criterion) {
     let mut buf = BytesMut::with_capacity(64);
@@ -35,5 +39,57 @@ fn types_benchmark(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, types_benchmark);
+fn calculate_token_bench(c: &mut Criterion) {
+    let simple_pk = ("I'm prepared!!!",);
+    let serialized_simple_pk = simple_pk.serialized().unwrap().into_owned();
+    let simple_pk_long_column = (
+        17_i32,
+        16_i32,
+        String::from_iter(std::iter::repeat('.').take(2000)),
+    );
+    let serialized_simple_pk_long_column = simple_pk_long_column.serialized().unwrap().into_owned();
+
+    let complex_pk = (17_i32, 16_i32, "I'm prepared!!!");
+    let serialized_complex_pk = complex_pk.serialized().unwrap().into_owned();
+    let complex_pk_long_column = (
+        17_i32,
+        16_i32,
+        String::from_iter(std::iter::repeat('.').take(2000)),
+    );
+    let serialized_values_long_column = complex_pk_long_column.serialized().unwrap().into_owned();
+
+    c.bench_function("calculate_token_from_partition_key simple pk", |b| {
+        b.iter(|| calculate_token_for_partition_key(&serialized_simple_pk, &Murmur3Partitioner))
+    });
+
+    c.bench_function(
+        "calculate_token_from_partition_key simple pk long column",
+        |b| {
+            b.iter(|| {
+                calculate_token_for_partition_key(
+                    &serialized_simple_pk_long_column,
+                    &Murmur3Partitioner,
+                )
+            })
+        },
+    );
+
+    c.bench_function("calculate_token_from_partition_key complex pk", |b| {
+        b.iter(|| calculate_token_for_partition_key(&serialized_complex_pk, &Murmur3Partitioner))
+    });
+
+    c.bench_function(
+        "calculate_token_from_partition_key complex pk long column",
+        |b| {
+            b.iter(|| {
+                calculate_token_for_partition_key(
+                    &serialized_values_long_column,
+                    &Murmur3Partitioner,
+                )
+            })
+        },
+    );
+}
+
+criterion_group!(benches, types_benchmark, calculate_token_bench);
 criterion_main!(benches);