Add in-memory writer (#369)

This PR refactors the ingestion flow to allow optional,
memory only staging approach during ingestion. Memory
staging means many times faster ingestion at the risk of
loosing events (in memory) in case of a server crash / restart.
---------

Co-authored-by: Nitish Tiwari <[email protected]>

Author: trueleo

Cargo.lock

@@ -13,10 +13,10 @@ actix-cors = "0.6"
 actix-web-prometheus = { version = "0.1" }
 prometheus = { version = "0.13", features = ["process"] }
 anyhow = { version = "1.0", features = ["backtrace"] }
-arrow-schema = { version = "34.0.0", features = ["serde"] }
-arrow-array = { version = "34.0.0" }
-arrow-json = "34.0.0"
-arrow-ipc = "34.0.0"
+arrow-schema = { version = "36.0.0", features = ["serde"] }
+arrow-array = { version = "36.0.0" }
+arrow-json = "36.0.0"
+arrow-ipc = "36.0.0"
 async-trait = "0.1"
 base64 = "0.21"
 bytes = "1.4"
@@ -32,7 +32,7 @@ clap = { version = "4.1", default-features = false, features = [
   "error-context",
 ] }
 crossterm = "0.26"
-datafusion = "21.0.0"
+datafusion = "22.0.0"
 object_store = { version = "0.5.6", features = ["aws", "aws_profile"] }
 derive_more = "0.99"
 env_logger = "0.10"
@@ -70,7 +70,7 @@ xxhash-rust = { version = "0.8", features = ["xxh3"] }
 xz2 = { version = "*", features=["static"] }
 bzip2 = { version = "*", features=["static"] }
 once_cell = "1.17.1"
-parquet = "34.0.0"
+parquet = "36.0.0"
 pyroscope = { version = "0.5.3", optional = true }
 pyroscope_pprofrs = { version = "0.2", optional = true }
 uptime_lib = "0.2.2"

server/Cargo.toml

@@ -64,15 +64,15 @@ pub struct Alert {
 }
 
 impl Alert {
-    pub fn check_alert(&self, stream_name: String, events: RecordBatch) {
+    pub fn check_alert(&self, stream_name: &str, events: RecordBatch) {
         let resolves = self.rule.resolves(events.clone());
 
         for (index, state) in resolves.into_iter().enumerate() {
             match state {
                 AlertState::Listening | AlertState::Firing => (),
                 alert_state @ (AlertState::SetToFiring | AlertState::Resolved) => {
                     let context = self.get_context(
-                        stream_name.clone(),
+                        stream_name.to_owned(),
                         alert_state,
                         &self.rule,
                         events.slice(index, 1),

server/src/alerts/mod.rs

@@ -48,21 +48,25 @@ pub struct Event {
 impl Event {
     pub async fn process(self) -> Result<(), EventError> {
         let key = get_schema_key(&self.rb.schema().fields);
+        let num_rows = self.rb.num_rows() as u64;
 
         if self.is_first_event(metadata::STREAM_INFO.schema(&self.stream_name)?.as_ref()) {
             commit_schema(&self.stream_name, self.rb.schema())?;
         }
 
-        self.process_event(&key)?;
+        Self::process_event(&self.stream_name, &key, self.rb.clone())?;
 
         metadata::STREAM_INFO.update_stats(
             &self.stream_name,
             self.origin_format,
             self.origin_size,
-            self.rb.num_rows() as u64,
+            num_rows,
         )?;
 
-        if let Err(e) = metadata::STREAM_INFO.check_alerts(&self).await {
+        if let Err(e) = metadata::STREAM_INFO
+            .check_alerts(&self.stream_name, self.rb)
+            .await
+        {
             log::error!("Error checking for alerts. {:?}", e);
         }
 
@@ -90,8 +94,12 @@ impl Event {
 
     // event process all events after the 1st event. Concatenates record batches
     // and puts them in memory store for each event.
-    fn process_event(&self, schema_key: &str) -> Result<(), EventError> {
-        STREAM_WRITERS.append_to_local(&self.stream_name, schema_key, &self.rb)?;
+    fn process_event(
+        stream_name: &str,
+        schema_key: &str,
+        rb: RecordBatch,
+    ) -> Result<(), EventError> {
+        STREAM_WRITERS.append_to_local(stream_name, schema_key, rb)?;
         Ok(())
     }
 }

server/src/event.rs

@@ -214,6 +214,7 @@ fn valid_type(data_type: &DataType, value: &Value) -> bool {
                 false
             }
         }
+        DataType::Timestamp(_, _) => value.is_string() || value.is_number(),
         _ => unreachable!(),
     }
 }