rebase with main

Author: parmesant

Cargo.lock

@@ -26,6 +26,7 @@ object_store = { version = "0.13.1", features = [
     "azure",
     "gcp",
 ] }
+datafusion-proto = "53.1.0"
 parquet = "58.0.0"
 
 # Web server and HTTP-related

Cargo.toml

@@ -257,6 +257,18 @@ pub static TOTAL_QUERY_CALLS_BY_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
     .expect("metric can be created")
 });
 
+pub static TOTAL_FILES_SCANNED_IN_HOTTIER_BY_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
+        Opts::new(
+            "total_files_scanned_in_hottier_by_date",
+            "Total files scanned in hottier by date",
+        )
+        .namespace(METRICS_NAMESPACE),
+        &["stream", "date", "tenant_id"],
+    )
+    .expect("metric can be created")
+});
+
 pub static TOTAL_FILES_SCANNED_IN_QUERY_BY_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
     IntCounterVec::new(
         Opts::new(
@@ -683,6 +695,17 @@ pub fn increment_files_scanned_in_query_by_date(count: u64, date: &str, tenant_i
         .inc_by(count);
 }
 
+pub fn increment_files_scanned_in_hottier_by_date(
+    count: u64,
+    date: &str,
+    tenant_id: &str,
+    stream_name: &str,
+) {
+    TOTAL_FILES_SCANNED_IN_HOTTIER_BY_DATE
+        .with_label_values(&[stream_name, date, tenant_id])
+        .inc_by(count);
+}
+
 pub fn increment_bytes_scanned_in_query_by_date(bytes: u64, date: &str, tenant_id: &str) {
     TOTAL_BYTES_SCANNED_IN_QUERY_BY_DATE
         .with_label_values(&[date, tenant_id])

src/metrics/mod.rs

@@ -98,6 +98,12 @@ type QueryResult = Result<(Either<Vec<RecordBatch>, BoxedBatchStream>, Vec<Strin
 //     Lazy::new(|| Query::create_session_context(PARSEABLE.storage()));
 pub static SCHEMA_PROVIDER: OnceCell<Box<dyn ParseableSchemaProvider>> = OnceCell::new();
 
+/// Additional physical optimizer rules registered by enterprise/plugins.
+/// Must be populated BEFORE `QUERY_SESSION_STATE` is first accessed.
+pub static ADDITIONAL_PHYSICAL_OPTIMIZER_RULES: Lazy<
+    RwLock<Vec<Arc<dyn datafusion::physical_optimizer::PhysicalOptimizerRule + Send + Sync>>>,
+> = Lazy::new(|| RwLock::new(Vec::new()));
+
 pub static QUERY_SESSION_STATE: Lazy<SessionState> =
     Lazy::new(|| Query::create_session_state(PARSEABLE.storage()));
 
@@ -280,11 +286,19 @@ impl Query {
             .parquet
             .schema_force_view_types = true;
 
-        SessionStateBuilder::new()
+        let mut builder = SessionStateBuilder::new()
             .with_default_features()
             .with_config(config)
-            .with_runtime_env(runtime)
-            .build()
+            .with_runtime_env(runtime);
+
+        // Append any additional physical optimizer rules (e.g., enterprise partial agg pushdown)
+        if let Ok(rules) = ADDITIONAL_PHYSICAL_OPTIMIZER_RULES.read() {
+            for rule in rules.iter() {
+                builder = builder.with_physical_optimizer_rule(Arc::clone(rule));
+            }
+        }
+
+        builder.build()
     }
 
     /// this function returns the result of the query
@@ -316,6 +330,8 @@ impl Query {
             return Ok((Either::Left(vec![]), fields));
         }
 
+        let ctx = QUERY_SESSION.get_ctx();
+
         let plan = ctx.state().create_physical_plan(df.logical_plan()).await?;
 
         let results = if !is_streaming {

src/query/mod.rs

@@ -56,7 +56,10 @@ use crate::{
     },
     event::DEFAULT_TIMESTAMP_KEY,
     hottier::HotTierManager,
-    metrics::{QUERY_CACHE_HIT, increment_files_scanned_in_query_by_date},
+    metrics::{
+        QUERY_CACHE_HIT, increment_files_scanned_in_hottier_by_date,
+        increment_files_scanned_in_query_by_date,
+    },
     option::Mode,
     parseable::{DEFAULT_TENANT, PARSEABLE, STREAM_EXISTS},
     storage::{ObjectStorage, ObjectStoreFormat},
@@ -205,6 +208,13 @@ impl StandardTableProvider {
             .await
             .map_err(|err| DataFusionError::External(Box::new(err)))?;
 
+        increment_files_scanned_in_hottier_by_date(
+            hot_tier_files.len() as u64,
+            &chrono::Utc::now().date_naive().to_string(),
+            self.tenant_id.as_deref().unwrap_or(DEFAULT_TENANT),
+            &self.stream,
+        );
+
         let hot_tier_files: Vec<File> = hot_tier_files
             .into_iter()
             .map(|mut file| {
@@ -352,101 +362,108 @@ impl StandardTableProvider {
         &self,
         manifest_files: Vec<File>,
     ) -> (Vec<Vec<PartitionedFile>>, datafusion::common::Statistics) {
-        let target_partition: usize = num_cpus::get();
-        let mut partitioned_files = Vec::from_iter((0..target_partition).map(|_| Vec::new()));
-        let mut column_statistics = HashMap::<String, Option<TypedStatistics>>::new();
-        let mut count = 0;
-        let mut file_count = 0u64;
-        for (index, file) in manifest_files
-            .into_iter()
-            .enumerate()
-            .map(|(x, y)| (x % target_partition, y))
+        partitioned_files(&self.schema, &self.tenant_id, manifest_files)
+    }
+}
+
+#[inline(always)]
+pub fn partitioned_files(
+    schema: &SchemaRef,
+    tenant_id: &Option<String>,
+    manifest_files: Vec<File>,
+) -> (Vec<Vec<PartitionedFile>>, datafusion::common::Statistics) {
+    let target_partition: usize = num_cpus::get();
+    let mut partitioned_files = Vec::from_iter((0..target_partition).map(|_| Vec::new()));
+    let mut column_statistics = HashMap::<String, Option<TypedStatistics>>::new();
+    let mut count = 0;
+    let mut file_count = 0u64;
+    for (index, file) in manifest_files
+        .into_iter()
+        .enumerate()
+        .map(|(x, y)| (x % target_partition, y))
+    {
+        #[allow(unused_mut)]
+        let File {
+            mut file_path,
+            num_rows,
+            columns,
+            ..
+        } = file;
+
+        // Track billing metrics for files scanned in query
+        file_count += 1;
+
+        // object_store::path::Path doesn't automatically deal with Windows path separators
+        // to do that, we are using from_absolute_path() which takes into consideration the underlying filesystem
+        // before sending the file path to PartitionedFile
+        // the github issue- https://github.com/parseablehq/parseable/issues/824
+        // For some reason, the `from_absolute_path()` doesn't work for macos, hence the ugly solution
+        // TODO: figure out an elegant solution to this
+        #[cfg(windows)]
         {
-            #[allow(unused_mut)]
-            let File {
-                mut file_path,
-                num_rows,
-                columns,
-                ..
-            } = file;
-
-            // Track billing metrics for files scanned in query
-            file_count += 1;
-
-            // object_store::path::Path doesn't automatically deal with Windows path separators
-            // to do that, we are using from_absolute_path() which takes into consideration the underlying filesystem
-            // before sending the file path to PartitionedFile
-            // the github issue- https://github.com/parseablehq/parseable/issues/824
-            // For some reason, the `from_absolute_path()` doesn't work for macos, hence the ugly solution
-            // TODO: figure out an elegant solution to this
-            #[cfg(windows)]
-            {
-                if PARSEABLE.storage.name() == "drive" {
-                    file_path = object_store::path::Path::from_absolute_path(file_path)
-                        .unwrap()
-                        .to_string();
-                }
+            if PARSEABLE.storage.name() == "drive" {
+                file_path = object_store::path::Path::from_absolute_path(file_path)
+                    .unwrap()
+                    .to_string();
             }
-            let pf = PartitionedFile::new(file_path, file.file_size);
-            partitioned_files[index].push(pf);
-
-            columns.into_iter().for_each(|col| {
-                column_statistics
-                    .entry(col.name)
-                    .and_modify(|x| {
-                        if let Some((stats, col_stats)) = x.as_ref().cloned().zip(col.stats.clone())
-                        {
-                            // update() returns None on type mismatch (e.g. column
-                            // historically written as both Utf8 and Timestamp(ms)).
-                            // Dropping to None here makes the planner skip min/max
-                            // pushdown for this column instead of crashing the worker.
-                            *x = stats.update(col_stats);
-                        }
-                    })
-                    .or_insert_with(|| col.stats.as_ref().cloned());
-            });
-            count += num_rows;
         }
-        let statistics = self
-            .schema
-            .fields()
-            .iter()
-            .map(|field| {
-                column_statistics
-                    .get(field.name())
-                    .and_then(|stats| stats.as_ref())
-                    .and_then(|stats| stats.clone().min_max_as_scalar(field.data_type()))
-                    .map(|(min, max)| datafusion::common::ColumnStatistics {
-                        null_count: Precision::Absent,
-                        max_value: Precision::Exact(max),
-                        min_value: Precision::Exact(min),
-                        distinct_count: Precision::Absent,
-                        sum_value: Precision::Absent,
-                        byte_size: Precision::Absent,
-                    })
-                    .unwrap_or_default()
-            })
-            .collect();
+        let pf = PartitionedFile::new(file_path, file.file_size);
+        partitioned_files[index].push(pf);
+
+        columns.into_iter().for_each(|col| {
+            column_statistics
+                .entry(col.name)
+                .and_modify(|x| {
+                    if let Some((stats, col_stats)) = x.as_ref().cloned().zip(col.stats.clone()) {
+                        // update() returns None on type mismatch (e.g. column
+                        // historically written as both Utf8 and Timestamp(ms)).
+                        // Dropping to None here makes the planner skip min/max
+                        // pushdown for this column instead of crashing the worker.
+                        *x = stats.update(col_stats);
+                    }
+                })
+                .or_insert_with(|| col.stats.as_ref().cloned());
+        });
+        count += num_rows;
+    }
+    let statistics = schema
+        .fields()
+        .iter()
+        .map(|field| {
+            column_statistics
+                .get(field.name())
+                .and_then(|stats| stats.as_ref())
+                .and_then(|stats| stats.clone().min_max_as_scalar(field.data_type()))
+                .map(|(min, max)| datafusion::common::ColumnStatistics {
+                    null_count: Precision::Absent,
+                    max_value: Precision::Exact(max),
+                    min_value: Precision::Exact(min),
+                    distinct_count: Precision::Absent,
+                    sum_value: Precision::Absent,
+                    byte_size: Precision::Absent,
+                })
+                .unwrap_or_default()
+        })
+        .collect();
 
-        let statistics = datafusion::common::Statistics {
-            num_rows: Precision::Exact(count as usize),
-            total_byte_size: Precision::Absent,
-            column_statistics: statistics,
-        };
+    let statistics = datafusion::common::Statistics {
+        num_rows: Precision::Exact(count as usize),
+        total_byte_size: Precision::Absent,
+        column_statistics: statistics,
+    };
 
-        // Track billing metrics for query scan
-        let current_date = chrono::Utc::now().date_naive().to_string();
-        increment_files_scanned_in_query_by_date(
-            file_count,
-            &current_date,
-            self.tenant_id.as_deref().unwrap_or(DEFAULT_TENANT),
-        );
+    // Track billing metrics for query scan
+    let current_date = chrono::Utc::now().date_naive().to_string();
+    increment_files_scanned_in_query_by_date(
+        file_count,
+        &current_date,
+        tenant_id.as_deref().unwrap_or(DEFAULT_TENANT),
+    );
 
-        (partitioned_files, statistics)
-    }
+    (partitioned_files, statistics)
 }
 
-async fn collect_from_snapshot(
+pub async fn collect_from_snapshot(
     snapshot: &Snapshot,
     time_filters: &[PartialTimeFilter],
     filters: &[Expr],
@@ -683,7 +700,8 @@ impl TableProvider for StandardTableProvider {
     }
 }
 
-fn reversed_mem_table(
+#[inline(always)]
+pub fn reversed_mem_table(
     mut records: Vec<RecordBatch>,
     schema: Arc<Schema>,
 ) -> Result<MemTable, DataFusionError> {
@@ -863,7 +881,7 @@ pub fn is_within_staging_window(time_filters: &[PartialTimeFilter]) -> bool {
     !has_upper_bound
 }
 
-fn expr_in_boundary(filter: &Expr) -> bool {
+pub fn expr_in_boundary(filter: &Expr) -> bool {
     let Expr::BinaryExpr(binexpr) = filter else {
         return false;
     };
@@ -881,7 +899,7 @@ fn expr_in_boundary(filter: &Expr) -> bool {
         )
 }
 
-fn extract_timestamp_bound(
+pub fn extract_timestamp_bound(
     binexpr: &BinaryExpr,
     time_partition: &Option<String>,
 ) -> Option<(Operator, NaiveDateTime)> {