Merge pull request #33299 from petrosagg/refactor-upsert-value

upsert: refactor `UpsertValue` struct and box error to reduce memory usage

Author: petrosagg

src/storage/src/render/sources.rs

@@ -43,7 +43,7 @@ use crate::decode::{render_decode_cdcv2, render_decode_delimited};
 use crate::healthcheck::{HealthStatusMessage, StatusNamespace};
 use crate::source::types::{DecodeResult, SourceOutput, SourceRender};
 use crate::source::{self, RawSourceCreationConfig, SourceExportCreationConfig};
-use crate::upsert::UpsertKey;
+use crate::upsert::{UpsertKey, UpsertValue};
 
 /// _Renders_ complete _differential_ [`Collection`]s
 /// that represent the final source and its errors
@@ -479,7 +479,7 @@ fn append_metadata_to_value<G: Scope, FromTime: Timestamp>(
 fn upsert_commands<G: Scope, FromTime: Timestamp>(
     input: Collection<G, DecodeResult<FromTime>, Diff>,
     upsert_envelope: UpsertEnvelope,
-) -> Collection<G, (UpsertKey, Option<Result<Row, UpsertError>>, FromTime), Diff> {
+) -> Collection<G, (UpsertKey, Option<UpsertValue>, FromTime), Diff> {
     let mut row_buf = Row::default();
     input.map(move |result| {
         let from_time = result.from_time;
@@ -493,9 +493,15 @@ fn upsert_commands<G: Scope, FromTime: Timestamp>(
         // If we have a well-formed key we can continue, otherwise we're upserting an error
         let key = match key {
             Ok(key) => key,
-            err @ Err(_) => match result.value {
-                Some(_) => return (UpsertKey::from_key(err.as_ref()), Some(err), from_time),
-                None => return (UpsertKey::from_key(err.as_ref()), None, from_time),
+            Err(err) => match result.value {
+                Some(_) => {
+                    return (
+                        UpsertKey::from_key(Err(&err)),
+                        Some(Err(Box::new(err))),
+                        from_time,
+                    );
+                }
+                None => return (UpsertKey::from_key(Err(&err)), None, from_time),
             },
         };
 
@@ -584,10 +590,10 @@ fn upsert_commands<G: Scope, FromTime: Timestamp>(
                         packer.extend_by_row(&metadata);
                         Some(Ok(row_buf.clone()))
                     }
-                    _ => Some(Err(UpsertError::Value(UpsertValueError {
+                    _ => Some(Err(Box::new(UpsertError::Value(UpsertValueError {
                         for_key: key_row,
                         inner,
-                    }))),
+                    })))),
                 }
             }
             None => None,

src/storage/src/upsert.rs

@@ -48,7 +48,7 @@ use crate::metrics::upsert::UpsertMetrics;
 use crate::storage_state::StorageInstanceContext;
 use crate::upsert_continual_feedback;
 use types::{
-    BincodeOpts, StateValue, UpsertState, UpsertStateBackend, Value, consolidating_merge_function,
+    BincodeOpts, StateValue, UpsertState, UpsertStateBackend, consolidating_merge_function,
     upsert_bincode_opts,
 };
 
@@ -58,7 +58,7 @@ pub(crate) mod rocksdb;
 // TODO(aljoscha): Move next to upsert module, rename to upsert_types.
 pub(crate) mod types;
 
-pub type UpsertValue = Result<Row, UpsertError>;
+pub type UpsertValue = Result<Row, Box<UpsertError>>;
 
 #[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct UpsertKey([u8; 32]);
@@ -508,7 +508,7 @@ enum DrainStyle<'a, T> {
 async fn drain_staged_input<S, G, T, FromTime, E>(
     stash: &mut Vec<(T, UpsertKey, Reverse<FromTime>, Option<UpsertValue>)>,
     commands_state: &mut indexmap::IndexMap<UpsertKey, types::UpsertValueAndSize<T, FromTime>>,
-    output_updates: &mut Vec<(Result<Row, UpsertError>, T, Diff)>,
+    output_updates: &mut Vec<(UpsertValue, T, Diff)>,
     multi_get_scratch: &mut Vec<UpsertKey>,
     drain_style: DrainStyle<'_, T>,
     error_emitter: &mut E,
@@ -607,15 +607,15 @@ async fn drain_staged_input<S, G, T, FromTime, E>(
                 if let Some(old_value) =
                     existing_value.replace(StateValue::finalized_value(value.clone()))
                 {
-                    if let Value::FinalizedValue(old_value) = old_value.into_decoded() {
+                    if let Some(old_value) = old_value.into_decoded().finalized {
                         output_updates.push((old_value, ts.clone(), Diff::MINUS_ONE));
                     }
                 }
                 output_updates.push((value, ts, Diff::ONE));
             }
             None => {
                 if let Some(old_value) = existing_value.take() {
-                    if let Value::FinalizedValue(old_value) = old_value.into_decoded() {
+                    if let Some(old_value) = old_value.into_decoded().finalized {
                         output_updates.push((old_value, ts, Diff::MINUS_ONE));
                     }
                 }
@@ -691,12 +691,16 @@ where
         let value = match result {
             Ok(ok) => Ok(ok),
             Err(DataflowError::EnvelopeError(err)) => match *err {
-                EnvelopeError::Upsert(err) => Err(err),
+                EnvelopeError::Upsert(err) => Err(Box::new(err)),
                 _ => return None,
             },
             Err(_) => return None,
         };
-        Some((UpsertKey::from_value(value.as_ref(), &key_indices), value))
+        let value_ref = match value {
+            Ok(ref row) => Ok(row),
+            Err(ref err) => Err(&**err),
+        };
+        Some((UpsertKey::from_value(value_ref, &key_indices), value))
     });
     let (output_handle, output) = builder.new_output();
 
@@ -931,7 +935,7 @@ where
     (
         output.as_collection().map(|result| match result {
             Ok(ok) => Ok(ok),
-            Err(err) => Err(DataflowError::from(EnvelopeError::Upsert(err))),
+            Err(err) => Err(DataflowError::from(EnvelopeError::Upsert(*err))),
         }),
         health_stream,
         snapshot_stream,

src/storage/src/upsert/memory.rs

@@ -17,6 +17,7 @@
 use std::collections::HashMap;
 
 use itertools::Itertools;
+use mz_ore::cast::CastFrom;
 
 use super::UpsertKey;
 use super::types::{
@@ -93,7 +94,7 @@ where
             stats.processed_gets += 1;
             let value = self.state.get(&key).cloned();
             let metadata = value.as_ref().map(|v| ValueMetadata {
-                size: v.memory_size(),
+                size: u64::cast_from(v.memory_size()),
                 is_tombstone: v.is_tombstone(),
             });
             stats.processed_gets_size += metadata.map_or(0, |m| m.size);