first pass at implementing predicate pushdown, seems to work

Cargo.toml

@@ -30,9 +30,17 @@ arrow-cast = { version = "50.0.0" }
 arrow-schema = { version = "50.0.0" }
 arrow-data = { version = "50.0.0" }
 datafusion = { version = "36.0", optional = true }
+datafusion-expr = { version = "36.0", optional = true }
+datafusion-common = { version = "36.0", optional = true }
+datafusion-physical-expr = { version = "36.0", optional = true }
 
 [features]
-datafusion = ["dep:datafusion"]
+datafusion = [
+    "dep:datafusion",
+    "dep:datafusion-physical-expr",
+    "dep:datafusion-expr",
+    "dep:datafusion-common",
+]
 all = ["datafusion"]
 
 [dev-dependencies]

src/async_reader/mod.rs

@@ -469,7 +469,7 @@ impl<T: for<'a> ZarrReadAsync<'a> + Clone + Unpin + Send + 'static>
 
         let mut predicate_stream: Option<ZarrStoreAsync<T>> = None;
         if let Some(filter) = &self.filter {
-            let predicate_proj = filter.get_all_projections();
+            let predicate_proj = filter.get_all_projections()?;
             predicate_stream = Some(
                 ZarrStoreAsync::new(
                     self.zarr_reader_async.clone(),

src/datafusion/file_opener.rs

@@ -17,22 +17,26 @@
 
 use arrow_schema::ArrowError;
 use datafusion::{datasource::physical_plan::FileOpener, error::DataFusionError};
+use datafusion_physical_expr::PhysicalExpr;
 use futures::{StreamExt, TryStreamExt};
+use std::sync::Arc;
 
 use crate::{
-    async_reader::{ZarrPath, ZarrRecordBatchStreamBuilder},
+    async_reader::{ZarrPath, ZarrReadAsync, ZarrRecordBatchStreamBuilder},
     reader::ZarrProjection,
 };
 
 use super::config::ZarrConfig;
+use super::helpers::build_row_filter;
 
 pub struct ZarrFileOpener {
     config: ZarrConfig,
+    filters: Option<Arc<dyn PhysicalExpr>>,
 }
 
 impl ZarrFileOpener {
-    pub fn new(config: ZarrConfig) -> Self {
-        Self { config }
+    pub fn new(config: ZarrConfig, filters: Option<Arc<dyn PhysicalExpr>>) -> Self {
+        Self { config, filters }
     }
 }
 
@@ -43,15 +47,27 @@ impl FileOpener for ZarrFileOpener {
     ) -> datafusion::error::Result<datafusion::datasource::physical_plan::FileOpenFuture> {
         let config = self.config.clone();
 
+        let filters_to_pushdown = self.filters.clone();
         Ok(Box::pin(async move {
             let zarr_path = ZarrPath::new(config.object_store, file_meta.object_meta.location);
-
             let rng = file_meta.range.map(|r| (r.start as usize, r.end as usize));
-
             let projection = ZarrProjection::from(config.projection.as_ref());
 
-            let batch_reader = ZarrRecordBatchStreamBuilder::new(zarr_path)
-                .with_projection(projection)
+            let mut batch_reader_builder =
+                ZarrRecordBatchStreamBuilder::new(zarr_path.clone()).with_projection(projection);
+            if let Some(filters) = filters_to_pushdown {
+                let schema = zarr_path
+                    .get_zarr_metadata()
+                    .await
+                    .map_err(|e| DataFusionError::External(Box::new(e)))?
+                    .arrow_schema()
+                    .map_err(|e| DataFusionError::External(Box::new(e)))?;
+                let filters = build_row_filter(&filters, &schema)?;
+                if let Some(filters) = filters {
+                    batch_reader_builder = batch_reader_builder.with_filter(filters);
+                }
+            }
+            let batch_reader = batch_reader_builder
                 .build_partial_reader(rng)
                 .await
                 .map_err(|e| DataFusionError::External(Box::new(e)))?;
@@ -81,7 +97,7 @@ mod tests {
         let test_data = get_test_v2_data_path("lat_lon_example.zarr".to_string());
 
         let config = ZarrConfig::new(Arc::new(local_fs));
-        let opener = ZarrFileOpener::new(config);
+        let opener = ZarrFileOpener::new(config, None);
 
         let file_meta = FileMeta {
             object_meta: ObjectMeta {

src/datafusion/helpers.rs

@@ -0,0 +1,228 @@
+use crate::reader::{ZarrArrowPredicate, ZarrChunkFilter, ZarrProjection};
+use arrow::array::BooleanArray;
+use arrow::error::ArrowError;
+use arrow::record_batch::RecordBatch;
+use arrow_schema::Schema;
+use datafusion_common::cast::as_boolean_array;
+use datafusion_common::tree_node::{RewriteRecursion, TreeNode, TreeNodeRewriter, VisitRecursion};
+use datafusion_common::Result as DataFusionResult;
+use datafusion_common::{internal_err, DataFusionError};
+use datafusion_expr::{Expr, ScalarFunctionDefinition, Volatility};
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr::utils::reassign_predicate_columns;
+use datafusion_physical_expr::{split_conjunction, PhysicalExpr};
+use std::collections::BTreeSet;
+use std::sync::Arc;
+
+// Checks whether the given expression can be resolved using only the columns `col_names`.
+// Copied from datafusion, because it's not accessible from the outside.
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
+    let mut is_applicable = true;
+    expr.apply(&mut |expr| match expr {
+        Expr::Column(datafusion_common::Column { ref name, .. }) => {
+            is_applicable &= col_names.contains(name);
+            if is_applicable {
+                Ok(VisitRecursion::Skip)
+            } else {
+                Ok(VisitRecursion::Stop)
+            }
+        }
+        Expr::Literal(_)
+        | Expr::Alias(_)
+        | Expr::OuterReferenceColumn(_, _)
+        | Expr::ScalarVariable(_, _)
+        | Expr::Not(_)
+        | Expr::IsNotNull(_)
+        | Expr::IsNull(_)
+        | Expr::IsTrue(_)
+        | Expr::IsFalse(_)
+        | Expr::IsUnknown(_)
+        | Expr::IsNotTrue(_)
+        | Expr::IsNotFalse(_)
+        | Expr::IsNotUnknown(_)
+        | Expr::Negative(_)
+        | Expr::Cast { .. }
+        | Expr::TryCast { .. }
+        | Expr::BinaryExpr { .. }
+        | Expr::Between { .. }
+        | Expr::Like { .. }
+        | Expr::SimilarTo { .. }
+        | Expr::InList { .. }
+        | Expr::Exists { .. }
+        | Expr::InSubquery(_)
+        | Expr::ScalarSubquery(_)
+        | Expr::GetIndexedField { .. }
+        | Expr::GroupingSet(_)
+        | Expr::Case { .. } => Ok(VisitRecursion::Continue),
+
+        Expr::ScalarFunction(scalar_function) => match &scalar_function.func_def {
+            ScalarFunctionDefinition::BuiltIn(fun) => match fun.volatility() {
+                Volatility::Immutable => Ok(VisitRecursion::Continue),
+                Volatility::Stable | Volatility::Volatile => {
+                    is_applicable = false;
+                    Ok(VisitRecursion::Stop)
+                }
+            },
+            ScalarFunctionDefinition::UDF(fun) => match fun.signature().volatility {
+                Volatility::Immutable => Ok(VisitRecursion::Continue),
+                Volatility::Stable | Volatility::Volatile => {
+                    is_applicable = false;
+                    Ok(VisitRecursion::Stop)
+                }
+            },
+            ScalarFunctionDefinition::Name(_) => {
+                internal_err!("Function `Expr` with name should be resolved.")
+            }
+        },
+
+        Expr::AggregateFunction { .. }
+        | Expr::Sort { .. }
+        | Expr::WindowFunction { .. }
+        | Expr::Wildcard { .. }
+        | Expr::Unnest { .. }
+        | Expr::Placeholder(_) => {
+            is_applicable = false;
+            Ok(VisitRecursion::Stop)
+        }
+    })
+    .unwrap();
+    is_applicable
+}
+
+// Below is all the logic necessary (I think) to convert a PhysicalExpr into a ZarrChunkFilter.
+// The logic is mostly copied from datafusion, and is simplified here for the zarr use case.
+pub struct ZarrFilterCandidate {
+    expr: Arc<dyn PhysicalExpr>,
+    projection: Vec<usize>,
+}
+
+struct ZarrFilterCandidateBuilder<'a> {
+    expr: Arc<dyn PhysicalExpr>,
+    file_schema: &'a Schema,
+    required_column_indices: BTreeSet<usize>,
+}
+
+impl<'a> ZarrFilterCandidateBuilder<'a> {
+    pub fn new(expr: Arc<dyn PhysicalExpr>, file_schema: &'a Schema) -> Self {
+        Self {
+            expr,
+            file_schema,
+            required_column_indices: BTreeSet::default(),
+        }
+    }
+
+    pub fn build(mut self) -> DataFusionResult<Option<ZarrFilterCandidate>> {
+        let expr = self.expr.clone().rewrite(&mut self)?;
+
+        Ok(Some(ZarrFilterCandidate {
+            expr,
+            projection: self.required_column_indices.into_iter().collect(),
+        }))
+    }
+}
+
+impl<'a> TreeNodeRewriter for ZarrFilterCandidateBuilder<'a> {
+    type N = Arc<dyn PhysicalExpr>;
+
+    fn pre_visit(&mut self, node: &Arc<dyn PhysicalExpr>) -> DataFusionResult<RewriteRecursion> {
+        if let Some(column) = node.as_any().downcast_ref::<Column>() {
+            if let Ok(idx) = self.file_schema.index_of(column.name()) {
+                self.required_column_indices.insert(idx);
+            }
+        }
+
+        Ok(RewriteRecursion::Continue)
+    }
+
+    fn mutate(&mut self, expr: Arc<dyn PhysicalExpr>) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
+        Ok(expr)
+    }
+}
+
+#[derive(Clone)]
+pub struct ZarrDatafusionArrowPredicate {
+    physical_expr: Arc<dyn PhysicalExpr>,
+    projection_mask: ZarrProjection,
+    projection: Vec<String>,
+}
+
+impl ZarrDatafusionArrowPredicate {
+    pub fn new(candidate: ZarrFilterCandidate, schema: &Schema) -> DataFusionResult<Self> {
+        let cols: Vec<_> = candidate
+            .projection
+            .iter()
+            .map(|idx| schema.field(*idx).name().to_string())
+            .collect();
+
+        let schema = Arc::new(schema.project(&candidate.projection)?);
+        let physical_expr = reassign_predicate_columns(candidate.expr, &schema, true)?;
+
+        Ok(Self {
+            physical_expr,
+            projection_mask: ZarrProjection::keep(cols.clone()),
+            projection: cols,
+        })
+    }
+}
+
+impl ZarrArrowPredicate for ZarrDatafusionArrowPredicate {
+    fn projection(&self) -> &ZarrProjection {
+        &self.projection_mask
+    }
+
+    fn evaluate(&mut self, batch: &RecordBatch) -> Result<BooleanArray, ArrowError> {
+        let index_projection = self
+            .projection
+            .iter()
+            .map(|col| batch.schema().index_of(col))
+            .collect::<Result<Vec<_>, _>>()?;
+        let batch = batch.project(&index_projection[..])?;
+
+        match self
+            .physical_expr
+            .evaluate(&batch)
+            .and_then(|v| v.into_array(batch.num_rows()))
+        {
+            Ok(array) => {
+                let bool_arr = as_boolean_array(&array)?.clone();
+                Ok(bool_arr)
+            }
+            Err(e) => Err(ArrowError::ComputeError(format!(
+                "Error evaluating filter predicate: {e:?}"
+            ))),
+        }
+    }
+}
+
+pub(crate) fn build_row_filter(
+    expr: &Arc<dyn PhysicalExpr>,
+    file_schema: &Schema,
+) -> DataFusionResult<Option<ZarrChunkFilter>> {
+    let predicates = split_conjunction(expr);
+    let candidates: Vec<ZarrFilterCandidate> = predicates
+        .into_iter()
+        .flat_map(|expr| {
+            if let Ok(candidate) =
+                ZarrFilterCandidateBuilder::new(expr.clone(), file_schema).build()
+            {
+                candidate
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    if candidates.is_empty() {
+        Ok(None)
+    } else {
+        let mut filters: Vec<Box<dyn ZarrArrowPredicate>> = vec![];
+        for candidate in candidates {
+            let filter = ZarrDatafusionArrowPredicate::new(candidate, file_schema)?;
+            filters.push(Box::new(filter));
+        }
+
+        let chunk_filter = ZarrChunkFilter::new(filters);
+
+        Ok(Some(chunk_filter))
+    }
+}

src/datafusion/mod.rs

@@ -17,6 +17,7 @@
 
 pub mod config;
 pub mod file_opener;
+mod helpers;
 pub mod scanner;
 pub mod table_factory;
 pub mod table_provider;

src/datafusion/scanner.rs

@@ -23,7 +23,7 @@ use datafusion::{
     datasource::physical_plan::{FileScanConfig, FileStream},
     physical_plan::{
         metrics::ExecutionPlanMetricsSet, DisplayAs, DisplayFormatType, ExecutionPlan,
-        Partitioning, SendableRecordBatchStream,
+        Partitioning, PhysicalExpr, SendableRecordBatchStream,
     },
 };
 
@@ -43,18 +43,22 @@ pub struct ZarrScan {
 
     /// The statistics for the scan.
     statistics: Statistics,
+
+    /// Filters that will be pushed down to the Zarr stream reader.
+    filters: Option<Arc<dyn PhysicalExpr>>,
 }
 
 impl ZarrScan {
     /// Create a new Zarr scan.
-    pub fn new(base_config: FileScanConfig) -> Self {
+    pub fn new(base_config: FileScanConfig, filters: Option<Arc<dyn PhysicalExpr>>) -> Self {
         let (projected_schema, statistics, _lex_sorting) = base_config.project();
 
         Self {
             base_config,
             projected_schema,
             metrics: ExecutionPlanMetricsSet::new(),
             statistics,
+            filters,
         }
     }
 }
@@ -100,9 +104,7 @@ impl ExecutionPlan for ZarrScan {
 
         let config =
             ZarrConfig::new(object_store).with_projection(self.base_config.projection.clone());
-
-        let opener = ZarrFileOpener::new(config);
-
+        let opener = ZarrFileOpener::new(config, self.filters.clone());
         let stream = FileStream::new(&self.base_config, partition, opener, &self.metrics)?;
 
         Ok(Box::pin(stream) as SendableRecordBatchStream)
@@ -169,7 +171,7 @@ mod tests {
             output_ordering: vec![],
         };
 
-        let scanner = ZarrScan::new(scan_config);
+        let scanner = ZarrScan::new(scan_config, None);
 
         let session = datafusion::execution::context::SessionContext::new();