Skip to content

Commit bd308d7

Browse files
committed
add union type
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 3e93048 commit bd308d7

50 files changed

Lines changed: 2798 additions & 1175 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

encodings/sparse/src/canonical.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort
122122
len,
123123
ctx,
124124
)?,
125+
DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
125126
DType::Decimal(decimal_dtype, nullability) => {
126127
let canonical_decimal_value_type =
127128
DecimalType::smallest_decimal_value_type(decimal_dtype);

fuzz/src/array/compare.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ pub fn compare_canonical_array(
173173
}))
174174
.into_array()
175175
}
176-
d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
176+
d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
177177
unreachable!("DType {d} not supported for fuzzing")
178178
}
179179
}

fuzz/src/array/filter.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ pub fn filter_canonical_array(
122122
}
123123
take_canonical_array_non_nullable_indices(array, indices.as_slice(), ctx)
124124
}
125-
d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
125+
d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
126126
unreachable!("DType {d} not supported for fuzzing")
127127
}
128128
}

fuzz/src/array/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -474,6 +474,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet<ActionType> {
474474
acc.intersection(&actions).copied().collect()
475475
})
476476
}
477+
DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
477478
DType::List(..) | DType::FixedSizeList(..) => {
478479
// List supports: Compress, Slice, Take, Filter, MinMax, Mask, ScalarAt
479480
// Does NOT support: SearchSorted, Compare, Cast, Sum, FillNull

fuzz/src/array/search_sorted.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ pub fn search_sorted_canonical_array(
148148
.collect::<VortexResult<Vec<_>>>()?;
149149
scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side)
150150
}
151-
d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
151+
d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
152152
unreachable!("DType {d} not supported for fuzzing")
153153
}
154154
}

fuzz/src/array/slice.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ pub fn slice_canonical_array(
123123
.into_array(),
124124
)
125125
}
126-
d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
126+
d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
127127
unreachable!("DType {d} not supported for fuzzing")
128128
}
129129
}

fuzz/src/array/sort.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ pub fn sort_canonical_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexR
101101
});
102102
take_canonical_array_non_nullable_indices(array, &sort_indices, ctx)
103103
}
104-
d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
104+
d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
105105
unreachable!("DType {d} not supported for fuzzing")
106106
}
107107
}

fuzz/src/array/take.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ pub fn take_canonical_array(
147147
}
148148
Ok(builder.finish())
149149
}
150-
d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
150+
d @ (DType::Null | DType::Union(..) | DType::Extension(_) | DType::Variant(_)) => {
151151
unreachable!("DType {d} not supported for fuzzing")
152152
}
153153
}

vortex-array/public-api.lock

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9346,6 +9346,8 @@ pub vortex_array::dtype::DType::Primitive(vortex_array::dtype::PType, vortex_arr
93469346

93479347
pub vortex_array::dtype::DType::Struct(vortex_array::dtype::StructFields, vortex_array::dtype::Nullability)
93489348

9349+
pub vortex_array::dtype::DType::Union(vortex_array::dtype::UnionVariants, vortex_array::dtype::Nullability)
9350+
93499351
pub vortex_array::dtype::DType::Utf8(vortex_array::dtype::Nullability)
93509352

93519353
pub vortex_array::dtype::DType::Variant(vortex_array::dtype::Nullability)
@@ -9376,6 +9378,8 @@ pub fn vortex_array::dtype::DType::as_struct_fields(&self) -> &vortex_array::dty
93769378

93779379
pub fn vortex_array::dtype::DType::as_struct_fields_opt(&self) -> core::option::Option<&vortex_array::dtype::StructFields>
93789380

9381+
pub fn vortex_array::dtype::DType::as_union_variants_opt(&self) -> core::option::Option<&vortex_array::dtype::UnionVariants>
9382+
93799383
pub fn vortex_array::dtype::DType::element_size(&self) -> core::option::Option<usize>
93809384

93819385
pub fn vortex_array::dtype::DType::eq_ignore_nullability(&self, &Self) -> bool
@@ -9396,6 +9400,8 @@ pub fn vortex_array::dtype::DType::into_struct_fields(self) -> vortex_array::dty
93969400

93979401
pub fn vortex_array::dtype::DType::into_struct_fields_opt(self) -> core::option::Option<vortex_array::dtype::StructFields>
93989402

9403+
pub fn vortex_array::dtype::DType::into_union_variants_opt(self) -> core::option::Option<vortex_array::dtype::UnionVariants>
9404+
93999405
pub fn vortex_array::dtype::DType::is_binary(&self) -> bool
94009406

94019407
pub fn vortex_array::dtype::DType::is_boolean(&self) -> bool
@@ -9422,6 +9428,8 @@ pub fn vortex_array::dtype::DType::is_signed_int(&self) -> bool
94229428

94239429
pub fn vortex_array::dtype::DType::is_struct(&self) -> bool
94249430

9431+
pub fn vortex_array::dtype::DType::is_union(&self) -> bool
9432+
94259433
pub fn vortex_array::dtype::DType::is_unsigned_int(&self) -> bool
94269434

94279435
pub fn vortex_array::dtype::DType::is_utf8(&self) -> bool
@@ -10506,6 +10514,66 @@ impl<T, V> core::iter::traits::collect::FromIterator<(T, V)> for vortex_array::d
1050610514

1050710515
pub fn vortex_array::dtype::StructFields::from_iter<I: core::iter::traits::collect::IntoIterator<Item = (T, V)>>(I) -> Self
1050810516

10517+
pub struct vortex_array::dtype::UnionVariants(_)
10518+
10519+
impl vortex_array::dtype::UnionVariants
10520+
10521+
pub fn vortex_array::dtype::UnionVariants::child_index_to_tag(&self, usize) -> i8
10522+
10523+
pub fn vortex_array::dtype::UnionVariants::empty() -> Self
10524+
10525+
pub fn vortex_array::dtype::UnionVariants::find(&self, impl core::convert::AsRef<str>) -> core::option::Option<usize>
10526+
10527+
pub fn vortex_array::dtype::UnionVariants::is_consecutive(&self) -> bool
10528+
10529+
pub fn vortex_array::dtype::UnionVariants::is_empty(&self) -> bool
10530+
10531+
pub fn vortex_array::dtype::UnionVariants::len(&self) -> usize
10532+
10533+
pub fn vortex_array::dtype::UnionVariants::names(&self) -> &vortex_array::dtype::FieldNames
10534+
10535+
pub fn vortex_array::dtype::UnionVariants::new_consecutive(vortex_array::dtype::FieldNames, alloc::vec::Vec<vortex_array::dtype::DType>) -> vortex_error::VortexResult<Self>
10536+
10537+
pub fn vortex_array::dtype::UnionVariants::nullability_constraints_satisfied(&self, vortex_array::dtype::Nullability) -> bool
10538+
10539+
pub fn vortex_array::dtype::UnionVariants::tag_to_child_index(&self, i8) -> core::option::Option<usize>
10540+
10541+
pub fn vortex_array::dtype::UnionVariants::try_new(vortex_array::dtype::FieldNames, alloc::vec::Vec<vortex_array::dtype::DType>, alloc::vec::Vec<i8>) -> vortex_error::VortexResult<Self>
10542+
10543+
pub fn vortex_array::dtype::UnionVariants::type_ids(&self) -> &[i8]
10544+
10545+
pub fn vortex_array::dtype::UnionVariants::variant(&self, impl core::convert::AsRef<str>) -> core::option::Option<vortex_array::dtype::DType>
10546+
10547+
pub fn vortex_array::dtype::UnionVariants::variant_by_index(&self, usize) -> core::option::Option<vortex_array::dtype::DType>
10548+
10549+
pub fn vortex_array::dtype::UnionVariants::variants(&self) -> impl core::iter::traits::exact_size::ExactSizeIterator<Item = vortex_array::dtype::DType> + '_
10550+
10551+
impl core::clone::Clone for vortex_array::dtype::UnionVariants
10552+
10553+
pub fn vortex_array::dtype::UnionVariants::clone(&self) -> vortex_array::dtype::UnionVariants
10554+
10555+
impl core::cmp::Eq for vortex_array::dtype::UnionVariants
10556+
10557+
impl core::cmp::PartialEq for vortex_array::dtype::UnionVariants
10558+
10559+
pub fn vortex_array::dtype::UnionVariants::eq(&self, &Self) -> bool
10560+
10561+
impl core::default::Default for vortex_array::dtype::UnionVariants
10562+
10563+
pub fn vortex_array::dtype::UnionVariants::default() -> Self
10564+
10565+
impl core::fmt::Debug for vortex_array::dtype::UnionVariants
10566+
10567+
pub fn vortex_array::dtype::UnionVariants::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
10568+
10569+
impl core::fmt::Display for vortex_array::dtype::UnionVariants
10570+
10571+
pub fn vortex_array::dtype::UnionVariants::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
10572+
10573+
impl core::hash::Hash for vortex_array::dtype::UnionVariants
10574+
10575+
pub fn vortex_array::dtype::UnionVariants::hash<__H: core::hash::Hasher>(&self, &mut __H)
10576+
1050910577
#[repr(transparent)] pub struct vortex_array::dtype::i256(_)
1051010578

1051110579
impl vortex_array::dtype::i256

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,7 @@ pub(crate) fn constant_uncompressed_size_in_bytes(
247247
let canonical = array.array().clone().execute::<Canonical>(ctx)?;
248248
return canonical_uncompressed_size_in_bytes(&canonical, ctx);
249249
}
250+
DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
250251
};
251252

252253
value_size
@@ -293,6 +294,9 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool {
293294
DType::Struct(fields, _) => fields
294295
.fields()
295296
.all(|field| supports_uncompressed_size_in_bytes(&field)),
297+
DType::Union(variants, _) => variants
298+
.variants()
299+
.all(|variant| supports_uncompressed_size_in_bytes(&variant)),
296300
DType::Extension(ext_dtype) => {
297301
supports_uncompressed_size_in_bytes(ext_dtype.storage_dtype())
298302
}

0 commit comments

Comments
 (0)