Skip to content

Commit 42adf29

Browse files
adriangbclaude
andcommitted
fix: keep global join-key bounds index-aligned across partitions
`compute_global_bounds` dropped columns whose min/max could not be ordered across partitions by calling `flatten()`, which compacted the column vector and shifted every later column down an index. `create_bounds_predicate` pairs bounds with `on_right` positionally, so a shifted column was filtered against the wrong probe expression — a silent wrong-results bug. Represent dropped columns as `None` slots in `PartitionBounds` so the surviving columns keep their join-key index. Add a regression test covering a dropped middle column. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
1 parent b90896a commit 42adf29

1 file changed

Lines changed: 96 additions & 36 deletions

File tree

‎datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs‎

Lines changed: 96 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -61,22 +61,29 @@ impl ColumnBounds {
6161
}
6262
}
6363

64-
/// Represents the bounds for all join key columns from a single partition.
65-
/// This contains the min/max values computed from one partition's build-side data.
64+
/// Min/max bounds for join key columns, indexed by join key expression index.
65+
///
66+
/// A `None` slot marks a column with no usable bounds — e.g. one dropped from
67+
/// a global envelope because it was not totally ordered across partitions. The
68+
/// slot is kept (rather than removed) so that every later column keeps its
69+
/// join-key index; [`create_bounds_predicate`] pairs bounds with `on_right`
70+
/// positionally. Per-partition bounds built via [`PartitionBounds::new`] are
71+
/// always fully populated.
6672
#[derive(Debug, Clone)]
6773
pub(crate) struct PartitionBounds {
68-
/// Min/max bounds for each join key column in this partition.
69-
/// Index corresponds to the join key expression index.
70-
column_bounds: Vec<ColumnBounds>,
74+
column_bounds: Vec<Option<ColumnBounds>>,
7175
}
7276

7377
impl PartitionBounds {
78+
/// Build from dense per-partition bounds, where every column has bounds.
7479
pub(crate) fn new(column_bounds: Vec<ColumnBounds>) -> Self {
75-
Self { column_bounds }
80+
Self {
81+
column_bounds: column_bounds.into_iter().map(Some).collect(),
82+
}
7683
}
7784

7885
pub(crate) fn get_column_bounds(&self, index: usize) -> Option<&ColumnBounds> {
79-
self.column_bounds.get(index)
86+
self.column_bounds.get(index).and_then(Option::as_ref)
8087
}
8188
}
8289

@@ -198,48 +205,53 @@ fn combine_membership_and_bounds(
198205

199206
/// Compute the global (envelope) min/max bounds across a set of partition bounds.
200207
///
201-
/// For each column index, returns the smallest min seen and the largest max seen.
202-
/// Columns where any partition is missing bounds, or where bounds are not totally
203-
/// ordered (e.g. mixed-type comparisons), are dropped from the global envelope.
208+
/// For each join-key column, returns the smallest min and largest max seen
209+
/// across partitions. A column that is not totally ordered across partitions
210+
/// (e.g. a mixed-type comparison where `partial_cmp` yields `None`) becomes a
211+
/// `None` slot: dropped from the envelope but keeping its index, so the
212+
/// surviving columns stay aligned with `on_right`. Returns `None` when there
213+
/// are no partitions, the partitions disagree on column count, or no column
214+
/// survives.
204215
fn compute_global_bounds(per_partition: &[&PartitionBounds]) -> Option<PartitionBounds> {
205216
let mut iter = per_partition.iter();
206217
let first = iter.next()?;
207-
let mut acc: Vec<Option<ColumnBounds>> = first
208-
.column_bounds
209-
.iter()
210-
.map(|cb| Some(cb.clone()))
211-
.collect();
218+
let mut acc: Vec<Option<ColumnBounds>> = first.column_bounds.clone();
212219

213220
for partition in iter {
214221
if partition.column_bounds.len() != acc.len() {
215222
return None;
216223
}
217-
for (slot, cb) in acc.iter_mut().zip(partition.column_bounds.iter()) {
224+
for (slot, incoming) in acc.iter_mut().zip(partition.column_bounds.iter()) {
218225
let Some(existing) = slot.as_mut() else {
226+
continue; // already dropped by an earlier partition
227+
};
228+
let Some(cb) = incoming else {
229+
*slot = None;
219230
continue;
220231
};
221-
match cb.min.partial_cmp(&existing.min) {
222-
Some(std::cmp::Ordering::Less) => existing.min = cb.min.clone(),
223-
Some(_) => {}
224-
None => {
225-
*slot = None;
226-
continue;
232+
match (
233+
cb.min.partial_cmp(&existing.min),
234+
cb.max.partial_cmp(&existing.max),
235+
) {
236+
(Some(min_ord), Some(max_ord)) => {
237+
if min_ord == std::cmp::Ordering::Less {
238+
existing.min = cb.min.clone();
239+
}
240+
if max_ord == std::cmp::Ordering::Greater {
241+
existing.max = cb.max.clone();
242+
}
227243
}
228-
}
229-
match cb.max.partial_cmp(&existing.max) {
230-
Some(std::cmp::Ordering::Greater) => existing.max = cb.max.clone(),
231-
Some(_) => {}
232-
None => *slot = None,
244+
// Not totally ordered: drop just this column, keep the slot
245+
// so later columns stay at their join-key index.
246+
_ => *slot = None,
233247
}
234248
}
235249
}
236250

237-
let merged: Vec<ColumnBounds> = acc.into_iter().flatten().collect();
238-
if merged.is_empty() {
239-
None
240-
} else {
241-
Some(PartitionBounds::new(merged))
251+
if acc.iter().all(Option::is_none) {
252+
return None;
242253
}
254+
Some(PartitionBounds { column_bounds: acc })
243255
}
244256

245257
/// Concatenate per-partition InList arrays into a single array. Returns `None`
@@ -985,10 +997,57 @@ mod tests {
985997
assert_eq!(cb.max, ScalarValue::Int32(Some(20)));
986998
}
987999

1000+
/// Regression test for the column-index misalignment bug.
1001+
///
1002+
/// When a *middle* join-key column is incomparable across partitions, it
1003+
/// must be dropped *in place* (its index left as a gap) so that every
1004+
/// later column keeps its join-key index. `create_bounds_predicate` pairs
1005+
/// each surviving column with `on_right` positionally; a column that
1006+
/// shifted index would be filtered against the wrong probe expression — a
1007+
/// silent wrong-results bug.
1008+
#[test]
1009+
fn compute_global_bounds_keeps_index_alignment_when_middle_column_dropped() {
1010+
let p1 = PartitionBounds::new(vec![
1011+
ColumnBounds::new(ScalarValue::Int32(Some(5)), ScalarValue::Int32(Some(10))),
1012+
ColumnBounds::new(ScalarValue::Int32(Some(1)), ScalarValue::Int32(Some(2))),
1013+
ColumnBounds::new(
1014+
ScalarValue::Int32(Some(100)),
1015+
ScalarValue::Int32(Some(200)),
1016+
),
1017+
]);
1018+
let p2 = PartitionBounds::new(vec![
1019+
ColumnBounds::new(ScalarValue::Int32(Some(3)), ScalarValue::Int32(Some(20))),
1020+
// Column 1 is incomparable with p1's column 1 (different
1021+
// ScalarValue variant -> partial_cmp returns None).
1022+
ColumnBounds::new(
1023+
ScalarValue::Utf8(Some("a".to_string())),
1024+
ScalarValue::Utf8(Some("z".to_string())),
1025+
),
1026+
ColumnBounds::new(
1027+
ScalarValue::Int32(Some(50)),
1028+
ScalarValue::Int32(Some(300)),
1029+
),
1030+
]);
1031+
let global = compute_global_bounds(&[&p1, &p2]).unwrap();
1032+
1033+
// Column 0 merges normally and stays at index 0.
1034+
let c0 = global.get_column_bounds(0).unwrap();
1035+
assert_eq!(c0.min, ScalarValue::Int32(Some(3)));
1036+
assert_eq!(c0.max, ScalarValue::Int32(Some(20)));
1037+
// Column 1 is incomparable and dropped.
1038+
assert!(global.get_column_bounds(1).is_none());
1039+
// Column 2 still merges, and critically stays at index 2 — the buggy
1040+
// `flatten()` shifted its bounds into index 1.
1041+
let c2 = global.get_column_bounds(2).unwrap();
1042+
assert_eq!(c2.min, ScalarValue::Int32(Some(50)));
1043+
assert_eq!(c2.max, ScalarValue::Int32(Some(300)));
1044+
}
1045+
9881046
/// When one partition's bounds for a column cannot be ordered against
9891047
/// another's (different ScalarValue variants → `partial_cmp` returns
990-
/// `None`), that column is dropped from the global envelope while
991-
/// orderable columns continue to widen normally.
1048+
/// `None`), that column is dropped from the global envelope (its index
1049+
/// becomes a `None` slot) while orderable columns continue to widen
1050+
/// normally.
9921051
#[test]
9931052
fn compute_global_bounds_drops_incomparable_column() {
9941053
let p1 = PartitionBounds::new(vec![
@@ -1006,8 +1065,9 @@ mod tests {
10061065
),
10071066
]);
10081067
let global = compute_global_bounds(&[&p1, &p2]).unwrap();
1009-
// Only the orderable column survives.
1010-
assert_eq!(global.column_bounds.len(), 1);
1068+
// The orderable column survives at its index; the incomparable one
1069+
// is a `None` slot.
1070+
assert!(global.get_column_bounds(1).is_none());
10111071
let cb = global.get_column_bounds(0).unwrap();
10121072
assert_eq!(cb.min, ScalarValue::Int32(Some(3)));
10131073
assert_eq!(cb.max, ScalarValue::Int32(Some(20)));

0 commit comments

Comments
 (0)