Skip to main content

datafusion_functions_nested/
array_subtract.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_subtract function.
19
20use crate::utils::{
21    array_math_binary_op, coerce_array_math_arg_types, make_scalar_function,
22};
23use arrow::array::ArrayRef;
24use arrow::datatypes::{
25    DataType,
26    DataType::{LargeList, List},
27};
28use datafusion_common::{Result, exec_err, utils::take_function_args};
29use datafusion_expr::{
30    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
31    Volatility,
32};
33use datafusion_macros::user_doc;
34
35make_udf_expr_and_func!(
36    ArraySubtract,
37    array_subtract,
38    array1 array2,
39    "returns the element-wise difference of two numeric arrays.",
40    array_subtract_udf
41);
42
43#[user_doc(
44    doc_section(label = "Array Functions"),
45    description = "Returns the element-wise difference of two numeric arrays of equal length, computed as `array1[i] - array2[i]` per position. NULL is propagated per element: if either input element at position `i` is NULL, the corresponding output element is NULL (positions are preserved). Returns NULL if either entire input array is NULL. Errors if the per-row lengths differ. Returns an empty array if both inputs are empty.",
46    syntax_example = "array_subtract(array1, array2)",
47    sql_example = r#"```sql
48> select array_subtract([10.0, 20.0, 30.0], [1.0, 2.0, 3.0]);
49+--------------------------------------------------------------+
50| array_subtract(List([10.0,20.0,30.0]),List([1.0,2.0,3.0]))   |
51+--------------------------------------------------------------+
52| [9.0, 18.0, 27.0]                                            |
53+--------------------------------------------------------------+
54```"#,
55    argument(
56        name = "array1",
57        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
58    ),
59    argument(
60        name = "array2",
61        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
62    )
63)]
64#[derive(Debug, PartialEq, Eq, Hash)]
65pub struct ArraySubtract {
66    signature: Signature,
67    aliases: Vec<String>,
68}
69
70impl Default for ArraySubtract {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76impl ArraySubtract {
77    pub fn new() -> Self {
78        Self {
79            signature: Signature::user_defined(Volatility::Immutable),
80            aliases: vec!["list_subtract".to_string()],
81        }
82    }
83}
84
85impl ScalarUDFImpl for ArraySubtract {
86    fn name(&self) -> &str {
87        "array_subtract"
88    }
89
90    fn signature(&self) -> &Signature {
91        &self.signature
92    }
93
94    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
95        Ok(arg_types[0].clone())
96    }
97
98    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
99        let [_, _] = take_function_args(self.name(), arg_types)?;
100        coerce_array_math_arg_types(self.name(), arg_types)
101    }
102
103    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
104        make_scalar_function(array_subtract_inner)(&args.args)
105    }
106
107    fn aliases(&self) -> &[String] {
108        &self.aliases
109    }
110
111    fn documentation(&self) -> Option<&Documentation> {
112        self.doc()
113    }
114}
115
116fn array_subtract_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
117    let [array1, array2] = take_function_args("array_subtract", args)?;
118    let sub = |a: f64, b: f64| a - b;
119    match (array1.data_type(), array2.data_type()) {
120        (List(_), List(_)) => {
121            array_math_binary_op::<i32, _>("array_subtract", array1, array2, sub)
122        }
123        (LargeList(_), LargeList(_)) => {
124            array_math_binary_op::<i64, _>("array_subtract", array1, array2, sub)
125        }
126        (arg_type1, arg_type2) => exec_err!(
127            "array_subtract received unexpected types after coercion: {arg_type1} and {arg_type2}"
128        ),
129    }
130}