feat: enhancement to allow custom partitions in ingestion (#795)

This PR allows custom partitions in ingestion to allow
adding an optional header - X-P-Custom-Partition in stream 
creation API with comma separated values the custom-partition 
will be stored in the stream.json file at the time of ingestion, below 
validations are put in place -

1. all events in the batch have the custom partitions
2. max of 3 partitions are allowed
3. custom partition values should not have a '.' in the log event
4. custom partition value should not be empty in the log event

Prefixes will be created in the alphabetical order of the custom partition fields

sort order in the parquet will be p_timestamp/time_partition and the custom 
partition fields same sort order will be reflected in the manifest file

Author: nikhilsinhaparseable

server/src/event.rs

@@ -23,13 +23,13 @@ mod writer;
 use arrow_array::RecordBatch;
 use arrow_schema::{Field, Fields, Schema};
 use itertools::Itertools;
-
 use std::sync::Arc;
 
 use self::error::EventError;
 pub use self::writer::STREAM_WRITERS;
 use crate::{handlers::http::ingest::PostError, metadata};
 use chrono::NaiveDateTime;
+use std::collections::HashMap;
 
 pub const DEFAULT_TIMESTAMP_KEY: &str = "p_timestamp";
 pub const DEFAULT_TAGS_KEY: &str = "p_tags";
@@ -44,6 +44,7 @@ pub struct Event {
     pub is_first_event: bool,
     pub parsed_timestamp: NaiveDateTime,
     pub time_partition: Option<String>,
+    pub custom_partition_values: HashMap<String, String>,
 }
 
 // Events holds the schema related to a each event for a single log stream
@@ -55,6 +56,14 @@ impl Event {
             key = format!("{key}{parsed_timestamp_to_min}");
         }
 
+        if !self.custom_partition_values.is_empty() {
+            let mut custom_partition_key = String::default();
+            for (k, v) in self.custom_partition_values.iter().sorted_by_key(|v| v.0) {
+                custom_partition_key = format!("{custom_partition_key}&{k}={v}");
+            }
+            key = format!("{key}{custom_partition_key}");
+        }
+
         let num_rows = self.rb.num_rows() as u64;
         if self.is_first_event {
             commit_schema(&self.stream_name, self.rb.schema())?;
@@ -65,6 +74,7 @@ impl Event {
             &key,
             self.rb.clone(),
             self.parsed_timestamp,
+            self.custom_partition_values,
         )?;
 
         metadata::STREAM_INFO.update_stats(
@@ -111,8 +121,15 @@ impl Event {
         schema_key: &str,
         rb: RecordBatch,
         parsed_timestamp: NaiveDateTime,
+        custom_partition_values: HashMap<String, String>,
     ) -> Result<(), EventError> {
-        STREAM_WRITERS.append_to_local(stream_name, schema_key, rb, parsed_timestamp)?;
+        STREAM_WRITERS.append_to_local(
+            stream_name,
+            schema_key,
+            rb,
+            parsed_timestamp,
+            custom_partition_values,
+        )?;
         Ok(())
     }
 }

server/src/event/format/json.rs

@@ -48,7 +48,7 @@ impl EventFormat for Event {
         static_schema_flag: Option<String>,
         time_partition: Option<String>,
     ) -> Result<(Self::Data, Vec<Arc<Field>>, bool, Tags, Metadata), anyhow::Error> {
-        let data = flatten_json_body(self.data, None, None, false)?;
+        let data = flatten_json_body(self.data, None, None, None, false)?;
         let stream_schema = schema;
 
         // incoming event may be a single json or a json array

server/src/event/writer.rs

@@ -53,6 +53,7 @@ impl Writer {
         schema_key: &str,
         rb: RecordBatch,
         parsed_timestamp: NaiveDateTime,
+        custom_partition_values: HashMap<String, String>,
     ) -> Result<(), StreamWriterError> {
         let rb = utils::arrow::replace_columns(
             rb.schema(),
@@ -61,8 +62,13 @@ impl Writer {
             &[Arc::new(get_timestamp_array(rb.num_rows()))],
         );
 
-        self.disk
-            .push(stream_name, schema_key, &rb, parsed_timestamp)?;
+        self.disk.push(
+            stream_name,
+            schema_key,
+            &rb,
+            parsed_timestamp,
+            custom_partition_values,
+        )?;
         self.mem.push(schema_key, rb);
         Ok(())
     }
@@ -84,6 +90,7 @@ impl WriterTable {
         schema_key: &str,
         record: RecordBatch,
         parsed_timestamp: NaiveDateTime,
+        custom_partition_values: HashMap<String, String>,
     ) -> Result<(), StreamWriterError> {
         let hashmap_guard = self.read().unwrap();
 
@@ -95,6 +102,7 @@ impl WriterTable {
                     schema_key,
                     record,
                     parsed_timestamp,
+                    custom_partition_values,
                 )?;
             }
             None => {
@@ -149,6 +157,7 @@ impl WriterTable {
                         schema_key,
                         record,
                         parsed_timestamp,
+                        custom_partition_values,
                     )?;
                 } else {
                     writer.lock().unwrap().push_mem(stream_name, record)?;
@@ -157,7 +166,13 @@ impl WriterTable {
             None => {
                 if CONFIG.parseable.mode != Mode::Query {
                     let mut writer = Writer::default();
-                    writer.push(stream_name, schema_key, record, parsed_timestamp)?;
+                    writer.push(
+                        stream_name,
+                        schema_key,
+                        record,
+                        parsed_timestamp,
+                        custom_partition_values,
+                    )?;
                     map.insert(stream_name.to_owned(), Mutex::new(writer));
                 } else {
                     let mut writer = Writer::default();

server/src/event/writer/file_writer.rs

@@ -45,6 +45,7 @@ impl FileWriter {
         schema_key: &str,
         record: &RecordBatch,
         parsed_timestamp: NaiveDateTime,
+        custom_partition_values: HashMap<String, String>,
     ) -> Result<(), StreamWriterError> {
         match self.get_mut(schema_key) {
             Some(writer) => {
@@ -56,8 +57,13 @@ impl FileWriter {
             // entry is not present thus we create it
             None => {
                 // this requires mutable borrow of the map so we drop this read lock and wait for write lock
-                let (path, writer) =
-                    init_new_stream_writer_file(stream_name, schema_key, record, parsed_timestamp)?;
+                let (path, writer) = init_new_stream_writer_file(
+                    stream_name,
+                    schema_key,
+                    record,
+                    parsed_timestamp,
+                    custom_partition_values,
+                )?;
                 self.insert(
                     schema_key.to_owned(),
                     ArrowWriter {
@@ -83,9 +89,10 @@ fn init_new_stream_writer_file(
     schema_key: &str,
     record: &RecordBatch,
     parsed_timestamp: NaiveDateTime,
+    custom_partition_values: HashMap<String, String>,
 ) -> Result<(PathBuf, StreamWriter<std::fs::File>), StreamWriterError> {
     let dir = StorageDir::new(stream_name);
-    let path = dir.path_by_current_time(schema_key, parsed_timestamp);
+    let path = dir.path_by_current_time(schema_key, parsed_timestamp, custom_partition_values);
     std::fs::create_dir_all(dir.data_path)?;
 
     let file = OpenOptions::new().create(true).append(true).open(&path)?;

server/src/handlers.rs

@@ -26,6 +26,7 @@ const STREAM_NAME_HEADER_KEY: &str = "x-p-stream";
 const LOG_SOURCE_KEY: &str = "x-p-log-source";
 const TIME_PARTITION_KEY: &str = "x-p-time-partition";
 const TIME_PARTITION_LIMIT_KEY: &str = "x-p-time-partition-limit";
+const CUSTOM_PARTITION_KEY: &str = "x-p-custom-partition";
 const STATIC_SCHEMA_FLAG: &str = "x-p-static-schema-flag";
 const AUTHORIZATION_KEY: &str = "authorization";
 const SEPARATOR: char = '^';