diff --git a/crates/polars-ops/src/frame/join/asof/mod.rs b/crates/polars-ops/src/frame/join/asof/mod.rs index 4b7ccd21b600..66eeb6d9808b 100644 --- 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 index 14bc6cfea729..d4eae68957c9 100644 --- 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)