time filter for all tables in query

nikhilsinhaparseable · nikhilsinhaparseable · commit 7af280d0737d · 2025-08-18T04:45:44.000-07:00
diff --git a/src/alerts/alerts_utils.rs b/src/alerts/alerts_utils.rs
@@ -107,7 +107,7 @@ async fn execute_local_query(
         filter_tag: None,
     };
 
-    let (records, _) = execute(query, &tables[0], false)
+    let (records, _) = execute(query, false)
         .await
         .map_err(|err| AlertError::CustomError(format!("Failed to execute query: {err}")))?;
 
diff --git a/src/handlers/airplane.rs b/src/handlers/airplane.rs
@@ -205,7 +205,7 @@ impl FlightService for AirServiceImpl {
             })?;
         let time = Instant::now();
 
-        let (records, _) = execute(query, &stream_name, false)
+        let (records, _) = execute(query, false)
             .await
             .map_err(|err| Status::internal(err.to_string()))?;
 
diff --git a/src/handlers/http/query.rs b/src/handlers/http/query.rs
@@ -92,12 +92,9 @@ pub async fn get_records_and_fields(
     let creds = extract_session_key_from_req(req)?;
     let permissions = Users.get_permissions(&creds);
 
-    let table_name = tables
-        .first()
-        .ok_or_else(|| QueryError::MalformedQuery("No table name found in query"))?;
     user_auth_for_datasets(&permissions, &tables).await?;
 
-    let (records, fields) = execute(query, table_name, false).await?;
+    let (records, fields) = execute(query, false).await?;
 
     let records = match records {
         Either::Left(vec_rb) => vec_rb,
@@ -121,28 +118,28 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpRespons
     let creds = extract_session_key_from_req(&req)?;
     let permissions = Users.get_permissions(&creds);
 
-    let table_name = tables
-        .first()
-        .ok_or_else(|| QueryError::MalformedQuery("No table name found in query"))?;
     user_auth_for_datasets(&permissions, &tables).await?;
     let time = Instant::now();
 
     // if the query is `select count(*) from <dataset>`
     // we use the `get_bin_density` method to get the count of records in the dataset
     // instead of executing the query using datafusion
     if let Some(column_name) = query.is_logical_plan_count_without_filters() {
-        return handle_count_query(&query_request, table_name, column_name, time).await;
+        let table = tables
+            .first()
+            .ok_or_else(|| QueryError::MalformedQuery("No table name found in query"))?;
+        return handle_count_query(&query_request, table, column_name, time).await;
     }
 
     // if the query request has streaming = false (default)
     // we use datafusion's `execute` method to get the records
     if !query_request.streaming {
-        return handle_non_streaming_query(query, table_name, &query_request, time).await;
+        return handle_non_streaming_query(query, tables, &query_request, time).await;
     }
 
     // if the query request has streaming = true
     // we use datafusion's `execute_stream` method to get the records
-    handle_streaming_query(query, table_name, &query_request, time).await
+    handle_streaming_query(query, tables, &query_request, time).await
 }
 
 /// Handles count queries (e.g., `SELECT COUNT(*) FROM <dataset-name>`)
@@ -211,11 +208,12 @@ async fn handle_count_query(
 /// - `HttpResponse` with the full query result as a JSON object.
 async fn handle_non_streaming_query(
     query: LogicalQuery,
-    table_name: &str,
+    table_name: Vec<String>,
     query_request: &Query,
     time: Instant,
 ) -> Result<HttpResponse, QueryError> {
-    let (records, fields) = execute(query, table_name, query_request.streaming).await?;
+    let first_table_name = table_name[0].clone();
+    let (records, fields) = execute(query, query_request.streaming).await?;
     let records = match records {
         Either::Left(rbs) => rbs,
         Either::Right(_) => {
@@ -228,7 +226,7 @@ async fn handle_non_streaming_query(
     let time = time.elapsed().as_secs_f64();
 
     QUERY_EXECUTE_TIME
-        .with_label_values(&[table_name])
+        .with_label_values(&[&first_table_name])
         .observe(time);
     let response = QueryResponse {
         records,
@@ -259,11 +257,12 @@ async fn handle_non_streaming_query(
 /// - `HttpResponse` streaming the query results as NDJSON, optionally prefixed with the fields array.
 async fn handle_streaming_query(
     query: LogicalQuery,
-    table_name: &str,
+    table_name: Vec<String>,
     query_request: &Query,
     time: Instant,
 ) -> Result<HttpResponse, QueryError> {
-    let (records_stream, fields) = execute(query, table_name, query_request.streaming).await?;
+    let first_table_name = table_name[0].clone();
+    let (records_stream, fields) = execute(query, query_request.streaming).await?;
     let records_stream = match records_stream {
         Either::Left(_) => {
             return Err(QueryError::MalformedQuery(
@@ -275,7 +274,7 @@ async fn handle_streaming_query(
     let total_time = format!("{:?}", time.elapsed());
     let time = time.elapsed().as_secs_f64();
     QUERY_EXECUTE_TIME
-        .with_label_values(&[table_name])
+        .with_label_values(&[&first_table_name])
         .observe(time);
 
     let send_null = query_request.send_null;
diff --git a/src/query/mod.rs b/src/query/mod.rs
@@ -25,7 +25,7 @@ use chrono::NaiveDateTime;
 use chrono::{DateTime, Duration, Utc};
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::catalog::resolve_table_references;
-use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::common::tree_node::Transformed;
 use datafusion::error::DataFusionError;
 use datafusion::execution::disk_manager::DiskManagerConfig;
 use datafusion::execution::{SendableRecordBatchStream, SessionState, SessionStateBuilder};
@@ -56,7 +56,7 @@ use crate::catalog::Snapshot as CatalogSnapshot;
 use crate::catalog::column::{Int64Type, TypedStatistics};
 use crate::catalog::manifest::Manifest;
 use crate::catalog::snapshot::Snapshot;
-use crate::event;
+use crate::event::{self, DEFAULT_TIMESTAMP_KEY};
 use crate::handlers::http::query::QueryError;
 use crate::option::Mode;
 use crate::parseable::PARSEABLE;
@@ -77,7 +77,6 @@ pub static QUERY_RUNTIME: Lazy<Runtime> =
 /// at a time and has access to the entire thread pool, enabling better concurrent processing, and thus quicker results.
 pub async fn execute(
     query: Query,
-    stream_name: &str,
     is_streaming: bool,
 ) -> Result<
     (
@@ -86,9 +85,8 @@ pub async fn execute(
     ),
     ExecuteError,
 > {
-    let time_partition = PARSEABLE.get_stream(stream_name)?.get_time_partition();
     QUERY_RUNTIME
-        .spawn(async move { query.execute(time_partition.as_ref(), is_streaming).await })
+        .spawn(async move { query.execute(is_streaming).await })
         .await
         .expect("The Join should have been successful")
 }
@@ -180,7 +178,6 @@ impl Query {
     /// if streaming is false, it returns a vector of record batches
     pub async fn execute(
         &self,
-        time_partition: Option<&String>,
         is_streaming: bool,
     ) -> Result<
         (
@@ -190,7 +187,7 @@ impl Query {
         ExecuteError,
     > {
         let df = QUERY_SESSION
-            .execute_logical_plan(self.final_logical_plan(time_partition))
+            .execute_logical_plan(self.final_logical_plan())
             .await?;
 
         let fields = df
@@ -214,19 +211,16 @@ impl Query {
         Ok((results, fields))
     }
 
-    pub async fn get_dataframe(
-        &self,
-        time_partition: Option<&String>,
-    ) -> Result<DataFrame, ExecuteError> {
+    pub async fn get_dataframe(&self) -> Result<DataFrame, ExecuteError> {
         let df = QUERY_SESSION
-            .execute_logical_plan(self.final_logical_plan(time_partition))
+            .execute_logical_plan(self.final_logical_plan())
             .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) -> 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
@@ -238,7 +232,6 @@ impl Query {
                     plan.plan.as_ref().clone(),
                     self.time_range.start.naive_utc(),
                     self.time_range.end.naive_utc(),
-                    time_partition,
                 );
                 LogicalPlan::Explain(Explain {
                     verbose: plan.verbose,
@@ -257,7 +250,6 @@ impl Query {
                     x,
                     self.time_range.start.naive_utc(),
                     self.time_range.end.naive_utc(),
-                    time_partition,
                 )
                 .data
             }
@@ -586,66 +578,69 @@ fn transform(
     plan: LogicalPlan,
     start_time: NaiveDateTime,
     end_time: NaiveDateTime,
-    time_partition: Option<&String>,
 ) -> Transformed<LogicalPlan> {
-    plan.transform(&|plan| match plan {
-        LogicalPlan::TableScan(table) => {
-            let mut new_filters = vec![];
-            if !table_contains_any_time_filters(&table, time_partition) {
-                let mut _start_time_filter: Expr;
-                let mut _end_time_filter: Expr;
-                match time_partition {
-                    Some(time_partition) => {
-                        _start_time_filter =
-                            PartialTimeFilter::Low(std::ops::Bound::Included(start_time))
-                                .binary_expr(Expr::Column(Column::new(
-                                    Some(table.table_name.to_owned()),
-                                    time_partition.clone(),
-                                )));
-                        _end_time_filter =
-                            PartialTimeFilter::High(std::ops::Bound::Excluded(end_time))
-                                .binary_expr(Expr::Column(Column::new(
-                                    Some(table.table_name.to_owned()),
-                                    time_partition,
-                                )));
-                    }
-                    None => {
-                        _start_time_filter =
-                            PartialTimeFilter::Low(std::ops::Bound::Included(start_time))
-                                .binary_expr(Expr::Column(Column::new(
-                                    Some(table.table_name.to_owned()),
-                                    event::DEFAULT_TIMESTAMP_KEY,
-                                )));
-                        _end_time_filter =
-                            PartialTimeFilter::High(std::ops::Bound::Excluded(end_time))
-                                .binary_expr(Expr::Column(Column::new(
-                                    Some(table.table_name.to_owned()),
-                                    event::DEFAULT_TIMESTAMP_KEY,
-                                )));
-                    }
+    plan.transform_up_with_subqueries(&|plan| {
+        match plan {
+            LogicalPlan::TableScan(table) => {
+                // Get the specific time partition for this stream
+                let time_partition = PARSEABLE
+                    .get_stream(&table.table_name.to_string())
+                    .ok()
+                    .and_then(|stream| stream.get_time_partition());
+
+                let mut new_filters = vec![];
+                if !table_contains_any_time_filters(&table, time_partition.as_ref()) {
+                    let default_timestamp = DEFAULT_TIMESTAMP_KEY.to_string();
+                    let time_column = time_partition.as_ref().unwrap_or(&default_timestamp);
+
+                    // Create time filters with table-qualified column names
+                    let start_time_filter = PartialTimeFilter::Low(std::ops::Bound::Included(
+                        start_time,
+                    ))
+                    .binary_expr(Expr::Column(Column::new(
+                        Some(table.table_name.to_owned()),
+                        time_column.clone(),
+                    )));
+
+                    let end_time_filter = PartialTimeFilter::High(std::ops::Bound::Excluded(
+                        end_time,
+                    ))
+                    .binary_expr(Expr::Column(Column::new(
+                        Some(table.table_name.to_owned()),
+                        time_column.clone(),
+                    )));
+
+                    new_filters.push(start_time_filter);
+                    new_filters.push(end_time_filter);
                 }
 
-                new_filters.push(_start_time_filter);
-                new_filters.push(_end_time_filter);
+                let new_filter = new_filters.into_iter().reduce(and);
+                if let Some(new_filter) = new_filter {
+                    let filter =
+                        Filter::try_new(new_filter, Arc::new(LogicalPlan::TableScan(table)))
+                            .unwrap();
+                    Ok(Transformed::yes(LogicalPlan::Filter(filter)))
+                } else {
+                    Ok(Transformed::no(LogicalPlan::TableScan(table)))
+                }
             }
-            let new_filter = new_filters.into_iter().reduce(and);
-            if let Some(new_filter) = new_filter {
-                let filter =
-                    Filter::try_new(new_filter, Arc::new(LogicalPlan::TableScan(table))).unwrap();
-                Ok(Transformed::yes(LogicalPlan::Filter(filter)))
-            } else {
-                Ok(Transformed::no(LogicalPlan::TableScan(table)))
+            _ => {
+                // For all other plan types, continue the transformation recursively
+                // This ensures that subqueries and other nested plans are also transformed
+                Ok(Transformed::no(plan))
             }
         }
-        x => Ok(Transformed::no(x)),
     })
-    .expect("transform only transforms the tablescan")
+    .expect("transform processes all plan nodes")
 }
 
 fn table_contains_any_time_filters(
     table: &datafusion::logical_expr::TableScan,
     time_partition: Option<&String>,
 ) -> bool {
+    let default_timestamp = DEFAULT_TIMESTAMP_KEY.to_string();
+    let time_column = time_partition.unwrap_or(&default_timestamp);
+
     table
         .filters
         .iter()
@@ -658,8 +653,7 @@ fn table_contains_any_time_filters(
         })
         .any(|expr| {
             matches!(&*expr.left, Expr::Column(Column { name, .. })
-            if (time_partition.is_some_and(|field| field == name) ||
-            (time_partition.is_none() && name == event::DEFAULT_TIMESTAMP_KEY)))
+            if name == time_column)
         })
 }
 
diff --git a/src/query/stream_schema_provider.rs b/src/query/stream_schema_provider.rs
@@ -487,10 +487,6 @@ impl TableProvider for StandardTableProvider {
             .map_err(|err| DataFusionError::Plan(err.to_string()))?;
         let time_partition = object_store_format.time_partition;
         let mut time_filters = extract_primary_filter(filters, &time_partition);
-        if time_filters.is_empty() {
-            return Err(DataFusionError::Plan("potentially unbounded query on time range. Table scanning requires atleast one time bound".to_string()));
-        }
-
         if is_within_staging_window(&time_filters) {
             self.get_staging_execution_plan(
                 &mut execution_plans,
