From d8fe64db9df8b850314ee070b5ce60d18e5c6dea Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Tue, 4 Aug 2026 16:27:58 +0300 Subject: [PATCH 1/2] preserve empty projection --- datafusion/proto/src/physical_plan/mod.rs | 27 +++++------ .../tests/cases/roundtrip_physical_plan.rs | 45 +++++++++++++++++++ 2 files changed, 57 insertions(+), 15 deletions(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 79c6394933eae..c7d5bc9c4f4e5 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -1179,15 +1179,11 @@ pub trait PhysicalPlanNodeExt: Sized { })?; let schema: SchemaRef = SchemaRef::new(proto_schema.try_into()?); - let projection = if !scan.projection.is_empty() { - Some( - scan.projection - .iter() - .map(|i| *i as usize) - .collect::>(), - ) - } else { - None + // Preserve the empty-projection sentinel written by `try_from_data_source_exec`. + let projection = match scan.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), }; let mut sort_information = vec![]; @@ -2364,12 +2360,13 @@ pub trait PhysicalPlanNodeExt: Sized { let proto_schema: protobuf::Schema = source_conf.original_schema().as_ref().try_into()?; - let proto_projection = source_conf - .projection() - .as_ref() - .map_or_else(Vec::new, |v| { - v.iter().map(|x| *x as u32).collect::>() - }); + // Proto3 can't tell `None` from `Some(vec![])`; encode the latter + // as the `[u32::MAX]` sentinel, matching the join/filter nodes. + let proto_projection = match source_conf.projection().as_ref() { + None => Vec::new(), + Some(v) if v.is_empty() => vec![u32::MAX], + Some(v) => v.iter().map(|x| *x as u32).collect(), + }; let proto_sort_information = source_conf .sort_information() diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 19a5ca337d7f6..aed45951a667a 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -3227,6 +3227,51 @@ async fn roundtrip_memory_source() -> Result<()> { roundtrip_test(plan) } +#[tokio::test] +async fn roundtrip_memory_source_empty_projection() -> Result<()> { + use datafusion::datasource::memory::MemorySourceConfig; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(arrow::array::Int64Array::from(vec![1, 2])), + Arc::new(arrow::array::StringArray::from(vec!["x", "y"])), + ], + )?; + + // `Some(vec![])` projects away every column, which is not the same as + // `None` (keep all columns) + let source = + MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), Some(vec![]))?; + let plan = DataSourceExec::from_data_source(source); + assert_eq!(plan.schema().fields().len(), 0); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let decoded = roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; + + let decoded_source = decoded + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("Expected DataSourceExec after roundtrip") + })? + .data_source() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("Expected MemorySourceConfig after roundtrip") + })? + .clone(); + + assert_eq!(decoded_source.projection(), &Some(vec![])); + assert_eq!(decoded.schema().fields().len(), 0); + Ok(()) +} + #[tokio::test] async fn roundtrip_listing_table_with_schema_metadata() -> Result<()> { let ctx = SessionContext::new(); From 6e5213fbdba12a3b1a8a8cc56b179aece5963106 Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Tue, 4 Aug 2026 16:35:41 +0300 Subject: [PATCH 2/2] use roundtrip_test_sql_with_context in test --- .../tests/cases/roundtrip_physical_plan.rs | 64 ++++++------------- 1 file changed, 19 insertions(+), 45 deletions(-) diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index aed45951a667a..e1592186ef49a 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -2678,6 +2678,25 @@ async fn roundtrip_empty_projection() -> Result<()> { roundtrip_test_sql_with_context(sql, &ctx).await } +#[tokio::test] +async fn roundtrip_memory_source_empty_projection() -> Result<()> { + // Memory scan: `Some(vec![])` must not decode back as `None` + let ctx = SessionContext::new(); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Int64, false), + ])), + vec![ + Arc::new(arrow::array::StringArray::from(vec!["Tom"])), + Arc::new(arrow::array::Int64Array::from(vec![18i64])), + ], + )?; + ctx.register_batch("tmem", batch)?; + let sql = "select 1 from tmem"; + roundtrip_test_sql_with_context(sql, &ctx).await +} + #[tokio::test] async fn roundtrip_physical_plan_node() { use datafusion::prelude::*; @@ -3227,51 +3246,6 @@ async fn roundtrip_memory_source() -> Result<()> { roundtrip_test(plan) } -#[tokio::test] -async fn roundtrip_memory_source_empty_projection() -> Result<()> { - use datafusion::datasource::memory::MemorySourceConfig; - - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Utf8, false), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(arrow::array::Int64Array::from(vec![1, 2])), - Arc::new(arrow::array::StringArray::from(vec!["x", "y"])), - ], - )?; - - // `Some(vec![])` projects away every column, which is not the same as - // `None` (keep all columns) - let source = - MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), Some(vec![]))?; - let plan = DataSourceExec::from_data_source(source); - assert_eq!(plan.schema().fields().len(), 0); - - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - let decoded = roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; - - let decoded_source = decoded - .downcast_ref::() - .ok_or_else(|| { - internal_datafusion_err!("Expected DataSourceExec after roundtrip") - })? - .data_source() - .downcast_ref::() - .ok_or_else(|| { - internal_datafusion_err!("Expected MemorySourceConfig after roundtrip") - })? - .clone(); - - assert_eq!(decoded_source.projection(), &Some(vec![])); - assert_eq!(decoded.schema().fields().len(), 0); - Ok(()) -} - #[tokio::test] async fn roundtrip_listing_table_with_schema_metadata() -> Result<()> { let ctx = SessionContext::new();