datafusion_functions_nested/
dimension.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 array_dims and array_ndims functions.
19
20use arrow::array::{Array, ArrayRef, ListArray, UInt64Array};
21use arrow::datatypes::{
22    DataType,
23    DataType::{FixedSizeList, LargeList, List, Null, UInt64},
24    UInt64Type,
25};
26
27use datafusion_common::cast::{
28    as_fixed_size_list_array, as_large_list_array, as_list_array,
29};
30use datafusion_common::{Result, exec_err, utils::take_function_args};
31
32use crate::utils::{compute_array_dims, make_scalar_function};
33use datafusion_common::utils::list_ndims;
34use datafusion_expr::{
35    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
36    Volatility,
37};
38use datafusion_macros::user_doc;
39use itertools::Itertools;
40use std::sync::Arc;
41
42make_udf_expr_and_func!(
43    ArrayDims,
44    array_dims,
45    array,
46    "returns an array of the array's dimensions.",
47    array_dims_udf
48);
49
50#[user_doc(
51    doc_section(label = "Array Functions"),
52    description = "Returns an array of the array's dimensions.",
53    syntax_example = "array_dims(array)",
54    sql_example = r#"```sql
55> select array_dims([[1, 2, 3], [4, 5, 6]]);
56+---------------------------------+
57| array_dims(List([1,2,3,4,5,6])) |
58+---------------------------------+
59| [2, 3]                          |
60+---------------------------------+
61```"#,
62    argument(
63        name = "array",
64        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
65    )
66)]
67#[derive(Debug, PartialEq, Eq, Hash)]
68pub struct ArrayDims {
69    signature: Signature,
70    aliases: Vec<String>,
71}
72
73impl Default for ArrayDims {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79impl ArrayDims {
80    pub fn new() -> Self {
81        Self {
82            signature: Signature::arrays(1, None, Volatility::Immutable),
83            aliases: vec!["list_dims".to_string()],
84        }
85    }
86}
87
88impl ScalarUDFImpl for ArrayDims {
89    fn name(&self) -> &str {
90        "array_dims"
91    }
92
93    fn signature(&self) -> &Signature {
94        &self.signature
95    }
96
97    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
98        Ok(DataType::new_list(UInt64, true))
99    }
100
101    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
102        make_scalar_function(array_dims_inner)(&args.args)
103    }
104
105    fn aliases(&self) -> &[String] {
106        &self.aliases
107    }
108
109    fn documentation(&self) -> Option<&Documentation> {
110        self.doc()
111    }
112}
113
114make_udf_expr_and_func!(
115    ArrayNdims,
116    array_ndims,
117    array,
118    "returns the number of dimensions of the array.",
119    array_ndims_udf
120);
121
122#[user_doc(
123    doc_section(label = "Array Functions"),
124    description = "Returns the number of dimensions of the array.",
125    syntax_example = "array_ndims(array)",
126    sql_example = r#"```sql
127> select array_ndims([[1, 2, 3], [4, 5, 6]]);
128+----------------------------------+
129| array_ndims(List([1,2,3,4,5,6])) |
130+----------------------------------+
131| 2                                |
132+----------------------------------+
133```"#,
134    argument(
135        name = "array",
136        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
137    )
138)]
139#[derive(Debug, PartialEq, Eq, Hash)]
140pub(super) struct ArrayNdims {
141    signature: Signature,
142    aliases: Vec<String>,
143}
144impl ArrayNdims {
145    pub fn new() -> Self {
146        Self {
147            signature: Signature::arrays(1, None, Volatility::Immutable),
148            aliases: vec![String::from("list_ndims")],
149        }
150    }
151}
152
153impl ScalarUDFImpl for ArrayNdims {
154    fn name(&self) -> &str {
155        "array_ndims"
156    }
157
158    fn signature(&self) -> &Signature {
159        &self.signature
160    }
161
162    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
163        Ok(UInt64)
164    }
165
166    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
167        make_scalar_function(array_ndims_inner)(&args.args)
168    }
169
170    fn aliases(&self) -> &[String] {
171        &self.aliases
172    }
173
174    fn documentation(&self) -> Option<&Documentation> {
175        self.doc()
176    }
177}
178
179fn array_dims_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
180    let [array] = take_function_args("array_dims", args)?;
181    let data: Vec<_> = match array.data_type() {
182        List(_) => as_list_array(&array)?
183            .iter()
184            .map(compute_array_dims)
185            .try_collect()?,
186        LargeList(_) => as_large_list_array(&array)?
187            .iter()
188            .map(compute_array_dims)
189            .try_collect()?,
190        FixedSizeList(..) => as_fixed_size_list_array(&array)?
191            .iter()
192            .map(compute_array_dims)
193            .try_collect()?,
194        arg_type => {
195            return exec_err!("array_dims does not support type {arg_type}");
196        }
197    };
198
199    let result = ListArray::from_iter_primitive::<UInt64Type, _, _>(data);
200    Ok(Arc::new(result))
201}
202
203fn array_ndims_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
204    let [array] = take_function_args("array_ndims", args)?;
205
206    fn general_list_ndims(array: &ArrayRef) -> Result<ArrayRef> {
207        let ndims = list_ndims(array.data_type());
208        let data = vec![ndims; array.len()];
209        let result = UInt64Array::new(data.into(), array.nulls().cloned());
210        Ok(Arc::new(result))
211    }
212
213    match array.data_type() {
214        Null => Ok(Arc::new(UInt64Array::new_null(array.len()))),
215        List(_) | LargeList(_) | FixedSizeList(..) => general_list_ndims(array),
216        arg_type => exec_err!("array_ndims does not support type {arg_type}"),
217    }
218}
datafusion_functions_nested/dimension.rs

datafusion_functions_nested/
dimension.rs
