recactor: query only expects time_partition

Author: Devdutt Shenoi

src/alerts/alerts_utils.rs

@@ -31,6 +31,7 @@ use tracing::trace;
 
 use crate::{
     alerts::AggregateCondition,
+    parseable::PARSEABLE,
     query::{TableScanVisitor, QUERY_SESSION},
     rbac::{
         map::SessionKey,
@@ -137,8 +138,9 @@ async fn execute_base_query(
         AlertError::CustomError(format!("Table name not found in query- {}", original_query))
     })?;
 
+    let time_partition = PARSEABLE.get_stream(&stream_name)?.get_time_partition();
     query
-        .get_dataframe(stream_name)
+        .get_dataframe(time_partition.as_ref())
         .await
         .map_err(|err| AlertError::CustomError(err.to_string()))
 }

src/alerts/mod.rs

@@ -37,7 +37,7 @@ use ulid::Ulid;
 pub mod alerts_utils;
 pub mod target;
 
-use crate::parseable::PARSEABLE;
+use crate::parseable::{StreamNotFound, PARSEABLE};
 use crate::query::{TableScanVisitor, QUERY_SESSION};
 use crate::rbac::map::SessionKey;
 use crate::storage;
@@ -514,17 +514,16 @@ impl AlertConfig {
 
         // for now proceed in a similar fashion as we do in query
         // TODO: in case of multiple table query does the selection of time partition make a difference? (especially when the tables don't have overlapping data)
-        let stream_name = if let Some(stream_name) = query.first_table_name() {
-            stream_name
-        } else {
+        let Some(stream_name) = query.first_table_name() else {
             return Err(AlertError::CustomError(format!(
                 "Table name not found in query- {}",
                 self.query
             )));
         };
 
+        let time_partition = PARSEABLE.get_stream(&stream_name)?.get_time_partition();
         let base_df = query
-            .get_dataframe(stream_name)
+            .get_dataframe(time_partition.as_ref())
             .await
             .map_err(|err| AlertError::CustomError(err.to_string()))?;
 
@@ -704,6 +703,8 @@ pub enum AlertError {
     CustomError(String),
     #[error("Invalid State Change: {0}")]
     InvalidStateChange(String),
+    #[error("{0}")]
+    StreamNotFound(#[from] StreamNotFound),
 }
 
 impl actix_web::ResponseError for AlertError {
@@ -717,6 +718,7 @@ impl actix_web::ResponseError for AlertError {
             Self::DatafusionError(_) => StatusCode::INTERNAL_SERVER_ERROR,
             Self::CustomError(_) => StatusCode::BAD_REQUEST,
             Self::InvalidStateChange(_) => StatusCode::BAD_REQUEST,
+            Self::StreamNotFound(_) => StatusCode::NOT_FOUND,
         }
     }
 

src/handlers/airplane.rs

@@ -216,11 +216,12 @@ impl FlightService for AirServiceImpl {
         })?;
         let time = Instant::now();
 
-        let stream = PARSEABLE
+        let time_partition = PARSEABLE
             .get_stream(&stream_name)
-            .map_err(|err| Status::internal(err.to_string()))?;
+            .map_err(|err| Status::internal(err.to_string()))?
+            .get_time_partition();
         let (records, _) = query
-            .execute(&stream)
+            .execute(time_partition.as_ref())
             .await
             .map_err(|err| Status::internal(err.to_string()))?;
 

src/handlers/http/query.rs

@@ -131,8 +131,8 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpRespons
         return Ok(HttpResponse::Ok().json(response));
     }
 
-    let stream = PARSEABLE.get_stream(&table_name)?;
-    let (records, fields) = query.execute(&stream).await?;
+    let time_partition = PARSEABLE.get_stream(&table_name)?.get_time_partition();
+    let (records, fields) = query.execute(time_partition.as_ref()).await?;
 
     let response = QueryResponse {
         records,
@@ -321,7 +321,7 @@ Description: {0}"#
     #[error("Error: {0}")]
     Anyhow(#[from] anyhow::Error),
     #[error("Error: {0}")]
-    StreamNotFound(#[from] StreamNotFound)
+    StreamNotFound(#[from] StreamNotFound),
 }
 
 impl actix_web::ResponseError for QueryError {

src/query/mod.rs

@@ -53,7 +53,7 @@ use crate::catalog::Snapshot as CatalogSnapshot;
 use crate::event;
 use crate::handlers::http::query::QueryError;
 use crate::option::Mode;
-use crate::parseable::{Stream, PARSEABLE};
+use crate::parseable::PARSEABLE;
 use crate::storage::{ObjectStorageProvider, ObjectStoreFormat, STREAM_ROOT_DIRECTORY};
 use crate::utils::time::TimeRange;
 
@@ -131,12 +131,10 @@ impl Query {
 
     pub async fn execute(
         &self,
-        stream: &Stream,
+        time_partition: Option<&String>,
     ) -> Result<(Vec<RecordBatch>, Vec<String>), ExecuteError> {
-        let time_partition = stream.get_time_partition();
-
         let df = QUERY_SESSION
-            .execute_logical_plan(self.final_logical_plan(&time_partition))
+            .execute_logical_plan(self.final_logical_plan(time_partition))
             .await?;
 
         let fields = df
@@ -155,18 +153,19 @@ impl Query {
         Ok((results, fields))
     }
 
-    pub async fn get_dataframe(&self, stream_name: String) -> Result<DataFrame, ExecuteError> {
-        let time_partition = PARSEABLE.get_stream(&stream_name)?.get_time_partition();
-
+    pub async fn get_dataframe(
+        &self,
+        time_partition: Option<&String>,
+    ) -> Result<DataFrame, ExecuteError> {
         let df = QUERY_SESSION
-            .execute_logical_plan(self.final_logical_plan(&time_partition))
+            .execute_logical_plan(self.final_logical_plan(time_partition))
             .await?;
 
         Ok(df)
     }
 
     /// return logical plan with all time filters applied through
-    fn final_logical_plan(&self, time_partition: &Option<String>) -> LogicalPlan {
+    fn final_logical_plan(&self, time_partition: Option<&String>) -> LogicalPlan {
         // see https://github.com/apache/arrow-datafusion/pull/8400
         // this can be eliminated in later version of datafusion but with slight caveat
         // transform cannot modify stringified plans by itself
@@ -486,7 +485,7 @@ fn transform(
     plan: LogicalPlan,
     start_time: NaiveDateTime,
     end_time: NaiveDateTime,
-    time_partition: &Option<String>,
+    time_partition: Option<&String>,
 ) -> Transformed<LogicalPlan> {
     plan.transform(&|plan| match plan {
         LogicalPlan::TableScan(table) => {
@@ -544,7 +543,7 @@ fn transform(
 
 fn table_contains_any_time_filters(
     table: &datafusion::logical_expr::TableScan,
-    time_partition: &Option<String>,
+    time_partition: Option<&String>,
 ) -> bool {
     table
         .filters
@@ -558,8 +557,8 @@ fn table_contains_any_time_filters(
         })
         .any(|expr| {
             matches!(&*expr.left, Expr::Column(Column { name, .. })
-            if ((time_partition.is_some() && name == time_partition.as_ref().unwrap()) ||
-            (!time_partition.is_some() && name == event::DEFAULT_TIMESTAMP_KEY)))
+            if (time_partition.is_some_and(|field| field == name) ||
+            (time_partition.is_none() && name == event::DEFAULT_TIMESTAMP_KEY)))
         })
 }