feat: add vector distance and array math functions#21371
Draft
crm26 wants to merge 4 commits intoapache:mainfrom
Draft
feat: add vector distance and array math functions#21371crm26 wants to merge 4 commits intoapache:mainfrom
crm26 wants to merge 4 commits intoapache:mainfrom
Conversation
Add 6 new scalar functions to datafusion-functions-nested: - cosine_distance(array, array) — cosine distance (1 - cosine similarity) - inner_product(array, array) — dot product - array_normalize(array) — L2 unit normalization - array_add(array, array) — element-wise addition - array_subtract(array, array) — element-wise subtraction - array_scale(array, float) — scalar multiplication Shared math primitives (dot_product, magnitude, sum_of_squares) extracted into vector_math.rs to avoid duplication across functions. Includes aliases (list_*, dot_product), 29 unit tests, and a sqllogictest file with vector search pattern coverage. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Adds cosine_distance, inner_product, array_normalize, array_add, array_subtract, and array_scale to datafusion-functions-nested. Shared primitives in vector_math.rs (dot_product_f64, magnitude_f64, sum_of_squares_f64, convert_to_f64_array) are reused across all functions and the existing array_distance. Consolidates the duplicate convert_to_f64_array from distance.rs into the shared module. Functions: cosine_distance(a, b) → float64 (aliases: list_cosine_distance) inner_product(a, b) → float64 (aliases: list_inner_product, dot_product) array_normalize(a) → list(float64) (aliases: list_normalize) array_add(a, b) → list(float64) (aliases: list_add) array_subtract(a, b) → list(float64) (aliases: list_subtract) array_scale(a, f) → list(float64) (aliases: list_scale) Enables vector search in standard SQL: SELECT id, cosine_distance(embedding, ARRAY[0.1, 0.2, ...]) as dist FROM documents ORDER BY dist LIMIT 10 79 tests, sqllogictest coverage, clippy clean.
github-actions
Bot
added
sqllogictest
Contributor
|
Hi -- thank you for this PR. I think it will be challenging to review a PR of this size To help review can you please:
Thank you so much 🙏 |
Contributor
Author
|
Done — filed tracking issue #21536 and splitting into one PR per function. First up: |
4 tasks
zzcclp
pushed a commit
to zzcclp/arrow-datafusion
that referenced
this pull request
Apr 23, 2026
## Summary - Adds `cosine_distance(array1, array2)` / `list_cosine_distance` — computes cosine distance (1 - cosine similarity) between two numeric arrays - Introduces shared `vector_math.rs` primitives (`dot_product_f64`, `magnitude_f64`, `convert_to_f64_array`) for reuse by follow-on vector functions - Returns NULL for zero-magnitude vectors; errors on mismatched lengths - Supports List, LargeList, and FixedSizeList with any numeric element type Part of apache#21536 — first in a series of split PRs (replacing apache#21371). ## Test plan - [x] Unit tests: identical, orthogonal, opposite, 45-degree, zero-magnitude, mismatched-length, NULL, multi-row - [x] sqllogictest: `cosine_distance.slt` covering all edge cases including empty arrays, LargeList, integer coercion, alias, return type - [x] Full slt suite (426/426 pass) - [x] `cargo clippy`, `cargo fmt`, `taplo`, `prettier`, `cargo machete` — all clean 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
pull Bot
pushed a commit
to buraksenn/datafusion
that referenced
this pull request
May 3, 2026
## Which issue does this PR close? Part of apache#21536 — split of apache#21371 into one-function-per-PR. ## Rationale for this change Adds `inner_product(array1, array2)` — the dot product of two equal-length numeric arrays, returning `Float64`. Computed as `sum(array1[i] * array2[i])`. ## What changes are included in this PR? Mirrors the structural pattern of merged apache#21542 (`cosine_distance`): - Same `coerce_types` for `List`/`LargeList`/`FixedSizeList` of any numeric inner type, with widening to `LargeList` when any input is `LargeList` (per the apache#21704 pattern) - Same NULL semantics: bare `NULL` → `NULL`, NULL row → NULL, NULL element in list → NULL - Same Arrow-idiomatic implementation: single `as_float64_array(list_array.values())` downcast, slice by `value_offsets()`, iterate via `ScalarBuffer<f64>` - No alias, no shared module — standalone, inline math The arithmetic is the only semantic divergence from `cosine_distance`: - `dot += a*b` (no magnitude or normalization) - Empty arrays return `0.0` (sum of empty set), not `NULL` - No zero-magnitude special case (`inner_product([0,0], [1,2])` returns `0`, which is well-defined for inner product) ## Are these changes tested? Yes. SLT covers: - Orthogonal, identical, opposite, general non-trivial vectors - Single zero vector, both zero vectors - Bare `NULL` in either or both positions - NULL element inside a list (returns NULL for that row) - Mismatched lengths (error) - `LargeList` inputs - Mixed `(List, LargeList)` in both orders - `(FixedSizeList, FixedSizeList)` and `(FixedSizeList, LargeList)` - `Float32` and `Int64` inner type coercion - Multi-row query with NULL row propagation - Empty arrays (returns `0`) - No-args error - Return-type assertion (`Float64`) ## Are there any user-facing changes? New scalar function `inner_product`, documented in `docs/source/user-guide/sql/scalar_functions.md`. --------- Co-authored-by: Jeffrey Vo <jeffrey.vo.australia@gmail.com> Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds vector distance and array math functions to
datafusion-functions-nested, enabling vector search and array algebra in standard SQL.Functions
cosine_distance(a, b)inner_product(a, b)array_normalize(a)array_add(a, b)array_subtract(a, b)array_scale(a, f)All have
list_*aliases.inner_productalso aliased asdot_product.Design
Shared primitives in
vector_math.rs:dot_product_f64(a, b)— used byinner_productandcosine_distancemagnitude_f64(a)— used bycosine_distanceandarray_normalizesum_of_squares_f64(a)— used bymagnitude_f64convert_to_f64_array(a)— shared with existingarray_distanceThe existing
distance.rsduplicateconvert_to_f64_arrayis consolidated into the shared module.Follows the exact pattern of the existing
array_distancefunction: same signature style,coerce_types, null handling, and type support (Float32, Float64, Int32, Int64, FixedSizeList, LargeList, List).Tests
79 tests including: normal inputs, null handling, zero vectors, orthogonal vectors, empty arrays, Float32/Float64, mismatched lengths, vector search ranking pattern. Sqllogictest coverage in
vector_functions.slt. Clippy clean.