apache / apache/datafusion

`recompute_schema()` is broken with `LogicalPlan::Union`

Open
#22,447 6 comments 0 reactions 1 assignee Claimed by @Brijesh-Thakkar View on GitHub
bug
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

### Describe the bug

`recompute_schema()` has a pretty well documented contract that if you change the schema that it should be recomputed.

However, the current check for `Union` only checks the *length* of the schema, not the types or names/aliases:

From `recompute_schema()`:

```rust
LogicalPlan::Union(Union { inputs, schema }) => {
let first_input_schema = inputs[0].schema();
if schema.fields().len() == first_input_schema.fields().len() {
// If inputs are not pruned do not change schema
Ok(LogicalPlan::Union(Union { inputs, schema }))
} else {
// A note on `Union`s constructed via `try_new_by_name`:
//
// At this point, the schema for each input should have
// the same width. Thus, we do not need to save whether a
// `Union` was created `BY NAME`, and can safely rely on the
// `try_new` initializer to derive the new schema based on
// column positions.
Ok(LogicalPlan::Union(Union::try_new(inputs)?))
}
}
```

### To Reproduce

```rust
#[test]
fn test_recompute_schema_union_type_mismatch() -> Result<()> {
let schema_i32 =
Schema::new(vec![Field::new("a", DataType::Int32, false)]);
let schema_i64 =
Schema::new(vec![Field::new("a", DataType::Int64, false)]);

// Build a Union whose schema starts out as Int32 (matching its inputs).
let original = Union::try_new(vec![
Arc::new(table_scan(Some("t1"), &schema_i32, None)?.build()?),
Arc::new(table_scan(Some("t2"), &schema_i32, None)?.build()?),
])?;
assert_eq!(
original.schema.field(0).data_type(),
&DataType::Int32,
"sanity: starting schema is Int32"
);

// Now simulate something like a type-coercion pass replacing the
// inputs with Int64-typed versions while leaving the Union's cached
// schema untouched. Same width, different types.
let stale = LogicalPlan::Union(Union {
inputs: vec![
Arc::new(table_scan(Some("t1"), &schema_i64, None)?.build()?),
Arc::new(table_scan(Some("t2"), &schema_i64, None)?.build()?),
],
schema: Arc::clone(&original.schema),
});

let recomputed = stale.recompute_schema()?;

assert_eq!(
recomputed.schema().field(0).data_type(),
&DataType::Int64,
"Union schema should track the new Int64 input types after \
recompute_schema(), but the width-only check left it stale"
);

Ok(())
}
```

### Expected behavior

The schema is updated

### Additional context

I'm trying to rewrite all exprs/aliases using a pattern for some optimizations. Rewriting exprs to use a new alias, even if it's consistent, you need to recompute the schema for *all* nodes.

Union is the only one that breaks that I've found but there might be others.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.