Skip to content

Commit ce2f153

Browse files
authored
fix(proto): preserve empty projection when ser/de MemoryScanExec (#24087)
## Which issue does this PR close? - Closes #24085 ## Rationale for this change Empty projection is not ser/de correctly in MemoryScanExec ## What changes are included in this PR? Use same approach as `FilterExec`, `HashJoinExec` etc for projection encode/decode ## Are these changes tested? Yes, new roundtrip test `roundtrip_memory_source_empty_projection`, which fails on main. ## Are there any user-facing changes? No
1 parent 373fab7 commit ce2f153

2 files changed

Lines changed: 31 additions & 15 deletions

File tree

datafusion/proto/src/physical_plan/mod.rs

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1179,15 +1179,11 @@ pub trait PhysicalPlanNodeExt: Sized {
11791179
})?;
11801180
let schema: SchemaRef = SchemaRef::new(proto_schema.try_into()?);
11811181

1182-
let projection = if !scan.projection.is_empty() {
1183-
Some(
1184-
scan.projection
1185-
.iter()
1186-
.map(|i| *i as usize)
1187-
.collect::<Vec<_>>(),
1188-
)
1189-
} else {
1190-
None
1182+
// Preserve the empty-projection sentinel written by `try_from_data_source_exec`.
1183+
let projection = match scan.projection.as_slice() {
1184+
[] => None,
1185+
[u32::MAX] => Some(Vec::new()),
1186+
indices => Some(indices.iter().map(|i| *i as usize).collect()),
11911187
};
11921188

11931189
let mut sort_information = vec![];
@@ -2364,12 +2360,13 @@ pub trait PhysicalPlanNodeExt: Sized {
23642360
let proto_schema: protobuf::Schema =
23652361
source_conf.original_schema().as_ref().try_into()?;
23662362

2367-
let proto_projection = source_conf
2368-
.projection()
2369-
.as_ref()
2370-
.map_or_else(Vec::new, |v| {
2371-
v.iter().map(|x| *x as u32).collect::<Vec<u32>>()
2372-
});
2363+
// Proto3 can't tell `None` from `Some(vec![])`; encode the latter
2364+
// as the `[u32::MAX]` sentinel, matching the join/filter nodes.
2365+
let proto_projection = match source_conf.projection().as_ref() {
2366+
None => Vec::new(),
2367+
Some(v) if v.is_empty() => vec![u32::MAX],
2368+
Some(v) => v.iter().map(|x| *x as u32).collect(),
2369+
};
23732370

23742371
let proto_sort_information = source_conf
23752372
.sort_information()

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2678,6 +2678,25 @@ async fn roundtrip_empty_projection() -> Result<()> {
26782678
roundtrip_test_sql_with_context(sql, &ctx).await
26792679
}
26802680

2681+
#[tokio::test]
2682+
async fn roundtrip_memory_source_empty_projection() -> Result<()> {
2683+
// Memory scan: `Some(vec![])` must not decode back as `None`
2684+
let ctx = SessionContext::new();
2685+
let batch = RecordBatch::try_new(
2686+
Arc::new(Schema::new(vec![
2687+
Field::new("a", DataType::Utf8, false),
2688+
Field::new("b", DataType::Int64, false),
2689+
])),
2690+
vec![
2691+
Arc::new(arrow::array::StringArray::from(vec!["Tom"])),
2692+
Arc::new(arrow::array::Int64Array::from(vec![18i64])),
2693+
],
2694+
)?;
2695+
ctx.register_batch("tmem", batch)?;
2696+
let sql = "select 1 from tmem";
2697+
roundtrip_test_sql_with_context(sql, &ctx).await
2698+
}
2699+
26812700
#[tokio::test]
26822701
async fn roundtrip_physical_plan_node() {
26832702
use datafusion::prelude::*;

0 commit comments

Comments
 (0)