env to collect stats, refactor

Author: nikhilsinhaparseable

src/cli.rs

@@ -397,6 +397,15 @@ pub struct Options {
         help = "Maximum number of field statistics to store"
     )]
     pub max_field_statistics: usize,
+
+    // collect dataset stats
+    #[arg(
+        long,
+        env = "P_COLLECT_DATASET_STATS",
+        default_value = "false",
+        help = "Enable/Disable collecting dataset stats"
+    )]
+    pub collect_dataset_stats: bool,
 }
 
 #[derive(Parser, Debug)]

src/handlers/http/prism_home.rs

@@ -39,9 +39,7 @@ pub async fn home_api(req: HttpRequest) -> Result<impl Responder, PrismHomeError
     let query_map = web::Query::<HashMap<String, String>>::from_query(req.query_string())
         .map_err(|_| PrismHomeError::InvalidQueryParameter(HOME_QUERY_PARAM.to_string()))?;
 
-    let include_internal = query_map
-        .get(HOME_QUERY_PARAM)
-        .map_or(false, |v| v == "true");
+    let include_internal = query_map.get(HOME_QUERY_PARAM).is_some_and(|v| v == "true");
 
     let res = generate_home_response(&key, include_internal).await?;
 

src/parseable/streams.rs

@@ -518,7 +518,9 @@ impl Stream {
 
         if let Some(mut schema) = schema {
             // calculate field stats for all user defined streams
-            if self.get_stream_type() != StreamType::Internal {
+            if self.get_stream_type() != StreamType::Internal
+                && PARSEABLE.options.collect_dataset_stats
+            {
                 let stats_schema = schema.clone();
                 let stream_name = self.stream_name.clone();
                 let stats_rbs = rbs.clone();
@@ -1199,11 +1201,10 @@ async fn query_distinct_stats(
     if let Ok(df) = ctx.sql(&sql).await {
         if let Ok(batches) = df.collect().await {
             for rb in batches {
-                let counts = rb
-                    .column(0)
-                    .as_any()
-                    .downcast_ref::<Int64Array>()
-                    .expect("Counts should be Int64Array");
+                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);

src/prism/home/mod.rs

@@ -127,7 +127,7 @@ pub async fn generate_home_response(
                 if !include_internal
                     && metadata
                         .iter()
-                        .any(|m| m.stream_type == StreamType::Internal)
+                        .all(|m| m.stream_type == StreamType::Internal)
                 {
                     continue;
                 }