Skip to content

Commit 215ebb8

Browse files
authored
Fix union equivalence schema rewrite with stale constants (#23375)
## Which issue does this PR close? - Closes #23374. ## Rationale for this change `UnionExec::try_new` can panic while computing equivalence properties if stale constant metadata is carried across a projection and then rewritten to the union output schema. In the observed shape, a filter such as `ticker = 'ESU6'` can leave a uniform string constant in equivalence properties. After a parent projection drops `ticker`, union property schema rewriting can see the remaining column slot as a timestamp column and attempt to cast `'ESU6'` to `Timestamp`, which fails during planning. Equivalence constants are optimizer metadata, so an unrepresentable constant after schema rewrite should be discarded rather than failing query planning. ## What changes are included in this PR? - Drops a uniform constant during `EquivalenceProperties::with_new_schema` if its value cannot be cast to the rewritten expression type. - Removes trivial equivalence classes after dropping such constants. - Propagates `UnionExec::compute_properties` errors from `UnionExec::try_new` instead of unwrapping. - Adds a regression test for union equivalence schema rewrite with an unrepresentable stale constant value. ## Are these changes tested? Yes: ## Are there any user-facing changes? No API change. This prevents a planner panic for affected `UNION ALL` + filter + projection query shapes.
1 parent b790763 commit 215ebb8

3 files changed

Lines changed: 44 additions & 5 deletions

File tree

datafusion/physical-expr/src/equivalence/properties/mod.rs

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1314,7 +1314,18 @@ impl EquivalenceProperties {
13141314
if let (Some(data_type), Some(AcrossPartitions::Uniform(Some(value)))) =
13151315
(data_type, &mut eq_class.constant)
13161316
{
1317-
*value = value.cast_to(&data_type)?;
1317+
match value.cast_to(&data_type) {
1318+
Ok(cast_value) => *value = cast_value,
1319+
Err(_) => {
1320+
// This is optimizer metadata. If a stale constant
1321+
// value cannot be represented after schema rewrite,
1322+
// drop the constant instead of failing planning.
1323+
eq_class.constant = None;
1324+
}
1325+
}
1326+
}
1327+
if eq_class.is_trivial() {
1328+
continue;
13181329
}
13191330
eq_classes.push(eq_class);
13201331
}

datafusion/physical-expr/src/equivalence/properties/union.rs

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ mod tests {
311311
use crate::equivalence::tests::{create_test_schema, parse_sort_expr};
312312
use crate::expressions::col;
313313

314-
use arrow::datatypes::{DataType, Field, Schema};
314+
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
315315
use datafusion_common::ScalarValue;
316316

317317
use itertools::Itertools;
@@ -899,6 +899,36 @@ mod tests {
899899
Ok(())
900900
}
901901

902+
#[test]
903+
fn test_union_drops_unrepresentable_constant_value_after_schema_rewrite() -> Result<()>
904+
{
905+
let input_schema = Arc::new(Schema::new(vec![Field::new(
906+
"ticker",
907+
DataType::Timestamp(TimeUnit::Nanosecond, None),
908+
true,
909+
)]));
910+
let output_schema = Arc::new(Schema::new(vec![Field::new(
911+
"timestamp",
912+
DataType::Timestamp(TimeUnit::Nanosecond, None),
913+
true,
914+
)]));
915+
916+
let ticker = col("ticker", &input_schema)?;
917+
let stale_value = ScalarValue::Utf8(Some("ESU6".to_owned()));
918+
let const_expr = ConstExpr::new(
919+
Arc::clone(&ticker),
920+
AcrossPartitions::Uniform(Some(stale_value)),
921+
);
922+
923+
let mut input = EquivalenceProperties::new(input_schema);
924+
input.add_constants(vec![const_expr])?;
925+
926+
let union_props = calculate_union(vec![input], output_schema)?;
927+
assert!(union_props.constants().is_empty());
928+
929+
Ok(())
930+
}
931+
902932
/// Return a new schema with the same types, but new field names
903933
///
904934
/// The new field names are the old field names with `text` appended.

datafusion/physical-plan/src/union.rs

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -130,9 +130,7 @@ impl UnionExec {
130130
// The schema of the inputs and the union schema is consistent when:
131131
// - They have the same number of fields, and
132132
// - Their fields have same types at the same indices.
133-
// Here, we know that schemas are consistent and the call below can
134-
// not return an error.
135-
let cache = Self::compute_properties(&inputs, schema).unwrap();
133+
let cache = Self::compute_properties(&inputs, schema)?;
136134
Ok(Arc::new(UnionExec {
137135
inputs,
138136
metrics: ExecutionPlanMetricsSet::new(),

0 commit comments

Comments
 (0)