fix(rust): Add Int128 path for join_asof (#21282)

pola-rs · Feb 16, 2025 · 4464d00 · 4464d00
1 parent 4eb4fb7
commit 4464d00
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/crates/polars-ops/src/frame/join/asof/mod.rs b/crates/polars-ops/src/frame/join/asof/mod.rs
@@ -297,6 +297,11 @@ pub trait AsofJoin: IntoDf {
                 let ca = left_key.i32().unwrap();
                 join_asof_numeric(ca, &right_key, strategy, tolerance, allow_eq)
             },
+            #[cfg(feature = "dtype-i128")]
+            DataType::Int128 => {
+                let ca = left_key.i128().unwrap();
+                join_asof_numeric(ca, &right_key, strategy, tolerance, allow_eq)
+            },
             DataType::UInt64 => {
                 let ca = left_key.u64().unwrap();
                 join_asof_numeric(ca, &right_key, strategy, tolerance, allow_eq)

diff --git a/py-polars/tests/unit/operations/test_join_asof.py b/py-polars/tests/unit/operations/test_join_asof.py
@@ -5,6 +5,7 @@
 import pytest
 
 import polars as pl
+from polars._typing import AsofJoinStrategy, PolarsIntegerType
 from polars.exceptions import InvalidOperationError
 from polars.testing import assert_frame_equal
 
@@ -1287,3 +1288,30 @@ def test_join_asof_not_sorted() -> None:
         df.join_asof(df, on="b", by="a")
     with pytest.raises(InvalidOperationError, match="is not sorted"):
         df.join_asof(df, on="b")
+
+
+@pytest.mark.parametrize("left_dtype", [pl.Int64, pl.UInt64])
+@pytest.mark.parametrize("right_dtype", [pl.Int64, pl.UInt64])
+@pytest.mark.parametrize("strategy", ["backward", "forward", "nearest"])
+def test_join_asof_large_int_21276(
+    left_dtype: PolarsIntegerType,
+    right_dtype: PolarsIntegerType,
+    strategy: AsofJoinStrategy,
+) -> None:
+    large_int64 = 1608129000134000123  # it only happen when "on" column is large
+    left = pl.DataFrame({"ts": pl.Series([large_int64 + 2], dtype=left_dtype)})
+    right = pl.DataFrame(
+        {
+            "ts": pl.Series([large_int64 + 1, large_int64 + 3], dtype=right_dtype),
+            "value": [111, 333],
+        }
+    )
+    result = left.join_asof(right, on="ts", strategy=strategy)
+    idx = 0 if strategy == "backward" else 1
+    expected = pl.DataFrame(
+        {
+            "ts": left["ts"],
+            "value": right["value"].gather(idx),
+        }
+    )
+    assert_frame_equal(result, expected)