Repository: scylladb/scylla-rust-driver
Pull request #758 (2023-07-07 15:58:36)

Stateful partitioners

## Motivation
Currently, partitioners require the whole input to be passed in a single contiguous slice, which often forces an allocation that could else be avoided.

## What's done
1. **Partitioners' API is refactored** so it resembles `std::hash` API:
    - `Partitioner` corresponds to `BuildHasher` as it becomes kind of a factory,
    - `PartitionerHasher` is introduced and corresponds to `Hasher`. It contains the hashing state, which enables feeding it with multiple chunks of data separately.
A **test** is added that **ensures consistent hashing result** no matter how the partition key is partitioned into chunks.
2. After the partitioners gain ability to partition by chunks, the codebase is modified so that **all needless allocations of partition key are avoided**, i.e., allocations that only served the purpose of calculating token.
To this end, `PartitionKey` struct is introduced, whose motivation is as follows: the algorithm for computing partition key for a given prepared statement and bound values can be divided into steps. The first of them involves extracting values that consistute the partition key and putting them in proper partition key order. `PartitionKey` performs this step on construction, and serves as an entry point for further steps, with token calculation among others. For those further steps to be done, an iterator over partition key values is accessible, which yields pairs _(value, spec)_ for each column that constitutes partition key, in partition key order.
3. As **there was one more place** where materialised partition key was used - by decoding it into the produced trace - an adaptor is introduced that operates on `PartitionKey`'s iterator provides values to be deserialised by the mechanism added recently in #766. **Allocation is then saved there**, too.
4. **An API cleanup is done** regarding partition key and token computation. `calculate_token_from_partition_key()` is moved to `partitioner.rs`, `calculate_token()` to `prepared_statement.rs`, and a number of helper functions are deleted as no longer needed.

## Pre-review checklist
- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [x] I added relevant tests for new features and bug fixes.
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- [x] I have provided docstrings for the public items that I want to introduce.
- ~~[ ] I have adjusted the documentation in `./docs/source/`.~~
- ~~[ ] I added appropriate `Fixes:` annotations to PR description.~~


mykaul: Any measurable performance impact? 
wprzytula: > Any measurable performance impact?

Number of allocations has dropped (6 -> 5 per insert, 19 -> 18 per select).

**Before:**
```
Args { mode: All, node: "127.0.0.1:9042", parallelism: 256, requests: 100000 }
Connecting to 127.0.0.1:9042 ...
Sending 100000 inserts, hold tight ..........
----------
Inserts:
----------
allocs/req:                 6.08
reallocs/req:               6.00
frees/req:                  6.07
bytes allocated/req:      272.75
bytes reallocated/req:     73.06
bytes freed/req:          267.77
----------
Sending 100000 selects, hold tight ..........
----------
Selects:
----------
allocs/req:                19.00
reallocs/req:               6.00
frees/req:                 19.00
bytes allocated/req:     1123.09
bytes reallocated/req:     73.00
bytes freed/req:         1123.03
----------
```

**After:**
```
Args { mode: All, node: "127.0.0.1:9042", parallelism: 256, requests: 100000 }
Connecting to 127.0.0.1:9042 ...
Sending 100000 inserts, hold tight ..........
----------
Inserts:
----------
allocs/req:                 5.08
reallocs/req:               5.00
frees/req:                  5.07
bytes allocated/req:      265.09
bytes reallocated/req:     77.06
bytes freed/req:          260.06
----------
Sending 100000 selects, hold tight ..........
----------
Selects:
----------
allocs/req:                18.00
reallocs/req:               5.00
frees/req:                 18.00
bytes allocated/req:     1115.19
bytes reallocated/req:     77.00
bytes freed/req:         1115.05
----------

```
piodul: > > Any measurable performance impact?
> 
> Number of allocations has dropped (6 -> 5 per insert, 19 -> 18 per select).

Please prepare a microbenchmark for this. While an allocation is now avoided, the new implementation has more branches, so it will be interesting to see whether the tradeoff was worth it.

We are basically interested in the `PreparedStatement::compute_partition_key` function. Interesting cases are single-column pks with different lengths and multi-column pks with different lengths.

We already have a small amount of benchmarks in `scylla` and `scylla-cql`, you can look at them for reference (and at the external `criterion` crate that those benchmarks use).

mykaul: > > Please prepare a microbenchmark for this. While an allocation is now avoided, the new implementation has more branches, so it will be interesting to see whether the tradeoff was worth it.
> 
> After benchmarks and some changes to the algorithm, the new way seems to outperform the old way by 15% to 40% in various cases. I believe this suffices to merge this.

NICE! Well done! Would be happy to see details of this tests.
wprzytula: CI fails due to use of `Option::unzip()`, which is stable since Rust 1.66, whereas our MSRV is 1.65. Can we bump it @piodul?
Apart from that, ready for review.
wprzytula: I reimplemented `Option::unzip()` (technically, copied its code from the standard library) temporarily, not to force MSRV bump yet. 
wprzytula: Rebased on `main`.

Patch:
diff --git a/examples/compare-tokens.rs b/examples/compare-tokens.rs
index 793aae2e..5c56aa5f 100644
--- 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
index 80e4aa9b..288baf91 100644
--- 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
index 08c13a3f..c6058923 100644
--- 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);
diff --git a/scylla/src/statement/prepared_statement.rs b/scylla/src/statement/prepared_statement.rs
index 3a1d573b..b57d5d4b 100644
--- a/scylla/src/statement/prepared_statement.rs
+++ b/scylla/src/statement/prepared_statement.rs
@@ -1,5 +1,5 @@
-use byteorder::{BigEndian, ReadBytesExt};
-use bytes::{BufMut, Bytes, BytesMut};
+use bytes::{Bytes, BytesMut};
+use scylla_cql::errors::{BadQuery, QueryError};
 use smallvec::{smallvec, SmallVec};
 use std::convert::TryInto;
 use std::sync::Arc;
@@ -7,8 +7,7 @@ use std::time::Duration;
 use thiserror::Error;
 use uuid::Uuid;
 
-use scylla_cql::frame::frame_errors::ParseError;
-use scylla_cql::frame::response::result::{deser_cql_value, CqlValue};
+use scylla_cql::frame::response::result::ColumnSpec;
 
 use super::StatementConfig;
 use crate::frame::response::result::PreparedMetadata;
@@ -16,8 +15,9 @@ use crate::frame::types::{Consistency, SerialConsistency};
 use crate::frame::value::SerializedValues;
 use crate::history::HistoryListener;
 use crate::retry_policy::RetryPolicy;
+use crate::routing::Token;
 use crate::transport::execution_profile::ExecutionProfileHandle;
-use crate::transport::partitioner::PartitionerName;
+use crate::transport::partitioner::{Partitioner, PartitionerHasher, PartitionerName};
 
 /// Represents a statement prepared on the server.
 #[derive(Debug)]
@@ -124,68 +124,76 @@ impl PreparedStatement {
 
     /// Computes the partition key of the target table from given values —
     /// it assumes that all partition key columns are passed in values.
-    /// Partition keys have a specific serialization rules.
+    /// Partition keys have specific serialization rules.
     /// Ref: <https://github.com/scylladb/scylla/blob/40adf38915b6d8f5314c621a94d694d172360833/compound_compat.hh#L33-L47>
+    ///
+    /// Note: this is no longer required to compute a token. This is because partitioners
+    /// are now stateful and stream-based, so they do not require materialised partition key.
+    /// Therefore, for computing a token, there are more efficient methods, such as
+    /// [Self::calculate_token()].
     pub fn compute_partition_key(
         &self,
         bound_values: &SerializedValues,
     ) -> Result<Bytes, PartitionKeyError> {
+        let partition_key = self.extract_partition_key(bound_values)?;
         let mut buf = BytesMut::new();
+        let mut writer = |chunk: &[u8]| buf.extend_from_slice(chunk);
 
-        // Single-value partition key case
-        if self.get_prepared_metadata().pk_indexes.len() == 1 {
-            let pk_index = self.get_prepared_metadata().pk_indexes[0].index;
-            if let Some(v) = bound_values
-                .iter()
-                .nth(pk_index as usize)
-                .ok_or_else(|| PartitionKeyError::NoPkIndexValue(pk_index, bound_values.len()))?
-            {
-                buf.extend_from_slice(v);
-            }
-            return Ok(buf.into());
-        }
+        partition_key.write_encoded_partition_key(&mut writer)?;
 
-        // Composite partition key case
+        Ok(buf.freeze())
+    }
 
-        // Iterate on values using sorted pk_indexes (see deser_prepared_metadata),
-        // and use PartitionKeyIndex.sequence to insert the value in pk_values with the correct order.
-        // At the same time, compute the size of the buffer to reserve it before writing in it.
-        let mut pk_values: SmallVec<[_; 8]> =
-            smallvec![None; self.get_prepared_metadata().pk_indexes.len()];
-        let mut values_iter = bound_values.iter();
-        let mut buf_size = 0;
-        // pk_indexes contains the indexes of the partition key value, so the current offset of the
-        // iterator must be kept, in order to compute the next position of the pk in the iterator.
-        // At each iteration values_iter.nth(0) will roughly correspond to values[values_iter_offset],
-        // so values[pk_index.index] will be retrieved with values_iter.nth(pk_index.index - values_iter_offset)
-        let mut values_iter_offset = 0;
-        for pk_index in self.get_prepared_metadata().pk_indexes.iter().copied() {
-            // Find value matching current pk_index
-            let next_val = values_iter
-                .nth((pk_index.index - values_iter_offset) as usize)
-                .ok_or_else(|| {
-                    PartitionKeyError::NoPkIndexValue(pk_index.index, bound_values.len())
-                })?;
-            // Add it in sequence order to pk_values
-            if let Some(v) = next_val {
-                pk_values[pk_index.sequence as usize] = Some(v);
-                buf_size += std::mem::size_of::<u16>() + v.len() + std::mem::size_of::<u8>();
-            }
-            values_iter_offset = pk_index.index + 1;
-        }
-        // Add values' bytes
-        buf.reserve(buf_size);
-        for v in pk_values.iter().flatten() {
-            let v_len_u16: u16 = v
-                .len()
-                .try_into()
-                .map_err(|_| PartitionKeyError::ValueTooLong(v.len()))?;
-
-            buf.put_u16(v_len_u16);
-            buf.extend_from_slice(v);
-            buf.put_u8(0);
+    /// Determines which values consistute the partition key and puts them in order.
+    ///
+    /// This is a preparation step necessary for calculating token based on a prepared statement.
+    pub(crate) fn extract_partition_key<'ps>(
+        &'ps self,
+        bound_values: &'ps SerializedValues,
+    ) -> Result<PartitionKey, PartitionKeyExtractionError> {
+        PartitionKey::new(self.get_prepared_metadata(), bound_values)
+    }
+
+    pub(crate) fn extract_partition_key_and_calculate_token<'ps>(
+        &'ps self,
+        partitioner_name: &'ps PartitionerName,
+        serialized_values: &'ps SerializedValues,
+    ) -> Result<Option<(PartitionKey<'ps>, Token)>, QueryError> {
+        if !self.is_token_aware() {
+            return Ok(None);
         }
-        Ok(buf.into())
+
+        let partition_key =
+            self.extract_partition_key(serialized_values)
+                .map_err(|err| match err {
+                    PartitionKeyExtractionError::NoPkIndexValue(_, _) => {
+                        QueryError::ProtocolError("No pk indexes - can't calculate token")
+                    }
+                })?;
+        let token = partition_key
+            .calculate_token(partitioner_name)
+            .map_err(|err| match err {
+                TokenCalculationError::ValueTooLong(values_len) => {
+                    QueryError::BadQuery(BadQuery::ValuesTooLongForKey(values_len, u16::MAX.into()))
+                }
+            })?;
+
+        Ok(Some((partition_key, token)))
+    }
+
+    /// Calculates the token for given prepared statement and serialized values.
+    ///
+    /// Returns the token that would be computed for executing the provided
+    /// prepared statement with the provided values.
+    // As this function creates a `PartitionKey`, it is intended rather for external usage (by users).
+    // For internal purposes, `PartitionKey::calculate_token()` is preferred, as `PartitionKey`
+    // is either way used internally, among others for display in traces.
+    pub fn calculate_token(
+        &self,
+        serialized_values: &SerializedValues,
+    ) -> Result<Option<Token>, QueryError> {
+        self.extract_partition_key_and_calculate_token(&self.partitioner_name, serialized_values)
+            .map(|opt| opt.map(|(_pk, token)| token))
     }
 
     /// Returns the name of the keyspace this statement is operating on.
@@ -329,100 +337,132 @@ impl PreparedStatement {
 }
 
 #[derive(Clone, Debug, Error, PartialEq, Eq, PartialOrd, Ord)]
-pub enum PartitionKeyError {
+pub enum PartitionKeyExtractionError {
     #[error("No value with given pk_index! pk_index: {0}, values.len(): {1}")]
     NoPkIndexValue(u16, i16),
+}
+
+#[derive(Clone, Debug, Error, PartialEq, Eq, PartialOrd, Ord)]
+pub enum TokenCalculationError {
     #[error("Value bytes too long to create partition key, max 65 535 allowed! value.len(): {0}")]
     ValueTooLong(usize),
 }
 
-// The PartitionKeyDecoder reverses the process of PreparedStatement::compute_partition_key:
-// it returns the consecutive values of partition key column that were encoded
-// by the function into the Bytes object, additionally decoding them to CqlValue.
-//
-// The format follows the description here:
-// <https://github.com/scylladb/scylla/blob/40adf38915b6d8f5314c621a94d694d172360833/compound_compat.hh#L33-L47>
-//
-// TODO: Currently, if there is a null value specified for a partition column,
-// it will be skipped when creating the serialized partition key. We should
-// not create such partition keys in the future, i.e. fail the request or
-// route it to a random replica instead and let the DB reject it. In the
-// meantime, this struct will return some garbage data while it tries to
-// decode the key, but nothing bad like a panic should happen otherwise.
-// The struct is currently only used for printing the partition key, so that's
-// completely fine.
-#[derive(Clone, Copy)]
-pub(crate) struct PartitionKeyDecoder<'pk> {
-    prepared_metadata: &'pk PreparedMetadata,
-    partition_key: &'pk [u8],
-    value_index: usize,
+#[derive(Clone, Debug, Error, PartialEq, Eq, PartialOrd, Ord)]
+pub enum PartitionKeyError {
+    #[error(transparent)]
+    PartitionKeyExtraction(PartitionKeyExtractionError),
+    #[error(transparent)]
+    TokenCalculation(TokenCalculationError),
 }
 
-impl<'pk> PartitionKeyDecoder<'pk> {
-    pub(crate) fn new(prepared_metadata: &'pk PreparedMetadata, partition_key: &'pk [u8]) -> Self {
-        Self {
-            prepared_metadata,
-            partition_key,
-            value_index: 0,
-        }
+impl From<PartitionKeyExtractionError> for PartitionKeyError {
+    fn from(err: PartitionKeyExtractionError) -> Self {
+        Self::PartitionKeyExtraction(err)
+    }
+}
+
+impl From<TokenCalculationError> for PartitionKeyError {
+    fn from(err: TokenCalculationError) -> Self {
+        Self::TokenCalculation(err)
     }
 }
 
-impl<'pk> Iterator for PartitionKeyDecoder<'pk> {
-    type Item = Result<CqlValue, ParseError>;
+pub(crate) type PartitionKeyValue<'ps> = (&'ps [u8], &'ps ColumnSpec);
+
+pub(crate) struct PartitionKey<'ps> {
+    pk_values: SmallVec<[Option<PartitionKeyValue<'ps>>; PartitionKey::SMALLVEC_ON_STACK_SIZE]>,
+}
+
+impl<'ps> PartitionKey<'ps> {
+    const SMALLVEC_ON_STACK_SIZE: usize = 8;
 
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.value_index >= self.prepared_metadata.pk_indexes.len() {
-            return None;
+    fn new(
+        prepared_metadata: &'ps PreparedMetadata,
+        bound_values: &'ps SerializedValues,
+    ) -> Result<Self, PartitionKeyExtractionError> {
+        // Iterate on values using sorted pk_indexes (see deser_prepared_metadata),
+        // and use PartitionKeyIndex.sequence to insert the value in pk_values with the correct order.
+        let mut pk_values: SmallVec<[_; PartitionKey::SMALLVEC_ON_STACK_SIZE]> =
+            smallvec![None; prepared_metadata.pk_indexes.len()];
+        let mut values_iter = bound_values.iter();
+        // pk_indexes contains the indexes of the partition key value, so the current offset of the
+        // iterator must be kept, in order to compute the next position of the pk in the iterator.
+        // At each iteration values_iter.nth(0) will roughly correspond to values[values_iter_offset],
+        // so values[pk_index.index] will be retrieved with values_iter.nth(pk_index.index - values_iter_offset)
+        let mut values_iter_offset = 0;
+        for pk_index in prepared_metadata.pk_indexes.iter().copied() {
+            // Find value matching current pk_index
+            let next_val = values_iter
+                .nth((pk_index.index - values_iter_offset) as usize)
+                .ok_or_else(|| {
+                    PartitionKeyExtractionError::NoPkIndexValue(pk_index.index, bound_values.len())
+                })?;
+            // Add it in sequence order to pk_values
+            if let Some(v) = next_val {
+                let spec = &prepared_metadata.col_specs[pk_index.index as usize];
+                pk_values[pk_index.sequence as usize] = Some((v, spec));
+            }
+            values_iter_offset = pk_index.index + 1;
         }
+        Ok(Self { pk_values })
+    }
 
-        // It should be safe to unwrap because PartitionKeyIndex::sequence
-        // fields of the pk_indexes form a permutation, but let's err
-        // on the safe side in case of bugs.
-        let pk_index = self
-            .prepared_metadata
-            .pk_indexes
-            .iter()
-            .find(|pki| pki.sequence == self.value_index as u16)?;
-
-        let col_idx = pk_index.index as usize;
-        let spec = &self.prepared_metadata.col_specs[col_idx];
-
-        let cell = if self.prepared_metadata.pk_indexes.len() == 1 {
-            Ok(self.partition_key)
-        } else {
-            self.parse_cell()
-        };
-        self.value_index += 1;
-        Some(cell.and_then(|mut cell| deser_cql_value(&spec.typ, &mut cell)))
+    pub(crate) fn iter(&self) -> impl Iterator<Item = PartitionKeyValue<'ps>> + Clone + '_ {
+        self.pk_values.iter().flatten().copied()
     }
-}
 
-impl<'pk> PartitionKeyDecoder<'pk> {
-    fn parse_cell(&mut self) -> Result<&'pk [u8], ParseError> {
-        let buf = &mut self.partition_key;
-        let len = buf.read_u16::<BigEndian>()? as usize;
-        if buf.len() < len {
-            return Err(ParseError::IoError(std::io::Error::new(
-                std::io::ErrorKind::UnexpectedEof,
-                "value too short",
-            )));
+    fn write_encoded_partition_key(
+        &self,
+        writer: &mut impl FnMut(&[u8]),
+    ) -> Result<(), TokenCalculationError> {
+        let mut pk_val_iter = self.iter().map(|(val, _spec)| val);
+        if let Some(first_value) = pk_val_iter.next() {
+            if let Some(second_value) = pk_val_iter.next() {
+                // Composite partition key case
+                for value in std::iter::once(first_value)
+                    .chain(std::iter::once(second_value))
+                    .chain(pk_val_iter)
+                {
+                    let v_len_u16: u16 = value
+                        .len()
+                        .try_into()
+                        .map_err(|_| TokenCalculationError::ValueTooLong(value.len()))?;
+                    writer(&v_len_u16.to_be_bytes());
+                    writer(value);
+                    writer(&[0u8]);
+                }
+            } else {
+                // Single-value partition key case
+                writer(first_value);
+            }
         }
-        let col = &buf[..len];
-        *buf = &buf[len..];
-        buf.read_u8()?;
-        Ok(col)
+        Ok(())
+    }
+
+    pub(crate) fn calculate_token(
+        &self,
+        partitioner_name: &PartitionerName,
+    ) -> Result<Token, TokenCalculationError> {
+        let mut partitioner_hasher = partitioner_name.build_hasher();
+        let mut writer = |chunk: &[u8]| partitioner_hasher.write(chunk);
+
+        self.write_encoded_partition_key(&mut writer)?;
+
+        Ok(partitioner_hasher.finish())
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use bytes::{BufMut, Bytes, BytesMut};
-    use scylla_cql::frame::response::result::{
-        ColumnSpec, ColumnType, CqlValue, PartitionKeyIndex, PreparedMetadata, TableSpec,
+    use scylla_cql::frame::{
+        response::result::{
+            ColumnSpec, ColumnType, PartitionKeyIndex, PreparedMetadata, TableSpec,
+        },
+        value::SerializedValues,
     };
 
-    use super::PartitionKeyDecoder;
+    use crate::prepared_statement::PartitionKey;
 
     fn make_meta(
         cols: impl IntoIterator<Item = ColumnType>,
@@ -458,53 +498,8 @@ mod tests {
         }
     }
 
-    fn make_key<'pk>(cols: impl IntoIterator<Item = &'pk [u8]>) -> Bytes {
-        let cols: Vec<_> = cols.into_iter().collect();
-        // TODO: Use compute_partition_key or one of the variants
-        // after they are moved to a more sensible place
-        // instead of constructing the PK manually
-        let mut b = BytesMut::new();
-        if cols.len() == 1 {
-            b.extend_from_slice(cols[0]);
-        } else {
-            for c in cols {
-                b.put_i16(c.len() as i16);
-                b.extend_from_slice(c);
-                b.put_u8(0);
-            }
-        }
-        b.freeze()
-    }
-
-    #[test]
-    fn test_pk_decoder_single_column() {
-        let meta = make_meta([ColumnType::Int], [0]);
-        let pk = make_key([0i32.to_be_bytes().as_ref()]);
-        let cells = PartitionKeyDecoder::new(&meta, &pk)
-            .collect::<Result<Vec<_>, _>>()
-            .unwrap();
-        assert_eq!(cells, vec![CqlValue::Int(0)]);
-    }
-
-    #[test]
-    fn test_pk_decoder_multiple_columns() {
-        let meta = make_meta(std::iter::repeat(ColumnType::Int).take(3), [0, 1, 2]);
-        let pk = make_key([
-            12i32.to_be_bytes().as_ref(),
-            34i32.to_be_bytes().as_ref(),
-            56i32.to_be_bytes().as_ref(),
-        ]);
-        let cells = PartitionKeyDecoder::new(&meta, &pk)
-            .collect::<Result<Vec<_>, _>>()
-            .unwrap();
-        assert_eq!(
-            cells,
-            vec![CqlValue::Int(12), CqlValue::Int(34), CqlValue::Int(56)],
-        );
-    }
-
     #[test]
-    fn test_pk_decoder_multiple_columns_shuffled() {
+    fn test_partition_key_multiple_columns_shuffled() {
         let meta = make_meta(
             [
                 ColumnType::TinyInt,
@@ -515,21 +510,22 @@ mod tests {
             ],
             [4, 0, 3],
         );
-        let pk = make_key([
-            &[1, 2, 3, 4, 5],
-            67i8.to_be_bytes().as_ref(),
-            89i64.to_be_bytes().as_ref(),
-        ]);
-        let cells = PartitionKeyDecoder::new(&meta, &pk)
-            .collect::<Result<Vec<_>, _>>()
-            .unwrap();
+        let mut values = SerializedValues::new();
+        values.add_value(&67i8).unwrap();
+        values.add_value(&42i16).unwrap();
+        values.add_value(&23i32).unwrap();
+        values.add_value(&89i64).unwrap();
+        values.add_value(&[1u8, 2, 3, 4, 5]).unwrap();
+
+        let pk = PartitionKey::new(&meta, &values).unwrap();
+        let pk_cols = Vec::from_iter(pk.iter());
         assert_eq!(
-            cells,
+            pk_cols,
             vec![
-                CqlValue::Blob(vec![1, 2, 3, 4, 5]),
-                CqlValue::TinyInt(67),
-                CqlValue::BigInt(89),
-            ],
+                ([1u8, 2, 3, 4, 5].as_slice(), &meta.col_specs[4]),
+                (67i8.to_be_bytes().as_ref(), &meta.col_specs[0]),
+                (89i64.to_be_bytes().as_ref(), &meta.col_specs[3]),
+            ]
         );
     }
 }
diff --git a/scylla/src/transport/cluster.rs b/scylla/src/transport/cluster.rs
index 0e6e9047..eed76acc 100644
--- a/scylla/src/transport/cluster.rs
+++ b/scylla/src/transport/cluster.rs
@@ -1,6 +1,7 @@
 /// Cluster manages up to date information and connections to database nodes
 use crate::frame::response::event::{Event, StatusChangeEvent};
 use crate::frame::value::ValueList;
+use crate::prepared_statement::TokenCalculationError;
 use crate::routing::Token;
 use crate::transport::host_filter::HostFilter;
 use crate::transport::{
@@ -13,7 +14,6 @@ use crate::transport::{
 };
 
 use arc_swap::ArcSwap;
-use bytes::{BufMut, Bytes, BytesMut};
 use futures::future::join_all;
 use futures::{future::RemoteHandle, FutureExt};
 use itertools::Itertools;
@@ -36,6 +36,7 @@ pub struct ContactPoint {
 }
 
 use super::locator::ReplicaLocator;
+use super::partitioner::calculate_token_for_partition_key;
 use super::topology::Strategy;
 
 /// Cluster manages up to date information and connections to database nodes.
@@ -409,31 +410,13 @@ impl ClusterData {
             .and_then(|t| t.partitioner.as_deref())
             .and_then(PartitionerName::from_str)
             .unwrap_or_default();
-        let serialized_values = partition_key.serialized()?;
-        // Null values are skipped in computation; null values in partition key are unsound,
-        // but it is consistent with computation of prepared statements token.
-        let serialized_pk = match serialized_values.len() {
-            0 => None,
-            1 => serialized_values
-                .iter()
-                .next()
-                .unwrap()
-                .map(Bytes::copy_from_slice),
-            _ => {
-                let mut buf = BytesMut::new();
-                for value in serialized_values.iter().flatten() {
-                    let value_size = value
-                        .len()
-                        .try_into()
-                        .map_err(|_| BadQuery::ValuesTooLongForKey(value.len(), u16::MAX.into()))?;
-                    buf.put_u16(value_size);
-                    buf.extend_from_slice(value);
-                    buf.put_u8(0);
+
+        calculate_token_for_partition_key(&partition_key.serialized().unwrap(), &partitioner)
+            .map_err(|err| match err {
+                TokenCalculationError::ValueTooLong(values_len) => {
+                    BadQuery::ValuesTooLongForKey(values_len, u16::MAX.into())
                 }
-                Some(buf.into())
-            }
-        };
-        Ok(partitioner.hash(&serialized_pk.unwrap_or_default()))
+            })
     }
 
     /// Access to replicas owning a given token
diff --git a/scylla/src/transport/iterator.rs b/scylla/src/transport/iterator.rs
index 0cc95569..8c2e554a 100644
--- a/scylla/src/transport/iterator.rs
+++ b/scylla/src/transport/iterator.rs
@@ -31,7 +31,6 @@ use crate::frame::{
     value::SerializedValues,
 };
 use crate::history::{self, HistoryListener};
-use crate::routing::Token;
 use crate::statement::Consistency;
 use crate::statement::{prepared_statement::PreparedStatement, query::Query};
 use crate::transport::cluster::ClusterData;
@@ -40,6 +39,7 @@ use crate::transport::load_balancing::{self, RoutingInfo};
 use crate::transport::metrics::Metrics;
 use crate::transport::retry_policy::{QueryInfo, RetryDecision, RetrySession};
 use crate::transport::{Node, NodeRef};
+use crate::utils::unzip_option;
 use tracing::{trace, trace_span, warn, Instrument};
 use uuid::Uuid;
 
@@ -76,8 +76,6 @@ struct ReceivedPage {
 pub(crate) struct PreparedIteratorConfig {
     pub(crate) prepared: PreparedStatement,
     pub(crate) values: SerializedValues,
-    pub(crate) partition_key: Option<Bytes>,
-    pub(crate) token: Option<Token>,
     pub(crate) execution_profile: Arc<ExecutionProfileInner>,
     pub(crate) cluster_data: Arc<ClusterData>,
     pub(crate) metrics: Arc<Metrics>,
@@ -241,19 +239,29 @@ impl RowIterator {
 
         let parent_span = tracing::Span::current();
         let worker_task = async move {
+            let prepared_ref = &config.prepared;
+            let values_ref = &config.values;
+
+            let (partition_key, token) = match prepared_ref
+                .extract_partition_key_and_calculate_token(
+                    prepared_ref.get_partitioner_name(),
+                    values_ref,
+                ) {
+                Ok(res) => unzip_option(res),
+                Err(err) => {
+                    let (proof, _res) = ProvingSender::from(sender).send(Err(err)).await;
+                    return proof;
+                }
+            };
+
             let statement_info = RoutingInfo {
                 consistency,
                 serial_consistency,
-                token: config.token,
+                token,
                 keyspace: config.prepared.get_keyspace_name(),
                 is_confirmed_lwt: config.prepared.is_confirmed_lwt(),
             };
 
-            let prepared_ref = &config.prepared;
-            let values_ref = &config.values;
-            let partition_key = config.partition_key;
-            let token = config.token;
-
             let choose_connection = |node: Arc<Node>| async move {
                 match token {
                     Some(token) => node.connection_for_token(token).await,
@@ -294,8 +302,7 @@ impl RowIterator {
 
             let span_creator = move || {
                 let span = RequestSpan::new_prepared(
-                    prepared_ref.get_prepared_metadata(),
-                    partition_key.as_ref(),
+                    partition_key.as_ref().map(|pk| pk.iter()),
                     token,
                     serialized_values_size,
                 );
diff --git a/scylla/src/transport/partitioner.rs b/scylla/src/transport/partitioner.rs
index 887acb2f..9c8a5423 100644
--- a/scylla/src/transport/partitioner.rs
+++ b/scylla/src/transport/partitioner.rs
@@ -1,7 +1,9 @@
 use bytes::Buf;
 use std::num::Wrapping;
 
-use crate::routing::Token;
+use crate::{
+    frame::value::SerializedValues, prepared_statement::TokenCalculationError, routing::Token,
+};
 
 #[allow(clippy::upper_case_acronyms)]
 #[derive(Clone, PartialEq, Debug, Default)]
@@ -21,87 +23,229 @@ impl PartitionerName {
             None
         }
     }
+}
+
+impl Partitioner for PartitionerName {
+    type Hasher = PartitionerHasherAny;
+
+    fn build_hasher(&self) -> Self::Hasher {
+        match self {
+            PartitionerName::Murmur3 => {
+                PartitionerHasherAny::Murmur3(Murmur3Partitioner.build_hasher())
+            }
+            PartitionerName::CDC => PartitionerHasherAny::CDC(CDCPartitioner.build_hasher()),
+        }
+    }
+}
+
+#[allow(clippy::upper_case_acronyms)]
+pub(crate) enum PartitionerHasherAny {
+    Murmur3(Murmur3PartitionerHasher),
+    CDC(CDCPartitionerHasher),
+}
+
+impl PartitionerHasher for PartitionerHasherAny {
+    fn write(&mut self, pk_part: &[u8]) {
+        match self {
+            PartitionerHasherAny::Murmur3(h) => h.write(pk_part),
+            PartitionerHasherAny::CDC(h) => h.write(pk_part),
+        }
+    }
 
-    pub(crate) fn hash(&self, pk: &[u8]) -> Token {
+    fn finish(&self) -> Token {
         match self {
-            PartitionerName::Murmur3 => Murmur3Partitioner::hash(pk),
-            PartitionerName::CDC => CDCPartitioner::hash(pk),
+            PartitionerHasherAny::Murmur3(h) => h.finish(),
+            PartitionerHasherAny::CDC(h) => h.finish(),
         }
     }
 }
 
+/// A trait for creating instances of `PartitionHasher`, which ultimately compute the token.
+///
+/// The Partitioners' design is based on std::hash design: `Partitioner`
+/// corresponds to `HasherBuilder`, and `PartitionerHasher` to `Hasher`.
+/// See their documentation for more details.
 pub trait Partitioner {
-    fn hash(pk: &[u8]) -> Token;
+    type Hasher: PartitionerHasher;
+
+    fn build_hasher(&self) -> Self::Hasher;
+
+    fn hash_one(&self, data: &[u8]) -> Token {
+        let mut hasher = self.build_hasher();
+        hasher.write(data);
+        hasher.finish()
+    }
+}
+
+/// A trait for hashing a stream of serialized CQL values.
+///
+/// Instances of this trait are created by a `Partitioner` and are stateful.
+/// At any point, one can call `finish()` and a `Token` will be computed
+/// based on values that has been fed so far.
+pub trait PartitionerHasher {
+    fn write(&mut self, pk_part: &[u8]);
+    fn finish(&self) -> Token;
 }
 
 pub struct Murmur3Partitioner;
-pub struct CDCPartitioner;
 
-impl Murmur3Partitioner {
-    // An implementation of MurmurHash3 ported from Scylla. Please note that this
-    // is not a "correct" implementation of MurmurHash3 - it replicates the same
-    // bugs made in the original Cassandra implementation in order to be compatible.
-    fn hash3_x64_128(mut data: &[u8]) -> i128 {
-        let length = data.len();
+impl Partitioner for Murmur3Partitioner {
+    type Hasher = Murmur3PartitionerHasher;
+
+    fn build_hasher(&self) -> Self::Hasher {
+        Self::Hasher {
+            total_len: 0,
+            buf: Default::default(),
+            h1: Wrapping(0),
+            h2: Wrapping(0),
+        }
+    }
+}
+
+pub struct Murmur3PartitionerHasher {
+    total_len: usize,
+    buf: [u8; Self::BUF_CAPACITY],
+    h1: Wrapping<i64>,
+    h2: Wrapping<i64>,
+}
 
-        let c1: Wrapping<i64> = Wrapping(0x87c3_7b91_1142_53d5_u64 as i64);
-        let c2: Wrapping<i64> = Wrapping(0x4cf5_ad43_2745_937f_u64 as i64);
+impl Murmur3PartitionerHasher {
+    const BUF_CAPACITY: usize = 16;
 
-        let mut h1 = Wrapping(0_i64);
-        let mut h2 = Wrapping(0_i64);
+    const C1: Wrapping<i64> = Wrapping(0x87c3_7b91_1142_53d5_u64 as i64);
+    const C2: Wrapping<i64> = Wrapping(0x4cf5_ad43_2745_937f_u64 as i64);
 
-        while data.len() >= 16 {
-            let mut k1 = Wrapping(data.get_i64_le());
-            let mut k2 = Wrapping(data.get_i64_le());
+    fn hash_16_bytes(&mut self, mut k1: Wrapping<i64>, mut k2: Wrapping<i64>) {
+        k1 *= Self::C1;
+        k1 = Self::rotl64(k1, 31);
+        k1 *= Self::C2;
+        self.h1 ^= k1;
 
-            k1 *= c1;
-            k1 = Self::rotl64(k1, 31);
-            k1 *= c2;
-            h1 ^= k1;
+        self.h1 = Self::rotl64(self.h1, 27);
+        self.h1 += self.h2;
+        self.h1 = self.h1 * Wrapping(5) + Wrapping(0x52dce729);
 
-            h1 = Self::rotl64(h1, 27);
-            h1 += h2;
-            h1 = h1 * Wrapping(5) + Wrapping(0x52dce729);
+        k2 *= Self::C2;
+        k2 = Self::rotl64(k2, 33);
+        k2 *= Self::C1;
+        self.h2 ^= k2;
 
-            k2 *= c2;
-            k2 = Self::rotl64(k2, 33);
-            k2 *= c1;
-            h2 ^= k2;
+        self.h2 = Self::rotl64(self.h2, 31);
+        self.h2 += self.h1;
+        self.h2 = self.h2 * Wrapping(5) + Wrapping(0x38495ab5);
+    }
+
+    fn fetch_16_bytes_from_buf(buf: &mut &[u8]) -> (Wrapping<i64>, Wrapping<i64>) {
+        let k1 = Wrapping(buf.get_i64_le());
+        let k2 = Wrapping(buf.get_i64_le());
+        (k1, k2)
+    }
+
+    #[inline]
+    fn rotl64(v: Wrapping<i64>, n: u32) -> Wrapping<i64> {
+        Wrapping((v.0 << n) | (v.0 as u64 >> (64 - n)) as i64)
+    }
 
-            h2 = Self::rotl64(h2, 31);
-            h2 += h1;
-            h2 = h2 * Wrapping(5) + Wrapping(0x38495ab5);
+    #[inline]
+    fn fmix(mut k: Wrapping<i64>) -> Wrapping<i64> {
+        k ^= Wrapping((k.0 as u64 >> 33) as i64);
+        k *= Wrapping(0xff51afd7ed558ccd_u64 as i64);
+        k ^= Wrapping((k.0 as u64 >> 33) as i64);
+        k *= Wrapping(0xc4ceb9fe1a85ec53_u64 as i64);
+        k ^= Wrapping((k.0 as u64 >> 33) as i64);
+
+        k
+    }
+}
+
+// The implemented Murmur3 algorithm is roughly as follows:
+// 1. while there are at least 16 bytes given:
+//      consume 16 bytes by parsing them into i64s, then
+//      include them in h1, h2, k1, k2;
+// 2. do some magic with remaining n < 16 bytes,
+//      include them in h1, h2, k1, k2;
+// 3. compute the token based on h1, h2, k1, k2.
+//
+// Therefore, the buffer of capacity 16 is used. As soon as it gets full,
+// point 1. is executed. Points 2. and 3. are exclusively done in `finish()`,
+// so they don't mutate the state.
+impl PartitionerHasher for Murmur3PartitionerHasher {
+    fn write(&mut self, mut pk_part: &[u8]) {
+        let mut buf_len = self.total_len % Self::BUF_CAPACITY;
+        self.total_len += pk_part.len();
+
+        // If the buffer is nonempty and can be filled completely, so that we can fetch two i64s from it,
+        // fill it and hash its contents, then make it empty.
+        if buf_len > 0 && Self::BUF_CAPACITY - buf_len <= pk_part.len() {
+            // First phase: populate buffer until full, then consume two i64s.
+            let to_write = Ord::min(Self::BUF_CAPACITY - buf_len, pk_part.len());
+            self.buf[buf_len..buf_len + to_write].copy_from_slice(&pk_part[..to_write]);
+            pk_part.advance(to_write);
+            buf_len += to_write;
+
+            debug_assert_eq!(buf_len, Self::BUF_CAPACITY);
+            // consume 16 bytes from internal buf
+            let mut buf_ptr = &self.buf[..];
+            let (k1, k2) = Self::fetch_16_bytes_from_buf(&mut buf_ptr);
+            debug_assert!(buf_ptr.is_empty());
+            self.hash_16_bytes(k1, k2);
+            buf_len = 0;
         }
 
+        // If there were enough data, now we have an empty buffer. Further data, if enough, can be hence
+        // hashed directly from the external buffer.
+        if buf_len == 0 {
+            // Second phase: fast path for big values.
+            while pk_part.len() >= Self::BUF_CAPACITY {
+                let (k1, k2) = Self::fetch_16_bytes_from_buf(&mut pk_part);
+                self.hash_16_bytes(k1, k2);
+            }
+        }
+
+        // Third phase: move remaining bytes to the buffer.
+        debug_assert!(pk_part.len() < Self::BUF_CAPACITY - buf_len);
+        let to_write = pk_part.len();
+        self.buf[buf_len..buf_len + to_write].copy_from_slice(&pk_part[..to_write]);
+        pk_part.advance(to_write);
+        buf_len += to_write;
+        debug_assert!(pk_part.is_empty());
+
+        debug_assert!(buf_len < Self::BUF_CAPACITY);
+    }
+
+    fn finish(&self) -> Token {
+        let mut h1 = self.h1;
+        let mut h2 = self.h2;
+
         let mut k1 = Wrapping(0_i64);
         let mut k2 = Wrapping(0_i64);
 
-        debug_assert!(data.len() < 16);
+        let buf_len = self.total_len % Self::BUF_CAPACITY;
 
-        if data.len() > 8 {
-            for i in (8..data.len()).rev() {
-                k2 ^= Wrapping(data[i] as i8 as i64) << ((i - 8) * 8);
+        if buf_len > 8 {
+            for i in (8..buf_len).rev() {
+                k2 ^= Wrapping(self.buf[i] as i8 as i64) << ((i - 8) * 8);
             }
 
-            k2 *= c2;
+            k2 *= Self::C2;
             k2 = Self::rotl64(k2, 33);
-            k2 *= c1;
+            k2 *= Self::C1;
             h2 ^= k2;
         }
 
-        if !data.is_empty() {
-            for i in (0..std::cmp::min(8, data.len())).rev() {
-                k1 ^= Wrapping(data[i] as i8 as i64) << (i * 8);
+        if buf_len > 0 {
+            for i in (0..std::cmp::min(8, buf_len)).rev() {
+                k1 ^= Wrapping(self.buf[i] as i8 as i64) << (i * 8);
             }
 
-            k1 *= c1;
+            k1 *= Self::C1;
             k1 = Self::rotl64(k1, 31);
-            k1 *= c2;
+            k1 *= Self::C2;
             h1 ^= k1;
         }
 
-        h1 ^= Wrapping(length as i64);
-        h2 ^= Wrapping(length as i64);
+        h1 ^= Wrapping(self.total_len as i64);
+        h2 ^= Wrapping(self.total_len as i64);
 
         h1 += h2;
         h2 += h1;
@@ -112,47 +256,122 @@ impl Murmur3Partitioner {
         h1 += h2;
         h2 += h1;
 
-        ((h2.0 as i128) << 64) | h1.0 as i128
+        Token {
+            value: (((h2.0 as i128) << 64) | h1.0 as i128) as i64,
+        }
     }
+}
 
-    #[inline]
-    fn rotl64(v: Wrapping<i64>, n: u32) -> Wrapping<i64> {
-        Wrapping((v.0 << n) | (v.0 as u64 >> (64 - n)) as i64)
-    }
+enum CDCPartitionerHasherState {
+    Feeding {
+        len: usize,
+        buf: [u8; CDCPartitionerHasher::BUF_CAPACITY],
+    },
+    Computed(Token),
+}
 
-    #[inline]
-    fn fmix(mut k: Wrapping<i64>) -> Wrapping<i64> {
-        k ^= Wrapping((k.0 as u64 >> 33) as i64);
-        k *= Wrapping(0xff51afd7ed558ccd_u64 as i64);
-        k ^= Wrapping((k.0 as u64 >> 33) as i64);
-        k *= Wrapping(0xc4ceb9fe1a85ec53_u64 as i64);
-        k ^= Wrapping((k.0 as u64 >> 33) as i64);
+pub struct CDCPartitioner;
 
-        k
+pub struct CDCPartitionerHasher {
+    state: CDCPartitionerHasherState,
+}
+
+impl Partitioner for CDCPartitioner {
+    type Hasher = CDCPartitionerHasher;
+
+    fn build_hasher(&self) -> Self::Hasher {
+        Self::Hasher {
+            state: CDCPartitionerHasherState::Feeding {
+                len: 0,
+                buf: Default::default(),
+            },
+        }
     }
 }
 
-impl Partitioner for Murmur3Partitioner {
-    fn hash(pk: &[u8]) -> Token {
-        Token {
-            value: Self::hash3_x64_128(pk) as i64,
+impl CDCPartitionerHasher {
+    const BUF_CAPACITY: usize = 8;
+}
+
+impl PartitionerHasher for CDCPartitionerHasher {
+    fn write(&mut self, pk_part: &[u8]) {
+        match &mut self.state {
+            CDCPartitionerHasherState::Feeding { len, buf } => {
+                // We feed the buffer until it's full.
+                let copied_len = Ord::min(pk_part.len(), Self::BUF_CAPACITY - *len);
+                buf[*len..*len + copied_len].copy_from_slice(&pk_part[..copied_len]);
+                *len += copied_len;
+
+                // If the buffer is full, we can compute and fix the token.
+                if *len == Self::BUF_CAPACITY {
+                    let token = Token {
+                        value: (&mut &buf[..]).get_i64(),
+                    };
+                    self.state = CDCPartitionerHasherState::Computed(token);
+                }
+            }
+            CDCPartitionerHasherState::Computed(_) => (),
+        }
+    }
+
+    fn finish(&self) -> Token {
+        match self.state {
+            CDCPartitionerHasherState::Feeding { .. } => Token { value: i64::MIN },
+            CDCPartitionerHasherState::Computed(token) => token,
         }
     }
 }
 
-impl Partitioner for CDCPartitioner {
-    fn hash(mut pk: &[u8]) -> Token {
-        let value = if pk.len() < 8 { i64::MIN } else { pk.get_i64() };
-        Token { value }
+/// Calculates the token for given partitioner and serialized partition key.
+///
+/// The ordinary way to calculate token is based on a PreparedStatement
+/// and values for that statement. However, if a user knows:
+/// - the order of the columns in the partition key,
+/// - the values of the columns of the partition key,
+/// - the partitioner of the table that the statement operates on,
+///
+/// then having a `PreparedStatement` is not necessary and the token can
+/// be calculated based on that information.
+///
+/// NOTE: the provided values must completely constitute partition key
+/// and be in the order defined in CREATE TABLE statement.
+pub fn calculate_token_for_partition_key<P: Partitioner>(
+    serialized_partition_key_values: &SerializedValues,
+    partitioner: &P,
+) -> Result<Token, TokenCalculationError> {
+    let mut partitioner_hasher = partitioner.build_hasher();
+
+    if serialized_partition_key_values.len() == 1 {
+        let val = serialized_partition_key_values.iter().next().unwrap();
+        if let Some(val) = val {
+            partitioner_hasher.write(val);
+        }
+    } else {
+        for val in serialized_partition_key_values.iter().flatten() {
+            let val_len_u16: u16 = val
+                .len()
+                .try_into()
+                .map_err(|_| TokenCalculationError::ValueTooLong(val.len()))?;
+            partitioner_hasher.write(&val_len_u16.to_be_bytes());
+            partitioner_hasher.write(val);
+            partitioner_hasher.write(&[0u8]);
+        }
     }
+
+    Ok(partitioner_hasher.finish())
 }
 
 #[cfg(test)]
 mod tests {
+    use rand::Rng;
+    use rand_pcg::Pcg32;
+
+    use crate::transport::partitioner::PartitionerHasher;
+
     use super::{CDCPartitioner, Murmur3Partitioner, Partitioner};
 
     fn assert_correct_murmur3_hash(pk: &'static str, expected_hash: i64) {
-        let hash = Murmur3Partitioner::hash(pk.as_bytes()).value;
+        let hash = Murmur3Partitioner.hash_one(pk.as_bytes()).value;
         assert_eq!(hash, expected_hash);
     }
 
@@ -169,10 +388,68 @@ mod tests {
     }
 
     fn assert_correct_cdc_hash(pk: &'static str, expected_hash: i64) {
-        let hash = CDCPartitioner::hash(pk.as_bytes()).value;
+        let hash = CDCPartitioner.hash_one(pk.as_bytes()).value;
         assert_eq!(hash, expected_hash);
     }
 
+    #[test]
+    fn partitioners_output_same_result_no_matter_how_input_is_partitioned() {
+        let inputs: &[&[u8]] = &[
+            b"",
+            b"0",
+            "Ala ma kota, a kota ma Ala.".as_bytes(),
+            "Zażółć gęślą jaźń. Wsiadł rycerz Szaławiła na bułanego konia. Litwo, ojczyzno moja, ...".as_bytes(),
+        ];
+
+        let seed = 0x2137;
+        let mut randgen = Pcg32::new(seed, 0);
+
+        // Splits the given data 2^n times and feeds partitioner with the chunks got.
+        fn split_and_feed(
+            randgen: &mut impl Rng,
+            partitioner: &mut impl PartitionerHasher,
+            data: &[u8],
+            n: usize,
+        ) {
+            if n == 0 {
+                partitioner.write(data);
+            } else {
+                let pivot = if !data.is_empty() {
+                    randgen.gen_range(0..data.len())
+                } else {
+                    0
+                };
+                let (data1, data2) = data.split_at(pivot);
+                for data in [data1, data2] {
+                    split_and_feed(randgen, partitioner, data, n - 1);
+                }
+            }
+        }
+
+        fn check_for_partitioner<P: Partitioner>(
+            partitioner: P,
+            randgen: &mut impl Rng,
+            input: &[u8],
+        ) {
+            let result_single_batch = partitioner.hash_one(input);
+
+            let results_chunks = (0..1000).map(|_| {
+                let mut partitioner_hasher = partitioner.build_hasher();
+                split_and_feed(randgen, &mut partitioner_hasher, input, 2);
+                partitioner_hasher.finish()
+            });
+
+            for result_chunk in results_chunks {
+                assert_eq!(result_single_batch, result_chunk)
+            }
+        }
+
+        for input in inputs {
+            check_for_partitioner(Murmur3Partitioner, &mut randgen, input);
+            check_for_partitioner(CDCPartitioner, &mut randgen, input);
+        }
+    }
+
     #[test]
     fn test_cdc_partitioner() {
         for s in [
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index abda176b..2f09f9fe 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -7,18 +7,16 @@ use crate::cloud::CloudConfig;
 use crate::frame::types::LegacyConsistency;
 use crate::history;
 use crate::history::HistoryListener;
-use crate::prepared_statement::PartitionKeyDecoder;
 use crate::utils::pretty::{CommaSeparatedDisplayer, CqlValueDisplayer};
+use crate::utils::unzip_option;
 use arc_swap::ArcSwapOption;
 use async_trait::async_trait;
-use bytes::BufMut;
 use bytes::Bytes;
-use bytes::BytesMut;
 use futures::future::join_all;
 use futures::future::try_join_all;
 use itertools::Either;
 pub use scylla_cql::errors::TranslationError;
-use scylla_cql::frame::response::result::{PreparedMetadata, Rows};
+use scylla_cql::frame::response::result::{deser_cql_value, ColumnSpec, Rows};
 use scylla_cql::frame::response::NonErrorResponse;
 use std::borrow::Borrow;
 use std::collections::HashMap;
@@ -43,9 +41,8 @@ use super::connection::NonErrorQueryResponse;
 use super::connection::QueryResponse;
 #[cfg(feature = "ssl")]
 use super::connection::SslConfig;
-use super::errors::{BadQuery, NewSessionError, QueryError};
+use super::errors::{NewSessionError, QueryError};
 use super::execution_profile::{ExecutionProfile, ExecutionProfileHandle, ExecutionProfileInner};
-use super::partitioner::Partitioner;
 use super::partitioner::PartitionerName;
 use super::topology::UntranslatedPeer;
 use super::NodeRef;
@@ -53,9 +50,9 @@ use crate::cql_to_rust::FromRow;
 use crate::frame::response::cql_to_rust::FromRowError;
 use crate::frame::response::result;
 use crate::frame::value::{
-    BatchValues, BatchValuesFirstSerialized, BatchValuesIterator, SerializedValues, ValueList,
+    BatchValues, BatchValuesFirstSerialized, BatchValuesIterator, ValueList,
 };
-use crate::prepared_statement::{PartitionKeyError, PreparedStatement};
+use crate::prepared_statement::PreparedStatement;
 use crate::query::Query;
 use crate::routing::Token;
 use crate::statement::{Consistency, SerialConsistency};
@@ -982,10 +979,11 @@ impl Session {
         let values_ref = &serialized_values;
         let paging_state_ref = &paging_state;
 
-        let partition_key = self.calculate_partition_key(prepared, &serialized_values)?;
-        let token = partition_key
-            .as_ref()
-            .map(|pk| prepared.get_partitioner_name().hash(pk));
+        let (partition_key, token) =
+            unzip_option(prepared.extract_partition_key_and_calculate_token(
+                prepared.get_partitioner_name(),
+                &serialized_values,
+            )?);
 
         let execution_profile = prepared
             .get_execution_profile_handle()
@@ -1007,8 +1005,7 @@ impl Session {
         };
 
         let span = RequestSpan::new_prepared(
-            prepared.get_prepared_metadata(),
-            partition_key.as_ref(),
+            partition_key.as_ref().map(|pk| pk.iter()),
             token,
             serialized_values.size(),
         );
@@ -1125,10 +1122,6 @@ impl Session {
     ) -> Result<RowIterator, QueryError> {
         let prepared = prepared.into();
         let serialized_values = values.serialized()?;
-        let partition_key = self.calculate_partition_key(&prepared, &serialized_values)?;
-        let token = partition_key
-            .as_ref()
-            .map(|pk| prepared.get_partitioner_name().hash(pk));
 
         let execution_profile = prepared
             .get_execution_profile_handle()
@@ -1138,8 +1131,6 @@ impl Session {
         RowIterator::new_for_prepared_statement(PreparedIteratorConfig {
             prepared,
             values: serialized_values.into_owned(),
-            partition_key,
-            token,
             execution_profile,
             cluster_data: self.cluster.get_data(),
             metrics: self.metrics.clone(),
@@ -1219,7 +1210,7 @@ impl Session {
                 RoutingInfo {
                     consistency,
                     serial_consistency,
-                    token: self.calculate_token(ps, first_serialized_value)?,
+                    token: ps.calculate_token(first_serialized_value)?,
                     keyspace: ps.get_keyspace_name(),
                     is_confirmed_lwt: false,
                 }
@@ -1890,76 +1881,6 @@ impl Session {
         .await
     }
 
-    fn calculate_partition_key(
-        &self,
-        prepared: &PreparedStatement,
-        serialized_values: &SerializedValues,
-    ) -> Result<Option<Bytes>, QueryError> {
-        if !prepared.is_token_aware() {
-            return Ok(None);
-        }
-        let partition_key = calculate_partition_key(prepared, serialized_values)?;
-        Ok(Some(partition_key))
-    }
-
-    /// Calculates the token for given prepared statement and serialized values.
-    ///
-    /// Returns the token that would be computed for executing the provided
-    /// prepared statement with the provided values.
-    pub fn calculate_token(
-        &self,
-        prepared: &PreparedStatement,
-        serialized_values: &SerializedValues,
-    ) -> Result<Option<Token>, QueryError> {
-        match self.calculate_partition_key(prepared, serialized_values) {
-            Ok(Some(partition_key)) => {
-                let partitioner_name = prepared.get_partitioner_name();
-                Ok(Some(partitioner_name.hash(&partition_key)))
-            }
-            Ok(None) => Ok(None),
-            Err(err) => Err(err),
-        }
-    }
-
-    /// Calculates the token for given partitioner and serialized partition key.
-    ///
-    /// The ordinary way to calculate token is based on a PreparedStatement
-    /// and values for that statement. However, if a user knows:
-    /// - the order of the columns in the partition key,
-    /// - the values of the columns of the partition key,
-    /// - the partitioner of the table that the statement operates on,
-    ///
-    /// then having a `PreparedStatement` is not necessary and the token can
-    /// be calculated based on that information.
-    ///
-    /// NOTE: the provided values must completely constitute partition key
-    /// and be in the order defined in CREATE TABLE statement.
-    pub fn calculate_token_for_partition_key<P: Partitioner>(
-        serialized_partition_key_values: &SerializedValues,
-        _partitioner: &P,
-    ) -> Result<Token, PartitionKeyError> {
-        let mut buf: BytesMut = BytesMut::new();
-
-        if serialized_partition_key_values.len() == 1 {
-            let val = serialized_partition_key_values.iter().next().unwrap();
-            if let Some(val) = val {
-                buf.extend_from_slice(val);
-            }
-        } else {
-            for val in serialized_partition_key_values.iter().flatten() {
-                let val_len_u16: u16 = val
-                    .len()
-                    .try_into()
-                    .map_err(|_| PartitionKeyError::ValueTooLong(val.len()))?;
-                buf.put_u16(val_len_u16);
-                buf.extend_from_slice(val);
-                buf.put_u8(0);
-            }
-        }
-
-        Ok(P::hash(&buf.freeze()))
-    }
-
     /// Retrieves the handle to execution profile that is used by this session
     /// by default, i.e. when an executed statement does not define its own handle.
     pub fn get_default_execution_profile_handle(&self) -> &ExecutionProfileHandle {
@@ -1967,21 +1888,6 @@ impl Session {
     }
 }
 
-fn calculate_partition_key(
-    stmt: &PreparedStatement,
-    values: &SerializedValues,
-) -> Result<Bytes, QueryError> {
-    match stmt.compute_partition_key(values) {
-        Ok(key) => Ok(key),
-        Err(PartitionKeyError::NoPkIndexValue(_, _)) => Err(QueryError::ProtocolError(
-            "No pk indexes - can't calculate token",
-        )),
-        Err(PartitionKeyError::ValueTooLong(values_len)) => Err(QueryError::BadQuery(
-            BadQuery::ValuesTooLongForKey(values_len, u16::MAX.into()),
-        )),
-    }
-}
-
 // Resolve the given hostname using a DNS lookup if necessary.
 // The resolution may return multiple IPs and the function returns one of them.
 // It prefers to return IPv4s first, and only if there are none, IPv6s.
@@ -2111,9 +2017,8 @@ impl RequestSpan {
         }
     }
 
-    pub(crate) fn new_prepared(
-        prepared_metadata: &PreparedMetadata,
-        partition_key: Option<&Bytes>,
+    pub(crate) fn new_prepared<'ps>(
+        partition_key: Option<impl Iterator<Item = (&'ps [u8], &'ps ColumnSpec)> + Clone>,
         token: Option<Token>,
         request_size: usize,
     ) -> Self {
@@ -2136,10 +2041,9 @@ impl RequestSpan {
         if let Some(partition_key) = partition_key {
             span.record(
                 "partition_key",
-                tracing::field::display(format_args!(
-                    "{}",
-                    partition_key_displayer(prepared_metadata, partition_key),
-                )),
+                tracing::field::display(
+                    format_args!("{}", partition_key_displayer(partition_key),),
+                ),
             );
         }
         if let Some(token) = token {
@@ -2238,12 +2142,16 @@ impl Drop for RequestSpan {
     }
 }
 
-fn partition_key_displayer<'pk>(
-    prepared_metadata: &'pk PreparedMetadata,
-    partition_key: &'pk [u8],
-) -> impl Display + 'pk {
+fn partition_key_displayer<'ps, 'res>(
+    mut pk_values_iter: impl Iterator<Item = (&'ps [u8], &'ps ColumnSpec)> + 'res + Clone,
+) -> impl Display + 'res {
     CommaSeparatedDisplayer(
-        PartitionKeyDecoder::new(prepared_metadata, partition_key).map(|c| match c {
+        std::iter::from_fn(move || {
+            pk_values_iter
+                .next()
+                .map(|(mut cell, spec)| deser_cql_value(&spec.typ, &mut cell))
+        })
+        .map(|c| match c {
             Ok(c) => Either::Left(CqlValueDisplayer(c)),
             Err(_) => Either::Right("<decoding error>"),
         }),
diff --git a/scylla/src/utils/mod.rs b/scylla/src/utils/mod.rs
index c83adbf7..9399e332 100644
--- a/scylla/src/utils/mod.rs
+++ b/scylla/src/utils/mod.rs
@@ -2,3 +2,11 @@ pub(crate) mod parse;
 
 pub(crate) mod pretty;
 pub mod test_utils;
+
+// FIXME: replace this with `Option::unzip()` once MSRV is bumped to at least 1.66
+pub(crate) fn unzip_option<T, U>(opt: Option<(T, U)>) -> (Option<T>, Option<U>) {
+    match opt {
+        Some((a, b)) => (Some(a), Some(b)),
+        None => (None, None),
+    }
+}


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 35e2f7bf..5b9914f9 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -10,7 +10,9 @@ use crate::statement::Consistency;
 use crate::tracing::TracingInfo;
 use crate::transport::cluster::Datacenter;
 use crate::transport::errors::{BadKeyspaceName, BadQuery, DbError, QueryError};
-use crate::transport::partitioner::{Murmur3Partitioner, Partitioner, PartitionerName};
+use crate::transport::partitioner::{
+    calculate_token_for_partition_key, Murmur3Partitioner, Partitioner, PartitionerName,
+};
 use crate::transport::topology::Strategy::NetworkTopologyStrategy;
 use crate::transport::topology::{
     CollectionType, ColumnKind, CqlType, NativeType, UserDefinedType,
@@ -229,7 +231,7 @@ async fn test_prepared_statement() {
                 .as_bigint()
                 .unwrap(),
         };
-        let prepared_token = Murmur3Partitioner::hash(
+        let prepared_token = Murmur3Partitioner.hash_one(
             &prepared_statement
                 .compute_partition_key(&serialized_values)
                 .unwrap(),
@@ -255,7 +257,7 @@ async fn test_prepared_statement() {
                 .as_bigint()
                 .unwrap(),
         };
-        let prepared_token = Murmur3Partitioner::hash(
+        let prepared_token = Murmur3Partitioner.hash_one(
             &prepared_complex_pk_statement
                 .compute_partition_key(&serialized_values)
                 .unwrap(),
@@ -527,7 +529,7 @@ async fn test_token_calculation() {
                 .as_bigint()
                 .unwrap(),
         };
-        let prepared_token = Murmur3Partitioner::hash(
+        let prepared_token = Murmur3Partitioner.hash_one(
             &prepared_statement
                 .compute_partition_key(&serialized_values)
                 .unwrap(),
@@ -2797,15 +2799,13 @@ async fn test_manual_primary_key_computation() {
             .await
             .unwrap();
 
-        let token_by_prepared = session
-            .calculate_token(prepared, &serialized_values_in_query_order)
+        let token_by_prepared = prepared
+            .calculate_token(&serialized_values_in_query_order)
             .unwrap()
             .unwrap();
-        let token_by_hand = Session::calculate_token_for_partition_key(
-            &serialized_values_in_pk_order,
-            &Murmur3Partitioner,
-        )
-        .unwrap();
+        let token_by_hand =
+            calculate_token_for_partition_key(&serialized_values_in_pk_order, &Murmur3Partitioner)
+                .unwrap();
         println!(
             "by_prepared: {}, by_hand: {}",
             token_by_prepared.value, token_by_hand.value