diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index c7de4161d..5eea55760 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5333,6 +5333,11 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray else: values_has_nan = pc.is_nan(values) values_indices = pc.list_parent_indices(arr) + if arr.null_count: + # Match flatten(), which excludes child spans under null parent lists. + values_indices = pc.filter( + values_indices, pc.take(pc.is_valid(arr), values_indices) + ) has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan)) indices = pa.array(range(len(arr)), type=pa.uint32()) return pc.is_in(indices, has_nan_indices) diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 4671109b8..4fb16bb31 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -21,6 +21,7 @@ from lancedb.table import ( _infer_target_schema, _merge_metadata, _sanitize_data, + has_nan_values, sanitize_create_table, ) import pyarrow as pa @@ -489,6 +490,20 @@ def test_handle_bad_vectors_nan(on_bad_vectors): assert output["vector"].combine_chunks() == expected +@pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()]) +@pytest.mark.parametrize("large", [False, True]) +@pytest.mark.parametrize("chunked", [False, True]) +def test_has_nan_values_skips_hidden_children_of_null_lists(dtype, large, chunked): + array_type = pa.LargeListArray if large else pa.ListArray + values = pa.array([float("nan"), None, float("nan"), 3.0, 4.0, 5.0], type=dtype) + arr = array_type.from_arrays( + [0, 2, 4, 6], values, mask=pa.array([True, False, False]) + ) + if chunked: + arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2, 1)]) + assert has_nan_values(arr).to_pylist() == [False, True, False] + + def test_handle_bad_vectors_noop(): # ChunkedArray should be preserved as-is vector = pa.chunked_array(