Skip to content

Commit ce312cc

Browse files
crm26claude
crm26
and
claude
committed
feat: handle mixed-type and NULL inputs for cosine_distance
Addresses round-2 review comments on #21542: - Widen container variant in coerce_types when inputs mix List and LargeList (or FixedSizeList), so mixed-type calls like `cosine_distance([1.0, 0.0], arrow_cast([0.0, 1.0], 'LargeList(Float64)'))` succeed. Follows the pattern from PR #21704 (ArrayConcat). - Coerce bare NULL inputs to a matching list variant so `cosine_distance(NULL, [1.0, 2.0])` returns NULL instead of erroring. - Drop the `list_cosine_distance` alias — the base name is not `array_cosine_distance`, so the `array_X` -> `list_X` convention does not apply. - Expand SLT coverage: mixed-type variants, FixedSizeList inputs, Float32 and Int64 inner types, bare NULL in each position, NULL row in a multi-row VALUES, and an unsupported-type plan error case. Dispatch fallthrough in cosine_distance_inner is now unreachable after the coerce_types widening, changed from exec_err! to internal_err!. Part of #21536. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
1 parent fc3ee90 commit ce312cc

3 files changed

Lines changed: 105 additions & 45 deletions

File tree

‎datafusion/functions-nested/src/cosine_distance.rs‎

Lines changed: 40 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,18 @@ use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait};
2222
use arrow::datatypes::{
2323
DataType,
2424
DataType::{FixedSizeList, LargeList, List, Null},
25+
Field,
2526
};
2627
use datafusion_common::cast::{as_float64_array, as_generic_list_array};
2728
use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only};
28-
use datafusion_common::{Result, exec_err, plan_err, utils::take_function_args};
29+
use datafusion_common::{
30+
Result, exec_err, internal_err, plan_err, utils::take_function_args,
31+
};
2932
use datafusion_expr::{
3033
ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
3134
Volatility,
3235
};
3336
use datafusion_macros::user_doc;
34-
use itertools::Itertools;
3537
use std::sync::Arc;
3638

3739
make_udf_expr_and_func!(
@@ -66,7 +68,6 @@ make_udf_expr_and_func!(
6668
#[derive(Debug, PartialEq, Eq, Hash)]
6769
pub struct CosineDistance {
6870
signature: Signature,
69-
aliases: Vec<String>,
7071
}
7172

7273
impl Default for CosineDistance {
@@ -79,7 +80,6 @@ impl CosineDistance {
7980
pub fn new() -> Self {
8081
Self {
8182
signature: Signature::user_defined(Volatility::Immutable),
82-
aliases: vec!["list_cosine_distance".to_string()],
8383
}
8484
}
8585
}
@@ -100,29 +100,48 @@ impl ScalarUDFImpl for CosineDistance {
100100
fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
101101
let [_, _] = take_function_args(self.name(), arg_types)?;
102102
let coercion = Some(&ListCoercion::FixedSizedListToList);
103-
let arg_types = arg_types.iter().map(|arg_type| {
104-
if matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) {
105-
Ok(coerced_type_with_base_type_only(
103+
104+
for arg_type in arg_types {
105+
if !matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) {
106+
return plan_err!("{} does not support type {arg_type}", self.name());
107+
}
108+
}
109+
110+
// If any input is `LargeList`, both sides must be widened to `LargeList`
111+
// so the runtime dispatch in `cosine_distance_inner` sees a homogeneous
112+
// pair. Follows the pattern in `ArrayConcat::coerce_types`.
113+
let any_large_list = arg_types.iter().any(|t| matches!(t, LargeList(_)));
114+
115+
let coerced = arg_types
116+
.iter()
117+
.map(|arg_type| {
118+
if matches!(arg_type, Null) {
119+
let field = Arc::new(Field::new_list_field(DataType::Float64, true));
120+
return if any_large_list {
121+
LargeList(field)
122+
} else {
123+
List(field)
124+
};
125+
}
126+
let coerced = coerced_type_with_base_type_only(
106127
arg_type,
107128
&DataType::Float64,
108129
coercion,
109-
))
110-
} else {
111-
plan_err!("{} does not support type {arg_type}", self.name())
112-
}
113-
});
114-
115-
arg_types.try_collect()
130+
);
131+
match coerced {
132+
List(field) if any_large_list => LargeList(field),
133+
other => other,
134+
}
135+
})
136+
.collect();
137+
138+
Ok(coerced)
116139
}
117140

118141
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
119142
make_scalar_function(cosine_distance_inner)(&args.args)
120143
}
121144

122-
fn aliases(&self) -> &[String] {
123-
&self.aliases
124-
}
125-
126145
fn documentation(&self) -> Option<&Documentation> {
127146
self.doc()
128147
}
@@ -133,11 +152,9 @@ fn cosine_distance_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
133152
match (array1.data_type(), array2.data_type()) {
134153
(List(_), List(_)) => general_cosine_distance::<i32>(args),
135154
(LargeList(_), LargeList(_)) => general_cosine_distance::<i64>(args),
136-
(arg_type1, arg_type2) => {
137-
exec_err!(
138-
"cosine_distance does not support types {arg_type1} and {arg_type2}"
139-
)
140-
}
155+
(arg_type1, arg_type2) => internal_err!(
156+
"cosine_distance received unexpected types after coercion: {arg_type1} and {arg_type2}"
157+
),
141158
}
142159
}
143160

‎datafusion/sqllogictest/test_files/cosine_distance.slt‎

Lines changed: 65 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -41,13 +41,23 @@ select round(cosine_distance([1.0, 0.0], [1.0, 1.0]), 3);
4141
----
4242
0.293
4343

44-
# NULL input (bare NULL is not a list type, errors at planning)
45-
query error cosine_distance does not support type
44+
# Bare NULL input returns NULL
45+
query R
4646
select cosine_distance(NULL, [1.0, 2.0]);
47+
----
48+
NULL
4749

48-
# NULL in second position
49-
query error cosine_distance does not support type
50+
# NULL in second position returns NULL
51+
query R
5052
select cosine_distance([1.0, 2.0], NULL);
53+
----
54+
NULL
55+
56+
# Both NULL returns NULL
57+
query R
58+
select cosine_distance(NULL, NULL);
59+
----
60+
NULL
5161

5262
# Zero vector returns NULL (undefined cosine similarity)
5363
query R
@@ -74,29 +84,71 @@ select cosine_distance(
7484
----
7585
1
7686

77-
# Integer arrays (coerced to Float64)
87+
# Mixed List + LargeList: widens to LargeList
88+
query R
89+
select cosine_distance([1.0, 0.0], arrow_cast([0.0, 1.0], 'LargeList(Float64)'));
90+
----
91+
1
92+
93+
# Reverse order: LargeList + List also widens
94+
query R
95+
select cosine_distance(arrow_cast([1.0, 0.0], 'LargeList(Float64)'), [0.0, 1.0]);
96+
----
97+
1
98+
99+
# FixedSizeList inputs (coerced to List)
100+
query R
101+
select cosine_distance(
102+
arrow_cast([1.0, 0.0], 'FixedSizeList(2, Float64)'),
103+
arrow_cast([0.0, 1.0], 'FixedSizeList(2, Float64)')
104+
);
105+
----
106+
1
107+
108+
# FixedSizeList + LargeList: widens to LargeList
109+
query R
110+
select cosine_distance(
111+
arrow_cast([1.0, 0.0], 'FixedSizeList(2, Float64)'),
112+
arrow_cast([0.0, 1.0], 'LargeList(Float64)')
113+
);
114+
----
115+
1
116+
117+
# Float32 inner type (coerced to Float64)
118+
query R
119+
select cosine_distance(arrow_cast([1.0, 0.0], 'List(Float32)'), [0.0, 1.0]);
120+
----
121+
1
122+
123+
# Int64 inner type (coerced to Float64)
124+
query R
125+
select cosine_distance(arrow_cast([1, 0], 'List(Int64)'), arrow_cast([0, 1], 'List(Int64)'));
126+
----
127+
1
128+
129+
# Integer literals (coerced to Float64)
78130
query R
79131
select cosine_distance([1, 0], [0, 1]);
80132
----
81133
1
82134

83-
# Multi-row query
135+
# Unsupported non-list input (plan error)
136+
query error cosine_distance does not support type
137+
select cosine_distance(1, 2);
138+
139+
# Multi-row query with NULL row propagation
84140
query R
85141
select cosine_distance(column1, column2) from (values
86142
(make_array(1.0, 0.0), make_array(0.0, 1.0)),
87143
(make_array(1.0, 1.0), make_array(1.0, 1.0)),
88-
(make_array(1.0, 0.0), make_array(-1.0, 0.0))
144+
(make_array(1.0, 0.0), make_array(-1.0, 0.0)),
145+
(make_array(1.0, 0.0), NULL)
89146
) as t(column1, column2);
90147
----
91148
1
92149
0
93150
2
94-
95-
# list_cosine_distance alias
96-
query R
97-
select list_cosine_distance([1.0, 0.0], [0.0, 1.0]);
98-
----
99-
1
151+
NULL
100152

101153
# Empty arrays return NULL (magnitude = 0)
102154
query R

‎docs/source/user-guide/sql/scalar_functions.md‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3288,7 +3288,6 @@ _Alias of [current_date](#current_date)._
32883288
- [list_cat](#list_cat)
32893289
- [list_concat](#list_concat)
32903290
- [list_contains](#list_contains)
3291-
- [list_cosine_distance](#list_cosine_distance)
32923291
- [list_dims](#list_dims)
32933292
- [list_distance](#list_distance)
32943293
- [list_distinct](#list_distinct)
@@ -4467,10 +4466,6 @@ cosine_distance(array1, array2)
44674466
+-----------------------------------------------+
44684467
```
44694468

4470-
#### Aliases
4471-
4472-
- list_cosine_distance
4473-
44744469
### `empty`
44754470

44764471
Returns 1 for an empty array or 0 for a non-empty array.
@@ -4573,10 +4568,6 @@ _Alias of [array_concat](#array_concat)._
45734568

45744569
_Alias of [array_has](#array_has)._
45754570

4576-
### `list_cosine_distance`
4577-
4578-
_Alias of [cosine_distance](#cosine_distance)._
4579-
45804571
### `list_dims`
45814572

45824573
_Alias of [array_dims](#array_dims)._

0 commit comments

Comments
 (0)