perf: batch staging arrow writes

  Batch disk staging writes per output filename before writing to
  Arrow IPC. This reduces per-request writer mutex hold time and cuts
  the number of DiskWriter::write calls during high-volume OTEL metrics
  ingest.

  Also keep targeted hotpath probes around ingest, JSON conversion,
  staging, and parquet conversion paths, and skip object-store sync work
  for streams without local parquet/schema files.

Author: parmesant

Cargo.toml

@@ -90,6 +90,14 @@ tokio = { version = "^1.43", default-features = false, features = [
 tokio-stream = { version = "0.1.17", features = ["fs"] }
 tokio-util = { version = "0.7" }
 
+# perf
+hotpath = { version = "0.16.0", features = [
+    "hotpath",
+    "hotpath-cpu",
+    "hotpath-alloc",
+    "tokio"
+] }
+
 # # Logging and Metrics
 # opentelemetry-proto = { version = "0.30.0", features = [
 #     "gen-tonic",

src/event/format/json.rs

@@ -60,6 +60,7 @@ impl EventFormat for Event {
 
     // convert the incoming json to a vector of json values
     // also extract the arrow schema, tags and metadata from the incoming json
+    #[hotpath::measure]
     fn to_data(
         self,
         schema: &HashMap<String, Arc<Field>>,
@@ -156,7 +157,7 @@ impl EventFormat for Event {
                     infer_schema.clone(),
                 ]).map_err(|err| anyhow!("Could not merge schema of this event with that of the existing stream. {:?}", err))?;
                 is_first = true;
-                let schema = infer_schema
+                let schema: Vec<Arc<Field>> = infer_schema
                     .fields
                     .iter()
                     .filter(|field| !field.data_type().is_null())
@@ -327,6 +328,10 @@ fn rename_json_keys(values: Vec<Value>) -> Result<Vec<Value>, anyhow::Error> {
         .into_iter()
         .map(|value| {
             if let Value::Object(map) = value {
+                if !map.keys().any(|key| key.starts_with('@')) {
+                    return Ok(Value::Object(map));
+                }
+
                 // Collect original keys to check for collisions
                 let original_keys: HashSet<String> = map.keys().cloned().collect();
 

src/event/format/mod.rs

@@ -162,6 +162,7 @@ pub trait EventFormat: Sized {
     /// Returns the UTC time at ingestion
     fn get_p_timestamp(&self) -> DateTime<Utc>;
 
+    #[hotpath::measure]
     fn into_recordbatch(
         self,
         storage_schema: &HashMap<String, Arc<Field>>,
@@ -606,6 +607,28 @@ pub fn rename_per_record_type_mismatches(
             let Value::Object(map) = value else {
                 return value;
             };
+            let needs_rename = map.iter().any(|(key, val)| {
+                if val.is_null() {
+                    return false;
+                }
+                let target_type = existing_schema
+                    .get(key)
+                    .map(|f| f.data_type())
+                    .or_else(|| inferred_types.get(key.as_str()).copied());
+                let Some(target_type) = target_type else {
+                    return false;
+                };
+                if (val.is_array() || val.is_object())
+                    && (target_type.is_list()
+                        || matches!(target_type, DataType::Struct(_) | DataType::Map(_, _)))
+                {
+                    return false;
+                }
+                !value_compatible_with_type(val, target_type, schema_version)
+            });
+            if !needs_rename {
+                return Value::Object(map);
+            }
             let new_map: serde_json::Map<String, Value> = map
                 .into_iter()
                 .map(|(key, val)| {

src/event/mod.rs

@@ -72,6 +72,7 @@ impl Event {
             is_first_event = self.is_first_event
         )
     )]
+    #[hotpath::measure]
     pub fn process(self) -> Result<(), EventError> {
         let mut key = get_schema_key(&self.rb.schema().fields);
         if self.time_partition.is_some() {

src/handlers/http/modal/utils/ingest_utils.rs

@@ -156,6 +156,7 @@ pub async fn flatten_and_push_logs(
     skip(json, log_source, p_custom_fields, time_partition, telemetry_type, tenant_id),
     fields(stream_name, record_count = tracing::field::Empty)
 )]
+#[hotpath::measure]
 pub fn push_logs(
     stream_name: &str,
     json: Value,

src/lib.rs

@@ -66,8 +66,8 @@ use once_cell::sync::Lazy;
 pub use openid;
 use parseable::PARSEABLE;
 use reqwest::{Client, ClientBuilder};
+pub use {hotpath, tracing_actix_web, tracing_opentelemetry, tracing_subscriber};
 pub use {opentelemetry, opentelemetry_otlp, opentelemetry_proto, opentelemetry_sdk};
-pub use {tracing_actix_web, tracing_opentelemetry, tracing_subscriber};
 
 // It is very unlikely that panic will occur when dealing with locks.
 pub const LOCK_EXPECT: &str = "Thread shouldn't panic while holding a lock";

src/otel/metrics.rs

@@ -640,6 +640,7 @@ fn process_resource_metrics<T, S, M>(
 
 /// this function performs the custom flattening of the otel metrics
 /// and returns a `Vec` of `Value::Object` of the flattened json
+#[hotpath::measure]
 pub fn flatten_otel_metrics(message: MetricsData, tenant_id: &str) -> Vec<Value> {
     process_resource_metrics(
         &message.resource_metrics,

src/parseable/staging/writer.rs

@@ -19,7 +19,7 @@
 
 use std::{
     collections::{HashMap, HashSet},
-    fs::{File, OpenOptions},
+    fs::{self, File, OpenOptions},
     io::BufWriter,
     path::PathBuf,
     sync::Arc,
@@ -41,10 +41,83 @@ use crate::{
 
 use super::StagingError;
 
+const DISK_WRITE_BATCH_ROWS: usize = 16_384;
+
 #[derive(Default)]
 pub struct Writer {
     pub mem: MemWriter<16384>,
     pub disk: HashMap<String, DiskWriter>,
+    disk_pending: HashMap<String, PendingDiskBatch>,
+}
+
+impl Writer {
+    #[hotpath::measure]
+    pub fn push_disk(
+        &mut self,
+        filename: String,
+        rb: &RecordBatch,
+        file_path: PathBuf,
+        range: TimeRange,
+    ) -> Result<(), StagingError> {
+        let pending = self.disk_pending.entry(filename.clone()).or_default();
+        pending.rows += rb.num_rows();
+        pending.range.get_or_insert(range);
+        pending.batches.push(rb.clone());
+
+        if pending.rows >= DISK_WRITE_BATCH_ROWS {
+            self.flush_pending_disk(&filename, file_path)?;
+        }
+
+        Ok(())
+    }
+
+    pub fn flush_pending_disk(
+        &mut self,
+        filename: &str,
+        file_path: PathBuf,
+    ) -> Result<(), StagingError> {
+        let Some(pending) = self.disk_pending.remove(filename) else {
+            return Ok(());
+        };
+        if pending.batches.is_empty() {
+            return Ok(());
+        }
+
+        let schema = pending.batches[0].schema();
+        let batch = concat_batches(&schema, pending.batches.iter())?;
+        match self.disk.get_mut(filename) {
+            Some(writer) => writer.write(&batch)?,
+            None => {
+                let range = pending.range.expect("pending disk batch must have range");
+                if let Some(parent) = file_path.parent() {
+                    fs::create_dir_all(parent)?;
+                }
+                let mut writer = DiskWriter::try_new(file_path, &schema, range)?;
+                writer.write(&batch)?;
+                self.disk.insert(filename.to_owned(), writer);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn flush_all_pending_disk(
+        &mut self,
+        data_path: &std::path::Path,
+    ) -> Result<(), StagingError> {
+        let filenames = self.disk_pending.keys().cloned().collect_vec();
+        for filename in filenames {
+            self.flush_pending_disk(&filename, data_path.join(&filename))?;
+        }
+        Ok(())
+    }
+}
+
+#[derive(Default)]
+struct PendingDiskBatch {
+    rows: usize,
+    batches: Vec<RecordBatch>,
+    range: Option<TimeRange>,
 }
 
 pub struct DiskWriter {
@@ -77,6 +150,7 @@ impl DiskWriter {
     }
 
     /// Write a single recordbatch into file
+    #[hotpath::measure]
     pub fn write(&mut self, rb: &RecordBatch) -> Result<(), StagingError> {
         self.inner.write(rb).map_err(StagingError::Arrow)
     }