Feature: Add FixedSizeList support in vortex-scalar (#4380)

Adds `vortex-scalar` support for `FixedSizeList`.

We reuse `ListScalar` to represent **both** `List` and `FixedSizeList`
logical types, as there is no effective difference between a single
scalar value of a list or fixed-size list.

Also refactors the unit `tests` module to be a bit more clear in what
each submodule is actually testing, plus adds several tests for
`FixedSizeList`.

Also fixed some `no_coercion` tests that checked the wrong things.

For reviewers: it might be a good idea to review each commit separately

---------

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-scalar/src/list.rs

@@ -3,19 +3,26 @@
 
 use std::fmt::{Display, Formatter};
 use std::hash::Hash;
-use std::ops::Deref;
 use std::sync::Arc;
 
 use itertools::Itertools as _;
 use vortex_dtype::{DType, Nullability};
-use vortex_error::{VortexError, VortexExpect as _, VortexResult, vortex_bail, vortex_panic};
+use vortex_error::{
+    VortexError, VortexExpect as _, VortexResult, vortex_bail, vortex_err, vortex_panic,
+};
 
 use crate::{InnerScalarValue, Scalar, ScalarValue};
 
-/// A scalar value representing a list (array) of elements.
+/// A scalar value representing a list or fixed-size list (array) of elements.
 ///
-/// This type provides a view into a list scalar value, which can contain
-/// zero or more elements of the same type, or be null.
+/// We use the same [`ListScalar`] to represent both variants since a single list scalar's data is
+/// identical to a single fixed-size list scalar.
+///
+/// This type provides a view into a list or fixed-size list scalar value which can contain zero or
+/// more elements of the same type, or be null. If the `dtype` is a [`FixedSizeList`], then the
+/// number of `elements` is equal to the `size` field of the [`FixedSizeList`].
+///
+/// [`FixedSizeList`]: DType::FixedSizeList
 #[derive(Debug)]
 pub struct ListScalar<'a> {
     dtype: &'a DType,
@@ -28,9 +35,16 @@ impl Display for ListScalar<'_> {
         match &self.elements {
             None => write!(f, "null"),
             Some(elems) => {
+                let fixed_size_list_str: &dyn Display =
+                    if let DType::FixedSizeList(_, size, _) = self.dtype {
+                        &format!("fixed_size<{size}>")
+                    } else {
+                        &""
+                    };
+
                 write!(
                     f,
-                    "[{}]",
+                    "{fixed_size_list_str}[{}]",
                     elems
                         .iter()
                         .map(|e| Scalar::new(self.element_dtype().clone(), e.clone()))
@@ -101,10 +115,10 @@ impl<'a> ListScalar<'a> {
 
     /// Returns the data type of the list's elements.
     pub fn element_dtype(&self) -> &DType {
-        let DType::List(element_type, _) = self.dtype() else {
-            unreachable!();
-        };
-        (*element_type).deref()
+        self.dtype
+            .as_list_element_opt()
+            .unwrap_or_else(|| vortex_panic!("`ListScalar` somehow had dtype {}", self.dtype))
+            .as_ref()
     }
 
     /// Returns the element at the given index as a scalar.
@@ -129,14 +143,35 @@ impl<'a> ListScalar<'a> {
         })
     }
 
+    /// Casts the list to the target [`DType`].
+    ///
+    /// # Panics
+    ///
+    /// Panics if the target [`DType`] is not a [`List`]: or [`FixedSizeList`], or if trying to cast
+    /// to a [`FixedSizeList`] with the incorrect number of elements.
+    ///
+    /// [`List`]: DType::List
+    /// [`FixedSizeList`]: DType::FixedSizeList
     pub(crate) fn cast(&self, dtype: &DType) -> VortexResult<Scalar> {
-        let DType::List(element_dtype, ..) = dtype else {
+        let target_element_dtype = dtype
+            .as_list_element_opt()
+            .ok_or_else(|| {
+                vortex_err!(
+                    "Cannot cast {} to {}: list can only be cast to a list or fixed-size list",
+                    self.dtype(),
+                    dtype
+                )
+            })?
+            .as_ref();
+
+        if let DType::FixedSizeList(_, size, _) = dtype
+            && *size as usize != self.len()
+        {
             vortex_bail!(
-                "Cannot cast {} to {}: list can only be cast to list",
-                self.dtype(),
-                dtype
+                "tried to cast to a `FixedSizeList[{size}]` but had {} elements",
+                self.len()
             )
-        };
+        }
 
         Ok(Scalar::new(
             dtype.clone(),
@@ -146,61 +181,101 @@ impl<'a> ListScalar<'a> {
                     .vortex_expect("nullness handled in Scalar::cast")
                     .iter()
                     .map(|element| {
+                        // Recursively cast the elements of the list.
                         Scalar::new(DType::clone(self.element_dtype), element.clone())
-                            .cast(element_dtype)
+                            .cast(target_element_dtype)
                             .map(|x| x.value().clone())
                     })
-                    .process_results(|iter| iter.collect())?,
+                    .collect::<VortexResult<Arc<[ScalarValue]>>>()?,
             )),
         ))
     }
 }
 
+/// A helper enum for creating a [`ListScalar`].
+enum ListKind {
+    Variable,
+    FixedSize,
+}
+
+/// Helper functions to create a [`ListScalar`] as a [`Scalar`].
 impl Scalar {
+    fn create_list(
+        element_dtype: impl Into<Arc<DType>>,
+        children: Vec<Scalar>,
+        nullability: Nullability,
+        list_kind: ListKind,
+    ) -> Self {
+        let element_dtype = element_dtype.into();
+
+        let children: Arc<[ScalarValue]> = children
+            .into_iter()
+            .map(|child| {
+                if child.dtype() != &*element_dtype {
+                    vortex_panic!(
+                        "tried to create list of {} with values of type {}",
+                        element_dtype,
+                        child.dtype()
+                    );
+                }
+                child.value
+            })
+            .collect();
+        let size: u32 = children
+            .len()
+            .try_into()
+            .vortex_expect("tried to create a list that was too large");
+
+        let dtype = match list_kind {
+            ListKind::Variable => DType::List(element_dtype, nullability),
+            ListKind::FixedSize => DType::FixedSizeList(element_dtype, size, nullability),
+        };
+
+        Self::new(dtype, ScalarValue(InnerScalarValue::List(children)))
+    }
+
     /// Creates a new list scalar with the given element type and children.
     ///
     /// # Panics
     ///
-    /// Panics if any child scalar has a different type than the element type.
+    /// Panics if any child scalar has a different type than the element type, or if there are too
+    /// many children.
     pub fn list(
         element_dtype: impl Into<Arc<DType>>,
         children: Vec<Scalar>,
         nullability: Nullability,
     ) -> Self {
-        let element_dtype = element_dtype.into();
-        for child in &children {
-            if child.dtype() != &*element_dtype {
-                vortex_panic!(
-                    "tried to create list of {} with values of type {}",
-                    element_dtype,
-                    child.dtype()
-                );
-            }
-        }
-        Self::new(
-            DType::List(element_dtype, nullability),
-            ScalarValue(InnerScalarValue::List(
-                children.into_iter().map(|x| x.value).collect(),
-            )),
-        )
+        Self::create_list(element_dtype, children, nullability, ListKind::Variable)
     }
 
     /// Creates a new empty list scalar with the given element type.
     pub fn list_empty(element_dtype: Arc<DType>, nullability: Nullability) -> Self {
-        Self::new(
-            DType::List(element_dtype, nullability),
-            ScalarValue(InnerScalarValue::List(vec![].into())),
-        )
+        Self::create_list(element_dtype, vec![], nullability, ListKind::Variable)
+    }
+
+    /// Creates a new fixed-size list scalar with the given element type and children.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any child scalar has a different type than the element type, or if there are too
+    /// many children.
+    pub fn fixed_size_list(
+        element_dtype: impl Into<Arc<DType>>,
+        children: Vec<Scalar>,
+        nullability: Nullability,
+    ) -> Self {
+        Self::create_list(element_dtype, children, nullability, ListKind::FixedSize)
     }
 }
 
 impl<'a> TryFrom<&'a Scalar> for ListScalar<'a> {
     type Error = VortexError;
 
     fn try_from(value: &'a Scalar) -> Result<Self, Self::Error> {
-        let DType::List(element_dtype, ..) = value.dtype() else {
-            vortex_bail!("Expected list scalar, found {}", value.dtype())
-        };
+        let element_dtype = value
+            .dtype()
+            .as_list_element_opt()
+            .ok_or_else(|| vortex_err!("Expected list scalar, found {}", value.dtype()))?;
 
         Ok(Self {
             dtype: value.dtype(),

vortex-scalar/src/scalar.rs

@@ -125,8 +125,7 @@ impl Scalar {
             DType::Utf8(_) => self.as_utf8().cast(target),
             DType::Binary(_) => self.as_binary().cast(target),
             DType::Struct(..) => self.as_struct().cast(target),
-            DType::List(..) => self.as_list().cast(target),
-            DType::FixedSizeList(..) => unimplemented!("TODO(connor)[FixedSizeList]"),
+            DType::List(..) | DType::FixedSizeList(..) => self.as_list().cast(target),
             DType::Extension(..) => self.as_extension().cast(target),
         }
     }
@@ -163,12 +162,11 @@ impl Scalar {
                 .fields()
                 .map(|fields| fields.into_iter().map(|f| f.nbytes()).sum::<usize>())
                 .unwrap_or_default(),
-            DType::List(..) => self
+            DType::List(..) | DType::FixedSizeList(..) => self
                 .as_list()
                 .elements()
                 .map(|fields| fields.into_iter().map(|f| f.nbytes()).sum::<usize>())
                 .unwrap_or_default(),
-            DType::FixedSizeList(..) => unimplemented!("TODO(connor)[FixedSizeList]"),
             DType::Extension(_ext_dtype) => self.as_extension().storage().nbytes(),
         }
     }
@@ -199,7 +197,12 @@ impl Scalar {
                 Self::struct_(DType::Struct(sf, nullability), fields)
             }
             DType::List(edt, nullability) => Self::list(edt, vec![], nullability),
-            DType::FixedSizeList(..) => unimplemented!("TODO(connor)[FixedSizeList]"),
+            DType::FixedSizeList(edt, size, nullability) => {
+                let elements = (0..size)
+                    .map(|_| Scalar::default_value(edt.as_ref().clone()))
+                    .collect();
+                Self::list(edt, elements, nullability)
+            }
             DType::Extension(dt) => {
                 let scalar = Self::default_value(dt.storage_dtype().clone());
                 Self::extension(dt, scalar)
@@ -296,6 +299,9 @@ impl Scalar {
 
     /// Returns a view of the scalar as a list scalar.
     ///
+    /// Note that we use [`ListScalar`] to represent **both** [`DType::List`] and
+    /// [`DType::FixedSizeList`].
+    ///
     /// # Panics
     ///
     /// Panics if the scalar is not a list type.
@@ -304,8 +310,11 @@ impl Scalar {
     }
 
     /// Returns a view of the scalar as a list scalar if it has a list type.
+    ///
+    /// Note that we use [`ListScalar`] to represent **both** [`DType::List`] and
+    /// [`DType::FixedSizeList`].
     pub fn as_list_opt(&self) -> Option<ListScalar<'_>> {
-        matches!(self.dtype, DType::List(..)).then(|| self.as_list())
+        matches!(self.dtype, DType::List(..) | DType::FixedSizeList(..)).then(|| self.as_list())
     }
 
     /// Returns a view of the scalar as an extension scalar.
@@ -397,8 +406,7 @@ impl PartialEq for Scalar {
             DType::Utf8(_) => self.as_utf8() == other.as_utf8(),
             DType::Binary(_) => self.as_binary() == other.as_binary(),
             DType::Struct(..) => self.as_struct() == other.as_struct(),
-            DType::List(..) => self.as_list() == other.as_list(),
-            DType::FixedSizeList(..) => unimplemented!("TODO(connor)[FixedSizeList]"),
+            DType::List(..) | DType::FixedSizeList(..) => self.as_list() == other.as_list(),
             DType::Extension(_) => self.as_extension() == other.as_extension(),
         }
     }
@@ -447,8 +455,9 @@ impl PartialOrd for Scalar {
             DType::Utf8(_) => self.as_utf8().partial_cmp(&other.as_utf8()),
             DType::Binary(_) => self.as_binary().partial_cmp(&other.as_binary()),
             DType::Struct(..) => self.as_struct().partial_cmp(&other.as_struct()),
-            DType::List(..) => self.as_list().partial_cmp(&other.as_list()),
-            DType::FixedSizeList(..) => unimplemented!("TODO(connor)[FixedSizeList]"),
+            DType::List(..) | DType::FixedSizeList(..) => {
+                self.as_list().partial_cmp(&other.as_list())
+            }
             DType::Extension(_) => self.as_extension().partial_cmp(&other.as_extension()),
         }
     }
@@ -464,8 +473,7 @@ impl Hash for Scalar {
             DType::Utf8(_) => self.as_utf8().hash(state),
             DType::Binary(_) => self.as_binary().hash(state),
             DType::Struct(..) => self.as_struct().hash(state),
-            DType::List(..) => self.as_list().hash(state),
-            DType::FixedSizeList(..) => unimplemented!("TODO(connor)[FixedSizeList]"),
+            DType::List(..) | DType::FixedSizeList(..) => self.as_list().hash(state),
             DType::Extension(_) => self.as_extension().hash(state),
         }
     }