diff --git a/python/cudf_polars/cudf_polars/dsl/expr.py b/python/cudf_polars/cudf_polars/dsl/expr.py index 13e496136b5..377a905aed6 100644 --- a/python/cudf_polars/cudf_polars/dsl/expr.py +++ b/python/cudf_polars/cudf_polars/dsl/expr.py @@ -801,7 +801,7 @@ def do_evaluate( obj = plc.replace.replace_nulls( indices.obj, plc.interop.from_arrow( - pa.scalar(n, type=plc.interop.to_arrow(indices.obj.data_type())) + pa.scalar(n, type=plc.interop.to_arrow(indices.obj.type())) ), ) else: diff --git a/python/cudf_polars/tests/expressions/test_gather.py b/python/cudf_polars/tests/expressions/test_gather.py index df33e19a0b6..6bffa3e252c 100644 --- a/python/cudf_polars/tests/expressions/test_gather.py +++ b/python/cudf_polars/tests/expressions/test_gather.py @@ -2,8 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +import pytest + import polars as pl +from cudf_polars import execute_with_cudf from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -17,3 +20,31 @@ def test_gather(): query = ldf.select(pl.col("a").gather(pl.col("b"))) assert_gpu_result_equal(query) + + +def test_gather_with_nulls(): + ldf = pl.LazyFrame( + { + "a": [1, 2, 3, 4, 5, 6, 7], + "b": [0, None, 1, None, 6, 1, 0], + } + ) + + query = ldf.select(pl.col("a").gather(pl.col("b"))) + + assert_gpu_result_equal(query) + + +@pytest.mark.parametrize("negative", [False, True]) +def test_gather_out_of_bounds(negative): + ldf = pl.LazyFrame( + { + "a": [1, 2, 3, 4, 5, 6, 7], + "b": [0, -10 if negative else 10, 1, 2, 6, 1, 0], + } + ) + + query = ldf.select(pl.col("a").gather(pl.col("b"))) + + with pytest.raises(pl.exceptions.ComputeError): + query.collect(post_opt_callback=execute_with_cudf)