Minor: Add more projection pushdown tests, clarify comments (#14963)

* Minor: Add more projection pushdown tests

* Improve comment and indenting

Author: alamb

datafusion/core/tests/physical_optimizer/projection_pushdown.rs

@@ -31,7 +31,7 @@ use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_expr::{Operator, ScalarUDF, ScalarUDFImpl, Signature, Volatility};
 use datafusion_physical_expr::expressions::{
-    binary, col, BinaryExpr, CaseExpr, CastExpr, Column, Literal, NegativeExpr,
+    binary, cast, col, BinaryExpr, CaseExpr, CastExpr, Column, Literal, NegativeExpr,
 };
 use datafusion_physical_expr::ScalarFunctionExpr;
 use datafusion_physical_expr::{
@@ -1383,28 +1383,29 @@ fn test_union_after_projection() -> Result<()> {
     Ok(())
 }
 
-#[test]
-fn test_partition_col_projection_pushdown() -> Result<()> {
+/// Returns a DataSourceExec that scans a file with (int_col, string_col)
+/// and has a partitioning column partition_col (Utf8)
+fn partitioned_data_source() -> Arc<DataSourceExec> {
     let file_schema = Arc::new(Schema::new(vec![
         Field::new("int_col", DataType::Int32, true),
         Field::new("string_col", DataType::Utf8, true),
     ]));
 
-    let partitioned_schema = Arc::new(Schema::new(vec![
-        Field::new("int_col", DataType::Int32, true),
-        Field::new("string_col", DataType::Utf8, true),
-        Field::new("partition_col", DataType::Utf8, true),
-    ]));
-
-    let source = FileScanConfig::new(
+    FileScanConfig::new(
         ObjectStoreUrl::parse("test:///").unwrap(),
         file_schema.clone(),
         Arc::new(CsvSource::default()),
     )
     .with_file(PartitionedFile::new("x".to_string(), 100))
     .with_table_partition_cols(vec![Field::new("partition_col", DataType::Utf8, true)])
     .with_projection(Some(vec![0, 1, 2]))
-    .build();
+    .build()
+}
+
+#[test]
+fn test_partition_col_projection_pushdown() -> Result<()> {
+    let source = partitioned_data_source();
+    let partitioned_schema = source.schema();
 
     let projection = Arc::new(ProjectionExec::try_new(
         vec![
@@ -1427,8 +1428,50 @@ fn test_partition_col_projection_pushdown() -> Result<()> {
     let after_optimize =
         ProjectionPushdown::new().optimize(projection, &ConfigOptions::new())?;
 
-    let expected = ["ProjectionExec: expr=[string_col@1 as string_col, partition_col@2 as partition_col, int_col@0 as int_col]",
-        "  DataSourceExec: file_groups={1 group: [[x]]}, projection=[int_col, string_col, partition_col], file_type=csv, has_header=false"];
+    let expected = [
+        "ProjectionExec: expr=[string_col@1 as string_col, partition_col@2 as partition_col, int_col@0 as int_col]",
+        "  DataSourceExec: file_groups={1 group: [[x]]}, projection=[int_col, string_col, partition_col], file_type=csv, has_header=false"
+    ];
+    assert_eq!(get_plan_string(&after_optimize), expected);
+
+    Ok(())
+}
+
+#[test]
+fn test_partition_col_projection_pushdown_expr() -> Result<()> {
+    let source = partitioned_data_source();
+    let partitioned_schema = source.schema();
+
+    let projection = Arc::new(ProjectionExec::try_new(
+        vec![
+            (
+                col("string_col", partitioned_schema.as_ref())?,
+                "string_col".to_string(),
+            ),
+            (
+                // CAST(partition_col, Utf8View)
+                cast(
+                    col("partition_col", partitioned_schema.as_ref())?,
+                    partitioned_schema.as_ref(),
+                    DataType::Utf8View,
+                )?,
+                "partition_col".to_string(),
+            ),
+            (
+                col("int_col", partitioned_schema.as_ref())?,
+                "int_col".to_string(),
+            ),
+        ],
+        source,
+    )?);
+
+    let after_optimize =
+        ProjectionPushdown::new().optimize(projection, &ConfigOptions::new())?;
+
+    let expected = [
+        "ProjectionExec: expr=[string_col@1 as string_col, CAST(partition_col@2 AS Utf8View) as partition_col, int_col@0 as int_col]",
+        "  DataSourceExec: file_groups={1 group: [[x]]}, projection=[int_col, string_col, partition_col], file_type=csv, has_header=false"
+    ];
     assert_eq!(get_plan_string(&after_optimize), expected);
 
     Ok(())

datafusion/datasource/src/file_scan_config.rs

@@ -264,35 +264,35 @@ impl DataSource for FileScanConfig {
         &self,
         projection: &ProjectionExec,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-        // If there is any non-column or alias-carrier expression, Projection should not be removed.
         // This process can be moved into CsvExec, but it would be an overlap of their responsibility.
 
+        // Must be all column references, with no table partition columns (which can not be projected)
         let partitioned_columns_in_proj = projection.expr().iter().any(|(expr, _)| {
             expr.as_any()
                 .downcast_ref::<Column>()
                 .map(|expr| expr.index() >= self.file_schema.fields().len())
                 .unwrap_or(false)
         });
 
-        Ok(
-            (all_alias_free_columns(projection.expr()) && !partitioned_columns_in_proj)
-                .then(|| {
-                    let file_scan = self.clone();
-                    let source = Arc::clone(&file_scan.file_source);
-                    let new_projections = new_projections_for_columns(
-                        projection,
-                        &file_scan
-                            .projection
-                            .clone()
-                            .unwrap_or((0..self.file_schema.fields().len()).collect()),
-                    );
-                    file_scan
-                        // Assign projected statistics to source
-                        .with_projection(Some(new_projections))
-                        .with_source(source)
-                        .build() as _
-                }),
-        )
+        // If there is any non-column or alias-carrier expression, Projection should not be removed.
+        let no_aliases = all_alias_free_columns(projection.expr());
+
+        Ok((no_aliases && !partitioned_columns_in_proj).then(|| {
+            let file_scan = self.clone();
+            let source = Arc::clone(&file_scan.file_source);
+            let new_projections = new_projections_for_columns(
+                projection,
+                &file_scan
+                    .projection
+                    .clone()
+                    .unwrap_or((0..self.file_schema.fields().len()).collect()),
+            );
+            file_scan
+                // Assign projected statistics to source
+                .with_projection(Some(new_projections))
+                .with_source(source)
+                .build() as _
+        }))
     }
 }