diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 9a358138bac..8503350cabe 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -23,8 +23,11 @@ use vortex_array::TypedArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::array_slots; use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::ListViewArray; use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -502,19 +505,64 @@ pub(super) fn run_end_canonicalize( .execute_as::("values", ctx)?; runend_decode_varbinview(pends, values, array.offset(), array.len(), ctx)?.into_array() } + DType::List(..) => { + let values = array + .values() + .clone() + .execute_as::("values", ctx)?; + runend_decode_listview(pends, values, array.offset(), array.len())?.into_array() + } _ => vortex_bail!("Unsupported RunEnd value type: {}", array.dtype()), }) } +fn runend_decode_listview( + ends: PrimitiveArray, + values: ListViewArray, + offset: usize, + length: usize, +) -> VortexResult { + let validity = match values.validity()? { + Validity::NonNullable => Validity::NonNullable, + Validity::AllValid => Validity::AllValid, + Validity::AllInvalid => Validity::AllInvalid, + Validity::Array(validity) => Validity::Array(unsafe { + RunEnd::new_unchecked(ends.clone().into_array(), validity, offset, length).into_array() + }), + }; + + // SAFETY: the `RunEndArray`s re-express valid per-run ListView metadata over the logical output + // length. The original `elements` child is reused, so every view still points at a valid range. + Ok(unsafe { + ListViewArray::new_unchecked( + values.elements().clone(), + RunEnd::new_unchecked( + ends.clone().into_array(), + values.offsets().clone(), + offset, + length, + ) + .into_array(), + RunEnd::new_unchecked(ends.into_array(), values.sizes().clone(), offset, length) + .into_array(), + validity, + ) + }) +} + #[cfg(test)] mod tests { + use std::sync::Arc; use std::sync::LazyLock; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::DictArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::builders::VarBinBuilder; use vortex_array::dtype::DType; @@ -522,6 +570,7 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::i256; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -605,6 +654,37 @@ mod tests { assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } + #[test] + fn test_runend_list_i64() { + let mut ctx = SESSION.create_execution_ctx(); + let values = ListArray::from_iter_slow::( + vec![vec![1i64, 2], vec![3], vec![4, 5, 6]], + Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)), + ) + .unwrap() + .into_array(); + let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); + + let expected = ListArray::from_iter_slow::( + vec![ + vec![1i64, 2], + vec![1, 2], + vec![3], + vec![3], + vec![3], + vec![4, 5, 6], + vec![4, 5, 6], + vec![4, 5, 6], + vec![4, 5, 6], + vec![4, 5, 6], + ], + Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)), + ) + .unwrap() + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + } + #[test] fn test_runend_nullable_decimal() { let mut ctx = SESSION.create_execution_ctx(); @@ -638,6 +718,95 @@ mod tests { assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } + #[test] + fn test_runend_list_bool() { + let mut ctx = SESSION.create_execution_ctx(); + let values = ListArray::from_iter_slow::( + vec![vec![true, false], vec![false], vec![true, true, false]], + Arc::new(DType::Bool(Nullability::NonNullable)), + ) + .unwrap() + .into_array(); + let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); + + let expected = ListArray::from_iter_slow::( + vec![ + vec![true, false], + vec![true, false], + vec![false], + vec![false], + vec![false], + vec![true, true, false], + vec![true, true, false], + vec![true, true, false], + vec![true, true, false], + vec![true, true, false], + ], + Arc::new(DType::Bool(Nullability::NonNullable)), + ) + .unwrap() + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + } + + #[test] + fn test_runend_list_utf8() { + let mut ctx = SESSION.create_execution_ctx(); + let values = ListArray::try_new( + VarBinViewArray::from_iter_str(["a", "b", "c", "d", "e", "f"]).into_array(), + buffer![0u32, 2, 3, 6].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); + + let expected = ListArray::try_new( + VarBinViewArray::from_iter_str([ + "a", "b", "a", "b", "c", "c", "c", "d", "e", "f", "d", "e", "f", "d", "e", "f", + "d", "e", "f", "d", "e", "f", + ]) + .into_array(), + buffer![0u32, 2, 4, 5, 6, 7, 10, 13, 16, 19, 22].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + } + + #[test] + fn test_runend_list_canonicalizes_to_runend_listview_slots() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = ListArray::try_new( + buffer![1i64, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 3, 6].into_array(), + Validity::from_iter([true, false, true]), + )? + .into_array(); + let arr = RunEnd::try_new(buffer![2u32, 5, 6].into_array(), values, &mut ctx)?; + + let listview = arr + .clone() + .into_array() + .execute::(&mut ctx)?; + assert!(listview.offsets().is::()); + assert!(listview.sizes().is::()); + match listview.validity()? { + Validity::Array(validity) => assert!(validity.is::()), + validity => panic!("expected array-backed validity, got {validity:?}"), + } + + let expected = ListArray::try_new( + buffer![1i64, 2, 1, 2, 3, 3, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 4, 5, 6, 7, 10].into_array(), + Validity::from_iter([true, true, false, false, false, true]), + )? + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + Ok(()) + } + #[test] fn test_runend_dict() { let mut ctx = SESSION.create_execution_ctx();