Support FixedSizeList in list_length; use AnyList matcher

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

vortex-array/src/expr/exprs.rs

@@ -754,11 +754,9 @@ pub fn ext_storage(input: Expression) -> Expression {
 
 // ---- ListLength ----
 
-/// Creates an expression that computes the number of elements in each list.
-///
-/// This is akin to ANSI SQL `CARDINALITY()`, or DuckDB's `len()`/`array_length()`. The result is
-/// a `U64` array; a null list yields a null length. It reads only the list's offsets, not the
-/// element values.
+/// Creates an expression that computes the number of elements in each list
+/// for `List` and `FixedSizeList` inputs. This is akin to ANSI SQL `CARDINALITY()`,
+/// or DuckDB's `len()`/`array_length()`.
 ///
 /// ```rust
 /// # use vortex_array::expr::{list_length, root};

vortex-array/src/scalar_fn/fns/list_length.rs

@@ -10,30 +10,37 @@ use vortex_session::registry::CachedId;
 use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
+use crate::array::ArrayView;
 use crate::arrays::ConstantArray;
 use crate::arrays::List;
 use crate::arrays::ListView;
-use crate::arrays::ListViewArray;
 use crate::arrays::list::ListArrayExt;
 use crate::arrays::listview::ListViewArrayExt;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::PType;
 use crate::expr::Expression;
+use crate::matcher::Matcher;
 use crate::scalar::Scalar;
 use crate::scalar_fn::Arity;
 use crate::scalar_fn::ChildName;
 use crate::scalar_fn::EmptyOptions;
 use crate::scalar_fn::ExecutionArgs;
+use crate::scalar_fn::ReduceCtx;
+use crate::scalar_fn::ReduceNode;
+use crate::scalar_fn::ReduceNodeRef;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::ScalarFnVTable;
+use crate::scalar_fn::ScalarFnVTableExt;
+use crate::scalar_fn::fns::literal::Literal;
 use crate::scalar_fn::fns::operators::Operator;
 
-/// Number of elements in each list of a `List` typed array.
+/// Number of elements in each list of a `List` or `FixedSizeList` typed array.
 ///
-/// This is computed purely from the list's offsets/sizes (the per-list length), never reading the
-/// element *values*. Validity is carried over from the original array.
+/// This is computed purely from the list's offsets (`ListArray`), sizes (`ListViewArray`), or
+/// dtype (`FixedSizeListArray`) without reading the element *values*. Validity is carried over
+/// from the original array.
 #[derive(Clone)]
 pub struct ListLength;
 
@@ -70,8 +77,10 @@ impl ScalarFnVTable for ListLength {
 
     fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
         match &arg_dtypes[0] {
-            DType::List(_, nullable) => Ok(DType::Primitive(PType::U64, *nullable)),
-            other => vortex_bail!("list_length() requires List, got {other}"),
+            DType::List(_, nullable) | DType::FixedSizeList(_, _, nullable) => {
+                Ok(DType::Primitive(PType::U64, *nullable))
+            }
+            other => vortex_bail!("list_length() requires List or FixedSizeList, got {other}"),
         }
     }
 
@@ -89,10 +98,23 @@ impl ScalarFnVTable for ListLength {
             return Ok(ConstantArray::new(len_scalar, args.row_count()).into_array());
         }
 
-        match input.dtype() {
-            DType::List(..) => list_length(&input, nullability, ctx),
-            other => vortex_bail!("list_length() requires List, got {other}"),
+        list_length(&input, nullability, ctx)
+    }
+
+    fn reduce(
+        &self,
+        _options: &Self::Options,
+        node: &dyn ReduceNode,
+        ctx: &dyn ReduceCtx,
+    ) -> VortexResult<Option<ReduceNodeRef>> {
+        // The length of nonnullable fixed-size list is constant
+        if let DType::FixedSizeList(_, size, Nullability::NonNullable) =
+            node.child(0).node_dtype()?
+        {
+            let length = Scalar::primitive(size as u64, Nullability::NonNullable);
+            return Ok(Some(ctx.new_node(Literal.bind(length), &[])?));
         }
+        Ok(None)
     }
 
     fn validity(
@@ -126,26 +148,33 @@ pub(crate) fn list_length(
     nullability: Nullability,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
-    // TODO(mk): short-circuit when array is all null
-
-    let (lengths, validity) = if let Some(list) = array.as_opt::<ListView>() {
-        // Length array is exactly size child
-        (list.sizes().clone(), list.listview_validity())
-    } else if let Some(list) = array.as_opt::<List>() {
-        // `length[i] = offsets[i + 1] - offsets[i]`
-        let offsets = list.offsets();
-        let n = offsets.len().saturating_sub(1);
-        let lengths = offsets
-            .slice(1..offsets.len())?
-            .binary(offsets.slice(0..n)?, Operator::Sub)?;
-        (lengths, list.list_validity())
-    } else {
-        // Otherwise execute to `ListViewArray` and take size child
-        let list = array.clone().execute::<ListViewArray>(ctx)?;
-        (list.sizes().clone(), list.listview_validity())
+    let (lengths, validity) = match array.dtype() {
+        // The length of fixed-size list is constant, so just need to carry over validity
+        DType::FixedSizeList(_, size, _) => {
+            let lengths = ConstantArray::new(
+                Scalar::primitive(*size as u64, Nullability::NonNullable),
+                array.len(),
+            )
+            .into_array();
+            (lengths, array.validity()?)
+        }
+        DType::List(..) => {
+            let list = array.clone().execute_until::<AnyList>(ctx)?;
+
+            if let Some(list) = list.as_opt::<List>() {
+                let lengths = list_length_from_offsets(list)?;
+                (lengths, list.list_validity())
+            } else if let Some(list_view) = list.as_opt::<ListView>() {
+                // Length array is exactly the sizes child
+                (list_view.sizes().clone(), list_view.listview_validity())
+            } else {
+                unreachable!("AnyList matcher guarantees List or ListView")
+            }
+        }
+        other => vortex_bail!("list_length() requires List or FixedSizeList, got {other}"),
     };
 
-    // Cast to the declared `U64` result
+    // Cast to `U64`
     let len = lengths.len();
     let lengths = lengths.cast(DType::Primitive(PType::U64, nullability))?;
 
@@ -157,6 +186,28 @@ pub(crate) fn list_length(
     }
 }
 
+/// Calculate the lengths of `ListArray` elements via the `offsets` child:
+/// `length[i] = offsets[i + 1] - offsets[i]`.
+fn list_length_from_offsets(list: ArrayView<'_, List>) -> VortexResult<ArrayRef> {
+    let offsets = list.offsets();
+    let n = offsets.len().saturating_sub(1);
+
+    offsets
+        .slice(1..offsets.len())?
+        .binary(offsets.slice(0..n)?, Operator::Sub)
+}
+
+/// Matches an `Array<List>` or `Array<ListView>`.
+struct AnyList;
+
+impl Matcher for AnyList {
+    type Match<'a> = ();
+
+    fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
+        (array.as_opt::<List>().is_some() || array.as_opt::<ListView>().is_some()).then_some(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
@@ -171,16 +222,20 @@ mod tests {
     use crate::VortexSessionExecute;
     use crate::arrays::BoolArray;
     use crate::arrays::ConstantArray;
+    use crate::arrays::FixedSizeListArray;
     use crate::arrays::ListArray;
     use crate::arrays::ListViewArray;
     use crate::arrays::PrimitiveArray;
+    use crate::arrays::ScalarFn;
+    use crate::arrays::scalar_fn::ScalarFnArrayExt;
     use crate::assert_arrays_eq;
     use crate::dtype::DType;
     use crate::dtype::Nullability;
     use crate::dtype::PType;
     use crate::expr::list_length;
     use crate::expr::root;
     use crate::scalar::Scalar;
+    use crate::scalar_fn::fns::literal::Literal;
     use crate::validity::Validity;
 
     fn create_list_elements() -> ArrayRef {
@@ -299,6 +354,45 @@ mod tests {
         Ok(())
     }
 
+    fn create_fixed_size_list(validity: Validity) -> ArrayRef {
+        // 4 lists of size 2 over 8 primitive elements.
+        let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8]).into_array();
+        FixedSizeListArray::new(elements, 2, validity, 4).into_array()
+    }
+
+    #[test]
+    fn test_fixed_size_list_length() -> VortexResult<()> {
+        let fsl = create_fixed_size_list(Validity::NonNullable);
+        let result = fsl.apply(&list_length(root()))?;
+
+        // A non-nullable fixed-size list reduces to a constant literal length, never touching the
+        // `ListLength` execution path.
+        assert!(
+            result
+                .as_opt::<ScalarFn>()
+                .is_some_and(|f| f.scalar_fn().as_opt::<Literal>().is_some()),
+            "list_length over a non-nullable FixedSizeList must reduce to a constant literal"
+        );
+        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 2, 2]));
+        Ok(())
+    }
+
+    #[test]
+    fn test_fixed_size_list_length_nullable() -> VortexResult<()> {
+        let fsl = create_fixed_size_list(Validity::Array(
+            BoolArray::from_iter([true, false, true, false]).into_array(),
+        ));
+        let result = fsl.apply(&list_length(root()))?;
+
+        let session = VortexSession::empty();
+        let mut ctx = session.create_execution_ctx();
+        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
+
+        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(2), None, Some(2), None]);
+        assert_arrays_eq!(result, expected);
+        Ok(())
+    }
+
     #[test]
     fn test_display() {
         let expr = list_length(root());