feat: add array_subtract scalar function

Author: SubhamSinghal

datafusion/functions-nested/src/array_subtract.rs

@@ -0,0 +1,130 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`ScalarUDFImpl`] definitions for array_subtract function.
+
+use crate::utils::{
+    array_math_binary_op, coerce_array_math_arg_types, make_scalar_function,
+};
+use arrow::array::ArrayRef;
+use arrow::datatypes::{
+    DataType,
+    DataType::{LargeList, List},
+};
+use datafusion_common::{Result, exec_err, utils::take_function_args};
+use datafusion_expr::{
+    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+    Volatility,
+};
+use datafusion_macros::user_doc;
+
+make_udf_expr_and_func!(
+    ArraySubtract,
+    array_subtract,
+    array1 array2,
+    "returns the element-wise difference of two numeric arrays.",
+    array_subtract_udf
+);
+
+#[user_doc(
+    doc_section(label = "Array Functions"),
+    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.",
+    syntax_example = "array_subtract(array1, array2)",
+    sql_example = r#"```sql
+> select array_subtract([10.0, 20.0, 30.0], [1.0, 2.0, 3.0]);
++--------------------------------------------------------------+
+| array_subtract(List([10.0,20.0,30.0]),List([1.0,2.0,3.0]))   |
++--------------------------------------------------------------+
+| [9.0, 18.0, 27.0]                                            |
++--------------------------------------------------------------+
+```"#,
+    argument(
+        name = "array1",
+        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
+    ),
+    argument(
+        name = "array2",
+        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
+    )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ArraySubtract {
+    signature: Signature,
+    aliases: Vec<String>,
+}
+
+impl Default for ArraySubtract {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ArraySubtract {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::user_defined(Volatility::Immutable),
+            aliases: vec!["list_subtract".to_string()],
+        }
+    }
+}
+
+impl ScalarUDFImpl for ArraySubtract {
+    fn name(&self) -> &str {
+        "array_subtract"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        Ok(arg_types[0].clone())
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        let [_, _] = take_function_args(self.name(), arg_types)?;
+        coerce_array_math_arg_types(self.name(), arg_types)
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        make_scalar_function(array_subtract_inner)(&args.args)
+    }
+
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+
+    fn documentation(&self) -> Option<&Documentation> {
+        self.doc()
+    }
+}
+
+fn array_subtract_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let [array1, array2] = take_function_args("array_subtract", args)?;
+    let sub = |a: f64, b: f64| a - b;
+    match (array1.data_type(), array2.data_type()) {
+        (List(_), List(_)) => {
+            array_math_binary_op::<i32, _>("array_subtract", array1, array2, sub)
+        }
+        (LargeList(_), LargeList(_)) => {
+            array_math_binary_op::<i64, _>("array_subtract", array1, array2, sub)
+        }
+        (arg_type1, arg_type2) => exec_err!(
+            "array_subtract received unexpected types after coercion: {arg_type1} and {arg_type2}"
+        ),
+    }
+}

datafusion/functions-nested/src/lib.rs

@@ -47,6 +47,7 @@ pub mod array_filter;
 pub mod array_has;
 pub mod array_normalize;
 pub mod array_scale;
+pub mod array_subtract;
 pub mod array_transform;
 pub mod arrays_zip;
 pub mod cardinality;
@@ -99,6 +100,7 @@ pub mod expr_fn {
     pub use super::array_has::array_has_any;
     pub use super::array_normalize::array_normalize;
     pub use super::array_scale::array_scale;
+    pub use super::array_subtract::array_subtract;
     pub use super::array_transform::array_transform;
     pub use super::arrays_zip::arrays_zip;
     pub use super::cardinality::cardinality;
@@ -176,6 +178,7 @@ pub fn all_default_nested_functions() -> Vec<Arc<ScalarUDF>> {
         array_normalize::array_normalize_udf(),
         array_add::array_add_udf(),
         array_scale::array_scale_udf(),
+        array_subtract::array_subtract_udf(),
         cosine_distance::cosine_distance_udf(),
         inner_product::inner_product_udf(),
         distance::array_distance_udf(),

datafusion/functions-nested/src/utils.rs

@@ -22,12 +22,13 @@ use std::sync::Arc;
 use arrow::datatypes::{DataType, Field, Fields};
 
 use arrow::array::{
-    Array, ArrayRef, BooleanArray, GenericListArray, OffsetSizeTrait, Scalar,
+    Array, ArrayRef, BooleanArray, Float64Array, GenericListArray, NullBufferBuilder,
+    OffsetBufferBuilder, OffsetSizeTrait, Scalar,
 };
-use arrow::buffer::OffsetBuffer;
+use arrow::buffer::{NullBuffer, OffsetBuffer};
 use datafusion_common::cast::{
-    as_fixed_size_list_array, as_large_list_array, as_large_list_view_array,
-    as_list_array, as_list_view_array,
+    as_fixed_size_list_array, as_float64_array, as_generic_list_array,
+    as_large_list_array, as_large_list_view_array, as_list_array, as_list_view_array,
 };
 use datafusion_common::{Result, ScalarValue, exec_err, internal_err, plan_err};
 
@@ -327,6 +328,90 @@ pub(crate) fn coerce_array_math_arg_types(
     Ok(coerced)
 }
 
+/// Element-wise binary operation kernel for two `Float64` lists of equal per-row
+/// length. The caller is responsible for type-dispatching on `O` (`i32` for
+/// `List`, `i64` for `LargeList`).
+///
+/// Semantics:
+/// - whole-row NULL on either side → NULL output row, length 0
+/// - per-element NULL on either side → NULL at that output position
+/// - per-row length mismatch → exec error tagged with `op_name`
+///
+/// `op_name` flows into the error message; `op` is the per-element scalar op
+/// (e.g. `|a, b| a + b` for `array_add`, `|a, b| a - b` for `array_subtract`).
+pub(crate) fn array_math_binary_op<O, F>(
+    op_name: &str,
+    lhs: &ArrayRef,
+    rhs: &ArrayRef,
+    op: F,
+) -> Result<ArrayRef>
+where
+    O: OffsetSizeTrait,
+    F: Fn(f64, f64) -> f64,
+{
+    let lhs = as_generic_list_array::<O>(lhs)?;
+    let rhs = as_generic_list_array::<O>(rhs)?;
+
+    let lhs_values = as_float64_array(lhs.values())?;
+    let rhs_values = as_float64_array(rhs.values())?;
+    let lhs_offsets = lhs.value_offsets();
+    let rhs_offsets = rhs.value_offsets();
+
+    let row_nulls = NullBuffer::union(lhs.nulls(), rhs.nulls());
+
+    let mut out_values: Vec<f64> = Vec::with_capacity(lhs_values.len());
+    let mut out_inner_nulls = NullBufferBuilder::new(lhs_values.len());
+    let mut out_offsets = OffsetBufferBuilder::<O>::new(lhs.len());
+
+    for row in 0..lhs.len() {
+        if row_nulls.as_ref().is_some_and(|nb| nb.is_null(row)) {
+            out_offsets.push_length(0);
+            continue;
+        }
+
+        let start1 = lhs_offsets[row].as_usize();
+        let len1 = lhs.value_length(row).as_usize();
+        let start2 = rhs_offsets[row].as_usize();
+        let len2 = rhs.value_length(row).as_usize();
+
+        if len1 != len2 {
+            return exec_err!(
+                "{op_name} requires both list inputs to have the same length per row, got {len1} and {len2} at row {row}"
+            );
+        }
+
+        let l_slice = lhs_values.slice(start1, len1);
+        let r_slice = rhs_values.slice(start2, len2);
+
+        let l_vals = l_slice.values();
+        let r_vals = r_slice.values();
+
+        for i in 0..len1 {
+            out_values.push(op(l_vals[i], r_vals[i]));
+        }
+
+        match NullBuffer::union(l_slice.nulls(), r_slice.nulls()) {
+            Some(nb) => out_inner_nulls.append_buffer(&nb),
+            None => out_inner_nulls.append_n_non_nulls(len1),
+        }
+
+        out_offsets.push_length(len1);
+    }
+
+    let values_array = Arc::new(Float64Array::new(
+        out_values.into(),
+        out_inner_nulls.finish(),
+    ));
+    let field = Arc::new(Field::new_list_field(DataType::Float64, true));
+
+    Ok(Arc::new(GenericListArray::<O>::try_new(
+        field,
+        out_offsets.finish(),
+        values_array,
+        row_nulls,
+    )?))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;