Push down projection expressions into ParquetOpener (#19111)

Closes #14993

Once this is merged I think we can say we support projection expression
pushdown into scans and it is implemented for Parquet.

Remaining TODOs which I think should be tracked in other issues (I'll
find them or create them later):
- Support for pushing down struct field access into the Parquet scan
(#11745)
- Refactor other file sources that could push down expressions to use
the same approach as Parquet (Avro?)
- Replace remaining uses of `SchemaAdapter` with `PhysicalExprAdapter`
and decide if we want to actually deprecate `SchemaAdapter`

Author: adriangb

datafusion/datasource-parquet/src/opener.rs

@@ -23,7 +23,7 @@ use crate::{
     apply_file_schema_type_coercions, coerce_int96_to_resolution, row_filter,
     ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory,
 };
-use arrow::array::RecordBatch;
+use arrow::array::{RecordBatch, RecordBatchOptions};
 use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
 use datafusion_physical_expr::projection::ProjectionExprs;
 use datafusion_physical_expr::utils::reassign_expr_columns;
@@ -66,8 +66,8 @@ use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader};
 pub(super) struct ParquetOpener {
     /// Execution partition index
     pub partition_index: usize,
-    /// Column indexes in `table_schema` needed by the query
-    pub projection: Arc<[usize]>,
+    /// Projection to apply on top of the table schema (i.e. can reference partition columns).
+    pub projection: ProjectionExprs,
     /// Target number of rows in each output RecordBatch
     pub batch_size: usize,
     /// Optional limit on the number of rows to read
@@ -104,7 +104,7 @@ pub(super) struct ParquetOpener {
     #[cfg(feature = "parquet_encryption")]
     pub file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
     /// Rewrite expressions in the context of the file schema
-    pub(crate) expr_adapter_factory: Option<Arc<dyn PhysicalExprAdapterFactory>>,
+    pub(crate) expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,
     /// Optional factory to create file decryption properties dynamically
     #[cfg(feature = "parquet_encryption")]
     pub encryption_factory:
@@ -137,7 +137,7 @@ impl FileOpener for ParquetOpener {
 
         let batch_size = self.batch_size;
 
-        // Replace any partition column references in the predicate
+        // Build partition values map for replacing partition column references
         // with their literal values from this file's partition values.
         //
         // For example, given
@@ -157,29 +157,39 @@ impl FileOpener for ParquetOpener {
         // there are other cases where partition columns may appear in more
         // complex predicates that cannot be simplified until we are about to
         // open the file (such as dynamic predicates)
-        let mut predicate = {
-            // Only replace partition columns if there are partition values.
-            // For non-partitioned tables, skip the clone and replacement traversal.
-            if partitioned_file.partition_values.is_empty() {
-                self.predicate.clone()
-            } else {
-                let partition_values: HashMap<&str, &ScalarValue> = self
-                    .table_schema
-                    .table_partition_cols()
-                    .iter()
-                    .zip(partitioned_file.partition_values.iter())
-                    .map(|(field, value)| (field.name().as_str(), value))
-                    .collect();
-
-                self.predicate
-                    .clone()
-                    .map(|p| replace_columns_with_literals(p, &partition_values))
-                    .transpose()?
-            }
-        };
+        let partition_values: HashMap<&str, &ScalarValue> = self
+            .table_schema
+            .table_partition_cols()
+            .iter()
+            .zip(partitioned_file.partition_values.iter())
+            .map(|(field, value)| (field.name().as_str(), value))
+            .collect();
+
+        // Calculate the output schema from the original projection (before literal replacement)
+        // so we get correct field names from column references
+        let logical_file_schema = Arc::clone(self.table_schema.file_schema());
+        let output_schema = Arc::new(
+            self.projection
+                .project_schema(self.table_schema.table_schema())?,
+        );
 
-        let projection = Arc::clone(&self.projection);
-        let logical_file_schema = Arc::clone(self.table_schema.table_schema());
+        // Apply partition column replacement to projection expressions
+        let mut projection = self.projection.clone();
+        if !partition_values.is_empty() {
+            projection = projection.try_map_exprs(|expr| {
+                replace_columns_with_literals(Arc::clone(&expr), &partition_values)
+            })?;
+        }
+
+        // Apply partition column replacement to predicate
+        let mut predicate = if partition_values.is_empty() {
+            self.predicate.clone()
+        } else {
+            self.predicate
+                .clone()
+                .map(|p| replace_columns_with_literals(p, &partition_values))
+                .transpose()?
+        };
         let reorder_predicates = self.reorder_filters;
         let pushdown_filters = self.pushdown_filters;
         let force_filter_selections = self.force_filter_selections;
@@ -191,8 +201,7 @@ impl FileOpener for ParquetOpener {
         let predicate_creation_errors = MetricBuilder::new(&self.metrics)
             .global_counter("num_predicate_creation_errors");
 
-        let expr_adapter_factory = self.expr_adapter_factory.clone();
-        let mut predicate_file_schema = Arc::clone(self.table_schema.file_schema());
+        let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory);
 
         let enable_page_index = self.enable_page_index;
         #[cfg(feature = "parquet_encryption")]
@@ -311,30 +320,29 @@ impl FileOpener for ParquetOpener {
                 }
             }
 
-            // Adapt the predicate to the physical file schema.
+            // Adapt the projection & filter predicate to the physical file schema.
             // This evaluates missing columns and inserts any necessary casts.
-            if let Some(expr_adapter_factory) = expr_adapter_factory.as_ref() {
-                predicate = predicate
-                    .map(|p| {
-                        let expr = expr_adapter_factory
-                            .create(
-                                Arc::clone(&logical_file_schema),
-                                Arc::clone(&physical_file_schema),
-                            )
-                            .rewrite(p)?;
-                        // After rewriting to the file schema, further simplifications may be possible.
-                        // For example, if `'a' = col_that_is_missing` becomes `'a' = NULL` that can then be simplified to `FALSE`
-                        // and we can avoid doing any more work on the file (bloom filters, loading the page index, etc.).
-                        PhysicalExprSimplifier::new(&physical_file_schema).simplify(expr)
-                    })
-                    .transpose()?;
-                predicate_file_schema = Arc::clone(&physical_file_schema);
-            }
+            // After rewriting to the file schema, further simplifications may be possible.
+            // For example, if `'a' = col_that_is_missing` becomes `'a' = NULL` that can then be simplified to `FALSE`
+            // and we can avoid doing any more work on the file (bloom filters, loading the page index, etc.).
+            // Additionally, if any casts were inserted we can move casts from the column to the literal side:
+            // `CAST(col AS INT) = 5` can become `col = CAST(5 AS <col type>)`, which can be evaluated statically.
+            let rewriter = expr_adapter_factory.create(
+                Arc::clone(&logical_file_schema),
+                Arc::clone(&physical_file_schema),
+            );
+            let simplifier = PhysicalExprSimplifier::new(&physical_file_schema);
+            predicate = predicate
+                .map(|p| simplifier.simplify(rewriter.rewrite(p)?))
+                .transpose()?;
+            // Adapt projections to the physical file schema as well
+            projection = projection
+                .try_map_exprs(|p| simplifier.simplify(rewriter.rewrite(p)?))?;
 
             // Build predicates for this specific file
             let (pruning_predicate, page_pruning_predicate) = build_pruning_predicates(
                 predicate.as_ref(),
-                &predicate_file_schema,
+                &physical_file_schema,
                 &predicate_creation_errors,
             );
 
@@ -358,15 +366,6 @@ impl FileOpener for ParquetOpener {
                 reader_metadata,
             );
 
-            let mut projection =
-                ProjectionExprs::from_indices(&projection, &logical_file_schema);
-            if let Some(expr_adapter_factory) = expr_adapter_factory {
-                let adapter = expr_adapter_factory.create(
-                    Arc::clone(&logical_file_schema),
-                    Arc::clone(&physical_file_schema),
-                );
-                projection = projection.try_map_exprs(|expr| adapter.rewrite(expr))?;
-            }
             let indices = projection.column_indices();
 
             let mask = ProjectionMask::roots(builder.parquet_schema(), indices);
@@ -376,7 +375,6 @@ impl FileOpener for ParquetOpener {
                 let row_filter = row_filter::build_row_filter(
                     &predicate,
                     &physical_file_schema,
-                    &predicate_file_schema,
                     builder.metadata(),
                     reorder_predicates,
                     &file_metrics,
@@ -507,6 +505,9 @@ impl FileOpener for ParquetOpener {
             let predicate_cache_records = file_metrics.predicate_cache_records.clone();
 
             let stream_schema = Arc::clone(stream.schema());
+            // Check if we need to replace the schema to handle things like differing nullability or metadata.
+            // See note below about file vs. output schema.
+            let replace_schema = !stream_schema.eq(&output_schema);
 
             // Rebase column indices to match the narrowed stream schema.
             // The projection expressions have indices based on physical_file_schema,
@@ -517,13 +518,34 @@ impl FileOpener for ParquetOpener {
             let projector = projection.make_projector(&stream_schema)?;
 
             let stream = stream.map_err(DataFusionError::from).map(move |b| {
-                b.and_then(|b| {
+                b.and_then(|mut b| {
                     copy_arrow_reader_metrics(
                         &arrow_reader_metrics,
                         &predicate_cache_inner_records,
                         &predicate_cache_records,
                     );
-                    projector.project_batch(&b)
+                    b = projector.project_batch(&b)?;
+                    if replace_schema {
+                        // Ensure the output batch has the expected schema.
+                        // This handles things like schema level and field level metadata, which may not be present
+                        // in the physical file schema.
+                        // It is also possible for nullability to differ; some writers create files with
+                        // OPTIONAL fields even when there are no nulls in the data.
+                        // In these cases it may make sense for the logical schema to be `NOT NULL`.
+                        // RecordBatch::try_new_with_options checks that if the schema is NOT NULL
+                        // the array cannot contain nulls, amongst other checks.
+                        let (_stream_schema, arrays, num_rows) = b.into_parts();
+                        let options =
+                            RecordBatchOptions::new().with_row_count(Some(num_rows));
+                        RecordBatch::try_new_with_options(
+                            Arc::clone(&output_schema),
+                            arrays,
+                            &options,
+                        )
+                        .map_err(Into::into)
+                    } else {
+                        Ok(b)
+                    }
                 })
             });
 
@@ -826,7 +848,8 @@ mod test {
     use datafusion_datasource::{file_stream::FileOpener, PartitionedFile, TableSchema};
     use datafusion_expr::{col, lit};
     use datafusion_physical_expr::{
-        expressions::DynamicFilterPhysicalExpr, planner::logical2physical, PhysicalExpr,
+        expressions::DynamicFilterPhysicalExpr, planner::logical2physical,
+        projection::ProjectionExprs, PhysicalExpr,
     };
     use datafusion_physical_expr_adapter::DefaultPhysicalExprAdapterFactory;
     use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
@@ -910,7 +933,7 @@ mod test {
         let make_opener = |predicate| {
             ParquetOpener {
                 partition_index: 0,
-                projection: Arc::new([0, 1]),
+                projection: ProjectionExprs::from_indices(&[0, 1], &schema),
                 batch_size: 1024,
                 limit: None,
                 predicate: Some(predicate),
@@ -929,7 +952,7 @@ mod test {
                 coerce_int96: None,
                 #[cfg(feature = "parquet_encryption")]
                 file_decryption_properties: None,
-                expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
+                expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
                 max_predicate_cache_size: None,
@@ -978,7 +1001,7 @@ mod test {
         let make_opener = |predicate| {
             ParquetOpener {
                 partition_index: 0,
-                projection: Arc::new([0]),
+                projection: ProjectionExprs::from_indices(&[0], &file_schema),
                 batch_size: 1024,
                 limit: None,
                 predicate: Some(predicate),
@@ -1000,7 +1023,7 @@ mod test {
                 coerce_int96: None,
                 #[cfg(feature = "parquet_encryption")]
                 file_decryption_properties: None,
-                expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
+                expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
                 max_predicate_cache_size: None,
@@ -1065,7 +1088,7 @@ mod test {
         let make_opener = |predicate| {
             ParquetOpener {
                 partition_index: 0,
-                projection: Arc::new([0]),
+                projection: ProjectionExprs::from_indices(&[0], &file_schema),
                 batch_size: 1024,
                 limit: None,
                 predicate: Some(predicate),
@@ -1087,7 +1110,7 @@ mod test {
                 coerce_int96: None,
                 #[cfg(feature = "parquet_encryption")]
                 file_decryption_properties: None,
-                expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
+                expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
                 max_predicate_cache_size: None,
@@ -1155,7 +1178,7 @@ mod test {
         let make_opener = |predicate| {
             ParquetOpener {
                 partition_index: 0,
-                projection: Arc::new([0]),
+                projection: ProjectionExprs::from_indices(&[0], &file_schema),
                 batch_size: 1024,
                 limit: None,
                 predicate: Some(predicate),
@@ -1177,7 +1200,7 @@ mod test {
                 coerce_int96: None,
                 #[cfg(feature = "parquet_encryption")]
                 file_decryption_properties: None,
-                expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
+                expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
                 max_predicate_cache_size: None,
@@ -1253,7 +1276,7 @@ mod test {
         let make_opener = |predicate| {
             ParquetOpener {
                 partition_index: 0,
-                projection: Arc::new([0]),
+                projection: ProjectionExprs::from_indices(&[0], &file_schema),
                 batch_size: 1024,
                 limit: None,
                 predicate: Some(predicate),
@@ -1275,7 +1298,7 @@ mod test {
                 coerce_int96: None,
                 #[cfg(feature = "parquet_encryption")]
                 file_decryption_properties: None,
-                expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
+                expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
                 max_predicate_cache_size: None,