switch from serde_json to postcard (#3170)

* switch from json to postcard

switch to postcard for better datastructure support when sending intermediate results between nodes

* Apply suggestions from code review

Co-authored-by: Paul Masurel <[email protected]>

* Update quickwit/quickwit-proto/protos/quickwit/search_api.proto

Co-authored-by: Paul Masurel <[email protected]>

---------

Co-authored-by: Paul Masurel <[email protected]>

Author: PSeitz

quickwit/Cargo.lock

@@ -225,7 +225,7 @@ quickwit-serve = { version = "0.5.0", path = "./quickwit-serve" }
 quickwit-storage = { version = "0.5.0", path = "./quickwit-storage" }
 quickwit-telemetry = { version = "0.5.0", path = "./quickwit-telemetry" }
 
-tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "80df1d9", default-features = false, features = [
+tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "c599bf3", default-features = false, features = [
   "mmap",
   "lz4-compression",
   "zstd-compression",

quickwit/Cargo.toml

@@ -27,8 +27,7 @@ use quickwit_proto::opentelemetry::proto::common::v1::any_value::Value as OtlpVa
 use quickwit_proto::opentelemetry::proto::common::v1::{
     AnyValue as OtlpAnyValue, KeyValue as OtlpKeyValue,
 };
-use serde;
-use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
+use serde::{self, de, Deserialize, Deserializer, Serialize, Serializer};
 use serde_json::{Number as JsonNumber, Value as JsonValue};
 use tracing::warn;
 

quickwit/quickwit-opentelemetry/src/otlp/mod.rs

@@ -240,8 +240,12 @@ message LeafSearchResponse {
   // num_attempted_splits = num_successful_splits + num_failed_splits.
   uint64 num_attempted_splits = 4;
 
-  // json serialized intermediate aggregation_result.
-  optional string intermediate_aggregation_result = 5;
+  // Deprecated json serialized intermediate aggregation_result.
+  reserved 5;
+
+  // postcard serialized intermediate aggregation_result.
+  optional bytes intermediate_aggregation_result = 6;
+ 
 
 }
 

quickwit/quickwit-proto/protos/quickwit/search_api.proto

@@ -199,10 +199,10 @@ pub struct LeafSearchResponse {
     /// num_attempted_splits = num_successful_splits + num_failed_splits.
     #[prost(uint64, tag = "4")]
     pub num_attempted_splits: u64,
-    /// json serialized intermediate aggregation_result.
-    #[prost(string, optional, tag = "5")]
+    /// postcard serialized intermediate aggregation_result.
+    #[prost(bytes = "vec", optional, tag = "5")]
     pub intermediate_aggregation_result: ::core::option::Option<
-        ::prost::alloc::string::String,
+        ::prost::alloc::vec::Vec<u8>,
     >,
 }
 #[derive(Serialize, Deserialize, utoipa::ToSchema)]

quickwit/quickwit-proto/src/quickwit.rs

@@ -21,6 +21,7 @@ itertools = { workspace = true }
 lru = { workspace = true }
 mockall = { workspace = true }
 once_cell = { workspace = true }
+postcard = { workspace = true }
 rayon = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }

quickwit/quickwit-search/Cargo.toml

@@ -168,17 +168,18 @@ fn merge_leaf_search_results(
                 .append(&mut retry_response.partial_hits);
             let intermediate_aggregation_result = initial_response
                 .intermediate_aggregation_result
-                .map::<crate::Result<_>, _>(|res1_str| {
+                .map::<crate::Result<_>, _>(|res1_bytes| {
                     if let Some(res2_str) = retry_response.intermediate_aggregation_result.as_ref()
                     {
                         let mut res1: IntermediateAggregationResults =
-                            serde_json::from_str(&res1_str)?;
-                        let res2: IntermediateAggregationResults = serde_json::from_str(res2_str)?;
+                            postcard::from_bytes(res1_bytes.as_slice())?;
+                        let res2: IntermediateAggregationResults =
+                            postcard::from_bytes(res2_str.as_slice())?;
                         res1.merge_fruits(res2)?;
-                        let res = serde_json::to_string(&res1)?;
-                        Ok(res)
+                        let serialized = postcard::to_allocvec(&res1)?;
+                        Ok(serialized)
                     } else {
-                        Ok(res1_str)
+                        Ok(res1_bytes)
                     }
                 })
                 .transpose()?;

quickwit/quickwit-search/src/cluster_client.rs

@@ -31,7 +31,7 @@ use tantivy::aggregation::{AggregationLimits, AggregationSegmentCollector};
 use tantivy::collector::{Collector, SegmentCollector};
 use tantivy::columnar::ColumnType;
 use tantivy::fastfield::Column;
-use tantivy::{DocId, Score, SegmentOrdinal, SegmentReader};
+use tantivy::{DocId, Score, SegmentOrdinal, SegmentReader, TantivyError};
 
 use crate::filters::{create_timestamp_filter_builder, TimestampFilter, TimestampFilterBuilder};
 use crate::find_trace_ids_collector::{FindTraceIdsCollector, FindTraceIdsSegmentCollector};
@@ -267,15 +267,16 @@ impl SegmentCollector for QuickwitSegmentCollector {
             .collect();
 
         let intermediate_aggregation_result = match self.aggregation {
-            Some(AggregationSegmentCollectors::FindTraceIdsSegmentCollector(collector)) => Some(
-                serde_json::to_string(&collector.harvest())
-                    .expect("Collector fruit should be JSON serializable."),
-            ),
+            Some(AggregationSegmentCollectors::FindTraceIdsSegmentCollector(collector)) => {
+                let fruit = collector.harvest();
+                let serialized =
+                    postcard::to_allocvec(&fruit).expect("Collector fruit should be serializable.");
+                Some(serialized)
+            }
             Some(AggregationSegmentCollectors::TantivyAggregationSegmentCollector(collector)) => {
-                Some(
-                    serde_json::to_string(&collector.harvest()?)
-                        .expect("Collector fruit should be JSON serializable."),
-                )
+                let serialized = postcard::to_allocvec(&collector.harvest()?)
+                    .expect("Collector fruit should be serializable.");
+                Some(serialized)
             }
             None => None,
         };
@@ -438,15 +439,19 @@ impl Collector for QuickwitCollector {
     }
 }
 
+fn map_error(err: postcard::Error) -> TantivyError {
+    TantivyError::InternalError(format!("Merge Result Postcard Error: {}", err))
+}
+
 /// Merges a set of Leaf Results.
 fn merge_leaf_responses(
     aggregations_opt: &Option<QuickwitAggregations>,
-    leaf_responses: Vec<LeafSearchResponse>,
+    mut leaf_responses: Vec<LeafSearchResponse>,
     max_hits: usize,
 ) -> tantivy::Result<LeafSearchResponse> {
     // Optimization: No merging needed if there is only one result.
     if leaf_responses.len() == 1 {
-        return Ok(leaf_responses.into_iter().next().unwrap_or_default()); //< default is actually never called
+        return Ok(leaf_responses.pop().unwrap());
     }
     let merged_intermediate_aggregation_result = match aggregations_opt {
         Some(QuickwitAggregations::FindTraceIdsAggregation(collector)) => {
@@ -457,21 +462,24 @@ fn merge_leaf_responses(
                 .filter_map(|leaf_response| {
                     leaf_response.intermediate_aggregation_result.as_ref().map(
                         |intermediate_aggregation_result| {
-                            serde_json::from_str(intermediate_aggregation_result)
+                            postcard::from_bytes(intermediate_aggregation_result.as_slice())
+                                .map_err(map_error)
                         },
                     )
                 })
                 .collect::<Result<_, _>>()?;
             let merged_fruit = collector.merge_fruits(fruits)?;
-            Some(serde_json::to_string(&merged_fruit)?)
+            let serialized = postcard::to_allocvec(&merged_fruit).map_err(map_error)?;
+            Some(serialized)
         }
         Some(QuickwitAggregations::TantivyAggregations(_)) => {
             let fruits: Vec<IntermediateAggregationResults> = leaf_responses
                 .iter()
                 .filter_map(|leaf_response| {
                     leaf_response.intermediate_aggregation_result.as_ref().map(
                         |intermediate_aggregation_result| {
-                            serde_json::from_str(intermediate_aggregation_result)
+                            postcard::from_bytes(intermediate_aggregation_result.as_slice())
+                                .map_err(map_error)
                         },
                     )
                 })
@@ -483,7 +491,9 @@ fn merge_leaf_responses(
                 for fruit in fruit_iter {
                     merged_fruit.merge_fruits(fruit)?;
                 }
-                Some(serde_json::to_string(&merged_fruit)?)
+                let serialized = postcard::to_allocvec(&merged_fruit).map_err(map_error)?;
+
+                Some(serialized)
             } else {
                 None
             }

quickwit/quickwit-search/src/collector.rs

@@ -75,6 +75,12 @@ impl From<TantivyError> for SearchError {
     }
 }
 
+impl From<postcard::Error> for SearchError {
+    fn from(error: postcard::Error) -> Self {
+        SearchError::InternalError(format!("Postcard error: {error}"))
+    }
+}
+
 impl From<serde_json::Error> for SearchError {
     fn from(serde_error: serde_json::Error) -> Self {
         SearchError::InternalError(format!("Serde error: {serde_error}"))

quickwit/quickwit-search/src/error.rs

@@ -347,6 +347,7 @@ async fn leaf_search_single_split(
     .map_err(|_| {
         crate::SearchError::InternalError(format!("Leaf search panicked. split={split_id}"))
     })??;
+
     Ok(leaf_search_response)
 }
 

quickwit/quickwit-search/src/leaf.rs

@@ -46,9 +46,8 @@ mod tests;
 pub use collector::QuickwitAggregations;
 use metrics::SEARCH_METRICS;
 use quickwit_doc_mapper::DocMapper;
-use root::validate_request;
+use root::{finalize_aggregation, validate_request};
 use service::SearcherContext;
-use tantivy::aggregation::AggregationLimits;
 use tantivy::query::Query as TantivyQuery;
 use tantivy::schema::NamedFieldDocument;
 
@@ -66,8 +65,6 @@ use quickwit_doc_mapper::tag_pruning::extract_tags_from_query;
 use quickwit_metastore::{ListSplitsQuery, Metastore, SplitMetadata, SplitState};
 use quickwit_proto::{Hit, PartialHit, SearchRequest, SearchResponse, SplitIdAndFooterOffsets};
 use quickwit_storage::StorageUriResolver;
-use tantivy::aggregation::agg_result::AggregationResults;
-use tantivy::aggregation::intermediate_agg_result::IntermediateAggregationResults;
 use tantivy::DocAddress;
 
 pub use crate::client::{create_search_service_client, SearchServiceClient};
@@ -193,6 +190,7 @@ pub async fn single_node_search(
     // Validates the query by effectively building it against the current schema.
     doc_mapper.query(doc_mapper.schema(), search_request)?;
     let searcher_context = Arc::new(SearcherContext::new(SearcherConfig::default()));
+
     let leaf_search_response = leaf_search(
         searcher_context.clone(),
         search_request,
@@ -229,30 +227,17 @@ pub async fn single_node_search(
         })
         .collect();
     let elapsed = start_instant.elapsed();
-    let aggregation = if let Some(intermediate_aggregation_result) =
-        leaf_search_response.intermediate_aggregation_result
-    {
-        let aggregations: QuickwitAggregations =
-            serde_json::from_str(search_request.aggregation_request.as_ref().expect(
-                "Aggregation should be present since we are processing an intermediate \
-                 aggregation result.",
-            ))?;
-        match aggregations {
-            QuickwitAggregations::FindTraceIdsAggregation(_) => {
-                // There is nothing to merge here because there is only one leaf response.
-                Some(intermediate_aggregation_result)
-            }
-            QuickwitAggregations::TantivyAggregations(aggregations) => {
-                let res: IntermediateAggregationResults =
-                    serde_json::from_str(&intermediate_aggregation_result)?;
-                let res: AggregationResults =
-                    res.into_final_result(aggregations, &AggregationLimits::default())?;
-                Some(serde_json::to_string(&res)?)
-            }
-        }
-    } else {
-        None
-    };
+
+    let aggregations: Option<QuickwitAggregations> = search_request
+        .aggregation_request
+        .as_ref()
+        .map(|agg| serde_json::from_str(agg))
+        .transpose()?;
+
+    let aggregation = finalize_aggregation(
+        leaf_search_response.intermediate_aggregation_result,
+        aggregations,
+    )?;
     Ok(SearchResponse {
         aggregation,
         num_hits: leaf_search_response.num_hits,