Re-implement cardinality() on arrays, replacing the one from Datafusion.

assets/trino/functions.sdf.yml

@@ -1432,7 +1432,7 @@ function:
   optional-parameters: []
   returns:
     datatype: bigint
-  implemented-by: !datafusion
+  implemented-by: !rust
   section: array
   cross-link: https://trino.io/docs/current/functions/array.html#cardinality
   description: >

src/trino/cardinality_impl.rs

@@ -16,27 +16,39 @@
 // under the License.
 
 #![allow(non_camel_case_types)]
+use arrow::array::{ArrayRef, UInt64Array};
 use arrow::datatypes::DataType;
+use datafusion::common::cast::as_list_array;
 use datafusion::common::Result;
 use datafusion::error::DataFusionError;
 use datafusion::logical_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
 use datafusion::logical_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility};
 use std::any::Any;
-
-fn cardinality_array_3_invoke(_args: &[ColumnarValue]) -> Result<ColumnarValue> {
-    Err(DataFusionError::NotImplemented(format!(
-        "Not implemented {}:{}",
-        file!(),
-        line!()
-    )))
+use std::sync::Arc;
+
+use crate::utils::apply_unary_kernel;
+
+fn cardinality_array_3_invoke(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+    apply_unary_kernel(&args[0], move |array: &ArrayRef| {
+        let list_array = as_list_array(array)?;
+        let result = list_array
+            .iter()
+            .map(|item| match item {
+                Some(arr) => Some(arr.len() as u64),
+                None => None, // NULL
+            })
+            .collect::<UInt64Array>();
+        Ok(Arc::new(result) as ArrayRef)
+    })
 }
 
-fn cardinality_array_3_return_type(_arg_types: &[DataType]) -> Result<DataType> {
-    Err(DataFusionError::NotImplemented(format!(
-        "Not implemented {}:{}",
-        file!(),
-        line!()
-    )))
+fn cardinality_array_3_return_type(arg_types: &[DataType]) -> Result<DataType> {
+    match arg_types[0] {
+        DataType::List(_) => Ok(DataType::UInt64),
+        _ => Err(DataFusionError::Plan(format!(
+            "The cardinality function only accepts List."
+        ))),
+    }
 }
 
 fn cardinality_array_3_simplify(

src/utils.rs

@@ -106,6 +106,8 @@ pub(super) fn apply_datepart_kernel(arg: &ColumnarValue, part: DatePart) -> Resu
     })
 }
 
+/// Use Arrow `date_part` kernel to compute day of week,
+/// with Monday being 1 through Sunday being 7, per Trino convention..
 pub(super) fn apply_dow_kernel(arg: &ColumnarValue) -> Result<ColumnarValue> {
     apply_unary_kernel(arg, move |arr| {
         let monday0 = date_part(arr, DatePart::DayOfWeekMonday0)?;
@@ -132,7 +134,7 @@ pub(super) fn array_to_columnar(array: ArrayRef) -> ColumnarValue {
     if array.len() == 1 {
         match ScalarValue::try_from_array(&array, 0) {
             Ok(scalar) => ColumnarValue::Scalar(scalar),
-            Err(_) => ColumnarValue::Array(array),
+            Err(_) => ColumnarValue::Array(array), // Should not happen; can just panic instead
         }
     } else {
         ColumnarValue::Array(array)