datafusion_functions_nested/
repeat.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_repeat function.
19
20use crate::utils::make_scalar_function;
21use arrow::array::{
22    Array, ArrayRef, BooleanBufferBuilder, GenericListArray, Int64Array, OffsetSizeTrait,
23    UInt64Array,
24};
25use arrow::buffer::{NullBuffer, OffsetBuffer};
26use arrow::compute;
27use arrow::datatypes::DataType;
28use arrow::datatypes::{
29    DataType::{LargeList, List},
30    Field,
31};
32use datafusion_common::cast::{as_int64_array, as_large_list_array, as_list_array};
33use datafusion_common::types::{NativeType, logical_int64};
34use datafusion_common::{DataFusionError, Result};
35use datafusion_expr::{
36    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
37    Volatility,
38};
39use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
40use datafusion_macros::user_doc;
41use std::sync::Arc;
42
43make_udf_expr_and_func!(
44    ArrayRepeat,
45    array_repeat,
46    element count, // arg name
47    "returns an array containing element `count` times.", // doc
48    array_repeat_udf // internal function name
49);
50
51#[user_doc(
52    doc_section(label = "Array Functions"),
53    description = "Returns an array containing element `count` times.",
54    syntax_example = "array_repeat(element, count)",
55    sql_example = r#"```sql
56> select array_repeat(1, 3);
57+---------------------------------+
58| array_repeat(Int64(1),Int64(3)) |
59+---------------------------------+
60| [1, 1, 1]                       |
61+---------------------------------+
62> select array_repeat([1, 2], 2);
63+------------------------------------+
64| array_repeat(List([1,2]),Int64(2)) |
65+------------------------------------+
66| [[1, 2], [1, 2]]                   |
67+------------------------------------+
68```"#,
69    argument(
70        name = "element",
71        description = "Element expression. Can be a constant, column, or function, and any combination of array operators."
72    ),
73    argument(
74        name = "count",
75        description = "Value of how many times to repeat the element."
76    )
77)]
78#[derive(Debug, PartialEq, Eq, Hash)]
79pub struct ArrayRepeat {
80    signature: Signature,
81    aliases: Vec<String>,
82}
83
84impl Default for ArrayRepeat {
85    fn default() -> Self {
86        Self::new()
87    }
88}
89
90impl ArrayRepeat {
91    pub fn new() -> Self {
92        Self {
93            signature: Signature::coercible(
94                vec![
95                    Coercion::new_exact(TypeSignatureClass::Any),
96                    Coercion::new_implicit(
97                        TypeSignatureClass::Native(logical_int64()),
98                        vec![TypeSignatureClass::Integer],
99                        NativeType::Int64,
100                    ),
101                ],
102                Volatility::Immutable,
103            ),
104            aliases: vec![String::from("list_repeat")],
105        }
106    }
107}
108
109impl ScalarUDFImpl for ArrayRepeat {
110    fn name(&self) -> &str {
111        "array_repeat"
112    }
113
114    fn signature(&self) -> &Signature {
115        &self.signature
116    }
117
118    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
119        let element_type = &arg_types[0];
120        match element_type {
121            LargeList(_) => Ok(LargeList(Arc::new(Field::new_list_field(
122                element_type.clone(),
123                true,
124            )))),
125            _ => Ok(List(Arc::new(Field::new_list_field(
126                element_type.clone(),
127                true,
128            )))),
129        }
130    }
131
132    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
133        make_scalar_function(array_repeat_inner)(&args.args)
134    }
135
136    fn aliases(&self) -> &[String] {
137        &self.aliases
138    }
139
140    fn documentation(&self) -> Option<&Documentation> {
141        self.doc()
142    }
143}
144
145fn array_repeat_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
146    let element = &args[0];
147    let count_array = as_int64_array(&args[1])?;
148
149    match element.data_type() {
150        List(_) => {
151            let list_array = as_list_array(element)?;
152            general_list_repeat::<i32>(list_array, count_array)
153        }
154        LargeList(_) => {
155            let list_array = as_large_list_array(element)?;
156            general_list_repeat::<i64>(list_array, count_array)
157        }
158        _ => general_repeat::<i32>(element, count_array),
159    }
160}
161
162/// For each element of `array[i]` repeat `count_array[i]` times.
163///
164/// Assumption for the input:
165///     1. `count[i] >= 0`
166///     2. `array.len() == count_array.len()`
167///
168/// For example,
169/// ```text
170/// array_repeat(
171///     [1, 2, 3], [2, 0, 1] => [[1, 1], [], [3]]
172/// )
173/// ```
174fn general_repeat<O: OffsetSizeTrait>(
175    array: &ArrayRef,
176    count_array: &Int64Array,
177) -> Result<ArrayRef> {
178    let total_repeated_values: usize = (0..count_array.len())
179        .map(|i| get_count_with_validity(count_array, i))
180        .sum();
181
182    let mut take_indices = Vec::with_capacity(total_repeated_values);
183    let mut offsets = Vec::with_capacity(count_array.len() + 1);
184    offsets.push(O::zero());
185    let mut running_offset = 0usize;
186
187    for idx in 0..count_array.len() {
188        let count = get_count_with_validity(count_array, idx);
189        running_offset = running_offset.checked_add(count).ok_or_else(|| {
190            DataFusionError::Execution(
191                "array_repeat: running_offset overflowed usize".to_string(),
192            )
193        })?;
194        let offset = O::from_usize(running_offset).ok_or_else(|| {
195            DataFusionError::Execution(format!(
196                "array_repeat: offset {running_offset} exceeds the maximum value for offset type"
197            ))
198        })?;
199        offsets.push(offset);
200        take_indices.extend(std::iter::repeat_n(idx as u64, count));
201    }
202
203    // Build the flattened values
204    let repeated_values = compute::take(
205        array.as_ref(),
206        &UInt64Array::from_iter_values(take_indices),
207        None,
208    )?;
209
210    // Construct final ListArray
211    Ok(Arc::new(GenericListArray::<O>::try_new(
212        Arc::new(Field::new_list_field(array.data_type().to_owned(), true)),
213        OffsetBuffer::new(offsets.into()),
214        repeated_values,
215        count_array.nulls().cloned(),
216    )?))
217}
218
219/// Handle List version of `general_repeat`
220///
221/// For each element of `list_array[i]` repeat `count_array[i]` times.
222///
223/// For example,
224/// ```text
225/// array_repeat(
226///     [[1, 2, 3], [4, 5], [6]], [2, 0, 1] => [[[1, 2, 3], [1, 2, 3]], [], [[6]]]
227/// )
228/// ```
229fn general_list_repeat<O: OffsetSizeTrait>(
230    list_array: &GenericListArray<O>,
231    count_array: &Int64Array,
232) -> Result<ArrayRef> {
233    let list_offsets = list_array.value_offsets();
234
235    // calculate capacities for pre-allocation
236    let mut outer_total = 0usize;
237    let mut inner_total = 0usize;
238    for i in 0..count_array.len() {
239        let count = get_count_with_validity(count_array, i);
240        if count > 0 {
241            outer_total += count;
242            if list_array.is_valid(i) {
243                let len = list_offsets[i + 1].to_usize().unwrap()
244                    - list_offsets[i].to_usize().unwrap();
245                inner_total += len * count;
246            }
247        }
248    }
249
250    // Build inner structures
251    let mut inner_offsets = Vec::with_capacity(outer_total + 1);
252    let mut take_indices = Vec::with_capacity(inner_total);
253    let mut inner_nulls = BooleanBufferBuilder::new(outer_total);
254    let mut inner_running = 0usize;
255    inner_offsets.push(O::zero());
256
257    for row_idx in 0..count_array.len() {
258        let count = get_count_with_validity(count_array, row_idx);
259        let list_is_valid = list_array.is_valid(row_idx);
260        let start = list_offsets[row_idx].to_usize().unwrap();
261        let end = list_offsets[row_idx + 1].to_usize().unwrap();
262        let row_len = end - start;
263
264        for _ in 0..count {
265            inner_running = inner_running.checked_add(row_len).ok_or_else(|| {
266                DataFusionError::Execution(
267                    "array_repeat: inner offset overflowed usize".to_string(),
268                )
269            })?;
270            let offset = O::from_usize(inner_running).ok_or_else(|| {
271                DataFusionError::Execution(format!(
272                    "array_repeat: offset {inner_running} exceeds the maximum value for offset type"
273                ))
274            })?;
275            inner_offsets.push(offset);
276            inner_nulls.append(list_is_valid);
277            if list_is_valid {
278                take_indices.extend(start as u64..end as u64);
279            }
280        }
281    }
282
283    // Build inner ListArray
284    let inner_values = compute::take(
285        list_array.values().as_ref(),
286        &UInt64Array::from_iter_values(take_indices),
287        None,
288    )?;
289    let inner_list = GenericListArray::<O>::try_new(
290        Arc::new(Field::new_list_field(list_array.value_type().clone(), true)),
291        OffsetBuffer::new(inner_offsets.into()),
292        inner_values,
293        Some(NullBuffer::new(inner_nulls.finish())),
294    )?;
295
296    // Build outer ListArray
297    Ok(Arc::new(GenericListArray::<O>::try_new(
298        Arc::new(Field::new_list_field(
299            list_array.data_type().to_owned(),
300            true,
301        )),
302        OffsetBuffer::<O>::from_lengths(
303            count_array
304                .iter()
305                .map(|c| c.map(|v| if v > 0 { v as usize } else { 0 }).unwrap_or(0)),
306        ),
307        Arc::new(inner_list),
308        count_array.nulls().cloned(),
309    )?))
310}
311
312/// Helper function to get count from count_array at given index
313/// Return 0 for null values or non-positive count.
314#[inline]
315fn get_count_with_validity(count_array: &Int64Array, idx: usize) -> usize {
316    if count_array.is_null(idx) {
317        0
318    } else {
319        let c = count_array.value(idx);
320        if c > 0 { c as usize } else { 0 }
321    }
322}
datafusion_functions_nested/repeat.rs

datafusion_functions_nested/
repeat.rs
