datafusion_functions_nested/
array_scale.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_scale function.
19
20use crate::utils::make_scalar_function;
21use arrow::array::{Array, ArrayRef, Float64Array, GenericListArray, OffsetSizeTrait};
22use arrow::buffer::{NullBuffer, OffsetBuffer};
23use arrow::datatypes::{
24    DataType,
25    DataType::{FixedSizeList, LargeList, List, Null},
26    Field,
27};
28use datafusion_common::cast::{as_float64_array, as_generic_list_array};
29use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only};
30use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args};
31use datafusion_expr::{
32    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
33    Volatility,
34};
35use datafusion_macros::user_doc;
36use std::sync::Arc;
37
38make_udf_expr_and_func!(
39    ArrayScale,
40    array_scale,
41    array scalar,
42    "scales each element of a numeric array by a scalar.",
43    array_scale_udf
44);
45
46#[user_doc(
47    doc_section(label = "Array Functions"),
48    description = "Returns a new array with each element of the input array multiplied by a scalar value, computed as `array[i] * scalar`. Returns NULL if the input row is NULL or the scalar is NULL. If a NULL element appears in the input array at position `i`, the result element at position `i` is NULL. Returns an empty array for an empty input array.",
49    syntax_example = "array_scale(array, scalar)",
50    sql_example = r#"```sql
51> select array_scale([1.0, 2.0, 3.0], 2.0);
52+----------------------------------+
53| array_scale(List([1.0,2.0,3.0]),Float64(2.0)) |
54+----------------------------------+
55| [2.0, 4.0, 6.0]                  |
56+----------------------------------+
57```"#,
58    argument(
59        name = "array",
60        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
61    ),
62    argument(
63        name = "scalar",
64        description = "Numeric scalar to multiply each element by. Can be a constant or column expression."
65    )
66)]
67#[derive(Debug, PartialEq, Eq, Hash)]
68pub struct ArrayScale {
69    signature: Signature,
70    aliases: Vec<String>,
71}
72
73impl Default for ArrayScale {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79impl ArrayScale {
80    pub fn new() -> Self {
81        Self {
82            signature: Signature::user_defined(Volatility::Immutable),
83            aliases: vec!["list_scale".to_string()],
84        }
85    }
86}
87
88impl ScalarUDFImpl for ArrayScale {
89    fn name(&self) -> &str {
90        "array_scale"
91    }
92
93    fn signature(&self) -> &Signature {
94        &self.signature
95    }
96
97    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
98        // After `coerce_types`, `arg_types[0]` is one of List(Float64) or LargeList(Float64).
99        Ok(arg_types[0].clone())
100    }
101
102    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
103        let [array_type, scalar_type] = take_function_args(self.name(), arg_types)?;
104        let coercion = Some(&ListCoercion::FixedSizedListToList);
105
106        if !matches!(
107            array_type,
108            Null | List(_) | LargeList(_) | FixedSizeList(..)
109        ) {
110            return plan_err!(
111                "{} first argument must be a list type, got {array_type}",
112                self.name()
113            );
114        }
115
116        if !scalar_type.is_numeric() && !matches!(scalar_type, Null) {
117            return plan_err!(
118                "{} second argument must be numeric, got {scalar_type}",
119                self.name()
120            );
121        }
122
123        let coerced_array = if matches!(array_type, Null) {
124            List(Arc::new(Field::new_list_field(DataType::Float64, true)))
125        } else {
126            coerced_type_with_base_type_only(array_type, &DataType::Float64, coercion)
127        };
128
129        Ok(vec![coerced_array, DataType::Float64])
130    }
131
132    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
133        make_scalar_function(array_scale_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_scale_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
146    let [array, scalar] = take_function_args("array_scale", args)?;
147    match array.data_type() {
148        List(_) => general_array_scale::<i32>(array, scalar),
149        LargeList(_) => general_array_scale::<i64>(array, scalar),
150        arg_type => internal_err!(
151            "array_scale received unexpected type after coercion: {arg_type}"
152        ),
153    }
154}
155
156fn general_array_scale<O: OffsetSizeTrait>(
157    array: &ArrayRef,
158    scalar: &ArrayRef,
159) -> Result<ArrayRef> {
160    let list_array = as_generic_list_array::<O>(array)?;
161    let scalar_array = as_float64_array(scalar)?;
162
163    let values = as_float64_array(list_array.values())?;
164    let offsets = list_array.value_offsets();
165
166    // A row is null whenever either input row is null. The scalar applies
167    // uniformly across the array, so a null scalar makes the whole row
168    // undefined; union the two row-level null buffers in a single pass
169    // rather than tracking row nulls inside the value loop.
170    let row_nulls = NullBuffer::union(list_array.nulls(), scalar_array.nulls());
171
172    let mut value_builder = Float64Array::builder(values.len());
173    let mut new_offsets = Vec::<O>::with_capacity(list_array.len() + 1);
174    new_offsets.push(O::zero());
175
176    for row in 0..list_array.len() {
177        if row_nulls.as_ref().is_some_and(|nb| nb.is_null(row)) {
178            new_offsets.push(new_offsets[row]);
179            continue;
180        }
181
182        let start = offsets[row].as_usize();
183        let end = offsets[row + 1].as_usize();
184        let len = end - start;
185        let scalar_val = scalar_array.value(row);
186
187        let slice = values.slice(start, len);
188
189        // Per-element NULL propagation for NULL elements inside the array.
190        for i in 0..len {
191            if slice.is_null(i) {
192                value_builder.append_null();
193            } else {
194                value_builder.append_value(slice.value(i) * scalar_val);
195            }
196        }
197
198        new_offsets.push(new_offsets[row] + O::usize_as(len));
199    }
200
201    let values_array = Arc::new(value_builder.finish());
202
203    // Preserve the inner field from the input array (including any user
204    // metadata). After `coerce_types` the inner type is Float64, but the
205    // input may still carry field-level annotations worth keeping.
206    let field = match list_array.data_type() {
207        List(f) | LargeList(f) => Arc::clone(f),
208        other => {
209            return internal_err!("array_scale unexpected list type: {other}");
210        }
211    };
212
213    Ok(Arc::new(GenericListArray::<O>::try_new(
214        field,
215        OffsetBuffer::new(new_offsets.into()),
216        values_array,
217        row_nulls,
218    )?))
219}
datafusion_functions_nested/array_scale.rs

datafusion_functions_nested/
array_scale.rs
