Fast pushdown projection

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 {

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<String>,
+    ) -> Result<Self, DataFusionError> {
+        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<String> = 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::<Vec<_>>();
+
+            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]

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<String>>,
     ) -> datafusion::error::Result<LogicalPlanBuilder> {
         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::<FileCollection>())
+                    {
+                        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))
             }

beacon-query/src/lib.rs

@@ -94,6 +94,23 @@ impl Literal {
 }
 
 impl Select {
+    pub fn collect_columns(&self, columns: &mut Vec<String>) {
+        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<Expr> {
         match self {
             Select::ColumnName(name) => Ok(column_name(name)),

beacon-query/src/parser.rs

@@ -73,11 +73,26 @@ impl Parser {
         session: &SessionContext,
         data_lake: &DataLake,
     ) -> anyhow::Result<LogicalPlan> {
-        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();