sync thresholds in env, stats sync on shutdown, streaming response

nikhilsinhaparseable · nikhilsinhaparseable · commit 43926be8e54f · 2025-06-18T06:45:48.000-04:00
diff --git a/src/cli.rs b/src/cli.rs
@@ -406,6 +406,23 @@ pub struct Options {
         help = "Enable/Disable collecting dataset stats"
     )]
     pub collect_dataset_stats: bool,
+
+    // the duration during which local sync should be completed
+    #[arg(
+        long,
+        env = "P_LOCAL_SYNC_THRESHOLD",
+        default_value = "30",
+        help = "Local sync threshold in seconds"
+    )]
+    pub local_sync_threshold: u64,
+    // the duration during which object store sync should be completed
+    #[arg(
+        long,
+        env = "P_OBJECT_STORE_SYNC_THRESHOLD",
+        default_value = "60",
+        help = "Object store sync threshold in seconds"
+    )]
+    pub object_store_sync_threshold: u64,
 }
 
 #[derive(Parser, Debug)]
diff --git a/src/handlers/http/health_check.rs b/src/handlers/http/health_check.rs
@@ -56,12 +56,35 @@ pub async fn check_shutdown_middleware(
 
 // This function is called when the server is shutting down
 pub async fn shutdown() {
-    // Set the shutdown flag to true
-    let mut shutdown_flag = SIGNAL_RECEIVED.lock().await;
-    *shutdown_flag = true;
+    // Set shutdown flag to true
+    set_shutdown_flag().await;
 
     //sleep for 5 secs to allow any ongoing requests to finish
     tokio::time::sleep(std::time::Duration::from_secs(5)).await;
+
+    // Perform sync operations
+    perform_sync_operations().await;
+
+    // If collect_dataset_stats is enabled, perform sync operations
+    // This is to ensure that all stats data is synced before the server shuts down
+    if PARSEABLE.options.collect_dataset_stats {
+        perform_sync_operations().await;
+    }
+}
+
+async fn set_shutdown_flag() {
+    let mut shutdown_flag = SIGNAL_RECEIVED.lock().await;
+    *shutdown_flag = true;
+}
+
+async fn perform_sync_operations() {
+    // Perform local sync
+    perform_local_sync().await;
+    // Perform object store sync
+    perform_object_store_sync().await;
+}
+
+async fn perform_local_sync() {
     let mut local_sync_joinset = JoinSet::new();
 
     // Sync staging
@@ -76,7 +99,9 @@ pub async fn shutdown() {
             Err(err) => error!("Failed to join async task: {err}"),
         }
     }
+}
 
+async fn perform_object_store_sync() {
     // Sync object store
     let mut object_store_joinset = JoinSet::new();
     sync_all_streams(&mut object_store_joinset);
diff --git a/src/lib.rs b/src/lib.rs
@@ -72,12 +72,6 @@ pub const OBJECT_STORE_DATA_GRANULARITY: u32 = LOCAL_SYNC_INTERVAL.as_secs() as
 /// Describes the duration at the end of which parquets are pushed into objectstore.
 pub const STORAGE_UPLOAD_INTERVAL: Duration = Duration::from_secs(30);
 
-/// Describes the duration during which local sync should be completed
-pub const LOCAL_SYNC_THRESHOLD: Duration = Duration::from_secs(30); // 30 secs
-
-/// Describes the duration during which object store sync should be completed
-pub const OBJECT_STORE_SYNC_THRESHOLD: Duration = Duration::from_secs(15); // 15 secs
-
 // A single HTTP client for all outgoing HTTP requests from the parseable server
 pub static HTTP_CLIENT: Lazy<Client> = Lazy::new(|| {
     ClientBuilder::new()
diff --git a/src/parseable/mod.rs b/src/parseable/mod.rs
@@ -362,6 +362,11 @@ impl Parseable {
             return Ok(true);
         }
 
+        // validate custom partition if provided
+        if let Some(partition) = custom_partition {
+            validate_custom_partition(partition)?;
+        }
+
         // For distributed deployments, if the stream not found in memory map,
         //check if it exists in the storage
         //create stream and schema from storage
diff --git a/src/storage/object_storage.rs b/src/storage/object_storage.rs
@@ -830,16 +830,6 @@ pub trait ObjectStorage: Debug + Send + Sync + 'static {
         let custom_partition = stream.get_custom_partition();
         let schema = stream.get_schema();
         for path in stream.parquet_files() {
-            if stream.get_stream_type() != StreamType::Internal
-                && PARSEABLE.options.collect_dataset_stats
-            {
-                if let Err(err) = calculate_field_stats(stream_name, &path, &schema).await {
-                    warn!(
-                        "Error calculating field stats for stream {}: {}",
-                        stream_name, err
-                    );
-                }
-            }
             let filename = path
                 .file_name()
                 .expect("only parquet files are returned by iterator")
@@ -889,6 +879,18 @@ pub trait ObjectStorage: Debug + Send + Sync + 'static {
             let manifest = catalog::create_from_parquet_file(absolute_path.clone(), &path)?;
             catalog::update_snapshot(store, stream_name, manifest).await?;
 
+            // If the stream is not internal and stats collection is enabled, calculate field stats
+            // before removing the parquet file
+            if stream.get_stream_type() != StreamType::Internal
+                && PARSEABLE.options.collect_dataset_stats
+            {
+                if let Err(err) = calculate_field_stats(stream_name, &path, &schema).await {
+                    warn!(
+                        "Error calculating field stats for stream {}: {}",
+                        stream_name, err
+                    );
+                }
+            }
             if let Err(e) = remove_file(path) {
                 warn!("Failed to remove staged file: {e}");
             }
@@ -968,7 +970,13 @@ async fn calculate_field_stats(
         )
         .await?;
     let ctx = SessionContext::new_with_state(QUERY_SESSION_STATE.clone());
-    let ctx_table_name = format!("{}_{}", stream_name, parquet_path.display());
+    let parquet_file_name = parquet_path
+        .file_name()
+        .expect("only parquet files are returned by iterator")
+        .to_str()
+        .expect("filename is valid string");
+    let parquet_file_name = str::replace(parquet_file_name, ".", "_");
+    let ctx_table_name = format!("{}_{}", stream_name, parquet_file_name);
     ctx.register_parquet(
         &ctx_table_name,
         parquet_path.to_str().expect("valid path"),
@@ -1062,81 +1070,118 @@ async fn calculate_single_field_stats(
 /// This is used for fetching record count for a field and distinct count.
 async fn query_single_i64(ctx: &SessionContext, sql: &str) -> Option<i64> {
     let df = ctx.sql(sql).await.ok()?;
-    let mut stream = df.execute_stream().await.ok()?;
-    let mut count = 0;
-    while let Some(batch_result) = stream.next().await {
-        let batch = batch_result.ok()?;
-        if batch.num_rows() == 0 {
-            return None;
-        }
-        let array = batch.column(0).as_any().downcast_ref::<Int64Array>()?;
-        count += array.value(0);
+    let batches = df.collect().await.ok()?;
+    let batch = batches.first()?;
+    if batch.num_rows() == 0 {
+        return None;
     }
-    Some(count)
+    let array = batch.column(0).as_any().downcast_ref::<Int64Array>()?;
+
+    Some(array.value(0))
+}
+
+macro_rules! try_downcast {
+    ($ty:ty, $arr:expr, $body:expr) => {
+        if let Some(arr) = $arr.as_any().downcast_ref::<$ty>() {
+            $body(arr)
+        } else {
+            warn!(
+                "Expected {} for {:?}, but found {:?}",
+                stringify!($ty),
+                $arr.data_type(),
+                $arr.data_type()
+            );
+            "UNSUPPORTED".to_string()
+        }
+    };
 }
 
-/// Helper function to format an Arrow value at a given index into a string.
-/// Handles null values and different data types like String, Int64, Float64, Timestamp, Date32, and Boolean.
+/// Function to format an Arrow value at a given index into a string.
+/// Handles null values and different data types by downcasting the array to the appropriate type.
 fn format_arrow_value(array: &dyn Array, idx: usize) -> String {
     if array.is_null(idx) {
         return "NULL".to_string();
     }
 
     match array.data_type() {
-        DataType::Utf8 => array
-            .as_any()
-            .downcast_ref::<StringArray>()
-            .unwrap()
+        DataType::Utf8 => try_downcast!(StringArray, array, |arr: &StringArray| arr
             .value(idx)
-            .to_string(),
-        DataType::Utf8View => array
-            .as_any()
-            .downcast_ref::<StringViewArray>()
-            .unwrap()
+            .to_string()),
+        DataType::Utf8View => try_downcast!(StringViewArray, array, |arr: &StringViewArray| arr
             .value(idx)
-            .to_string(),
-        DataType::Binary => {
-            let arr = array.as_any().downcast_ref::<BinaryArray>().unwrap();
+            .to_string()),
+        DataType::Binary => try_downcast!(BinaryArray, array, |arr: &BinaryArray| {
             String::from_utf8_lossy(arr.value(idx)).to_string()
-        }
-        DataType::BinaryView => {
-            let arr = array.as_any().downcast_ref::<BinaryViewArray>().unwrap();
+        }),
+        DataType::BinaryView => try_downcast!(BinaryViewArray, array, |arr: &BinaryViewArray| {
             String::from_utf8_lossy(arr.value(idx)).to_string()
-        }
-        DataType::Int64 => array
-            .as_any()
-            .downcast_ref::<Int64Array>()
-            .unwrap()
+        }),
+        DataType::Int64 => try_downcast!(Int64Array, array, |arr: &Int64Array| arr
             .value(idx)
-            .to_string(),
-        DataType::Float64 => array
-            .as_any()
-            .downcast_ref::<Float64Array>()
-            .unwrap()
+            .to_string()),
+        DataType::Int32 => try_downcast!(
+            arrow_array::Int32Array,
+            array,
+            |arr: &arrow_array::Int32Array| arr.value(idx).to_string()
+        ),
+        DataType::Int16 => try_downcast!(
+            arrow_array::Int16Array,
+            array,
+            |arr: &arrow_array::Int16Array| arr.value(idx).to_string()
+        ),
+        DataType::Int8 => try_downcast!(
+            arrow_array::Int8Array,
+            array,
+            |arr: &arrow_array::Int8Array| arr.value(idx).to_string()
+        ),
+        DataType::UInt64 => try_downcast!(
+            arrow_array::UInt64Array,
+            array,
+            |arr: &arrow_array::UInt64Array| arr.value(idx).to_string()
+        ),
+        DataType::UInt32 => try_downcast!(
+            arrow_array::UInt32Array,
+            array,
+            |arr: &arrow_array::UInt32Array| arr.value(idx).to_string()
+        ),
+        DataType::UInt16 => try_downcast!(
+            arrow_array::UInt16Array,
+            array,
+            |arr: &arrow_array::UInt16Array| arr.value(idx).to_string()
+        ),
+        DataType::UInt8 => try_downcast!(
+            arrow_array::UInt8Array,
+            array,
+            |arr: &arrow_array::UInt8Array| arr.value(idx).to_string()
+        ),
+        DataType::Float64 => try_downcast!(Float64Array, array, |arr: &Float64Array| arr
             .value(idx)
-            .to_string(),
-        DataType::Timestamp(TimeUnit::Millisecond, _) => {
-            let arr = array
-                .as_any()
-                .downcast_ref::<TimestampMillisecondArray>()
-                .unwrap();
-            let timestamp = arr.value(idx);
-            DateTime::from_timestamp_millis(timestamp)
-                .map(|dt| dt.to_string())
-                .unwrap_or_else(|| "INVALID_TIMESTAMP".to_string())
-        }
-        DataType::Date32 => array
-            .as_any()
-            .downcast_ref::<Date32Array>()
-            .unwrap()
+            .to_string()),
+        DataType::Float32 => try_downcast!(
+            arrow_array::Float32Array,
+            array,
+            |arr: &arrow_array::Float32Array| arr.value(idx).to_string()
+        ),
+        DataType::Timestamp(TimeUnit::Millisecond, _) => try_downcast!(
+            TimestampMillisecondArray,
+            array,
+            |arr: &TimestampMillisecondArray| {
+                let timestamp = arr.value(idx);
+                chrono::DateTime::from_timestamp_millis(timestamp)
+                    .map(|dt| dt.to_string())
+                    .unwrap_or_else(|| "INVALID_TIMESTAMP".to_string())
+            }
+        ),
+        DataType::Date32 => try_downcast!(Date32Array, array, |arr: &Date32Array| arr
             .value(idx)
-            .to_string(),
-        DataType::Boolean => array
-            .as_any()
-            .downcast_ref::<BooleanArray>()
-            .unwrap()
+            .to_string()),
+        DataType::Boolean => try_downcast!(BooleanArray, array, |arr: &BooleanArray| if arr
             .value(idx)
-            .to_string(),
+        {
+            "true".to_string()
+        } else {
+            "false".to_string()
+        }),
         DataType::Null => "NULL".to_string(),
         _ => {
             warn!(
@@ -1162,20 +1207,20 @@ async fn query_distinct_stats(
     );
     let mut distinct_stats = Vec::new();
     if let Ok(df) = ctx.sql(&sql).await {
-        if let Ok(batches) = df.collect().await {
-            for rb in batches {
-                let Some(counts) = rb.column(0).as_any().downcast_ref::<Int64Array>() else {
-                    warn!("Unexpected type for count column in stats query");
-                    continue;
-                };
-                let values = rb.column(1).as_ref();
-                for i in 0..rb.num_rows() {
-                    let value = format_arrow_value(values, i);
-                    distinct_stats.push(DistinctStat {
-                        distinct_value: value,
-                        count: counts.value(i),
-                    });
-                }
+        let mut stream = df.execute_stream().await.expect("Failed to execute stream");
+        while let Some(batch_result) = stream.next().await {
+            let rb = batch_result.expect("Failed to execute stream");
+            let Some(counts) = rb.column(0).as_any().downcast_ref::<Int64Array>() else {
+                warn!("Unexpected type for count column in stats query");
+                continue;
+            };
+            let values = rb.column(1).as_ref();
+            for i in 0..rb.num_rows() {
+                let value = format_arrow_value(values, i);
+                distinct_stats.push(DistinctStat {
+                    distinct_value: value,
+                    count: counts.value(i),
+                });
             }
         }
     }
diff --git a/src/sync.rs b/src/sync.rs
