mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 20:22:37 +00:00
fix(python): align vector parent indices after flattening null lists (#4572)
Arrow's `flatten()` skips child values hidden under null list parents, but `list_parent_indices()` can retain their indices. This makes `has_nan_values()` raise `ArrowInvalid` because the parent indices and NaN mask have different lengths. Filter out indices belonging to null parents before applying the mask. The regression test covers `ListArray` and `LargeListArray`, float16/32/64, and chunked inputs: 12 cases. NaN and null-component policies stay the same. The crash also affects #4503. This patch fixes it separately from that PR's changes to bad-vector policies. Validation: - All 12 regression cases failed before the fix. The utility suite now passes with 87 passed and 1 skipped. - Ruff lint and formatting pass for both changed files, as does `git diff --check`. - The patch applies cleanly to main at `ca275b6d`. The patched helper from that revision also passes all 12 cases. - A separate AI reviewer checked commit `343082a6` and found no blocking issues. They reran the utility suite and compared 4,059 generated cases plus 12 explicit null-parent cases with a row-by-row reference implementation. Those checks include slices, chunked inputs, and fixed-size lists. SDK tests used base `07e626ec` with the existing compiled extension. The broader Python run stopped after 474 passed and 23 skipped at `test_local_module_outside_code_warns_at_registration` with `ModuleNotFoundError`. That test also fails with the original vector helper restored. Repository-wide Ruff reports 49 lint errors and 11 files needing formatting, all in untouched documentation examples. Full SDK integration tests against current main require a newer native build and have not been run.
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user