Add faster specialization for deserializing vector<float>

scylladb · Aug 13, 2024 · dc4bc3d · dc4bc3d
1 parent de99af7
commit dc4bc3d
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/scylla-cql/src/frame/response/result.rs b/scylla-cql/src/frame/response/result.rs
@@ -7,6 +7,7 @@ use crate::frame::{frame_errors::ParseError, types};
 use crate::types::deserialize::result::{RowIterator, TypedRowIterator};
 use crate::types::deserialize::value::{
     mk_deser_err, BuiltinDeserializationErrorKind, DeserializeValue, MapIterator, UdtIterator,
+    VectorIterator,
 };
 use crate::types::deserialize::{DeserializationError, FrameSlice};
 use bytes::{Buf, Bytes};
@@ -829,9 +830,15 @@ pub fn deser_cql_value(
                 .collect::<StdResult<_, _>>()?;
             CqlValue::Tuple(t)
         }
-        Vector(_type_name, _) => {
-            let l = Vec::<CqlValue>::deserialize(typ, v)?;
-            CqlValue::Vector(l)
+        // Specialization for faster deserialization of vectors of floats, which are currently
+        // the only type of vector
+        Vector(elem_type, _) if matches!(elem_type.as_ref(), Float) => {
+            let v = VectorIterator::<CqlValue>::deserialize_vector_of_float_to_vec_of_cql_value(typ, v)?;
+            CqlValue::Vector(v)
+        }
+        Vector(_, _) => {
+            let v = Vec::<CqlValue>::deserialize(typ, v)?;
+            CqlValue::Vector(v)
         }
     })
 }

diff --git a/scylla-cql/src/types/deserialize/value.rs b/scylla-cql/src/types/deserialize/value.rs
@@ -1,12 +1,11 @@
 //! Provides types for dealing with CQL value deserialization.
 
+use bytes::Bytes;
 use std::{
     collections::{BTreeMap, BTreeSet, HashMap, HashSet},
     hash::{BuildHasher, Hash},
     net::IpAddr,
 };
-
-use bytes::Bytes;
 use uuid::Uuid;
 
 use std::fmt::{Display, Pointer};
@@ -759,7 +758,7 @@ pub struct VectorIterator<'frame, T> {
 }
 
 impl<'frame, T> VectorIterator<'frame, T> {
-    fn new(
+    pub fn new(
         coll_typ: &'frame ColumnType,
         elem_typ: &'frame ColumnType,
         count: usize,
@@ -774,6 +773,25 @@ impl<'frame, T> VectorIterator<'frame, T> {
             phantom_data: std::marker::PhantomData,
         }
     }
+
+    /// Faster specialization for deserializing a `vector<float>` into `Vec<CqlValue>`.  
+    /// The generic code `Vec<CqlValue>::deserialize(...)` is much slower because it has to 
+    /// match on the element type for every item in the vector. 
+    /// Here we just hardcode `f32` and we can shortcut a lot of code.
+    ///
+    /// This could be nicer if Rust had generic type specialization in stable,
+    /// but for now we need a separate method.    
+    pub fn deserialize_vector_of_float_to_vec_of_cql_value(
+        typ: &'frame ColumnType,
+        v: Option<FrameSlice<'frame>>,
+    ) -> Result<Vec<CqlValue>, DeserializationError> {
+        VectorIterator::<'frame, f32>::deserialize(typ, v)
+            .and_then(|it| {
+                it.map(|e| e.map(CqlValue::Float))
+                    .collect::<Result<Vec<CqlValue>, DeserializationError>>()
+            })
+            .map_err(deser_error_replace_rust_name::<Self>)
+    }
 }
 
 impl<'frame, T> DeserializeValue<'frame> for VectorIterator<'frame, T>
@@ -828,6 +846,7 @@ where
 {
     type Item = Result<T, DeserializationError>;
 
+    #[inline]
     fn next(&mut self) -> Option<Self::Item> {
         let raw = self.raw_iter.next()?.map_err(|err| {
             mk_deser_err::<Self>(
@@ -883,6 +902,7 @@ where
                     .and_then(|it| it.collect::<Result<_, DeserializationError>>())
                     .map_err(deser_error_replace_rust_name::<Self>)
             }
+
             ColumnType::Vector(_, _) => VectorIterator::<'frame, T>::deserialize(typ, v)
                 .and_then(|it| it.collect::<Result<_, DeserializationError>>())
                 .map_err(deser_error_replace_rust_name::<Self>),
@@ -1435,6 +1455,7 @@ impl<'frame> VectorBytesSequenceIterator<'frame> {
 impl<'frame> Iterator for VectorBytesSequenceIterator<'frame> {
     type Item = Result<Option<FrameSlice<'frame>>, LowLevelDeserializationError>;
 
+    #[inline]
     fn next(&mut self) -> Option<Self::Item> {
         self.remaining = self.remaining.checked_sub(1)?;
         Some(self.slice.read_subslice(self.elem_len))