diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index f6b30b830c7..e0401460f65 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -43,17 +43,20 @@ pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexRes // function might not expect the output `ListViewArray` to have a bunch of leading and trailing // garbage data when they turn it back into a `ListArray`. let list = list.reset_offsets(false, ctx)?; - let list_offsets = list.offsets().clone(); + + // `reset_offsets` leaves a lazy subtraction in the offsets slot, which both the `sizes` below + // and the view's own offsets read. Execute it once here so that it does not run per reader. + let list_offsets = list.offsets().clone().execute::(ctx)?; // Create `sizes` array by computing differences between consecutive offsets. // We use the same `DType` for the sizes as the `offsets` array to ensure compatibility. - let sizes = match_each_integer_ptype!(list_offsets.dtype().as_ptype(), |O| { - build_sizes_from_offsets::(&list, ctx)? + let sizes = match_each_integer_ptype!(list_offsets.ptype(), |O| { + build_sizes_from_offsets::(&list_offsets)? }); // We need to slice the `offsets` to remove the last element (`ListArray` has `n + 1` offsets). debug_assert_eq!(list_offsets.len(), list.len() + 1); - let adjusted_offsets = list_offsets.slice(0..list.len())?; + let adjusted_offsets = list_offsets.into_array().slice(0..list.len())?; // SAFETY: Since everything came from an existing valid `ListArray`, and the `sizes` were // derived from valid and in-order `offsets`, we know these fields are valid. @@ -69,22 +72,19 @@ pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexRes }) } -/// Builds a sizes array from a `ListArray` by computing differences between consecutive offsets. -fn build_sizes_from_offsets( - list: &ListArray, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let len = list.len(); +/// Builds a sizes array by computing differences between consecutive offsets. +/// +/// `offsets` **must** be the `n + 1` sorted offsets of a non-empty `ListArray` of `n` rows. +fn build_sizes_from_offsets(offsets: &PrimitiveArray) -> VortexResult { + let offsets_slice = offsets.as_slice::(); + debug_assert!(offsets_slice.is_sorted()); + + let len = offsets_slice.len() - 1; let mut sizes_builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, len); // Create `UninitRange` for direct memory access. let mut sizes_range = sizes_builder.uninit_range(len); - let offsets = list.offsets().clone().execute::(ctx)?; - let offsets_slice = offsets.as_slice::(); - debug_assert_eq!(len + 1, offsets_slice.len()); - debug_assert!(offsets_slice.is_sorted()); - // Compute sizes as the difference between consecutive offsets. for i in 0..len { let size = offsets_slice[i + 1] - offsets_slice[i]; @@ -101,13 +101,12 @@ fn build_sizes_from_offsets( // TODO(connor)[ListView]: Note that it is not exactly zero-copy because we have to add a single // offset at the end, but it is fast enough. -/// Creates a `ListArray` from a `ListViewArray`. The resulting `ListArray` will not have any -/// leading or trailing garbage data. +/// Creates a `ListArray` from a `ListViewArray`. /// -/// If `ListViewArray::is_zero_copy_to_list` is `true`, then this operation is fast +/// The resulting `ListArray` will not have any leading or trailing garbage data. /// -/// Otherwise, this function fall back to the (very) expensive path and will rebuild the -/// `ListArray` from scratch. +/// This operation is fast when `ListViewArray::is_zero_copy_to_list` is `true`. Otherwise it +/// falls back to the (very) expensive path and rebuilds the `ListArray` from scratch. pub fn list_from_list_view( list_view: ListViewArray, ctx: &mut ExecutionCtx, @@ -136,8 +135,9 @@ pub fn list_from_list_view( // TODO(connor)[ListView]: We can optimize this by always keeping extra memory in `ListViewArray` // offsets for an `n+1`th offset. -/// Builds a `ListArray` offsets array from a `ListViewArray` by constructing `n+1` offsets. -/// The last offset is computed as `last_offset + last_size`. +/// Builds a `ListArray` offsets array from a `ListViewArray` by constructing `n + 1` offsets. +/// +/// The last offset is computed as `last_offset + last_size`, and is zero for an empty array. /// /// # Safety /// @@ -350,6 +350,24 @@ mod tests { Ok(()) } + #[test] + fn test_list_to_listview_resets_nonzero_offsets() -> VortexResult<()> { + let elements = buffer![0i32, 1, 2, 3, 4].into_array(); + let offsets = buffer![2u16, 4, 5].into_array(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?; + + let mut ctx = SESSION.create_execution_ctx(); + let list_view = list_view_from_list(list.clone(), &mut ctx)?; + + assert_arrays_eq!( + buffer![0u16, 2].into_array(), + list_view.offsets().clone(), + &mut ctx + ); + assert_arrays_eq!(list, list_view, &mut ctx); + Ok(()) + } + #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx();