fix: concat at once

Author: Devdutt Shenoi

src/event/format/json.rs

@@ -20,7 +20,6 @@
 #![allow(deprecated)]
 
 use anyhow::anyhow;
-use arrow::compute::concat_batches;
 use arrow_array::RecordBatch;
 use arrow_json::reader::{infer_json_schema_from_iterator, ReaderBuilder};
 use arrow_schema::{DataType, Field, Fields, Schema};
@@ -281,14 +280,13 @@ impl EventFormat for Event {
             }
 
             match partitions.get_mut(&key) {
-                Some(PartitionEvent { rb, .. }) => {
-                    *rb = concat_batches(&schema, [rb, &batch])?;
-                }
+                Some(PartitionEvent { rbs, .. }) => rbs.push(batch),
                 _ => {
                     partitions.insert(
                         key,
                         PartitionEvent {
-                            rb: batch,
+                            rbs: vec![batch],
+                            schema,
                             parsed_timestamp,
                             custom_partition_values,
                         },

src/event/mod.rs

@@ -19,20 +19,26 @@
 
 pub mod format;
 
+use arrow::compute::concat_batches;
 use arrow_array::RecordBatch;
-use arrow_schema::Field;
+use arrow_schema::{Field, Schema};
 use itertools::Itertools;
 use std::sync::Arc;
 
 use self::error::EventError;
-use crate::{metadata::update_stats, parseable::Stream, storage::StreamType};
+use crate::{
+    metadata::update_stats,
+    parseable::{StagingError, Stream},
+    storage::StreamType,
+};
 use chrono::NaiveDateTime;
 use std::collections::HashMap;
 
 pub const DEFAULT_TIMESTAMP_KEY: &str = "p_timestamp";
 
 pub struct PartitionEvent {
-    pub rb: RecordBatch,
+    pub rbs: Vec<RecordBatch>,
+    pub schema: Arc<Schema>,
     pub parsed_timestamp: NaiveDateTime,
     pub custom_partition_values: HashMap<String, String>,
 }
@@ -50,14 +56,15 @@ pub struct Event {
 impl Event {
     pub fn process(self, stream: &Stream) -> Result<(), EventError> {
         for (key, partition) in self.partitions {
+            let rb =
+                concat_batches(&partition.schema, &partition.rbs).map_err(StagingError::Arrow)?;
             if self.is_first_event {
-                let schema = partition.rb.schema().as_ref().clone();
-                stream.commit_schema(schema)?;
+                stream.commit_schema(partition.schema.as_ref().clone())?;
             }
 
             stream.push(
                 &key,
-                &partition.rb,
+                &rb,
                 partition.parsed_timestamp,
                 &partition.custom_partition_values,
                 self.stream_type,
@@ -67,20 +74,22 @@ impl Event {
                 &stream.stream_name,
                 self.origin_format,
                 self.origin_size,
-                partition.rb.num_rows(),
+                rb.num_rows(),
                 partition.parsed_timestamp.date(),
             );
 
-            crate::livetail::LIVETAIL.process(&stream.stream_name, &partition.rb);
+            crate::livetail::LIVETAIL.process(&stream.stream_name, &rb);
         }
         Ok(())
     }
 
     pub fn process_unchecked(&self, stream: &Stream) -> Result<(), EventError> {
         for (key, partition) in &self.partitions {
+            let rb =
+                concat_batches(&partition.schema, &partition.rbs).map_err(StagingError::Arrow)?;
             stream.push(
                 key,
-                &partition.rb,
+                &rb,
                 partition.parsed_timestamp,
                 &partition.custom_partition_values,
                 self.stream_type,

src/handlers/http/ingest.rs

@@ -236,18 +236,20 @@ pub async fn post_event(
 }
 
 pub async fn push_logs_unchecked(
-    rb: RecordBatch,
+    batch: RecordBatch,
     stream: &Stream,
 ) -> Result<event::Event, PostError> {
+    let schema = batch.schema();
     let unchecked_event = event::Event {
         origin_format: "json",
         origin_size: 0,
         time_partition: None,
         is_first_event: true, // NOTE: Maybe should be false
         partitions: [(
-            get_schema_key(&rb.schema().fields),
+            get_schema_key(&schema.fields),
             PartitionEvent {
-                rb,
+                rbs: vec![batch],
+                schema,
                 parsed_timestamp: Utc::now().naive_utc(),
                 custom_partition_values: HashMap::new(), // should be an empty map for unchecked push
             },