Merge branch 'main' into fix-timepartition

Author: Devdutt Shenoi

src/event/format/mod.rs

@@ -171,7 +171,6 @@ pub trait EventFormat: Sized {
             update_field_type_in_schema(new_schema, None, time_partition, None, schema_version);
 
         let rb = Self::decode(data, new_schema.clone())?;
-
         let rb = add_parseable_fields(rb, p_timestamp, p_custom_fields)?;
 
         Ok((rb, is_first))

src/handlers/http/ingest.rs

@@ -395,7 +395,6 @@ mod tests {
     use crate::{
         event::format::{json, EventFormat},
         metadata::SchemaVersion,
-        utils::json::{convert_array_to_object, flatten::convert_to_array},
     };
 
     trait TestExt {
@@ -560,21 +559,6 @@ mod tests {
         assert_eq!(rb.num_columns(), 1);
     }
 
-    #[test]
-    fn non_object_arr_is_err() {
-        let json = json!([1]);
-
-        assert!(convert_array_to_object(
-            json,
-            None,
-            None,
-            None,
-            SchemaVersion::V0,
-            &crate::event::format::LogSource::default()
-        )
-        .is_err())
-    }
-
     #[test]
     fn array_into_recordbatch_inffered_schema() {
         let json = json!([
@@ -770,28 +754,17 @@ mod tests {
             {
                 "a": 1,
                 "b": "hello",
-                "c": [{"a": 1}]
+                "c_a": [1],
             },
             {
                 "a": 1,
                 "b": "hello",
-                "c": [{"a": 1, "b": 2}]
+                "c_a": [1],
+                "c_b": [2],
             },
         ]);
-        let flattened_json = convert_to_array(
-            convert_array_to_object(
-                json,
-                None,
-                None,
-                None,
-                SchemaVersion::V0,
-                &crate::event::format::LogSource::default(),
-            )
-            .unwrap(),
-        )
-        .unwrap();
 
-        let (rb, _) = json::Event::new(flattened_json)
+        let (rb, _) = json::Event::new(json)
             .into_recordbatch(
                 &HashMap::default(),
                 false,
@@ -859,28 +832,17 @@ mod tests {
             {
                 "a": 1,
                 "b": "hello",
-                "c": [{"a": 1}]
+                "c_a": 1,
             },
             {
                 "a": 1,
                 "b": "hello",
-                "c": [{"a": 1, "b": 2}]
+                "c_a": 1,
+                "c_b": 2,
             },
         ]);
-        let flattened_json = convert_to_array(
-            convert_array_to_object(
-                json,
-                None,
-                None,
-                None,
-                SchemaVersion::V1,
-                &crate::event::format::LogSource::default(),
-            )
-            .unwrap(),
-        )
-        .unwrap();
 
-        let (rb, _) = json::Event::new(flattened_json)
+        let (rb, _) = json::Event::new(json)
             .into_recordbatch(
                 &HashMap::default(),
                 false,

src/parseable/streams.rs

@@ -67,6 +67,18 @@ use super::{
     LogStream, ARROW_FILE_EXTENSION,
 };
 
+/// Returns the filename for parquet if provided arrows file path is valid as per our expectation
+fn arrow_path_to_parquet(path: &Path, random_string: &str) -> Option<PathBuf> {
+    let filename = path.file_stem()?.to_str()?;
+    let (_, front) = filename.split_once('.')?;
+    assert!(front.contains('.'), "contains the delim `.`");
+    let filename_with_random_number = format!("{front}.{random_string}.parquet");
+    let mut parquet_path = path.to_owned();
+    parquet_path.set_file_name(filename_with_random_number);
+
+    Some(parquet_path)
+}
+
 #[derive(Debug, thiserror::Error)]
 #[error("Stream not found: {0}")]
 pub struct StreamNotFound(pub String);
@@ -182,7 +194,10 @@ impl Stream {
         let paths = dir
             .flatten()
             .map(|file| file.path())
-            .filter(|file| file.extension().is_some_and(|ext| ext.eq("arrows")))
+            .filter(|file| {
+                file.extension()
+                    .is_some_and(|ext| ext.eq(ARROW_FILE_EXTENSION))
+            })
             .sorted_by_key(|f| f.metadata().unwrap().modified().unwrap())
             .collect();
 
@@ -225,12 +240,13 @@ impl Stream {
                     &arrow_file_path, self.stream_name
                 );
                 remove_file(&arrow_file_path).unwrap();
-            } else {
-                let key = Self::arrow_path_to_parquet(&arrow_file_path, &random_string);
+            } else if let Some(key) = arrow_path_to_parquet(&arrow_file_path, &random_string) {
                 grouped_arrow_file
                     .entry(key)
                     .or_default()
                     .push(arrow_file_path);
+            } else {
+                warn!("Unexpected arrows file: {}", arrow_file_path.display());
             }
         }
         grouped_arrow_file
@@ -289,17 +305,6 @@ impl Stream {
         }
     }
 
-    fn arrow_path_to_parquet(path: &Path, random_string: &str) -> PathBuf {
-        let filename = path.file_stem().unwrap().to_str().unwrap();
-        let (_, filename) = filename.split_once('.').unwrap();
-        assert!(filename.contains('.'), "contains the delim `.`");
-        let filename_with_random_number = format!("{filename}.{random_string}.arrows");
-        let mut parquet_path = path.to_owned();
-        parquet_path.set_file_name(filename_with_random_number);
-        parquet_path.set_extension("parquet");
-        parquet_path
-    }
-
     /// Converts arrow files in staging into parquet files, does so only for past minutes when run with `!shutdown_signal`
     pub fn prepare_parquet(&self, shutdown_signal: bool) -> Result<(), StagingError> {
         info!(
@@ -831,7 +836,7 @@ impl Streams {
 
 #[cfg(test)]
 mod tests {
-    use std::{sync::Barrier, thread::spawn, time::Duration};
+    use std::{io::Write, sync::Barrier, thread::spawn, time::Duration};
 
     use arrow_array::{Int32Array, StringArray, TimestampMillisecondArray};
     use arrow_schema::{DataType, Field, TimeUnit};
@@ -1207,6 +1212,57 @@ mod tests {
         assert_eq!(staging.arrow_files().len(), 1);
     }
 
+    fn create_test_file(dir: &TempDir, filename: &str) -> PathBuf {
+        let file_path = dir.path().join(filename);
+        let mut file = File::create(&file_path).expect("Failed to create test file");
+        // Write some dummy content
+        file.write_all(b"test content")
+            .expect("Failed to write to test file");
+        file_path
+    }
+
+    #[test]
+    fn test_valid_arrow_path_conversion() {
+        let temp_dir = TempDir::new().expect("Failed to create temp dir");
+        let filename = "12345abcde&key1=value1.date=2020-01-21.hour=10.minute=30.key1=value1.key2=value2.ee529ffc8e76.data.arrows";
+        let file_path = create_test_file(&temp_dir, filename);
+        let random_string = "random123";
+
+        let result = arrow_path_to_parquet(&file_path, random_string);
+
+        assert!(result.is_some());
+        let parquet_path = result.unwrap();
+        assert_eq!(
+                parquet_path.file_name().unwrap().to_str().unwrap(),
+                "date=2020-01-21.hour=10.minute=30.key1=value1.key2=value2.ee529ffc8e76.data.random123.parquet"
+            );
+    }
+
+    #[test]
+    fn test_complex_path() {
+        let temp_dir = TempDir::new().expect("Failed to create temp dir");
+        let nested_dir = temp_dir.path().join("nested/directory/structure");
+        std::fs::create_dir_all(&nested_dir).expect("Failed to create nested directories");
+
+        let filename = "20200201T1830f8a5fc1edc567d56&key1=value1&key2=value2.date=2020-01-21.hour=10.minute=30.region=us-west.ee529ffc8e76.data.arrows";
+        let file_path = nested_dir.join(filename);
+
+        let mut file = File::create(&file_path).expect("Failed to create test file");
+        file.write_all(b"test content")
+            .expect("Failed to write to test file");
+
+        let random_string = "random456";
+
+        let result = arrow_path_to_parquet(&file_path, random_string);
+
+        assert!(result.is_some());
+        let parquet_path = result.unwrap();
+        assert_eq!(
+            parquet_path.file_name().unwrap().to_str().unwrap(),
+            "date=2020-01-21.hour=10.minute=30.region=us-west.ee529ffc8e76.data.random456.parquet"
+        );
+    }
+
     #[test]
     fn get_or_create_returns_existing_stream() {
         let streams = Streams::default();

src/utils/arrow/mod.rs

@@ -17,37 +17,12 @@
  *
  */
 
-//! example function for concat recordbatch(may not work)
-//! ```rust
-//! # use arrow::record_batch::RecordBatch;
-//! # use arrow::error::Result;
-//!
-//! fn concat_batches(batch1: RecordBatch, batch2: RecordBatch) -> Result<RecordBatch> {
-//!     let schema = batch1.schema();
-//!     let columns = schema
-//!         .fields()
-//!         .iter()
-//!         .enumerate()
-//!         .map(|(i, _)| -> Result<_> {
-//!             let array1 = batch1.column(i);
-//!             let array2 = batch2.column(i);
-//!             let array = arrow::compute::concat(&[array1.as_ref(), array2.as_ref()])?;
-//!             Ok(array)
-//!         })
-//!         .collect::<Result<Vec<_>>>()?;
-//!
-//!     RecordBatch::try_new(schema.clone(), columns)
-//! }
-//! ```
-
 use std::{
     collections::{HashMap, HashSet},
     sync::Arc,
 };
 
-use arrow_array::{
-    Array, ArrayRef, RecordBatch, StringArray, TimestampMillisecondArray, UInt64Array,
-};
+use arrow_array::{ArrayRef, RecordBatch, StringArray, TimestampMillisecondArray, UInt64Array};
 use arrow_schema::{ArrowError, DataType, Field, Schema, TimeUnit};
 use arrow_select::take::take;
 use chrono::{DateTime, Utc};
@@ -62,31 +37,6 @@ use serde_json::{Map, Value};
 
 use crate::event::DEFAULT_TIMESTAMP_KEY;
 
-/// Replaces columns in a record batch with new arrays.
-///
-/// # Arguments
-///
-/// * `schema` - The schema of the record batch.
-/// * `batch` - The record batch to modify.
-/// * `indexes` - The indexes of the columns to replace.
-/// * `arrays` - The new arrays to replace the columns with.
-///
-/// # Returns
-///
-/// The modified record batch with the columns replaced.
-pub fn replace_columns(
-    schema: Arc<Schema>,
-    batch: &RecordBatch,
-    indexes: &[usize],
-    arrays: &[Arc<dyn Array + 'static>],
-) -> RecordBatch {
-    let mut batch_arrays = batch.columns().iter().map(Arc::clone).collect_vec();
-    for (&index, arr) in indexes.iter().zip(arrays.iter()) {
-        batch_arrays[index] = Arc::clone(arr);
-    }
-    RecordBatch::try_new(schema, batch_arrays).unwrap()
-}
-
 /// Converts a slice of record batches to JSON.
 ///
 /// # Arguments
@@ -213,40 +163,11 @@ pub fn reverse(rb: &RecordBatch) -> RecordBatch {
 mod tests {
     use std::sync::Arc;
 
-    use arrow_array::{Array, Int32Array, RecordBatch};
-    use arrow_schema::{DataType, Field, Schema};
+    use arrow_array::RecordBatch;
+    use arrow_schema::Schema;
 
     use super::*;
 
-    #[test]
-    fn check_replace() {
-        let schema = Schema::new(vec![
-            Field::new("a", DataType::Int32, false),
-            Field::new("b", DataType::Int32, false),
-            Field::new("c", DataType::Int32, false),
-        ]);
-
-        let schema_ref = Arc::new(schema);
-
-        let rb = RecordBatch::try_new(
-            schema_ref.clone(),
-            vec![
-                Arc::new(Int32Array::from_value(0, 3)),
-                Arc::new(Int32Array::from_value(0, 3)),
-                Arc::new(Int32Array::from_value(0, 3)),
-            ],
-        )
-        .unwrap();
-
-        let arr: Arc<dyn Array + 'static> = Arc::new(Int32Array::from_value(0, 3));
-
-        let new_rb = replace_columns(schema_ref.clone(), &rb, &[2], &[arr]);
-
-        assert_eq!(new_rb.schema(), schema_ref);
-        assert_eq!(new_rb.num_columns(), 3);
-        assert_eq!(new_rb.num_rows(), 3)
-    }
-
     #[test]
     fn check_empty_json_to_record_batches() {
         let r = RecordBatch::new_empty(Arc::new(Schema::empty()));