Skip to content

Commit 695e068

Browse files
authored
Add Map compute and compression support (#9110)
## Rationale for this change <!-- Why are you proposing this change, and what is its impact? Is it part of a long term effort, or a bigger change? If this PR is related to a tracked effort or an open issue, please link to the relevant issue. --> - Closes: #000 ## What changes are included in this PR? <!-- No need to duplicate information from the previous section, but if you're touching many parts of the code base, its worth explicitly noting the important changes or how they are tested. --> ## What APIs are changed? Are there any user-facing changes? <!-- Are there any user-facing changes that might require documentation updates Is any public API changed? --> --------- Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent c6629ac commit 695e068

39 files changed

Lines changed: 1554 additions & 125 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
}
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
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::ListView;
9+
use crate::arrays::MapArray;
10+
use crate::arrays::map::MapArrayExt;
11+
use crate::arrays::map::MapArraySlotsExt;
12+
13+
pub(super) fn check_map_identical(
14+
lhs: &MapArray,
15+
rhs: &MapArray,
16+
ctx: &mut ExecutionCtx,
17+
) -> VortexResult<bool> {
18+
if lhs.map_dtype() != rhs.map_dtype() {
19+
return Ok(false);
20+
}
21+
22+
check_list_identical(
23+
&lhs.entries().as_::<ListView>().into_owned(),
24+
&rhs.entries().as_::<ListView>().into_owned(),
25+
ctx,
26+
)
27+
}

‎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
}
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
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::ListView;
9+
use crate::arrays::MapArray;
10+
use crate::arrays::map::MapArraySlotsExt;
11+
12+
pub(super) fn check_map_constant(map: &MapArray, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
13+
check_listview_constant(&map.entries().as_::<ListView>().into_owned(), ctx)
14+
}

‎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/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ use crate::aggregate_fn::EmptyOptions;
4545
use crate::array::ArrayView;
4646
use crate::arrays::Constant;
4747
use crate::arrays::ConstantArray;
48-
use crate::arrays::map::MapArrayExt;
48+
use crate::arrays::ListView;
49+
use crate::arrays::map::MapArraySlotsExt;
4950
use crate::arrays::varbinview::BinaryView;
5051
use crate::dtype::DType;
5152
use crate::dtype::DecimalType;
@@ -200,9 +201,10 @@ pub(crate) fn canonical_uncompressed_size_in_bytes(
200201
Canonical::Decimal(array) => decimal_uncompressed_size_in_bytes(array, ctx),
201202
Canonical::VarBinView(array) => varbinview_uncompressed_size_in_bytes(array, ctx),
202203
Canonical::List(array) => list_view_uncompressed_size_in_bytes(array, ctx),
203-
Canonical::Map(array) => {
204-
list_view_uncompressed_size_in_bytes(&array.entries().into_owned(), ctx)
205-
}
204+
Canonical::Map(array) => list_view_uncompressed_size_in_bytes(
205+
&array.entries().as_::<ListView>().into_owned(),
206+
ctx,
207+
),
206208
Canonical::FixedSizeList(array) => fixed_size_list_uncompressed_size_in_bytes(array, ctx),
207209
Canonical::Struct(array) => struct_uncompressed_size_in_bytes(array, ctx),
208210
Canonical::Union(array) => union_uncompressed_size_in_bytes(array, ctx),

0 commit comments

Comments
 (0)