refactor: direct mutation vs reconstruction

Author: Devdutt Shenoi

src/event/format/json.rs

@@ -30,7 +30,7 @@ use serde_json::Value;
 use std::{collections::HashMap, sync::Arc};
 use tracing::error;
 
-use super::{update_field_type_in_schema, EventFormat};
+use super::EventFormat;
 use crate::{metadata::SchemaVersion, storage::StreamType, utils::arrow::get_field};
 
 static TIME_FIELD_NAME_PARTS: [&str; 11] = [
@@ -74,7 +74,6 @@ impl EventFormat for Event {
     fn to_data(
         self,
         schema: &HashMap<String, Arc<Field>>,
-        time_partition: Option<&String>,
         schema_version: SchemaVersion,
         static_schema_flag: bool,
     ) -> Result<(Self::Data, Vec<Arc<Field>>, bool), anyhow::Error> {
@@ -101,12 +100,11 @@ impl EventFormat for Event {
                     .map_err(|err| {
                         anyhow!("Could not infer schema for this event due to err {:?}", err)
                     })?;
-                let new_infer_schema = update_field_type_in_schema(
-                    Arc::new(infer_schema),
-                    Some(stream_schema),
-                    time_partition,
-                );
-                infer_schema = Schema::new(new_infer_schema.fields().clone());
+
+                for log_record in value_arr.iter() {
+                    override_inferred_data_type(&mut infer_schema, log_record, schema_version);
+                }
+
                 Schema::try_merge(vec![
                     Schema::new(stream_schema.values().cloned().collect::<Fields>()),
                     infer_schema.clone(),
@@ -237,14 +235,14 @@ fn extract_and_parse_time(
 // From Schema v1 onwards, convert json fields with name containig "date"/"time" and having
 // a string value parseable into timestamp as timestamp type and all numbers as float64.
 pub fn override_inferred_data_type(
-    inferred_schema: Schema,
+    schema: &mut Schema,
     log_record: &Value,
     schema_version: SchemaVersion,
-) -> Schema {
+) {
     let Value::Object(map) = log_record else {
-        return inferred_schema;
+        return;
     };
-    let updated_fields = inferred_schema
+    schema.fields = schema
         .fields()
         .iter()
         .map(|field| {
@@ -277,9 +275,7 @@ pub fn override_inferred_data_type(
                 _ => Field::new(field_name, field.data_type().clone(), true),
             }
         })
-        .collect_vec();
-
-    Schema::new(updated_fields)
+        .collect();
 }
 
 // Returns arrow schema with the fields that are present in the request body

src/event/format/mod.rs

@@ -115,7 +115,6 @@ pub trait EventFormat: Sized {
     fn to_data(
         self,
         schema: &HashMap<String, Arc<Field>>,
-        time_partition: Option<&String>,
         schema_version: SchemaVersion,
         static_schema_flag: bool,
     ) -> Result<(Self::Data, EventSchema, bool), AnyError>;
@@ -134,12 +133,8 @@ pub trait EventFormat: Sized {
         p_custom_fields: &HashMap<String, String>,
     ) -> Result<(RecordBatch, bool), AnyError> {
         let p_timestamp = self.get_p_timestamp();
-        let (data, schema, is_first) = self.to_data(
-            storage_schema,
-            time_partition,
-            schema_version,
-            static_schema_flag,
-        )?;
+        let (data, schema, is_first) =
+            self.to_data(storage_schema, schema_version, static_schema_flag)?;
 
         if get_field(&schema, DEFAULT_TIMESTAMP_KEY).is_some() {
             return Err(anyhow!(
@@ -149,21 +144,22 @@ pub trait EventFormat: Sized {
         };
 
         // prepare the record batch and new fields to be added
-        let mut new_schema = Arc::new(Schema::new(schema));
-        if !Self::is_schema_matching(new_schema.clone(), storage_schema, static_schema_flag) {
+        let mut new_schema = Schema::new(schema);
+        if !Self::is_schema_matching(&new_schema, storage_schema, static_schema_flag) {
             return Err(anyhow!("Schema mismatch"));
         }
-        new_schema = update_field_type_in_schema(new_schema, None, time_partition);
 
-        let rb = Self::decode(data, new_schema.clone())?;
+        update_field_type_in_schema(&mut new_schema, Some(storage_schema), time_partition);
+        let updated_schema = Arc::new(new_schema);
 
+        let rb = Self::decode(data, updated_schema)?;
         let rb = add_parseable_fields(rb, p_timestamp, p_custom_fields)?;
 
         Ok((rb, is_first))
     }
 
     fn is_schema_matching(
-        new_schema: Arc<Schema>,
+        new_schema: &Schema,
         storage_schema: &HashMap<String, Arc<Field>>,
         static_schema_flag: bool,
     ) -> bool {
@@ -200,7 +196,7 @@ pub trait EventFormat: Sized {
 }
 
 pub fn get_existing_field_names(
-    inferred_schema: Arc<Schema>,
+    inferred_schema: &Schema,
     existing_schema: Option<&HashMap<String, Arc<Field>>>,
 ) -> HashSet<String> {
     let mut existing_field_names = HashSet::new();
@@ -219,8 +215,8 @@ pub fn get_existing_field_names(
 
 pub fn override_existing_timestamp_fields(
     existing_schema: &HashMap<String, Arc<Field>>,
-    inferred_schema: Arc<Schema>,
-) -> Arc<Schema> {
+    inferred_schema: &mut Schema,
+) {
     let timestamp_field_names: HashSet<String> = existing_schema
         .values()
         .filter_map(|field| {
@@ -231,7 +227,8 @@ pub fn override_existing_timestamp_fields(
             }
         })
         .collect();
-    let updated_fields: Vec<Arc<Field>> = inferred_schema
+
+    inferred_schema.fields = inferred_schema
         .fields()
         .iter()
         .map(|field| {
@@ -246,28 +243,24 @@ pub fn override_existing_timestamp_fields(
             }
         })
         .collect();
-
-    Arc::new(Schema::new(updated_fields))
 }
 
 pub fn update_field_type_in_schema(
-    inferred_schema: Arc<Schema>,
+    inferred_schema: &mut Schema,
     existing_schema: Option<&HashMap<String, Arc<Field>>>,
     time_partition: Option<&String>,
-) -> Arc<Schema> {
-    let mut updated_schema = inferred_schema.clone();
-    let existing_field_names = get_existing_field_names(inferred_schema.clone(), existing_schema);
-
+) {
+    let existing_field_names = get_existing_field_names(inferred_schema, existing_schema);
     if let Some(existing_schema) = existing_schema {
         // overriding known timestamp fields which were inferred as string fields
-        updated_schema = override_existing_timestamp_fields(existing_schema, updated_schema);
+        override_existing_timestamp_fields(existing_schema, inferred_schema);
     }
 
     let Some(time_partition) = time_partition else {
-        return updated_schema;
+        return;
     };
 
-    let new_schema: Vec<Field> = updated_schema
+    inferred_schema.fields = inferred_schema
         .fields()
         .iter()
         .map(|field| {
@@ -283,5 +276,4 @@ pub fn update_field_type_in_schema(
             }
         })
         .collect();
-    Arc::new(Schema::new(new_schema))
 }

src/handlers/http/logstream.rs

@@ -116,7 +116,7 @@ pub async fn detect_schema(Json(json): Json<Value>) -> Result<impl Responder, St
 
     let mut schema = infer_json_schema_from_iterator(log_records.iter().map(Ok)).unwrap();
     for log_record in log_records.iter() {
-        schema = override_inferred_data_type(schema, log_record, SchemaVersion::V1);
+        override_inferred_data_type(&mut schema, log_record, SchemaVersion::V1);
     }
 
     Ok((web::Json(schema), StatusCode::OK))