datafusion_functions_nested/
reverse.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_reverse function.
19
20use crate::utils::make_scalar_function;
21use arrow::array::{
22    Array, ArrayRef, Capacities, FixedSizeListArray, GenericListArray, MutableArrayData,
23    OffsetSizeTrait,
24};
25use arrow::buffer::OffsetBuffer;
26use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null};
27use arrow::datatypes::{DataType, FieldRef};
28use datafusion_common::cast::{
29    as_fixed_size_list_array, as_large_list_array, as_list_array,
30};
31use datafusion_common::{exec_err, utils::take_function_args, Result};
32use datafusion_expr::{
33    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
34};
35use datafusion_macros::user_doc;
36use std::any::Any;
37use std::sync::Arc;
38
39make_udf_expr_and_func!(
40    ArrayReverse,
41    array_reverse,
42    array,
43    "reverses the order of elements in the array.",
44    array_reverse_udf
45);
46
47#[user_doc(
48    doc_section(label = "Array Functions"),
49    description = "Returns the array with the order of the elements reversed.",
50    syntax_example = "array_reverse(array)",
51    sql_example = r#"```sql
52> select array_reverse([1, 2, 3, 4]);
53+------------------------------------------------------------+
54| array_reverse(List([1, 2, 3, 4]))                          |
55+------------------------------------------------------------+
56| [4, 3, 2, 1]                                               |
57+------------------------------------------------------------+
58```"#,
59    argument(
60        name = "array",
61        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
62    )
63)]
64#[derive(Debug)]
65pub struct ArrayReverse {
66    signature: Signature,
67    aliases: Vec<String>,
68}
69
70impl Default for ArrayReverse {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76impl ArrayReverse {
77    pub fn new() -> Self {
78        Self {
79            signature: Signature::any(1, Volatility::Immutable),
80            aliases: vec!["list_reverse".to_string()],
81        }
82    }
83}
84
85impl ScalarUDFImpl for ArrayReverse {
86    fn as_any(&self) -> &dyn Any {
87        self
88    }
89
90    fn name(&self) -> &str {
91        "array_reverse"
92    }
93
94    fn signature(&self) -> &Signature {
95        &self.signature
96    }
97
98    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
99        Ok(arg_types[0].clone())
100    }
101
102    fn invoke_with_args(
103        &self,
104        args: datafusion_expr::ScalarFunctionArgs,
105    ) -> Result<ColumnarValue> {
106        make_scalar_function(array_reverse_inner)(&args.args)
107    }
108
109    fn aliases(&self) -> &[String] {
110        &self.aliases
111    }
112
113    fn documentation(&self) -> Option<&Documentation> {
114        self.doc()
115    }
116}
117
118/// array_reverse SQL function
119pub fn array_reverse_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
120    let [input_array] = take_function_args("array_reverse", arg)?;
121
122    match &input_array.data_type() {
123        List(field) => {
124            let array = as_list_array(input_array)?;
125            general_array_reverse::<i32>(array, field)
126        }
127        LargeList(field) => {
128            let array = as_large_list_array(input_array)?;
129            general_array_reverse::<i64>(array, field)
130        }
131        FixedSizeList(field, _) => {
132            let array = as_fixed_size_list_array(input_array)?;
133            fixed_size_array_reverse(array, field)
134        }
135        Null => Ok(Arc::clone(input_array)),
136        array_type => exec_err!("array_reverse does not support type '{array_type:?}'."),
137    }
138}
139
140fn general_array_reverse<O: OffsetSizeTrait + TryFrom<i64>>(
141    array: &GenericListArray<O>,
142    field: &FieldRef,
143) -> Result<ArrayRef> {
144    let values = array.values();
145    let original_data = values.to_data();
146    let capacity = Capacities::Array(original_data.len());
147    let mut offsets = vec![O::usize_as(0)];
148    let mut nulls = vec![];
149    let mut mutable =
150        MutableArrayData::with_capacities(vec![&original_data], false, capacity);
151
152    for (row_index, offset_window) in array.offsets().windows(2).enumerate() {
153        // skip the null value
154        if array.is_null(row_index) {
155            nulls.push(false);
156            offsets.push(offsets[row_index] + O::one());
157            mutable.extend(0, 0, 1);
158            continue;
159        } else {
160            nulls.push(true);
161        }
162
163        let start = offset_window[0];
164        let end = offset_window[1];
165
166        let mut index = end - O::one();
167        let mut cnt = 0;
168
169        while index >= start {
170            mutable.extend(0, index.to_usize().unwrap(), index.to_usize().unwrap() + 1);
171            index = index - O::one();
172            cnt += 1;
173        }
174        offsets.push(offsets[row_index] + O::usize_as(cnt));
175    }
176
177    let data = mutable.freeze();
178    Ok(Arc::new(GenericListArray::<O>::try_new(
179        Arc::clone(field),
180        OffsetBuffer::<O>::new(offsets.into()),
181        arrow::array::make_array(data),
182        Some(nulls.into()),
183    )?))
184}
185
186fn fixed_size_array_reverse(
187    array: &FixedSizeListArray,
188    field: &FieldRef,
189) -> Result<ArrayRef> {
190    let values = array.values();
191    let original_data = values.to_data();
192    let capacity = Capacities::Array(original_data.len());
193    let mut nulls = vec![];
194    let mut mutable =
195        MutableArrayData::with_capacities(vec![&original_data], false, capacity);
196    let value_length = array.value_length() as usize;
197
198    for row_index in 0..array.len() {
199        // skip the null value
200        if array.is_null(row_index) {
201            nulls.push(false);
202            mutable.extend(0, 0, 1);
203            continue;
204        } else {
205            nulls.push(true);
206        }
207        let start = row_index * value_length;
208        let end = start + value_length;
209        for idx in (start..end).rev() {
210            mutable.extend(0, idx, idx + 1);
211        }
212    }
213
214    let data = mutable.freeze();
215    Ok(Arc::new(FixedSizeListArray::try_new(
216        Arc::clone(field),
217        array.value_length(),
218        arrow::array::make_array(data),
219        Some(nulls.into()),
220    )?))
221}
datafusion_functions_nested/reverse.rs

datafusion_functions_nested/
reverse.rs
