Skip to content

Commit 66e354f

Browse files
connortsui20connortsui20
committed
Fix UnionArray canonicalization edge cases
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 092651f commit 66e354f

3 files changed

Lines changed: 226 additions & 21 deletions

File tree

‎vortex-array/src/arrays/constant/vtable/canonical.rs‎

Lines changed: 74 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use vortex_buffer::Buffer;
88
use vortex_buffer::buffer;
99
use vortex_error::VortexExpect;
1010
use vortex_error::VortexResult;
11+
use vortex_error::vortex_err;
1112

1213
use crate::Canonical;
1314
use crate::ExecutionCtx;
@@ -149,16 +150,22 @@ pub(crate) fn constant_canonicalize(
149150
.collect(),
150151
None => {
151152
assert!(matches!(validity, Validity::AllInvalid));
152-
// The struct is entirely null, so fields just need placeholder values with the
153-
// correct dtype. We use `default_value` which returns a zero for non-nullable
154-
// dtypes and null for nullable dtypes, preserving each field's nullability.
155153
struct_dtype
156154
.fields()
157155
.map(|dt| {
158-
let scalar = Scalar::default_value(&dt);
159-
ConstantArray::new(scalar, array.len()).into_array()
156+
if array.is_empty() {
157+
return Ok(Canonical::empty(&dt).into_array());
158+
}
159+
160+
let scalar = try_placeholder_scalar(&dt).ok_or_else(|| {
161+
vortex_err!(
162+
"cannot canonicalize null constant struct: field dtype {dt} \
163+
has no placeholder value"
164+
)
165+
})?;
166+
Ok(ConstantArray::new(scalar, array.len()).into_array())
160167
})
161-
.collect()
168+
.collect::<VortexResult<Vec<_>>>()?
162169
}
163170
};
164171
// SAFETY: Fields are constructed from the same struct scalar, all have same
@@ -202,25 +209,32 @@ fn constant_canonical_union(
202209
nullability: Nullability,
203210
len: usize,
204211
) -> VortexResult<UnionArray> {
212+
if len == 0 {
213+
return Ok(UnionArray::empty(variants.clone(), nullability));
214+
}
215+
205216
let union = scalar.as_union();
206217
let selected = union.child_index().zip(union.value());
207218

208219
let children = variants
209220
.variants()
210221
.enumerate()
211222
.map(|(index, dtype)| {
212-
let value = match selected {
213-
Some((selected_child, selected_value)) if index == selected_child => {
223+
let value = match &selected {
224+
Some((selected_child, selected_value)) if index == *selected_child => {
214225
selected_value.clone()
215226
}
216-
_ if matches!(&dtype, DType::Union(_, Nullability::NonNullable)) => {
217-
Scalar::zero_value(&dtype)
218-
}
219-
_ => Scalar::default_value(&dtype),
227+
_ => try_placeholder_scalar(&dtype).ok_or_else(|| {
228+
vortex_err!(
229+
"cannot canonicalize constant union: unselected variant {} ({dtype}) has \
230+
no placeholder value",
231+
variants.names()[index]
232+
)
233+
})?,
220234
};
221-
ConstantArray::new(value, len).into_array()
235+
Ok(ConstantArray::new(value, len).into_array())
222236
})
223-
.collect::<Vec<_>>();
237+
.collect::<VortexResult<Vec<_>>>()?;
224238

225239
let type_id = match union.type_id() {
226240
Some(type_id) => Scalar::primitive(type_id, nullability),
@@ -234,6 +248,52 @@ fn constant_canonical_union(
234248
)
235249
}
236250

251+
/// Builds an arbitrary valid scalar for array positions ignored by a parent validity or type ID.
252+
/// This does not give the dtype a semantic default value.
253+
fn try_placeholder_scalar(dtype: &DType) -> Option<Scalar> {
254+
if let Some(default) = Scalar::try_default_value(dtype) {
255+
return Some(default);
256+
}
257+
258+
match dtype {
259+
DType::FixedSizeList(element_dtype, size, nullability) => {
260+
let element = try_placeholder_scalar(element_dtype)?;
261+
Some(Scalar::fixed_size_list(
262+
Arc::clone(element_dtype),
263+
vec![element; *size as usize],
264+
*nullability,
265+
))
266+
}
267+
DType::Struct(fields, _) => {
268+
let children = fields
269+
.fields()
270+
.map(|field| try_placeholder_scalar(&field))
271+
.collect::<Option<Vec<_>>>()?;
272+
Some(Scalar::struct_(dtype.clone(), children))
273+
}
274+
DType::Union(variants, nullability) => {
275+
variants
276+
.variants()
277+
.enumerate()
278+
.find_map(|(index, child_dtype)| {
279+
let child = try_placeholder_scalar(&child_dtype)?;
280+
Scalar::union(
281+
variants.clone(),
282+
variants.child_index_to_tag(index),
283+
child,
284+
*nullability,
285+
)
286+
.ok()
287+
})
288+
}
289+
DType::Extension(ext_dtype) => {
290+
let storage = try_placeholder_scalar(ext_dtype.storage_dtype())?;
291+
Some(Scalar::extension_ref(ext_dtype.clone(), storage))
292+
}
293+
_ => None,
294+
}
295+
}
296+
237297
fn constant_canonical_byte_view(
238298
scalar_bytes: Option<&[u8]>,
239299
dtype: &DType,

‎vortex-array/src/arrays/union/tests.rs‎

Lines changed: 130 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,13 @@
44
use vortex_buffer::ByteBufferMut;
55
use vortex_buffer::buffer;
66
use vortex_error::VortexResult;
7+
use vortex_error::vortex_err;
78
use vortex_mask::Mask;
89
use vortex_session::registry::ReadContext;
910

1011
use crate::ArrayContext;
12+
use crate::Canonical;
13+
use crate::CanonicalValidity;
1114
use crate::IntoArray;
1215
use crate::VortexSessionExecute;
1316
use crate::array_session;
@@ -302,12 +305,18 @@ fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> {
302305
vec![DType::Primitive(PType::I64, Nullability::NonNullable)],
303306
vec![3],
304307
)?;
308+
let nested_placeholder = Scalar::union(
309+
nested_variants.clone(),
310+
3,
311+
0_i64.into(),
312+
Nullability::NonNullable,
313+
)?;
305314
let nested_dtype = DType::Union(nested_variants, Nullability::NonNullable);
306315
let outer_variants = UnionVariants::try_new(
307316
["number", "nested"].into(),
308317
vec![
309318
DType::Primitive(PType::I32, Nullability::NonNullable),
310-
nested_dtype.clone(),
319+
nested_dtype,
311320
],
312321
vec![5, 9],
313322
)?;
@@ -324,7 +333,7 @@ fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> {
324333
.execute::<UnionArray>(&mut ctx)?;
325334
assert_eq!(
326335
array.child_by_name("nested")?.execute_scalar(0, &mut ctx)?,
327-
Scalar::zero_value(&nested_dtype)
336+
nested_placeholder
328337
);
329338

330339
let outer_null = Scalar::null(DType::Union(outer_variants, Nullability::Nullable));
@@ -333,12 +342,130 @@ fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> {
333342
.execute::<UnionArray>(&mut ctx)?;
334343
assert_eq!(
335344
array.child_by_name("nested")?.execute_scalar(0, &mut ctx)?,
336-
Scalar::zero_value(&nested_dtype)
345+
nested_placeholder
337346
);
338347

339348
Ok(())
340349
}
341350

351+
#[test]
352+
fn constant_union_builds_deeply_nested_union_placeholders() -> VortexResult<()> {
353+
let nested_variants = UnionVariants::try_new(
354+
["value"].into(),
355+
vec![DType::Primitive(PType::I64, Nullability::NonNullable)],
356+
vec![3],
357+
)?;
358+
let nested_placeholder = Scalar::union(
359+
nested_variants.clone(),
360+
3,
361+
0_i64.into(),
362+
Nullability::NonNullable,
363+
)?;
364+
let nested_dtype = DType::Union(nested_variants, Nullability::NonNullable);
365+
let wrapper_dtype = DType::struct_([("nested", nested_dtype)], Nullability::NonNullable);
366+
let outer_variants = UnionVariants::try_new(
367+
["number", "wrapper"].into(),
368+
vec![
369+
DType::Primitive(PType::I32, Nullability::NonNullable),
370+
wrapper_dtype,
371+
],
372+
vec![5, 9],
373+
)?;
374+
let selected_number =
375+
Scalar::union(outer_variants, 5, 42_i32.into(), Nullability::NonNullable)?;
376+
let mut ctx = array_session().create_execution_ctx();
377+
let array = ConstantArray::new(selected_number, 2)
378+
.into_array()
379+
.execute::<UnionArray>(&mut ctx)?;
380+
let wrapper = array
381+
.child_by_name("wrapper")?
382+
.execute_scalar(0, &mut ctx)?;
383+
384+
assert_eq!(
385+
wrapper.as_struct().field("nested"),
386+
Some(nested_placeholder)
387+
);
388+
389+
Ok(())
390+
}
391+
392+
#[test]
393+
fn constant_union_rejects_uninhabited_placeholders_without_panicking() -> VortexResult<()> {
394+
let uninhabited = DType::Union(
395+
UnionVariants::new(Default::default(), vec![])?,
396+
Nullability::NonNullable,
397+
);
398+
let outer_variants = UnionVariants::try_new(
399+
["number", "uninhabited"].into(),
400+
vec![
401+
DType::Primitive(PType::I32, Nullability::NonNullable),
402+
uninhabited,
403+
],
404+
vec![5, 9],
405+
)?;
406+
let selected_number =
407+
Scalar::union(outer_variants, 5, 42_i32.into(), Nullability::NonNullable)?;
408+
let mut ctx = array_session().create_execution_ctx();
409+
410+
assert!(
411+
ConstantArray::new(selected_number.clone(), 1)
412+
.into_array()
413+
.execute::<UnionArray>(&mut ctx)
414+
.is_err()
415+
);
416+
assert_eq!(
417+
ConstantArray::new(selected_number, 0)
418+
.into_array()
419+
.execute::<UnionArray>(&mut ctx)?
420+
.len(),
421+
0
422+
);
423+
424+
Ok(())
425+
}
426+
427+
#[test]
428+
fn empty_union_supports_variant_children() -> VortexResult<()> {
429+
let variants = UnionVariants::try_new(
430+
["dynamic"].into(),
431+
vec![DType::Variant(Nullability::NonNullable)],
432+
vec![5],
433+
)?;
434+
let array = Canonical::empty(&DType::Union(variants, Nullability::NonNullable)).into_union();
435+
436+
assert_eq!(array.len(), 0);
437+
assert_eq!(
438+
array.child(0).map(|child| child.dtype()),
439+
Some(&DType::Variant(Nullability::NonNullable))
440+
);
441+
442+
Ok(())
443+
}
444+
445+
#[test]
446+
fn canonical_validity_canonicalizes_union_type_ids() -> VortexResult<()> {
447+
let array = UnionArray::try_new(
448+
ConstantArray::new(Scalar::primitive(5_u8, Nullability::Nullable), 2).into_array(),
449+
variants()?,
450+
vec![
451+
PrimitiveArray::from_iter([10_i32, 20]).into_array(),
452+
BoolArray::from_iter([false, true]).into_array(),
453+
],
454+
)?;
455+
let mut ctx = array_session().create_execution_ctx();
456+
let Canonical::Union(array) = array.into_array().execute::<CanonicalValidity>(&mut ctx)?.0
457+
else {
458+
return Err(vortex_err!(
459+
"UnionArray must remain canonical Union storage"
460+
));
461+
};
462+
463+
assert!(array.type_ids().is::<crate::arrays::Primitive>());
464+
assert_eq!(array.dtype().nullability(), Nullability::Nullable);
465+
466+
Ok(())
467+
}
468+
342469
#[test]
343470
fn chunked_union_packs_components() -> VortexResult<()> {
344471
let first = union_array()?.into_array().slice(0..1)?;

‎vortex-array/src/canonical.rs‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use crate::array::ArrayView;
2121
use crate::array::child_to_validity;
2222
use crate::arrays::Bool;
2323
use crate::arrays::BoolArray;
24+
use crate::arrays::ChunkedArray;
2425
use crate::arrays::Decimal;
2526
use crate::arrays::DecimalArray;
2627
use crate::arrays::Extension;
@@ -237,9 +238,15 @@ impl Canonical {
237238
DType::Union(variants, nullability) => {
238239
Canonical::Union(UnionArray::empty(variants.clone(), *nullability))
239240
}
240-
DType::Variant(_) => {
241-
vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant")
242-
}
241+
DType::Variant(_) => Canonical::Variant(
242+
VariantArray::try_new(
243+
ChunkedArray::try_new(vec![], dtype.clone())
244+
.vortex_expect("empty Variant core storage must be valid")
245+
.into_array(),
246+
None,
247+
)
248+
.vortex_expect("empty VariantArray must be valid"),
249+
),
243250
DType::Extension(ext_dtype) => Canonical::Extension(ExtensionArray::new(
244251
ext_dtype.clone(),
245252
Canonical::empty(ext_dtype.storage_dtype()).into_array(),
@@ -678,7 +685,18 @@ impl Executable for CanonicalValidity {
678685
StructArray::new_unchecked(fields, struct_fields, len, validity.execute(ctx)?)
679686
})))
680687
}
681-
union @ Canonical::Union(_) => Ok(CanonicalValidity(union)),
688+
Canonical::Union(union) => {
689+
let UnionDataParts {
690+
variants,
691+
type_ids,
692+
children,
693+
} = union.into_data_parts();
694+
let type_ids = type_ids.execute::<CanonicalValidity>(ctx)?.0.into_array();
695+
696+
Ok(CanonicalValidity(Canonical::Union(unsafe {
697+
UnionArray::new_unchecked(type_ids, variants, children)
698+
})))
699+
}
682700
Canonical::Extension(ext) => Ok(CanonicalValidity(Canonical::Extension(
683701
ExtensionArray::new(
684702
ext.ext_dtype().clone(),

0 commit comments

Comments
 (0)