add test

ion-elgreco · Nov 17, 2024 · 65f98d5 · 65f98d5
1 parent 9c22733
commit 65f98d5
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/polars_distance/src/expressions.rs b/polars_distance/src/expressions.rs
@@ -39,7 +39,7 @@ fn elementwise_str_u32(
     y: &ChunkedArray<StringType>,
     f: fn(&str, &str) -> u32,
 ) -> UInt32Chunked {
-    let (x, y) = if x.len() < y.len() { (x, y) } else { (y, x) };
+    let (x, y) = if x.len() < y.len() { (y, x) } else { (x, y) };
     match y.len() {
         1 => match unsafe { y.get_unchecked(0) } {
             Some(y_value) => arity::unary_elementwise(x, |x| x.map(|x| f(x, y_value))),
@@ -59,7 +59,7 @@ fn elementwise_str_f64(
     y: &ChunkedArray<StringType>,
     f: fn(&str, &str) -> f64,
 ) -> Float64Chunked {
-    let (x, y) = if x.len() < y.len() { (x, y) } else { (y, x) };
+    let (x, y) = if x.len() < y.len() { (y, x) } else { (x, y) };
     match y.len() {
         1 => match unsafe { y.get_unchecked(0) } {
             Some(y_value) => arity::unary_elementwise(x, |x| x.map(|x| f(x, y_value))),

diff --git a/polars_distance/tests/test_distance_arr.py b/polars_distance/tests/test_distance_arr.py
@@ -359,3 +359,16 @@ def test_broadcast():
         ]
     )
     assert_frame_equal(result, expected)
+
+    df = pl.DataFrame(
+        {
+            "a1": ["test1", "hello", "test1", "hello", "test1", "hello"],
+        }
+    )
+    result = df.select(d=pld.col("a1").dist_str.levenshtein(pl.lit("testaa")))
+    expected = pl.DataFrame(
+        [
+            pl.Series("d", [2, 5, 2, 5, 2, 5], dtype=pl.UInt32),
+        ]
+    )
+    assert_frame_equal(result, expected)