Skip to content

Commit cec8d20

Browse files
committed
Support runend canonicalization for lists
Signed-off-by: Daniel King <dan@spiraldb.com>
1 parent 97953be commit cec8d20

1 file changed

Lines changed: 143 additions & 0 deletions

File tree

encodings/runend/src/array.rs

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,13 @@ use vortex_array::IntoArray;
2222
use vortex_array::TypedArrayRef;
2323
use vortex_array::VortexSessionExecute;
2424
use vortex_array::array_slots;
25+
use vortex_array::arrays::BoolArray;
2526
use vortex_array::arrays::DecimalArray;
27+
use vortex_array::arrays::ListViewArray;
2628
use vortex_array::arrays::Primitive;
29+
use vortex_array::arrays::PrimitiveArray;
2730
use vortex_array::arrays::VarBinViewArray;
31+
use vortex_array::arrays::listview::ListViewArraySlotsExt;
2832
use vortex_array::buffer::BufferHandle;
2933
use vortex_array::dtype::DType;
3034
use vortex_array::dtype::Nullability;
@@ -502,18 +506,68 @@ pub(super) fn run_end_canonicalize(
502506
.execute_as::<VarBinViewArray>("values", ctx)?;
503507
runend_decode_varbinview(pends, values, array.offset(), array.len(), ctx)?.into_array()
504508
}
509+
DType::List(..) => {
510+
let values = array
511+
.values()
512+
.clone()
513+
.execute_as::<ListViewArray>("values", ctx)?;
514+
runend_decode_listview(pends, values, array.offset(), array.len(), ctx)?.into_array()
515+
}
505516
_ => vortex_bail!("Unsupported RunEnd value type: {}", array.dtype()),
506517
})
507518
}
508519

520+
fn runend_decode_listview(
521+
ends: PrimitiveArray,
522+
values: ListViewArray,
523+
offset: usize,
524+
length: usize,
525+
ctx: &mut ExecutionCtx,
526+
) -> VortexResult<ListViewArray> {
527+
let offsets = values.offsets().clone().execute_as("offsets", ctx)?;
528+
let decoded_offsets =
529+
runend_decode_primitive(ends.clone(), offsets, offset, length, ctx)?.into_array();
530+
531+
let sizes = values.sizes().clone().execute_as("sizes", ctx)?;
532+
let decoded_sizes =
533+
runend_decode_primitive(ends.clone(), sizes, offset, length, ctx)?.into_array();
534+
535+
let validity = match values.validity()? {
536+
Validity::NonNullable => Validity::NonNullable,
537+
Validity::AllValid => Validity::AllValid,
538+
Validity::AllInvalid => Validity::AllInvalid,
539+
Validity::Array(validity) => Validity::Array(runend_decode_bools(
540+
ends,
541+
validity.execute_as::<BoolArray>("validity", ctx)?,
542+
offset,
543+
length,
544+
ctx,
545+
)?),
546+
};
547+
548+
// SAFETY: `decoded_offsets`, `decoded_sizes`, and `validity` are expanded from valid ListView
549+
// metadata for each run. The original `elements` child is reused, so every expanded view still
550+
// points at the same valid element ranges.
551+
Ok(unsafe {
552+
ListViewArray::new_unchecked(
553+
values.elements().clone(),
554+
decoded_offsets,
555+
decoded_sizes,
556+
validity,
557+
)
558+
})
559+
}
560+
509561
#[cfg(test)]
510562
mod tests {
563+
use std::sync::Arc;
511564
use std::sync::LazyLock;
512565

513566
use vortex_array::IntoArray;
514567
use vortex_array::VortexSessionExecute;
515568
use vortex_array::arrays::DecimalArray;
516569
use vortex_array::arrays::DictArray;
570+
use vortex_array::arrays::ListArray;
517571
use vortex_array::arrays::VarBinViewArray;
518572
use vortex_array::assert_arrays_eq;
519573
use vortex_array::builders::VarBinBuilder;
@@ -522,6 +576,7 @@ mod tests {
522576
use vortex_array::dtype::Nullability;
523577
use vortex_array::dtype::PType;
524578
use vortex_array::dtype::i256;
579+
use vortex_array::validity::Validity;
525580
use vortex_buffer::buffer;
526581
use vortex_error::VortexResult;
527582
use vortex_session::VortexSession;
@@ -583,6 +638,94 @@ mod tests {
583638
assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
584639
}
585640

641+
#[test]
642+
fn test_runend_list_i64() {
643+
let mut ctx = SESSION.create_execution_ctx();
644+
let values = ListArray::from_iter_slow::<u32, _>(
645+
vec![vec![1i64, 2], vec![3], vec![4, 5, 6]],
646+
Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
647+
)
648+
.unwrap()
649+
.into_array();
650+
let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
651+
652+
let expected = ListArray::from_iter_slow::<u32, _>(
653+
vec![
654+
vec![1i64, 2],
655+
vec![1, 2],
656+
vec![3],
657+
vec![3],
658+
vec![3],
659+
vec![4, 5, 6],
660+
vec![4, 5, 6],
661+
vec![4, 5, 6],
662+
vec![4, 5, 6],
663+
vec![4, 5, 6],
664+
],
665+
Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
666+
)
667+
.unwrap()
668+
.into_array();
669+
assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
670+
}
671+
672+
#[test]
673+
fn test_runend_list_bool() {
674+
let mut ctx = SESSION.create_execution_ctx();
675+
let values = ListArray::from_iter_slow::<u32, _>(
676+
vec![vec![true, false], vec![false], vec![true, true, false]],
677+
Arc::new(DType::Bool(Nullability::NonNullable)),
678+
)
679+
.unwrap()
680+
.into_array();
681+
let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
682+
683+
let expected = ListArray::from_iter_slow::<u32, _>(
684+
vec![
685+
vec![true, false],
686+
vec![true, false],
687+
vec![false],
688+
vec![false],
689+
vec![false],
690+
vec![true, true, false],
691+
vec![true, true, false],
692+
vec![true, true, false],
693+
vec![true, true, false],
694+
vec![true, true, false],
695+
],
696+
Arc::new(DType::Bool(Nullability::NonNullable)),
697+
)
698+
.unwrap()
699+
.into_array();
700+
assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
701+
}
702+
703+
#[test]
704+
fn test_runend_list_utf8() {
705+
let mut ctx = SESSION.create_execution_ctx();
706+
let values = ListArray::try_new(
707+
VarBinViewArray::from_iter_str(["a", "b", "c", "d", "e", "f"]).into_array(),
708+
buffer![0u32, 2, 3, 6].into_array(),
709+
Validity::NonNullable,
710+
)
711+
.unwrap()
712+
.into_array();
713+
let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
714+
715+
let expected = ListArray::try_new(
716+
VarBinViewArray::from_iter_str([
717+
"a", "b", "a", "b", "c", "c", "c", "d", "e", "f", "d", "e", "f", "d", "e", "f",
718+
"d", "e", "f", "d", "e", "f",
719+
])
720+
.into_array(),
721+
buffer![0u32, 2, 4, 5, 6, 7, 10, 13, 16, 19, 22].into_array(),
722+
Validity::NonNullable,
723+
)
724+
.unwrap()
725+
.into_array();
726+
assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
727+
}
728+
586729
#[test]
587730
fn test_runend_dict() {
588731
let mut ctx = SESSION.create_execution_ctx();

0 commit comments

Comments
 (0)