datafusion_functions_nested/
array_product.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_product function.
19
20use crate::utils::make_scalar_function;
21use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait};
22use arrow::datatypes::{
23    DataType,
24    DataType::{FixedSizeList, LargeList, List, Null},
25    Field,
26};
27use datafusion_common::cast::{as_float64_array, as_generic_list_array};
28use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only};
29use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args};
30use datafusion_expr::{
31    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
32    Volatility,
33};
34use datafusion_macros::user_doc;
35use std::sync::Arc;
36
37make_udf_expr_and_func!(
38    ArrayProduct,
39    array_product,
40    array,
41    "returns the product of the elements of a numeric array.",
42    array_product_udf
43);
44
45#[user_doc(
46    doc_section(label = "Array Functions"),
47    description = "Returns the product of the elements in the input numeric array. \
48                   NULL elements inside the array are skipped (matching SQL aggregate \
49                   convention). Returns NULL if the input is NULL, every element is  \
50                   NULL, or the array is empty. The result is always returned as \
51                   `Float64`.",
52    syntax_example = "array_product(array)",
53    sql_example = r#"```sql
54> select array_product([1.0, 2.0, 3.0]);
55+------------------------------------+
56| array_product(List([1.0,2.0,3.0])) |
57+------------------------------------+
58| 6.0                                |
59+------------------------------------+
60```"#,
61    argument(
62        name = "array",
63        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
64    )
65)]
66#[derive(Debug, PartialEq, Eq, Hash)]
67pub struct ArrayProduct {
68    signature: Signature,
69    aliases: Vec<String>,
70}
71
72impl Default for ArrayProduct {
73    fn default() -> Self {
74        Self::new()
75    }
76}
77
78impl ArrayProduct {
79    pub fn new() -> Self {
80        Self {
81            signature: Signature::user_defined(Volatility::Immutable),
82            aliases: vec!["list_product".to_string()],
83        }
84    }
85}
86
87impl ScalarUDFImpl for ArrayProduct {
88    fn name(&self) -> &str {
89        "array_product"
90    }
91
92    fn signature(&self) -> &Signature {
93        &self.signature
94    }
95
96    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
97        Ok(DataType::Float64)
98    }
99
100    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
101        let [arg_type] = take_function_args(self.name(), arg_types)?;
102        let coercion = Some(&ListCoercion::FixedSizedListToList);
103
104        if !matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) {
105            return plan_err!("{} does not support type {arg_type}", self.name());
106        }
107
108        let coerced = if matches!(arg_type, Null) {
109            List(Arc::new(Field::new_list_field(DataType::Float64, true)))
110        } else {
111            coerced_type_with_base_type_only(arg_type, &DataType::Float64, coercion)
112        };
113
114        Ok(vec![coerced])
115    }
116
117    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
118        make_scalar_function(array_product_inner)(&args.args)
119    }
120
121    fn aliases(&self) -> &[String] {
122        &self.aliases
123    }
124
125    fn documentation(&self) -> Option<&Documentation> {
126        self.doc()
127    }
128}
129
130fn array_product_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
131    let [array] = take_function_args("array_product", args)?;
132    match array.data_type() {
133        List(_) => general_array_product::<i32>(args),
134        LargeList(_) => general_array_product::<i64>(args),
135        arg_type => internal_err!(
136            "array_product received unexpected type after coercion: {arg_type}"
137        ),
138    }
139}
140
141fn general_array_product<O: OffsetSizeTrait>(arrays: &[ArrayRef]) -> Result<ArrayRef> {
142    let list_array = as_generic_list_array::<O>(&arrays[0])?;
143    let values = as_float64_array(list_array.values())?;
144    let offsets = list_array.value_offsets();
145
146    let mut builder = Float64Array::builder(list_array.len());
147
148    for row in 0..list_array.len() {
149        if list_array.is_null(row) {
150            builder.append_null();
151            continue;
152        }
153
154        let start = offsets[row].as_usize();
155        let end = offsets[row + 1].as_usize();
156
157        let mut prod = 1.0_f64;
158        let mut any_valid = false;
159        for i in start..end {
160            if values.is_valid(i) {
161                prod *= values.value(i);
162                any_valid = true;
163            }
164        }
165
166        if any_valid {
167            builder.append_value(prod);
168        } else {
169            builder.append_null();
170        }
171    }
172
173    Ok(Arc::new(builder.finish()))
174}
datafusion_functions_nested/array_product.rs

datafusion_functions_nested/
array_product.rs
