Skip to content

Commit 4e6c30d

Browse files
AdamGSAdamGS
committed
Add Map compute and compression support
Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent 364b5d2 commit 4e6c30d

34 files changed

Lines changed: 1456 additions & 73 deletions

File tree

‎fuzz/src/array/filter.rs‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use vortex_array::arrays::StructArray;
1111
use vortex_array::arrays::VarBinViewArray;
1212
use vortex_array::arrays::bool::BoolArrayExt;
1313
use vortex_array::arrays::struct_::StructArrayExt;
14+
use vortex_array::builders::builder_with_capacity;
1415
use vortex_array::dtype::DType;
1516
use vortex_array::match_each_decimal_value_type;
1617
use vortex_array::match_each_native_ptype;
@@ -121,11 +122,17 @@ pub fn filter_canonical_array(
121122
)
122123
.map(|a| a.into_array())
123124
}
124-
d @ (DType::Null
125-
| DType::Map(..)
126-
| DType::Union(..)
127-
| DType::Variant(_)
128-
| DType::Extension(_)) => {
125+
DType::Map(..) => {
126+
let mut builder =
127+
builder_with_capacity(array.dtype(), filter.iter().filter(|b| **b).count());
128+
for (idx, keep) in filter.iter().enumerate() {
129+
if *keep {
130+
builder.append_scalar(&array.execute_scalar(idx, ctx)?)?;
131+
}
132+
}
133+
Ok(builder.finish())
134+
}
135+
d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
129136
unreachable!("DType {d} not supported for fuzzing")
130137
}
131138
}

‎fuzz/src/array/mask.rs‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
2121
use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
2222
use vortex_array::arrays::listview::ListViewArraySlotsExt;
2323
use vortex_array::arrays::struct_::StructArrayExt;
24+
use vortex_array::builders::builder_with_capacity;
2425
use vortex_array::dtype::Nullability;
2526
use vortex_array::match_each_decimal_value_type;
2627
use vortex_array::validity::Validity;
@@ -139,6 +140,18 @@ pub fn mask_canonical_array(
139140
.vortex_expect("StructArray creation should succeed in fuzz test")
140141
.into_array()
141142
}
143+
Canonical::Map(array) => {
144+
let result_dtype = array.dtype().as_nullable();
145+
let mut builder = builder_with_capacity(&result_dtype, array.len());
146+
for idx in 0..array.len() {
147+
if mask.value(idx) {
148+
builder.append_scalar(&array.execute_scalar(idx, ctx)?.cast(&result_dtype)?)?;
149+
} else {
150+
builder.append_null();
151+
}
152+
}
153+
builder.finish()
154+
}
142155
Canonical::Extension(array) => {
143156
// Recursively mask the storage array
144157
let storage_canonical = array.storage_array().clone().execute::<Canonical>(ctx)?;
@@ -153,7 +166,6 @@ pub fn mask_canonical_array(
153166
Canonical::Union(_) => {
154167
todo!("TODO(connor)[Union]: support Union arrays in the mask fuzzer")
155168
}
156-
Canonical::Map(_) => unreachable!("Map arrays are not fuzzed"),
157169
Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
158170
})
159171
}

‎fuzz/src/array/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -518,7 +518,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet<ActionType> {
518518
acc.intersection(&actions).copied().collect()
519519
})
520520
}
521-
DType::Map(..) => HashSet::new(),
521+
DType::Map(..) => [Compress, Slice, Take, Filter, Mask, ScalarAt].into(),
522522
DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
523523
// Currently, no support at all
524524
DType::Variant(_) => unreachable!("Variant dtype shouldn't be fuzzed"),

‎fuzz/src/array/slice.rs‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
1616
use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
1717
use vortex_array::arrays::listview::ListViewArraySlotsExt;
1818
use vortex_array::arrays::struct_::StructArrayExt;
19+
use vortex_array::builders::builder_with_capacity;
1920
use vortex_array::dtype::DType;
2021
use vortex_array::match_each_decimal_value_type;
2122
use vortex_array::match_each_native_ptype;
@@ -125,11 +126,14 @@ pub fn slice_canonical_array(
125126
)
126127
.map(|a| a.into_array())
127128
}
128-
d @ (DType::Null
129-
| DType::Map(..)
130-
| DType::Union(..)
131-
| DType::Variant(_)
132-
| DType::Extension(_)) => {
129+
DType::Map(..) => {
130+
let mut builder = builder_with_capacity(array.dtype(), stop - start);
131+
for idx in start..stop {
132+
builder.append_scalar(&array.execute_scalar(idx, ctx)?)?;
133+
}
134+
Ok(builder.finish())
135+
}
136+
d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
133137
unreachable!("DType {d} not supported for fuzzing")
134138
}
135139
}

‎fuzz/src/array/take.rs‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -148,11 +148,20 @@ pub fn take_canonical_array(
148148
)
149149
.map(|a| a.into_array())
150150
}
151-
d @ (DType::Null
152-
| DType::Map(..)
153-
| DType::Union(..)
154-
| DType::Variant(_)
155-
| DType::Extension(_)) => {
151+
DType::Map(..) => {
152+
let result_dtype = array.dtype().union_nullability(nullable);
153+
let mut builder = builder_with_capacity(&result_dtype, indices.len());
154+
for idx in indices {
155+
if let Some(idx) = idx {
156+
builder
157+
.append_scalar(&array.execute_scalar(*idx, ctx)?.cast(&result_dtype)?)?;
158+
} else {
159+
builder.append_null();
160+
}
161+
}
162+
Ok(builder.finish())
163+
}
164+
d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
156165
unreachable!("DType {d} not supported for fuzzing")
157166
}
158167
}

‎vortex-array/src/aggregate_fn/fns/all_non_distinct/map.rs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use vortex_error::VortexResult;
5+
6+
use super::list::check_list_identical;
7+
use crate::ExecutionCtx;
8+
use crate::arrays::MapArray;
9+
use crate::arrays::map::MapArrayExt;
10+
11+
pub(super) fn check_map_identical(
12+
lhs: &MapArray,
13+
rhs: &MapArray,
14+
ctx: &mut ExecutionCtx,
15+
) -> VortexResult<bool> {
16+
if lhs.map_dtype() != rhs.map_dtype() {
17+
return Ok(false);
18+
}
19+
20+
check_list_identical(
21+
&lhs.entries().into_owned(),
22+
&rhs.entries().into_owned(),
23+
ctx,
24+
)
25+
}

‎vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ mod extension;
77
mod filter;
88
mod fixed_size_list;
99
mod list;
10+
mod map;
1011
mod primitive;
1112
mod struct_;
1213
#[cfg(test)]
@@ -27,6 +28,7 @@ use self::extension::check_extension_identical;
2728
use self::filter::shared_validity_mask;
2829
use self::fixed_size_list::check_fixed_size_list_identical;
2930
use self::list::check_list_identical;
31+
use self::map::check_map_identical;
3032
use self::primitive::check_primitive_identical;
3133
use self::struct_::check_struct_identical;
3234
use self::varbin::check_varbinview_identical;
@@ -262,6 +264,7 @@ fn check_canonical_identical(
262264
}
263265
(Canonical::Struct(lhs), Canonical::Struct(rhs)) => check_struct_identical(lhs, rhs, ctx),
264266
(Canonical::List(lhs), Canonical::List(rhs)) => check_list_identical(lhs, rhs, ctx),
267+
(Canonical::Map(lhs), Canonical::Map(rhs)) => check_map_identical(lhs, rhs, ctx),
265268
(Canonical::FixedSizeList(lhs), Canonical::FixedSizeList(rhs)) => {
266269
check_fixed_size_list_identical(lhs, rhs, ctx)
267270
}

‎vortex-array/src/aggregate_fn/fns/is_constant/map.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use vortex_error::VortexResult;
5+
6+
use super::list::check_listview_constant;
7+
use crate::ExecutionCtx;
8+
use crate::arrays::MapArray;
9+
use crate::arrays::map::MapArrayExt;
10+
11+
pub(super) fn check_map_constant(map: &MapArray, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
12+
check_listview_constant(&map.entries().into_owned(), ctx)
13+
}

‎vortex-array/src/aggregate_fn/fns/is_constant/mod.rs‎

Lines changed: 65 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ mod decimal;
66
mod extension;
77
mod fixed_size_list;
88
mod list;
9+
mod map;
910
pub mod primitive;
1011
mod struct_;
1112
mod varbin;
@@ -20,6 +21,7 @@ use self::decimal::check_decimal_constant;
2021
use self::extension::check_extension_constant;
2122
use self::fixed_size_list::check_fixed_size_list_constant;
2223
use self::list::check_listview_constant;
24+
use self::map::check_map_constant;
2325
use self::primitive::check_primitive_constant;
2426
use self::struct_::check_struct_constant;
2527
use self::varbin::check_varbinview_constant;
@@ -402,9 +404,7 @@ impl AggregateFnVTable for IsConstant {
402404
Canonical::Struct(s) => check_struct_constant(s, ctx)?,
403405
Canonical::Extension(e) => check_extension_constant(e, ctx)?,
404406
Canonical::List(l) => check_listview_constant(l, ctx)?,
405-
Canonical::Map(_) => {
406-
vortex_bail!("Map arrays don't support IsConstant")
407-
}
407+
Canonical::Map(m) => check_map_constant(m, ctx)?,
408408
Canonical::FixedSizeList(f) => check_fixed_size_list_constant(f, ctx)?,
409409
Canonical::Null(_) => true,
410410
Canonical::Union(_) => {
@@ -456,14 +456,54 @@ mod tests {
456456
use crate::arrays::ListArray;
457457
use crate::arrays::PrimitiveArray;
458458
use crate::arrays::StructArray;
459+
use crate::builders::MapBuilder;
459460
use crate::dtype::DType;
460461
use crate::dtype::DecimalDType;
461462
use crate::dtype::FieldNames;
463+
use crate::dtype::MapDType;
462464
use crate::dtype::Nullability;
463465
use crate::dtype::PType;
464466
use crate::expr::stats::Stat;
467+
use crate::scalar::Scalar;
465468
use crate::validity::Validity;
466469

470+
type MapEntryFixture<'a> = (i32, Option<&'a str>);
471+
type MapRowFixture<'a> = Option<Vec<MapEntryFixture<'a>>>;
472+
473+
fn map_array_from_rows(rows: &[MapRowFixture<'_>]) -> VortexResult<crate::ArrayRef> {
474+
let map_dtype = MapDType::try_new(
475+
DType::Primitive(PType::I32, Nullability::NonNullable),
476+
DType::Utf8(Nullability::Nullable),
477+
false,
478+
)?;
479+
let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable);
480+
let mut builder =
481+
MapBuilder::<u64, u64>::with_capacity(map_dtype, Nullability::Nullable, rows.len());
482+
483+
for row in rows {
484+
let scalar = match row {
485+
Some(entries) => {
486+
let entries = entries
487+
.iter()
488+
.map(|(key, value)| {
489+
let key = Scalar::primitive(*key, Nullability::NonNullable);
490+
let value = value.map_or_else(
491+
|| Scalar::null(DType::Utf8(Nullability::Nullable)),
492+
|value| Scalar::utf8(value, Nullability::Nullable),
493+
);
494+
(key, value)
495+
})
496+
.collect::<Vec<_>>();
497+
Scalar::try_map(dtype.clone(), entries)?
498+
}
499+
None => Scalar::null(dtype.clone()),
500+
};
501+
builder.append_value(scalar.as_map())?;
502+
}
503+
504+
Ok(builder.finish_into_map().into_array())
505+
}
506+
467507
// Tests migrated from compute/is_constant.rs
468508
#[test]
469509
fn is_constant_min_max_no_nan() -> VortexResult<()> {
@@ -687,4 +727,26 @@ mod tests {
687727
assert_eq!(is_constant(&list_array.into_array(), &mut ctx)?, expected);
688728
Ok(())
689729
}
730+
731+
#[test]
732+
fn test_map_is_constant() -> VortexResult<()> {
733+
let mut ctx = array_session().create_execution_ctx();
734+
735+
let identical = map_array_from_rows(&[
736+
Some(vec![(1, Some("one")), (2, None)]),
737+
Some(vec![(1, Some("one")), (2, None)]),
738+
])?;
739+
assert!(is_constant(&identical, &mut ctx)?);
740+
741+
let different = map_array_from_rows(&[
742+
Some(vec![(1, Some("one")), (2, None)]),
743+
Some(vec![(1, Some("one")), (3, None)]),
744+
])?;
745+
assert!(!is_constant(&different, &mut ctx)?);
746+
747+
let all_null = map_array_from_rows(&[None, None])?;
748+
assert!(is_constant(&all_null, &mut ctx)?);
749+
750+
Ok(())
751+
}
690752
}

‎vortex-array/src/arrays/arbitrary.rs‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,10 @@ use crate::builders::ArrayBuilder;
2828
use crate::builders::DecimalBuilder;
2929
use crate::builders::FixedSizeListBuilder;
3030
use crate::builders::ListViewBuilder;
31+
use crate::builders::MapBuilder;
3132
use crate::dtype::DType;
3233
use crate::dtype::IntegerPType;
34+
use crate::dtype::MapDType;
3335
use crate::dtype::NativePType;
3436
use crate::dtype::Nullability;
3537
use crate::dtype::OffsetBuilderPType;
@@ -158,7 +160,9 @@ fn random_array_chunk(
158160
DType::FixedSizeList(elem_dtype, list_size, null) => {
159161
random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len)
160162
}
161-
DType::Map(..) => Err(IncorrectFormat),
163+
DType::Map(map_dtype, nullability) => {
164+
random_map(u, map_dtype.clone(), *nullability, chunk_len)
165+
}
162166
DType::Struct(sdt, n) => {
163167
let first_array = sdt
164168
.fields()
@@ -199,6 +203,41 @@ fn random_array_chunk(
199203
}
200204
}
201205

206+
fn random_map(
207+
u: &mut Unstructured,
208+
map_dtype: MapDType,
209+
nullability: Nullability,
210+
chunk_len: Option<usize>,
211+
) -> Result<ArrayRef> {
212+
let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
213+
let key_dtype = map_dtype.key_dtype();
214+
let value_dtype = map_dtype.value_dtype();
215+
let dtype = DType::Map(map_dtype.clone(), nullability);
216+
let mut builder = MapBuilder::<u64, u64>::with_capacity(map_dtype, nullability, array_length);
217+
218+
for _ in 0..array_length {
219+
if nullability == Nullability::Nullable && u.arbitrary::<bool>()? {
220+
builder.append_null();
221+
} else {
222+
let entry_count = u.int_in_range(0..=20)?;
223+
let entries = (0..entry_count)
224+
.map(|_| {
225+
let key = random_scalar(u, &key_dtype)?;
226+
let value = random_scalar(u, &value_dtype)?;
227+
Ok((key, value))
228+
})
229+
.collect::<Result<Vec<_>>>()?;
230+
let scalar = Scalar::try_map(dtype.clone(), entries)
231+
.vortex_expect("generated map scalar should be valid");
232+
builder
233+
.append_scalar(&scalar)
234+
.vortex_expect("generated map scalar should append");
235+
}
236+
}
237+
238+
Ok(builder.finish_into_map().into_array())
239+
}
240+
202241
/// Creates a random fixed-size list array.
203242
///
204243
/// If the `chunk_len` is specified, the length of the array will be equal to the chunk length.

0 commit comments

Comments
 (0)