From b2b81d3f169f64cb3afe49608faffe49f4421791 Mon Sep 17 00:00:00 2001 From: Radu Stoenescu Date: Tue, 14 Jul 2026 15:47:07 +0300 Subject: [PATCH] [HSTACK] fix: reassign predicate column indices by name before building pruning predicates A filter's Column can carry an index that is only valid for the schema it was originally built against (e.g. a filter inferred across a join keeps the other side's index). ParquetSource::try_pushdown_filters and other producers store such predicates as-is, so build_pruning_predicates (used both for real per-file pruning and for ParquetSource::fmt_extra's Display output) could receive a Column whose index no longer matches its own name in file_schema. In debug builds this can panic: PhysicalExprSimplifier::simplify runs a debug-only type-preservation assertion that calls Column::data_type() using the raw index, which can resolve to an unrelated column (e.g. a struct-typed one), producing a type-coercion error that gets unwrapped. This is most visible via fmt_extra/Display (EXPLAIN, logging, or any code that displays the plan), since it does not go through per-file schema adaptation the way the real execution path does. Fix build_pruning_predicates to reassign every Column's index by name against file_schema before use, via the existing reassign_expr_columns helper (already used elsewhere in this crate for the same purpose). Adds a regression test reproducing the panic with real parquet-backed ParquetSources joined via HashJoinExec, using a DynamicFilterPhysicalExpr to match the shape of the predicate that triggers PruningPredicate::try_new's snapshot/simplify path. --- datafusion/datasource-parquet/src/opener.rs | 4 + datafusion/datasource-parquet/src/source.rs | 141 ++++++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/datafusion/datasource-parquet/src/opener.rs b/datafusion/datasource-parquet/src/opener.rs index 2e67658478800..f802653718b76 100644 --- a/datafusion/datasource-parquet/src/opener.rs +++ b/datafusion/datasource-parquet/src/opener.rs @@ -1100,6 +1100,10 @@ pub(crate) fn build_pruning_predicates( let Some(predicate) = predicate.as_ref() else { return (None, None); }; + // Column indices may not match file_schema (e.g. filters inferred across a + // join keep the other side's index); reassign by name before use. + let predicate = &reassign_expr_columns(Arc::clone(predicate), file_schema.as_ref()) + .unwrap_or_else(|_| Arc::clone(predicate)); let pruning_predicate = build_pruning_predicate( Arc::clone(predicate), file_schema, diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 6a16152bb5838..75b46db7df2be 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -923,4 +923,145 @@ mod tests { assert!(source.reverse_row_groups()); assert!(source.filter().is_some()); } + + /// Reproduces a crash where `Display`ing a plan (as `EXPLAIN` or any logging + /// code would) panics inside `ParquetSource::fmt_extra`. + /// + /// The scenario: two parquet-backed tables share a leading struct-typed column + /// (e.g. an audit/system column) followed by a join key of the same name on + /// both sides. A join's filter-pushdown/inference can produce, for one side, a + /// filter whose `Column` has the right *name* but an index that only matches + /// the *other* side's schema layout - here, index 0, which is `meta` (a + /// struct) on this side, not `MERGE_POLICY_ID`. `fmt_extra` used to build a + /// `PruningPredicate` straight from that mismatched `Column` and panic + /// comparing `Struct(..) = Float64`. + #[tokio::test] + async fn test_display_does_not_panic_on_predicate_with_mismatched_column_index() { + use arrow::array::{ArrayRef, Float64Array, StringArray, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use arrow::record_batch::RecordBatch; + use bytes::{BufMut, BytesMut}; + use datafusion_common::{JoinType, NullEquality}; + use datafusion_datasource::PartitionedFile; + use datafusion_datasource::file_scan_config::FileScanConfigBuilder; + use datafusion_datasource::source::DataSourceExec; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, binary, lit}; + use datafusion_physical_plan::displayable; + use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; + use datafusion_physical_plan::ExecutionPlan; + use object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path}; + use parquet::arrow::ArrowWriter; + + let struct_field = Field::new( + "meta", + DataType::Struct(Fields::from(vec![Field::new( + "id", + DataType::Utf8, + true, + )])), + true, + ); + let key_field = Field::new("MERGE_POLICY_ID", DataType::Float64, true); + let schema = Arc::new(Schema::new(vec![struct_field, key_field])); + + let make_batch = |ids: Vec| -> RecordBatch { + let meta_id_field = Arc::new(Field::new("id", DataType::Utf8, true)); + let meta_ids: ArrayRef = + Arc::new(StringArray::from(vec!["a".to_string(); ids.len()])); + let meta: ArrayRef = Arc::new(StructArray::from(vec![(meta_id_field, meta_ids)])); + let keys: ArrayRef = Arc::new(Float64Array::from(ids)); + RecordBatch::try_new(Arc::clone(&schema), vec![meta, keys]).unwrap() + }; + + async fn write_parquet( + store: Arc, + filename: &str, + batch: RecordBatch, + ) -> u64 { + let mut out = BytesMut::new().writer(); + { + let mut writer = ArrowWriter::try_new(&mut out, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + let data = out.into_inner().freeze(); + let len = data.len() as u64; + store.put(&Path::from(filename), data.into()).await.unwrap(); + len + } + + let store: Arc = Arc::new(InMemory::new()); + let left_len = write_parquet(Arc::clone(&store), "left.parquet", make_batch(vec![42.0])).await; + let right_len = + write_parquet(Arc::clone(&store), "right.parquet", make_batch(vec![42.0])).await; + let object_store_url = ObjectStoreUrl::local_filesystem(); + + // Left side: the literal `WHERE MERGE_POLICY_ID = 42`, correctly indexed (1). + let left_predicate = binary( + Arc::new(Column::new("MERGE_POLICY_ID", 1)), + Operator::Eq, + lit(42.0f64), + &schema, + ) + .unwrap(); + let left_source = + Arc::new(ParquetSource::new(Arc::clone(&schema)).with_predicate(left_predicate)); + let left_config = FileScanConfigBuilder::new(object_store_url.clone(), left_source) + .with_file(PartitionedFile::new("left.parquet", left_len)) + .build(); + let left_exec: Arc = DataSourceExec::from_data_source(left_config); + + // Right side: what a join's transitive-filter inference produces from + // `L.MERGE_POLICY_ID = R.MERGE_POLICY_ID AND L.MERGE_POLICY_ID = 42` - same + // column name, but index 0, which is only valid for `L`'s layout (here it + // means `meta`, a struct, on `R`'s own schema). + let right_predicate = binary( + Arc::new(Column::new("MERGE_POLICY_ID", 0)), + Operator::Eq, + lit(42.0f64), + &schema, + ) + .unwrap(); + // Wrapping in a `DynamicFilterPhysicalExpr` matches the real bug shape (the + // predicate carried a join-side dynamic filter alongside the mismatched + // comparison) and is what makes `PruningPredicate::try_new` take the + // snapshot path that runs the vulnerable, unrewritten simplifier pass. + let right_predicate = Arc::new(DynamicFilterPhysicalExpr::new( + right_predicate.children().into_iter().map(Arc::clone).collect(), + right_predicate, + )) as Arc; + let right_source = + Arc::new(ParquetSource::new(Arc::clone(&schema)).with_predicate(right_predicate)); + let right_config = FileScanConfigBuilder::new(object_store_url, right_source) + .with_file(PartitionedFile::new("right.parquet", right_len)) + .build(); + let right_exec: Arc = DataSourceExec::from_data_source(right_config); + + let on = vec![( + Arc::new(Column::new("MERGE_POLICY_ID", 1)) as Arc, + Arc::new(Column::new("MERGE_POLICY_ID", 1)) as Arc, + )]; + let join = Arc::new( + HashJoinExec::try_new( + left_exec, + right_exec, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + + // This must not panic: Displaying the plan (as EXPLAIN, logging, or any + // bookkeeping code would) used to crash inside `ParquetSource::fmt_extra` + // while building a pruning predicate from the right side's mismatched + // `Column`. + let _ = displayable(join.as_ref()).indent(true).to_string(); + } }