Skip to content

Commit 7ea9240

Browse files
davidlghellindavidlghellin
committed
fix revision
1 parent 4b9f517 commit 7ea9240

3 files changed

Lines changed: 186 additions & 55 deletions

File tree

‎datafusion/spark/src/function/string/levenshtein.rs‎

Lines changed: 170 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use std::any::Any;
1918
use std::sync::Arc;
2019

2120
use arrow::array::{Array, ArrayRef, Int32Array, Int64Array, OffsetSizeTrait};
@@ -83,10 +82,6 @@ impl SparkLevenshtein {
8382
}
8483

8584
impl ScalarUDFImpl for SparkLevenshtein {
86-
fn as_any(&self) -> &dyn Any {
87-
self
88-
}
89-
9085
fn name(&self) -> &str {
9186
"levenshtein"
9287
}
@@ -122,14 +117,40 @@ impl ScalarUDFImpl for SparkLevenshtein {
122117
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
123118
let ScalarFunctionArgs { args, .. } = args;
124119

125-
// Determine the coerced string type (handles mixed Utf8 + LargeUtf8)
120+
// Arity is validated in `return_type` at planning time, but invoke can
121+
// also be reached via direct UDF calls (tests, programmatic invocation)
122+
// so we re-check here before indexing `args[0]` / `args[1]`.
123+
if args.len() != 2 && args.len() != 3 {
124+
return exec_err!("levenshtein expects 2 or 3 arguments, got {}", args.len());
125+
}
126+
127+
// Determine the coerced string type (handles mixed Utf8 + LargeUtf8).
128+
// If neither string nor binary→string coercion applies, error out
129+
// explicitly — mirrors `return_type`'s behaviour instead of silently
130+
// defaulting to Utf8.
126131
let coerced_type = string_coercion(&args[0].data_type(), &args[1].data_type())
127132
.or_else(|| {
128133
binary_to_string_coercion(&args[0].data_type(), &args[1].data_type())
129134
})
130-
.unwrap_or(DataType::Utf8);
131-
132-
// Spark returns NULL when any scalar argument is NULL.
135+
.ok_or_else(|| {
136+
datafusion_common::DataFusionError::Execution(format!(
137+
"Unsupported data types for levenshtein. Expected Utf8, \
138+
LargeUtf8 or Utf8View, got {:?} and {:?}",
139+
args[0].data_type(),
140+
args[1].data_type()
141+
))
142+
})?;
143+
144+
// Spark NULL behaviour for `levenshtein`:
145+
// - NULL string argument (scalar OR per-row) → result is NULL
146+
// - NULL threshold scalar → result is NULL
147+
// - NULL threshold value PER ROW → treated as `0`
148+
// (handled inside `apply_threshold_*` in the kernel, NOT here)
149+
//
150+
// The per-row asymmetry (string-null propagates, threshold-null does
151+
// not) is verified Spark behaviour — see the Sail Gherkin scenario
152+
// "null threshold per row treated as zero" in
153+
// `python/pysail/tests/spark/function/features/levenshtein.feature`.
133154
let null_int = |dt: &DataType| match dt {
134155
DataType::LargeUtf8 => ColumnarValue::Scalar(ScalarValue::Int64(None)),
135156
_ => ColumnarValue::Scalar(ScalarValue::Int32(None)),
@@ -154,6 +175,50 @@ impl ScalarUDFImpl for SparkLevenshtein {
154175
}
155176
}
156177

178+
/// Applies the optional per-row threshold to a computed Levenshtein distance,
179+
/// returning an `i32` result. Used by the `Utf8` / `Utf8View` paths.
180+
///
181+
/// Spark semantics — verified against PySpark and codified in the Sail
182+
/// test suite (`spark/function/features/levenshtein.feature`, scenario
183+
/// "null threshold per row treated as zero"):
184+
///
185+
/// - No threshold column at all → return the raw distance.
186+
/// - Non-null threshold value → return `-1` if `dist > threshold`, else `dist`.
187+
/// - NULL threshold value at this row → treat as `0`, so any non-zero distance
188+
/// returns `-1` and identical strings (distance 0) return `0`.
189+
///
190+
/// This matches Spark's behaviour where a NULL threshold per row is NOT
191+
/// null-propagating; only NULL string arguments produce a NULL result.
192+
#[inline]
193+
fn apply_threshold_i32(dist: i32, threshold: Option<(&Int32Array, usize)>) -> i32 {
194+
match threshold {
195+
Some((arr, i)) => {
196+
let t = if arr.is_null(i) { 0 } else { arr.value(i) };
197+
if dist > t { -1 } else { dist }
198+
}
199+
None => dist,
200+
}
201+
}
202+
203+
/// `i64` counterpart of [`apply_threshold_i32`], used by the `LargeUtf8` path
204+
/// where the distance is materialised as `i64`. The threshold column itself
205+
/// remains `Int32` (Spark's signature), so the value is widened to `i64`
206+
/// before comparison.
207+
#[inline]
208+
fn apply_threshold_i64(dist: i64, threshold: Option<(&Int32Array, usize)>) -> i64 {
209+
match threshold {
210+
Some((arr, i)) => {
211+
let t = if arr.is_null(i) {
212+
0
213+
} else {
214+
arr.value(i) as i64
215+
};
216+
if dist > t { -1 } else { dist }
217+
}
218+
None => dist,
219+
}
220+
}
221+
157222
/// Spark-compatible Levenshtein distance with optional threshold.
158223
fn spark_levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
159224
if args.len() < 2 || args.len() > 3 {
@@ -197,18 +262,10 @@ fn spark_levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef>
197262
let dist = datafusion_strsim::levenshtein_with_buffer(
198263
string1, string2, &mut cache,
199264
) as i32;
200-
match &threshold {
201-
Some(t) => {
202-
let thresh =
203-
if t.is_null(i) { 0 } else { t.value(i) };
204-
if dist > thresh {
205-
Some(-1i32)
206-
} else {
207-
Some(dist)
208-
}
209-
}
210-
None => Some(dist),
211-
}
265+
Some(apply_threshold_i32(
266+
dist,
267+
threshold.as_ref().map(|t| (*t, i)),
268+
))
212269
}
213270
_ => None,
214271
})
@@ -229,18 +286,10 @@ fn spark_levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef>
229286
let dist = datafusion_strsim::levenshtein_with_buffer(
230287
string1, string2, &mut cache,
231288
) as i32;
232-
match &threshold {
233-
Some(t) => {
234-
let thresh =
235-
if t.is_null(i) { 0 } else { t.value(i) };
236-
if dist > thresh {
237-
Some(-1i32)
238-
} else {
239-
Some(dist)
240-
}
241-
}
242-
None => Some(dist),
243-
}
289+
Some(apply_threshold_i32(
290+
dist,
291+
threshold.as_ref().map(|t| (*t, i)),
292+
))
244293
}
245294
_ => None,
246295
})
@@ -261,18 +310,10 @@ fn spark_levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef>
261310
let dist = datafusion_strsim::levenshtein_with_buffer(
262311
string1, string2, &mut cache,
263312
) as i64;
264-
match &threshold {
265-
Some(t) => {
266-
let thresh =
267-
if t.is_null(i) { 0 } else { t.value(i) as i64 };
268-
if dist > thresh {
269-
Some(-1i64)
270-
} else {
271-
Some(dist)
272-
}
273-
}
274-
None => Some(dist),
275-
}
313+
Some(apply_threshold_i64(
314+
dist,
315+
threshold.as_ref().map(|t| (*t, i)),
316+
))
276317
}
277318
_ => None,
278319
})
@@ -591,4 +632,88 @@ mod tests {
591632
assert_eq!(&Int64Array::from(vec![3i64]), result);
592633
Ok(())
593634
}
635+
636+
// ── Utf8View path ───────────────────────────────────────────
637+
638+
#[test]
639+
fn test_utf8view_two_arg() -> Result<()> {
640+
use arrow::array::StringViewArray;
641+
let s1: ArrayRef = Arc::new(StringViewArray::from(vec!["kitten"]));
642+
let s2: ArrayRef = Arc::new(StringViewArray::from(vec!["sitting"]));
643+
let res = spark_levenshtein::<i32>(&[s1, s2])?;
644+
let result = as_int32_array(&res)?;
645+
assert_eq!(&Int32Array::from(vec![3]), result);
646+
Ok(())
647+
}
648+
649+
#[test]
650+
fn test_utf8view_with_threshold_exceeds() -> Result<()> {
651+
use arrow::array::StringViewArray;
652+
let s1: ArrayRef = Arc::new(StringViewArray::from(vec!["kitten"]));
653+
let s2: ArrayRef = Arc::new(StringViewArray::from(vec!["sitting"]));
654+
let t: ArrayRef = Arc::new(Int32Array::from(vec![2]));
655+
// distance 3 > threshold 2 → -1
656+
let res = spark_levenshtein::<i32>(&[s1, s2, t])?;
657+
let result = as_int32_array(&res)?;
658+
assert_eq!(&Int32Array::from(vec![-1]), result);
659+
Ok(())
660+
}
661+
662+
#[test]
663+
fn test_utf8view_with_threshold_within() -> Result<()> {
664+
use arrow::array::StringViewArray;
665+
let s1: ArrayRef = Arc::new(StringViewArray::from(vec!["kitten"]));
666+
let s2: ArrayRef = Arc::new(StringViewArray::from(vec!["sitting"]));
667+
let t: ArrayRef = Arc::new(Int32Array::from(vec![5]));
668+
// distance 3 <= threshold 5 → 3
669+
let res = spark_levenshtein::<i32>(&[s1, s2, t])?;
670+
let result = as_int32_array(&res)?;
671+
assert_eq!(&Int32Array::from(vec![3]), result);
672+
Ok(())
673+
}
674+
675+
// ── LargeUtf8 + threshold path ──────────────────────────────
676+
677+
#[test]
678+
fn test_largeutf8_with_threshold_exceeds() -> Result<()> {
679+
use arrow::array::LargeStringArray;
680+
use datafusion_common::cast::as_int64_array;
681+
let s1: ArrayRef = Arc::new(LargeStringArray::from(vec!["kitten"]));
682+
let s2: ArrayRef = Arc::new(LargeStringArray::from(vec!["sitting"]));
683+
let t: ArrayRef = Arc::new(Int32Array::from(vec![2]));
684+
// distance 3 > threshold 2 → -1 (as i64)
685+
let res = spark_levenshtein::<i64>(&[s1, s2, t])?;
686+
let result = as_int64_array(&res)?;
687+
assert_eq!(&Int64Array::from(vec![-1i64]), result);
688+
Ok(())
689+
}
690+
691+
#[test]
692+
fn test_largeutf8_with_threshold_within() -> Result<()> {
693+
use arrow::array::LargeStringArray;
694+
use datafusion_common::cast::as_int64_array;
695+
let s1: ArrayRef = Arc::new(LargeStringArray::from(vec!["kitten"]));
696+
let s2: ArrayRef = Arc::new(LargeStringArray::from(vec!["sitting"]));
697+
let t: ArrayRef = Arc::new(Int32Array::from(vec![5]));
698+
// distance 3 <= threshold 5 → 3 (as i64)
699+
let res = spark_levenshtein::<i64>(&[s1, s2, t])?;
700+
let result = as_int64_array(&res)?;
701+
assert_eq!(&Int64Array::from(vec![3i64]), result);
702+
Ok(())
703+
}
704+
705+
#[test]
706+
fn test_largeutf8_with_null_threshold_row_treated_as_zero() -> Result<()> {
707+
use arrow::array::LargeStringArray;
708+
use datafusion_common::cast::as_int64_array;
709+
// Row 0: "abc" vs "def" with NULL threshold → treated as 0 → -1
710+
// Row 1: "xyz" vs "xyz" with NULL threshold → distance 0 → 0
711+
let s1: ArrayRef = Arc::new(LargeStringArray::from(vec!["abc", "xyz"]));
712+
let s2: ArrayRef = Arc::new(LargeStringArray::from(vec!["def", "xyz"]));
713+
let t: ArrayRef = Arc::new(Int32Array::from(vec![None, None]));
714+
let res = spark_levenshtein::<i64>(&[s1, s2, t])?;
715+
let result = as_int64_array(&res)?;
716+
assert_eq!(&Int64Array::from(vec![-1i64, 0i64]), result);
717+
Ok(())
718+
}
594719
}

‎datafusion/spark/src/function/string/mod.rs‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -94,11 +94,17 @@ pub mod expr_fn {
9494
"Returns the character length of string data or number of bytes of binary data. The length of string data includes the trailing spaces. The length of binary data includes binary zeros.",
9595
arg1
9696
));
97-
export_functions!((
98-
levenshtein,
99-
"Returns the Levenshtein distance between two strings. Optionally accepts a threshold; returns -1 if the distance exceeds it.",
100-
str1 str2 threshold
101-
));
97+
// `levenshtein` has both 2-arg (`str1`, `str2`) and 3-arg
98+
// (`str1`, `str2`, `threshold`) forms, so the fixed-arity
99+
// `export_functions!` macro would force one shape and make the other
100+
// unreachable from the `expr_fn` API. Mirror how the core
101+
// `datafusion-functions` crate exposes variadic helpers
102+
// (e.g. `concat_ws`, `trim`) and ship a manual `Vec<Expr>` wrapper.
103+
#[doc = "Returns the Levenshtein distance between two strings. Optionally accepts a threshold; returns -1 if the distance exceeds it."]
104+
pub fn levenshtein(args: Vec<datafusion_expr::Expr>) -> datafusion_expr::Expr {
105+
super::levenshtein().call(args)
106+
}
107+
102108
export_functions!((
103109
like,
104110
"Returns true if str matches pattern (case sensitive).",

‎datafusion/sqllogictest/test_files/spark/string/levenshtein.slt‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -147,35 +147,35 @@ SELECT levenshtein(s1, s2, 2) AS result FROM VALUES ('abc', 'abc'), ('abc', 'def
147147

148148
## Different threshold per row
149149
query I
150-
SELECT levenshtein(s1, s2, t) AS result FROM VALUES ('abc', 'def', 2), ('abc', 'def', 5), ('abc', 'def', 3) AS t(s1, s2, t);
150+
SELECT levenshtein(s1, s2, threshold) AS result FROM VALUES ('abc', 'def', 2), ('abc', 'def', 5), ('abc', 'def', 3) AS v(s1, s2, threshold);
151151
----
152152
-1
153153
3
154154
3
155155

156156
## Null threshold per row — Spark treats null threshold as 0
157157
query I
158-
SELECT levenshtein(s1, s2, t) AS result FROM VALUES ('kitten', 'sitting', CAST(NULL AS INT)), ('abc', 'def', 1) AS t(s1, s2, t);
158+
SELECT levenshtein(s1, s2, threshold) AS result FROM VALUES ('kitten', 'sitting', CAST(NULL AS INT)), ('abc', 'def', 1) AS v(s1, s2, threshold);
159159
----
160160
-1
161161
-1
162162

163163
## Null threshold per row with identical strings — dist=0 <= 0 → returns 0
164164
query I
165-
SELECT levenshtein(s1, s2, t) AS result FROM VALUES ('abc', 'abc', CAST(NULL AS INT)) AS t(s1, s2, t);
165+
SELECT levenshtein(s1, s2, threshold) AS result FROM VALUES ('abc', 'abc', CAST(NULL AS INT)) AS v(s1, s2, threshold);
166166
----
167167
0
168168

169169
## Null strings with threshold in columns
170170
query I
171-
SELECT levenshtein(s1, s2, t) AS result
171+
SELECT levenshtein(s1, s2, threshold) AS result
172172
FROM VALUES
173173
(CAST(NULL AS STRING), 'hello', 5),
174174
('hello', CAST(NULL AS STRING), 5),
175175
(CAST(NULL AS STRING), CAST(NULL AS STRING), 5),
176176
('abc', 'abc', CAST(NULL AS INT)),
177177
('abc', 'def', CAST(NULL AS INT))
178-
AS t(s1, s2, t);
178+
AS v(s1, s2, threshold);
179179
----
180180
NULL
181181
NULL

0 commit comments

Comments
 (0)