Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 130 additions & 0 deletions datafusion/functions-nested/src/array_subtract.rs
Original file line number Diff line number Diff line change
@@ -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}"
),
}
}
3 changes: 3 additions & 0 deletions datafusion/functions-nested/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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(),
Expand Down
93 changes: 89 additions & 4 deletions datafusion/functions-nested/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down Expand Up @@ -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::*;
Expand Down
Loading
Loading