prepared: reuse prepared statement's result metadata

Author: muzarski

scylla-rust-wrapper/src/cass_types.rs

@@ -138,6 +138,7 @@ pub enum MapDataType {
     KeyAndValue(Arc<CassDataType>, Arc<CassDataType>),
 }
 
+#[derive(Debug)]
 pub struct CassColumnSpec {
     pub name: String,
     pub data_type: Arc<CassDataType>,

scylla-rust-wrapper/src/prepared.rs

@@ -5,6 +5,7 @@ use crate::{
     argconv::*,
     cass_error::CassError,
     cass_types::{get_column_type, CassDataType},
+    query_result::CassResultData,
     statement::{CassStatement, Statement},
     types::size_t,
 };
@@ -14,11 +15,10 @@ use scylla::prepared_statement::PreparedStatement;
 pub struct CassPrepared {
     // Data types of columns from PreparedMetadata.
     pub variable_col_data_types: Vec<Arc<CassDataType>>,
-    // Data types of columns from ResultMetadata.
-    //
-    // Arc<CassDataType> -> to share each data type with other structs such as `CassValue`
-    // Arc<Vec<...>> -> to share the whole vector with `CassResultData`.
-    pub result_col_data_types: Arc<Vec<Arc<CassDataType>>>,
+
+    // Cached result metadata. Arc'ed since we want to share it
+    // with result metadata after execution.
+    pub result_metadata: Arc<CassResultData>,
     pub statement: PreparedStatement,
 }
 
@@ -30,17 +30,13 @@ impl CassPrepared {
             .map(|col_spec| Arc::new(get_column_type(col_spec.typ())))
             .collect();
 
-        let result_col_data_types: Arc<Vec<Arc<CassDataType>>> = Arc::new(
-            statement
-                .get_result_set_col_specs()
-                .iter()
-                .map(|col_spec| Arc::new(get_column_type(col_spec.typ())))
-                .collect(),
-        );
+        let result_metadata = Arc::new(CassResultData::from_column_specs(
+            statement.get_result_set_col_specs(),
+        ));
 
         Self {
             variable_col_data_types,
-            result_col_data_types,
+            result_metadata,
             statement,
         }
     }

scylla-rust-wrapper/src/query_result.rs

@@ -23,6 +23,7 @@ pub struct CassResult {
     pub paging_state_response: PagingStateResponse,
 }
 
+#[derive(Debug)]
 pub struct CassResultData {
     pub col_specs: Vec<CassColumnSpec>,
 }

scylla-rust-wrapper/src/session.rs

@@ -300,20 +300,20 @@ pub unsafe extern "C" fn cass_session_execute(
                 // Since `query.query` is consumed, we cannot match the statement
                 // after execution, to retrieve the cached metadata in case
                 // of prepared statements.
-                Option<Arc<Vec<Arc<CassDataType>>>>,
+                Option<Arc<CassResultData>>,
             ),
             QueryError,
         >;
         let query_res: QueryRes = match statement {
             Statement::Simple(query) => {
                 // We don't store result metadata for Queries - return None.
-                let maybe_result_col_data_types = None;
+                let maybe_result_metadata = None;
 
                 if paging_enabled {
                     session
                         .query_single_page(query.query, bound_values, paging_state)
                         .await
-                        .map(|(qr, psr)| (qr, psr, maybe_result_col_data_types))
+                        .map(|(qr, psr)| (qr, psr, maybe_result_metadata))
                 } else {
                     session
                         .query_unpaged(query.query, bound_values)
@@ -322,21 +322,21 @@ pub unsafe extern "C" fn cass_session_execute(
                             (
                                 result,
                                 PagingStateResponse::NoMorePages,
-                                maybe_result_col_data_types,
+                                maybe_result_metadata,
                             )
                         })
                 }
             }
             Statement::Prepared(prepared) => {
-                // Clone vector of the Arc<CassDataType>, so we don't do additional allocations when constructing
-                // CassDataTypes in `CassResultData::from_result_payload`.
-                let maybe_result_col_data_types = Some(prepared.result_col_data_types.clone());
+                // Clone result metadata, so we don't need to construct it from scratch in
+                // `CassResultData::from_column_specs` - it requires a lot of allocations for complex types.
+                let maybe_result_metadata = Some(Arc::clone(&prepared.result_metadata));
 
                 if paging_enabled {
                     session
                         .execute_single_page(&prepared.statement, bound_values, paging_state)
                         .await
-                        .map(|(qr, psr)| (qr, psr, maybe_result_col_data_types))
+                        .map(|(qr, psr)| (qr, psr, maybe_result_metadata))
                 } else {
                     session
                         .execute_unpaged(&prepared.statement, bound_values)
@@ -345,16 +345,21 @@ pub unsafe extern "C" fn cass_session_execute(
                             (
                                 result,
                                 PagingStateResponse::NoMorePages,
-                                maybe_result_col_data_types,
+                                maybe_result_metadata,
                             )
                         })
                 }
             }
         };
 
         match query_res {
-            Ok((result, paging_state_response, _maybe_col_data_types)) => {
-                let metadata = Arc::new(CassResultData::from_column_specs(result.col_specs()));
+            Ok((result, paging_state_response, maybe_result_metadata)) => {
+                // maybe_result_metadata is:
+                // - Some(_) for prepared statements
+                // - None for unprepared statements
+                let metadata = maybe_result_metadata.unwrap_or_else(|| {
+                    Arc::new(CassResultData::from_column_specs(result.col_specs()))
+                });
                 let cass_rows = result
                     .rows
                     .map(|rows| create_cass_rows_from_rows(rows, &metadata));