Skip to content

Commit fc4d111

Browse files
authored
Fix struct cast panic in rule and add tests (#6962)
## Summary Closes: #6933 ## Testing Adds a regression test as well as a few other tests to confirm this is correct. Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 8ac7234 commit fc4d111

1 file changed

Lines changed: 128 additions & 2 deletions

File tree

  • vortex-array/src/arrays/struct_/compute

vortex-array/src/arrays/struct_/compute/rules.rs

Lines changed: 128 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,9 @@ impl ArrayParentReduceRule<Struct> for StructCastPushDownRule {
5151
parent: ScalarFnArrayView<Cast>,
5252
_child_idx: usize,
5353
) -> VortexResult<Option<ArrayRef>> {
54-
let target_fields = parent.options.as_struct_fields();
54+
let Some(target_fields) = parent.options.as_struct_fields_opt() else {
55+
return Ok(None);
56+
};
5557
let mut new_fields = Vec::with_capacity(target_fields.nfields());
5658

5759
for (target_name, target_dtype) in target_fields.names().iter().zip(target_fields.fields())
@@ -136,6 +138,8 @@ impl ArrayParentReduceRule<Struct> for StructGetItemRule {
136138

137139
#[cfg(test)]
138140
mod tests {
141+
use vortex_buffer::buffer;
142+
139143
use crate::IntoArray;
140144
use crate::arrays::StructArray;
141145
use crate::arrays::VarBinViewArray;
@@ -146,6 +150,7 @@ mod tests {
146150
use crate::dtype::DType;
147151
use crate::dtype::FieldNames;
148152
use crate::dtype::Nullability;
153+
use crate::dtype::PType;
149154
use crate::dtype::StructFields;
150155
use crate::scalar::Scalar;
151156
use crate::validity::Validity;
@@ -173,7 +178,8 @@ mod tests {
173178
Nullability::NonNullable,
174179
);
175180

176-
// Use ArrayBuiltins::cast which goes through the optimizer and applies StructCastPushDownRule
181+
// Use `ArrayBuiltins::cast` which goes through the optimizer and applies
182+
// `StructCastPushDownRule`.
177183
let result = source.into_array().cast(target).unwrap().to_struct();
178184
assert_arrays_eq!(
179185
result.unmasked_field_by_name("a").unwrap(),
@@ -188,4 +194,124 @@ mod tests {
188194
ConstantArray::new(Scalar::null(utf8_null), 1)
189195
);
190196
}
197+
198+
/// Regression test: casting a struct to a non-struct DType must not panic. Previously,
199+
/// `StructCastPushDownRule` called `as_struct_fields()` which panics on non-struct types.
200+
#[test]
201+
fn cast_struct_to_non_struct_does_not_panic() {
202+
let source = StructArray::try_new(
203+
FieldNames::from(["x"]),
204+
vec![buffer![1i32, 2, 3].into_array()],
205+
3,
206+
Validity::NonNullable,
207+
)
208+
.unwrap();
209+
210+
// Casting a struct to a primitive type should not panic. Before the fix,
211+
// `StructCastPushDownRule` would panic via `as_struct_fields()` on the non-struct target.
212+
let result = source
213+
.into_array()
214+
.cast(DType::Primitive(PType::I32, Nullability::NonNullable));
215+
// Whether this errors or succeeds depends on execution, but the key invariant is that the
216+
// optimizer rule does not panic.
217+
if let Ok(arr) = &result {
218+
assert_eq!(
219+
arr.dtype(),
220+
&DType::Primitive(PType::I32, Nullability::NonNullable)
221+
);
222+
}
223+
}
224+
225+
#[test]
226+
fn cast_struct_drop_field() {
227+
// Casting to a struct with a subset of fields should succeed.
228+
let source = StructArray::try_new(
229+
FieldNames::from(["a", "b", "c"]),
230+
vec![
231+
buffer![1i32, 2, 3].into_array(),
232+
buffer![10i64, 20, 30].into_array(),
233+
buffer![100u8, 200, 255].into_array(),
234+
],
235+
3,
236+
Validity::NonNullable,
237+
)
238+
.unwrap();
239+
240+
let target = DType::Struct(
241+
StructFields::new(
242+
FieldNames::from(["a", "c"]),
243+
vec![
244+
DType::Primitive(PType::I32, Nullability::NonNullable),
245+
DType::Primitive(PType::U8, Nullability::NonNullable),
246+
],
247+
),
248+
Nullability::NonNullable,
249+
);
250+
251+
let result = source.into_array().cast(target).unwrap().to_struct();
252+
assert_eq!(result.unmasked_fields().len(), 2);
253+
assert_arrays_eq!(
254+
result.unmasked_field_by_name("a").unwrap(),
255+
buffer![1i32, 2, 3].into_array()
256+
);
257+
assert_arrays_eq!(
258+
result.unmasked_field_by_name("c").unwrap(),
259+
buffer![100u8, 200, 255].into_array()
260+
);
261+
}
262+
263+
#[test]
264+
fn cast_struct_field_type_widening() {
265+
// Casting struct fields to wider types (i32 -> i64).
266+
let source = StructArray::try_new(
267+
FieldNames::from(["val"]),
268+
vec![buffer![1i32, 2, 3].into_array()],
269+
3,
270+
Validity::NonNullable,
271+
)
272+
.unwrap();
273+
274+
let target = DType::Struct(
275+
StructFields::new(
276+
FieldNames::from(["val"]),
277+
vec![DType::Primitive(PType::I64, Nullability::NonNullable)],
278+
),
279+
Nullability::NonNullable,
280+
);
281+
282+
let result = source.into_array().cast(target).unwrap().to_struct();
283+
assert_eq!(
284+
result.unmasked_field_by_name("val").unwrap().dtype(),
285+
&DType::Primitive(PType::I64, Nullability::NonNullable)
286+
);
287+
assert_arrays_eq!(
288+
result.unmasked_field_by_name("val").unwrap(),
289+
buffer![1i64, 2, 3].into_array()
290+
);
291+
}
292+
293+
#[test]
294+
fn cast_struct_add_non_nullable_field_fails() {
295+
// Adding a non-nullable field via cast should fail.
296+
let source = StructArray::try_new(
297+
FieldNames::from(["a"]),
298+
vec![buffer![1i32].into_array()],
299+
1,
300+
Validity::NonNullable,
301+
)
302+
.unwrap();
303+
304+
let target = DType::Struct(
305+
StructFields::new(
306+
FieldNames::from(["a", "b"]),
307+
vec![
308+
DType::Primitive(PType::I32, Nullability::NonNullable),
309+
DType::Primitive(PType::I32, Nullability::NonNullable),
310+
],
311+
),
312+
Nullability::NonNullable,
313+
);
314+
315+
assert!(source.into_array().cast(target).is_err());
316+
}
191317
}

0 commit comments

Comments
 (0)