feat: Added a new trait to expose SchemaProvider

Cargo.toml

@@ -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

src/lib.rs

@@ -58,7 +58,12 @@ pub mod validator;
 use std::time::Duration;
 
 // Public re-exports of crates being used in enterprise
+pub use arrow_array;
+pub use arrow_flight;
+pub use arrow_ipc;
+pub use catalog as parseable_catalog;
 pub use datafusion;
+pub use datafusion_proto;
 pub use handlers::http::modal::{
  ParseableServer, ingest_server::IngestServer, query_server::QueryServer, server::Server,
 };
@@ -68,6 +73,7 @@ use once_cell::sync::Lazy;
 pub use openid;
 use parseable::PARSEABLE;
 use reqwest::{Client, ClientBuilder};
+pub use utils as parseable_utils;
 pub use {opentelemetry, opentelemetry_otlp, opentelemetry_proto, opentelemetry_sdk};
 pub use {tracing_actix_web, tracing_opentelemetry, tracing_subscriber};
 

src/metrics/mod.rs

@@ -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(
@@ -716,6 +728,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/query/mod.rs

@@ -25,6 +25,7 @@ use arrow_schema::SchemaRef;
 use chrono::NaiveDateTime;
 use chrono::{DateTime, Duration, Utc};
 use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::catalog::SchemaProvider;
 use datafusion::common::tree_node::Transformed;
 use datafusion::execution::disk_manager::DiskManager;
 use datafusion::execution::{
@@ -45,7 +46,7 @@ use datafusion::sql::sqlparser::dialect::PostgreSqlDialect;
 use futures::Stream;
 use futures::StreamExt;
 use itertools::Itertools;
-use once_cell::sync::Lazy;
+use once_cell::sync::{Lazy, OnceCell};
 use serde::{Deserialize, Serialize};
 use serde_json::{Value, json};
 use std::ops::Bound;
@@ -59,7 +60,6 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
 use tracing::Instrument;
 
 use self::error::ExecuteError;
-use self::stream_schema_provider::GlobalSchemaProvider;
 pub use self::stream_schema_provider::PartialTimeFilter;
 use crate::alerts::alert_structs::Conditions;
 use crate::alerts::alerts_utils::get_filter_string;
@@ -72,7 +72,8 @@ use crate::handlers::http::query::QueryError;
 use crate::metrics::increment_bytes_scanned_in_query_by_date;
 use crate::option::Mode;
 use crate::parseable::{DEFAULT_TENANT, PARSEABLE};
-use crate::storage::{ObjectStorageProvider, ObjectStoreFormat};
+use crate::query::stream_schema_provider::GlobalSchemaProvider;
+use crate::storage::{ObjectStorage, ObjectStorageProvider, ObjectStoreFormat};
 use crate::utils::time::{DATE_BIN_EPOCH_ANCHOR, TimeRange, count_api_bin_interval};
 
 /// Boxed record-batch stream used as the streaming half of query results.
@@ -81,8 +82,13 @@ type BoxedBatchStream = SendableRecordBatchStream;
 /// Result type returned by query execution: either collected batches or a streaming adapter, plus field names.
 type QueryResult = Result<(Either<Vec<RecordBatch>, BoxedBatchStream>, Vec<String>), ExecuteError>;
 
-// pub static QUERY_SESSION: Lazy<SessionContext> =
-// 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()));
@@ -98,6 +104,15 @@ pub static QUERY_SESSION: Lazy<InMemorySessionContext> = Lazy::new(|| {
  }
 });
 
+/// Trait to enable implementation of SchemaProvider
+pub trait ParseableSchemaProvider: Send + Sync {
+ fn new_provider(
+ &self,
+ storage: Option<Arc<dyn ObjectStorage>>,
+ tenant_id: &Option<String>,
+ ) -> Box<dyn SchemaProvider>;
+}
+
 pub struct InMemorySessionContext {
  session_context: Arc<RwLock<SessionContext>>,
 }
@@ -112,18 +127,23 @@ impl InMemorySessionContext {
  }
 
  pub fn add_schema(&self, tenant_id: &str) {
+ let schema_provider = if let Some(provider) = SCHEMA_PROVIDER.get() {
+ provider.new_provider(
+ Some(PARSEABLE.storage().get_object_store()),
+ &Some(tenant_id.to_owned()),
+ )
+ } else {
+ Box::new(GlobalSchemaProvider {
+ storage: PARSEABLE.storage().get_object_store(),
+ tenant_id: Some(tenant_id.to_owned()),
+ })
+ };
  self.session_context
  .write()
  .expect("SessionContext should be writeable")
  .catalog("datafusion")
  .expect("Default catalog should be available")
- .register_schema(
- tenant_id,
- Arc::new(GlobalSchemaProvider {
- storage: PARSEABLE.storage().get_object_store(),
- tenant_id: Some(tenant_id.to_owned()),
- }),
- )
+ .register_schema(tenant_id, schema_provider.into())
  .expect("Should be able to register new schema");
  }
 
@@ -179,29 +199,41 @@ impl Query {
  // register multiple schemas
  if let Some(tenants) = PARSEABLE.list_tenants() {
  for t in tenants.iter() {
- let schema_provider = Arc::new(GlobalSchemaProvider {
- storage: storage.get_object_store(),
- tenant_id: Some(t.clone()),
- });
- let _ = catalog.register_schema(t, schema_provider);
+ let schema_provider = if let Some(provider) = SCHEMA_PROVIDER.get() {
+ provider.new_provider(
+ Some(PARSEABLE.storage().get_object_store()),
+ &Some(t.to_owned()),
+ )
+ } else {
+ Box::new(GlobalSchemaProvider {
+ storage: PARSEABLE.storage().get_object_store(),
+ tenant_id: Some(t.to_owned()),
+ })
+ };
+ let _ = catalog.register_schema(t, schema_provider.into());
  }
  }
  } else {
  // register just one schema
- let schema_provider = Arc::new(GlobalSchemaProvider {
- storage: storage.get_object_store(),
- tenant_id: None,
- });
+ let schema_provider = if let Some(provider) = SCHEMA_PROVIDER.get() {
+ provider.new_provider(Some(PARSEABLE.storage().get_object_store()), &None)
+ } else {
+ Box::new(GlobalSchemaProvider {
+ storage: PARSEABLE.storage().get_object_store(),
+ tenant_id: None,
+ })
+ };
+
  let _ = catalog.register_schema(
  &state.config_options().catalog.default_schema,
- schema_provider,
+ schema_provider.into(),
  );
  }
 
  SessionContext::new_with_state(state)
  }
 
- fn create_session_state(storage: Arc<dyn ObjectStorageProvider>) -> SessionState {
+pub fn create_session_state(storage: Arc<dyn ObjectStorageProvider>) -> SessionState {
  let runtime_config = storage
  .get_datafusion_runtime()
  .with_disk_manager_builder(DiskManager::builder());
@@ -260,11 +292,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
@@ -296,14 +336,12 @@ impl Query {
  return Ok((Either::Left(vec![]), fields));
  }
 
- let plan = QUERY_SESSION
- .get_ctx()
- .state()
- .create_physical_plan(df.logical_plan())
- .await?;
+ let ctx = QUERY_SESSION.get_ctx();
+
+ let plan = ctx.state().create_physical_plan(df.logical_plan()).await?;
 
  let results = if !is_streaming {
- let task_ctx = QUERY_SESSION.get_ctx().task_ctx();
+ let task_ctx = ctx.task_ctx();
 
  let batches = collect_partitioned(plan.clone(), task_ctx.clone())
  .await?
@@ -319,7 +357,7 @@ impl Query {
 
  Either::Left(batches)
  } else {
- let task_ctx = QUERY_SESSION.get_ctx().task_ctx();
+ let task_ctx = ctx.task_ctx();
 
  let output_partitions = plan.output_partitioning().partition_count();