datafusion_functions_nested/
cardinality.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`ScalarUDFImpl`] definitions for cardinality function.
19
20use crate::utils::make_scalar_function;
21use arrow::array::{
22    Array, ArrayRef, GenericListArray, MapArray, OffsetSizeTrait, UInt64Array,
23};
24use arrow::datatypes::{
25    DataType,
26    DataType::{
27        FixedSizeList, LargeList, LargeListView, List, ListView, Map, Null, UInt64,
28    },
29};
30use datafusion_common::Result;
31use datafusion_common::cast::{
32    as_fixed_size_list_array, as_large_list_array, as_large_list_view_array,
33    as_list_array, as_list_view_array, as_map_array,
34};
35use datafusion_common::exec_err;
36use datafusion_common::utils::{ListCoercion, take_function_args};
37use datafusion_expr::{
38    ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation,
39    ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
40};
41use datafusion_macros::user_doc;
42use std::sync::Arc;
43
44make_udf_expr_and_func!(
45    Cardinality,
46    cardinality,
47    array,
48    "returns the total number of elements in the array or map.",
49    cardinality_udf
50);
51
52impl Cardinality {
53    pub fn new() -> Self {
54        Self {
55            signature: Signature::one_of(
56                vec![
57                    TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
58                        arguments: vec![ArrayFunctionArgument::Array],
59                        array_coercion: Some(ListCoercion::FixedSizedListToList),
60                    }),
61                    TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray),
62                ],
63                Volatility::Immutable,
64            ),
65        }
66    }
67}
68
69#[user_doc(
70    doc_section(label = "Array Functions"),
71    description = "Returns the total number of elements in the array.",
72    syntax_example = "cardinality(array)",
73    sql_example = r#"```sql
74> select cardinality([[1, 2, 3, 4], [5, 6, 7, 8]]);
75+--------------------------------------+
76| cardinality(List([1,2,3,4,5,6,7,8])) |
77+--------------------------------------+
78| 8                                    |
79+--------------------------------------+
80```"#,
81    argument(
82        name = "array",
83        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
84    )
85)]
86#[derive(Debug, PartialEq, Eq, Hash)]
87pub struct Cardinality {
88    signature: Signature,
89}
90
91impl Default for Cardinality {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96impl ScalarUDFImpl for Cardinality {
97    fn name(&self) -> &str {
98        "cardinality"
99    }
100
101    fn signature(&self) -> &Signature {
102        &self.signature
103    }
104
105    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
106        Ok(UInt64)
107    }
108
109    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
110        make_scalar_function(cardinality_inner)(&args.args)
111    }
112
113    fn documentation(&self) -> Option<&Documentation> {
114        self.doc()
115    }
116}
117
118fn cardinality_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
119    let [array] = take_function_args("cardinality", args)?;
120    match array.data_type() {
121        Null => Ok(Arc::new(UInt64Array::new_null(array.len()))),
122        List(_) => {
123            let list_array = as_list_array(array)?;
124            generic_list_cardinality::<i32>(list_array)
125        }
126        LargeList(_) => {
127            let list_array = as_large_list_array(array)?;
128            generic_list_cardinality::<i64>(list_array)
129        }
130        Map(_, _) => {
131            let map_array = as_map_array(array)?;
132            generic_map_cardinality(map_array)
133        }
134        arg_type => {
135            exec_err!("cardinality does not support type {arg_type}")
136        }
137    }
138}
139
140fn generic_map_cardinality(array: &MapArray) -> Result<ArrayRef> {
141    let result: UInt64Array = array
142        .iter()
143        .map(|opt_arr| opt_arr.map(|arr| arr.len() as u64))
144        .collect();
145    Ok(Arc::new(result))
146}
147
148fn generic_list_cardinality<O: OffsetSizeTrait>(
149    array: &GenericListArray<O>,
150) -> Result<ArrayRef> {
151    let result = array
152        .iter()
153        .map(|arr| match arr {
154            Some(arr) => value_cardinality(&arr).map(Some),
155            None => Ok(None),
156        })
157        .collect::<Result<UInt64Array>>()?;
158    Ok(Arc::new(result) as ArrayRef)
159}
160
161fn value_cardinality(array: &ArrayRef) -> Result<u64> {
162    match array.data_type() {
163        List(_) => {
164            let list = as_list_array(&array)?;
165            sum_list_cardinality(list.iter())
166        }
167        LargeList(_) => {
168            let list = as_large_list_array(&array)?;
169            sum_list_cardinality(list.iter())
170        }
171        ListView(_) => {
172            let list = as_list_view_array(&array)?;
173            sum_list_cardinality(list.iter())
174        }
175        LargeListView(_) => {
176            let list = as_large_list_view_array(&array)?;
177            sum_list_cardinality(list.iter())
178        }
179        FixedSizeList(..) => {
180            let list = as_fixed_size_list_array(&array)?;
181            sum_list_cardinality(list.iter())
182        }
183        _ => Ok(array.len() as u64),
184    }
185}
186
187fn sum_list_cardinality<I>(mut iter: I) -> Result<u64>
188where
189    I: Iterator<Item = Option<ArrayRef>>,
190{
191    iter.try_fold(0u64, |total, arr| {
192        let value_count = match arr {
193            Some(arr) => value_cardinality(&arr)?,
194            None => 0,
195        };
196        total.checked_add(value_count).ok_or_else(|| {
197            datafusion_common::exec_datafusion_err!("cardinality overflowed u64")
198        })
199    })
200}
datafusion_functions_nested/cardinality.rs

datafusion_functions_nested/
cardinality.rs
