fix in list round trip in df proto (#16744)

Co-authored-by: Andrew Lamb <[email protected]>

Author: XiangpengHao

datafusion/proto/src/physical_plan/mod.rs

@@ -39,7 +39,7 @@ use crate::protobuf::{
 use crate::{convert_required, into_required};
 
 use datafusion::arrow::compute::SortOptions;
-use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::arrow::datatypes::{Schema, SchemaRef};
 use datafusion::datasource::file_format::csv::CsvSink;
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
 use datafusion::datasource::file_format::json::JsonSink;
@@ -611,14 +611,38 @@ impl protobuf::PhysicalPlanNode {
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let input: Arc<dyn ExecutionPlan> =
             into_physical_plan(&filter.input, registry, runtime, extension_codec)?;
+        let projection = if !filter.projection.is_empty() {
+            Some(
+                filter
+                    .projection
+                    .iter()
+                    .map(|i| *i as usize)
+                    .collect::<Vec<_>>(),
+            )
+        } else {
+            None
+        };
+
+        // Use the projected schema if projection is present, otherwise use the full schema
+        let predicate_schema = if let Some(ref proj_indices) = projection {
+            // Create projected schema for parsing the predicate
+            let projected_fields: Vec<_> = proj_indices
+                .iter()
+                .map(|&i| input.schema().field(i).clone())
+                .collect();
+            Arc::new(Schema::new(projected_fields))
+        } else {
+            input.schema()
+        };
+
         let predicate = filter
             .expr
             .as_ref()
             .map(|expr| {
                 parse_physical_expr(
                     expr,
                     registry,
-                    input.schema().as_ref(),
+                    predicate_schema.as_ref(),
                     extension_codec,
                 )
             })
@@ -629,17 +653,6 @@ impl protobuf::PhysicalPlanNode {
                 )
             })?;
         let filter_selectivity = filter.default_filter_selectivity.try_into();
-        let projection = if !filter.projection.is_empty() {
-            Some(
-                filter
-                    .projection
-                    .iter()
-                    .map(|i| *i as usize)
-                    .collect::<Vec<_>>(),
-            )
-        } else {
-            None
-        };
         let filter =
             FilterExec::try_new(predicate, input)?.with_projection(projection)?;
         match filter_selectivity {
@@ -727,11 +740,31 @@ impl protobuf::PhysicalPlanNode {
         {
             let schema =
                 parse_protobuf_file_scan_schema(scan.base_conf.as_ref().unwrap())?;
+
+            // Check if there's a projection and use projected schema for predicate parsing
+            let base_conf = scan.base_conf.as_ref().unwrap();
+            let predicate_schema = if !base_conf.projection.is_empty() {
+                // Create projected schema for parsing the predicate
+                let projected_fields: Vec<_> = base_conf
+                    .projection
+                    .iter()
+                    .map(|&i| schema.field(i as usize).clone())
+                    .collect();
+                Arc::new(Schema::new(projected_fields))
+            } else {
+                schema
+            };
+
             let predicate = scan
                 .predicate
                 .as_ref()
                 .map(|expr| {
-                    parse_physical_expr(expr, registry, schema.as_ref(), extension_codec)
+                    parse_physical_expr(
+                        expr,
+                        registry,
+                        predicate_schema.as_ref(),
+                        extension_codec,
+                    )
                 })
                 .transpose()?;
             let mut options = TableParquetOptions::default();
@@ -745,7 +778,7 @@ impl protobuf::PhysicalPlanNode {
                 source = source.with_predicate(predicate);
             }
             let base_config = parse_protobuf_file_scan_config(
-                scan.base_conf.as_ref().unwrap(),
+                base_conf,
                 registry,
                 extension_codec,
                 Arc::new(source),

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

@@ -1739,7 +1739,7 @@ async fn roundtrip_physical_plan_node() {
 }
 
 // Failing due to https://github.com/apache/datafusion/pull/16662
-#[ignore]
+// Fixed: Column index mismatch during protobuf deserialization
 #[tokio::test]
 async fn test_tpch_part_in_list_query_with_real_parquet_data() -> Result<()> {
     // Test the specific query: SELECT p_size FROM part WHERE p_size IN (14, 6, 5, 31)