diff --git a/beacon-config/src/lib.rs b/beacon-config/src/lib.rs index 25aaad08..25972f81 100644 --- a/beacon-config/src/lib.rs +++ b/beacon-config/src/lib.rs @@ -66,6 +66,8 @@ pub struct Config { pub allowed_credentials: bool, #[envconfig(from = "BEACON_CORS_MAX_AGE", default = "3600")] pub max_age: u64, + #[envconfig(from = "BEACON_ENABLE_PUSHDOWN_PROJECTION", default = "false")] + pub enable_pushdown_projection: bool, } impl Config { diff --git a/beacon-data-lake/src/files/collection.rs b/beacon-data-lake/src/files/collection.rs index a1a21330..642395d5 100644 --- a/beacon-data-lake/src/files/collection.rs +++ b/beacon-data-lake/src/files/collection.rs @@ -1,6 +1,6 @@ -use std::{any::Any, borrow::Cow, sync::Arc}; +use std::{any::Any, borrow::Cow, collections::HashSet, sync::Arc}; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{Schema, SchemaRef}; use datafusion::{ catalog::{Session, TableProvider}, common::{Constraints, Statistics}, @@ -51,6 +51,39 @@ impl FileCollection { Ok(Self { inner_table: table }) } + + pub fn with_pushdown_projection( + &self, + projection: Vec, + ) -> Result { + let mut schema = self.inner_table.schema(); + let options = self.inner_table.options().clone(); + let table_paths = self.inner_table.table_paths().to_vec(); + + let maybe_set_projection: HashSet = HashSet::from_iter(projection.iter().cloned()); + + if !maybe_set_projection.is_empty() { + // Only keep fields that are in the projection + let filtered_fields = schema + .fields() + .iter() + .filter(|f| projection.contains(f.name())) + .map(|f| f.as_ref().clone()) + .collect::>(); + + if !filtered_fields.is_empty() { + schema = Arc::new(Schema::new(filtered_fields)); + } + } + + let table_config = ListingTableConfig::new_with_multi_paths(table_paths) + .with_listing_options(options) + .with_schema(schema); + + let table = ListingTable::try_new(table_config)?; + + Ok(Self { inner_table: table }) + } } #[async_trait::async_trait] diff --git a/beacon-query/src/from.rs b/beacon-query/src/from.rs index 22539a47..9ee442dc 100644 --- a/beacon-query/src/from.rs +++ b/beacon-query/src/from.rs @@ -10,6 +10,7 @@ use beacon_data_lake::{ use beacon_formats::{ arrow::ArrowFormat, csv::CsvFormat, odv_ascii::OdvFormat, parquet::ParquetFormat, }; +use datafusion::catalog::SchemaProvider; use datafusion::{ datasource::{file_format::FileFormat, listing::ListingTableUrl, provider_as_source}, logical_expr::{LogicalPlanBuilder, TableSource}, @@ -55,10 +56,25 @@ impl From { &self, session_context: &SessionContext, data_lake: &DataLake, + projection: Option<&Vec>, ) -> datafusion::error::Result { match self { From::Table(name) => { // Use a registered table. + let table = data_lake.table(name).await?; + if let Some(mut table) = table { + if let (Some(projection), Some(file_collection)) = + (projection, table.as_any().downcast_ref::()) + { + let projected_table = + file_collection.with_pushdown_projection(projection.clone())?; + table = Arc::new(projected_table); + } + let source = provider_as_source(table); + + return LogicalPlanBuilder::scan(name, source, None); + } + let table = session_context.table(name).await?; Ok(LogicalPlanBuilder::new(table.into_parts().1)) } diff --git a/beacon-query/src/lib.rs b/beacon-query/src/lib.rs index 92c785c3..6c6bd634 100644 --- a/beacon-query/src/lib.rs +++ b/beacon-query/src/lib.rs @@ -94,6 +94,23 @@ impl Literal { } impl Select { + pub fn collect_columns(&self, columns: &mut Vec) { + match self { + Select::ColumnName(name) => { + columns.push(name.clone()); + } + Select::Column { column, .. } => { + columns.push(column.clone()); + } + Select::Literal { .. } => {} + Select::Function { args, .. } => { + for arg in args { + arg.collect_columns(columns); + } + } + } + } + pub fn to_expr(&self, session_state: &SessionState) -> anyhow::Result { match self { Select::ColumnName(name) => Ok(column_name(name)), diff --git a/beacon-query/src/parser.rs b/beacon-query/src/parser.rs index 073d0fd8..41b0bf40 100644 --- a/beacon-query/src/parser.rs +++ b/beacon-query/src/parser.rs @@ -73,11 +73,26 @@ impl Parser { session: &SessionContext, data_lake: &DataLake, ) -> anyhow::Result { - let mut builder = query_body - .from - .unwrap_or_default() - .init_builder(&session, data_lake) - .await?; + let mut builder = if beacon_config::CONFIG.enable_pushdown_projection { + let mut all_columns = vec![]; + for select in &query_body.select { + let mut select_cols = vec![]; + select.collect_columns(&mut select_cols); + all_columns.extend(select_cols); + } + + query_body + .from + .unwrap_or_default() + .init_builder(session, data_lake, Some(&all_columns)) + .await? + } else { + query_body + .from + .unwrap_or_default() + .init_builder(session, data_lake, None) + .await? + }; let session_state = session.state();