Config datafusion query with schema (#99)

Instead of pre listing all the valid prefixes before 
query execution and letting datafusion infer the 
schema, it is better to pass the schema we already 
have for a given stream and let datafusion do all 
the heavy lifting. In case stream info does not 
have a schema for a given stream then we return 
early with Error suggesting they need to post 
events to this logstream first.

Author: trueleo

server/src/query.rs

@@ -17,6 +17,7 @@
  */
 
 use chrono::{DateTime, Utc};
+use datafusion::arrow::datatypes::Schema;
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion::datasource::listing::ListingOptions;
@@ -44,6 +45,7 @@ fn get_value<'a>(value: &'a Value, key: &'static str) -> Result<&'a str, Error>
 pub struct Query {
     pub query: String,
     pub stream_name: String,
+    pub schema: Arc<Schema>,
     pub start: DateTime<Utc>,
     pub end: DateTime<Utc>,
 }
@@ -123,10 +125,25 @@ impl Query {
 #[cfg(test)]
 mod tests {
     use super::Query;
+    use crate::{alerts::Alerts, metadata::STREAM_INFO};
+    use datafusion::arrow::datatypes::Schema;
+    use datafusion::arrow::datatypes::{DataType, Field};
     use rstest::*;
     use serde_json::Value;
     use std::str::FromStr;
 
+    #[fixture]
+    fn schema() -> Schema {
+        let field_a = Field::new("a", DataType::Int64, false);
+        let field_b = Field::new("b", DataType::Boolean, false);
+        Schema::new(vec![field_a, field_b])
+    }
+
+    fn clear_map() {
+        STREAM_INFO.write().unwrap().clear();
+    }
+
+    // A query can only be performed on streams with a valid schema
     #[rstest]
     #[case(
         r#"{
@@ -144,12 +161,38 @@ mod tests {
         }"#,
         &["stream_name/date=2022-10-15/hour=10/minute=00/", "stream_name/date=2022-10-15/hour=10/minute=01/"]
     )]
-    fn query_parse_prefix(#[case] prefix: &str, #[case] right: &[&str]) {
+    #[serial_test::serial]
+    fn query_parse_prefix_with_some_schema(#[case] prefix: &str, #[case] right: &[&str]) {
+        clear_map();
+        STREAM_INFO
+            .add_stream("stream_name".to_string(), Some(schema()), Alerts::default())
+            .unwrap();
+
         let query = Value::from_str(prefix).unwrap();
         let query = Query::parse(query).unwrap();
         assert_eq!(&query.stream_name, "stream_name");
         let prefixes = query.get_prefixes();
         let left = prefixes.iter().map(String::as_str).collect::<Vec<&str>>();
         assert_eq!(left.as_slice(), right);
     }
+
+    // If there is no schema for this stream then parsing a Query should fail
+    #[rstest]
+    #[case(
+        r#"{
+            "query": "SELECT * FROM stream_name",
+            "startTime": "2022-10-15T10:00:00+00:00",
+            "endTime": "2022-10-15T10:01:00+00:00"
+        }"#
+    )]
+    #[serial_test::serial]
+    fn query_parse_prefix_with_no_schema(#[case] prefix: &str) {
+        clear_map();
+        STREAM_INFO
+            .add_stream("stream_name".to_string(), None, Alerts::default())
+            .unwrap();
+
+        let query = Value::from_str(prefix).unwrap();
+        assert!(Query::parse(query).is_err());
+    }
 }

server/src/s3.rs

@@ -44,8 +44,6 @@ const S3_URL_ENV_VAR: &str = "P_S3_URL";
 
 // max concurrent request allowed for datafusion object store
 const MAX_OBJECT_STORE_REQUESTS: usize = 1000;
-// max concurrent request allowed for prefix checks during listing
-const MAX_CONCURRENT_PREFIX_CHECK: usize = 1000;
 
 lazy_static::lazy_static! {
     #[derive(Debug)]
@@ -276,6 +274,7 @@ impl S3 {
         Ok(body_bytes)
     }
 
+    #[allow(dead_code)]
     async fn prefix_exists(&self, prefix: &str) -> Result<bool, AwsSdkError> {
         // TODO check if head object is faster compared to
         // list objects
@@ -418,42 +417,15 @@ impl ObjectStorage for S3 {
             SessionContext::with_config_rt(SessionConfig::default(), Arc::clone(&STORAGE_RUNTIME));
 
         // Get all prefix paths and convert them into futures which yeilds ListingTableUrl
-        let handles = query.get_prefixes().into_iter().map(|prefix| async move {
-            let t: Result<Option<ListingTableUrl>, ObjectStorageError> = {
-                if self.prefix_exists(&prefix).await? {
-                    let path = format!("s3://{}/{}", &S3_CONFIG.s3_bucket_name, prefix);
-                    Ok(Some(ListingTableUrl::parse(path)?))
-                } else {
-                    Ok(None)
-                }
-            };
-            t
-        });
-
-        // Poll futures but limit them to concurrency of MAX_CONCURRENT_PREFIX_CHECK
-        let list: Vec<_> = futures::stream::iter(handles)
-            .buffer_unordered(MAX_CONCURRENT_PREFIX_CHECK)
-            .collect()
-            .await;
-
-        // Collect all available prefixes
-        let prefixes: Vec<_> = list
+        let prefixes = query
+            .get_prefixes()
             .into_iter()
-            .filter_map(|fetch| {
-                fetch
-                    .map_err(|err| {
-                        log::error!("prefix lookup failed due to {}", err);
-                        err
-                    })
-                    .ok()
+            .map(|prefix| {
+                let path = format!("s3://{}/{}", &S3_CONFIG.s3_bucket_name, prefix);
+                ListingTableUrl::parse(path).unwrap()
             })
-            .flatten()
             .collect();
 
-        if prefixes.is_empty() {
-            return Ok(());
-        }
-
         let file_format = ParquetFormat::default().with_enable_pruning(true);
         let listing_options = ListingOptions {
             file_extension: format!("{}.data.parquet", hostname_unchecked()),
@@ -465,8 +437,7 @@ impl ObjectStorage for S3 {
 
         let config = ListingTableConfig::new_with_multi_paths(prefixes)
             .with_listing_options(listing_options)
-            .infer(&ctx.state())
-            .await?;
+            .with_schema(Arc::clone(&query.schema));
 
         let table = ListingTable::try_new(config)?;
         ctx.register_table(query.stream_name.as_str(), Arc::new(table))?;

server/src/validator.rs

@@ -16,10 +16,13 @@
  *
  */
 
+use std::sync::Arc;
+
 use chrono::{DateTime, Utc};
 use serde_json::json;
 
 use crate::alerts::Alerts;
+use crate::metadata::STREAM_INFO;
 use crate::query::Query;
 use crate::Error;
 
@@ -133,10 +136,18 @@ pub fn query(query: &str, start_time: &str, end_time: &str) -> Result<Query, Err
         return Err(Error::StartTimeAfterEndTime());
     }
 
+    let stream_name = tokens[stream_name_index].to_string();
+
+    let schema = match STREAM_INFO.schema(&stream_name)? {
+        Some(schema) => Arc::new(schema),
+        None => return Err(Error::MissingRecord),
+    };
+
     Ok(Query {
         stream_name: tokens[stream_name_index].to_string(),
         start,
         end,
         query: query.to_string(),
+        schema,
     })
 }