fix and add docs on schemas

Author: adriangb

datafusion-examples/examples/struct_field_rewrite.rs

@@ -353,8 +353,6 @@ impl TreeNodeRewriter for StructFieldRewriterImpl {
                                 .file_schema
                                 .field_with_name(&expected_flattened_column_name)
                             {
-                                println!("source_field: {:?}", source_field);
-                                println!("shredded_field: {:?}", shredded_field);
                                 if source_field.data_type() == shredded_field.data_type()
                                 {
                                     // Rewrite the expression to use the flattened column

datafusion/datasource-parquet/src/opener.rs

@@ -17,6 +17,7 @@
 
 //! [`ParquetOpener`] for opening Parquet files
 
+use std::collections::BTreeSet;
 use std::sync::Arc;
 
 use crate::file_format::{
@@ -28,14 +29,16 @@ use crate::{
     row_filter, should_enable_page_index, ParquetAccessPlan, ParquetFileMetrics,
     ParquetFileReaderFactory,
 };
+use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
 use datafusion_datasource::file_expr_rewriter::FileExpressionRewriter;
 use datafusion_datasource::file_meta::FileMeta;
 use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
 use datafusion_datasource::schema_adapter::SchemaAdapterFactory;
 
-use arrow::datatypes::SchemaRef;
+use arrow::datatypes::{Field, Schema, SchemaRef};
 use arrow::error::ArrowError;
-use datafusion_common::{exec_err, Result};
+use datafusion_common::{exec_err, Result, SchemaError};
+use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_optimizer::pruning::PruningPredicate;
 use datafusion_physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder};
@@ -90,6 +93,17 @@ pub(super) struct ParquetOpener {
 
 impl FileOpener for ParquetOpener {
     fn open(&self, file_meta: FileMeta) -> Result<FileOpenFuture> {
+        // Note about schemas: we are actually dealing with _4_ different schemas here:
+        // - The table schema as defined by the TableProvider. This is what the user sees, what they get when they `SELECT * FROM table`, etc.
+        // - The "virtual" file schema: this is the table schema minus any hive partition columns. This is what the file schema is coerced to.
+        // - The physical file schema: this is the schema as defined by the parquet file. This is what the parquet file actually contains.
+        // - The filter schema: a hybrid of the virtual file schema and the physical file schema.
+        //   If a filter is rewritten to reference columns that are in the physical file schema but not the virtual file schema, we need to add those columns to the filter schema so that the filter can be evaluated.
+        //   This schema is generated by taking any columns from the virtual file schema that are referenced by the filter and adding any columns from the physical file schema that are referenced by the filter but not in the virtual file schema.
+        //   Columns from the virtual file schema are added in the order they appear in the virtual file schema.
+        //   The columns from the physical file schema are always added to the end of the schema, in the order they appear in the physical file schema.
+        //
+        // I think it might be wise to do some renaming of parameters where possible, e.g. rename `file_schema` to `table_schema_without_partition_columns` and `physical_file_schema` or something like that.
         let file_range = file_meta.range.clone();
         let extensions = file_meta.extensions.clone();
         let file_name = file_meta.location().to_string();
@@ -110,13 +124,12 @@ impl FileOpener for ParquetOpener {
 
         let projected_schema =
             SchemaRef::from(self.table_schema.project(&self.projection)?);
-        let schema_adapter = self
-            .schema_adapter_factory
-            .create(projected_schema, Arc::clone(&self.table_schema));
         let mut predicate = self.predicate.clone();
         let mut pruning_predicate = self.pruning_predicate.clone();
         let mut page_pruning_predicate = self.page_pruning_predicate.clone();
+        let schema_adapter_factory = self.schema_adapter_factory.clone();
         let table_schema = Arc::clone(&self.table_schema);
+        let mut filter_schema = table_schema.clone();
         let filter_expression_rewriter = self.filter_expression_rewriter.clone();
         let reorder_predicates = self.reorder_filters;
         let pushdown_filters = self.pushdown_filters;
@@ -165,14 +178,6 @@ impl FileOpener for ParquetOpener {
 
             let file_schema = Arc::clone(builder.schema());
 
-            let (schema_mapping, adapted_projections) =
-                schema_adapter.map_schema(&file_schema)?;
-
-            let mask = ProjectionMask::roots(
-                builder.parquet_schema(),
-                adapted_projections.iter().cloned(),
-            );
-
             // Try to rewrite the predicate using our filter expression rewriter
             // This must happen before build_row_filter and other subsequent steps
             // so that they all use the rewritten predicate if it was rewritten.
@@ -181,12 +186,18 @@ impl FileOpener for ParquetOpener {
                     if let Ok(rewritten) = filter_rewriter
                         .rewrite(Arc::clone(&file_schema), original_predicate)
                     {
+                        let mut filter_schema_builder = FilterSchemaBuilder::new(
+                            &file_schema,
+                            &table_schema,
+                        );
+                        rewritten.visit(&mut filter_schema_builder)?;
+                        filter_schema = filter_schema_builder.build();
                         // If we rewrote the filter we need to recompute the pruning predicates to match the new filter.
                         page_pruning_predicate =
-                            Some(build_page_pruning_predicate(&rewritten, &file_schema));
+                            Some(build_page_pruning_predicate(&rewritten, &filter_schema));
                         pruning_predicate = build_pruning_predicate(
                             rewritten.clone(),
-                            &file_schema,
+                            &filter_schema,
                             &predicate_creation_errors,
                         );
                         // Update the predicate to the rewritten version
@@ -195,12 +206,23 @@ impl FileOpener for ParquetOpener {
                 }
             }
 
+            let schema_adapter = schema_adapter_factory
+                .create(projected_schema, Arc::clone(&filter_schema));
+
+            let (schema_mapping, adapted_projections) =
+                schema_adapter.map_schema(&file_schema)?;
+
+            let mask = ProjectionMask::roots(
+                builder.parquet_schema(),
+                adapted_projections.iter().cloned(),
+            );
+
             // Filter pushdown: evaluate predicates during scan
             if let Some(predicate) = pushdown_filters.then_some(predicate).flatten() {
                 let row_filter = row_filter::build_row_filter(
                     &predicate,
                     &file_schema,
-                    &table_schema,
+                    &filter_schema,
                     builder.metadata(),
                     reorder_predicates,
                     &file_metrics,
@@ -370,3 +392,68 @@ pub(crate) fn build_page_pruning_predicate(
         Arc::clone(file_schema),
     ))
 }
+
+/// A vistor for a PhysicalExpr that collects all column references to determine what columns the expression needs to be evaluated.
+struct FilterSchemaBuilder<'schema> {
+    filter_schema_fields: BTreeSet<Arc<Field>>,
+    file_schema: &'schema Schema,
+    table_schema: &'schema Schema,
+}
+
+impl<'schema> FilterSchemaBuilder<'schema> {
+    fn new(file_schema: &'schema Schema, table_schema: &'schema Schema) -> Self {
+        Self {
+            filter_schema_fields: BTreeSet::new(),
+            file_schema,
+            table_schema,
+        }
+    }
+
+    fn sort_fields(fields: &mut Vec<Arc<Field>>, table_schema: &Schema, file_schema: &Schema) {
+        fields.sort_by_key(|f| f.name().to_string());
+        fields.dedup_by_key(|f| f.name().to_string());
+        fields.sort_by_key(|f| {
+            let table_schema_index = table_schema.index_of(f.name()).unwrap_or(std::usize::MAX);
+            let file_schema_index = file_schema.index_of(f.name()).unwrap_or(std::usize::MAX);
+            (table_schema_index, file_schema_index)
+        });
+    }
+
+    fn build(self) -> SchemaRef {
+        let mut fields = self.filter_schema_fields.into_iter().collect::<Vec<_>>();
+        FilterSchemaBuilder::sort_fields(&mut fields, self.table_schema, self.file_schema);
+        Arc::new(Schema::new(fields))
+    }
+}
+
+
+impl<'node> TreeNodeVisitor<'node> for FilterSchemaBuilder<'_> {
+    type Node = Arc<dyn PhysicalExpr>;
+
+    fn f_down(
+        &mut self,
+        node: &'node Arc<dyn PhysicalExpr>,
+    ) -> Result<TreeNodeRecursion> {
+        if let Some(column) = node.as_any().downcast_ref::<Column>() {
+            if let Ok(field) = self.table_schema.field_with_name(column.name()) {
+                self.filter_schema_fields.insert(Arc::new(field.clone()));
+            } else if let Ok(field) = self.file_schema.field_with_name(column.name()) {
+                self.filter_schema_fields.insert(Arc::new(field.clone()));
+            } else {
+                // valid fields are the table schema's fields + the file schema's fields, preferring the table schema's fields when there is a conflict
+                let mut valid_fields = self.table_schema.fields().iter().chain(self.file_schema.fields().iter()).cloned().collect::<Vec<_>>();
+                FilterSchemaBuilder::sort_fields(&mut valid_fields, self.table_schema, self.file_schema);
+                let valid_fields = valid_fields.into_iter().map(|f| datafusion_common::Column::new_unqualified(f.name())).collect();
+                let field = datafusion_common::Column::new_unqualified(column.name());
+                return Err(
+                    datafusion_common::DataFusionError::SchemaError(
+                        SchemaError::FieldNotFound { field: Box::new(field), valid_fields },
+                        Box::new(None),
+                    )
+                )
+            }
+        }
+
+        Ok(TreeNodeRecursion::Continue)
+    }
+}

datafusion/datasource-parquet/src/row_filter.rs

@@ -274,7 +274,7 @@ impl<'a> FilterCandidateBuilder<'a> {
 
         let required_bytes = size_of_columns(&required_indices, metadata)?;
         let can_use_index = columns_sorted(&required_indices, metadata)?;
-
+    
         Ok(Some(FilterCandidate {
             expr: rewritten_expr,
             required_bytes,

datafusion/datasource/src/schema_adapter.rs

@@ -395,7 +395,7 @@ impl SchemaMapper for SchemaMapping {
             .fields()
             .iter()
             .zip(batch_cols.iter())
-            .map(|(field, batch_col)| {
+            .flat_map(|(field, batch_col)| {
                 self.table_schema
                     // try to get the same field from the table schema that we have stored in self
                     .field_with_name(field.name())
@@ -413,10 +413,6 @@ impl SchemaMapper for SchemaMapping {
                             // and if that works, return the field and column.
                             .map(|new_col| (new_col, table_field.clone()))
                     })
-                    .unwrap_or_else(|| {
-                        // If we don't have the field in the table schema return it as-is from the file schema
-                        Ok((batch_col.clone(), Arc::unwrap_or_clone(field.clone())))
-                    })
             })
             .collect::<Result<Vec<_>, _>>()?
             .into_iter()