You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/05/20 02:02:55 UTC

(datafusion) 01/01: feat: add array_normalize scalar function (#22013)

This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22013-4f2b4e3d88eeafb085cf72cf6d2d3cdf5e6bb7c4
in repository https://gitbox.apache.org/repos/asf/datafusion.git

commit 821260fdd8e9a4e5e539a02310d08fc2601a7670
Author: crm26 <58...@users.noreply.github.com>
AuthorDate: Tue May 19 22:02:31 2026 -0400

    feat: add array_normalize scalar function (#22013)
    
    ## Which issue does this PR close?
    
    Part of #21536 — split of #21371 into one-function-per-PR. Third in the
    series after #21542 (cosine_distance) and #21861 (inner_product).
    
    ## Rationale for this change
    
    Adds `array_normalize(array)` — the L2-normalized version of a numeric
    input vector. Computed as `array[i] / sqrt(sum(array[i]^2))` per
    element. Returns the same shape as the input (`List<Float64>` or
    `LargeList<Float64>`).
    
    Aliased as `list_normalize` to match the `array_X`/`list_X` convention
    used across the crate.
    
    ## What changes are included in this PR?
    
    Coercion shell mirrors the merged cosine_distance/inner_product pattern:
    - `coerce_types` accepts `List`/`LargeList`/`FixedSizeList` of any
    numeric inner type, plus bare `NULL`. After coercion the inner function
    only sees `List(Float64)` or `LargeList(Float64)`.
    - Per-row L2 norm computed inline (no shared module), using a single
    `as_float64_array(list_array.values())` downcast plus `value_offsets()`
    slicing — no per-row downcasts.
    - Manual list builder: `Vec<f64>` for values, `Vec<O>` for offsets,
    `NullBuffer` for row validity.
    
    Per-row semantics:
    - NULL row → NULL output
    - NULL element in list → NULL row
    - Empty list → empty list (no division-by-zero hazard)
    - Zero magnitude → NULL row (consistent with cosine_distance's
    zero-magnitude → NULL)
    - Otherwise → divide each element by `sqrt(sum-of-squares)`
    
    ## Are these changes tested?
    
    Yes. SLT covers:
    - 3-4-5 right triangle, 3D vector, already-unit-axis, single non-zero
    component, negative components
    - Bare `NULL` input, NULL element in list, zero vector, empty array
    - `LargeList`, `FixedSizeList` (via coercion), `Float32` and `Int64`
    inner types, integer literals
    - Multi-row query mixing normal / NULL row / zero-vector row /
    null-element row
    - Plan error for non-list input
    - No-args error
    - Return-type assertion (`List(Float64)`)
    - `list_normalize` alias coverage (constant + multi-row with NULL)
    
    ## Are there any user-facing changes?
    
    New scalar function `array_normalize` (alias `list_normalize`),
    documented in `docs/source/user-guide/sql/scalar_functions.md`.
---
 datafusion/functions-nested/src/array_normalize.rs | 207 +++++++++++++++++++++
 datafusion/functions-nested/src/lib.rs             |   3 +
 .../sqllogictest/test_files/array_normalize.slt    | 146 +++++++++++++++
 docs/source/user-guide/sql/scalar_functions.md     |  33 ++++
 4 files changed, 389 insertions(+)

diff --git a/datafusion/functions-nested/src/array_normalize.rs b/datafusion/functions-nested/src/array_normalize.rs
new file mode 100644
index 0000000000..0ff7674032
--- /dev/null
+++ b/datafusion/functions-nested/src/array_normalize.rs
@@ -0,0 +1,207 @@
+// 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_normalize function.
+
+use crate::utils::make_scalar_function;
+use arrow::array::{
+    Array, ArrayRef, Float64Array, GenericListArray, NullBufferBuilder,
+    OffsetBufferBuilder, OffsetSizeTrait,
+};
+use arrow::datatypes::{
+    DataType,
+    DataType::{FixedSizeList, LargeList, List, Null},
+    Field,
+};
+use datafusion_common::cast::{as_float64_array, as_generic_list_array};
+use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only};
+use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args};
+use datafusion_expr::{
+    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+    Volatility,
+};
+use datafusion_macros::user_doc;
+use std::sync::Arc;
+
+make_udf_expr_and_func!(
+    ArrayNormalize,
+    array_normalize,
+    array,
+    "returns the L2-normalized vector for a numeric array.",
+    array_normalize_udf
+);
+
+#[user_doc(
+    doc_section(label = "Array Functions"),
+    description = "Returns the L2-normalized vector for the input numeric array, computed as `array[i] / sqrt(sum(array[i]^2))` per element. Returns NULL if the input is NULL, contains NULL elements, or has zero magnitude (all elements are zero). Returns an empty array for an empty input array.",
+    syntax_example = "array_normalize(array)",
+    sql_example = r#"```sql
+> select array_normalize([3.0, 4.0]);
++-----------------------------+
+| array_normalize(List([3.0,4.0])) |
++-----------------------------+
+| [0.6, 0.8]                  |
++-----------------------------+
+```"#,
+    argument(
+        name = "array",
+        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
+    )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ArrayNormalize {
+    signature: Signature,
+    aliases: Vec<String>,
+}
+
+impl Default for ArrayNormalize {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ArrayNormalize {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::user_defined(Volatility::Immutable),
+            aliases: vec!["list_normalize".to_string()],
+        }
+    }
+}
+
+impl ScalarUDFImpl for ArrayNormalize {
+    fn name(&self) -> &str {
+        "array_normalize"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        // After `coerce_types`, `arg_types[0]` is one of List(Float64) or LargeList(Float64).
+        Ok(arg_types[0].clone())
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        let [arg_type] = take_function_args(self.name(), arg_types)?;
+        let coercion = Some(&ListCoercion::FixedSizedListToList);
+
+        if !matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) {
+            return plan_err!("{} does not support type {arg_type}", self.name());
+        }
+
+        let coerced = if matches!(arg_type, Null) {
+            List(Arc::new(Field::new_list_field(DataType::Float64, true)))
+        } else {
+            coerced_type_with_base_type_only(arg_type, &DataType::Float64, coercion)
+        };
+
+        Ok(vec![coerced])
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        make_scalar_function(array_normalize_inner)(&args.args)
+    }
+
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+
+    fn documentation(&self) -> Option<&Documentation> {
+        self.doc()
+    }
+}
+
+fn array_normalize_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let [array] = take_function_args("array_normalize", args)?;
+    match array.data_type() {
+        List(_) => general_array_normalize::<i32>(args),
+        LargeList(_) => general_array_normalize::<i64>(args),
+        arg_type => internal_err!(
+            "array_normalize received unexpected type after coercion: {arg_type}"
+        ),
+    }
+}
+
+fn general_array_normalize<O: OffsetSizeTrait>(arrays: &[ArrayRef]) -> Result<ArrayRef> {
+    let list_array = as_generic_list_array::<O>(&arrays[0])?;
+    let values = as_float64_array(list_array.values())?;
+    let offsets = list_array.value_offsets();
+
+    let mut new_values: Vec<f64> = Vec::with_capacity(values.len());
+    let mut new_offsets = OffsetBufferBuilder::<O>::new(list_array.len());
+    let mut nulls = NullBufferBuilder::new(list_array.len());
+
+    for row in 0..list_array.len() {
+        if list_array.is_null(row) {
+            nulls.append_null();
+            new_offsets.push_length(0);
+            continue;
+        }
+
+        let start = offsets[row].as_usize();
+        let end = offsets[row + 1].as_usize();
+        let len = end - start;
+
+        let slice = values.slice(start, len);
+        if slice.null_count() != 0 {
+            nulls.append_null();
+            new_offsets.push_length(0);
+            continue;
+        }
+
+        let vals = slice.values();
+
+        // Empty array: return empty array (no normalization needed, no division by zero risk)
+        if len == 0 {
+            nulls.append_non_null();
+            new_offsets.push_length(0);
+            continue;
+        }
+
+        // Compute squared magnitude.
+        let mut sq_sum = 0.0;
+        for i in 0..len {
+            sq_sum += vals[i] * vals[i];
+        }
+
+        // Zero magnitude: undefined normalization. Emit NULL row.
+        if sq_sum == 0.0 {
+            nulls.append_null();
+            new_offsets.push_length(0);
+            continue;
+        }
+
+        let mag = sq_sum.sqrt();
+        for i in 0..len {
+            new_values.push(vals[i] / mag);
+        }
+        nulls.append_non_null();
+        new_offsets.push_length(len);
+    }
+
+    let values_array = Arc::new(Float64Array::from(new_values));
+    let field = Arc::new(Field::new_list_field(DataType::Float64, true));
+
+    Ok(Arc::new(GenericListArray::<O>::try_new(
+        field,
+        new_offsets.finish(),
+        values_array,
+        nulls.finish(),
+    )?))
+}
diff --git a/datafusion/functions-nested/src/lib.rs b/datafusion/functions-nested/src/lib.rs
index e4d08a9a5e..6c4556902b 100644
--- a/datafusion/functions-nested/src/lib.rs
+++ b/datafusion/functions-nested/src/lib.rs
@@ -43,6 +43,7 @@ pub mod macros_lambda;
 pub mod array_any_match;
 pub mod array_compact;
 pub mod array_has;
+pub mod array_normalize;
 pub mod array_transform;
 pub mod arrays_zip;
 pub mod cardinality;
@@ -90,6 +91,7 @@ pub mod expr_fn {
     pub use super::array_has::array_has;
     pub use super::array_has::array_has_all;
     pub use super::array_has::array_has_any;
+    pub use super::array_normalize::array_normalize;
     pub use super::array_transform::array_transform;
     pub use super::arrays_zip::arrays_zip;
     pub use super::cardinality::cardinality;
@@ -164,6 +166,7 @@ pub fn all_default_nested_functions() -> Vec<Arc<ScalarUDF>> {
         array_has::array_has_any_udf(),
         empty::array_empty_udf(),
         length::array_length_udf(),
+        array_normalize::array_normalize_udf(),
         cosine_distance::cosine_distance_udf(),
         inner_product::inner_product_udf(),
         distance::array_distance_udf(),
diff --git a/datafusion/sqllogictest/test_files/array_normalize.slt b/datafusion/sqllogictest/test_files/array_normalize.slt
new file mode 100644
index 0000000000..ba4711d02c
--- /dev/null
+++ b/datafusion/sqllogictest/test_files/array_normalize.slt
@@ -0,0 +1,146 @@
+# 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.
+
+## array_normalize
+
+# 3-4-5 right triangle: [3,4] / 5 = [0.6, 0.8]
+query ?
+select array_normalize([3.0, 4.0]);
+----
+[0.6, 0.8]
+
+# Already-unit vector along axis: [1,0] -> [1,0]
+query ?
+select array_normalize([1.0, 0.0]);
+----
+[1.0, 0.0]
+
+# 3D vector: [1,2,2] / 3 = [0.333..., 0.666..., 0.666...]
+query ?
+select array_normalize([1.0, 2.0, 2.0]);
+----
+[0.3333333333333333, 0.6666666666666666, 0.6666666666666666]
+
+# Negative components preserved
+query ?
+select array_normalize([-3.0, 4.0]);
+----
+[-0.6, 0.8]
+
+# Bare NULL input returns NULL
+query ?
+select array_normalize(NULL);
+----
+NULL
+
+# NULL element inside a list returns NULL for that row
+query ?
+select array_normalize([1.0, NULL, 2.0]);
+----
+NULL
+
+# Zero vector returns NULL (undefined normalization)
+query ?
+select array_normalize([0.0, 0.0]);
+----
+NULL
+
+# Single non-zero component: [5] -> [1]
+query ?
+select array_normalize([5.0]);
+----
+[1.0]
+
+# LargeList support
+query ?
+select array_normalize(arrow_cast([3.0, 4.0], 'LargeList(Float64)'));
+----
+[0.6, 0.8]
+
+# FixedSizeList input (coerced to List)
+query ?
+select array_normalize(arrow_cast([3.0, 4.0], 'FixedSizeList(2, Float64)'));
+----
+[0.6, 0.8]
+
+# Float32 inner type (coerced to Float64)
+query ?
+select array_normalize(arrow_cast([3.0, 4.0], 'List(Float32)'));
+----
+[0.6, 0.8]
+
+# Int64 inner type (coerced to Float64)
+query ?
+select array_normalize(arrow_cast([3, 4], 'List(Int64)'));
+----
+[0.6, 0.8]
+
+# Integer literals (coerced to Float64)
+query ?
+select array_normalize([3, 4]);
+----
+[0.6, 0.8]
+
+# Unsupported non-list input (plan error)
+query error array_normalize does not support type
+select array_normalize(1);
+
+# Multi-row query: normal row, NULL row, zero-vector row, NULL-element row
+query ?
+select array_normalize(column1) from (values
+    (make_array(3.0, 4.0)),
+    (NULL),
+    (make_array(0.0, 0.0)),
+    (make_array(1.0, NULL))
+) as t(column1);
+----
+[0.6, 0.8]
+NULL
+NULL
+NULL
+
+# Empty array: returns empty array (no normalization needed, no division by zero)
+query ?
+select array_normalize(arrow_cast(make_array(), 'List(Float64)'));
+----
+[]
+
+# No arguments error
+query error array_normalize function requires 1 argument, got 0
+select array_normalize();
+
+# Return type matches input variant (List of Float64)
+query ?T
+select array_normalize([3.0, 4.0]), arrow_typeof(array_normalize([3.0, 4.0]));
+----
+[0.6, 0.8] List(Float64)
+
+# list_normalize alias produces the same result
+query ?
+select list_normalize([3.0, 4.0]);
+----
+[0.6, 0.8]
+
+# list_normalize alias multi-row with NULL row
+query ?
+select list_normalize(column1) from (values
+    (make_array(3.0, 4.0)),
+    (NULL)
+) as t(column1);
+----
+[0.6, 0.8]
+NULL
diff --git a/docs/source/user-guide/sql/scalar_functions.md b/docs/source/user-guide/sql/scalar_functions.md
index 4786163cb5..70216e3e9b 100644
--- a/docs/source/user-guide/sql/scalar_functions.md
+++ b/docs/source/user-guide/sql/scalar_functions.md
@@ -3268,6 +3268,7 @@ _Alias of [current_date](#current_date)._
 - [array_max](#array_max)
 - [array_min](#array_min)
 - [array_ndims](#array_ndims)
+- [array_normalize](#array_normalize)
 - [array_pop_back](#array_pop_back)
 - [array_pop_front](#array_pop_front)
 - [array_position](#array_position)
@@ -3321,6 +3322,7 @@ _Alias of [current_date](#current_date)._
 - [list_length](#list_length)
 - [list_max](#list_max)
 - [list_ndims](#list_ndims)
+- [list_normalize](#list_normalize)
 - [list_pop_back](#list_pop_back)
 - [list_pop_front](#list_pop_front)
 - [list_position](#list_position)
@@ -3893,6 +3895,33 @@ array_ndims(array, element)
 
 - list_ndims
 
+### `array_normalize`
+
+Returns the L2-normalized vector for the input numeric array, computed as `array[i] / sqrt(sum(array[i]^2))` per element. Returns NULL if the input is NULL, contains NULL elements, or has zero magnitude (all elements are zero). Returns an empty array for an empty input array.
+
+```sql
+array_normalize(array)
+```
+
+#### Arguments
+
+- **array**: Array expression. Can be a constant, column, or function, and any combination of array operators.
+
+#### Example
+
+```sql
+> select array_normalize([3.0, 4.0]);
++-----------------------------+
+| array_normalize(List([3.0,4.0])) |
++-----------------------------+
+| [0.6, 0.8]                  |
++-----------------------------+
+```
+
+#### Aliases
+
+- list_normalize
+
 ### `array_pop_back`
 
 Returns the array without the last element.
@@ -4778,6 +4807,10 @@ _Alias of [array_max](#array_max)._
 
 _Alias of [array_ndims](#array_ndims)._
 
+### `list_normalize`
+
+_Alias of [array_normalize](#array_normalize)._
+
 ### `list_pop_back`
 
 _Alias of [array_pop_back](#array_pop_back)._


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org