map schema keys to recordbatches

Author: Devdutt Shenoi

src/event/format/json.rs

@@ -39,7 +39,7 @@ use tracing::error;
 
 use super::{EventFormat, LogSource};
 use crate::{
-    event::PartitionEvent,
+    event::{get_schema_key, PartitionEvent},
     kinesis::{flatten_kinesis_logs, Message},
     metadata::SchemaVersion,
     otel::{logs::flatten_otel_logs, metrics::flatten_otel_metrics, traces::flatten_otel_traces},
@@ -246,7 +246,7 @@ impl EventFormat for Event {
             log_source,
         )?;
 
-        let mut partitions = vec![];
+        let mut partitions = HashMap::new();
         for json in data {
             let custom_partition_values = match custom_partitions.as_ref() {
                 Some(custom_partitions) => {
@@ -271,11 +271,25 @@ impl EventFormat for Event {
                 schema_version,
             )?;
 
-            partitions.push(PartitionEvent {
-                rb,
-                parsed_timestamp,
-                custom_partition_values,
-            });
+            let schema = rb.schema();
+            let mut key = get_schema_key(&schema.fields);
+            if time_partition.is_some() {
+                let parsed_timestamp_to_min = parsed_timestamp.format("%Y%m%dT%H%M").to_string();
+                key.push_str(&parsed_timestamp_to_min);
+            }
+
+            for (k, v) in custom_partition_values.iter().sorted_by_key(|v| v.0) {
+                key.push_str(&format!("&{k}={v}"));
+            }
+
+            partitions.insert(
+                key,
+                PartitionEvent {
+                    rb,
+                    parsed_timestamp,
+                    custom_partition_values,
+                },
+            );
         }
 
         Ok(super::Event {

src/event/mod.rs

@@ -42,29 +42,14 @@ pub struct Event {
     pub origin_size: u64,
     pub is_first_event: bool,
     pub time_partition: Option<String>,
-    pub partitions: Vec<PartitionEvent>,
+    pub partitions: HashMap<String, PartitionEvent>,
     pub stream_type: StreamType,
 }
 
 // Events holds the schema related to a each event for a single log stream
 impl Event {
     pub fn process(self, stream: &Stream) -> Result<(), EventError> {
-        for partition in self.partitions {
-            let mut key = get_schema_key(&partition.rb.schema().fields);
-            if self.time_partition.is_some() {
-                let parsed_timestamp_to_min =
-                    partition.parsed_timestamp.format("%Y%m%dT%H%M").to_string();
-                key.push_str(&parsed_timestamp_to_min);
-            }
-
-            for (k, v) in partition
-                .custom_partition_values
-                .iter()
-                .sorted_by_key(|v| v.0)
-            {
-                key.push_str(&format!("&{k}={v}"));
-            }
-
+        for (key, partition) in self.partitions {
             if self.is_first_event {
                 let schema = partition.rb.schema().as_ref().clone();
                 stream.commit_schema(schema)?;
@@ -92,11 +77,9 @@ impl Event {
     }
 
     pub fn process_unchecked(&self, stream: &Stream) -> Result<(), EventError> {
-        for partition in &self.partitions {
-            let key = get_schema_key(&partition.rb.schema().fields);
-
+        for (key, partition) in &self.partitions {
             stream.push(
-                &key,
+                key,
                 &partition.rb,
                 partition.parsed_timestamp,
                 &partition.custom_partition_values,

src/handlers/http/ingest.rs

@@ -28,7 +28,7 @@ use serde_json::Value;
 
 use crate::event::error::EventError;
 use crate::event::format::LogSource;
-use crate::event::{self, PartitionEvent};
+use crate::event::{self, get_schema_key, PartitionEvent};
 use crate::handlers::{LOG_SOURCE_KEY, STREAM_NAME_HEADER_KEY};
 use crate::option::Mode;
 use crate::parseable::{Stream, StreamNotFound, PARSEABLE};
@@ -239,21 +239,26 @@ pub async fn post_event(
 }
 
 pub async fn push_logs_unchecked(
-    batches: RecordBatch,
+    rb: RecordBatch,
     stream: &Stream,
 ) -> Result<event::Event, PostError> {
-    let unchecked_event = event::Event {
+    let mut unchecked_event = event::Event {
         origin_format: "json",
         origin_size: 0,
         time_partition: None,
         is_first_event: true, // NOTE: Maybe should be false
-        partitions: vec![PartitionEvent {
-            rb: batches,
-            parsed_timestamp: Utc::now().naive_utc(),
-            custom_partition_values: HashMap::new(), // should be an empty map for unchecked push
-        }],
+        partitions: HashMap::new(),
         stream_type: StreamType::UserDefined,
     };
+    unchecked_event.partitions.insert(
+        get_schema_key(&rb.schema().fields),
+        PartitionEvent {
+            rb,
+            parsed_timestamp: Utc::now().naive_utc(),
+            custom_partition_values: HashMap::new(), // should be an empty map for unchecked push
+        },
+    );
+
     unchecked_event.process_unchecked(stream)?;
 
     Ok(unchecked_event)