Add lists to vortex (#1524)

I add a simple implementation of lists based off varbin. I hope to unify
parts of both of these.

Author: joseph-isaacs

vortex-array/src/array/arbitrary.rs

@@ -81,6 +81,7 @@ fn random_array(u: &mut Unstructured, dtype: &DType, len: Option<usize>) -> Resu
                     .vortex_unwrap()
                     .into_array())
                 }
+                // TOOD(joe): add arbitrary list
                 DType::List(..) => {
                     todo!("List arrays are not implemented")
                 }

vortex-array/src/array/chunked/canonical.rs

@@ -1,16 +1,18 @@
 use arrow_buffer::{BooleanBufferBuilder, Buffer, MutableBuffer, ScalarBuffer};
-use vortex_dtype::{DType, PType, StructDType};
+use vortex_dtype::{DType, Nullability, PType, StructDType};
 use vortex_error::{vortex_bail, vortex_err, ErrString, VortexExpect, VortexResult};
 
 use crate::array::chunked::ChunkedArray;
 use crate::array::extension::ExtensionArray;
 use crate::array::null::NullArray;
 use crate::array::primitive::PrimitiveArray;
 use crate::array::struct_::StructArray;
-use crate::array::{BinaryView, BoolArray, VarBinViewArray};
+use crate::array::{BinaryView, BoolArray, ListArray, VarBinViewArray};
+use crate::compute::{scalar_at, slice, try_cast};
 use crate::validity::Validity;
 use crate::{
-    ArrayDType, ArrayData, ArrayValidity, Canonical, IntoArrayData, IntoArrayVariant, IntoCanonical,
+    ArrayDType, ArrayData, ArrayLen, ArrayValidity, Canonical, IntoArrayData, IntoArrayVariant,
+    IntoCanonical,
 };
 
 impl IntoCanonical for ChunkedArray {
@@ -88,9 +90,11 @@ pub(crate) fn try_canonicalize_chunks(
             )))
         }
 
-        // TODO(aduffy): better list support
         DType::List(..) => {
-            todo!()
+            // TODO(joe): improve performance, use a listview, once it exists
+
+            let list = pack_lists(chunks.as_slice(), validity, dtype)?;
+            Ok(Canonical::List(list))
         }
 
         DType::Bool(_) => {
@@ -117,6 +121,50 @@ pub(crate) fn try_canonicalize_chunks(
     }
 }
 
+fn pack_lists(chunks: &[ArrayData], validity: Validity, dtype: &DType) -> VortexResult<ListArray> {
+    let len: usize = chunks.iter().map(|c| c.len()).sum();
+    let mut offsets = Vec::with_capacity(len + 1);
+    offsets.push(0);
+    let mut elements = Vec::new();
+    let elem_dtype = dtype
+        .as_list_element()
+        .vortex_expect("ListArray must have List dtype");
+
+    for chunk in chunks {
+        let chunk = chunk.clone().into_list()?;
+        // TODO: handle i32 offsets if they fit.
+        let offsets_arr = try_cast(
+            chunk.offsets(),
+            &DType::Primitive(PType::I64, Nullability::NonNullable),
+        )?
+        .into_primitive()?;
+
+        let first_offset_value: usize = usize::try_from(&scalar_at(offsets_arr.as_ref(), 0)?)?;
+        let last_offset_value: usize =
+            usize::try_from(&scalar_at(offsets_arr.as_ref(), offsets_arr.len() - 1)?)?;
+        elements.push(slice(
+            chunk.elements(),
+            first_offset_value,
+            last_offset_value,
+        )?);
+
+        let adjustment_from_previous = *offsets
+            .last()
+            .ok_or_else(|| vortex_err!("List offsets must have at least one element"))?;
+        offsets.extend(
+            offsets_arr
+                .maybe_null_slice::<i64>()
+                .iter()
+                .skip(1)
+                .map(|off| off + adjustment_from_previous - first_offset_value as i64),
+        );
+    }
+    let chunked_elements = ChunkedArray::try_new(elements, elem_dtype.clone())?.into_array();
+    let offsets = PrimitiveArray::from_vec(offsets, Validity::NonNullable);
+
+    ListArray::try_new(chunked_elements, offsets.into_array(), validity)
+}
+
 /// Swizzle the pointers within a ChunkedArray of StructArrays to instead be a single
 /// StructArray, where the Array for each Field is a ChunkedArray.
 ///
@@ -238,12 +286,17 @@ fn pack_views(
 
 #[cfg(test)]
 mod tests {
-    use vortex_dtype::{DType, Nullability};
+    use std::sync::Arc;
+
+    use vortex_dtype::DType;
+    use vortex_dtype::DType::{List, Primitive};
+    use vortex_dtype::Nullability::NonNullable;
+    use vortex_dtype::PType::I32;
 
     use crate::accessor::ArrayAccessor;
     use crate::array::chunked::canonical::pack_views;
-    use crate::array::{ChunkedArray, StructArray, VarBinViewArray};
-    use crate::compute::slice;
+    use crate::array::{ChunkedArray, ListArray, StructArray, VarBinViewArray};
+    use crate::compute::{scalar_at, slice};
     use crate::validity::Validity;
     use crate::variants::StructArrayTrait;
     use crate::{ArrayDType, ArrayLen, IntoArrayData, IntoArrayVariant, ToArrayData};
@@ -258,7 +311,7 @@ mod tests {
         let array2 = slice(stringview_array().as_ref(), 2, 4).unwrap();
         let packed = pack_views(
             &[array1, array2],
-            &DType::Utf8(Nullability::NonNullable),
+            &DType::Utf8(NonNullable),
             Validity::NonNullable,
         )
         .unwrap();
@@ -308,4 +361,37 @@ mod tests {
             .unwrap();
         assert_eq!(orig_values, canon_values);
     }
+
+    #[test]
+    pub fn pack_nested_lists() {
+        let l1 = ListArray::try_new(
+            vec![1, 2, 3, 4].into_array(),
+            vec![0, 3].into_array(),
+            Validity::NonNullable,
+        )
+        .unwrap();
+
+        let l2 = ListArray::try_new(
+            vec![5, 6].into_array(),
+            vec![0, 2].into_array(),
+            Validity::NonNullable,
+        )
+        .unwrap();
+
+        let chunked_list = ChunkedArray::try_new(
+            vec![l1.clone().into_array(), l2.clone().into_array()],
+            List(Arc::new(Primitive(I32, NonNullable)), NonNullable),
+        );
+
+        let canon_values = chunked_list.unwrap().into_list().unwrap();
+
+        assert_eq!(
+            scalar_at(l1, 0).unwrap(),
+            scalar_at(canon_values.clone(), 0).unwrap()
+        );
+        assert_eq!(
+            scalar_at(l2, 0).unwrap(),
+            scalar_at(canon_values, 1).unwrap()
+        );
+    }
 }

vortex-array/src/array/list/compute/mod.rs

@@ -0,0 +1,39 @@
+use std::sync::Arc;
+
+use itertools::Itertools;
+use vortex_error::VortexResult;
+use vortex_scalar::Scalar;
+
+use crate::array::{ListArray, ListEncoding};
+use crate::compute::{scalar_at, slice, ComputeVTable, ScalarAtFn, SliceFn};
+use crate::{ArrayDType, ArrayData, IntoArrayData};
+
+impl ComputeVTable for ListEncoding {
+    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
+        Some(self)
+    }
+
+    fn slice_fn(&self) -> Option<&dyn SliceFn<ArrayData>> {
+        Some(self)
+    }
+}
+
+impl ScalarAtFn<ListArray> for ListEncoding {
+    fn scalar_at(&self, array: &ListArray, index: usize) -> VortexResult<Scalar> {
+        let elem = array.elements_at(index)?;
+        let scalars: Vec<Scalar> = (0..elem.len()).map(|i| scalar_at(&elem, i)).try_collect()?;
+
+        Ok(Scalar::list(Arc::new(elem.dtype().clone()), scalars))
+    }
+}
+
+impl SliceFn<ListArray> for ListEncoding {
+    fn slice(&self, array: &ListArray, start: usize, stop: usize) -> VortexResult<ArrayData> {
+        Ok(ListArray::try_new(
+            array.elements(),
+            slice(array.offsets(), start, stop + 1)?,
+            array.validity().slice(start, stop)?,
+        )?
+        .into_array())
+    }
+}